论文阅读 [精读]-DIETCODE: AUTOMATIC OPTIMIZATION FOR DYNAMIC TENSOR PROGRAMS

昨天阅读 survey 之后，今天继续阅读编译优化论文一篇。是 MLSys 会议上的论文。这个会上全是一些深度学习 + 系统的文章，好像国内知名度不是很高，甚至不再 CCF 推荐列表上 ……

说回本文，本文最后一个作者是陈天奇，TVM 的作者，之前我看了他的经历，挺受鼓舞的，推荐大家都可以看一下～

陈天奇：机器学习科研的十年

由于这个方向我了解甚少，这篇论文还是精度，笔记也做得详细一点。

摘要

在计算密集的负载下得到高效率的计算是很难的问题。目前解决办法：

• vendor libraries：需要大量工程开发
• auto-schedulers：只支持确定形状。否则 scheduler 速度会变的非常慢。

作者观察到只支持确定形状的瓶颈在于：搜索空间是和形状有关的。作者提出新方法，可以构建一个形状无关的搜索空间，在同一空间中、用同样的 cost model 来同时搜索所有形状。

实验表明，新方法的速度和效果均大幅超越 SOTA。

Introduction

随着 DL 大规模应用，成本变得很高。$5\mathrm{\%}$ 的性能提升可能要多花 4000$。

然而，提升所有算子的效果很难，需要详细的实验、开发。很多研究者应用 oneDNN，cuDNN 等 vendor library。后者的开发时间和更新周期都很长。（一年一更）

为了解决这个问题，多种 auto-schedule 负责进行 High-level IR 计算定义和 Low-level IR 实现之间的过程。然而，现有方法都需要在计算时知道所有位置的形状，这样才能正确地构建搜索空间，考虑所有情况。

然而，很多工作需要形状不确定。一个方法是对所有形状都遍历优化一遍：

• 慢，一个 BERT 层需要 42h 来优化

作者发现，如果把想同类型、不同形状的算子看成一个可变形状的 workload，然后只规划一次，可以显著减少计算量。已有方法不能这么做，是因为他们要对不同形状开不用的搜索空间。本文提出一种新的方法 DietCode：

构建一个由 micro-kernel 构成的 shape-generic 搜索空间，每个 micro-kernel 都是形状不固定的。这样就能让 DietCode 有一个统一的搜索空间。

DietCode 构建的 cost-model 包括：

• 复杂的 shape-generic 的部分。需要从 micro-kernel 中提取 loop feature 等信息，需要从硬件上获取参数
• 简单的依赖形状的部分。不需要额外信息

本文的贡献可以归纳为：

• 提出通过构建 shape-generic 的搜索空间解决动态形状的规划问题
• 提出 DietCode，可以同时规划所有形状，用同样的 shape-generic cost model
• 在 BERT 模型编译中，比 SOTA 快 5.88 倍，如果考虑所有形状，快 100 倍。效果上比 auto-schedule 好 $69\mathrm{\%}$, 比库方法好 $18.6\mathrm{\%}$

Background

现有 auto-scheduler 工作流

现有的 vendor library 成本很高，有很多 auto-scheduler 来替代。

其输入如上图：包含一个 tensor 的表达式，同时对于 tensor 的形状有附加的说明。

现有方法如上图左:

• 1. 通过表达式和形状信息，与硬件的可能优化方式一起构建搜索空间
• 2. 通过硬件信息构建一个 cost model
• 3. 用某种算法进行学习

最后搜出来的方法大约比 vendor-library 方法快 3.88 倍

动态形状 tensor 的程序

现有方法需要：在优化时知道所有的形状信息。因此不适用于以下需求：

• NAS 中，需要搜索网络超参。每组超参的实现，网络的结构都不同 (举例，12800 种)
• 在 NLP 等序列任务中，输入的长度可能运行时才能确定，比如 BERT 就是长度 1-128
• 统一模型，不同层的维度可能都不一样：比如 BERT 的 hidden state 在不同层可能是 $(768,2304,3072)$

为什么需要新的 auto-scheduler 框架

现有方法加一些代码不行吗？

• 对于 vendor-library：输入性状改变可能带来 13 倍的性能衰减。这是因为代码里有好多 hard-code。改这些东西的代价很大，码量爆炸。
• 已有 auto-scheduler：现有的工作流，上图左，形状信息都是焊死的。不改代码的话，只能对所有形状遍历，这个速度肯定不可行。

已有一些 auto-scheduler 的动态形状改进有一些问题：

• Selective tuning：把不同形状聚成一个个 cluster，然后对每个 cluster 优化。需要额外的知识，不能完全自动化
• Nimble：在一个大形状上预计算，然后在别的动态形状应用。问题是在一个大形状上的最优不一定可以泛化到别的最优。
• Bucketing：把形状变化区间分成子区间，然后用每个子区间的最大值作为值来计算一次。
  • 需要额外计算形状信息，并且由于化简，带来计算不准确。
  • 由于 padding，会出现不必要的计算开销。

综上，已有方法都不好。

DietCode 原理、关键因素

Shape-Generic Search Space

观察下面这个图：

左边是一个含参数 T 的代码，右边是已有 auto-scheduler 的做法：

• 可见，他们只尝试了 T 的因子，比如 $T=49,t\in [1,7,49]$, 显然比如 t=10 时就不准确。

本文，则考虑硬件约束下，给出了一些 micro-kernel。这是符合硬件条件的一些算子，可以一个动态形状的运算可以视为一些 micro-kernel 的组合，比如下列运算： $$\begin{gathered} Y=X{W}^T \\ X:[16\ast T,768],W:[2304,768] \end{gathered}$$ 其中 T 是可变的参数。如果我们有一个运算器可以运算 $128\ast 128$ 的矩阵 micro-kernel：dense_128x128。

当 T=64 时，由于 $16·T=8\times 128,2304=18\times 128$ 可以把他们视为 $8\ast 18$ 的 kernel 的排列在一起进行运算，如图：

通过排布使用 micro-kernel dense_128x128, 我们对于 $T=1,2,...128$ 都可以很方便的实现上面的运算。也就是说，通过这种设计，我们只需要使用形状无关的 micro-kernel，就可以实现任意形状输入的优化

这种方法要解决几个问题：

• 如果 T=60，那么就需要一个 padding，如上图的下面。也就是说，会带来一些额外计算。要怎样高效利用所有 micro-kernel？
• 怎样准确评判这种由 micro-kernel 构成的程序的效率？

Micro-Kernel-based Cost Model

上面说的第一个问题通过 padding 可以一定程度上解决。对于后一个问题，需要构建一个新的 cost model。

已有的 cost-model 需要输入整个程序，从中抽取一些特征： $$\begin{gathered} \text{Cost}(P)=f(\text{FeatureExtractor}(P)) \\ f:\text{cost function (e.g., XGBoost)} \\ P:\text{complete program } \\ \text{Cost}:\text{ compute throughput} \end{gathered}$$ 但是，由于 micro-kernel 是程序的一种切片，不是整个程序，因此不适用传统 cost model。

设计新的 cost model，发现如果一个动态形状程序 P 被切片成很多 micro-kernel M，那么 cost 可以拆解成两部分：

• $f_{MK}$：这些 micro-kernel 的损失。和形状无关（多个 micro-kernel 可以并行，时间是常数）
• $f_{adapt}$ ：把 P 切片成 M 带来的损失。和形状有关，但是计算很简单 (比如上图就是算 padding 的比例)

$${\text{Cost}}_M(P)=f_{MK}(\text{FeatureExtractor}(M))·f_{\text{adapt}}(P,M)$$

式子的左边可以通过已有方法的公式计算，更重要的是，由于 M 和程序 P 无关，可以事先预计算。式子的右边计算速度很快。综上，新模型的计算效率很高。

Joint Learning with DietCode

本文提出了 DietCode，基于以下要素：

• 一个 Shape-Generic Search Space
• 一个 Micro-Kernel-based Cost Model
• 在运行时通过上面的 cost 公式对确定的输入形状派发一种 micro-kernel 排列

和已有方法对比如下图：

可以对所有 categories 同时学习，因为有统一的 cost 公式，复杂度是 $O(1)$

实现细节

这种方法作为 TVM 的 auto-scheduler 方法

local padding

上面提到的由于使用 padding，可能会带来性能下降，高达 17 倍。这是因为引入了很多的 branch 指令，而且每个 branch 都要计算数据。

作者研究了三种已有解决方法，举例子如下图，这个例子有一个数据读取，一个数据修改，一个数据写回

• global padding：如 b。提前算好，不需要在里面做 branch，但是需要引入额外的存储空间和 padding 结果的运算
• Loop partitioning：如 c。把程序切成两片，其中上一片无论如何都不会爆。当 $t\ll T$ 时效果不错，但这个条件很难保证。
• local padding: 在取数据和写回数据时保证正确性，但是在计算时不保证。这个有两个好处：
  • 只有计算阶段的 branch 会大幅影响表现。fetch 和 write back 的时间可以被已有流水线技术优化掉
  • 如果计算阶段算了额外算了 padding 的部分，写回阶段会把它忽略，因此不怎么影响表现

本工作使用 local padding 方法

Micro-Kernel-based Cost Model 具体实现

cost 的计算应该考虑以下部分：

• 1、micro-kernel 的表现
• 2、硬件占有率带来的惩罚，用的 micro-kernel 越多，瞬时的占有率高，吞吐量大，就是好。
• 3、padding 带来的额外惩罚，padding 的比例越低越好

这三部分可以这样理解：

最终公式如下：

其中 $f_{PAD}$ 和硬件无关，之和 micro-kernel 设计、程序 P 有关

其中 $f_{OCC}$ 是占有率惩罚， 用一个回归模型预测

k，b 是可学习的参数，由于每个 micro-kernel 被派发到不同的核，$f_{OCC}$ 也可以看做占了多少的核。如果做一个归一化的话：

可以节省一个参数 b

Automatic Dispatching

在同时计算完所有的 M 以后，我们根据： $$vote(S)=argma{x}_M(CostM(P(S,M)))$$ 这个得分是和吞吐率成正比的，因此越高越好。用 argmax

来对与特定的形状 S 派发一组 micro-kernel 的排列 M，这个可以遍历。结束以后，我们把 $<S,vote(S)>$ 做成一颗决策树，在运行时通过决策树决定 M 即可。

​

EVALUATION

这一部分就不详细讲了。简单说一下，就是测试 BERT 模型的表现，baseline 是：

• vendor: cuDNN
• Auto-scheduler: Ansor
• dynamic code generation: Nimble。进一步优化 Ansor，计算最大形状，然后应用在所有形状上

评测指标，全都是越小越好：

• end-to-end latency on the entire model, measured as msec。
• the runtime cost of the generated tensor programs on a single operator, measured as μsec and averaged over 100 runs。
• the total auto-scheduling time used, measured as hours

上面这个图可以看出，对于不同的输入维度，DietCode 的效果基本最好。同时和 vendor 的差距小，但是 vendor 需要人工大码量、长时间的开发。

上面这个图是编译时间，虚线代表理论推导，实际根本不可行。可以看出，DietCode 非常快。

上面这个图是指选定某种结构以后，真正的跑的速度。更快。作者测试了 dense 层 (一个动态量)，batchmatmul 层 (2 个动态量)，效果都更好。

我的思考

• 其实我之前也不知道这种 “动态形状” 的输入是怎么规划的。但我确实感觉到动态形状的支持应该是必要的，基本所有的 NLP 问题都需要这个。
• 看作者的意思，这个功能现在已经上 TVM 了，我在想这个是不是和操作系统一样：

先有实现，后有理论

• 后续我可能会继续调研一些 TVM 的原理，感觉还是挺神奇的。