CuAsmRL: Optimizing GPU SASS Schedules via Deep Reinforcement Learning

📄 arXiv: 2501.08071v1 📥 PDF

作者: Guoliang He, Eiko Yoneki

分类: cs.AR, cs.LG

发布日期: 2025-01-14

备注: cgo 2025

💡 一句话要点

CuAsmRL：利用深度强化学习优化GPU SASS指令调度

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: GPU优化 SASS调度 强化学习 CUDA内核 自动优化

📋 核心要点

  1. 现有CUDA内核优化依赖人工，耗时且易出错，存在性能提升瓶颈。
  2. 提出CuAsmRL，将SASS调度优化建模为强化学习问题，自动探索更优调度。
  3. 实验表明，CuAsmRL能显著提升现有CUDA内核性能，最高提升26%，平均提升9%。

📝 摘要（中文）

大型语言模型(LLMs)对计算资源的需求巨大。为了降低成本，研究人员开发了专门的CUDA内核，通常融合多个张量操作，以最大限度地利用GPU。然而，这些专门的内核仍然可能存在性能提升空间，CUDA汇编专家表明，手动优化GPU SASS调度可以带来更好的性能，但通常需要大量的试错来手动找到最佳的GPU SASS调度。本文提出了一种自动优化GPU SASS调度的方法，可以集成到现有的编译器框架中。自动优化的关键是训练一个强化学习(RL)智能体，模仿人类专家执行手动调度。为此，我们构建了一个汇编游戏，RL智能体可以在其中找到最佳的GPU SASS调度。汇编游戏从一个 extit{-O3}优化的SASS调度开始，RL智能体可以迭代地应用动作来改变当前的调度。如果改变后的调度通过在GPU上执行获得更高的吞吐量，则会产生正向奖励。实验表明，CuAsmRL可以透明地进一步提高现有专用CUDA内核的性能，最高可达26%，平均提高9%。此外，它还可以作为一个工具来揭示自动学习到的潜在优化策略。

🔬 方法详解

问题定义：论文旨在解决GPU SASS指令调度优化问题。现有方法主要依赖人工手动调整，需要专家经验，耗时且容易出错，难以充分挖掘GPU的性能潜力。即使经过编译器优化（如-O3），仍然存在进一步优化的空间。

核心思路：论文的核心思路是将SASS指令调度优化问题建模为一个强化学习任务。通过训练一个RL智能体，使其能够学习人类专家的优化策略，自动探索更优的指令调度方案。这种方法避免了人工试错的低效性，并有可能发现人工难以找到的优化方案。

技术框架：CuAsmRL的整体框架可以概括为：首先，从一个经过-O3优化的SASS调度开始；然后，RL智能体通过执行动作（例如指令移动、交换等）来修改当前的调度；修改后的调度在GPU上执行，并根据吞吐量变化获得奖励；RL智能体根据奖励更新策略，不断学习和改进。这个过程构成了一个“汇编游戏”。

关键创新：该论文的关键创新在于将GPU SASS指令调度问题转化为一个强化学习问题，并设计了一个有效的汇编游戏环境。通过强化学习，可以自动地学习和优化SASS指令调度，而无需人工干预。这种方法能够发现传统编译器优化无法找到的优化机会。

关键设计：汇编游戏的设计至关重要，包括状态表示、动作空间、奖励函数等。状态表示需要能够充分描述当前的SASS调度状态，动作空间需要包含所有可能的调度操作，奖励函数需要能够准确反映调度方案的性能。具体参数设置和网络结构等细节在论文中可能有所描述，但摘要中未提及，因此未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CuAsmRL在实验中表现出色，能够透明地提升现有专用CUDA内核的性能，最高提升幅度达到26%，平均提升幅度为9%。这些结果表明，即使是经过人工优化的CUDA内核，仍然存在很大的优化空间。CuAsmRL能够自动发现这些优化机会，并显著提高性能。这些实验结果验证了CuAsmRL的有效性和实用性。

🎯 应用场景

CuAsmRL具有广泛的应用前景，可以应用于各种需要高性能计算的领域，例如深度学习、科学计算、图形渲染等。它可以作为编译器优化的一部分，自动优化生成的CUDA内核，提高GPU的利用率和性能。此外，CuAsmRL还可以作为一个分析工具，帮助开发人员理解GPU的底层行为，并发现潜在的优化机会。该研究有望推动GPU自动优化技术的发展，降低高性能计算的开发成本。

📄 摘要（原文）

Large language models (LLMs) are remarked by their substantial computational requirements. To mitigate the cost, researchers develop specialized CUDA kernels, which often fuse several tensor operations to maximize the utilization of GPUs as much as possible. However, those specialized kernels may still leave performance on the table as CUDA assembly experts show that manual optimization of GPU SASS schedules can lead to better performance, and trial-and-error is largely employed to manually find the best GPU SASS schedules. In this work, we employ an automatic approach to optimize GPU SASS schedules, which thus can be integrated into existing compiler frameworks. The key to automatic optimization is training an RL agent to mimic how human experts perform manual scheduling. To this end, we formulate an assembly game, where RL agents can play to find the best GPU SASS schedules. The assembly game starts from a \textit{-O3} optimized SASS schedule, and the RL agents can iteratively apply actions to mutate the current schedules. Positive rewards are generated if the mutated schedules get higher throughput by executing on GPUs. Experiments show that CuAsmRL can further improve the performance of existing specialized CUDA kernels transparently by up to $26\%$, and on average $9\%$. Moreover, it is used as a tool to reveal potential optimization moves learned automatically.