Zhangzhe's Blog

The projection of my life.

0%

机器学习编译(3)——端到端模型执行

Posted on Edited on In Valine:

URL

机器学习编译的本质与关注点

  • MLC 的本质：张量函数之间的转换
  • MLC 的关注点：
    • 所以可能的张量函数抽象表达
    • 所有可能的张量函数转换

构造 IR_Module

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
@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def relu0(X: T.Buffer[(1, 128), "float32"], Y: T.Buffer[(1, 128), "float32"]):
        # function attr dict
        T.func_attr({"global_symbol": "relu0", "tir.noalias": True})
        for i, j in T.grid(1, 128):
            with T.block("Y"):
                vi, vj = T.axis.remap("SS", [i, j])
                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))
    @T.prim_func
    def linear0(
        X: T.Buffer[(1, 784), "float32"],
        W: T.Buffer[(128, 784), "float32"],
        B: T.Buffer[(128,), "float32"],
        Z: T.Buffer[(1, 128), "float32"],
    ):
        T.func_attr({"global_symbol": "linear0", "tir.noalias": True})
        Y = T.alloc_buffer((1, 128), "float32")
        for i, j, k in T.grid(1, 128, 784):
            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] + X[vi, vk] * W[vj, vk]
        for i, j in T.grid(1, 128):
            with T.block("Z"):
                vi, vj = T.axis.remap("SS", [i, j])
                Z[vi, vj] = Y[vi, vj] + B[vj]
    @T.prim_func
    def linear1(
        X: T.Buffer[(1, 128), "float32"],
        W: T.Buffer[(10, 128), "float32"],
        B: T.Buffer[(10,), "float32"],
        Z: T.Buffer[(1, 10), "float32"],
    ):
        T.func_attr({"global_symbol": "linear1", "tir.noalias": True})
        Y = T.alloc_buffer((1, 10), "float32")
        for i, j, k in T.grid(1, 10, 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] + X[vi, vk] * W[vj, vk]
        for i, j in T.grid(1, 10):
            with T.block("Z"):
                vi, vj = T.axis.remap("SS", [i, j])
                Z[vi, vj] = Y[vi, vj] + B[vj]
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        with R.dataflow():
            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir(relu0, (lv0,), (1, 128), dtype="float32")
            out = R.call_tir(linear1, (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out
  • @tvm.script.ir_module 装饰 IR_Module
  • @T.prim_func 装饰 主张量函数
  • @R.function 装饰 高层神经网络执行的抽象，（将整个 IR_Module 中的主张量函数串起来组成一个计算图）
  • R.dataflow() 用于标记程序计算图区域
  • R.call_tir(prim_func, inputs, output_shape, dtype) 分配输出内存并和输入一起输入主张量函数

构建并运行模型

1
2
3
4
5
ex = relax.vm.build(MyModule, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
nd_res = vm["main"](
    data_nd, nd_params["w0"], nd_params["b0"], nd_params["w1"], nd_params["b1"]
)
  • 可执行文件 = build(IR_Module)
  • 虚拟机执行器 = 虚拟机(可执行文件)
  • 运行结果 = 虚拟机执行器(模型输入)

使用现有库避免重复造轮子

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
@tvm.script.ir_module
class MyModuleWithExternCall:
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        # block 0
        with R.dataflow():
            lv0 = R.call_tir("env.linear", (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir("env.relu", (lv0,), (1, 128), dtype="float32")
            out = R.call_tir("env.linear", (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out
  • 上图中的 env.linear 是库张量函数，同一个 IR_Module 中可使用库张量函数，也可使用自定义张量函数，也可以二者混用。
1
2
3
@tvm.register_func("env.linear", override=True)
def torch_linear(x: tvm.nd.NDArray, w: tvm.nd.NDArray, b: tvm.nd.NDArray, out: tvm.nd.NDArray):
    ...
  • @tvm.register_func，注册库张量函数

绑定参数到 IR_Module

  • 以上的所有 IR_Module 都是在调用时才传入 数据 和 权重参数，但 权重参数 是不变的，所以可以将 权重参数 提前绑定到 IR_Module 中，调用时只传入输入数据。
1
MyModuleWithParams = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
  • relax.transform.BindParams(计算图入口函数，模型参数)(IR_Module) 将模型参数绑定到 IR_Module 的计算图入口函数中，返回一个绑定好模型参数的 IR_Module。

总结

  • 计算图抽象（一般表示 main 函数） 有助于将元张量函数拼接在一起以进行端到端执行。
  • Relax 抽象的关键要素包括
    • call_tir 构造，将目标传递规范的元函数嵌入到计算图中
    • Dataflow block
  • 计算图允许调用环境库函数和 TensorIR 函数。