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

跟踪细节

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

with PyCallGraph(output=graphviz,config=config):

# Build with Relay

    with relay.build_config(opt_level=0):

• 各函数之间的调用关系，如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

else:

先是寻找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):

def build(self, func, target=None, target_host=None, params=None):

# Setup the params.

self._set_params(params)

# Build the function

# Get artifacts

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

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

runtime::Module RelayBuildCreate() {

}

TVM_REGISTER_GLOBAL("relay.build_module._BuildModule")

});

就是注册了一个RelayBuildModule供调用，由于主要用的是build函数，因此到RelayBuildModule中找对应的函数。这里TVM用PackedFunc做了一层封装。

PackedFunc GetFunction(const std::string& name,

const ObjectPtr<Object>& sptr_to_self) final {

// ...

});

// ...

}

就是调用的是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

// Get the updated function.

// Generate code for the updated function.

graph_codegen_->Codegen(func);

lowered_funcs,

target_host_,

BuildConfig::Current());

}

}

经过多番跳转，终于到达build的核心模块，再来看TVM逐步做的工作。

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

优化

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

Function func,

const TargetsMap& targets,

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

// BindParamsByName(func, params)

// Perform Module->Module optimizations.

// 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.

// ... judge & do

// Handle heterogeneous compilation.

}

// Fuse the operations if it is needed.

CHECK(relay_module.defined());

}

计算图生成

对应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) {

break;

}

}

}

}

// Get the module for a certain target.

fhost_all.push_back(it);

}

device_modules.push_back(mdev);

}

// Import all modules

mhost.Import(it);

}

}

}

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

// do something

// the build function.

}

以生成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。

调用tvm.build后首先跳转到python/tvm/build_module.py，其中的build函数主要做两个步骤：

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

代码变换

lower高层次代码对应

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

args,

name="default_function",

binds=None,

simple_mode=False):

# initialization

# Phase 0

# Phase 1

# Phase 2

else:

stmt,

cfg.auto_unroll_max_step,

cfg.auto_unroll_max_depth,

cfg.auto_unroll_max_extent,

cfg.unroll_explicit)

# Phase 3

# Instrument BoundCheckers

优化Pass的主体实施都在src/api/api_pass.cc中，以tvm.ir_pass进行注册（由于C++函数中已经在tvm的命名空间里，搜索时直接搜ir_pass，出来对应的API）。

代码生成

lower完之后就进入到后端代码生成，对应build函数中的

同样，跳转至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

# Tensor Expression

# args: (shape, label)

# args: (shape, function, label)

# function represented in lambda expression (element-wise)

#     lambda axis1, axis2, ... : f(axis1, axis2, ...)

# generate schedule

# 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后的代码验证。

}

可以对原始schedule做一些优化，如将循环分割为内外两层循环，可以调用split方法实现。

# split(parent[, factor, nparts])

# Split the stage either by factor providing outer scope, or both. Return outer, inner vaiable of iteration.

print(tvm.lower(s,[A,B,C],simple_mode=True))

打印lower后的代码，可以发现确实将原来的单层循环分割为内外循环。

  }

这个schedule变换没有带来任何好处，只是为了说明Tensor Expression应该怎么用。就可以调用build生成目标代码了，可以设置target和target_host，参考代码如下。

# Change it to respective GPU if gpu is enabled Ex: cuda, opencl, rocm

# test

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优化。

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

target_input_mod = {'target': [LoweredFunc]}

lower高层次代码

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

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

后端代码生成

调用codegen.build_module，跳转至tvm/python/tvm/target/codegen.py文件，通过FFI对C++函数Build进行调用，命名空间是"target"。Build函数的C++实现在tvm/src/target/http://codegen.cc文件，后续流程就与relay.build一致了，根据不同的硬件平台生成代码。

  // the build function.

}

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