Search-Based LLMs for Code Optimization

TL;DR

This paper reframes code optimization as a search problem and introduces SBLLM, a framework that iteratively guides LLMs to discover and refine efficient optimization methods using execution feedback, adaptive pattern retrieval, and GO-COT prompting. Empirical evaluations on the PIE Python/C++ benchmark across multiple LLMs show SBLLM achieving notable speedups and higher correctness, outperforming four prompting baselines. Ablation and parameter analyses confirm the contribution of each component and provide guidance on settings, such as an optimal 3 representative samples and up to 4 iterations. The work advances practical, scalable code optimization with LLMs and suggests broader applicability to other software engineering tasks.

Abstract

The code written by developers usually suffers from efficiency problems and contain various performance bugs. These inefficiencies necessitate the research of automated refactoring methods for code optimization. Early research in code optimization employs rule-based methods and focuses on specific inefficiency issues, which are labor-intensive and suffer from the low coverage issue. Recent work regards the task as a sequence generation problem, and resorts to deep learning (DL) techniques such as large language models (LLMs). These methods typically prompt LLMs to directly generate optimized code. Although these methods show state-of-the-art performance, such one-step generation paradigm is hard to achieve an optimal solution. First, complex optimization methods such as combinatorial ones are hard to be captured by LLMs. Second, the one-step generation paradigm poses challenge in precisely infusing the knowledge required for effective code optimization within LLMs, resulting in under-optimized code.To address these problems, we propose to model this task from the search perspective, and propose a search-based LLMs framework named SBLLM that enables iterative refinement and discovery of improved optimization methods. SBLLM synergistically integrate LLMs with evolutionary search and consists of three key components: 1) an execution-based representative sample selection part that evaluates the fitness of each existing optimized code and prioritizes promising ones to pilot the generation of improved code; 2) an adaptive optimization pattern retrieval part that infuses targeted optimization patterns into the model for guiding LLMs towards rectifying and progressively enhancing their optimization methods; and 3) a genetic operator-inspired chain-of-thought prompting part that aids LLMs in combining different optimization methods and generating improved optimization methods.

Figures (11)

• Figure 1: The overview of SBLLM.
• Figure 2: Examples for illustrating the two kinds of retrieved patterns in the adaptive optimization pattern retrieval module.
• Figure 3: The illustration of the GO-COT prompt. Contents in "[]" will be substituted by the corresponding data. The complete prompt can be found in our GitHub repository SBLLM.
• Figure 4: Venn diagram of optimized code provided by SBLLM and baseline methods on ChatGPT.
• Figure 5: The proportion of top-1 prediction based on ChatGPT with different optimization levels on Python and C++, respectively. "NC", "NO", "LH", and "FH" indicate the code is not correct, correct but not optimized, optimized but lower than human reference and faster than human reference, respectively.