基于推导的人工智能编译器-v3-翟季东

1 . 第十届开源操作系统年度技术会议-OS2ATC 基于表达式推导的人工智能编译器 翟季冬 清华大学 2023.3.26 PACMAN pacman.cs.tsinghua.edu.cn 清华大学 Tsinghua University 

2 . 背景和意义  人工智能芯片成为深度学习应用的主要计算力  摩尔定律发展逐渐放缓  深度学习对计算巨大需求  为了降低编程复杂度，深度学习编程框架及人工智能编译器应运而生 编程框架 人工智能芯片 ? 人工智能编译器 Caffe 清华大学 Tsinghua University 2 

3 . 研究背景和意义 为了提高人工智能应用的优化与部署效 率，现有的张量编译器将优化过程解耦  深度学习模型的编译过程 为高层次的图优化和低层次的算子优化 体系结构无关 的图层优化 (图融合、替换、 中间数据流图 删冗) 深度学习代码 数学算子库， 人工智能芯片 张量编译等 体系结构相关 的算子优化 清华大学 Tsinghua University 3 

4 . 研究现状-1：图层优化  传统的优化框架以基于规则的图优化为主，如 TensorFlow中的 Grappler， 需要领域专家对优化规则进行人工设计  为了提高优化效率，以 TASO、PET 为代表的自动优化框架成为当前图优化领 域发展的主流趋势 优点 缺点  可自动地对图优化空间进行搜索  优化结果只能包含算子库中的算子  利用形式化验证保证正确性  需尝试大量错误状态，优化时间长 清华大学 Tsinghua University 4 

5 . 研究现状-2：算子层优化  以 cuDNN、cuBLAS、oneDNN 等为代表的人工智能算子库在特定硬件平台上提 供极高的性能，但需要大量人力进行维护  近年来以 TVM、Ansor 为代表的算子自动生成框架成为了发展趋势 优点 a = tvm.nd.array(numpy.random.rand(M, K) .astype(dtype), ctx)  基于原语的编程接口相对简单 b = tvm.nd.array(numpy.random.rand(K, N) .astype(dtype), ctx)  具有一定的跨平台能力 k = te.reduce_axis((0, K), "k") A = te.placeholder((M, K), name="A") 行优先 列优先 B = te.placeholder((K, N), name="B") C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C") 缺点 s = te.create_schedule(C.op) xo, yo, xi, yi = s[C].tile( 向量化 分块 C.op.axis[0], C.op.axis[1], bn, bn)  很难达到算子库的性能 (k,) = s[C].op.reduce_axis ko, ki = s[C].split(k, factor=4)  缺少对图层信息的感知 s[C].reorder(xo, yo, ko, ki, xi, yi) s[C].vectorize(yi) 并行 并行分块 s[packedB].parallel(x) 缺少一套可以融合图层与算子层优化的方法 用户手写的调度代码 典型的调度原语 清华大学 Tsinghua University 5 

6 . 现有图层优化编译器 - TASO  TASO：自动的等价计算图搜索和替换  搜索由后端支持的算子所组成的计算图  利用形式化验证证明计算图等价性  用等价的高效计算图替换低效计算图 利用所有可用的算子，遍历所有算子个数 不大于某个固定阈值的计算图 计算图 搜索 输入子程序 后端支持的 算子 … 发现并利用等价的高效计算图 清华大学 Tsinghua University 6 

7 . 基于局部等价优化的人工智能应用编译优化 Pet: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections 通过构造算子库中没有的内存布局重排算子，引入局部等价优 化，为张量程序的优化带来新的机会 清华大学 Tsinghua University 7 

8 . 局部等价优化 ∀𝒑. 𝒀[𝒑] = 𝒁[𝒑] Y Z Add 张量中有成千上万个元 Conv Conv Conv 素，是否可以在优化时 Add 只考虑“大部分”元素 W1 X W1 W2 X 的等价性？ W2 完全等价优化 优点：保证功能上的完全等价性 局部等价优化 缺点：会错过一些优化机会 清华大学 Tsinghua University 8 

9 . 局部等价优化 ∀𝒑. 𝒀[𝒑] = 𝒁[𝒑] ∃𝒑. 𝒀[𝒑] ≠ 𝒁[𝒑] Y Z Z Y Add Conv Dilated Conv Conv Conv Conv Add W1 W2 X W1 W2 X X W X W 完全等价优化 局部等价优化 优点：保证功能上的完全等价性 优点：可以取得更好的性能 • 更高效的算子 • 更高效的张量内存布局 • 更多的系统优化机会 缺点：会错过一些优化机会 缺点：可能造成精度降低、功能不等价 清华大学 Tsinghua University 9 

10 . 局部等价优化 ∀𝒑. 𝒀[𝒑] = 𝒁[𝒑] ∃𝒑. 𝒀[𝒑] ≠ 𝒁[𝒑] Y Z Z Y Add 利用局部等价变形所带来的性能提升的同时 Conv Dilated Conv Conv Conv Conv 维持功能的等价性 Add W1 W2 X W1 W2 X X W X W 完全等价优化 局部等价优化 优点：保证功能上的完全等价性 优点：可以取得更好的性能 • 更高效的算子 • 更高效的张量内存布局 • 更多的系统优化机会 缺点：会错过一些优化机会 缺点：可能造成精度降低、功能不等价 清华大学 Tsinghua University 10 

11 . 突变 突变 程序 突变生成器 生成器 纠错器 优化器 利用所有可用的算子，遍历所有算子个数 不大于某个固定阈值的子程序 突变 生成器 输入子程序 后端支持的 所有算子 … 可发现完全等价与局部等价的 优化变形 内存布局 重排算子 … 清华大学 Tsinghua University 11 

12 . PET 整体架构 突变 … 突变 … 程序 生成器 纠错器 优化器 输入程序 优化后的 突变程序 完成纠错 程序 的突变 清华大学 Tsinghua University 12 

13 . 输入程序 突变 突变 程序 程序优化器 生成器 纠错器 优化器 • 利用非线性算子将程序划分为 子程序 多重线性 • 对每个子程序搜索其候选突变 张量程序 基于搜索的 突变生成器与 程序优化器 突变纠错器 … • 通过后处理优化来保证优化后 程序的高效性 已完成纠错的突变 优化后的程序 清华大学 Tsinghua University 13 

14 . 突变 突变 程序 挑战：如何检测和纠错？ 生成器 纠错器 优化器 Program f Program g 利用多重线性张量程序良好的代数性质，快速定位不 等价结果并生成纠错内核 具体见 "PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections." OSDI. 2021. 清华大学 Tsinghua University 14 

15 . 优化案例 – dilated conv  下图中的例子通过将输入张量进行变形把 dilated conv 的计算变为普通 conv 的计算 清华大学 Tsinghua University 15 

16 . 实验结果  在 batch size 分别为 1 和 16 的情况下，在常用的 DNN 模型上可取 得最高 2.5 倍的加速比 清华大学 Tsinghua University 16 

17 . 小结：基于内存布局重排的局部等价优化  提出了基于局部等价变换的张量程序编译优化技术  结合张量程序应用特征，创新性地提出了张量重排与局部等价的优化策略  对常见的模型和算子可以有最高2.5倍的性能提升  但该方法依然存在一定局限性  优化空间受限：仅能利用用户指定的算子进行搜索和优化  无效搜索较多：遍历搜索需尝试大量错误状态，增加优化开销 清华大学 Tsinghua University 17 

18 . 基于表达式推导的图算融合优化 EinNet: Optimizing Tensor Programs with Derivation-Based Transformations 以表达式替代算子表示，利用代数运算规则对表达式进行自动 推导，以发掘更多的图算融合优化潜力 清华大学 Tsinghua University 18 

19 . 基于张量表达式的形式表达  张量表达式  可表示遍历顺序、维度取值空间、内存布局等张量程序的重要特征  人工智能应用中常用的算子均可表示，如卷积、矩阵乘等  卷积算子的张量表达式示例 引入了循环符号，显式 以索引方式表示内存布 表示出元素遍历的顺序 显式标示出循环和求和 局，以支持不同的内存 及取值空间 变量的取值空间，以支 布局变量 持取值空间的映射变换 在保证完备性的同时，准确表示了张量程序的特征 清华大学 Tsinghua University 19 

20 . 在表达式中融合图层信息  通过表达式嵌套的方法将图层信息融合到表达式中 A W1 W2 Conv Conv ➢ 将不同算子以嵌套方式放入同一 Add 个表达式，不同子表达式间可以 B 进行联合优化 清华大学 Tsinghua University 20 

21 . 利用表达式替代基于计算图的算子表示 基于算子的计算图表示 基于表达式的表示 A K Conv B 优点：➢ 表示方法简洁 优点：➢ 保留图层信息的同时包含算子计算逻辑 ➢ 底层硬件无关 ➢ 底层硬件无关的同时摆脱算子库束缚 ➢ 可支撑图算融合优化 缺点：➢ 优化受限于算子库 缺点：➢ 缺少现有优化技术的支持 ➢ 只能进行图层优化 ➢ 优化的搜索空间变大 清华大学 Tsinghua University 21 

22 . 表达式推导及算子映射的基本流程 输入 计算图 张量 表达式 输出 计算图 推导过程利用基本 代数运算规则，如 交换律、结合律等 将算子库没有的算子进 将已有算子映射到算子库 行自动生成 清华大学 Tsinghua University 22 

23 . 优点：扩展优化空间  相对现有自动图优化工作的优化空间局限性  仅能利用用户预定义的算子进行搜索和优化  局限于POR变换(predefined operator representable)  利用张量表达式将优化空间扩展至非预定义算子  探索涵盖任意张量表达式的变换  融合图层和算子层信息，发掘更多优化机会 清华大学 Tsinghua University 23 

24 . 推导规则汇总 推导规则 形式化描述* 描述 表达式拆分 将一个表达式拆分为多个无依赖的表达式 算子间 表达式合并 将多个无依赖的表达式合并为一个表达式 规则 表达式融合 将多个有依赖的表达式合并为一个表达式 求和拆分 将一个求和作用域拆分为两个 变量替换 将遍历操作的迭代变量进行替换 算子内 遍历合并 将两个作用域下的遍历操作合并为一个 规则 边界放松 将迭代变量的取值范围扩大 边界收紧 将迭代变量的取值范围缩小 * 由于篇幅原因，此处只保留了推导规则的形式化描述的公式，省略了其限制条件 推导规则可方便地进行扩展，以支持更多类型的优化 清华大学 Tsinghua University 24 

25 . 自动推导过程 ➢ 代码生成部分以访存 密集型算子为主，复 杂度较低，调度简单， 搜索时间较短 ➢ 本研究的形式化表示 方法与TVM的表达式 有兼容，可以方便地 通过代码映射生成 TVM代码 利用 TVM 进 行代码生成 清华大学 Tsinghua University 25 

26 . 新变换的机会  单个卷积算子的部分推导方案  Conv3x3计算可变换为Matmul、Conv1x1、Conv2x2、Group conv等计算 输入程序 表达式推导过程 可行变换方案 (a) T0 W (c) (d) (e) T0 W (f) T0 W … VS+BR+… … Conv3x3 Matmul Transpose Σ(𝑐) Σ(𝑐) Σ(𝑐) Σ(𝑐) 转换为表达式 图算融合的搜索空间带来了大量新的变换机会 SS (g) Σ(𝑟𝑠) … Σ(𝑟𝑠) (i) OffsetReducers T0 Conv1x1 OffsetReducek (h) W VS+TM+… … VS ... ... OffsetConcat Im2col变换 Σ(𝑘) (n) T W (b) Matmul 0 … Σ(𝑘) VS+TM+… (j) f (k) T0 W Pad+Split Σ(c𝑟𝑠) (m) … Conv2x2 Op BR … Duplicate 预定义算子 Σ(𝑐′𝑟′𝑠′) eOp 张量表达式算子 Σ(𝑐𝑟1 𝑠1 ) Σ(𝑐𝑟1 𝑠1 ) Group Conv3x3 OffsetReducer2s2 沿循环变量𝑥求和 (l) VS+SS+… Σ(𝑥) Σ(𝑟2 𝑠2 ) 公式推导 BR (p) T0 W Σ(c𝑟𝑠) … 表达式实例化 ( … Split 推导规则 o) TM 遍历合并 Σ(𝑐𝑟1 𝑠1 ) Σ(𝑐𝑟1 𝑠1 ) Conv2x2 Conv2x1 Conv1x2 Conv1x1 SS 求和拆分 BR 边界放松 VS 变量替换 BT 边界收紧 Σ(𝑟2 𝑠2 ) OffsetAdd 清华大学 Tsinghua University 26 

27 . 新优化空间的挑战  探索图算融合优化空间面临的挑战 1. 自动算子生成开销高  利用TVM在GPU上生成典型矩阵乘算子需要30分钟  无法对所有表达式都进行自动算子生成 2. 图算融合搜索空间大 1. 预定义算子数量较少，但存在无穷多张量表达式 EinNet搜索空间 清华大学 Tsinghua University 27 

28 . 高效探索图算融合空间 – 算子库映射  将计算密集部分映射至算子库，访存密集部分映射至算子生成框架  同时利用算子库的高效性和算子生成框架的灵活性  基于循环变量映射表，实现表达式至算子库算子的匹配  根据预定义算子的张量表达式，自动生成相关信息 循环变量映射表 清华大学 Tsinghua University 28 

29 . 高效探索图算融合空间 – 搜索方法  两阶段搜索算法  探索式推导阶段：遍历所有推导规则，探索有效的变化方案  利用搜索深度限制搜索范围  收敛式推导阶段：选择性执行推导规则，减小搜索开销  比较当前表达式和目标算子（如Matmul）表达式，计算表达式距离  执行使距离减小的推导规则，加速实现到目标算子的匹配 清华大学 Tsinghua University 29