YCL AI計算庫在resnet50上的優(yōu)化

1. 背景介紹

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

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

1.1 問題

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

1.2 策略

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

2. YCL計算庫架構(gòu)

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

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

圖2 YCL計算庫架構(gòu)

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

3. 優(yōu)化方法

本文從以下4個方面針對倚天架構(gòu)做優(yōu)化：

3.1 子任務(wù)劃分，利用倚天710各級cache提升數(shù)據(jù)吞吐

圖3 YCL中矩陣運(yùn)算子任務(wù)劃分

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

YCL子子任務(wù)劃分也采用上述方法，但是在設(shè)計子任務(wù)是考慮倚天Cache結(jié)構(gòu)與物理核優(yōu)勢，首先根據(jù)分配的倚天core數(shù)與任務(wù)大小，確定最終分配的線程數(shù)，如果計算任務(wù)較小，則考慮少分配線程數(shù)n，可以降低線程調(diào)度產(chǎn)生的開銷。

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

3.2 任務(wù)調(diào)度

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

3.3 底層算子融合

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

圖4 底層算子融合

3.4 BF16算子計算

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

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

3.5 性能評估

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

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

其中g(shù)8i為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

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

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

圖5 倚天710 resnet50優(yōu)化前后性能對比

4. 安裝與使用方法

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