Reinforcement Learning for Graph Coloring: Understanding the Power and Limits of Non-Label Invariant Representations

📄 arXiv: 2401.12470v1 📥 PDF

作者: Chase Cummins, Richard Veras

分类: cs.LG, cs.AI

发布日期: 2024-01-23

💡 一句话要点

提出基于强化学习的图着色方法以解决寄存器分配问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 寄存器分配 图着色 强化学习 近端策略优化 机器学习 编译器优化 图表示

📋 核心要点

  1. 寄存器分配问题复杂，现有方法在处理大量用户变量与有限寄存器时容易出现冲突。
  2. 本文提出将寄存器分配视为图着色问题，利用强化学习模型进行求解，强调图的标记对模型性能的重要性。
  3. 实验结果显示，模型在不同标记下的表现差异显著，强调了无标记重排序不变表示的必要性。

📝 摘要（中文）

寄存器分配是现代编译器中的重要问题，涉及将大量用户变量分配到有限的CPU寄存器中以避免冲突。本文将寄存器分配问题转化为图着色问题，利用PyTorch和OpenAI Gymnasium环境展示了近端策略优化模型在解决图着色问题上的有效性。研究表明，图的标记对模型性能至关重要，通过对图的矩阵表示进行置换，测试模型在不同标记下的效果，发现模型在相同图的重新标记下效果不佳。本文的主要贡献在于展示了机器学习模型需要无标记重排序不变的图表示，以实现一致的性能。

🔬 方法详解

问题定义：本文旨在解决寄存器分配问题，该问题在现代编译器中至关重要，现有方法在处理大量变量与有限寄存器时存在性能瓶颈。

核心思路：通过将寄存器分配问题转化为图着色问题，利用强化学习中的近端策略优化模型进行求解，探索图的标记对模型性能的影响。

技术框架：整体架构包括图的矩阵表示、标记置换、模型训练与评估等主要模块。首先构建图的表示，然后进行标记置换，最后训练模型并评估其在不同标记下的表现。

关键创新：最重要的创新在于揭示了图的标记重排序对模型性能的影响，强调了无标记重排序不变表示在机器学习中的必要性。

关键设计：在模型设计中，采用了近端策略优化算法，设置了适当的损失函数以优化图着色效果，网络结构经过精心设计以适应图的表示与标记变化。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，模型在不同标记下的表现差异显著，尤其是在相同图的重新标记下，模型的有效性显著降低。这一发现强调了无标记重排序不变表示的必要性，为未来的研究提供了新的方向。

🎯 应用场景

该研究的潜在应用领域包括编译器优化、图算法设计及其他需要高效资源分配的计算场景。通过改进寄存器分配策略，可以显著提升程序执行效率，降低计算资源消耗，具有重要的实际价值和未来影响。

📄 摘要（原文）

Register allocation is one of the most important problems for modern compilers. With a practically unlimited number of user variables and a small number of CPU registers, assigning variables to registers without conflicts is a complex task. This work demonstrates the use of casting the register allocation problem as a graph coloring problem. Using technologies such as PyTorch and OpenAI Gymnasium Environments we will show that a Proximal Policy Optimization model can learn to solve the graph coloring problem. We will also show that the labeling of a graph is critical to the performance of the model by taking the matrix representation of a graph and permuting it. We then test the model's effectiveness on each of these permutations and show that it is not effective when given a relabeling of the same graph. Our main contribution lies in showing the need for label reordering invariant representations of graphs for machine learning models to achieve consistent performance.