TVM基础编程示例分析

TVM基础编程示例分析

一．TVM编程基础示例

前言

继前图灵奖获得者Hennessy和Patterson在ISCA 2018提出“A New Golden Age for Computer Architecture”，编译器界大神Chris Lattner在ASPLOS 2021提出了“The Golden Age of Compiler Design”。另一方面，2020年图灵奖授予了编译器“龙书”作者Jeffrey Ullman和Alfred Aho。编译器技术在新的时代背景下似乎又再次焕发了新的活力，成为了业界的热点。

作为现在最热门的AI计算场景，与编译器技术的结合自然成为了大家不约而同的技术路线。机器学习跨入深度学习时代后，比较老一代的计算框架基本将神经网络建模为计算图，其中算子为节点，张量为边。然后以拓扑序执行，辅以并行优化等。这种范式下，为了达到好的性能，一般需要对网络中的算子深度优化。但是，今天的神经网络结构日益复杂，算子种类也更加繁多。不同的算子参数、输入配置以及算子间的融合，使得需要优化的算子数量组合爆炸，一一硬扛不切实际，而且很多时候也缺乏专家经验和开发时间。为了挖掘极致的性能，同时使得新算子实现更为方便，基于编译技术的方法成为了主流。像TVM，XLA，Glow，nGraph，MindSpore，Jittor，MegEngine，ONNC，Tiramisu等等用到或是基于编译技术的计算框架层出不穷。

在这个方向上，TVM可以说是先驱者，一个端到端的深度学习编译器，在平台兼容性和性能等方面都有很好的表现，社区也非常活跃。但TVM代码读起来不太容易理解（编译器的代码好像都不太好读…）。TVM经过几年的快速演进，今天已是一个比较复杂的系统了，里边的功能很多。可以通过过一个最简单的例子来看看其大致处理流程。本文通过官方教程Working with Operators Using Tensor Expressions中的例程vecadd为例，介绍TVM的流程示例。

import tvm

import os

n = 1024

A = tvm.te.placeholder((n,), name='A')

B = tvm.te.placeholder((n,), name='B')

C = tvm.te.compute(A.shape, lambda i: A[i] + B[i], name="C")

s = tvm.te.create_schedule(C.op)

# outer, inner = s[C].split(C.op.axis[0], factor=64)

# s[C].parallel(outer)

tgt = tvm.target.Target(target="llvm", host="llvm")

fadd = tvm.build(s, [A, B, C], tgt, name="vecadd")

dev = tvm.device(tgt.kind.name, 0)

a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev)

b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev)

c = tvm.nd.array(np.zeros(n, dtype=C.dtype), dev)

fadd(a, b, c)

程序做的事就是两个向量的逐元素相加。这个case中不考虑复杂算子，不考虑Relay，不考虑复杂pass，不考虑复杂的schedule，不考虑auto-tuning机制，不考虑graph runtime等。也正是因为简单，分析处理流程可以抓住主干，避免陷入复杂的细节。麻省虽小，五脏俱全。包含了TVM主要流程中的几个关键要素。

整个过程会分量部分介绍。第一部分主要涉及计算定义与schedule的创建。TVM是基于Halide中algorithm与schedule分离的思想。简单而言，前者指定算什么，后者指定怎么算。下面两节就是分别对应计算的定义与schedule的构建。

定义计算

现实使用当中，多数情况下会通过前端的解析器，从已有的机器学习模型中导入。如from_onnx.py中的relay.frontend.from_onnx()函数，可以从onnx模型导入。但上面例子是单个算子的例子，直接通过TE（Tensor expression）定义的。

先来看下例子中的计算定义部分：

A = tvm.te.placeholder((n,), name='A')

B = tvm.te.placeholder((n,), name='B')

C = tvm.te.compute(A.shape, lambda i: A[i] + B[i], name="C")

通过TEDD，可构建可视化图如下：

上面语句中，首先通过placeholder()函数创建tensor对象。调用_ffi_api.Placeholder()函数，从Python调到C++层构建PlaceholderOpNode对象，然后输出tensor返回。主要流程如下：

te.placeholder() # operation.py

return _ffi_api.Placeholder() # placeholder_op.cc

return placeholder()

return PlaceholderOp(...).output(0) # tensor.cc

n = make_object<PlaceholderOpNode>();

...

data_ = std::move(n);

这里的返回类型，或者说上面的A，B类型为tvm.te.tensor.Tensor。C++层对应TensorNode类。TensorNode中关联的Operation对象，代表通过什么操作计算得到的。Operation的output()函数可以得到输出tensor。OperationNode的InputTensors()函数（纯虚函数，在各继承类中会实现，如ComputeOpNode::InputTensors()）得到输入tensor。通过这样的方式在逻辑上形成计算图，表示了相互间的依赖关系。

接下去的compute()函数（实现在operation.py），主要用于根据给定用TE描述的计算，构建一个新的tensor。主要流程如下：

compute(shape, fcompute, ...) # operation.py

...

dim_var = [tvm.tir.IterVar((0, s), x, 0) for x, s in zip(arg_names, shape[:out_ndim])] # expr.py

body = fcompute(*[v.var for v in dim_var])

body = convert(body)

op_node = _ffi_api.ComputeOp(name, tag, attrs, dim_var, body)

outputs = tuple(op_node.output(i) for i in range(num))

return outputs[0] if num == 1 else outputs

其中有几个关键步骤：

1. 为每个axis创建tvm.tir.IterVar，对应循环变量。如上例中就只有一个axis，范围为[0,1024)。对应的C++层的IterVar类定义在var.h文件中。

2. 语句body = fcompute(*[v.var for v in dim_var])最为关键，调用传入的lambda函数，返回的body类型为tvm.tir.expr.Add（继承关系：->BinaryOpExpr->PrimExprWithOp->ExprOp & PrimExpr）。lambda函数中的A[i]类型为TensorSlice（继承自ObjectGeneric与ExprOp），代表Tensor的切片。调用下面的函数前会使用TensorSlice::asobject()函数，转成ProducerLoad（expr.py和expr.h）对象，继承自PrimExpr。这里由于是加操作，因此会调用ExprOp的操作符重载函数__add__()。继而调用add()函数（定义在tir/generic.py）。该函数调用到C++层，相应的函数在tir/op/op.cc中，通过下面的宏注册：REGISTER_MAKE_BINARY_OP(_OpAdd, add);。实现如下：

PrimExpr add(PrimExpr a, PrimExpr b, Span span) {

BinaryOpMatchTypes(a, b, span);

PrimExpr ret = arith::TryConstFold<tir::Add>(a, b);

if (ret.defined()) return ret;

return tir::Add(a, b, span);

}

返回的是tir::Add对象，对应Python中的Add对象（定义在tir/expr.py）。

调用convert()函数（实现在object_generic.py），对body对象进行转换，转化为TVM对象。经过转换后body类型为tvm.ir.container.Array。

创建C++层的ComputeOp对象（实现在compute_op.cc）。这个对象中包含ComputeOpNode对象的引用。C++层中ComputeOp（继承自Operaton），对应Python中的对象类型为te.tensor.ComputeOp。Python层中ComputeOp（继承关系：ComputeOp->BaseComputeOp->Operation）。最后返回output张量对象，类型为te.tensor.Tensor。

对于上面的例子，构建的数据结构大体如下：

二．TVM调用llvm编译

前面基于一个最基本的case，介绍了TVM中计算的定义与schedule的构建。这里继续介绍接下去的一个重点部分，就是编译。

有了前面构建的schedule后，接着就需要编译生成目标代码了。这个工作主要由tvm.build()和relay.build()两个函数完成。区别在于应用目标的范围，前者用于单个算子，后者用于整个网络。由于网络可看作由算子组成，后者会调用前者。本例中是针对单个算子的，因此这里使用的是前者：

tgt = tvm.target.Target(target="llvm", host="llvm")

fadd = tvm.build(s, [A, B, C], tgt, name="vecadd")

其中最主要的build()函数定义在driver/build_module.py文件中。该函数基于给定参数构建出可调用的目标函数。按照官方介绍里的说法，主要做两个工作：

l Lowering：将high-level的循环嵌套结构，转换成最终的low-level的IR。

l Codegen：从low-level的IR生成目标机器代码。

该函数的第一个参数是前面构建出来的schedule，第二个参数是函数的参数列表，第三个参数是target。提供用于lowering和codegen所需的目标平台信息。代码中对应的Target对象定义在target.*文件中。构造函数有两个参数，第一个参数target指示目标平台的配置。配置项如：

kind: 平台类型，基本决定了生成的代码是在什么处理器上运行。注册的target kind详细见target_kind.cc，有llvm, c, cuda, nvptx, romc, opencl, metal, vulkan, hexagon等。

keys: 如kind是opencl的话，key可以是mali, opencl, gpu。

device：对应实际运行的设备，会添加到keys后面。

libs：外部库，如cblas, cudnn, cublas, mkl这些。

…

另外，参数host与target类似，但用于指示host平台。如果taret平台为cuda，毕竟GPU还是不能完全脱离CPU运行，因此还需要host的代码做胶水，如内存分配，kernel启动这些。默认为llvm。

Lowering过程可以单独用tvm.lower()函数完成，如：

m = tvm.lower(s, [A, B, C], name="vecadd")

rt_mod = tvm.build(m, target="llvm")

也可以通过tvm.build()函数完成（因为一进去就会先调用lower()函数）。lower()函数的主要流程相关代码：

lower(sch, args, name="main", ...) // driver/build_module.py

// Handle add_lower_pass, if any.

lower_phases0 = ...

...

// According to the given schedule, form a function (in IRModule).

mod = form_irmodule(sch, args, ...) // build_module.py

sch.normalize()

Schedule::normalize() // schedule_dataflow_rewrite.cc

InjectInline()

RebaseNonZeroMinLoop()

LegalizeInvalidAttach()

bounds = schedule.InferBound(sch)

InferBound() // bound.cc

stmt = schedule.ScheduleOps(sch, bounds)

ScheduleOps() // schedule_ops.cc

body = Stmt()

// scan init and scan updates

...

for each stage in schedule: // in reverse order

body = MakePipeline(stage, dom_map, body, ...)

SchedulePostProc post_proc

post_proc.Init(sch)

return post_proc(body)

compact = schedule.VerifyCompactBuffer(stmt)

binds, arg_list = get_binds(args, compact, binds)

stmt = schedule.SchedulePostProcRewriteForTensorCore(stmt, sch, ...)

// func type: PrimFunc

func = schedule.SchedulePostProcToPrimFunc(arg_list, stmt, ...) // schedule_postproc_to_primfunc.cc

// Prepare parameters

...

return tie::PrimFunc(params, body, ...)

// name: vecadd

func = func.with_attr("global_symbol", name)

// Set functions

return tvm.IRModule({name: func})

// Phase 0: InjectPrefetch, StorageFlatten, BF16Legalize, NarrowDataType, Simplify

pass_list = lower_phase0

// Phase 1: LoopPartition, VectorizeLoop, InjectVirtualThread, InjectDoubleBuffer, StorageRewrite, UnrollLoop

pass_list += lower_phase1

// Phase 3: Simplify, RemoveNoOp, RewriteUnsafeSelect, HoistIfThenElse

pass_list += lower_phase2

// Apply the above passes.

optimize = tvm.transform.Sequential(pass_list)

mod = optimize(mod)

// mod type: tvm.ir.module.IRModule

return mod

主要根据参数给的schedule与参数生成对应的IRModule对象（定义在ir/module.h中）。IRModule是软件栈中所有IR变换的基础单元。维护函数与类型定义。这里的各种pass就是在IRModule上进行并吐出IRModule。

其中几个主要数据结构关系如下：

lower()函数中有四个阶段，第一个阶段中通过form_irmodule()函数，根据给定的schedule生成IRModule对象，然后在这个IRModule对象上，应用4轮的pass。这些pass主要分为几个阶段，分别是：

Phase 0：使用者自定义的pass。

Phase 1：使用者自定义的pass。以及：

InjectPrefetch

StorageFlatten

BF16Legalize

NarrowDataType

Simplify

Phase 2：使用者自定义的pass。以及：

LoopPartition

VectorizeLoop

InjectVirtualThread

InjectDoubleBuffer

StorageRewrite

UnrollLoop

Phase 3：使用者自定义的pass。以及：

Simplify

RemoveNoOp

RewriteUnsafeSelect

HoistIfThenElse

InstrumentBoundCheckers

这里pass其实是编译构建过程中的精华之一。

lower()函数的最后返回经过上面多轮pass优化后的IRModule对象。其中form_irmodule()函数是相对比较复杂的一部分，主要负责生成最初的IRModule对象。几个关键步骤如下：

l Schedule::normalize()函数规范化给定的schedule。主要实现在schedule_dataflow_rewrite.cc文件中。调用以下三个函数。

InjectInline()函数处理算子内联。用到调度原语 compute_inline的话会用到。

RebaseNonZeroMinLoop()函数将循环迭代的最小界置为0。感觉有点canonicalization的意思。

LegalizeInvalidAttach()函数处理在使用调度原语compute_at时且目标迭代又被split或fuse情况下的合法化。

l InferBound()函数顾名思义就是边界推导（Bound inference），主要用于推导循环边界。更具体地，就是确定每个IterVar的范围，返回IterVar到Range的映射，即每个循环变量的范围。这个信息在后面的MakeLoopNest()函数中，用于确定for循环的范围，在BuildRealize()函数中设置缓冲的大小。具体可参见官方文档 InferBound Pass。

l ScheduleOps()函数基于前面经过一些处理后的Schedule对象和推导出来的循环边界产生Stmt对象。表示一个初始的循环嵌套结构。C++层中的Stmt为所有语句（Statement）的容器。子类有LetStmt，AttrStmt，AssertStmt，Store，Allocate，SeqStmt，IfThenElse，Evaluate，For，While等等。该函数会处理schedule的依赖，核心部分是逆向遍历Schedule当中的Stage（对于上面例子中就是先Compute Op，再两个Placeholder Op）。对于每个stage（PlaceholderOp除外），根据attach type调用相应的逻辑。

l 对于上面的例子，Compute Op没有attach在其它计算中，因此对应Stage的attach type为kGroupRoot，因此这里调用MakePipeline()函数产生Stmt。这步比较关键比较复杂，后面再展开。

l 然后通过SchedulePostProc对象（继承自StmtExprMutator），对前面生成的Stmt进行后处理。

l get_binds()函数用于绑定buffer。给每个参数张量分配buffer。如对于上面例子中的A, B, C三个张量，分别通过tvm.tir.decl_buffer()，创建buffer并绑定张量。

l SchedulePostProcToPrimFunc()函数基于ScheduleOps()产生的Stmt创建PrimFunc对象，可以用于TIR优化。PrimFunc代表包含了TIR statement的primitive function，是low-level的代码表示。

l 创建IRModule对象。基于上面生成的对象封装成IRModule对象并返回。一个IRModule可以有多个函数，比较简单的情况下就一个。

上面第ScheduleOps()函数中，会调用MakePipeline()函数，针对ComputeOp对应Stage，返回一条由Stmt组成的pipeline，大体流程相关代码如下：

MakePipeline(Stage, unordered_map<IterVar, Range>, Stmt, ...) // schedule_ops.cc

producer = s->op->BuildProvide(stage, ...) // ComputeOpNode::BuildProvide() in compute_op.cc

ComputeType ctype = DetectComputeType(this, stage)

MakeComputeStmt(...) // compute_op.cc

ComputeLoopNest n = ComputeLoopNest::Create(...) // compute_op.cc

ComputeLoopNest ret

// make main loop nest

ret.main_nest = MakeLoopNest(stage, dom_map, ...) // op_utils.cc

vector<vector<Stmt>> nest

nest.resize(leaf_iter_vars.size() + 1)

for iter_var in leaf_iter_vars:

nest[i + 1].emplace_back(For(var, 0, dom->extent, kind, no_op))

nest[i + 1].emplace_back(AttrStmt(iv, tir::attr::loop_scope, iv->var, no_op))

...

n.init_nest.emplace_back(MakeIfNest(n.init_predicates))

n.main_nest.emplace_back(MakeIfNest(n.main_predicates))

if has reduce_axis:

...

else:

vector<Stmt> provides

...

// Array<Stmt> -> SeqStmt

Stmt provide = SeqStmt::Flatten(provides) // stmt.h

provide = MergeNest(n.main_nest, provide) // ir_utils.cc

return Substitute(provide, n.main_vmap) // stmt_functor.cc

Stmt pipeline = producer

pipeline = s->op->BuildRealize(stage, dom_map, pipeline)

// set the sizes of allocated buffers

BaseComputeOpNode::BuildRealize(stage, realize_map, body) // compute_op.cc

Stmt realize = body

realize = tir::ProducerRealize(...)

pipeline = AttrStmt(s->op, tir::attr::realize_scope, ..., pipeline)

return pipeline

MakePipeline()函数主要步骤如下：

ComputeOpNode::BuildProvide()函数主要创建ComputeOp对应的循环嵌套，对应的那些Stmt对象并串成pipeline。

首先用DetectComputeType()函数检测计算类型。遍历当前Stage的所有当前有效IterVar对象，根据属性判断计算类型，对于上面的简单例子这里为ComputeType::kNormal。

然后根据类型调用相应函数创建Stmt对象。这里对应地是调用MakeComputeStmt()函数。

根据Stage对象和边界推导的结果，通过ComputeLoopNest::Create()函数，创建ComputeLoopNest对象。该对象表示循环嵌套，几个主要成员：

init_predicates与main_predicates：类型为vector<PrimExpr>。表示每个循环的边界判断，调用MakeBoundCheck()函数来生成。

init_nest与main_nest：类型为vector<vector<Stmt>>。其中main_nest是最主要的表示循环嵌套的对象，对于上面的例子，经过split后，包含两个for循环。

根据main_predicates创建对应的Stmt（如有），用于在循环中判断该predicate是否成立，添加到main_nest结构中。

根据有无reduce axis走不同的path。如果没有的话（如本例），对于ComputeOp的body中的每一个输出，创建ProducerStore对象，再通过MergeNest()函数将之与主嵌套main_nest合并。

通过Substitute()函数，基于main_vmap（在MakeLoopNest()函数中准备）进行替换。

如schedule中设置了double buffer（如s[A].double_buffer），添加对应的AttrStmt。通过增大额外的buffer，达到达到计算与访存的重叠。本例中没用到。

如传入的consumer有定义且不是no op（指无定义、const init的EvaluateNode，或者是长度为0的SeqStmtNode），添加SeqStmt，将producer与consumer串连。本例中也不适用。

调用BuildRealize()函数。对于每个输出的张量，在pipeline中加入ProducerRealize节点。

最后，在pipeline中添加AttrStmt节点，标注操作的范围，返回该pipeline。

对于前面vecadd的例子，得到的pipeline大致如下示意图：

整个lower()函数后完成后的IR（TIR），打印出来如下：

primfn(A_1: handle, B_1: handle, C_1: handle) -> ()

attr = {"global_symbol": "main", "tir.noalias": True}

buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024], []),

B: Buffer(B_2: Pointer(float32), float32, [1024], []),

A: Buffer(A_2: Pointer(float32), float32, [1024], [])}

buffer_map = {A_1: A, B_1: B, C_1: C} {

for (i.outer: int32, 0, 16) {

for (i.inner: int32, 0, 64) {

C_2[((i.outer*64) + i.inner)] = ((float32*)A_2[((i.outer*64) + i.inner)] + (float32*)B_2[((i.outer*64) + i.inner)])

}

Lowering完成后，接下去就是build了。Build的主要流程相关代码如下：

build() # driver/build_module.py

input_mod = lower(inputs, args, ...)

mod_host_all = tvm.IRModule()

for tar, input_mod in target_input_mod.items():

# build the lowered functions for a device with the given compilation

mod_host, mdev = _build_for_device(input_mod, tar, target_host)

# input_mod type: IRModule

mod_mixed = input_mod

# Apply passes: ThreadSync, InferFragment, LowerThreadAllreduce, MakePackedAPI, SplitHostDevice

...

# Device optimizations: Filter, LowerWarpMemory, ,Simplify, LowerDeviceStorageAccessInfo, LowerIntrin

...

mod_dev = opt_device(mod_mixed) # IRModule

# Host optimization: LowerTVMBuiltin, LowerDeviceStorageAccessInfo, CustomDataType, LowerIntrin, CombineContextCall

...

mod_host = opt_host(mod_mixed) # IRModule

# Build IRModule into Module

# If there are dev functions

rt_mod_dev = codegen.build_module(mod_dev, target) # target/codegen.py

_ffi_api.Build(mod, target) # codegen.py

# mod_host type: IRModule, rt_mod_dev type: Module

return mod_host, rt_mod_dev

mod_host_all.update(mod_host)

# Insert functions in another Module to current one

_ffi_api.Module_Update()

IRModuleNode::Update() # ir/module.cc

device_modules.append(mdev)

# Generate a unified host module (type: runtime.Module)

rt_mod_host = codegen.build_module(mod_host_all, target_host)

# Create LLVMModuleNode and return the corresponding Module

_ffi_api.Build(mod, target) # target/codegen.cc

# Import all modules

for mdev in device_modules:

rt_mod_host.import_module(mdev)

_LIB.TVMModImport(mod, dep) # c_runtime_api.cc

GetModuleNode(mod)->Import(...) # runtime/module.cc

imports_.emplace_back(...)

return rt_mod_host # runtime.module.Module

target_input_mod包含了前面lowering输出的需要编译的IRModule及相应的target信息。如LLVM(CPU)为target，就是：{"llvm -keys=cpu -link-params=0", IRModule}。如cuda为target，可能就是{“cuda -keys=cuda,gpu -max_num_threads=1024 -thread_warp_size=32"， IRModule}。对于简单的case，target_input_mod只包含一个元素，_build_for_device()函数返回host端的IRModule，以及target端的Module（如cuda平台C++层对应CUDAModuleNode对象）。然后将host端IRModule生成一个统一的host模块，再将前面生成的对应target的Module导入其中。

这里mod_host_all与mod_host的类型为tvm.ir.module.IRModule。rt_mod_host与mdev的类型为tvm.runtime.module.Module。注意mdev只有当目标为非CPU（如GPU等）平台时才会有，当target为llvm（即for CPU）时mdev为空。

这个流程大体示意图如下：

其中比较核心和重要的部分是Build()函数，实现在codegen.cc文件中。会调用到具体后端的编译函数，进行目标代码生成。如cuda平台对应函数定义在build_cuda_on.cc文件中，llvm在llvm_module.cc文件中。以llvm后端为例，主要流程相关代码为：

TVM_REGISTER_GLOBAL("target.build.llvm")

.set_body_typed([](IRModule mod, Target target) -> runtime::Module {

auto n = make_object<LLVMModuleNode>();

n->Init(mod, target); // llvm_module.cc

InitializeLLVM();

llvm::InitializeAllTargetInfos();

llvm::InitializeAllTargets();

...

unique_ptr<CodeGenLLVM> cg = CodeGenLLVM::Create(...) // codegen_llvm.cc

// Call the corresponding codegen backend according to the target.

const PackedFunc* f = runtime::Registry::Get("tvm.codegen.llvm.target_" + target);

handle = (*f)()

return unique_ptr<CodeGenLLVM>(handle);

vector<PrimFunc> funcs;

for kv : mod->functions:

...

f = Downcast<PrimFunc>(kv.second);

if (f->HasNonzeroAttr(tir::attr::kIsEntryFunc))

entry_func = global_symbol.value();

funcs.push_back(f);

cg->Init("TVMMod", ...);

CodeGenCPU::Init() // codegen_cpu.cc

CodeGenLLVM::Init() // codegen_llvm.cc

for f in funcs:

cg->AddFunction(f); // codegen_cpu.cc

CodeGenLLVM::AddFunction();

AddFunctionInternal(f);

llvm::FunctionType* ftype = llvm::FunctionType::get(...);

// kGlobalSymbol: "global_symbol"

global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);

function_ = llvm::Function::Create(...);

llvm::BasicBlock* entry = llvm::BasicBlock::Create(..., function_);

IRBuilder::SetInsertPoint(entry);

this->VisitStmt(f->body);

builder_->CreateRet(ConstInt32(0));

if entry_func.length() != 0:

cg->AddMainFunction(entry_func); // codegen_cpu.cc

// tvm_module_main : "__tvm_main__"

llvm::GlobalVariable* global = new llvm::GlobalVariable(*module_, ..., tvm_module_main);

global->setInitializer(llvm::ConstantDataArray::getString(*ctx_, entry_func_name))

global->setDLLStorageClass(llvm::GlobalVariable::DLLExportStorageClass);

module_ = cg->Finish(); // CodeGenCPU::Finish() in codegen_cpu.cc

CodeGenLLVM::Finish(); // codegen_llvm.cc

CodeGenCPU::AddStartupFunction();

function_ = llvm::Function::Create(ftype, llvm::Function::InternalLinkage,"__tvm_module_startup", module_.get());

llvm::BasicBlock* startup_entry = llvm::BasicBlock::Create(*ctx_, "entry", function_);

llvm::appendToGlobalCtors(*module_, function_, 65535);

builder_->CreateRet(nullptr);

CodeGenLLVM::Optimize(); // codegen_llvm.cc

// Function pass manager

FPassManager fpass(module_.get());

// Module pass manager

MPassManager mpass;

mpass.add(llvm::createTargetTransformInfoWrapperPass(getTargetIRAnalysis()));

fpass.add(llvm::createTargetTransformInfoWrapperPass(getTargetIRAnalysis()));

llvm::PassManagerBuilder builder;

builder.Inliner = llvm::createFunctionInliningPass(builder.OptLevel, ...);

builder.LoopVectorize = true;

builder.SLPVectorize = true;

...

// Run the function passes

for mod in module_:

fpass.run(mod);

fpass.doFinalization();

// Run the module passes.

mpass.run(*module_);

return runtime::Module(n);

});

该函数中先创建LLVMModuleNode对象，然后调用Init()函数进行初始化，最后封装成Module对象返回。其中的Init()函数主要是将生成的TIR转为LLVM IR。主要分几步：

InitializeLLVM()函数初始化LLVM环境。这里边主要是例行调用LLVM的一大堆初始化函数。

创建用于代码生成的CodeGenLLVM对象。这里由于target字符串为x86-64，因此工厂函数名为tvm.codegen.llvm.target_x86-64。该工厂函数中创建CodeGenX86_64对象。因为继承关系为CodeGenX86_64 -> CodeGenCPU -> CodeGenLLVM，所以返回的是CodeGenLLVM的指针。

类型为IRModule的参数mod中的functions成员包含了该模块中的函数。这一步中将这些函数存于类型PrimFunc的数组funcs中。对于标为入口函数（kIsEntryFunc）的函数，记录在entry_func变量中。

接下来初始化前面创建的CodeGenX86_64对象。先调用CodeGenCPU::Init()，里边又会调用到CodeGenLLVM::Init()。前者主要创建一堆TVM运行时类型与函数。后者创建一些llvm中用于codegen的对象，如IRBuilder、llvm::Module和llvm::MDBuilder。

对前面放入funcs数组的每个函数，调用CodeGenCPU::AddFunction()函数，进行代码生成。对本文涉及的case只有一个函数就是vecadd()。

首先产生llvm::Function和llvm::BasicBlock对象，分别对应函数与基本块。前面在loewr()函数中将函数的名为global_symbol的属性设为相应的函数名（如vecadd）。这里将该属性取出，作为生成函数的链接时的symbol。

通过VisitStmt()函数遍历IRModule中的各节点并转为LLVM中对应的数据结构，生成LLVM IR。这是最关键的一步了。前面构建起的TIR主要就是为了这里的转换。例如，对于ForNode就会调用CodeGenLLVM::VisitStmt_(ForNode *op)函数。会调用CreateSerialFor()函数，产生相应的LLVM IR。在优化pass中的MakePackedAPI（make_packed_api.cc）会添加一个AttrStmt，对应一个值为目标函数名加_compute_后缀的compute_scope。这样，在code generation时，CodeGenCPU::CreateComputeScope()函数（为什么加compute_scope在该函数的注释中有提到）调用。

因此，最终的binary（可通过fadd.export_library("vecadd.so")语句导出）中大概会是这个样子：

AddMainFunction()函数设置主函数。如上面的例子中只有一个函数vecadd()，主函数。这个symbol会放在runtime::symbol::tvm_module_main（即__tvm_main__)这个全局变量中。可以拿编译好binary验证这一点。用objdump命令dump导出的so文件，可以看到如下这段。如果将里边的0x766563616464的16进制转为ASCII，就是主函数的symbol名：vecadd。

0000000000003c87 <__tvm_main__>:

3c87: 76 65 jbe 3cee <__GNU_EH_FRAME_HDR+0x5e>

3c89: 63 61 64 movslq 0x64(%rcx),%esp

3c8c: 64 fs

最后，调用CodeGenCPU::Finish()函数将LLVM IR生成后端代码。实际调用CodeGenLLVM::Finish()函数，会调用CodeGenLLVM::Finish()函数。主要调用CodeGenCPU::AddStartupFunction()函数和CodeGenLLVM::Optimize()函数。前者创建_tvm_module_startup函数，然后将一些需要启动时调用的函数填入。后者主要利用LLVM pass做一些优化。主要是向量化和函数内联。llvm中两种自动向量化。具体可参见Auto-Vectorization in LLVM。

其实，到这里编译还没有完全结束，只是构建好了LLVM的module。剩下的事情就是交给LLVM来编译生成可执行的binary了。真正生成可执行的binary是在第一次运行时通过LazyInitJIT()函数完成。运行时会调用到LLVMModuleNode::GetFunction()函数。当发现还未生成可执行binary时，会调用LazyInitJIT()函数。该函数通过llvm::ExecutionEngine将前面产生的llvm::Module编译成真正的（能在机器上跑的）binary。然后GetFunctionAddr()函数从中获得相应的函数指针，用于执行。

参考链接：

https://blog.csdn.net/jinzhuojun/article/details/117135551

https://blog.csdn.net/jinzhuojun/article/details/119696091

https://releases.llvm.org/12.0.0/docs/Vectorizers.html#the-slp-vectorizer

posted @ 2021-11-07 07:19 吴建明wujianming 阅读(1352) 评论(0) 编辑收藏举报

刷新页面返回顶部

吴建明

TVM基础编程示例分析

二．TVM调用llvm编译

公告