PaperlayerView exact revision on arXiv

Source-linked AI summary

GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks

Guyue Huang, Guohao Dai, Yu Wang, Huazhong Yang

arXiv:2007.03179v1cs.DC

TL;DR

GNN workloads rely on SpMM, but existing GPU approaches struggle with general SpMM-like operations, preprocessing overhead, and inefficient sparse-data access. GE-SpMM provides a CSR-based design using Coalesced Row Caching and Coarse-grained Warp Merging. It achieves speedups over cuSPARSE and GraphBLAST and reduces CUDA time in embedded GNN models.

  • Problem

    GNN SpMM acceleration must support general SpMM-like operations without preprocessing overhead while addressing redundant and uncoalesced GPU memory access.

  • Method

    GE-SpMM is a CSR-based GPU SpMM design using Coalesced Row Caching and Coarse-grained Warp Merging to improve sparse-data access and reuse.

  • Results

    GE-SpMM achieves up to 1.41× speedup over Nvidia cuSPARSE, 1.81× over GraphBLAST, and 3.67× CUDA time reduction in embedded GNN models.

  • Takeaways & Limitations

    GE-SpMM provides a CSR-compatible SpMM-like kernel that can be embedded in GNN frameworks without data conversion overhead.

Abstract

from arXiv · show

Graph Neural Networks (GNNs) have achieved significant improvements in various domains. Sparse Matrix-Matrix multiplication (SpMM) is a fundamental operator in GNNs, which performs a multiplication between a sparse matrix and a dense matrix. Accelerating SpMM on parallel hardware like GPUs can face the following challenges: From the GNN application perspective, the compatibility needs to be considered. General GNN algorithms require SpMM-like operations (e.g., pooling) between matrices, which are not supported in current high-performance GPU libraries (e.g., Nvidia cuSPARSE). Moreover, the sophisticated preprocessing in previous implementations will lead to heavy data format conversion overheads in GNN frameworks. From the GPU hardware perspective, optimizations in SpMV (Sparse Matrix-Vector) designs on GPUs do not apply well to SpMM. SpMM exposes the column-wise parallelism in the dense output matrix, but straightforward generalization from SpMV leads to inefficient, uncoalesced access to sparse matrix in global memory. The sparse row data can be reused among GPU threads, which is neither possible in SpMM designs inherited from SpMV. To tackle these challenges, we propose GE-SpMM. GE-SpMM performs SpMM-like operation on sparse matrices represented in the most common Compressed Sparse Row (CSR) format, so it can be embedded in GNN frameworks with no preprocessing overheads and support general GNN algorithms. We introduce the Coalesced Row Caching method to process columns in parallel and ensure coalesced access to sparse matrix data. We also present the Coarse-grained Warp Merging to reduce redundant data loading among GPU warps. Experiments on a real-world graph dataset show that GE-SpMM achieves up to 1.41X speedup over Nvidia cuSPARSE and up to 1.81X over GraphBLAST. We also embed GE-SpMM in GNN frameworks and get up to 3.67X speedup over popular GNN models like GCN and GraphSAGE.

I. INTRODUCTION

SpMM is a time-consuming GNN operation whose GPU acceleration must support general SpMM-like operators without costly preprocessing while avoiding redundant and uncoalesced memory access. GE-SpMM addresses these requirements with CSR-based caching and warp-merging techniques, achieving speedups over established libraries and GNN models.

  • Motivation: GNNs propagate vertex feature vectors across graph edges, making neighbor aggregation a fundamental operation expressible as sparse adjacency–dense feature matrix multiplication.The aggregation operator can be general-purpose rather than simple summation.
  • Motivation: SpMM accounts for approximately 30% of CUDA time during the cited GCN training example, motivating targeted acceleration.Dense matrix multiplications account for approximately 10%, while other operators each account for less than 10%.
  • Challenges: GPU SpMM acceleration must support general SpMM-like operators and avoid substantial data-format conversion overhead in GNN frameworks.Existing frameworks use cuSPARSE for standard SpMM but may fall back to separate implementations for unsupported SpMM-like operations.
  • Challenges: SpMM exposes column-wise output parallelism, but straightforward SpMV generalization produces uncoalesced sparse-matrix accesses and redundant data loading.Unlike SpMV, SpMM can reach high bandwidth utilization while still suffering from excessive data movement.
  • GE-SpMM: GE-SpMM uses a CSR-based GPU kernel that integrates with GNN frameworks without data conversion overhead and supports SpMM-like operations.CSR is presented as a common format suitable for both SpMM and other sparse matrix operations.
  • GE-SpMM: Coalesced Row Caching and Coarse-grained Warp Merging improve sparse-data access and reuse, with up to 1.25× and 1.51× speedups respectively.CRC caches sparse rows in shared memory, while CWM merges warp workloads to reuse loaded sparse data.
  • Evaluation: GE-SpMM achieves up to 1.41× speedup over Nvidia cuSPARSE, 1.81× over GraphBLAST, and up to 3.67× CUDA-time reduction in embedded GNN models.The reported GNN models include GCN and GraphSAGE.

II. BACKGROUNDS AND RELATED WORKS

This section introduces the background context for SpMM and GNNs on GPUs, including the notation used throughout the paper.

  • The paper introduces background information about SpMM and GNNs on GPUs.The notation used in the paper is summarized in Table III.

A. GPU Preliminaries

GPU execution and memory organization shape SpMM performance, while existing GNN frameworks and SpMM implementations face compatibility, layout, and preprocessing limitations.

  • GPU Preliminaries: GPU blocks own shared memory, which supports cross-thread data reuse and reduces global-memory transactions.GPU threads execute in warps of 32 under CUDA's SIMT model.
  • GPU Preliminaries: Coalesced memory access occurs when warp threads access consecutive, aligned memory locations in one instruction.GPUs merge warp memory requests into as few transactions as possible.
  • Related Work: Sequentially applying SpMV fails to exploit parallelism across the dense output columns in SpMM.GraphBLAST instead assigns multiple threads to process one output row in parallel.
  • Related Work: Preprocessed sparse formats can outperform cuSPARSE but add storage, maintenance, and preprocessing costs that hinder GNN-framework adoption.The cited preprocessing time can reach 5× the actual SpMM computation time.
  • GNN Frameworks: DGL and PyG provide GNN abstractions, with DGL and PyG implementing custom SpMM-related kernels.DGL internally calls cuSPARSE for standard SpMM and uses its own kernels for unsupported SpMM-like operations.
  • GNN Frameworks: DGL transposes cuSPARSE output because csrmm2 produces column-major output while GNN feature matrices are row-major.This layout conversion is performed using cuBLAS.
  • GNN Frameworks: PyG's MessagePassing abstraction enables flexible user-defined operations but prevents fusion of edge-message generation and reduction into one SpMM kernel.The passage states that this generality leaves less room for optimizing specific operations such as SpMM.

III. GE-SPMM DESIGN

GE-SpMM is designed around CSR input to preserve compatibility with GNN frameworks while organizing SpMM computation over sparse rows and dense output columns.

  • Data Organization in GE-SpMM: GE-SpMM uses CSR because it is a widely used format across vendor libraries, data-science toolkits, and GNN frameworks.CSR stores row offsets, column indices, and nonzero values in row order.
  • Data Organization in GE-SpMM: Algorithm 2 assigns a thread block to a sparse row and a thread to an output column before loading row data in warp-sized tiles.The CRC algorithm uses shared-memory arrays sm k and sm v for the loaded column indices and values.
  • Data Organization in GE-SpMM: CSR-based SpMM computes each C[i,j] by traversing sparse row i and multiplying its nonzeros with corresponding entries from dense matrix B.The algorithm accesses A.rowPtr, then A.colInd and A.val, before writing the output element.

B. Coalesced Row Caching

Coalesced Row Caching improves sparse-row loading by staging warp-sized tiles in shared memory, then reusing them while computing dense output columns.

  • Coalesced Row Caching: The baseline parallelization coalesces dense-matrix accesses but not accesses to sparse-matrix arrays.Threads sharing a sparse row issue inefficient accesses to A.colInd and A.val.
  • Coalesced Row Caching: CRC partially unrolls the sparse-row loop by the warp size, loading each tile cooperatively into shared memory.A warp then consumes the staged elements sequentially from shared memory.
  • Coalesced Row Caching: CRC uses a two-phase load-and-consume strategy so sparse-row accesses are coalesced while each warp accesses the same dense-matrix row.The procedure repeats for additional tiles when a sparse row exceeds one tile.
  • Coarse-grained Warp Merging: In CWM, merging workloads across warps allows a sparse-matrix element to be loaded once instead of twice.The example uses N = 64 and CF = 2.
  • Coalesced Row Caching: CRC can reduce sparse-row load requests by nearly the warp size, although alignment and short rows reduce the realized reduction.The resulting algorithm still reduces load transactions and improves global-bandwidth efficiency.

C. Coarse-grained Warp Merging

Coarse-grained Warp Merging (CWM) reduces redundant sparse-row loading by assigning each thread multiple output values, trading memory reuse against parallelism and resource usage. The method is integrated with CRC in the GE-SpMM kernel.

  • C. Coarse-grained Warp Merging: CRC shares sparse-matrix data only within a warp, so different warps still perform redundant loading; CWM is introduced to address this limitation.Warp-level synchronization avoids the larger overhead of block-level synchronization.
  • C. Coarse-grained Warp Merging: Algorithm 3 combines CRC and CWM for CSR SpMM, loading sparse indices and values before producing two output values per thread when CF = 2.The algorithm takes A.rowPtr[], A.colInd[], A.val[], and B[] as inputs and writes C[].
  • C. Coarse-grained Warp Merging: CWM merges workloads across warps so threads produce multiple output values and redundant sparse-row loads are reduced.Threads in different warps cannot share data under CRC alone; merging their workloads addresses this redundancy.
  • C. Coarse-grained Warp Merging: The coarsening factor (CF) controls thread reduction: CF = 2 assigns each thread two output values and can reduce sparse-row load transactions by CF.Thread coarsening also exposes independent memory-loading instructions for instruction-level parallelism.
  • C. Coarse-grained Warp Merging: Increasing CF reduces memory loading but can hurt performance by lowering parallelism and increasing per-thread resource usage.The design must balance data reuse against the parallel threads needed to hide memory latency.
  • C. Coarse-grained Warp Merging: GE-SpMM embeds these techniques in GNN frameworks through its CSR-based kernel and an integration flow for existing frameworks.The framework integration is described as part of applying GE-SpMM to accelerate GNN applications.

A. GE-SpMM for Different Matrix Sizes

GE-SpMM supports arbitrary dense-matrix column sizes and user-defined SpMM-like operations, enabling integration into existing GNN frameworks without changing their general graph-operation workflow. Its evaluation uses GNN classification graphs and SNAP benchmark graphs.

  • A. GE-SpMM for Different Matrix Sizes: GE-SpMM applies CRC and CWM for arbitrary dense-matrix column size N, using both techniques when N > 32.The techniques target cases where loading the dense matrix creates a bandwidth bottleneck.
  • A. GE-SpMM for Different Matrix Sizes: User-defined initialization and associative, commutative reduction functions extend GE-SpMM from standard SpMM to general SpMM-like GNN operations.Common reductions such as sum and maximum satisfy the stated reduction requirements.
  • A. GE-SpMM for Different Matrix Sizes: GE-SpMM is integrated into DGL and PyG by wrapping the kernel as a PyTorch custom autograd function or replacing relevant graph-operation calls.DGL substitutes its CUDA kernel, while PyG uses the operator in place of MessagePassing calls in training code.
  • A. GE-SpMM for Different Matrix Sizes: Experiments evaluate GE-SpMM on Cora, Citeseer, Pubmed, and 66 valid graphs from the SNAP group in SuiteSparse.The selected graphs include node-classification workloads and a broader sparse-matrix benchmark.

2) SpMM Baselines:

The evaluation compares GE-SpMM with vendor, open-source, and graph-processing baselines under specified GPU and profiling conditions. CRC reduces global-load transactions and improves efficiency, while its performance benefit varies by GPU.

  • 2) SpMM Baselines:: GE-SpMM is compared with cuSPARSE csrmm2, GraphBLAST rowsplit, and the GunRock GPU graph-processing engine.These baselines represent vendor, open-source CSR, and graph-processing approaches.
  • 2) SpMM Baselines:: cuSPARSE csrmm2 produces column-major output, so GNN applications requiring row-major output must perform a matrix transpose in real applications.The reported kernel comparison excludes this transpose overhead.
  • 2) SpMM Baselines:: Experiments run on Nvidia GTX 1080Ti and RTX 2080 systems with different GPU architectures, memory capacities, and theoretical bandwidths.The machines use Pascal and Turing GPUs, respectively.
  • 2) SpMM Baselines:: CRC evaluation profiles global-load transactions and global-load efficiency to measure coalescing and bandwidth utilization.The tests use nvprof on three synthetic random graphs with N = 512.
  • 2) SpMM Baselines:: CRC significantly reduces total load transactions and improves memory-load efficiency through coalesced memory access.These effects are reported for the three synthetic random graphs used in the profiling tests.
  • 2) SpMM Baselines:: 1.246× average performance gain comes from CRC on GTX 1080Ti, whereas RTX 2080 shows an average of 1.011× from CRC alone.CRC remains the foundation for combining with CWM on RTX 2080.

2) Benefits of Coarse-grained Warp Merging:

CWM reduces global transactions through workload coarsening but trades lower parallelism and occupancy against data reuse. Across SNAP graphs, CF = 2 is the practical default, and combining CRC with CWM improves performance on RTX 2080.

  • 2) Benefits of Coarse-grained Warp Merging:: CWM profiling measures global-load transactions, global-load throughput, and achieved occupancy to expose its memory and parallelism trade-offs.The tested matrix has M = 65K, nnz = 650K, and N = 512.
  • 2) Benefits of Coarse-grained Warp Merging:: CWM reduces global data loading as CF increases, but its benefit becomes smaller when dense-matrix loading dominates transactions.Larger CF also lowers occupancy, indicating reduced parallelism.
  • 2) Benefits of Coarse-grained Warp Merging:: CF = 2 works well for most SNAP matrices, while CF > 4 shows an obvious performance drop.The comparison plots relative speedup over not using CWM for each SNAP graph.
  • 2) Benefits of Coarse-grained Warp Merging:: The reported SNAP comparisons include GraphBLAST, cuSPARSE, and GE-SpMM across 64 graphs, with omitted bars indicating out-of-memory cases.Matrix identifiers follow the alphabetical order of matrix names.
  • 2) Benefits of Coarse-grained Warp Merging:: 1.51× average improvement is achieved on RTX 2080 by combining CRC and CWM with CF = 2.CRC alone does not provide a large gain on RTX 2080 in the preceding tests.

C. Overall Performance of SpMM Kernel

GE-SpMM is evaluated against GPU SpMM baselines on three real-world GNN graphs, where its two optimizations improve kernel performance.

  • Kernel tests use single-precision matrices and dense-matrix column counts N from 128 through 512.
  • 1.62× maximum speedup over cuSPARSE is achieved on three graphs used in GNN models.The results indicate potential acceleration for real GNN models when GE-SpMM is integrated into GNN frameworks.

2) Graphs from SNAP:

On SNAP graphs, GE-SpMM becomes more competitive as the dense-matrix column count grows and can outperform both cuSPARSE and GraphBLAST. Against ASpT, preprocessing overhead materially affects the comparison, while GE-SpMM also outperforms a GunRock-based implementation.

  • Graphs from SNAP: 1.43× and 1.81× are the maximum speedups over cuSPARSE and GraphBLAST, respectively, on the SNAP dataset.The advantage increases with larger N, making the techniques especially relevant for applications with large dense-matrix column counts.
  • Comparison with Graph Engines on GPUs: 18.27× average speedup over the GunRock-based implementation shows that SpMM and GNN workloads require feature-dimension parallelism beyond traditional graph-processing primitives.GunRock does not provide parallelization along the feature dimension, which harms SpMM performance.
  • Comparison with ASpT: GE-SpMM reaches 0.93X, 0.97X, and 1.00X against ASpT for N=128, 256, and 512 on GTX1080Ti, respectively.On RTX2080, the corresponding averages are 0.85X, 0.93X, and 0.98X.
  • Comparison with ASpT: Preprocessing overhead ranges from 0.01× to 64.53× of actual SpMM execution time, averaging 0.47× on GTX1080Ti and 0.34× on RTX2080.Including one preprocessing step and one run, GE-SpMM is 1.43×∼2.06× against ASpT on average.

1) GNNs based on SpMM Operators:

GE-SpMM is embedded in DGL and PyG to accelerate GNN workloads involving SpMM or SpMM-like operations. Its benefits vary with the framework and feature dimension, while its flexibility also covers operations such as GraphSAGE pooling that cuSPARSE does not provide.

  • GNNs based on SpMM Operators: GE-SpMM is not competitive for output-layer SpMM when the feature length equals the number of classification classes and N is small.This limitation explains four GTX1080Ti tests in which GE-SpMM did not accelerate the original DGL implementation.
  • GNNs based on SpMM Operators: GE-SpMM brings up to 3.67× and 2.10× CUDA time reduction on two GPUs when integrated with PyG.PyG’s explicit message generation and reduction leave more room for improvement than DGL’s fused SpMM kernel.
  • SpMM-like Operators: GraphSAGE-pool uses maximum aggregation over neighbor features, an SpMM-like operation that cuSPARSE does not provide.GE-SpMM is intended to support such flexible, user-defined reductions in new GNN models.
  • Conclusion: GE-SpMM is a CSR-based design that can be embedded in GNN frameworks without preprocessing overhead.
Loading 2007.03179v1…