# Compilers 2.0: AI as stochastic optimizer > Source: > Published: 2026-09-02 00:16:33+00:00 There has been a lot of discussion following the presentation of the Jalapeño MLA kernel at HotChips and subsequent commentary by SemiAnalysis. As OpenAI’s hardware team, we just barely touched on this little gold nugget: the fact that AI is writing our kernels, and that, when it does, we don’t really need to understand what the kernel does line by line. We glaringly left out: how is such a thing possible? What is the right way to think about this, versus a more traditional method of code generation? Is the optimized kernel as sound as the unoptimized one? For my background, I’ve worked on compilers for accelerators for well over a decade. I started XLA, which is an excellent compiler infrastructure with an industry-leading cross-company team and effort working on it. For the past 2+ years at OpenAI I’ve been trying to reconceptualize how compilers should work in the age of AI. New compiler formulations will draw on existing strengths, but it is impossible to deny that there is a powerful new tool to leverage in the toolkit. This will be a bit of a journey, but I hope to illuminate how AI is being used for the automation of computer program improvement; i.e. optimizing compilation. I do think, by way of AI, we may experience something we think of as “compilers 2.0”. AI is less fundamentally constrained in what it can propose, and what it proposes is a result of the model’s training and context, this leads me to classify it as a “stochastic optimizer” – this can pose challenges but, as we will see, is also a source of great strengths… A great deal of academic research and industry application is already headed in this direction, and rapidly uncovering the potential for AI’s involvement in the optimizing compiler realm, but we are at a point where it warrants a broad strokes explanation. Background Compilers take in programs and spit out translated or improved versions of those programs. Programs, on both the input and output side, have semantics that tell us what the programs mean, what they could possibly do, and how to reason about those things it could do. Those of us who work on compilers think of them much like pure functions – they take in a data structure and spit out a data structure that should have corresponding semantics. Sometimes our compilers focus on “lowering” or “translating”. For example, they may take in C and spit out x86-64 assembly, which we would often consider to be “lower level”. But often they are doing more than just translation as a sub-portion of that process… Our compilers, in practice, focus on “optimizing”. They may take in a data structure that represents the program – in our parlance an “Intermediate Representation” (IR) – and they try to produce a better version of that program. Sometimes “better” means it takes fewer cycles to run, sometimes it means it’ll have less unnecessary code, sometimes it means specializing for things that we can prove “must be true” about the program (partial evaluation). Now, briefly, consider that LLMs were originally created to translate human text from one language to another. Clearly translation is in their wheelhouse. And we can see through our use of LLMs on day to day tasks that they can also write new solutions and improve existing solutions. Many of us coders also have experience asking an LLM “optimize this snippet of code” and they remarkably can. (However, we need to know that they optimized the code correctly, which we will get to!) This is simply to highlight that LLMs have the capabilities that we look for in an optimizing compiler. Optimization and Optimality Optimizing compilers are, unsurprisingly, trying to increase optimality of the program they’re working on, by some objective (usually execution time). That is so difficult to do in the general case, for an arbitrary program, that there is a theorem called the full employment theorem for compiler engineers. (I only found this out after I chose to be a compiler engineer, but it still brought me comfort!) “Superoptimizers” are an amazing little sub-field of optimizing compilers. Imagine there is a given program, and we can say what it does via semantics. What is the most optimal program that has those same semantics? That’s what superoptimizers attempt to tackle, and it’s effectively a search problem… Imagine I’m trying to find the shortest program that had those same semantics, and I had a way to ask if a candidate program had the same semantics. I could, hypothetically, enumerate every program in objective order, and pick the smallest one that had the same semantics. However, enumerating every program in objective order sounds pretty intractable. One of my favorite academic papers, made in 2013 titled “STOKE” (Stochastic Superoptimization), asked: “well, what if we just randomly tweak programs over and over, do we then eventually observe the best program?” They proposed that via a random walk (and with our OG machine learning friend Markov Chain Monte Carlo / Metropolis-Hastings), eventually you’d see that optimal program. Monte Carlo tweaking is typically dumb (you randomly pick a tweak), but also fast. LLMs are very smart (many reasoning tokens), but comparatively slow. What if, instead of the dumb/fast Monte Carlo tweaking, we had LLMs figure out the directions in which to take the programs? We’d have a stochastic optimizer that was very intelligent, walking our program through the optimized program space. Intuitions for Optimization Let’s take a step back. Consider the person you know that best personifies “optimizes the heck out of snippets of code”. For short let’s call them “optimizin’ Ollie”. Ollie probably has a gut intuition for what kinds of code tweaks could bear fruit. Ollie probably tries some things to see if they work, and if they don’t work out, rolls it back and tries something else. But they have some intuition for what kinds of things are possible, and how they might be able to beat the compiler. These intuitions that Ollie has are often beyond what compilers do. Although modern optimizing compilers are quite impressive in their results, they are based on fairly simple rules and heuristics. In technical jargon, they are based on the idea of a local dataflow transform that is run to fixed point. We also phase order the considerations; i.e. we build compiler pipelines to consider A and then B, but not the composite AB problem. Schedulers and register allocators are a notorious example of this, many PhDs have been attempted on the composite scheduler-register-allocator (to get the benefits of collapsing the phase ordering), but they have been challenging to make work in practice. This is why Ollie’s expertise is valuable. Often Ollie knows how to balance several NP-complete problems with heuristics that are bespoke to the situation. So there is more bespoke context awareness and sensitivity. Ollie is also able to employ techniques that optimizing compilers may not apply profitably, especially in combination, things like outlining or crafting custom ABIs or transforms to enable vectorization, or the other slew of things that make us grumble “I wish the compiler had a way to just do this…” Now consider that AI, through whatever reasoning facilities it has, may be able to act as a mini Ollie. It may not have the matched intuition in terms of what will come to fruition, but it has an inkling of what can be profitable, and it can take many, many shots on goal. With this approach, unlike in the STOKE paper, we cannot guarantee that as time goes to infinity we can see the optimal program, but because the AI has “more human like” reasoning facilities, it can actually get significant human-like traction per unit time. Tying it Back: MLA Kernel Let me start by saying: I don’t know what low level code the AI spat out for our Jalapeño MLA kernel, but I do know how to type in the numpy for MLA. In the XLA compiler I previously worked on, we would fuse those numpy operations together into clumps, and then use a metaprogram called an “emitter” to lower it down to loops, instructions, and lower-level primitives. When the XLA compiler / emitter program did that, I didn’t need to care what assembly came out the back. For our stochastic optimizer, AI conceptually takes the place of the emitter meta-program – it both lowers down and optimizes, and we can ask it to optimize further and further towards roofline. I hope this makes it clear where the AI slots in and how it is analogous to a component in an existing optimizing compiler system. It’s also helpful to think: what layer we consider to be “assembly code” is now moving up. When you type in normal C++ and compile it at -O3 (the highest typical optimization level) you don’t expect to understand the assembly that comes out, even if you understood the C++ you had typed in. We’re doing the analogous thing here, but with a higher-level and more mathematical input specification. Now, a key question is how we check that the program we get out of the AI is indeed equivalent to the higher level description / numpy. That checking mechanism establishes the soundness of the AI stochastic optimization process. I expect a future blog post may go into more detail on this, but for now, suffice it to say that testing for semantic equivalence is possible and we do it. Accelerator programs are particularly amenable to strong, complete contracts that we can verify “are exactly what the AI optimized program does”, as they are quite mathematical and data flow oriented in their broad context. Note that many relevant techniques in this area were pioneered by efforts in the sub-field of program synthesis. Whereas optimizing compilers say, “here is a program with semantics, make it better but with equivalent semantics!”, program synthesis says, “I believe there exists a program with these semantics, please try to find the best one you can”. Program synthesis is a harder problem than optimizing compilation, but it is also less fundamentally constrained. It is effectively what humans like Ollie do when they do better than the optimizing compiler, and it is something that AI can now help us to automate. The AI can draw “inspiration” from the original program, but it need not just perform minor local transforms on it. Classic optimizing compilers won’t see “oh, you wrote a bubble sort” and, by understanding the contract, switch it to a quick-sort, but both Ollie and the AI are able to do that. This is what puts us more in the program synthesis regime with stochastic optimization than classical optimizing-compiler regime. This all comes together in the fact that you can start with something that is “not very far from the numpy”, wait 48 hours, and have an optimized kernel with the same semantics, as we showed in our HotChips talk: As the slide also notes, on our machine we’re often able to observe the AI climbing performance past our human experts even on the kernels we felt were fairly well tuned. Often there is a decent achievable percentage still left just due to the many varieties of combinations / permutations that may need to be explored. These are often intractably tedious for a human performance engineer. Recap & Conclusion A compiler, at the end of the day, is just a function. We give our program to that function, and we get back a better version of our program. The program we get out and the program we put in we expect to have the same semantics. Traditional optimizing compilers make programs better via dataflow rules and heuristics. These are fully understandable in their provenance, but also can be more limited in what moves they can make. By contrast AI, as a stochastic optimizer, just has to “think hard” and spit something out. Its moves are not as fundamentally limited, making them more analogous to our human expert optimizer. We do need ways to check that the programs that it spits out are sound and implement the same semantics we put in, and we do have those in place. And this kind of AI optimization is particularly well suited to mathematical operations which have very strong contracts. The contracts avoid the need to understand what the kernel does line by line. That’s how we got the AI generated MLA kernel!