机器学习编译笔记

机器学习编译

前言

• 已经有官方笔记了，但还是简略地做一下个人笔记加深理解，需要参考着看

概述

• 问题：机器学习应用和部署之间存在鸿沟
• 怎么解决：用软硬件的桥梁——编译，特别是IR
• 机器学习编译的目标
  • 集成与最小化依赖
  • 利用硬件加速
  • 通用优化
• 机器学习编译的关键要素
  • 机器学习：张量和张量函数
  • 编译：IR 抽象

张量程序抽象

张量程序

• 机器学习程序可以表示为张量的变换
• 元张量函数：机器学习模型计算中的单个不可分的单元计算 e.g. add
  • 机器学习编译需要实现灵活地转换元张量函数，方便后续优化
• 张量程序是一个表示元张量函数的有效抽象 e.g. TensorIR
  • 有效指张量程序能够被一系列有效的程序变换所改变
    • 足够灵活
    • 额外的结构能够为程序变换提供更多的信息 e.g. TVMScript 中 axis.spatial 提供并行化信息
  • 关键成分
    • 多维数组
    • 循环嵌套
    • 计算语句
• TVMScript 是 TensorIR 的一种直接实现
  • 可以理解为底层 numpy （指开空间手动加减乘除这种）的 IR ，这种 IR 看起来比 pytorch/高级numpy 要复杂，但实际上是因为它们做了太多封装，屏蔽了实现细节
  • TVMScript 也有一些针对性的语法糖

import numpy as np
import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T

# 高级 numpy
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)

# 底层 numpy
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)

# TVMScript
@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))

• 张量表达式 TE 相关 API 也可以生成 TensorIR 函数，往往是为给定的更高级别的输入服务
• 此外，算子融合也是生成 TensorIR 函数的常见方式

张量程序变换

• 机器学习编译核心工作流
• 主要是调 API 变换循环以实现一些优化，可以显示调用，但在 TVM 中是在编译阶段隐式做的！
  • 这样就把算子库的优化统一起来了
  • 和传统编译一样，不排斥进一步人工调优

端到端模型整合

• 详细介绍 TVMScript 语法和工作流，详见 notebook 和官方笔记

自动程序优化

• 介绍 TVM 中如何自动变换张量程序
• 关键思想：使用随机调度变换来指定好程序的搜索空间，在搜索空间内搜索并找到最优的调度变换
• 随机调度变换：瞎猜循环如何拆分获取搜索空间
• 优化点：多级循环转换，向量化，循环展开等

机器学习框架整合

• 机器学习模型很复杂，手写 TVMScript 太累了相当于手写IR能不累吗，还是看看前面众多的机器学习框架吧
• 途径：通过 Builder 构造 IRModule
  • 通过张量表达式 TE
  • 通过 TVM 中的另外一种计算图级 IR Relax（之前叫 relay ）
• Pytorch 计算图转化为 TVM 中的 IRModule
  • 通过 Pytorch 接口拓扑顺序迭代遍历 Pytorch 的计算图，调用 Builder 构造 IRModule
  • 实际上是 Pytorch 的计算图到 Relax 的 codegen，遍历树节点生成这套完全就是编译
• TVM 的两种 IR
  • Relax
  • TensorIR/TVMScript
• TVM 中其它的表达式
  • 深度学习框架中的计算图
  • 张量表达式
• 允许各种 IR 共同出现同时作用，IRModule 是沟通桥梁

GPU 硬件加速

• 详细介绍 TVMScript 针对 GPU 的额外优化处理，其中有几个 CUDA 优化经典例子，详见 notebook 和官方笔记
• TVM 能利用自动程序优化自动寻找这些优化

领域加速器和特殊硬件支持

• 陈天奇自己最喜欢的章节
• MLC 需要为不同的硬件平台生成不同的目标代码，这需要 MLC 实现领域加速
• 第一步，循环变换生成新的循环层级，使满足硬件张量化条件，这同时也需要符合异构计算器的内存模型 e.g. 最内层循环变为$16\times16$
• 第二步，张量化，将 TensorIR 中的块映射到硬件加速指令中去

计算图优化

• 之前的变换都是关注单元算子的变换，我们还可以通过改写计算图数据结构来优化模型
• 详细介绍 TVM 中如何实现一个 relu 和 add 算子融合的 Pass，详见 notebook 和官方笔记
• 如果接触过编译优化，那么 vistor 模式和 Pass 这一套会很熟悉

未来展望

• 跨层级优化
• 在领域加速器上自动化地实现张量化映射
• HPC 优化的更多应用

参考

• colab
• 课程中文主页