基于国产众核处理器的深度神经网络算子加速库优化

基于不同硬件设备的算子加速库已经成为深度学习框架不可或缺的一部分,能够为大规模训练或者推理任务提供数倍的性能加速.当前的主流算子库都是基于GPU架构开发的,与其他异构设计并不兼容;SWDNN算子库是基于申威26010开发的,无法充分发挥升级后的申威26010 pro处理器的性能,也不能满足当前GPT-3等大型神经网络模型对大容量内存和高访存带宽的需求.文中面向申威26010 pro处理器体系结构的特点和大型神经网络模型的训练需求,提出了基于多核组的三级并行和神经网络算子任务调度方案,在满足大型模型训练内存需求的同时,提高了并行效率和整体计算性能;提出了三级异步流水机制和计算访存重叠的访存优化方法,显著缓解了神经网络算子的访存性能瓶颈.基于以上方法,文中构建了基于申威26010 pro处理器的SWTensor多核组算子加速库,在自然语言处理模型GPT-2上进行了实验,结果表明,其典型计算密集型算子和访存密集型算子在单精度浮点计算性能和访存带宽上分别达到了理论峰值的90.4％和88.7％.