利用GPU异构计算加速稀疏矩阵乘法,性能提升超2.9倍。
Libra: Unleashing GPU Heterogeneity for High-Performance Sparse Matrix Multiplication
- 设计2D感知任务分配,协同利用Tensor Core与CUDA核心优势
- 在H100和RTX 4090上实现平均1.77倍至2.9倍加速
- 适合需要高性能稀疏计算的图神经网络开发者
稀疏矩阵乘法(SpMM和SDDMM)广泛应用于深度学习与科学计算。现代加速器通常配备张量核心单元(TCUs)和CUDA核心,前者擅长结构化计算,后者更具编程灵活性。然而如何结合二者以最大化稀疏算子性能尚不明确。本文首先分析混合计算的性能来源,系统揭示其互补优势。基于此,提出Libra框架,通过2D感知的工作负载分配方法,精准识别最优任务映射,同时利用TCUs的数据重用能力与CUDA核心的灵活性,减少计算冗余。Libra还引入混合负载均衡、占用率感知调度与高效内核实现,全面提升执行效率。在H100与RTX 4090上的实验表明,相比12个最新基线,其平均速度提升1.77倍(优于FlashSparse)、1.73倍(优于RoDe)、2.9倍(优于DGL),适用于端到端图神经网络应用。Libra为稀疏算子加速提供了全新视角。
原文摘要 · Abstract (English)
Sparse matrix multiplication operators (i.e., SpMM and SDDMM) are widely used in deep learning and scientific computing. Modern accelerators are commonly equipped with Tensor Core Units (TCUs) and CUDA cores to accelerate sparse operators. The former excels at structured matrix computations, whereas the latter offers greater programming flexibility. However, how to combine these two resources to maximize sparse-operator performance remains unclear. In this work, we first identify the source of performance gains in hybrid computation and systematically analyze their complementary strengths. Motivated by this, we propose Libra, a holistic framework that efficiently leverages heterogeneous computing resources to accelerate both SpMM and SDDMM operators. Specifically, Libra introduces a 2D-aware (locality and utilization) workload distribution method to precisely identify the optimal task mapping, simultaneously leveraging the data reuse capabilities of TCUs and the flexibility of CUDA cores to minimize computational redundancy. Libra further incorporates hybrid load balancing, occupancy-aware task scheduling, and efficient kernel implementations to maximize execution efficiency. Extensive experiments on H100 and RTX 4090 GPUs demonstrate that Libra surpasses all the 12 up-to-date baselines significantly, e.g., on average 1.77x speedup over FlashSparse, 1.73x over RoDe, and 2.9x over DGL for end-to-end GNN applications. Libra opens up a new perspective for sparse operator acceleration by fully unleashing the power of heterogeneous GPU resources.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。