# Operator Fusion in XLA: Analysis and Evaluation

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

## Key facts

- Cited by 9 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 20 references that are in the map.
- First posted on arXiv 30 Jan 2023.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems](https://mlsanity.com/paper/1512.01274/) (2015)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions](https://mlsanity.com/paper/1802.04730/) (2018)
- [TVM: An Automated End-to-End Optimizing Compiler for Deep Learning](https://mlsanity.com/paper/1802.04799/) (2018)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [The Deep Learning Compiler: A Comprehensive Survey](https://mlsanity.com/paper/2002.03794/) (2020)
- [Horizontally Fused Training Array: An Effective Hardware Utilization Squeezer for Training Novel Deep Learning Models](https://mlsanity.com/paper/2102.02344/) (2021)

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

- [QiMeng-Kernel: Macro-Thinking Micro-Coding Paradigm for LLM-Based High-Performance GPU Kernel Generation](https://mlsanity.com/paper/2511.20100/) (2025)
- [Faster Predictive Coding Networks via Better Initialization](https://mlsanity.com/paper/2601.20895/) (2026)
- [FastCHGNet: Training one Universal Interatomic Potential to 1.5 Hours with 32 GPUs](https://mlsanity.com/paper/2412.20796/) (2024)
- [Survey on Computer Vision Techniques for Internet-of-Things Devices](https://mlsanity.com/paper/2308.02553/) (2023)
- [Graph neural networks with configuration cross-attention for tensor compilers](https://mlsanity.com/paper/2405.16623/) (2024)
- [Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization](https://mlsanity.com/paper/2605.25954/) (2026)
- [Nova: An End-to-End MLIR Compiler for Deep Learning](https://mlsanity.com/paper/2608.00029/) (2026)
- [Measuring and Reducing WebGPU Dispatch Overhead for LLM Inference](https://mlsanity.com/paper/2608.08730/) (2026)
- [KernelZero: Co-Evolving Proposer and Coder for Continuously Improved GPU Kernel Generation](https://mlsanity.com/paper/2609.33074/) (2026)

## Abstract (arXiv)

Machine learning (ML) compilers are an active area of research because they offer the potential to automatically speedup tensor programs. Kernel fusion is often cited as an important optimization performed by ML compilers. However, there exists a knowledge gap about how XLA, the most common ML compiler, applies this nuanced optimization, what kind of speedup it can afford, and what low-level effects it has on hardware. Our paper aims to bridge this knowledge gap by studying key compiler passes of XLA's source code. Our evaluation on a reinforcement learning environment Cartpole shows how different fusion decisions in XLA are made in practice. Furthermore, we implement several XLA kernel fusion strategies that can achieve up to 10.56x speedup compared to our baseline implementation.