Zhangzhe's Blog

The projection of my life.

0%

机器学习编译(2)——张量程序抽象

Posted on Edited on In Valine:

URL

https://mlc.ai/zh/chapter_tensor_program/index.html

元张量函数

1.png

  • 元张量函数 表示机器学习模型计算中的单个单元计算。
    • 一个机器学习编译过程可以有选择地转换元张量函数的实现。

张量程序

2.png

  • 张量程序 是一个表示元张量函数的有效抽象。
    • 关键成分包括: 多维数组，循环嵌套，计算语句。
    • 程序变换可以被用于加速张量程序的执行。
    • 张量程序中额外的结构能够为程序变换提供更多的信息。

TensorIR: 张量程序抽象案例研究

  • TensorIR 是标准机器学习编译框架 Apache TVM 中使用的张量程序抽象。

目标

  • 使用 TensorIR 张量程序抽象 ReLU(A @ B) 张量函数。
  • 数学表示：
    • Yi,j=∑kAi,k×Bk,jY_{i,j}=\sum_k A_{i,k}\times B_{k,j}
    • Ci,j=ReLU(Yi,j)=max(Yi,j,0)C_{i,j}=ReLU(Y_{i,j})=max(Y_{i,j}, 0)

不同实现方法

使用 Numpy 实现

1
2
3
4
5
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)

使用 Low Level Numpy 实现

Low Level Numpy 是指只使用 Numpy 的数据结构而不调用 Numpy 的 API

1
2
3
4
5
6
7
8
9
10
11
12
# Use low level numpy to implement matmal ReLU
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)

使用 TensorIR 实现

TensorIR 是 TVMScript 中的一种 Python 方言

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

TensorIR 代码与 Low Level Numpy 代码对比

函数参数

1
2
3
4
5
6
7
8
# 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):
    ...

buffer

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

循环

1
2
3
4
5
6
# 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):

计算块

1
2
3
4
5
6
7
8
9
10
11
12
13
# 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]

块（Block） 是 TensorIR 中的基本计算单位。

  • 值得注意的是，对于一组固定的 vi 和 vj，计算块在 Y 的空间位置 (Y[vi, vj]) 处生成一个点值，该点值独立于 Y 中的其他位置（具有不同的vi, vj 值的位置）。我们可以称 vi、vj 为 空间轴，因为它们直接对应于块写入的缓冲区空间区域的开始。 涉及归约的轴（vk）被命名为 归约轴。
  • 空间轴上的每个点都独立于其他点。
1
2
3
4
5
6
vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
# 使用语法糖可等价写成如下形式
# SSR means the properties of each axes are "spatial", "spatial", "reduce"
vi, vj, vk = T.axis.remap("SSR", [i, j, k])

函数属性

1
T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})

其中：

  • global_symbol 对应函数名。
  • tir.noalias 是一个属性，表示所有的缓冲存储器不重叠。

装饰器

  • @tvm.script.ir_module 表示被装饰的类是一个 IRModule。
  • @T.prim_func 表示被装饰的函数是一个张量函数。