TensorIR_张量程序抽象案例研究

!pip install apache-tvm

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

• \(Y_{ij} = \sum_kA_{ik}B_{kj}\)
• \(C_{ij} = \mathbb{relu}(Y_{ij}) = \mathbb{max}(Y_{ij}, 0)\)

dtype = "float32"
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)

%%time
# a @ b is equivalent to np.matmul(a, b)
c_mm_relu = np.maximum(a_np @ b_np, 0)

我们使用low-level Numpy

• 尽量不使用numpy提供的函数，使用循环
• 尽量显式开对应的数组

def llnp_mm_relu(A: np.array, B: np.array, C: np.array):
    Y = np.empty([128, 128], dtype="float32")
    # 外层循环遍历矩阵A的行
    for i in range(128):
        # 中层循环遍历矩阵B的列
        for j in range(128):
            # 内层循环遍历矩阵A的列（同时也是矩阵B的行）
            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])

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

这里的代码是用一种名为 TVMScript 的语言实现的，它是一种嵌入在 Python AST 中的特定领域方言。

@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")):
        # global_symbol: 最后的函数名; noalias: 相当于Cpp的restrict
        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))

块轴绑定的语法糖

@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"):
                # 块轴绑定的语法糖 SSR: Sepatial Sepatial Reduce
                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"):
                # 块轴绑定的语法糖 SS: Sepatial Sepatial
                vi, vj = T.axis.remap("SS", [i, j])
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

type(MyModule)

type(MyModule["mm_relu"])

机器学习编译过程中的一个 IRModule 可以包含多个张量函数

@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))

元张量函数的变换

def llnp_mm_relu_v2(A: np.array, B: np.array, C: np.array):
    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)
llnp_mm_relu_v2(a_np, b_np, c_np)
np.testing.assert_allclose(c_mm_relu, c_np, rtol=1e-5)

import IPython

def code2html(code):
    """
    Helper function to use pygments to turn the code string into hightlight huml.
    """
    import pygments
    from pygments.lexers import Python3Lexer
    from pygments.formatters import HtmlFormatter
    formatter = HtmlFormatter(style='monokai')
    html = pygments.highlight(code, Python3Lexer(), formatter)
    return "<style>%s</style>%s\n" % (formatter.get_style_defs(".highlight"), html)

IPython.display.HTML(code2html(MyModule.script()))

我们首先创建一个以给定的 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。

请注意，变换是程序性的，因此如果你不小心执行了两次该代码块，我们将得到“变量 j 不再存在”的错误。如果发生这种情况，你可以从头（创建 sch 的位置）开始再次运行

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

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

IPython.display.HTML(code2html(MyModule.script()))

IPython.display.HTML(code2html(sch.mod.script()))

sch.reorder(j0, k, j1)
IPython.display.HTML(code2html(sch.mod.script()))

我们使用名为 reverse_compute_at 的原语将块 C 移动到 Y 的内循环里

block_C = sch.get_block("C", "mm_relu")
sch.reverse_compute_at(block_C, j0)
IPython.display.HTML(code2html(sch.mod.script()))

到目前为止，我们将归约初始化和更新放在一个块体中。这种组合形式为循环变换带来了便利（因为初始化和更新的外循环 i、j 通常需要彼此保持同步）

在循环变换之后，我们可以将 Y 元素的初始化与归约更新分开。我们可以通过 decompose_reduction 原语来做到这一点。（注意：这也是 TVM 在以后编译的时候隐式做的，所以这一步的主要目的是让它显式，看看最终效果）

sch.decompose_reduction(block_Y, k)
IPython.display.HTML(code2html(sch.mod.script()))

最终变换后的代码类似于以下低级 NumPy 代码

def llnp_mm_relu_v3(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):
            # Y_init
            for j1 in range(4):
                j = j0 * 4 + j1
                Y[i, j] = 0
            # Y_update
            for k in range(128):
                for j1 in range(4):
                    j = j0 * 4 + j1
                    Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
            # C
            for j1 in range(4):
                j = j0 * 4 + j1
                C[i, j] = max(Y[i, j], 0)

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

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

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)

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)

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)

IPython.display.HTML(code2html(MyModule.script()))

IPython.display.HTML(code2html(sch.mod.script()))

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

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

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

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

我们当前的示例主要是为了证明不同的代码变体可以导致不同的性能。更多的变换步骤可以帮助我们获得更好的性能，我们将在以后的章节中介绍。本练习的主要目标是首先让我们获得程序变换工具，并首先体验通过变换可能实现的功能。

练习

作为练习，尝试不同的 j_factor 选择，看看它们如何影响代码的性能

def transform(mod, jfactor):
    sch = tvm.tir.Schedule(mod)
    block_Y = sch.get_block("Y", func_name="mm_relu")
    i, j, k = sch.get_loops(block_Y)
    j0, j1 = sch.split(j, factors=[None, jfactor])
    sch.reorder(j0, k, j1)
    block_C = sch.get_block("C", "mm_relu")
    sch.reverse_compute_at(block_C, j0)
    return sch.mod

mod_transformed = transform(MyModule, jfactor=8)

rt_lib_transformed = tvm.build(mod_transformed, "llvm")
f_timer_transformed = rt_lib_transformed.time_evaluator("mm_relu", tvm.cpu())
print("Time cost of transformed mod_transformed %g sec" % f_timer_transformed(a_nd, b_nd, c_nd).mean)
# display the code below
IPython.display.HTML(code2html(mod_transformed.script()))

创建 TensorIR 并与之交互的方法

通过 TVMScript 创建 TensorIR

获取 TensorIR 函数的第一种方法是直接在 TVMScript 中编写函数，这也是我们在上一节中使用的方法。 TVMScript 还允许我们在必要时跳过某些信息部分。 例如，T.axis.remap 使我们能够缩短迭代器大小注释。

TVMScript 也是一种在变换过程中检查张量函数的有用方法。在某些情况下，一种很有帮助的做法是打印出 TVMScript，进行一些手动编辑，然后将其反馈给机器学习编译流程以调试和尝试可能的（手动）变换，然后将变换后的程序重新应用到 MLC 流程中。
使用张量表达式生成 TensorIR 代码

在许多情况下，我们的开发形式是不在循环级别的更高级别的抽象。 所以另一种常见的获取 TensorIR 的方式是务实地生成相关代码。

张量表达式 (TE) 是一种特定领域的语言，它通过 API 之类的表达式描述一系列计算。

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 函数描述了对于给定的索引 (i, j) 我们要如何计算元素 Y[i, j] 的值。

lambda i, j: te.sum(A[i, k] * B[k, j], axis=k)

上面的 lambda 表达式描述了计算 \(Y_{ij} = \sum_kA_{ik}B_{kj}\)。在描述计算之后，我们可以通过传递我们感兴趣的相关参数来创建一个 TensorIR 函数。在这种特殊情况下，我们想要创建一个具有两个输入参数（A，B）和一个输出参数（C）的函数。

te_func = te.create_prim_func([A, B, C]).with_attr({"global_symbol": "mm_relu"})
MyModuleFromTE = tvm.IRModule({"mm_relu": te_func})
IPython.display.HTML(code2html(MyModuleFromTE.script()))

张量表达式 API 作为一个有用的工具，帮助为给定的更高级别的输入生成 TensorIR 函数。

讨论

我们了解到，一个常见的机器学习编译过程遵循一系列程序变换。将 TensorIR 变换过程与低级 NumPy 参考开发过程进行比较是很有趣的。

上图显示了标准的开发过程。我们需要重复开发不同程序变体的过程，然后（如果它是编译语言，则构建）在感兴趣的平台上运行它们。

机器学习编译流程（如下图所示）的主要区别在于 IRModule（程序）之间的程序变换。所以我们不仅可以通过开发（通过手动编写代码或生成代码）提出程序变体，还可以通过变换张量程序来获得变体。

变换是一个非常强大的工具，可以帮助我们简化开发成本并为流程引入更多自动化。本节介绍了通过 TensorIR 对元张量函数的特定视角，我们将在未来涵盖更多视角。

总结

• TensorIR 抽象
• 包含循环、多维缓冲区等常用元素
  • 引入了一个封装循环计算要求的新结构块。
  • 可以在 Python AST 中构建（通过 TVMScript）
• 我们可以使用变换来创建不同的 TensorIR 变体。
• 通用 MLC 流程：开发、变换、构建。