Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs

GPU架构的复杂性持续增加，近年来的新架构引入了功能日益强大的专用硬件单元，用于矩阵乘法和数据搬运，同时辅以高度并行的通用计算核心。为了充分释放这些硬件的性能，软件必须采用复杂的调度策略，以最大限度地利用所有硬件资源。由于实现此类调度十分复杂，程序员和编译器通常依赖程序变换技术，例如软件流水（SWP）和线程束专精（WS），在实践中达成这一目标。然而，如何最佳地组合使用SWP与WS仍是一个难题，目前主要依靠脆弱的编译启发式规则和易出错的人工经验来解决，对可能的解决方案空间缺乏系统性的理解。为改善这一状况，我们提出了一种将SWP与WS联合建模为统一优化问题的新方法，该问题可由现成的约束求解器进行整体求解。我们基于此方法实现了Twill系统，这是首个能够自动为一大类迭代程序生成最优SWP和WS调度的系统。Twill不依赖任何启发式规则，易于扩展至新型GPU架构，并能保证生成的调度方案是最优的。实验表明，Twill能够重新发现专家为NVIDIA Hopper和Blackwell GPU架构上的Flash Attention手动设计的SWP与WS调度方案，从而证明了这些人工调度方案的最优性。

论文试图解决如何在现代复杂的GPU架构上自动发现最优的软件流水（SWP）和线程束专用化（WS）调度策略的问题。随着GPU引入更多固定功能单元，最大化硬件资源利用率变得极为复杂，当前依赖脆弱的编译启发式和人工经验来组合SWP与WS，缺乏系统性探索调度空间的方法。这是一个尚未被充分解决的新问题，尤其是在联合优化SWP与WS方面。

论文提出将SWP和WS的调度问题形式化为一个统一的约束优化问题，利用现成的约束求解器进行全局求解，从而实现全自动、最优的调度生成。关键创新在于对两种程序变换进行了联合建模，摆脱了传统启发式方法的限制。这是首次将两者视为整体优化目标的工作，具有较强的理论清晰性和实用性。

作者实现了名为Twill的系统，据称是首个能自动推导出最优SWP+WS调度的框架。Twill无需任何启发式规则，可扩展至新GPU架构，并保证最优性。实验部分成功复现了专家手工为Flash Attention在NVIDIA Hopper和Blackwell架构上设计的调度方案，并证明其最优性。代码已开源，具备较强可复现性。未来可进一步研究将其扩展到更广泛的算子类别和异构计算平台。

1. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning 2. Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines 3. TACO: A Tool for Compiling Tensor Contraction Expressions 4. FlexTensor: An Automatic Schedule Explorer for High-Performance Tensor Computation on FPGAs and GPUs 5. Ansor: Generating High-Performance Tensor Programs for Deep Learning