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。
對於上面的例子，構建的數據結構大體如下：
相關主要類簡圖：
圖中也可以看到，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<ScheduleNode>();
data_ = n;
n->outputs = ops;
auto g = te::CreateReadGraph(n->outputs); # graph.cc
Array<Operation> 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對象。構造函數中幾個主要步驟：
- 創建相應的ScheduleNode對象，將參數中傳入的Operation數組，設置到成員outputs中。對於上面的例子，Schedule()函數傳入的參數中Operation數組的size為1，即ComputeOp。
- CreateReadGraph()函數返回ReadGraph對象，包含了輸出依賴的所有操作及對應的張量。實質是一個Operation到該Operation的輸入tensor的數組Array<Tensor>的映射。構建過程主要是以輸入節點為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類中。其中主要的幾個相關成員：
l relations（類型Array<IterVarRelation>）：如這里創建的SplitNode繼承自IterVarRelationNode，幾個成員（parent, outer, inner, factor, nparts）描述了split的參數及前后計算軸變量。
l all_vars（類型為Array<IterVar>）：所有的循環變量。包括split過程中所有新老循環變量。
l leaf_vars（類型為Array<IterVar>）：當前生效的循環變量。如在這個例子中，只有經過split后的兩個循環變量。
經過split過后，循環變量關系通過TEDD可視化如下：
主要工作在SplitHelper()函數中完成。主要步驟：
- 原循環變量（用IterVar表示）按照給定因子，經過切分成為兩個，分別為外循環和內循環兩個。如示例中，外循環范圍為[0,16)，內循環范圍范圍為[0,64)。
- 通過FindLeafVar()函數找到父循環變量（即split前）在leaf_vars數組中的位置，一會split后的新循環變量會插在這個位置。
- 創建Split對象並存入成員relations中，對應SplitNode類。保存了使用了何種調度原語（這里是split），以及應用調度原語前后的循環變量間的關系。
- 更新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<IterVarAttrNode> n = make_object<IterVarAttrNode>();
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<IterVar>({*p_x_outer, *p_y_outer, *p_x_inner, *p_y_inner}));
return Array<IterVar>({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文件中。該函數基於給定參數構建出可調用的目標函數。按照官方介紹里的說法，主要做兩個工作 ：
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