526互联

机器学习编译（三）：张量程序案例 TensorIR

发布时间 2023-08-13 12:21:25作者: machine_gun_lin

使用张量程序抽象的目的是为了表示循环和相关的硬件加速选择，如多线程、特殊硬件指令的使用和内存访问。

1. 一个例子

使用张量程序抽象，我们可以在较高层的抽象制定一些与特定硬件无关的较通用的 IR 优化（计算优化）。

比如，

对于两个大小为 128×128 的矩阵 A 和 B，我们进行如下两步的张量计算：

也就是做了一个 matmul 操作和一个 relu 操作。

使用 numpy 做这个操作，大致代码如下；

dtype = "float32"
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
# a @ b is equivalent to np.matmul(a, b)
c_mm_relu = np.maximum(a_np @ b_np, 0)

在底层，NumPy 调用库（例如 OpenBLAS）和它自己在低级 C 语言中的一些实现来执行这些计算。

类似的一个代码是这样：

def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    Y = np.empty((128, 128), dtype="float32")
    for i in range(128):
        for j in range(128):
            for k in range(128):
                if k == 0:
                    Y[i, j] = 0
                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
    for i in range(128):
        for j in range(128):
            C[i, j] = max(Y[i, j], 0)

这里相比原来的 numpy 代码做了内存分配和循环展开，比较接近底层语言做的事了。

验证一下两种实现的结果，发现是等价的：

c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)

但第二种实现更接近实际的框架在目标硬件上的调度和计算。这里的抽象展示了我们在实际实现中需要的元素：

  • 多维缓冲区（数组）。
  • 在数组维度上的循环。
  • 在循环下执行的计算语句。

2. TensorIR

接下来介绍 TensorIR，TensorIR 是 TVM 里的 Tensor 抽象，是最接近硬件的 IR。

在 TVM 中，可以用一种叫做 TVMScript DSL（Domain Specified Language，特定领域语言）来用 python 写一些 IR 优化，比如上面的例子可以写成：

@tvm.script.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")):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                vk = T.axis.reduce(128, 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 = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

对比一下用 numpy 指定的优化和用 TVMScript 的优化：

从这个程序出发我们看一下 TVMScript 的语法和他们对应的 Tensor 抽象：

  • 函数参数与缓冲区

    TVMScript 用 T.Buffer 指定数组大小：

    # TensorIR
def mm_relu(A: T.Buffer((128, 128), "float32"),
            B: T.Buffer((128, 128), "float32"),
            C: T.Buffer((128, 128), "float32")):
    ...
# numpy
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    ...

    这里 A、B、C 都是 128 × 128、元素类型为 float32 的 Tensor。制定数组大小有助于编译器对计算过程的调度和计算做更多的优化。

    计算过程中会产生中间结果 Y（matmul 之后，relu 之前），TVMScript 这里也对中间结果分配了缓冲区：

    # TensorIR
Y = T.alloc_buffer((128, 128), dtype="float32")
# numpy
Y = np.empty((128, 128), dtype="float32")

  • for 循环

    TVMScript 用 T.grid 来表示循环迭代：

    # TensorIR
for i, j, k in T.grid(128, 128, 128):
# numpy
for i in range(128):
    for j in range(128):
        for k in range(128):

  • 计算块

    CUDA 中也有 block 的概念，这里也是类似的，对于循环中的变量需要映射到一个块上做计算：

    # TensorIR
with T.block("Y"):
    vi = T.axis.spatial(128, i)
    vj = T.axis.spatial(128, j)
    vk = T.axis.reduce(128, k)
    with T.init():
        Y[vi, vj] = T.float32(0)
    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

# coressponding numpy code
vi, vj, vk = i, j, k
if vk == 0:
    Y[vi, vj] = 0
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

    这里的 matmul 映射到了一个块 Y 上得到计算结果，第 8 行给出了计算规则。

    块 是 TensorIR 中的基本计算单位。值得注意的是，该块包含比普通 NumPy 代码更多的信息。一个块包含一组块轴（vi、vj、vk）和围绕它们定义的计算。

    vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)

    上面三行声明了关于块轴的关键性质，语法如下。

    [block_axis] = T.axis.[axis_type]([axis_range], [mapped_value])

    这三行包含以下信息：

    • 定义了 vi、vj、vk 应被绑定到的位置（在本例中为 i、j 和 k）；
    • 声明了 vi、vj、vk 的原始范围（T.axis.spatial(128, i) 中的 128）；
    • 声明了块轴的属性（spatial, reduce）。

    我们一一浏览这些属性。首先，就边界关系而言，vi = T.axis.spatial(128, i) 有效地蕴含了 vi = i。[axis_range] 值提供了 [block_axis] 的预期范围。例如，vi = T.axis.spatial(128, i) 中的 128 表明 vi 应该在 range(0, 128) 中。

  • 块轴绑定的语法糖

    在每个块轴直接映射到外部循环迭代器的情况下，我们可以使用 T.axis.remap 在一行中声明所有块轴：

    # SSR means the properties of each axes are "spatial", "spatial", "reduce"
vi, vj, vk = T.axis.remap("SSR", [i, j, k])

    这里就相当于块 vi, vj, vk 分别和循环迭代器 i, j, k 绑定上了，类似 T.grid，只需要一行代码就可以完成多个绑定。

    使用这个语法之后的程序如下：

    @tvm.script.ir_module
class MyModuleWithAxisRemapSugar:
    @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({"global_symbol": "mm_relu", "tir.noalias": True})
        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))

  • 函数属性和装饰器

    @tvm.script.ir_module 表示 MyModule 是一个 IRModule。IRModule 是 TVM 中的最小编译单位（可以用 tvm.build 来编译）。

    一个 IRModule 中可以包含多个元张量函数 prim_func，比如这里 MyModuleWithTwoFunctions 有两个 prim_func：

    @tvm.script.ir_module
class MyModuleWithTwoFunctions:
    @T.prim_func
    def mm(A: T.Buffer((128, 128), "float32"),
           B: T.Buffer((128, 128), "float32"),
           Y: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "mm", "tir.noalias": True})
        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]

    @T.prim_func
    def relu(A: T.Buffer((128, 128), "float32"),
             B: T.Buffer((128, 128), "float32")):
        T.func_attr({"global_symbol": "relu", "tir.noalias": True})
        for i, j in T.grid(128, 128):
            with T.block("B"):
                vi, vj = T.axis.remap("SS", [i, j])
                B[vi, vj] = T.max(A[vi, vj], T.float32(0))

3. 变换

对于之前的 numpy 的 matmul + relu 的例子，在指定内存分配和循环展开的基础上还可以做一个等价的计算变换：

def lnumpy_mm_relu_v2(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    Y = np.empty((128, 128), dtype="float32")
    for i in range(128):
        for j0 in range(32):
            for k in range(128):
                for j1 in range(4):
                    j = j0 * 4 + j1
                    if k == 0:
                        Y[i, j] = 0
                    Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
    for i in range(128):
        for j in range(128):
            C[i, j] = max(Y[i, j], 0)

c_np = np.empty((128, 128), dtype=dtype)
lnumpy_mm_relu_v2(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)

这里对于第二层循环 0 ~ 127 做了变换，用两个循环(32× 4) j0 和 j1 替换了 j 循环；

TVMScript 中提供了 Schedule 类可以帮助我们做类似的循环变换。

这是原先的未做循环变量的 TVM 程序：

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({"global_symbol": "mm_relu", "tir.noalias": 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"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, 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, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(Y[vi, vj])
                T.writes(C[vi, vj])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

我们首先创建一个以给定的 MyModule 作为输入的 Schedule 辅助类。

sch = tvm.tir.Schedule(MyModule)

然后我们执行以下操作以获得对块 Y 和相应循环的引用。

block_Y = sch.get_block("Y", func_name="mm_relu")
i, j, k = sch.get_loops(block_Y)

我们把对 j 的循环变换成两个循环，内部循环的长度是 4：

j0, j1 = sch.split(j, factors=[None, 4])

我们可以查看存储在 sch.mod 中的变换结果。

IPython.display.Code(sch.mod.script(), language="python")

变换后的 TVM 程序：

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({"global_symbol": "mm_relu", "tir.noalias": True})
        # with T.block("root"):
        Y = T.alloc_buffer((128, 128))
        for i, j_0, j_1, k in T.grid(128, 32, 4, 128):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j_0 * 4 + j_1)
                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, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(Y[vi, vj])
                T.writes(C[vi, vj])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

变换之后，我们创建了两个新的循环，j_0 和 j_1，分别对应范围 32 和 4。下一步是重新排序这两个循环。

sch.reorder(j0, k, j1)
IPython.display.Code(sch.mod.script(), language="python")

重新排序之后的 TVM 程序：

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({"global_symbol": "mm_relu", "tir.noalias": True})
        # with T.block("root"):
        Y = T.alloc_buffer((128, 128))
        for i, j_0, k, j_1 in T.grid(128, 32, 128, 4):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j_0 * 4 + j_1)
                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, j in T.grid(128, 128):
            with T.block("C"):
                vi, vj = T.axis.remap("SS", [i, j])
                T.reads(Y[vi, vj])
                T.writes(C[vi, vj])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

可以发现这里 i, j_0, k, j_1 的迭代顺序和上面重新排序之前的 i, j_0, j_1, k 是不一样的。

这样，我们就用 TVMScript 对计算过程手动做了循环展开上的优化。

4. 编译和运行

前面我们说了 TensorIR 是 TVM 里最接近硬件的 IR，而 IRModule 是 TVM 中可以编译的最小单位。

也就是说，对于 MyModule，我们可以对它进行编译和运行。

TVM 中可以用 tvm.build 来编译指定的 IRModule，并且指定 target（目标硬件平台）：

rt_lib = tvm.build(MyModule, target="llvm")

我们创建三个用于保存输入和输出的 TVM NDArray：

a_nd = tvm.nd.array(a_np)
b_nd = tvm.nd.array(b_np)
c_nd = tvm.nd.empty((128, 128), dtype="float32")
type(c_nd)       # tvm.runtime.ndarray.NDArray

编译后的 rt_lib 可以运行，我们把输入扔进去：

func_mm_relu = rt_lib["mm_relu"]
func_mm_relu(a_nd, b_nd, c_nd)

np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)

可以发现运行结果和我们手写的 numpy 是一样的。

我们还可以构建经过变换后的程序：

rt_lib_after = tvm.build(sch.mod, target="llvm")
rt_lib_after["mm_relu"](a_nd, b_nd, c_nd)
np.testing.assert_allclose(c_mm_relu, c_nd.numpy(), rtol=1e-5)

最后，我们可以比较一下两者的时间差，看一下我们的变换的效果。 time_evaluator 是一个辅助的测试函数，可用于比较不同生成函数的运行性能：

f_timer_before = rt_lib.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of MyModule %g sec" % f_timer_before(a_nd, b_nd, c_nd).mean)
f_timer_after = rt_lib_after.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of transformed sch.mod %g sec" % f_timer_after(a_nd, b_nd, c_nd).mean)

Time cost of MyModule 0.00214587 sec
Time cost of transformed sch.mod 0.00101764 sec

看来我们的变换还是有很大的作用的，那么原因是什么呢？

如下图，现代 CPU 带有多级缓存，需要先将数据提取到缓存中，然后 CPU 才能访问它。

访问已经在缓存中的数据要快得多。CPU 采用的一种策略是获取彼此更接近的数据（局部性原理）。 当我们读取内存中的一个元素时，它会尝试将附近的元素（更正式的名称为“缓存行”）获取到缓存中。 因此，当你读取下一个元素时，它已经在缓存中。 因此，具有连续内存访问的代码通常比随机访问内存不同部分的代码更快。

现在让我们看看上面的迭代可视化，分析一下是怎么回事。 在这个分析中，让我们关注最里面的两个循环：k 和 j1。高亮的地方显示了当我们针对 k 的一个特定实例迭代 j1 时迭代触及的 Y、A 和 B 中的相应区域。

我们可以发现，j1 这一迭代产生了对 B 元素的连续访问。具体来说，它意味着在 j1=0 和 j1=1 时我们读取的值彼此相邻。这可以让我们拥有更好的缓存访问行为。此外，我们使 C 的计算更接近 Y，从而实现更好的缓存行为。

改变计算的循环迭代使之更接近 CPU 执行策略是编译器的常见优化手段之一。

除了 TensorIR，TVM 中还可以用张量表示（Tensor Expression, TE）来做张量的抽象，这里略过，在介绍 TVM 的文章中再详细介绍。