# A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in RL training of language and agent models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 20 of its 69 references that are in the map.
- First posted on arXiv 17 Sep 2024.

## Placed in

- Deep reinforcement learning methods / [RL training of language and agent models](https://mlsanity.com/deeprl/rl-training-of-language-and-agent-models/)

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Pointer Networks](https://mlsanity.com/paper/1506.03134/) (2015)
- [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)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (2018)
- [Ithemal: Accurate, Portable and Fast Basic Block Throughput Estimation using Deep Neural Networks](https://mlsanity.com/paper/1808.07412/) (2018)
- [CompilerGym: Robust, Performant Compiler Optimization Environments for AI Research](https://mlsanity.com/paper/2109.08267/) (2021)
- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://mlsanity.com/paper/2104.04955/) (2021)
- [Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation](https://mlsanity.com/paper/2001.08743/) (2020)
- [AutoPhase: Juggling HLS Phase Orderings in Random Forests with Deep Reinforcement Learning](https://mlsanity.com/paper/2003.00671/) (2020)
- [Compiling ONNX Neural Network Models Using MLIR](https://mlsanity.com/paper/2008.08272/) (2020)
- [Reinforced Genetic Algorithm Learning for Optimizing Computation Graphs](https://mlsanity.com/paper/1905.02494/) (2019)
- [AutoPhase: Compiler Phase-Ordering for High Level Synthesis with Deep Reinforcement Learning](https://mlsanity.com/paper/1901.04615/) (2019)
- [MLGOPerf: An ML Guided Inliner to Optimize Performance](https://mlsanity.com/paper/2207.08389/) (2022)
- [TIRAMISU: A Polyhedral Compiler for Dense and Sparse Deep Learning](https://mlsanity.com/paper/2005.04091/) (2020)
- [LOOPer: A Learned Automatic Code Optimizer For Polyhedral Compilers](https://mlsanity.com/paper/2403.11522/) (2024)
- [X-RLflow: Graph Reinforcement Learning for Neural Network Subgraphs Transformation](https://mlsanity.com/paper/2304.14698/) (2023)

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

- [Enhancing Code LLMs with Reinforcement Learning in Code Generation: A Survey](https://mlsanity.com/paper/2412.20367/) (2024)
- [FasterPy: An LLM-based Code Execution Efficiency Optimization Framework](https://mlsanity.com/paper/2512.22827/) (2025)
- [Curriculum-Guided Reinforcement Learning for Synthesizing Gas-Efficient Financial Derivatives Contracts](https://mlsanity.com/paper/2509.23976/) (2025)

## Abstract (arXiv)

Code optimization is a crucial task that aims to enhance code performance. However, this process is often tedious and complex, highlighting the necessity for automatic code optimization techniques. Reinforcement Learning (RL) has emerged as a promising approach for tackling such complex optimization problems. In this project, we introduce MLIR RL, an RL environment for the MLIR compiler, dedicated to facilitating MLIR compiler research and enabling automatic code optimization. We propose a multi-discrete formulation of the action space where the action space is the Cartesian product of simpler action subspaces. We also propose a new method, called level pointers, to reduce the size of the action space related to the loop interchange transformation. This enables more efficient and effective learning of the policy. To demonstrate the effectiveness of MLIR RL, we train an RL agent to optimize MLIR Linalg code, targeting CPU. The code is generated from two domain-specific frameworks: deep-learning models generated from PyTorch, and LQCD (Lattice Quantum Chromodynamics) code generated from an LQCD compiler. The result of this work is a research environment that allows the community to experiment with novel ideas in RL-driven loop-nest optimization.