计算机科学 ›› 2026, Vol. 53 ›› Issue (6): 163-170.doi: 10.11896/jsjkx.251000116

• 高性能计算 • 上一篇    下一篇

利用即时代码生成加速ARM平台的SPMM

施俊, 王庆林, 田飞扬, 王志成, 李润华, 刘杰   

  1. 国防科技大学高端装备数字化软件湖南省重点实验室 长沙 410073
    国防科技大学并行与分布计算全国重点实验室 长沙 410073
    国防科技大学计算机学院 长沙 410073
  • 收稿日期:2025-10-27 修回日期:2025-12-22 出版日期:2026-06-15 发布日期:2026-06-09
  • 通讯作者: 王庆林(wangqinglin@nudt.edu.cn)
  • 作者简介:(shijun@nudt.edu.cn)
  • 基金资助:
    天河高性能计算机系统优化与升级(2021YFB0300101)

Optimizing SPMM on ARM Architectures with JIT Instruction Generation

SHI Jun, WANG Qinglin, TIAN Feiyang, WANG Zhicheng, LI Runhua, LIU Jie   

  1. Laboratory of Digitizing Software for Frontier Equipment,National University of Defense Technology,Changsha 410073,China
    National Key Laboratory of Parallel and Distributed Computing,National University of Defense Technology,Changsha 410073,China
    College of Computer Science and Technology,National University of Defense Technology,Changsha 410073,China
  • Received:2025-10-27 Revised:2025-12-22 Published:2026-06-15 Online:2026-06-09
  • About author:SHI Jun,born in 2000,MS candidate.His main research interests include high performance computing and parallel computing.
    WANG Qinglin,born in 1987,Ph.D,associate research fellow.His main research interests includes high perfor-mance computing and artificial intelligence.
  • Supported by:
    Optimization and Upgrading of the Tianhe High-Performance Computing System(2021YFB0300101).

摘要: 近年来,随着ARM架构处理器在边缘计算设备和云服务器领域的快速普及,以及稀疏矩阵乘稠密矩阵(SPMM)在深度学习等计算密集型应用中的关键作用日益凸显,针对ARM平台的稀疏计算优化研究已成为学术热点。当前ARM平台的SPMM解决方案主要采用提前编译(AOT)模式,即在程序执行前完成全部编译工作。然而,SPMM的AOT解决方案面临3个主要限制:不必要的内存访问、额外的分支开销以及冗余指令。因此,提出ASPJIT,其是一个ARM平台的即时(JIT)汇编代码生成框架,以加速ARM多核CPU上的SPMM计算。ASPJIT基于运行时稀疏特征分析,采用基于稀疏判断的列转行算法动态优化计算序列,显著提升指令级并行度(ILP);此外,ASPJIT通过寄存器分配策略缓存高频访问数据以减少内存访问时延,并精选单指令多数据SIMD指令集最大化算术吞吐量。对ASPJIT进行了性能评估,并将其与两个AOT基线进行了比较,第一个涉及使用ARM gcc编译器自动向量化编译的现有SPMM实现,第二个使用了ARM Eigen提供的经过优化的SPMM例程。结果表明,ASPJIT分别提供了平均3.8倍和5.6倍的提升。

关键词: 稀疏矩阵乘稠密矩阵, 提前编译, 即时汇编代码生成, 单指令多数据, ARM多核CPU

Abstract: In recent years,with the rapid adoption of ARM architecture processors in edge computing devices and cloud servers,as well as the increasingly critical role of sparse matrix multiplication(SPMM)in compute-intensive applications such as deep learning,research on sparse computing optimization for ARM platforms has become an academic hotspot.However,current SPMM solutions for ARM platforms primarily employ the ahead-of-time(AOT)compilation model,where all compilation is completed before program execution.Nevertheless,AOT solutions for SPMM face three major limitations:unnecessary memory accesses,additional branch overhead,and redundant instructions.This paper proposes ASPJIT,a just-in-time(JIT)assembly code generation framework for ARM platforms,designed to accelerate SPMM computation on ARM multi-core CPUs.ASPJIT dynamically optimizes the computation sequence using a sparse judgment-based column-to-row algorithm,leveraging runtime sparse feature analysis to significantly improve instruction-level parallelism(ILP).Additionally,ASPJIT reduces memory access latency by employing a register allocation strategy to cache frequently accessed data and maximizes arithmetic throughput by carefully selecting SIMD instruction sets.A performance evaluation of ASPJIT is conducted and compared it with two AOT baselines.The first involves existing SPMM implementations compiled with automatic vectorization using the ARM gcc compiler.The second utilizes optimized SPMM routines provided by ARM Eigen.The results demonstrate that ASPJIT delivers average speedups of 3.8x and 5.6x,respectively.

Key words: Sparse matrix-matrix multiplication(SPMM), Ahead-of-time(AOT), Just-in-time(JIT), Single instruction multiple data(SIMD), ARM multi-core CPUs

中图分类号: 

  • TP391
[1]YANG C,AYDIN B,OWENS J D.Design principles for sparse matrix multiplication on the gpu[C]//European Conference on Parallel Processing.Cham:Springer,2018.
[2]WU T,WANG B,SHAN Y,et al.Efficient pagerank and spmv computation on amd gpus[C]//2010 39th International Conference on Parallel Processing.IEEE,2010:81-89.
[3]FU Q,ROLINGER T B,HUANG H H.JITSPMM:Just-in-time instruction generation for accelerated sparse matrix-matrix multiplication[C]//2024 IEEE/ACM International Symposium on Code Generation and Optimization(CGO).IEEE,2024:448-459.
[4]LANGVILLE A N,MEYER C D.Google's PageRank and beyond:The science of search engine rankings[M].Princeton University Press,2006.
[5]KOREN Y,BELL R,VOLINSKY C.Matrix factorization techniques for recommender systems[J].Computer,2009,42(8):30-37.
[6]SCHAEFFER S E.Graph clustering[J].Computer Science Review,2007,1(1):27-64.
[7]WANG E,ZHANG Q,SHEN B,et al.Intel Math Kernel Library[J].2014.DOI:10.1007/978-3-319-06486-4_7.
[8]VIRTANEN P,GOMMERS R,OLIPHANT T E,et al.Fundamental algorithms for scientific computing in python and SciPy 1.0 contributors.SciPy 1.0[J].Nature Methods,2020,17:261-272.
[9]WANG M Y.Deep graph library:Towards efficient and scalable deep learning on graphs[C]//ICLR Workshop on Representation Learning on Graphs and Manifolds.2019.
[10]FEY M,LENSSEN J E.Fast graph representation learning with PyTorch Geometric[J].arXiv:1903.02428,2019.
[11]SELVITOPI O,BROCK B,NISA I,et al.Distributed-memoryparallel algorithms for sparse times tall-skinny-dense matrix multiplication[C]//Proceedings of the 35th ACM International Conference on Supercomputing.2021:431-442.
[12]ZHANG Y,YANG W,LI K,et al.Performance analysis and optimization for SpMV based on aligned storage formats on an ARM processor[J].Journal of Parallel and Distributed Computing,2021,158:126-137.
[13]ZHANG W,JIANG Z,CHEN Z,et al.NUMA-Aware DGEMM based on 64-bit ARMv8 multicore processors architecture[J].Electronics,2021,10(16):1984.
[14]ZHENG J,JIANG J,DU J,et al.Optimizing massively parallel sparse matrix computing on ARM many-core processor[J].Parallel Computing,2023,117:103035.
[15]SATO M,ISHIKAWA Y,TOMITA H,et al.Co-design fora64fx manycore processor and” fugaku”[C]//SC20:International Conference for High Performance Computing,Networking,Storage and Analysis.IEEE,2020:1-15.
[16]HUANG G,DAI G,WANG Y,et al.Ge-spmm:General-purpose sparse matrix-matrix multiplication on gpus for graph neural networks[C]//SC20:International Conference for High Performance Computing,Networking,Storage and Analysis.IEEE,2020:1-12.
[17]INOUE H,OHARA M,TAURA K.Faster set intersection with SIMD instructions by reducing branch mispredictions[J].Proceedings of the VLDB Endowment,2014,8(3):293-304.
[18]SHAYLOR N.A {Just-in-Time} Compiler for {Memory-Constrained}{Low-Power} Devices[C]//2nd Java Virtual Machine Research and Technology Symposium(Java VM 02).2002.
[19]FU Q,JI Y,HUANG H H.TLPGNN:A lightweight two-level parallelism paradigm for graph neural network computation on GPU[C]//Proceedings of the 31st International Symposium on High-Performance Parallel and Distributed Computing.2022:122-134.
[20]MERRILL D,GARLAND M.Merge-based parallel sparse matrix-vector multiplication[C]//Proceedings of the International Conference for High Performance Computing(SC'16).Networking,Storage and Analysis.IEEE,2016:678-689.
[21]PAI V S,RANGANATHAN P,ADVE S V.The impact of instruction-level parallelism on multiprocessor performance and simulation methodology[C]//Proceedings Third International Symposium on High-Performance Computer Architecture.IEEE,1997:72-83.
[22]HENNESSY J L,PATTERSON D A.Computer architecture:a quantitative approach[M].Elsevier,2011.
[23]DAVIS T A,HU Y.The University of Florida sparse matrix collection[J].ACM Transactions on Mathematical Software,2011,38(1):1-25.
[24]CHANDRA R.Parallel programming in OpenMP[M].Morgan Kaufmann,2001.
[25]ACER S,SELVITOPI,O,AYKANAT C.Improving performance of sparse matrix dense matrix multiplication on large-scale parallel systems[J].Parallel Computing,2016,59:71-96.
[26]HU Y,YE Z,WANG M,et al.Featgraph:A flexible and efficient backend for graph neural network systems[C]//SC20:International Conference for High Performance Computing,Networking,Storage and Analysis.IEEE,2020:1-13.
[27]GUO M,WANG Y,GU Y,et al.Bs-SpMM:Accelerate Sparse Matrix-Matrix Multiplication by Balanced Split Strategy on the GPU[C]//IEEE INFOCOM 2023-IEEE Conference on Computer Communications Workshops(INFOCOM WKSHPS).IEEE,2023:1-6.
[28]CAO L,WANG Q,YANG S,et al.LSSM-SpMM:A Long-Row Splitting and Short-Row Merging Approach for Parallel SpMM on PEZY-SC3s[C]//International Conference on Algorithms and Architectures for Parallel Processing.Singapore:Springer,2024:78-97.
Viewed
Full text


Abstract

Cited

  Shared   
  Discussed   
No Suggested Reading articles found!