基于Tensor Core的高性能非结构化稀疏矩阵-矩阵乘法计算
分布式、并行与集群计算
2024-08-22 v1
摘要
高性能稀疏矩阵-矩阵（SpMM）乘法对于科学和工业应用至关重要，因为数据规模不断增大，已不适合使用稠密数据结构。然而，现有硬件（如Tensor Core）不适用于SpMM，因为其对数据结构施加了严格限制，而这些限制无法满足许多应用中遇到的非结构化稀疏性。为此，我们引入(S)parse(M)a(trix)Matrix(T)ensor Core加速(SMaT)：一种新颖的SpMM库，可为非结构化稀疏矩阵利用Tensor Core。我们的块稀疏库利用低级CUDA MMA（矩阵-矩阵-累加）API，最大化现代GPU的性能。此外，稀疏矩阵置换等算法优化可通过最小化非零块的数量来进一步提高性能。在NVIDIA A100上的评估显示，SMaT相较于现有SOTA库（DASP、cuSPARSE和Magicube）可实现最高125倍提升（平均提升2.6倍）。SMaT可用于加速科学计算、大规模模型训练、推理等众多工作负载。
引用
@article{arxiv.2408.11551,
title = {High Performance Unstructured SpMM Computation Using Tensor Cores},
author = {Patrik Okanovic and Grzegorz Kwasniewski and Paolo Sylos Labini and Maciej Besta and Flavio Vella and Torsten Hoefler},
journal= {arXiv preprint arXiv:2408.11551},
year = {2024}
}
备注
Accepted by 2024 International Conference on High Performance Computing, Networking, Storage and Analysis, 2023 (SC'24)