Skip to content

T
tic
Unverified Commit e63e08fe by Tianqi Chen Committed by GitHub
Options

[REFACTOR][TIR] Migrate all low-level passes to the Pass Manager. (#5233) 

* [REFACTOR][TIR] Migrate all low-level passes to the Pass Manager.

This PR migrates the tvm.lower to return IRModule of PrimFuncs
instead of the LoweredFuncs.

* Remove LoweredFunc.
parent fd9ce583
Showing with 598 additions and 756 deletions
apps/lldb/tvm.py
......@@ -46,7 +46,6 @@ def __lldb_init_module(debugger, _):
"tvm::IterVarAttr",
"tvm::IterVarRelation",
"tvm::Layout",
"tir::LoweredFunc",
"tvm::Map",
"tvm::Map",
"tvm::MemoryInfo",
......
docs/dev/codebase_walkthrough.rst
......@@ -145,15 +145,6 @@ After lowering is done, ``build()`` function generates target machine code from
Code generation is done by ``build_module()`` function, defined in ``python/tvm/target/codegen.py``. On the C++ side, code generation is implemented in ``src/target/codegen`` subdirectory. ``build_module()`` Python function will reach ``Build()`` function below in ``src/target/codegen/codegen.cc``:
::
runtime::Module Build(const Array<LoweredFunc>& funcs,
const std::string& target) {
std::string build_f_name = "codegen.build_" + target;
const PackedFunc* bf = runtime::Registry::Get(build_f_name);
runtime::Module m = (*bf)(funcs, target);
return m;
}
The ``Build()`` function looks up the code generator for the given target in the ``PackedFunc`` registry, and invokes the function found. For example, ``codegen.build_cuda`` function is registered in ``src/codegen/build_cuda_on.cc``, like this:
......
include/tvm/driver/driver_api.h
......@@ -32,8 +32,8 @@
#include <tvm/runtime/packed_func.h>
#include <tvm/target/target.h>
#include <tvm/support/with.h>
#include <tvm/ir/module.h>
#include <tvm/te/schedule_pass.h>
#include <tvm/tir/lowered_func.h>
#include <string>
#include <vector>
......@@ -43,15 +43,15 @@
namespace tvm {
/*!
* \brief Build a LoweredFunc given a schedule, args and binds
* \brief Build an IRModule given a schedule, args and binds
* \param sch The schedule to lower.
* \param args The arguments to the function.
* \param name The name of the lowered function.
* \param binds Buffer assignments.
* \param config The build configuration.
* \return The lowered function.
* \return The result module.
*/
TVM_DLL Array<tir::LoweredFunc> lower(
TVM_DLL IRModule lower(
te::Schedule sch,
const Array<te::Tensor>& args,
const std::string& name,
......@@ -59,44 +59,43 @@ TVM_DLL Array<tir::LoweredFunc> lower(
const BuildConfig& config);
/*!
* \brief Build a device and host module for a specific target from an array of lowered functions.
* \brief Build a device and host module for a specific target from an IRModule.
* \param funcs The functions to be built.
* \param target The target device to build for.
* \param target_host The target for building host code. To use the default, pass Target()
* \param config The build configuration.
* \return The built module.
*/
TVM_DLL runtime::Module build(const Array<tir::LoweredFunc>& funcs,
TVM_DLL runtime::Module build(const IRModule& funcs,
const Target& target,
const Target& target_host,
const BuildConfig& config);
/*!
* \brief Build a device and host module for a specific target from a map
* contains target to a list of lowered functions pairs. This function is used
* contains target to IRModule. This function is used
* for heterogeneous build.
* \param input The map contains target to a list of lowered functions pairs.
* \param input The map contains target to an IRModule.
* \param target_host The target for building host code. To use the default,
* pass Target().
* \param config The build configuration.
* \return The built module that contains code for different processors.
*/
TVM_DLL runtime::Module build(const Map<Target, Array<tir::LoweredFunc>>& input,
TVM_DLL runtime::Module build(const Map<Target, IRModule>& input,
const Target& target_host,
const BuildConfig& config);
/*!
* \brief Build a device and host module for a specific target from a map
* contains target to a list of lowered functions pairs. This function is used
* contains target to IRModule. This function is used
* for heterogeneous build.
* \param input The map contains target string to a list of lowered functions
* pairs.
* \param input The map contains target string to an IRModule.
* \param target_host The target for building host code. To use the default,
* pass Target().
* \param config The build configuration.
* \return The built module that contains code for different processors.
*/
TVM_DLL runtime::Module build(const Map<std::string, Array<tir::LoweredFunc>>& input,
TVM_DLL runtime::Module build(const Map<std::string, IRModule>& input,
const Target& target_host,
const BuildConfig& config);
} // namespace tvm
......
include/tvm/ir/module.h
......@@ -297,6 +297,15 @@ class IRModule : public ObjectRef {
CHECK(ptr != nullptr);
return static_cast<IRModuleNode*>(ptr);
}
/*!
* \brief Construct an empty module.
*
* \returns The constructed module
*/
static IRModule Empty() {
return IRModule(Map<GlobalVar, BaseFunc>());
}
/*!
* \brief Construct a module from a standalone expression.
*
......
include/tvm/target/codegen.h
......@@ -27,7 +27,6 @@
#include <tvm/runtime/packed_func.h>
#include <tvm/ir/module.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/target/target.h>
#include <string>
......@@ -42,17 +41,6 @@ using runtime::TVMArgs;
using runtime::TVMRetValue;
/*!
* \brief Temporary backward compatible function to convert a list
* of LoweredFunc to a IRModule of PrimfFuncs
* \param funcs The input lowered function.
* \return The IRModule.
*
* \note This function is only used for code refactor and will be
* removed once the refactor completes.
*/
IRModule ToIRModule(const Array<tir::LoweredFunc>& funcs);
/*!
* \brief Build a module from array of lowered function.
* \param mod The Module to be built
* \param target The target to be built.
......
include/tvm/tir/analysis.h
......@@ -24,9 +24,12 @@
#ifndef TVM_TIR_ANALYSIS_H_
#define TVM_TIR_ANALYSIS_H_
#include <tvm/ir/module.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/function.h>
#include <tvm/tir/stmt.h>
namespace tvm {
namespace tir {
......@@ -59,6 +62,18 @@ struct ExprDeepEqual {
*/
Array<Var> UndefinedVars(const Stmt& stmt, const Array<Var>& defs);
/*!
* \brief Verify if memory accesses are legal for a specific target device type.
*
* In the case that tgt is cuda, if not all workload is bound with
* threads, CPU code is generated that tries to access GPU memory,
* which is illegal. This pass performs verification for this case.
*
* \param mod The module to be verified.
* \return Success of memory verification.
*/
void VerifyMemory(const IRModule& mod);
} // namespace tir
} // namespace tvm
#endif // TVM_TIR_ANALYSIS_H_
include/tvm/tir/ir_pass.h
......@@ -31,7 +31,6 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/function.h>
#include <tvm/tir/lowered_func.h>
#include <unordered_map>
#include <unordered_set>
......@@ -367,60 +366,6 @@ Stmt HoistIfThenElse(Stmt stmt);
Stmt NarrowDataType(Stmt stmt, int target_bits);
/*!
* \brief Make an user callable API LoweredFunc.
*
* The main task of this function is to create code to :
* - Map the values in the api_args to Var that is required by body.
* - Insert assertions to check type/value of the passed arguments.
*
* \param body The body of the function.
* \param name The name of the function.
* \param api_args Arguments to the function, can be either Var, or Buffer
* \param num_unpacked_args Number of arguments that
* are processed in plain form instead of packed form.
* \param is_restricted Whether the caller can guarantee that each buffer argument do not overlap.
* It is recommended to set to true for optimized code if such invariant holds.
*
* \return a LoweredFunc with the specified signiture.
*
* \note
* The function signature have two cases
*
* let num_packed_args = len(api_args) - num_unpacked_args;
*
* if num_packed_args is zero:
* f(api_arg_0, api_arg_1, .., api_arg_n) where n == len(api_args)
*
* if num_packed_args is not zero:
* f(TVMArg* packed_args, int* packed_arg_type_ids, int num_packed_args,
* api_arg_k, api_arg_k+1, ... api_arg_n,
* TVMValue* out_ret_val, int* out_ret_tcode)
*
* where n == len(api_args), k == num_packed_args
*
* There is no thread_axis in generated function.
*/
LoweredFunc MakeAPI(Stmt body,
std::string name,
Array<ObjectRef> api_args,
int num_unpacked_args,
bool is_restricted);
/*!
* \brief Remap the thread axis
*
* This can be used to get equivalent program which uses
* threadIdx.y in place of threadIdx.x by passing
* {"threadIdx.x": thread_axis("threadIdx.y")}
*
*
* \param f The device function to be lowered.
* \param axis_map The map from StringImm -> ItrVar
* \return Transformed function.
*/
LoweredFunc RemapThreadAxis(LoweredFunc f, Map<PrimExpr, IterVar> axis_map);
/*!
* \brief Rewrite the pointer content type of arguments,
* as well as Alloc internal to the function to use
* the most frequently accessed type for load/store
......@@ -433,31 +378,6 @@ LoweredFunc RemapThreadAxis(LoweredFunc f, Map<PrimExpr, IterVar> axis_map);
PrimFunc PointerValueTypeRewrite(PrimFunc f);
/*!
* \brief Lower custom datatypes.
*
* See tvm::datatypes::Registry for more information on adding custom datatypes.
*
* \param f The device function to be lowered.
* \param target The target device.
* \return Transformed function.
*/
LoweredFunc LowerCustomDatatypes(LoweredFunc f, const std::string& target);
/*!
* \brief Verify if memory accesses are legal for a specific target device type.
*
* In the case that tgt is cuda, if not all workload is bound with
* threads, CPU code is generated that tries to access GPU memory,
* which is illegal. This pass performs verification for this case.
*
* \param func The function to be verified.
* \param device_type The target device type.
* \return Success of memory verification.
*/
bool VerifyMemory(LoweredFunc func, int device_type);
/*!
* \brief Verify the correctness of a GPU code
* It will check the whether the amount of memory usage or the number of threads
* in a block exceeds the limit
......
include/tvm/tir/lowered_func.h deleted 100644 → 0
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file tvm/tir/lowered_func.h
* \brief Information about a lowered TVM function.
* This data structure is final step toward codegen.
*/
#ifndef TVM_TIR_LOWERED_FUNC_H_
#define TVM_TIR_LOWERED_FUNC_H_
#include <tvm/node/container.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt.h>
#include <string>
namespace tvm {
namespace tir {
// Internal node container of lowered function.
class LoweredFuncNode;
/*!
* \brief LoweredFunc represents function after lowering.
* This is the final IR representation before codegen.
*/
class LoweredFunc : public FunctionRef {
public:
LoweredFunc() {}
explicit LoweredFunc(ObjectPtr<Object> n) : FunctionRef(n) {}
/*!
* \brief access the internal node container
* \return the pointer to the internal node container
*/
inline const LoweredFuncNode* operator->() const;
/*! \brief specify container node */
using ContainerType = LoweredFuncNode;
};
/*! \brief specific type of lowered function */
enum LoweredFuncType : int {
/*! \brief Function that can mix device and host calls */
kMixedFunc = 0,
/*! \brief Only contains host code */
kHostFunc = 1,
/*! \brief Only contains device code */
kDeviceFunc = 2
};
/*! \brief Node container of LoweredFunc */
class LoweredFuncNode : public tir::FunctionBaseNode {
public:
/*! \brief The name of the function */
std::string name;
/*!
* \brief The arguments of the function
* This function can only take pod type(int, float) and void* as arguments.
*/
Array<Var> args;
/*!
* \brief The IterVar axis of threads
* Each axis need host function to specify a size.
* \note Calling convention into LoweredFunc
*
* Assume we have a LoweredFunc f, a call into f
* Call(f, arg1, arg2, ..., arg_n,
* size_axis_1, size_axis_2, ... size_axis_m)
*
* Here n = len(args), m = len(thread_axis)
*
* The CodeGen should take this and translate this call
* to corresponding API specific kernel launchs or function calls.
*/
Array<IterVar> thread_axis;
/*!
* \brief The hint data type of Var handles defined in LetStmt
* Can be used as hint when generating type signiture.
* The creation rule is given by
* handle_data_type[var_handle] = make_const(the_type, 0);
*
* \note Expr is used instead Type, because Type cannot be hold by Map.
* constant Expr of given type is used.
*/
Map<Var, PrimExpr> handle_data_type;
/*! \brief The type of the function */
LoweredFuncType func_type{kMixedFunc};
/*! \brief Whether this function is packed function */
bool is_packed_func{true};
/*!
* \brief Whether function ensures that argument pointers do not alias.
* This corresponds to restrict keyword in C.
*/
bool is_restricted{true};
/*! \brief The body statment of the function */
Stmt body;
/*! \return name of the operation */
const std::string& func_name() const final {
return name;
}
// there is no return value, but return 1
// to enable Call into this function.
int num_outputs() const final {
return 1;
}
void VisitAttrs(AttrVisitor* v) {
v->Visit("name", &name);
v->Visit("args", &args);
v->Visit("thread_axis", &thread_axis);
v->Visit("handle_data_type", &handle_data_type);
v->Visit("func_type", &func_type);
v->Visit("is_packed_func", &is_packed_func);
v->Visit("is_restricted", &is_restricted);
v->Visit("body", &body);
}
static constexpr const char* _type_key = "LoweredFunc";
TVM_DECLARE_FINAL_OBJECT_INFO(LoweredFuncNode, Object);
};
// Implementations of inline functions
inline const LoweredFuncNode* LoweredFunc::operator->() const {
return static_cast<const LoweredFuncNode*>(get());
}
} // namespace tir
} // namespace tvm
namespace std {
template <>
struct hash<::tvm::tir::LoweredFunc> : public tvm::ObjectHash {
};
}
#endif // TVM_TIR_LOWERED_FUNC_H_
include/tvm/tir/transform.h
......@@ -59,6 +59,61 @@ TVM_DLL Pass CreatePrimFuncPass(const runtime::TypedPackedFunc<
const tvm::Array<tvm::PrimExpr>& required);
/*!
* \brief Transform the high-level PrimFunc to a low-level version
* that can be used as an API function.
*
*
* The main task of this function is to create code to :
* - Map the values in the api_args to Var that is required by body.
* - Insert assertions to check type/value of the passed arguments.
*
* \param num_unpacked_args Number of arguments that
* are processed in plain form instead of packed form.
*
* \note
* The function signature have two cases
*
* let num_packed_args = len(api_args) - num_unpacked_args;
*
* if num_packed_args is zero:
* f(api_arg_0, api_arg_1, .., api_arg_n) where n == len(api_args)
*
* if num_packed_args is not zero:
* f(TVMArg* packed_args, int* packed_arg_type_ids, int num_packed_args,
* api_arg_k, api_arg_k+1, ... api_arg_n,
* TVMValue* out_ret_val, int* out_ret_tcode)
*
* where n == len(api_args), k == num_packed_args
*
* \return The pass.
*/
TVM_DLL Pass MakePackedAPI(int num_unpacked_args);
/*!
* \brief Remap the thread axis
*
* This can be used to get equivalent program which uses
* threadIdx.y in place of threadIdx.x by passing
* {"threadIdx.x": thread_axis("threadIdx.y")}
*
*
* \return The pass.
*/
TVM_DLL Pass RemapThreadAxis(Map<PrimExpr, IterVar> axis_map);
/*!
* \brief Lower custom datatypes.
*
* See tvm::datatypes::Registry for more information on adding custom datatypes.
*
* \return The pass.
*/
TVM_DLL Pass LowerCustomDatatypes();
/*!
* \brief Bind the device type ofthe function to be
* the device_type specified in the target attribute.
*
......
python/tvm/driver/build_module.py
......@@ -17,9 +17,6 @@
# pylint: disable=invalid-name
"""The build utils in python.
This module provides the functions to transform schedule to
LoweredFunc and compiled Module.
"""
import warnings
......@@ -30,7 +27,6 @@ from tvm.ir import container
from tvm.ir import CallingConv
from tvm.target import codegen, BuildConfig
from tvm.tir import ir_pass
from tvm.tir.stmt import LoweredFunc
from tvm.te import tensor
from tvm.te import schedule
from tvm import target as _target
......@@ -136,8 +132,8 @@ def lower(sch,
Returns
-------
f : LoweredFunc or Stmt
The result function, if with_api_wrapper=False
m : IRModule or Stmt
The result IRModule, if simple_mode=False
Then the Stmt before make api is returned.
"""
cfg = BuildConfig.current()
......@@ -199,16 +195,21 @@ def lower(sch,
if simple_mode:
return stmt
return ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)
f = tvm.tir.PrimFunc(arg_list, stmt).with_attr(
"global_symbol", tvm.runtime.String(name))
if cfg.restricted_func:
f = f.with_attr("tir.no_alias", True)
mod = tvm.IRModule({name: f})
return tvm.tir.transform.MakePackedAPI()(mod)
def _build_for_device(flist, target, target_host):
def _build_for_device(input_mod, target, target_host):
"""Build the lowered functions for a device with the given compilation
target.
Parameters
----------
flist : list of LoweredFunc
input_mod : IRModule
The schedule to be built.
target : str or :any:`tvm.target.Target`
......@@ -219,8 +220,8 @@ def _build_for_device(flist, target, target_host):
Returns
-------
fhost : list of LoweredFunc
A list of lowered functions for the host.
fhost : IRModule
The host IRModule.
mdev : tvm.module
A module that contains device code.
......@@ -229,14 +230,13 @@ def _build_for_device(flist, target, target_host):
target_host = _target.create(target_host)
device_type = ndarray.context(target.target_name, 0).device_type
for func in flist:
if not ir_pass.VerifyMemory(func, device_type):
raise ValueError(
"Direct host side access to device memory is detected in %s. "
"Did you forget to bind?" % func.name)
mod_mixed = input_mod
mod_mixed = tvm.tir.transform.Apply(lambda f: f.with_attr("target", target))(mod_mixed)
tvm.tir.analysis.verify_memory(mod_mixed)
mod_mixed = tvm.testing.LoweredFuncsToIRModule(flist)
opt_mixed = [tvm.tir.transform.Apply(lambda f: f.with_attr("target", target))]
opt_mixed = []
if len(mod_mixed.functions) == 1:
opt_mixed += [tvm.tir.transform.Apply(lambda f: f.with_attr("tir.is_entry_func", True))]
if BuildConfig.current().detect_global_barrier:
opt_mixed += [tvm.tir.transform.ThreadSync("global")]
opt_mixed += [tvm.tir.transform.ThreadSync("shared"),
......@@ -292,7 +292,7 @@ def build(inputs,
Parameters
----------
inputs : tvm.te.Schedule, LoweredFunc, or dict of target to LoweredFunc list
inputs : tvm.te.Schedule, IRModule, or dict of target to IRModule
The schedule to be built
args : list of Buffer or Tensor or Var, optional
......@@ -326,7 +326,7 @@ def build(inputs,
________
There are two typical example uses of this function depending on the type
of the argument `inputs`:
1. it is a list of lowered functions:
1. it is an IRModule.
.. code-block:: python
......@@ -335,10 +335,10 @@ def build(inputs,
B = te.placeholder((n,), name='B')
C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
s = tvm.te.create_schedule(C.op)
f = tvm.lower(s, [A, B, C], name="test_add")
m = tvm.build(f, target="llvm")
m = tvm.lower(s, [A, B, C], name="test_add")
rt_mod = tvm.build(m, target="llvm")
2. it is a dict of compilation target to list of lowered functions:
2. it is a dict of compilation target to IRModule.
.. code-block:: python
......@@ -349,9 +349,9 @@ def build(inputs,
s1 = tvm.te.create_schedule(C.op)
with tvm.target.cuda() as cuda_tgt:
s2 = topi.cuda.schedule_injective(cuda_tgt, [C])
f1 = tvm.lower(s1, [A, B, C], name="test_add1")
f2 = tvm.lower(s2, [A, B, C], name="test_add2")
m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm")
m1 = tvm.lower(s1, [A, B, C], name="test_add1")
m2 = tvm.lower(s2, [A, B, C], name="test_add2")
rt_mod = tvm.build({"llvm": m1, "cuda": m2}, target_host="llvm")
Note
----
......@@ -360,45 +360,36 @@ def build(inputs,
if isinstance(inputs, schedule.Schedule):
if args is None:
raise ValueError("args must be given for build from schedule")
flist = lower(inputs, args,
input_mod = lower(inputs, args,
name=name,
binds=binds)
if isinstance(flist, LoweredFunc):
flist = [flist]
elif isinstance(inputs, LoweredFunc):
if args:
raise ValueError("args must be done when build from LoweredFunc.")
flist = [inputs]
elif isinstance(inputs, (list, tuple, container.Array)):
flist = inputs
merged_mod = tvm.IRModule({})
for x in inputs:
merged_mod.update(x)
input_mod = merged_mod
elif isinstance(inputs, tvm.IRModule):
input_mod = inputs
elif not isinstance(inputs, (dict, container.Map)):
raise ValueError("inputs must be Schedule, LoweredFunc, list of "
"LoweredFunc, or dict of target to list of "
"LoweredFunc.")
raise ValueError("inputs must be Schedule, IRModule or dict of target to IRModule")
if not isinstance(inputs, (dict, container.Map)):
target = _target.Target.current() if target is None else target
target = target if target else "llvm"
target_flist = {target: flist}
target_input_mod = {target: input_mod}
else:
target_flist = inputs
target_input_mod = inputs
for tar, flist in target_flist.items():
for tar, mod in target_input_mod.items():
if not isinstance(tar, (str, _target.Target)):
raise ValueError("The key of inputs must be str or "
"_target.Target when inputs is dict.")
fname_set = set()
for x in flist:
if not isinstance(x, LoweredFunc):
raise ValueError("inputs must be Schedule, LoweredFunc, list "
"of LoweredFunc, or dict of str to list of "
"LoweredFunc.")
if x.name in fname_set:
raise ValueError("Duplicate function name %s" % x.name)
fname_set.add(x.name)
if not isinstance(mod, tvm.IRModule):
raise ValueError("inputs must be Schedule, IRModule,"
"or dict of str to IRModule.")
if not target_host:
for tar, _ in target_flist.items():
for tar, _ in target_input_mod.items():
tar = _target.create(tar)
device_type = ndarray.context(tar.target_name, 0).device_type
if device_type == ndarray.cpu(0).device_type:
......@@ -410,8 +401,8 @@ def build(inputs,
mod_host_all = tvm.IRModule({})
device_modules = []
for tar, flist in target_flist.items():
mod_host, mdev = _build_for_device(flist, tar, target_host)
for tar, input_mod in target_input_mod.items():
mod_host, mdev = _build_for_device(input_mod, tar, target_host)
mod_host_all.update(mod_host)
device_modules.append(mdev)
......
python/tvm/relay/backend/_backend.py
......@@ -17,7 +17,6 @@
"""The interface of expr function exposed from C++."""
import tvm._ffi
import tvm.driver
from tvm.ir import container as _container
@tvm._ffi.register_func("relay.backend.lower")
......@@ -40,7 +39,7 @@ def lower(sch, inputs, func_name, source_func):
Returns
-------
lowered_funcs : List[tvm.LoweredFunc]
mod : tvm.IRModule
The result of lowering.
"""
# pylint: disable=broad-except, import-outside-toplevel
......@@ -56,20 +55,17 @@ def lower(sch, inputs, func_name, source_func):
msg += "-----------------------------\n"
msg += source_func.astext()
raise RuntimeError(msg)
return f if isinstance(
f, (_container.Array, tuple, list)) else [f]
return f
@tvm._ffi.register_func("relay.backend.build")
def build(funcs, target, target_host=None):
def build(mod, target, target_host=None):
"""Backend build function.
Parameters
----------
funcs : List[tvm.LoweredFunc] or Dict[str, List[tvm.LoweredFunc]]
A list of lowered functions or dictionary mapping from targets to
lowered functions.
mod : tvm.IRModule or Dict[str, tvm.IRModule]
Input module
target : tvm.Target
The target to run the code on.
......@@ -84,7 +80,7 @@ def build(funcs, target, target_host=None):
"""
if target_host == "":
target_host = None
return tvm.driver.build(funcs, target=target, target_host=target_host)
return tvm.driver.build(mod, target=target, target_host=target_host)
@tvm._ffi.register_func("relay._tensor_value_repr")
......
python/tvm/relay/backend/graph_runtime_codegen.py
......@@ -48,7 +48,7 @@ class GraphRuntimeCodegen(object):
self._get_graph_json = self._mod["get_graph_json"]
self._list_params_name = self._mod["list_params_name"]
self._get_param_by_name = self._mod["get_param_by_name"]
self._get_lowered_funcs = self._mod["get_lowered_funcs"]
self._get_irmodule = self._mod["get_irmodule"]
self._setup(mod, target)
def _setup(self, mod, target):
......@@ -74,14 +74,14 @@ class GraphRuntimeCodegen(object):
-------
graph_json : str
The graph json that can be consumed by runtime.
lowered_funcs : List[tvm.LoweredFunc] or Dict[str, List[tvm.LoweredFunc]]
mod : IRModule or Dict[str, IRModule]
The lowered functions.
params : Dict[str, tvm.nd.NDArray]
Additional constant parameters.
"""
self._codegen(func)
graph_json = self._get_graph_json()
lowered_func = self._get_lowered_funcs()
lowered_func = self._get_irmodule()
param_names = self._list_params_name()
params = {}
for name in param_names:
......
python/tvm/runtime/__init__.py
......@@ -28,3 +28,4 @@ from .object_generic import convert_to_object, convert, const
from .ndarray import context, cpu, gpu, opencl, cl, vulkan, metal, mtl
from .ndarray import vpi, rocm, opengl, ext_dev, micro_dev
from .module import load_module, enabled, system_lib
from .container import String
python/tvm/target/build_config.py
......@@ -20,9 +20,7 @@ import tvm._ffi
import tvm.ir
from tvm.runtime import Object
from tvm.ir import container
from tvm.tir import Stmt
from tvm.tir.stmt import LoweredFunc
from . import _ffi_api
......@@ -48,17 +46,13 @@ class DumpIR(object):
def dump(*args, **kwargs):
"""dump function"""
retv = func(*args, **kwargs)
if not isinstance(retv, (Stmt, LoweredFunc, container.Array)):
if not isinstance(retv, (Stmt,)):
return retv
fname = func.func_name if hasattr(func, 'func_name') else func.__name__
pname = str(self._pass_id) + "_" + fname + "_ir.cc"
with open(pname, "a") as f:
out = retv.body if isinstance(retv, LoweredFunc) else retv
out = retv
f.write(str(out))
if isinstance(retv, container.Array):
for x in retv:
out = x.body if isinstance(x, LoweredFunc) else x
f.write("---------%s\n%s\n-----------\n"%(x.name, str(out)))
self._pass_id += 1
return retv
return dump
......
python/tvm/testing.py
......@@ -14,9 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=invalid-name
""" TVM testing utilities """
import logging
import numpy as np
import tvm
import tvm._ffi
......@@ -165,4 +168,40 @@ def check_numerical_grads(function, input_values, grad_values, function_value=No
x_name, grad.shape, dist, max_diff, avg_diff)
def MakeAPILegacy(stmt, name, args, num_unpacked_args, noalias):
"""Legacy adapter to build a Module from statement.
Used for migrating existing test cases only.
Parameters
----------
stmt: Stmt
The input statement.
name: str
The name of the funciton.
args: list of Buffer or Vars
The function arguments
num_unpacked_args: int
Number of unpacked arguments.
nolias: bool
Whether allow noalias.
Returns
-------
mod : IRModule
The created IRModule.
"""
f = tvm.tir.PrimFunc(args, stmt).with_attr(
"global_symbol", tvm.runtime.String(name))
f = f.with_attr("tir.is_entry_func", True)
if noalias:
f = f.with_attr("tir.no_alias", True)
mod = tvm.IRModule({name: f})
return tvm.tir.transform.MakePackedAPI(num_unpacked_args)(mod)
tvm._ffi._init_api("testing", __name__)
python/tvm/tir/__init__.py
......@@ -29,7 +29,7 @@ from .expr import IterVar, Any
from .stmt import Stmt, LetStmt, AssertStmt, ProducerConsumer, For
from .stmt import BufferStore, Store, Provide, Allocate, AttrStmt, Free, Realize, SeqStmt
from .stmt import IfThenElse, Evaluate, Prefetch, LoweredFunc, stmt_seq, stmt_list
from .stmt import IfThenElse, Evaluate, Prefetch, stmt_seq, stmt_list
from .function import PrimFunc
......
python/tvm/tir/analysis/analysis.py
......@@ -55,3 +55,14 @@ def expr_deep_equal(lhs, rhs):
tvm.ir.structural_equal
"""
return _ffi_api.expr_deep_equal(lhs, rhs)
def verify_memory(mod):
"""Verify if module contains illegal host side direct memory access.
Parameters
----------
mod: tvm.IRModule
The module to be verified.
"""
_ffi_api.verify_memory(mod)
python/tvm/tir/function.py
......@@ -18,6 +18,7 @@
import tvm._ffi
import tvm.runtime
from tvm.runtime import Object
from tvm.ir import BaseFunc
from .buffer import Buffer
from .expr import Var
......@@ -54,6 +55,7 @@ class PrimFunc(BaseFunc):
param_list = []
buffer_map = {} if buffer_map is None else buffer_map
for x in params:
x = tvm.runtime.convert(x) if not isinstance(x, Object) else x
if isinstance(x, Buffer):
var = Var(x.name, dtype="handle")
param_list.append(var)
......
python/tvm/tir/stmt.py
......@@ -385,14 +385,6 @@ class Prefetch(Stmt):
_ffi_api.Prefetch, func, value_index, dtype, bounds)
@tvm._ffi.register_object
class LoweredFunc(Object):
"""Represent a LoweredFunc in TVM."""
MixedFunc = 0
HostFunc = 1
DeviceFunc = 2
def stmt_seq(*args):
"""Make sequence of statements
......
python/tvm/tir/transform/transform.py
......@@ -60,6 +60,36 @@ def Filter(fcond):
return _fpass.prim_func_pass(_transform, opt_level=0)
def LowerCustomDatatypes():
"""Lower custom datatypes.
See tvm::datatypes::Registry for more information on adding custom datatypes.
Returns
-------
fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.LowerCustomDatatypes()
def MakePackedAPI(num_unpacked_params=0):
"""Transform the PrimFuncs in the module to a packed func API.
Parameters
----------
num_unpacked_params : int
Number of parameters that we hope to directly pass via normal arguments
following the PackedFunc input signature.
Returns
-------
fpass : tvm.ir.transform.Pass
The result pass
"""
return _ffi_api.MakePackedAPI(num_unpacked_params)
def BindDeviceType():
"""Bind the device type of the function to be
the device_type specified in the target attribute.
......
src/contrib/hybrid/codegen_hybrid.h
......@@ -27,7 +27,6 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/target/codegen.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/te/schedule.h>
#include <map>
#include <string>
......
src/driver/driver_api.cc
......@@ -26,8 +26,10 @@
#include <tvm/te/operation.h>
#include <tvm/tir/transform.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/ir_pass.h>
#include <tvm/target/codegen.h>
#include <tvm/runtime/container.h>
#include <tvm/runtime/registry.h>
#include <algorithm>
......@@ -39,7 +41,6 @@ namespace tvm {
using runtime::TVMArgs;
using runtime::TVMRetValue;
using runtime::PackedFunc;
using tir::LoweredFunc;
bool LLVMEnabled() {
const runtime::PackedFunc* pf = runtime::Registry::Get("target.build.llvm");
......@@ -166,17 +167,6 @@ tir::Stmt BuildStmt(te::Schedule sch,
return stmt;
}
Array<LoweredFunc> lower(te::Schedule sch,
const Array<te::Tensor>& args,
const std::string& name,
const std::unordered_map<te::Tensor, tir::Buffer>& binds,
const BuildConfig& config) {
Array<ObjectRef> out_arg_list;
auto stmt = BuildStmt(sch, args, binds, true, &out_arg_list, config);
return Array<LoweredFunc>({ tir::MakeAPI(stmt, name, out_arg_list, 0, config->restricted_func) });
}
transform::Pass BindTarget(Target target) {
auto fpass = [target](tir::PrimFunc f, IRModule m, transform::PassContext ctx) {
return WithAttr(std::move(f), tvm::attr::kTarget, target);
......@@ -198,18 +188,46 @@ transform::Pass FilterBy(FCond fcond) {
}
IRModule lower(te::Schedule sch,
const Array<te::Tensor>& args,
const std::string& name,
const std::unordered_map<te::Tensor, tir::Buffer>& binds,
const BuildConfig& config) {
Array<ObjectRef> out_arg_list;
auto stmt = BuildStmt(sch, args, binds, true, &out_arg_list, config);
Array<tir::Var> params;
Map<tir::Var, tir::Buffer> buffer_map;
for (auto var : out_arg_list) {
if (auto* n = var.as<tir::VarNode>()) {
params.push_back(GetRef<tir::Var>(n));
} else {
tir::Buffer buffer = Downcast<tir::Buffer>(var);
tir::Var bptr(buffer->name, DataType::Handle());
params.push_back(bptr);
buffer_map.Set(bptr, buffer);
}
}
auto f = tir::PrimFunc(params, stmt, VoidType(), buffer_map);
f = WithAttr(std::move(f), "global_symbol", runtime::String(name));
if (config->restricted_func) {
f = WithAttr(std::move(f), "tir.no_alias", Integer(1));
}
auto mod = IRModule(Map<GlobalVar, BaseFunc>({{GlobalVar(name), f}}));
return tir::transform::MakePackedAPI(0)(mod);
}
std::pair<IRModule, IRModule>
split_dev_host_funcs(const Array<LoweredFunc>& funcs,
split_dev_host_funcs(IRModule mod_mixed,
const Target& target,
const Target& target_host,
const BuildConfig& config) {
for (const auto& x : funcs) {
CHECK(tir::VerifyMemory(x, target->device_type))
<< "Direct host side access to device memory is detected in "
<< x->func_name() << ". Did you forget to bind?";
}
IRModule mod_mixed = codegen::ToIRModule(funcs);
mod_mixed = BindTarget(target)(std::move(mod_mixed));
tir::VerifyMemory(mod_mixed);
Array<tvm::transform::Pass> mixed_pass_list = {BindTarget(target)};
if (config->detect_global_barrier) {
......@@ -274,10 +292,9 @@ split_dev_host_funcs(const Array<LoweredFunc>& funcs,
// Build for heterogeneous execution.
runtime::Module build(const Map<Target, Array<LoweredFunc>>& inputs,
runtime::Module build(const Map<Target, IRModule>& inputs,
const Target& target_host,
const BuildConfig& config) {
Array<LoweredFunc> fhost_all;
std::vector<runtime::Module> device_modules;
Target target_host_val = target_host;
......@@ -319,10 +336,10 @@ runtime::Module build(const Map<Target, Array<LoweredFunc>>& inputs,
}
// Build for heterogeneous execution when target is a string.
runtime::Module build(const Map<std::string, Array<LoweredFunc>>& inputs,
runtime::Module build(const Map<std::string, IRModule>& inputs,
const Target& target_host,
const BuildConfig& config) {
Map<Target, Array<LoweredFunc>> updated_input;
Map<Target, IRModule> updated_input;
for (const auto& it : inputs) {
auto target = Target::Create(it.first);
if (target->device_name == "vta") {
......@@ -334,11 +351,11 @@ runtime::Module build(const Map<std::string, Array<LoweredFunc>>& inputs,
}
// Build for homogeneous execution.
runtime::Module build(const Array<LoweredFunc>& funcs,
runtime::Module build(const IRModule& funcs,
const Target& target,
const Target& target_host,
const BuildConfig& config) {
Map<Target, Array<LoweredFunc>> inputs = {{target, funcs}};
Map<Target, IRModule> inputs = {{target, funcs}};
return build(inputs, target_host, config);
}
......
src/relay/backend/build_module.cc
......@@ -38,7 +38,6 @@ namespace tvm {
namespace relay {
namespace backend {
using tir::LoweredFunc;
using TargetsMap = Map<tvm::Integer, tvm::Target>;
using namespace tvm::relay::transform;
......@@ -78,16 +77,16 @@ struct GraphCodegen {
}
Array<tvm::runtime::Module> GetExternalModules() {
return CallFunc<Array<tvm::runtime::Module> >("get_external_modules", nullptr);
return CallFunc<Array<tvm::runtime::Module>>("get_external_modules", nullptr);
}
Map<std::string, Array<LoweredFunc> > GetLoweredFunc() {
return CallFunc<Map<std::string, Array<LoweredFunc> > >("get_lowered_funcs", nullptr);
Map<std::string, IRModule> GetIRModule() {
return CallFunc<Map<std::string, IRModule>>("get_irmodule", nullptr);
}
std::unordered_map<std::string, tvm::runtime::NDArray> GetParams() {
std::unordered_map<std::string, tvm::runtime::NDArray> ret;
auto names = CallFunc<Array<tvm::PrimExpr> >("list_params_name", nullptr);
auto names = CallFunc<Array<tvm::PrimExpr>>("list_params_name", nullptr);
for (auto expr : names) {
auto key = expr.as<tir::StringImmNode>()->value;
ret[key] = CallFunc<runtime::NDArray>("get_param_by_name", key);
......@@ -152,9 +151,9 @@ class RelayBuildModule : public runtime::ModuleNode {
this->SetParam(kv.first, kv.second->data);
}
});
} else if (name == "get_lowered_funcs") {
} else if (name == "get_irmodule") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
*rv = this->graph_codegen_->GetLoweredFunc();
*rv = this->graph_codegen_->GetIRModule();
});
} else if (name == "get_external_modules") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
......@@ -452,7 +451,7 @@ class RelayBuildModule : public runtime::ModuleNode {
ret_.graph_json = graph_codegen_->GetJSON();
ret_.params = graph_codegen_->GetParams();
auto lowered_funcs = graph_codegen_->GetLoweredFunc();
auto lowered_funcs = graph_codegen_->GetIRModule();
// When there is no lowered_funcs due to reasons such as optimization.
if (lowered_funcs.size() == 0) {
......
src/relay/backend/compile_engine.h
......@@ -27,7 +27,6 @@
#include <tvm/node/structural_equal.h>
#include <tvm/node/structural_hash.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/runtime/module.h>
#include <tvm/relay/analysis.h>
#include <tvm/relay/expr.h>
......@@ -82,7 +81,8 @@ struct CachedFuncNode : public Object {
/*! \brief The schedule to the function */
te::Schedule schedule;
/*! \brief The lowered functions to support the function. */
tvm::Array<tir::LoweredFunc> funcs;
IRModule funcs = IRModule::Empty();
/*! \brief Parameter usage states in the shape function. */
tvm::Array<Integer> shape_func_param_states;
......
src/relay/backend/graph_runtime_codegen.cc
......@@ -55,7 +55,7 @@ using TargetsMap = std::unordered_map<int, Target>;
/*! \brief Lowered outputs */
struct LoweredOutput {
std::string graph_json;
Map<std::string, Array<tir::LoweredFunc> > lowered_funcs;
Map<std::string, IRModule> lowered_funcs;
Array<tvm::runtime::Module> external_mods;
std::unordered_map<std::string, tvm::runtime::NDArray> params;
};
......@@ -214,19 +214,14 @@ class GraphRuntimeCodegen
LoweredOutput ret;
ret.graph_json = os.str();
ret.params = params_;
for (auto& kv : lowered_funcs_) {
if (ret.lowered_funcs.count(kv.first) == 0) {
ret.lowered_funcs.Set(kv.first, Array<tir::LoweredFunc>());
}
auto& vec = ret.lowered_funcs[kv.first];
Array<tir::LoweredFunc> tmp;
for (auto f : kv.second) {
tmp.push_back(f);
}
for (auto f : vec) {
tmp.push_back(f);
ret.lowered_funcs.Set(kv.first, IRModule::Empty());
}
ret.lowered_funcs.Set(kv.first, tmp);
auto& mod = ret.lowered_funcs[kv.first];
mod->Update(kv.second);
ret.lowered_funcs.Set(kv.first, mod);
}
ret.external_mods = compile_engine_->LowerExternalFunctions();
return ret;
......@@ -457,12 +452,9 @@ class GraphRuntimeCodegen
CCacheKey key = (*pf0)(func, target);
CachedFunc lowered_func = (*pf1)(compile_engine_, key);
if (!lowered_funcs_.count(target->str())) {
lowered_funcs_[target->str()] = {};
lowered_funcs_[target->str()] = IRModule::Empty();
}
for (auto f : lowered_func->funcs) {
lowered_funcs_[target->str()].insert(f);
}
lowered_funcs_[target->str()]->Update(lowered_func->funcs);
return GraphAddCallNode(op,
_GetUniqueName(lowered_func->func_name),
lowered_func->func_name);
......@@ -602,8 +594,7 @@ class GraphRuntimeCodegen
/*! \brief plan memory of device result */
Map<Expr, Array<IntegerArray>> storage_device_map_;
/*! \brief lowered funcs */
std::unordered_map<std::string, std::unordered_set<tir::LoweredFunc, ObjectHash, ObjectEqual>>
lowered_funcs_;
std::unordered_map<std::string, IRModule> lowered_funcs_;
/*! \brief name map */
std::unordered_map<std::string, size_t> name_map_;
/*! \brief compile engine */
......@@ -655,7 +646,7 @@ class GraphRuntimeCodegenModule : public runtime::ModuleNode {
CHECK_GT(this->output_.params.count(key), 0);
*rv = this->output_.params[key];
});
} else if (name == "get_lowered_funcs") {
} else if (name == "get_irmodule") {
return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
*rv = this->output_.lowered_funcs;
});
......
src/relay/backend/vm/compiler.cc
......@@ -226,6 +226,7 @@ std::vector<int64_t> ToAllocTensorShape32(NDArray shape) {
return raw_shape;
}
class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
public:
VMFunctionCompiler(VMCompilerContext* context, TargetsMap targets, Target target_host)
......@@ -407,12 +408,15 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
CCacheKey key(func, target_host_);
auto cfunc = engine_->LowerShapeFunc(key);
int op_index = -1;
if (context_->seen_funcs.count(cfunc->funcs[0]) == 0) {
// pick the only function inside the context
CHECK_EQ(cfunc->funcs->functions.size(), 1);
auto pfunc = Downcast<tir::PrimFunc>((*cfunc->funcs->functions.begin()).second);
if (context_->seen_funcs.count(pfunc) == 0) {
op_index = context_->cached_funcs.size();
context_->cached_funcs.push_back(cfunc);
context_->seen_funcs[cfunc->funcs[0]] = op_index;
context_->seen_funcs[pfunc] = op_index;
} else {
op_index = context_->seen_funcs[cfunc->funcs[0]];
op_index = context_->seen_funcs[pfunc];
}
// Prepare input and output registers
......@@ -494,13 +498,14 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
context_->cached_funcs.push_back(cfunc);
} else {
// TODO(jroesch): support lowered funcs for multiple targets
CHECK_EQ(cfunc->funcs.size(), 1);
if (context_->seen_funcs.find(cfunc->funcs[0]) == context_->seen_funcs.end()) {
CHECK_EQ(cfunc->funcs->functions.size(), 1);
auto pfunc = Downcast<tir::PrimFunc>((*cfunc->funcs->functions.begin()).second);
if (context_->seen_funcs.find(pfunc) == context_->seen_funcs.end()) {
op_index = context_->cached_funcs.size();
context_->cached_funcs.push_back(cfunc);
context_->seen_funcs[cfunc->funcs[0]] = op_index;
context_->seen_funcs[pfunc] = op_index;
} else {
op_index = context_->seen_funcs[cfunc->funcs[0]];
op_index = context_->seen_funcs[pfunc];
}
}
......@@ -862,11 +867,7 @@ void VMCompiler::Lower(IRModule mod,
// update primitive function map
size_t primitive_index = 0;
for (const auto& cfunc : context_.cached_funcs) {
if (cfunc->target->str() == "ext_dev") {
exec_->primitive_map.insert({cfunc->func_name, primitive_index++});
} else {
exec_->primitive_map.insert({cfunc->funcs[0]->name, primitive_index++});
}
}
}
......@@ -961,8 +962,6 @@ void VMCompiler::PopulateGlobalMap() {
}
void VMCompiler::Codegen() {
using tir::LoweredFunc;
if (!context_.module.defined()) {
LOG(WARNING) << "Did you forget to call VMCompiler::Lower?";
return;
......@@ -971,15 +970,21 @@ void VMCompiler::Codegen() {
if (cached_funcs.size() == 0) {
return;
}
std::unordered_map<std::string, Array<LoweredFunc>> funcs;
std::unordered_map<std::string, IRModule> funcs;
for (auto& cfunc : cached_funcs) {
std::string target_str = cfunc->target->str();
// NOTE: because module, is mutable, we need to make an
// explicit copy of the IRModule.
IRModule mod = cfunc->funcs;
mod.CopyOnWrite();
if (target_str == "ext_dev") {
continue;
} else if (funcs.count(target_str) == 0) {
funcs.emplace(target_str, Array<LoweredFunc>{cfunc->funcs[0]});
funcs.emplace(target_str, mod);
} else {
funcs[target_str].push_back(cfunc->funcs[0]);
funcs[target_str]->Update(mod);
}
}
......
src/relay/backend/vm/compiler.h
......@@ -76,7 +76,7 @@ struct VMCompilerContext {
// List of cached functions
std::vector<CachedFunc> cached_funcs;
// The functions that have been lowered.
std::unordered_map<tir::LoweredFunc, size_t, ObjectHash, ObjectEqual> seen_funcs;
std::unordered_map<tir::PrimFunc, size_t, ObjectHash, ObjectEqual> seen_funcs;
};
......
src/relay/transforms/gradient.cc
......@@ -22,7 +22,6 @@
* \brief API for Automatic Differentiation for the Relay IR.
*/
#include <tvm/ir/type_functor.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/te/operation.h>
#include <tvm/relay/expr_functor.h>
#include <tvm/relay/analysis.h>
......
src/target/build_common.h
......@@ -31,29 +31,12 @@
#include <tvm/tir/function.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt.h>
#include <tvm/tir/lowered_func.h>
#include <unordered_map>
#include <string>
#include "../runtime/meta_data.h"
namespace tvm {
namespace codegen {
// Extract function information from device function.
inline std::unordered_map<std::string, runtime::FunctionInfo>
ExtractFuncInfo(const Array<tir::LoweredFunc>& funcs) {
std::unordered_map<std::string, runtime::FunctionInfo> fmap;
for (tir::LoweredFunc f : funcs) {
runtime::FunctionInfo info;
for (size_t i = 0; i < f->args.size(); ++i) {
info.arg_types.push_back(f->args[i].dtype());
}
for (size_t i = 0; i < f->thread_axis.size(); ++i) {
info.thread_axis_tags.push_back(f->thread_axis[i]->thread_tag);
}
fmap[f->name] = info;
}
return fmap;
}
inline std::unordered_map<std::string, runtime::FunctionInfo>
ExtractFuncInfo(const IRModule& mod) {
......
src/target/codegen.cc
......@@ -43,50 +43,6 @@
namespace tvm {
namespace codegen {
// convert legacy LoweredFunc to PrimFunc.
tir::PrimFunc ToPrimFunc(tir::LoweredFunc from) {
// remap args to attach type annotations.
Array<tir::Var> args;
Map<tir::Var, PrimExpr> remap_vars;
for (auto var : from->args) {
auto it = from->handle_data_type.find(var);
if (it != from->handle_data_type.end()) {
tir::Var new_var(var->name_hint,
PointerType(PrimType((*it).second->dtype)));
args.push_back(new_var);
remap_vars.Set(var, new_var);
} else {
args.push_back(var);
}
}
tir::PrimFunc func(args, Substitute(from->body, remap_vars));
func = WithAttr(std::move(func), attr::kGlobalSymbol, runtime::String(from->name));
func = WithAttr(std::move(func), tir::attr::kDeviceThreadAxis, from->thread_axis);
if (from->func_type == tir::LoweredFuncType::kDeviceFunc) {
func = WithAttr(std::move(func),
attr::kCallingConv, Integer(CallingConv::kDeviceKernelLaunch));
}
if (from->is_restricted) {
func = WithAttr(std::move(func), tir::attr::kNoAlias, Integer(1));
}
return func;
}
IRModule ToIRModule(const Array<tir::LoweredFunc>& funcs) {
Map<GlobalVar, BaseFunc> functions;
for (size_t i = 0; i < funcs.size(); ++i) {
auto f = funcs[i];
tir::PrimFunc pf = ToPrimFunc(f);
if (i == 0) {
pf = WithAttr(std::move(pf), tir::attr::kIsEntryFunc, Integer(1));
}
functions.Set(GlobalVar(f->name), pf);
}
return IRModule(functions);
}
runtime::Module Build(IRModule mod, const Target& target) {
if (BuildConfig::Current()->disable_assert) {
mod = tir::transform::SkipAssert()(mod);
......@@ -284,9 +240,6 @@ runtime::Module PackImportsToLLVM(const runtime::Module& mod,
TVM_REGISTER_GLOBAL("target.Build")
.set_body_typed(Build);
TVM_REGISTER_GLOBAL("testing.LoweredFuncsToIRModule")
.set_body_typed(ToIRModule);
// Export two auxiliary function to the runtime namespace.
TVM_REGISTER_GLOBAL("runtime.ModulePackImportsToC")
.set_body_typed(PackImportsToC);
......
src/target/llvm/codegen_llvm.cc
......@@ -448,7 +448,7 @@ CodeGenLLVM::CreateDebugInfo(llvm::Module* module) {
auto debug_info = llvm::make_unique<CodeGenLLVM::DebugInfo>();
debug_info->di_builder_ = llvm::make_unique<llvm::DIBuilder>(*module);
#endif
// TODO(tulloch): pass this information through relay::Span classes to the LoweredFunc instance?
// TODO(tulloch): pass this information through relay::Span classes to the IRModule instance?
debug_info->file_ = debug_info->di_builder_->createFile("model.tvm", "/tmp/");
debug_info->compilation_unit_ = debug_info->di_builder_->createCompileUnit(
llvm::dwarf::DW_LANG_C, debug_info->file_, "TVM", 0, "", 0, "",
......
src/target/llvm/llvm_module.cc
......@@ -67,20 +67,23 @@ class LLVMModuleNode final : public runtime::ModuleNode {
} else if (name == "_get_target_triple") {
std::string target_triple = tm_->getTargetTriple().str();
return PackedFunc([target_triple](TVMArgs args, TVMRetValue *rv) {
* rv = target_triple;
*rv = target_triple;
});
}
if (ee_ == nullptr) LazyInitJIT();
// This LLVMModule is empty and no function can be retrieved.
if (entry_func_.empty()) return nullptr;
std::lock_guard<std::mutex> lock(mutex_);
const std::string& fname = (name == runtime::symbol::tvm_module_main ?
entry_func_ : name);
TVMBackendPackedCFunc faddr =
reinterpret_cast<TVMBackendPackedCFunc>(GetFunctionAddr(fname));
TVMBackendPackedCFunc faddr;
if (name == runtime::symbol::tvm_module_main) {
const char* entry_name = reinterpret_cast<const char*>(
GetGlobalAddr(runtime::symbol::tvm_module_main));
CHECK(entry_name != nullptr)
<< "Symbol " << runtime::symbol::tvm_module_main << " is not presented";
faddr = reinterpret_cast<TVMBackendPackedCFunc>(GetFunctionAddr(entry_name));
} else {
faddr = reinterpret_cast<TVMBackendPackedCFunc>(GetFunctionAddr(name));
}
if (faddr == nullptr) return PackedFunc();
return WrapPackedFunc(faddr, sptr_to_self);
}
......@@ -205,6 +208,7 @@ class LLVMModuleNode final : public runtime::ModuleNode {
std::unique_ptr<CodeGenLLVM> cg = CodeGenLLVM::Create(tm_.get());
std::vector<PrimFunc> funcs;
std::string entry_func;
for (auto kv : mod->functions) {
CHECK(kv.second->IsInstance<PrimFuncNode>())
<< "Can only lower IR Module with PrimFuncs";
......@@ -212,7 +216,7 @@ class LLVMModuleNode final : public runtime::ModuleNode {
if (f->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
auto global_symbol = f->GetAttr<runtime::String>(tvm::attr::kGlobalSymbol);
CHECK(global_symbol.defined());
entry_func_ = global_symbol;
entry_func = global_symbol;
}
funcs.push_back(f);
}
......@@ -225,8 +229,8 @@ class LLVMModuleNode final : public runtime::ModuleNode {
cg->AddFunction(f);
}
if (entry_func_.length() != 0) {
cg->AddMainFunction(entry_func_);
if (entry_func.length() != 0) {
cg->AddMainFunction(entry_func);
}
module_ = cg->Finish();
......@@ -321,13 +325,7 @@ class LLVMModuleNode final : public runtime::ModuleNode {
CHECK(ee_ != nullptr)
<< "Failed to initialize jit engine for " << mptr_->getTargetTriple();
ee_->runStaticConstructorsDestructors(false);
// setup context address.
// we will skip context setup if this LLVMModule is empty.
if (GetGlobalAddr(runtime::symbol::tvm_module_main) == 0)
return;
entry_func_ =
reinterpret_cast<const char*>(GetGlobalAddr(runtime::symbol::tvm_module_main));
if (void** ctx_addr = reinterpret_cast<void**>(
GetGlobalAddr(runtime::symbol::tvm_module_ctx))) {
*ctx_addr = this;
......@@ -356,8 +354,6 @@ class LLVMModuleNode final : public runtime::ModuleNode {
// The target configuration string
std::string target_;
// Name of entry function.
std::string entry_func_;
// JIT lock
std::mutex mutex_;
// execution engine
......
src/target/source/codegen_c.h
......@@ -29,7 +29,6 @@
#include <tvm/tir/function.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/target/codegen.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/runtime/container.h>
#include <string>
#include <vector>
......
src/target/spirv/codegen_spirv.h
......@@ -27,7 +27,6 @@
#include <tvm/arith/analyzer.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/lowered_func.h>
#include <vector>
#include <memory>
......
src/target/stackvm/codegen_stackvm.h
......@@ -26,7 +26,6 @@
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/lowered_func.h>
#include <tvm/target/codegen.h>
#include <string>
#include <vector>
......
src/tir/pass/verify_memory.cc → src/tir/analysis/verify_memory.cc
......@@ -22,8 +22,10 @@
* \brief Pass to check if memory accesses are legal.
*/
#include <tvm/tir/expr.h>
#include <tvm/tir/ir_pass.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/target/target.h>
#include <tvm/runtime/registry.h>
namespace tvm {
......@@ -44,7 +46,7 @@ class MemoryAccessVerifier final : protected StmtExprVisitor {
public:
/// Special member functions
//@{
explicit MemoryAccessVerifier(LoweredFunc f, int device_type)
explicit MemoryAccessVerifier(PrimFunc f, int device_type)
: func_(f), dev_type_(device_type) {}
virtual ~MemoryAccessVerifier() = default;
MemoryAccessVerifier(const MemoryAccessVerifier &) = delete;
......@@ -116,7 +118,7 @@ class MemoryAccessVerifier final : protected StmtExprVisitor {
CHECK(V) << "Invalid Variable\n";
// Variable is from function args. Return true.
if (V == func_->args[0].get()) return true;
if (V == func_->params[0].get()) return true;
// The value is expected to come from a tvm_struct_get Call.
// Get the first argument of tvm_struct_get, and continue.
......@@ -179,18 +181,33 @@ class MemoryAccessVerifier final : protected StmtExprVisitor {
const ProducerConsumerNode *pc_{nullptr};
bool failure_{false}; ///< If the verification fails (i.e. has illegal access)
//@}
LoweredFunc func_{nullptr}; ///< Function to be verified.
tir::PrimFunc func_{nullptr}; ///< Function to be verified.
int dev_type_{kDLCPU}; ///< Device type
std::unordered_map<const VarNode *, PrimExpr> defs_; ///< Variable definitions
};
} // namespace
/// Interface of VerifyMemory pass
bool VerifyMemory(LoweredFunc func, int device_type) {
MemoryAccessVerifier v(func, device_type);
void VerifyMemory(const IRModule& mod) {
for (auto kv : mod->functions) {
if (auto* n = kv.second.as<PrimFuncNode>()) {
PrimFunc func = GetRef<PrimFunc>(n);
auto target = func->GetAttr<Target>(tvm::attr::kTarget);
CHECK(target.defined())
<< "LowerWarpMemory: Require the target attribute";
MemoryAccessVerifier v(func, target->device_type);
v.Run();
return !v.Failed();
if (v.Failed()) {
LOG(FATAL)
<< "ValueError: Direct host side access to device memory is detected."
<< " Did you forget to bind?\n"
<< func;
}
}
}
}
TVM_REGISTER_GLOBAL("tir.analysis.verify_memory")
.set_body_typed(VerifyMemory);
} // namespace tir
} // namespace tvm
src/tir/ir/buffer.cc
......@@ -48,7 +48,7 @@ Buffer decl_buffer(Array<PrimExpr> shape,
DataType dtype,
std::string name) {
return BufferNode::make(
Var(name, DataType::Handle()),
Var(name, PointerType(PrimType(dtype))),
dtype,
shape,
Array<PrimExpr>(),
......
src/tir/ir/lowered_func.cc deleted 100644 → 0
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/*!
* \file lowered_func.cc
*/
#include <tvm/tir/lowered_func.h>
namespace tvm {
namespace tir {
TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
.set_dispatch<LoweredFuncNode>([](const ObjectRef& node, ReprPrinter* p) {
auto* op = static_cast<const LoweredFuncNode*>(node.get());
p->stream << "LoweredFunc(" << op->name << ", " << op << ")";
});
TVM_REGISTER_NODE_TYPE(LoweredFuncNode);
} // namespace tir
} // namespace tvm
src/tir/pass/ffi_api.cc
......@@ -105,13 +105,6 @@ TVM_REGISTER_GLOBAL("ir_pass.PostOrderVisit")
});
});
TVM_REGISTER_GLOBAL("ir_pass.LowerStorageAccess")
.set_body([](TVMArgs args, TVMRetValue *ret) {
LoweredFunc f = args[0];
auto n = make_object<LoweredFuncNode>(*f.operator->());
n->body = LowerStorageAccessInfo(f->body);
*ret = LoweredFunc(n);
});
// make from two arguments
#define REGISTER_PASS(PassName) \
......@@ -128,7 +121,6 @@ REGISTER_PASS(VectorizeLoop);
REGISTER_PASS(SkipVectorize);
REGISTER_PASS(UnrollLoop);
REGISTER_PASS(InjectCopyIntrin);
REGISTER_PASS(MakeAPI);
REGISTER_PASS(StorageRewrite);
REGISTER_PASS(CoProcSync);
REGISTER_PASS(LowerStorageAccessInfo);
......@@ -138,9 +130,6 @@ REGISTER_PASS(InjectDoubleBuffer);
REGISTER_PASS(LoopPartition);
REGISTER_PASS(RemoveNoOp);
REGISTER_PASS(LiftAttrScope);
REGISTER_PASS(RemapThreadAxis);
REGISTER_PASS(LowerCustomDatatypes);
REGISTER_PASS(VerifyMemory);
REGISTER_PASS(VerifyGPUCode);
REGISTER_PASS(DecorateDeviceScope);
REGISTER_PASS(InstrumentBoundCheckers);
......
src/tir/pass/storage_rewrite.cc
......@@ -994,29 +994,6 @@ class VectorAllocRewriter : public StmtExprMutator {
};
LoweredFunc PointerValueTypeRewrite(LoweredFunc f) {
auto n = make_object<LoweredFuncNode>(*f.operator->());
VectorAllocRewriter rewriter;
n->body = rewriter(n->body);
for (Var arg : f->args) {
if (arg.dtype().is_handle()) {
const auto& tvec = rewriter.acc_map_[arg.get()];
if (tvec.size() == 1) {
PrimExpr dtype = make_const(tvec[0], 0);
n->handle_data_type.Set(arg, dtype);
} else {
// always set data type to be non vectorized so
// load/store can still work via scalarization
if (tvec.size() != 0 && !n->handle_data_type.count(arg)) {
PrimExpr dtype = make_const(tvec[0].with_lanes(1), 0);
n->handle_data_type.Set(arg, dtype);
}
}
}
}
return LoweredFunc(n);
}
PrimFunc PointerValueTypeRewrite(PrimFunc f) {
auto* n = f.CopyOnWrite();
VectorAllocRewriter rewriter;
......
src/tir/pass/lower_custom_datatypes.cc → src/tir/transforms/lower_custom_datatypes.cc
......@@ -22,7 +22,9 @@
*/
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/ir_pass.h>
#include <tvm/tir/transform.h>
#include <tvm/target/target.h>
#include <tvm/runtime/registry.h>
#include "../../target/datatype/registry.h"
namespace tvm {
......@@ -129,11 +131,26 @@ class CustomDatatypesLowerer : public StmtExprMutator {
std::string target_;
};
LoweredFunc LowerCustomDatatypes(LoweredFunc f, const std::string& target) {
auto n = make_object<LoweredFuncNode>(*f.operator->());
n->body = CustomDatatypesLowerer(target)(n->body);
return LoweredFunc(n);
namespace transform {
Pass LowerCustomDatatypes() {
auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
auto* n = f.CopyOnWrite();
auto target = f->GetAttr<Target>(tvm::attr::kTarget);
CHECK(target.defined())
<< "LowerCustomDatatypes: Require the target attribute";
n->body = CustomDatatypesLowerer(target->target_name)(std::move(n->body));
return f;
};
return CreatePrimFuncPass(pass_func, 0, "tir.LowerCustomDatatypes", {});
}
TVM_REGISTER_GLOBAL("tir.transform.LowerCustomDatatypes")
.set_body_typed(LowerCustomDatatypes);
} // namespace transform
} // namespace tir
} // namespace tvm
src/tir/pass/make_api.cc → src/tir/transforms/make_packed_api.cc
......@@ -18,20 +18,24 @@
*/
/*!
* \file make_api.cc Build API function.
* \file make_packed_api.cc Lower PrimFunc to use the packed function API.
*/
#include <tvm/tir/ir_pass.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/analysis.h>
#include <tvm/tir/transform.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/buffer.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/registry.h>
#include <tvm/runtime/container.h>
#include <vector>
#include <utility>
#include <unordered_set>
#include "ir_util.h"
#include "arg_binder.h"
#include "../pass/ir_util.h"
#include "../pass/arg_binder.h"
namespace tvm {
namespace tir {
......@@ -40,14 +44,18 @@ inline Stmt MakeAssertEQ(PrimExpr lhs, PrimExpr rhs, std::string msg) {
return AssertStmtNode::make(lhs == rhs, msg, EvaluateNode::make(0));
}
LoweredFunc MakeAPI(Stmt body,
std::string name,
Array<ObjectRef> api_args,
int num_unpacked_args,
bool is_restricted) {
PrimFunc MakePackedAPI(PrimFunc&& func,
int num_unpacked_args) {
auto global_symbol = func->GetAttr<runtime::String>(tvm::attr::kGlobalSymbol);
CHECK(global_symbol.defined())
<< "MakePackedAPI: Expect PrimFunc to have the global_symbol attribute";
std::string name_hint = global_symbol;
auto* func_ptr = func.CopyOnWrite();
const Stmt nop = EvaluateNode::make(0);
int num_args = static_cast<int>(api_args.size());
int num_args = static_cast<int>(func_ptr->params.size());
CHECK_LE(num_unpacked_args, num_args);
int num_packed_args = num_args - num_unpacked_args;
// Data field definitions
// The packed fields
......@@ -69,7 +77,8 @@ LoweredFunc MakeAPI(Stmt body,
// local function definitions
// load i-th argument as type t
auto f_arg_value = [&](DataType t, int i) {
Array<PrimExpr> call_args{v_packed_args,
Array<PrimExpr> call_args{
v_packed_args,
IntImm(DataType::Int(32), i),
IntImm(DataType::Int(32), intrinsic::kTVMValueContent)};
// load 64 bit version
......@@ -83,13 +92,7 @@ LoweredFunc MakeAPI(Stmt body,
}
return res;
};
// get declaration of argument i
auto f_arg_decl = [&](int i) {
std::ostringstream os;
os << "arg" << i;
const VarNode* v = api_args[i].as<VarNode>();
return Var(os.str(), v ? v->dtype: DataType::Handle());
};
// ---------------------------
// start of logics
// add signiture for packed arguments.
......@@ -99,16 +102,25 @@ LoweredFunc MakeAPI(Stmt body,
args.push_back(v_num_packed_args);
std::ostringstream os;
os << name << ": num_args should be " << num_packed_args;
os << name_hint << ": num_args should be " << num_packed_args;
seq_init.emplace_back(
MakeAssertEQ(v_num_packed_args, num_packed_args, os.str()));
}
// Save the input variables and buffers that will be bound later.
std::vector<std::pair<Var, Var> > var_defs;
std::vector<std::pair<Buffer, Var> > buf_defs;
for (int i = 0; i < static_cast<int>(api_args.size()); ++i) {
Var v_arg = f_arg_decl(i);
// Need to re-declare vars, in case some arguments also appears in the buffer.
std::vector<std::pair<Var, Var> > var_def;
std::vector<std::pair<Var, Buffer> > buffer_def;
for (int i = 0; i < static_cast<int>(func_ptr->params.size()); ++i) {
Var param = func_ptr->params[i];
Var v_arg = Var("arg" + std::to_string(i), param->dtype);
auto it = func_ptr->buffer_map.find(param);
if (it != func_ptr->buffer_map.end()) {
buffer_def.emplace_back(v_arg, (*it).second);
} else {
var_def.emplace_back(v_arg, param);
}
if (i < num_packed_args) {
// Value loads
seq_init.emplace_back(LetStmtNode::make(
......@@ -123,7 +135,7 @@ LoweredFunc MakeAPI(Stmt body,
DataType t = v_arg.dtype();
if (t.is_handle()) {
std::ostringstream msg;
msg << name << ": Expect arg[" << i << "] to be pointer";
msg << name_hint << ": Expect arg[" << i << "] to be pointer";
seq_check.emplace_back(
AssertStmtNode::make(tcode == kTVMOpaqueHandle ||
tcode == kTVMNDArrayHandle ||
......@@ -131,27 +143,18 @@ LoweredFunc MakeAPI(Stmt body,
tcode == kTVMNullptr, msg.str(), nop));
} else if (t.is_int() || t.is_uint()) {
std::ostringstream msg;
msg << name << ": Expect arg[" << i << "] to be int";
msg << name_hint << ": Expect arg[" << i << "] to be int";
seq_check.emplace_back(AssertStmtNode::make(tcode == kDLInt, msg.str(), nop));
} else {
CHECK(t.is_float());
std::ostringstream msg;
msg << name << ": Expect arg[" << i << "] to be float";
msg << name_hint << ": Expect arg[" << i << "] to be float";
seq_check.emplace_back(
AssertStmtNode::make(tcode == kDLFloat, msg.str(), nop));
}
} else {
args.push_back(v_arg);
}
// add checks for functions.
if (api_args[i].as<VarNode>()) {
var_defs.emplace_back(std::make_pair(Downcast<Var>(api_args[i]), v_arg));
} else {
// Buffer checks
CHECK(api_args[i].as<BufferNode>())
<< "api_args can only be Buffer or Var";
buf_defs.emplace_back(std::make_pair(Downcast<Buffer>(api_args[i]), v_arg));
}
}
// allow return value if the function is packed.
......@@ -170,24 +173,22 @@ LoweredFunc MakeAPI(Stmt body,
// either 0 or the original stride will be correctly used. Checks here have
// to use the args that may have no let bining yet. Therefore, hoisting let
// binding for args before buffer declaration is needed.
for (const auto& arg : var_defs) {
binder.Bind(arg.first, arg.second, arg.second->name_hint, true);
for (const auto& kv : var_def) {
binder.Bind(kv.second, kv.first, kv.first->name_hint, true);
}
for (const auto& kv : buffer_def) {
binder.BindDLTensor(kv.second, device_type, device_id,
kv.first, kv.first->name_hint);
}
for (const auto& buf_arg : buf_defs) {
binder.BindDLTensor(buf_arg.first, device_type, device_id,
buf_arg.second, buf_arg.second->name_hint);
if (num_unpacked_args == 0) {
func = WithAttr(std::move(func), tvm::attr::kCallingConv, Integer(CallingConv::kCPackedFunc));
}
ObjectPtr<LoweredFuncNode> n = make_object<LoweredFuncNode>();
n->name = name;
n->args = args;
n->handle_data_type = binder.def_handle_dtype();
n->is_packed_func = num_unpacked_args == 0;
n->is_restricted = is_restricted;
body = AttrStmtNode::make(
auto body = AttrStmtNode::make(
make_zero(DataType::Int(32)), attr::compute_scope,
StringImmNode::make(name + "_compute_"), body);
StringImmNode::make(name_hint + "_compute_"), func_ptr->body);
// Set device context
if (vmap.count(device_id.get())) {
PrimExpr node = StringImmNode::make("default");
......@@ -203,21 +204,59 @@ LoweredFunc MakeAPI(Stmt body,
device_type, device_id}, CallNode::Intrinsic)));
body = SeqStmt({set_device, body});
}
n->body = MergeNest(
func_ptr->body = MergeNest(
{seq_init, binder.init_nest(), seq_check, binder.asserts()}, body);
LoweredFunc f(n);
Array<Var> undefined = UndefinedVars(f->body, f->args);
func_ptr->params = args;
Array<Var> undefined = UndefinedVars(func_ptr->body, func_ptr->params);
if (undefined.size() != 0) {
std::ostringstream os;
for (Var v : undefined) {
os << " \'" << v->name_hint << "\' ";
}
os << " does not appear in api_args";
os << " is not bound to any variables";
LOG(FATAL) << "Not all Vars are passed in api_args: " << os.str();
}
return f;
func_ptr->buffer_map = Map<Var, Buffer>();
func_ptr->checked_type_ = func_ptr->func_type_annotation();
func_ptr->ret_type = PrimType(DataType::Int(32));
// return the function.
return std::move(func);
}
namespace transform {
Pass MakePackedAPI(int num_unpacked_args) {
auto pass_func = [num_unpacked_args](IRModule m, PassContext ctx) {
IRModuleNode* mptr = m.CopyOnWrite();
std::vector<std::pair<GlobalVar, PrimFunc> > updates;
for (const auto& kv : mptr->functions) {
if (auto* n = kv.second.as<PrimFuncNode>()) {
PrimFunc func = GetRef<PrimFunc>(n);
if (func->GetAttr<Integer>(tvm::attr::kCallingConv, 0)->value
== static_cast<int>(CallingConv::kDefault)) {
auto updated_func = MakePackedAPI(std::move(func), num_unpacked_args);
updates.push_back({kv.first, updated_func});
}
}
}
for (const auto& pair : updates) {
mptr->AddUnchecked(pair.first, pair.second);
}
return m;
};
return tvm::transform::CreateModulePass(
pass_func, 0, "tir.MakePackedAPI", {});
}
TVM_REGISTER_GLOBAL("tir.transform.MakePackedAPI")
.set_body_typed(MakePackedAPI);
} // namespace transform
} // namespace tir
} // namespace tvm
src/tir/pass/remap_thread_axis.cc → src/tir/transforms/remap_thread_axis.cc
......@@ -22,7 +22,8 @@
*/
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/ir_pass.h>
#include <tvm/tir/transform.h>
#include <tvm/runtime/registry.h>
#include <unordered_map>
......@@ -74,8 +75,8 @@ class ThreadAxisRewriter : private StmtExprMutator {
std::unordered_map<const VarNode*, Var> vmap_;
};
LoweredFunc
RemapThreadAxis(LoweredFunc f, Map<PrimExpr, IterVar> thread_map) {
PrimFunc RemapThreadAxis(PrimFunc&& f, Map<PrimExpr, IterVar> thread_map) {
std::unordered_map<std::string, IterVar> tmap;
for (const auto& kv : thread_map) {
const StringImmNode* str = kv.first.as<StringImmNode>();
......@@ -83,18 +84,33 @@ RemapThreadAxis(LoweredFunc f, Map<PrimExpr, IterVar> thread_map) {
tmap[str->value] = kv.second;
}
CHECK_EQ(f->func_type, kDeviceFunc);
auto n = make_object<LoweredFuncNode>(*f.operator->());
auto thread_axis = f->GetAttr<Array<IterVar> >(tir::attr::kDeviceThreadAxis);
auto* n = f.CopyOnWrite();
// replace the thread axis
for (size_t i = 0; i < n->thread_axis.size(); ++i) {
auto it = tmap.find(n->thread_axis[i]->thread_tag);
for (size_t i = 0; i < thread_axis.size(); ++i) {
auto it = tmap.find(thread_axis[i]->thread_tag);
if (it != tmap.end()) {
n->thread_axis.Set(i, it->second);
thread_axis.Set(i, it->second);
}
}
n->body = ThreadAxisRewriter(tmap).Rewrite(n->body);
return LoweredFunc(n);
n->body = ThreadAxisRewriter(tmap).Rewrite(std::move(n->body));
return WithAttr(std::move(f), tir::attr::kDeviceThreadAxis, thread_axis);
}
namespace transform {
Pass RemapThreadAxis(Map<PrimExpr, IterVar> thread_map) {
auto pass_func = [thread_map](PrimFunc f, IRModule m, PassContext ctx) {
return RemapThreadAxis(std::move(f), thread_map);
};
return CreatePrimFuncPass(pass_func, 0, "tir.RemapThreadAxis", {});
}
TVM_REGISTER_GLOBAL("tir.transform.RemapThreadAxis")
.set_body_typed(RemapThreadAxis);
} // namespace transform
} // namespace tir
} // namespace tvm
src/tir/transforms/split_host_device.cc
......@@ -264,7 +264,6 @@ class HostDeviceSplitter : public StmtMutator {
std::string name_prefix_;
// Number of device functions.
int device_func_counter_{0};
std::vector<LoweredFunc> device_funcs_;
std::unordered_map<const VarNode*, PrimExpr> handle_data_type_;
};
......
tests/cpp/build_module_test.cc
......@@ -117,7 +117,7 @@ TEST(BuildModule, Heterogeneous) {
std::unordered_map<Tensor, Buffer> binds;
auto lowered_s1 = lower(s1, args1, "elemwise_add", binds, config);
auto lowered_s2 = lower(s2, args2, "elemwise_sub", binds, config);
Map<tvm::Target, Array<LoweredFunc>> inputs = {{target_cuda, lowered_s1},
Map<tvm::Target, IRModule> inputs = {{target_cuda, lowered_s1},
{target_llvm, lowered_s2}};
auto module = build(inputs, Target(), config);
......
tests/python/integration/test_dot.py
......@@ -18,29 +18,6 @@ import tvm
from tvm import te
import numpy as np
def lower(s, args, name="mydot"):
binds = {}
arg_list = []
for x in args:
assert isinstance(x, te.tensor.Tensor)
buf = tvm.tir.decl_buffer(x.shape, dtype=x.dtype, name=x.op.name)
binds[x] = buf
arg_list.append(buf)
s = s.normalize()
bounds = tvm.te.schedule.InferBound(s)
stmt = tvm.te.schedule.ScheduleOps(s, bounds)
stmt = tvm.tir.ir_pass.StorageFlatten(stmt, binds, 16)
stmt = tvm.tir.ir_pass.CanonicalSimplify(stmt)
stmt = tvm.tir.ir_pass.Simplify(stmt)
fapi = tvm.tir.ir_pass.MakeAPI(stmt, name, arg_list, 0, True)
fapi = tvm.tir.ir_pass.LowerTVMBuiltin(fapi)
return fapi
def mybuild(fapi, target="llvm"):
return
def test_dot():
nn = 12
......
tests/python/unittest/test_runtime_extension.py
......@@ -38,8 +38,9 @@ def test_dltensor_compatible():
with ib.for_range(0, n - 1, "i") as i:
A[i + 1] = A[i] + 1
stmt = ib.get()
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "arange", [Ab], 0, True)
mod = tvm.testing.LoweredFuncsToIRModule([fapi])
mod = tvm.testing.MakeAPILegacy(stmt, "arange", [Ab], 0, True)
mod = tvm.tir.transform.LowerTVMBuiltin()(mod)
f = tvm.target.codegen.build_module(mod, "stackvm")
a = tvm.nd.array(np.zeros(10, dtype=dtype))
......
tests/python/unittest/test_runtime_heterogeneous.py
......@@ -156,7 +156,7 @@ def test_simplex_data_transferring():
elemwise_sub],
name="elemwise_sub")
target_flist = {target_device: [lower_add], target_host: [lower_sub]}
target_flist = {target_device: lower_add, target_host: lower_sub}
mhost = tvm.build(target_flist, target_host=target_host)
ctx = [host_ctx, device_ctx]
mod = graph_runtime.create(graph, mhost, ctx)
......@@ -354,8 +354,9 @@ def test_duplex_data_transferring():
elemwise_sub],
name="elemwise_sub")
target_flist = {target_device: [lower_add0, lower_add1], target_host:
[lower_sub]}
lower_add0.update(lower_add1)
target_flist = {target_device: lower_add0, target_host:
lower_sub}
mhost = tvm.build(target_flist, target_host=target_host)
ctx = [host_ctx, device_ctx]
params = {}
......
tests/python/unittest/test_runtime_module_load.py
......@@ -57,8 +57,8 @@ def test_dso_module_load():
tvm.tir.Store(Ab.data,
tvm.tir.Load(dtype, Ab.data, i) + 1,
i + 1))
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
m = tvm.driver.build(fapi, target="llvm")
m = tvm.testing.MakeAPILegacy(stmt, "ramp", [Ab], 0, True)
m = tvm.driver.build(m, target="llvm")
for name in names:
m.save(name)
......
tests/python/unittest/test_target_codegen_llvm.py
......@@ -22,6 +22,7 @@ import numpy as np
import ctypes
import math
def test_llvm_intrin():
ib = tvm.tir.ir_builder.create()
n = tvm.runtime.convert(4)
......@@ -34,7 +35,8 @@ def test_llvm_intrin():
tvm.tir.Call(
"int32", "prefetch", args, tvm.tir.Call.Intrinsic, None, 0)))
body = ib.get()
func = tvm.tir.ir_pass.MakeAPI(body, "prefetch", [A], 0, True)
func = tvm.testing.MakeAPILegacy(body, "prefetch", [A], 0, True)
fcode = tvm.build(func, None, "llvm")
......@@ -85,7 +87,7 @@ def test_llvm_lookup_intrin():
x = tvm.tir.call_llvm_intrin("uint8x8", "llvm.ctpop.i8", tvm.tir.const(1, 'uint32'), A)
ib.emit(x)
body = ib.get()
func = tvm.tir.ir_pass.MakeAPI(body, "ctpop", [A], 1, True)
func = tvm.testing.MakeAPILegacy(body, "ctpop", [A], 1, True)
fcode = tvm.build(func, None, "llvm")
......@@ -307,8 +309,9 @@ def test_multiple_func():
f2 = tvm.lower(s, [A, B, C], name="fadd1")
f1 = tvm.lower(s, [A, B, C], name="fadd2")
m = tvm.build([f1, f2], "llvm")
fadd1 = m['fadd1']
fadd2 = m['fadd2']
fadd1 = m['fadd1']
ctx = tvm.cpu(0)
# launch the kernel.
n = nn
......@@ -665,6 +668,7 @@ def test_llvm_shuffle():
tvm.testing.assert_allclose(c_.asnumpy(), (a_.asnumpy() * 2).astype('int32'))
if __name__ == "__main__":
test_multiple_func()
test_llvm_large_uintimm()
test_llvm_import()
test_alignment()
......@@ -676,7 +680,6 @@ if __name__ == "__main__":
test_llvm_vadd_pipeline()
test_llvm_add_pipeline()
test_llvm_intrin()
test_multiple_func()
test_llvm_flip_pipeline()
test_llvm_madd_pipeline()
test_llvm_temp_space()
......
tests/python/unittest/test_target_codegen_static_init.py
......@@ -19,6 +19,18 @@ from tvm import te
import ctypes
import numpy as np
def MakeAPILegacy(stmt, name, args, num_unpacked_args, noalias):
"""Legacy adapter to create a API"""
f = tvm.tir.PrimFunc(args, stmt).with_attr(
"global_symbol", tvm.runtime.String(name))
f = f.with_attr("tir.is_entry_func", True)
if noalias:
f = f.with_attr("tir.no_alias", True)
mod = tvm.IRModule.from_expr(f)
return tvm.tir.transform.MakePackedAPI()(mod)
def test_static_callback():
dtype = 'int64'
n = te.size_var('n')
......@@ -32,7 +44,7 @@ def test_static_callback():
with ib.for_range(0, n, "i", for_type="parallel") as i:
A[i] = A[i] + 1
stmt = ib.get()
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
fapi = tvm.testing.MakeAPILegacy(stmt, "ramp", [Ab], 0, True)
f = tvm.driver.build(fapi, target="llvm")
a = tvm.nd.array(np.zeros(10, dtype=dtype))
f(a)
......@@ -55,7 +67,7 @@ def test_static_init():
return sh
stmt = ib.get()
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
fapi = tvm.testing.MakeAPILegacy(stmt, "ramp", [Ab], 0, True)
f = tvm.driver.build(fapi, target="llvm")
a = tvm.nd.array(np.zeros(10, dtype=dtype))
f(a)
......
tests/python/unittest/test_target_codegen_vm_basic.py
......@@ -26,6 +26,18 @@ def run_jit(fapi, check):
s = f.get_source()
check(f)
def MakeAPILegacy(stmt, name, args, num_unpacked_args, noalias):
"""Legacy adapter to create a API"""
f = tvm.tir.PrimFunc(args, stmt).with_attr(
"global_symbol", tvm.runtime.String(name))
f = f.with_attr("tir.is_entry_func", True)
if noalias:
f = f.with_attr("tir.no_alias", True)
mod = tvm.IRModule.from_expr(f)
return tvm.tir.transform.MakePackedAPI()(mod)
def test_stack_vm_basic():
a = tvm.nd.array(np.zeros(10, dtype='float32'))
@tvm.register_func
......@@ -36,7 +48,7 @@ def test_stack_vm_basic():
n = te.size_var('n')
Ab = tvm.tir.decl_buffer((n, ), "float32")
stmt = tvm.tir.Evaluate(tvm.tir.call_packed("tvm_call_back_get_shape", Ab.shape[0]))
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "print_shape", [Ab], 0, True)
fapi = tvm.testing.MakeAPILegacy(stmt, "print_shape", [Ab], 0, True)
run_jit(fapi, lambda f: f(a))
......@@ -57,7 +69,7 @@ def test_stack_vm_loop():
ib.emit(tvm.tir.call_packed("tvm_stack_vm_print", i))
stmt = ib.get()
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
fapi = tvm.testing.MakeAPILegacy(stmt, "ramp", [Ab], 0, True)
a = tvm.nd.array(np.zeros(10, dtype=dtype))
def check(f):
f(a)
......@@ -79,7 +91,7 @@ def test_stack_vm_cond():
A[i + 1] = A[i] + 2
stmt = ib.get()
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "test", [Ab], 0, True)
fapi = tvm.testing.MakeAPILegacy(stmt, "test", [Ab], 0, True)
def check(f):
a = tvm.nd.array(np.zeros(10, dtype=dtype))
f(a)
......@@ -98,7 +110,7 @@ def test_vm_parallel():
with ib.for_range(0, n, "i", for_type="parallel") as i:
A[i] = A[i] + 1
stmt = ib.get()
fapi = tvm.tir.ir_pass.MakeAPI(stmt, "ramp", [Ab], 0, True)
fapi = tvm.testing.MakeAPILegacy(stmt, "ramp", [Ab], 0, True)
def check(f):
a = tvm.nd.array(np.zeros(10, dtype=dtype))
f(a)
......
tests/python/unittest/test_target_custom_datatypes.py
......@@ -19,7 +19,6 @@ import tvm
from tvm import te
from ctypes import *
import topi
import tvm.tir.ir_pass as ir_pass
import numpy as np
tgt = "llvm"
......@@ -51,10 +50,12 @@ def lower_datatypes_and_build(schedule, args):
Once datatype lowering is integrated directly into TVM's lower/build
process, we won't need to do this manually.
TODO(gus) integrate datatype lowering into build process; change this test"""
flist = tvm.lower(schedule, args)
flist = [flist]
flist = [ir_pass.LowerCustomDatatypes(func, tgt) for func in flist]
return tvm.build(flist[0], target=tgt)
mod = tvm.lower(schedule, args)
target = tvm.target.create(tgt)
mod = tvm.tir.transform.Apply(lambda f: f.with_attr("target", target))(mod)
mod = tvm.tir.transform.LowerCustomDatatypes()(mod)
return tvm.build(mod, target=tgt)
def test_bfloat_add_and_cast_1():
X = te.placeholder((3, ), name="X")
......
tests/python/unittest/test_tir_pass_verify_memory.py → tests/python/unittest/test_tir_analysis_verify_memory.py
......@@ -15,12 +15,13 @@
# specific language governing permissions and limitations
# under the License.
import tvm
import pytest
from tvm import te
# The following DLDeviceType/TVMDeviceExtType values
# are originally defined in dlpack.h and c_runtime_api.h.
gpu_devices = [2, 4, 7, 8, 10, 11]
other_devices = [1, 3, 9, 12]
gpu_devices = ["cuda", "opencl", "metal", "vulkan"]
other_devices = ["llvm", "ext_dev"]
def lower(sch, args):
......@@ -39,8 +40,11 @@ def lower(sch, args):
stmt = tvm.te.schedule.ScheduleOps(sch, bounds)
stmt = tvm.tir.ir_pass.LoopPartition(stmt, False)
stmt = tvm.tir.ir_pass.StorageFlatten(stmt, binds, 64)
func = tvm.tir.ir_pass.MakeAPI(stmt, "myadd", arg_list, 0, True)
return func
f = tvm.tir.PrimFunc(arg_list, stmt).with_attr(
"global_symbol", tvm.runtime.String("test"))
mod = tvm.IRModule({"test": f})
return tvm.tir.transform.MakePackedAPI()(mod)
# All computations are bound.
......@@ -57,10 +61,13 @@ def test_verify_memory_all_bind():
s[B].bind(bx, te.thread_axis("blockIdx.x"))
s[B].bind(tx, te.thread_axis("threadIdx.x"))
func = lower(s, [A, B])
mod = lower(s, [A, B])
for dev_type in gpu_devices + other_devices:
assert tvm.tir.ir_pass.VerifyMemory(func, dev_type)
binded_mod = tvm.tir.transform.Apply(
lambda f: f.with_attr("target", tvm.target.create(dev_type)))(mod)
tvm.tir.analysis.verify_memory(binded_mod)
# Computations are not bound.
......@@ -74,12 +81,18 @@ def test_verify_memory_not_bind():
# B is not bound to threads.
s = te.create_schedule(B.op)
func = lower(s, [A, B])
mod = lower(s, [A, B])
for dev_type in gpu_devices:
assert not tvm.tir.ir_pass.VerifyMemory(func, dev_type)
binded_mod = tvm.tir.transform.Apply(
lambda f: f.with_attr("target", tvm.target.create(dev_type)))(mod)
with pytest.raises(ValueError):
tvm.tir.analysis.verify_memory(binded_mod)
for dev_type in other_devices:
assert tvm.tir.ir_pass.VerifyMemory(func, dev_type)
binded_mod = tvm.tir.transform.Apply(
lambda f: f.with_attr("target", tvm.target.create(dev_type)))(mod)
tvm.tir.analysis.verify_memory(binded_mod)
# Computations are partially bound.
......@@ -98,16 +111,22 @@ def test_verify_memory_partially_bind():
s[C].bind(bx, te.thread_axis("blockIdx.x"))
s[C].bind(tx, te.thread_axis("threadIdx.x"))
func = lower(s, [A, B, C, D])
mod = lower(s, [A, B, C, D])
for dev_type in gpu_devices:
assert not tvm.tir.ir_pass.VerifyMemory(func, dev_type)
binded_mod = tvm.tir.transform.Apply(
lambda f: f.with_attr("target", tvm.target.create(dev_type)))(mod)
with pytest.raises(ValueError):
tvm.tir.analysis.verify_memory(binded_mod)
for dev_type in other_devices:
assert tvm.tir.ir_pass.VerifyMemory(func, dev_type)
binded_mod = tvm.tir.transform.Apply(
lambda f: f.with_attr("target", tvm.target.create(dev_type)))(mod)
tvm.tir.analysis.verify_memory(binded_mod)
if __name__ == "__main__":
test_verify_memory_all_bind()
test_verify_memory_not_bind()
test_verify_memory_partially_bind()
tests/python/unittest/test_tir_pass_bound_checkers.py
......@@ -118,7 +118,6 @@ def test_in_bounds_vectorize_llvm():
s[B].vectorize(xi)
# build and invoke the kernel.
lowered_func = tvm.lower (s, [A, C], "llvm", simple_mode=False)
print (lowered_func.body)
f = tvm.build(s, [A, C], "llvm")
ctx = tvm.cpu(0)
# launch the kernel.
......@@ -137,7 +136,6 @@ def test_in_bounds_loop_partition_basic_llvm():
s = te.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=4)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -156,7 +154,6 @@ def test_out_of_bounds_loop_partition_basic_llvm(index_a, index_b):
s = te.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=4)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -205,12 +202,11 @@ def test_in_bounds_const_loop_partition_ir():
# after instrumentation
assert_bound_instrumentation(stmt, check_attr_stmt, 2 * 3)
assert_bound_instrumentation(stmt, check_branch_stmt, 2)
print (stmt)
branch_collector = list()
collect_visit(stmt, collect_branch_stmt)
assert(len(branch_collector) == 2)
print (branch_collector[0].condition)
print (branch_collector[1].condition)
def test_in_bounds_const_loop_partition_llvm():
with tvm.target.build_config(instrument_bound_checkers=True, partition_const_loop=True):
......@@ -222,7 +218,6 @@ def test_in_bounds_const_loop_partition_llvm():
s = te.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=4)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -242,7 +237,6 @@ def test_out_of_bounds_const_loop_partition_llvm(index_a, index_b):
s = te.create_schedule(T.op)
xo, xi = s[T].split(T.op.axis[0], factor=4)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -276,7 +270,6 @@ def test_in_bounds_conv_llvm(loop_tiling=False):
if loop_tiling:
oho, owo, ohi, owi = s[conv].tile(oh, ow, 16, 16)
lowered_func = tvm.lower(s, [data, kernel, conv], simple_mode=True)
print (lowered_func.body)
ctx = tvm.cpu (0)
f = tvm.build(s, [data, kernel, conv], "llvm")
......@@ -320,7 +313,6 @@ def test_out_of_bounds_conv_llvm(data_offsets, kernel_offsets, loop_tiling=False
if loop_tiling:
oho, owo, ohi, owi = s[conv].tile(oh, ow, 16, 16)
lowered_func = tvm.lower(s, [data, kernel, conv], simple_mode=True)
print (lowered_func.body)
ctx = tvm.cpu (0)
f = tvm.build(s, [data, kernel, conv], "llvm")
......@@ -341,7 +333,6 @@ def test_in_bounds_tensors_with_same_shapes1D_llvm():
T = te.compute((m, ), lambda i: A[i]*B[i])
s = te.create_schedule(T.op)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -361,7 +352,6 @@ def test_out_of_bounds_tensors_with_diff_shapes1D_llvm(a_shape, b_shape, c_shape
T = te.compute((m, ), lambda i: A[i]*B[i])
s = te.create_schedule(T.op)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -380,7 +370,6 @@ def test_in_bounds_tensors_with_same_shapes2D_llvm():
T = te.compute((m, m), lambda i, j: A[i][j]*B[i][j])
s = te.create_schedule(T.op)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -400,7 +389,6 @@ def test_out_of_bounds_tensors_with_diff_shapes2D_llvm(a_shape, b_shape, c_shape
T = te.compute((m, m), lambda i, j: A[i][j]*B[i][j])
s = te.create_schedule(T.op)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -419,7 +407,7 @@ def test_in_bounds_tensors_with_same_shapes3D_llvm():
T = te.compute((m, m, m), lambda i, j, p: A[i][j][p]*B[i][j][p])
s = te.create_schedule(T.op)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -439,7 +427,7 @@ def test_out_of_bounds_tensors_with_diff_shapes3D_llvm(a_shape, b_shape, c_shape
T = te.compute((m, m, m), lambda i, j, p: A[i][j][p]*B[i][j][p])
s = te.create_schedule(T.op)
lowered_func = tvm.lower (s, [A, B, T], "llvm", simple_mode=False)
print (lowered_func.body)
ctx = tvm.cpu(0)
f = tvm.build(s, [A, B, T], "llvm")
......@@ -460,7 +448,7 @@ def test_out_of_bounds_tensors_with_zero_shape_op_with_not_zero_shape_llvm():
D = te.compute((), lambda : C + 1)
s = te.create_schedule(D.op)
stmt = tvm.lower (s, [A, scale, D], simple_mode=True)
print (stmt)
# build and invoke the kernel.
f = tvm.build(s, [A, scale, D], "llvm")
ctx = tvm.cpu(0)
......
tests/python/unittest/test_tir_pass_inject_double_buffer.py
......@@ -40,8 +40,7 @@ def test_double_buffer():
stmt = tvm.tir.ir_pass.Simplify(stmt)
assert isinstance(stmt.body.body, tvm.tir.Allocate)
assert stmt.body.body.extents[0].value == 2
f = tvm.tir.ir_pass.MakeAPI(stmt, "db", [A.asobject(), C.asobject()], 2, True)
mod = tvm.testing.LoweredFuncsToIRModule([f])
mod = tvm.testing.MakeAPILegacy(stmt, "db", [A.asobject(), C.asobject()], 2, True)
f = tvm.tir.transform.ThreadSync("shared")(mod)["db"]
count = [0]
......
tests/python/unittest/test_tir_pass_loop_partition.py
......@@ -381,7 +381,7 @@ def test_multilevel_splitting_with_indivisble_factors():
## But this does the right thing.
with tvm.target.build_config(partition_const_loop=True):
lowered_body = tvm.lower(s, [A, B]).body
lowered_body = tvm.lower(s, [A, B], name="x")["x"].body
def visit_stmt(op):
return(isinstance(op, tvm.tir.Max))
num_max = collect_visit(lowered_body, visit_stmt)
......@@ -407,7 +407,7 @@ def test_double_splitting_with_indivisible_factors():
# Find the beginning of the Halide IR corresponding to kernel code
# and make sure it doesn't have an if statements left
top_produce = find_top_produce(f.body)
top_produce = find_top_produce(f["fadd1"].body)
assert(not any(collect_visit(top_produce, lambda x: isinstance(x, tvm.tir.IfThenElse))))
# check functional correctness of generated code
......
tests/python/unittest/test_tir_pass_storage_flatten.py
......@@ -92,9 +92,7 @@ def test_flatten_double_buffer():
stmt = tvm.tir.ir_pass.Simplify(stmt)
assert isinstance(stmt.body.body, tvm.tir.Allocate)
assert stmt.body.body.extents[0].value == 2
f = tvm.tir.ir_pass.MakeAPI(stmt, "db", [A.asobject(), C.asobject()], 2, True)
f = tvm.tir.ir_pass.MakeAPI(stmt, "db", [A.asobject(), C.asobject()], 2, True)
mod = tvm.testing.LoweredFuncsToIRModule([f])
mod = tvm.testing.MakeAPILegacy(stmt, "db", [A.asobject(), C.asobject()], 2, True)
f = tvm.tir.transform.ThreadSync("shared")(mod)["db"]
count = [0]
......
tests/python/unittest/test_tir_transform_combine_context_call.py
......@@ -36,12 +36,7 @@ def test_for():
ib.emit(tvm.tir.call_extern
("int32", "fadd", device_context(0), A))
body = ib.get()
f = tvm.tir.ir_pass.MakeAPI(body, "func", [dev_type, n], 2, True)
# temp adapter to convert loweredFunc to IRModule
# to test passes in the new style.x
mod = tvm.testing.LoweredFuncsToIRModule([f])
mod = tvm.testing.MakeAPILegacy(body, "func", [dev_type, n], 2, True)
mod = tvm.tir.transform.CombineContextCall()(mod)
assert mod["func"].body.value.dtype == "handle"
......
tests/python/unittest/test_tir_transform_lower_warp_memory.py
......@@ -35,10 +35,8 @@ def test_lower_warp_mem():
cuda_target = tvm.target.create("cuda")
assert cuda_target.thread_warp_size == 32
f = tvm.lower(s, [A, B], name="f")
mod = tvm.lower(s, [A, B], name="f")
mod = tvm.testing.LoweredFuncsToIRModule([f])
mod = tvm.tir.transform.Apply(lambda f: f.with_attr("target", cuda_target))(mod)
fdevice = tvm.tir.transform.SplitHostDevice()(mod)["f_kernel0"]
mod = tvm.IRModule.from_expr(fdevice)
......
tests/python/unittest/test_tir_pass_makeapi.py → tests/python/unittest/test_tir_transform_make_packed_api.py
......@@ -35,11 +35,11 @@ def test_makeapi():
stmt = tvm.tir.ir_pass.StorageFlatten(stmt, {A: Ab, B:Bb, C:Cb}, 64)
num_unpacked_args = 2
f = tvm.tir.ir_pass.MakeAPI(
stmt, "myadd", [n, Ab, Bb, Cb], num_unpacked_args, True)
assert(f.handle_data_type[Ab.data].dtype == Ab.dtype)
assert(len(f.args) == 7)
output_ssa = False
f = tvm.tir.PrimFunc([n, Ab, Bb, Cb], stmt).with_attr(
"tir.no_alias", True).with_attr("global_symbol", tvm.runtime.String("myadd"))
mod = tvm.IRModule.from_expr(f)
f = tvm.tir.transform.MakePackedAPI(num_unpacked_args)(mod)["main"]
assert(len(f.params) == 7)
if __name__ == "__main__":
......
tests/python/unittest/test_tir_transform_thread_sync.py
......@@ -37,10 +37,9 @@ def test_thread_storage_sync():
Ab = tvm.tir.decl_buffer(A.shape, A.dtype, name='A')
A2b = tvm.tir.decl_buffer(A2.shape, A2.dtype, name='A2')
stmt = tvm.tir.ir_pass.StorageFlatten(stmt, {A: Ab, A2: A2b}, 64)
f = tvm.tir.ir_pass.MakeAPI(stmt, "test", [Ab, A2b], 0, True)
cuda_target = tvm.target.create("cuda")
mod = tvm.testing.LoweredFuncsToIRModule([f])
cuda_target = tvm.target.create("cuda")
mod = tvm.testing.MakeAPILegacy(stmt, "test", [Ab, A2b], 0, True)
mod = tvm.tir.transform.Apply(lambda f: f.with_attr("target", cuda_target))(mod)
fdevice = tvm.tir.transform.SplitHostDevice()(mod)["test_kernel0"]
mod = tvm.IRModule.from_expr(fdevice)
......
tutorials/dev/low_level_custom_pass.py
......@@ -36,7 +36,7 @@ Before reading this tutorial, we assume readers have already known these topics
- Visitor design pattern. Otherwise, check the
`Python AST module <https://docs.python.org/3/library/ast.html>`_ to see how an AST
visitor is implemented.
- How a HalideIR/Schedule is lowered to either a LoweredFunc class or a LLVM module. Otherwise,
- How a Schedule is lowered to either an IRModule class or a LLVM module. Otherwise,
take a look at ``python/tvm/build_module.py`` to get some basics.
"""
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment