MegEngine計算圖、MatMul優化解析
本文針對天元在推理優化過程中所涉及的計算圖優化與 MatMul 優化進行深度解讀，希望能夠幫助廣大開發者在利用天元 MegEngine「深度學習，簡單開發」的同時，也能夠了解 CPU 優化的相關知識。從而幫助大家在模型部署的整體流程中更好地進行加速；在實際模型部署時能夠評估模型在特定平臺上運行所能達到的性能以及內存使用情況；以及在算法設計時可以設計出更利于 CPU 優化加速的卷積 Opr 等。
本文針對曠視天元深度學習框架在推理優化過程中所涉及的計算圖優化與 MatMul 優化進行深度解讀。
背景及引言
在深度學習大規模落地邊緣端場景的今天，如何最大程度降本增效，是企業與開發者共同關注的話題。其中，模型的訓練與推理是兩個關鍵環節。
天元（MegEngine）深度學習框架憑借「訓練與推理一體化」的獨特范式，能夠極大程度上（90%）節省模型從研發到部署的整體成本，降低轉換難度，真正實現小時級轉化；同時，天元（MegEngine）在 CPU 推理方面所做的大量優化工作，也使得開發者在推理時能夠發揮出處理器的最佳性能。
在之前對天元的極致推理優化進行了綜述《 工程之道，MegEngine 推理性能極致優化之綜述篇》。本文則針對天元在推理優化過程中所涉及的計算圖優化與 MatMul 優化進行深度解讀，希望能夠幫助廣大開發者在利用天元 MegEngine「深度學習，簡單開發」的同時，也能夠了解 CPU 優化的相關知識。
從而幫助大家在模型部署的整體流程中更好地進行加速；在實際模型部署時能夠評估模型在特定平臺上運行所能達到的性能以及內存使用情況；以及在算法設計時可以設計出更利于 CPU 優化加速的卷積 Opr 等。
CPU 推理優化概覽
對于產業應用而言，CPU 推理的性能優化至關重要，如下表所示，經過優化的推理性能可以較未經優化的原始性能提升數十倍。

在進行模型推理階段的優化時，首先需要對模型的計算圖進行優化，以避免冗余的計算與訪存，確保計算圖在推理時是最優的；其次，在大多 CV 相關模型中，卷積的計算比重最高，達到模型總計算量 90% 以上，因此對模型推理的優化主要聚焦在對卷積計算的優化上。
通常而言，卷積計算的實現方式有 3 種：direct 卷積，Im2col 卷積以及 Winograd 卷積。
? direct 卷積：根據卷積的計算公式直接對 FeatureMap 上進行滑窗計算；
? Im2col 卷積：根據卷積計算需要在輸入通道上進行 reduce sum 的特點，將卷積運行轉化為 MatMul 計算；
? winograd：在保證計算無誤的前提下，使用加法替代乘法，達到優化卷積乘法計算量的目的，在中間過程需要使用 MatMul 進行計算。
由上文可知，以 Im2col 或 Winograd 方法進行的卷積計算會頻繁使用到 MatMul，因此對 MatMul 這種基礎算子的優化就顯得尤為重要。
基于上述考量，本文將首先介紹模型優化中的圖優化，然后介紹基礎算子 MatMul 在 CPU 上的優化方法。
推理計算圖優化
在訓練階段定義模型的計算圖，主要是為了滿足模型參數的訓練需求。當訓練結束，模型參數固定后，對計算圖的進一步優化能夠幫助模型推理更加高效。
推理和訓練的計算圖是一張有向無環圖 (DAG)，在天元中，開發者能夠以類似 LLVM 的方式對 DAG 計算圖定義許多優化方法，這里簡稱 OptPass。OptPass 可以根據用戶的配置有選擇性的加入到圖優化的 OptPass 列表中，從而幫助用戶靈活地為圖優化定義 OptPass。
計算圖優化
天元定義了多個為推理進行計算圖優化的 OptPass，開發者使用這些 OptPass 之后，將得到一張用于推理的最優計算圖。下面以 MobileNetV1 中， Convolution+Batch Norm+Relu 這樣的典型結構經過計算圖優化之后 Fuse 為 ConvBias 的過程為例，介紹天元的計算圖優化過程。
由于 Batch Norm 在推理階段除了輸入 Tensor 外其他都是常數，因此可以簡化為多個 elemwise 的組合，天元實現了一個 ConvertBatchNormToElemwisePass 的 Optpass，這個 OptPass 將模型中的所有 Batch Norm 轉化為 Elemwise，具體轉化原理如下：

緊接這一過程，天元將運行 OptPass ParamRedistributePass，該 OptPass 會將上述 Batch Norm 轉化而來的 Elemwise 中的 scale 融合到 convolution 的權重中，具體實現原理如下：

最后天元將運行 OptPass FuseConvBiasNonlinPass，該 OptPass 會將計算圖中的 Convolution+Elemwise 轉化為天元內部實現的 ConvBias Op 中，同時，它還會設置 ConvBias Op 中 NonelineaMode 參數。
如此便完成了從 Convolution+Batch Norm+Relu 到 ConvBias 的轉換，整體轉換過程如下圖：

實驗驗證圖優化之后的性能
在推理之前，如果要對完成訓練的模型進行圖優化，則需要在模型 dump 的期間進行，當然，也可以在 SDK 運行模型之前進行。下面是在模型 dump 時進行圖優化的代碼：
from megengine.jit import trace
@trace(symbolic=True)
def pred_fun(data, *, net):
net.eval()
pred = net(data)
pred_normalized = F.softmax(pred)
return pred_normalized

使用 trace 類的 trace 接口無需運行直接編譯

pred_fun.trace(data, net=xor_net)

使用 trace 類的 dump 接口進行部署

pred_fun.dump(“xornet_deploy.mge”, arg_names=[“data”], optimize_for_inference=True, enable_fuse_conv_bias_nonlinearity=True)
上面的 optimize_for_inference=True 將在 dump 模型時候針對 inference 進行優化， enable_fuse_conv_bias_nonlinearity=True 將在模型中進行 Op Fuse 的優化，此外天元還支持其他的優化參數，具體可見天元的文檔。
下表展示的，是在實驗過程中對模型進行圖優化前、后，模型運行性能的測試對比。

可以看出圖優化對模型性能的提升效果顯著，具體提升比例由模型自身決定。
MatMul 優化
如前文所述，MatMul 作為卷積運算的基礎算子，會頻繁地被 Im2col、Winograd 以及 FullyConnect 使用。因此在天元中，MatMul 既被封裝為單獨的 Op，也可以作為單獨的 algo 供卷積的實現使用。
此外，由于 MatMul 也是計算最為密集的算子，因此天元對它進行了極致的優化。
優化
MatMul 是線性代數中的矩陣乘，假設矩陣 A 大小為 MK，矩陣 B 大小為 KN，則得到矩陣 C 大小為 M*N，其中 C 的每個元素的計算公式如下：

可以發現，在 MatMul 的計算中乘法和加法的計算量為 2MNK (計算 C 中每個元素時，加法和乘法計算量分別為 K，C 的總元素個數為 MN)，訪存量為 2MNK (計算每個 C 中元素需要 2K 訪存)+ 2MN(整個 C 矩陣讀一次和寫一次)。由于計算量固定 (排除 Strassen)，所以只能優化訪存，使得乘法和加法運算達到處理器的極限性能，從而實現 MatMul 的最佳性能。
MatMul 分塊
關于減少 MatMul 計算時的訪存量，最有效的方法是對 MatMul 的計算進行分塊，并將分塊之后的數據保存在 CPU 的 Cache 中。
如下圖所示，將 A 按照 mr 進行行分塊，將 B 按照 nr 進行列分塊，計算時將需要用到的分塊保存在 CPU 的各級 Cache 中，從而極大的減少了內存的讀寫。

進一步細化上面的分塊計算過程如下圖所示，A 中的一個行塊都要重復地和 B 中的每一個列塊進行小分塊的 MatMul，寫入到 C 的小塊中，為了使得分塊盡量的大 (如上所述，能夠減少內存訪存量)，Cache miss 率盡量低，因此需要根據 CPU 的 Cache 結構特點 (速度：L1D>L2>L3，容量 L1D<L2<L3) 來分配 Cache 和數據塊之間的存放關系，天元中進行的分配如下：
? 將下圖中紅色 (訪問重復次數最多的 A 的行塊，計算時需要的 B 的一個列塊以及計算結果的 C 的小塊) 部分都保存在 L1 中。
? 由于計算完每一個 C 的行塊，都需要重復遍歷一次整個 B 矩陣，因此將 B 存放在 L2 中使得每次讀取 B 的一個列塊都是從 L2 中讀取，代價最小。
? 將重復訪問率最高的 C 的累加中間結果保存在速度最快的 CPU 處理器的寄存器中。

通過上面的分配策略，并結合 CPU 中資源 (寄存器數量，L1D 和 L2 的大小)，便可以確定最佳的 MatMul 計算中的 Nr，Kr：
? 可以根據 CPU 處理器的寄存器數量得到 mr 和 nr 的具體大小，寄存器容量 > mrnr
? 根據 L1D Cache 的大小結合 mr 和 nr 計算出 Kr，Kr=L1D/(mr+nr)
? 再根據 L2 的大小計算出 B 矩陣中的 Nr，Nr=(L2-L1D)/Kr
在上面計算 N 時，用 L2-L1D 的原因是，由于當前 CPU 使用的 Cache 是 Inclusive 的，且 L2 是指令緩存和數據緩存合并的，另外上面 M 沒有明確限制。
在得到上面最佳 Nr、Kr、mr 和 nr 之后，進一步便可以首先對 MatMul 計算中的 N、K 進行 Nr 和 Kr 分塊，然后在 Nr、Kr 的基礎上再進行 mr 和 nr 分塊。如此，MatMul 計算中的計算訪存比達到最高，且 CPU 處理器的資源也得到充分利用。
經過分塊之后，由于在最內層的計算 Kernel 為 A(mrKr)\B(Krnr)=C(mrnr) 的分塊矩陣乘，決定了整個 MatMul 的計算性能，因此需要極致地優化。
Kernel 計算
最核心的計算 Kernel 進行的是尺寸為 (mrKr)x(Krnr)–>(mrnr) 的小尺寸 MatMul，其計算示意圖如下：

如上圖所示，Kernel 在計算時會讀取 A 中一列， B 中一行，進行矩陣乘，得到 大小為 mr*nr 的 C，然后和原來 C 中的值相加，如此循環 Kr 次，完成該 Kernel 的計算。
在該 Kernel 的計算過程中乘法和加法的計算量為 2mrnrKr，訪存量為 (mr+nr)Kr+2mrnr，可以根據處理器來判斷該計算能否隱藏數據的訪存。下面以 ARM cortex A76 為例進行分析，根據 A76 的數據手冊得到：
? FP32 SIMD Load throughput=2，即單周期可以 load 8 個 float 數據
? FP32 SIMD FMLA throughput=2，即單周期可以進行 16 個乘加運算
因此，當 Kernel 的尺寸 mr=8、nr=12、Kr=256，計算量為 49152 次乘加運算，訪存量為 5312 個 float 數據時，該計算訪存量為 9.25，大于處理器的計算訪存比 2。因此可以得出結論，如果 A 和 B 均在 L1 中，則該 Kernel 的計算不會因為數據的 Load 被阻塞，所以計算單元能夠發揮出處理器的最佳性能。
雖然在對 MatMul 進行分塊時，已經計劃將 A 和 B 這樣的小矩陣置于 L1D Cache 中，但是數據在真正運行時卻不一定都在 L1D 中，原因在于，B 矩陣的列塊在原來的大矩陣中內存并不連續，其次 A 中的一列由于內存不連續，也不能使用 SIMD 進行 load，為此需要對 A 和 B 進行數據 PACK。
MatMul 數據 PACK
上文 kernel 在計算過程中，需要同時讀取 A 矩陣 mr 行數據，而每行數據之間內存不連續，因此這種訪問對 Cache 不友好，同樣，在讀取矩陣 B 中 nr 列的時候也存在數據讀取不連續的問題，加之 A 的所有行塊和 B 中的所有列塊將被讀取多次，因此可以通過對 A 和 B 提前進行數據 PACK， 實現在后續計算中頻繁多次訪問數據時是連續的（只在第一次 PACK 時對 Cache 不友好），進而獲得巨大收益。

對矩陣 A 進行數據 PACK 是將 A 中 mr 行數據的相同列拷貝到一起，如上圖中將 A PACK 到 A’ 的步驟。重復完所有 A 中的行塊便完成了 A 矩陣的數據 PACK。B 矩陣的 PACK 操作是，將 nr 列數據拷貝到連續的內存地址中，它對應上圖 B PACK 到 B’ 的過程。
實 驗
按照文介紹方式方式，天元在 X86 和 ARM 上分別對 MatMul 進行了優化。下表展示了 ARM64 上的性能測試結果，實驗平臺為 kirin 980。
首先，對該處理器進行分析可以看到，其主頻為 2.6 GHz，每個周期能夠進行 16 次乘加計算，因此其理論計算峰值為 16*2.6=41.6 Gflops。

可以看到，經天元優化的 MatMul 計算，發揮出了該處理器 90% 以上的計算性能。
總 結
本文以 Batch Norm 為例介紹了推理計算圖的具體實現，以及 MatMul 在 CPU 上的優化細節。作為 CPU 推理優化的基石，最優的推理計算圖是實現高性能 CPU 推理的前提條件，極致性能的 MatMul 計算基礎算子將為實現卷積計算中的 Im2col 和 Winograd 提供性能保障。
在后面的文章中，將在詳細介紹卷積計算中 Im2col 和 Winograd 的優化細節。
參考文獻

1. Anatomy of High-Performance Matrix Multiplication
2. Fusing batch normalization and convolution in runtime
3. Arm? Cortex?-A76 Software Optimization Guide
   了解更多信息可訪問：
   ? MegEngine WebSite： https://megengine.org.cn
   ? MegEngine GitHub（歡迎 Star）： https://github.com/MegEngine/MegEngine

總結

以上是生活随笔為你收集整理的MegEngine计算图、MatMul优化解析的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。