摘要：Host排程模式下，GE將模型中運算元的執行單元劃分為Host CPU執行與Device（昇騰AI處理器）執行兩大類。
本文分享自華為雲社群《深度解讀昇騰CANN小shape運算元計算最佳化技術，進一步減少排程開銷》，作者：昇騰CANN。
GE（Graph Engine）將模型的排程分為Host排程與下沉排程兩種模式。經過上期的介紹我們知道，在模型為靜態shape時，由於其輸入tensor shape固定不變，在編譯時就能確定所有運算元的輸入輸出shape，並能提前完成模型級記憶體編排、tiling計算等Host排程工作，因此採用模型下沉排程方式可以將整個模型下沉到Device側執行，從而提升模型排程效能。
與之對應的，在模型為動態shape的情況下，由於輸入tensor shape不確定，需要在上一個運算元完成shape推導後，才能確定下一個運算元的輸入shape等資訊，因此無法將整個模型下沉執行，只能採用Host排程模式。
1 Host排程簡介及最佳化背景
所謂Host排程，是指模型的排程主體位於Host CPU，由CPU完成逐運算元排程。一個運算元的排程任務為kernel執行準備必要引數，通常包含shape推導、tiling、記憶體分配、launch等。
Host排程模式下，GE將模型中運算元的執行單元劃分為Host CPU執行與Device（昇騰AI處理器）執行兩大類。對於卷積、MatMul等對算力要求高的運算元，會被劃分到Device執行；而由於shape資訊在Host CPU維護，Shape、Reshape等運算元更適合被劃分到Host CPU執行；除此之外，還有一些運算元，在shape較小時，計算量也很小，排程開銷往往大於運算元的實際計算開銷，就需要考慮如何儘可能減少排程開銷帶來的效能影響。
圖1 網路拓撲片段
圖1是一段網路拓撲片段示例，按照一般的排程機制，Gather、Concat運算元會下沉到Device側計算，Shape、Unsqueeze、Reshape運算元在Host側計算。其執行時序如圖2所示，模型E2E執行耗時除了包含運算元計算的時間外，還包含Host與Device之間的資料複製、運算元下沉排程、Stream同步等開銷，整體執行E2E耗時在毫秒級別。
圖2 最佳化前執行時序
而對於小shape（如shape size小於8）的Gather、Concat，運算元本身在Host側CPU的計算開銷上僅微秒級別，與Device側計算的效能相差無幾。此時下發帶來的額外開銷就顯得比較明顯。針對上述這種shape較小且輸入Tensor記憶體在Host的場景，GE識別將這部分運算元保留在Host側執行，可有效減少排程開銷帶來的效能影響。
2 小shape運算元計算最佳化實現
在圖編譯流程執行到引擎選擇之後，GE選擇在Host側執行的運算元並將其作為錨點，然後向後遞迴查詢計算資料個數小於8的運算元，並將這些運算元的執行引擎修改為Host CPU。針對圖1所示的網路片段，假設shape運算元的輸出的shape size小於8，則Gather、Concat運算元的執行引擎都會被重新整理成Host CPU。最佳化後執行時序如圖3所示，此時模型執行只有運算元計算帶來的開銷，經測試約為10微秒（3ms –> 10us），顯著的提高了E2E執行效能。
圖3 最佳化前後執行前後時序對比
3 最佳化效果
以LLaMA2大語言推理模型為例，符合上述執行引擎重新整理的運算元有Pack、Gather、Concat等約650+個，重新整理前模型E2E耗時約1.062S，重新整理後執行時間最佳化到了1.009S，吞吐提升5%。
4 更多介紹
GE小shape運算元計算最佳化技術的相關介紹就到這裡，歡迎大家關注後續技術分享。如需獲取更多學習資源請登入昇騰社群。
點選關注，第一時間瞭解華為雲新鮮技術~