Machine learning?guided compiler optimization has recently attracted significant attention, with applications such as optimization pass prediction in LLVM IR and heuristic tuning in domain-specific language (DSL) compilers. In this work, we extend this paradigm to the Multi-Level Intermediate Representation (MLIR), a compiler infrastructure for DSL development that is widely used in AI compiler stacks. We present MLIRCompilerEnv, a reinforcement learning (RL) environment for predicting optimization passes and pass options in MLIR-based compiler pipelines, and evaluate it in the context of AI compilation. Our design targets the linalg, affine, and scf dialect levels, while remaining extensible to other MLIR-based compiler stacks for meta-optimization. We formulate the compiler optimization passes as a multi-stage RL problem, where each dialect stage is associated with a dedicated agent responsible for predicting optimization passes and their corresponding options. A message-passing graph neural network serves as a shared backbone across agents extracting structural features from the intermediate representations (IR). Agents operate sequentially to mirror the ordering of stages in the MLIR lowering pipeline, and end-to-end execution runtime is used as the reward signal to guide learning. To support structured learning over compiler IRs, we introduce a graph construction framework that fuses control-flow and data-flow information into a unified graph representation, covering MLIR dialects including linalg, scf, affine, arith, and math. We illustrate how this design naturally generalizes beyond AI compilers and facilitates extension to other MLIR-based compiler stacks such as stencil compilers. Initial results demonstrate the feasibility of the approach and we discuss performance trade-offs and compilation overhead introduced by the learning-based framework. The framework is currently being further developed to include matching optimization passes to the MLIR based IREE compiler for a comprehensive like-for-like comparison and analysis.

Speaker: Prakanth Thilakaraj (University of Warwick)