# GeoT: Tensor Centric Library for Graph Neural Network via Efficient Segment Reduction on GPU

arXiv: https://arxiv.org/abs/2404.03019 · HTML: https://mlsanity.com/paper/2404.03019/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 44 of its 88 references that are in the map.
- First posted on arXiv 3 Apr 2024.

## Builds on (most cited first, 30 of 44)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://mlsanity.com/paper/1706.02413/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Simplifying Graph Convolutional Networks](https://mlsanity.com/paper/1902.07153/) (2019)
- [Geometric deep learning: going beyond Euclidean data](https://mlsanity.com/paper/1611.08097/) (2016)
- [Pitfalls of Graph Neural Network Evaluation](https://mlsanity.com/paper/1811.05868/) (2018)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://mlsanity.com/paper/1810.02244/) (2018)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking](https://mlsanity.com/paper/1707.03815/) (2017)
- [GraphCast: Learning skillful medium-range global weather forecasting](https://mlsanity.com/paper/2212.12794/) (2022)
- [Graph Neural Network for Traffic Forecasting: A Survey](https://mlsanity.com/paper/2101.11174/) (2021)
- [MeshGPT: Generating Triangle Meshes with Decoder-Only Transformers](https://mlsanity.com/paper/2311.15475/) (2023)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [MLPerf Inference Benchmark](https://mlsanity.com/paper/1911.02549/) (2019)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Predicting multicellular function through multi-layer tissue networks](https://mlsanity.com/paper/1707.04638/) (2017)
- [PinnerSage: Multi-Modal User Embedding Framework for Recommendations at Pinterest](https://mlsanity.com/paper/2007.03634/) (2020)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs](https://mlsanity.com/paper/2006.06608/) (2020)
- [SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning](https://mlsanity.com/paper/2207.04606/) (2022)
- [GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks](https://mlsanity.com/paper/2007.03179/) (2020)

## Abstract (arXiv)

In recent years, Graph Neural Networks (GNNs) have ignited a surge of innovation, significantly enhancing the processing of geometric data structures such as graphs, point clouds, and meshes. As the domain continues to evolve, a series of frameworks and libraries are being developed to push GNN efficiency to new heights. While graph-centric libraries have achieved success in the past, the advent of efficient tensor compilers has highlighted the urgent need for tensor-centric libraries. Yet, efficient tensor-centric frameworks for GNNs remain scarce due to unique challenges and limitations encountered when implementing segment reduction in GNN contexts. We introduce GeoT, a cutting-edge tensor-centric library designed specifically for GNNs via efficient segment reduction. GeoT debuts innovative parallel algorithms that not only introduce new design principles but also expand the available design space. Importantly, GeoT is engineered for straightforward fusion within a computation graph, ensuring compatibility with contemporary tensor-centric machine learning frameworks and compilers. Setting a new performance benchmark, GeoT marks a considerable advancement by showcasing an average operator speedup of 1.80x and an end-to-end speedup of 1.68x.