机器学习编译笔记

机器学习编译

前言

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

概述

  • 问题:机器学习应用和部署之间存在鸿沟
  • 怎么解决:用软硬件的桥梁——编译,特别是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 中是在编译阶段隐式做的!
    • 这样就把算子库的优化统一起来了
    • 和传统编译一样,不排斥进一步人工调优

端到端模型整合

自动程序优化

  • 介绍 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 优化的更多应用

参考


机器学习编译笔记
http://example.com/2022/09/25/机器学习编译笔记/
作者
zty
发布于
2022年9月25日
许可协议