# A Deep Learning Model for Predicting Transformation Legality

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 28 references that are in the map.
- First posted on arXiv 8 Nov 2025.

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

- [Learning to Optimize Tensor Programs](https://mlsanity.com/paper/1805.08166/) (2018)
- [Device Placement Optimization with Reinforcement Learning](https://mlsanity.com/paper/1706.04972/) (2017)
- [Neural Code Comprehension: A Learnable Representation of Code Semantics](https://mlsanity.com/paper/1806.07336/) (2018)
- [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)
- [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)
- [AutoPhase: Compiler Phase-Ordering for High Level Synthesis with Deep Reinforcement Learning](https://mlsanity.com/paper/1901.04615/) (2019)
- [LOOPer: A Learned Automatic Code Optimizer For Polyhedral Compilers](https://mlsanity.com/paper/2403.11522/) (2024)
- [LoopTune: Optimizing Tensor Computations with Reinforcement Learning](https://mlsanity.com/paper/2309.01825/) (2023)

## Abstract (arXiv)

Compilers must check the legality of code transformations to guarantee the correctness of applying a sequence of code transformations to a given code. While such a legality check needs to be precisely computed in general, we can use an approximate legality prediction model in certain cases, such as training a reinforcement learning (RL) agent for schedule prediction. In this paper, we propose an approximate method for legality checks. We propose a novel DL model for predicting the legality of transformations. The model takes the code representation and a list of transformations as input and predicts whether applying those transformations to the code is legal. We implement and evaluate the proposed model, demonstrating its effectiveness. Our evaluation shows an F1 score of 0.91 on a test set of randomly generated programs. To further evaluate the model in a practical scenario, we used the model to replace the legality check used during the training of an RL agent designed for automatic code optimization. We demonstrate that such a replacement enables the agent to train on twice as many steps, resulting in faster training and reducing resource usage by approximately 80% for CPU and 35% for RAM. The agent trained using this approach maintains comparable performance, with only a 4% reduction on benchmarks from the Polybench suite compared to the traditional method.