Meta Large Language Model Compiler: Foundation Models of Compiler Optimization

LLM Compiler is trained predominantly on compiler intermediate representations and assembly code generated by LLVM (Lattner & Adve, 2004) version 17.0.6. These are derived from the same dataset of publicly available code used to train Code Llama. We summarize this dataset in Table 2. As in Code Llama, we also source a small proportion of training batches from natural language datasets.

Given unoptimized LLVM-IR (as emitted by the clang frontend), a list of optimization passes, and a starting code size, generate the resulting code after those optimizations have been applied and the resulting code size.

There are two flavors of this task: in the first the model is expected to output compiler IR, in the second the model is expected to output assembly code. The input IR, optimization passes, and code size are the same for both flavors. The prompt dictates the required output format. Examples of each prompt are provided in Appendices Listings 10 and 11.

We use two metrics for code size: the number of IR instructions, and binary size. Binary size is computed by summing the size of the .TEXT and .DATA sections of the IR or assembly after lowering to an object file; we exclude .BSS section from our binary size metric since it does not affect on-disk size.

In this work we target LLVM 17.0.6 and use the New Pass Manager (PM, 2021) which classifies passes for different levels such as module, function, loop, etc. as well as transformation and analysis passes. Transformation passes change given input IR while analysis passes generate information that influence subsequent transformations.

Of the 346 possible pass arguments for opt, we select 167 to use. This includes each of the default optimization pipelines (e.g. module(default<Oz>)), individual optimization transform passes (e.g. module(constmerge)), but excludes non-optimization utility passes (e.g. module(dot-callgraph)) and transformations passes that are not semantics preserving (e.g. module(internalize)). We exclude analysis passes since they have no side effects and we rely on the pass manager to inject dependent analysis passes as needed. For passes that accept parameter arguments we use the default values (e.g. module(licm<allowspeculation>)). Table 9 contains a list of all passes used. We used LLVM’s opt tool to apply pass lists and clang to lower the resulting IR to object file. Listing 3 shows the commands used.

We generated the compiler emulation dataset by applying random lists of between 1 and 50 optimization passes to unoptimized programs summarized in Table 2. The length of each pass list was selected uniformly at random. Pass lists were generated by uniformly sampling from the set of 167 passes described above. Pass lists which resulted in compiler crashes or timed out after 120 seconds were excluded.

We present the LLM Compiler FTD models with an unoptimized LLVM-IR (as emitted by the clang frontend) and ask it to produce a list of opt flags that should be applied to it, the binary size before and after these optimizations are applied, and the output code. If no improvement can be made over the input code, a short output message is generated that contains only the unoptimized binary size. Listings 12 and 13 provide the prompt and output templates for this task.

We used the same constrained set of optimization passes as in the compiler emulation task, and compute binary size in the same manner.

Figure 4 illustrates the process used to generate training data (described below) and how the model is used for inference. Only the generated pass list is needed at evaluation time. We extract the pass list from the model output and run opt using the given arguments. We can then evaluate the accuracy of the model predicted binary sizes and optimized output code, but those are auxiliary learning tasks not required for use.

LLVM’s optimizer is not free from bugs and running optimization passes in unexpected or untested orders may expose subtle correctness errors that undermine the utility of the model. To mitigate this risk we developed PassListEval, a tool to help in automatically identifying pass lists that break program semantics or cause compiler crashes. An overview of the tool is shown in Figure 5. PassListEval accepts as input a candidate pass list and evaluates it over a suite of 164 self-testing C++ programs, taken from HumanEval-X (Zheng et al., 2023). Each program contains a reference solution for a programming challenge, e.g. “Check if in given vector of numbers, are any two numbers closer to each other than given threshold”, and a suite of unit tests that validate correctness. We apply the candidate pass lists to the reference solution, and then link them against the test suites to produce a binary. When executed, the binary will crash if any of the tests fail. If any binary crashes, or if any of the compiler invocations fail, we reject the candidate pass list.

We trained LLM Compiler FTD models on a dataset of flag tuning examples derived from 4.5M of the unoptimized IRs used for pretraining. To generate the example optimal pass list for each program we ran an extensive iterative compilation process depicted in Figure 4 and outlined below:

1. 1.

   We used large-scale random search to generate an initial candidate best pass list for the programs. For each program we independently generated random lists of up to 50 passes by uniformly sampling from the set of 167 searchable passes described previously. Every time we evaluated a pass list on a program we recorded the resulting binary size. We then pick the per-program pass lists that produced the lowest binary size. We ran 22 billion unique compilations for an average 4,877 per program.

2. 2.

   The pass lists generated by random search may contain redundant passes that have no effect on the final outcome. Further, some pass orderings are commutative such that reordering then does not affect the final outcome. Since these would introduce noise in our training data, we developed a minimization process which we applied to each pass list. Minimization comprises three steps: redundant pass elimination, bubble sort, and insertion search. In redundant pass elimination we minimize the best pass list by iteratively removing individual passes to see if they contribute to the binary size. If not, they are discarded. This is repeated until no further passes can be discarded. Bubble sort then attempts to provide a uniform ordering for pass subsequences by sorting passes based on a key. Finally, insertion sort performs a local search by iterating over each pass in the pass list and attempting to insert each of the 167 search passes before it. If doing so improves the binary size, this new pass list is kept. The entire minimization pipeline loops until a fixed point is reached. The distribution of minimized pass list lengths is shown in Figure 17. The average pass list length is 3.84.

3. 3.

   We apply PassListEval, described previously, to the candidate best pass lists. Through this we identified 167,971 of 1,704,443 unique pass lists (9.85%) as causing compile time or runtime errors.

4. 4.

   We broadcast the top 100 most frequently optimal pass lists across all programs, updating the per-program best pass lists if improvements are found. After this the total number of unique best pass lists decreases from 1,536,472 to 581,076.

The autotuning pipeline outlined above produced a geometric mean 7.1% reduction in binary size over -Oz. Figure 18 shows the frequency of individual passes. For our purposes, this autotuning serves as a gold standard for the optimization of each program. While the binary size savings discovered are significant, this required 28 billion additional compilations at a computational cost of over 21,000 CPU days. The goal of instruction fine-tuning LLM Compiler FTD to perform the flag tuning task is to achieve some fraction of the performance of the autotuner without requiring running the compiler thousands of times.

Using an LLM for disassembly causes problems of correctness. The lifted code must be verified by an equivalence checker which is not always feasible or manually verified for correctness or subjected to sufficient test cases to give confidence. However, a lower bound on correctness can be found by round-tripping. That is to say by compiling the lifted IR back into assembly, if the assembly is identical then the IR is correct. This gives an easy route to using the results of the LLM and an easy way to measure the utility of a disassembly model.

We provide the model with assembly code and train it to emit the corresponding disassembled IR. Listing 15 shows the prompt format. The context length for this task is set to 8k tokens for the input assembly code and 8k tokens for the output IR.

We derive the assembly codes and IR pairs from the same dataset used in previous tasks. Our fine-tuning dataset consists in 4.7M samples. The input IR has been optimized with -Oz before being lowered to x86 assembly.

We evaluate LLM Compiler FTD on the task of optimization flag tuning for unseen programs and compare to GPT-4 Turbo and Code Llama - Instruct. We run inference on each model and extract from the model output the optimization pass list. We then use this pass list to optimize the particular program and record the binary size. The baseline is the binary size of the program when optimized using -Oz.

For GPT-4 Turbo and Code Llama - Instruct we append a suffix to the prompt with additional context to further describe the problem and expected output format. After some experimentation we found that the prompt suffix shown in Listing 14 provides the best performance.

All model-generated pass lists are validated using PassListEval, and -Oz is used as substitute if validation fails. To further validate correctness of model-generated pass lists we link the final program binaries and differential test their outputs against the outputs of the benchmark when optimized using a conservative -O2 optimization pipeline.

We evaluate on 2,398 test prompts extracted from the MiBench benchmark suite (Guthaus et al., 2001). To generate these prompts we take all of the 713 translation units that make up the 24 MiBench benchmarks and generate unoptimized IRs from each. We then format them as prompts as per Listing 12. If the resulting prompt exceeds 15k tokens we split the LLVM module representing that translation unit into smaller modules, one for each function, using llvm-extract. This results in 1,985 prompts which fit within the 15k token context window, leaving 443 translation units which do not fit. We use -Oz when for the 443 excluded translation units when computing performance scores. Table 10 summarizes the benchmarks.

Table 3 shows zero-shot performance of all models on the flag tuning task. Only LLM Compiler FTD models provide an improvement over -Oz, with the 13B parameter model marginally outperforming the smaller model, generating smaller object files than -Oz in 61% of cases.

In some cases the model-generated pass list causes a larger object file size than -Oz. For example, LLM Compiler FTD 13B regresses in 12% of cases. These regressions can be avoided by simply compiling the program twice: once using the model-generated pass list, once using -Oz, and selecting the pass list which produces the best result. By eliminating regressions wrt -Oz, these -Oz backup scores raise the overall improvement over -Oz to 5.26% for LLM Compiler FTD 13B, and enable modest improvements over -Oz for Code Llama - Instruct and GPT-4 Turbo. Figure 7 shows the performance of each model broken down by individual benchmark.

While the model-generated binary size predictions have no effect on actual compilation, we can evaluate the performance of the models at predicting binary sizes before and after optimization to give an indication of each model’s understanding of optimization. Figure 8 shows the results. LLM Compiler FTD binary size predictions correlate well with ground truth, with the 7B parameter model achieving MAPE values of 0.083 and 0.225 for unoptimized and optimized binary sizes respectively. The 13B parameter model improved has similar MAPE values of 0.082 and 0.225. Code Llama - Instruct and GPT-4 Turbo binary size predictions show little correlation with ground truth. We note that the LLM Compiler FTD errors are slightly higher for optimized code than unoptimized code. In particular, there is an occasional tendency for LLM Compiler FTD to overestimate the effectiveness of optimization, resulting in a lower predicted binary size than actual.

Table 4 ablates the performance of models on a small holdout validation set of 500 prompts taken from the same distribution as our training data (though not used during training). We trained for flag tuning at each stage of the training pipeline from Figure 1 to compare performance. As shown, disassembly training causes a slight regression in performance from average 5.15% to 5.12% improvement over -Oz. We also show performance of the autotuner used for generating the training data described in Section 2. LLM Compiler FTD achieves 77% of the performance of the autotuner.

We evaluate the functional correctness of LLM-generated code when disassembling assembly code to LLVM-IR. As in Section 5.1 we evaluate LLM Compiler FTD and compare to Code Llama - Instruct and GPT-4 Turbo, and find that an additional prompt suffix, shown in Listing 16, is required to extract the best performance from these models. The suffix provides additional context about the task and the expected output format. To evaluate the performance of models we round-trip the model-generated disassembled IR back down to assembly. This enables us to evaluate accuracy of the disassembly by comparing the BLEU score (Papineni et al., 2002) of the original assembly against the round-trip result. A lossless and perfect disassembly from assembly to IR will have a round-trip BLEU score of 1.0 (exact match).

We evaluate on 2,015 test prompts extracted from the MiBench benchmark suite. We took the 2,398 translation units used for the flag tuning evaluation above and generated disassembly prompts. We then filtered the prompts on a maximum 8k token length, allowing 8k tokens for the model output, leaving 2,015. Table 11 summarizes the benchmarks.

Table LABEL:tab:disassembly-perf shows performance of the models on the disassembly task. LLM Compiler FTD 7B has a slightly higher round-trip success rate than LLM Compiler FTD 13B, but LLM Compiler FTD 13B has the highest accuracy of round-tripped assembly (round trip BLEU) and most frequently produces a perfect disassembly (round trip exact match). Code Llama - Instruct and GPT-4 Turbo struggle with generating syntactically correct LLVM-IR. Figure 9 shows the distribution of round-trip BLEU scores for all models.

Table 6 ablates the performance of models on a small holdout validation set of 500 prompts taken from the MiBench dataset used previously. We trained for disassembly at each stage of the training pipeline from Figure 1 to compare performance. Round trip rate is highest when going through the whole stack of training data and drops consistently with every training stage, though round trip BLEU varies little with each stage.

We ablate LLM Compiler models on the two foundation model tasks of next-token prediction and compiler emulation. We perform this evaluation at each stage of the training pipeline to see how training for each successive task affects performance. For next-token prediction we compute perplexity on a small sample of LLVM-IR and assembly code from all optimization levels. We evaluate compiler emulation using two metrics: whether the generated IR or assembly code compiles, and whether the generated IR or assembly code is an exact match for what the compiler would produce.

For next-token prediction we use a small holdout set of validation data that is drawn from the same distribution as our training data but has not been used for training. We use a mixture of optimization levels including unoptimized code, code optimized with -Oz, and randomly generated pass lists. For compiler emulatino we evaluate using 500 prompts generated from MiBench using randomly pass lists generated in the manner described in Section 2.2.

Table LABEL:tab:compiler-eval-tasks shows performance of LLM Compiler FTD across all training stages on the two foundation model training tasks of next-token prediction and compiler emulation. Next-token prediction performance jumps sharply after Code Llama, which has seen very little IR and assembly, and declines slightly with each subsequent stage of fine-tuning.

For compiler emulation, the Code Llama base model and the pre-trained models perform poorly since they have not been trained on this task. The highest performance is achieved directly after compiler emulation training where 95.6% of IR and assembly generated by LLM Compiler FTD 13B compiles, and 20% of it matches the compiler exactly. Performance declines after fine-tuning for flag tuning and disassembly.

While the purpose of LLM Compiler FTD is to provide foundation models for code optimization, it builds upon base Code Llama models which were trained for software engineering tasks. To evaluate how the additional training of LLM Compiler FTD has affected the performance of code generation we use the same benchmark suites as in Code Llama that evaluate the ability of LLMs to generate Python code from natural language prompts, such as “Write a function to find the longest chain which can be formed from the given set of pairs.”.

We use the HumanEval (Chen et al., 2021) and MBPP (Austin et al., 2021) benchmarks as in Code Llama.

Table LABEL:tab:human-eval-mbpp-scores shows the greedy decoding performance (pass@1) of all model training stages and model sizes starting at the Code Llama base model. It also shows the models’ scores on pass@10 and pass@100 which were generated with p=0.95 and temperature=0.6. Each stage of compiler-centric training causes a slight regression in Python programming ability. pass@1 performance on HumanEval and MBPP declines by up to 18% and 5% for LLM Compiler and by up to 29% and 22% for LLM Compiler FTD after the additional flag tuning and disassembly fine-tuning. All models still outperform Llama 2 on both tasks.