第三章：静态分析与计时
前两章在研究每种操作的同时，已经用设备上的周期计数器给它们计时：第一章第 3 节介绍了 LCC 的读法和 TC VMEM 计时边界，此后每一节都给出了周期数。那里只需要会用这套测法。本章回答它背后的问题：这样测为什么是对的，主机上看到的时间与设备上的周期数是什么关系，怎样把各个单元的参数组合成一个可以从清单推算周期数的模型，以及不同 TensorCore、不同芯片上的时间怎样比较。
时间可以在三处测：主机的时钟，以及设备上的两个计数器。
- 主机的时钟测到的是调用方等待的时间，包含提交、启动和完成通知，与 kernel 本身关系不大。
- LCC（local cycle counter）是每个 TensorCore 自己的周期计数器，每个周期加 1。在一个 TensorCore 上数一段代码要多少周期，用它：读数就是周期，可以放在程序的任何位置。本教程中的周期数都是这样测的。
- GTC（global time counter）是整个切片共用的时间轴，各芯片通过一棵时钟同步树对齐到同一个时间源。比较不同 TensorCore、不同芯片上的事件，或者把一段时间换成秒，用它。JAX 自带的 profiler XProf 的设备时间就来自 GTC。
两个计数器不能互相替代：LCC 的读数只在同一个 TensorCore 之内可比；GTC 的分辨率是 1.5 个周期，在时间源之外的芯片上还会被不断校正。
本章从最粗的主机计时开始，说明它包含什么、为什么不够；再说明比较两个实现时边界怎样固定；然后给出两级发射的确切规则，把第一章测得的参数整理成发射模型；接着介绍 GTC 和时钟同步树；最后介绍建立在 GTC 上的 XProf：它的事件怎样产生、怎样解析，以及哪些信息只有它能给出。
读完本章，应当能够：
- 说清一个“时间”从哪里开始、到哪里结束，包含哪些与 kernel 无关的开销；
- 公平地比较两个实现，并从编译结果判断两边是否做了同一件事；说清 TC VMEM 计时边界的外壳各部分防住了什么；
- 理解 TensorCore 的两级发射，判断何时需要
sfence，知道读数本身怎样干扰被测的程序； - 判断一段代码受哪一类资源限制，从清单推算它的周期数，并用 LCC 验证；模型与实测不符时，逐个 bundle 找出缺少的规则；
- 说清 GTC 的读数由哪两部分组成，时间源与跟随者读到的 GTC 有什么不同，并用 GTC 比较不同 TensorCore、不同芯片上的事件；
- 读懂 XProf 的事件来自哪条指令、时间来自哪个计数器，需要时自己写出这条指令。
本章目录
优化一个 kernel 的顺序
前言说过，TensorCore 按编译器排好的指令包顺序执行，一段代码要多少周期可以从清单算出来。把前面各章连起来，优化一个 kernel 的工作顺序是：
- 先固定边界。 比较的是完整的程序（HBM 到 HBM），还是数据已经在片上之后的一段计算（TC VMEM 到 TC VMEM）；两边的输入、结果和布局是否相同（第一章第 3 节、本章第 2 节）。边界没有固定，周期数再精确也不能比较。
- 读两边的清单。 原生 XLA 是底线：先弄清它生成了哪些指令、数据走了哪些通路，再看自己的 kernel（第一章第 2 节）。
- 判断受哪一类资源限制。 是提交—取回的往返延迟、单元的发射间隔，还是指令的条数与内存的等待（第一章第 15 节、本章第 4 节）。这决定了该换算法、该重排，还是该减少指令。
- 只改一个量，在设备上读周期数。 区间两端各读一次 LCC，中间不放读数（第一章第 3 节、本章第 3 节）。区间内只有指令时读数每次相同，差一个周期也是确定的差别；有 DMA 时报告中位数和范围。需要只变一条指令、一个顺序时，用 tpuasm 直接改写已编译的程序，不要改上层的源码：源码的任何改动都会让编译器把别的指令也重排一遍，前后两个程序相差的就不止一个量。
- 用模型解释周期数。 实测与模型一致，说明时间花在哪里已经清楚；不一致时逐个 bundle 定位，补上规则（第 4 节）。
- 在完整的程序上核对结果。 一段手写的片段在载体里正确，放进完整的程序仍可能出错：载体可能恰好在某个寄存器里留着片段需要的值，掩盖了漏写的依赖。小形状上正确也不够，寄存器紧张、编译器开始复用和溢出寄存器之后才出现的错误，要用更大的输入才能发现。改写前后的程序都用同样的输入核对全部输出。
这个顺序里没有“先采集一段时间线，再看哪里慢”。XProf 只在 LCC 做不到时使用（第 6 节）：主机一侧的事件，或者程序不能用 tpuasm 改写的场合。
本章的工具
计时用 tpuasm_tools.py 中的工具，都在第一章第 3 节介绍。LccProbe 用 Pallas 编译一个载体 kernel，再用 tpuasm 在其中插入手写的 bundle，读出 LCC 或 GTC；插入的片段不经过编译器，设备执行的就是写下的指令。KernelClock 给任意已编译的程序计时：在指定的位置插入 LCC 读数，读数留在 SMEM 中，程序运行之后取回。vmem_shell、time_in_vmem 把输入和结果固定在 TC VMEM，用到 JAX 的内部接口 jax._src.state.primitives（pin、unpin）和 jax.experimental.layout（第 2 节）。第 4 节的发射模型在仓库根目录的 issue_model.py 中。第 5 节用同一个 LccProbe 读 GTC，并从主机只读地读回同步树的角色寄存器。第 6 节用 xprof_tools.py 采集并解析 XProf 的 XPlane 文件。
本章得到的参数
各个硬件单元的延迟与发射间隔在第一章测得，汇总在第一章的首页和本章第 4 节。本章得到的是与计时方法、发射规则和时钟有关的参数：
| 参数 | 值 | 位置 |
|---|---|---|
| 主机调用并等到结果，比设备时间多出的部分 | 约 130–136 µs | 第 1 节 |
| TensorCore 0 在 kernel 之前的 runtime 代码 | 2–5 万个周期，每次不同 | 第 1 节 |
KernelClock 一次读数的开销 |
20 个周期（两次读数紧挨时）；按此扣除后，n 个互不等待的 bundle 读作 n − 2 | 第 1、3 节 |
| 同一份程序、区间内只有指令时的读数 | 每次运行相同；区间内有 DMA 时有几个周期的波动 | 第一章第 3 节，本章第 2 节 |
sfence 的开销（VIF 为空时） |
11 个周期 | 第一章第 3 节，本章第 3 节 |
| VIF 容量 / 释放延迟 | 20 项 / 向量发射后 10 个周期 | 第 3、4 节 |
CMEM cld → vpop |
53 个周期；串行每组 54 个周期 | 第 3、4 节 |
TC VMEM 写后读：vst → 读同一批行的 vld |
7 个周期 | 第 4 节 |
| LCC 的频率 | 1050 MHz（时间轴上的包络 1049.98–1050.01 MHz），一个周期 0.952 ns | 第 5 节 |
| GTC 的结构 | 高 60 位是时间计数，每步 1/0.7 ns（1.5 个周期）；低 4 位是本地序号，到 15 饱和 | 第 5 节 |
| GTC 与 LCC 的比例 | 时间源上恰好 32 : 3，每周期的增量按 1、15、16 循环；跟随者上偏离百万分之几，并有成批的校正 | 第 5 节 |
| GTC 的对齐精度 | 同一颗芯片的两个 TensorCore 读到同一个值；跨芯片能证明的界是 ±0.57 µs（相邻）到 ±1.4 µs（最远），按对称时延估计落后时间源 1–32 ns，随树上的深度增加 | 第 5 节 |
| XProf 设备时间的分辨率 | GTC 高位的一步，1.43 ns | 第 6 节 |