JITSPMM：面向加速稀疏矩阵-矩阵乘法的即时指令生成
分布式、并行与集群计算
2023-12-12 v1 性能
编程语言
摘要
由于图神经网络（GNN）等应用中输入数据量庞大，实现高性能的稀疏矩阵-矩阵乘法（SpMM）受到了越来越多的研究关注，尤其是在多核 CPU 上。大多数现有的 SpMM 计算解决方案遵循提前编译（AOT）方法，即在程序执行前将其完全编译。SpMM 的 AOT 编译面临三个关键局限：不必要的内存访问、额外的分支开销和冗余指令。这些局限源于与 SpMM 相关的关键信息直到运行时才可知这一事实。在本文中，我们提出了 JITSPMM，这是一个即时（JIT）汇编代码生成框架，用于在具有 SIMD 扩展的多核 CPU 上加速 SpMM 计算。首先，JITSPMM 将 JIT 汇编代码生成技术集成到三种广泛使用的 SpMM 工作负载划分方法中，以实现 CPU 线程间的均衡工作负载分配。其次，借助运行时信息的可用性，JITSPMM 采用了一种新技术——粗粒度列合并，通过展开性能关键循环来最大化指令级并行性。此外，JITSPMM 智能分配寄存器以缓存频繁访问的数据从而最小化内存访问，并采用特定的 SIMD 指令以增强算术吞吐量。我们对 JITSPMM 进行了性能评估，并将其与两个 AOT 基准进行了比较。第一个基准涉及使用 Intel icc 编译器通过自动向量化编译的现有 SpMM 实现。第二个基准利用了 Intel MKL 提供的高度优化的 SpMM 例程。结果表明，JITSPMM 分别实现了平均 3.8 倍和 1.4 倍的性能提升。
引用
@article{arxiv.2312.05639,
title = {JITSPMM: Just-in-Time Instruction Generation for Accelerated Sparse Matrix-Matrix Multiplication},
author = {Qiang Fu and Thomas B. Rolinger and H. Howie Huang},
journal= {arXiv preprint arXiv:2312.05639},
year = {2023}
}