有幸参与了MICRO2020,见识到了很多优秀的论文,其中最让我惊艳的是华为的在多面体优化上做优化的文章 <Optimizing the Memory Hierarchy by Compositing Automatic Transformations on Computations and Data>(https://www.di.ens.fr/~zhaojie/micro2020-paper),再看到video之后立刻读了原文,觉得其中很多思想和Halide提出的Pipeline很像。因此我想在这里发表一下自己对深度学习与编译器的结合的看法,抛砖引玉,和各位大佬讨论一下。
传统的多面体优化(Polyhedral model)会将给定的程序直接当成一大坨进行分析,在一大段复杂的数学和工程实现后,将模型完成转换,得到一个相对不错的(自动挖掘并行度)程序。
我们都知道,在HPC领域我们关注的主要是循环结构,对循环结构的调整可以影响locality, reuse distance等等。多面体优化就是可以在保证正确性的情况下自动对循环进行调整,并行,如下所示:
当然,多面体优化本身是一个极复杂的算法,我会在之后的文章单独来写,这次只是给大家一个直观的认识。
华为的这篇文章做了什么事情?首先考虑下面的代码(例子来自论文原文)
对这个代码做变化的话,可能会得到两种结果:
这种方法得到的循环结构比较整齐,都是嵌套循环,利于并行
这种方法虽然尽可能的消掉了循环,但是计算逻辑复杂,有很多的if判断。
对于GPU,我们自然希望模型并行度高,而且diverge少,也就是分支判断少,那么显然第一种变换优于第二种。
因此我们的目标也比较明确:就是将程序做变换,使变换后得到的模型并行度好。
一般来说,GPU程序都是对output做并行,因此我们也希望可以整齐的切分output。这里可以介绍一下tiling的概念,对于一个普通的二重循环:
for(int i = 0; i < 16; i++)for(int j = 0; j < 16; j++)fn(i,j);
如果我们将两重循环分别做切割,得到四重循环:
for(int i_outer = 0; i_outer<4;i_outer++)for(int i_inner = 0; i_inner<4; i_inner++)for(int j_outer = 0; j_outer < 4; j_outer++)for(int j_inner = 0; j_inner < 4; j_inner++)int i = i_outer * 4 + i_inner;int j = j_outer * 4 + j_inner;fn(i, j);
再对四重循环的二三层做交换,即顺序变成i outer, j outer, i inner, j inner,那么这个变换就叫做tiling
tiling往往和棋盘格格式紧密相连,下面图中矩阵乘法的例子就用了tiling,把C切成了9*16块
tiling的好处在此不做展开,网上的资料比较多。有兴趣的同学直接搜CUDA的矩阵乘法优化即可。
介绍完了我们的目标,还要提一下另一个重要的概念:Pipeline。这个概念最开始在Halide的论文中(http://people.csail.mit.edu/jrk/halide-pldi13.pdf)被提出。Halide是一个图像处理编译器,做过CV的同学都知道,我们处理图像一般都是走一个Pipeline。比如正则化->高斯模糊->调整对比度。在这里面有一个偏序关系,那就是Consumer和Producer(其实就是上下游两个op)
生产者生产的值只会被消费者使用,也就是说消费者的使用决定了生产者的生产。
举一个最经典的例子,假设我有一个2D矩阵,我希望计算每个点和相邻8个点的和。
那我可以把这个算法拆成一个Pipeline:
输入矩阵->blurx矩阵(每个矩阵的元素存储输入矩阵每个点和左右两个点的和)->output(每个矩阵的元素存储blurx矩阵每个点和上下两个点的和)
实际上就是干了下面的事情:
然而根据消费者(out)不同的使用情况,生产者(blurx)的生产也不同,比如下面两种例子:
第一个例子很正常,一般人都会这么写,体现不出来啥。。。
这种写法就很奇葩:几乎不把blurx存到内存里面,用到的时候当场算一遍。
算法1 | 算法2 | |
---|---|---|
内存占用 | 2048*3072 | 3 |
计算量 | 2048*3072 | 2046*3072*3 |
可以发现前者有最大的内存,最小的计算量,而后者正好相反。这两种不同的producer-consumer模式可以看作是computation和memory的tradeoff
华为的文章利用了这种思想:你不是想并行度好么?行,那我直接就把最后的输出切一下,每个输出需要哪些producer用多面体优化技术算一遍,这样我又可以做tiling(因为output被切了),又可以利用多面体优化,对每个output小块做优化
论文中给的例子也很直观:先把右下角的output均匀切成4块,然后算每块output需要的producer(左下角)。
我在深度学习编译器方面也有一些了解,觉得目前看来,这种producer-consumer的优化很popular,TVM也在用(compute_at原语)。这种通过输出推导输入可以尽可能避免冗余计算,同时生成对于output并行的结构有良好性质的代码。
总结一下,之前的多面体优化是把模型当成一坨来优化,而在华为的工作中,把模型看作了层级很清楚的Pipeline,对每层分别做优化,这样生成的程序就有更好的并行性。当然,这篇论文更可贵的地方是提供了详细的完整的代码实现,真的可以落地到实际应用场景。不过鉴于篇幅有限(主要是我懒。。。),就简单的把思想说一下。还是强烈推荐有兴趣的读者看看原文!