TVM算子自动生成：TE与TIR Script联合使用教程
项目地址: https://gitcode.com/gh_mirrors/tvm/tvm 引言
你是否在深度学习模型部署过程中遇到过算子性能优化的难题？手动编写高效算子代码不仅耗时费力，还难以适配不同硬件架构。TVM（Tensor Virtual Machine）作为一款开源深度学习编译器栈，提供了强大的算子自动生成能力，能够帮助开发者轻松应对这一挑战。本文将详细介绍如何结合TE（Tensor Expression）与TIR（Tensor Intermediate Representation）Script，实现高效算子的自动生成。读完本文，你将掌握：
- TE与TIR Script的核心概念及协同工作原理
- 使用TE定义算子计算逻辑的方法
- 通过TIR Script进行底层优化的技巧
- 完整的算子开发流程与最佳实践
TE与TIR Script概述
TE（Tensor Expression）
TE是TVM中用于定义张量计算的高级接口，它允许开发者以声明式的方式描述算子的计算逻辑，而无需关心具体的硬件实现细节。TE提供了丰富的张量操作原语，如compute、reduce、scan等，能够轻松构建复杂的算子表达式。
TIR（Tensor Intermediate Representation）Script
TIR是TVM的中间表示形式，它是一种类C的领域特定语言，用于精确描述算子的底层实现。TIR Script允许开发者直接操作循环、内存布局、数据类型等底层细节，实现算子的深度优化。
TE与TIR Script的协同工作流程
TE与TIR Script并非相互替代，而是相辅相成。典型的协同工作流程如下：
TE基础：定义张量计算
张量与计算表达式
在TE中，张量（Tensor）是最基本的数据结构。通过te.compute函数，我们可以定义张量之间的计算关系。以下是一个简单的矩阵加法示例：
import tvm
from tvm import te
# 定义输入张量形状
n = te.var("n")
m = te.var("m")
A = te.placeholder((n, m), name="A")
B = te.placeholder((n, m), name="B")
# 定义计算表达式
C = te.compute((n, m), lambda i, j: A[i, j] + B[i, j], name="C")
在上述代码中，te.placeholder用于定义输入张量，te.compute则定义了输出张量C与输入张量A、B之间的计算关系。lambda函数lambda i, j: A[i, j] + B[i, j]描述了计算逻辑，其中i和j是循环变量。
循环嵌套与迭代变量
TE会自动根据计算表达式生成相应的循环嵌套。对于二维张量的计算，将生成两层嵌套循环。我们可以通过te.schedule对象获取并优化这些循环：
# 创建调度对象
s = te.create_schedule(C.op)
# 获取计算的循环嵌套
i, j = s[C].op.axis
# 对循环进行拆分
j0, j1 = s[C].split(j, factor=8)
# 重新排序循环
s[C].reorder(i, j0, j1)
# 打印TIR代码
print(tvm.lower(s, [A, B, C], simple_mode=True))
上述代码中，我们对j轴进行了拆分，并重新排序了循环，以提高数据局部性。tvm.lower函数可以将TE计算表达式转换为TIR代码，便于我们查看和进一步优化。
TIR Script进阶：底层优化
TIR基础语法
TIR Script提供了精确控制算子实现的能力。以下是一个简单的TIR函数示例，实现了向量加法：
from tvm.script import tir as T
@T.prim_func
def add(a: T.handle, b: T.handle, c: T.handle) -> None:
# 绑定缓冲区
A = T.match_buffer(a, (1024,), dtype="float32")
B = T.match_buffer(b, (1024,), dtype="float32")
C = T.match_buffer(c, (1024,), dtype="float32")
# 定义循环
for i in T.serial(1024):
with T.block("add"):
vi = T.axis.spatial(1024, i)
C[vi] = A[vi] + B[vi]
在TIR中，T.match_buffer用于绑定内存缓冲区，T.serial定义了串行循环，T.block则标记了计算块。
TE与TIR Script的联合使用
TE生成的初始TIR代码可以通过TIR Script进行进一步优化。以下是一个结合TE和TIR Script的完整示例，实现一个优化的矩阵乘法算子：
import tvm
from tvm import te, tir
from tvm.script import tir as T
# 使用TE定义矩阵乘法
def te_matmul(n, m, l):
A = te.placeholder((n, l), name="A")
B = te.placeholder((l, m), name="B")
k = te.reduce_axis((0, l), name="k")
C = te.compute((n, m), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="C")
s = te.create_schedule(C.op)
# TE级别的优化
i, j = s[C].op.axis
k = s[C].op.reduce_axis[0]
# 循环分块
i0, i1 = s[C].split(i, factor=32)
j0, j1 = s[C].split(j, factor=32)
k0, k1 = s[C].split(k, factor=4)
s[C].reorder(i0, j0, k0, i1, k1, j1)
return s, [A, B, C]
# 生成TE计算并转换为TIR
s, args = te_matmul(1024, 1024, 1024)
mod = tvm.lower(s, args, simple_mode=False)
# 使用TIR Script进行进一步优化
@tvm.script.ir_module
class OptimizedMatmul:
@T.prim_func
def main(A: T.handle, B: T.handle, C: T.handle) -> None:
T.func_attr({"global_symbol": "main", "tir.noalias": True})
A = T.match_buffer(A, (1024, 1024), dtype="float32")
B = T.match_buffer(B, (1024, 1024), dtype="float32")
C = T.match_buffer(C, (1024, 1024), dtype="float32")
# 定义共享内存
A_shared = T.alloc_buffer((32, 32), dtype="float32", scope="shared")
B_shared = T.alloc_buffer((32, 32), dtype="float32", scope="shared")
for i0, j0 in T.grid(32, 32):
for k0 in T.serial(32):
with T.block("load_A"):
vi0, vk0 = T.axis.remap("SS", [i0, k0])
A_shared[vi0, vk0] = A[vi0 * 32, vk0 * 4]
with T.block("load_B"):
vk0, vj0 = T.axis.remap("SS", [k0, j0])
B_shared[vk0, vj0] = B[vk0 * 4, vj0 * 32]
for i1, j1, k1 in T.grid(32, 32, 4):
with T.block("compute"):
vi, vj, vk = T.axis.remap("SSR", [i0 * 32 + i1, j0 * 32 + j1, k0 * 4 + k1])
with T.init():
C[vi, vj] = T.float32(0)
C[vi, vj] = C[vi, vj] + A_shared[i1, k1] * B_shared[k1, j1]
# 替换原函数
mod["main"] = OptimizedMatmul["main"]
# 编译优化后的算子
target = "cuda"
lib = tvm.build(mod, target=target)
在上述示例中，我们首先使用TE定义了矩阵乘法的基本计算逻辑，并进行了初步的循环分块优化。然后，我们将TE生成的TIR代码转换为TIR Script表示，并添加了共享内存优化，进一步提升了算子性能。
完整算子开发流程
步骤1：问题分析与需求定义
在开发新算子之前，首先需要明确算子的功能需求和性能目标。例如，我们需要开发一个适用于移动端CPU的高效卷积算子，要求支持3x3卷积核和任意输入通道数。
步骤2：使用TE定义计算逻辑
根据需求，使用TE定义算子的计算逻辑。以下是一个简化的卷积算子示例：
def conv2d(n, h, w, c, k, r, s, padding, stride):
# 计算输出形状
oh = (h + 2 * padding - r) // stride + 1
ow = (w + 2 * padding - s) // stride + 1
# 定义输入和权重张量
data = te.placeholder((n, c, h, w), name="data")
weight = te.placeholder((k, c, r, s), name="weight")
# 定义卷积计算
di = te.reduce_axis((0, c), name="di")
dr = te.reduce_axis((0, r), name="dr")
ds = te.reduce_axis((0, s), name="ds")
# 计算输出张量
output = te.compute(
(n, k, oh, ow),
lambda b, f, i, j: te.sum(
data[b, di, i*stride + dr - padding, j*stride + ds - padding] *
weight[f, di, dr, ds],
axis=[di, dr, ds]
),
name="output"
)
return te.create_schedule(output.op), [data, weight, output]
步骤3：TE级优化
对TE生成的计算进行初步优化，如循环分块、重排序等：
s, args = conv2d(1, 224, 224, 3, 64, 3, 3, 1, 1)
output = args[-1]
# 获取输出轴
b, f, i, j = s[output].op.axis
di, dr, ds = s[output].op.reduce_axis
# 循环分块优化
f0, f1 = s[output].split(f, factor=16)
i0, i1 = s[output].split(i, factor=8)
j0, j1 = s[output].split(j, factor=8)
s[output].reorder(b, f0, i0, j0, f1, i1, j1, di, dr, ds)
步骤4：TIR级优化
将TE生成的TIR代码转换为TIR Script，并进行底层优化，如数据布局调整、向量化等：
mod = tvm.lower(s, args, simple_mode=False)
# 使用TIR Script进行向量化优化
@tvm.script.ir_module
class OptimizedConv2d:
@T.prim_func
def main(data: T.handle, weight: T.handle, output: T.handle) -> None:
# 函数属性和缓冲区定义
T.func_attr({"global_symbol": "main", "tir.noalias": True})
data = T.match_buffer(data, (1, 3, 224, 224), dtype="float32")
weight = T.match_buffer(weight, (64, 3, 3, 3), dtype="float32")
output = T.match_buffer(output, (1, 64, 224, 224), dtype="float32")
# 循环定义和向量化
for b in T.serial(1):
for f0 in T.serial(4):
for i0 in T.serial(28):
for j0 in T.serial(28):
for f1 in T.serial(16):
with T.block("compute"):
# 轴定义和向量化标记
vb, vf0, vi0, vj0, vf1 = T.axis.remap("SSSSS", [b, f0, i0, j0, f1])
T.reads(data[b, :, i0*8:(i0+1)*8+2, j0*8:(j0+1)*8+2],
weight[vf0*16+vf1, :, :, :])
T.writes(output[b, vf0*16+vf1, i0*8:(i0+1)*8, j0*8:(j0+1)*8])
# 向量化计算
output[b, vf0*16+vf1, i0*8:(i0+1)*8, j0*8:(j0+1)*8] = T.vectorized(
T.sum(data[b, di, i0*8 + dr : (i0+1)*8 + dr, j0*8 + ds : (j0+1)*8 + ds] *
weight[vf0*16+vf1, di, dr, ds],
axis=[di, dr, ds])
)
mod["main"] = OptimizedConv2d["main"]
步骤5：编译与验证
编译优化后的算子，并进行正确性和性能验证：
# 编译算子
target = "llvm -mcpu=armv8.2-a"
lib = tvm.build(mod, target=target)
# 生成测试数据
import numpy as np
data_np = np.random.uniform(size=(1, 3, 224, 224)).astype("float32")
weight_np = np.random.uniform(size=(64, 3, 3, 3)).astype("float32")
output_np = np.zeros((1, 64, 224, 224), dtype="float32")
# 运行算子
dev = tvm.cpu()
data_tvm = tvm.nd.array(data_np, device=dev)
weight_tvm = tvm.nd.array(weight_np, device=dev)
output_tvm = tvm.nd.array(output_np, device=dev)
lib(data_tvm, weight_tvm, output_tvm)
# 验证结果
def numpy_conv2d(data, weight, padding=1, stride=1):
# 简单的numpy实现，用于验证
n, c, h, w = data.shape
k, _, r, s = weight.shape
oh = (h + 2*padding - r) // stride + 1
ow = (w + 2*padding - s) // stride + 1
output = np.zeros((n, k, oh, ow))
data_pad = np.pad(data, ((0,0), (0,0), (padding,padding), (padding,padding)), mode='constant')
for b in range(n):
for f in range(k):
for i in range(oh):
for j in range(ow):
for di in range(c):
for dr in range(r):
for ds in range(s):
output[b, f, i, j] += data_pad[b, di, i*stride+dr, j*stride+ds] * weight[f, di, dr, ds]
return output
output_np = numpy_conv2d(data_np, weight_np)
np.testing.assert_allclose(output_tvm.numpy(), output_np, rtol=1e-4)
# 性能测试
evaluator = lib.time_evaluator("main", dev, number=10)
print(f"Conv2d performance: {evaluator(data_tvm, weight_tvm, output_tvm).mean * 1e3:.2f} ms")
TE与TIR Script联合使用最佳实践
1. 合理划分优化层级
TE适合进行高层的计算逻辑定义和循环分块，而TIR Script则更适合底层的内存优化和指令级优化。合理划分优化层级可以提高开发效率。
2. 充分利用TVM内置调度原语
TVM提供了丰富的调度原语，如split、reorder、tile、vectorize等，应优先使用这些原语进行优化，只有在需要特殊优化时才手动编写TIR Script。
3. 利用自动调优工具
TVM的AutoTVM和AutoScheduler工具可以自动搜索最优的调度参数，结合TE使用可以大幅提升算子性能，减少手动优化的工作量。
4. 注意数据局部性
在进行循环优化时，应尽量提高数据的空间局部性和时间局部性，如通过分块将数据缓存在CPU缓存中，减少内存访问延迟。
5. 验证与性能分析相结合
优化过程中应不断进行正确性验证和性能分析，使用TVM提供的profile工具定位性能瓶颈，有针对性地进行优化。
结论
TE与TIR Script的联合使用为算子开发提供了强大的工具链，既可以通过TE快速定义计算逻辑，又可以通过TIR Script进行深度优化。本文详细介绍了TE和TIR Script的基本用法、协同工作流程以及完整的算子开发流程，并总结了最佳实践。通过掌握这些技术，开发者可以高效地开发出适配各种硬件平台的高性能算子。
未来，随着TVM生态的不断完善，TE和TIR Script的功能将更加强大，算子开发流程也将更加自动化和智能化。建议开发者持续关注TVM社区的最新进展，及时掌握新的优化技术和工具。
最后，鼓励大家积极参与TVM社区贡献，分享自己的算子开发经验和优化技巧，共同推动深度学习编译器技术的发展。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
