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

TL;DR

LOOPRAG tackles the difficulty of finding profitable loop transformations by using retrieval-augmented generation to guide LLMs with loop-property-aware demonstrations. It introduces a parameter-driven dataset synthesis to produce diverse, legal SCoP examples, a loop-aware retrieval algorithm to select informative demonstrations, and a feedback-driven iterative generation loop that uses compilation, testing, and performance rankings to refine results. Evaluated on PolyBench, TSVC, and LORE, LOOPRAG achieves substantial speedups over production compilers and base LLMs, and outperforms several LLM-based optimization approaches, demonstrating the practical impact of retrieval-guided, feedback-informed code optimization. While promising, the approach acknowledges limitations in formal equivalence guarantees and scalability, outlining future work to broaden optimization techniques and improve soundness. Overall, LOOPRAG represents a significant step toward automated, high-quality loop transformations guided by structured demonstrations and runtime feedback.

Abstract

Loop transformations are semantics-preserving optimization techniques, widely used to maximize objectives such as parallelism. Despite decades of research, applying the optimal composition of loop transformations remains challenging due to inherent complexities, including cost modeling for optimization objectives. Recent studies have explored the potential of Large Language Models (LLMs) for code optimization. However, our key observation is that LLMs often struggle with effective loop transformation optimization, frequently leading to errors or suboptimal optimization, thereby missing opportunities for performance improvements. To bridge this gap, we propose LOOPRAG, a novel retrieval-augmented generation framework designed to guide LLMs in performing effective loop optimization on Static Control Part. We introduce a parameter-driven method to harness loop properties, which trigger various loop transformations, and generate diverse yet legal example codes serving as a demonstration source. To effectively obtain the most informative demonstrations, we propose a loop-aware algorithm based on loop features, which balances similarity and diversity for code retrieval. To enhance correct and efficient code generation, we introduce a feedback-based iterative mechanism that incorporates compilation, testing and performance results as feedback to guide LLMs. Each optimized code undergoes mutation, coverage and differential testing for equivalence checking. We evaluate LOOPRAG on PolyBench, TSVC and LORE benchmark suites, and compare it against compilers (GCC-Graphite, Clang-Polly, Perspective and ICX) and representative LLMs (DeepSeek and GPT-4). The results demonstrate average speedups over base compilers of up to 11.20$\times$, 14.34$\times$, and 9.29$\times$ for PolyBench, TSVC, and LORE, respectively, and speedups over base LLMs of up to 11.97$\times$, 5.61$\times$, and 11.59$\times$.

Figures (14)

• Figure 1: Percentage of code optimization results generated by GPT-4 compared to PLuTo on PolyBench pouchet2023PolyBench and TSVC 44642. All codes are compiled using GCC. $\uparrow$ and $\downarrow$ indicate performance that outperforms and underperforms PLuTo, respectively. $\neq$ represents code that is not equivalent in semantics to the original.
• Figure 2: Original code syrk.
• Figure 3: LOOPRAG overview.
• Figure 4: The relationship between parameters and loop properties. The directed line segment from a parameter to a loop property indicates that the parameter can control the loop property. A loop property may be configured by one or more parameters.
• Figure 5: Implementation of Code Synthesis and Retrieval.