Google XProf 为 TPU 自定义内核引入寄存器级性能分析，实测缩短矩阵乘法执行时间

一句话结论

Google 在 XProf 中新增针对 Pallas 等自定义内核的周期级硬件计数器分析能力，实证可定位内存瓶颈并优化内核性能。

关键要点

• XProf 新增内核性能分析套件，支持在 TPU v7（Ironwood）上对 Pallas、Mosaic 或 Triton 编写的自定义内核进行周期级硬件性能计数器采样。
• 在分块矩阵乘法示例中，通过三重缓冲优化内存停顿，将内核执行时间从 125.5µs 缩短至 88µs，降幅约为 30%。
• 开发者可通过 `--xla_enable_custom_call_region_trace=true` 和 `--xla_xprof_register_llo_debug_info=true` 两个编译器标志启用自定义调用区域的 MLIR 检查。
• 系统支持最多四个 SparseCore 上配置最多 28 个计数器，形成 4×28 矩阵，采样分辨率下限为 1µs，可通过外部事件触发模式实现亚微秒级捕获。
• 新的 Perf Counters View 提供超过 16000 个原始硬件计数器的表格视图，其轨道高度反映时间窗口内计数器的最大值而非归一化百分比。

背景与事实

Google AI Infra 团队的 Yogesh SY 介绍，XProf 是面向 TPU 工作负载的开源性能分析器，也是 OpenXLA 项目的一部分并与 JAX 性能分析集成。此次更新的核心价值在于解决自定义编译路径导致的性能分析盲区。此前，使用 Pallas、Mosaic 或 Triton 创建的内核会跳过标准 XLA Pass，导致编译期静态成本模型产生偏差。静态工具可能错误地将某个 MXU 指令块标记为已充分利用，而实际上该单元正在空闲等待 HBM，因为静态分析忽略了时间因素。这种偏差使得“optimal FLOPs”和有效吞吐效率等指标在自定义内核场景中可能不准确甚至无法使用。

该套件分三个层级工作。第一层是编译器检查，开发者传入指定标志后，Graph Viewer 会显示“Custom Call Text”面板，展示每个自定义调用降低后的 MLIR，供工程师检查操作融合与内存分块结构。第二层是静态执行分析，Trace Viewer 展示用于低级别操作（LLO）的包数据，包括每个时钟周期执行的机器指令，并为 MXU、标量与向量 ALU、向量填充、加载、溢出、存储及跨通道单元（XLU）提供时间对齐的轨道。第三层是运行时遥测，XProf 会定期对硬件计数器进行采样，基础分辨率下限为 1µs。

在 TPU v7（Ironwood）上，该系统采用了一种新的外部事件触发模式，取消了 1µs 的采样限制。采样器会捕获 TPU 跟踪指令和边界触发器，包括自定义调用作用域的进入和退出，从而实现亚微秒级捕获并提高归因准确性。配置上，开发者可在最多四个 SparseCore 上，为每个核心配置最多 28 个计数器。采集通过 `jax.profiler.ProfileOptions` 启用，需设置 `tpu_enable_periodic_counter_sampling`、`tpu_tc_perf_counter_sampling_options` 并将 `is_external_trigger` 设为 true；若使用周期采样模式，则以 `interval_us` 替代触发器标志。

文章以受内存限制的矩阵乘法变体为例，发现 `sync_wait` 计数器出现大量峰值。通过实施三重缓冲，使 HBM 加载与 MXU 计算重叠，减少了停顿事件，执行时间从 125.5µs 降至 88µs。需要强调的是，该数据来自 Google 编写的单个演示内核，而非广泛基准测试，其意义在于验证分析流程的有效性，而非预示所有场景均可获得相同收益。

影响分析

对于中文开发者与 TPU 从业者而言，这一更新直接改变了自定义内核调优的方法论。分析判断认为，工程师不应再将 XLA 得出的效率数据视为自定义内核的可靠指标，而应转而采用寄存器级计数器作为优化依据。XLA 成本模型是为标准编译路径设计的，在面对 Pallas 等绕过标准 Pass 的自定义代码时，其静态估算存在固有盲区。开发者必须建立新的性能信任层级：直接从硬件寄存器读取的值（如 HBM 利用率和 TPO 指标）应作为事实依据，而编译期静态估算需保持谨慎。这意味着调优工作流将从“依赖工具报告”转向“主动配置计数器并解读原始硬件数据”，对工程师的底层硬件理解能力和性能分析技能提出了更高要求。

适用边界

该结论的适用范围有明确限制。首先，计数器采样功能的文档明确面向 TPU v7（Ironwood），使用早期 TPU 世代的团队不应假定其具备相同的计数器覆盖范围和功能支持。其次，4×28 的计数器预算是固定的，团队必须针对每次具体的性能调查预先选择最相关的计数器，无法同时监控所有指标。最后，文中的性能收益（30% 降幅）源自单一演示内核，直接将其推广到其他类型或规模的自定义内核是不成立的，不同工作负载的瓶颈分布和优化空间存在显著差异。

孤本观察

本文编辑观察认为，Google 将超过 16000 个原始计数器直接暴露给开发者，而非提供经过归一化的友好指标，这是一种典型的“信任开发者”设计哲学，既赋予了极大的诊断自由度，也相应转移了数据解读的认知负担给最终用户。

来源：InfoQ 中文 AI