稀疏张量算子的硬件加速

摘要

工业界针对稠密张量的硬件设计已经接近成熟，一些具有代表性的DSA已经达到较高的运算效率，如tensor core，systolic array。而稀疏张量的带宽和算力要求与稠密张量不同，有着自己独特的坐标表示，负载差异较大，导致硬件加速设计难度较大或者效果欠佳。 

本篇文章首先介绍稀疏算子的基本算法，然后分析稀疏张量的压缩格式，最后以ExTensor为例，介绍稀疏张量处理器的基本结构和设计难点。

引言

随着神经网络规模的不断增长，深度神经网络的稀疏加速是近两年的一个重要研究方向。权重剪枝（pruning）、低位宽量化（quantization）、ReLU激活函数、Drop-out等技术的广泛使用，都在神经网络的权重或特征值中引入了大量的零元素。而一些特定领域的科学问题，其数据本身就是高度稀疏化的。零元素与任意元素相乘的结果总是零，零元素出现的乘法项对多项式的结果没有贡献，因此，以矩阵或向量乘法为核心的稀疏张量算子中，存在大量的冗余存储和无效计算。这为我们开发专门的稀疏张量加速硬件，或是优化现有的硬件以高效地支持稀疏算法，提供了契机。

不同领域的稀疏张量，其稀疏性有很大的不同。一般的深度神经网络，经过权重剪枝、低位宽量化等手段稀疏化之后，非零元素比例多在1~100%，其中不同层的稀疏率一般也有所不同。而在推荐系统、社交媒体等场景下，数据中非零元素比例一般低于10^-5%（不难想象，社交媒体上与自己是好友的账号只占所有账号非常小的一个比例）。下图展示了不同任务场景中，张量的稀疏率从10^-6%到100%不等。

图1 不同任务场景中张量数据的稀疏率[1]

从后文中的讨论可以看到，硬件针对迥异的稀疏率，无论是存储方式还是计算模式都有着完全不同的优化方向。此外，在特定的应用场景中，稀疏加速往往可以和该领域数据的其他分布特点（比如非零元素的聚团、数据位宽的长尾分布等）结合起来，一起作为硬件优化的考量因素。

常见的稀疏张量算子

所谓张量，一般指高维数组。向量和矩阵可以认为是维度的数目分别为1和2的特殊张量。一般计算机领域的稀疏张量算子中张量的维度都在3维、4维或更低（即使是高维数据，单个算子一般也只会操作几个维度）。这里我列举了一些常见的张量算子，更多的算子可以参考BLAS[2]、TACO[3]等线性代数加速库的文档。

通用矩阵-向量乘法（GEMV），将矩阵A中的每一行A与向量b做内积，再将获得的向量与另一个向量c加权求和，获得输出向量y。数学表达式为：

在有的论文中（尤其是很多和稀疏加速相关的论文中），GEMV算子不包括与向量c加权求和的部分（即alpha=1和beta=0）。如果A矩阵是稀疏的，我们也可以将该算子记为SpMV。如果向量b也是稀疏的，则记为SpMSpV。

通用矩阵-矩阵乘法（GEMM），将矩阵A与矩阵B相乘，再将获得的矩阵与另一个矩阵C加权求和，获得输出矩阵Y。数学表达式为：

如果矩阵A是稀疏的，则一般记为SpMM。如果矩阵B也是稀疏的，则记为SpMSpM。

GEMV和GEMM在深度神经网络中极为常见，它们占据了神经网络的主要工作负载。前者包括全连接、RNN等，后者包括CONV（IMG2COL展平后）。

张量-向量乘法（TTV），将d阶张量A在某个维度Mk上的每个纤维A(m1,…,:,md)与向量b做内积，可以获得一个（d-1）阶张量y。学术论文中，比较常见的是k=d的情形。对于3阶张量A，它与向量b在第三个维度上的TTV（mode-3 TTV），数学表达式为：

如果张量A稀疏，则记为SpTTV。除非A非常稀疏（以至于大量参与计算的纤维全部都是零元素），输出张量y总是密集的。

张量-矩阵乘法（TTM），将TTV中的向量b扩展成矩阵B，此时输出张量y也从(d-1)阶扩展成d阶。同样地，我们给出k=d=3时的定义，mode-3 TTM数学表达式为：

如果张量A是稀疏的，则记为SpTTM。

更多的稀疏张量算子，例如CONV（卷积，大家都很熟悉，一般可展开成GEMM），SDDMM（稀疏采样的密集矩阵乘法，多见于Sparse Attention网络）、MTTKRP（矩阵化张量Khatri Rao积，多见于张量分解）等，限于篇幅，这里就不展开了。

无论是什么算子，实现硬件稀疏加速的关键在于高效地定位对最终计算有贡献的非零元素在输入张量中的坐标，并将其从存储器中索引并搬运到计算单元上，以避免不必要的数据传输与计算。

稀疏张量的压缩表示

稀疏张量为了高效地存储，一般都被以压缩格式存储。很多时候，压缩格式的选择需要结合应用场景、硬件的数据流、存储器带宽等综合考虑。常见的一些压缩格式如下图所示：

图2 稀疏张量的常见压缩表示方法（a为矩阵，b为张量）[4]

DIA等一些特定的压缩表示对数据的分布有一定要求，否则压缩效率会显著下降。也有相当多的硬件稀疏加速的论文会在常见压缩格式的基础上改造，设计符合自己需求的压缩格式。

数据的压缩效率和索引效率需要权衡。我们对比上图中的CSF（字面意思是压缩稀疏纤维，也有译作压缩稀疏张量）格式和CSR格式。后者CSR大概是最常见的稀疏矩阵压缩方法，而前者是最近几年提出的一种新的压缩格式[5]。但本质上两者是统一的。CSR对矩阵的每一行进行了去零压缩，保存了每个元素的列坐标（以相对前一个元素的偏移量的形式）。因此，每行有几个非零元素，就会有几个坐标。而每一行的起始地址指针，都保存在row_ptr中。与行内压缩不同，即使某一行是空行，CSR也会保存该行的地址指针（通过和前一个指针没有相对偏移表示这是空的指针）。所以一个M行的矩阵，CSR需要M+1个行指针。就这个意义上来讲，我们可以说CSR在第一个维度（行）上没有做压缩，在第二个维度（列）上做了压缩。而CSF与CSR的本质区别在于，CSF在每一个维度上进行了压缩。也就是说，如果某一个行是空行，CSF不会存储这一行的地址。很显然，如果数据非常稀疏，CSF可以比CSR提供更高的压缩效率。但反过来，CSR的行指针是密集排布的，意味着更高的索引效率。

计算性能和存储器带宽需要权衡。上文中的ZVC格式（有时也叫Binary Map、Bit Map、Bit Mask等）是一种很适合计算的压缩格式。两个ZVC格式的矢量做乘法，我们可以将它们的mask做逻辑AND，就可以立刻获得发生有效计算的元素的坐标。但如果矩阵稀疏度过高，ZVC中mask中大量的零（尽管每个0只有1bit）占用的存储器空间并不能省略，因此存储效率较低。

案例介绍：ExTensor

ExTensor是Kartik Hedge等来自NVIDIA和UIUC的研究者在MICRO-2019上发表的一个稀疏张量算子加速器架构。ExTensor并不是第一个针对稀疏算子优化的架构，但它是第一个面向通用（而非某一个特定算子）的张量架构，它基于坐标的相交（intersection）提出了一种对稀疏算子的普遍抽象，并开发了对应的硬件架构。ExTensor并不局限于上述的某一种压缩表示方法，它是一种普遍的抽象，可以用于CSR、CSF、COO等（尽管作者只在CSF上开展了研究）。

正如我前面所说的，稀疏算子的计算中重点是找出有效的计算，并定位相关非零元素的索引。基于此，ExTensor把稀疏算子的计算分成了两个部分：1）坐标的相交，2）对应非零元素的计算。下图是一个矩阵乘的例子：

图3 矩阵乘法的例子[1]

我们以这个例子为例，介绍intersection的过程。对于上述矩阵乘，K维度最后被归约到了一点，是发生intersection的维度。下图是intersection的示意图。

图4 矩阵乘法的intersection示意图[1]

其中左边和右边是矩阵乘的两种计算流，分别是Output Stationary和A Stationary。两种数据流都可以提取出intersection，我们就以左边的Output Stationary为例。此时I、J两个维度在外侧，K在内侧，因此我们先对于每一对I、J的K纤维进行intersection。对于矩阵A，I=0时的K纤维为[1,2]，而矩阵B中J=1时的K纤维为[0]，两束纤维的交集为空集，因此在输出矩阵的Z(0, 1)处没有需要有效计算。矩阵A的I=0时的K纤维[1,2]与矩阵B中J=2时的K纤维[0,1]的交集非空，为[1]，意味着A(I=0,K=1)与B(K=1,J=2)会发生计算。此时计算单元需要取得对应的元素a和h，并进行相应的计算。以此类推，我们遍历矩阵A和B中的I和J，并对对应的每一束K纤维进行集合的相交运算，就可以获得所有参与有效计算的元素坐标。

特别的是，该抽象不要求非零元素必须是标量，因为我们这一步比较、计算的是坐标。即使非零元素是一个子张量（即上图中的a、b、c等均为一棵子树），我们仍然可以进行intersection，从而跳过不必要的子张量的读取与计算。这个特点有助于我们在Tile层面（较高的维度上）进行稀疏优化，也使得我们可以对张量的多个维度进行层次化的intersection。

对应到硬件上，ExTensor提出了对应的intersection硬件模型，如下图所示。

图5 Intersection的硬件模型[1]

其中Scanner从压缩表示的张量的meta信息中获得需要进行intersection的坐标流（比如上面矩阵乘的例子中的矩阵A的不同I时的每个K纤维），Intersect负责将两个Scanner产生的坐标流求交集，其算法原理如上图右侧框中所示。当然该算法实际实现起来是低效的（因为实际坐标不需要连续取），研究者还进行了一些优化，感兴趣的读者可以阅读ExTensor的论文了解。

图6 Coordinator的硬件结构[1]

基于上述硬件模型，研究者设计了一个较实际的硬件结构Coordinator，如上图所示。其中Scanner A和B从对应的meta信息（存储在MD Storage中）不断地产生张量A和张量B的坐标流（Stream A和Stream B）发送给Intersect。后者按照前述算法进行Intersect，丢弃掉无效计算，输出有效计算对应的非零元素的坐标（pos A和pos B）。坐标用于Data Storage中获取非零元素的数据，此外Intersect需要产生计算结果的写回坐标（图中coord，可以理解成前面的例子中的I和J）。最后，非零元素被发射到相应的计算单元进行计算（或者是下一级的Coordinator进行另一个维度的intersection）。

研究者在Coordinator的基础上构建了ExTensor加速器，并在Near-DRAM、Last Level Buffer、Inside PE三个层面上进行了intersection。该加速器同典型的服务器CPU进行了性能的对比，在SpMSpM、SpMM、TTV、TTM、SDDMM等测试基准上获得了平均3.4x，1.3x，2.8x，24.9x，2.7x的提速。

总结与展望

ExTensor以及其他很多已经发表的硬件加速相关工作，共同推动稀疏张量计算朝着更高性能、更高能效前进。对稀疏性的探索并不容易，但机会与挑战总是并存。展望未来，在算法-硬件协同设计优化，针对存算一体、模拟技术等新型器件的稀疏优化，更通用的稀疏张量加速器架构等，也许会是研究的新方向。 

参考文献

[1] Austin Derrow-Pinion, Jennifer She, David Wong, et al. ETA Predictionwith Graph Neural Networks in Google Maps. 2021

[1] K. Hedge, et al., ExTensor: An Accelerator for Sparse Tensor Algebra. https://dl.acm.org/doi/pdf/10.1145/3352460.3358275 . 

[2] BLAS：BLAS (Basic Linear Algebra Subprograms).

[3] TACO: Home - Documentation - The Tensor Algebra Compiler (TACO) 

[4] E. Qin, et al., Extending Sparse Tensor Accelerators to Support Multiple Compression Formats. https://arxiv.org/ftp/arxiv/papers/2103/2103.10452.pdf 

[5] S. Smith, et al., Tensor-Matrix Products with a Compressed Sparse Tensor. https://dl.acm.org/doi/pdf/10.1145/2833179.2833183 .

 往期推荐

关于壁仞科技研究院

壁仞科技研究院作为壁仞科技的前沿研究部门，旨在研究新型智能计算系统的关键技术，重点关注新型架构，先进编译技术和设计方法学，并将逐渐拓展研究方向，探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则，将积极投入各类产学研合作并参与开源社区的建设，为相关领域的技术进步做出自己的贡献。

扫码关注我们