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

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

緊接這一過程，天元將運(yùn)行 OptPass ParamRedistributePass，該 OptPass 會(huì)將上述 Batch Norm 轉(zhuǎn)化而來的 Elemwise 中的 scale 融合到 convolution 的權(quán)重中，具體實(shí)現(xiàn)原理如下：

最后天元將運(yùn)行 OptPass FuseConvBiasNonlinPass，該 OptPass 會(huì)將計(jì)算圖中的 Convolution+Elemwise 轉(zhuǎn)化為天元內(nèi)部實(shí)現(xiàn)的 ConvBias Op 中，同時(shí)，它還會(huì)設(shè)置 ConvBias Op 中 NonelineaMode 參數(shù)。
如此便完成了從 Convolution+Batch Norm+Relu 到 ConvBias 的轉(zhuǎn)換，整體轉(zhuǎn)換過程如下圖：

實(shí)驗(yàn)驗(yàn)證圖優(yōu)化之后的性能
在推理之前，如果要對(duì)完成訓(xùn)練的模型進(jìn)行圖優(yōu)化，則需要在模型 dump 的期間進(jìn)行，當(dāng)然，也可以在 SDK 運(yùn)行模型之前進(jìn)行。下面是在模型 dump 時(shí)進(jìn)行圖優(yōu)化的代碼：
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 接口無需運(yùn)行直接編譯

pred_fun.trace(data, net=xor_net)

使用 trace 類的 dump 接口進(jìn)行部署

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

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

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

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

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

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

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

可以看到，經(jīng)天元優(yōu)化的 MatMul 計(jì)算，發(fā)揮出了該處理器 90% 以上的計(jì)算性能。
總 結(jié)
本文以 Batch Norm 為例介紹了推理計(jì)算圖的具體實(shí)現(xiàn)，以及 MatMul 在 CPU 上的優(yōu)化細(xì)節(jié)。作為 CPU 推理優(yōu)化的基石，最優(yōu)的推理計(jì)算圖是實(shí)現(xiàn)高性能 CPU 推理的前提條件，極致性能的 MatMul 計(jì)算基礎(chǔ)算子將為實(shí)現(xiàn)卷積計(jì)算中的 Im2col 和 Winograd 提供性能保障。
在后面的文章中，將在詳細(xì)介紹卷積計(jì)算中 Im2col 和 Winograd 的優(yōu)化細(xì)節(jié)。
參考文獻(xiàn)

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

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。