自动优化GPU矩阵计算的流水线与线程组调度，提升性能并保证最优。
Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs
- 将流水线与线程组特化联合建模为约束优化问题，用求解器自动求解。
- 在Hopper和Blackwell架构上复现专家手动设计的Flash Attention调度，证明其最优性。
- 系统无启发式规则，可扩展至新GPU架构，适合高性能计算开发者使用。
GPU架构日益复杂，近年引入了强大的固定功能单元用于矩阵乘法和数据搬运，以配合高度并行的通用核心。要充分发挥这类硬件性能，软件需采用复杂的调度策略以最大化利用所有资源。由于实现这些调度非常困难，程序员和编译器常依赖软件流水线（SWP）和线程组特化（WS）等程序变换。然而，如何最优地结合使用SWP与WS仍是难题，目前主要依赖脆弱的编译启发式和不可靠的人类直觉，缺乏对解空间的深入理解。为此，我们首次将SWP与WS联合建模为一个可通过现成约束求解器求解的优化问题。我们实现了Twill系统，这是首个能为一大类迭代程序自动推导最优SWP与WS调度的工具。Twill无需启发式规则，易于扩展至新GPU架构，并保证生成最优调度。实验表明，Twill能复现并在NVIDIA Hopper和Blackwell架构上证明专家手动设计的Flash Attention调度的最优性。
原文摘要 · Abstract (English)
GPU architectures have continued to grow in complexity, with recent incarnations introducing increasingly powerful fixed-function units for matrix multiplication and data movement to accompany highly parallel general-purpose cores. To fully leverage these machines, software must use sophisticated schedules that maximally utilize all hardware resources. Since realizing such schedules is complex, both programmers and compilers routinely employ program transformations, such as software pipelining (SWP) and warp specialization (WS), to do so in practice. However, determining how best to use SWP and WS in combination is a challenging problem that is currently handled through a mix of brittle compilation heuristics and fallible human intuition, with little insight into the space of solutions. To remedy this situation, we introduce a novel formulation of SWP and WS as a joint optimization problem that can be solved holistically by off-the-shelf constraint solvers. We reify our approach in Twill, the first system that automatically derives optimal SWP and WS schedules for a large class of iterative programs. Twill is heuristic-free, easily extensible to new GPU architectures, and guaranteed to produce optimal schedules. We show that Twill can rediscover, and thereby prove optimal, the SWP and WS schedules manually developed by experts for Flash Attention on both the NVIDIA Hopper and Blackwell GPU architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。