TVM自定义修改代码示例
TVM自定义修改代码示例
一．TVM设备添加代码
因为要添加的设备是一种类似于GPU的加速卡，TVM中提供了对GPU编译器的各种支持，有openCl，OpenGL和CUDA等，这里选取比较熟悉的CUDA进行模仿生成。总体上看，TVM是一个多层的结构。
TVM在python这一层提供了相关的设备接口，然后使用tvm.build真正的编译，然后调用get_source函数来获得想要的源码（或者IR，如llvm选项提供的是LLVM的IR，或者PTX选项提供的就是NVPTX类型的IR）。
因此，添加新设备(device)推测的步骤就是：
- 补全相应的python接口
- 找到python和C交互的接口
- 正确维护中间代码的IR pass变换中新设备引入的特性
- 代码生成对新设备和新特性的支持
- 添加编译选项支持(非必须)
以下就分别就这4个步骤进行介绍。
1. 补全相应的python接口
我之前给的那个测试代码中使用的是字符串解析的方式，但是从其它tutorial中发现，还存在一种tvm.target.cuda()的设备建立方式，这个很明显比字符串解析，相对找起来容易（字符串最终对应的也是这种方式）。按照这种方式找到了tvm/python/tvm/target.py文件中，这个类中定义了现在能支持的target。添加新的target叫做dpu。
def dpu(model='unknown', options=None):
"""Returns a dpu target.
Parameters
----------
model: str
The model of dpu device
options : str or list of str
Additional options
"""
opts = _merge_opts(['-model=%s' % model], options)
return _api_internal._TargetCreate("dpu", *opts)
每个设备都包括硬件自身的上下文信息和硬件上运行软件运行时，就是runtime，在TVM中相关的软件运行时信息在tvm/python/tvm/_ffi/runtime_ctypes.py文件中，添加对dpu的支持
在class TVMContext的两个掩码MASK2STR和STR2MASK中分别添加：
13: 'dpu',
和
'dpu':13,
2. 找到python和C交互的接口
回到刚才的target.py文件中，核心的代码只有两句
opts = _merge_opts(['-model=%s' % model], options)
return _api_internal._TargetCreate("dpu", *opts)
第一句是将model和相关的options组合在一起，就是个字符串相关的拼接，没有特别多需要关注的内容，后边有一个_api_internel._TargetCreate的函数调用，从名字上看起来非常的重要，是创建真正的Target的，但是，在tvm/python文件中，无论如何都找不到该函数的实现。
前边已经提到过TVM中使用的是python提供接口，真正的实现都是在C++中，因此，这里猜测是调用了C语言的实现。下面列一下TVM相关的文件夹
3rdparty是很多第三方库的实现
build 目录是建立的编译后的.so文件所在的位置
docs 是相关的文档
include C++代码的include的主目录
jvm 是java相关的文件夹
nnvm 是中间的nnvm算子所在的目录
python 是python文件所在的目录，所有与python相关的都在该目录中
rust apps conda docker golang web verilog都是特有领域中的内容，对一般项目没有影响
tests 是测试文件，中间包含了作者写的很多测试，是学习TVM的另一个手段
Tutorial是官网上相关的历程
vta 是TVM的软件栈
cmake包含了所有的编译配置文件，和CmakeLists.txt共同工作
src 是全部的C++代码
topi 是Tensor Operator Index Library，后续进行详细介绍
在src目录下搜索_TargetCreate，得到src/codegen/build_module.cc:116中有相关的内容
TVM_REGISTER_API("_TargetCreate")
.set_body([](TVMArgs args, TVMRetValue* ret) {
std::string target_name = args[0];
std::vector<std::string> options;
for (int i = 1; i < args.num_args; ++i) {
std::string arg = args[i];
options.push_back(arg);
}
*ret = CreateTarget(target_name, options);
});
这段代码就是通过一种TVM_REGISTER_API的注册机制，注册_TargetCreate函数，真正的函数体是.set_body内执行的，实际上C++中tvm::CreateTarget函数。TVM_REGISTER_API的注册机制在TVM项目中非常普遍，其实现在项目中也有，不是主要的研究内容，不需要改，所以不另行赘述。
3. 正确维护中间代码的IR pass变换中新设备引入的特性
在src/codegen/build_module.cc文件中的tvm::CreateTarget函数中添加对dpu的支持
else if (target_name == "dpu") {
t->device_type = kDLDPU;
}
这里边的kDLDPU是一个DLDeviceType类型值，实现是在3rdparty/dlpack/include/dlpack/dlpack.h中添加的
kDLDPU =13,
在include/tvm/runtime/device_api.h：200补充对kDLDPU的支持
case kDLDPU: return "dpu";
Target部分添加完了，还需要补充运行时的内容。
运行时的内容在src/runtime/目录下，需要在module.cc中添加对dpu 的支持。
在RuntimeEnabled函数中，添加
else if (target == "dpu") {
f_name = "device_api.dpu";
}
这只是添加了一个名字的支持，需要新建一个dpu目录，里边存放DPUModuleNode、DPUWorkspace等支持，测试代码的getSource函数的真正实现，存放在这里边，主要模仿CUDA和openCl的实现进行。目前存放有dpu_common.h、dpu_device_api.cc、dpu_module.cc、dpu_module.h四个文件，大概1K行代码，实现逻辑不是很复杂。
4. 代码生成对新设备和新特性的支持
上边准备好了module部分，也就是运行时，但是这里第一步想要实现的是一个能在dpu编译器上运行的C代码。需要在codegen部分添加对dpu这个设备的支持。
codegen是在tvm.build（Python）中形成的，在其对应的C++实现上是codegen/build_module.cc文件，之前添加了名字的支持，现在还需要添加这个真正的Target调用点
Target DPU(const std::vector<std::string>& options ) {
return CreateTarget("dpu", options);
}
最主要的codegen对DPU的支持是新建CodeGenDPU类，这个类的实现在该目录的codegen_dpu.h和codegen_dpu.cc文件内。其它的函数可以不实现，有两个函数必须实现
runtime::Module BuildDPU(Array<LoweredFunc> funcs) {
using tvm::runtime::Registry;
bool output_ssa = false;
CodeGenDPU cg;
cg.Init(output_ssa);
for (LoweredFunc f : funcs) {
cg.AddFunction(f);
}
std::string code = cg.Finish();
if (const auto* f = Registry::Get("tvm_callback_dpu_postproc")) {
code = (*f)(code).operator std::string();
}
return DPUModuleCreate(code, "dpu", ExtractFuncInfo(funcs), code);
}
TVM_REGISTER_API("codegen.build_dpu")
.set_body([](TVMArgs args, TVMRetValue* rv) {
*rv = BuildDPU(args[0]);
});
5. 添加编译选项支持
上边可以说是完成了从设备添加到代码生成的部分，但是如果只有上边，新添加的设备一直无法运行。但如果仅是对一个设备进行修改的话，这部分没有必要。后来排查发现是部分代码未编译进去导致的。所以开始修改cmake配置。
在上一个TVM调试文档中提到，编译需要打开LLVM和CUDA选项，这里新添加了dpu的设备，需要增加一个新的编译选项，在cmake/config.cmake中添加
|
1 2 |
|
cmake目录下存在着modules和util目录，modules是指定了相关设备的目录等配置，util文件夹下的内容，用来寻找如CUDA等的配置。暂时只需要modules下添加DPU.cmake，这部分的配置代码相对比较简单，就是指定runtime对应的目录。
|
1 2 3 4 5 6 7 8 9 10 |
|
这里修改完config.cmake，需要重新拷贝到build目录下，以使下次配置生效。编译tvm时是cmake目录下的config.cmake和CMakeLists.txt共同工作生效。在CMakeLists.txt中添加
|
1 2 |
|
然后在build目录下，运行cmake命令，重新编译生效。
|
1 2 |
|
这里不加-DCMAKE_BUILD_TYPE=Debug，C++代码无法进行调试。
二．TVM代码生成
本节主要介绍TVM的代码生成流程，即调用relay.build或tvm.build后发生了什么，将深入到TVM的源代码进行剖析。（这里采用的依然是TVM v0.6）
首先区分两个build的区别：tvm.build主要针对单一算子（参照Tensor Expression一文），relay.build是针对整个模型进行编译（参照GCN优化一文），Relay最后也会调用到tvm::build做代码生成。
relay.build
通常的模型编译，由以下两条语句完成。
# Build with Relay
withrelay.build_config(opt_level=0):
graph,lib,params=relay.build(func,target,params=params)
跟踪细节
如何进行代码跟踪，一方面可以直接通过VS Code在函数上方Alt+单击跳转，另一方面如果想有更直观的印象，可以利用pycallgraph进行可视化（需先用pip安装），代码如下，用GCN的代码编译模块。
frompycallgraphimportPyCallGraph
frompycallgraph.outputimportGraphvizOutput
frompycallgraphimportConfig
graphviz=GraphvizOutput()
graphviz.output_file='relay_callgraph.png'
config=Config(max_depth=5)
with PyCallGraph(output=graphviz,config=config):
# Build with Relay
with relay.build_config(opt_level=0):
graph, lib, params=relay.build(func,target,params=params)
- 各函数之间的调用关系，如
tvm.relay.build_module.build->tvm.relay.build_module.BuildModule.build - FFI的打包调用关系，C++和Python在哪些函数上实现互调。
- 深色标注的结点（执行时间长）实际上也是核心的执行步骤，即关键路径
- 结点的调用次数，如
tvm.build_module.lower调用了14次，对应的正是14个Relay算子，可见Relay IR计算图可视化。
对relay.build进行跟踪，跳转进来是python/tvm/relay/build_module.py（这里是因为在relay/__init__.py中，将build函数直接import到relay的命名空间，跳过了build_module这一层），其中的build函数是build_module内的全局函数(helper)。
def build(mod, target=None, target_host=None, params=None):
# do somthing
ifisinstance(autotvm.DispatchContext.current,autotvm.FallbackContext):
tophub_context=autotvm.tophub.context(list(target.values()))
else:
tophub_context=autotvm.util.EmptyContext()
withtophub_context:
bld_mod=BuildModule()
graph_json,mod,params=bld_mod.build(func,target,target_host,params)
returngraph_json,mod,params
先是寻找AutoTVM是否有预先tune好的参数记录，然后构造tophub_context，在内部构建了BuildModule后，跳转到BuildModule.build，然后返回BuildModule.__init__中的内容。
class BuildModule(object):
"""Build a Relay function to run on TVM graph runtime. This class is used
to expose the `RelayBuildModule` APIs implemented in C++.
"""
def __init__(self):
self.mod=_build_module._BuildModule()
self._get_graph_json=self.mod["get_graph_json"]
self._get_module=self.mod["get_module"]
self._build=self.mod["build"]
self._optimize=self.mod["optimize"]
self._set_params_func=self.mod["set_params"]
self._get_params_func=self.mod["get_params"]
def build(self, func, target=None, target_host=None, params=None):
target=_update_target(target)
# Setup the params.
ifparams:
self._set_params(params)
# Build the function
self._build(func,target,target_host)
# Get artifacts
graph_json=self.get_json()
mod=self.get_module()
params=self.get_params()
returngraph_json,mod,params
_build_module._BuildModule()通过FFI，在python/tvm/relay/_build_module.py中，与C++函数建立联系（tvm._ffi._cytpes.function.Function.__call__）。
fromtvm._ffi.functionimport_init_api
_init_api("relay.build_module", __name__)
对应的C++函数在src/relay/backend/build_module.cc
runtime::Module RelayBuildCreate() {
autoexec=make_object<RelayBuildModule>();
returnruntime::Module(exec);
}
TVM_REGISTER_GLOBAL("relay.build_module._BuildModule")
.set_body([](TVMArgsargs,TVMRetValue*rv){
*rv=RelayBuildCreate();
});
就是注册了一个RelayBuildModule供调用，由于主要用的是build函数，因此到RelayBuildModule中找对应的函数。这里TVM用PackedFunc做了一层封装。
PackedFunc GetFunction(const std::string& name,
const ObjectPtr<Object>& sptr_to_self) final {
// ...
if(name=="build"){
returnPackedFunc([sptr_to_self,this](TVMArgsargs,TVMRetValue*rv){
CHECK_EQ(args.num_args,3);
this->Build(args[0],args[1],args[2]);
});
// ...
}
就是调用的是this->Build，跳转过去会指向BuildRelay。
void BuildRelay(Function func,
const std::unordered_map<std::string, tvm::runtime::NDArray>& params) {
// Optimize input Relay Function and returns Relay Module
relay::Modulerelay_module=Optimize(func,targets_,params);
// Get the updated function.
func=relay_module->Lookup("main");
// Generate code for the updated function.
graph_codegen_=std::unique_ptr<GraphCodegen>(newGraphCodegen());
graph_codegen_->Init(nullptr,targets_);
graph_codegen_->Codegen(func);
ret_.graph_json=graph_codegen_->GetJSON();
ret_.params=graph_codegen_->GetParams();
autolowered_funcs=graph_codegen_->GetLoweredFunc();
if(lowered_funcs.size()==0){
LOG(WARNING)<<"no lowered funcs exist in the compiled module";
}else{
ret_.mod=tvm::build(
lowered_funcs,
target_host_,
BuildConfig::Current());
}
}
经过多番跳转，终于到达build的核心模块，再来看TVM逐步做的工作。
- 优化
- 计算图生成
- 后端代码生成
优化
先是优化Optimize，可以看到这里的优化主要是设备无关的优化，graph-level针对tensor运算的优化。（这里的优化pass都已经在C++中实现，先前版本的NNVM似乎还是在Python中调用）
relay::ModuleOptimize(
Function func,
const TargetsMap& targets,
const std::unordered_map<std::string, runtime::NDArray>& params) {
// BindParamsByName(func, params)
// Perform Module->Module optimizations.
relay::Modulerelay_module=relay::ModuleNode::FromExpr(func);
Array<Pass>pass_seqs;
// Run all dialect legalization passes.
// ...
pass_seqs.push_back(transform::SimplifyInference());
//
// ...fskip
//
pass_seqs.push_back(transform::EliminateCommonSubexpr(fskip));
pass_seqs.push_back(transform::CombineParallelConv2D(3));
pass_seqs.push_back(transform::CombineParallelDense(3));
pass_seqs.push_back(transform::FoldConstant());
pass_seqs.push_back(transform::FoldScaleAxis());
pass_seqs.push_back(transform::CanonicalizeCast());
pass_seqs.push_back(transform::CanonicalizeOps());
// ...AlterOpLayout
pass_seqs.push_back(transform::FoldConstant());
// Create a sequential pass and perform optimizations.
transform::Passseq=transform::Sequential(pass_seqs);
// ... judge & do
relay_module=seq(relay_module);
// Handle heterogeneous compilation.
transform::PassContextpass_ctx=PassContext::Current();
if(targets_.size()>1){
relay_module=
RunDeviceAnnotationPass(relay_module,pass_ctx->fallback_device);
}
// Fuse the operations if it is needed.
relay_module=transform::FuseOps()(relay_module);
relay_module=transform::InferType()(relay_module);
CHECK(relay_module.defined());
returnrelay_module;
}
计算图生成
对应GraphCodegen类，同样调用src/relay/backend/build_module.cc中的relay.build_module._GraphRuntimeCodegen（一样是FFI），然后跳转至src/relay/backend/graph_runtime_codegen.cc，其中已经用TVM_REGISTER_GLOBAL注册了对应函数，即用GraphRuntimeCodegenModule生成对应Object。
因此实际graph_codegen_->Codegen的函数是一个PackedFunc，定义在GraphRuntimeCodegen.Codegen，将relay::Function func进行遍历，然后生成计算图。
后端代码生成
Relay得到lower后的函数，最后一步则是交给tvm::build做代码生成，跳转到src/codegen/build_module.cc中的build函数（注意这里重载了几个版本），然后跳转到核心build，注意这里的build函数支持异构编译，只要再inputs划分好不同硬件设施即可。
// Build for heterogeneous execution.
runtime::Module build(const Map<Target, Array<LoweredFunc>>& inputs,
const Target& target_host,
const BuildConfig& config) {
Array<LoweredFunc>fhost_all;
std::vector<runtime::Module>device_modules;
Targettarget_host_val=target_host;
if(!target_host.defined()){
for(constauto&it:inputs){
if(it.first->device_type==kDLCPU){
target_host_val=it.first;
break;
}
}
}
if(!target_host_val.defined()){
target_host_val=DefaultTargetHost(target_host_val);
}
for(constauto&it:inputs){
autohost_dev_funcs=
split_dev_host_funcs(it.second,it.first,target_host_val,config);
auto&fhost=host_dev_funcs[0];
auto&fdevice=host_dev_funcs[1];
// Get the module for a certain target.
runtime::Modulemdev=DeviceBuild(fdevice,it.first);
for(constauto&it:fhost){
fhost_all.push_back(it);
}
device_modules.push_back(mdev);
}
runtime::Modulemhost=codegen::Build(fhost_all,target_host_val->str());
// Import all modules
for(constauto&it:device_modules){
if(it.operator->()){
mhost.Import(it);
}
}
returnmhost;
}
最核心是mhost = codegen::Build，最后跳转过去就开始调用代码生成模块了（src/codegen/codegen.cc）。
runtime::ModuleBuild(constArray<LoweredFunc>&funcs,
conststd::string&target){
// do something
std::stringbuild_f_name="codegen.build_"+mode;
// the build function.
constPackedFunc*bf=runtime::Registry::Get(build_f_name);
runtime::Modulem=transformed_funcs.empty()?
(*bf)(funcs,target):
(*bf)(transformed_funcs,target);
returnm;
}
以生成LLVM IR为例，codegen.build_llvm会在src/codegen/llvm/llvm_module.cc注册，然后调用同个文件中的LLVMModuleNode->Init。这时会跳转到src/codegen/llvm/codegen_llvm.cc中的CodeGenLLVM类进行代码生成。
tvm.build
用tvm.build对算子进行编译，按照以下方式进行调用，例子来自Tensor Expression。
s=tvm.create_schedule(C.op)
tgt="llvm"# "cuda"
fadd=tvm.build(s,[A,B,C],target=tgt,name="myadd")
调用tvm.build后首先跳转到python/tvm/build_module.py，其中的build函数主要做两个步骤：
- lower高层次代码
- 后端代码生成
代码变换
lower高层次代码对应
flist=lower(inputs,args,name=name,binds=binds)
lower函数同样在python/tvm/build_module.py中，类似于relay.build中的Optimize，但这里执行的是operator-level的优化，主要针对循环变换。
deflower(sch,
args,
name="default_function",
binds=None,
simple_mode=False):
# initialization
# Phase 0
ifisinstance(sch, schedule.Schedule):
stmt=form_body(sch)
forf in lower_phase0:
stmt=f(stmt)
compact=ir_pass.VerifyCompactBuffer(stmt)
binds,arg_list = get_binds(args, compact, binds)
# Phase 1
stmt=ir_pass.RewriteForTensorCore(stmt, sch, binds)
stmt=ir_pass.StorageFlatten(stmt, binds, 64, cfg.instrument_bound_checkers)
stmt=ir_pass.CanonicalSimplify(stmt)
forf in lower_phase1:
stmt=f(stmt)
# Phase 2
ifnot simple_mode:
stmt=ir_pass.LoopPartition(stmt, cfg.partition_const_loop)
ifcfg.disable_vectorize:
stmt=ir_pass.SkipVectorize(stmt)
else:
stmt=ir_pass.VectorizeLoop(stmt)
stmt=ir_pass.InjectVirtualThread(stmt)
stmt=ir_pass.InjectDoubleBuffer(stmt, cfg.double_buffer_split_loop)
stmt=ir_pass.StorageRewrite(stmt)
stmt=ir_pass.UnrollLoop(
stmt,
cfg.auto_unroll_max_step,
cfg.auto_unroll_max_depth,
cfg.auto_unroll_max_extent,
cfg.unroll_explicit)
forf in lower_phase2:
stmt=f(stmt)
# Phase 3
stmt=ir_pass.Simplify(stmt)
stmt=ir_pass.RemoveNoOp(stmt)
ifnot cfg.disable_select_rewriting:
stmt=ir_pass.RewriteUnsafeSelect(stmt)
forf in lower_phase3:
stmt=f(stmt)
# Instrument BoundCheckers
ifcfg.instrument_bound_checkers:
stmt=ir_pass.InstrumentBoundCheckers(stmt)
ifsimple_mode:
returnstmt
returnir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)
优化Pass的主体实施都在src/api/api_pass.cc中，以tvm.ir_pass进行注册（由于C++函数中已经在tvm的命名空间里，搜索时直接搜ir_pass，出来对应的API）。
代码生成
lower完之后就进入到后端代码生成，对应build函数中的
mhost=codegen.build_module(fhost_all,str(target_host))
同样，跳转至tvm/codegen.py，初始化tvm.codegen的API codegen._Build，调用FFI，跳转至src/api/api_codegen.cc，最后跳转至src/codegen/codegen.cc中的tvm::Build，后端代码生成与relay.build相同。
三．TVM代码生成分析
本节介绍一种生成代码方式tvm.build的流程。先以官方提供的一个向量相加程序为例，简单熟悉下Tensor Expression和lower后的代码，tvm版本都是0.7dev。
简单例子
Get Started with Tensor
Expressiontvm.apache.org/docs/tutorials/get_started/tensor_expr_get_started.html#sphx-glr-tutorials-get-started-tensor-expr-get-started-py
importtvm
importnumpyasnp
# Tensor Expression
# args: (shape, label)
A=tvm.placeholder((10,),name='A')
B=tvm.placeholder((10,),name='B')
# args: (shape, function, label)
# function represented in lambda expression (element-wise)
# lambda axis1, axis2, ... : f(axis1, axis2, ...)
C=tvm.compute((10,),lambdai:A[i]+B[i],name="C")
# generate schedule
s=tvm.create_schedule(C.op)
# print low level codes
print(tvm.lower(s,[A,B,C],simple_mode=True))
上述代码先定义了两个维度为10的张量A、B，C，一个维度为10的张量，每一个维度的值等于A、B对应值和，使用tvm接口定义了一个compute。s为tvm创建的计算调度schedule，按照默认情况，s会生成朴素的嵌套循环形式，实际可以通过打印输出lower后的代码验证。
for(i:int32,0,10){
C_2[i]=((float32*)A_2[i]+(float32*)B_2[i])
}
可以对原始schedule做一些优化，如将循环分割为内外两层循环，可以调用split方法实现。
# split(parent[, factor, nparts])
# Split the stage either by factor providing outer scope, or both. Return outer, inner vaiable of iteration.
bx,tx=s[C].split(C.op.axis[0],factor=2)
print(tvm.lower(s,[A,B,C],simple_mode=True))
打印lower后的代码，可以发现确实将原来的单层循环分割为内外循环。
for(i.outer:int32,0,5){
for(i.inner:int32,0,2){
C_2[((i.outer*2)+i.inner)]=((float32*)A_2[((i.outer*2)+i.inner)]+(float32*)B_2[((i.outer*2)+i.inner)])
}
}这个schedule变换没有带来任何好处，只是为了说明Tensor Expression应该怎么用。就可以调用build生成目标代码了，可以设置target和target_host，参考代码如下。
tgt_host="llvm"
# Change it to respective GPU if gpu is enabled Ex: cuda, opencl, rocm
tgt="llvm"# cuda llvm
n=10
fadd=tvm.build(s,[A,B,C],tgt,target_host=tgt_host,name="myadd")
ctx=tvm.context(tgt,0)
a=tvm.nd.array(np.random.uniform(size=n).astype(A.dtype),ctx)
b=tvm.nd.array(np.random.uniform(size=n).astype(B.dtype),ctx)
c=tvm.nd.array(np.zeros(n,dtype=C.dtype),ctx)
fadd(a,b,c)# run
# test
tvm.testing.assert_allclose(c.asnumpy(),a.asnumpy()+b.asnumpy())
print(fadd.get_source())
接下来重点梳理下调用tvm.build后代码生成的流程。
tvm.build
调用tvm.build后，先跳转到tvm/python/tvm/driver/build_module.py，其中的build函数主要做两个步骤：
- lower高层次代码
- 后端代码生成
build函数的输入类型可以是
ScheduleLoweredFunc[LoweredFunc]{target: [LoweredFunc]}
如果输入是schedule.Schedule，调用文件中lower函数进行Schedule优化。
defbuild(inputs,args=None,target=None,target_host=None,name="default_function",binds=None):
ifisinstance(inputs,schedule.Schedule):
ifargsisNone:
raiseValueError("args must be given for build from schedule")
input_mod=lower(inputs,args,name=name,binds=binds)
//skipsomecode.....
最终所有的输入，都整理成如下形式：
target_input_mod = {'target': [LoweredFunc]}
lower高层次代码
lower函数类似relay.build流程中的Optimize函数，但lower函数执行的是operator-level的优化，主要针对循环变换。
deflower(sch,args,name="main",binds=None,simple_mode=False):
# config setup
pass_ctx=PassContext.current()
instrument_bound_checkers=bool(pass_ctx.config.get(
"tir.instrument_bound_checkers",False))
disable_vectorize=bool(pass_ctx.config.get(
"tir.disable_vectorize",False))
add_lower_pass=pass_ctx.config.get("tir.add_lower_pass",[])
lower_phase0=[x[1]forxinadd_lower_passifx[0]==0]
lower_phase1=[x[1]forxinadd_lower_passifx[0]==1]
lower_phase2=[x[1]forxinadd_lower_passifx[0]==2]
lower_phase3=[x[1]forxinadd_lower_passifx[0]>2]
# Phase 0
ifisinstance(sch,schedule.Schedule):
mod=form_irmodule(sch,args,name,binds)
else:
mod=sch
pass_list=lower_phase0
# Phase 1
pass_list+=[
tvm.tir.transform.InjectPrefetch(),
tvm.tir.transform.StorageFlatten(64,instrument_bound_checkers),
tvm.tir.transform.BF16Legalize(),
tvm.tir.transform.NarrowDataType(32),
tvm.tir.transform.Simplify(),
]
pass_list+=lower_phase1
# Phase 2
ifnotsimple_mode:
pass_list+=[(tvm.tir.transform.LoopPartition())]
pass_list+=[
tvm.tir.transform.VectorizeLoop(notdisable_vectorize),
tvm.tir.transform.InjectVirtualThread(),
tvm.tir.transform.InjectDoubleBuffer(),
tvm.tir.transform.StorageRewrite(),
tvm.tir.transform.UnrollLoop()
]
pass_list+=lower_phase2
# Phase 3
pass_list+=[
tvm.tir.transform.Simplify(),
tvm.tir.transform.RemoveNoOp(),
]
pass_list+=[tvm.tir.transform.RewriteUnsafeSelect()]
pass_list+=[tvm.tir.transform.HoistIfThenElse()]
pass_list+=lower_phase3
# Instrument BoundCheckers
ifinstrument_bound_checkers:
pass_list+=[tvm.tir.transform.InstrumentBoundCheckers()]
optimize=tvm.transform.Sequential(pass_list)
mod=optimize(mod)
returnmod
lower函数后，对target device和target host分别生成代码，调用的代码如下所示。
defbuild(inputs,args=None,target=None,target_host=None,name="default_function",binds=None):
# skip some code.....
device_modules=[]
fortar,input_modintarget_input_mod.items():
# build for device module
mod_host,mdev=_build_for_device(input_mod,tar,target_host)
mod_host_all.update(mod_host)
device_modules.append(mdev)
# Generate a unified host module.
rt_mod_host=codegen.build_module(mod_host_all,target_host)
# Import all modules.
formdevindevice_modules:
ifmdev:
rt_mod_host.import_module(mdev)
returnrt_mod_host
后端代码生成
调用codegen.build_module，跳转至tvm/python/tvm/target/codegen.py文件，通过FFI对C++函数Build进行调用，命名空间是"target"。Build函数的C++实现在tvm/src/target/http://codegen.cc文件，后续流程就与relay.build一致了，根据不同的硬件平台生成代码。
runtime::ModuleBuild(IRModulemod,Targettarget){
if(transform::PassContext::Current()
->GetConfig<Bool>("tir.disable_assert",Bool(false))
.value()){
mod=tir::transform::SkipAssert()(mod);
}
std::stringbuild_f_name;
if(target->kind->name=="micro_dev"){
build_f_name="target.build.c";
}else{
build_f_name="target.build."+target->kind->name;
}
// the build function.constPackedFunc*bf=runtime::Registry::Get(build_f_name);
CHECK(bf!=nullptr)<<build_f_name<<" is not enabled";
return(*bf)(mod,target);
}
TVM_REGISTER_GLOBAL("target.Build").set_body_typed(Build);
参考链接：
https://www.cnblogs.com/jourluohua/p/10191269.html
https://www.h5w3.com/128623.html
https://zhuanlan.zhihu.com/p/258432371
References
- TVM Codebase Walkthrough by Example, https://docs.tvm.ai/dev/codebase_walkthrough.html
- TVM图编译器Relay简单探究 – 郑思泽的文章 – 知乎, https://zhuanlan.zhihu.com/p/91283238
- 谢睿峰, TVM/VTA代码生成流程, https://krantz-xrf.github.io/2019/10/24/tvm-workflow.html
- https://discuss.tvm.ai/t/relationship-between-tvm-build-and-relay-build/4166
浙公网安备 33010602011771号