AI-guided compiler optimization: search, semantics, and trust

A search loop with an independent acceptance test

Let P be a reference program, K a candidate kernel, and H the target hardware. An optimizer searches for a candidate that reduces an execution cost while preserving the required behavior:

minimize:    cost(K, H)
subject to:  K is legal on H
             K satisfies the behavioral contract of P

The cost might be latency, throughput, energy, memory use, or code size. These objectives can conflict. A larger specialized kernel may run faster but consume more instruction memory; a configuration that maximizes throughput may worsen individual request latency. The objective and workload therefore belong in the specification of the search.

A generative model proposes implementations. Parsing, compilation, resource checks, correctness validation, and benchmarking determine whether a proposal is worth retaining. A compiler error becomes feedback for another attempt. An incorrect result produces a counterexample. A slower implementation supplies evidence about an unprofitable transformation.

best = qualified_baseline
best_cost = measure(best)
history = []

while search_budget_remains():
    candidate = propose(specification, target, best, history)
    artifact, diagnostics = compile_candidate(candidate)
    if artifact is absent:
        history.append(diagnostics)
        continue

    result = validate(artifact, fixed_contract, reference)
    if not result.accepted:
        history.append(result.counterexamples)
        continue

    cost = measure(artifact)
    history.append(cost)
    if confirmed_improvement(cost, best_cost):
        best, best_cost = artifact, cost

return best, contract, qualification_record

This sketch is a design pattern, not a description of a particular production system. Its essential property is that the proposal mechanism cannot relax the contract or modify the benchmark to improve its own score. Search inputs and acceptance infrastructure must remain separate.

Only qualified candidates enter timing comparisons. The eventual winner also needs fresh validation and repeated measurement: repeatedly selecting the smallest noisy timing can reward luck rather than a better implementation.

The contract survives lowering

A compiler translates between program representations. High-level tensor operations may become tiled loops, vector instructions, explicit memory transfers, and finally device instructions. This progression is called lowering. Each step commits to implementation details while retaining the behavior promised by the source.

An intermediate representation, or IR, makes the relevant facts available to a compiler pass. Tensor IR exposes shapes, reductions, and contractions. Loop IR exposes iteration spaces and memory accesses. Machine IR exposes registers and instruction constraints. StableHLO, for example, specifies tensor operations so that framework producers and compiler consumers can agree on their meaning.

For a pure numerical function, semantics begins with a mapping from input arrays to output arrays. A deployable kernel also has rules for types, shapes, strides, alignment, aliasing, mutation, synchronization, and exceptional values. An implementation that works for one contiguous matrix is not thereby valid for a transposed view or overlapping input and output buffers.

Semantic preservation may mean exact equivalence, numerical agreement within a stated error bound, or refinement. Refinement permits an implementation to choose among behaviors allowed by the specification without introducing forbidden ones. None of these relations permits the optimizer to silently discard a caller-visible guarantee.

Specialization narrows the problem. If a dimension, layout, or data type is known before execution, partial evaluation can remove branches, simplify indexing, and precompute static work. The resulting program handles a smaller input domain. Its caller must enforce that domain or select another implementation when the assumptions do not hold.

Why optimization choices interact

Conventional compilers combine rewrite rules, analyses, heuristics, and cost models. Constant folding, dead-code elimination, loop transformations, and vectorization can each improve a program. Running a transformation to a fixed point means that another application makes no further change under that rule set. It does not mean that the program is globally optimal.

Pass ordering matters because one transformation changes the opportunities available to another. Fusion may eliminate an intermediate array while increasing register pressure. Unrolling may expose independent work while lengthening live ranges. A layout change may enable vector loads but make another consumer's accesses less convenient.

Instruction scheduling and register allocation illustrate the coupling. A schedule that overlaps more operations can keep more temporary values live. Allocating those values may require spills, which introduce loads and stores that change the scheduling problem. Solving the choices together can be computationally expensive; production compilers use decomposition and heuristics to keep compilation affordable.

Some bounded compiler subproblems have NP-hard formulations. That is different from the undecidability of general program equivalence or termination. The first concerns the cost of solving finite problems; the second rules out a universal decision procedure for arbitrary programs. Neither prevents useful optimization within a restricted language and workload.

Learned decisions already appear inside compiler infrastructure. LLVM's MLGO applies machine learning to optimization policies. A generative kernel search broadens the possible changes: instead of selecting a decision within an existing pass, it may propose a different decomposition, storage strategy, or schedule.

From program mutations to learned proposals

Superoptimization treats implementation choice as a search problem. Given a specification and a candidate language, it seeks an equivalent program with a better cost. Exhaustive enumeration can establish optimality in a sufficiently small, well-defined space. Realistic instruction sequences quickly make that space too large.

STOKE explored stochastic search for loop-free x86-64 code. Its cost function combined correctness and performance terms, and an MCMC sampler explored program mutations. The research explicitly sacrificed completeness in exchange for a broader useful search space and strong results.

A random mutation might replace an instruction, alter an operand, or rearrange operations. A language model can propose coordinated changes using the specification, target documentation, current candidate, and accumulated diagnostics. Each proposal costs more to generate, but may traverse a larger useful part of the implementation space.

Calling the process stochastic identifies a role for probability in exploration. It does not establish that the optimizer implements Markov Chain Monte Carlo or inherits its convergence properties. Metropolis-Hastings requires a defined target distribution, proposal probabilities, and an acceptance rule. A history-dependent sequence of model calls with heuristic selection does not automatically meet those conditions.

Even a mathematically justified asymptotic result would not promise that a finite search budget finds the best kernel. The practical question is whether the search produces repeatable improvements within an acceptable time and compute budget. Cheap parameter tuning, structural model proposals, and conventional compiler passes can coexist in the same system.

What a candidate can change

Accelerator performance often depends on the arrangement of work and data as much as on the arithmetic expression. A kernel implements a small operation such as reduction, matrix multiplication, normalization, or attention. Its host supplies buffers and launch parameters; the kernel distributes the work across the device.

These choices interact. Increasing a tile can improve reuse while consuming enough registers to reduce resident workgroups. A fused operation can save bandwidth but spill temporary values. Occupancy is useful for hiding latency, but maximizing it is not itself the objective; the objective is faster execution of the required work.

Vectorization also depends on the execution model. The comparison of SVE and AVX-512 follows how vector width, predicates, and tail handling shape a program. A rewrite that helps one target may need a different form on another.

More invasive changes can alter an internal calling convention: how arguments, results, registers, and storage are shared between generated components. This is viable only within a controlled boundary where every participant agrees. It cannot silently change an externally visible ABI.

A kernel language makes the search tractable

A mathematical reference and a low-level kernel serve different purposes. Array notation makes the operation easy to inspect; a kernel language exposes the layouts, memory operations, and synchronization needed to execute it efficiently. A NumPy-like reference is a useful specification technique, but an analogy to NumPy does not identify a production compiler's actual input language.

Gluon is a lower-level GPU language on the Triton compiler stack. It exposes control over tile layouts, storage, data movement, and asynchronous execution that Triton often delegates to the compiler. That control creates more optimization opportunities and more obligations for the implementation.

Restricting proposals to a kernel language makes both generation and checking more manageable. Known primitives, explicit tensor shapes, bounded effects, and a defined target rule out large parts of general application behavior. The language's restrictions do not automatically guarantee memory safety or numerical correctness, but they make those properties more concrete.

XLA's GPU emitters provide a useful comparison: compiler code constructs a lower-level implementation for an operation or fused region. A model can explore implementation choices in that part of a toolchain while still relying on conventional parsing, type checks, lowering, assembly, and runtime integration.

Attention makes memory traffic part of the algorithm

For a single attention head, a simplified expression is:

attention(Q, K, V) = softmax(Q * transpose(K) / sqrt(d)) * V

The query Q is compared with keys K; normalized scores weight the values V. The scale uses the query/key dimension d. During autoregressive generation, previously computed keys and values are retained in a KV cache so that each new token can reuse them.

As sequence lengths and batches grow, reading that state can become expensive. DeepSeek-V2's Multi-head Latent Attention, or MLA, reduces persistent attention state through a compressed latent representation. Its design also handles positional information; treating MLA as simply compressing two ordinary arrays misses part of the implementation.

A latent representation changes the balance between stored state, projection work, and reuse. An implementation need not materialize every expanded key and value: algebraic rearrangement and absorbed projections can change which intermediates exist. The profitable choice depends on the attention phase, tensor shapes, precision, and hardware.

An optimizer must decide how to tile the work, retain useful state, schedule reductions, manage precision, and overlap transfers with computation. The same formula can admit several very different memory schedules. This makes attention a useful example of joint optimization across mathematical structure and device execution.

OpenAI's Jalapeño results report that AI-generated implementations of selected GPT-OSS attention and mixture-of-experts blocks ran 1.5 to 1.8 times faster than existing expert implementations. These are reported results for selected blocks on that system. They are not a general speedup for generated code, a whole-model multiplier, or an independent measurement performed here.

Roofline separates bandwidth from compute limits

Arithmetic intensity is the amount of arithmetic performed per byte transferred from a specified memory level. The roofline model combines that intensity with bandwidth and peak arithmetic throughput:

I = operations / bytes transferred
P <= min(P_peak, bandwidth * I)

At low intensity, bandwidth limits the attainable rate. Reusing a loaded tile or avoiding an intermediate write can improve performance more than selecting a faster arithmetic instruction. At high intensity, arithmetic throughput may become the dominant limit.

Consider a hypothetical device with a 100 TFLOP/s arithmetic ceiling and 2 TB/s memory bandwidth. A kernel performing 10 floating-point operations per byte has a bandwidth roof of 20 TFLOP/s. Raising its intensity to 40 operations per byte raises that roof to 80 TFLOP/s. These are illustrative bounds, not measurements of Jalapeño or another specific accelerator.

The byte count must correspond to the memory level being analyzed. Cache traffic, external-memory traffic, and inter-device communication have different ceilings. A measured point below the simple roof may be limited by dependencies, instruction issue, synchronization, occupancy, or launch overhead. The model identifies questions to investigate; it does not certify that every kernel can reach the line.

Fusion and tiling can change intensity, while scheduling and vectorization can improve utilization at a given intensity. This distinction helps explain why a successful optimization may reduce traffic rather than increase instruction throughput. The M4 cache-latency investigation similarly shows why the memory level exercised by a benchmark matters when interpreting its result.

Numerical equivalence must be specified

Real-number identities are not sufficient rules for floating-point optimization. Rounding makes addition non-associative; changing a reduction tree can change its output. Fused multiply-add, mixed-precision accumulation, and approximate device instructions introduce further choices.

A concrete example uses binary64 values that are exactly representable individually:

a = 2**54
b = -(2**54)
c = 1.0

(a + b) + c = 1.0
a + (b + c) = 0.0   # b + c rounds back to b

The example assumes round-to-nearest, ties-to-even binary64 arithmetic at each addition. It demonstrates why an optimizer needs permission to reassociate expressions. Similar issues arise when a parallel reduction changes the order of many additions.

A contract can require bitwise equality, a bound in units of last place, or an absolute/relative tolerance. One common comparison for finite values has the form:

abs(candidate - reference) <= atol + rtol * abs(reference)

The absolute term handles values near zero, where relative error becomes unstable. The contract must separately define how to treat NaNs, infinities, signed zero, underflow, and subnormal values. A tolerance should follow from the algorithm's accuracy requirements, not be widened until a promising candidate passes.

Aggregate accuracy can also hide localized failures. Depending on the operation, validation may need both elementwise bounds and a meaningful norm or downstream accuracy measure. Goldberg's floating-point analysis provides the background for these distinctions.

Testing and proof answer different questions

Differential testing runs the candidate and an independent reference on the same inputs. Useful cases include irregular shapes, alignment boundaries, extreme magnitudes, cancellation, masks, and permitted exceptional values. Random testing broadens coverage, while preserved counterexamples prevent a later proposal from reintroducing an earlier error.

Property-based and metamorphic checks add another perspective. Permuting independent batch elements should permute their outputs. Adding masked padding should preserve unmasked results if the mask contract says so. Such relations can expose mistakes shared by a candidate and a superficially similar reference.

Numerical agreement alone does not establish safety. A kernel can return correct values while reading beyond a buffer, using uninitialized storage, or relying on a race that happens not to fail. Static analysis, sanitizers, simulation, and target-specific checks address properties that output comparisons miss.

Formal verification establishes properties within a mathematical model. Translation validation checks a particular generated result against its source, rather than requiring a proof of the entire generating compiler. Alive2 applies this approach to LLVM optimizations; its supported operations and analysis limits remain part of the guarantee. Its existence does not imply that an arbitrary accelerator kernel is covered.

Soundness means that an accepted claim is true under the checker's model and assumptions. Completeness means that every true claim in the relevant class can be established. A useful verifier may be sound but incomplete, reporting an unknown result when it cannot finish a proof. An unknown result must not silently become an acceptance.

Finite tests generally do not prove equivalence across an unrestricted domain. Conversely, a proof under restricted assumptions does not justify deploying the result outside those assumptions. Confidence comes from knowing precisely what each check establishes and ensuring that the deployed invocation remains within that scope.

Benchmark the artifact that will run

Device execution is often asynchronous. Timing a host launch without waiting for completion can measure submission overhead instead of the kernel. Compilation, initialization, transfers, and synchronization must be included or excluded deliberately, according to the objective.

Warm-up, repeated trials, randomized candidate order, and matched power and thermal conditions reduce misleading comparisons. The distribution matters: a lower median accompanied by severe tail regressions may be unacceptable. Finalists should be remeasured against the baseline in fresh trials rather than accepted from their best observed run.

A search can overfit its benchmark just as it can overfit tests. Keep qualification workloads separate from the examples used to guide generation, cover the intended deployment shapes, and compare end-to-end behavior when kernel changes affect surrounding transfers or synchronization.

Preserve the exact binary that was validated and timed. Recompiling the same candidate with different flags, a new backend, or different target features creates a new artifact whose evidence may no longer apply. Generating source deterministically is useful, but reproducible deployment also requires recording the toolchain and the executable output.

Broader transformations require stronger boundaries

Optimizing compilation begins with an implementation to improve. Program synthesis begins with a description of required behavior and searches for an implementation. Generative optimization can cross that boundary when a candidate changes the algorithm rather than merely rewriting its instructions.

For example, replacing one sorting method with another may preserve sorted values while changing stability, storage use, comparator calls, or exception timing. If callers can observe those differences, the replacement may violate the contract. An apparently obvious algorithmic improvement is legal only after the required behavior has been made explicit.

Restriction helps. A kernel language can limit available primitives, memory effects, synchronization, and supported shapes. A reference can state the mathematical operation independently of the candidate. Acceptance tools can reject unsafe or numerically invalid implementations before any timing result influences deployment.

Reading generated code remains valuable for debugging, maintenance, portability, and incident response. It is not interchangeable with verification: a readable explanation does not prove equivalence, and a valid artifact can be difficult to inspect. The relevant decision is which guarantees the surrounding system establishes strongly enough to let engineers work above the instruction level.

A deployable result therefore consists of three things: the executable candidate, its contract, and its qualification record. That record should identify the source and binary hashes, generation inputs and budget, compiler and runtime versions, target configuration, correctness checks, benchmark distributions, and supported invocation domain. A known-good fallback and explicit requalification triggers complete the operational boundary.

The compiler's responsibility persists throughout: translate a defined computation into an implementation that callers may safely use. Stochastic search expands how that implementation is found. Its value depends on the precision of the contract and the independence of the evidence used to accept it.