# Graph neural networks with configuration cross-attention for tensor compilers

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 24 references that are in the map.
- First posted on arXiv 26 May 2024.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Quantifying the Carbon Emissions of Machine Learning](https://mlsanity.com/paper/1910.09700/) (2019)
- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs](https://mlsanity.com/paper/2308.13490/) (2023)
- [Operator Fusion in XLA: Analysis and Evaluation](https://mlsanity.com/paper/2301.13062/) (2023)
- [Learning Large Graph Property Prediction via Graph Segment Training](https://mlsanity.com/paper/2305.12322/) (2023)
- [TensorFlow: A system for large-scale machine learning](https://mlsanity.com/paper/1605.08695/) (2016)

## Abstract (arXiv)

With the recent popularity of neural networks comes the need for efficient serving of inference workloads. A neural network inference workload can be represented as a computational graph with nodes as operators transforming multidimensional tensors. The tensors can be transposed and/or tiled in a combinatorially large number of ways, some configurations leading to accelerated inference. We propose TGraph, a neural graph architecture that allows screening for fast configurations of the target computational graph, thus representing an artificial intelligence (AI) tensor compiler in contrast to the traditional heuristics-based compilers. The proposed solution improves mean Kendall's τ across layout collections of TpuGraphs from 29.8% of the reliable baseline to 67.4% of TGraph. We estimate the potential CO2 emission reduction associated with our work to be equivalent to over 50% of the total household emissions in the areas hosting AI-oriented data centers.