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源码编译安装,相关的编译与集成方式可以参考《倚天AI实践》该部分在后续文章中发出。