从零手动创建TensorIR：深入TVM底层编译原理与实践

1. 项目概述从TVM到TensorIR的跨越如果你正在深度学习模型部署的深水区里扑腾对PyTorch或TensorFlow导出的模型如何在各种硬件上跑得更快一筹莫展那你大概率已经听说过TVMTensor Virtual Machine这个名字。它不像一些框架那样“开箱即用”更像是一个给资深工程师的“编译器工具箱”让你能深入到计算图的最底层去扣性能。而TensorIR正是这个工具箱里近年来最锋利的一把新锉刀。它不是另一个前端API而是TVM中用于表示和变换中间表示IR的核心抽象。简单来说以前我们可能更关注“用什么算子”而TensorIR让我们能精细地控制“这个算子在硬件上具体怎么算”。这次我们不谈高屋建瓴的架构就扎扎实实地聊一个最基础、但也是所有高级优化的起点如何从零开始手动创建一段TensorIR。这听起来像是“从拧螺丝开始造火箭”但我的经验是不理解TensorIR的构造后续的循环变换、算子融合、自动调度AutoScheduler都像是空中楼阁。你会用te.compute写表达式但不知道为什么它最后变成了那样的循环嵌套你会调用schedule的方法但不清楚它底层到底操作了什么数据结构。通过亲手创建你能直观地理解TVM是如何将你的计算意图一步步分解、映射到可执行代码的。这对于调试性能瓶颈、实现自定义算子、乃至理解整个编译栈都至关重要。2. TensorIR核心概念与设计哲学解析在动手之前我们必须统一“语言”。TensorIR的设计哲学深深植根于现代深度学习编译的需求显式地表达数据并行与任务并行并为硬件优化提供充分的变换空间。2.1 TensorIR是什么超越计算图的抽象你可以把TensorIR理解为一种“可编程的计算模板”。常见的计算图IR如Relay主要描述“做什么”——算子之间的数据流关系。而TensorIR则深入到“怎么做”的层面它显式地定义了循环迭代空间计算在哪些维度上循环循环的边界是什么。数据读取模式在循环的每一步从哪里全局内存、共享内存、寄存器读取数据。计算语句在循环体内执行的具体算术或逻辑运算。内存作用域数据存放在哪个层次的内存中这直接关系到访存效率。这种显式性带来了巨大的灵活性。例如一个矩阵乘法在TensorIR中你可以选择用三重循环i, j, k的朴素实现也可以表达成分块tile、向量化vectorize、循环展开unroll等多种优化后的形态。TensorIR本身不规定哪种形态是最好的它只是为你提供了描述所有这些形态的能力并将选择权交给了调度器Scheduler或搜索算法。2.2 核心数据结构PrimFunc与Block在TVM的现代体系中一个可被编译的最小单元通常是一个PrimFunc原始函数。而一个PrimFunc的核心组成部分就是Block。PrimFunc 这是一个包含计算逻辑和缓冲区Buffer声明的函数。它可以直接对应一个GPU kernel或者CPU函数。在TensorIR中我们构建的最终就是一个或多个PrimFunc。Block 这是TensorIR中组织和承载计算的主体。一个Block可以理解为一个“计算块”它包含迭代器IterVar 定义了这个块的计算空间例如i,j,k。读写关系Read/Write Regions 声明了这个块会读取和写入哪些缓冲区Buffer的哪些区域。这是TVM进行依赖分析和算子融合的关键依据。计算体Body 包含具体的计算语句例如C[i, j] A[i, k] * B[k, j]。创建TensorIR的过程很大程度上就是在构建和组装这些Block。2.3 与Tensor Expression (TE) 的关系你可能更熟悉TVM的Tensor ExpressionTE接口。确实在大多数教程里我们这样写一个矩阵乘法import tvm from tvm import te A te.placeholder((m, n), name‘A’) B te.placeholder((n, k), name‘B’) k te.reduce_axis((0, n), name‘k’) C te.compute((m, k), lambda i, j: te.sum(A[i, k] * B[k, j], axisk), name‘C’)te.compute会为我们自动生成一个最基础的TensorIR通常是一个多重循环。手动创建TensorIR本质上是在做te.compute在背后自动完成的事情但给了我们百分百的控制权。当你需要实现一个非常规的、TE无法简洁表达的计算模式或者你想极精细地控制初始循环结构以配合特定的调度策略时直接操作TensorIR就变得必要了。3. 手动创建TensorIR的完整实操流程理论说得再多不如一行代码。让我们从一个最简单的例子开始创建一个用于执行逐元素相加C[i, j] A[i, j] B[i, j]的TensorIR。我们将使用TVM的Python接口tvm.script来以编程方式构建。3.1 环境准备与模块导入首先确保你的环境已安装TVM。推荐从源码编译以获取完整功能。我们的创建将主要依赖tvm.script模块下的ir_module和tirTensorIR相关类。import tvm from tvm.script import tir as T import numpy as np这里的关键是from tvm.script import tir as T它引入了所有用于构建TensorIR的装饰器和类。T模块下的类如T.BufferT.iter_var是我们构建IR的“砖块”。3.2 构建一个简单的逐元素加法PrimFunc我们将分步构建一个完整的PrimFunc。步骤1定义函数签名与缓冲区T.prim_func def elementwise_add( A: T.Buffer((1024, 1024), “float32”), B: T.Buffer((1024, 1024), “float32”), C: T.Buffer((1024, 1024), “float32”) ) - None: # 函数体将在后续步骤中填充T.prim_func装饰器表明这是一个TensorIR原始函数。参数ABC被声明为T.Buffer这是TensorIR中表示数据缓冲区的类型。(1024, 1024)指定了缓冲区的形状“float32”是数据类型。- None表示这个函数没有返回值结果通过缓冲区C输出。步骤2在函数体内添加Block一个PrimFunc的主体由多个Block组成。对于这个简单的逐点操作我们只需要一个计算块。T.prim_func def elementwise_add(A: T.Buffer((1024, 1024), “float32”), B: T.Buffer((1024, 1024), “float32”), C: T.Buffer((1024, 1024), “float32”)): # 在函数作用域下我们创建一个Block with T.block(“root”): # 为这个Block定义两个数据并行迭代器i和j i T.axis.spatial(1024, T.axis.remap(“S”, [0])) j T.axis.spatial(1024, T.axis.remap(“S”, [0])) # 上面的定义是“占位符”实际范围由调度或上下文决定。 # 在Block内部我们使用T.reads和T.writes声明数据访问区域 T.reads(A[i, j], B[i, j]) T.writes(C[i, j]) # 核心计算语句 C[i, j] A[i, j] B[i, j]但上面的写法有一个问题i和j只是单个索引而我们需要它们循环覆盖整个0到1023的空间。在TensorIR中我们通常通过调度来将这种“单点计算”的Block绑定到具体的循环迭代器上。更直观的创建方式是从循环开始。步骤3通过循环嵌套构建完整计算更符合直觉的方式是先创建循环然后在循环体内实例化Block。我们可以使用T.axis.range来定义循环迭代器。T.prim_func def elementwise_add_loops(A: T.Buffer((1024, 1024), “float32”), B: T.Buffer((1024, 1024), “float32”), C: T.Buffer((1024, 1024), “float32”)): # 使用T.grid生成两层嵌套循环 for i in T.grid(1024): for j in T.grid(1024): with T.block(“C”): # 将循环变量i, j绑定到Block的迭代器上 vi T.axis.spatial(1024, i) vj T.axis.spatial(1024, j) T.reads(A[vi, vj], B[vi, vj]) T.writes(C[vi, vj]) C[vi, vj] A[vi, vj] B[vi, vj]T.grid(1024)创建了一个从0到1023的迭代序列。在循环体内每次迭代都创建一个名为“C”的Block实例。vi T.axis.spatial(1024, i)是关键。它声明了一个spatial类型的迭代器表示数据并行维度其定义域dom是0到1023并将当前循环变量i的值绑定给它。vj同理。T.reads和T.writes明确了每个Block实例访问的数据区域这对于后续分析至关重要。步骤4验证与打印创建完成后我们可以打印并检查这个TensorIR的结构。print(elementwise_add_loops.script())输出将是一个结构化的文本表示清晰地展示了循环、块和计算语句。你还可以使用tvm.lower来将其编译为底层IR看看它是否合法。sch tvm.tir.Schedule(elementwise_add_loops) print(type(sch)) # 确认成功创建了调度对象3.3 创建更复杂的案例矩阵乘法理解了基础我们来挑战一个经典案例矩阵乘法C[i, j] sum_k(A[i, k] * B[k, j])。这里引入了归约操作。T.prim_func def matmul(A: T.Buffer((1024, 1024), “float32”), B: T.Buffer((1024, 1024), “float32”), C: T.Buffer((1024, 1024), “float32”)): # 三层循环i, j 是空间维度k 是归约维度 for i in T.grid(1024): for j in T.grid(1024): with T.block(“init”): vi_init T.axis.spatial(1024, i) vj_init T.axis.spatial(1024, j) T.writes(C[vi_init, vj_init]) C[vi_init, vj_init] T.float32(0) # 初始化C为0 for k in T.grid(1024): with T.block(“update”): vi T.axis.spatial(1024, i) vj T.axis.spatial(1024, j) vk T.axis.reduce(1024, k) # 注意这里是reduce迭代器 T.reads(A[vi, vk], B[vk, vj], C[vi, vj]) T.writes(C[vi, vj]) C[vi, vj] C[vi, vj] A[vi, vk] * B[vk, vj]关键点解析分离初始化我们创建了一个独立的“init”块来将C的每个元素初始化为0。这比在归约循环中判断k0更清晰也给了调度器更多优化空间例如初始化可以和计算并行。归约迭代器vk T.axis.reduce(1024, k)声明vk是一个归约迭代器。这告诉TVM沿着k维度的循环是用于归约求和的在后续优化中如循环交换、分块需要保持归约的语义正确性。读写声明在“update”块中T.reads包含了C[vi, vj]因为我们要读取它当前的值并累加。T.writes同样指向C[vi, vj]表示会更新它。实操心得在手动创建涉及归约的TensorIR时清晰地区分spatial和reduce迭代器是重中之重。如果错误地将归约轴标记为spatial调度器可能会进行非法的变换导致计算结果错误。在不确定时多使用print(prim_func.script())来检查IR结构。4. 从TensorIR到可执行代码调度与构建创建出原始的TensorIR只是第一步。原始的、多重循环的IR通常性能很差。接下来我们需要通过调度来优化它然后将其构建为特定后端如LLVM、CUDA的可执行模块。4.1 对创建的TensorIR进行基础调度我们以elementwise_add_loops为例演示几个基础的调度操作。# 从我们创建的PrimFunc创建调度对象 sch tvm.tir.Schedule(elementwise_add_loops) # 获取名为“C”的Block block_c sch.get_block(“C”) # 获取包围这个Block的循环 i, j sch.get_loops(block_c) # 调度操作1循环分块 (Tiling) # 将i轴拆分为外循环i_outer和内循环i_inner i_outer, i_inner sch.split(i, factors[None, 64]) # None表示自动推导外循环大小 j_outer, j_inner sch.split(j, factors[None, 64]) # 重新组织循环顺序i_outer, j_outer, i_inner, j_inner sch.reorder(i_outer, j_outer, i_inner, j_inner) print(“After tiling and reorder:”) print(sch.mod.script()) # 调度操作2向量化 (Vectorization) # 对最内层的j_inner循环进行向量化假设硬件支持256位宽即8个float32 sch.vectorize(j_inner) print(“After vectorization:”) print(sch.mod.script()) # 调度操作3并行化 (Parallel) # 将最外层的i_outer循环标记为并行执行 sch.parallel(i_outer)通过调度我们将一个简单的两层循环转换成了分块、重排序、向量化、并行化的复杂但高性能的循环结构。这就是TVM威力的体现分离计算描述与调度优化。4.2 编译与运行测试调度完成后我们需要将其编译为具体硬件上的代码。# 1. 构建运行时模块 target tvm.target.Target(“llvm -mcpucore-avx2”) # 指定CPU目标 # 也可以使用 “cuda” 以GPU为目标 # target tvm.target.Target(“cuda”) # 通过build接口将调度后的模块sch.mod编译 rt_mod tvm.build(sch.mod, targettarget) # 2. 准备数据并运行 n 1024 A_np np.random.uniform(size(n, n)).astype(“float32”) B_np np.random.uniform(size(n, n)).astype(“float32”) C_np np.zeros((n, n)).astype(“float32”) # 创建TVM运行时需要的数据容器NDArray A_nd tvm.nd.array(A_np) B_nd tvm.nd.array(B_np) C_nd tvm.nd.array(C_np) # 运行编译好的函数 rt_mod(A_nd, B_nd, C_nd) # 3. 验证结果 C_ref A_np B_np tvm.testing.assert_allclose(C_nd.numpy(), C_ref, rtol1e-5) print(“Test passed!”)如果一切顺利你将看到“Test passed!”。这意味着你手动创建的TensorIR经过TVM编译成功在CPU上执行并得到了正确结果。4.3 性能对比分析为了体会手动创建并调度的价值我们可以做一个简单的性能对比。import timeit # 测试我们优化后的版本 timer rt_mod.time_evaluator(rt_mod.entry_name, devicetvm.cpu(), number100, repeat3) optimized_time timer(A_nd, B_nd, C_nd).mean print(f“Optimized version average time: {optimized_time * 1000:.2f} ms”) # 作为对比用TVM的TE生成一个朴素版本并编译 from tvm import te A_te te.placeholder((n, n), name‘A’) B_te te.placeholder((n, n), name‘B’) C_te te.compute((n, n), lambda i, j: A_te[i, j] B_te[i, j]) s_naive te.create_schedule(C_te.op) func_naive tvm.build(s_naive, [A_te, B_te, C_te], target“llvm”) C_nd_naive tvm.nd.array(np.zeros((n, n)).astype(“float32”)) timer_naive func_naive.time_evaluator(func_naive.entry_name, devicetvm.cpu(), number100, repeat3) naive_time timer_naive(A_nd, B_nd, C_nd_naive).mean print(f“Naive TE version average time: {naive_time * 1000:.2f} ms”) print(f“Speedup: {naive_time / optimized_time:.2f}x”)在我的测试环境中AVX2指令集经过基础调度的版本通常能比完全朴素的版本快上数倍。这个差距在矩阵乘法等计算密集型操作中会更加惊人。5. 常见问题、调试技巧与高级模式手动创建TensorIR就像直接操作汇编强大但容易出错。下面分享一些我踩过坑后总结的经验。5.1 常见错误与排查清单Buffer维度或数据类型不匹配症状在tvm.build阶段报错提示Buffer形状或dtype错误。排查仔细检查T.Buffer声明中的形状元组和数据类型字符串是否与后续访问的索引一致。例如声明为(1024, 1024)的缓冲区索引必须是[i, j]形式且i, j的值必须在0-1023范围内。迭代器绑定错误症状运行结果不正确或者调度变换时报语义错误。排查确认在Block内部T.axis.spatial/reduce的第一个参数定义域是否正确第二个参数绑定值是否确实来自外层的循环变量。使用sch.mod.script()打印IR检查每个Block的迭代器定义。T.reads/T.writes区域声明错误症状TVM在进行算子融合等优化时行为异常或者无法正确分析依赖。排查确保T.reads包含了Block内所有被读取的缓冲区区域T.writes包含了所有被写入的区域。对于归约操作输出缓冲区既在reads中也在writes中。区域声明应尽可能精确例如A[vi, vk]比A[0:1024, 0:1024]更好。调度变换破坏了计算语义症状调度后如reordercompute_at计算结果变成NaN或与预期不符。排查这是最棘手的问题。对于复杂的调度建议逐步进行每做一个变换就验证一次计算结果。TVM的调度器会尽量保证合法性但并非万能。特别是手动创建复杂IR时需要你对数据流依赖有清晰的理解。5.2 调试工具与技巧print(sch.mod.script())是你的最佳朋友在任何步骤后都打印IR直观地看到变换效果。使用tvm.tir.stmt_functor进行IR遍历对于非常复杂的IR可以编写访问者模式函数来打印特定节点信息辅助调试。从简到繁逐步构建不要试图一次性写出完美的复杂IR如带共享内存的GPU矩阵乘法。先写出功能正确的朴素版本然后逐步添加优化结构如分块、缓存每步都进行测试。与TE生成的IR进行对比当你对某个计算不确定如何用TensorIR表达时先用te.compute写出来然后用tvm.lower或sch tvm.tir.Schedule(func)得到其TensorIR打印出来作为参考模板。5.3 探索高级模式引入中间缓存在GPU编程或优化CPU缓存时我们经常需要手动管理数据局部性。TensorIR可以很好地表达这一点。以下是一个概念性的示例展示如何在矩阵乘法中为A矩阵的块引入缓存T.prim_func def matmul_with_cache(A: T.Buffer((1024, 1024), “float32”), B: T.Buffer((1024, 1024), “float32”), C: T.Buffer((1024, 1024), “float32”)): # 在共享内存或寄存器级别声明一个缓存块 A_shared T.alloc_buffer((64, 64), dtype“float32”, scope“shared”) # 假设是GPU共享内存 # 或者 A_local T.alloc_buffer((64, 64), dtype“float32”, scope“local”) # 寄存器 for i_outer in T.grid(16): # 1024 / 64 16 for j_outer in T.grid(16): # 阶段1: 将A的一个块从全局内存加载到共享内存缓存 for x in T.grid(64): for y in T.grid(64): with T.block(“load_A”): vx T.axis.spatial(64, x) vy T.axis.spatial(64, y) T.reads(A[i_outer * 64 vx, j_outer * 64 vy]) T.writes(A_shared[vx, vy]) A_shared[vx, vy] A[i_outer * 64 vx, j_outer * 64 vy] # 阶段2: 使用缓存的数据进行计算 (这里省略了B和C的类似逻辑) # ... 计算块 ...这个例子展示了T.alloc_buffer的用法它可以声明一个指定形状、数据类型和内存作用域scope的临时缓冲区。scope可以是“global”“shared”“local”“warp”等对应不同的硬件内存层次。通过精细控制数据的缓存和移动可以极大化利用硬件带宽和计算单元这也是手动创建TensorIR在极致性能优化中的核心价值所在。手动创建TensorIR是一个由浅入深的过程。从最简单的逐点操作开始理解Block、迭代器、读写区域这些基本概念然后尝试表达归约、分块、缓存等复杂模式。这个过程会让你对TVM的编译逻辑有脱胎换骨的理解。当你再使用te.compute或者AutoScheduler时你就能清晰地知道它们在你写的计算描述之下究竟构建和变换了什么。这不仅仅是学习一个工具更是掌握一种对计算进行系统性思考和优化的思维方式。