A Deep Learning Model for Predicting Transformation Legality

TL;DR

This work addresses the high cost of exact legality checks when exploring code transformation schedules in polyhedral compilers. It introduces a DL-based legality predictor that ingests a program representation and a sequence of transformations to output a legal/illegal decision, achieving an F1 score of $0.91$ on a synthetic test set. The model employs a three-layer architecture with computation embedding, a recursive loop embedding via dual LSTMs, and a binary classifier, trained on $5.8$ million training examples and $1.21$ million test examples generated from synthetic TiramisU programs. Integrating the predictor into RL training yields roughly a twofold increase in training steps and substantial resource reductions (CPU ~80%, RAM ~35%) with only about a 4% drop in Polybench performance, illustrating the practical benefits of differentiable legality guidance for search and optimization. The work also provides open-source data and demonstrates potential extensions to gradient-based search and broader applicability in compiler optimization.

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.

Figures (9)

• Figure 1: Illustration of Program Representation with Loop and Assignment Vectors
• Figure 2: DL Model Architecture for Transformation Legality Prediction
• Figure 3: Confusion matrix for the randomly generated test data
• Figure 4: Comparison of mean reward between the Legality-Model-Based Agent, the Dependency-Analysis-Based Agent,and Random-Legality-Based Agent over training steps.
• Figure 5: Geometric means of the speedup of the Legality-Model-Based Agent and Random-Legality-Based Agent over the dependency-analysis agent on the Polybench suite