Melhorando o Desempenho de Código Assembly com Modelos de Linguagem de Grande Escala via Aprendizado por Reforço

Improving Assembly Code Performance with Large Language Models via Reinforcement Learning

May 16, 2025

Autores: Anjiang Wei, Tarun Suresh, Huanmi Tan, Yinglun Xu, Gagandeep Singh, Ke Wang, Alex Aiken

cs.AI

Resumo

Modelos de linguagem de grande escala (LLMs) têm demonstrado um desempenho robusto em uma ampla gama de tarefas de programação, mas seu potencial para otimização de código ainda é pouco explorado. Este trabalho investiga se LLMs podem otimizar o desempenho de código assembly, onde o controle refinado sobre a execução permite melhorias que são difíceis de expressar em linguagens de alto nível. Apresentamos uma estrutura de aprendizado por reforço que treina LLMs usando Otimização de Política Proximal (PPO), orientada por uma função de recompensa que considera tanto a correção funcional, validada por casos de teste, quanto o desempenho de execução em relação ao compilador padrão da indústria gcc -O3. Para apoiar este estudo, introduzimos um benchmark de 8.072 programas do mundo real. Nosso modelo, Qwen2.5-Coder-7B-PPO, alcança uma taxa de aprovação de 96,0% nos testes e uma aceleração média de 1,47x em relação à linha de base gcc -O3, superando todos os outros 20 modelos avaliados, incluindo Claude-3.7-sonnet. Esses resultados indicam que o aprendizado por reforço pode desbloquear o potencial dos LLMs para atuarem como otimizadores eficazes do desempenho de código assembly.

English

Large language models (LLMs) have demonstrated strong performance across a wide range of programming tasks, yet their potential for code optimization remains underexplored. This work investigates whether LLMs can optimize the performance of assembly code, where fine-grained control over execution enables improvements that are difficult to express in high-level languages. We present a reinforcement learning framework that trains LLMs using Proximal Policy Optimization (PPO), guided by a reward function that considers both functional correctness, validated through test cases, and execution performance relative to the industry-standard compiler gcc -O3. To support this study, we introduce a benchmark of 8,072 real-world programs. Our model, Qwen2.5-Coder-7B-PPO, achieves 96.0% test pass rates and an average speedup of 1.47x over the gcc -O3 baseline, outperforming all 20 other models evaluated, including Claude-3.7-sonnet. These results indicate that reinforcement learning can unlock the potential of LLMs to serve as effective optimizers for assembly code performance.