LOOPRAG: Enhancing Loop Transformation Optimization with Retrieval-Augmented Large Language Models

Yijie Zhi, Yayu Cao, Jianhua Dai, Xiaoyang Han, Jingwen Pu, Qinran Wu and 2 others

Rewriting a program's loops — reordering, tiling, fusing them without changing what it computes — is one of the biggest levers for performance, but finding the best combination is so hard that even decades-old polyhedral compilers fall short and large language models asked to do it often produce slower or outright wrong code; this paper (Zhejiang University and collaborators — Yijie Zhi, Ming Cai, and colleagues) builds LoopRAG, which teaches an LLM to optimize loops by synthesizing a library of diverse legal transformed examples with a polyhedral compiler, retrieving the most relevant ones as demonstrations, and iterating with compile-test-performance feedback plus rigorous equivalence checking — beating both compilers and base LLMs by up to roughly 11 to 14 times.

The argument in one glance

Problem

Rewriting loops is a huge performance lever, and nobody does it optimally

A large share of a program's time is spent in loops, and how those loops are structured — the order iterations run in, how they are tiled into blocks, whether adjacent loops are fused — decides whether the hardware can vectorize and parallelize the work. Loop transformations rewrite that structure at the source level without changing what the program computes, and they are one of the most powerful optimizations a compiler has. The catch is that finding the best combination of transformations is genuinely hard: it requires analyzing the data dependencies to know which reorderings are legal, and modeling the cost of each candidate to know which is fastest, over a combinatorial space of possibilities. Despite decades of research, even the most sophisticated compilers apply only a fraction of the available transformations and leave real speedups unclaimed. The optimal composition of loop transformations remains an open problem.
What people did

Ask a large language model to optimize the loops

Large language models have shown a knack for code, and a wave of recent work asks them to optimize programs directly — feed the model the loop nest and prompt it to rewrite it for speed, letting the model's learned knowledge of transformations stand in for a compiler's analysis.
Why it fell short

Left to themselves, LLMs get loop transformations wrong or make them slower

The paper's key observation is that LLMs struggle badly at loop transformation specifically. Asked to optimize a loop nest, a strong model frequently produces code that fails to compile or errors out, or that applies a transformation but misses the large speedups a real optimizer would find, or — most dangerous of all — that runs but computes a different result, silently violating the semantics-preserving requirement that makes an optimization an optimization. Compared against a polyhedral compiler on standard benchmarks, an off-the-shelf model both underperforms on many cases and produces a meaningful fraction of semantically wrong outputs. An unguided LLM is neither reliable enough to trust nor good enough to beat the compilers.
This paper's idea

Teach the LLM with synthesized, retrieved demonstrations and verified feedback

LoopRAG is a retrieval-augmented framework that guides an LLM to transform loops well and safely. First, it synthesizes its own library of demonstrations: a parameter-driven method generates diverse but legal loop programs, optimizes each with a polyhedral compiler, and stores the before-and-after pairs plus data-flow information. Second, for a target loop it uses a loop-aware retrieval algorithm — matching on loop features and balancing similarity against diversity — to pull the most informative example pairs and show them to the LLM as demonstrations. Third, it runs a feedback-driven iterative loop: the LLM's output is compiled, tested, and timed, and those results are fed back into the prompt to steer the next attempt, while every candidate is checked for semantic equivalence with mutation, coverage, and differential testing. The result beats both compilers and base LLMs by up to roughly 11 to 14 times across three benchmark suites, and even surpasses the compiler whose outputs taught it.

Why this problem is worth solving

Loop optimization is where the biggest source-level speedups live, and it is exactly the part compilers do imperfectly.

A loop transformation rewrites how a loop nest executes — permuting the loop order, tiling iterations into cache-friendly blocks, fusing or fissioning loops, skewing them to expose parallelism — without changing the program’s results. These transformations are what let a compiler map a naive loop onto vector units and many cores, so they are among the highest-leverage optimizations available. But applying them well is a long-standing challenge for two reasons. First, legality: a reordering is only valid if it respects the program’s data dependencies, so the compiler must analyze which statements depend on which, and a transformation that breaks a dependence produces wrong answers. Second, cost: among the legal transformations, which composition is fastest depends on the target hardware and the interaction of the transformations with each other, a combinatorial space that is hard to model and search. The polyhedral model gives a mathematical framework for the legal region — the Static Control Part, a loop nest with affine bounds and accesses — and tools built on it can find good schedules, but even they do not find the optimum, and much of the potential speedup goes unrealized. So there is a large, real gap between what loop transformations could deliver and what compilers actually extract, which is the gap worth closing.

Language models seemed like a way to close it, but they fail in a way that is uniquely unacceptable for a compiler task.

Optimizing code with a large language model is attractive because the model has absorbed enormous amounts of code and can propose transformations without hand-built analysis. But an optimization has a hard requirement that ordinary code generation does not: it must preserve semantics exactly — the transformed program must compute the same thing as the original, always. A language model generates by plausibility, not by proof, so it has no guarantee of legality; it can apply a reordering that looks right but violates a data dependence, yielding code that runs and returns a different, wrong answer. This is far worse than merely being slow: a slow optimization is a missed opportunity, but an incorrect one silently corrupts results, which no compiler can ship. The paper finds that a strong model, asked to transform loops, does exactly this on a nontrivial fraction of cases, alongside outright failures and suboptimal rewrites. So the model brings useful intuition but three failure modes — errors, suboptimality, and semantic violation — and any system that wants to use it for loop optimization must supply what the model lacks: legal, high-quality guidance about what transformations to make, and a rigorous check that whatever it produces is actually equivalent. Providing both is the design problem LoopRAG solves.

What people did before, and where it broke

Polyhedral compilers (PLuTo, GCC-Graphite, Clang-Polly). Use the polyhedral model to analyze dependencies and search for a good schedule of loop transformations. Where it broke. They are sound and effective but do not find the optimal composition; the search and cost models are limited, so substantial speedups remain unrealized on many loop nests.
Prompting an LLM to optimize code directly. Give the model the loop nest and ask it to rewrite it faster, relying on its learned knowledge. Where it broke. Unguided, the model frequently errors, misses the large speedups a real optimizer finds, and sometimes emits code that is not semantically equivalent — an unacceptable correctness failure for an optimization.
Retrieval-augmented generation over existing code datasets. Retrieve similar code examples to condition the model. Where it broke. Existing loop-code datasets lack the diversity and the paired before-and-after optimization information needed to demonstrate the full range of legal transformations, so the retrieved examples are not informative enough to teach effective loop optimization.
Trusting generated optimizations without strong verification. Accept the model’s output if it looks correct or passes light testing. Where it broke. Light checks miss the semantic violations LLMs produce; without rigorous equivalence checking, wrong optimizations slip through, which is fatal for a compiler-grade tool.
The gap. No approach combined a purpose-built, diverse, legally-transformed demonstration library, a loop-aware retrieval that surfaces the most informative examples, and a feedback-plus-equivalence-checking loop that makes the LLM both fast and provably semantics-preserving.

How LoopRAG works

Guiding an LLM to transform loops with synthesized demonstrations and verified feedback

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Do not ask the LLM to invent transformations; teach it with a compiler’s legal examples. The central move is to change the LLM’s role from an unaided optimizer to a guided one, because the model’s weakness is exactly a lack of legal, high-quality knowledge about which loop transformations apply, while its strength is composing and adapting patterns it is shown. LoopRAG supplies the missing knowledge by building its own demonstration library rather than trusting the model’s memory or a generic code corpus: a parameter-driven method synthesizes diverse but legal loop programs, and a polyhedral compiler transforms each, producing paired before-and-after examples of correct optimizations that span the space of transformations. Retrieved and shown as few-shot demonstrations, these examples give the model concrete, legal templates to follow, so it proposes transformations grounded in what a real optimizer does. This exploits a striking fact the paper leans on: the LLM, taught by the compiler’s outputs, can then go beyond them. Recognizing that mature compilers leave optimization potential unrealized, and treating their output as a starting point to exceed rather than a ceiling, is the same spirit the work on finding missed optimizations covered elsewhere on this site is built on, systematically exposing the speedups compilers overlook so they can be captured. LoopRAG turns that gap into training signal for a model that then surpasses the very compiler that taught it.

Retrieve demonstrations that are both like the target and unlike each other. A retrieval-augmented system is only as good as the examples it retrieves, and for loop optimization the naive choice — the most similar examples — is not enough, because a handful of near-identical demonstrations teach the model little more than one does. LoopRAG’s retrieval is loop-aware: it extracts features from the loop’s properties to measure both syntactic and semantic closeness, and it deliberately balances similarity against diversity, so the selected demonstrations are relevant to the target yet cover different transformation patterns, maximizing the information the model receives in a limited prompt. This is the difference between showing the model five variations of the same idea and showing it five complementary ideas that bracket the target. Carefully navigating a large space of candidate optimizations by balancing exploiting what looks promising against exploring what is different is the same search discipline the Bayesian compiler-optimization work covered elsewhere on this site applies to autotuning, trading off around a huge configuration space so the effort spent yields the most information. LoopRAG applies that exploration-exploitation balance to which demonstrations to retrieve, not which configurations to try.

Close the loop with feedback and prove equivalence, so the model is fast and correct. Even well-guided, an LLM’s first attempt may not compile, may be wrong, or may be beatable, so LoopRAG refuses to accept a single blind generation. It iterates, feeding the concrete results of compiling, testing, and timing each candidate back into the prompt, so the model fixes its own compilation errors and climbs toward faster code over a few rounds. And because the one failure a compiler can never tolerate is a silent change in behavior, LoopRAG verifies every candidate’s semantic equivalence to the original using mutation, coverage, and differential testing drawn from compiler-testing research, discarding anything that computes differently and outputting only the fastest that provably matches. This is what converts an unreliable generator into a trustworthy optimizer: the feedback makes it good, the equivalence checking makes it safe. Subjecting generated or transformed code to rigorous differential testing to catch behavioral divergences that plausibility alone would miss is the same correctness stance the compiler-driven differential-testing work covered elsewhere on this site takes, using systematic testing to surface the cases where a transformation quietly breaks. LoopRAG makes that verification an inline gate on every optimization the model proposes.

Following a loop nest through LoopRAG

Optimizing a loop nest with retrieved demonstrations and verified iteration

A trace, drawn from the paper's design, of how LoopRAG transforms a target loop safely and effectively.

Step 1 of 5

Offline: synthesize and optimize a demonstration library

Before any target is seen, LoopRAG uses its parameter-driven method to generate a diverse set of legal loop programs by configuring loop properties, runs a polyhedral compiler to optimize each, and stores each original alongside its compiler-optimized version and data-flow information as a demonstration dataset.

Synthesizing the examples gives coverage and paired before-and-after optimizations that existing code datasets lack.

Why it works

Compiler-generated demonstrations give the LLM the legality it lacks. An unguided model has no guarantee its transformations are legal → but paired before-and-after examples from a polyhedral compiler show concrete legal transformations → so the model proposes valid, high-quality rewrites instead of plausible-but-wrong ones.
Synthesized diversity covers transformations real datasets miss. Existing loop datasets lack breadth and paired optimizations → but a parameter-driven generator spans loop properties and each is compiler-optimized → so the demonstration library actually teaches the full range of transformations.
Similarity-plus-diversity retrieval maximizes prompt information. The most-similar examples are redundant and teach little → but balancing similarity against diversity selects relevant yet complementary demonstrations → so the model learns more per limited-size prompt and generalizes to the target.
Feedback iteration turns first-draft attempts into good ones. A single blind generation may not compile or may be slow → but feeding compile, test, and timing results back into the prompt → so the model repairs errors and improves speed across rounds, reaching results neither a compiler nor a one-shot LLM attains.
Equivalence checking makes the optimizer trustworthy. LLMs sometimes emit semantically-wrong code that light checks miss → but mutation, coverage, and differential testing verify each candidate against the original → so only provably-equivalent transformations are output, meeting the compiler-grade correctness bar.

What it costs (trade-offs and scope)

It targets Static Control Parts, the polyhedral-friendly loops. LoopRAG operates on loop nests with affine bounds and accesses that the polyhedral model and its demonstration-generating compiler can handle; loops outside this class — irregular control flow, non-affine indexing — are outside its scope, as they are for the compilers it builds on.
It depends on a compiler as its teacher. The demonstration library is generated by optimizing synthesized programs with a polyhedral compiler (PLuTo), so the quality and range of demonstrations are bounded by what that compiler can produce; LoopRAG surpasses it by composing and adapting, but the teacher still shapes the starting point.
Equivalence checking is testing, not proof. Mutation, coverage, and differential testing are strong and catch the semantic violations LLMs produce, but they are dynamic testing rather than formal verification, so they provide high confidence rather than a mathematical guarantee that a transformation is equivalent on all inputs.
The pipeline is heavier than a single prompt. Synthesizing and optimizing a demonstration library, running feature-based retrieval, iterating with compile-test-time feedback, and verifying each candidate is far more machinery and compute than a one-shot LLM call, so the speedups come at the cost of a substantial offline and per-query process.
The speedups are best cases over specific baselines. The up-to-11-to-14-times figures are maxima across PolyBench, TSVC, and LORE against particular compilers (GCC-Graphite, Clang-Polly, Perspective, ICX) and models (DeepSeek, GPT-4); average gains and behavior on other loops, hardware, or models will differ, and the results depend on the underlying LLM’s capability.
It optimizes performance, not the model. LoopRAG improves the code an LLM produces for loops but does not change the model itself; its value is a framework around the model, so it inherits the model’s limits and must be re-run per target rather than producing a general optimizer.

The evaluation

Applying loop transformations to speed up a loop nest

Property	Polyhedral compiler	Unguided LLM	LoopRAG
Guarantees legal transformations
Finds strong (non-obvious) transformations
Handles its own errors
Rejects semantically-wrong output
Reported speedup

A comparison drawn from the paper. LoopRAG is evaluated on PolyBench, TSVC, and LORE against compilers (GCC-Graphite, Clang-Polly, Perspective, ICX) and LLMs (DeepSeek, GPT-4); the up-to-11.20/14.34/9.29-times gains over compilers and up-to-11.97/5.61/11.59-times gains over base LLMs are the paper's reported speedups.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

The full paper text was available (the authors’ arXiv version, 2512.15766), so this account of the mechanism — the difficulty of optimal loop transformation, the three LLM failure modes (errors, suboptimality, semantic violation), and LoopRAG’s three-part pipeline (parameter-driven synthesis of a diverse legal demonstration library optimized by a polyhedral compiler, loop-aware retrieval balancing similarity and diversity, and feedback-based iteration with mutation, coverage, and differential equivalence checking) — follows the paper’s own description. The headline numbers are the paper’s reported results: average speedups over compilers of up to 11.20 times on PolyBench, 14.34 times on TSVC, and 9.29 times on LORE, and over base LLMs of up to 11.97, 5.61, and 11.59 times, using DeepSeek and GPT-4 and comparing against GCC-Graphite, Clang-Polly, Perspective, and ICX, with the framework surpassing its demonstration source PLuTo. Those are best-case figures over specific benchmarks, baselines, and models; the results depend on the underlying LLM and the polyhedral-friendly (Static Control Part) scope, and the equivalence checking is rigorous testing rather than formal proof. Exact per-benchmark numbers, the synthesis and retrieval algorithms, and the feedback details should be read from the paper itself.

Background

Why LLMs need scaffolding to do compiler work

It is worth seeing why a language model, impressive as it is at code, cannot simply replace a compiler pass. A compiler optimization is held to a standard ordinary programming is not: it must be correct on every input, every time, because it runs silently on code the programmer trusts, and a single wrong transformation corrupts results with no warning. Language models are fundamentally probabilistic — they generate the most plausible continuation, which is often but not always right, and they offer no proof of correctness. That mismatch means an LLM used naively for optimization is a liability, however good its suggestions look, because its inevitable mistakes are exactly the catastrophic, silent kind. The productive response, which this line of work embodies, is not to trust the model but to scaffold it: give it high-quality, legal examples so its suggestions are grounded, and wrap its output in rigorous checking so its mistakes are caught before they matter. The model supplies creative reach the compiler lacks; the scaffolding supplies the correctness the model lacks. This pattern of pairing a powerful but unreliable learned component with a verification layer that enforces the hard guarantees is the same architecture the ML-based program-tuning work covered elsewhere on this site relies on, using a learned model to propose and a rigorous evaluator to decide. LoopRAG is a careful instance of it for the specific, high-stakes task of transforming loops.

If you remember only three things

1. Loop transformations (reordering, tiling, fusing, skewing loop iterations without changing results) are among the highest-leverage source-level optimizations because they unlock vectorization and parallelism, but choosing the optimal composition is a decades-old hard problem of dependency analysis and cost modeling that even mature polyhedral compilers do not solve optimally. Large language models applied to it directly do worse: they frequently error, miss big speedups, or — worst — emit code that is not semantically equivalent, an unacceptable silent correctness failure for an optimization.
2. LoopRAG is a retrieval-augmented framework that guides an LLM to transform loops well and safely, in three parts: a parameter-driven method synthesizes diverse but legal loop programs, optimizes each with a polyhedral compiler, and stores the before-and-after pairs plus data-flow information as a demonstration library; a loop-aware retrieval algorithm balances similarity and diversity to pull the most informative demonstrations for a target loop; and a feedback-based iterative loop compiles, tests, and times each candidate, feeding results back into the prompt, while mutation, coverage, and differential testing verify semantic equivalence and reject wrong transformations.
3. Across PolyBench, TSVC, and LORE, LoopRAG achieves average speedups over compilers (GCC-Graphite, Clang-Polly, Perspective, ICX) of up to 11.20, 14.34, and 9.29 times, and over base LLMs (DeepSeek, GPT-4) of up to 11.97, 5.61, and 11.59 times, and it surpasses PLuTo, the compiler whose outputs formed its demonstrations. It targets polyhedral-friendly Static Control Parts, depends on that compiler as its teacher, uses rigorous testing (not formal proof) for equivalence, and is a heavier pipeline than a single prompt whose gains depend on the underlying model.

Check yourself

Check your understanding

1. Why is applying the optimal composition of loop transformations hard even for mature compilers?
2. What is the most dangerous failure mode of using an unguided LLM for loop optimization?
3. How does LoopRAG guide the LLM to make good, legal transformations?
4. What are LoopRAG's main results and the boundaries of the approach?

Sources

The paper is “LOOPRAG: Enhancing Loop Transformation Optimization with Retrieval-Augmented Large Language Models” by Yijie Zhi, Yayu Cao, Xiaoyang Han, Jingwen Pu, Qinran Wu, Sheng Cheng, and Ming Cai (Zhejiang University; Ming Cai also Beijing ShenZhou Aerospace Software Technology Ltd.) and Jianhua Dai (Zhejiang Institute of Administration), published at ASPLOS 2026, DOI 10.1145/3779212.3790183. It is licensed CC BY.

This explainer is based on the paper’s full text, obtained from the authors’ arXiv preprint (arXiv:2512.15766). The account of the mechanism — the difficulty of applying the optimal composition of loop transformations (legality via dependency analysis and cost modeling), the observation that LLMs struggle with loop transformation and produce errors, suboptimal results, or semantically inequivalent code; and LoopRAG’s three-part pipeline of dataset synthesis (a parameter-driven method configuring loop properties to generate diverse legal loop programs, optimized by a polyhedral compiler such as PLuTo, stored with data-flow information as a demonstration source), loop-aware retrieval (balancing similarity and diversity using loop features to select the most informative demonstrations for a target Static Control Part), and feedback-based iterative generation (four iterations using compilation, testing, and performance feedback, with mutation, coverage, and differential testing for semantic equivalence checking) — follows the paper’s own description. The reported results cited here are the paper’s own: average speedups over compilers (GCC-Graphite, Clang-Polly, Perspective, ICX) of up to 11.20 times on PolyBench, 14.34 times on TSVC, and 9.29 times on LORE, and over base LLMs (DeepSeek, GPT-4) of up to 11.97, 5.61, and 11.59 times, with LoopRAG surpassing its demonstration source PLuTo. Exact per-benchmark numbers, the synthesis and retrieval algorithms, and the feedback and equivalence-checking details should be read from the paper itself.