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

Tirichine, Mohammed

doi:10.1109/CGO68049.2026.11394838

2026-01-01

PRE

Tirichine, Mohammed *

DOI：10.1109/CGO68049.2026.11394838

原文链接

原文求助

摘要

En 中文

代码优化是一项旨在提升代码性能的关键任务。然而，这一过程往往繁琐且复杂，凸显了自动代码优化技术的必要性。强化学习（RL）已作为一种有前景的方法，用于解决此类复杂的优化问题。在本项目中，我们介绍了MLIR RL，这是一个针对MLIR编译器的RL环境，旨在促进MLIR编译器研究并实现自动代码优化。我们提出了一个多离散的动作空间表述，其中动作空间是更简单动作子空间的笛卡尔积。我们还提出了一种新方法，称为“层级指针”（level pointers），用于减少与循环交换变换相关的动作空间大小。这使策略学习更加高效和有效。为证明MLIR RL的有效性，我们训练了一个RL代理来优化面向CPU的MLIR Linalg代码。代码来源于两个特定领域的框架：由PyTorch生成的深度学习模型，以及由LQCD编译器生成的LQCD（格点量子色动力学）代码。本工作的成果是一个研究环境，允许社区实验RL驱动的循环嵌套优化中的新思想。

Keyword:

Automatic Code Optimization

Reinforcement Learning

MLIR

Deep Learning

Machine Learning

Compiler

期刊

2026 IEEE/ACM INTERNATIONAL SYMPOSIUM ON CODE GENERATION AND OPTIMIZATION, CGO

IF:

论文数:

被引数:

机构

new york university

学者数:

6.2K

论文数: 3.0K

被引数: 1

引用论文

A loop transformation theory and an algorithm to maximize parallelism

IEEE Transactions on Parallel and Distributed Systems

IF0

1991-01-01

PREAI

M.E. Wolf; M.S. Lam

Loop transformations

Proceedings of the 38th annual ACM SIGPLAN-SIGACT symposium on Principles of programming languages

IF0

2011-01-26

PREAI

Louis-Noël Pouchet; Uday Bondhugula; Cédric Bastoul; Albert Cohen; J. Ramanujam; P. Sadayappan; Nicolas Vasilache

PolyGym: Polyhedral Optimizations as an Environment for Reinforcement Learning

2021 30th International Conference on Parallel Architectures and Compilation Techniques (PACT)

IF0

2021-09-01

PREAI

Alexander Brauckmann; Andres Goens; Jeronimo Castrillon

A Deep Learning Model for Loop Interchange

Proceedings of the 32nd ACM SIGPLAN International Conference on Compiler Construction

IF0

2023-02-17

PREAI

Lina Mezdour; Khadidja Kadem; Massinissa Merouani; Amina Selma Haichour; Saman Amarasinghe; Riyadh Baghdadi

Learning to Optimize Halide with Tree Search and Random Programs 学习用树搜索和随机程序优化Halide

ACM TRANSACTIONS ON GRAPHICS

IF9.5

2019-07-12

147

PREAI

Adams, Andrew; Ma, Karima; Anderson, Luke; Baghdad, Riyadh; Li, Tzu-Mao; Gharbi, Michael; Steiner, Benoit; Johnson, Steven; Fatahalian, Kayvon; Durand, Fredo; Ragan-Kelley, Jonathan

Optimizing memory usage in the polyhedral model

ACM Transactions on Programming Languages and Systems

IF0

2000-09-01

PREAI

Fabien Quilleré; Sanjay Rajopadhye

A Hybrid Machine Learning Model for Code Optimization 代码优化的混合机器学习模型

International Journal of Parallel Programming

IF0

2023-09-22

PREAI

Yacine Hakimi; Riyadh Baghdadi; Yacine Challal