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

转载地址：https://bbs.huaweicloud.com/f...

作者：搬砖小强abc

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

达芬奇架构

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

图1 AI Core架构图

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

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

计算单元
  计算单元是AI Core中提供强大算力的核心单元，主要包含矩阵计算单元、向量计算单元、标量计算单元和累加器。

  矩阵计算单元 & 累加器      

   矩阵计算单元和累加器主要完成矩阵相关计算。常见的深度神经网络算法大量的使用了矩阵计算，达芬奇架构对矩阵计算进行了深度的优化与定制。如图2表示矩阵A和矩阵B相乘C=A×B，其中M是矩阵A的行数，K是A的列数以及B的行数，N是矩阵B的列数。

图2 矩阵乘法示意图
(1) CPU计算

传统CPU中计算矩阵乘法的典型代码如代码1所示:

for (int m=0;n<;m++)

for (int n=0;n

for (int k=0;k

cm += Am*Bk;

代码1 CPU计算矩阵乘法

该程序用三个循环进行一次完整的矩阵乘法计算，如果在一个CPU上执行需要M×K×N个时钟周期才能完成，当矩阵非常庞大时执行过程极为耗时。

而且在CPU计算过程中，矩阵A是按照行扫描，矩阵B按照列扫描，而矩阵存储一般是按行连续存储，而内存读取是具有很强的数据局部特征的，即当读取内存中某个数时会打开内存中相应的一整行并且把同一行中所有数读取出来，这种内存读取方式对矩阵A很友好而对B却显得很不友好，为此需要将B的存储方式转成按列存储。因此在矩阵计算中往往通过改变某个矩阵的存储方式来提高效率。

（2）GPU计算
在深度神经网络中实现计算卷积过程，关键的步骤是将卷积运算转化为矩阵计算。在CPU中大规模矩阵计算往往成为性能瓶颈，而矩阵计算在深度学习中又极为重要。为了解决这个矛盾，GPU采用通用矩阵乘法(GEMM)的方式来实现矩阵乘法。例如实现一个16×16矩阵与另一个16×16矩阵的乘法，因为需要同时计算16×16组向量的乘加运算，需要安排256个并行线程，每个线程都可以独立计算完成结果矩阵中的一个输出点。假设每个线程在一个时钟周期内可以完成一次乘加运算，而每组向量乘运算需要16次乘加运算，则GPU完成整个矩阵计算需要16个时钟周期，这个时延是传统GPU无法避免的瓶颈。而昇腾AI处理器针对这个问题做了深度优化。

（3）昇腾计算
达芬奇架构在AI Core中特意设计了矩阵计算单元作为昇腾AI处理器的核心计算模块。通过精巧设计的定制电路和极致的后端优化手段，矩阵计算单元可以用一条指令完成两个16×16矩阵的相乘运算(因M=K=N=16，标记为16^3，这也是Cube名称的由来)，等同于在极端时间内进行了16^3=4096个乘加运算，并且可实现FP16的运算精度。如图3所示，矩阵计算单元完成C=A×B的矩阵运算时，会事先将矩阵A按行存放在IB中，同时将矩阵B按列存放在IB中，通过矩阵计算单元计算后得到的结果矩阵C按行存放在OB中。在矩阵相乘运算中，矩阵C的第一元素由矩阵A的第一行的16个元素和矩阵B的第一列的16个元素通过矩阵计算单元子电路进行16次乘法和15次加法运算得出。

图3 矩阵计算单元计算示意图

矩阵计算单元中共有256个计算子电路，可以由一条指令并行完成矩阵C的256个元素计算，即大致一个时钟周期就可以完成一个矩阵乘法。笼统的说，N个矩阵乘法运算只需N条指令，这正比例资源关系自然就解释了MindSpore在昇腾处理起上的训练性能成线性提升的原因。

在有关矩阵的处理上，通常在进行完一次矩阵乘法后还需要和上次结果进行累加，以实现类似C=A×B+C的运算。矩阵计算单元的设计上也考虑到了这种情况，为此专门在计算单元后面增加了一组累加器单元，可以实现将上次中间结果与当前结果累加，总共累加次数可由软件控制，并在累加完成之后将结果写入OB中，在卷积计算过程中，累加器可以完成加偏置的累加计算。

矩阵计算单元这以空间换时间的设计理念，可以快速完成16×16的矩阵相乘，但当超过16×16大小的矩阵利用该单元进行计算时，则需要市事先按照特定数据格式进行矩阵的存储，并在计算过程中以特定的分块方式进行数据的读取。而这种切割分块操作，昇腾从软硬件两方面都做了设计，MindSpore从软件架构上实现了自动最优分割的接口化设计(具体是如何设计的这里先挖个坑，后续再填坑)，而在硬件上的优化设计，如图4所示。矩阵A展示的切割和排序方式称作“大Z小Z”，直观地看就是矩阵A地各个分块之间按照行的顺序排序，称为“大Z”方式；而每个块的内部数据也是按行排列，称为“小Z”方式。与之形成对比的是矩阵B的各个分块之间按行排序，而每个块的内部按列排序，称为“大Z小N”的排序方式。按照矩阵计算的一般法则，昇腾AI处理器内部专用电路可以实现将如此排列的A、B矩阵相乘之后得到矩阵C，而矩阵C会呈现出各个分块之间按照列排序，而每个块内部按照行排序的格式，称为“大N小Z”的排列方式。

图4 存储格式的要求

  向量计算单元

  未完待续......