Libra: Unleashing GPU Heterogeneity for High-Performance Sparse Matrix Multiplication

TL;DR

Libra is a comprehensive framework that optimally combines Tensor Core Units and CUDA cores to significantly accelerate sparse matrix multiplication, enhancing performance in deep learning and scientific computing applications.

Contribution

This work introduces Libra, a novel framework that systematically leverages GPU heterogeneity for high-performance sparse matrix multiplication, combining workload distribution, load balancing, and efficient kernel design.

Findings

Achieves up to 2.9x speedup over state-of-the-art baselines.

Effectively accelerates end-to-end GNN applications.

Demonstrates significant performance improvements on H100 and RTX 4090 GPUs.

Abstract

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.