🤖 AI Summary
Manual optimization of MLIR-based compilers is labor-intensive and inefficient. Method: This paper introduces the first reinforcement learning (RL)-driven automated optimization framework tailored for MLIR. We propose a multi-action RL modeling approach that decomposes the high-dimensional optimization action space into the Cartesian product of lower-dimensional subspaces, preserving expressiveness while significantly improving policy search efficiency and training stability. Furthermore, we establish the first end-to-end RL-driven compilation optimization loop in MLIR, integrating policy network training, environment interaction, and performance feedback. Contribution/Results: Experiments demonstrate that our method achieves performance comparable to or exceeding hand-tuned TensorFlow optimizations across multiple MLIR operator optimization tasks, thereby validating the feasibility, effectiveness, and practical potential of RL-enabled automated compiler optimization.
📝 Abstract
Code optimization is a crucial task aimed at enhancing code performance. However, this process is often tedious and complex, highlighting the necessity for automatic code optimization techniques. Reinforcement Learning (RL), a machine learning technique, has emerged as a promising approach for tackling such complex optimization problems. In this project, we introduce the first RL environment for the MLIR compiler, dedicated to facilitating MLIR compiler research, and enabling automatic code optimization using Multi-Action Reinforcement Learning. We also propose a novel formulation of the action space as a Cartesian product of simpler action subspaces, enabling more efficient and effective optimizations. Experimental results demonstrate that our proposed environment allows for an effective optimization of MLIR operations, and yields comparable performance to TensorFlow, surpassing it in multiple cases, highlighting the potential of RL-based optimization in compiler frameworks.