# A Reinforcement Learning Environment for Polyhedral Optimizations

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 37 references that are in the map.
- First posted on arXiv 28 Apr 2021.

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

- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [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)
- [Tiramisu: A Polyhedral Compiler for Expressing Fast and Portable Code](https://mlsanity.com/paper/1804.10694/) (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)
- [A Deep Learning Based Cost Model for Automatic Code Optimization](https://mlsanity.com/paper/2104.04955/) (2021)
- [ProGraML: Graph-based Deep Learning for Program Optimization and Analysis](https://mlsanity.com/paper/2003.10536/) (2020)
- [Static Neural Compiler Optimization via Deep Reinforcement Learning](https://mlsanity.com/paper/2008.08951/) (2020)
- [Machine Learning in Compiler Optimisation](https://mlsanity.com/paper/1805.03441/) (2018)
- [TDO-CIM: Transparent Detection and Offloading for Computation In-memory](https://mlsanity.com/paper/2007.00060/) (2020)

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

- [CompilerGym: Robust, Performant Compiler Optimization Environments for AI Research](https://mlsanity.com/paper/2109.08267/) (2021)
- [Target-independent XLA optimization using Reinforcement Learning](https://mlsanity.com/paper/2308.14364/) (2023)
- [LoopTune: Optimizing Tensor Computations with Reinforcement Learning](https://mlsanity.com/paper/2309.01825/) (2023)
- [CHEHAB RL: Learning to Optimize Fully Homomorphic Encryption Computations](https://mlsanity.com/paper/2601.19367/) (2026)
- [DRIVE: Dynamic Rule Inference and Verified Evaluation for Constraint-Aware Autonomous Driving](https://mlsanity.com/paper/2508.04066/) (2025)

## Abstract (arXiv)

The polyhedral model allows a structured way of defining semantics-preserving transformations to improve the performance of a large class of loops. Finding profitable points in this space is a hard problem which is usually approached by heuristics that generalize from domain-expert knowledge. Existing problem formulations in state-of-the-art heuristics depend on the shape of particular loops, making it hard to leverage generic and more powerful optimization techniques from the machine learning domain. In this paper, we propose PolyGym, a shape-agnostic formulation for the space of legal transformations in the polyhedral model as a Markov Decision Process (MDP). Instead of using transformations, the formulation is based on an abstract space of possible schedules. In this formulation, states model partial schedules, which are constructed by actions that are reusable across different loops. With a simple heuristic to traverse the space, we demonstrate that our formulation is powerful enough to match and outperform state-of-the-art heuristics. On the Polybench benchmark suite, we found transformations that led to a speedup of 3.39x over LLVM O3, which is 1.83x better than the speedup achieved by ISL. Our generic MDP formulation enables using reinforcement learning to learn optimization policies over a wide range of loops. This also contributes to the emerging field of machine learning in compilers, as it exposes a novel problem formulation that can push the limits of existing methods.