1. 背景介紹

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

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

1.1 問(wèn)題

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

1.2 策略

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

2. YCL計(jì)算庫(kù)架構(gòu)

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

圖1 YCL計(jì)算庫(kù)在tensorflow中的位置

圖2 YCL計(jì)算庫(kù)架構(gòu)

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

3. 優(yōu)化方法

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

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

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

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

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

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

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

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

3.3 底層算子融合

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

圖4 底層算子融合

3.4 BF16算子計(jì)算

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

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

3.5 性能評(píng)估

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

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

其中g(shù)8i為intel SPR實(shí)例，其tensorflow安裝方式與python依賴(lài)如下（通過(guò) 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來(lái)加速矩陣運(yùn)算

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

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

4. 安裝與使用方法

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