high performance computing on graphics processing units: hgpu.org

This work focuses on the use of deep reinforcement learning (DRL) to automate code optimization within modern compiler infrastructures. Code optimization is a critical step in program transformation that aims to improve performance and reduce resource consumption while preserving correctness. Traditional approaches to code optimization rely on manual or heuristic-based methods, which are often time-consuming and require deep knowledge of low-level programming and hardware. To overcome these limitations, this project proposes a DRL-based system capable of automatically learning and applying optimization strategies without human intervention. The system targets the MLIR (Multi-Layer Intermediate Representation) compiler framework and is designed to generalize across different application domains. To evaluate this generality, the system is tested on two distinct domains: neural network operators and Lattice Quantum Chromodynamics (LQCD) simulations. These domains differ significantly in computational patterns and optimization needs. The DRL environment is carefully constructed to represent code states, define legal transformation actions, and incorporate a reward function that reflects the effectiveness of applied optimizations. Additionally, The project includes two new proposed methods: a new architecture inspired from Pointer Networks for the interchange action and an adaptation of the AlphaZero algorithm for the code optimization problem. The proposed approach demonstrates promising results in learning optimization sequences that improve code performance across varied computational workloads.