机器学习编译总结

• TVM 是什么：是 Tensor Virtual Machine 的缩写，是一个 Open Deep Learning Compiler Stack（深度学习开源编译栈）

• TVM 想干什么：将机器学习算法从开发阶段，通过变换和优化算法，使其变成部署状态

• TVM 的原则：

  • 集成与最小化依赖
  • 利用硬件加速
  • 通用优化

• TVM Module 层次结构：

  • IRModule：包含一个或多个 元张量函数 和一个 高层神经网络执行的抽象。通常用 @tvm.script.ir_module 装饰器装饰
  • tensorIR: 元张量函数。通常表示一个算子实例的计算过程，包含多个 计算块。通常用 @T.prim_func 装饰器装饰
  • 高层神经网络执行的抽象：IRModule 的程序入口。通常用 @R.function
  • block: 计算块。张量的基本计算单位，通常包含多个 计算轴 上的循环。通常用 with T.block(block_name) 来标明作用域
  • 计算轴：
    • 空间轴（spatial axis）：空间轴上循环的每个位置的计算独立于其他位置
    • 规约轴（reduce axis）：规约轴上的位置不会反映到最后的计算输出上

• TVM Module 变换过程：

  • 自动程序优化
  • cuda 多线程优化
  • 内存优化
  • 图优化
  • 等等

• TVM Module 执行过程：

  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)
  • 虚拟机执行器 = 虚拟机(可执行程序)
  • 运行结果 = 虚拟机执行器(模型输入 + 模型权重)