JITSPMM: Just-in-Time Instruction Generation for Accelerated Sparse Matrix-Matrix Multiplication

文章标题：JITSPMM：通过即时指令生成加速稀疏矩阵-矩阵乘法
作者/机构：
- Qiang Fu, Advanced Micro Devices Inc.
- Thomas B. Rolinger, NVIDIA
- H. Howie Huang, George Washington University

A1 主要贡献

本文旨在解决在多核CPU上进行稀疏矩阵-矩阵乘法（SpMM）的性能瓶颈问题。SpMM在图神经网络（GNNs）等应用中因输入数据规模庞大而备受关注。

核心问题：现有的SpMM解决方案大多采用提前编译（Ahead-of-Time, AOT）方法。AOT编译在处理SpMM时面临三大限制：
1. 不必要的内存访问：由于寄存器分配方案并非针对SpMM的内存访问模式进行优化。
2. 额外的分支开销：由于缺乏运行时信息，无法避免为处理不同输入数据而引入的分支指令，从而导致分支预测错误的开销。
3. 冗余指令：为处理不必要的内存访问（如寄存器溢出和内存加载）和分支控制（如比较和条件跳转）而引入了多余的指令。

研究目标：为了克服AOT方法的局限性，本文提出采用即时（Just-in-Time, JIT）汇编代码生成技术。JIT技术能够在运行时直接生成针对特定SpMM计算实例的汇编代码，从而实现更优化的寄存器分配、利用运行时信息规避分支开销，并减少冗余指令。

创新点与主要贡献：
本文提出了JITSPMM，一个用于在带SIMD扩展的多核CPU上加速SpMM计算的JIT汇编代码生成框架。
- 识别了AOT方法在SpMM计算中 impeding 性能提升的局限性。
- 提出了JITSPMM框架，通过采用JIT汇编代码生成技术来解决这些局限性。该框架包含以下关键技术：
- 将JIT汇编代码生成技术集成到三种广泛使用的工作负载划分方法（row-split, nnz-split, merge-split）中，以实现CPU线程间的负载均衡。
- 提出一种新颖的粗粒度列合并（coarse-grain column merging）技术，利用运行时信息展开性能关键循环，最大化指令级并行。
- 智能分配寄存器以缓存频繁访问的数据，最小化内存访问，并采用精选的SIMD指令来提升算术吞吐量。
- 实现了JITSPMM并进行了全面的性能评估。实验结果表明，与使用Intel icc编译器自动向量化的AOT基准相比，JITSPMM平均性能提升3.8倍；与高度优化的Intel MKL库相比，平均性能提升1.4倍。

A3 背景知识与关键观察

稀疏矩阵-矩阵乘法（SpMM）

SpMM的通用形式与CSR存储格式。通常，SpMM被表述为 $Y = AX$，其中A是一个维度为 $m \times n$ 的稀疏矩阵，X是一个维度为 $n \times d$ 的稠密矩阵，Y是维度为 $m \times d$ 的稠密结果矩阵。稀疏矩阵有多种表示格式，但本文专注于压缩稀疏行（Compress Sparse Row, CSR）格式，因为它被广泛用于供应商库（如Intel MKL 【18，E. Wang, Q. Zhang, B. Shen, G. Zhang, X. Lu, Q. Wu, Y. Wang, E. Wang, Q. Zhang, B. Shen et al., “Intel math kernel library,” HighPerformance Computing on the Intel® Xeon Phi™: How to Fully Exploit MIC Architectures, pp. 167–188, 2014.】）、数据科学工具包（如SciPy 【22，P. Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski, P. Peterson, W. Weckesser, J. Bright et al., “Scipy 1.0: fundamental algorithms for scientific computing in python,” Nature methods, vol. 17, no. 3, pp. 261–272, 2020.】）和GNN框架（如DGL 【9，M. Y. Wang, “Deep graph library: Towards efficient and scalable deep learning on graphs,” in ICLR workshop on representation learning on graphs and manifolds, 2019.】、PyG 【23，M. Fey and J. E. Lenssen, “Fast graph representation learning with pytorch geometric,” arXiv preprint arXiv:1903.02428, 2019.】）。CSR格式将稀疏矩阵存储在三个一维数组中：row_ptr、col_indices 和 vals，如图2所示。为了优化存储空间，CSR有效地将所有非零元素的列索引和值按行顺序分别打包到 col_indices 和 vals 数组中。row_ptr 数组存储了每一行的第一个非零元素在另外两个数组中的偏移量。此外，值得注意的是，在像GNN这样的实际应用中，输入稠密矩阵X的形状通常倾向于“高瘦”型，即n远大于d 【24，O. Selvitopi, B. Brock, I. Nisa, A. Tripathy, K. Yelick, and A. Buluc¸, “Distributed-memory parallel algorithms for sparse times tall-skinnydense matrix multiplication,” in Proceedings of the ACM International Conference on Supercomputing, 2021, pp. 431–442.】。

SpMM的计算过程。算法1展示了SpMM的计算过程。输出矩阵 Y[i][j] 的每个元素是稀疏矩阵A的第i行与稠密矩阵X的第j列的点积。该算法访问 A.row_ptr 数组来确定第i行的起始位置，并遍历 A.col_indices 和 A.vals 数组的一个片段，以处理稀疏行i中的所有非零元素。对于每个非零元素，算法利用 A.col_indices 中相应的列索引（在第5行表示为k）来定位稠密矩阵X中的相应行。此时，算法检索 X[k][j] 的值，将其与 A.vals 中的非零元素值相乘，并将结果加到累加值 ret 中。这个过程一直持续到稀疏行i中的所有非零元素都处理完毕，从而得到 Y[i][j] 的最终值。

Input : Spare matrix $A$ of size $m \times n$, i.e., $A.row\_ptr$, $A.col\_indices$, $A.vals$, and dense matrix $X$ of size $n \times d$.
Output: Dense matrix $Y = AX$.
1 for $i = 0$ to $m$ do
2 for $j = 0$ to $d$ do
3 $ret = 0$
4 for $idx = A.row\_ptr[i]$ to $A.row\_ptr[i+1]$ do
5 $k = A.col\_indices[idx]$
6 $ret += A.vals[idx] * X[k][j]$
7 $Y[i][j] = ret$

现代多核CPU与SIMD指令扩展

在CPU上优化SpMM的价值。尽管图形处理单元（GPUs）和领域特定加速器已被用于执行SpMM等稀疏操作，利用其巨大的计算能力和高内存带宽【19，C. Yang, A. Buluc¸, and J. D. Owens, “Design principles for sparse matrix multiplication on the gpu,” in Euro-Par 2018: Parallel Processing: 24th International Conference on Parallel and Distributed Computing, Turin, Italy, August 27-31, 2018, Proceedings. Springer, 2018, pp. 672–687.】、【25，M. Zhu, T. Zhang, Z. Gu, and Y. Xie, “Sparse tensor core: Algorithm and hardware co-design for vector-wise sparse neural networks on modern gpus,” in Proceedings of the 52nd Annual IEEE/ACM International Symposium on Microarchitecture, 2019, pp. 359–371.】–【27，G. Ortega, F. Vazquez, I. Garc ´ ´ıa, and E. M. Garzon, “Fastspmm: An ´ efficient library for sparse matrix matrix product on gpus,” The Computer Journal, vol. 57, no. 7, pp. 968–979, 2014.】，探索在通用CPU上优化SpMM的技术仍然具有重要价值。这是因为CPU仍然是许多系统中的主要计算资源，提供了巨大的内存容量【28，Z. Gong, H. Ji, Y. Yao, C. W. Fletcher, C. J. Hughes, and J. Torrellas, “Graphite: optimizing graph neural networks on cpus through cooperative software-hardware techniques,” in Proceedings of the 49th Annual International Symposium on Computer Architecture, 2022, pp. 916–931.】、【29，Y. Liu, Y. Wang, R. Yu, M. Li, V. Sharma, and Y. Wang, “Optimizing CNN model inference on CPUs,” in 2019 USENIX Annual Technical Conference (USENIX ATC 19). Renton, WA: USENIX Association, Jul. 2019, pp. 1025–1040. [Online]. Available: https://www.usenix.org/conference/atc19/presentation/liu-yizhi】。例如，在com-Friendster数据集上执行SpMM，输入稠密矩阵有32列时 ，Intel MKL 【18，E. Wang, Q. Zhang, B. Shen, G. Zhang, X. Lu, Q. Wu, Y. Wang, E. Wang, Q. Zhang, B. Shen et al., “Intel math kernel library,” HighPerformance Computing on the Intel® Xeon Phi™: How to Fully Exploit MIC Architectures, pp. 167–188, 2014.】会消耗87.2 GB内存，这超过了当今许多GPU的内存容量（例如Nvidia A100的80 GB内存）。因此，本工作的重点是提升SpMM在处理通常超过加速器内存容量的大型矩阵时的性能。

CPU架构与SIMD指令。现代CPU通过多核架构利用线程级并行来提升整体性能，以解决因晶体管预算限制而导致的单核处理器扩展瓶颈【30，K. Olukotun, B. A. Nayfeh, L. Hammond, K. Wilson, and K. Chang, “The case for a single-chip multiprocessor,” ACM Sigplan Notices, vol. 31, no. 9, pp. 2–11, 1996.】。在软硬件接口方面，许多现代CPU使用x86指令集架构（ISA）。从执行程序的角度看，一个64位x86处理器的内部架构可以概念性地分为几个不同单元，包括指令指针（RIP寄存器）、通用寄存器、状态和控制标志（RFLAGS寄存器）、浮点寄存器、控制和状态（MXCSR）以及SIMD扩展，如图3所示。与我们工作特别相关的是，SIMD（单指令多数据）指令通过使用向量寄存器同时处理多个数据元素来提升性能。SSE2扩展引入了16个128位寄存器（XMM0–15），而AVX/AVX2扩展将寄存器大小扩展到256位（YMM0–15）。最新的AVX512扩展提供了32个更大的512位寄存器（ZMM0–31），允许并行处理多达16个不同的32位浮点数。例如，指令vmulps zmm2, zmm1, zmm0 对ZMM0和ZMM1内的每对单精度值执行乘法操作，并将结果存储在ZMM2中。

采用即时汇编代码生成的动机

AOT编译方法的局限性。大多数现有的SpMM计算解决方案采用提前编译（AOT）方法【9，M. Y. Wang, “Deep graph library: Towards efficient and scalable deep learning on graphs,” in ICLR workshop on representation learning on graphs and manifolds, 2019.】、【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】、【23，M. Fey and J. E. Lenssen, “Fast graph representation learning with pytorch geometric,” arXiv preprint arXiv:1903.02428, 2019.】，如图4(a)所示。这个过程始于用C/C++等传统语言编写的通用SpMM计算的源代码实现。下一步是使用主流编译器（如gcc）将源代码编译成可执行二进制文件。最后，用输入数据执行该二进制文件以计算结果。然而，尽管在SpMM的源代码实现上投入了细致的优化工作【18，E. Wang, Q. Zhang, B. Shen, G. Zhang, X. Lu, Q. Wu, Y. Wang, E. Wang, Q. Zhang, B. Shen et al., “Intel math kernel library,” HighPerformance Computing on the Intel® Xeon Phi™: How to Fully Exploit MIC Architectures, pp. 167–188, 2014.】、【26，G. Huang, G. Dai, Y. Wang, and H. Yang, “Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks,” in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020, pp. 1–12.】，AOT方法仍然无法生成最高效的可执行二进制文件，主要原因有三：1) C/C++编译器依赖于基于启发式规则的寄存器分配方案【31，G. J. Chaitin, “Register allocation & spilling via graph coloring,” ACM Sigplan Notices, vol. 17, no. 6, pp. 98–101, 1982.】，这些方案不足以捕捉SpMM计算的内存访问模式特征，导致不必要的内存访问。2) AOT方法的一个固有局限是无法利用运行时信息，这无意中导致为处理不同输入数据而引入额外的分支指令，从而可能因分支预测错误【32，H. Inoue, M. Ohara, and K. Taura, “Faster set intersection with simd instructions by reducing branch mispredictions,” Proceedings of the VLDB Endowment, vol. 8, no. 3, pp. 293–304, 2014.】带来开销。3) 由于不必要的内存访问（如寄存器溢出和内存加载）和分支控制操作（如比较和条件跳转），引入了冗余的汇编指令，导致执行的指令过多。

JIT汇编代码生成的优势。在本文中，我们提出为SpMM采用即时（JIT）汇编代码生成技术【33，N. Shaylor, “A just-in-time compiler for memory-constrained low-power devices.” in Java Virtual Machine Research and Technology Symposium, 2002, pp. 119–126.】、【34，M. Arnold, S. J. Fink, D. Grove, M. Hind, and P. F. Sweeney, “A survey of adaptive optimization in virtual machines,” Proceedings of the IEEE, vol. 93, no. 2, pp. 449–466, 2005.】，其过程如图4(b)所示。它始于开发一个JIT代码生成器，该生成器不产生“通用”的SpMM代码。相反，当数据在运行时可用时（例如，矩阵的维度），JIT代码生成器可以产生针对正在执行的特定SpMM实例量身定制的汇编代码。与AOT解决方案相比，这种方法可以更好地利用寄存器，减少分支开销，并减少指令数量，如下一小节所示。

单线程标量SpMM性能对比

JIT相比AOT的性能优势验证。为了展示JIT汇编代码生成相对于AOT方法的优势，我们使用三种主流C编译器（gcc、clang和icc）编译了算法1中的SpMM顺序C实现，未使用SIMD指令或多线程。然后，我们将其性能与单线程JIT实现进行比较，后者是我们提出的解决方案（见第四节）的简化版，未使用SIMD指令。我们在uk-2005稀疏矩阵数据集上运行所有实现，该数据集是一个方阵，有3900万行和9.36亿个非零元素。该稀疏矩阵与一个具有3900万行和8列的随机值稠密矩阵相乘。结果如表II所示。可以观察到，即使包含了代码生成开销，JIT方法的执行时间也比三种AOT解决方案快2.1-3倍。

JIT优势的剖析数据。表II中呈现的剖析数据证明了JIT汇编代码生成有效缓解了与AOT解决方案相关的三个限制。1) JIT解决方案有效地利用寄存器存储频繁访问的数据，与AOT编译相比，内存加载次数减少了2.4-2.7倍。例如，算法1中ret的中间结果存储在标量寄存器XMM0-XMM7中，第6行的A.vals[idx]值存储在寄存器XMM31中。这使得JIT程序可以在不多次访问内存的情况下对这些变量进行读写操作。2) 借助运行时信息（稠密矩阵的列数d为8），JIT解决方案能够展开算法1中第2行的for循环。因此，JIT解决方案减少了每次迭代结束时发出的分支指令数量，与AOT方法相比，分支预测错误次数减少了1.2-4.1倍。最小化分支预测错误至关重要，因为它们会导致流水线刷新和重填，从而严重降低整体性能【35，A. Farcy, O. Temam, R. Espasa, and T. Juan, “Dataflow analysis of branch mispredictions and its application to early resolution of branch outcomes,” in Proceedings. 31st Annual ACM/IEEE International Symposium on Microarchitecture. IEEE, 1998, pp. 59–68.】。3) 通过移除与不必要内存访问和分支控制相关的指令，JIT解决方案以减少3.4-4.4倍的已执行指令数完成了相同的计算。

A2 方法细节

JITSPMM框架概述

JITSPMM框架的整体工作流程。在第三节B中，我们展示了JIT汇编代码生成技术在提升单线程标量SpMM性能方面的显著优势。然而，将此技术应用于具有SIMD指令扩展的现代多线程CPU面临若干挑战。第一个挑战是如何在可用的CPU线程间有效划分SpMM操作的工作负载，而不会导致可能影响性能的负载不平衡。第二个挑战是调度生成的指令以充分利用每个CPU核心的指令级并行性。最后，有效地将计算映射到SIMD指令并利用SIMD寄存器需要仔细的考虑和优化。为了应对这些挑战，我们提出了JITSPMM，一个专为在多核CPU上利用SIMD指令扩展进行高性能SpMM计算而设计的即时汇编代码生成框架。JITSPMM的工作流程如图5所示。给定稀疏和稠密矩阵的输入数据，我们的JITSPMM框架遵循一个三步过程。首先，它为计算生成汇编代码，并将其封装为一个jit-function。然后，根据硬件配置创建特定数量的线程。每个线程独立确定其分配的工作负载，并调用jit-function对其分配的数据部分执行计算。一旦所有线程完成各自的工作负载，它们会同步并合并成单个线程。最后，合并后的线程整合各个结果并返回计算的最终输出。在接下来的小节中，我们将分三部分详细介绍JITSPMM框架，每部分解决前面确定的一个特定挑战。

工作负载划分

JITSPMM中的三种工作负载划分策略。在本小节中，我们简要概述JITSPMM中使用的工作负载划分方法，并说明如何将我们的JIT汇编代码生成技术集成到这些方法中。由于输入稀疏矩阵中非零元素在各行之间分布不均，实现高效的工作负载分配对于获得良好的SpMM性能至关重要【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】。在JITSPMM的背景下，我们包含了三种广泛使用的策略【19，C. Yang, A. Buluc¸, and J. D. Owens, “Design principles for sparse matrix multiplication on the gpu,” in Euro-Par 2018: Parallel Processing: 24th International Conference on Parallel and Distributed Computing, Turin, Italy, August 27-31, 2018, Proceedings. Springer, 2018, pp. 672–687.】、【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】：row-split、nnz-split和merge-split，如图6所示。这三种方法中的每一种都选择一个特定的对象在所有CPU线程之间平均分配。具体来说，row-split方法将稀疏矩阵的行平均分配给每个线程，如图6(a)所示。相反，nnz-split方法将相同数量的非零元素分配给每个线程，如图6(b)所示。merge-split策略旨在均衡每个线程的行数和非零元素总数，如图6(c)所示。

动态行分派方法。值得注意的是，row-split策略可能导致工作负载不平衡，因为特定线程可能接收到显著更多或更少的非零元素【10，Q. Fu, Y. Ji, and H. H. Huang, “Tlpgnn: A lightweight two-level parallelism paradigm for graph neural network computation on gpu,” in Proceedings of the 31st International Symposium on High-Performance Parallel and Distributed Computing, 2022, pp. 122–134.】。考虑图6(a)，第二个线程没有接收到任何非零元素，而第四个线程被分配了四个。为了解决这个问题，我们引入了一种基于row-split方法的动态行分派方法。列表1提供了在SpMM计算中为每个线程实现动态行分派的x86实现。首先，分配一个名为NEXT的全局整型变量来跟踪下一个未处理的行ID，并在计算开始前初始化为零。每个线程通过原子地读取NEXT的值并加上批处理大小，使用xadd指令【36，Intel, “Xadd - exchange and add,” 2023. [Online]. Available: https://www.felixcloutier.com/x86/xadd】动态请求一批行 。

JIT与工作负载划分的集成。考虑到这三种工作负载分配方法的AOT实现也容易受到第三节中强调的限制的影响，我们通过为每种方法生成不同的汇编代码来应用我们的JIT代码生成技术。这三种策略中的每一种都为每个线程分配了不同的工作负载分布。因此，它们需要对线程执行的计算进行不同的实现。我们的JITSPMM代码生成过程始于计算单行的汇编代码，这在所有三种方法中都是相同的。对于采用动态行分派的row-split策略，我们加入了列表1中的代码片段用于工作负载请求。这些指令促进了向每个线程动态分配工作。对于nnz-split和merge-split方法，我们引入额外的指令来实现循环，这些循环在一系列连续的行上执行计算，这个范围由这两种方法采用的二分搜索确定【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】。这使得线程能够在分配的段上执行矩阵乘法。值得注意的是，我们的JIT代码生成框架能够无缝集成这些不同的工作负载分配技术。通过动态生成量身定制的汇编代码，我们为每种策略优化SpMM计算，确保计算资源的有效利用并实现高性能并行处理。在接下来的两个小节中，我们的注意力将转向为计算稀疏矩阵的单行生成优化代码，特别是针对算法1的第2-7行。

粗粒度列合并

粗粒度列合并（CCM）技术概述。在本小节中，我们介绍我们的粗粒度列合并（Coarse-grain Column Merging, CCM）技术，该技术旨在借助运行时信息，最大化为计算单个矩阵行生成的汇编代码的指令级并行性（ILP）。通过允许多条指令同时执行，ILP有助于保持处理器的计算资源（如流水线和浮点单元（FPUs））处于繁忙状态【26，G. Huang, G. Dai, Y. Wang, and H. Yang, “Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks,” in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020, pp. 1–12.】、【37，V. S. Pai, P. Ranganathan, and S. V. Adve, “The impact of instructionlevel parallelism on multiprocessor performance and simulation methodology,” in Proceedings Third International Symposium on HighPerformance Computer Architecture. IEEE, 1997, pp. 72–83.】。

\begin{algorithm}
\caption{Coarse-grain Column Merging (CCM) for computing a single row in SpMM with $d=45$}
\textbf{Input}: Row index $i$, $A.row\_ptr$, $A.col\_indices$, $A.vals$, and dense matrix $X$.
\textbf{Output}: Result of the $i$-th row $Y[i][0: 45]$.
\begin{algorithmic}[1]
\State $ret[0: 45] := 0$ \Comment{All elements initialized to zero}
\For{$idx = A.row\_ptr[i]$ to $A.row\_ptr[i+1]$}
\State $k = A.col\_indices[idx]$
\State $ret[0: 45] += A.vals[idx] * X[k][0: 45]$ \Comment{Mul-and-add}
\EndFor
\State $Y[i][0: 45] = ret[0: 45]$ \Comment{Write result back}
\end{algorithmic}
\end{algorithm}

CCM的高层实现。CCM技术旨在通过展开算法1中第2行的for循环，并将整行的所有列合并为单个向量，来优化单行SpMM计算。这里以稠密矩阵X的列数d=45为例。算法2概述了实现CCM的步骤，它将在运行时替换算法1的第2-7行。首先，我们分配并初始化一个长度为45的向量，即ret[0:45]，用于存储结果行的45个元素，如第1行所示，其中:=表示将一个标量值赋给向量中的每个元素。接下来，我们遍历行i的非零（nz）元素列表（第2行）。对于每个nz元素，我们加载其值，并将其与输入稠密矩阵X中相应行X[k][0:45]的每个元素相乘。结果向量被加到累加向量ret[0:45]中（第4行）。最后，累加的向量被写回到结果矩阵Y的第i行（第5行）。必须强调的是，只有当输入稠密矩阵的列数d=45这个运行时信息已知时，CCM才变得可行。相比之下，所有AOT解决方案都受限于使用一个从0到d的循环，从而错失了下文阐述的好处。

CCM带来的益处。CCM技术提供了几个增强ILP的优点。首先，它实现了独立的工作负载处理。由于不同列索引的计算是独立的，用于单向量计算的指令（如算法2中的第1、4或5行）也是独立的。指令之间缺乏数据依赖性，消除了可能的流水线停顿，促进了更好的ILP。其次，CCM通过展开算法1中遍历每个列索引的原始循环（第2行）来消除分支开销，这是通过在运行时知道列数d来实现的。这使我们能够避免执行与循环相关的分支指令。此外，这也消除了与分支预测错误相关的惩罚，分支预测错误会因流水线刷新和指令重填而严重影响流水线效率并降低ILP【38，S. Eyerman, J. E. Smith, and L. Eeckhout, “Characterizing the branch misprediction penalty,” in 2006 IEEE International Symposium on Performance Analysis of Systems and Software. IEEE, 2006, pp. 48–58.】。最后一个优点在于改进的内存访问模式。在SpMM计算期间，大量的内存操作涉及为每个nz元素加载稠密输入矩阵中的相应行，表示为X[k][0:45]，如算法2的第5行所示。如果没有粗粒度列合并，程序将需要遍历所有列索引。在每次迭代中，它会读取所有相应行在相同列索引处的值，然后再移动到下一个列索引。这导致了非顺序的内存访问，因为不同行的相同列索引处的值在内存中不是连续存储的（如图7(a)所示），考虑到矩阵的行主序格式。相比之下，我们的粗粒度列合并技术一步处理单行的所有列，如图7(b)所示。这使得内存访问是顺序的，因为值在内存中是连续存储的。因此，这导致了缓存未命中减少，最小化了内存加载的流水线停顿，并最终增强了ILP【39，J. L. Hennessy and D. A. Patterson, Computer architecture: a quantitative approach. Elsevier, 2011.】。

寄存器分配与指令选择

JITSPMM的寄存器分配策略。在本小节中，我们深入研究JITSPMM框架生成的汇编代码，以展示我们如何利用SIMD寄存器和指令来减少内存访问并提高算术吞吐量。列表2提供了一个为单行计算生成的汇编代码示例，其中列数设置为45，数值类型为32位浮点数。在x86架构中，有三类SIMD寄存器：XMM、YMM和ZMM。这些寄存器的大小分别为128位、256位和512位。这意味着它们可以分别容纳4个、8个和16个单精度浮点值。ZMM寄存器的低256位与对应的YMM寄存器别名，YMM寄存器的低128位与对应的XMM寄存器别名。这允许在使用不同寄存器大小时具有兼容性和灵活性。对于支持AVX512指令扩展的CPU，总共有32个ZMM/YMM/XMM寄存器可用。这些寄存器为并行计算提供了大量的存储空间，可以极大地提高SIMD操作的性能。值得一提的是，标量浮点操作指令也可以使用XMM寄存器作为操作数。在这种情况下，只使用这些寄存器的低32位或64位，这适用于不需要SIMD并行的标量操作。

寄存器分配的具体实现。我们在寄存器分配中的目标是通过改善数据在寄存器中的保留来减少内存访问。为实现这一点，我们采用SIMD寄存器的组合来存储算法2中的整个向量ret。为了处理任意大小的向量，我们将其大小分解为可以用不同类型SIMD寄存器（ZMM/YMM/XMM）存储的大小线性组合，同时使用尽可能少的寄存器。例如，考虑d=45且值为32位浮点数的情况。我们将其分解为16(ZMM0)+16(ZMM1)+8(YMM2)+4(XMM3)+1(XMM4)。最后一个寄存器（XMM4）用于存储一个标量值。通过将向量ret[0:45]存储在SIMD寄存器中，我们可以避免在更新过程中（算法2的第5行）进行内存访问。此外，我们将第4行中的nz值A.vals[idx]广播到寄存器ZMM31中，从而无需多次内存加载。为了跟踪行的nz列表及其相应的索引，我们利用三个通用寄存器（r10、r11和r12）来存储行的nz列表的起始和结束位置，以及相应的行索引（A.row_ptr[i]、A.row_ptr[i+1]和A.col_indices[idx]）。策略性地分配寄存器使我们能够将数据在寄存器中保留更长时间，减少了频繁且昂贵的内存访问操作的需求。因此，程序可以使用存储在寄存器中的数据执行计算，从而实现更快的处理时间和更高的整体性能。

JITSPMM的指令选择策略。x86架构的SIMD扩展为程序员提供了广泛的指令，以同时对多个数据项执行各种操作，包括算术计算、逻辑运算、数据重排和数据移动【40，Intel, “Intel® 64 and ia-32 architectures software developer manuals,” 2023. [Online]. Available: https://www.intel.com/content/www/us/en/ developer/articles/technical/intel-sdm.html】。在为单行计算生成汇编代码时，我们利用了三类关键的SIMD指令。首先，我们使用打包浮点数按位异或（vxorps）指令将代表结果向量ret[0:45]的寄存器清零，如列表2的第3-6行所示。这些指令优于像vmovups这样的数据移动指令，因为它们避免了修改程序状态和控制寄存器【41，D. Kusswurm, Modern X86 Assembly Language Programming. Springer, 2014.】。其次，我们利用打包融合乘加（FMA）指令（vfmadd231ps）来累加nz值与稠密矩阵中相应行的乘法结果（第20-23行）。FMA指令能够通过单次舍入操作执行浮点乘法和加法，相比于两个独立的指令，其算术吞吐量更高【41，D. Kusswurm, Modern X86 Assembly Language Programming. Springer, 2014.】。最后，我们使用打包浮点数移动指令（vmovups）将结果存回内存（第30-33行）。总之，通过利用这些SIMD指令，我们能够在生成的二进制文件中利用数据并行的力量。

; 为处理的行索引加载
mov rdi, #i 
; 初始化存储结果的寄存器
vxorps zmm0, zmm0, zmm0
vxorps zmm1, zmm1, zmm1
vxorps ymm2, ymm2, ymm2
vxorps xmm3, xmm3, xmm3
vxorps xmm4, xmm4, xmm4
; 加载nz列表的起始和结束位置
mov r10, #row_ptr[rdi]
mov r11, #row_ptr[rdi+1]

.nnzloop_start: ; 遍历nz列表
; 边界检查
cmp r10, r11
jge .nnzloop_end
; 加载相应的行id
mov r12, #col_indices[r10]
; 加载nz值并将其广播到zmm31
vbroadcastss zmm31, #vals[r12]
; 累加结果
vfmadd231ps zmm0, zmm31, #X[r12,0:16]
vfmadd231ps zmm1, zmm31, #X[r12,16:32]
vfmadd231ps ymm2, ymm31, #X[r12,32:40]
vfmadd231ps xmm3, zmm31, #X[r12,40:44]
vfmadd231ss xmm4, xmm31, #X[r12,44]
; 下一个nz元素
inc r10
jmp .nnzloop_start

.nnzloop_end:
; 将结果写入内存
vmovups #Y[rdi,0:16], zmm0
vmovups #Y[rdi,16:32], zmm1
vmovups #Y[rdi,32:40], ymm2
vmovups #Y[rdi,40:44], xmm3
vmovss #Y[rdi,44], xmm4

列表 2. 为计算SpMM的第i行生成的汇编代码，其中稠密矩阵列数为45，数值类型为32位浮点数。

A4 实验

实验环境

• 数据集：使用了SuiteSparse矩阵集合【42，T. A. Davis and Y. Hu, “The university of florida sparse matrix collection,” ACM Transactions on Mathematical Software (TOMS), vol. 38, no. 1, pp. 1–25, 2011.】中14个最大的稀疏矩阵（按非零元素数量排序），所有矩阵均为方阵。稠密输入矩阵由随机32位浮点值生成，列数分别为16或32。
  
• 硬件配置：
  • CPU: 24核 Intel Xeon(R) Gold 6126 CPU
  • 内存: 1.5 TB DRAM
  • 实验使用48个线程（利用超线程技术）。
• 软件配置：
  • JITSPMM实现: C/C++ 与 OpenMP【43，R. Chandra, Parallel programming in OpenMP. Morgan kaufmann, 2001.】，约2000行代码。
  • JIT汇编代码生成库: AsmJit【44，“Asmjit project: Low-latency machine code generation,” 2023. [Online]. Available: https://asmjit.com/】 。
  • 操作系统: Centos 8，内核版本 4.18.0。
  • 编译器: GCC 11.3.1。
  • 性能分析工具: Linux Perf【45，A. C. De Melo, “The new linux’perf’tools,” in Slides from Linux Kongress, vol. 18, 2010, pp. 1–42.】。
• 基准解决方案：
  1. Auto-vectorization：基于Merrill和Garland的工作【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】的C/C++ SpMM实现，使用Intel icc编译器并开启 -O3 -mavx512f 标志进行自动向量化。
  2. Intel MKL：使用Intel数学核心函数库【18，E. Wang, Q. Zhang, B. Shen, G. Zhang, X. Lu, Q. Wu, Y. Wang, E. Wang, Q. Zhang, B. Shen et al., “Intel math kernel library,” HighPerformance Computing on the Intel® Xeon Phi™: How to Fully Exploit MIC Architectures, pp. 167–188, 2014.】中高度优化的mkl_sparse_spmm例程。

实验结果

代码生成开销

JITSPMM的执行时间包括代码生成和代码执行两部分。如表IV所示，代码生成时间相对于总执行时间非常小，平均仅占总执行时间的0.0074%，最低可达0.0003%，表明其开销可以忽略不计。

性能比较

1. 与Auto-Vectorization的比较：
   JITSPMM在所有三种工作负载分配方法、不同数据集和列数上均显著优于icc编译器的自动向量化方案。

   • 当列数为16时，JITSPMM在row-split, nnz-split, merge-split方法上分别取得了平均3.5倍（最高9.3倍）、3.5倍（最高7.8倍）和3.3倍（最高7.6倍）的加速比。
   • 当列数增加到32时，加速比更为显著，分别达到平均4.1倍（最高10.0倍）、4.2倍（最高8.0倍）和4.1倍（最高7.9倍）。
     

2. 与Intel MKL的比较：
   JITSPMM同样优于高度优化的Intel MKL库。

   • 当列数为16时，row-split, nnz-split, merge-split方法分别取得了平均1.4倍（最高1.9倍）、1.5倍（最高2.3倍）和1.4倍（最高2.1倍）的加速比。
   • 当列数为32时，三种方法分别取得了平均1.4倍（最高1.8倍）、1.3倍（最高1.8倍）和1.3倍（最高1.8倍）的加速比。
     

总结：实验证明，JITSPMM框架相比AOT编译解决方案（无论是自动向量化还是高度优化的手工库）具有显著的性能优势，最高可达10倍和2.3倍。

剖析分析

为了深入理解性能差异的原因，论文对JITSPMM和两个AOT基准（列数为16时）进行了剖析，如图11所示。
- 内存加载：JITSPMM的内存加载次数远少于AOT方法，平均比auto-vectorization和MKL分别少2.8倍和2倍。这验证了其寄存器分配策略在通过SIMD寄存器保留数据方面非常有效。
- 分支指令：通过利用运行时信息和循环展开，JITSPMM的分支指令数平均比auto-vectorization和MKL分别减少了3.8倍和2.9倍。
- 分支预测错误：JITSPMM在减少分支预测错误方面没有显著优势，比auto-vectorization少1.4倍，与MKL相当。这可能是因为处理器内置的分支预测器对AOT引入的额外分支指令预测准确率较高。
- 执行指令总数：通过大幅减少不必要的内存访问和分支控制，JITSPMM完成相同计算所需的指令数平均比auto-vectorization和MKL分别少7.9倍和2.0倍，显示了其在优化指令执行方面的有效性。

A7 补充细节

SpMM优化的相关工作。SpMM的优化研究已在多核CPU和通用GPU等多种架构上进行了数十年【19，C. Yang, A. Buluc¸, and J. D. Owens, “Design principles for sparse matrix multiplication on the gpu,” in Euro-Par 2018: Parallel Processing: 24th International Conference on Parallel and Distributed Computing, Turin, Italy, August 27-31, 2018, Proceedings. Springer, 2018, pp. 672–687.】、【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】、【26，G. Huang, G. Dai, Y. Wang, and H. Yang, “Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks,” in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020, pp. 1–12.】、【46，X. Liu, M. Smelyanskiy, E. Chow, and P. Dubey, “Efficient sparse matrix-vector multiplication on x86-based many-core processors,” in Proceedings of the 27th international ACM conference on International conference on supercomputing, 2013, pp. 273–282.】–【49，C. Hong, A. Sukumaran-Rajam, B. Bandyopadhyay, J. Kim, S. E. Kurt, I. Nisa, S. Sabhlok, U. V. C¸ ataly ¨ urek, S. Parthasarathy, and ¨ P. Sadayappan, “Efficient sparse-matrix multi-vector product on gpus,” in Proceedings of the 27th International Symposium on High-Performance Parallel and Distributed Computing, 2018, pp. 66–79.】。这些研究的主要目标是通过解决工作负载平衡【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】、优化内存访问【19，C. Yang, A. Buluc¸, and J. D. Owens, “Design principles for sparse matrix multiplication on the gpu,” in Euro-Par 2018: Parallel Processing: 24th International Conference on Parallel and Distributed Computing, Turin, Italy, August 27-31, 2018, Proceedings. Springer, 2018, pp. 672–687.】、【26，G. Huang, G. Dai, Y. Wang, and H. Yang, “Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks,” in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020, pp. 1–12.】和创新的数据表示技术【48，C. Hong, A. Sukumaran-Rajam, I. Nisa, K. Singh, and P. Sadayappan, “Adaptive sparse tiling for sparse matrix multiplication,” in Proceedings of the 24th Symposium on Principles and Practice of Parallel Programming, 2019, pp. 300–314.】、【49，C. Hong, A. Sukumaran-Rajam, B. Bandyopadhyay, J. Kim, S. E. Kurt, I. Nisa, S. Sabhlok, U. V. C¸ ataly ¨ urek, S. Parthasarathy, and ¨ P. Sadayappan, “Efficient sparse-matrix multi-vector product on gpus,” in Proceedings of the 27th International Symposium on High-Performance Parallel and Distributed Computing, 2018, pp. 66–79.】来提升SpMM性能。Merill和Garland为SpMV提出了基于合并的工作负载划分【20，D. Merrill and M. Garland, “Merge-based parallel sparse matrix-vector multiplication,” in SC’16: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2016, pp. 678–689.】，以解决row-split【50，J. L. Greathouse and M. Daga, “Efficient sparse matrix-vector multiplication on gpus using the csr storage format,” in SC’14: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2014, pp. 769–780.】、【51，A. Ashari, N. Sedaghati, J. Eisenlohr, S. Parthasarath, and P. Sadayappan, “Fast sparse matrix-vector multiplication on gpus for graph applications,” in SC’14: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2014, pp. 781–792.】和nnz-split【52，S. Dalton, S. Baxter, D. Merrill, L. Olson, and M. Garland, “Optimizing sparse matrix operations on gpus using merge path,” in 2015 IEEE International Parallel and Distributed Processing Symposium. IEEE, 2015, pp. 407–416.】方法的负载不平衡问题，使用2-D二分搜索来找到平衡的工作负载分解。Yang等人【19，C. Yang, A. Buluc¸, and J. D. Owens, “Design principles for sparse matrix multiplication on the gpu,” in Euro-Par 2018: Parallel Processing: 24th International Conference on Parallel and Distributed Computing, Turin, Italy, August 27-31, 2018, Proceedings. Springer, 2018, pp. 672–687.】将row-split和merge-split方法推广到GPU上的SpMM，并实现了合并内存访问以获得更好性能。Ge-spmm【26，G. Huang, G. Dai, Y. Wang, and H. Yang, “Ge-spmm: General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks,” in SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020, pp. 1–12.】是一种GPU核设计，用于对以常见CSR格式表示的稀疏矩阵执行类SpMM操作，其技术确保了对GPU全局内存的高效合并访问并减少了GPU warp间的冗余数据加载。通过深入分析SpMV和SpMM的对比，Hong等人【49，C. Hong, A. Sukumaran-Rajam, B. Bandyopadhyay, J. Kim, S. E. Kurt, I. Nisa, S. Sabhlok, U. V. C¸ ataly ¨ urek, S. Parthasarathy, and ¨ P. Sadayappan, “Efficient sparse-matrix multi-vector product on gpus,” in Proceedings of the 27th International Symposium on High-Performance Parallel and Distributed Computing, 2018, pp. 66–79.】开发了一种新的稀疏矩阵表示和计算方法，适用于实现高数据移动效率和SpMM在GPU上的有效并行化。硬件供应商库，如Intel MKL【18，E. Wang, Q. Zhang, B. Shen, G. Zhang, X. Lu, Q. Wu, Y. Wang, E. Wang, Q. Zhang, B. Shen et al., “Intel math kernel library,” HighPerformance Computing on the Intel® Xeon Phi™: How to Fully Exploit MIC Architectures, pp. 167–188, 2014.】和NVIDIA cuSPARSE【53，M. Naumov, L. Chien, P. Vandermersch, and U. Kapasi, “Cusparse library,” in GPU Technology Conference, 2010.】，也为CPU和GPU提供了高性能（非开源）的SpMM实现。

与先前工作的区别。虽然这些先前的研究在加速SpMM方面取得了显著进展，但需要注意的是，它们都遵循提前编译（AOT）方法。因此，它们在不同程度上遇到了本工作中阐述的局限性。

A5 结论

本文介绍了JITSPMM框架，该框架利用即时（JIT）汇编代码生成来提升在具有SIMD扩展的现代多核CPU上稀疏矩阵-矩阵乘法（SpMM）操作的效率。通过解决传统提前编译（AOT）方法的局限性，JITSPMM显著减少了不必要的内存访问、分支操作和执行的指令。我们的实验评估在各种数据集上证明了JITSPMM的有效性，与基准AOT解决方案相比，性能提升高达10倍。