Large Language Models for Compiler Optimization
arxiv.org
arxiv.org
Writing code to unroll a loop is trivial. The limitations of compilers are that almost all currently existing languages are too low level for optimization. ML has the potential to extract back this lost information. Basically the opposite of lowering an IR.
Heuristic replacement (like loop unrolling) is another big one. For the specific case of loop unrolling, I would think lower level elements like how much iCache pressure the unrolling creates/whether or not the loop could fit in the DSB buffer would matter more.
For your point about existing IRs being too low-level, there has been a large push to try and work on that. MLIR has been used pretty extensively for that problem in ML applications, and languages like Rust have multiple higher level IRs. There's also a preliminary implementation of a Clang-IR for C/C++, and there's even be some work on higher level representations within LLVM-IR itself.
You could still do that, you'd just also need to ask the model for a proof. (But I guess that's much harder than heuristically picking which passes to apply.)
My comment was describing the wacky idea of asking the model to come up with a formal, machine-checkable mathematical proof of correctness, too. That's hard in general.
The idea in the article of just letting the model pick between different, already proven-correct, optimization passes is much saner most of the time.
Hundreds of passes, sure. Guaranteed that they preserve correctness is a bit more dubious, that's pretty hard to establish for most transforms. Passes that make no assumptions about prior passes are tricky too since compilers tend to work in terms of a lowering pipeline.
If the compiler has N correct passes that can be combined in arbitrary order without compromising compiler termination, exponentially increasing code size or generally making the output much worse, then you've already built a really good compiler. The subtask of then shuffling the order of passes to see if you missed anything is trivial, using machine learning to control your sort & test loop doesn't seem very compelling here.
My hunch is that the low hanging fruit in compiler dev using LLM is driving a fuzz tester with one. Other things seem worthwhile but difficult.
This is called the "phase ordering problem", and it's neither trivial nor solved.
LIMA: Less Is More for Alignment https://arxiv.org/abs/2305.11206
AlpaGasus: Training A Better Alpaca with Fewer Data https://arxiv.org/abs/2307.08701
Textbooks Are All You Need II: phi-1.5 technical report https://arxiv.org/abs/2309.05463
It's possible, nay, mandatory to constrain the outputs of the model at each step of generation in order to guarantee that a given structure or grammar is adhered to. If you can fine-tune the model with these constraints in place you can offload a lot of the effort that the LLM otherwise has to perform in comprehending correctness so it has more capacity for generating good content. To be sure, quality and quantity of data are important, but it's all too easy to introduce subtle bugs that take years to tease out if you don't adhere to the right constraints.
This specific paper focuses on phase ordering, which should guarantee correctness, assuming the underlying transformations are correct. They do train the model to perform compilation, but as an auxiliary task.
- 3.0% improvement in reducing instruction counts over the compiler
- generating compilable code 91% of the time
- perfectly emulating the output of the compiler 70% of the time.
I read through the paper to see what perfectly emulating the output means. In this case, I think it's that it's also possible to get the same code out of the compiler using a different pass order. I was hoping for still passes original test suite or similar.
The authors are aware that the generated code has different semantics to the input in some cases but don't seem to consider that particularly important. The section is "Evaluation of Generated Code".
So - using machine learning, it is possible to delete a small percentage of the instructions emitted by a compiler while breaking the semantics. A tool which miscompiles programs while making them slightly smaller doesn't seem to be progress in compiler optimization.
Does anyone see some value add here that I'm missing?
We also train the model to generate what it thinks the optimized code will look like. We find that this helps the model choose better pass lists, but obviously the code cannot be trusted and semantics are not guaranteed. It only compiles in 91% of cases. "Perfectly emulating the output of the compiler" means the model spat out code that is character-for-character identical to what the compiler generates with the given pass list (even choosing the same variable names etc). IMO this is no mean feat, but is still a long way to go before using LLMs for codegen. We provide a bunch of examples in the paper of things that LLMs can and cannot do.
I agree that what you say is true of compilation as a whole, but that doesn't seem to be the focus here (rather, it's used as a sort of crutch to help the LLM learn)
Whether LLMs are the right approach is a separate question.
In SQL optimization, the problem is a bit trickier (IMO) because compilation is in the query path. One successful approach I know of is Bao: https://arxiv.org/abs/2004.03814
Current compilers use a complex set of heuristics, a more holistic approach the kind neural networks do might outperform.
Additionally, there are other factors to consider when productionizing these systems. Compile time is important (which LLMs will almost certainly explode), and anyone concerned about code size will probably be doing (Thin)LTO which would require feeding a lot more context into a LLM making inlining decisions.
1. https://arxiv.org/abs/2101.04808 2. https://dl.acm.org/doi/10.1145/3503222.3507744
1. https://github.com/google/ml-compiler-opt/pull/109 2. https://youtu.be/0uUKDQyn1Z4?si=PHrx9RICJIiA3E6C
But I agree, as of now I haven't seen good uses where LLMs produce reliable output. Not only do you need that guarantee that whatever output always generates a correct program, you need something where an LLM is considerably better than a simple or random algorithm, and you need a lot of training data (severely restricting how creative you can be with the output).
- Challenge: https://codalab.lisn.upsaclay.fr/competitions/15096
- Paper describing the challenge: https://arxiv.org/abs/2308.07899
(I am one of the authors, AMA)
Good luck on the challenge though, this seems like an interesting and valuable area of research.
Feel free to submit something! A simple submission is probably just a few lines of code.
3% code size reduction is really good. The challenge will be having codegen like this that someone is willing to support. And for that they'd want to be able to reason about why the compiler made this decision or that one. IIUC that's an outstanding problem for AI in general.
Good luck finding such a bug, because you will be looking on correct code, but computer will be executing invalid output.
Most work in ML for compilers focuses on replacing heuristics and phase ordering precisely because they don't impact correctness. There is some work being done on neural compilation [1], but I'm not sure that's going to be a viable approach anytime soon.
lol let's say they're less likely to impact correctness than an arbitrary new optimization.
At this point this is like (usefully) fuzzing your optimizer, which long term is going to be great for correctness.
3% code size reduction without changing semantics would be more interesting but might still be a bad thing for performance.
*fast-math etc is a thing, where similar-enough output is fine
Does that ever actually happen? I've only heard of it happening to people who forced the AI's hand by including the comments for said code in the prompt.
clippy in rust, for example, will tell you about removing needless allocations, etc, which a backend compiler might have a harder time doing.
#include <stdio.h>
int main(void) {
int i;
int a[2000] ;
for(i = 0; i < 2000; i++){
a[i] = i ;
}
printf("%d\n", a[5]);
}
Just as with the Python program, we can optimize this program significantly. Given that we're only interested in the sixth element (index 5) of the array, we don't need to construct and populate the entire array.
Here's an optimized version of the code:
#include <stdio.h>
int main(void) {
printf("%d\n", 5);
}LLMs are just good at searching and transforming between representations, but they really are not good at logical inferences.
I have used ChatGPT on some fairly awful code, asking it to add comments, rename variables and functions, etc. I find the outputs useful, but a couple of times it's broken the code. I imagine for many (not all) languages, you could ask the LLM to produce suggested changes, then use a code-rewriting tool to apply them in a way you could be 100% sure they weren't going to change the behaviour of the progrma.
Where the T can be anything from Ruby unit tests to Coq proofs.
Certified Reasoning with Language Models https://github.com/gpoesia/certified-reasoning
It's based on Peano, a theorem proving environment Peano: Learning Formal Mathematical Reasoning https://arxiv.org/abs/2211.15864
(https://github.com/kyegomez/LOGICGUIDE claims to implement the same paper as the first repo but it is fake)
https://discourse.llvm.org/t/pre-llvm-dev23-ml-guided-compil...
With that said, it is not doing neural compilation as others have mentioned, it’s only about ordering/enabling different phases of the compiler based on ML, over the current, simpler heuristics.
We haven’t experimented with model size yet, we just used the same configuration as the smallest Code Llama. We did play with dataset size and found thah performance tracks the usual scaling laws. Details in the paper
Attempting to convert assembler to C using an LLM would be interesting too, but the results would probably be poor since there is just so much information that gets lost when compiling, so the C code would be pretty statistical. I guess I could improve it by somehow adding "C code is or is not compilable because of line n" and whether the result is correct to the lost function.
What do you guys think about this?
Shouldn't this mean that it is safe to ask GPTs about something proprietary if it was just once because rare examples should just disappear in weights of everything else. And this also means that even GPT4 won't be able to answer any queries about obscure or rare knowledge it has seen in its training dataset.
Rather than just picking compiler arguments, an LLM could parse, and then transform AST into its more optimized form in a way that does not rely on a very formal way it is treated by compilers of today, apply guesstimate-based branch reordering and inlining heuristics not dissimilar to how a programmer would do so manually.
Once done, a compiler could perform final AST validation and either route it back to LLM to fix it or auto-fix most common cases.
For branch reordering, techniques like BOLT [5] are pretty effectively able to reorder code layout at the binary level for big performance gains by using profile information. ML models can sometimes synthesize that information [3], but if I recall correctly, the performance of those models wasn't as good.
Neural compilation (like what you're describing) has been tried with LLMs [4], but has a lot of correctness problems currently, and I don't think it's going to be feasible anytime soon to do reinforcement learning for performance/code-size improvements.
1. https://arxiv.org/abs/2101.04808 2. https://arxiv.org/abs/2207.08389 3. https://arxiv.org/abs/2112.14679 4. https://ieeexplore.ieee.org/document/9926313 5. https://arxiv.org/abs/1807.06735
If that's their target, then what's the point? LLVM's optimizations are not done to minimize instructions but to maximize performance. On modern processors, these can be very different things.
But LLVM targets both depending upon what optimization pipeline you select. (-Oz/-Os are targeting minimum code size, -O1,-O2,-O3 are optimization focussed).
Code size reduction is critical in some use cases like embedded environments and mobile apps and it is a significant area of research.
Theoretically changing the order of the passes in the optimization pipeline shouldn't cause any correctness issues, but the fact is that the ordering in the default compilation pipelines is the one that is most tested, so there will probably be bugs exposed when fuzzing the pass ordering.
Additionally, this work focuses on phase ordering, which produces correct code regardless of what the LLM puts out, assuming there aren't any bugs in the passes being used (which could crop up as random orderings aren't as well tested as the standard orderings in the commonly used pipelines).
⸻
1. Although I also find myself thinking about my kids as they were developing language where they initially correctly conjugated some common irregular verbs, then they started conjugating them as if they were regular and then finally returned to correctly conjugating them, which might be what’s happening with ChatGPT and math.