A Deep Learning Model for Predicting Transformation Legality

TL;DR

This paper introduces a deep learning model that predicts the legality of code transformations, enabling faster and more resource-efficient training of code optimization agents with minimal impact on performance.

Contribution

The paper presents a novel DL-based legality prediction model that improves training efficiency for code optimization without sacrificing accuracy.

Findings

F1 score of 0.91 on test set

Training speed doubled with the model

Resource usage reduced by 80% CPU and 35% RAM

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.