TVM addresses these challenges with three key modules.
(1) We introduce a tensor expression languageto build operators and provide program transformation primitives that generate different versions of the program with various optimizations.
tensor expression 是 tvm 引入的一门面向张量计算的函数式编程语言。前面我们知道 topi 就是基于 te 实现了一系列常见的张量算子库
在 tvm 里面，每个算子都可以表达为一个 te 表达式，te 接受张量作为输入，并返回一个新的张量
这个语言本身是非常简单的：
- Type：因为处理的输入输出，都是 tensor，所以只有一种数据类型（不像 relay ir，有 Var、Constant、Function、Call 等各种类型）
- Op：操作符也很简单，最常见的就2个 placeholder_op，compute_op
te 的代码主要在 include/tvm/te/ src/te/ 这2个目录，全部加起来才1.3w行，所以说是并不复杂
这篇文章，我们从编译器的视角来理解一下 tensor expression
参考资料：
1. te 的设计目标
te 可能应该有其他更多的设计目标
但从我目前了解到的来看，te 最大的作用应该是简化tensor 算子的编写难度，并支持进一步的调度和优化
比如向量加法
1）TVMScript 版本
在没有 te 之前，比如只有 tir 的时候，用户只能用 TVM script 来编写tensor 算子。如下：
@tvm.script.ir_module
class MyModule:
@T.prim_func
def main(a: T.handle, b: T.handle):
# 我们通过 T.handle 进行数据交换，类似于内存指针
T.func_attr({"global_symbol": "main", "tir.noalias": True})
# 通过 handle 创建 Buffer
A = T.match_buffer(a, (8,), dtype="float32")
B = T.match_buffer(b, (8,), dtype="float32")
for i in range(8):
# block 是针对计算的抽象
with T.block("B"):
# 定义一个空间（可并行）block 迭代器，并且将它的值绑定成 i
vi = T.axis.spatial(8, i)
B[vi] = A[vi] + 1.0
ir_module = MyModule
2）tensor expression 版本
但是如果用 te 来描述向量加法的话，只需要2行，并且通过script() 打印出来会发现和 TVM script 是一样的
from tvm import te
A = te.placeholder((8,), dtype="float32", name="A")
B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B")
func = te.create_prim_func([A, B])
ir_module_from_te = IRModule({"main": func})
print(ir_module_from_te.script())
2. te 表达式示例
目前我并没有在官方文档上，找到 tensor expression 的语法描述。只能通过官方文档，以及代码里的测试例子，还有网上收集的一些例子，来推测 te 的一些基本用法
2.1. 向量加法
比如使用 tensor expression 实现向量加法，向量是一维的数组
用 np.arange 函数可以生成一维数组
A = te.placeholder((8,), dtype="float32", name="A")
B = te.placeholder((8,), dtype="float32", name="A")
C = te.compute((8,), lambda i: A[i] + B[i], name="B")
s = te.create_schedule(C.op)
func = tvm.build(s, [A, B, C], target="llvm")
a = tvm.nd.array(np.arange(8).astype("float32"))
b = tvm.nd.array(np.arange(8).astype("float32"))
c = tvm.nd.array(np.zeros((8,)).astype("float32"))
func(a, b, c)
print(a, b, c)
placeholder 是一个占位符，一般用来声明一个输入变量
调用 create_prim_func 创建 func 的时候，必须把ABC都作为输入，也就是说，函数的返回值也是函数的一个参数
所以后面在调用这个算子的时候，返回值参数需要事先创建好。
2.2. 矩阵加法
用 np.random.randint(0, 10, (n, m)) 可以随机生成一个10以内的二维矩阵
我们用 tensor expression 实现2个矩阵的加法
A = te.placeholder((8,4,), dtype="float32", name="A")
B = te.placeholder((8,4,), dtype="float32", name="B")
C = te.compute((8,4,), lambda i,j: A[i,j] + B[i,j], name="C")
s = te.create_schedule(C.op)
func = tvm.build(s, [A, B, C], target="llvm")
a = tvm.nd.array(np.random.randint(0, 10, (8,4)).astype("float32"))
b = tvm.nd.array(np.random.randint(0, 10, (8,4)).astype("float32"))
c = tvm.nd.array(np.zeros((8,4)).astype("float32"))
func(a, b, c)
2.3. GEMM 矩阵乘法
大学线性代数的内容，行和列的乘积
用 tensor expression 代码表示如下：
import numpy as np
import tvm
from tvm import te
n = te.var("n")
k = te.var("k")
m = te.var("m")
A = te.placeholder((n, k), dtype="float32")
B = te.placeholder((k, m), dtype="float32")
l = te.reduce_axis((0, k), name='l')
C = te.compute((n, m),
lambda i, j: te.sum(A[i, l] * B[l, j], axis = l),)
func = tvm.build(te.create_schedule(C.op), [A, B, C], target="llvm")
a = tvm.nd.array(np.random.randint(0, 3, (2, 4)).astype("float32"))
b = tvm.nd.array(np.random.randint(0, 3, (4, 2)).astype("float32"))
c = tvm.nd.array(np.zeros((2, 2)).astype("float32"))
func(a, b, c)
print(a)
print(b)
print(c)
这里面涉及到一个新的 reduce_axis 函数，中文叫规约，作用是用来减少轴，比如把一个2维的矩阵转换成一个1维的向量，就叫规约
3. 计算 & 调度分离
我们看到上面 tensor expression 的例子里，都有一个 te.create_schedule() 的过程，这是在给当前的计算创建一个默认的调度逻辑（当涉及到更深度的优化时候，可以自行实现更复杂的调度逻辑）
计算 & 调度，是 te 或者 tvm 里一个非常重要的2个概念，而计算和调度分离，是 tvm 里优化性能的一个重要手段
我举个最简单的例子：矩阵加法
1） 什么是计算？
计算：描述算法本身
首先你不要把下面的代码想象成具体的实现，它只是一段伪代码，一段描述算法的伪代码
在这段伪代码里，它所表达的算法，是把两个矩阵里的每个相同位置的元素相加，其值作为结果向量里对应位置的值
for (i = 0; i < 8; i++) {
c[i] = a[i] + b[i]
}
2）什么是调度？
调度：决定计算如何在硬件上执行
到这里就很关键了，调度决定计算如何在具体的硬件上执行，调度的目标是要结合硬件特性实现执行效率的优化
比如分块执行，假设硬件具备并行计算的能力，我可以把上述代码，分成2块并行执行
调度过程就是要描述怎么把这个计算过程分成多少块，块如何并行的在硬件上执行
for (i = 0; i < 4; i++) {
c[i] = a[i] + b[i]
}
for (i = 4; i < 8; i++) {
c[i] = a[i] + b[i]
}
除了分块，还有更多调度的优化技术
4. te 的一些核心概念
Tensor 和 Operation 是 tensor expression 最主要的2个数据结构，后者是最关键的
4.1. Tensor
tensor 是 te 里面唯一的一种数据类型，不管是变量，还是一维向量，还是二维矩阵，还是三维xx，统统用 tensor 来标识
4.2. Operation
xx
4.2.1. PlaceholderOp
xx
4.2.2. ComputeOp
xx
4.3. Stmt
Stmt 可以包含多种类型的语句，如赋值语句、控制流语句（如 if-else 和循环语句）、函数调用语句等
Stmt 的各种具体定义，可以在 include/tvm/tir/stmt.h 里找到
4.4. Block & Stage
block 是一个抽象的概念，可以理解为是一组 stmt 的集合
block 基本块的引入可以实现更细粒度和更高级别的优化。例如，可以针对基本块内的操作进行循环展开、指令重排、寄存器优化等操作，以提高计算的效率。基本块的划分和优化是基于计算图的特性和性能需求进行的，为了更好地利用硬件资源和降低执行开销
一个 stage 包含多个计算步骤（Operation），stage的主要作用是提供一种粗粒度的优化和调度，它可以将多个计算步骤组合成一个较小的实体，从而更好地进行优化和调度。
5. te 的编译过程
tvm 中有2个地方会把 te 编译成 tir PrimFunc
一个就是前面我们讲 topi 的时候了解到，tvm 在编译 relay ir 的时候，如果发现这个算子是 topi 的算子，会直接调用 topi 算子的实现，而 topi 的算子实现里面，基本上都会调用 tvm::te::compute() 函数，把 te 编译成 tir 的 PrimFunc
另一个就是用户直接编写的 te 代码，调用 te.create_prim_func() 函数把 te 编译成 tir
5.1. tvm::te::compute()
上面的代码，te.compute 其实是 tvm::te::compute() 的一个 Python 前端
te::compute 主要在2个地方会被调用到：
- 处理 topi 的时候
- python 显式调用 te.compute()
topi 会把所有算子实现，通过 TVM_REGISTER_GLOBAL 注册到全局的 ffi 里面
relay ir 在编译的时候，如果发现是 topi 算子，转而调用对应注册的算子计算函数，这个地方在 tvm::relay::tec::TECompilerImpl::Lower
Lower 会调用 PrimFuncFor 寻找上面注册好的 topi 算子实现函数
而 topi 的算子实现，最终都会调用 tvm::te::compute()，这里就实现了 te -> tir 的编译，得到 tir 的 PrimFunc
需要注意的是，tvm::te::compute() 返回的仍然还是 tensor 结构，fcompute(args) 返回的是一个 PrimFunc
Tensor compute(Array shape, FCompute fcompute, std::string name, std::string tag,
Map<String, ObjectRef> attrs) {
// compute dimension.
size_t ndim = shape.size();
std::vector axis;
std::vector args;
for (size_t i = 0; i < ndim; ++i) {
std::ostringstream os;
os << "ax" << i; axis.emplace_back(IterVar(Range(0, shape[i]), Var(os.str(), shape[i].dtype()), kDataPar)); args.push_back(axis.back()->var);
}
return ComputeOp(name, tag, attrs, axis, {fcompute(args)}).output(0);
}
5.2. te.create_prim_func
如果我们直接使用 te 来编写算子
可以通过 te.create_prim_func 来讲 te 编译成 PrimFunc，这是一个 python 前端，对应 c++ 的 CreatePrimFunc 函数，我们直接看 c++ 实现好了
CreatePrimFunc 函数分位4个阶段：
- CollectOrderedOps，后序遍历（递归），得到 AST graph
- InitializeBufferBinds 这里好像是初始化 ExternOp 的什么东西
- 得到 AST graph 之后，对每个 op 调用 RewriteStageToBlock 转换成 Stmt
- 最后封装成一个 PrimFunc
PrimFunc CreatePrimFunc(const Array& arg_list) {
// Infomations used in CreatePrimFunc and its sub-functions.
CreateFuncInfo info(arg_list);
// Root body stmts.
Array root_stmts;
// Analyzer
arith::Analyzer analyzer;
// Step 1. Create ordered array of operations and validate they are supported.
Array order = CollectOrderedOps(arg_list);
// Step 2. Initialize buffer binds map
InitializeBufferBinds(order, &info);
// Step 3. Rewrite compute stages into blocks.
for (const te::Operation& op : order) {
RewriteStageToBlock(op, &info, &root_stmts, &analyzer);
}
// Step 4. Create func and complete prim func.
return GenerateAndCompletePrimFunc(arg_list, root_stmts, &info);
}
5.2.1. post DFS 遍历 AST graph
这里是在 CollectOrderedOps 函数完成的
Array CollectOrderedOps(const Array& arg_list) {
Array arg_ops;
for (const te::Tensor& arg : arg_list) {
arg_ops.push_back(arg->op);
}
te::ReadGraph g = te::CreateReadGraph(arg_ops);
Array order = te::PostDFSOrder(arg_ops, g);
return order;
}
te::ReadGraph 是一个 type 别名，真实的类型是 using ReadGraph = Map<Operation, Array>;
CreateReadGraph 函数通过递归遍历，建立 Operation 到输入 InputTensors 的映射关系。
// construct a read graph that gives readers of each operation
// that the root depend on
ReadGraph CreateReadGraph(const Array& roots) {
ReadGraph rmap;
std::vector stack;
std::unordered_set visited;
// initialize the roots
for (Operation op : roots) {
stack.push_back(op);
visited.insert(op.get());
}
while (!stack.empty()) {
Operation op = stack.back();
stack.pop_back();
Array deps = op->InputTensors();
rmap.Set(op, deps);
for (Tensor t : deps) {
if (t->op.defined() && visited.count(t->op.get()) == 0) {
visited.insert(t->op.get());
stack.push_back(t->op);
}
}
}
return rmap;
}
te::PostDFSOrder 和上面的遍历其实有点类似。本质上是对 AST graph 做了一次后序遍历
void PostDFSOrder(const Operation& op, const ReadGraph& g, std::unordered_set* visited,
Array* post_order) {
if (visited->count(op)) return;
visited->insert(op);
for (const auto& t : g.at(op)) {
PostDFSOrder(t->op, g, visited, post_order);
}
post_order->push_back(op);
}
Array PostDFSOrder(const Array& roots, const ReadGraph& g) {
std::unordered_set visited;
Array post_order;
for (Operation op : roots) {
PostDFSOrder(op, g, &visited, &post_order);
}
return post_order;
}
5.2.2. Stage Rewriting
阶段重写（Stage Rewriting），就是把阶段（Stage）会根据一些策略被重写为基本块（block）的形式
RewriteStageToBlock 就是干这个的，在这里它把每个 Opration（及其依赖）译成 Stmt 计算语句
如下：
void RewriteStageToBlock(const te::Operation& op, CreateFuncInfo* info, Array* root_stmts,
arith::Analyzer* analyzer) {
if (const auto* placeholder = op.as()) {
// ...
if (info->tensor2buffers.count(tensor) == 0) {
const Buffer& buffer =
decl_buffer(placeholder->shape, placeholder->dtype, placeholder->name, "global");
info->tensor2buffers[tensor] = buffer;
}
} else if (const auto* compute_op = op.as()) {
// Case 2. ComputeOp (te.compute)
root_stmts->push_back(
GenerateStmtFromCompute(GetRef(compute_op), info, analyzer));
} else if (const auto extern_op = op.as()) {
// Case 3. ExternOp (te.extern)
root_stmts->push_back(GenerateStmtFromExternOp(GetRef(extern_op), info));
}
}
这里我们其实只需要关注 PlaceholderOp 和 ComputeOp 的处理即可
Placeholder 是一个占位符，属于计算图里面最叶子的节点了，它自己是没有输入的（InputTensors）
tensor2buffers 的作用是将每个张量映射到预先分配的存储缓冲区上。这样可以实现对张量数据进行内存分配和管理，方便后面的代码生成
对于 ComputeOp，调用 GenerateStmtFromCompute 生成一些列的 Stmt，然后组合成一个 Block
5.2.3. 生成 PrimFunc
xxx