0%

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

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

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

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

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

摘要

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

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

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

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

Introduction

随着DL大规模应用,成本变得很高。$5\%$的性能提升可能要多花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\%$,比库方法好$18.6\%$

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的组合,比如下列运算:

其中T是可变的参数。如果我们有一个运算器可以运算$128*128$的矩阵micro-kernel:dense_128x128。

当T=64时,由于$16·T = 8×128, 2304 = 18×128$可以把他们视为$8*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需要输入整个程序,从中抽取一些特征:

但是,由于micro-kernel是程序的一种切片,不是整个程序,因此不适用传统cost model。

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

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

式子的左边可以通过已有方法的公式计算,更重要的是,由于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以后,我们根据:

这个得分是和吞吐率成正比的,因此越高越好。用argmax

来对与特定的形状S派发一组micro-kernel的排列M,这个可以遍历。结束以后,我们把$$做成一颗决策树,在运行时通过决策树决定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的原理,感觉还是挺神奇的。