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();
…
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::TryConstFoldtir::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。
   对于上面的例子，构建的数据结构大体如下：
   

相关主要类简图：

图中也可以看到，Python与C++层中的对象有对应关系。这便于Python与C++间的调用，这也是TVM的特色之一。一般名为XXX的是相应XXXNode的引用（如ComputeOp与ComputeOpNode）。前者继承自ObjectRef，后者继承自Object。主要的内容是在XXXNode中，XXX中的->操作符重载了，将操作及访问会应用到XXXNode上。
Operation代表操作，如PlaceholderOp和ComputeOp。Tensor代表张量，TensorSlice表示Tensor的切片，如例子中A[i]。PrimExpr主要用于low-level的表示，是所有primitive expression的基类。Primitive expression处理POD数据类型。这里表示计算的Add和包含了张量的ProducerLoad都是PrimExpr。
稍微复杂些的常见例子是矩阵乘matmul：
k = tvm.te.reduce_axis((0, l), name=‘k’)
A = tvm.te.placeholder((n, l), name=‘A’)
B = tvm.te.placeholder((l, m), name=‘B’)
C = tvm.te.compute((n, m), lambda x, y: tvm.te.sum(A[x, k] * B[k, y], axis=k), name=‘C’)
与上例有所区别的是这里操作数都是二维的，且有reduce轴（计算过程中约减，因此输入中有，输出中没有的轴）。计算中使用了tvm.te.sum()（实现在python/tvm/tir/op.py）函数来reduce中间轴。函数的定义为：
sum = comm_reducer(lambda x, y: x + y, lambda t: const(0, dtype=t), name=“sum”) # tir/op.py
tvm.te.sum(A[x, k] * B[k, y], axis=k)
tvm.tir.Reduce(…) # expr.py
return Reduce(…); # expr.cc
生成的数据结构与上面vecadd例子中是类似的，其中Add换成了Reduce。
构建schedule
TVM中继承了Halide中algorithm与schedule分离的思想。上面定义好了算什么，接下来就需要确定怎么算了。这就是schedule要定义的事。首先，需要创建一个schedule：
s = tvm.te.create_schedule(C.op)
其中C.op类型为te.tensor.ComputeOp，返回的变量s类型为te.schedule.Schedule。基本流程如下：
create_schedule(ops) # in schedule.py
return ffi_api.CreateSchedule(ops)
create_schedule(ops) // schedule.h
return Schedule(ops) // schedule_lang.cc
auto n = make_object();
data = n;
n->outputs = ops;
auto g = te::CreateReadGraph(n->outputs); # graph.cc
Array post_order = te::PostDFSOrder(n->outputs, g); // graph.cc
for op in post_order:
Stage stage(op);
n->stages.push_back(stage);
n->stage_map.Set(op, stage);
…
这里从Python调用到C++，主要作用是创建Schedule对象。构造函数中几个主要步骤：

1. 创建相应的ScheduleNode对象，将参数中传入的Operation数组，设置到成员outputs中。对于上面的例子，Schedule()函数传入的参数中Operation数组的size为1，即ComputeOp。
2. CreateReadGraph()函数返回ReadGraph对象，包含了输出依赖的所有操作及对应的张量。实质是一个Operation到该Operation的输入tensor的数组Array的映射。构建过程主要是以输入节点为root，然后通过Operation的InputTensors()函数，找出对应的输入tensor。上面例子就是：
   

调用PostDFSOrder()函数得到后序的Operation数组。对于该例子便是A, B, C。表示了各个Operation之间的依赖关系。
按照上面得到的后序数组，对每个Operation创建相应的Stage对象。Schedule对象包含一系列Stage。每个Stage对象对应一个Operation。如上面的例子，就有三个Stage。每个Stage保存了一个循环嵌套（Loop nest）结构的信息，及每个循环的类型（如parallel, vectorized, unrolled）等。
创建了Schedule及对应的Stage对象后，接下来就可以进行一些操作。对于该schedule，可以应用一些调度原语（Schedule primitive）。详细可见官方文档Schedule Primitives in TVM 。下面是一个很常用的split的简单例子：
outer, inner = s[C].split(C.op.axis[0], factor=64)
上面的语句中，s[C]从schedule中得到对应的Stage对象，类型为tvm.te.schedule.Stage。split()函数第一个参数和返回值的类型都是tir.expr.IterVar，对应相应的循环变量（或者说计算轴）。将操作C的计算中的轴，以64为因子进行分割，将一重循环分成二重循环。例如，如果原来的循环次数为1024，分割后就是外循环16次，内循环64次。大体流程如下：
Stage::split() // schedule.py
outer, inner = _ffi_api.StageSplitByFactor(…) // schedule_lang.cc
IterVar outer, inner;
Stage::split(parent, factor, &outer, &inner);
SplitHelper(opertor->(), parent, factor, PrimExpr(), p_outer, p_inner);
IterVar outer = IterVar(…);
IterVar inner = IterVar(…);

            size_t pos = FindLeafVar(...);self->relations.push_back(Split(parent, outer, inner, factor, nparts))auto n = make_object<SplitNode>();...data_ = std::move(n);all_vars.push_back(outer);all_vars.push_back(inner);leaf_vars.erase(leaf_vars.begin() + pos);leaf_vars.insert(leaf_vars.begin() + pos, inner);leaf_vars.insert(leaf_vars.begin() + pos, outer);return Array<IterVar>({outer, inner});
return outer, inner;

前面提到，循环结构表示在StageNode类中。其中主要的几个相关成员：
 relations（类型Array）：如这里创建的SplitNode继承自IterVarRelationNode，几个成员（parent, outer, inner, factor, nparts）描述了split的参数及前后计算轴变量。
 all_vars（类型为Array）：所有的循环变量。包括split过程中所有新老循环变量。
 leaf_vars（类型为Array）：当前生效的循环变量。如在这个例子中，只有经过split后的两个循环变量。
经过split过后，循环变量关系通过TEDD可视化如下：

主要工作在SplitHelper()函数中完成。主要步骤：

1. 原循环变量（用IterVar表示）按照给定因子，经过切分成为两个，分别为外循环和内循环两个。如示例中，外循环范围为[0,16)，内循环范围范围为[0,64)。
2. 通过FindLeafVar()函数找到父循环变量（即split前）在leaf_vars数组中的位置，一会split后的新循环变量会插在这个位置。
3. 创建Split对象并存入成员relations中，对应SplitNode类。保存了使用了何种调度原语（这里是split），以及应用调度原语前后的循环变量间的关系。
4. 更新all_vars与leaf_vars两个IterVar数组。前者表示所有的（即split前后）循环变量，后者表示split后循环变量，可以理解为目前生效的循环变量。添加新产生的循环变量到all_vars和leaf_vars中，同时删除leaf_vars中的原有循环变量。
   主要数据结构如下：
   

相关主要类简图：

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

经过split后，让外循环并行提高性能。可以用下面的调度原语：
s[C].parallel(outer)
调用大体流程如下：

Stage::paralle() // schedule.py
_ffi_api.StageParallel(self, var)
Stage::parallel() // schedule_lang.cc
SetAttrIterType(operator->(), var, kParallelized);
UpdateIterVarAttr(self, var, …);
ObjectPtr n = make_object();
n->iter_type = kParallelized;
self->iter_var_attrs.Set(var, IterVarAttr(n));
与上面类似，也是从Python层调用到C++层，完成实质的工作。只要设置循环变量属性就行，因此比较简单，函数UpdateIterVarAttr()中，主要就是创建相应的IterVarAttrNode对象，根据参数设置属性，最后保存到StageNode的iter_var_attrs成员中。

例如，对于常见的矩阵乘计算，通常会应用tile这个调度原语做tiling：
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], 32, 32)
对于两个计算轴做tiling，对每个轴都分成外循环与内循环，然后返回总共4个新的计算轴。大体流程如下：
Stage::tile() // schedule.py
x_outer, y_outer, x_inner, y_inner = _ffi_api.StageTile(…) // schedule_lange.cc
IterVar x_outer, y_outer, x_inner, y_inner;
stage.tile(x_parent, y_parent, x_factor, y_factor, &x_outer, &y_outer, &x_inner, &y_inner);
split(x_parent, x_factor, p_x_outer, p_x_inner);
split(y_parent, y_factor, p_x_outer, p_y_inner);
…
reorder(Array({*p_x_outer, *p_y_outer, *p_x_inner, *p_y_inner}));
return Array({x_outer, y_outer, x_inner, y_inner);
return x_outer, y_outer, x_inner, y_inner;
其实主要的工作就是在两个维度上做split，然后对切分后的循环变量，按指定顺序做reorder。
计算的定义与schedule的构建基本就完成了。下面一下编译部分。

二．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文件中。该函数基于给定参数构建出可调用的目标函数。按照官方介绍里的说法，主要做两个工作 ：
 Lowering：将high-level的循环嵌套结构，转换成最终的low-level的IR。
 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对象。几个关键步骤如下：

 Schedule::normalize()函数规范化给定的schedule。主要实现在schedule_dataflow_rewrite.cc文件中。调用以下三个函数。
 InjectInline()函数处理算子内联。用到调度原语 compute_inline的话会用到。
 RebaseNonZeroMinLoop()函数将循环迭代的最小界置为0。感觉有点canonicalization的意思。
 LegalizeInvalidAttach()函数处理在使用调度原语compute_at时且目标迭代又被split或fuse情况下的合法化。
 InferBound()函数顾名思义就是边界推导（Bound inference），主要用于推导循环边界。更具体地，就是确定每个IterVar的范围，返回IterVar到Range的映射，即每个循环变量的范围。这个信息在后面的MakeLoopNest()函数中，用于确定for循环的范围，在BuildRealize()函数中设置缓冲的大小。具体可参见官方文档 InferBound Pass。
 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调用相应的逻辑。
 对于上面的例子，Compute Op没有attach在其它计算中，因此对应Stage的attach type为kGroupRoot，因此这里调用MakePipeline()函数产生Stmt。这步比较关键比较复杂，后面再展开。
 然后通过SchedulePostProc对象（继承自StmtExprMutator），对前面生成的Stmt进行后处理。
 get_binds()函数用于绑定buffer。给每个参数张量分配buffer。如对于上面例子中的A, B, C三个张量，分别通过tvm.tir.decl_buffer()，创建buffer并绑定张量。
 SchedulePostProcToPrimFunc()函数基于ScheduleOps()产生的Stmt创建PrimFunc对象，可以用于TIR优化。PrimFunc代表包含了TIR statement的primitive function，是low-level的代码表示。
 创建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> 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 provides
…
// Array -> 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。表示每个循环的边界判断，调用MakeBoundCheck()函数来生成。
 init_nest与main_nest：类型为vector<vector>。 其中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.outer64) + i.inner)] = ((float32)A_2[((i.outer64) + 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 compilationmod_host, mdev = _build_for_device(input_mod, tar, target_host)# input_mod type: IRModulemod_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 functionsrt_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: Modulereturn 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.ccdevice_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.ccGetModuleNode(mod)->Import(...) # runtime/module.ccimports_.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();
n->Init(mod, target); // llvm_module.cc
InitializeLLVM();
llvm::InitializeAllTargetInfos();
llvm::InitializeAllTargets();
…
unique_ptr 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(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.ccCodeGenLLVM::Init() // codegen_llvm.ccfor f in funcs:cg->AddFunction(f); // codegen_cpu.ccCodeGenLLVM::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.ccCodeGenLLVM::Finish(); // codegen_llvm.ccCodeGenCPU::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 managerFPassManager fpass(module_.get());// Module pass managerMPassManager 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 passesfor 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