用强化学习优化编译器代码，提升性能并降低人工成本。
A Reinforcement Learning Environment for Automatic Code Optimization in the MLIR Compiler
- 设计多离散动作空间与层级指针法，缩小优化搜索范围。
- 在PyTorch和LQCD生成的代码上实现自动优化，提升CPU性能。
- 为研究者提供可复现的强化学习编译优化实验环境。
代码优化对提升程序性能至关重要，但过程繁琐复杂，亟需自动化方法。强化学习（RL）为解决此类复杂优化问题提供了新思路。本文提出MLIR RL，一个面向MLIR编译器的强化学习环境，旨在推动编译器研究并实现自动代码优化。我们采用多离散动作空间设计，将动作空间构建为更简单子空间的笛卡尔积；同时引入“层级指针”方法，显著压缩循环交换变换相关动作空间规模，从而提升策略学习效率。为验证有效性，我们在目标为CPU的MLIR Linalg代码上训练了强化学习智能体，这些代码分别来自PyTorch生成的深度学习模型和由LQCD编译器生成的格点量子色动力学代码。本工作成果是一个可供社区探索新型强化学习驱动循环嵌套优化方案的研究平台。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。