TC_IPDPS25_01_st-bi.pdf

Transcript

1. BRP-SpMM: Block-Row Partition Based Sparse Matrix Multiplication with Tensor and CUDA Cores IPDPS 2025 Yukang Dong, Wenbin Jiang, Xinhai Shen, Haihong Guo, Zhiyuan Shao, Hai Jin Presenter: BI SHUTAN 2026/4/17 1
2. Overview Background - SpMM is a fundamental operation in scientific computing and GNNs - GPUs feature specialized Tensor Cores for dense math and general-purpose CUDA Cores Challenges - Tensor Cores offer massive throughput but struggle with high-sparsity matrices (>90%) - Existing methods suffer from unaligned memory access Proposed Method: BRP-SpMM - Partitions sparse matrix into TC Block and Residual Row parts - Customized storage format to manage these two parts efficiently - Leverages Tensor Cores and CUDA Cores separately for optimal utilization Results - Average speedup of 2.1x (up to 2.9x) over cuSPARSE on NVIDIA A800 - Accelerates GNN training by up to 1.9x compared to DGL and PyG 3 2026/4/17
3. Preliminary 1: SpMM Operation Sparse-Dense Matrix Multiplication (SpMM) • Computes 𝐶𝑀×𝑁 = 𝐴𝑀×𝐾 × 𝐵𝐾×𝑁 • Matrix A is sparse, Matrices B and C are dense Sparse Matrix Characteristics • A matrix where the number of non-zero elements is significantly fewer than the total elements. • Matrix A: 𝑁𝑁𝑍 ≪ 𝑀 × 𝐾 • High Sparsity (>90%) or Extreme Sparsity (Graph adjacency matrices >99%) 4 𝐶𝑀×𝑁 𝐴𝑀×𝐾 𝐵𝐾×𝑁 2026/4/17
4. Preliminary 2: Standard Sparse Formats(𝑁 × 𝑁 matrix) COO (Coordinate) • Stores (𝑟𝑜𝑤, 𝑐𝑜𝑙, 𝑣𝑎𝑙𝑢𝑒) triples for every non-zero element. • Storage Space: 3 × 𝑁𝑁𝑍. CSR (Compressed Sparse Row) • Data Structure: Uses three arrays to minimize space: • 𝑅𝑜𝑤𝑃𝑡𝑟𝑠: Pointers to the start of each row (size: 𝑛 + 1). • 𝐶𝑜𝑙𝐼𝑑𝑥: Column index for each value. • 𝑉𝑎𝑙𝑢𝑒: All non-zero elements. • Computation: Number of elements in row 𝑖 = 𝑅𝑜𝑤𝑃𝑡𝑟𝑠[𝑖 + 1] − 𝑅𝑜𝑤𝑃𝑡𝑟𝑠[𝑖] • Storage Space: 2 × 𝑁𝑁𝑍 + 𝑛. 5 CSR COO 2026/4/17
5. Preliminary 3: Tensor Core CUDA Cores: Support standard data types (FP16, FP32, FP64, INT32) Flexible with irregular storage formats like CSR and COO Tensor Cores: Hardware designed to accelerate dense GEMM op. Deliver significantly higher throughput than CUDA Cores. Require strict blocked data formats (e.g., m8n8k4) 6 Ampere GPU SM (A100) Example: m4n4k4 D(m,n) = A(m,k) × B(k,n) + C(m,n) 2026/4/17
6. Related Works: TCGNN-SpMM A pioneering study focusing on SpMM operations within Graph Neural Networks. Core Innovation: Sparse Graph Translation (SGT) (On CPU) • Condense sparse matrices by grouping rows into row windows. • These blocks serve as the basic processing unit for GPU Tensor Cores. • Non-zero elements within each window are condensed into 16×8 TC Blocks. 7 2026/4/17
7. Related Works: DTC-SpMM A state-of-the-art SpMM implementation built directly upon the TC-GNN framework. Shifts from high-level APIs to low-level MMA PTX instructions for finer hardware control. Advanced Runtime Optimizations • Offloading SGT from CPU to GPU to significantly reduce preprocessing overhead. • Implements vectorized global memory accesses and double buffering. • Introduces pipeline strategies to overlap data loading with Tensor Core computation. • Utilizes register remapping and load balancing to improve throughput Critical Shortcomings • Extensive optimizations result in a significantly large kernel function (~300 lines). • Large program sizes risk overflowing the GPU instruction cache, reducing efficiency. • Does not improve upon the original SGT technique, inheriting its memory alignment flaws. 8 2026/4/17
8. Challenge 1: Inefficient Memory Access Hardware Alignment Constraints • GPU memory access efficiency is maximized when: • Starting addresses are aligned to 128 bytes. • Dense block sizes are fixed (e.g., 128B). • GPU L2 cache line size is 32 128 bytes, managed in 32 bytes sectors • Cache utilization peaks only when data loads are multiples of 32 bytes The Limitation of Sparse Graph Translation (SGT) • Non-zero elements in each TC block are unpredictable in existing methods. • This leads to unaligned memory access and inefficient L2 cache usage. 9 2026/4/17
9. Challenge Challenge 2: Compute Unit Mismatch • Some TC blocks generated by SGT contain only a single non-zero element • Tensor Cores are high-throughput units designed for dense operations • Processing "almost empty" blocks reduces computational efficiency 10 Challenge 3: Oversized Runtime GPU Kernels TCGNN-SpMM: Uses high-level WMMA C API but lacks advanced optimizations. DTC-SpMM: Incorporates low-level MMA PTX and numerous runtime optimizations. The Instruction Cache Bottleneck • Extensive optimizations result in significantly large kernel functions (~300 lines). • Large program sizes can lead to Instruction Cache Overflow on the GPU. 2026/4/17
10. Motivations 1. Memory Alignment & Cache Efficiency Problem: Existing SGT methods produce blocks with unpredictable non-zero counts. Proposed Method: Implement a Fixed Non-zero Strategy to ensure 100% memory alignment 2. Optimal Compute Unit Utilization Problem: Processing nearly empty blocks on Tensor Cores wastes high-throughput hardware. Proposed Method: Adopt a hybrid strategy by offloading residuals to flexible CUDA cores. 3. Instruction Cache & Kernel Complexity Problem: SOTA kernels (e.g., DTC-SpMM) are excessively bloated and complex (~300 lines). Proposed Solution: Design Compact Kernels (1/3 size of SOTA) to maximize I-Cache efficiency. 11 2026/4/17
11. Method: design overview of BRP-SpMM 12 2026/4/17
12. Part 1: Block-Row Partition 13 The Fixed Non-zero Strategy (𝒇 = 𝟏) • BRP-SpMM fixes the number of non-zero elements per column, denoted as f. • By setting 𝑓 = 1, every TC Block(h×w) contains exactly w (width) non-zero elements. • Columns with multiple elements are split to maintain this strict density Benefits • Eliminates the variable non-zero counts per block • Minimizes TC block growth while ensuring consistent memory access patterns Note: Diagram uses simplified 4×4 blocks (g=2). Real-world settings uses 16×8 blocks (g=4). 2026/4/17
13. Part 1: Block-Row Partition 14 TC Block Group • Grouping Logic: Every g TC blocks are organized into a single TC Block Group. • TC Block Part: Complete groups are assigned to Tensor Cores. • Residual Row Part: Elements failing to form a full group are offloaded to CUDA Cores. Benefits • 32 floats (128Byte) ensure 100% alignment with memory access widths. • 128B is a direct multiple of the 32B L2 cache line. Note: Diagram uses simplified 4×4 blocks (g=2). Real-world settings uses 16×8 blocks (g=4). 2026/4/17
14. Part 2: Customized Storage Format — TC Block Part 16 Field Description group_window_id Index of the row window the group belongs to. local_id Local element index within the TC block(0~ℎ × 𝑤 − 1) old_col_id Original column index in the sparse matrix values Actual non-zero data values (e.g., weights or features). 2026/4/17
15. 17 Field Description res_rows_id Indices of rows containing residual non-zero elements res_row_offset Start and end positions of non-zeros for each residual row. res_col_id Original column index for each residual non-zero element. values Actual data values in the residual sparse matrix. 0 1 2 3 8 9 10 11 Part 2: Customized Storage Format — Residual Row part 2026/4/17
16. 18 Part 2: Customized Storage Format — Storage Space Matrix Dimensions: 𝑛 × 𝑛 with 𝑛𝑛𝑧 total non-zero elements. Workload Split: 𝑛𝑛𝑧 = 𝑛𝑛𝑧𝑡 𝑇𝐶 𝑃𝑎𝑟𝑡 + 𝑛𝑛𝑧𝑟 𝑅𝑒𝑠𝑖𝑑𝑢𝑎𝑙 𝑃𝑎𝑟𝑡 . Configuration: TC Block size ℎ × 𝑤 = 16 × 8; Group size 𝑔 = 4 TC Block Part 4 arrays: group_window_id, local_id, old_col_id, and values TC Block Part Count = nnzt g×w Total Size = nnzt g×w + 3 × nnzt = nnzt 4×8 + 3 × nnzt ≈ 3.03nnzt ≈ 3nnzt Residual Row Part Upper Limit: Worst case assumes every row is a Residual Row. Total Size = 𝑛 + 𝑛 + 1 + 2 × 𝑛𝑛𝑧𝑟 = 2𝑛 + 1 + 2𝑛𝑛𝑧𝑟 BRP-SPMM Total Total = 3𝑛𝑛𝑧𝑡 + 2𝑛 + 1 + 2𝑛𝑛𝑧𝑟 = 2(𝑛𝑛𝑧𝑡 +𝑛𝑛𝑧𝑟 ) + 𝑛𝑛𝑧𝑡 + 2𝑛 + 1 ≈ 2𝑛𝑛𝑧 + 𝑛𝑛𝑧𝑡 + 2𝑛 2026/4/17
17. 19 Part 2: Customized Storage Format — Storage Space BRP-SPMM ≈ 2𝑛𝑛𝑧 + 𝑛𝑛𝑧𝑡 + 2𝑛 Standard CSR ≈ 2𝑛𝑛𝑧 + 𝑛 Overhead: Approximately 𝑛𝑛𝑧𝑡 + 𝑛 additional storage, roughly 50% more compared to CSR? Trade-off: The 50% increase in storage space is a minimal cost for the significant gains in memory access efficiency and L2 cache hit rates. Standard CSR 2026/4/17
18. 21 Hidden Memory Overhead in BRP-SpMM Preprocessing • Paper Claim: Custom format uses only ~50% more space than standard CSR. • This only accounts for the final static storage, ignoring the massive peak memory footprint during preprocessing (Algorithm 1). Overhead 1: Redundant Input Formats • Algorithm 1 requires row_offset (CSR) AND row_idx (COO) simultaneously. • Input alone takes 3 × 𝑛𝑛𝑧 + 𝑛 before any processing begins. Overhead 2: Temporary array nnz_window_id • requires 𝑛𝑛𝑧 allocation. 2026/4/17
19. 22 Hidden Memory Overhead in BRP-SpMM Preprocessing Overhead 3: Out-of-Place GPU Sorting • MultiKeySort on GPU (e.g., CUB/Thrust) cannot be executed in-place. • Sorting tuples requires Double Buffering, adding an implicit 3 × 𝑛𝑛𝑧 temporary allocation. Peak Memory Footprint Calculation: • Input Data: ∼ 3𝑛𝑛𝑧 + 𝑛 • Sort Buffers & Temp Arrays: ∼ 4𝑛𝑛𝑧 • Output Allocation: ∼ 3𝑛𝑛𝑧 + 2𝑛 • Total Peak Memory ≈ 10𝑛𝑛𝑧 (Nearly 3x to 4x the size of standard CSR). 2026/4/17
20. 24 Part 3: TC Block Kernel - Load Balance 2026/4/17 Thread Block Mapping (Load Balanced) • 1 Thread Block → Computes 1 TC Block Group(4 TC Block). Warp Division • 1 Thread Block(128 threads) = 4 Warps. • 1 Warp → Computes 8 columns of Matrix B.(m16n8k8) • Multiply TC Block Group (4 blocks of 16 × 8) by 8 × 8 𝐵𝑡𝑖𝑙𝑒 Addressing Matrix B Issue: Compression(Fixed Non-zero) changes A's original column layout. Solution: Use old_col_id to find matching rows in B. Fetching Example • If condensed A has 𝑜𝑙𝑑_𝑐𝑜𝑙_𝑖𝑑 = [2, 5, 11, 13]... • Fetch rows 2, 5, 11, 13 from Matrix B. • Each warp fetches 8 elements from these rows
21. 25 Part 3: TC Block Kernel - Register Remapping 2026/4/17 The Bottleneck • High overhead from atomic write-backs to Matrix C. DTC-SpMM • 4 independent m16n8k4 instructions → 4 𝐶𝑓𝑟𝑎𝑔 outputs. • Requires 4 atomic writes per thread. BRP-SpMM • Use 4 m16n8k8 instructions to compute 1 single 𝐶𝑡𝑖𝑙𝑒 . • Accumulate intermediate sums directly in registers(RC). Result & Benefit • Only 1 atomic write per thread. • Reduces atomic operations by 75%. TC Block Group(𝐴𝑓𝑟𝑎𝑔 ) 16 32(4×8) B𝑓𝑟𝑎𝑔 32(4×8) 8 × = 16 8 C𝑓𝑟𝑎𝑔
22. 27 Part 3: Residual Row Kernel Computation Goal • 1 × 𝐾 (Residual Row) × 𝐾 × 𝑁 (Matrix B) → 1 × 𝑁 (Row in Matrix C) 1-D Tiling Strategy • set 𝑏𝑙𝑜𝑐𝑘𝐷𝑖𝑚. 𝑥 = 1, 𝑏𝑙𝑜𝑐𝑘𝐷𝑖𝑚. 𝑦 = 𝑁 • 1 Thread Block→Handles 1 residual row(Computes 1 row in Matrix C). • 1 Thread→Computes 1 element in Matrix C(Handles 1 column in Matrix B) Execution Flow & Write-Back • Executes serially after the TC Block kernel finishes. • Each thread writes to a unique element in Matrix C. • No Atomics: Safe to use += directly to aggregate final results (𝐶[𝑥 ∗ 𝑁 + 𝑡𝑖𝑑] += 𝑟𝑒𝑠𝑢𝑙𝑡). 2026/4/17 Residual Row Matrix B Row in Matrix C
23. Experiment: Setup Hardware Environment • GPU: NVIDIA A800 GPU (Ampere, 80GB memory). • Compute Capability: 8.6. • Software Stack: CUDA 12.4 • Measurement: Each test is repeated 100 times, reporting average results. Dataset Source: Eight widely-used graph datasets from Deep Graph Library (DGL). Key Metric: All datasets exhibit extreme sparsity, exceeding 99%. Scale Range: 20K rows to 717K rows 28 2026/4/17
24. Experiment: Baseline CUDA Core Based Methods: cuSPARSE: NVIDIA's standard vendor library (v12.4), used as the primary baseline. Sputnik: A SOTA sparse library designed for deep learning, optimized for CUDA cores. Tensor Core Based Methods: TCGNN-SpMM: Pioneering work bringing SpMM computation to Tensor Cores via SGT. DTC-SpMM: SOTA implementation focused on bridging the gap for SpMM with Tensor Cores Evaluate normalized speedup relative to cuSPARSE across varying dense matrix columns (N=32~256). 29 2026/4/17
25. Experiment: SpMM Performance Comparison(N=32) • Vs. cuSPARSE: Achieves an average speedup of 2.1x when N=32. • Peak Performance: Reaches a maximum speedup of 2.9x on the PubMed dataset. • Vs. DTC-SpMM: Delivers an average speedup of 1.2x across all scenarios. 30 2026/4/17
26. Experiment: SpMM Performance Comparison • Peak performance at N=32 is attributed to the 32 × 4 thread block configuration(4 warps). • Each warp computes 8 columns via m16n8k8, maximizing thread efficiency. • For larger N, redundant memory reads of 𝐴𝑡𝑖𝑙𝑒 degrade relative performance. 31 2026/4/17
27. Experiment: Preprocessing Overhead Analysis • Comparison: Consistently the lowest overhead across all eight datasets. • Scalability: On the largest dataset, BRP-SpMM is 212.9x faster than TCGNN-SpMM (CPU). • Speedup vs SOTA: 1.9x faster than DTC-SpMM on the Yelp dataset. • Scalability: As the dataset size increases, the advantage becomes more significant. 32 2026/4/17
28. Experiment: Detailed comparison of two parts • TC Block Dominance: Accounts for 88.5% of non-zero elements on average. • Execution Time: Consumes 85.9% of total kernel execution time. • Edge Case: In massive datasets, the Residual Row proportion drops significantly (~5%). • Future Direction: Fixed window constraints (𝑔 = 4, 𝑤 = 8) limit residual size. 33 2026/4/17
29. Experiment: L2 Cache Hit Rate • Average Gains: +13.2% over cuSPARSE and +8.6% over DTC-SpMM. • Peak Performance: Over 20% improvement on the Yelp dataset. • The 128-byte aligned memory access significantly improves hardware efficiency. 34 2026/4/17
30. Experiment: Storage Analysis • Comparison: Memory footprint is nearly identical to the standard COO format. • Requires at most 50% more space compared to the highly compressed CSR format. • Peak Memory Footprint during preprocessing: 10𝑛𝑛𝑧 35 2026/4/17
31. Case Study: End-to-End GNN Training(500 epochs) • Experimental Setup • Model: 2-layer Graph Convolutional Network (GCN). • Frameworks: Compared against DGL, PyG, TC-GNN, and DTC-GCN. • Average Speedup: BRP-GCN achieves 1.13x to 1.51x speedup over all baselines. • Maximum Gain: Reaches 1.87x speedup over DGL on the Yelp dataset when N=128. • Format conversion time is included in the totals and proven negligible 36 2026/4/17
32. Conclusion Motivation & Method • Issue: Inefficient memory access and underutilized Tensor Cores in high-sparsity SpMM. • Solution: Partition matrices into a TC Block part and a Residual Row part. Key Results • vs. cuSPARSE: Up to 2.9× (avg 2.1×) kernel speedup. • vs. SOTA (DTC-SpMM): Up to 1.8× (avg 1.2×) kernel speedup. • End-to-End GNN Training: Up to 1.9× faster in GNN training compared to DGL/PyG. Limitations & Future Work • The "Fixed Non-zero(𝑓 = 1)" strategy is sub-optimal for moderate sparsity (50%-80%). • Residual part becomes overly small in massive graphs, limiting GPU thread utilization. • Future Work: Needs an adaptive approach to dynamically adjust f and g parameters. 37 2026/4/17
33. Questions and Thoughts 38 2026/4/17 • Extra attention must be paid to the actual memory cost during the preprocessing phase. • Even with 99% sparsity, local sparsity patterns can trigger excessive TC Blocks and reduce efficiency. • The reduction in atomic writes is less significant than claimed; • The primary source of improvement (BRP or reduced instructions) remains unclear. • For Krylov iterations, SpMM primarily appears in specialized variants such as Block-GMRES.
34. References [1] Yuke Wang, Boyuan Feng, Zheng Wang, Guyue Huang, and Yufei Ding. TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs. In Proceedings of the 2023 USENIX Annual Technical Conference (ATC 2023), pages 149–164. USENIX, 2023. [2] Ruibo Fan, Wei Wang, and Xiaowen Chu. DTC-SpMM: Bridging the Gap in Accelerating General Sparse Matrix Multiplication with Tensor Cores. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS 2024), pages 253–267. ACM, 2024. 39 2026/4/17