A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler

TL;DR

This paper introduces MLIR RL, an RL environment for the MLIR compiler that facilitates automatic code optimization, demonstrating its effectiveness by training an agent to optimize MLIR Linalg code from deep learning and scientific computing domains.

Contribution

The paper presents a novel RL environment for MLIR, including a multi-discrete action space formulation and level pointers method to improve optimization efficiency.

Findings

RL agent successfully optimized MLIR Linalg code for CPU

Environment supports code from deep learning and scientific computing

Enables research in RL-driven loop-nest optimization

Abstract

Code optimization is a crucial task that aims to enhance code performance. However, this process is often tedious and complex, highlighting the necessity for automatic code optimization techniques. Reinforcement Learning (RL) has emerged as a promising approach for tackling such complex optimization problems. In this project, we introduce MLIR RL, an RL environment for the MLIR compiler, dedicated to facilitating MLIR compiler research and enabling automatic code optimization. We propose a multi-discrete formulation of the action space where the action space is the Cartesian product of simpler action subspaces. We also propose a new method, called level pointers, to reduce the size of the action space related to the loop interchange transformation. This enables more efficient and effective learning of the policy. To demonstrate the effectiveness of MLIR RL, we train an RL agent to optimize MLIR Linalg code, targeting CPU. The code is generated from two domain-specific frameworks: deep-learning models generated from PyTorch, and LQCD (Lattice Quantum Chromodynamics) code generated from an LQCD compiler. The result of this work is a research environment that allows the community to experiment with novel ideas in RL-driven loop-nest optimization.