机器学习编译入门课程学习笔记第二讲 张量程序抽象

  本节课的slides链接如下：https://mlc.ai/summer22-zh/slides/2-TensorProgram.pdf；notes链接如下：https://mlc.ai/zh/chapter_tensor_program/。

1. 本节课内容大纲

1. 元张量函数(张量算子)
2. 张量程序抽象

  通过本节课的学习让我对批处理有了更深的理解。批处理和小批量梯度下降有异曲同工之妙(前提是计算具有独立性)。将在后文中进行更加详细的介绍。

2. 元张量函数

  元张量函数指的是单个单元(最小颗粒度)的张量函数。以PyTorch中的加法为例：

import torch
C = torch.empty((128,), dtype=torch.float32)
a = torch.tensor(np.arange(128, dtype="float32"))
b = torch.tensor(np.ones(128, dtype="float32"))

torch.add(a，b，out=c) # torch.add can be viewed as a primitive tensor function

  那么torch.add是如何实现的呢？或者在实际部署中，如何将不同的步骤(多个元张量函数)合并在一起？

  抽象与实现是本课程的核心概念之一。对于同一个张量算子而言，可以使用不同的方式来实现：

  第二种表示是对第一种表示的细化，第三种表示是使用低级语言(C语言，甚至是使用汇编语言)来进行更高效的实现。

  One most common MLC process that many frameworks offer is to transform the implementations of primitive functions(or dispatch them in runtime) to more optimized ones based on the environment.

  张量算子变换中比较常见的一种做法是直接把它映射到对应的算子库上。比如在CUDA环境中，就可以使用cudaAdd。另外一种比较常见的是更细粒度的程序变换。前者直接用符号表示即可，而后者需要表示成循环的形式。

3. 张量程序抽象

  上文提到的张量算子变换包括两种常见的形式，一种是直接映射到对应的算子库上，另一种是进行更细粒度的程序变换。将后者称之为是张量函数抽象。

  下图是TVM中的一个样例：

  张量函数抽象包括三大关键要素：

1. 多维的输入、输出的数据。
2. 循环。
3. 在每个循环中，进行相应的计算。
   

  涉及到一个问题：即为什么我们不直接使用C语言等低级语言进行编写程序。主要在于机器学习编译的核心不仅在于程序猿能够完成代码的编写，更重要的是完成张量算子变换，即给定一个张量将它从原始的形态变换到更加优化的形态(尝试不同的变换，最终找到一个最为优化的形态)。如果可以通过程序来探索可能的张量函数对应的等价空间，就能够大大提升效率。

3.1 为什么需要进行张量程序抽象

Initial state

Program-based transformnations

Transformed Program

  在Program-based transformnations中进行了循环拆分，本质上是将单层循环变换为双层循环(循环拆分是是比较常见的一种变换)，而Transformed program涉及到了并行化与向量化的操作，其中向量化在这里指的是每一步都以长度为4的单元(批处理)来进行运算。

3.2 常见的张量变换

3.2.1 循环拆分

其中等价替换包括：$$x=x_o\ast 4+x_i$$

3.2.2 循环重排

  循环重排的本质是将两个循环的顺序进行更换。

3.2.3 线程绑定

  如果在GPU上进行并行化，可以把xi和xo对应到GPU的两个线程上。通过绑定GPU的线程得到最终的结果(后续章节中会进行更加深入的讲解)。

  在不同的硬件环境上尝试不同的变换组合，然后根据测试在张量函数对应的等价空间中得到最优的变换。需要注意的一点是，在这里进行的变换必须是等价变换。下面举个反例，本质上可以表示成$C=A+C$，C会对其他index下的C的结果存在依赖(在xo = 4, xi = 0的情况会对 xo = 0, xi = 1的结果造成依赖， 而如果我们改变了循环顺序之后 xo = 0, xi = 1会先被循环到)，如果只是单纯的reorder，循环的顺序会有变化导致依赖关系被打破，会出现等式左边的C计算时等式右边的C还没有结果的情况。所以在这里不能进行循环重排。

for xo in range(32):
    for xi in range(4):
        C[xo * 4 + xi] = A[xo * 4 + xi] + C[max(xo * 4 + xi – 1, 0)]

3.3 张量程序抽象中的其它结构

  根据上文所述，我们不能任意地对程序进行变换(比如部分计算会依赖于循环之间的顺序)。但幸运的是，我们所感兴趣的大多数元张量函数都具有良好的属性（例如循环迭代之间的独立性）。

  张量程序可以将这些额外的信息合并为程序的一部分，以使程序变换更加便利。

  其中vi = T.axis.spatial(128, i)不仅等价于vi=i，并且表明vi表示了一个迭代器，并且迭代器中所有的迭代都是独立的。这个信息对于执行这个程序而言并非必要，但会使得我们在变换这个程序时更加方便。在这个例子中，我们知道我们可以安全地并行或者重新排序所有与 vi 有关的循环，只要实际执行中 vi 的值按照从 0 到 128 的顺序变化。
23:20

4. 张量程序变换实践

  首先需要说明的是：如果已经有Colab账号，更建议直接使用Colab环境，那么可直接使用下文中的pip安装包。

4.1 安装包

  官方安装命令如下所示，但由于https://mlc.ai/wheels需要代理才能连通，如果不是Colab环境则不建议使用下列命令进行安装：

python3 -m  pip install mlc-ai-nightly -f https://mlc.ai/wheels

  比较推荐的方法是手动下载适合环境的wheel文件，然后进行安装：

  个人环境为Linux+Python3.8环境，对应文件的下载链接为：https://github.com/mlc-ai/utils/releases/download/v0.9.dev0/mlc_ai_nightly-0.9.dev1661%2Bg2b34ced6c-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl。

4.2 构造张量程序

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

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def main(A: T.Buffer[128, "float32"],
             B: T.Buffer[128, "float32"],
             C: T.Buffer[128, "float32"]):
        # extra annotations for the function
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i in range(128):
            with T.block("C"):
                # declare a data parallel iterator on spatial domain
                vi = T.axis.spatial(128, i)
                C[vi] = A[vi] + B[vi]

  TVMScript 是一种让我们能以Python抽象语法树的形式来表示张量程序的方式。注意到这段代码并不实际对应一个 Python 程序，而是对应一个机器学习编译过程中的张量程序。TVMScript 的语言设计是为了与 Python 语法所对应，并在 Python 语法的基础上增加了能够帮助程序分析与变换的额外结构。

print(type(MyModule)) # tvm.ir.module.IRModule

  MyModule 是 IRModule 数据结构的一个实例，是一组张量函数的集合。

  我们可以通过 script 函数得到这个 IRModule 的 TVMScript 表示。这个函数对于在一步步程序变换间检查 IRModule 而言非常有帮助。

print(MyModule.script())

@tvm.script.ir_module
class Module:
    @tir.prim_func
    def func(A: tir.Buffer[128, "float32"], B: tir.Buffer[128, "float32"], C: tir.Buffer[128, "float32"]) -> None:
        # function attr dict
        tir.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with tir.block("root")
        for i in tir.serial(128):
            with tir.block("C"):
                vi = tir.axis.spatial(128, i)
                tir.reads(A[vi], B[vi])
                tir.writes(C[vi])
                C[vi] = A[vi] + B[vi]

4.3 编译与运行

  使用build函数将一个 IRModule 转化为可以执行的函数，其中llvm表示CPU作为硬件执行环境。

rt_mod = tvm.build(MyModule, target="llvm")  # The module for CPU backends.
type(rt_mod)

tvm.driver.build_module.OperatorModule

在编译后，mod 包含了一组可以执行的函数。我们可以通过这些函数的名字拿到对应的函数。

func = rt_mod["main"]
func

<tvm.runtime.packed_func.PackedFunc>

先构建三个张量，其中c用来接收输出：

a = tvm.nd.array(np.arange(128, dtype="float32"))
b = tvm.nd.array(np.ones(128, dtype="float32"))
c = tvm.nd.empty((128,), dtype="float32")

func(a, b, c)
print(c)

[  1.   2.   3.   4.   5.   6.   7.   8.   9.  10.  11.  12.  13.  14.
  15.  16.  17.  18.  19.  20.  21.  22.  23.  24.  25.  26.  27.  28.
  29.  30.  31.  32.  33.  34.  35.  36.  37.  38.  39.  40.  41.  42.
  43.  44.  45.  46.  47.  48.  49.  50.  51.  52.  53.  54.  55.  56.
  57.  58.  59.  60.  61.  62.  63.  64.  65.  66.  67.  68.  69.  70.
  71.  72.  73.  74.  75.  76.  77.  78.  79.  80.  81.  82.  83.  84.
  85.  86.  87.  88.  89.  90.  91.  92.  93.  94.  95.  96.  97.  98.
  99. 100. 101. 102. 103. 104. 105. 106. 107. 108. 109. 110. 111. 112.
 113. 114. 115. 116. 117. 118. 119. 120. 121. 122. 123. 124. 125. 126.
 127. 128.]

4.4 张量程序变换

  现在我们开始变换张量程序。一个张量程序可以通过一个辅助的名为调度(schedule)的数据结构得到变换。

sch = tvm.tir.Schedule(MyModule)
type(sch)

tvm.tir.schedule.schedule.Schedule

我们首先尝试拆分循环。

# Get block by its name
block_c = sch.get_block("C")
# Get loops surrounding the block,拿到block之外对应的循环
(i,) = sch.get_loops(block_c)
# Tile the loop nesting.
i_0, i_1, i_2 = sch.split(i, factors=[None, 4, 4])
print(sch.mod.script())

  将for i in tir.serial(128)中的i变换成了i_0, i_1和i_2。另外需要注意的是factors=[None, 4, 4]和tir.grid(8, 4, 4)中之间一一的对应关系以及vi等价于tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)。后者是通过三重循环来表示一重循环。

@tvm.script.ir_module
class Module:
    @tir.prim_func
    def func(A: tir.Buffer[128, "float32"], B: tir.Buffer[128, "float32"], C: tir.Buffer[128, "float32"]) -> None:
        # function attr dict
        tir.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with tir.block("root")
        for i_0, i_1, i_2 in tir.grid(8, 4, 4):
            with tir.block("C"):
                vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)
                tir.reads(A[vi], B[vi])
                tir.writes(C[vi])
                C[vi] = A[vi] + B[vi]

  另外可以对循环进行重排，比如将 i_2 移动到 i_1 的外侧。

sch.reorder(i_0, i_2, i_1)
print(sch.mod.script())

@tvm.script.ir_module
class Module:
    @tir.prim_func
    def func(A: tir.Buffer[128, "float32"], B: tir.Buffer[128, "float32"], C: tir.Buffer[128, "float32"]) -> None:
        # function attr dict
        tir.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with tir.block("root")
        for i_0, i_2, i_1 in tir.grid(8, 4, 4):
            with tir.block("C"):
                vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)
                tir.reads(A[vi], B[vi])
                tir.writes(C[vi])
                C[vi] = A[vi] + B[vi]

  另外比较常见的变换是并行化，比如对最外层的循环进行并行操作。

sch.parallel(i_0)
print(sch.mod.script())

@tvm.script.ir_module
class Module:
    @tir.prim_func
    def func(A: tir.Buffer[128, "float32"], B: tir.Buffer[128, "float32"], C: tir.Buffer[128, "float32"]) -> None:
        # function attr dict
        tir.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with tir.block("root")
        for i_0 in tir.parallel(8):
            for i_2, i_1 in tir.grid(4, 4):
                with tir.block("C"):
                    vi = tir.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)
                    tir.reads(A[vi], B[vi])
                    tir.writes(C[vi])
                    C[vi] = A[vi] + B[vi]

schedule

build

get function

run

IRModule

IRModule_new

build_module

function

result

5. 总结

2.4. 总结

• 元张量函数表示机器学习模型计算中的单个单元计算。

  • 一个机器学习编译过程可以有选择地转换元张量函数的实现。

• 张量程序是一个表示元张量函数的有效抽象。

  • 关键成分包括：多维数组、循环嵌套、计算语句。
  • 程序变换可以被用于加速张量程序的执行。
  • 张量程序中额外的结构能够为程序变换提供更多的信息。