KPerfIR: Towards an Open and Compiler-centric Ecosystem for GPU Kernel Performance Tooling on Modern AI Workloads

Yue Guan1†, Yuanwei Fang2‡, Keren Zhou3,4, Corbin Robeck2‡, Manman Ren2‡, Zhongkai Yu1†, Yufei Ding1,2†, Adnan Aziz2‡ 1University of California, San Diego, 2Meta, 3George Mason University, 4OpenAI †{yueguan, zhy055, yufeiding}@ucsd.edu ‡{fywkevin, robeck, mren, adnanaziz}@meta.com kzhou6@gmu.edu

A1 主要贡献

本工作提出KPerfIR，一种新型的多级编译器中心基础设施，用于支持在现代GPU上针对现代人工智能（AI）工作负载的自定义、可扩展和可移植的性能分析工具开发。核心问题在于现有编译器难以超越手工调优实现（如cuBLAS、rocBLAS和CUTLASS），这是由于GPU架构的快速演进和新算子变体的不断出现，例如Nvidia Hopper架构的第5代Tensor Cores和Tensor Memory Accelerators，以及Flash-Attention-3（FA3）内核的复杂平铺和流水线技术。研究目标是桥接编译器和分析器之间的差距，通过将性能分析功能集成到编译器工作流中，实现作为编译器通路的分析功能，提供可编程和可重用的性能分析框架，从而获得对复杂优化挑战的细粒度洞察，如GPU上细粒度功能单元的执行重叠。创新点包括：将多级编译器IR仪器化融入分析工作流，使性能工具作为编译器通路组成；集成到Triton基础设施中，支持AMD和Nvidia平台；开发首个基于区域的GPU定时工具，提供内核内行为的洞察，并通过IR语义恢复仪器化引起的 inaccuracy。评估显示工具开销低（8.2%），测量准确（2%相对误差），并为复杂GPU内核内优化提供可行动洞察。

关键贡献总结如下：
- 提出KPerfIR，一种编译器中心基础设施，支持自定义性能工具，集成到Triton编译器中，支持AMD和Nvidia平台，为优化AI编译器和算子的开发者提供关键支持（第4节）。
- 作为基础设施的演示，开发新型基于区域的定时工具，体现编译器中心方法，提供细粒度GPU分析能力（第5节）。
- 使用拟议定时工具进行深入案例研究，分析GPU重叠，揭示内核内重叠技术（如warp specialization）的宝贵洞察（第6节）。

A3 背景知识/关键Observation/设计原则

GPU编译器背景。 AI编译器的设计已显著进步，MLIR的引入标志着向模块化、可扩展性和多层抽象的关键步骤。MLIR通过可重用框架简化现代计算工作负载的复杂性，降低编译器开发的入门门槛，同时支持自定义组件和第三方工具集成以驱动创新。在MLIR基础上，Triton是一种突破性编译器，旨在桥接AI工作负载和GPU硬件。通过专为GPU编程定制的方言，Triton提供高层抽象，简化高效GPU内核的开发，满足计算密集型AI应用的需求。Triton利用MLIR的模块化架构实现可重用性、无缝集成和多级优化，示例说明如何在统一框架内处理领域特定需求。Triton的中间表示（IR）涵盖各种方言，包括TTIR、TTGIR、Triton、TritonGPU、TritonNvidiaGPU和TritonAMDGPU，每种针对GPU编程的具体阶段。

GPU分析器背景。 GPU性能工具对于分析工作负载执行特征和识别性能瓶颈至关重要。这些工具使开发者能够理解应用行为和硬件利用率，这对于在GPU平台上实现最优性能至关重要。GPU性能工具可分为两大类。现成分析器包括NVIDIA的Nsight Compute（NCU）和Nsight Systems（NSys），以及AMD的RocTracer和RocProfiler。NCU为CUDA应用提供详细的内核级分析，提供各种指标和可行动优化建议。NSys提供系统级性能的全景视图，捕获CPU和GPU之间的交互以诊断延迟和同步问题。类似地，RocTracer支持AMD GPU的性能优化，通过在异构环境中跟踪和分析应用。一个重要方面是定时启动和停止调用（如cudaProfilerStart和cudaProfilerStop），这些作为主机侧API实现，围绕内核启动调用，而非直接在内核代码中。

MosaicGPU分析器通过使用高层Python绑定插入PTX指令，代表与我们方法最接近的想法。然而，它在汇编代码级别操作，而非高层IR，限制其提供全面和可重用分析能力。ThunderKitten（TK）是一种有前景的Nvidia平台DSL，最近开发了自定义定时接口用于基于区域的跟踪。然而，此工具与DSL设计耦合，仅支持Nvidia平台。分析基础设施如CUPTI、NVBit、RocProfiler提供一定程度的仪器化以进行自定义。这些允许开发者在运行时插入自定义指令，提供灵活的分析选项。然而，它们与程序和仪器化框架紧密耦合。这种方法缺乏可移植性，无法无缝集成到编译器工作流中，限制其支持跨平台开发的能力。

表1: GPU性能工具比较

GPU重叠技术背景。 为了有效利用GPU设备的资源并释放其全部潜力，GPU内核必须安排多样执行和内存单元以重叠其执行，实现更好并发。随着专用片上加速单元的引入，如Tensor Core和Tensor Memory Accelerator，它们的重叠变得更加重要和困难。波优先级和指令调度是最直接的重叠解决方案，通过显式优先级设置扩展调度器。这作为波优先级技术实现，主要由AMD平台采用，通过编译器内在函数（如LLVM的__builtin_amdgcn_s_setprio()），在编译时控制波前在计算单元（CU）上的调度和执行。CU将在运行时选择优先级更高的波前，当调度冲突指令类型时。通过设置合适的波优先级，程序员可以显式重叠一个波上的内存加载和另一个波上的计算操作，以充分利用执行流水线。重叠可以通过编译器调度内在函数如__builtin_amdgcn_sched_barrier()进一步细化，以控制矩阵操作的聚类和流水线。

软件流水线（SWP）将独立循环迭代操作（即内存和计算）的执行转换为多个阶段，以在迭代之间重叠，如图2-(b)所示。通过使用异步指令进行数据加载和计算，不同操作可以并发执行以重叠延迟。然而，SWP需要在加载的内存层次中需要额外的暂存存储和寄存器来存储每个阶段的中间结果。软件流水线通常在多个级别采用，如全局内存到共享内存和共享内存到寄存器文件，以有效隐藏延迟。除了用于重叠数据加载和计算外，SWP还用于重叠由单独硬件单元执行的不同计算部分。VALU、VMEM和SMEM指令可以独立调度——目标是让这些完全并行于张量/矩阵核心操作。由于其提升利用率的关键优势，SWP已在高性能算子库和编译器优化中广泛接受【[22], ALCOP: Automatic LoadCompute Pipelining in Deep Learning Compiler for AIGPUs+2023+Proceedings of Machine Learning and Systems, volume 5】。

Warp Specialization（WS）是另一种基于生产者-消费者模型的重叠技术【[8], CudaDMA: optimizing GPU memory bandwidth via warp specialization+2011+Proceedings of 2011 International Conference for High Performance Computing, Networking, Storage and Analysis】【[23], Alcop: Automatic loadcompute pipelining in deep learning compiler for aigpus+2023+Proceedings of Machine Learning and Systems, volume 5】。在将执行分割成阶段后，WS将阶段分配给专用warp作为生产者和消费者。不同于在不同循环迭代阶段之间重叠，WS在不同warp的生产者和消费者阶段之间重叠，如图2-(c)所示。WS先前被提出作为不规则和大问题的解决方案，以重叠不平衡工作负载并减少寄存器压力【[17], WASP: Exploiting GPU Pipeline Parallelism with Hardware-Accelerated Automatic Warp Specialization+2024+2024 IEEE International Symposium on High-Performance Computer Architecture (HPCA)】。与SWP相比，其中每个warp充当生产者和消费者，WS分离角色以便正确分配寄存器。例如，数据加载阶段消耗的寄存器远少于重型计算阶段。然而，由于缺乏硬件warp级寄存器分配支持，此特性未被充分利用。最近，在Nvidia Hopper架构上【[15], Nvidia hopper h100 gpu: Scaling performance+2023+IEEE Micro】，新型特性如寄存器重新分配和异步事务屏障使WS实现优越的重叠性能【[42], Flashattention3: Fast and accurate attention with asynchrony and lowprecision+2024+arXiv preprint arXiv:2407.08608】【[45], CUTLASS+2023+January 2023】。

关键观察和动机。 随着异构执行单元复杂性和异步数据流编程范式的增加，直观理解执行流水线变得具有挑战性。用户可能选择SWP或WS来编排示例中的数据加载和计算。然而，为了最大化硬件利用率，他们必须仔细确定SWP中的阶段分区和阶段数量，或决定WS中的执行顺序和同步屏障。现有分析器提供聚合结果，缺乏将分析指标与程序语义相关联所需的程序语义。要真正理解重叠行为，必须解析循环结构、分析跨warp依赖，并跟踪执行全程的细粒度指标。用户难以有效识别和解决低效问题。关键 takeaway 1：性能分析工具需要编译器的IR来提供细粒度性能指标。获得执行行为的全面理解使手动优化成为可能。然而，在编译器通路内实现自动优化时，挑战更大。编译器需要与现有IR设计绑定的自定义性能反馈来确定最优变换。如示例所示，选择WS和SWP的适当重叠设计需要每个方法的稳定性能估计。编译器通路必须依赖精确分析数据来指导其优化选择。嵌入编译器基础设施的可编程性能工具提供了启用此工作流的关键第一步。关键 takeaway 2：编译器优化通路需要可编程性能分析工具来有效指导其优化决策。

A2 方法细节

编译器中心性能工具概述。 虽然它受益于编译器和内核开发者，但当前编译器设计缺乏对性能工具的支持。现有MLIR方言主要关注静态编译时变换和优化通路，几乎不考虑捕获动态跟踪或分析语义。这一差距使开发者缺乏原生框架来分析运行时行为或在编译器生态内微调内核性能。将性能工具基础设施设计为MLIR方言是一项非平凡任务。例如，在Triton中，此类方言必须桥接高层张量IR和低层目标特定IR之间的抽象差距，同时保持干净抽象以与现有多级程序IR语义一致和兼容。理想基础设施应能处理多样分析需求，并具有针对不断演进GPU架构的极大可扩展性。为解决此复杂性，我们引入KPerfIR方言，这是首创，将仪器化通路直接集成到编译器IR中作为性能工具基础设施。该系统构建在主流AI编译器Triton之上【[47], Triton: an intermediate language and compiler for tiled neural network computations+2019+Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】，它采用MLIR作为其编译器基础设施。我们演示KPerfIR基础设施的多级IR设计和运行时约定，如图4所示。我们进一步讨论KPerfIR的几个代表性用例。

IR设计。 在Triton的多级IR结构中，TTIR和TTGIR是开发者交互的主要程序表示。我们将KPerfIR实现为多层MLIR方言，与Triton的IR结构交互，并将高层MLIR操作链接到LLVM-IR级分析。如图4所示，我们在MLIR方言中插入记录操作，并将它们降低到架构特定的TritonGPU操作。我们还在LLVM IR级别降低和仪器化一些脚手架函数，如从本地缓冲区到主机内存的数据拷贝。此架构将编程语言结构链接到硬件相关特征以进行分析。具体而言，我们在不同级别定义新型操作来仪器化分析语义并收集测量数据。重写的Triton程序然后降低到LLVM IR以进行进一步分析和代码生成。此设计通过高层IR设计和低层代码生成平衡可编程性和复杂性。仪器化可以在MLIR或LLVM IR级别进行，作为灵活性和通用性的权衡。MLIR级别仪器化可以访问高层结构，如循环和数据对象（例如张量、矩阵等）。然而，在编译器通路管道中插入仪器化指令越高，准确性和可靠性越受编译器后端指令调度和重排序的影响。相反，LLVM IR更接近低层汇编代码，并提供更多关于GPU设备的信息。然而，LLVM IR级别的仪器化可能丢失与高层程序语义的连接。因此，我们通过将高层MLIR方言链接到LLVM IR级分析来设计编译器中心性能工具，如表2所示。我们在TTIR和TTGIR级别提供分析仪器化接口，因为它们是隐藏低层供应商特定细节并捕获关键优化（如软件流水线、代码移动、循环融合和warp specialization）的通用表示。KPerfIR是我们最高抽象级别，包括关键RecordOp操作符，代表一般程序标记，其语义解释完全依赖于KPerfIR到KPerfGPUIR降低通路的配置。RecordOp的输入包括name和isStart，指定程序中的注释位置和标识符。在程序中使用RecordOp的示例如下。

表2: KPerfIR的主要操作

KPerfIR级别的RecordOp抽象掉硬件细节，并降低到GPU特定的KPerfGPUIR操作。在KPerfGPUIR级别，我们获得供应商独立但GPU特定的硬件特征（例如共享内存抽象）。在降低到KPerfGPUIR时，给出各种MLIR通路选项来确定转换，如MetricType（指定性能计数器）和Granularity（例如warp-group、warp、thread等）。例如，在分析GPU周期时，每个RecordOp降低到ReadCounterOp和StoreCounterOp，这些可以由编译器分别调度到适当的程序位置。ReadCounterOp收集精确的GPU周期寄存器值。StoreCounterOp将计数器值存储到具有偏移的分析缓冲区中。相应的缓冲区和簿记资源在降低通路中根据可配置约束确定和分配。表2概述了KPerfIR和KPerfGPUIR中具有输入、输出和属性的关键操作符。仪器化的具体操作由降低/转换通路中的各种MLIR通路选项控制，包括BufferType、BufferStrategy、MetricType、Granularity和资源约束（例如缓冲区大小）。给定这些控制旋钮，编译器生成特定KPerfGPUIR操作，用于由RecordOp标记指定分析位置的性能工具。例如，指定BufferStrategy为circular、MetricType为clock和Granularity为warp-group，导致内核内分析器使用循环缓冲区存储每个warp-group的时钟周期（参见第5节对分析器的详细讨论）。在这种情况下，StoreCounterOp操作转换为CircularStoreOp来处理循环缓冲区事件记录。此外，编译器还在KPerfGPUIR中生成脚手架操作来设置和清理分析阶段的环境。InitOp初始化分析状态（例如缓冲区索引）。我们使用栈分配来启用LLVM工具链对缓冲区索引应用寄存器提升，该索引位于时钟测量的关键路径中。此操作返回分配索引的地址，并且内存加载/存储操作将被LLVM后端的优化通过将索引值保持在寄存器中而优化掉。FinalizeOp执行清理并写回分析记录。类似地，精确缓冲区内存分配操作（LocalAllocOp、GlobalScratchAllocOp和StackAllocOp）由从KPerfIR到KPerfGPUIR降低期间的配置确定。最后，我们有LLVM级别标记startInstrumentationOp和stopInstrumentationOp来控制低层基于库的仪器化。startInstrumentationOp告诉LLVM级别插入起始内存分析/时间戳（包括warp ID、SM ID）。stopInstrumentationOp插入结束内存分析/时间戳并从本地缓冲区冲刷数据到暂存缓冲区/主机固定内存。这允许低层IR级别的仪器化链接到高层数据对象和流结构（例如在仅存在于高层IR的数据对象和流结构周围的LLVM IR级别插入时间戳）。将所有仪器化包装在MLIR级别使派生的工具无法扩展到ML框架，由于它们在支持的操作集上的分歧。相反，在LLVM IR级别插入指令允许灵活和可重用工具。分析代码可以用任何语言编写，编译到LLVM字节码，并插入到MLIR框架代码中。此特性特别吸引那些不链接到Clang或其库但包括C++仪器化函数的MLIR框架。换言之，分析函数可以用HIP/CUDA编写，并在没有C++代码的MLIR框架中使用（如Triton、PyTorch）。在内核函数入口处，我们分配分析缓冲区、全局暂存内存和簿记资源。我们推迟第5.2节中内存管理的讨论，并在这里专注于降低通路的解释。分配的分析缓冲区与每个ReadCounterOp关联，以收集硬件性能计数器（例如GPU周期）。收集的原始数据由FinalizeOp以预定义内存布局写回到GPU全局内存中，FinalizeOp插入在内核函数末尾。FinalizeOp的输入包括分析记录的数量、分析缓冲区的基地址和分析器全局内存的基地址。Triton将TTGIR降低到LLVM后端以生成供应商特定代码（例如ptx和amdgcn）并执行低层优化，如冗余代码消除【[18], Eliminating redundant object code+1982+Proceedings of the 9th ACM SIGPLAN-SIGACT Symposium on Principles of Programming Languages, POPL ’82】和寄存器提升【[28], Register promotion in c programs+1997+SIGPLAN Not.】。我们开发一套LLVM转换模式来处理仪器化分析操作的代码生成。具体而言，InitOp降低到缓冲区索引的栈分配（llvm.alloca）。ReadCounterOp降低到性能计数器（例如%clock）的读取到寄存器，StoreCounterOp降低到带有标签创建和缓冲区索引管理的寄存器值存储。对于FinalizeOp，我们计算当前线程块的全局内存偏移，并将第一个线程分配为工作者来将整个分析缓冲区写回到全局内存。我们提供关于后端的主要支持，以减少高层仪器化在优化如指令重排序中的干扰。对于Nvidia GPU，影响微不足道，因为硬件负责指令调度，PTX大多遵循仪器化程序与KPerfIR的操作符关联。指令不会围绕关键指令如WGMMA调度。AMD GPU将指令调度暴露给软件，甚至对于amdgcm中的SMEM加载和MFMA指令，使仪器化操作符成为关键因素。因此，我们提供用户可调整的三级配置：1. 使用KPerfIR提示的手动调整；2. 在amdgcn上的直接仪器化；3. 使用屏障掩码显式指定指令调度窗口。

编译器通路。 我们然后利用编译器通路将RecordOp仪器化到目标程序并将其降低到底层IR。有两个主要级别的编译器通路处理它们的插入和转换，包括将KPerfIR降低到KPerfGPUIR和将KPerfGPUIR降低到LLVM。KPerfIR到KPerfGPUIR降低通路将RecordOp转换为ReadCounterOp和StoreCounterOp，并插入资源分配（如LocalAlloc）和设置/清理操作（如InitOp和FinalizeOp）的操作。

与第三方工具协作。 我们然后演示用户开发第三方工具的接口，如图4右侧所示，包括处理重写的仪器化API和必要的内存管理来解码分析的原始数据。仪器化API。要使用KPerfIR方言构建自定义性能工具，用户必须显式修改内核签名并使用仪器化通路更新IR。我们为用户提供两个API来与KPerfIR交互，命令行API和Python API。使用命令行API，在运行时遇到的所有Triton函数都使用指定的分析通路仪器化。这允许方便操纵目标工作负载，但缺乏灵活性，因为用户无法跳过他们不感兴趣的内核。Python API让用户指定仪器化入口和要仪器化的目标函数。这里的仪器化入口指编译过程中的插入位置（即在哪个通路之前/之后）。另一方面，由RecordOp指定的仪器化点指IR中的位置。它包含核心API作为KPerfIR.patch(instrumentation_obj, fn)，其中instrumentation_obj标识要仪器化的编译器通路（之前或之后），fn是可选的以选择仪器化内核。如果instrumentation_obj指定MLIR仪器化，我们采用MLIR的通路基础设施PassManagement。我们将在编译通路结束时仪器化LLVM。此外，我们提供KPerfIR.unpatch() API来恢复所有现有仪器化。这涉及在KPerfIR运行时中维护内核的原始和仪器化版本。运行时内存管理。第三方工具的另一个重要部分涉及缓冲区存储和分析结果的管理。在顶层，我们使用额外最后一个参数作为设备缓冲区指针重写内核签名并修改调用约定。累积的分析数据然后返回到由KPerfIR运行时管理的主机缓冲区。这里，我们采用基于丢弃的循环缓冲区设计，将在第5.2节中与实际工具设计一起阐述。我们有两种管理GPU设备缓冲区和主机缓冲区通信的方式，遵循先前研究【[53], Gvprof: A value profiler for gpu-based clusters+2020+SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】【[54], Valueexpert: Exploring value patterns in gpu-accelerated applications+2022+Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems】。第一种方法是使用CUDA管理分配，以便设备缓冲区分配在GPU和CPU可见的内存空间中。第二种方法是分配优先级高于PyTorch运行时使用的流，并使用该流将数据拷贝回CPU。一旦数据拷贝回CPU，它被解码成类似于CUPTI Activity API的C/C++结构。第三方工具将在运行时初始化时注册回调来处理此数据。回调在主机缓冲区解码后触发，以允许这些工具处理数据，如CUPTI活动。因此，工具可以将数据后处理成与它们自己的前端可视化器兼容的格式，如Chrome Trace【[19], Chrome trace format+2023】或Hatchet【[9], Hatchet: Pruning the overgrowth in parallel profiles+2019+Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】。

用例。 我们然后展示由KPerfIR基础设施促成的几个新型性能工具，如图6所示。基于迭代的定时。理解warp组跨迭代边界的同步行为对于分析SWP至关重要。例如，观察特定迭代步骤的重叠需要将循环归纳变量i与定时记录中的版本参数关联。然而，传统分析工具缺乏确定迭代位置所需的循环语义，使理解资源跨迭代的精确重叠行为具有挑战性。通过将分析通路融入编译器，我们可以仪器化IR来收集时钟tick与循环迭代索引语义一起，启用后处理来构建细粒度时间线，如图6-(a)所示。虽然MLIR的Python绑定等工具可以实现此功能，但它们是临时性的并限制于简单循环结构，强调需要将编译器语义更深入集成到分析工作流中。关键路径分析。运行时反馈对于编译器通路特别宝贵，尤其是在涉及WS的场景中识别执行关键路径。例如，在FA3内核中，生产者warp组加载输入张量，而消费者warp组执行多个GEMM操作和softmax计算。关键路径——由数据移动和计算延迟等因素决定——随程序平铺和访问模式变化，使静态分析不足以安排。对于FA3案例，有三个常见关键路径由不同操作主导，如图6-(b)所示。通过将分析通路嵌入编译器，我们可以仪器化定时记录来捕获运行时阶段延迟，并动态识别关键路径。这使编译器能够策略性地插入异步屏障，减少等待气泡时间并改善资源利用率。此类优化通路突显了将性能工具与编译器通路集成的价值。程序相关内存分析。细粒度内存分析从编译器中心设计中显著受益，允许开发者将内存性能指标（如访问模式、热图和银行冲突）与高层程序对象相关联。以图6-(c)为例，分析特定数据对象如张量的内存访问模式需要解析程序IR来插入分析内在函数，将高层结构映射到低层硬件寄存器。此方法对于优化L2级调度特别关键，如AMD的芯片设计【[1], AMD CDNA 3 Architecture+2024】。分析warp级内存访问模式并将其与程序对象相关联，使能够设计优化的L2访问混洗策略。此类能力仅通过编译器中心分析实现，展示了将内存分析工具集成到编译器基础设施的必要性。

基于区域的定时工具概述。 在本节中，我们深入探讨构建在KPerfIR基础设施上的强大第三方性能工具，它进行基于区域的内核内定时，以协助理解和优化GPU硬件资源的重叠。要使用KPerfIR构建自定义工具，用户必须实现仪器化通路、分析数据管理和后处理过程，如第4.3节所述。这些与基于区域的定时工具一起阐述，其仪器化工作流（第5.1节）、内存系统（第5.2节）和跟踪重放后处理（第5.3节）。

工作流。 我们首先介绍定时工具的工作流，有两种接口，用户接口和编译器接口，如图7所示。根据目标程序如何解析和指定分析区域，调用不同接口来仪器化分析记录。KPerfIR然后编译和执行仪器化程序来收集原始性能计数器。最后，我们使用原始元数据重放跟踪并产生最终时间线跟踪。此跟踪输出并反馈到不同接口用于可视化或分析器引导的编译器通路。用户接口是开发者分析目标GPU程序感兴趣区域的基本用法。我们在Triton内核中提供PythonDSL绑定，并允许手动重写和覆盖转储的中层IR以关注区域。这很简单，因为开发者手动重写转储的中层IR以关注区域。这在开发高性能内核和调试效率问题时极其有用。KPerfIR将使用重写程序IR仪器化原始程序并产生输出跟踪。开发者可以使用可视化工具，如Chrome Tracer【[19], Chrome trace format+2023】，获得程序时间线的直观视图。编译器接口是编译器与分析器交互的方式。与用户接口相比，编译器通路负责解析中层IR并在关注区域插入分析记录。类似地，仪器化程序被编译和执行，分析跟踪反馈到编译器内的编译器通路。

分析数据管理。 我们然后介绍工具的具体管理和解码技术。容纳每个内存层次存储的系统如图8所示。我们自底向上演示数据结构和布局。Warp级区域数据结构。对于每个细粒度分析区域，插入两个分析记录，区域开始记录和区域结束记录。每个记录有8字节，包含4字节标签和4字节负载，如图9所示。标签包含1位标志作为START或END，和31位控制位，这些依赖于后端，将在后面解释。负载包含从硬件计数器捕获的32位时间戳指标（Nvidia的%clock周期和AMD的S_MEMTIME指令的LSB 32位）。此设计在空间容量和写入速度之间实现良好权衡。因为存储消耗低，并且存储采用向量化存储指令。我们使用32位时钟捕获当前记录的周期tick，这可能导致值溢出。我们在后处理过程中解决此问题，其中我们检测并抛出溢出异常。跟踪重放可以补偿时钟回绕溢出，只要每个迭代运行少于40亿周期（在1GHz下为4秒）。由于大多数执行花费在循环上，限制放松到循环迭代级别，对于大多数情况小于1毫秒，使32位周期数据成为安全选择。如果必要，KPerfIR可以通过在IR中添加专用操作符和属性轻松扩展到64位时钟。SM级数据布局。然后记录从本地寄存器文件使用存储指令拷贝到共享内存中的数据缓冲区。数据缓冲区由warp组分割成分析空间，具有非重叠记录槽，如图8所示。因此，编译时记录索引确定共享内存中的存储位置偏移。例如，如果我们分配64槽（0.5KB）并且每个线程块包含2个warp组（8个warp，256线程），则每个warp组有32槽。在编译时，我们预计算每个分析空间的基地址并生成代码来管理每个warp组的索引。每个warp组缓冲区的内存布局如图9所示。开始和结束记录在共享内存中交错，以支持嵌套区域和多迭代。图的下部说明三种分析模式：常见、嵌套和多迭代。记录存储时不配对，并在跟踪重放期间对齐以根据分析模式确保正确性。共享内存循环缓冲区。因为共享内存容量有限，它可能无法容纳所有分析数据，尤其是在长多迭代记录模式下。例如，使用2个warp组、4个分析区域和512迭代的循环，需要多达4096记录槽，转换为16KB共享内存存储。对于生产级内核，我们通常获得剩余的可用共享内存容量从1KB到4KB，用于内核内分析的数据缓冲区，仅占H100设备总228KB的1.75%【[15], Nvidia hopper h100 gpu: Scaling performance+2023+IEEE Micro】。我们提出每个warp组数据缓冲区的循环缓冲区结构，作为平衡准确性和可用性的解决方案。循环缓冲区在达到共享内存限制时仅循环保持跟踪的尾部记录。我们循环覆盖最旧结果，而不是将它们冲刷到低层内存。洞察是可视化几个最近迭代足以识别瓶颈。我们使用预计算的缓冲区容量回绕索引来实现循环缓冲区。如果捕获的事件溢出缓冲区，它被重定向到占用的槽。这导致运行时索引管理的额外指令，这在编译时解决。我们仅需要在运行时使用轻量级模指令来回绕索引。此外，此循环缓冲区是运行时的默认内存管理行为。用户可以切换到朴素冲刷策略，其中捕获的事件在缓冲区满时立即写回到全局内存。此策略可以保持所有分析事件，但开销是更频繁的内存写入指令。在KPerfIR基础设施中，我们提供全面操作来支持每个层次的内存分配，包括Stack、Shared、Global，以及管理策略，包括Circular和Flush。用户可以利用这些抽象来开发工具，结合缓冲区和管理策略的使用。具体而言，我们为AMD后端实现协作存储策略以最小化分析开销。由于每个warp或warp组仅分析一个时间戳记录，AMD GPU利用带线程掩码的分支来使用谓词发出指令。然而，此线程分歧可能导致意外的指令缓存未命中，导致高达600周期的开销。为了缓解此，我们强制所有线程写入相同的共享内存位置，仅保留最后一个结果。为了对齐记录，我们包括额外的12位签名，从HARDWARE_ID REG的LSB派生，如图9所示。此签名编码wave_slot_id、SIMD_id和pipe_id来注释分析线程。在warp组级别，我们确保所有四个warp写入相同的共享内存位置。全局程序跟踪。最后，我们在从内核函数返回之前，从所有SM收集分析数据回GPU全局内存。我们保持临时分析数据在共享内存中，并在内核结束时一次性拷贝整个数据缓冲区。附加几个额外元数据用于重建跟踪，如图8所示，包括线程块索引、warp级索引和每个warp组中记录槽的数量。Triton前端自动使用额外参数（profile_mem）修补内核函数，分配GPU暂存内存用于此类分析数据。

跟踪重放。 内核内分析的基于区域定时场景中的主要挑战是由分析指令引起的扰动。后处理技术，如周期采样【[21], Profile inference revisited+2022+Proc. ACM Program. Lang.】，通常用于缓解不准确性，通过周期采样性能指标并随后总结结果以减少与数据收集相关的开销。在基于区域的定时工具中，我们提出类似的后处理方法跟踪重放，根据KPerfIR的设计利用程序语义重新缩放分析区域。我们可以通过偏移常数来考虑时钟测量开销，从而获得准确的可视化结果。对于仅具有同步指令的区域，我们可以通过减去记录操作的开销轻松重建区域。对于异步指令，我们有两种常见不准确性，如图10所示。在香草执行时间线示例中，我们有一序列连续指令在张量核心上发出异步MMA操作后执行。在程序到达指令流中的屏障后，在功能单元产生结果之前有一段时间等待时间。为了理解这里的异步行为，用户可以在连续执行中插入分析记录以更好地对齐操作。然而，此类记录的引入导致不准确的等待时间测量。在减少气泡情况下，由于分析引起的额外开销，等待时间被低估。在意外空闲情况下，情况更糟，其中功能单元的执行时间无法覆盖分析开销。将观察到意外空闲时间，并且优化将被误导。为了解决此，我们提出在跟踪重放过程中扣除准确等待时间，如图10-(b)所示。我们不是在异步指令后放置一个END记录，而是插入两个START记录：在异步启动之前和等待屏障之后，以及一个END记录就在屏障之前。这保证准确等待时间Twait测量为Twait = (CLK2 − Ta) − (CLK1 −Ta) = CLK2 −CLK1，其中Ta和Tb分别是时钟读取指令之前和之后的记录区域中的经过时间。在后处理中，分析开销被抵消，使用两个仔细放置的记录。我们可以类似地推导Texe时间。此方法要求TMMA − Texe > Ta + Tb。由于分析开销对于大多数情况小于25周期，如第6.4节所示，功能单元的执行时间可以覆盖仅插入在这里的记录，这通常约为1000周期。

A4 实验环境

数据集名称、规模及用途：基准测试使用关键AI工作负载，包括GEMM和Triton的实验Flash Attention算子。这些算子是流行AI模型的基础，代表训练和推理任务中的大部分工作负载。选择提供最先进性能的主流实现，并进行敏感性实验以评估方法的有效性和可扩展性。具体配置包括头维度128、16头、批大小16、序列长度4096，用于FA3内核基准。

模型架构关键参数：GEMM算子分为2或3阶段SWP设计，标记为GEMM-SWP-2/3；FA3算子使用vanilla重叠和改进版本，标记为FA3-WS-a/b。FA3涉及多个GEMM、softmax和加载阶段，关键路径包括4个GEMM和2个加载阶段。

硬件配置：服务器配备最先进GPU卡，包括NVIDIA H100-HBM3（GPU型号/数量未指定，但用于Nvidia平台）和AMD MI300X（GPU型号/数量未指定，但用于AMD平台）。连接关系未详细说明，但评估针对现代GPU架构。

软件配置：基于Triton 3.0.0（代码实现/语言/依赖代码库）和LLVM 19.1，确保与最新编译器技术和GPU架构兼容。操作系统未指定。

A4 实验结果

实验内容、结果和分析：评估KPerfIR在NVIDIA H100-HBM3和AMD MI300X上的性能，使用GEMM和Flash Attention算子基准。案例研究针对FA3内核，使用基于区域的定时工具识别vanilla实现中的空闲气泡区域，并提取关键优化洞察。实施新型编译器通路以改进重叠策略，减少瓶颈。性能模型基于分析结果评估优化重叠设计的效率。对于头维度128、16头、批大小16、序列长度4096的注意力内核，2阶段SWP计算为467.07 TFLOPs，vanilla Triton FA3达到526.97 TFLOPs，优化后达到582.44 TFLOPs预测。优化FA3内核实现24.1%改进（图12），优于最佳手动FA3内核7.6%【[41], Flashattention3: Fast and accurate attention with asynchrony and lowprecision+2024+The Thirty-eighth Annual Conference on Neural Information Processing Systems】。分析显示关键路径包括区域22、12、25、15、6和3，通过推进区域16的到达屏障释放GEMM1区域，实现更紧凑时间线（图11）。

表3: FA3内核的分析区域（出处：表3）

表4: 性能模型（出处：表4）

分析开销实验：测量仪器化内核的端到端延迟，并归一化到原始执行时间。对于大多数情况，开销低于10%；对于3阶段SWP GEMM，开销在15%以内。平均8%开销，通过后处理跟踪重放缓解干扰。共享内存使用循环缓冲区设计，适应可用空间（1KB-4KB），对于3阶段SWP GEMM，剩余10.9KB未用空间，可跟踪16迭代（图13、14）。

低层深入分析：在SASS ISA级别，每个KPerfIR记录节点降低到三条指令：时钟读取、整数移动和谓词存储，平均开销33周期（图15）。对于循环定时，循环入口生成五条额外指令。性能影响在2%以内，使用模型$T_{instrumented} = T_{vanilla} + N_{record} \times cycles$评估（表5）。

表5: 性能退化评估（出处：表5）

A5 结论

随着AI编译器和GPU架构的持续演进，性能工具必须推进以满足现代编译器的增长需求。为解决此需求，我们引入KPerfIR，一种构建在Triton编译器上的新型编译器中心分析基础设施。KPerfIR桥接编译器和分析器设计之间的差距，支持可编程、可重用工具的开发，并解锁性能分析和编译器优化的新可能性。我们将KPerfIR视为开放、编译器中心生态的基础，赋能社区构建多样和创新性能工具，以适应AI工作负载的动态需求。未来工作包括扩展到HPC和科学模拟领域，以及支持分布式GPU工作负载，包括融合计算-通信内核。

A6 附录

摘要。KPerfIR是Triton编译器的性能工具基础设施【[47], Triton: an intermediate language and compiler for tiled neural network computations+2019+Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】，OSDI’25提交的结果源自一些特性分支。虽然核心概念稳定，但实现仍在演进并可能变化。正式文档可在https://triton-lang.org/main/dialects/ProtonOps.html找到。

范围。我们展示性能工具的可用性和论文中讨论的改进结果，通过重现图11和图12。具体而言，使用基于区域的性能工具分析FA3内核演示编译器中心设计的使用。vanilla和改进FA3内核的分析跟踪比较显示内核内区域分析的方法论。评估结果显示实际性能改进。

内容。我们提供包含Triton编译器和评估脚本的Docker镜像用于artifact评估。请遵循以下部分的安装说明设置环境。在Docker镜像中，内容组织如下：workspace/ triton % Triton编译器源代码 tritonbench % Triton FA3基准套件 – kperfir_artifact % 其他脚本和文件

托管。此artifact在https://github.com/ChandlerGuan/kperfir_artifact开源，主分支和提交版本e45891d。

要求。Docker环境需要安装NVIDIA Container Toolkit的基于Linux系统。虽然KPerfIR项目设计为跨平台，但FA3内核的artifact评估需要Nvidia H100 GPU。