# Target-independent XLA optimization using Reinforcement Learning

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 52 references that are in the map.
- First posted on arXiv 28 Aug 2023.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [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)
- [The Deep Learning Compiler: A Comprehensive Survey](https://mlsanity.com/paper/2002.03794/) (2020)
- [CompilerGym: Robust, Performant Compiler Optimization Environments for AI Research](https://mlsanity.com/paper/2109.08267/) (2021)
- [A Survey on Compiler Autotuning using Machine Learning](https://mlsanity.com/paper/1801.04405/) (2018)
- [MLGO: a Machine Learning Guided Compiler Optimizations Framework](https://mlsanity.com/paper/2101.04808/) (2021)
- [NeuroVectorizer: End-to-End Vectorization with Deep Reinforcement Learning](https://mlsanity.com/paper/1909.13639/) (2019)
- [Beyond variance reduction: Understanding the true impact of baselines on policy optimization](https://mlsanity.com/paper/2008.13773/) (2020)
- [Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation](https://mlsanity.com/paper/2001.08743/) (2020)
- [Transferable Graph Optimizers for ML Compilers](https://mlsanity.com/paper/2010.12438/) (2020)
- [AutoPhase: Juggling HLS Phase Orderings in Random Forests with Deep Reinforcement Learning](https://mlsanity.com/paper/2003.00671/) (2020)
- [Mind Mappings: Enabling Efficient Algorithm-Accelerator Mapping Space Search](https://mlsanity.com/paper/2103.01489/) (2021)
- [A2C is a special case of PPO](https://mlsanity.com/paper/2205.09123/) (2022)
- [Static Neural Compiler Optimization via Deep Reinforcement Learning](https://mlsanity.com/paper/2008.08951/) (2020)
- [A new Potential-Based Reward Shaping for Reinforcement Learning Agent](https://mlsanity.com/paper/1902.06239/) (2019)
- [A Reinforcement Learning Environment for Polyhedral Optimizations](https://mlsanity.com/paper/2104.13732/) (2021)
- [A Collaborative Filtering Approach for the Automatic Tuning of Compiler Optimisations](https://mlsanity.com/paper/2005.04092/) (2020)

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

- [Per-example gradients: a new frontier for understanding and improving optimizers](https://mlsanity.com/paper/2510.00236/) (2025)

## Abstract (arXiv)

An important challenge in Machine Learning compilers like XLA is multi-pass optimization and analysis. There has been recent interest chiefly in XLA target-dependent optimization on the graph-level, subgraph-level, and kernel-level phases. We specifically focus on target-independent optimization XLA HLO pass ordering: our approach aims at finding the optimal sequence of compiler optimization passes, which is decoupled from target-dependent optimization. However, there is little domain specific study in pass ordering for XLA HLO. To this end, we propose introducing deep Reinforcement Learning (RL) based search for optimal XLA HLO pass ordering. We also propose enhancements to the deep RL algorithms to further improve optimal search performance and open the research direction for domain-specific guidance for RL. We create an XLA Gym experimentation framework as a tool to enable RL algorithms to interact with the compiler for passing optimizations and thereby train agents. Overall, in our experimentation we observe an average of 13.3% improvement in operation count reduction on a benchmark of GPT-2 training graphs and 10.4% improvement on a diverse benchmark including GPT-2, BERT, and ResNet graphs using the proposed approach over the compiler's default phase ordering.