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)推测的步骤就是：

1. 补全相应的python接口
2. 找到python和C交互的接口
3. 正确维护中间代码的IR pass变换中新设备引入的特性
4. 代码生成对新设备和新特性的支持
5. 添加编译选项支持(非必须)

以下就分别就这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	#Build DPU set(USE_DPU ON)

cmake目录下存在着modules和util目录，modules是指定了相关设备的目录等配置，util文件夹下的内容，用来寻找如CUDA等的配置。暂时只需要modules下添加DPU.cmake，这部分的配置代码相对比较简单，就是指定runtime对应的目录。

1 2 3 4 5 6 7 8 9 10

# DPU Module     if(USE_DPU)   message(STATUS "Build with DPU support")   file(GLOB RUNTIME_DPU_SRCS src/runtime/dpu/*.cc)   list(APPEND RUNTIME_SRCS ${RUNTIME_DPU_SRCS}) else()   message(STATUS "NOT BUILD DPU SUPPORT") endif(USE_DPU)

这里修改完config.cmake，需要重新拷贝到build目录下，以使下次配置生效。编译tvm时是cmake目录下的config.cmake和CMakeLists.txt共同工作生效。在CMakeLists.txt中添加

1 2	tvm_option(USE_DPU "Build with DPU" ON) include(cmake/modules/DPU.cmake)

然后在build目录下，运行cmake命令，重新编译生效。

1 2	cmake  -DCMAKE_BUILD_TYPE=Debug ../ make

这里不加-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

with relay.build_config(opt_level=0):

graph, lib, params = relay.build(func, target, params=params)

跟踪细节

如何进行代码跟踪，一方面可以直接通过VS Code在函数上方Alt+单击跳转，另一方面如果想有更直观的印象，可以利用pycallgraph进行可视化（需先用pip安装），代码如下，用GCN的代码编译模块。

from pycallgraph import PyCallGraph

from pycallgraph.output import GraphvizOutput

from pycallgraph import Config

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

if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):

tophub_context = autotvm.tophub.context(list(target.values()))

else:

tophub_context = autotvm.util.EmptyContext()

with tophub_context:

bld_mod = BuildModule()

graph_json, mod, params = bld_mod.build(func, target, target_host, params)

return graph_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.

        if params:

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()

return graph_json, mod, params

_build_module._BuildModule()通过FFI，在python/tvm/relay/_build_module.py中，与C++函数建立联系（tvm._ffi._cytpes.function.Function.__call__）。

from tvm._ffi.function import _init_api

_init_api("relay.build_module", __name__)

对应的C++函数在src/relay/backend/build_module.cc

runtime::Module RelayBuildCreate() {

auto exec = make_object<RelayBuildModule>();

return runtime::Module(exec);

}

TVM_REGISTER_GLOBAL("relay.build_module._BuildModule")

.set_body([](TVMArgs args, 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") {

return PackedFunc([sptr_to_self, this](TVMArgs args, 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::Module relay_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>(new GraphCodegen());

graph_codegen_->Init(nullptr, targets_);

graph_codegen_->Codegen(func);

ret_.graph_json = graph_codegen_->GetJSON();

ret_.params = graph_codegen_->GetParams();

auto lowered_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逐步做的工作。

1. 优化
2. 计算图生成
3. 后端代码生成

优化

先是优化Optimize，可以看到这里的优化主要是设备无关的优化，graph-level针对tensor运算的优化。（这里的优化pass都已经在C++中实现，先前版本的NNVM似乎还是在Python中调用）

  relay::Module Optimize(

Function func,

const TargetsMap& targets,

const std::unordered_map<std::string, runtime::NDArray>& params) {

// BindParamsByName(func, params)

// Perform Module->Module optimizations.

relay::Module relay_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::Pass seq = transform::Sequential(pass_seqs);

// ... judge & do

relay_module = seq(relay_module);

// Handle heterogeneous compilation.

transform::PassContext pass_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());

return relay_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;

Target target_host_val = target_host;

if (!target_host.defined()) {

for (const auto& 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 (const auto& it : inputs) {

auto host_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::Module mdev = DeviceBuild(fdevice, it.first);

for (const auto& it : fhost) {

fhost_all.push_back(it);

}

device_modules.push_back(mdev);

}

runtime::Module mhost = codegen::Build(fhost_all, target_host_val->str());

// Import all modules

for (const auto& it : device_modules) {

if (it.operator->()) {

mhost.Import(it);

}

}

return mhost;

}

最核心是mhost = codegen::Build，最后跳转过去就开始调用代码生成模块了（src/codegen/codegen.cc）。

runtime::Module Build(const Array<LoweredFunc>& funcs,

const std::string& target) {

// do something

std::string build_f_name = "codegen.build_" + mode;

// the build function.

const PackedFunc* bf = runtime::Registry::Get(build_f_name);

runtime::Module m = transformed_funcs.empty() ?

(*bf)(funcs, target) :

(*bf)(transformed_funcs, target);

return m;

}

以生成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函数主要做两个步骤：

1. lower高层次代码
2. 后端代码生成

代码变换

lower高层次代码对应

flist = lower(inputs,args,name=name,binds=binds)

lower函数同样在python/tvm/build_module.py中，类似于relay.build中的Optimize，但这里执行的是operator-level的优化，主要针对循环变换。

def lower(sch,

args,

name="default_function",

binds=None,

simple_mode=False):

# initialization

# Phase 0

    if isinstance(sch, schedule.Schedule):

stmt = form_body(sch)

for f 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)

for f in lower_phase1:

stmt = f(stmt)

# Phase 2

    if not simple_mode:

stmt = ir_pass.LoopPartition(stmt, cfg.partition_const_loop)

if cfg.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)

for f in lower_phase2:

stmt = f(stmt)

# Phase 3

    stmt = ir_pass.Simplify(stmt)

stmt = ir_pass.RemoveNoOp(stmt)

if not cfg.disable_select_rewriting:

stmt = ir_pass.RewriteUnsafeSelect(stmt)

for f in lower_phase3:

stmt = f(stmt)

# Instrument BoundCheckers

    if cfg.instrument_bound_checkers:

stmt = ir_pass.InstrumentBoundCheckers(stmt)

if simple_mode:

return stmt

return ir_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

import tvm

import numpy as np

# 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,), lambda i: 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函数的输入类型可以是

• Schedule
• LoweredFunc
• [LoweredFunc]
• {target: [LoweredFunc]}

如果输入是schedule.Schedule，调用文件中lower函数进行Schedule优化。

def build(inputs,args=None,target=None,target_host=None,name="default_function",binds=None):

    if isinstance(inputs, schedule.Schedule):

        if args is None:

            raise ValueError("args must be given for build from schedule")

        input_mod = lower(inputs, args,name=name,binds=binds)

    // skip some code.....

最终所有的输入，都整理成如下形式：

target_input_mod = {'target': [LoweredFunc]}

lower高层次代码

lower函数类似relay.build流程中的Optimize函数，但lower函数执行的是operator-level的优化，主要针对循环变换。

def lower(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] for x in add_lower_pass if x[0] == 0]

    lower_phase1 = [x[1] for x in add_lower_pass if x[0] == 1]

    lower_phase2 = [x[1] for x in add_lower_pass if x[0] == 2]

    lower_phase3 = [x[1] for x in add_lower_pass if x[0] > 2]

    # Phase 0

    if isinstance(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

    if not simple_mode:

        pass_list += [(tvm.tir.transform.LoopPartition())]

    pass_list += [

        tvm.tir.transform.VectorizeLoop(not disable_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

    if instrument_bound_checkers:

        pass_list += [tvm.tir.transform.InstrumentBoundCheckers()]

    optimize = tvm.transform.Sequential(pass_list)

    mod = optimize(mod)

    return mod

lower函数后，对target device和target host分别生成代码，调用的代码如下所示。

def build(inputs,args=None,target=None,target_host=None,name="default_function",binds=None):

    # skip some code.....

    device_modules = []

    for tar, input_mod in target_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.

    for mdev in device_modules:

        if mdev:

            rt_mod_host.import_module(mdev)

    return rt_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::Module Build(IRModule mod, Target target) {

  if (transform::PassContext::Current()

          ->GetConfig<Bool>("tir.disable_assert", Bool(false))

          .value()) {

    mod = tir::transform::SkipAssert()(mod);

  }

  std::string build_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.

  const PackedFunc* 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