Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Agentic Auto-Scheduling: An Experimental Study of LLM-Guided Loop Optimization

Massinissa Merouani, Islem Kara Bernou, Riyadh Baghdadi

TL;DR

This work investigates whether off-the-shelf Large Language Models can effectively drive complex loop optimizations by interacting with a compiler in a closed loop. The proposed ComPilot framework treats the LLM as an optimization agent that proposes high-level loop transformations, which the compiler validates for legality and evaluates via measured performance, feeding results back to the agent. On the PolyBench suite, ComPilot achieves a geometric mean speedup of $2.66\times$ in a single run and $3.54\times$ in a best-of-5 strategy, and outperforms the Pluto polyhedral optimizer in many cases, highlighting the viability of learning-driven, feedback-based performance engineering. The results underscore the potential of agentic AI in code optimization, with practical implications for building adaptive, hardware-aware optimizers that rely on empirical measurements rather than static cost models.

Abstract

Automatic code optimization remains a difficult challenge, particularly for complex loop nests on modern hardware. This paper investigates a novel approach to code optimization where Large Language Models (LLMs) guide the process through a closed-loop interaction with a compiler. We present ComPilot, an experimental framework that leverages off-the-shelf LLMs, without any task-specific fine-tuning, as interactive optimization agents. ComPilot establishes a feedback loop where an LLM proposes transformations for a given loop nest to a compiler. The compiler attempts the transformations, reporting back legality status and measured speedup or slowdown. The LLM utilizes this concrete feedback to iteratively refine its optimization strategy. Our extensive evaluation across the PolyBench benchmark suite demonstrates the effectiveness of this zero-shot approach. ComPilot achieves geometric mean speedups of 2.66x (single run) and 3.54x (best-of-5 runs) over the original code. Furthermore, ComPilot demonstrates competitive performance against the state-of-the-art Pluto polyhedral optimizer, outperforming it in many cases. This experimental study demonstrates that general-purpose LLMs can effectively guide the code optimization process when grounded by compiler feedback, opening promising research directions for agentic AI in code optimization.

Agentic Auto-Scheduling: An Experimental Study of LLM-Guided Loop Optimization

TL;DR

This work investigates whether off-the-shelf Large Language Models can effectively drive complex loop optimizations by interacting with a compiler in a closed loop. The proposed ComPilot framework treats the LLM as an optimization agent that proposes high-level loop transformations, which the compiler validates for legality and evaluates via measured performance, feeding results back to the agent. On the PolyBench suite, ComPilot achieves a geometric mean speedup of in a single run and in a best-of-5 strategy, and outperforms the Pluto polyhedral optimizer in many cases, highlighting the viability of learning-driven, feedback-based performance engineering. The results underscore the potential of agentic AI in code optimization, with practical implications for building adaptive, hardware-aware optimizers that rely on empirical measurements rather than static cost models.

Abstract

Automatic code optimization remains a difficult challenge, particularly for complex loop nests on modern hardware. This paper investigates a novel approach to code optimization where Large Language Models (LLMs) guide the process through a closed-loop interaction with a compiler. We present ComPilot, an experimental framework that leverages off-the-shelf LLMs, without any task-specific fine-tuning, as interactive optimization agents. ComPilot establishes a feedback loop where an LLM proposes transformations for a given loop nest to a compiler. The compiler attempts the transformations, reporting back legality status and measured speedup or slowdown. The LLM utilizes this concrete feedback to iteratively refine its optimization strategy. Our extensive evaluation across the PolyBench benchmark suite demonstrates the effectiveness of this zero-shot approach. ComPilot achieves geometric mean speedups of 2.66x (single run) and 3.54x (best-of-5 runs) over the original code. Furthermore, ComPilot demonstrates competitive performance against the state-of-the-art Pluto polyhedral optimizer, outperforming it in many cases. This experimental study demonstrates that general-purpose LLMs can effectively guide the code optimization process when grounded by compiler feedback, opening promising research directions for agentic AI in code optimization.

Paper Structure

This paper contains 25 sections, 22 figures, 5 tables.

Table of Contents

  1. Introduction
  2. System Design and Methodology
  3. Context Initialization Phase
  4. Iterative Optimization Phase
  5. Results and Analysis
  6. Main Results
  7. Ablation Study
  8. Supplementary Analyses
  9. Related Work
  10. Prior Work on LLM-Based Code Optimization
  11. Domain-Specialized LLMs for Compiler Code Optimization.
  12. High-Level Code Optimization
  13. Positioning of Our Work
  14. Automatic Code Optimization Methods
  15. Polyhedral Compilers and Autoschedulers
  16. ...and 10 more sections

Figures (22)

  • Figure 1: Overview of the ComPilot framework, depicting the interaction between the LLM, the different modules, and the compiler.
  • Figure 2: Structure of the Context Prompt
  • Figure 3: Example of the message sent to the LLM showing the loop nest corresponding to the doitgen_XLARGE kernel.
  • Figure 4: Example program analysis generated by the LLM.
  • Figure 5: An example of the LLM's schedule proposition message.
  • ...and 17 more figures