提示：文章有不对的地方，麻烦大佬们指出下

---

前言

提示：本文是个人上课的一些提炼笔记，需要更全面的可以看陈天奇老师课程

课程：https://www.bilibili.com/video/BV15v4y1g7EU
课程主页： https://mlc.ai/summer22-zh
课程笔记：https://mlc.ai/zh/

---

一、什么是机器学习编译？

机器学习编译，个人简单理解就是，把机器学习模型经过一定转换，使其能以更接近底层API的形式部署在各种不同的平台上。其中值得关注的问题是集成和最小化依赖，机器学习模型依赖的包可能只占pytorch
的一小部分，比如只用到Conv2d, ReLU之类，从节省资源考虑，在部署时只打包必要的依赖。部署时也分为开发环境和部署环境，但二者很多时候是相同的。

二、张量抽象

举个例子，抽象可以把linear和relu操作合并表示为循环操作的矩阵元素操作

1.什么是算子

张量算子函数指的是张量计算过程中的最基本操作，比如加法add和减法等
，正常的加法在编译过后可能是逐元素的加法操作，但加法是可并行，可以根据底层支持优化成并行的向量加法。

# 注意运行要去掉注释，要不编译过不了
# pip install mlc-ai-nightly -f https://mlc.ai/wheels
# 张量算子计算示例
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

class MyModule:
    @T.prim_func
    def main(A: T.Buffer[128, 'float32'],
             B: T.Buffer[128, 'float32'],
             C: T.Buffer[128, 'float32']):
        T.func_attr({'global_symbol':"main", "tir.noalian":True})
        for i in range(128):
            with T.block("C"):
                vi = T.axis.spatial(128, i)
                C[vi] = A[vi] + B[vi]

sch = tvm.tir.Schedule(MyModule) #辅助类
block_c = sch.get_block("C") # 拿到上面C部分的代码
i, = sch.get_loops(block_c)
i0, i1, i2 = sch.split(i, factors=[None, 4, 4]) #一个迭代循环分解长成3个
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.noalian": True})
        # body
        # with tir.block("root")
        for i_0, i_1, i_2 in tir.grid(8, 4, 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]
    
'''

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

2.张量实践

矩阵乘法和Relu对应代码如下（示例）：

from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np
# https://tvm.apache.org/docs tvm官方文档，但课程用的好像不一样

dtype = 'float32'
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
cmm_relu = np.maximum(a_np @ b_np, 0) #矩阵乘法+ReLU的numpy实现

# 矩阵乘法的low-level numpy写法，尽量贴近C
def lnumpy_mm_relu(A:np.ndarray,
                   B:np.ndarray,
                   C:np.ndarray):
    # https://blog.csdn.net/artorias123/article/details/86527456
    # 矩阵乘法按行计算，内存优化
    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)

c_np = np.empty((128,128), dtype=dtype)
lnumpy_mm_relu(a_np, b_np, c_np)
# 检测两个矩阵的差值是否满足一定范围，不满足就报错
np.testing.assert_allclose(cmm_relu, c_np, rtol=1e-5)
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

# tvm版本，注意运行要去掉注释，要不编译过不了
@tvm.script.ir_module
class MyModule:
    @T.prim_func #元张量函数
    def main(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.noalian":True}) # 内存指针不重复？ 
        Y = T.alloc_buffer((128, 128), dtype='float32')  #模拟申请内存 
        '''
        mm 部分
        '''
        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) #每次调用block, 0..127
                #vi，vj, vk = T.axis.remap("SSR", [i, j, k]) 课程说的简化写法, SSR是spatial, spatial, reduction
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk,vj]
        '''
        relu部分 可以与上面分成两个函数
        '''
        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))
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

def lnumpy_mm_relu_v2(A:np.ndarray,
                   B:np.ndarray,
                   C:np.ndarray):
    # https://blog.csdn.net/artorias123/article/details/86527456
    # 矩阵乘法按行计算，内存优化
    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,j] * B[i,k]
    
    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(cmm_relu, c_np, rtol=1e-5)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

# mm_relu_v2版本的对应操作
sch = tvm.tir.Schedule(MyModule)
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, 4]) 

block_C = sch.get_block("C", "mm_relu")
sch.reverse_compute_at(block_C, j0) # 算子融合
print(sch.mod.script()) 

# 比较计算内存优化前后的运行时间
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')

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}")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

上述优化的意思是，由于内存读写的特性，RAM和CPU缓存的速率差很多，CPU计算时的数据是以一行行的形式从RAM放进CPU的各级(L1~L3)Cache的，也就是A[1,3]元素在缓存时，A[1]行的元素或者其左右的元素都在缓存，所以在计算的时候尽量把连续的元素放一起算。

图的意思是，矩阵乘法计算是Y = A行 x B列，正常三重循环遍历k次算一个, Y[i, j] = A[i, k] * B[k, j]， 这里B要行指针K需移动多次，但B[k, j]读进缓存时B[K]行也在了，所以可以多算几个同一行的Y[i, j]，提高效率。

持续更新中--------

提示：这里对文章进行总结：
例如：以上就是今天要讲的内容，本文仅仅简单介绍了pandas的使用，而pandas提供了大量能使我们快速便捷地处理数据的函数和方法。