Leveraging LLMs to Automate Energy-Aware Refactoring of Parallel Scientific Codes (2505.02184v1)

Abstract: While LLMs are increasingly used for generating parallel scientific code, most current efforts emphasize functional correctness, often overlooking performance and energy considerations. In this work, we propose LASSI-EE, an automated LLM-based refactoring framework that generates energy-efficient parallel code on a target parallel system for a given parallel code as input. Through a multi-stage, iterative pipeline process, LASSI-EE achieved an average energy reduction of 47% across 85% of the 20 HeCBench benchmarks tested on NVIDIA A100 GPUs. Our findings demonstrate the broader potential of LLMs, not only for generating correct code but also for enabling energy-aware programming. We also address key insights and limitations within the framework, offering valuable guidance for future improvements.

• The paper introduces LASSI-EE, a novel framework leveraging LLMs for iterative, energy-aware refactoring of parallel scientific codes.
• Experimental results demonstrate significant energy savings, achieving up to a 47.36% average reduction and supporting 85% of evaluated codes.
• This research highlights LLMs' potential in automating complex HPC tasks, improving energy efficiency and bridging performance optimization with correctness.

Leveraging LLMs for Energy-Aware Code Refactoring

The paper "Leveraging LLMs to Automate Energy-Aware Refactoring of Parallel Scientific Codes" presents a paper on employing LLMs to automate the refactoring of parallel scientific codes with an emphasis on energy efficiency. The researchers propose LASSI-EE, a framework uniquely structured to harness the capabilities of LLMs beyond generating functionally correct code. The system addresses the energy aspects of code execution, aiming to develop optimizations that reduce energy consumption without compromising correctness.

Framework Structure and Methodology

LASSI-EE is built upon a multi-stage pipeline designed for iterative refinement and performance profiling of parallel codes. The framework emphasizes a thorough initial profiling stage, where baseline performance and power metrics of a code are established. This is followed by an LLM-invoked zero-shot refactoring attempt, which initializes a framework to automatically correct any generated code discrepancies through a self-correcting loop.

Subsequent stages involve building comprehensive contextual knowledge of optimizations and the execution platform. LASSI-EE leverages the LLM to generate self-prompts and context summaries, followed by developing a code refactoring plan. This plan informs iterative regeneration attempts, ensuring continual performance improvement and energy reduction. The final stage evaluates and compares energy metrics to ascertain the best refactoring outcome. A specialized LLM-as-a-Judge agent ensures functional validity across iterations, marking valid outputs accordingly.

Results and Analysis

The framework demonstrates significant potential, achieving up to a 47.36% average reduction in energy consumption across successful trials, supporting 85% of the evaluated codes. Notably, the power and execution time were considerably improved, with average reductions of 31.34% and 51.93%, respectively. Key strategies identified in refactored code include improved memory management, utilizing CUDA streams, and efficient communication practices. These highlight the versatility of LLMs in automating sophisticated coding practices tailored for performance optimization.

Some cases showed extremely high reductions, over 90%, which while flagged valid, suggest the need for human verification due to potential oversimplifications or misinterpretations of code functionality.

Implications and Future Directions

This research underscores the growing potential for leveraging LLMs in HPC, particularly in automating the labor-intensive task of code refactoring for energy efficiency. LASSI-EE effectively integrates context-driven enhancements and LLM-mediated validation to iterate on code transformations autonomously. This approach could transform code development processes, significantly reducing the energy footprint of HPC operations.

Future research could explore extending LASSI-EE to various programming languages and architectures beyond CUDA C++, enhancing adaptability. Further refinement of the validation module and incorporation of more sophisticated energy profiling techniques could improve the accuracy and reliability of agentic decisions in code refactoring. This foundational work paves the way for more robust AI integration as co-scientists in computational disciplines, further bridging the gap between performance optimization and code correctness in sustainable computing.