Behavioral Embeddings of Programs: A Quasi-Dynamic Approach for Optimization Prediction

TL;DR

This work tackles the challenge of learning program representations that effectively predict compiler optimization outcomes. It introduces the Program Behavior Spectrum, a quasi-dynamic representation built by probing LLVM IR with diverse optimization sequences and measuring scale-invariant changes in a $56$-dimensional Autophase feature space. The spectrum is discretized via Product Quantization into a compositional code and learned with a multi-task Transformer named PQ-BERT, enabling deep modeling of the behavioral grammar. Empirically, Behavioral-PQ outperforms state-of-the-art static baselines on Best Pass Prediction and $-Oz$ Benefit Prediction, with strong ablations and embedding-space analyses supporting the value of scale-invariant, compositional encoding. The approach offers a scalable, interpretable path toward more accurate and robust compiler optimization predictions, with public code to support reproducibility and further exploration.

Abstract

Learning effective numerical representations, or embeddings, of programs is a fundamental prerequisite for applying machine learning to automate and enhance compiler optimization. Prevailing paradigms, however, present a dilemma. Static representations, derived from source code or intermediate representation (IR), are efficient and deterministic but offer limited insight into how a program will behave or evolve under complex code transformations. Conversely, dynamic representations, which rely on runtime profiling, provide profound insights into performance bottlenecks but are often impractical for large-scale tasks due to prohibitive overhead and inherent non-determinism. This paper transcends this trade-off by proposing a novel quasi-dynamic framework for program representation. The core insight is to model a program's optimization sensitivity. We introduce the Program Behavior Spectrum, a new representation generated by probing a program's IR with a diverse set of optimization sequences and quantifying the resulting changes in its static features. To effectively encode this high-dimensional, continuous spectrum, we pioneer a compositional learning approach. Product Quantization is employed to discretize the continuous reaction vectors into structured, compositional sub-words. Subsequently, a multi-task Transformer model, termed PQ-BERT, is pre-trained to learn the deep contextual grammar of these behavioral codes. Comprehensive experiments on two representative compiler optimization tasks -- Best Pass Prediction and -Oz Benefit Prediction -- demonstrate that our method outperforms state-of-the-art static baselines. Our code is publicly available at https://github.com/Panhaolin2001/PREP/.

Figures (4)

• Figure 1: Example of machine learning for compiler optimization task.
• Figure 2: The overall architecture of our program representation learning framework. It transforms LLVM IR into a Behavioral Spectrum, encodes it using Product Quantization, and learns its underlying grammar via a pre-trained Transformer model.
• Figure 3: Performance comparison across two prediction tasks: Best Pass Prediction, evaluated in terms of accuracy (higher values indicate better performance), and -Oz Benefit Prediction, evaluated in terms of Mean Absolute Error (lower values indicate better performance).
• Figure 4: t-SNE visualization of embedding spaces for the test set. Each point is a program, colored by its true best optimization pass.