本文详细介绍了如何将基础的TensorIR实现通过系统化的变换方法优化为高性能版本。以批处理矩阵乘法接ReLU激活函数为例，我们将展示完整的变换流程，包括循环分割、重排序、计算位置调整、归约分解以及向量化等关键技术。
目录
-
1. 问题背景与原始实现 -
2. 目标优化代码分析 -
3. TensorIR变换步骤详解 -
4. 完整变换过程实现 -
5. 性能优化原理分析 -
6. 总结与展望
问题背景与原始实现
批处理矩阵乘法（Batched Matrix Multiplication）后接ReLU激活函数是深度学习中的常见操作。在NumPy中，这一操作可以直观地表示为：
def lnumpy_mm_relu_v2(A: np.ndarray, B: np.ndarray, C: np.ndarray):
Y = np.empty((16, 128, 128), dtype="float32")
for n in range(16):
for i in range(128):
for j in range(128):
for k in range(128):
if k == 0:
Y[n, i, j] = 0
Y[n, i, j] = Y[n, i, j] + A[n, i, k] * B[n, k, j]
for n in range(16):
for i in range(128):
for j in range(128):
C[n, i, j] = max(Y[n, i, j], 0)
这段NumPy代码展示了一个直观但低效的实现方式。它包含两个主要阶段：首先进行批处理矩阵乘法，然后应用ReLU激活函数。这种实现的主要问题在于：
-
• 四层嵌套循环导致高时间复杂度 -
• 内存访问模式不连续，缓存利用率低 -
• 缺乏并行化机会 -
• 中间结果存储开销大
将这一NumPy实现转换为TensorIR后，我们得到以下基础实现：
@tvm.script.ir_module
class MyBmmRule:
@T.prim_func
def bmm_relu(A: T.Buffer[(16, 128, 128), "float32"],
W: T.Buffer[(16, 128, 128), "float32"],
Y: T.Buffer[(16, 128, 128), "float32"]):
T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
# 必须在此分配缓冲区
Y_ = T.alloc_buffer([16, 128, 128], dtype="float32")
for n, i, j, k in T.grid(16, 128, 128, 128):
with T.block("M"):
vn = T.axis.spatial(16, n)
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
with T.init():
Y_[vn, vi, vj] = T.float32(0)
Y_[vn, vi, vj] += A[vn, vi, vk] * W[vn, vk, vj]
for n, i, j in T.grid(16, 128, 128):
with T.block("R"):
vn = T.axis.spatial(16, n)
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
Y[vn, vi, vj] = T.max(Y_[vn, vi, vj], T.float32(0))
这个TensorIR实现虽然功能正确，但在性能上存在明显不足。它采用了简单的四层嵌套循环结构，缺乏现代处理器所需的数据局部性优化和并行化支持。T.grid函数创建了平面的循环迭代空间，而T.block定义了计算块，其中通过T.axis明确了各维度的角色：空间维度（spatial）和归约维度（reduce）。
目标优化代码分析
我们的目标是将基础TensorIR转换为以下高性能版本：
@tvm.script.ir_module
class TargetModule:
@T.prim_func
def bmm_relu(A: T.Buffer[(16, 128, 128), "float32"],
B: T.Buffer[(16, 128, 128), "float32"],
C: T.Buffer[(16, 128, 128), "float32"]) -> None:
T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
Y = T.alloc_buffer([16, 128, 128], dtype="float32")
for i0 in T.parallel(16):
for i1, i2_0 in T.grid(128, 16):
for ax0_init in T.vectorized(8):
with T.block("M_init"):
n, i = T.axis.remap("SS", [i0, i1])
j = T.axis.spatial(128, i2_0 * 8 + ax0_init)
Y[n, i, j] = T.float32(0)
for ax1_0 in T.serial(32):
for ax1_1 in T.unroll(4):
for ax0 in T.serial(8):
with T.block("M_update"):
n, i = T.axis.remap("SS", [i0, i1])
j = T.axis.spatial(128, i2_0 * 8 + ax0)
k = T.axis.reduce(128, ax1_0 * 4 + ax1_1)
Y[n, i, j] = Y[n, i, j] + A[n, i, k] * B[n, k, j]
for i2_1 in T.vectorized(8):
with T.block("R"):
n, i = T.axis.remap("SS", [i0, i1])
j = T.axis.spatial(128, i2_0 * 8 + i2_1)
C[n, i, j] = T.max(Y[n, i, j], T.float32(0))
这个优化版本引入了多种高级优化技术。外层循环使用T.parallel(16)实现了批处理维度的并行化，充分利用多核处理器。初始化阶段采用T.vectorized(8)进行向量化初始化，一次性处理8个元素。矩阵乘法阶段通过循环分割和重排序优化数据局部性，其中T.unroll(4)指示编译器展开内层循环以减少控制开销。ReLU计算与矩阵乘法融合，避免了额外的数据遍历开销。
优化技术对比分析
|
|
|
|
|
|---|---|---|---|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
TensorIR变换步骤详解
要实现从基础版本到优化版本的转换，我们需要遵循系统化的变换流程：
变换步骤框架
|
|
|
|
|
|---|---|---|---|
|
|
|
|
sch.get_block() |
|
|
|
|
sch.get_loops() |
|
|
|
|
split()
reorder(), compute_at()
|
|
|
|
|
decompose_reduction() |
|
|
|
|
vectorize()
parallel(), unroll()
|
关键变换技术说明
循环分割(Split)
循环分割将大循环分解为多个小循环，便于后续的重排序和优化。这种技术允许我们更精细地控制数据访问模式，为向量化和缓存优化创造条件。
循环重排序(Reorder)
改变循环嵌套顺序可以显著影响数据局部性。正确的重排序模式应该使内存访问模式更加连续，减少缓存失效，提高内存带宽利用率。
计算位置调整(Compute_at/Reverse_compute_at)
将计算移动到合适的位置，可以减少中间结果的存储和加载，实现计算融合。这种技术特别适用于消除不必要的中间缓冲区，减少内存占用。
归约分解(Decompose Reduction)
将归约操作的初始化和计算分离，为后续优化创造机会。这种分解使得我们可以对初始化阶段应用不同的优化策略，如向量化。
完整变换过程实现
下面是完整的TensorIR变换代码实现：
# 步骤1: 创建调度器并获取块
sch = tvm.tir.Schedule(MyBmmRule)
block_M = sch.get_block("M", func_name="bmm_relu")
# 步骤2: 获取循环
n, i, j, k = sch.get_loops(block_M)
# 步骤3: 组织循环 - 分割
k0, k1 = sch.split(k, factors=[32, 4])
j0, j1 = sch.split(j, factors=[16, 8])
# 步骤3: 组织循环 - 重排序
sch.reorder(j0, k0, k1, j1)
# 步骤3: 组织循环 - 调整计算位置
block_R = sch.get_block("R", func_name="bmm_relu")
sch.reverse_compute_at(block_R, j0)
# 步骤4: 分解归约
block_M = sch.get_block("M", func_name="bmm_relu")
M_init = sch.decompose_reduction(block_M, k0)
# 步骤5: 应用优化原语
# 向量化初始化
n, i, j_0, j_1_init = sch.get_loops(M_init)
sch.vectorize(j_1_init)
# 向量化ReLU计算
n, i, j_0, i2_1 = sch.get_loops(block_R)
sch.vectorize(i2_1)
# 获取更新块并准备后续优化
block_M_update = sch.get_block("M_update", func_name="bmm_relu")
n, i, j_0, k_0, k_1, j_1 = sch.get_loops(block_M_update)
# 应用循环展开
sch.unroll(k_1)
# 应用并行化
sch.parallel(n)
这个变换流程展示了如何系统地将基础实现转换为高性能版本。首先通过获取计算块和循环结构了解代码结构，然后通过循环分割和重排序优化数据访问模式，接着调整计算位置实现操作融合，最后应用硬件特定的优化原语。
让我们逐步分析每个变换阶段的结果：
初始变换后的中间结果
在执行循环分割和重排序后，我们得到以下中间代码：
@tvm.script.ir_module
class Module:
@tir.prim_func
def bmm_relu(A: tir.Buffer[(16, 128, 128), "float32"],
B: tir.Buffer[(16, 128, 128), "float32"],
C: tir.Buffer[(16, 128, 128), "float32"]) -> None:
tir.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
Y = tir.alloc_buffer([16, 128, 128], dtype="float32")
for n in tir.parallel(16):
for i, j_0, k_0, k_1, j_1 in tir.grid(128, 16, 32, 4, 8):
with tir.block("M"):
vn, vi = tir.axis.remap("SS", [n, i])
vj = tir.axis.spatial(128, j_0 * 8 + j_1)
vk = tir.axis.reduce(128, k_0 * 4 + k_1)
tir.reads(A[vn, vi, vk], B[vn, vk, vj])
tir.writes(Y[vn, vi, vj])
with tir.init():
Y[vn, vi, vj] = tir.float32(0)
Y[vn, vi, vj] = Y[vn, vi, vj] + A[vn, vi, vk] * B[vn, vk, vj]
for n, i, j in tir.grid(16, 128, 128):
with tir.block("R"):
vn, vi, vj = tir.axis.remap("SSS", [n, i, j])
tir.reads(Y[vn, vi, vj])
tir.writes(C[vn, vi, vj])
C[vn, vi, vj] = tir.max(Y[vn, vi, vj], tir.float32(0))
这个阶段主要完成了循环结构的重组。通过将j维度分割为[16, 8]，k维度分割为[32, 4]，我们为后续的向量化和循环展开创造了条件。循环重排序确保内存访问模式更加连续，提高了缓存利用率。
归约分解后的结果
执行归约分解后，初始化部分被分离出来：
@tvm.script.ir_module
class Module:
@tir.prim_func
def bmm_relu(A: tir.Buffer[(16, 128, 128), "float32"],
B: tir.Buffer[(16, 128, 128), "float32"],
C: tir.Buffer[(16, 128, 128), "float32"]) -> None:
tir.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
Y = tir.alloc_buffer([16, 128, 128], dtype="float32")
for n in tir.parallel(16):
for i, j_0 in tir.grid(128, 16):
for j_1_init in tir.serial(8):
with tir.block("M_init"):
vn, vi = tir.axis.remap("SS", [n, i])
vj = tir.axis.spatial(128, j_0 * 8 + j_1_init)
tir.reads()
tir.writes(Y[vn, vi, vj])
Y[vn, vi, vj] = tir.float32(0)
for k_0, k_1, j_1 in tir.grid(32, 4, 8):
with tir.block("M_update"):
vn, vi = tir.axis.remap("SS", [n, i])
vj = tir.axis.spatial(128, j_0 * 8 + j_1)
vk = tir.axis.reduce(128, k_0 * 4 + k_1)
tir.reads(Y[vn, vi, vj], A[vn, vi, vk], B[vn, vk, vj])
tir.writes(Y[vn, vi, vj])
Y[vn, vi, vj] = Y[vn, vi, vj] + A[vn, vi, vk] * B[vn, vk, vj]
for ax0 in tir.serial(8):
with tir.block("R"):
vn, vi = tir.axis.remap("SS", [n, i])
vj = tir.axis.spatial(128, j_0 * 8 + ax0)
tir.reads(Y[vn, vi, vj])
tir.writes(C[vn, vi, vj])
C[vn, vi, vj] = tir.max(Y[vn, vi, vj], tir.float32(0))
归约分解是优化过程中的关键步骤。它将矩阵乘法的初始化阶段与计算阶段分离，使得我们可以对初始化应用向量化优化。同时，通过reverse_compute_at将ReLU计算融合到主循环中，避免了额外的数据遍历，减少了内存访问开销。
最终优化结果
应用所有优化原语后，我们得到最终的TensorIR代码：
@tvm.script.ir_module
class TargetModule:
@T.prim_func
def bmm_relu(A: T.Buffer[(16, 128, 128), "float32"],
B: T.Buffer[(16, 128, 128), "float32"],
C: T.Buffer[(16, 128, 128), "float32"]) -> None:
T.func_attr({"global_symbol": "bmm_relu", "tir.noalias": True})
Y = T.alloc_buffer([16, 128, 128], dtype="float32")
for i0 in T.parallel(16):
for i1, i2_0 in T.grid(128, 16):
for ax0_init in T.vectorized(8):
with T.block("M_init"):
n, i = T.axis.remap("SS", [i0, i1])
j = T.axis.spatial(128, i2_0 * 8 + ax0_init)
Y[n, i, j] = T.float32(0)
for ax1_0 in T.serial(32):
for ax1_1 in T.unroll(4):
for ax0 in T.serial(8):
with T.block("M_update"):
n, i = T.axis.remap("SS", [i0, i1])
j = T.axis.spatial(128, i2_0 * 8 + ax0)
k = T.axis.reduce(128, ax1_0 * 4 + ax1_1)
Y[n, i, j] = Y[n, i, j] + A[n, i, k] * B[n, k, j]
for i2_1 in T.vectorized(8):
with T.block("R"):
n, i = T.axis.remap("SS", [i0, i1])
j = T.axis.spatial(128, i2_0 * 8 + i2_1)
C[n, i, j] = T.max(Y[n, i, j], T.float32(0))
最终优化版本融合了多种高级优化技术。外层循环使用T.parallel实现批处理维度的并行化，充分利用多核处理器。初始化阶段采用T.vectorized(8)进行向量化初始化，一次性处理8个元素。矩阵乘法阶段通过T.unroll(4)实现循环展开，减少控制开销。ReLU计算同样使用向量化优化，进一步提高性能。
性能优化原理分析
数据局部性优化
原始实现的四层嵌套循环在数据访问上存在明显的问题。通过循环分割和重排序，我们确保了连续的内存访问模式，这对现代处理器至关重要：
-
1. 连续内存访问：通过将 j1维度保持为最内层循环，确保对Y、A和B的访问在内存中是连续的，这显著提高了缓存行利用率 -
2. 缓存重用：适当的分块大小确保计算所需的数据块能够保留在L1/L2缓存中，减少主内存访问 -
3. 数据预取：连续的访问模式使得硬件预取器能够有效工作，提前加载需要的数据
并行化策略
优化实现采用了多层次的并行化策略：
|
|
|
|
|
|---|---|---|---|
|
|
n
|
|
T.parallel(16) |
|
|
|
|
T.vectorized(8) |
|
|
|
|
T.unroll(4) |
这种分层并行化策略确保了计算任务在多个维度上同时进行，最大限度地利用了现代处理器的计算资源。
内存层次结构利用
优化后的代码更好地利用了现代处理器的内存层次结构：
-
1. 寄存器级别：通过循环展开，编译器可以将更多变量保存在寄存器中，减少内存访问 -
2. 缓存级别：分块计算确保数据块大小适合L1/L2缓存，提高缓存命中率 -
3. 内存级别：连续访问模式充分利用内存带宽，减少访问延迟
计算与内存访问平衡
在原始实现中，计算与内存访问的比例不理想。优化后的实现通过多种技术改善了这种平衡：
-
• 循环融合：将矩阵乘法和ReLU计算融合，减少中间结果存储 -
• 数据重用：通过循环重排序增加计算/内存访问比 -
• 向量化：同时提高计算效率和内存访问效率
总结与展望
本文通过一个具体的批处理矩阵乘法接ReLU操作的例子，详细展示了TensorIR从基础实现到高性能优化的完整变换流程。关键的收获包括：
核心变换技术
-
1. 系统性方法：遵循获取块→获取循环→组织循环→分解归约→应用优化的系统流程 -
2. 循环变换：分割、重排序、计算位置调整是优化数据局部性的关键 -
3. 硬件原语：向量化、并行化、循环展开是释放硬件性能的重要手段
实际应用价值
这种变换方法不仅适用于矩阵乘法，还可以推广到各种计算密集型操作。在实际的深度学习模型优化中，类似的变换技术可以带来显著的性能提升。
未来发展方向
随着硬件不断发展，TensorIR等中间表示将继续演化，支持更多自动化优化和针对特定硬件的优化策略。理解这些基础变换原理，对于适应未来的深度学习编译器发展至关重要。
通过掌握TensorIR变换技术，开发者可以在保持代码可维护性的同时，充分发挥现代硬件的计算潜力，为高效的深度学习推理和训练奠定坚实基础。