登录
当前位置： 首页 > 工具软件 > tvm > 使用案例 >

TVM代码流程分析

谷梁宝
2023-12-01

TVM代码流程分析
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)
跟踪细节
那么对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();
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& 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_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>& inputs,
    const Target& target_host,
    const BuildConfig& config) {
    Array fhost_all;
    std::vectorruntime::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& 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 - Tensor Expression
    本节以向量加法为例，记录TVM最最基本的Tensor Expression的使用，以及简单的编译运行流程。
    下面的代码为简单的向量加法，参考自Tensor Expression官方教程，在TVM v0.6下执行（注意与v0.7dev的模块有区别)。
    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))
其中placeholder代表特定维度的张量，最后生成的代码会要求用户输入两个tensor，如果是C++代码，则要求用户输入两个float*。注意，会发现这个过程实际上是没有计算发生的，而只是定义了计算如何进行。
输出的low-level代码如下所示，还是相当好理解的，即i从0到10循环，循环内每次计算C[i]的值。
produce C {
for (i, 0, 10) {
C[i] = (A[i] + B[i])
}
}
一些常用的循环优化API可以在这里找到。这里使用循环分割split作为尝试。
split(parent[, factor, nparts])
Split the stage either by factor providing outer scope, or both. Return outer, innervaiable of iteration.
bx, tx = s[C].split(C.op.axis[0],factor=2)
print(tvm.lower(s,[A,B,C],simple_mode=True))
由于对schedule的操作是原地变换，因此可以直接输出lower后的代码，发现确实已经改变了，原来的循环体变成52的循环。
produce C {
for (i.outer, 0, 5) {
for (i.inner, 0, 2) {
C[((i.outer2) + i.inner)] = (A[((i.outer2) + i.inner)] + B[((i.outer2) + i.inner)])
}
}
}
当然这一个schedule变换并没有带来任何好处，只是为了说明Tensor Expression应该怎么用。
之后就可以调用build生成目标代码了，可以设置target和target_host。
tgt = “c” # “llvm”, “cuda”
fadd = tvm.build(s,[A,B,C],target=tgt,name=“myadd”)
然后可以创造运行时环境，进行运行测试。
n = 10
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())
生成的C代码如下
for (int32_t i_outer = 0; i_outer < 5; ++i_outer) {
for (int32_t i_inner = 0; i_inner < 2; ++i_inner) {
C[((i_outer * 2) + i_inner)] = (A[((i_outer * 2) + i_inner)] + B[((i_outer * 2) + i_inner)]);
}
}
生成的myadd.c完整代码如下
最后通过fadd.save(“myadd.c”)保存文件。
TVM - Relay IR Pass
本节介绍Relay IR Pass的构造。
Relay IR Pass核心依然是在C++中实现，但提供了Python接口，方便上层直接调用并对计算流图进行变换优化。
Pass管理器在include/tvm/relay/transform.h中，里面包含所有Pass的声明，希望做到
• 管理调度不同的优化pass
• 收集需要的分析信息，并且保持是最新的
• 减少程序员实现新pass的麻烦
Python的接口函数声明在python/tvm/relay/transform.py中，在python/tvm/relay/_transform.py中通过FFI对C++函数进行调用，命名空间为relay._transform。
具体C++的实现则分为两个部分：
• 高层IR图变换，源码在src/relay/pass中，集中变换则是在src/relay/backend/build_module.cc中的relay::Module Optimize
• 后端代码的图变换，源码在src/relay/backend/vm中，集中变换在python/tvm/build_module.py中的lower函数
Pass的构造
• PassInfo
• class PassInfoNode : public RelayNode {
• std::string name;
• int opt_level;
• std::vectorstd::string required;
• };
• PassContext
• class PassContextNode : public RelayNode {
• public:
• ErrorReporter err_reporter;
• int opt_level{2};
• int fallback_device{static_cast(kDLCPU)};
• tvm::Arraytvm::Expr required_pass;
• tvm::Arraytvm::Expr disabled_pass;
• };
•
• class PassContext : public NodeRef {
• public:
• TVM_DLL static PassContext Create();
• TVM_DLL static PassContext Current();
• /* Other fields are omitted. /
•
• private:
• // The entry of a pass context scope.
• TVM_DLL void EnterWithScope();
• // The exit of a pass context scope.
• TVM_DLL void ExitWithScope();
•
• // Classes to get the Python with like syntax.
• friend class tvm::With;
• };
•
• struct RelayPassContextThreadLocalEntry {
• /! \brief The default pass context. /
• PassContext default_context;
• /! \brief The current pass context. /
• std::stack context_stack;
• RelayPassContextThreadLocalEntry() {
• default_context = PassContext(make_node());
• }
• };
•
• /! \brief The thread-local store to hold the pass context. */
• typedef dmlc::ThreadLocalStore
• RelayPassContextThreadLocalStore;
• Pass Constructs：提供基类
• class PassNode : RelayNode {
• virtual PassInfo Info() const = 0;
• virtual Module operator()(const IRModule& mod
• const PassContext& pass_ctx) const = 0;
• };
也就是说，一个Pass一定是作用在特定context下的IRModule，所有Pass都设计成Module到Module的映射，完整Pass的定义在src/relay/ir/transform.cc和src/ir/transform.cc中。
Module-Level
class ModulePassNode : PassNode {
PassInfo pass_info;
runtime::TypedPackedFunc<Module(Module, PassContext)> pass_func;
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
// Other members/methods are omitted
};
Function-Level
class FunctionPassNode : PassNode {
PassInfo pass_info;
runtime::TypedPackedFunc<Function(Function, Module, PassContext)> pass_func;
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
bool SkipFunction(const Function& func) const;
// Other members/methods are omitted…
};
Sequential
类似于PyTorch中的nn.Sequential，顺序执行多个Pass
class SequentialPassNode : PassNode {
PassInfo pass_info;
// Passes need to be executed.
Array passes;
bool PassEnabled(const PassInfo& info) const;
Module operator()(const Module& mod, const PassContext& pass_ctx) const final;
};

References
• TVM内置Pass索引，https://docs.tvm.ai/api/python/relay/transform.html
• Relay Pass Infrastructure, https://tvm.apache.org/docs/dev/relay_pass_infra.html
• 初识TVM - 立交桥跳水冠军的文章 - 知乎，https://zhuanlan.zhihu.com/p/88188955
• 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
• https://blog.csdn.net/qq_33287871/article/details/113898181
• https://www.cnblogs.com/wangtianning1223/p/14662970.html

 类似资料：

相关阅读

Python实现线程池代码分享thinkphp3.2.3 分页代码分享Mybatis源码分析之存储过程调用和运行流程设置流程退出代码的正确方法？RocketMq事务消息发送代码流程详解

相关文章

MySQL的逻辑架构和SQL执行流程分析YAML流程8.30携程前端笔试代码题4.15 携程笔试 AK，附C++代码Bootstrap 代码

相关问答

WebPack4与代码拆分分流码验证表达式OAuth 2.0身份验证代码授予流程vs客户端凭据授予流程如何将部分php代码放入wordpress短代码在区分源代码、目标代码、汇编代码和机器代码时，我感到困惑

相关文档

Tendermint 源码分析Memcached 源码分析工作流程OpenStack Neutron 源码分析OpenStack HeatClient 源码分析
快捷导航： 新手教程 算法原理 架构设计 Java进阶 数据库进阶 大厂专栏 面试经验 编程笔记 编程问答 所有专题 文档资料 工具软件 电子书籍 小牛导航
在线工具： 房贷计算器 个税计算器 Linux命令查询 Json格式化 正则表达式 颜色转换 AES加解密 SHA1加密 MD5加密 毒鸡汤 字数统计 随机密码生成 进制转换 Base64编解码 励志句子
Copyright © 2019-2024 小牛知识库@xnip.cn. All Rights Reserved.