TVM代码流程分析
TVM代碼流程分析
TVM - 代碼生成流程
本節(jié)主要介紹TVM的代碼生成流程，即調(diào)用relay.build或tvm.build之后發(fā)生了什么，將深入到TVM的源代碼進(jìn)行剖析。（這里采用的依然是TVM v0.6）
首先區(qū)分兩個build的區(qū)別：tvm.build主要針對單一算子（參照Tensor Expression一文），而relay.build是針對整個模型進(jìn)行編譯（參照GCN優(yōu)化一文），而Relay最后也會調(diào)用到tvm::build做代碼生成。
relay.build
通常的模型編譯由以下兩條語句完成。
Build with Relay
with relay.build_config(opt_level=0):
graph, lib, params = relay.build(func, target, params=params)
跟蹤細(xì)節(jié)
那么對relay.build進(jìn)行跟蹤，跳轉(zhuǎn)進(jìn)來是python/tvm/relay/build_module.py（這里是因為在relay/init.py中將build函數(shù)直接import到relay的命名空間，因此跳過了build_module這一層），其中的build函數(shù)是build_module內(nèi)的全局函數(shù)(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是否有預(yù)先tune好的參數(shù)記錄，然后構(gòu)造tophub_context，在其內(nèi)部構(gòu)建了BuildModule之后，才跳轉(zhuǎn)到BuildModule.build，然后返回BuildModule.__init__中的內(nèi)容。
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 functionself._build(func, target, target_host)# Get artifactsgraph_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++函數(shù)建立聯(lián)系（tvm._ffi._cytpes.function.Function.call）。
from tvm._ffi.function import _init_api
_init_api(“relay.build_module”, name)
對應(yīng)的C++函數(shù)在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供調(diào)用，由于主要用的是build函數(shù)，因此到RelayBuildModule中找對應(yīng)的函數(shù)。這里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]);
});
// …
}
也就是調(diào)用的是this->Build，再跳轉(zhuǎn)過去會指向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());
}
}
經(jīng)過多番跳轉(zhuǎn)，終于到達(dá)build的核心模塊，再來看TVM逐步做的工作。
-
優(yōu)化
-
計算圖生成
-
后端代碼生成
優(yōu)化
先是優(yōu)化Optimize，可以看到這里的優(yōu)化主要是設(shè)備無關(guān)的優(yōu)化，是graph-level的針對tensor運(yùn)算的優(yōu)化。（這里的優(yōu)化pass都已經(jīng)在C++中實現(xiàn)，先前版本的NNVM似乎還是在Python中調(diào)用）
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;
}
計算圖生成
對應(yīng)GraphCodegen類，以同樣的方式調(diào)用src/relay/backend/build_module.cc中的relay.build_module.GraphRuntimeCodegen（一樣是FFI），然后跳轉(zhuǎn)至src/relay/backend/graph_runtime_codegen.cc，其中已經(jīng)用TVM_REGISTER_GLOBAL注冊了對應(yīng)函數(shù)，即用GraphRuntimeCodegenModule生成對應(yīng)Object。
因此實際graph_codegen->Codegen的函數(shù)是一個PackedFunc，定義在GraphRuntimeCodegen.Codegen，用來將relay::Function func進(jìn)行遍歷，然后生成計算圖。
后端代碼生成
Relay得到lower后的函數(shù)，最后一步則是交給tvm::build做代碼生成，跳轉(zhuǎn)到src/codegen/build_module.cc中的build函數(shù)（注意這里重載了幾個版本），然后跳轉(zhuǎn)到核心build，注意這里的build函數(shù)支持異構(gòu)編譯，只要再inputs劃分好不同硬件設(shè)施即可。
// 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;
}
當(dāng)中最最核心的則是mhost = codegen::Build，最后跳轉(zhuǎn)過去就開始調(diào)用代碼生成模塊了（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注冊，然后調(diào)用同個文件中的LLVMModuleNode->Init。這時會跳轉(zhuǎn)到src/codegen/llvm/codegen_llvm.cc中的CodeGenLLVM類進(jìn)行代碼生成。
tvm.build
用tvm.build對算子進(jìn)行編譯則是按照以下方式進(jìn)行調(diào)用，例子來自Tensor Expression。
s = tvm.create_schedule(C.op)
tgt = “l(fā)lvm” # “cuda”
fadd = tvm.build(s,[A,B,C],target=tgt,name=“myadd”)
調(diào)用tvm.build后首先跳轉(zhuǎn)到python/tvm/build_module.py，其中的build函數(shù)主要做兩個步驟：
-
lower高層次代碼
-
后端代碼生成
代碼變換
lower高層次代碼對應(yīng)的是
flist = lower(inputs,args,name=name,binds=binds)
而lower函數(shù)同樣在python/tvm/build_module.py中，類似于relay.build中的Optimize，但這里執(zhí)行的是operator-level的優(yōu)化，主要針對循環(huán)變換。
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 stmtreturn ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)
優(yōu)化Pass的主體實施都在src/api/api_pass.cc中，以tvm.ir_pass進(jìn)行注冊（注意由于C++函數(shù)中已經(jīng)在tvm的命名空間里，故搜索時直接搜ir_pass才會出來對應(yīng)的API）。
代碼生成
lower完之后就進(jìn)入到后端代碼生成，對應(yīng)build函數(shù)中的
mhost = codegen.build_module(fhost_all, str(target_host))
同樣的原理，跳轉(zhuǎn)至tvm/codegen.py，初始化tvm.codegen的API codegen._Build，調(diào)用FFI，跳轉(zhuǎn)至src/api/api_codegen.cc，最后跳轉(zhuǎn)至src/codegen/codegen.cc中的tvm::Build，之后的后端代碼生成則與relay.build相同。
TVM - Tensor Expression
本節(jié)以向量加法為例，記錄TVM最最基本的Tensor Expression的使用，以及簡單的編譯運(yùn)行流程。
下面的代碼為簡單的向量加法，參考自Tensor Expression官方教程，在TVM v0.6下執(zhí)行（注意與v0.7dev的模塊有區(qū)別)。
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*。注意，會發(fā)現(xiàn)這個過程實際上是沒有計算發(fā)生的，而只是定義了計算如何進(jìn)行。
輸出的low-level代碼如下所示，還是相當(dāng)好理解的，即i從0到10循環(huán)，循環(huán)內(nèi)每次計算C[i]的值。
produce C {
for (i, 0, 10) {
C[i] = (A[i] + B[i])
}
}
一些常用的循環(huán)優(yōu)化API可以在這里找到。這里使用循環(huán)分割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后的代碼，發(fā)現(xiàn)確實已經(jīng)改變了，原來的循環(huán)體變成52的循環(huán)。
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)])
}
}
}
當(dāng)然這一個schedule變換并沒有帶來任何好處，只是為了說明Tensor Expression應(yīng)該怎么用。
之后就可以調(diào)用build生成目標(biāo)代碼了，可以設(shè)置target和target_host。
tgt = “c” # “l(fā)lvm”, “cuda”
fadd = tvm.build(s,[A,B,C],target=tgt,name=“myadd”)
然后可以創(chuàng)造運(yùn)行時環(huán)境，進(jìn)行運(yùn)行測試。
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
本節(jié)介紹Relay IR Pass的構(gòu)造。
Relay IR Pass核心依然是在C++中實現(xiàn)，但提供了Python接口，方便上層直接調(diào)用并對計算流圖進(jìn)行變換優(yōu)化。
Pass管理器在include/tvm/relay/transform.h中，里面包含所有Pass的聲明，希望做到
? 管理調(diào)度不同的優(yōu)化pass
? 收集需要的分析信息，并且保持是最新的
? 減少程序員實現(xiàn)新pass的麻煩
Python的接口函數(shù)聲明在python/tvm/relay/transform.py中，在python/tvm/relay/_transform.py中通過FFI對C++函數(shù)進(jìn)行調(diào)用，命名空間為relay._transform。
具體C++的實現(xiàn)則分為兩個部分：
? 高層IR圖變換，源碼在src/relay/pass中，集中變換則是在src/relay/backend/build_module.cc中的relay::Module Optimize
? 后端代碼的圖變換，源碼在src/relay/backend/vm中，集中變換在python/tvm/build_module.py中的lower函數(shù)
Pass的構(gòu)造
? 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都設(shè)計成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，順序執(zhí)行多個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內(nèi)置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
總結(jié)
- 上一篇： TVM实现hardware backen
- 下一篇： 如何将自定义代码生成TVM