在CPU上加速卷積神經網絡

在AWS AI，我們發現雖然GPU在深度學習裡扮演著重要角色，實際操作中有很多卷積神經網絡推理其實都由CPU上完成，無論是在雲上還是在終端。經過一番評測，我們認識到沒有一款深度學習框架可以在各種主流CPU（包括Intel, AMD和ARM）都高效地進行卷積神經網絡推理，主要原因是大家的做法都過於依賴第三方庫，再加上框架本身不可避免地帶來了一些overhead。為了做一個面向不同CPU，依賴度儘量少的高效模型推理方案，我們把目光投向了深度學習編譯器。

當時（現在也是）功能最全社區活躍度最高的開源深度學習編譯器是TVM，它可以接受不同框架（Keras/MXNet/TensorFlow/...）的模型並編譯到多種設備（CPU/GPU/FPGA/...）上。由於TVM的編譯做了不同粒度的優化，在端到端（end-to-end）性能上有時甚至能得到幾十倍的加速。TVM發起於華盛頓大學（當時的）在讀博士生陳天奇，AWS AI從一開始就在活躍地參與這個開源項目，一直都是TVM社區最大一股的工業界力量。

最早TVM的工作比較注重端到端功能的實現，以及在這之上的GPU調優。對於CPU的性能，從基本的卷積算子優化，到更複雜的計算圖與算子之間的聯合優化，以及運行時的多線程並行都沒有仔細做。我們對上面幾點經過一頓操作後如願以償地得到了上面說的「一個面向不同CPU，依賴度儘量少的高效模型推理方案」。這個工作被我們發表於計算機系統頂會之一USENIX ATC 』19，論文Optimizing CNN Model Inference on CPUs具體闡述了優化方案和結果。下圖節選了一些結果，展示了我們方案對經典卷積神經網絡模型推理相比深度學習框架的性能提升（數字是以Tensorflow為基準的加速比，越大越好），測試平臺涵蓋Intel, AMD和ARM的主流CPU，數據類型都是fp32。注意下圖裡的Tensorflow和MXNet在x86 CPU上都調用了Intel MKL-DNN，在ARM CPU上也調用了OpenBlas和Eigen這樣的高性能庫，而我們的方案完全沒有藉助任何第三方加速庫，所以運行時的依賴和佔用內存量也都做到了最小。這一系列優化已經被廣泛應用在亞馬遜從雲到端的各種服務和產品中，比如亞馬遜智能音箱Echo的wakeword model就是用我們的方案優化的。換句話說，在遍布全球的Echo上，當你說Alexa喚醒它們時，本質上都是我們優化過的模型輸出了一個正例。我們的整套解決方案完全開源，已經集成在TVM裡面，歡迎大家參照TVM網站上的Auto-tuning a convolutional network for x86 CPU教程試用。




如果你好奇從技術上說為什麼TVM跑得這麼快，請繼續往下讀，我們嘗試從下面幾個方面來極簡略地一窺究竟。

基本算子優化

以Convolution為例（卷積神經網絡絕大多數的運算都消耗在這裡），TVM裡可以簡潔地定義如下

C = tvm.compute(
(batch, out_channel, out_height, out_width),
lambda nn, ff, yy, xx: tvm.sum(
Input[nn, rc, yy * stride_h + ry * dilation_h,
xx * stride_w + rx * dilation_w] *

Filter[ff, rc, ry, rx].astype(out_dtype),

axis=[rc, ry, rx]), tag="conv2d")

TVM編譯系統會把它轉化為一系列的循環乘加。根據Roofline模型，計算性能要麼受制於獲取數據的帶寬（可通過tile增加數據的locality），要麼受制於硬體本身的並行計算能力（可通過vectorize等方法充分利用晶片的性能）。TVM提供了另一套稱作schedule的接口，讓我們可以用一兩行Python代碼實現複雜的優化操作。

s = tvm.create_schedule(C.op)

什麼是tile？

由於內存成本的限制，處理器通常採用多級cache，越靠近處理器的cache速度越快，容量越小。如果數據小到足夠放入cache，並且能儘可能長時間的被使用，那麼就能節省從DRAM裡拉取數據的開銷。TVM scheduling提供了對tile, reorder, cache_write等原語來把數據切成小塊以便放到cache裡。

CC = s.cache_write(C, 『global』)
ow_chunk, ow_block = s[C].split(ow, factor=16)
s[C].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block)

什麼是vectorize？

現代處理器的SIMD指令用於一次處理多個數據。比如下面這個for循環就可以用一條指令完成。

for (int i = 0; i < 16; ++i) {
c[i] = c[i] + a[i] * b[i];
}

TVM裡使用SIMD指令只需要一行代碼：

s[C].vectorize(oc_block)

計算優化的方法還有很多，感興趣的同學可以參考MIT的Performance Engineering of Software Systems公開課。

通過巧妙地組織TVM提供的這些scheduling primitive，在Intel Skylake CPU上我們實現了面向AVX-512這一特定指令集的特殊優化。舉例來說，我們用tile將小塊數據放入最內層循環，用tile+unroll+cache_read/cache_write生成連續的乘加操作，從而充分使用到AVX-512指令集提供的寄存器zmm，等等。總言之，我們的方案在不藉助彙編的情況下，也能夠有效利用晶片的底層特性，從而做到最大化性能。具體的細節可以參考上文提到的ATC』19 paper。

讓機器自己優化算子（Operator-Level Optimization）

人的經驗和精力畢竟是有限的，深度學習算子那麼多，即便只是一個convolution，在神經網絡中的大小形狀也有千千萬（不同形狀的輸入其性能差別非常大，可以對比一下MKL裡矩陣乘法在正方形和長方形輸入上的表現）。那麼，能不能讓機器針對特定workload，窮舉tile的大小，循環的順序，vectorize的寬度，等等？進一步說，如果計算資源有限，能不能藉助機器學習的方法，縮小這個搜索空間？

這就是AutoTVM嘗試去做的事情。我們在實際應用中，確實得到了比人工調優更好的性能，在一些尺寸的輸入上，甚至可以超過MKL-DNN，ACL，cuDNN這類官方加速器。並且我們觀察到，AutoTVM會針對不同的卷積層產生不一樣的算法，這一點是人工靜態代碼很難做到的。

計算圖的優化（Graph/Function-Level Optimization）

深度學習系統現在普遍採用了圖結構來描述網絡（TVM中的Relay IR更進一步用Programming Language的方式來描述Deep Learning Network，大大增強了框架的表達能力，本文並不展開討論這一點）。TVM會針對inference做大量圖結構的優化，比如把多個算子fuse在一起，減少數據移動和內存分配的開銷；在編譯期預先完成模型參數的計算；把BatchNorm摺疊進Convolution，等等。

交叉優化算子和計算圖

除此之外，在ATC』19 paper中，我們還提出算子和計算圖的聯合優化。我們的方案讓operator-level和graph-level配合，使得算子的優化可以反映到計算圖上；反過來，計算圖的tuning又能指導算子層面做出更好的調優策略。

舉例來說，為了優化convolution的性能，我們通常會對輸入數據的layout (NCHW四維，分別對應batch size, input channel, height, width)做處理，比如tile input channel這一維去做vectorize（一個常見辦法是把NCHW轉換成NCHW16c，這個16就是vectorize的大小），在計算結束後再把layout轉回去（NCHW16c → NCHW）。這個layout的轉換操作是有性能成本的，每個convolution都這麼轉來轉去肯定造成一些浪費。一個自然的想法是，如果網絡中的convolution都採用NCHW16c這個layout，是不是只需要在整個網絡的輸入輸出各轉換一次就可以了？進一步說，這裡的參數16也是可調的，網絡中有些conv可能用16比較好，有些用32比較好，還有些形狀奇怪的可能要用7，14……這樣一來，我們就需要在layout transform產生的性能損失，和使用不同layout產生的性能提升之間做trade-off。具體算法可以參考我們paper中的相關描述。

通過算子和計算圖的自動優化，TVM可以生成非常高效的計算代碼。我們最終希望，給定一個網絡結構和目標硬體，使用TVM可以自動生成最優的計算邏輯。就像如今程式設計師只需要專注於業務邏輯，而把性能調優交給高級語言編譯器來完成；深度學習科學家也只需要專注於模型結構本身，把部署到生產環境（可以是CPU/GPU伺服器，Edge device，瀏覽器，FPGA等等）這個任務交給深度學習編譯器。

CPU上的運行時多線程

上面說的各種優化思想其實可以用到各種不同的設備上，專門對於CPU來說，我們還優化了TVM運行時的多線程並行操作。你可能會問為什麼不直接用OpenMP, 簡單的答案是我們實測的OpenMP擴展性和穩定性並不理想（具體參見paper）。另外OpenMP這套接口的在不同平臺上的實現各不相同，這也帶來了性能上的不確定性。再者，TVM裡遇到的多線程並行都是embarassingly parallel的，我們並不需要像OpenMP那樣處理各種複雜的並行同步。基於以上種種，我們決定在TVM用pthread和C++11以後提供的std::atomic實現一套簡單實用的CPU運行時多線程。具體一點說，我們用細粒度的原子操作來實現單生產者單消費者的lockfree queue以便調度作業的主線程和各個工作線程高效通信；我們還把不同的工作線程綁在不用的CPU物理核上，最大程度地消除資源競爭；最後我們在全局變量上加了相應的cache line padding以消除線程間的false sharing。