RL4ReAl: Reinforcement Learning for Register Allocation

TL;DR

This paper introduces RL4ReAl, a reinforcement learning-based approach integrated into LLVM for register allocation, achieving results comparable or superior to existing highly optimized allocators across multiple architectures.

Contribution

It presents a novel multi-agent reinforcement learning framework for register allocation that is architecture-independent and maintains semantic correctness.

Findings

Matches or outperforms LLVM's production register allocators

Demonstrates architecture independence with x86 and ARM AArch64

Uses a modular gRPC-based framework for training and inference

Abstract

We aim to automate decades of research and experience in register allocation, leveraging machine learning. We tackle this problem by embedding a multi-agent reinforcement learning algorithm within LLVM, training it with the state of the art techniques. We formalize the constraints that precisely define the problem for a given instruction-set architecture, while ensuring that the generated code preserves semantic correctness. We also develop a gRPC based framework providing a modular and efficient compiler interface for training and inference. Our approach is architecture independent: we show experimental results targeting Intel x86 and ARM AArch64. Our results match or out-perform the heavily tuned, production-grade register allocators of LLVM.