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源码编译安装。