First-Principle Mojo: KGEN, MLIR, and What a Compiler Chooses to Remember

One thing I have come to admire about MLIR, and I say this with all due respect to LLVM, is how willing it is to move away from the rigidity of asking one intermediate representation to carry the entire burden of compilation. LLVM IR became enormously successful because it chose a particular level of abstraction and became extremely good at reasoning there, but that strength is also a boundary. By the time a program has been lowered into LLVM IR, some of the higher-level facts that once made the program easy to understand may already be gone, and trying to reconstruct those facts afterward can range from difficult to practically impossible.

MLIR approaches the problem differently. Instead of asking one IR to represent every useful idea, it provides a common compiler infrastructure in which multiple dialects can coexist. Each dialect can preserve a different kind of information about the program while still participating in the same broader system of operations, types, regions, blocks, SSA values, rewriting, and compiler passes.

Modular’s now open-source KGEN compiler makes this especially interesting because Mojo does not simply move from source code into one intermediate representation and then down to machine code. It passes through several representations, and sometimes several dialects exist together before one has completely disappeared. LIT, KGEN, POP, HLCF, and other MLIR dialects can coexist because they are not separate compilers lined up waiting for their turn. They are different vocabularies describing different aspects of the same program at the points where those aspects are still useful.

A simplified view of the journey looks something like this:

Mojo source
     ↓
    LIT
     ↓
KGEN / POP / HLCF
     ↓
elaboration
     ↓
concrete KGEN
     ↓
LLVM dialect
     ↓
LLVM IR
     ↓
machine code

What interested me almost immediately is that this makes lowering feel like something more nuanced than simply moving a program closer to the machine. The compiler is continually deciding what it still needs to know about the program, which representation makes that information easiest to reason about, and when a particular concept has finally done enough work that it can safely disappear.

That starts hella early in Mojo.

The conventional beginner model of a compiler usually begins with source code being tokenized, parsed into an Abstract Syntax Tree, transformed into some intermediate representation, and eventually lowered into machine code. An AST is essentially a structured representation of the program’s syntax. If we wrote a + b * 2, the tree would preserve the fact that the multiplication belongs underneath the addition, giving the compiler something much easier to reason about than the original flat sequence of characters.

Mojo does not center its compiler around a traditional full AST in quite the same way. It still has lightweight expression nodes and supporting declaration and type structures where they are useful, but the KGEN walkthrough describes the parser as beginning to emit MLIR operations into the LIT dialect very early. Rather than spending a long period living inside one large syntax tree before eventually entering the IR world, Mojo becomes compiler-shaped pretty damn quickly.

That decision becomes more interesting when you look at the parser itself. Mojo uses a lazy three-phase approach consisting broadly of name resolution, signature resolution, and body resolution. “Lazy” here does not mean the compiler is being slow or careless. It means it deliberately refuses to solve certain problems before it has the information required to solve them properly.

During name resolution, the compiler can walk through the program establishing that declarations exist without needing to understand their entire implementation yet. A structure named Foo exists. A function named bar exists. Their bodies can wait. During signature resolution, the compiler starts asking more specific questions about what those declarations are. Perhaps Foo is parameterized by some T, or bar accepts a Foo[Int]. Only later does body resolution dig into the actual implementations.

One practical benefit is that Mojo can handle forward references without forcing the programmer to litter source code with explicit forward declarations simply to keep the compiler happy. The programmer gets more freedom over how code is organized, while the compiler takes responsibility for discovering the names and signatures it will need later.

More than the convenience itself, I like the principle underneath it. The compiler does not insist on fully understanding a declaration the first moment it encounters something related to it. It gathers enough information to continue, then comes back once the surrounding context has become clearer. I keep seeing variations of that idea throughout KGEN: do not commit earlier than you need to.

Once parsing is complete, Mojo is represented through the LIT dialect, which the walkthrough describes as its source-level IR. LIT is where the program has clearly left human syntax behind, but has not yet left Mojo meaning very far behind.

Its operations still look recognizably connected to the language:

lit.fn
lit.call
lit.var.decl
lit.ref.load
lit.ref.store
lit.struct.decl
lit.trait.decl
lit.return

The types preserve similarly rich information, including things such as structures, traits, functions, generators, and particularly references with origin information:

!lit.ref

That last one becomes important because the compiler still has to reason about lifetime and borrowing rules. At this stage, treating a reference as merely a raw pointer would throw away information that the compiler has not finished using.

This is where I started thinking of LIT as Mojo, but compiler-shaped.

If the source contains something like:

def foo(arg: Int):
    pass

the resulting representation might contain a lit.fn and a lit.return. We are no longer looking at indentation, punctuation, or the human syntax of the language, but a function is still recognizably a function in Mojo terms. A trait declaration is still a trait declaration. A reference still carries information about its origins and lifetime semantics.

That matters because semantic checking comes next.

The compiler now has to determine whether the program merely parses or whether it actually makes sense under Mojo’s rules. Lifetime analysis and borrow checking happen while the representation still contains the information necessary to perform them. Destructor calls can be inserted, and uses of references that would outlive the things they depend upon can be rejected before those references are lowered into something less expressive.

This provides one of the cleanest examples I have found so far for understanding what lowering actually does. Imagine starting with:

!lit.ref

At that point, the compiler knows considerably more than “this is an address.” It knows this is a Mojo reference carrying lifetime-related information. Once the compiler has used that information to perform the relevant semantic checks, LowerLIT can eventually move the program toward something such as a KGEN pointer.

Conceptually:

reference + lifetime meaning
            ↓
borrow/lifetime analysis
            ↓
those semantics have done their job
            ↓
lower-level pointer representation

Nothing about this suggests that information should be preserved forever. In fact, a compiler eventually has to discard enormous amounts of information. The CPU does not need to know that some address originated as a Mojo reference governed by a particular source-level lifetime rule. What matters is whether that information was preserved long enough for the compiler to do everything useful with it before throwing it away.

That changed how I think about lowering. My earlier mental model was mostly directional: take something high-level and gradually translate it into something closer to the machine. That is not wrong, but it leaves out the more interesting part. Lowering is also a process of deciding which information has finished being useful.

Once LIT has done enough of that work, KGEN becomes the more important representation. Modular describes KGEN as the canonical parametric IR, and the word parametric is where things start getting especially fun.

Imagine a function conceptually represented as:

add(x)

where N is something that will be known at compile time, but has not yet been made concrete. KGEN can preserve this as a generator rather than immediately producing one fixed function. This has nothing to do with Python-style generators and yield; a KGEN generator is closer to a recipe capable of creating concrete code once its parameters become known.

A C++ template is a useful, if imperfect, analogy. Before specialization, there is a family of possible programs. Once a particular parameter arrives, the compiler can create the specific version it actually needs.

For example, a generator representing:

add(lhs)

might eventually be elaborated with rhs = 42, producing a concrete function in which 42 has become an actual constant. The compiler has moved from reasoning about possibility to reasoning about a particular program instance.

That transition is called elaboration, but KGEN does something interesting before it gets there: it optimizes the parametric program first.

At first this felt slightly backwards to me. Why not instantiate everything and optimize the results afterward? The answer becomes much clearer when you remember that elaboration can multiply code. One parametric generator may eventually produce many specialized concrete functions. If the original generator contains unnecessary complexity, every specialization risks carrying some version of that complexity with it.

It makes more sense to clean up the recipe before using it over and over again.

That is why pre-elaboration passes simplify generators, eliminate things that are not needed, promote suitable memory-backed values into SSA, propagate constants when possible, and perform carefully limited inlining. The qualifier carefully limited is important because even something we normally think of as an obvious optimization can become counterproductive depending on where it occurs. Excessive inlining before elaboration can make generators larger and hurt the elaborator’s ability to cache useful pieces of work.

That is the sort of compiler nuance I enjoy because it strips away the idea that an optimization is simply a magic button marked “make program faster.” An optimization changes the shape of a program, and that changed shape becomes somebody else’s input. Whether the transformation is useful therefore depends on what the next stage needs.

The presence of the POP and HLCF dialects reinforces that point. POP represents parametric operations around arithmetic, memory, and SIMD, while HLCF preserves higher-level structured control-flow concepts such as if, for, loops, break, and continue.

Preserving something like hlcf.for may initially seem unnecessary when the compiler could eventually express the same behavior using branches and basic blocks. But “this is a for loop” is useful information. While that structure remains explicit, the compiler can ask loop-shaped questions. It can reason about the bounds, consider unrolling, or perform transformations that are much easier when the loop is still visibly a loop rather than a collection of branch targets.

Later, that same structured loop can be lowered into lower-level control flow once the higher-level representation has finished helping.

There is even a RaiseForLoops transformation in the broader pipeline, which I find particularly interesting because it breaks the overly simple idea that compilation only moves downward. Sometimes a compiler can recognize structure in a lower-level representation and deliberately raise it into something richer because that richer form makes the next optimization easier. Once the optimization is finished, it can lower it again.

Compilation starts looking less like a one-way elevator and more like moving between different vantage points depending on which one gives you the clearest view of the problem.

Elaboration is where much of that parametric uncertainty finally collapses. Generators become concrete functions, generic structures become concrete structure instances, compile-time expressions can be evaluated, and static constraints can finally be resolved against known values.

Before elaboration, the compiler might be reasoning about add. After elaboration, it might be reasoning about add<42>.

That is a surprisingly meaningful transition. The compiler has exchanged generality for knowledge, and knowledge tends to create new opportunities for optimization. A branch whose condition depended on a compile-time parameter may now be constant. A generic structure may now have a fixed layout. An operation whose exact behavior was unclear before specialization may suddenly become trivial to simplify.

This is also why KGEN optimizes again after elaboration. Some passes appear on both sides of the boundary because they are not seeing the same program twice. Elaboration itself has changed what is knowable.

A transformation that could not prove something before may now have enough information to do so.

That makes the optimization pipeline feel less like:

optimize program
done

and more like an ongoing conversation with the representation. Transform the program, discover new information, simplify something else, expose another opportunity, and keep going until the next abstraction level becomes more useful.

KGEN also includes an interpreter for compile-time evaluation, which is one of those details that makes modern compiler machinery wonderfully strange. Not every compile-time value is simply sitting there as a literal waiting to be substituted. Sometimes the compiler actually has to execute enough program logic to determine what a parameter or expression evaluates to. KGEN can interpret control flow, function calls, and memory behavior within an emulated environment during compilation, allowing the result of that computation to become part of the concrete program being produced.

The compiler is, in a limited and deliberate sense, running some of the program in order to figure out what program it should ultimately compile.

Once elaboration has produced a concrete program and another round of optimization has done its work, the journey begins moving much more decisively toward LLVM. KGEN and related dialect types are progressively converted into lower-level forms until the program reaches MLIR’s LLVM dialect.

There is a distinction here that confused me initially and seems worth making explicit: the MLIR LLVM dialect is not LLVM IR.

The path is closer to:

KGEN / POP / HLCF
        ↓
MLIR LLVM dialect
        ↓
MLIR-to-LLVM translation
        ↓
LLVM IR
        ↓
LLVM optimization and target backend
        ↓
machine code

The LLVM dialect remains an MLIR dialect. It models concepts that map closely enough to LLVM that the program can eventually be translated into actual LLVM IR, after which LLVM’s existing optimizer and target infrastructure can continue carrying the program toward whatever CPU or GPU eventually has to execute it.

I find this especially interesting in the context of the LLVM work I have been reading recently. LLVM does not disappear because MLIR exists. Instead, LLVM becomes one particularly important place in a longer chain of representations. The program reaches LLVM once LLVM becomes the right level of abstraction for the questions that remain.

Even Mojo’s precompiled package format reflects this reluctance to rush toward final machine code. A .mojoc package can preserve MLIR bytecode before elaboration rather than flattening everything immediately into a final native binary. Function bodies remain available so later users of that package can provide concrete parameters and allow specialization to happen when those parameters actually exist.

Again, the compiler postpones commitment because later information may still change what the best program looks like.

One part of the KGEN walkthrough that I originally assumed would be a side issue ended up reinforcing this entire idea from a different direction: debug information.

Optimization spends much of its time changing the program into something that looks less and less like what the human originally wrote. Mem2Reg might take a local variable that once clearly lived in a stack slot and turn it into an SSA value. SROA might take one source-level structure and split it into several independent scalar values. Inlining can erase function boundaries. Constant propagation may cause entire branches to disappear.

All of that is great until someone opens a debugger and asks what happened to x.

The debugger still needs some connection back to the source-level world even though the optimized machine-oriented program may no longer contain anything resembling the original variable in one neat location. KGEN therefore has debug infrastructure that carries information about source scopes, variables, types, and how transformed IR values correspond to things the programmer remembers writing.

There is an interesting tension there. The optimizer is continually deciding which structure no longer matters for execution, while the debugging infrastructure has to preserve enough lineage that a human can still make sense of the transformed result.

That made me realize that there is not even one definition of “useful information” inside a compiler. Something can be irrelevant to code generation and still essential to debugging. Different consumers of the representation need different truths about the same program.

Before digging through KGEN, I still pictured compilation mostly as a path from rich human meaning toward increasingly concrete machine instructions. That direction remains true, but I think it misses what is actually interesting about modern compiler infrastructure.

The program keeps changing representation because the questions being asked about the program keep changing.

At the LIT level, Mojo semantics, references, traits, lifetimes, and source-level structure still matter. Semantic checking uses those facts to determine whether the program is valid. KGEN can then concentrate on parametric computation once some source-level semantics have finished their job. POP and HLCF preserve SIMD, memory, and structured control-flow information while those structures remain useful. Elaboration turns families of possible programs into the concrete instances actually required. Post-elaboration optimization takes advantage of the new information revealed by specialization. Eventually LLVM becomes useful because the remaining problem has moved much closer to low-level optimization and target-specific execution.

By the time a CPU or GPU receives the result, an enormous amount of meaning has disappeared, but ideally none of it disappeared before the compiler had finished using it.

That is probably the biggest thing KGEN has changed in how I think about lowering. The goal is not to preserve everything, nor is it to rush everything toward the machine as quickly as possible. The interesting problem is deciding which representation exposes the information needed right now, and when the compiler has finally earned the right to let that information go.

LLVM demonstrated how powerful a carefully chosen intermediate representation could become. MLIR loosens one assumption further by allowing several representations to occupy the space between source and silicon, rather than insisting that one of them be responsible for every useful abstraction along the way.

KGEN makes that idea unusually visible. LIT can remain Mojo-like while Mojo semantics still matter. KGEN can remain parametric while compile-time possibilities matter. HLCF can remember that something is a structured loop while that structure still helps optimization. Eventually those concepts can disappear into LLVM and then machine instructions without the compiler needing to pretend they were never useful in the first place.