CompilerGym: Robust, Performant Compiler Optimization Environments for AI Research (2024)

Frontend

The CompilerGym frontend is a Python library that exposescompiler optimization tasks using the OpenAI Gym[7] environmentinterface. Figure1 shows the interaction loop for the Gymenvironments. This allows researchers to interact with important compileroptimization problems in a familiar language and vocabulary with which many arecomfortable. The frontend is described in SectionIII.

Backend

CompilerGym uses a client-server architecture, shown inFigure2. This design provides separation of concernsas systems developers can easily add support for new compiler problems byimplementing a simple Compilation Session interface that comprises only fourmethods. The backend is described in SectionIV.

III-A OpenAI Gym Environments

We formulate compiler optimization tasks as Markov Decision Processes (MDPs) andexpose them as environments using the popular OpenAI Gym[7] interface. AGym environment comprises five ingredients:

1) An Action Space defines the set of possible actions that can be takenfrom a given MDP state. In CompilerGym, action spaces can be composed of discretechoices (e.g. selecting an optimization pass from a finite set), continuouschoices (e.g. selecting a function inlining threshold), or any combination ofthe two. The action space can change between states, such as in the case whereone optimization precludes another.

2) An Observation Space from which observations of the MDP state aredrawn. CompilerGym environments support multiple observation types such as numericfeature vectors generated by compiler analyses, control flow graphs, and stringsof compiler IR. Each environment exposes multiple observation spaces that can beselected from or composed.

3) A Reward Space defines the range of values generated by the rewardfunction, used to provide feedback on the quality of a chosen action, eitherpositive or negative. In CompilerGym, reward spaces can be nondeterministic (e.g.change in program runtime), platform specific (e.g. change in the size of acompiled binary), or entirely deterministic.

4) A Step operator applies an action at the current state and respondswith a new observation, a reward, and a signal that indicates whether the MDPhas reached a terminal state. Not all compiler optimization problems haveterminal states.

5) A Reset operator resets the environment to an initial state andreturns an initial observation.

Listing1 demonstrates how the core CompilerGym API is used.A make() function instantiates a subclass of the gym.Env environment that represents a particular compiler optimization task.The Gym interface is self describing: the action space and observation spacesare described by action_space and observation_space attributes, respectively. This enables CompilerGym environments tobe integrated directly with techniques that are compatible with other Gymenvironments. Listing2 shows one such integration.

In interacting with an environment the user’s goal is to select the sequence ofactions that maximizes the cumulative reward. Although Gym is designed primarilyfor reinforcement learning research, it makes no assumptions about the structureof user code and therefore can be used with a wide range of approaches. For asingle environment, the best sequence of actions may be found through search. Togeneralize a solution that works for unseen environments, a policy islearned to map from observation to optimal actions, or a $Q$-function islearned to give expected cumulative rewards for state-action pairs.

III-B API Extensions for Compiler Optimization

The advantage of the Gym interface is that it is simple and can be used across arange of domains. We supplement this interface with additional APIs that arespecific to compilers.

III-B6 Lightweight deep copy operator

CompilerGym environments provide a fork() operator that efficiently createsindependent deep copies of environment states. This can be used to optimizebacktracking or other techniques that require frequently evaluating a commonsubsequence of actions. For example, a greedy search can be implemented bycreating $n$ forks of an environment with an $n$-dimensional action space,running a single action in each fork, and selecting the one which produced thegreatest reward. Backtracking is especially expensive in compilers because mostactions have no “undo”.

III-C Customizing Environment Behavior

The Gym[7] library defines environment wrapper classes to mutate the MDPformulation of a wrapped environment. CompilerGym provides an additional suite ofenvironment wrappers for a broad range of compiler research uses. These includespecifying a subset of command line flags to use in an action space, iteratingover a suite of benchmarks, and defining derived observation spaces such asusing custom compiler analyses on compiler IR. These wrappers can be composed.Listing2 shows integration with the popularRLlib[9] library using two of these wrappers.

III-D Command Line Tools

We include a complete set of command line tools for CompilerGym, including scripts torun parallelized searches, replay and validate results from past runs, and aninteractive shell that includes inline documentation and tab completion,enabling users to interact with the compiler optimization environments withoutwriting any code.

III-E Web Service and CompilerGym Explorer

We designed a REST API to enable CompilerGym environments to be used over a network,and CompilerGym Explorer³³3Available at:https://compilergym.ai/explorer, a web frontend that makes it easy tonavigate compiler optimization spaces, implemented using React. CompilerGym Explorerpresents a visualization of the search tree, shown inFigure3, and asynchronously calls the RESTAPI to update the tree in real time as the user interacts with it.

A key feature of the tool is to visualize not only the current state, but alsohistorical trends of the rewards and observation metrics. This allows users toeasily pinpoint interesting actions in a large search tree and trigger newexplorations. We expect this to be valuable for feature engineering, debuggingthe behavior of agents, and as a general educational tool.

III-F State Transition Dataset

We designed a relational database schema to log the state transitions of CompilerGymenvironments for later offline analysis, shown inFigure4. A Steps tablerecords every unique action sequence for a particular benchmark and a hash ofthe environment state. An Observations table stores variousrepresentations of each unique state, indexed by state hash. A StateTransitions table encodes the unique transitions between states and therewards received for each.

We implemented a wrapper class for CompilerGym environments that asynchronouslypopulates the Steps and Observations tables of astate transition database upon every step of an environment. A post-processingscript de-duplicates and populates the StateTransitions table.

We are releasing a large instance of this database ($50+$GB) which contains over1M unique LLVM environment states, suitable for a range of offline supervisedand unsupervised learning tasks. We evaluate an example usage inSectionVII-F.

IV-A The CompilationSession Interface

CompilerGym is designed for seamless compiler integration. The integration centersaround implementing a state machine to interact with the compiler called aCompilationSession. A CompilationSession exposes actions and observations usinga simple schema and must implement two methods, apply_actionand get_observation, as shown inFigure5. We provide CompilationSessioninterfaces for Python and C++.Listing3 demonstrates an exampleimplementation.

IV-B Compiler Service Runtime

A common runtime maps implementations of the CompilationSession interface(Listing3) to the Gym API(Listing1). This runtime is shared by all compilerintegrations and is architected to be performant and scalable. The design isresilient to failures, crashes, infinite loops, and nondeterministic behavior inbackend compiler services. All compiler service operations have appropriatetimeouts, graceful error handling, or retry loops. Improvements to the runtimecan be made without changing compiler integration or user code.

A key design point of the CompilerGym runtime is that the service that provides thecompiler integration is isolated in a separate process to the user’s Pythoninterpreter. The Python interpreter invokes operations on the compiler servicethrough Remote Procedure Calls (RPCs). The benefits of this are fault toleranceand recovery in cases where the compiler crashes or terminates abruptly; supportfor compiling on a different system architecture than the host by running thecompiler service on a remote machine; and scalability as the expensive computework is offloaded, enabling many user threads to interact with separate compilerenvironments without contention on Python’s global interpreter lock.

V-A LLVM Phase Ordering

LLVM[11] is a modular compiler infrastructure used throughout academiaand industry. After parsing an input source program to a language-agnosticIntermediate Representation (IR), the LLVM optimizer applies a configurablepipeline of optimization passes to the IR. The selection and ordering ofcompiler optimizations – known as phase ordering – greatly impacts thequality of the final binary and has been the focus of muchresearch[1, 12].

We include a phase ordering environment in CompilerGym as an example of a challenging,high-dimensional optimization problem in which significant gains can beachieved.

V-B GCC Flag Tuning

We include an environment that exposes the optimization space defined byGCC’s command line flags. The environment works with any version of GCC from5 up to and including the current version at time of writing, 11.2. Theenvironment uses Docker images to enable hassle free install and consistencyacross machines. Alternatively, any local installation of the compiler can beused. This selection is made by simple string specifier of the path or dockerimage name. The only change that an RL agent needs to make to work with GCC instead of LLVM is to call env=gym.make("gcc-v0"), instead ofusing "llvm-v0".

V-C CUDA Loop Nest Code Generation

Manually tuning CUDA code requires sweeping over many parameters. Due to thesheer size of the tunable space, the problem of generating fast CUDA is wellsuited for automated techniques[24, 25]. As a flexible compilationenvironment, CompilerGym is well equipped to handle compilers for tuning GPUworkloads. We integrated loop_tool, a simple dense linear algebracompiler[26]. loop_tool takes a minimalist approach tolinear algebra representations by decomposing standard BLAS-like routines into aDAG of $n$-dimensional applications of arithmetic primitives. The DAG is thenannotated with three pieces of information about loop ordering: the order inwhich loops are emitted, the nesting structure of each loop, and the reuse ofloops by subsequent operations. This is lowered to a loop tree that can beannotated with which loop should be run in parallel. These four annotationsacross a slew of point-wise operations represent a large optimization space.

Binary Releases

Periodic versioned releases are made from a stable branch. We ship pre-compiledrelease binaries for macOS and Linux (Ubuntu 18.04, Fedora 28, Debian 10 ornewer equivalents) that can be installed as Python wheels.

Documentation

Our public facing documentation includes full API references for Python and C++,a getting started guide, FAQ, and code samples demonstrating integration withRLlib[9], implementations of exhaustive, random, and greedy searches,and Q-learning[27] and Actor Critic[28].

Testing

We have a comprehensive unittest suite with 85.8% branch coverage that is runon every code change across a test matrix of all supported operating systems andPython versions. Additionally, a suite of fuzz and stress tests are ran daily bycontinuous integration services to proactively identify issues.

Experimental Platforms

Results in this section are obtainedfrom shared compute servers equipped with Intel Xeon 8259CL CPUs, NVIDIA GP100GPUs, and flash storage.

VII-A Computational Efficiency

A key design goal of CompilerGym is to provide the best performance possible, enablingresearchers to train larger models, try more configurations, and get betterresults in less time. We evaluate the computational efficiency of CompilerGym’s LLVMphase ordering environment and compare to two prior works:Autophase[4] and OpenTuner[29].

We use code size rewards signals for all three platforms and the observationspace used in[4] for Autophase and CompilerGym; OpenTuner is a black boxsearch framework and so does not provide observation spaces. We measure thecomputational efficiencies of each environment by measuring the wall times ofoperations during 1M random trajectories. For CompilerGym, which uses a client-serverarchitecture, we also measure the initial server startup time.

TableII shows the results. CompilerGym achieves a much higherthroughput than Autophase while offering the same interface, observation space,and reward signal. This is enabled by CompilerGym’s client-server architecture. Afterinitially reading and parsing the bitcode file from disk, the CompilerGym serverincrementally applies an individual optimization pass at each step. In contrast,Autophase and OpenTuner must, at each step, read and parse the IR, apply theentire sequence of passes, and then serialize the result. OpenTuner, which wasdesigned for uses where the search time is dominated by compilation time, hasthe highest environment initialization cost, as it requires several diskoperations and the creation of a database. The CompilerGym server maintains a cache ofparsed unoptimized bitcodes that enables an amortized $\mathcal{O}(1)$ cost ofenvironment initialization.

The distribution of operation wall times depends on the action being performedand the program being optimized. Figure6 showsthe wide distribution of wall times within the benchmarks of a single dataset.

VII-B Computational Efficiency of Observation Spaces

This experiment evaluates the computational efficiency of the LLVM environmentobservation and reward spaces. We recorded 1M wall times of each using randomtrajectories.

TableIII summarizes the results. There is a$192\times$ range in observation space times, demonstrating a tradeoff betweenobservation space computational cost and fidelity; and $4727\times$ range inreward space times, motivating the development of fast approximate proxy rewardsand cost models[30, 31].

VII-C Autotuning LLVM Phase Ordering

We evaluate various autotuning techniques on the LLVM phase ordering task todemonstrate the ease and speed of CompilerGym. We use the following autotuningtechniques: Greedy search, which at each step evaluates all possible actions andselects the action which provides the greatest reward, terminating once nopositive reward can be achieved by any action; LaMCTS[32], anextension of Monte Carlo Tree Search[33] that partitions the searchspace on the fly to focus on important search regions[33];Nevergrad[34] and OpenTuner[29], two black boxoptimization frameworks that contain ensembles of techniques; and Random Search,which selects actions randomly until a configurable number of steps have elapsedwithout a positive reward.

We run single threaded versions of each autotuning technique on each benchmarkin the cBench[15] suite for one hour. Hyperparameters for alltechniques were tuned on a validation set of 50 Csmith[8] benchmarks.We evaluate each technique when optimizing for three different targets: codesize, binary size, and runtime. For runtime we use the median of threemeasurements to provide the reward signals during search, and the median of 30measurements for final reported values. Each experiment was repeated 10 times.

The standard interface exposed by CompilerGym makes it simple to integrate with thirdparty autotuning libraries or to develop new autotuning approaches.TableIV shows the number of lines of code required tointegrate each search technique, and the performance achieved.

Phase ordering is challenging because the optimization space is unbounded,high-dimensional, and contains sparse rewards. Nevertheless, autotuning – whenfurnished with a sufficiently generous search budget – outperforms the defaultcompiler heuristics by tailoring the configuration to each benchmark. We notethat the optimal configuration differs between all benchmarks and optimizationtargets.

VII-D Autotuning GCC Command Line Flags

For GCC we show a different aspect of the CompilerGym. For these experiments we explorethe GCC environment’s high-dimensional action space using a number of simplesearch techniques. These experiments are performed using GCC version 11.2.0 inDocker. That version of GCC has 502 optimization settings that can be selected.We evaluate three search techniques on the the CHstone[18] suite:

1) Random search. A random list of 502 integers from the allowable rangeis selected at each step.

2) Hill climbing search. At each step a small number of random changesare made to the current choices. If this improves the objective then the currentstate is accepted and future steps modify from there.

3) Genetic algorithm (GA). A population of 100 random choices ismaintained. We use the Python library geneticalgorithm[35] with its default parameters.

TableV shows the geometric mean of the object codesize objective across the benchmarks in CHstone[18], averaged over 3searches. Each search was allowed 1000 compilations.

VII-E Autotuning CUDA Loop Nests

The loop_tool environment provides an easily accessible interface tobeing exploring the landscape of GPU optimizations. Tuning a simple space bysearching threading and then sizing the inner loop reaches 73.5% of theoreticalpeak performance on our GP100 test hardware ($\sim$6e10 FLOPs or $\sim$750GB/sfor two 4-byte floating point reads and one write), and parity with PyTorchperformance on the same operation across a variety of problem sizes.Figure7 shows the results for different loopconfigurations, demonstrating potentially useful hardware and compilercharacteristics, notably a drop in performance near 100k threads.

VII-F Learning a Cost Model using the State Transition Dataset

Auxiliary tasks are commonly used in reinforcement learning to produce betterrepresentation learning and help with downstream tasks[36, 37]. This experimentdemonstrates using the State Transition Dataset(SectionIII-F) to learn a cost model ofinstruction count from a graph representation of program state.

We implemented a Gated Graph Neural Network[38] in PyTorch[39] and used Mean Squared Error loss to train a regressor to predictthe instruction count of a program after two rounds of message passing on theProGraML[14] graph representation built into CompilerGym. We trained on 80%of the State Transition Database by iterating over pairs of (graph, instructioncount) from the database. We used the remaining 20% of the database as avalidation set. Figure8 shows the convergence ofthe neural network. The network achieves a relative error of 0.025, while anaive mean prediction scores 1.393.

VII-G Reinforcement Learning for LLVM Phase Ordering

CompilerGym offers seamless integration with third party reinforcement learningframeworks. For example, by changing a single parameter value inListing2 we can use any of the 26 reinforcement learningalgorithms included in RLlib[9].

We use CompilerGym to replicate the LLVM phase ordering environment usedin[4]. Specifically: we fix episode lengths to 45 steps, use thesame observation space comprising a feature vector concatenated with a histogramof the agent’s previous actions, and we use a subset of the full actionspace⁵⁵5We use 42 actions (out of 124 total) rather than the 45 actionsused in[4] as three of the actions have been removed in recentversions of LLVM.. We note that each of these modifications to the base LLVMenvironment can be achieved using the wrapper classes built into CompilerGym(SectionIII-C). Our environment differsfrom[4] in that we use a code size reward signal rather thansimulated cycle counts.

We train three different reinforcement learning algorithms for 100k episodes andperiodically evaluate performance on a holdout validation set. We use Csmith togenerate both training and validation sets, as in[4].

TableVI shows the performance of the trained agents whenevaluated on a random 50 programs from each of the datasets available out of thebox in CompilerGym. 3 of the 4 algorithms achieve positive results when generalizingto programs within the same domain (Csmith), but only PPO[43] is able toachieve a positive score on two of the 13 other datasets. This highlights thechallenge of generalization across program domains.

VII-H Effect of Training Set on RL

The generalization of reinforcement learning agents across domains is thesubject of activeresearch[44, 45, 46].As demonstrated in the previous experiment, the performance of agents trained onone dataset can differ wildly on datasets from other domains. We evaluate theeffect of training set on generalization by training a PPO[43] agent ondifferent training sets and then evaluating their generalization performance ontest sets from different domains. All other experimental parameters are as perthe previous experiment.

TableVII shows the results. As can be seen, eachalgorithm performs best when generalizing to benchmarks from within the samedataset, suggesting the importance of training on benchmarks across a wide rangeof program domains.

VII-I Effect of Program Representation on Learning

Representation learning and feature engineering is an area of muchresearch[47, 48, 2].CompilerGym environments provide multiple state representations for each environment.We evaluate the performance of two different program representations, and theirperformance when concatenated with a histogram of the agent’s previous actions,as used in[4]. We use the same experimental setup as in the priorsections.

The results are shown in Figure9. In bothcases stronger performance is achieved when coupling the program representationwith a histogram of the agent’s previous actions. The Autophase representationencodes more attributes of the structure of programs than InstCount and achievesgreater performance. We believe that representation learning is one of the mostexciting areas for future research, and CompilerGym provides the supportinginfrastructure for this research.