TensorIR 创建
在本节中，我们将介绍在 Apache TVM Unity 中编写 TensorIR 函数的方法。本教程假定您熟悉 TensorIR 的基本概念。如果尚不熟悉，请先参考理解 TensorIR 抽象。
注意
本教程侧重于构建**独立的** TensorIR 函数。此处介绍的技术并非终端用户编译 Relax 模型所必需的。
使用 TVMScript 创建 TensorIR
通过 TVMScript 创建 TensorIR 函数的最直接方法。TVMScript 是一种 TVM Python 方言，用于表示 TVM 中的 TensorIR。
重要提示
虽然 TVMScript 采用 Python 语法和 AST，确保与自动补全和语法检查等 Python 工具的完全兼容性，但它不是原生 Python 语言，不能由 Python 解释器执行。
更准确地说，装饰器 **@tvm.script** 从被装饰的函数中提取 Python AST，随后将其解析为 TensorIR。
标准格式
让我们以mm_relu 为例，它来自理解 TensorIR 抽象。以下是 ir_module 的完整格式以及 TVMScript 中的格式
import numpy as np
import tvm
from tvm.script import ir as I
from tvm.script import tir as T
@I.ir_module
class MyModule:
@T.prim_func
def mm_relu(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
):
Y = T.alloc_buffer((128, 128), dtype="float32")
for i in range(128):
for j in range(128):
for k in range(128):
with T.block("Y"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
T.reads(A[vi, vk], B[vk, vj])
T.writes(Y[vi, vj])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i in range(128):
for j in range(128):
with T.block("C"):
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
T.reads(Y[vi, vj])
T.writes(C[vi, vj])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
使用语法糖简化
为了便于编写，我们可以使用以下语法糖来简化代码
使用
T.grid来简化嵌套循环；使用
T.axis.remap来缩写块迭代器注释；对于内容可以从块体推断出来的块，排除
T.reads和T.writes；
@I.ir_module
class ConciseModule:
@T.prim_func
def mm_relu(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32"),
):
Y = T.alloc_buffer((128, 128), dtype="float32")
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
Y[vi, vj] = T.float32(0)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(128, 128):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = T.max(Y[vi, vj], T.float32(0))
我们可以使用以下代码来验证这两个模块是等效的
print(tvm.ir.structural_equal(MyModule, ConciseModule))
True
与 Python 变量交互
尽管 TVMScript 不由 Python 解释器执行，但与 Python 的有限交互是可行的。例如，Python 变量可用于确定 TensorIR 的形状和数据类型。
# Python variables
M = N = K = 128
dtype = "float32"
# IRModule in TVMScript
@I.ir_module
class ConciseModuleFromPython:
@T.prim_func
def mm_relu(
A: T.Buffer((M, K), dtype),
B: T.Buffer((K, N), dtype),
C: T.Buffer((M, N), dtype),
):
Y = T.alloc_buffer((M, N), dtype)
for i, j, k in T.grid(M, N, K):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
Y[vi, vj] = T.cast(T.float32(0), dtype)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(M, N):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))
检查等效性
print(tvm.ir.structural_equal(ConciseModule, ConciseModuleFromPython))
True
具有动态形状的 TensorIR 函数
尽管 TVMScript 不由 Python 解释器执行，但与 Python 的有限交互是可行的。例如，Python 变量可用于确定 TensorIR 的形状和数据类型。
@I.ir_module
class DynamicShapeModule:
@T.prim_func
def mm_relu(a: T.handle, b: T.handle, c: T.handle):
# Dynamic shape definition
M, N, K = T.int32(), T.int32(), T.int32()
# Bind the input buffers with the dynamic shapes
A = T.match_buffer(a, [M, K], dtype)
B = T.match_buffer(b, [K, N], dtype)
C = T.match_buffer(c, [M, N], dtype)
Y = T.alloc_buffer((M, N), dtype)
for i, j, k in T.grid(M, N, K):
with T.block("Y"):
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
with T.init():
Y[vi, vj] = T.cast(T.float32(0), dtype)
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
for i, j in T.grid(M, N):
with T.block("C"):
vi, vj = T.axis.remap("SS", [i, j])
C[vi, vj] = T.max(Y[vi, vj], T.cast(T.float32(0), dtype))
现在让我们检查运行时动态形状推断
def evaluate_dynamic_shape(lib: tvm.runtime.Module, m: int, n: int, k: int):
A = tvm.nd.array(np.random.uniform(size=(m, k)).astype("float32"))
B = tvm.nd.array(np.random.uniform(size=(k, n)).astype("float32"))
C = tvm.nd.array(np.zeros((m, n), dtype="float32"))
lib(A, B, C)
return C.numpy()
# Compile lib only once
dyn_shape_lib = tvm.compile(DynamicShapeModule, target="llvm")
# Able to handle different shapes
print(evaluate_dynamic_shape(dyn_shape_lib, m=4, n=4, k=4))
print(evaluate_dynamic_shape(dyn_shape_lib, m=64, n=64, k=128))
[[0.7870551 0.40384182 0.4803333 1.3946644 ]
[0.30299586 0.58826745 0.989751 0.756879 ]
[1.2259218 0.8456589 0.66306597 1.341947 ]
[1.1576492 0.8520223 0.8902653 1.3287237 ]]
[[38.987442 32.73539 33.6178 ... 33.693363 30.839447 35.58607 ]
[36.944706 34.4549 33.324024 ... 31.845808 32.066383 35.31651 ]
[37.251682 33.683674 34.065304 ... 33.9451 31.355616 33.2658 ]
...
[31.92583 26.470911 27.017 ... 29.027529 26.014395 28.261797]
[36.585373 33.189007 33.31166 ... 32.758415 30.030182 35.012505]
[35.34669 28.41992 31.260956 ... 29.160038 27.674332 30.496407]]
使用张量表达式创建 TensorIR
通常，为了更简洁地表达计算，TensorIR 的具体细节被忽略，从而促成了 TensorIR 的务实生成。这就是张量表达式 (TE) 发挥作用的地方。
张量表达式 (TE) 是一种领域特定语言，通过类似表达式的 API 描述一系列计算。
注意
张量表达式在 TVM 堆栈中包含两个组件：表达式和调度。表达式是体现计算模式的领域特定语言，这正是我们在本节中要讨论的内容。相反，TE 调度是传统的调度方法，在 TVM Unity 堆栈中已被 TensorIR 调度取代。
创建静态形状函数
我们使用上小节中相同的 mm_relu 示例来演示 TE 创建方法。
from tvm import te
A = te.placeholder((128, 128), "float32", name="A")
B = te.placeholder((128, 128), "float32", name="B")
k = te.reduce_axis((0, 128), "k")
Y = te.compute((128, 128), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = te.compute((128, 128), lambda i, j: te.max(Y[i, j], 0), name="C")
这里 te.compute 接受签名 te.compute(output_shape, fcompute)。而 fcompute 函数描述了我们希望如何计算给定索引的每个元素 Y[i, j] 的值
上述 lambda 表达式封装了计算：\(Y_{i, j} = \sum_k A_{i, k} \times B_{k, j}\)。在定义计算之后，我们可以通过结合感兴趣的相关参数来制定 TensorIR 函数。在这个特定实例中，我们的目标是构建一个具有两个输入参数 **A、B** 和一个输出参数 **C** 的函数。
te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
TEModule = tvm.IRModule({"mm_relu": te_func})
TEModule.show()
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(A: T.Buffer((128, 128), "float32"), B: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")):
T.func_attr({"tir.noalias": T.bool(True)})
# with T.block("root"):
Y = T.alloc_buffer((128, 128))
for i, j, k in T.grid(128, 128, 128):
with T.block("Y"):
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k])
T.reads(A[v_i, v_k], B[v_k, v_j])
T.writes(Y[v_i, v_j])
with T.init():
Y[v_i, v_j] = T.float32(0.0)
Y[v_i, v_j] = Y[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
for i, j in T.grid(128, 128):
with T.block("C"):
v_i, v_j = T.axis.remap("SS", [i, j])
T.reads(Y[v_i, v_j])
T.writes(C[v_i, v_j])
C[v_i, v_j] = T.max(Y[v_i, v_j], T.float32(0.0))
创建动态形状函数
我们还可以使用张量表达式创建动态形状函数。唯一的区别是我们需要将输入张量的形状指定为符号变量。
# Declare symbolic variables
M, N, K = te.var("m"), te.var("n"), te.var("k")
A = te.placeholder((M, N), "float32", name="A")
B = te.placeholder((K, N), "float32", name="B")
k = te.reduce_axis((0, K), "k")
Y = te.compute((M, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="Y")
C = te.compute((M, N), lambda i, j: te.max(Y[i, j], 0), name="C")
dyn_te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
DynamicTEModule = tvm.IRModule({"mm_relu": dyn_te_func})
DynamicTEModule.show()
# from tvm.script import ir as I
# from tvm.script import tir as T
@I.ir_module
class Module:
@T.prim_func
def mm_relu(var_A: T.handle, var_B: T.handle, var_C: T.handle):
T.func_attr({"tir.noalias": T.bool(True)})
m, n = T.int32(), T.int32()
A = T.match_buffer(var_A, (m, n))
k = T.int32()
B = T.match_buffer(var_B, (k, n))
C = T.match_buffer(var_C, (m, n))
# with T.block("root"):
Y = T.alloc_buffer((m, n))
for i, j, k_1 in T.grid(m, n, k):
with T.block("Y"):
v_i, v_j, v_k = T.axis.remap("SSR", [i, j, k_1])
T.reads(A[v_i, v_k], B[v_k, v_j])
T.writes(Y[v_i, v_j])
with T.init():
Y[v_i, v_j] = T.float32(0.0)
Y[v_i, v_j] = Y[v_i, v_j] + A[v_i, v_k] * B[v_k, v_j]
for i, j in T.grid(m, n):
with T.block("C"):
v_i, v_j = T.axis.remap("SS", [i, j])
T.reads(Y[v_i, v_j])
T.writes(C[v_i, v_j])
C[v_i, v_j] = T.max(Y[v_i, v_j], T.float32(0.0))