YCL AI計算庫在resnet50上的優化

1. 背景介紹

英特爾第四代至強可擴展處理器（代號Sapphire Rapids，簡稱SPR）上引入了全新的加速引擎AMX（Advanced Matrix Extensions)，通過指令集層面的支持來顯著加速深度學習算法中的Tensor計算。AMX針對廣泛的硬件和軟件優化，進一步增強了前一代矢量神經網絡指令VNNI和BF16，推出了AMX_INT8和AMX_BF16指令，從一維向量計算發展到二維矩陣計算，最大限度的利用計算資源。由于神經網絡訓練推理涉及大量的矩陣運算，AMX的引入將大幅提高AI性能。

2022年的云棲大會上，阿里云推出了搭載倚天710芯片的ECS服務器，受到了業界的廣泛關注。該服務器CPU芯片基于ARM Neoverse N2架構，支持ARM v9 指令集，最高支持128核。業界權威性能報告指出，其并行計算能力在CPU服務器中非常搶眼，并且極具性價比，有潛力作為昂貴的GPU服務器的替代品。然而，Yitian710 作為平頭哥第一代ARM通用芯片，在AI場景與X86相比，軟件生態與推理性能都存在一定的短板，本文旨在通過倚天AI計算庫的優化，打造適合ARM架構的軟件平臺，提升倚天性能。

1.1 問題

倚天710目前主要依賴開源社區提供對AI場景的軟件支持，存在以下幾個問題：
（1）目前主流AI軟件生態對X86架構適配更好，各種推理場景性能表現更優，倚天缺乏相關的軟件生態，推廣依靠具體業務場景定制優化與ARM生態支持，效率低下
（2）倚天SIMD位寬受限，與X86 ICL，SPR相比有較大劣勢，需要任務調度充分發揮倚天物理核算力優勢
（3）ARM AI生態演進考慮不同架構兼容，迭代速度慢，且缺乏倚天微架構針對性調優，不利于充分發揮倚天在AI場景優勢

1.2 策略

YCL（YiTian Compute Library）為平頭哥數據中心解決方案團隊開發的一款高性能AI計算庫，該庫基于ARM開源的ACL(ARM Compute Library)實現。ACL是一個用于機器學習和計算機視覺的高性能 C++ 庫。它提供了一系列優化的算法和操作，可在 ARM CPU、GPU 和 DSP 上執行。YCL在ACL的基礎上，針對倚天硬件架構的特性，做了深入的適配與優化，通過調度算法優化、GEMM拆分、底層算子融合、BF16精度優化等方法，實現包括Core、Support、Graph、Backends等不同層次模塊的性能優化，并通過oneDNN標準接口對接上層推理框架如tensorflow，pytorch，實現上層計算任務不感知。經測試，集成優化版本的tensorlfow在mlperf resnet50評測中性能提升超過40%，目前該版本已集成到cap2自動化測試系統。

2. YCL計算庫架構

自Tensorflow 2.5版本開始，已經有了對 oneDNN 的實驗性支持，此后ARM開源社區在oneDNN的backend增加了ARM實現，來加速 AArch64 CPU 的性能。如下圖1所示，Tensorflow framework 將上層的計算任務分解成各個算子，調用底層實現以提升性能。Tensorflow默認調用Eigen實現各算子如GEMM。Intel提供oneDNN加速庫用于實現基于X86 backend的kernel實現，對于arm backend，oneDNN調用ACL來使用ARM向量指令以提升性能。YCL即為 arm backend替代ACL的計算庫，專門針對倚天SoC架構特性做出優化。

圖1 YCL計算庫在tensorflow中的位置

圖2 YCL計算庫架構

YCL在框架架構上與ACL基本相同，如圖2所示,綠色部分為倚天710軟件架構，接口層實現了常用的AI算子，在使用每個算子之前通過配置(configure)接口設置輸入數據、數據類型、計算模式、算子評估、權重數據packing、調度方法等，配置完成即可啟動運算過程(run)，該過程首先將計算任務劃分成子任務，并為不同子任務分配線程并發計算，最后各線程調用計算kernel完成各自計算任務，主線程合并計算結果完成最終的計算。

3. 優化方法

本文從以下4個方面針對倚天架構做優化：

3.1 子任務劃分，利用倚天710各級cache提升數據吞吐

圖3 YCL中矩陣運算子任務劃分

矩陣運算(GEMM)一般為當前AI推理任務中的主要計算來源，很多加速庫也是重點優化提升GEMM計算性能。當前學術上提升在CPU上提升GEMM性能的主要思路為：將A矩陣在M方向劃分為寬度為Lvh的子塊，將B矩陣在N方向上劃分寬度為Lvw的子塊，然后根據L1 cache大小確定K方向(Kc)的值，然后確定每個子塊計算順序，使用多核完成計算。

YCL子子任務劃分也采用上述方法，但是在設計子任務是考慮倚天Cache結構與物理核優勢，首先根據分配的倚天core數與任務大小，確定最終分配的線程數，如果計算任務較小，則考慮少分配線程數n，可以降低線程調度產生的開銷。

然后根據任務大小與計算單元緩存確定子任務數，原則是劃分后的子任務可以一次性存入緩存，提升數據存取速度；然后如圖3所示，將矩陣A與B分別劃分成Akj(Lvh x kc), Bki(kc x Lvw)子矩陣，每個線程分別計算Ck=Akj x Bki ，通過調節倚天SIMD寄存器布局，降低數據重復訪問，獲得最優性能。

3.2 任務調度

設計兩級線程與子任務對應表，如圖4所示，其中level 0子任務平均分配到各線程上執行，level 1為多余的子任務首先緩存在buffer中，等到有線程空閑時執行。該方法有三個好處，第一，子任務劃分利用了多核系統緩存，子任務在單核中執行效率最高；第二，線程與任務對應，充分利用線程資源，先完成的線程繼續執行level 1子任務，減少線程長尾效應影響；第三，各子任務在整體任務中數據連續存儲，提升cache命中率。

3.3 底層算子融合

在tensorflow中有大量的eltwise計算，然后結果輸入激活函數的操作，該部分在獨立計算，不依賴其他操作，可以在底層將eltwise計算的中間結果保存在寄存器中，然后緊接中做ACT，以eltwise(sum) + ReLU為例，可以在oneDNN與YCL中將這部分功能合并，如下圖4所示。

圖4 底層算子融合

3.4 BF16算子計算

倚天710采用armv9架構，指令集支持bfloat16矩陣計算，單個bfmmla指令可以計算一個2x2大小的矩陣，理論性能相比float指令可以提升4倍，下表為倚天710不同精度下指令的理論算力。因此，使用BF16指令可以在保證精度的前提下大幅提升性能。

YCL計算庫在不改變tensoflow框架的前提下，實現了從float到bfloat的簡單切換。在oneDNN層面，將卷積算子做了改造，首先將輸入tensor配置為bfloat16格式，然后將輸入數據從float格式轉換成bfloat16格式數據，改轉化可能會有overhead，最終實現采用simd 匯編實現，將轉換完成的數據導入oneDNN原始的memory中，并釋放臨時buffer。

3.5 性能評估

倚天710單個SoC有128個core，且都是物理核，有獨立的L1與L2cache，我們使用阿里云ecs.c8y.8xlarge來測試YCL計算庫的性能，為了充分發揮倚天物理核算力，測試采用MLperf resnet 0ffline模式將CPU壓力打到最大，測試開啟BF16，具體測試命令為：

./run_local.sh tf resnet50 cpu --scenario Offline

其中g8i為intel SPR實例，其tensorflow安裝方式與python依賴如下（通過 pip install tensorflow==2.11.0安裝）

使能BF16方法如下：

export DNNL_VERBOSE=1
export TF_ENABLE_ONEDNN_OPTS=1
export ONEDNN_DEFAULT_FPMATH_MODE=BF16

運行benchmark，查看log，如果存在avx512_core_amx_bf16，代表使能AMX_BF16來加速矩陣運算

測試均在32c下進行，如下圖5所示，使用優化后的YCL計算庫resnet50性能提升45%。

圖5 倚天710 resnet50優化前后性能對比

4. 安裝與使用方法

目前YCL計算庫已經適配了tensorflow 1.15與2.9兩個版本，通過打patch方式支持tensorflow源碼編譯安裝。