发射模型与静态分析

第一章每认识一个硬件单元，就测出它的延迟和发射间隔；第 3 节给出了两级发射的规则。本节把它们整理成一个可以计算的模型：给定一段实际执行的 bundle 序列，逐个 bundle 推算它的标量发射和向量发射时刻，从而预测 LCC 读数。模型需要的参数只有几类：VIF 的容量和释放延迟、每种结果的可用延迟、每个单元的发射间隔。本节先把这些参数汇总起来，用模型预测前面各节全部手写片段的读数，与真机比较；然后说明模型与实测不符时怎样找出缺少的规则，以及在动手推算之前怎样先判断一段代码受哪一类资源限制。

模型

VIF 的两个参数来自 tpu-v4-latency-numbers 第 0 节，它用 vdelay 人为堵住向量发射，直接测出了 VIF 的容量与释放延迟；各个单元的参数来自第一章。规则如下：

1. 标量发射按顺序进行，没有阻塞时每周期一个 bundle。标量指令（包括读 LCC）在此执行。
2. VIF：含向量或 DMA 指令的 bundle 标量发射后进入 VIF。一项在向量发射后 10 个周期才在标量一侧释放；未释放的项达到 20 个时，任何 bundle 都不能标量发射。
3. 向量发射按顺序进行，每周期最多一个 bundle，最早在标量发射后 1 个周期；含访问内存的 vld、vst、cld 的 bundle 最早在 2 个周期后。此外要等计分板放行：源操作数就绪、对应的单元和队列允许；vld 要读的行刚被 vst 写过时，还要等到那条 vst 向量发射后 7 个周期。模型中的 ready、queues、last 三张表就是对计分板的模拟。
4. sfence 所在 bundle 在向量一侧占一个位置，下一 bundle 要等 VIF 中全部项释放之后才标量发射。

单元与队列的参数：

mrf2、mrf3 的结果比 mrf0、mrf1 晚 18 个周期，即 101（研究报告 55）。

模型写在仓库根目录的 issue_model.py 中，一百多行。核心是每个 bundle 的向量发射时刻：

issue = max(vector_free, scalar + 1 + any(decode(item)[0].startswith(MEMORY) for item in vector))
for item in vector:
    mnemonic, operands = decode(item)
    if mnemonic.startswith('vpop'):
        events.append(('pop', operands[1]))
        issue = max(issue, queues[operands[1]][0])          # 等结果进入队列
    elif operands and kind(operands[0]) in PUSH_LATENCY:
        events.append(('push', operands[0]))
    unit = next((prefix for prefix in UNIT_INTERVAL if mnemonic.startswith(prefix)), None)
    if unit:
        events.append(('unit', unit))
    issue = max([issue, *(ready.get(source, 0) for source in operands[1:])])   # 等源操作数
for event, queue in events:
    issue = max(issue, last.get(f'{event} {queue}', -99) + (UNIT_INTERVAL[queue] if event == 'unit' else INTERVAL.get(f'{event} {kind(queue)}', 1)))   # 等单元

TWO_CYCLES 列出结果 2 个周期后才可用的向量 ALU 指令，其余按 1 个周期处理。UNIT_INTERVAL 记录不经过队列、但同一单元相邻两条有最小间隔的指令，目前只有一项：第四章第 6 节测得的 vrng，每 8 个周期一条。

模型的输入是动态执行的 bundle 序列：循环要按实际迭代次数展开。模型不处理 DMA 和 vwait，它们的时间取决于数据通路，由第二章的开销模型给出。标量一侧加入了第一章第 8 节测得的 sld 规则：

if mnemonic == 'sld':
    scalar = max(scalar, last_sld + SLD_INTERVAL)       # 相邻两条 sld 至少相隔 4 个周期
# 读取尚未就绪的标量寄存器时，整个 bundle 等待。
scalar = max([scalar, *(scalar_ready.get(source, 0) for source in sources)])

sld 写入的寄存器在 4 个周期后就绪（SLD_LATENCY）；其余标量指令的结果都按下一周期可用处理。有一点与预期不同：紧接在 sld 之后的 vmov.8x128 v11, s24 读同一个寄存器，读数里看不到任何等待，模型不加等待时恰好与真机一致（下表的“sld 与依赖 sld 结果的向量运算”）。第一章第 8 节的实验确认，此时 vmov 读到的就是 sld 载入的新值，而不是旧值；vmov 读标量操作数的时刻本节没有进一步确定。

用模型预测，再与真机比较

本小节实验源码、输出。

脚本把每段片段同时交给模型和真机：

measured = sorted({tuple(row) for row in probe.run(program, setup=setup).tolist()})
reads = issue_model.replay(issue_model.parse(trace))
predicted = (reads[21] - reads[20], reads[22] - reads[20])

片段取自前面各节的时序实验：独立与相互依赖的向量指令、EUP、MXU、转置、归约、循环移位、vpush/spop、sld、CMEM 的 cld，以及计数循环。51 段片段中 50 段的 R1 − R0、R2 − R0 与模型完全一致，包括 VIF 积压、挡住标量发射的情况。例如 16 组 cld + vpop：模型和真机都是 (337, 877)。R1 = 337 的来历是：R1 前面有 32 个进入 VIF 的 bundle，R1 要等到只剩 19 项未释放，也就是第 13 项（第 7 条 cld）释放；它在第 327 个周期向量发射，第 337 个周期释放。

不一致的只有 1 段：1 组 cld + vpop，模型给出 (3, 67)，实测在 (3, 66) 与 (3, 67) 之间变化。

编译器生成的代码同样可以预测。第一章第 13、15 节把 kernel 的计算部分原样插进载体，模型与实测逐个相同：

这张表的后几行在加入“f32 的加、减和类型转换也是 2 个周期”这条规则之前，模型比实测少 1 到 7 个周期。差几个周期也说明缺了规则，下一小节讲怎样把它找出来。

模型与实测不符时：逐个 bundle 定位

本小节实验源码、输出。

上面那张参数表不是一次写全的。模型预测与实测不符时，说明缺了一条规则，要找的是它在哪一个 bundle 上起作用。办法是第 3 节末尾的那一条：起点的读数不动，把终点的读数逐个 bundle 往后移，每个位置各运行一次，得到“前 p 个 bundle 全部发射完”的时刻；模型对同样的前 p 个 bundle 也给出一个预测。两者之差在哪个 bundle 上跳变，缺的规则就与那个 bundle 有关。

下面这段清单把中间结果存进 TC VMEM，稍后再读回，是寄存器不够用时常见的写法：

SECTION = (
    'va0: vadd.8x128.s32 v11, 1, v10',
    'vst: vst.8x128 [vmem:0x80], v11',
    'va0: vadd.8x128.s32 v12, 2, v10',
    'va0: vadd.8x128.s32 v13, v12, v11',
    'vld: vld.8x128 v14, [vmem:0x80]',
    ...
)
for count in range(1, len(SECTION) + 1):
    body = ''.join(bundle(text) for text in SECTION[:count])
    measured, without = measure(probe, body), predict(body, 0)

predict(body, 0) 是去掉写后读规则的模型，也就是这条规则被发现之前的模型。括号中是差值比上一行增加了多少，每行末尾是新加进来的那个 bundle：

p =  1：实测 14，没有写后读规则的模型 14，差 0（+0），加上规则的模型 14    va0: vadd.8x128.s32 v11, 1, v10
p =  2：实测 15，没有写后读规则的模型 15，差 0（+0），加上规则的模型 15    vst: vst.8x128 [vmem:0x80], v11
p =  3：实测 16，没有写后读规则的模型 16，差 0（+0），加上规则的模型 16    va0: vadd.8x128.s32 v12, 2, v10
p =  4：实测 17，没有写后读规则的模型 17，差 0（+0），加上规则的模型 17    va0: vadd.8x128.s32 v13, v12, v11
p =  5：实测 22，没有写后读规则的模型 18，差 4（+4），加上规则的模型 22    vld: vld.8x128 v14, [vmem:0x80]
p =  6：实测 23，没有写后读规则的模型 19，差 4（+0），加上规则的模型 23    va0: vadd.8x128.s32 v15, v14, v13
p =  7：实测 24，没有写后读规则的模型 20，差 4（+0），加上规则的模型 24    vst: vst.8x128 [vmem:0x88], v15
p =  8：实测 25，没有写后读规则的模型 21，差 4（+0），加上规则的模型 25    vld: vld.8x128 v16, [vmem:0x90]
p =  9：实测 26，没有写后读规则的模型 22，差 4（+0），加上规则的模型 26    va0: vadd.8x128.s32 v17, v16, v15
p = 10：实测 31，没有写后读规则的模型 23，差 8（+4），加上规则的模型 31    vld: vld.8x128 v18, [vmem:0x88]
p = 11：实测 32，没有写后读规则的模型 24，差 8（+0），加上规则的模型 32    va0: vadd.8x128.s32 v19, v18, v17
p = 12：实测 33，没有写后读规则的模型 25，差 8（+0），加上规则的模型 33    vst: vst.8x128 [vmem:0x80], v19

差值只在第 5、第 10 个 bundle 上跳变，每次 4 个周期。这两个 bundle 都是 vld，读的都是 3 个 bundle 之前刚由 vst 写过的 tile；第 8 个 bundle 也是 vld，读的是没有被写过的 0x90，差值不变。所以缺的规则是：读刚写过的 tile 要等。

接着像前面测通路参数那样，单独测出这条规则的参数。vst 之后隔 d − 1 个 vnop 再 vld：

读相邻的 tile 时 R2 − R0 = 14 + d，没有等待。读同一个 tile 时，d ≤ 7 都是 21，d = 8 起才随 d 增加：vld 最早在写同一批行的 vst 向量发射之后 7 个周期发射。 判断的单位是行，不是 tile：0x84 起的 8 行与 0x80 起的 8 行有 4 行重叠，同样要等。另外几种排列：

同一个 bundle 中的 vst 与 vld 同时发射，不等。交替地写一个 tile、读一个 tile，每组 8 个周期；先写完再读，只有第一条 vld 等一次。把这条规则加进模型（issue_model.STORE_LOAD = 7，地址是立即数时才能判断读写的是哪几行）之后，上面逐 bundle 的 12 个时刻与表中的 5 种排列全部与实测一致，此前各段片段的结果不变。第一章第 6 节中 Pallas 的 int8 转换比 XLA 慢一倍，就是这条规则的后果。

这条规则对写 kernel 的含义是：存进 TC VMEM 的中间结果不要马上读回，中间至少隔 7 个周期的其他工作；编译器把寄存器溢出到 TC VMEM 时也受它约束。top-k 论文的 3.3 节用同样的办法测出了按索引读写（第一章第 9 节）的等待：vsetiar 之后的 vld.iar 要等，任何一次写内存之后的按索引读也要等，不论写的是哪个 tile。后一条正是用逐 bundle 的办法发现的：实测与模型之差只在含 vld.iar 的 bundle 上跳变。本节的模型没有包含按索引读写。

先判断受哪一类资源限制

逐个 bundle 推算之前，往往只要数清楚一段计算占用哪些资源，就能知道它的时间由什么决定、换哪一种写法才有用。上面的参数分成三类：

• 提交—取回的往返延迟。 一个结果要送进某个单元、等它出来才能接着用时，时间是依赖链上往返的次数乘以延迟：XLU 的归约 79 个周期，MXU 83 个周期，vpush/spop 43 个周期。链上的往返不能重叠，指令再少也快不了。
• 单元的发射间隔。 往返很多但互不依赖时，延迟被流水掩盖，时间是提交的条数乘以发射间隔：XLU、MXU 每 8 个周期一条，EUP 每 2 个周期一条。
• 逐元素指令的条数。 完全不经过这些单元时，时间是向量 ALU 指令的条数，每个周期可以发射两条（va0、va1 两个槽），再加上内存读写的等待。

这就是第一章第 15 节归纳的三类资源。那里的 top-k 是第一类：k 轮 argmax 串成一条链，每轮 82 个周期，其中 79 个在等 XLU，所以 f32[8,128] 取前 8 名的 685 个周期与指令条数几乎无关。要让它更快，减少指令没有用，只能缩短链，或者让链不经过 XLU。第一章第 12 节的累加链是第三类里“依赖链的长度”的例子，拆成两条之后瓶颈换成了发射槽。top-k 论文按这三类逐一给出了写法，并把这种估计给出的下限与实测放在一起（论文表 19）：

链的估计最准；XLU 不再是瓶颈之后，剩下的时间分散在写后读、按索引读这类等待和向量 ALU 的工作上，只数一类资源就不够了，要回到逐个 bundle 的模型。这种分类用来决定换哪一类写法，精确的周期数仍然要实测。

指令相同，先后不同，周期数也不同。 向量一侧按序发射，一个 bundle 在等，后面的都在等。同一组指令，把要等的那条往后挪、把不依赖它的工作提到前面，等待就被填上了。论文的 7.2 节给出了一个数量：同一个算法、同一组指令，编译器排出的顺序是 714 个周期，按实测的延迟重排之后是 595 个周期，一条指令也没有改。差别来自编译器的排程器不知道某些等待（转置之后 XLU 队列的占用、按索引读之前的等待），把指令排在一起，硬件只好一条条地等。手写或改写一段清单时，周期数取决于每种资源被谁占用、占多久，以及这些占用怎样错开，而不只是指令的条数。

编译器自己也有一张延迟表。 libtpu 的排程器按一张内部的延迟表决定两条指令之间隔几个周期。论文的 3.3 节在编译 top-k 时记录了它的回答：归约从提交到取回 79，循环移位和重排从提交到取回 69，同一个 XLU 上相邻两次提交 8，与本节实测的参数一致。这是一个独立的旁证。但这张表并不完整：转置之后队列的占用，表里是 104 个周期，实测约 223 个；与索引寄存器有关的等待，表里没有。编译器生成的清单大体上已经按延迟排好，上面那 119 个周期的差别就出在这张表缺失和偏小的地方；参数以真机实测为准。

由参数得到的写法规则

第一章各节在测出参数的同时，已经给出了相应的写法。把它们放在一起，就是排一段代码时要遵守的规则，每一条都对应模型中的一个参数：

• 循环。 周期数是循环体的 bundle 数（含延迟槽）乘以迭代次数，再加上循环体内的等待；分支本身没有额外开销。展开让不同迭代的指令互相填补（第一章第 8 节）。
• 提交—取回的单元要让足够多的提交同时在途。 单独一次往返，延迟是空等；连续提交时延迟被流水掩盖，时间由发射间隔决定。EUP 的延迟是 7、间隔是 2，要 4 条在途；XLU 的归约是 79 比 8，MXU 是 83 比 8，都要约 10 条在途（第一章第 7、11、12 节）。
• 取回有自己的间隔。 MXU 和 XLU 的取回每 8 个周期一条，积压的结果不能一口气取完，提交与取回要交错排列；两个 XLU 的提交和取回要成对放进同一个 bundle（第一章第 10、11 节）。
• 依赖链按延迟计。 f32 的加、减、乘和类型转换每条 2 个周期，整数运算每条 1 个周期。一条长的累加链可以拆成几条互不依赖的链（第一章第 12 节）。
• 向量 → 标量很贵。 spop 最早在 vpush 标量发射之后 43 个周期执行，期间标量一侧停止发射。把向量的计算结果变成循环次数或分支条件，每次都要付这个开销（第一章第 12 节）。
• TC VMEM 写后读。 vst 之后 7 个周期之内，读同一批行的 vld 不能发射。暂存的中间结果隔 7 个周期以上再读回，或者在这 7 个周期里安排别的工作。
• CMEM。 cld 之后 53 个周期才能 vpop，cld 每 2 个周期可以发射一次。串行地“读一个、取一个”时每组恰好 54 个周期，提前发出多个 cld 才能接近每 2 个周期一个 TC VREG（第二章第 3 节）。

模型的边界

模型可以预测不含 DMA 的直线代码，输入是动态执行的 bundle 序列。它没有包含按索引读写（vsetiar、vld.iar、vst.iar）、分段转置和转置之后 XLU 队列的占用；用到这些指令的清单，按上面逐个 bundle 的办法先补上规则。对整个 kernel，还需要把 DMA 的开销（第二章第 2 节）和各 TensorCore 之间的等待加进来。TC VMEM 读写、向量 ALU、EUP、MXU 各类指令实际执行了多少条，可以用 TensorCore 的硬件性能计数器从主机读出来，用于核对展开后的序列（研究报告 61）。研究报告 49 讨论了把这样的模型扩展成读取真实清单的静态计时器还缺什么：DMA 与向量工作的组合、并发 DMA、完整的 MXU 规则。本节的参数可以作为这项工作的输入。