图像 pipeline_多面体优化，Pipeline与深度学习编译器

有幸参与了MICRO2020，见识到了很多优秀的论文，其中最让我惊艳的是华为的在多面体优化上做优化的文章 （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，对每层分别做优化，这样生成的程序就有更好的并行性。当然，这篇论文更可贵的地方是提供了详细的完整的代码实现，真的可以落地到实际应用场景。不过鉴于篇幅有限（主要是我懒。。。），就简单的把思想说一下。还是强烈推荐有兴趣的读者看看原文！