1. 项目概述为稀疏计算“瘦身”与“提速”在深度神经网络推理、推荐系统、科学计算和图分析这些现代计算的核心场景里矩阵乘法Matrix Multiplication, MM是当之无愧的“算力吞噬者”。一个有趣且普遍的现象是这些工作负载中的数据往往极其“稀疏”——矩阵中超过90%甚至99%的元素都是零。想象一下你面前有一张巨大的网格上面星星点点地分布着几个有意义的数字其余全是空白。传统的计算硬件比如我们熟知的CPU、GPU甚至是专为矩阵计算优化的脉动阵列如Google的TPU在处理这种稀疏矩阵乘法Sparse MM, SpMM时就像让一台重型卡车去运送几件小包裹绝大部分动力都浪费在了搬运“空气”零值上。它们无法智能地跳过这些无效计算导致性能低下能耗却居高不下。因此设计一个能“看见”并“忽略”零值的专用SpMM硬件加速器就成了提升整个系统效能的关键。这不仅仅是做一个更快的计算单元更是要设计一套全新的计算哲学只对有意义的数据非零元素做功。本次分享的项目正是基于这一理念提出并实现了一个基于行积Row-Wise Product算法的稀疏矩阵乘法硬件加速器并配套了一个精巧的基于操作计数的负载均衡优化方案。我们的核心武器是业界通用的压缩稀疏行CSR格式它就像为稀疏矩阵量身定做的“压缩包”只存储非零元素及其位置信息极大地节省了存储和带宽。简单来说这个项目解决了两个核心痛点第一如何高效地执行稀疏计算我们选择了行积算法它相比内积和外积在稀疏场景下避免了复杂的索引匹配减少了对大容量片上存储的依赖数据访问模式也更友好。第二如何让多个计算核心PE高效协同工作我们提出了一种智能的矩阵分块策略不是简单地把矩阵切成大小相等的块而是根据每个块实际需要进行的乘法累加MAC操作数量来划分确保每个PE的活儿差不多一样多从而避免“忙的忙死闲的闲死”的局面。经过实测在我们设计的32个处理单元PE的加速器上相比传统的128x128和256x256规模脉动阵列平均取得了13.6倍到47.9倍的性能提升。而我们的负载均衡方案也比简单的固定分块或仅按非零元素数量分块带来了最高8.5%和6.3%的额外性能收益且预处理开销几乎可以忽略不计平均仅占加速器计算时间的0.06%。接下来我将深入拆解这个加速器的设计思路、硬件架构、负载均衡算法的细节并分享在实现过程中踩过的坑和总结出的实战经验。2. 核心思路与方案选型为什么是“行积”“CSR”在设计稀疏矩阵乘法加速器时摆在我们面前的有几条主流技术路径内积Inner-Product、外积Outer-Product和行积Row-Wise Product。每种方法都有其适用的场景和固有的优缺点。我们的选择并非凭空而来而是基于稀疏计算的特性和硬件实现的成本效益深思熟虑的结果。2.1 传统方法的困境内积与外积在稀疏场景的短板内积法是最直观的矩阵乘法理解方式结果矩阵C的每个元素C[i][j]是矩阵A的第i行与矩阵B的第j列的点积。但在稀疏世界里问题来了A的第i行和B的第j列可能都只有零星几个非零元素。为了计算它们的点积硬件必须进行复杂的“索引匹配”Index Matching即快速找出A行和B列中那些列索引与行索引相等的非零元素对。这个过程需要大量的比较和查找操作在硬件上实现既复杂又耗时特别是当矩阵非常稀疏时为了一次有效的乘法可能要做很多次无效的索引比对。外积法则是将计算视为一系列秩-1矩阵的叠加矩阵A的每一列与矩阵B的对应行做外积得到一个个部分结果矩阵最后将它们累加起来。这种方法避免了内积的索引匹配问题因为A的列和B的行直接相乘所有配对都会产生结果。但它的代价是会产生大量的中间部分结果。这些部分结果矩阵本身也可能是稀疏的但为了累加它们必须被暂存在片上或片外存储器中。对于大规模矩阵这部分存储开销会成为巨大的瓶颈严重制约了硬件的可扩展性和能效。2.2 行积法的优势为稀疏计算量身定制行积法采取了不同的视角。计算A x B C时对于矩阵A中的每一个非零元素A[m][n]位于第m行第n列我们将其视为一个标量。然后取出矩阵B的第n行一个行向量用标量A[m][n]去乘这个行向量。这样一次“标量-向量”乘法会产生结果矩阵C中第m行的一部分结果。遍历A中所有非零元素并对同一位置C[i][j]的结果进行累加就得到了最终的C。这种方法在稀疏计算中展现出三大显著优势无需索引匹配计算直接由A元素的列索引n驱动去取B的第n行。这是一个简单的、确定性的寻址操作没有内积法中的复杂匹配逻辑。中间结果存储压力小部分结果直接累加到最终结果矩阵C的对应行中不需要像外积法那样生成并存储大量的中间矩阵。片上只需要维护正在计算和累加的C行数据存储需求大幅降低。数据局部性友好计算过程按A的行顺序进行对A的访问是顺序的CSR格式存储。同时对B的访问也是以行为单位。这种访问模式与CSR格式的存储方式按行压缩天然契合有利于充分利用缓存和内存带宽减少随机访问。注意行积法的一个潜在缺点是对于A中同一行的不同非零元素它们会多次读写C的同一行进行累加可能带来一定的读写冲突或同步开销。但在我们的架构中通过合理的PE任务划分和片上缓冲设计可以有效管理这一问题。2.3 数据格式的选择为什么是CSR确定了行积算法接下来需要选择稀疏矩阵的存储格式。常见的格式有坐标格式COO、压缩稀疏行CSR、压缩稀疏列CSC等。我们选择CSR主要基于以下几点考量广泛支持CSR是科学计算和众多软件库如SciPy、CuSPARSE中的标准稀疏格式之一生态成熟便于与现有软件栈集成。行访问高效CSR格式由三个数组构成values存储非零值column_indices存储每个非零值所在的列号row_ptr存储每行第一个非零值在values数组中的起始位置。这种结构使得按行遍历非零元素变得极其高效通过row_ptr快速定位完美匹配行积算法。存储紧凑相比COO格式需要存储每个非零元的行号和列号CSR通过row_ptr压缩了行索引信息对于行数很多的矩阵节省了可观的存储空间。相比位图Bitmap格式需要为每个元素分配1bit标记是否为零在极高稀疏度下CSR的压缩率更高。在我们的设计中输入矩阵A和B都以CSR格式存储输出矩阵C也动态生成CSR格式。这形成了一个从输入到输出的全CSR流水线最大化地减少了数据转换的开销和存储压力。3. 硬件加速器架构深度解析有了行积算法和CSR格式的理论基础我们来看看如何将它们落地为一个高效的硬件电路。我们的核心是一个可扩展的多PEProcessing Element架构每个PE都围绕行积-CSR计算模型精心设计。3.1 单PE核心计算单元设计单个PE是执行稀疏行积计算的基本单元。它的设计目标是高效地实现图5所示的算法。图6展示了其硬件架构框图我们可以将其分解为几个关键模块来理解数据供给与索引管理PE内部设有多个缓冲区用于存放输入矩阵A、B的CSR数据NVA,CIA,RPA和NVB,CIB,RPB以及正在构建的输出矩阵C的CSR数据NVC,CIC,RPC。两个关键的索引寄存器idxA和idxB是计算的“指挥棒”。idxA指向当前正在处理的矩阵A的非零元素在NVA/CIA数组中的位置。idxB的初始值由CIA[idxA]决定它指向矩阵B中对应行行号CIA[idxA]的第一个非零元素在NVB/CIB中的位置。乘法器从NVA[idxA]和NVB[idxB]读取操作数进行乘法运算。结果写入与CSR动态构建 这是设计中最精妙也最具挑战性的部分。由于输出矩阵C也是CSR格式且非零元素在每行内必须按列索引升序排列我们不能简单地将部分结果写到任意位置。PE需要动态地在NVC和CIC数组中为新的结果元素“插入”一个条目或找到已存在的条目进行累加。写入控制单元Writing Control Unit这是该模块的大脑。它持续比较CIB[idxB]当前B元素的目标列号和CIC[idx_search]当前在C的CSR数组中搜索位置的列号。四状态机与处理基于比较结果写入控制单元触发图4所示的四种情况之一情况a (CIC CIB)当前搜索位置的列号小于目标列号说明目标位置还在后面。只需递增搜索索引idx_search继续向后查找。情况b (CIC CIB)当前搜索位置的列号大于目标列号说明找到了应插入的位置。此时需要触发右移单元将NVC和CIC数组中从idx_search到idx_empty-1的所有元素向右移动一位腾出一个空位然后写入新结果。情况c (CIC CIB)找到了匹配的列索引直接将乘积累加到NVC[idx_search]上即可。情况d (idx_search idx_empty)搜索指针已到达数组末尾的空位无需移动直接在此空位写入新结果。这个动态插入/累加机制保证了输出CSR数组始终有序是行积算法能高效生成CSR格式输出的关键。累加与更新 当写入控制单元确定好写入位置对应情况b, c, d后乘法器的结果会被送入累加器与NVC中对应位置的值如果是情况c或初始值0情况b, d相加然后写回。同时CIC对应位置被更新为CIB[idxB]。完成一次标量-标量乘加后idxB递增以处理B当前行的下一个非零元素。当B的当前行处理完毕idxA递增处理A的下一个非零元素。3.2 多PE扩展与共享缓冲单个PE的性能有限。为了提升吞吐量我们可以实例化多个相同的PE。在多PE设计中共享片上缓冲输入矩阵A、B的CSR数据通常存储在较大的共享片上缓冲区SRAM中所有PE通过一个互联网络如交叉开关或总线访问这些数据。这避免了数据在多个PE间的重复存储。私有计算资源每个PE拥有自己独立的乘法器、累加器、索引寄存器组和写入控制单元。这样多个PE可以并行处理矩阵的不同部分。关键挑战——写冲突由于所有PE共同构建同一个输出矩阵C的CSR数组如果多个PE试图同时写入或修改C的同一行就会发生冲突。我们的解决方案是通过矩阵分块和任务调度确保在任何一个调度周期内不同的PE负责计算C中完全不同的行区域从而从根本上避免写冲突。这正是下一节负载均衡要解决的核心问题之一。4. 负载均衡优化从“均分数据”到“均分工作”使用多个PE并行计算最理想的状态是大家同时干完活。但如果任务分配不均整体完成时间就会由最慢的那个PE决定。传统的矩阵并行计算往往简单地将矩阵按行或按列等分成大小相同的块固定分块Fixed Tiling。但在稀疏矩阵乘法中这行不通因为一个块里非零元素多计算量就大非零元素少计算量就小。更复杂的是计算量并不单纯取决于非零元素的数量。4.1 问题本质为什么非零元素数量不等于计算量考虑一个简单的例子矩阵A的一个非零元素A[i][j]需要与矩阵B的第j行所有非零元素相乘。如果B的第j行有10个非零元那么A[i][j]就会引发10次乘加运算MAC。如果B的第j行只有1个非零元则只引发1次运算。因此分配给一个PE的计算负载应该是它负责的A的子块中每个非零元素A[m][n]所对应的B的第n行中非零元素数量的总和。我们称之为操作计数Operation Count。基于非零元素数量Element Count的分块只考虑了A子块的“重量”却忽略了与之配对的B子块的“密度”必然导致负载不均。我们的目标是基于精确的操作计数来进行分块。4.2 基于操作计数的矩阵分块算法我们的负载均衡方案分为两步水平划分和垂直划分最终将矩阵A切分成一个网格状的瓦片Tile并为每个瓦片配对B矩阵的一个行块。水平划分Horizontal Division of Matrix A目标将矩阵A的行划分成若干组使得每组一个水平条带包含的非零元素总数尽可能相等。方法直接利用CSR格式中的RPA数组。RPA[i1] - RPA[i]就是第i行的非零元数量。我们按行顺序累加这个数量当累计值达到“总非零元数/PE数”的阈值时就切一刀。这样确保了每个水平条带承载的“数据量”大致均衡。垂直划分Vertical Division of Matrix A目标在水平划分的基础上对每个水平条带再进行垂直划分使得最终每个小瓦片Tile的操作计数尽可能相等。方法这是算法的核心。我们需要计算每个垂直划分候选点的操作计数。操作计数预计算在主机CPU上我们预先对矩阵A进行采样例如采样10%的行快速估算出opstotal。对于采样到的A的每一行中的每个非零元素A[i][j]我们查看RPB数组得到RPB[j1] - RPB[j]即B矩阵第j行的非零元数量。这个数量就是A[i][j]将引发的MAC操作数。对所有采样元素求和并放大即可估算出总操作数opstotal。划分决策目标是让每个PE分配到的操作数接近⌈opstotal / N_PE⌉。我们在垂直方向上移动划分线计算划分线左侧所有A瓦片对应的操作计数之和当该和值最接近目标值时确定划分点。由于A是CSR格式垂直划分实际上是在CIA数组上确定一个列索引范围使得落入该范围的A的非零元素对应的操作计数之和满足要求。矩阵B的划分矩阵B的划分是跟随矩阵A的垂直划分而定的。A的每个垂直条带对应一个列索引范围[col_start, col_end)。那么矩阵B中行号落在这个范围内的所有行就被划归为同一个B的行块Tile Row。这样保证了计算的一致性A瓦片(i,j)只与B的行块j相乘。4.3 瓦片配对与调度策略划分完成后我们得到了一个N x N的A瓦片网格假设有N个PE以及N个B的行块。如何将它们分配给N个PE并安排执行顺序呢我们采用了一种对角线调度策略其伪代码如图8所示。对于第k个调度轮次sched_round我们将瓦片对(tile_A(i, (ik-1) % N), tile_B((ik-1) % N))分配给第i个PEi从0到N-1。这种策略的精妙之处在于避免资源争用在同一个调度轮次内任何两个PE所访问的A瓦片来自A的不同行区域和B的行块B的不同行集合都是互不重叠的。这意味着它们不会同时读写输出矩阵C的同一行彻底消除了PE间的写冲突。同时它们访问的输入数据区域也不同减少了对共享输入缓冲区的访问竞争。负载均衡由于分块是基于操作计数的每个瓦片对的计算量理论上相近。再经过这种轮转调度每个PE在各轮次中承担的工作量也趋于平均。图9用一个2PE的简单例子清晰展示了整个过程先水平划分A使非零元数相等再垂直划分使操作计数相等各为4最后通过两轮调度让PE0和PE1分别计算(A00, B0)、(A01, B1)和(A11, B1)、(A10, B0)完美避免了冲突。4.4 预处理开销与实战考量你可能会问这个分块调度算法需要在主机CPU上执行会不会带来很大的额外开销我们的评估结果给出了令人放心的答案开销极低几乎可忽略。即使在32个PE的配置下矩阵分块和调度所带来的预处理延迟平均也只占加速器本身计算时间的0.055%。对于4PE和16PE配置这个开销更是低至0.00033%和0.01015%。这是因为采样而非全量计算我们只对A矩阵采样10%的行来估算操作计数大大减少了预计算量。计算简单分块算法主要涉及的是对CSR数组RPA,CIA,RPB的扫描和简单算术运算计算复杂度是O(nnz)对于稀疏矩阵来说非常高效。一次计算多次使用对于固定的矩阵A和B分块方案只需计算一次便可以反复用于多次SpMM计算例如在迭代算法中。实操心得在实际硬件设计时主机CPU计算出的分块信息每个PE负责的A瓦片的行/列范围、B的行块范围、以及各PE的初始索引寄存器值可以通过控制寄存器或DMA描述符的方式传递给加速器。加速器上电或任务开始前由驱动加载这些配置PE即可开始并行计算。这套方案将复杂的负载均衡逻辑放在软件端保持了硬件架构的简洁和高效。5. 性能评估与对比分析理论再优美也需要实战检验。我们构建了一个周期精确的模拟器其行为基于实际的FPGA原型实现并与广泛用于DNN加速的脉动阵列128x128和256x256规模进行了性能对比。5.1 实验设置与基准测试对比平台脉动阵列模拟了128x128和256x256两种规模采用权重静止Weight Stationary数据流时钟频率设为1 GHz。这是对标Google TPU的典型配置。我们的SpMM加速器评估了4PE、16PE、32PE三种配置。其中16PE和32PE的设计是与TPUv4i和TPUv1中的一个矩阵乘法单元MXU进行等功耗Iso-power对比的尽管我们的设计使用32位浮点数而TPU使用低精度整数这使我们的对比条件更为保守。测试集采用了来自SuiteSparse矩阵集的14个具有不同维度和稀疏度的真实矩阵作为基准见表2。我们根据这些矩阵的维度和非零密度合成了方阵A和B进行乘法测试。评估指标主要对比执行完整SpMM运算的延迟Latency。5.2 加速效果碾压性的优势表3的对比结果令人振奋。我们的32PE-SpMM加速器相比128x128和256x256的脉动阵列平均取得了47.9倍和13.6倍的加速比。即使是16PE版本也分别有8.8倍和2.5倍的平均加速。4PE版本在面对小矩阵时可能不占优但在处理大规模矩阵如wg,m2,az等时性能依然优于脉动阵列。核心原因在于计算效率的本质不同脉动阵列其固化的数据流无法跳过零值操作。无论矩阵多稀疏它都需要遍历整个稠密矩阵的空间进行计算计算量是M*N*K。随着矩阵增大延迟线性增长。我们的SpMM加速器只对非零元素执行MAC操作计算量正比于nnz(A) * avg_nnz_per_row(B)。在稀疏度很高例如95%的场景下有效计算量远小于稠密计算量。因此我们的加速器性能随矩阵规模增长较慢可扩展性远优于脉动阵列。当然对于非常小的稠密矩阵如wv,p3脉动阵列因其极高的时钟频率和规整的数据流可能仍有优势。但这恰恰说明了专用稀疏加速器的必要性——它针对的是脉动阵列不擅长的主流稀疏负载。5.3 负载均衡方案的有效性验证我们进一步比较了三种分块策略的性能固定分块FT简单地将矩阵按行、列数均等分割。基于非零元素数量的分块ET仅考虑使每个块的非零元素数量均衡。我们的基于操作计数的分块OT。表4的结果证实了我们的分析OT方案在绝大多数情况下优于FT和ET。在32PE配置下OT相比FT和ET平均带来了8.5%和6.3%的性能提升在个别矩阵如pg上提升甚至超过25%。提升幅度随着PE数量增加而变大因为PE越多负载不均衡的负面影响越容易被放大。ET方案虽然比FT好但因为它忽略了B矩阵行稀疏度的差异其分块仍然不是最优的。我们的OT方案通过精确预估计算负载实现了更精细的均衡。注意事项分块策略的性能也受数据分布模式影响。在极少数特定分布下如测试中的f3矩阵4PE配置FT可能偶然取得更好效果。这是因为固定分块有时会巧合地匹配计算负载分布。但OT方案的优势在于其稳定性和普适性它不依赖于数据的特定分布模式而是通过计算来保证均衡因此在各种真实、未知的稀疏矩阵上都能提供可靠的高性能。6. 硬件实现细节与FPGA原型为了让设计落地我们在Xilinx ZCU106 FPGA开发平台上实现了4PE版本的SpMM加速器原型。6.1 实现结果表1总结了实现的关键数据时钟频率综合后达到214.27 MHz。对于包含复杂控制逻辑如动态CSR构建的FPGA设计来说这是一个合理的频率。资源利用率在ZCU106平台包含大量DSP、BRAM、LUT资源上我们的设计占用了中等水平的资源。这证明了设计的可实现性。功耗实测功耗约为6.18W。需要强调的是这是在FPGA上的实现。如果采用ASIC工艺进行流片通过定制化设计、优化关键路径、使用更先进的存储单元时钟频率可以大幅提升达到GHz级别同时功耗会显著降低。FPGA原型的主要价值在于验证功能正确性和架构可行性。6.2 关键硬件设计技巧CSR缓冲区的设计输入A、B的CSR缓冲区通常设计为多bank SRAM以支持多个PE并行访问。输出C的CSR缓冲区设计更为关键因为它需要支持随机插入右移操作。一种高效的设计是将其组织为多个独立的行缓冲Row Buffer每个PE被分配一组互不相交的行这样就能并行写入而无冲突。行缓冲内部可以采用双端口SRAM或小型的寄存器文件配合地址生成逻辑来管理有序插入。写入控制单元的优化图4中的四种情况判断和对应的操作递增搜索、右移、累加、直接写入需要在一个或几个周期内完成。可以通过并行比较器和简单的状态机来实现。右移操作可以通过一个多路选择器MUX链或一个小的移位寄存器来实现对于CSR行缓冲来说移动的数据量一行内的非零元通常不大开销可控。数据通路与流水线乘法器和累加器可以流水化。关键路径往往在写入控制逻辑和CSR缓冲区的访问上。需要仔细平衡流水线级数确保在目标频率下能稳定工作。7. 常见问题、挑战与未来展望在设计和实现这个加速器的过程中我们遇到并克服了一系列挑战也看到了未来可以继续优化的方向。7.1 实战中遇到的典型问题与解决思路输出CSR动态构建的硬件开销问题在硬件中动态维护一个有序的、不断增长的CSR数组特别是右移操作是面积和时序上的挑战。解决思路我们采用了“行缓冲”架构。每个PE或一组PE负责输出矩阵C的一个连续行区间。这样每个行缓冲是独立的规模较小右移操作只发生在一个缓冲区内硬件实现简单。通过负载均衡分块我们确保了行区间的划分从算法上避免了跨行缓冲的写冲突。不规则访存与带宽压力问题稀疏计算导致对输入矩阵B的访问是间接的、不规则的通过CIA索引跳转这可能造成缓存命中率低内存带宽利用率不佳。解决思路我们在硬件中设计了较大的输入缓冲区并采用预取Prefetching机制。主机CPU在传输数据时可以根据分块信息将每个B的行块连续地放入加速器的片上缓冲。在计算时一个B的行块会被反复使用与A瓦片中的多个元素相乘这提供了良好的数据局部性缓解了带宽压力。负载均衡预计算的精度与开销权衡问题100%精确计算操作计数需要扫描整个A矩阵开销大。解决思路采用采样法。我们发现只需对A矩阵随机采样约10%的行就能足够准确地估算出总操作计数和分布从而做出高质量的分块决策。这使预处理开销降低了近一个数量级而性能损失微乎其微。7.2 未来可能的优化方向支持更广泛的数据类型与稀疏格式当前设计针对32位浮点数和CSR格式。未来可以扩展支持INT8、BFLOAT16等低精度格式以适配DNN推理以及支持块稀疏Block Sparse、结构化稀疏等格式进一步提升存储和计算效率。更自适应的负载均衡当前的分块是静态的、在计算前确定的。对于动态稀疏模式如在迭代法中稀疏模式变化可以探索轻量级的运行时负载监测与动态任务迁移机制。与片上网络NoC集成当PE数量扩展到成百上千时PE间以及PE与共享缓冲间的通信将成为瓶颈。设计一个高效的、支持不规则稀疏数据访问模式的NoC是下一步的研究重点。软件栈与编译器支持要让硬件发挥最大效能需要强大的软件支持。开发能够自动将稠密矩阵运算转换为稀疏格式、调用本加速器、并进行高效分块调度的编译器中间件和运行时库是推向实际应用的关键。回过头看这个项目的核心价值在于它提供了一套从算法、架构到系统优化的完整稀疏计算解决方案。它不只是一个更快的计算单元而是一个考虑了数据格式、并行效率、负载均衡和实际硬件约束的协同设计。在数据稀疏性日益普遍的时代这类专为稀疏性而生的加速器设计思路或许比单纯追求更高峰值算力的稠密加速器更能代表未来高效计算的方向。