[Paper Review] TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs

Yuke Wang, Boyuan Feng+3 more|arXiv (Cornell University)|Dec 3, 2021

Advanced Graph Neural Networks7 citations

TL;DR

TC-GNN is the first GNN acceleration framework that bridges sparse graph neural network (GNN) computation with dense GPU Tensor Core Units (TCUs). It introduces a novel sparse graph translation (SGT) technique to condense irregular sparse neighbor data into dense tiles compatible with TCUs, while jointly optimizing CUDA cores and TCUs for high-performance GNN inference and training, achieving an average 1.70× speedup over state-of-the-art frameworks like DGL across diverse models and datasets.

ABSTRACT

Recently, graph neural networks (GNNs), as the backbone of graph-based machine learning, demonstrate great success in various domains (e.g., e-commerce). However, the performance of GNNs is usually unsatisfactory due to the highly sparse and irregular graph-based operations. To this end, we propose TC-GNN, the first GNN acceleration framework based on GPU Tensor Core Units (TCUs). The core idea is to reconcile the "Sparse" GNN computation with the high-performance "Dense" TCUs. Specifically, we conduct an in-depth analysis of the sparse operations in mainstream GNN computing frameworks. We introduce a novel sparse graph translation technique to facilitate TCU processing of the sparse GNN workload. We implement an effective CUDA core and TCU collaboration design to fully utilize GPU resources. We integrate TC-GNN with the PyTorch framework for high programmability. Rigorous experiments show an average of 1.70X speedup over the state-of-the-art DGL framework across various models and datasets.

Motivation & Objective

• To address the performance bottleneck in GNNs caused by highly sparse and irregular graph operations, which dominate runtime (>80%) in training and inference.
• To enable effective utilization of GPU Tensor Core Units (TCUs), which are optimized for dense matrix operations but underused in existing sparse GNN frameworks.
• To overcome the limitations of cuSPARSE and cuBLAS by designing a system that reconciles sparse GNN computation with dense TCU execution.
• To provide a high-programmability solution integrated with PyTorch, minimizing user learning overhead while maximizing hardware utilization.
• To explore hardware-aware optimizations that can inspire future AI accelerators with dynamic TCU input shapes and support for structural sparsity.

Proposed method

• Introduces a novel Sparse Graph Translation (SGT) technique that reorganizes sparse neighbor data into dense tiles by exploiting frequent neighbor sharing across nodes, reducing redundant memory access.
• Designs a hybrid CUDA core and TCU kernel architecture: CUDA cores manage fine-grained memory access and data movement, while TCUs accelerate dense GEMM operations on condensed tiles.
• Employs memory and data flow optimizations to minimize idle time and maximize data reuse, particularly for high-dimensional node embeddings.
• Enables seamless integration with PyTorch via custom APIs, allowing users to leverage TC-GNN without changing their existing programming workflow.
• Supports dynamic configuration of TCU block sizes and precision formats via compile-time parameters, ensuring adaptability across different GPU architectures.
• Leverages recent CUDA features that support blocked SpMM with TCU, but extends it to handle irregular sparse graphs through SGT pre-processing.

Experimental results

Research questions

• RQ1 Can Tensor Core Units (TCUs) be effectively utilized for sparse GNN computation, despite their design for dense matrix operations?
• RQ2 How can the inherent irregularity and sparsity of real-world graphs be transformed into a form compatible with fixed-size TCU input tiles?
• RQ3 What system-level design enables efficient collaboration between CUDA cores and TCUs for GNN workloads, avoiding performance degradation from naive TCU application?
• RQ4 To what extent can a framework like TC-GNN outperform existing state-of-the-art GNN frameworks such as DGL in terms of speedup across diverse models and datasets?
• RQ5 How can future GPU hardware be enhanced to natively support dynamic TCU tile shapes and structural sparsity for even greater performance gains?

Key findings

• TC-GNN achieves an average of 1.70× speedup over the state-of-the-art DGL framework across multiple GNN models and real-world datasets, demonstrating significant performance gains.
• The throughput of TC-GNN scales proportionally with increasing node embedding dimensions, indicating effective utilization of GPU resources and strong memory bandwidth scaling.
• The sparse graph translation (SGT) technique successfully reduces redundant memory access by condensing shared neighbors into fewer dense tiles, improving data locality and reducing overhead.
• The hybrid CUDA core and TCU kernel design effectively balances memory-intensive operations (handled by CUDA cores) and compute-intensive GEMM (handled by TCUs), maximizing hardware utilization.
• TC-GNN maintains high programmability through native PyTorch integration, enabling end-to-end GNN workloads with minimal code changes and high developer productivity.
• The framework is portable across modern GPUs (e.g., A100, H100, RTX4090) and supports configurable TCU parameters, ensuring adaptability to future GPU generations.