AI的张量世界，直面维度灾难

16并不是一个很大的数字。那么，在64个时钟周期内，计算一个卷积神经网络（Convolutional Neural Network，简称CNN），使其能够从16个输入通道、3*3张量卷积中生成具有16*16张瓦片图（tile）的16个输出通道，将需要多少MAC（Multiply-Accumulate Unit，乘加器）呢？

如果不使用快速算法，那么答案将会是至少9216个。通常，9216个MAC可以用来构建一个96*96的脉冲阵列，但是计算一个96*96的矩阵乘法（matrix multiplication，简称MM）将延迟至少96个时钟周期。也就是说，构建这个脉动阵列需要96个96*96矩阵乘法。

欢迎来到AI的张量世界。是时候直面维度灾难了。

张量vs矩阵

本文标题灵感来自Charles F.Van Loan教授在2010年所作的关于张量的模展开矩阵的演讲中的一句话—— 所有张量暗地里都希望它们是矩阵！

这一陈述表明了张量研究领域想要先把张量展开成矩阵，然后利用成熟的矩阵理论框架对其进行研究。同时，尽管主流AI框架认为张量是最基本的数据类型，但在行业标准实践中，也常将张量展开为矩阵，以便利用高度优化的矩阵乘法（MM）库和矩阵乘法加速器（MMAs）库。通常，AI领域认为矩阵是特殊的张量。

一些主流的看法认为：

1. 可以利用MM中的并行性和数据共享模式。

2. 和张量相比，矩阵更适用于可编程硬件。

3. MM有原生硬件加速方法，即脉动阵列。

然而，

1. 从结构上来说，CNN和MM是相同的。展开张量来利用和MM相同的并行性和数据共享模式是完全没有必要的。

2. 在从存储层次的下层向上层运行的过程中，由于时间局部性，矩阵会递归式分块；由于空间局部性，矩阵会压缩打包。最终，矩阵会变成微面板，即小块行或列，并为软件微内核或GPU着色内核所用。

3. 在谷歌TPU v1：256*256脉冲阵列概述中，笔者解决了方形脉动阵列的可扩展性问题。自此，使用多个小型脉动阵列似乎成了主流。因此，在矩阵成为能够为脉动阵列所用的方形矩阵前，仍然需要被分块和压缩打包。

所以，由展开张量得到的矩阵实际上被分块和压缩打包成合适的结构，用于高性能运行（如下图所示）。前者被称为平滑展开，后者为分块展开。由于数十年的研究、构建和利用分块矩阵框架用于高性能MM执行，矩阵成为CNN和AI中最常见的默认数据类型。

MM的分块展开

主流看法认为，CNN的特征图可看作是某个矩阵中的一列，卷积核（filter）则可展开成为一列中的某些连续矩阵元素。由于平滑展开，特征图中的相邻像素值被空间和时间重复引用的功能被抹去了。

拯救分块张量

MM中的递归矩阵分块展开有一个重要特征，即当矩阵逐渐接近硬件的裸金属（bare metal）时，高水平的时空位置得以保留。那么，考虑CNN中的数据位置会不会像张量形式中的位置一样被保留将会非常有趣？

将特征图分成瓦片图，并保持其平铺结构，以便于重复利用卷积核、将快速算法用于瓦片图以及达到细粒度的SIMD（Single Instruction Multiple Data，单指令多数据流）并行是一个合理的做法。当接近裸金属时，张量应该保持张量的形式，并保持平铺结构和完整特征图中的数据局部性。

此外，必须解决输入特征图和输出特征图中都存在的局部模式。前者使得多个输入特征图的数据能够在计算多个输出特征图时实现共享。后者则使得更多方面能够共享输入特征图。不同之处在于，由于生成输出像素值并不需要用到所有的输入特征图，所以不能共享所有的特征图。更何况，把所有特征图放在芯片上并不实际。总的来说，所有维度都要按一定程度进行分块，这使得张量能够被分成更小的单元，即分块张量。

分块张量即接口为张量的张量，和分块矩阵相对应。和分块矩阵被高性能计算（high-performance computing，简称HPC）机群用于MM相似，分块张量可用于解决维度灾难和保留CNN的数据局部性。张量包，相当于微通道或MM中的方形子矩阵，是最基本的张量单元。它必须按照原子级运行，以利用所有维度的空间局部性。由张量包构成的张量块也是一种张量单元。它必须在整体计算单位和外部记忆之间转移，以促进张量包之间的时间局部性。

原子级张量包运行可根据最小充分输入通道量来生成具有最小充分大小瓦片图的最小充分输出通道量。由于张量中的维度灾难，即使在每个维度的张量包都很小时，上述张量包运行也能发挥很大作用。它可以在张量块中迭代或并行运行，来解决更严峻的问题。该方法将在下文中半正式地详细阐述。

CNN中的分块张量

输入为A、输出为C的CNN包括多个输入特征图（input feature maps，简称IFMs）和多个输出特征图（output feature maps，简称OFMs）。这些可以看作是三维张量，用x，y表示特征图，用输入深度w索引IFMs，用输出深度z索引OFMs。为了达到细粒度SIMD并行和利用具有特殊局部性的快速算法，每个特征图将会进一步沿着x维和y维分成瓦片图。相应的索引(tx, ty)分别表示输入瓦片图和输出瓦片图。对于每一对唯一的IFM w和OFM z，有卷积核B(w, z)，通常由3*3的卷积核参数构成。输入张量和输出张量成为瓦片图的分块张量，如下图所示：

CNN中的瓦片图分块张量

利用张量理论符号，可以更加完整、准确地表示平铺CNN：

冒号表示采用该维度中的所有数据。如A(:, :, w)表示采用了IFM w中的所有瓦片图，可表示IFM w。

CNNs在结构上和MMs相同

主流想法认为必须把张量平滑展开为矩阵，以便于利用MMs的并行性和数据共享模式。然而，这实际上会适得其反。如下图所示，从结构上来说，CNNs和MMs的并行性和数据共享模式是相同的。这也是MMs在CNNs中广泛使用的原因。

CNNs在结构上和MMs相同

因为C行是独立计算的，所以A行可以被划分。同理，因为像素值可以被单独计算，所以特征图可以被分成瓦片图。

瓦片图中和MM相同的并行性和数据共享模式保持完整。

在同一特征图中，输出瓦片图的卷积核相同，但输入瓦片图不同。在输出维度中，输出瓦片图的输入瓦片图相同，但卷积核不同。下表利用张量理论符号更加完整地解释了MMs和CNNs的相同之处。

深度也需要被划分

现在，假设芯片的记忆容量是有限的，并且需要基于芯片上的一大堆数据，在IFMs、OFMs以及卷积核中充分发挥CNNs的与MM相同的数据共享模式。这些数据的几何结构会是怎样的呢？如下图所示，这需要有充分的(x, y)来重复利用卷积核，充分的w使得有足够的输入数据被共享，以及充分的z来有效共享输入数据：

芯片缓冲器必须有充分的信息范围涵盖所有维度

从上图中可以观察到，需要拆分w维，z维，x维以及y维。这一步已经完成了。这一步是为了保证大量的输入数据和输出数据会分别涵盖最小充分IFMs量和最小充分OFMs量。引入新的索引tw和tz，分别用于定义IFMs组和OFMs组。在IFM组tw中，定义一个输入张量包A(tx, ty, tw)，作为一组对齐的瓦片图(tx, ty)。同理，在OFM组tz中，定义C(tx, ty, tz)作为一组对齐的瓦片图(tx, ty)。B(tw, tz)表示一组卷积核，每个卷积核用于由一个来自组tw的IFM和一个来自组tz的OFM构成的一对。这就是卷积核张量包。输入，输出和卷积核张量构成张量包中的分块张量，如下图所示：

CNN中张量包中的分块张量

张量包中和MM相同的并行性和数据共享模式保持完整。

软件中的张量微核或者硬件中的张量包引擎可以用于处理卷积一个输入张量包A(tx, ty, tw)和一个卷积核张量包B(tw, tz)的原子级操作，如下图所示：

原子级张量包操作

假设输入和输出瓦片图分别为6*6和4*4，并把IFM组和OFM组大小设置为8。那么在不使用快速算法的情况下，运行一次3*3卷积，需要2304个MAC和4个时钟周期。这2304条并行运行平均分布在所有维度中，包括特征图中的x维和y维、输入深度w和输出深度z。而如果将特征图分为瓦片图，如4*4规模，那么快速算法，如Winograd就可以得到应用，则这2304条并行运行仅需576个MAC就可以完成了。

张量递归划分

在计算单元中，张量包是基本单元。为了保留tile函数间的数据位置和平铺结构，在完整张量、张量包和张量块之间引入一个中间块，以完整涵盖想要放到芯片上的张量包。张量包是tile函数间的时间位置的基本单元，也是当有充分计算单元、芯片记忆容量和供应带宽时的空间局部性的基本单元。在x，y和z维上，一个张量块可定义为(bx, by, bz)。

下文展示了一个张量递归划分的例子。整个张量是一个分块张量，该分块张量可分成4*4*2个分块张量，每一块可再进一步分成1*1*8个分块张量，每一块包含4*4个瓦片图。其中(dx, dy)是分支，用于找到每个瓦片图中的像素值，dz是索引，用于定义每个组中的特殊特征图。

张量递归划分，用于保留特征图中的局部性

张量块在存储器中线性存储。但其存储顺序并不需要在意。对于一个张量块，有两种不同的分块展开顺序，一种不断优化将张量块存储到DRAM（Dynamic Random Access Memory，动态随机存取存储器）中，另一种则不断优化将张量块展示到计算单元中。软件和/或硬件中的换位机制需要在运行过程中从一种格式转换到另一种格式，如下图所示。

张量块换位

由换位导致的延迟可以由双缓冲来解决。 当一个张量块被展示到计算单元中时，它被展开成张量包，并按(tx, ty)排序，如此一来，张量包就可以按相应的计算单元在(x, y)域并行运行。

所有张量暗地里都想做自己

Van Loan教授在一场关于分块张量的分块展开的谈话中讲到：

分块展开保留了数据结构和局部性……我认为，分块在张量计算中的影响力最终将会和其在矩阵计算中的影响力一样大。

在多维时间局部性中，张量被分块展开成张量块，以便于芯片上不同维度的数据实现共享。张量块则可以进一步分块展开成张量包，每个张量包有各个维度的最小充分信息范围。在本文所举的例子中，运行这样的张量包将会需要2304条并行运行。在张量块中，和MM相同的并行性和数据共享模式应用为张量包。

就本文所了解的来看，本文可能在计算机史上第一次讨论了将分块张量的概念应用到裸金属中。跨x维和y维的更好的、规模化的并行性可用于处理高清视频/图像。基于CNN革命性的能力和其在张量中的深刻根基，CNN可能成为第一个分块张量杀手锏。

所有张量暗地里都想要成为自己，即分块张量，并释放潜能，在AI硬件领域实现突破。