Investigating Execution-Aware Language Models for Code Optimization

1. Problem Statement

Code optimization is critical for software performance but traditionally relies on human developers. While Large Language Models (LLMs) have shown promise in automating code optimization, they are typically trained on static source code representations. Consequently, they lack an intrinsic understanding of run-time execution behavior (e.g., how often a line runs, variable states, or branch coverage).

Although recent research suggests that integrating execution information (execution traces) improves LLM performance in tasks like vulnerability detection and program repair, the specific impact of "execution-awareness" on code optimization remains unexplored. The authors hypothesize that teaching models to understand run-time behavior should theoretically help them identify inefficiencies, but this has not been empirically validated for optimization tasks.

2. Methodology

The study employs an empirical approach using CodeT5+ (a 220M parameter model) as the base architecture. The researchers investigated how incorporating four specific code execution aspects affects optimization capabilities using three distinct training strategies.

A. Code Execution Aspects Investigated

1. Line Executions (LE): Frequency of execution for specific code lines (identifying hotspots).
2. Line Coverage (LC): Whether a line is executed at least once.
3. Branch Coverage (BC): Whether a line belongs to a covered or uncovered conditional branch.
4. Variable States (VS): The type and value of variables at specific execution points.

B. Training Strategies

The authors defined three strategies to integrate these aspects into CodeT5+:

1. S1 (Execution-aware Pre-training + Fine-tuning):
   • Pre-training: The model learns to predict execution aspects (e.g., "given this code and input, what is the line coverage?") without running the code.
   • Fine-tuning: The model is fine-tuned on the code optimization task (slow $\to$ fast).
2. S2 (Execution-aware Pre-training + Masked Language Modeling + Fine-tuning):
   • Similar to S1, but the pre-training phase alternates between execution prediction and standard Masked Language Modeling (MLM) to preserve general code understanding.
3. S3 (Execution-aware Fine-tuning):
   • No pre-training. Execution information (traces) is injected directly into the input code (via comments) during the fine-tuning phase. This assumes the execution trace for the specific input is available.

C. Dataset and Evaluation

• Dataset: Derived from the PIE (Performance Improving Edits) dataset (77k C++ slow-fast pairs from IBM CodeNet).
• Trace Collection: Used gdb to generate execution traces for slow programs, limited to 500 seconds to prevent timeouts.
• Quantization: Raw numerical execution data (e.g., loop counts) was quantized into special tokens (e.g., <e> for 1 execution, <E+> for 21+ executions) to make them learnable by the transformer.
• Metrics:
  • Correctness (%): Percentage of generated code that compiles, runs, and passes all test cases.
  • Speedup: Ratio of input execution time to generated execution time ($T_I / T_G$).
  • %Opt: Percentage of generated code that is both correct and at least 10% faster.
• Baselines: Standard CodeT5+ fine-tuned on the same data (without execution info).

3. Key Contributions

• Comprehensive Evaluation: The first study to systematically evaluate the impact of execution-awareness specifically on code optimization.
• Twelve Model Variants: Evaluation of 12 distinct execution-aware models (4 aspects $\times$ 3 strategies) against standard baselines.
• Replication Package: Public release of code, datasets, and scripts to reproduce the findings.
• Counter-Intuitive Findings: Provided evidence that execution-awareness does not necessarily improve optimization and may even degrade performance in certain metrics.

4. Results

The study yielded surprising results that contradict the intuition that "more runtime knowledge equals better optimization."

• Overall Performance: Execution-aware models generally underperformed compared to the standard CodeT5+ baseline across all metrics (Correctness, Speedup, and %Opt).
  • Example: The baseline (BLS12) achieved 18.75% correctness and 1.79x speedup. The best execution-aware pre-training model (LCS2) achieved only 14.45% correctness and 1.76x speedup.
• Correctness Degradation: Execution-aware models produced significantly fewer correct programs. While they often generated syntactically valid code that compiled, they failed to preserve the intended functionality (semantic correctness) as often as the baseline.
• Strategy Comparison:
  • S1 & S2 (Pre-training): Consistently worse than the baseline.
  • S3 (Direct Fine-tuning): Showed marginal improvements over its specific baseline (BLS3) in some cases (e.g., Line Executions), but these improvements were statistically insignificant (p > 0.05).
• Nuanced Finding (Speedup on Correct Code): When analyzing only the subset of generated programs that were correct, execution-aware models sometimes achieved higher speedups than the baseline (e.g., LCS1 achieved 9.96x speedup on correct programs vs. 5.29x for baseline). This suggests that while these models struggle to generate correct code, when they do succeed, they can produce highly optimized solutions.

5. Significance and Conclusion

• Main Conclusion: Simply injecting run-time execution information into language models does not enhance their ability to optimize code. In fact, it often hinders the model's ability to generate semantically correct code.
• Implications for Future Research:
  • The "execution-aware" paradigm, while successful for bug detection or repair, may not be directly transferable to optimization without addressing the correctness gap.
  • Future work should focus on improving the semantic understanding of execution traces to ensure functional correctness before optimizing for speed.
  • Researchers should explore larger models or alternative training strategies, as the current approach (CodeT5+ with specific quantization) appears insufficient for this specific task.

In summary, the paper challenges the assumption that execution-awareness is a "silver bullet" for code optimization, highlighting a critical trade-off between understanding run-time behavior and maintaining functional correctness.