Zhangzhe's Blog

The projection of my life.

0%

机器学习编译(4)——自动程序优化

Posted on Edited on In Valine:

URL

自动程序优化的原因

  • MLC 的本质是张量函数之间的转换，但我们不知道哪种转换是让模型运行更快的，所以需要使用自动程序优化，去自动搜索最有转换。

自动程序优化过程

end-to-end 构建模型

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
# IR_Module 使用自定义主张量函数和库张量函数
@tvm.script.ir_module
class MyModuleMixture:
    @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})
        ...
    @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("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
# 注册库张量函数
@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("env.relu", override=True)
def lnumpy_relu(x: tvm.nd.NDArray, out: tvm.nd.NDArray):
    ...
# 绑定模型权重参数（nd_params 是模型权重），作用类似于 functools.partial()
MyModuleWithParams = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
# IR_Module -> 可执行程序 -> 虚拟机执行器
ex = relax.vm.build(MyModuleWithParams, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
# 执行
nd_res = vm["main"](data_nd)
# 测速
ftimer = vm.module.time_evaluator("main", tvm.cpu(), number=100)

自动优化 linear0 主张量函数

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
# 调优 API 只接受一个带有一个 main 函数的 IRModule，所以需要将原始 IRModule 中的 linear0 转成新 IRModule 的 main 函数
mod_linear = tvm.IRModule.from_expr(MyModuleMixture["linear0"].with_attr("global_symbol", "main"))
# 打印新IRModule
IPython.display.HTML(code2html(mod_linear.script()))
# 打印输出
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(
        X: T.Buffer[(1, 784), "float32"],
        W: T.Buffer[(128, 784), "float32"],
        B: T.Buffer[(128,), "float32"],
        Z: T.Buffer[(1, 128), "float32"],
    ):
        # 函数中内容是 MyModuleMixture.linear0
# 自动调优 API，input 是需要调优的 IRModule，output 是调优后的 schedule，schedule.mod 是调优后的 IRModule
sch_tuned_linear = ms.tune_tir(
    mod=mod_linear,                 # 待调优 IRModule
    target="llvm --num-cores=1",    # 调优目标
    config=ms.TuneConfig(           # 自动调优配置
      max_trials_global=64,
      num_trials_per_iter=64,
    ),
    work_dir="./tune_tmp",
    task_name="main",
)
# 将返回的 IRModule 中的 main 函数更新到原 IRModule 的 linear0 中
# 绑定参数
MyModuleWithParams2 = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
# 获取调优后的 main 函数
new_func = sch_tuned_linear.mod["main"].with_attr("global_symbol", "linear0")
# 获取原 IRModule 的 linear0 张量函数
gv = MyModuleWithParams2.get_global_var("linear0")
# 更新调优后的 main 函数到原 IRModule 的 linear0 张量函数
MyModuleWithParams2.update_func(gv, new_func)
# 重新测速，速度变快