张量程序抽象

在我们深入 TensorIR 的细节之前，首先介绍什么是原始张量函数。原始张量函数是对应于计算操作的单个“单元”的函数。例如，卷积操作可以是一个原始张量函数，融合的卷积 + relu 操作也可以是一个原始张量函数。通常，原始张量函数实现的典型抽象包含以下要素：多维缓冲区、驱动张量计算的循环嵌套，以及最终的计算语句本身。

from tvm.script import tir as T

@T.prim_func
def main(
    A: T.Buffer((128,), "float32"),
    B: T.Buffer((128,), "float32"),
    C: T.Buffer((128,), "float32"),
) -> None:
    for i in range(128):
        with T.block("C"):
            vi = T.axis.spatial(128, i)
            C[vi] = A[vi] + B[vi]

张量程序的主要元素

演示的原始张量函数计算两个向量的元素wise和。该函数

• 接受三个多维缓冲区作为参数，并生成一个多维缓冲区作为输出。

• 包含一个单独的循环嵌套 i，用于促进计算。

• 具有一个单一的计算语句，用于计算两个向量的元素wise和。

TensorIR 中的额外结构

关键的是，我们无法对程序执行任意转换，因为某些计算依赖于循环的顺序。幸运的是，我们关注的大多数原始张量函数都具有有利的属性，例如循环迭代之间的独立性。例如，上述程序包括块和迭代注释

• 块注释 with T.block("C") 表示该块是指定用于调度的基本计算单元。一个块可能包含单个计算语句、带有循环的多个计算语句或不透明的内部函数（如 Tensor Core 指令）。

• 迭代注释 T.axis.spatial，指示变量 vi 映射到 i，并且所有迭代都是独立的。

虽然此信息对于执行特定程序不是至关重要的，但在转换程序时证明很有用。因此，只要我们遍历从 0 到 128 的所有索引元素，我们就可以自信地并行化或重新排序与 vi 关联的循环。