机器学习编译

一、机器学习概述

1.1 什么是机器学习编译

将机器学习算法从开发形态通过变换和优化算法使其变成部署形态。即将训练好的机器学习模型应用落地，部署在特定的系统环境之中的过程。

开发形态：开发机器学习模型时使用的形态。Pytorch,TensorFlow等通用框架编写的模型描述，以及与之相关的权重。

部署形态：指执行机器学习应用程序所需的形态。通常涉及机器学习模型的每个步骤的支持代码、管理资源的控制器以及应用程序开发环境的接口。

不用的AI应用对应的部署环境是互不相同的。

在进行机器学习模型部署时，不仅要考虑硬件系统的环境，与此同时也需要考虑软件环境(如操作系统环境、机器学习平台环境等)。

1.2 机器学习编译的目标

机器学习编译的直接目标包括两点：最小化依赖和利用硬件加速。最终目的是实现性能(时间复杂度、空间复杂度)优化。

最小化依赖可以认为是集成的一部分，提取出与应用相关的库(删除与应用无关的库)，从而减少应用的大小，达到节省空间的目的。

利用硬件加速指的是利用硬件本身的特性进行加速。可以通过构建调用原生加速库的部署代码或生成利用原生指令的代码来实现。

1.3 为什么要学习机器学习编译

构建机器学习部署的解决方案。

对现有机器学习框架形成更加深刻的理解。

1.4 机器学习编译的核心要素

张量(Tensor):是执行中最重要的元素。张量是表示神经网络模型执行的输入、输出和中间结果的多维数组。
张量函数(Tensor Function):神经网络的知识被编码在权重和接收张量和输出张量的计算序列中。将这些计算称为张量函数。注意：张量函数不需要对应于神经网络计算的单个步骤。部分计算或者整个端到端的计算也可以看作张量函数。

不仅是单个函数可以认作张量函数，也可以把其中一部分(或者整个整体)看作是张量函数。

举个栗子：

第一个linear层和relu层计算被折叠成一个linear_relu函数(算子融合)，这需要特定的linear_relu的详细实现。

def linear_relu(x,w,out):for i in range(1):for j in range(200):out[i,j]=0for k in range(3072):out[i,j]+=w[i,k]+w[j,k]out[i,j]=max(out[i,j],0)

1.5 抽象和实现

对同样的目标有不同颗粒度的表示。例如对linear_relu而言，既可以使用左边两个框图进行表示，也可以用右边的循环进行表示。

抽象：表示相同张量函数的方式。不同的抽象可能会指定一些细节，而忽略其他实现细节。例如,linear_relu可以使用另外一个不同的for循环来实现。
抽象和实现可能是所有计算机系统中最重要的关键字。抽象指定"做什么"，实现提供”如何做“。没有具体的界限。根据我们的看法，for循环本身可以被视作一种抽象，可以使用python解释器实现或编译为本地汇编代码，

机器学习编译实际上是在相同或不同抽象下转换和组装张量函数的过程。我们将研究张量函数的不同抽象类型，以及他们如何协同工作以解决机器学习部署中的挑战。

在后续文章中分别会涉及到四种类型的抽象：

计算图，张量程序，库(运行时)，硬件专用指令

二、张量程序抽象

2.1 元张量函数(张量算子)

元张量函数指的是单个单元(最小颗粒度)的张量函数。

以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语言，甚至是使用汇编语言)来进行更高效的实现。

许多框架提供的一个最常见的MLC过程是将基本函数的实现（或在运行时分派它们）转换为基于环境的更优化的实现。

例如当前环境可以四个四个读取程序数据，所以可以使用并行化来优化当前代码。

2.1.1 张量算子变换

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

2.2 张量程序抽象

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

下图是TVM中的一个样例：

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

1.多维的输入、输出的数据。

2.循环

3.每个循环中进行相应的计算。

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

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

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

2.2.2 常见张量变换

1.循环拆分

其中等价替换包括：

x=xo*4+xi

这样简化并行化处理，提高了内存访问效率，便于向量化操作，简化编程模型。

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

这里的xi需要依赖于xi-1所以不能进行线程绑定。

2.2.3 张量程序抽象中的其他结构

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

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

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

2.3 张量程序变换实践

建议直接使用Colab，可以直接使用pip安装包来安装tvm

1.配置环境

由于 https://mlc.ai/wheels需要代理才能连通。

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

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_funcdef main(A:T.Buffer[128,"float32"],B:T.Buffer[128,"float32"],C:T.Buffer[128,"float32"]):T.func_attr({"global_symbol":"main","tir.noalias":True})for i in range(128):with T.block("C"):vi=T.axis.spatial(128,i)C[vi]=A[vi]+B[vi]

TvmScript是一种我们能以python抽象语法树的形式来表示张量程序的方式。与python语法所对应，并在该语法的基础上增加了能够帮助程序分析与变换的额外结构。

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

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

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

print(MyModule.script())

3.编译与运行

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

rt_mod=tvm.build(MyModule,target="llvm")
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.张量程序变换

一个张量程序可以通过一个辅助的名为调度(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)。后者是通过三重循环来表示一重循环。

你可以通过这幅图片来了解这个拆分关系

# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:@T.prim_funcdef main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32"), C: T.Buffer((128,), "float32")):T.func_attr({"tir.noalias": T.bool(True)})# with T.block("root"):for i_0, i_1, i_2 in T.grid(8, 4, 4):with T.block("C"):vi = T.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)T.reads(A[vi], B[vi])T.writes(C[vi])C[vi] = A[vi] + B[vi]

循环重排

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

# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:@T.prim_funcdef main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32"), C: T.Buffer((128,), "float32")):T.func_attr({"tir.noalias": T.bool(True)})# with T.block("root"):for i_0, i_2, i_1 in T.grid(8, 4, 4):with T.block("C"):vi = T.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)T.reads(A[vi], B[vi])T.writes(C[vi])C[vi] = A[vi] + B[vi]

并行化

对最外层循环进行并行操作

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

# from tvm.script import ir as I
# from tvm.script import tir as T

@I.ir_module
class Module:@T.prim_funcdef main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32"), C: T.Buffer((128,), "float32")):T.func_attr({"tir.noalias": T.bool(True)})# with T.block("root"):for i_0 in T.parallel(8):for i_2, i_1 in T.grid(4, 4):with T.block("C"):vi = T.axis.spatial(128, i_0 * 16 + i_1 * 4 + i_2)T.reads(A[vi], B[vi])T.writes(C[vi])C[vi] = A[vi] + B[vi]

一下是张量程序创建->编译->运行

三、张量程序案例 TensorIR

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

3.1 举个例子

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

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

即一个matmul操作和一个relu操作。

使用numpy

dtype="float32"
a=np.random.rand(128,128).astype(dtype)
b=np.random.rand(128,128).astype(dtype)
c_mm_relu=np.maximum(a@b,0)dtype="float32"
a=np.random.rand(128,128).astype(dtype)
b=np.random.rand(128,128).astype(dtype)
c_mm_relu=np.maximum(a@b,0)

在底层，Numpy调用库和它自己在低级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):Y[i, j] = 0  # 初始化 Y[i, j] 为 0for k in range(128):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)

主程序对比

# 主程序
dtype = "float32"
a = np.random.rand(128, 128).astype(dtype)
b = np.random.rand(128, 128).astype(dtype)

# 使用 NumPy 的矩阵乘法和 ReLU 激活计算正确结果
c_mm_relu = np.maximum(a @ b, 0)

# 初始化输出矩阵 c
c = np.empty((128, 128), dtype=dtype)

# 使用 lnumpy_mm_relu 计算结果
lnumpy_mm_relu(a, b, c)

# 使用 np.testing.assert_allclose 验证结果是否一致
np.testing.assert_allclose(c_mm_relu, c, rtol=1e-5, atol=1e-8)
print("Validation passed: The results are consistent.")

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

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

3.2 Tensor IR

Tensor IR是TVM里的Tensor抽象，是最接近硬件的IR

在TVM中可以用一种叫做TVMScript DSL来用python写一些IR优化，比如上面的例子可以写成：

import numpy as np
from tvm.script import tir as T
import tvm

@tvm.script.ir_module
class Module:@T.prim_funcdef 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))
import numpy as np
from tvm.script import tir as T
import tvm@tvm.script.ir_module
class Module:@T.prim_funcdef 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指定优化对比：

计算块：

TensorIR 是 TVM 中用于表示张量计算的低级中间表示，它不仅描述了计算的数据流，还描述了计算之间的依赖关系。这些依赖关系对于理解数据如何在计算图中流动以及如何正确地执行计算至关重要。在 TensorIR 中，数据依赖关系体现在各个计算块（block）之间的相互作用上。

对于循环中的变量需要映射到一个块上做计算。

块是TensorIR中的基本计算单位，该块比numpy包含的信息多。一个块包含一组块轴(vi、vj、vk)和围绕他们定义的计算。

块轴绑定的语法糖：在每个块轴直接映射到外部循环迭代器的情况下我们可以使用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，只需要一行代码就可以完成多个绑定。

于是上述代码可以写成

3.3 变换

3.3.1 循环拆分

sch=tvm.tir.Schedule(Module)
block_T=sch.get_block("matmul",func_name="mm_relu")
i,j,k=sch.get_loops(block_T)
j0,j1=sch.split(j,factors=[None,4])
print(sch.mod.script())

3.3.2 循环重排

sch.reorder(j0,k,j1)
print(sch.mod.script())

3.3.3 算子融合

将 ReLU 块（block_C）的计算移动到矩阵乘法块（block_T）的 j0 循环中。

作用：

(1)减少内存访问：将Relu的计算移动到矩阵乘法的内层循环中，可以减少对中间缓冲区Y的访问次数，从而提高内存访问效率。

(2)提高计算效率：减少了循环嵌套的层数，降低了循环开销。可以利用CPU或GPU的缓存，减少数据在内存和缓存之间的传输。

3.3.4 规约分解

将规约操作(如矩阵乘法中的累加)分解为多个阶段，从而优化计算过程。

sch.decompose_reduction(block_T,k)
print(sch.mod.script())

3.3.5 优化对比

import numpy as np
# 主程序
dtype = "float32"
a = np.random.rand(128, 128).astype(dtype)
b = np.random.rand(128, 128).astype(dtype)

# 使用 NumPy 的矩阵乘法和 ReLU 激活计算正确结果
c_mm_relu = np.maximum(a @ b, 0)

mod=tvm.build(Module,target="llvm")
func=mod["mm_relu"]

# 初始化输出矩阵 c
c = np.empty((128, 128), dtype="float32")

a_tvm = tvm.nd.array(a)
b_tvm = tvm.nd.array(b)
c_tvm = tvm.nd.array(c)
# 使用 lnumpy_mm_relu 计算结果
func(a_tvm, b_tvm, c_tvm)

np.testing.assert_allclose(c_mm_relu, c_tvm.numpy(), rtol=1e-5, atol=1e-8)
print("Validation passed: The results are consistent.")

Validation passed: The results are consistent.

对比优化前后时间差,帮助评估调度优化(如循环分块、规约分解等)对性能的影响。

# 比较计算内存优化前后的运行时间
rt_lib = tvm.build(Module, target='llvm')
a_nd = tvm.nd.array(a)
b_nd = tvm.nd.array(b)
c_nd = tvm.nd.empty((128, 128), dtype='float32')rt_lib_after = tvm.build(sch.mod, target='llvm')
f_time_before = rt_lib.time_evaluator("mm_relu", tvm.cpu())
print(f"Time cost of Mymodule:{f_time_before(a_nd, b_nd, c_nd).mean}")
f_time_after =  rt_lib_after.time_evaluator("mm_relu", tvm.cpu())
print(f"Time cost after transform:{f_time_after(a_nd, b_nd, c_nd).mean}")