卷積操作的高速實現

使用深度學習庫可以大幅加速CNN模型運行，那么這些庫中的哪些具體的做法實現了這種高速度和高性能呢？佐治亞理工學院計算機科學碩士研究生Manas Sahni在自己的電腦上試驗了多種方法的策略，深入剖析高速卷積的實現過程。

我的筆記本電腦CPU還可以，在TensorFlow等庫的加持下，這台計算機可以在 10-100 毫秒內運行大部分常見CNN模型。2019年，即使是智能手機也能在不到半秒內運行「重量級」CNN模型。而當我自己做了一個簡單的卷積層實現，發現這一個層的運行時間竟然超過2秒時，我非常震驚。

大家都知道，現代深度學習庫對大部分運算具備高度優化的生產級實現。但是這些庫使用了哪些人類不具備的「黑魔法」呢？它們如何將性能提升100倍？當它們「優化」或加速神經網絡運算時，它們在做什么？當談及高性能/高效DNN時，我常常問（或被問及）這些問題。

本文嘗試介紹在DNN庫中如何實現一個卷積層。卷積不僅是很多DNN模型中最常見和最繁重的操作，它也是導致上述高性能實現的代表性trick：一點點算法智慧 + 大量調參 + 充分利用低級架構。

本文所涉及的很多內容來自這篇開創性論文《Anatomy of High-Performance Matrix Multiplication》，這篇論文形成了線性代數庫（如OpenBLAS）所用算法的基礎；本文還涵蓋 Stefan Hadjis 博士和 Chris Rose的教程。

論文地址：https://www.cs.utexas.edu/~flame/pubs/GotoTOMS_revision.pdf
教程地址：https://cs.stanford.edu/people/shadjis/blas.html

朴素卷積（Naive Convolution）

不成熟的優化是萬惡之源。——Donald Knuth

在探討優化之前，我們先來了解一下基線和瓶頸。這里是一個朴素 numpy/for-loop 卷積：

'''
Convolve `input` with `kernel` to generate `output`
    input.shape = [input_channels, input_height, input_width]
    kernel.shape = [num_filters, input_channels, kernel_height, kernel_width]
    output.shape = [num_filters, output_height, output_width]
'''
for filter in 0..num_filters
    for channel in 0..input_channels
        for out_h in 0..output_height
            for out_w in 0..output_width
                for k_h in 0..kernel_height
                    for k_w in 0..kernel_width
                        output[filter, channel, out_h, out_h] += 
                            kernel[filter, channel, k_h, k_w] * 
                            input[channel, out_h + k_h, out_w + k_w]

這個卷積包含6個嵌套的for loop，這里不涉及步幅（stride）、擴張（dilation）等參數。假如這是MobileNet第一層的規模，我們在純C中運行該層，花費的時間竟然高達22秒！在使用最強悍的編譯器優化后（如-O3 或 -Ofast），該卷積層的運行時間降至2.2秒。但是就一個層而言，這個速度仍然太慢了。

那么如果我使用Caffe運行這個層呢？在同一台計算機上使用Caffe運行同一個層所花費的時間僅為18毫秒，實現了100倍的加速！整個網絡運行時間才大約100毫秒。

那么「瓶頸」是什么？我們應該從哪兒開始優化呢？

最內的循環進行了兩次浮點運算（乘和加）。對於實驗所使用的卷積層規模，它執行了8516萬次，即該卷積需要1.7億次浮點運算（170MFLOPs）。我的電腦CPU的峰值性能是每秒800億FLOPs，也就是說理論上它可以在0.002秒內完成運行。但是很明顯我們做不到，同樣很明顯的是，電腦的原始處理能力是非常充足的。

理論峰值無法實現的原因在於，內存訪問同樣需要時間：無法快速獲取數據則無法快速處理數據。上述高度嵌套的for-loop使得數據訪問非常艱難，從而無法充分利用緩存。

貫穿本文的問題是：如何訪問正在處理的數據，以及這與數據存儲方式有何關聯。

先決條件

FLOP/s

性能或者速度的度量指標是吞吐量（throughput），即每秒浮點運算數（FLoating Point Operations per Second，FLOP/s）。使用較多浮點運算的大型運算自然運行速度較慢，因此FLOP/s是對比性能的一種較為可行的指標。

我們還可以用它了解運行性能與CPU峰值性能的差距。我的計算機CPU具備以下屬性：

2個物理內核；
每個內核的頻率為2.5 GHz，即每秒運行2.5×10^9個CPU循環；
每個循環可處理32 FLOPs（使用AVX & FMA）。

因此，該CPU的峰值性能為：

GFLOP/s。這就是我的CPU的理論峰值性能。類似地，我們可以得出單核CPU的峰值性能為80 GFLOP/s。

存儲順序和行優先

邏輯上我們將矩陣/圖像/張量看作是多維度的，但實際上它們存儲在線性、一維的計算機內存中。我們必須定義一個慣例，來規定如何將多個維度展開到線性一維存儲空間中，反之亦然。

大部分現代深度學習庫使用行主序作為存儲順序。這意味着同一行的連續元素被存儲在相鄰位置。對於多維度而言，行主序通常意味着：在線性掃描內存時第一個維度的變化速度最慢。

那么維度本身的存儲順序呢？通常4維張量（如CNN中的張量）的存儲順序是NCHW、NHWC等。本文假設CNN中的張量使用NCHW存儲順序，即如果HxW 圖像的block為N，通道數為C，則具備相同N的所有圖像是連續的，同一block內通道數C相同的所有像素是連續的，等等。

Halide語言

本文討論的優化有時需要使用C語法甚至匯編語言這類底層的低級語言。這會影響代碼的可讀性，還會使嘗試不同優化方法變得困難，因為它需要重寫全部代碼。Halide是一種嵌入到 C++ 中的語言，它可以幫助抽象概念，旨在幫助用戶寫出快速的圖像處理代碼。它可以分離算法（需要計算的東西）和調度策略（如何計算算法以及何時計算），因此使用Halide試驗不同優化方法會更加簡便。我們可以保持算法不變，試用不用的調度策略。

本文將使用Halide語言展示這些低級概念，但是你需要首先了解函數名稱。

從卷積到矩陣相乘

上文討論的朴素卷積已經夠慢了，本節要介紹的實現則更加復雜，它包含步幅、擴張、填充（padding）等參數。我們將繼續使用基礎卷積作為運行示例，但是最大程度利用計算機的性能需要一些技巧——在多個層次上精心調參，以及充分利用所用計算機架構的具體知識。換言之，解決所有難題是一項艱巨任務。

那么我們可以將它轉換成容易解決的問題嗎？比如矩陣相乘。

矩陣相乘（又稱matmul，或者Generalized Matrix Multiplication，GEMM）是深度學習的核心。全連接層、RNN等中都有它的身影，它還可用於實現卷積。

卷積是濾波器和輸入圖像塊（patch）的點乘。如果我們將濾波器展開為2-D矩陣，將輸入塊展開為另一個2-D矩陣，則將兩個矩陣相乘可以得到同樣的數字。與CNN不同，近幾十年來矩陣相乘已經得到廣泛研究和優化，成為多個科學領域中的重要問題。將圖像塊展開為矩陣的過程叫做im2col（image to column）。我們將圖像重新排列為矩陣的列，每個列對應一個輸入塊，卷積濾波器就應用於這些輸入塊上。

下圖展示了一個正常的3x3卷積：

下圖展示的是該卷積運算被實現為矩陣相乘的形式。右側的矩陣是im2col的結果，它需要從原始圖像中復制像素才能得以構建。左側的矩陣是卷積的權重，它們已經以這種形式儲存在內存中。

注意，矩陣乘積直接得出了卷積輸出，無需額外「轉換」至原始形式。

出於視覺簡潔考慮，此處將每個圖像塊作為獨立的個體進行展示。而在現實中，不同圖像塊之間通常會有重疊，因而im2col可能導致內存重疊。生成im2col 緩沖（im2col buffer）和過多內存（inflated memory）所花費的時間必須通過GEMM實現的加速來抵消。

使用im2col可以將卷積運算轉換為矩陣相乘。現在我們可以使用更加通用和流行的線性代數庫（如OpenBLAS、Eigen）對矩陣相乘執行高效計算，而這是基於幾十年的優化和微調而實現的。

如果要抵消im2col變換帶來的額外工作和內存，我們需要一些加速。接下來我們來看這些庫是如何實現此類加速的。本文還介紹了在系統級別進行優化時的一些通用方法。

加速GEMM

朴素

本文剩余部分將假設GEMM按照該公式執行：

我們首先計算基礎標准矩陣相乘的時間：

for i in 0..M:
    for j in 0..N:
        for k in 0..K:
            C[i, j] += A[i, k] * B[k, j]

使用Halide語言：

Halide::Buffer C, A, B;
Halide::Var x, y;

C(x,y) += A(k, x) *= B(y, k);  // loop bounds, dims, etc. are taken care of automatically

最內的代碼行做了2次浮點運算（乘和加），代碼執行了M∗N∗K次，因此該GEMM的FLOPs數為2MNK。

下圖展示了不同矩陣大小對應的性能：

僅僅達到峰值性能的10%！接下來我們將探索加速計算的方式，但不變的主題仍然是：如果無法快速獲取數據，則僅僅快速計算數據並不足夠。隨着矩陣規模越來越大，內存成為更加嚴重的問題，而性能也會繼續逐漸下降。你看到圖中的劇烈下跌了嗎？這表示當矩陣太大無法適應緩存時，吞吐量突然下跌。

緩存

RAM是大的存儲空間，但速度較慢。CPU緩存的速度要快得多，但其規模較小，因此恰當地使用CPU緩存至關重要。但是並不存在明確的指令：「將該數據加載到緩存」。該過程是由CPU自動管理的。

每一次從主內存中獲取數據時，CPU會將該數據及其鄰近的數據加載到緩存中，以便利用訪問局部性（locality of reference）。

你應該首先注意到我們訪問數據的模式。我們按照下圖A的形式逐行遍歷數據，按照下圖B的形式逐列遍歷數據。

它們的存儲也是行優先的，因此一旦我們找到 A[i, k]，則它在該行中的下一個元素A[i, k+1]已經被緩存了。接下來我們來看B中發生了什么：

列的下一個元素並未出現在緩存中，即出現了緩存缺失（cache miss）。這時盡管獲取到了數據，CPU也出現了一次停頓。
獲取數據后，緩存同時也被 B 中同一行的其他元素填滿。我們實際上並不會使用到它們，因此它們很快就會被刪除。多次迭代后，當我們需要那些元素時，我們將再次獲取它們。我們在用實際上不需要的值污染緩存。

我們需要重新修改loop，以充分利用緩存能力。如果數據被讀取，則我們要使用它。這就是我們所做的第一項改變：循環重排序（loop reordering）。

我們將i,j,k 循環重新排序為 i,k,j：

for i in 0..M:
    for k in 0..K:
        for j in 0..N:

答案仍然是正確的，因為乘/加的順序對結果沒有影響。而遍歷順序則變成了如下狀態：

循環重排序這一簡單的變化，卻帶來了相當可觀的加速：

平鋪（Tiling）

要想進一步改進重排序，我們需要考慮另一個緩存問題。

對於A中的每一行，我們針對B中所有列進行循環。而對於 B 中的每一步，我們會在緩存中加載一些新的列，去除一些舊的列。當到達A的下一行時，我們仍舊重新從第一列開始循環。我們不斷在緩存中添加和刪除同樣的數據，即緩存顛簸（cache thrashing）。

如果所有數據均適應緩存，則顛簸不會出現。如果我們處理的是小型矩陣，則它們會舒適地待在緩存里，不用經歷重復的驅逐。慶幸的是，我們可以將矩陣相乘分解為子矩陣。要想計算 C 的r×c平鋪，我們僅需要A的r行和B的c列。接下來，我們將 C 分解為6x16的平鋪：

C(x, y) += A(k, y) * B(x, k);

C.update()
 .tile(x, y, xo, yo, xi, yi, 6, 16)

/*
in pseudocode:
for xo in 0..N/16:
    for yo in 0..M/6:
        for yi in 6:
            for xi in 0..16:
                for k in 0..K:
                    C(...) = ...
*/

我們將x,y 維度分解為外側的xo,yo和內側的xi,yi。我們將為該6x16 block優化micro-kernel（即xi,yi），然后在所有block上運行micro-kernel（通過xo,yo進行迭代）。

向量化 & FMA

大部分現代CPU支持SIMD（Single Instruction Multiple Data，單指令流多數據流）。在同一個CPU循環中，SIMD可在多個值上同時執行相同的運算/指令（如加、乘等）。如果我們在4個數據點上同時運行SIMD指令，就會直接實現4倍的加速。

因此，當我們計算處理器的峰值速度時，我們其實有些作弊，把該向量化性能作為峰值性能。對於向量等數據而言，SIMD用處多多，在處理此類數據時，我們必須對每一個向量元素執行同樣的指令。但是我們仍然需要設計計算核心，以充分利用它。

計算峰值FLOPs時，我們所使用的第二個技巧是FMA（Fused Multiply-Add）。盡管乘和加是兩種獨立的浮點運算，但它們非常常見，有些專用硬件單元可以將二者融合為一，作為單個指令來執行。編譯器通常會管理FMA的使用。

在英特爾CPU上，我們可以使用SIMD（AVX & SSE）在單個指令中處理多達8個浮點數。編譯器優化通常能夠獨自識別向量化的時機，但是我們需要掌控向量化以確保無誤。

C.update()
 .tile(x, y, xo, yo, xi, yi, 6, 16)
 .reorder(xi, yi, k, xo, yo)
 .vectorize(xi, 8)

/*
in pseudocode:
for xo in 0..N/16:
    for yo in 0..M/6:
        for k in 0..K:
            for yi in 6:
                for vectorized xi in 0..16:
                    C(...) = ...
*/

多線程處理（Threading）

到現在為止，我們僅使用了一個CPU內核。我們擁有多個內核，每個內核可同時執行多個指令。一個程序可被分割為多個線程，每個線程在單獨的內核上運行。

C.update()
 .tile(x, y, xo, yo, xi, yi, 6, 16)
 .reorder(xi, yi, k, xo, yo)
 .vectorize(xi, 8)
 .parallel(yo)

/*
in pseudocode:
for xo in 0..N/16 in steps of 16:
    for parallel yo in steps of 6:
        for k in 0..K:
            for yi in 6:
                for vectorized xi in 0..16 in steps of 8:
                    C(...) = ...
*/

你可能注意到，對於非常小的規模而言，性能反而下降了。這是因為工作負載很小，線程花費較少的時間來處理工作負載，而花費較多的時間同步其他線程。多線程處理存在大量此類問題。

展開（Unrolling）

循環使我們避免重復寫同樣代碼的痛苦，但同時它也引入了一些額外的工作，如檢查循環終止、更新循環計數器、指針運算等。如果手動寫出重復的循環語句並展開循環，我們就可以減少這一開銷。例如，不對1個語句執行8次迭代，而是對4個語句執行2次迭代。

這種看似微不足道的開銷實際上是很重要的，最初意識到這一點時我很驚訝。盡管這些循環操作可能「成本低廉」，但它們肯定不是免費的。每次迭代2-3個額外指令的成本會很快累積起來，因為此處的迭代次數是數百萬。隨着循環開銷越來越小，這種優勢也在不斷減小。

展開是幾乎完全被編譯器負責的另一種優化方式，除了我們想要更多掌控的micro-kernel。

C.update()
 .tile(x, y, xo, yo, xi, yi, 6, 16)
 .reorder(xi, yi, k, xo, yo)
 .vectorize(xi, 8)
 .unroll(xi)
 .unroll(yi)

/*
in pseudocode:
for xo in 0..N/16:
    for parallel yo:
        for k in 0..K:
            C(xi:xi+8, yi+0)
            C(xi:xi+8, yi+1)
            ...
            C(xi:xi+8, yi+5)
            C(xi+8:xi+16, yi+0)
            C(xi+8:xi+16, yi+1)
            ...
            C(xi+8:xi+16, yi+5)
*/

現在我們可以將速度提升到接近60 GFLOP/s。

總結

上述步驟涵蓋一些性能加速最常用的變換。它們通常以不同方式組合，從而得到更加復雜的調度策略來計算同樣的任務。

下面就是用Halide語言寫的一個調度策略：

 matrix_mul(x, y) += A(k, y) * B(x, k);
    out(x, y) = matrix_mul(x, y);

    out.tile(x, y, xi, yi, 24, 32)
        .fuse(x, y, xy).parallel(xy)
        .split(yi, yi, yii, 4)
        .vectorize(xi, 8)
        .unroll(xi)
        .unroll(yii);

    matrix_mul.compute_at(out, yi)
        .vectorize(x, 8).unroll(y);

    matrix_mul.update(0)
        .reorder(x, y, k)
        .vectorize(x, 8)
        .unroll(x)
        .unroll(y)
        .unroll(k, 2);

它執行了：

將out分解為32x24的平鋪，然后將每個平鋪進一步分解為8x24的子平鋪。
使用類似的重排序、向量化和展開，在臨時緩沖區（matrix_mul）計算8x24 matmul。
使用向量化、展開等方法將臨時緩沖區matrix_mul 復制回out。
在全部32x24平鋪上並行化這一過程。

最后，我們將速度提升至超過120 GFLOPs，接近峰值性能160 GFLOPs，堪比OpenBLAS等生產級庫。使用im2col類似的微調代碼和矩陣相乘，同樣的卷積可以在大約20毫秒內完成運行。如想深入了解，你可以嘗試自行試驗不同的調度策略。作為一名工程師，我經常聽到關於緩存、性能等的語句，而親眼看到它們的實際效果是非常有趣和值得的。

注意：im2col+gemm方法是大部分深度學習庫中流行的通用方法。但是，對於特定的常用規模、不同的架構（GPU）和不同的運算參數（如擴張、分組等），專門化（specialization）是關鍵。這些庫可能也有更專門化的實現，這些實現利用類似的trick或具體的假設。經過不斷試錯的高度迭代過程之后，我們構建了這些micro-kernel。對於什么方法效果好/不好、必須基於結果考慮注釋，程序員通常只有直觀的想法。聽起來和深度學習研究差不多，不是嗎？

原文鏈接：

https://sahnimanas.github.io/post/anatomy-of-a-high-performance-convolution/

原文鏈接：

https://www.jiqizhixin.com/articles/2019-09-15-6

如何實現高速卷積？深度學習庫使用了這些「黑魔法」

免責聲明！