TensorIR 构建

TensorIR 构建#

在本节中，将介绍在 Apache TVM Unity 中编写 TensorIR 函数的方法。本教程假设您已经熟悉 TensorIR 的基本概念。

备注

本教程专注于构建 独立的 TensorIR 函数。这里介绍的技术对于最终用户编译 Relax 模型来说不是必需的。

使用 TVMScript 构建 TensorIR#

通过 TVMScript 创建 TensorIR 函数的最简单方法。TVMScript 是 TVM Python 方言，代表 TVM 中的 TensorIR。

标准格式#

以下是 ir_module 的完整格式，以及在TVMScript中：

使用语法糖的简洁写法#

为了便于编写，可以采用以下语法糖来简化代码：

利用 T.grid 来压缩嵌套循环；
使用 T.axis.remap 来简化块迭代器注释；
排除 T.reads 和 T.writes 对于可以从块体中推断出内容的块；

可以使用以下代码来验证这两个模块是等价的：

print(tvm.ir.structural_equal(MyModule, ConciseModule))

True

与 Python 变量交互#

尽管 TVMScript 不由 Python 解释器执行，但与 Python 的有限交互是可行的。例如，可以使用 Python 变量来确定 TensorIR 的形状和数据类型。

检查等价性：

print(tvm.ir.structural_equal(ConciseModule, ConciseModuleFromPython))

True

具有动态形状的 TensorIR 函数#

尽管 TVMScript 不由 Python 解释器执行，但与 Python 的有限交互是可行的。例如，可以使用 Python 变量来确定 TensorIR 的形状和数据类型。

现在让我们检查运行时动态形状推断：

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.build(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))

使用张量表达式创建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()

创建动态形状函数#

还可以使用张量表达式创建动态形状函数。唯一的区别是需要将输入张量的形状指定为符号变量。

# 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()