ParamSpMM: Adaptive and Efficient Sparse Matrix-Matrix Multiplication on GPUs for GNNs⋆
arXiv:2605.15695v1 [cs.DC] 15 May 2026
Lixing Zhang, Guanhua Ye, Hongzheng Li, Shigang Li, and Yingxia Shao Beijing University of Posts and Telecommunications, Beijing, China {zhanglixing, g.ye, Ethan_Lee, lishigang, shaoyx }@bupt.edu.cn
Abstract. Fueled by the ability to mine real-world graph data, GNN applications have experienced phenomenal growth. Sparse Matrix-Matrix Multiplication (SpMM) is a critical operator in GNNs. However, existing SpMM designs for GNNs struggle to adapt to diverse input characteristics. In this paper, we first conduct a comprehensive analysis of existing SpMM optimizations, revealing their limitations through statistical and empirical evidence. Based on this analysis, we introduce ParamSpMM, a parametric approach for highly adaptive and efficient SpMM computation in GNNs. It incorporates a new data structure, the Parameterized Compressed Sparse Row (PCSR), to flexibly integrate existing optimization techniques. ParamSpMM enables the configuration of these optimization techniques according to various input characteristics. Furthermore, we complement ParamSpMM with an ML-based SpMMdecider that predicts optimal configurations based on carefully crafted input features. Our evaluations demonstrate that ParamSpMM outperforms Nvidia cuSPARSE with an average speedup of 1.92×, significantly enhancing GNN training efficiency. Keywords: graph neural networks · sparse matrix-matrix multiplication · GPU
1
Introduction
Graph Neural Networks (GNNs) have emerged as a vital tool in graph data mining, showing advantages in traffic prediction [9], recommendation systems [27], chemistry [8], and biomedicine [17]. GNNs are renowned for their scatterand-gather paradigm, where each node aggregates the neighbors’ embeddings and updates its own embedding accordingly. This computation can be modeled as Sparse Matrix-Matrix Multiplication (SpMM), where the sparse matrix is the graph’s adjacency matrix, and the dense input matrix corresponds to the embedding matrix. ⋆
This manuscript is an author-posted version of a paper accepted by the 9th International Workshop on Graph Data Management and Analysis (GDMA 2025), held in conjunction with the 30th International Conference on Database Systems for Advanced Applications (DASFAA 2025).
2
Lixing et al.
SpMM often accounts for the majority of GNNs’ execution time, making it a primary bottleneck that requires optimization [13, 22]. However, the diversity of input graphs and embedding dimensions complicates SpMM optimization for GNNs in three aspects: Various data locality of graphs. If the neighboring nodes in a graph have close ID numbers, the graph will benefit from better data locality, resulting in a higher cache hit ratio in sparse linear computation [1,25]. However, the presence of diverse community structures in real-world graphs [7] leads to various data locality. Various degree distributions of graphs. There are Poisson degree distributions found in road traffic networks [3] and Power-law degree distributions observed in social networks [21]. Graphs with the latter distribution suffer more from workload imbalance in SpMM. Various embedding dimensions. GNNs have a wide range of network architectures with varying dimensions of node embeddings across layers. SpMM in GNNs thus involves dense input matrices with diverse column dimensions. However, existing SpMM designs lack the flexibility to address the above three aspects of diversity, offering limited customization of the SpMM kernel for specific inputs. Static kernels in popular vendor-provided frameworks like DGL [24] and PyG [6] cannot adapt to different inputs. On the other hand, existing dynamic designs [4,5,11,13,25] use fixed optimization strategies, providing limited adaptability by tuning a few parameters focusing on limited aspects of the whole performance issues. The limitations of existing dynamic designs are evident in two aspects: (1) Incomplete exploitation of optimization opportunities. Some works [11] employ blocking to utilize data locality but ignore workload imbalance caused by skewed degree distributions. Others [4, 5, 14, 25] focus on workload balancing without adequately utilizing data locality. (2) Poor adaptability to diverse SpMM inputs. Existing works struggle with unexpected inputs, demonstrating a mismatch between applied optimization techniques and the input characteristics. The blocking technique in ASpT [11] becomes inefficient on ultra-sparse matrices due to excessive zero padding. Some work [5,25] apply workload balancing by default, degrading performance on graphs with relatively balanced degree distribution. In this paper, we first provide an in-depth analysis of existing SpMM optimizations and reveal their limitations in handling diverse SpMM inputs. Based on the analysis, we introduce ParamSpMM, a parametric approach for SpMM computation in GNNs. The basic idea under the hood of ParamSpMM is to integrate the three existing optimization strategies—blocking, workload balancing, and thread coarsening—to enable high adaptability across diverse SpMM inputs, as each optimization strategy is tailored to a specific aspect of SpMM input diversity. We design an adaptive representation for the sparse matrix, called Parameterized Compressed Sparse Row (PCSR) format, to enable the seamless integration of the three optimizations. Based on the new format, ParamSpMM
ParamSpMM
3
provides adaptive and customizable high-performance SpMM kernels via flexible parameter configuration, that responds to the various input characteristics. While ParamSpMM offers comprehensive configuration options, manually identifying the optimal configuration remains challenging. To fully capitalize on the characteristics of diverse SpMM inputs, we introduce an ML-based SpMMdecider for ParamSpMM to predict the optimal configurations. With carefully crafted features across size, degree distribution, and data locality, SpMM-decider automates complicated configuration tasks, thereby augmenting ParamSpMM with self-driven capability. Extensive evaluations demonstrate that ParamSpMM secures an 1.92× average speedup over Nvidia cuSPARSE [20]. Embedded into GNNs training, ParamSpMM demonstrates an average of 1.96× speedup over DGL [24] across various GNN models and datasets. We outline contributions as follows: – We conduct an in-depth investigation into the limitations of existing SpMM optimization methods, revealing their inefficiencies in handling diverse data localities, degree distributions, and node embedding dimensions. – We develop ParamSpMM, a flexible parametric framework for SpMM computation that dynamically adjusts to the varying characteristics of SpMM inputs in GNNs. This approach enables high-performance SpMM kernel customization, offering significant adaptability and efficiency across diverse input scenarios. – We propose an ML-based SpMM-decider for the optimal configuration of ParamSpMM to efficiently leverage the SpMM input characteristics.
2
Background
In this section, we introduce the basic Compressed Sparse Row (CSR) based SpMM. Notably, SpMM is defined as An×n Bn×dim = Cn×dim , involving the sparse matrix A and dense matrices B and C. dim is B’s column dimension. CSR format. CSR represents a sparse matrix via three arrays—rowP tr, colIdx, and val—facilitating efficient traversal of nonzero elements within a sparse row. The traversal range of i-th row is specified by (rowP tr[i], rowP tr[i+ 1]). Based on CSR, a prevalent approach for computing SpMM on GPUs is in a row-wise manner, where a worker (e.g., a thread warp of ω threads) computes a segment of C[i][:] while traversing nonzeros in A[i][:]. Thread Mapping. A commonly employed thread mapping scheme arranges thread blocks into a 2D grid [30]. A thread block consists of W thread warps, each identified by (blk.x, blk.y, warpId). Warps with the same (blk.x, warpId) but different blk.y compute different ω-length segments in the same C row. CSR-based SpMM. The essence of SpMM computation is the multiplyaccumulate (MAC) operation, which multiplies a nonzero from A with a corresponding element in B and accumulates the result to C. Given that dim is a multiple of ω, Algorithm 1 describes the computing process from a thread’s perspective. The algorithm first determines the output position in C (Lines 15). Then, it establishes the traversal range (Lines 6-7), which also indicates the
4
Lixing et al.
Algorithm 1 Basic CSR-based SpMM on GPUs Input: rowP tr, colIdx, val, W Output: C 1: warpId = threadId/ω 2: laneId = threadId%ω 3: /*position of a computed segment in C*/ 4: Crow = blk.x ∗ W + warpId 5: seg = blk.y ∗ ω + laneId 6: head = rowP tr[Crow] 7: tail = rowP tr[Crow + 1] 8: for i in [head, tail) do 9: /*one multiply-accumulate (MAC) job*/ 10: Brow, v = colIdx[i], val[i] 11: res+ = v ∗ B[Brow][seg] 12: end for 13: C[Crow][seg + i ∗ ω] = res
current warp’s workload. For each iteration, a MAC operation is performed by retrieving B’s data based on the colIdx of the traversed nonzero (Lines 9-11). Finally, the results of the accumulated segments are written to C (Line 13).
3
Analysis of Existing SpMM Optimizations
In this section, we provide relevant statistics and empirical findings that highlight the limitations of existing optimizations. These observations underpin our design of ParamSpMM. Details on the experimental setup are provided in Section 6.1. 3.1
Analysis of Blocking Technique
Basic CSR-based SpMM has irregular memory access patterns for matrix B due to the non-uniform layout of nonzeros in A as demonstrated in Lines 10-11 of Algorithm 1. To mitigate this issue, blocking techniques are employed in sparse kernels [11, 18, 31, 32] to reduce irregular memory accesses. Nonzeros in the same column of A exhibit identical memory access patterns to B during MAC operations. By grouping these nonzeros for processing by a single thread, we can effectively exploit B’s data reuse through registers or shared memory. For graph datasets with extreme sparsity (99.99%) [2, 12, 16], vectorized blocking (e.g., 2 × 1) is more efficient than 2D blocking (e.g., 2 × 2) by minimizing zero padding and thus reducing unnecessary computations. Thus, we adopt vectorized blocking in this work, organizing A’s nonzeros into vectors based on a configurable Vector Size (V). Observation: The effectiveness of vectorized blocking varies with graph data locality. Inappropriate V choices can lead to excessive zero padding and performance degradation. Table 1 shows SpMM throughput with different V on four representative real-world graphs when dim = 32, with the corresponding zero padding ratios listed in parentheses. For V = 1, the padding ratio is 0%, as it introduces no zero padding. Notably, with V = 3, 97.5% of our 202 graphs show
ParamSpMM
5
padding ratios above 50%, indicating over half of the computation is redundant. Therefore, in this paper, the domain of V is limited to {1, 2}. The optimal V and padding ratios vary by graph, highlighting the necessity of adaptive V for blocking tailored to each graph’s data locality. Table 1. Throughput in GFLOPS under Various V, with zero padding ratios listed in parentheses. graphs coPapersCiteseer coPapersDBLP sx-askubuntu sx-mathoverflow V=1 V=2 V=3
3.2
1794 (0%) 2170 (26.8%) 1635 (39.6%)
1672 (0%) 1870 (30.6%) 1470 (44.6%)
386 (0%) 337 (47.8%) 318 (64.3%)
865 (0%) 767 (49.0%) 700 (65.6%)
Analysis of Workload Balancing Technique
The workload assigned to a thread warp is defined as l = tail−head, representing the traversal range of nonzeros (Lines 5-6 of Algorithm 1). Workload imbalance in SpMM stems from uneven l distribution across sparse rows. The nonzero-split approach for workload balancing redistributes workloads across warps to mitigate imbalances in sparse computations [4, 5, 19, 25, 30]. It redistributes the workload of a heavily loaded thread warp across multiple warps by splitting nonzeros within a sparse row. While this improves workload balance, it compromises memory efficiency as each warp must independently write its partial results to the same output segment, resulting in repeated writes of some results. In contrast, SpMM without balancing writes each result to C only once. Observation: The effectiveness of workload balancing varies with graph degree distribution. Applying balancing on graphs with relatively balanced degree distributions may not yield performance improvements, owing to increased memory writes and additional overhead from balancing-related bookkeeping instructions. Figure 1 demonstrates the throughputs (GFLOPS) of SpMM with and without workload balancing on 12 real-world graphs when
2000 Throughput (GFLOPS)
1750 1500
S False True
1250 1000 750 500 250 0
z ama
r ntu BLP 302 312 505 t-TX t-PA flow iki-RfA iki-Vote t-CA see on0 azon0 azon0 oadNe oadNe oadNe rsCite apersD askubu thover w w r r e r m m a P x a a s ap co sx-m coP Fig. 1. Throughputs of SpMM with or without workload balancing.
6
Lixing et al.
dim = 32. For simplicity, we use S, a boolean value, to represent whether to apply balancing. The effectiveness of balancing varies across graphs, necessitating adaptive balancing. This observation highlights the importance of the adaptive application of workload balancing based on graph degree distribution. 3.3
Analysis of Thread Coarsening Technique
Thread coarsening [5, 13, 23, 25, 30] allocates a set of independent MAC jobs to a single thread for each nonzero traversal iteration, and the size of the set is determined by Coarsening Factor (F ). Each warp reuses the fetched value and column index of a nonzero from A in one iteration to execute F ∗ ω MAC jobs, producing a length of F ∗ ω segment results in C, with F ∈ [1, CEIL( dim ω )]. Thread coarsening brings a theoretical F-fold reduction in the memory transactions for A’s data and improves instruction-level parallelism (ILP). However, it increases register usage and reduces active threads on the fly [13]. Therefore, an improper F potentially degrades performance. Additionally, when dim is not a multiple of ω ∗F, the residual warp (computing the last segment of a C row) may have a MAC-job gap compared to other warps with the same (blk.x, warpId) but different blk.y. This MAC-job gap occurs under an improper F and wastes computing resources, which is overlooked by existing works. We quantify it as gapF : tnF dim = min(dim, F · ω) F trdim = dim
mod (F · ω)
F F gapF dim = tndim − trdim
(1)
F where tnF dim is the ideal MAC-job number under chosen F and dim, and trdim is the actual MAC-job number for the residual warp. Observation: The effectiveness of thread coarsening varies with dim and graphs. Given dim ∈ {64, 96, 128, 160}, Table 2 presents the distribution of optimal F settings across 202 real-world graphs. N/A means that the value of F is out of range and is thus exempted for a given dim. The corresponding MAC-job gap for each dim is also presented in parentheses. The results reveal: (1) Improper F degrades performance due to MAC-job gaps for a given dim. F = 2 for dim = 96, F = 3 for dim = 128, and F = 2, 3, 4 for dim = 160 are less preferred due to their wider MAC-job gaps compared to other F options. (2) Despite some F demonstrating the same MAC-job gap (e.g., F = 1, 2 when dim = 64), identifying the optimal F for a specific sparse matrix (graph) is nontrivial, involving memory transactions, local resource usage, and parallelism. These findings necessitate adaptive F selection based on both dim and graph characteristics.
3.4
Summary
As analyzed above, existing optimizations inadequately accommodate diverse SpMM inputs. A comprehensive design that flexibly integrates these optimizations is crucial for constructing a robust SpMM kernel accommodating three
ParamSpMM
7
Table 2. The Distribution of Optimal F Setting dim F 1 2 3 4 5
64
96
128
160
20.30% (0) 24.26% (0) 21.29% (0) 25.74% (0) 79.70% (0) 1.98% (32) 36.14% (0) 5.94% (32) N/A 73.76% (0) 0% (64) 6.44% (32) N/A N/A 42.57% (0) 0% (96) N/A N/A N/A 61.88% (0)
aspects of input diversity. Therefore, we introduce ParamSpMM with key parameters: V for vectorized blocking, S for workload balancing, F for thread coarsening, and W to determine the thread block size. These parameters collectively enable fine-tuned optimization, adapting to various input characteristics.
4
Backbone Design of ParamSpMM
In this section, we first provide an overview of ParamSpMM. Then, we introduce Parameterized Compressed Sparse Row (PCSR), the data representation of a sparse matrix in ParamSpMM. Subsequently, we detail the computing engine and present graph reordering to further enhance ParamSpMM’s performance. 4.1
Overview of ParamSpMM
ParamSpMM consists of three phases: configuration prediction, PCSR generation, and SpMM computing (Figure 2). First, the ML-based SpMM-decider (Section 5) predicts optimal ⟨W, F , V, S⟩ using input features. Second, it applies vectorized blocking to extract nonzero vectors from A, followed by workload balancing if S is T rue. Workload balancing redistributes workload across warps and generates output positions for partial result accumulation in C. Finally, based on PCSR, the computing engine executes the SpMM kernel with customized optimization of blocking, workload balancing, and thread coarsening. DASFAA 4.2
Parameterized Compressed Sparse Row Format
To enable seamless integration of various optimizations in ParamSpMM, we propose Parameterized Compressed Sparse Row (PCSR), an adaptive representation for sparse matrix A.
0
2
Layer 𝑘 0
2
1
3
①Config Prediction
②PCSR Generation
③Computing
SpMM-decider
Vectorized Blocking rearrange nonzeros into vectors
Computing Engine
1
3
Input Layer 𝑘 + 1 Info. 0
2
GNN model
1
3
size
Features
Graph
degree distribution data locality
Config
Workload Balancing redistribute workload
PCSR Balancing Blocking
Thread Coarsening
Fig. 2. Overview of ParamSpMM: a three-phase workflow for SpMM in GNNs.
8
Lixing et al. 1 ⟨ 𝒮 = 𝐹𝑎𝑙𝑠𝑒, 𝒱 = 1 ⟩ 𝑎
𝑏
𝑒
𝑓
2 ⟨ 𝒮 = 𝐹𝑎𝑙𝑠𝑒, 𝒱 = 2 ⟩
𝑐
𝑔
𝑑
ℎ
4
PCSR
8 10
7
1
2
3
0
2
3
𝑒
𝑏
𝑓
𝑐
0
𝑑
1
2
3
0
1
0
2
3
𝑐𝑜𝑙𝐼𝑑𝑥
𝑏
𝑐
𝑑
𝑒
𝑓
𝑔
ℎ T-Row 𝑖 𝑗
𝑣𝑎𝑙
𝑎
TW1
TW2
Split
𝑓
𝑔
𝑔
𝑖
ℎ
0
0
𝑗
TW1
𝑇𝑅𝑜𝑤
3 ⟨ 𝒮 = 𝑇𝑟𝑢𝑒, 𝒱 = 1, 𝑆𝐺 = 2 ⟩
𝑒
0
TW0
TW3
𝑇𝑅𝑜𝑤
𝑏
PCSR
4
𝑣𝑎𝑙 𝑎
0
𝑗
0
𝑐𝑜𝑙𝐼𝑑𝑥 0
𝑎
𝑑
ℎ
𝑟𝑜𝑤𝑃𝑡𝑟
0
TW0
𝑐
𝑖
𝑗 6
𝑏 𝑓
𝑔
𝑖
𝑟𝑜𝑤𝑃𝑡𝑟 0
𝑎 𝑒
𝑐
Split
rowPtr Reassign
𝑑
4
0 ℎ
0
𝑖
4 ⟨ 𝒮 = 𝑇𝑟𝑢𝑒, 𝒱 = 2, 𝑆𝐺 = 3 ⟩
2
6 4
8 6
𝑎 𝑒
10 8
𝑏 𝑓
𝑑
ℎ
0
𝑖
𝑗
PCSR
𝑟𝑜𝑤𝑃𝑡𝑟 0
2
4
6
8 10
𝑐𝑜𝑙𝐼𝑑𝑥
0
1
2
3
0
1
0
2
3
𝑣𝑎𝑙 𝑎
𝑏
𝑐
𝑑
𝑒
𝑓
𝑔
ℎ T-Row 𝑖 𝑗
TW0
𝑇𝑅𝑜𝑤
TW1
TW2
TW3
0
1
3
0
2
0
TW4
4
0
𝑔
10
rowPtr Reassign
𝑐
7
3
4
7
0
0
𝑗
PCSR
𝑟𝑜𝑤𝑃𝑡𝑟 0
3
4
7
𝑐𝑜𝑙𝐼𝑑𝑥
0
1
2
3
0
2
3
𝑣𝑎𝑙
𝑎
𝑒
𝑏
𝑓
𝑐
0
𝑑
TW0
𝑇𝑅𝑜𝑤
0
TW1
0
0
𝑔
𝑖
ℎ
𝑗
TW2
1
Fig. 3. PCSR under four different configurations. TW: Thread Warp.
Data Representation of PCSR. PCSR represents a sparse matrix via four arrays: rowP tr, colIdx, val, and T Row, arranging matrix elements into V × 1 nonzero vectors. rowP tr defines the traversal range of nonzero vectors and indicates each thread warp’s workload l. val and colIdx store the values and column indices of nonzero vectors. T Row is for workload balancing: without balancing, a thread warp determines its output position in C using (blk.x, warpId) (Lines 3,11 of Algorithm 1), leaving T Row empty. With balancing, the workload is redistributed across multiple warps, using T Row to store target positions in C to accumulate partial results, and rowP tr is reassigned to redistribute the workload across thread warps. PCSR Generation. Given the configuration of ⟨W, F, V, S⟩, it begins by dividing A into multiple row panels of height V, followed by vectorized blocking to generate rowP tr, colIdx, and val for the nonzero vectors (size of V × 1) within row panels. Subsequently, if S = T rue, for workload balancing, rowP tr is reassigned to redistribute the workload across thread warps. T Row is generated to store the target positions in C for partial results accumulation. Examples of PCSR. Figure 3 illustrates the examples of PCSR under four different configurations, where different row panels are marked with different colors. In Examples (2) and (4), val contains flattened vectors with some necessary zero padding when the number of nonzeros in a vector is less than V. Elements such as a and e can thus reuse the data from dense matrix B when packed into a vector, attributed to their identical column index. Examples (3) and (4) with S = T rue are cases when workload balancing is applied. Example (3) differs from Example (1) by sharing the original workload of thread warp 0
ParamSpMM
9
with warp 1, thus improving workload balance. However, Example (4) reveals a worse workload distribution after balancing, where the original workload in Example (2) is already fairly balanced, highlighting the importance of proper S decisions as discussed in Section 3.2. Algorithm 2 ParamSpMM Computing Engine Input: PCSR, dim, ⟨W, F, V, S⟩ Output: C 1: /*initialization (Lines 1-6 of Algorithm 1)*/ 2: res[V][F] = {0} 3: t = min(F, CEILDIV (dim − seg, ω)) 4: for i in [head, tail) do 5: /*thread coarsening: prefetch val and Brow*/ 6: Brow = PCSR.colIdx[i] 7: v[0 : V) = PCSR.val[i ∗ V : (i + 1) ∗ V) 8: /*performing t MAC operations*/ 9: for j in [0, t) do 10: /*blocking: prefetch Bval for data reuse*/ 11: Bval = B[Brow][seg + j ∗ ω] 12: for k in [0, V) do 13: res[k][j]+ = v[k] ∗ Bval 14: end for 15: end for 16: end for 17: if S == F alse then 18: Crow = Crow ∗ V 19: for i in [0, t) do 20: for j in [0, V) do 21: C[Crow + j][seg + i ∗ ω] = res[j][i] 22: end for 23: end for 24: else 25: Crow = PCSR.T Row[Crow] ∗ V 26: for i in [0, t) do 27: for j in [0, V) do 28: atomicAdd(res[j][i], C[Crow + j][seg + i ∗ ω]) 29: end for 30: end for 31: end if
4.3 Computing Engine The computing engine executes SpMM using PCSR with a customized kernel based on the predicted ⟨W, F , V, S⟩ by SpMM-decider. Algorithm 2 outlines the core functionality from the perspective of a thread in a warp. The algorithm begins by defining the traversal ranges of nonzero vectors via PCSR.rowP tr, setting thread mapping with W, and initializing res buffer (Lines
10
Lixing et al.
1-2). The variable t (Line 3) determines the actual MAC operations per thread in each nonzero vector iteration (Line 4). t is crucial for handling the boundary cases when dim is not a multiple of ω ∗ F , since some threads may execute fewer than F MAC operations in one iteration. Subsequently, it iterates through the allocated nonzero vectors (Line 4). For each vector, it performs thread coarsening by prefetching PCSR.val and PCSR.colIdx (Lines 6-7) for subsequent t MAC operations within the nested loops (Lines 9-15). Through blocking, it reuses B’s values V times (Lines 11-14), reducing memory transactions. Lastly, if workload balancing is not applied (S = F alse), the results are directly written to C (Lines 18-23). Otherwise, when S = T rue, it uses the PCSR.T Row to locate the write-back positions and performs atomicAdd to handle potential conflicts in partial results accumulation (Lines 25-30). 4.4
Graph Reordering to Enhance Data Locality
Various graph reordering techniques [1,5,14,25,26] have been devised to improve data locality by rearranging nodes with similar neighbors, which share similar memory access patterns for B during computation, to be positioned closer. This complements ParamSpMM of V > 1 by creating more consecutive nonzeros in the same columns, leading to less zero padding. Appropriate application of blocking can further exploit the enhanced data locality from graph reordering. Practically speaking, Rabbit Reordering [1] is a default step in ParamSpMM, which is highly parallelized and time-efficient. The reordering cost is amortizable as reordered graphs can be reused in iterative applications.
5
Learning Optimal ParamSpMM Configuration
As the optimal ParamSpMM configuration varies across SpMM inputs and is challenging to determine manually, we propose a data-driven ML-based SpMMdecider to predict the optimal configuration. In this section, we introduce our carefully crafted sparse matrix features and the configuration prediction model. 5.1
Sparse Matrix Features
We classify sparse matrix features into three categories: (1) Size features: They decide the scale of threads and overall workload, guiding the setting of F and W; (2) Degree distribution features: They navigate the decision of workload balancing; (3) Data locality features: They help the selection of V in vectorized blocking. These features, detailed in Table 3, can be measured once to train models and configure ParamSpMM across different dim, which is amortizable in iterative applications. While most features are well-known, we specifically introduce two essential metrics, SR and P R, to guide ParamSpMM’s optimization of blocking and balancing: P R for Blocking. As discussed in Section 3.1, V in vectorized blocking significantly impacts zero padding levels, consequently affecting ParamSpMM performance. To quantify the zero padding level, we introduce P RV : nnz (2) P RV = 1 − nnzV ∗ V
ParamSpMM
11
Table 3. Summary of Features Category
size features
Feature Description n n b nnz δ d db dmax
degree CV d distribution CV features SRi data locality features
ρ b bmax P Ri
number of rows in a sparse matrix number of non-empty rows in a sparse matrix number of nonzeros in a sparse matrix the ratio of n b to n the average number of nonzeros per row d without considering empty rows max number of nonzeros across sparse rows coefficient of variation of node degrees CV without considering empty rows SR under ⟨V = i, S = true⟩ the density of nonzeros in a sparse matrix average bandwidth of all rows max bandwidth1 across all rows in a sparse matrix padding ratio under vectorized blocking of V = i
where nnz is the number of nonzeros and nnzV is the number of nonzero vectors. P RV ranges from [0, 1 − V1 ], with higher values indicating increased zero padding. Lower P RV values reflect better data reuse and reduced unnecessary computation, potentially leading to improved performance. Consequently, P RV serves as one of many input features for the ML-based SpMM-decider. SR for Balancing. Our workload balancing method sets a workload upper bound, Split Granularity (SG), ensuring each warp processes no more than SG nonzero vectors. We set SG as: SG = CEILDIV (dc (3) V , ω) ∗ ω where ω is the thread warp size. After balancing, multiple warps may accumulate partial results to the same segment in C, increasing memory writes as discussed in Section 3.2. To quantify this overhead, we introduce Split Ratio (SR): split_size (4) row_size where split_size and row_size are the lengths of reassigned and original rowP tr arrays, respectively. SR ∈ [1, +∞] measures the increased memory writes to C. For example, SR = 1.2 indicates 1.2× more writes than without balancing. Given the complexity of manually selecting the optimal strategy based on SR, the ML-based SpMM-decider is employed for this decision. SR =
5.2
SpMM-decider: Configuration Prediction Model
Despite different ParamSpMM configurations serving distinct optimization purposes, establishing a definitive boundary among their application contexts is challenging. An ML-based SpMM-decider is thus trained to predict optimal 1
We define the bandwidth of one row in the sparse matrix as the difference in column indices between the first nonzero element and the last one.
12
Lixing et al. GE-SpMM
GNNASpMM
DA-SpMM
ParamSpMM
speedup
4 3 2 1 0
16
32
48
64
80
96
112
128
144
160
176
192
208
224
240
256
Fig. 4. The performance of SpMM across various dim. The reported speedup is normalized to cuSPARSE.
ParamSpMM configurations based on the crafted input features, enabling systematic and data-driven SpMM optimization. With the sparse matrix features listed in Table 3 as input, ML-based SpMMdecider predicts the optimal configuration of ⟨W, F , V, S⟩. SpMM-decider is based on the random forests model, which is a lightweight ensemble learning model. This approach facilitates easier model training and deployment, as well as a lower risk of overfitting compared to other more complicated ML models.
6
Experimental Evaluation
In this section, we evaluate ParamSpMM’s performance, technique effectiveness, and application in GNN training. 6.1
Experimental Setup
Experiment Platform. We evaluate ParamSpMM and baselines on NVIDIA A6000, which has 84 Amper SMs and a compute capability of 8.6. We integrate ParamSpMM as an extension into PyTorch, leveraging its expressiveness in GNNs. The code is available at https://github.com/zzzlxhhh/ParamSpMM_AE, which is compiled with GCC 9.4.0 and NVCC 11.6 using -O3 optimization. Datasets. For SpMM datasets, we collect 202 matrices (1, 000 < n < 8, 000, 000) from SNAP [16] and DIMACS10 [2]. Too small matrices are exempted for the evaluation consistency concerns. Matrices that are too large are excluded due to possible out-of-memory errors. For GNN datasets, we take 6 graphs from OGB [12], which have not previously been used for training SpMM-decider. SpMM Baselines. We compare ParamSpMM with three categories of baselines: Static (Nvidia cuSPARSE [20] and GE-SpMM [13]), Heuristic-based (GNNAdvisor [25]), and ML-based (DA-SpMM [4]). All methods use Rabbit Reordering [1] as default preprocessing. GNN Evaluation Models. We evaluate ParamSpMM on two representative GNN models: Graph Convolutional Network (GCN) [15] and Graph Isomorphic Network (GIN) [29]. We use GNNAdivisor [25] (a GNN system with optimized SpMM kernel) and DGL [24] (a vendor-provided GNN library) as baselines.
ParamSpMM
6.2
13
Overall ParamSpMM Performance
We evaluate ParamSpMM’s performance across 202 matrices with varying dim, resulting in 3232 diverse SpMM inputs exhibiting a substantial diversity. These matrices exhibit diverse characteristics, including matrix size (n from 1005 ∼ 7, 733, 822), density (ρ from 2.73 × e−7 ∼ 0.025), coefficient of variation (CV from 0.00638 ∼ 58.097), zero padding ratio (P R2 from 0.247 ∼ 0.499). To visualize the adaptability of ParamSpMM and the baselines across various SpMM inputs, Figure 4 presents the boxplot of the speedups of each method over cuSPARSE on A6000, demonstrating the speedups distribution of each method. ParamSpMM significantly outperforms all baselines across a wide range of SpMM input diversity, demonstrating superior adaptability. While baseline methods occasionally underperform cuSPARSE, ParamSpMM maintains robust acceleration in most cases. Table 4 further summarizes the speedups of ParamSpMM over baselines across various dim. ParamSpMM achieves an average speedup of 1.92×, 2.41×, 1.55×, and 1.64× over cuSPARSE, GE-SpMM, GNNAdivisor, and DA-SpMM, demonstrating superior performance. Table 4. The speedups of ParamSpMM over baselines on A6000. Category
Baselines
Static
cuSPARSE GE-SpMM
16
32
48
64
80
96
112
128
144
160
176
192
208
224
240
256 Average
2.68× 2.07× 1.73× 1.81× 1.77× 1.90× 1.84× 1.97× 2.00× 1.90× 1.88× 1.87× 1.77× 1.75× 1.87× 1.93× 1.92× 4.32× 3.03× 3.30× 2.79× 3.00× 2.67× 2.78× 2.00× 2.24× 2.08× 1.96× 1.88× 1.98× 1.88× 1.84× 1.82× 2.41×
Heuristic GNNAdivisor 2.44× 1.70× 1.87× 1.67× 1.71× 1.57× 1.81× 1.29× 1.31× 1.28× 1.39× 1.54× 1.34× 1.33× 1.54× 1.34× 1.55× ML
DA-SpMM
2.10× 1.95× 1.64× 1.58× 1.69× 1.64× 1.66× 1.53× 1.57× 1.58× 1.58× 1.54× 1.57× 1.49× 1.56× 1.52× 1.64×
6.3 Effectiveness of SpMM-decider To evaluate SpMM-decider’s prediction accuracy, we split the datasets into 80% training and 20% testing sets. We compare the performance of predicted configurations (pre) against optimal configurations, with randomly configured ParamSpMM (rnd) serving as a baseline. Table 5 shows that SpMM-decider effectively provides suitable configurations, with most normalized performances exceeding 99%, whereas rnd shows significantly lower performance. Table 5. The normalized performance of SpMM-decider dim 16 32 48 64
pre
rnd
dim
pre
rnd
dim
pre
rnd
dim
pre
rnd
98.84% 82.26% 80 99.97% 70.58% 144 99.28% 74.04% 208 99.29% 68.60% 99.69% 73.47% 96 99.55% 75.94% 160 99.74% 72.13% 224 99.67% 79.42% 99.98% 70.34% 112 99.31% 70.45% 176 99.21% 73.49% 240 99.19% 70.56% 98.24% 76.65% 128 99.30% 78.14% 192 98.96% 74.18% 256 98.75% 70.33%
6.4 Effectiveness of Graph Reordering In this section, we evaluate the impact of graph reordering. With ParamSpMM_wor and cuSPARSE_wor to denote ParamSpMM and cuSPARSE without reordering, Table 6 reveals three key findings: (1) Graph reordering significantly enhances ParamSpMM’s performance over ParamSpMM_wor. (2) Even without reordering, ParamSpMM_wor still substantially outperforms cuSPARSE_wor due to our parametric approach. (3) While graph reordering provides cuSPARSE a modest 1.14× speedup over cuSPARSE_wor, ParamSpMM better leverages the improved data locality, achieving a 1.26× speedup over ParamSpMM_wor.
14 Lixing et al. Table 6. The speedups of cuSPARSE, ParamSpMM, and ParamSpMM_wor over cuSPARSE_wor. dim
16
32
48
64
80
96
112
128 Average
cuSPARSE 1.06× 1.08× 1.15× 1.16× 1.18× 1.17× 1.17× 1.15× 1.14× ParamSpMM_wor 1.91× 1.92× 1.66× 1.70× 1.66× 1.72× 1.68× 1.76× 1.75× ParamSpMM 2.80× 2.19× 1.95× 2.07× 2.06× 2.20× 2.13× 2.26× 2.21×
6.5
Application in GNNs
We evaluate ParamSpMM in 5-layer GCN and GIN with input/output sizes of 16 and hidden sizes of {32, 64, 128} on A6000. Figure 5 shows training speedups of ParamSpMM and GNNAdvisor over DGL. ParamSpMM consistently outperforms DGL, achieving average speedups of 1.60× (up to 2.19×) for GCN and 1.61× (up to 2.59×) for GIN. As SpMM is key in both models, ParamSpMM effectively accelerates both GCN and GIN. model = GCN | hidden_size = 32
DGL GNNAdvisor ParamSpMM model = GCN | hidden_size = 64 model = GCN | hidden_size = 128
2.5
speedup
2.0 1.5 1.0 0.5 0.0
model = GIN | hidden_size = 32
model = GIN | hidden_size = 64
model = GIN | hidden_size = 128
2.5
speedup
2.0 1.5 1.0 0.5
cit
ati on co 2 lla b pp a ve sse l arx iv pro du c cit ts ati on co 2 lla b pp a ve sse l arx iv pro du c cit ts ati on co 2 lla b pp a ve sse l arx iv pro du cts
0.0
Fig. 5. GCN and GIN performance speedup normalized to DGL.
7
Related Work
In this section, we discuss the related works for SpMM optimization. Blocking. Blocking techniques [11,18,28,32] exploit data locality. FlashLLM [28] and Magicube [18] utilize the high throughput tensor core units to facilitate SpMM in sparse DNN training [10], where the sparse matrices are relatively dense and uniform. Other studies [11, 31, 32] propose hybrid blocking strategies with better adaptability to various data locality. Workload Balancing. While merge-based balancing [4, 19, 30] accelerates SpMM for Power-law distributions, it requires costly binary searches for nonzeros’ row indices. GNNAdvisor [25] proposes a complex shared memory accumulation mechanism to reduce global memory atomic operations after balancing. Thread Coarsening. While several works [5,13,25] employ thread coarsening to reduce memory access, they fail to configure an appropriate F. GNNAdvisor [25] and GE-SpMM [13] simply increase F with dim without considering
ParamSpMM
15
MAC-job gaps. HP-SpMM [5] applies vectorized instructions (float2 and float4) to achieve thread coarsening, leading to limited F choices from 2 and 4. ML-guided Optimization. ML is used as an auto-tuning approach in some recent works [4, 32]. DA-SpMM [4] utilizes ML to optimize SpMM on GPUs, yet their strategy space overlooks blocking and thread coarsening. DDB [32] focuses on hardware with matrix-multiply units.
8
Conclusions
In this work, we presented ParamSpMM, a flexible parametric framework for optimizing GPU-based SpMM kernels against diverse inputs in GNNs. After analyzing existing works’ limitations in handling input diversity, we flexibly integrated blocking, workload balancing, and thread coarsening techniques, offering customizable optimization through parameter configuration. We introduced PCSR format to enable seamless cooperation among optimization techniques. We developed an ML-based SpMM-decider for automatic configuration prediction based on input characteristics. Extensive evaluations show ParamSpMM’s superior adaptability with an average 1.92× speedup over cuSPARSE [20] while achieving significant acceleration over DGL [24] in various GNN models.
9
Acknowledgments
This work is supported by National Natural Science Foundation of China (Nos. 62272054, 62192784, 62372055), Beijing Nova Program (No. 20230484319, 20250484968), State Key Laboratory of Multimedia Information Processing Open Fund (No. SKLMIP-KF-2025-07), and CAAI-CANN Open Fund, developed on OpenI Community (No. CAAIXSJLJJ2025CANN10). Yingxia Shao is the corresponding author.
References 1. Arai, J., Shiokawa, H., Yamamuro, T., Onizuka, M., Iwamura, S.: Rabbit order: Just-in-time parallel reordering for fast graph analysis. In: (IPDPS). pp. 22–31 (2016) 2. Bader, D.A., Meyerhenke, H., Sanders, P., Wagner, D.: 10th dimacs implementation challenge-graph partitioning and graph clustering (2011) 3. Chandrasekaran, B.: Survey of network traffic models. Waschington University in St. Louis CSE 567 (2009) 4. Dai, G., Huang, G., Yang, S., Yu, Z., Zhang, H., Ding, Y., Xie, Y., Yang, H., Wang, Y.: Heuristic adaptability to input dynamics for spmm on gpus. p. 595–600. DAC ’22 (2022) 5. Fan, R., Wang, W., Chu, X.: Fast sparse gpu kernels for accelerated training of graph neural networks. In: (IPDPS). pp. 501–511 (2023) 6. Fey, M., Lenssen, J.E.: Fast graph representation learning with pytorch geometric (2019) 7. Fortunato, S.: Community detection in graphs. Physics Reports 486(3–5), 75–174 (Feb 2010) 8. Gilmer, J., Schoenholz, S.S., Riley, P.F., Vinyals, O., Dahl, G.E.: Neural message passing for quantum chemistry. p. 1263–1272. ICML’17 (2017)
16
Lixing et al.
9. Guo, S., Lin, Y., Feng, N., Song, C., Wan, H.: Attention based spatialtemporal graph convolutional networks for traffic flow forecasting. AAAI’19/IAAI’19/EAAI’19 (2019) 10. Hoefler, T., Alistarh, D., Ben-Nun, T., Dryden, N., Peste, A.: Sparsity in deep learning: pruning and growth for efficient inference and training in neural networks. J. Mach. Learn. Res. 22(1) (jan 2021) 11. Hong, C., Sukumaran-Rajam, A., Nisa, I., Singh, K., Sadayappan, P.: Adaptive sparse tiling for sparse matrix multiplication. p. 300–314. PPoPP ’19 (2019) 12. Hu, W., Fey, M., Zitnik, M., Dong, Y., Ren, H., Liu, B., Catasta, M., Leskovec, J.: Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems 33, 22118–22133 (2020) 13. Huang, G., Dai, G., Wang, Y., Yang, H.: Ge-spmm: General-purpose sparse matrixmatrix multiplication on gpus for graph neural networks. In: SC20. pp. 1–12 (2020) 14. Huang, K., Zhai, J., Zheng, Z., Yi, Y., Shen, X.: Understanding and bridging the gaps in current gnn performance optimizations. p. 119–132. PPoPP ’21 (2021) 15. Kipf, T.N., Welling, M.: Semi-supervised classification with graph convolutional networks (2017) 16. Leskovec, J., Krevl, A.: Snap datasets: Stanford large network dataset collection (2014) 17. Li, M.M., Huang, K., Zitnik, M.: Graph representation learning in biomedicine and healthcare. Nature Biomedical Engineering 6(12), 1353–1369 (2022) 18. Li, S., Osawa, K., Hoefler, T.: Efficient quantized sparse matrix operations on tensor cores. SC ’22 (2022) 19. Merrill, D., Garland, M.: Merge-based sparse matrix-vector multiplication (spmv) using the csr storage format. PPoPP ’16 (2016) 20. Naumov, M., Chien, L., Vandermersch, P., Kapasi, U.: Cusparse library. In: GPU Technology Conference (2010) 21. Sala, A., Zheng, H., Zhao, B.Y., Gaito, S., Rossi, G.P.: Brief announcement: revisiting the power-law degree distribution for social graph analysis. p. 400–401. PODC ’10 (2010) 22. Shao, Y., Li, H., Gu, X., Yin, H., Li, Y., Miao, X., Zhang, W., Cui, B., Chen, L.: Distributed graph neural network training: A survey. ACM Comput. Surv. 56(8) (Apr 2024). https://doi.org/10.1145/3648358, https://doi.org/ 10.1145/3648358 23. Volkov, V., Demmel, J.W.: Benchmarking gpus to tune dense linear algebra. In: SC ’08. pp. 1–11 (2008) 24. Wang, M., Zheng, D., Ye, Z., Gan, Q., Li, M., Song, X., Zhou, J., Ma, C., Yu, L., Gai, Y., Xiao, T., He, T., Karypis, G., Li, J., Zhang, Z.: Deep graph library: A graph-centric, highly-performant package for graph neural networks (2020) 25. Wang, Y., Feng, B., Li, G., Li, S., Deng, L., Xie, Y., Ding, Y.: Gnnadvisor: An adaptive and efficient runtime system for gnn acceleration on gpus. In: (OSDI 21). pp. 515–531 (2021) 26. Wei, H., Yu, J.X., Lu, C., Lin, X.: Speedup graph processing by graph ordering. p. 1813–1828. SIGMOD ’16 (2016) 27. Wu, S., Sun, F., Zhang, W., Xie, X., Cui, B.: Graph neural networks in recommender systems: A survey. ACM Comput. Surv. 55(5) (dec 2022) 28. Xia, H., Zheng, Z., Li, Y., Zhuang, D., Zhou, Z., Qiu, X., Li, Y., Lin, W., Song, S.L.: Flash-llm: Enabling cost-effective and highly-efficient large generative model inference with unstructured sparsity. arXiv preprint arXiv:2309.10285 (2023) 29. Xu, K., Hu, W., Leskovec, J., Jegelka, S.: How powerful are graph neural networks? (2019)
ParamSpMM
17
30. Yang, C., Buluç, A., Owens, J.D.: Design principles for sparse matrix multiplication on the gpu. p. 672–687. Euro-Par ’18 (2018) 31. Yang, X., Parthasarathy, S., Sadayappan, P.: Fast sparse matrix-vector multiplication on gpus: Implications for graph mining. Proc. VLDB Endow. 4(4), 231–242 (jan 2011) 32. Yesil, S., Moreira, J.E., Torrellas, J.: Dense dynamic blocks: optimizing spmm for processors with vector and matrix units using machine learning techniques. ICS ’22, New York, NY, USA (2022)