# Using Graph Neural Networks to model the performance of Deep Neural Networks

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 20 references that are in the map.
- First posted on arXiv 27 Aug 2021.

## Builds on (most cited first, 8 of 8)

- [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)
- [Graph Neural Networks: A Review of Methods and Applications](https://mlsanity.com/paper/1812.08434/) (2018)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [TVM: An Automated End-to-End Optimizing Compiler for Deep Learning](https://mlsanity.com/paper/1802.04799/) (2018)
- [Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks](https://mlsanity.com/paper/1808.07412/) (2018)
- [A Survey on Compiler Autotuning using Machine Learning](https://mlsanity.com/paper/1801.04405/) (2018)
- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://mlsanity.com/paper/2104.04955/) (2021)

## Built on by (most cited first, 2 of 2)

- [The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding](https://mlsanity.com/paper/2110.10221/) (2021)
- [Geometric Learning in Black-Box Optimization: A GNN Framework for Algorithm Performance Prediction](https://mlsanity.com/paper/2506.16144/) (2025)

## Abstract (arXiv)

With the unprecedented proliferation of machine learning software, there is an ever-increasing need to generate efficient code for such applications. State-of-the-art deep-learning compilers like TVM and Halide incorporate a learning-based performance model to search the space of valid implementations of a given deep learning algorithm. For a given application, the model generates a performance metric such as the run time without executing the application on hardware. Such models speed up the compilation process by obviating the need to benchmark an enormous number of candidate implementations, referred to as schedules, on hardware. Existing performance models employ feed-forward networks, recurrent networks, or decision tree ensembles to estimate the performance of different implementations of a neural network. Graphs present a natural and intuitive way to model deep-learning networks where each node represents a computational stage or operation. Incorporating the inherent graph structure of these workloads in the performance model can enable a better representation and learning of inter-stage interactions. The accuracy of a performance model has direct implications on the efficiency of the search strategy, making it a crucial component of this class of deep-learning compilers. In this work, we develop a novel performance model that adopts a graph representation. In our model, each stage of computation represents a node characterized by features that capture the operations performed by the stage. The interaction between nodes is achieved using graph convolutions. Experimental evaluation shows a 7:75x and 12x reduction in prediction error compared to the Halide and TVM models, respectively.