LOw-cOst yet High-Performant Sparse Matrix-Matrix Multiplication on Arm SME Architectures
Kelun Lei, Hailong Yang, Kaige Zhang, Kejie Ma, Yiqing Wang, Xin You, Yufan Xu, Enrique S. Quintana-Orti, Zhongzhi Luan, Yi Liu, Depei Qian
稀疏矩阵密度矩阵乘法(SpMM)是科学计算和新兴图形学习工作负载的关键内核。 最近的 Armv9 架构引入了可扩展矩阵扩展 (SME),实现了基于图块的矩阵操作,吞吐量很高。 然而,有效地利用中小企业和传统SIMD资源来实现非结构化的稀疏工作负载仍然是一个悬而未决的挑战。 为了解决这个问题,我们提出了LOOPS,这是一个混合执行框架,将行明智的CSR部分与向量-明智的BCSR部分布局相结合,使矢量指令(NEON)和可扩展矩阵扩展(SME)资源的合作利用成为可能。 LOOPS通过轻量级性能模型引导的自适应两级并行化方案支持FP64、FP32和FP16的多精度SpMM。 苹果M4Pro CPU上整个SuiteSparse的实验结果表明,相对于Armadillo,LOOPS的平均速度为9.93×(FP32)1.44×(FP64),而CPU基线TACO为71.3×(FP32)/54.8×(FP64)。 在NVIDIA A100 GPU上执行的两种GPU方法(cuSPARSE,Magicube)在同一CPU上运行的LOOPS的比较显示LOOPS的平均速度在19.8×和33.5×之间,具体取决于精度。 值得注意的是,LOOPS提供的能源效率明显优于A100 GPU上的GPU代码。
Sparse matrix-dense matrix multiplication (SpMM) is a critical kernel in both scientific computing and emerging graph learning workloads. The recent Armv9 architecture introduces Scalable Matrix Extension (SME), enabling tile-based matrix operations with high throughput. However, effectively exploiting both SME and traditional SIMD resources for unstructured sparse workloads remains an open challenge. To address this, we propose LOOPS, a hybrid execution framework that combines row-wise CSR-part...