Supercompiler Code Optimization with Zero-Shot Reinforcement Learning
Jialong Wu Chaoyi Deng Jianmin Wang Mingsheng Long




1]Jialong Wu
These authors contributed equally to this work.

1]Chaoyi Deng
These authors contributed equally to this work.

1]Jianmin Wang

[1]Mingsheng Longmingsheng@tsinghua.edu.cn

[1] School of Software, BNRist, Tsinghua University

Effective code optimization in compilers plays a central role in computer and software engineering. While compilers can be made to automatically search the optimization space without the need for user interventions, this is not a standard practice since the search is slow and cumbersome. Here we present CodeZero, an artificial intelligence agent trained extensively on large data to produce effective optimization strategies instantly for each program in a single trial of the agent. To overcome the huge range of possible test programs, we prepare a large dataset of training programs that emphasize quality, naturalness, and diversity. To tackle the vast space of possible optimizations, we adapt deep reinforcement learning to train the agent in a sample-efficient manner through interacting with a world model of the compiler environment. Evaluation on both benchmark suites and production-level code optimization problems demonstrates our agent's supercompiler performances and zero-shot generalization abilities, outperforming built-in optimization options designed by compiler experts. Our methodology kindles the great potential of artificial intelligence for engineering and paves the way for scaling machine learning techniques in the realm of code optimization.

Supercompiler Code Optimization with Zero-Shot Reinforcement Learning
    *
    April 29, 2024
=====================================================================






§ INTRODUCTION


Entering the post-Moore’s law era, code optimization is crucial for computer and software engineering, which plays an important role in realizing the full potential of slow-growing hardware. Developers typically rely on a compiler's ability to transform input programs into semantically equivalent but more efficient versions, improving metrics like execution time, code size, and power consumption. For example, standard optimization options , , and  aim to reduce execution time, while the  and  options are crafted to reduce code size. Still, it is not common for users to explore beyond these conventional compiler options. Given the vast diversity of programs and platforms, coupled with the increasing number of optimization passes integrated into compiler frameworks, these off-the-shelf optimization strategies predefined heuristically by compiler experts may struggle to guarantee near-optimal performance across ever-changing scenarios <cit.>.

Automatic code optimization is therefore crucial in compilers. Autotuning <cit.> improves code by systematically searching the optimization space through iterative executing and profiling optimization strategies. This search technique can yield remarkable performance gains but must be rerun for each new program with thousands of compilations, which is too time-consuming to be practical for all but a few specialized use cases. Meanwhile, machine learning techniques hold the capabilities to generalize the optimization strategy of one program to other similar ones, thereby facilitating faster code optimization. A direct method is to use supervised learning to predict good optimization strategies of input programs <cit.>, which is impractical due to prohibitive computations to construct labeled training data by search. Another more promising routine, reinforcement learning, that successfully discovered faster sorting algorithm (AlphaDev <cit.>) and matrix multiplication algorithm (AlphaTensor <cit.>), can explore the optimization space from feedback on optimization metrics without requiring optimal labeled data. For both techniques, broad generalization across different programs, even out of the training samples, arises as a major bottleneck. The community has noted that in a range of machine learning applications <cit.>, training high-capacity models on large-scale datasets has yielded unprecedented performances. For example, large language models like GPT-4 have demonstrated impressive zero-shot generalization abilities <cit.>. However, in the area of code optimization, the current practice is to learn optimization heuristics in a per-program manner <cit.> or from a small training set with hundreds of programs <cit.>, lagging far behind the new era of solving challenging problems by scaling up machine learning models.



In this study, we focus on the LLVM <cit.> phase ordering problem, a longstanding challenge for compiler research, and propose CodeZero, a reinforcement-learned code optimization agent, capable of generating a sequence of optimization passes tailored to a particular input program. The problem is formulated such that a code optimization agent, upon encountering a program, selects a series of optimization passes, and receives feedback based on the outcomes of applying these transformations, from the compiler environment. To tackle the huge range of possible test programs, we aim to enhance the generalization ability of trained agents by assembling a large-scale training dataset of natural programs. Several previous works rely on randomly generated programs due to data scarcity but the significant distribution shift from random-generated and human-written programs can even hurt the generalization of learned agents when tested on real-world scenarios <cit.>. 
Our training data not only includes real-world programs sourced from GitHub <cit.> but also incorporates complex algorithmic solutions of competitive programming <cit.> and diverse programs generated by large language models (LLM) <cit.>.
To tackle the vast space of possible optimizations (∼ 10^73 sequences), we employ a state-of-the-art model-based reinforcement learning method <cit.> to train the code optimization agent sample-efficiently. This method not only learns a predictive world model of compilers to reduce the amount of real compiler executions but also benefits generalization by learning representations that better capture the structure of the compiler state transitions <cit.>. After training on massive programs via trial-and-error, the code optimization agent can generalize in a zero-shot way to previously unseen programs with superior optimization performance against off-the-shelf compiler heuristics.

We demonstrate the effectiveness of our trained CodeZero agent on a range of domains <cit.> from benchmark suites encompassing fundamental algorithms that are ubiquitously employed in everyday applications, to production-level open-source programs, including object files from C++ TensorFlow <cit.> and OpenCV <cit.> library. On six test datasets, our agent can produce optimization sequences in a single trial, yielding more efficient code size reduction compared to the  flag. Detailed analysis underscores the effectiveness of both the newly introduced program datasets and the model-based reinforcement learning technique within the realm of code optimization. We posit that the agents developed through our approach could be integrated into the existing toolkits of optimization strategy in compilers alongside other manually designed heuristics, such as  or .



§ CODEZERO AGENT




 §.§ Code Optimization as Decision Making



As illustrated in Figure <ref>, compilers consist of three main components: the front end which translates the source code into an intermediate representation (IR), the middle end, and the back end which converts IR to the binary code. The middle end is responsible for language- and platform-agnostic optimizations over the IR, implemented as passes to either collect information about the program or apply a transform on it, like function inlining, loop unrolling, and dead code elimination. For example, the LLVM- tool has more than 100 optimization passes available. A specific order of applying these passes forms an optimization sequence for an input program. This is critical as the right selection and ordering of passes can significantly boost the program’s performance. Despite that compiler developers have provided standard optimization sequences at various levels, e.g. , , , these preset sequences may not always yield optimal results, especially for emerging programs written in new frameworks such as TensorFlow.  Particularly due to the increasing number of optimization passes, it is an open challenge to determine the most effective sequence for each program.

This problem, known as phase ordering, can be naturally formulated as a partially observable Markov decision process (POMDP) <cit.>. In this formulation, a code optimization agent determines the optimization sequence of an input program through a series of interactions with the compiler environment, guided by its policy π. 
The process starts at the initial state s_0, representing the IR of the program to be optimized, which is randomly sampled from all IRs of interest. It is critical that aiming to capture the most important characteristics of the target optimization, the agent only receives partially observable information o_0 of the state. The observation space can vary widely, ranging from manually designed features (e.g. the number of basic blocks) <cit.> to more complex tree-based or graph-based program representations <cit.>, and even raw text strings of IR <cit.>. At each time step t=0,1,2,…, the agent takes an action a_t based on its policy, corresponding to selecting an optimization pass. Following this action, the agent receives an immediate reward r_t+1, reflecting changes in the optimization metrics, and an observation o_t+1 of the next state s_t+1 = p(s_t, a_t) which represents the IR transformed by the compiler using the selected pass. This process can be terminated either when the agent finds no positive gains can be achieved or a maximum number of steps is reached. The goal of the agent is to learn an optimization policy π(a_t | o_≤ t) that effectively maximizes the cumulative rewards, thereby optimizing the performance metrics of the final IR.



 §.§ Large-Scale Data Preparation


To ensure that our CodeZero agent can effectively generalize to unseen situations, a concept known as zero-shot generalization, we have identified three critical factors in preparing our training dataset. Firstly, the dataset must reflect naturalness. Training data should be within the distribution of human-written programs, otherwise, overfitting to programs that deviate significantly from this, such as those generated by tools like Csmith <cit.> and llvm-stress <cit.>, could provide no benefits or even hurt the generalization to real-world scenarios. Secondly, diversity in the dataset is essential. It should encompass a wide range of human-written program styles and structures, ensuring globally comprehensive coverage of possible scenarios. Lastly, the pursuit of high-quality training data is imperative. Existing large-scale program collections <cit.> have been proposed to serve as training data <cit.> but are proven unfruitful in our preliminary experiments. Instead, we focus on data that features moderate lengths of IR, complex algorithmic logic, and potential room for optimizations. This allows the agent to flexibly explore and understand the local transformation structure within the IR space. Collectively, these three properties – naturalness, diversity, and high quality – contribute to aligning the visited IR distribution during training with that will be encountered in real-world applications.

We construct our training dataset containing hundreds of thousands of programs by combining three distinct single-source program datasets, CodeContests <cit.>, FormAI <cit.>, and AnghaBench <cit.>. CodeContests comprises human-written solutions to competitive programming problems with complex and optimizable algorithmic logic. FormAI is a large collection of AI-generated C programs with various functionality types and coding styles, aiming to enrich the dataset's diversity. AnghaBench is a collection of real-world C programs extracted from GitHub. As shown by visualization in Figure <ref>, our training data has a broad coverage of evaluation programs in a variety of domains. Future expansions of the dataset can be conducted following the aforementioned principles.



 §.§ Agent Training with Model-based RL




While existing work on reinforcement learning for code optimization prevalently focuses on model-free RL methods, model-based RL can offer advantages in terms of both sample efficiency <cit.> and generalization <cit.>. Executing and profiling extensive optimization sequences, especially for the runtime metric, can be time-consuming. This is further compounded when constructing complex observations, such as control-data flow graphs <cit.>. Model-based RL addresses these challenges by learning a world model to approximate state transitions and reward signals of the environment. This allows the agent to learn its policy by simulating trajectories based on model predictions, rather than relying solely on trial-and-error interactions in the real compiler environment. This approach thus improves sample efficiency. Moreover, as the policy can share the representation with the world model, model learning can act as an auxiliary task and thus aid in learning representations that better capture the structure of the environment and manifest in better generalization of the policy <cit.>.

We train the CodeZero agent by adapting an advanced model-based RL method, Dreamer <cit.>, as depicted in Figure <ref>. This approach involves learning a predictive world model (p̂_θ, r̂_θ ) of the compiler environment by approximating the underlying transition dynamics `=12 p(o_t+1 | o_≤ t, a_≤ t) and reward function r(o_≤ t, a_≤ t). Through imaginary rollouts using this world model over a horizon H, the policy `=12 π_ψ(a_t | o_≤ t) can be learned using the REINFORCE policy gradients <cit.> with an entropy regularizer H. The training objective is formulated as

    ℒ(ψ) ≐𝔼_p̂_θ, r̂_θ, π_ψ[∑_t=1^H - (V_t -v_ξ(o_≤ t))logπ_ψ(a_t| o_≤ t) - ηH[π_ψ(o_≤ t)]],

where V_t is the imagined return estimation and the value prediction v_ξ serves as a baseline for variance reduction <cit.>. The value prediction is learned separately by regressing the cumulative rewards that the agent should expect to receive from the current state. All components are realized as deep neural networks.

After extensive training on a wide range of programs, our agent demonstrates the capability to generalize “zero-shot” — that is, to be effective without further training — to new, unseen programs across various domains.



§ RESULTS




 §.§ Evaluation Benchmarks


Our experiments focus on code size reduction, which benefits applications targeting low-resource hardware such as embedded systems. This focus is driven by the practical advantages of code size as a metric: it is both cost-effective and convenient to construct extensive compilable training and test datasets and to evaluate the optimization performance for code size.

We evaluate our method on benchmarks from the CompilerGym platform <cit.>: benchmark suites including cBench <cit.>, CHStone <cit.>, MiBench <cit.>, and NASA Parallel Benchmarks (NPB) <cit.>, in addition to kernels from open source projects including BLAS <cit.>, Linux, OpenCV <cit.>, and TensorFlow <cit.>. We disregard benchmarks generated by synthetic program generators <cit.> since they do not align with real-world scenarios. 










 §.§ Code Optimization Results


Figure <ref> reports CodeZero's performance in reducing code size, measured in terms of IR instruction count, compared to the expert-designed  flag in LLVM. Remarkably, CodeZero, without any specific training on in-domain programs, outperforms  across all but two test benchmark datasets in a single trial. The marginal performance variations between different methods on the BLAS and Linux datasets suggest that they are already highly optimized. CodeZero also demonstrates the ability to match or even slightly surpass a competitive autotuning method, random search, within a similar time budget (in the order of seconds), while the latter blindly aggregates hundreds of trials. This notable performance, coupled with the significant advantage over a single random trial, underscores the effectiveness of our agent's policy learning. Moreover, CodeZero's zero-shot generalization capability either matches or exceeds the performance of in-domain counterparts, which are trained using the train subset of each benchmark dataset. This is particularly significant in cases like the NPB dataset, which contains only 22 training programs. In such scenarios, where data sparsity presents a challenge for in-domain agents, the zero-shot CodeZero notably achieves an extra 3% reduction in code size. Notably, CodeZero's pre-trained policy can generalize to IRs compiled from novel programming languages beyond C and C++, as evidenced in its performances on the BLAS and NPB datasets of Fortran programs. Further, in an AI-generated benchmark dataset of Objective-C, another language supported by LLVM-Clang apart from C/C++, CodeZero successfully improves upon LLVM , achieving an average code size reduction of 1.027×, and reaching up to 2.87× in certain instances.






 §.§ Program Case Study


We analyze the internal behavior of our CodeZero agent to optimize IR code via recording the sequence of passes chosen by the policy on individual programs. Table <ref> presents the agent's top performance outcomes on various benchmark datasets. We observe that the agent indeed produces a specialized optimization strategy tailored for each program. Additionally, these results also highlight certain passes, such as -sroa (scalar replacement of aggregates), -gvn (global value numbering), and -simplifycfg, as particularly effective in code size optimization.

In Figure <ref>, we display a predicted optimization trajectory for an unseen program from cBench, as forecasted by our learned compiler world model. The model successfully forecasts numeric features of future IR, including the counts of branches and blocks, alongside future rewards that signify optimization outcomes. This instance exemplifies the capability of our learned compiler world model to serve as a viable alternative for a real compiler environment in training code optimization agents.





 §.§ Effect of Training Dataset




In our quest to understand the impact of training datasets on generalization, we explored how different datasets from various domains and sizes influence performance. As shown in Figure <ref>a, we discovered that the CodeContests dataset stands out as the most effective for generalizing to our test benchmarks. Although comparable outcomes were observed with other datasets like FormAI and AnghaBench, agents trained on these datasets notably lag in performance when tested on OpenCV and TensorFlow benchmarks. Moreover, combining these datasets does not yield performance enhancement. Consequently, we present in Figure <ref> the results based on the agent trained exclusively with the CodeContests dataset. These results underscore the significance of high-quality training data, a principle that is increasingly recognized as vital in other scalable machine learning applications, such as large language models. In Figure <ref>b, we compare the agent trained on the full CodeContests dataset versus the one trained on a 100-program subset sampled from the same dataset. This subset is comparable in size to the datasets used in previous research <cit.> that applies reinforcement learning to code optimization. The suboptimal performance on this smaller dataset highlights its inadequacy for training a strong zero-shot generalizable agent and emphasizes the significance of a larger and more diverse training set prepared by this study.




 §.§ Sample Efficiency and Zero-Shot Generalization


We further evaluate the sample efficiency and zero-shot generalization abilities of CodeZero based on the world model algorithm, Dreamer, against a range of model-free counterparts. These include notable algorithms such as DQN <cit.>, A2C <cit.>, APEX <cit.>, IMPALA <cit.>, and PPO <cit.> (see Appendix <ref> for more details). Figure <ref>c shows that while PPO is the most competitive among the model-free baselines, Dreamer outperforms it by learning an order of magnitude faster, requiring fewer interactions with the compiler environment to achieve excellent performance. Additionally, as Figure <ref>d indicates, Dreamer excels in generalizing to unseen test benchmarks, outstripping its counterparts. These findings verify our hypothesis that a world model-based agent holds a stronger capability to capture the environment's structure and can effectively enhance zero-shot generalization.



§ DISCUSSION


We tackled the major challenge of zero-shot generalization that arises when applying artificial intelligence techniques to improve code optimization. We introduce the CodeZero agent that leverages the strengths of a diverse, high-quality training program dataset, as well as the sample efficiency and generalization capabilities of world model-based reinforcement learning. Our results have shown that CodeZero, with its zero-shot generalization ability, attains supercompiler code optimization performance, in the challenging phase ordering problem for code size reduction. Our study provides an AI-centric framework and serves as a meaningful step towards scaling machine learning techniques for code optimization. There is a substantial scope for further exploration, including expansion of the training dataset, scaling up the compiler world models, optimizing multiple objectives like execution time, and enriching feature and action spaces with deeper expert knowledge or large language models.



§ METHODS


In this section, we provide the mathematical and computational details of our approach.



 §.§ LLVM Phase Ordering POMDP


We formulate the phase ordering problem of LLVM as a partially observable Markov decision process (POMDP) ℳ = (𝒮, 𝒜, r, p, μ, 𝒪, ϕ). The state space 𝒮 covers all possible Intermediate Representations (IRs), the action space 𝒜 comprises individual compiler optimization passes, and the reward function r is defined by the metric being optimized. The transition dynamics p: 𝒮×𝒜↦𝒮 represents the outcome of applied IR transformations. The initial state distribution μ∈Δ(𝒮) captures all IRs of interest, which can be approximated via uniform sampling from the training dataset. The observation function ϕ: 𝒮↦𝒪 maps the underlying IR into the observation space. Complex observation spaces offer comprehensive program information, while expert-designed features embed problem-specific knowledge, potentially enhancing the optimization policy's generalization by eliminating irrelevant details. 
The code optimization agent with a policy π interacts with the compiler environment according to the protocol described in Section <ref>. We fix the horizon of interactions to 45 steps.

Following Autophase <cit.> and CompilerGym <cit.>, the action space consists of 42 optimization passes out of over a hundred in LLVM, which effectively prunes the vast optimization space while not sacrificing the performance of the learned policy. The observation space is a concatenation of two numeric features: the 56-dimension Autophase feature summarizing the statistics of the IR, such as the number of specific basic blocks, branches, and instructions, and a 42-dimension histogram of the agent's previous actions. The reward function is defined as the normalized change of the optimization metric C(s):

    r_t+1 = C(s_t) - C(s_t+1)/C(s_0) -C(s_b),

where a lower C indicates better performance and C(s_b) stands for the performance of a baseline policy, such as built-in  or . A total reward greater than 1 means that the optimization sequence performs better than the baseline policy.



 §.§ Training Dataset


We construct our training datasets, by considering three datasets from distinct domains. The CodeContests dataset <cit.> consists of over 13,000 competitive problems, each on average having hundreds of solutions in three languages. We subsample up to ten C++ solutions for each training problem, resulting in 110,240 programs, as part of our training data, and sample one solution for each of 100 test problems from CodeContests as our validation data. We also include the full Form-AI dataset <cit.> with 112,000 AI-generated programs and the collection of the largest 9,998 single-function programs and 15,264 multiple-function programs from the AnghaBench <cit.> into our training data.



 §.§ Model-based RL Method


We utilize a deep model-based RL method, DreamerV3 <cit.>, to model the compiler environment upon which it learns an effective policy that maps observations to actions with parameterized deep neural networks.



  
World Model

The world model simulating the compiler environment is formulated as a latent dynamics model with the following four components: 

    Representation model:      z_t∼ q_θ(z_t| z_t-1,a_t-1, o_t) 
       Transition model:      ẑ_t∼ p_θ(ẑ_t| z_t-1, a_t-1) 
       Image decoder:      ô_t∼ p_θ(ô_t| z_t)
       Reward decoder:      r̂_t ∼ p_θ(r̂_t  |  z_t)

The representation model estimates the latent state z_t based on the previous state z_t-1, the previous action a_t-1 and the current observation o_t, while the transition model predicts it directly from z_t-1 and a_t-1. The overall models are jointly learned by minimizing the negative evidence lower bound (ELBO) <cit.>:

    ℒ_model(θ) ≐ 𝔼_q_θ(z_1:T| a_1:T, o_1:T)   [ ∑_t=1^T(
        -ln p_θ(o_t| z_t) -ln p_θ(r_t| z_t) 
           + KL[q_θ(z_t| z_t-1,a_t-1, o_t)  ‖ p_θ(ẑ_t| z_t-1, a_t-1) ])].


In practice, we employ a reward smoothing technique <cit.> to mitigate the sparsity and the long tail distribution of rewards during an episode. This is achieved through the following equation:

    r_t^'←α r_t-1^' + (1-α) r_t,   t=1,2,…

with α∈ [0,1). Consequently, we train a reward decoder p_θ(r̂_t^'|  z_t) to predict the smoothed rewards.



  
Actor-Critic Learning The actor and critic neural networks are parameterized on top of the latent representations:

    Actor: â_t ∼π_ψ(â_t |ẑ_t)    Critic:  v_ξ(ẑ_t) ≈𝔼_p_θ, π_ψ[∑_τ≥ tγ^τ-tr̂_τ].

The actor and critic are jointly trained on the same imagined trajectories {ẑ_τ, â_τ, r̂_τ} with horizon H, generated by the interactions between the transition model and reward model in Eq. (<ref>) and the actor: starting at the latent state ẑ_t = z_t, at each step τ=t, t+1, t+2, …, the policy takes an action â_τ∼π_ψ(â_τ|ẑ_τ), and transits to the next latent state ẑ_τ+1∼ p_θ(ẑ_τ+1| z_τ, a_τ) with a reward r̂_τ+1∼ p_θ(r̂_τ+1|ẑ_τ+1). The critic is trained to predict the λ-return <cit.> through a discrete regression loss <cit.>: 

    ℒ_critic(ξ)≐𝔼_p_θ, π_ψ[∑^t+H_τ=t - log v_ξ(V_τ^λ|ẑ_τ) ],
        
       V_τ^λ≐r̂_τ + γ
          (1 - λ)v_ξ(ẑ_τ+1)+λ V_τ+1^λ   if τ<t+H 
    
          v_ξ(ẑ_τ + 1)    if τ=t+H.

The actor, meanwhile, is trained to maximize the imagined return through the REINFORCE policy gradient <cit.>:

    ℒ_actor(ψ)≐𝔼_p_θ, π_ψ[∑^t+H_τ=t(-(V_τ^λ -v_ξ(ẑ_τ))logπ_ψ(â_τ|ẑ_τ) - η H[π_ψ (ẑ_τ) ] ) ],

where H[π_ψ (ẑ_τ) ] is an entropy regularization which encourages exploration, and η is a hyperparameter that adjusts the regularization strength.








§ RELATED WORK


In this section, we will elaborate a more extensive review of the development of machine learning methods for code optimization in compilers.

One of the key challenges for compilation is to determine which code transformations to apply, how to apply them (e.g., using suitable parameters), and in what order. This requires effectively searching and evaluating a massive number of possible options, known as iterative compilation <cit.> or autotuning <cit.>. However, this search-based approach only finds a good optimization for one specific program and does not generalize into a compiler heuristic. This limitation underscores the importance of integrating machine learning techniques.

Pioneering work has delved into supervised machine learning, adopting two main approaches. The first one learns a predictive model that can directly predict the best option. It involves iteratively compiling training programs to identify the most effective compilation strategy for each, which then serves as the labels of training data. An early example <cit.> used a neural network for branch prediction, and one more well-known work is MilepostGCC <cit.>, a practical attempt to integrate machine learning into a production compiler, GCC. It employs models trained on a large dataset of programs distributed over the Internet. The second approach aims to learn a cost or performance function capable of estimating the quality of various compiler options, which enables evaluation of a range of possible options without the need to compile and profile each one <cit.>.

Recent advancements have seen reinforcement learning (RL) techniques making strides in compiler optimization, circumventing the requirement for collecting optimal labeled data <cit.>. This technique has been applied to optimize individual compilation heuristics, such as inlining <cit.>, loop transformation <cit.>, and graph partitioning <cit.>. Several works relevant to us have explored the full optimization pipeline, i.e. the LLVM phase ordering problem, including AutoPhase <cit.>, CORL <cit.>, and POSET-RL <cit.>. These methods predominantly utilize model-free RL algorithms while our work pioneers the use of an advanced model-based RL approach to reduce real compiler interactions. Model-based RL <cit.> learns a simulation model of the compiler environment, which is related in spirit to the aforementioned approach of supervised performance models. However, it goes a step further by also learning a policy capable of directly determining the best optimization option, thereby eliminating the necessity for a guided search process. 

Both machine learning techniques require crafting high-quality features that capture the important characteristics of programs, a process known as feature engineering. The most prevalent feature vectors are based on the frequencies of various types of instructions within the programs <cit.>, designed by expert intuitions. Numerous studies have aimed to reduce the cost of feature design. Following the success of word2vec embeddings in natural language processing <cit.>, methods like code2vec <cit.>, inst2vec <cit.>, and IR2vec <cit.> represent programs as distributed vectors that capture syntactic and semantic information from the abstract syntax tree (AST) or intermediate representation (IR). The surge in deep learning has enabled feeding raw information such as AST <cit.>, control-data flow graphs (CFG) <cit.>, and code token sequences <cit.> into powerful deep neural networks, capable of learning useful representations end-to-end.

Platforms that expose the compiler as a playground for AI experiments have significantly reduced the entry barriers to intelligent compiler research. OpenTuner <cit.> and YaCoS <cit.> serve as autotuning frameworks with a range of compiler optimization techniques. Our experiments utilize CompilerGym <cit.>, which offers user-friendly interfaces for researchers to interact with compilers in a reinforcement learning manner. We are optimistic that the future release of our trained code optimization agents, in conjunction with these platforms, can have a democratizing effect on applying AI techniques to compiler optimizations.



§ IMPLEMENTATION DETAILS




 §.§ Compiler Environment


Our experiments are conducted on the CompilerGym platform <cit.>, version 0.2.5, with LLVM-10.0.0 integration.



 §.§ Features and Actions




We extract program features following the approach of Autophase <cit.>. As described in Section <ref>, we use a 56-dimension Autophase feature vector concatenated with a 42-dimension action histogram vector as the observation. The Autophase feature contains various statistics of the IR code, with each dimension thoroughly explained in Table <ref>. The action histogram vector contains the counts of actions already taken by the agent within the current episode. Both vectors are normalized to ensure that their values fall within a reasonable range. Specifically, each element of the Autophase vector is divided by the total instruction count of the program, whereas the action histogram vector is normalized by the total action count per episode, set as 45.

Our action space is derived from Autophase <cit.>, comprising 45 LLVM optimization passes. However, CompilerGym excludes 3 actions due to updates in the latest LLVM version. Thus, we utilize a total of 42 actions. Table <ref> presents the names of the LLVM optimization passes corresponding to these 42 actions.






 §.§ Hyperparameters


The hyperparameters for our DreamerV3 agent implementation are outlined in Table <ref>. For hyperparameters not specified, we use the same value as the original DreamerV3 <cit.>.





 §.§ Model-Free Baselines



We use RLlib <cit.> to train and test model-free reinforcement learning algorithms including PPO <cit.>, A2C <cit.>, IMPALA <cit.>, APEX <cit.>, and DQN <cit.>. All the experiments are conducted with 10 CPUs and an RTX-3090 GPU over a minimum duration of 10 hours. The training process uses 5 rollout workers to interact with the environment and use 4 evaluation workers to evaluate the checkpoint on the validation set. We use default hyperparameters of algorithms in RLlib following the CompilerGym platform <cit.>, except that we have carefully tuned the hyperparameters for our strongest baseline, PPO, as listed in Table <ref>.







§ BENCHMARKS




  
CompilerGym Benchmarks

In our study, we select eight benchmarks for zero-shot test and in-domain training: BLAS, cBench, CHStone, Linux, MiBench, NPB, OpenCV, and TensorFlow. These benchmarks are part of the built-in datasets provided by CompilerGym version 0.2.5. For benchmarks with a total number of programs more than 100, we use the first 50 programs as the test set, the following 50 programs as the validation set, and all of the rest as the training set. These training and validation sets are only used for in-domain training. The datasets comprising fewer than 100 programs are not applicable for in-domain training; instead, all their programs are allocated to the test set. The number of programs in each dataset after division is detailed in Table <ref>.






  
AI-Generated Benchmarks
To further test the generalization ability of our CodeZero agent on different programming languages, we borrow the method from FormAI <cit.> and generate a dataset containing 50 unique Objective-C programs using GPT-3.5. We use the same prompt as FormAI, except that we add an instruction to ask GPT to generate programs that can be directly compiled under Clang version 10.0.0 and do not use ARC (Automatic Reference Counting), to improve the compilation pass rate of generated programs. We compile the generated programs using Clang without including any third-party libraries, and all programs that cannot pass compilation are discarded. 



§ EXTENDED EXPERIMENTAL RESULTS




 §.§ Learning Curves


The zero-shot test performance of our CodeZero agents during training is shown in Figure <ref>. Note that we validate and test the agent every 10000 environment steps and report the test performance from the checkpoint that achieved the best validation results for comparison among various methods. 





 §.§ Quantitative Results


Quantitative results corresponding to Figure <ref> in the main text are provided in Table <ref>.





 §.§ Data Distribution Visualization


In Figure <ref>, we visualize the distribution of our training and test datasets. 
To accurately represent the dynamic behavior of programs, we randomly select 1000 action sequences, each with a length of 45, from our action space. These sequences are subsequently executed on each program, with the resulting Autophase features concatenated to form a feature vector with dimensions of 1000× 45× 56 for every program. These comprehensive feature vectors are finally processed using t-SNE <cit.> for dimensionality reduction and visualization.

Figure <ref> illustrates that our training data (denoted as triangles) has a broad coverage of test programs (denoted as circles). This contrasts with the Csmith dataset (denoted as crosses) employed in CompilerGym experiments <cit.>, which shows a significant deviation from real-world applications. Nonetheless, our visualization can still not perfectly capture the transferability across datasets. For instance, empirical evidence suggests that CodeContests are the most effective in generalizing to OpenCV and TensorFlow, while the visual analysis does not directly imply this.





