A Deep Learning Model for Predicting Transformation Legality

2025-11-08

AI Summary

This work addresses the high computational cost of verifying the legality of code transformations in compiler optimization. We propose an end-to-end deep learning model that directly learns the mapping between program representations and multi-step transformation sequences, replacing traditional precise but inefficient compiler-based legality checks. Our approach innovatively integrates abstract syntax tree encoding with sequential modeling of transformation histories to enable efficient, approximate legality prediction. Evaluated on a standard benchmark, the model achieves an F1 score of 0.91. When integrated into a reinforcement learning–driven automated optimization framework, it reduces training steps by 2×, decreases CPU and memory overhead by 80% and 35%, respectively, while incurring only a 4% degradation in end-to-end performance. These results demonstrate substantial acceleration of correctness-sensitive optimization search—without compromising soundness guarantees—thereby enabling scalable, reliable compiler auto-tuning.

Technology Category

• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP
• Machine Learning: Optimization

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Predicting legality of code transformations using deep learning
• Replacing precise compiler checks with approximate legality model
• Accelerating reinforcement learning training for automatic code optimization

Innovation

• Deep learning model predicts code transformation legality
• Model inputs are code representation and transformation list
• Approximate legality check accelerates reinforcement learning training