浅谈MindSpore之一：训练性能为何能够近似线性提升-一一网

【摘要】作为华为推出的全场景自研AI框架，MindSpore支持在x86&ARM64 CPU/GPU及昇腾处理器等多种平台训练推理。而其在昇腾上的训练性能，相比于TF/Torch在GPU上训练，能够有接近线性的提升，这背后蕴含着设计开发人员在软硬件协同优化的深思熟虑，那么具体又是如何去设计优化的呢？此篇文章着重从硬件层面浅谈了其个中缘由。

萌新一枚，之前断断续续使用了一段时间的MindSpore框架，在和伙伴们交流的过程中，很多人发现MindSpore相比于tensorflow/pytorch在训练性能上有着近乎线性的提升，而且是在昇腾芯片上才表现出来，遂一直好奇这其中的个中缘由。最近看了《昇腾AI处理器架构与编程》这本书后，茅塞顿开，恍然大悟。奥秘藏在MindSpore精妙设计与昇腾芯片达芬奇架构上面。今天就来简单探讨下。

达芬奇架构

昇腾AI处理器的计算核心主要由AI Core构成，负责执行标量、向量和张量相关的计算密集型算子。AI Core采用了达芬奇架构，其基本结构如图1所示。包括了三种基础计算资源：矩阵计算单元(Cube Unit)、向量计算单元(Vector Unit)和标量计算单元(Scalar Unit)。这三种计算单元从名字上就可以看出，分别对应了张量、向量和标量三种常见的计算模式。此外，在矩阵计算单元和向量计算单元内部还提供了不同精度、不同类型的计算模式。AI Core中的矩阵计算单元目前支持INT8和FP16的计算；向量计算单元目前支持FP16和FP32以及多种整型熟的计算。

AI Core架构图.jpg

图1 AI Core架构图

为了配合AI Core中数据的传输和搬运，围绕着三种计算资源还分布式地设置了一系列片上缓冲区，比如用来放置整体图像特征数据、网络参数以及中间结果地输入缓冲区(Input Buffer, IB)和输出缓冲区(Output Buffer, OB)，以及提供一些临时变量地高速寄存器单元。

在AI Core中，IB之后设置了一个存储转换单元(Memory Transfer Unit, MTE)。这是达芬奇架构地特色之一，主要目的是为了以极高地效率实现数据格式的转换。比如我们使用的GPU要通过矩阵计算来实现卷积，首先要通过Img2col的方法把输入的网络和特征数据以一定格式排列起来，这一步在GPU中是通过软件来实现的，效率比较低下。达芬奇架构采用了一个专门的存储转换单元来完成这一过程，通过固化的硬件电路可以在很短时间内完成整个转置过程。由于类似转置计算在深度神经网络中出现频繁，这种定制化电路模块设计可以提升AI Core的执行效率，从而实现不间断的卷积运算。