周期数:用 LCC 计时
上一节的清单回答了“TensorCore 执行了哪些指令”,这一节回答“这些指令要多少周期”。TensorCore 上的一切都按时钟周期进行:一个 bundle 的发射、一个结果要等多久、一次 DMA 的开销,原本都是周期计数。TPU v4 的时钟是 1050 MHz,一个周期 0.952 ns,100 个周期约 95 ns。本教程从这一节起,对每一种操作都同时给出它的指令和它的周期数。
周期数不用主机的时钟测。主机上调用一次 kernel 并等到结果,要一百多微秒,也就是十几万个周期,其中绝大部分是提交、启动和完成通知(第三章第 1 节);而本章的一段计算往往只有几十到几百个周期。周期数直接在设备上读:TensorCore 有一个本地周期计数器 LCC(local cycle counter),每个周期加 1,标量单元可以用一条指令读出它的值。在一段代码的两端各读一次,相减就是这段代码的周期数,没有换算,也不依赖 profiler。
本节先写出读 LCC 的指令,用它测几段手写的 bundle,得到“bundle 数不是周期数”的第一个例子;再给第 1 节的最小 kernel 计时,看它的时间花在哪里;然后把计时的边界放在 TC VMEM,得到本章之后各节都要用的测法;最后归纳周期数由哪几类开销组成。
本节实验:N 个 bundle 源码、输出;相互依赖的向量指令 源码、输出;最小 kernel 的周期数 源码、输出;TC VMEM 边界 源码、输出。
读 LCC 的指令
LCC 是 64 位的,标量寄存器是 32 位的,所以要分两半读:
{ s0: srdreg.lcclo s20 ; s1: srdreg.lcchi s25 }srdreg 读一个特殊寄存器(special register),lcclo 和 lcchi 是 LCC 的低、高 32 位,结果写进标量寄存器。两条指令放在同一个 bundle 的两个标量槽 s0、s1 中,在同一个周期执行,得到的两半属于同一个时刻,拼起来就是一致的 64 位读数。如果分在两个 bundle 读,两次读之间低位可能恰好从 0xFFFFFFFF 进位到 0,拼出的值会差 2³²。
Pallas 的公开接口没有读 LCC 的函数,所以读数用 tpuasm 直接插进机器程序。本教程把插入读数的步骤封装成 tpuasm_tools.py 中的两个工具:LccProbe 测手写的片段,KernelClock 测编译出的程序。下面依次介绍。
手写片段:LccProbe
LccProbe 把“编译载体、插入片段、取回读数”包装成一个类。它先用 Pallas 编译一个普通的 kernel 作为载体:
def kernel(x_hbm, out_hbm, data, sem) -> None:
pltpu.async_copy(x_hbm, data, sem).wait()
# 带唯一立即数的 vxor 标出插入位置。
data[:8, :] = data[:8, :] ^ jnp.uint32(0x13579BDF)
core = jax.lax.axis_index('tc')
pltpu.async_copy(data.at[:72, :], out_hbm.at[pl.ds(core * 72, 72)], sem).wait()输入是 u32[256,128] 的随机数,先整块搬进 TC VMEM 的 data;然后对前 8 行做一次异或,常数 0x13579BDF 在整个程序中只出现一次,find_bundles 由它找到这条 vxor 所在的 bundle,作为插入位置;最后把 data 的前 72 行(9 个 tile)写回 HBM。
run(body) 在插入位置之前放入三段清单:
- 前缀。 保存载体此后还要用到的标量寄存器,把输入的第 0 个 tile 读进
v10供片段当作数据使用,执行调用者给的setup,最后用一条sfence让此前的工作全部发射完(sfence的含义见下文)。 - 片段
body。 由调用者写出,其中用read_lcc(20)、read_lcc(21)、read_lcc(22)三次读 LCC,记作 R0、R1、R2。 - 后缀。 把读数用
vmov广播成 TC VREG,用vst存进 TC VMEM 的第 1–6 个 tile,随载体的输出 DMA 返回主机;再恢复标量寄存器。
主机把每对半字拼成 64 位,返回每次运行的 (R1 − R0, R2 − R0)。插入后的程序不再经过编译器,bundle 的内容和顺序就是设备实际执行的内容。片段用两个小函数写出:bundle(text) 生成一个 bundle 的清单文本,text 为空时是空 bundle;read_lcc(low) 生成上面那个读 LCC 的 bundle,低位写进 s{low}、高位写进 s{low + 5}。
实验:N 个 bundle 要几个周期
片段是“R0、N 个相同的 bundle、R1、一条 sfence、R2”:
body = read_lcc(20) + bundle(instruction) * count + read_lcc(21) + bundle('s0: sfence') + read_lcc(22)实验只改变两处:instruction 取四种 bundle 之一,count 取 0、1、4、16、64。
| 中间的 bundle | instruction |
|---|---|
| 空 bundle | '' |
| 相互依赖的标量加法 | 's0: sadd.s32 s24, 1, s24' |
| 独立的向量加法 | 'va0: vadd.8x128.s32 v11, 1, v10' |
| 从 TC VMEM 读一个 TC VREG | 'vld: vld.8x128 v11, [vmem:0x0]' |
四种 bundle 的结果完全相同,每种配置 8 次运行的读数也完全相同:
| N | 0 | 1 | 4 | 16 | 64 |
|---|---|---|---|---|---|
| R1 − R0 | 1 | 2 | 5 | 17 | 65 |
| R2 − R0 | 13 | 14 | 17 | 29 | 77 |
从中读出三条规则:
- 没有等待时,每个周期发射一个 bundle。 R1 − R0 = N + 1:两个相邻的读数相差 1,中间每多一个 bundle,差值加 1。标量加法链中每条依赖前一条的结果,也没有变慢:标量加法的结果在下一个周期就可用。
sfence加 11 个周期。 R2 − R1 总是 12:一个是相邻读数之间的 1,另外 11 来自sfence。- 读数没有波动。 同一段片段每次运行的读数都相同。指令什么时候发射只取决于程序本身,与哪一次运行无关。
所以 R2 − R0 的固定部分是 13 个周期。后文凡是写“R2 − R0”的地方,都包含这 13 个周期的读数与 sfence 的开销。
bundle 数不是周期数
上面的向量加法互不依赖。把它们换成相互依赖的:每一条用上一条写出的 v11。实验比较四种序列,每种 N 个 bundle:
'互不依赖的 vadd.f32': ('va1: vadd.8x128.f32 v11, v10, v10', 'va1: vadd.8x128.f32 v11, v10, v10'),
'相互依赖的 vadd.f32': ('va1: vadd.8x128.f32 v11, v10, v10', 'va1: vadd.8x128.f32 v11, v11, v10'),vmul.f32 的两种序列写法相同。结果:
| 序列 | R1 − R0(N = 1、4、16) | R2 − R0(N = 1、4、16) |
|---|---|---|
互不依赖的 vadd.f32 |
2、5、17 | 14、17、29 |
相互依赖的 vadd.f32 |
2、5、17 | 14、19、43 |
互不依赖的 vmul.f32 |
2、5、17 | 14、17、29 |
相互依赖的 vmul.f32 |
2、5、17 | 14、19、43 |
同样是 16 个 bundle,互不依赖时 R2 − R0 = N + 13,相互依赖时是 2N + 11。f32 的加法和乘法,结果要 2 个周期后才能被下一条指令使用;相互依赖的 16 条就要 32 个周期,而不是 16 个。清单上看不出这个差别:两种序列都是 16 个 bundle,每个 bundle 一条指令。
还有一点同样重要:四种序列的 R1 − R0 完全相同,只有 R2 − R0 不同。原因在于 TensorCore 的 bundle 要经过两次发射:
- 标量发射。 bundle 按顺序进入标量单元,没有阻塞时每个周期一个。标量槽中的指令(包括读 LCC 的
srdreg)在这时执行。 - 向量发射。 含向量指令的 bundle 随后在一个队列里排队,按顺序向量发射。一条向量指令要等它的源操作数就绪、要用的单元空闲才能发射;它在等,排在它后面的 bundle 都跟着等。
R1 紧跟在 16 条乘法后面,它记录的是标量一侧走到这里的时刻,这时向量一侧的乘法还没有发射完。sfence 是一条标量指令,它让下一个 bundle 等到此前所有的 bundle 都已向量发射之后才标量发射。所以 sfence 之后的 R2 才包含了向量一侧的等待。
暂且可以理解为:标量一侧可以领先向量一侧最多 20 个 bundle;
sfence等的是此前的指令全部“发射”,不是全部“完成”。第三章第 3 节详细介绍这个队列(VIF)、sfence的确切语义,以及它对计时和正确性的影响。
由此得到计时的两条规则,本教程的两个工具都已经按它们实现:
- 读数之前要有一条
sfence,否则读到的只是标量一侧的时刻。 - bundle 数只是周期数的下限。 每个周期最多发射一个 bundle;一条指令要等的东西没有就绪,它和它后面的指令都要等。一段代码要多少周期,取决于其中的指令各自要等什么。
给编译出的程序计时:KernelClock
LccProbe 测的是插进载体的手写片段。要测一个编译出的 kernel,读数可以照样插入,问题是读数怎样带出来:被测的程序没有多余的输出可用。办法来自 SMEM 的一个性质:它的内容在程序之间保留(第三章第 1 节用实验确认)。读数先留在 SMEM 中,等被测的程序运行结束,再用一个专门的读取程序取回。KernelClock 的每次读数是插入的 8 个 bundle:
{ s1: sst [smem:0x20100], s30 } # 借用 s30、s31,原值先存起来
{ s1: sst [smem:0x20101], s31 }
{ s0: sfence } # 等此前的向量工作全部发射
{ s0: srdreg.lcclo s30 ; s1: srdreg.lcchi s31 } # 读数
{ s1: sst [smem:0x20000], s30 } # 第 i 个读数存进 0x20000 + 2i、+ 2i + 1
{ s1: sst [smem:0x20001], s31 }
{ s1: sld s30, [smem:0x20100] } # 恢复
{ s1: sld s31, [smem:0x20101] }sst、sld 是标量单元写、读 SMEM 的指令(本章第 8 节)。被测程序的寄存器和输出都不受影响,所以它可以是任何已编译的程序,包括原生 XLA 的程序。用法是三步:
clock = tpuasm_tools.KernelClock(num_cores=2)
timed = clock.instrument(compiled, [pc0, pc1]) # 在原 bundle 编号 pc0、pc1 之前各插入一次读数
timed(x).block_until_ready()
readings = clock.read(2) # (TensorCore 数, 2) 的 64 位读数紧挨着的两次读数相差 20 个周期,这是读数自身的开销。本教程约定:区间的读数差减去 20,记作这段程序的周期数;clock.interval(compiled, first, last, call) 把这几步包在一起,返回区间 [first, last) 的周期数。这个约定带有一个固定的偏移:n 个互不等待的 bundle 读作 n − 2(第三章第 3 节说明这 2 个周期的来历)。两个实现用同一种读法比较时,偏移相互抵消。
最小 kernel 的时间花在哪里
第 1 节的最小 kernel 把 f32[8,128] 从 HBM 搬进 TC VMEM,乘以 2,再搬回 HBM。它的主体是 8 个 bundle:
509: { s0: dma.simple [vmem:s7], [hbm:s0], length=8, dst_flag=[sflag:52] }
510: { misc: vwait.ge [sflag:52], 8 }
511: { vld: vld.8x128 v0, [vmem:0x0] ; misc: vsyncadd.s32 [sflag:52], -8 }
512: { va0: vmul.8x128.f32 v1, 2.0, v0 }
513: { vst: vst.8x128 [vmem:0x0], v1 }
514: { s0: dma.simple [hbm:s1], [vmem:s7], length=8, dst_flag=[sflag:52] }
515: { misc: vwait.ge [sflag:52], 8 }
516: { misc: vsyncadd.s32 [sflag:52], -8 }实验对四个区间各测 8 次,每个区间都只在两端读数:
intervals = {
'整个 kernel(两个 vtrace 标记之间)': (start, end + 1),
'输入 DMA:发起到等待结束': (first_dma, first_dma + 2),
'计算:vld、vmul、vst': (first_dma + 2, second_dma),
'输出 DMA:发起到等待结束': (second_dma, second_dma + 3),
}
for name, (first, last) in intervals.items():
cycles = clock.interval(compiled, first, last, call, samples=8)| 区间 | 8 次运行的周期数 | 中位数 |
|---|---|---|
| 整个 kernel | 874、874、918、992、901、875、876、876 | 876 |
| 输入 DMA:发起到等待结束 | 445、445、617、446、445、445、446、446 | 445 |
计算:vld、vmul、vst |
1、1、1、1、1、1、1、1 | 1 |
| 输出 DMA:发起到等待结束 | 407、407、407、406、407、408、408、407 | 407 |
三个 bundle 的计算读作 1 个周期(3 − 2),两次 DMA 各四百多个周期。这个 kernel 的 876 个周期几乎全部在等 DMA:一次只搬 4 KiB 的 DMA,开销相当于四百多条向量指令。本章第 4 节测 DMA 的开销怎样随大小变化,第二章第 2 节讨论怎样安排 DMA 与计算的先后,使等待重叠起来。
两类读数的性质也不同。计算的读数 8 次完全相同;DMA 的读数每次相差一两个周期,偶尔多出一两百个。指令的发射是确定的;DMA 由 TensorCore 之外的搬运机构完成,什么时候完成取决于它当时的状态。所以区间内只有指令时,报告一个值;区间内有 DMA 时,报告中位数和范围。
把计时的边界放在 TC VMEM
上面的 kernel 以 HBM 为边界,周期数被搬运淹没了。本章研究的是每种操作本身的开销:数据已经在 TC VMEM 中,做完这一步之后结果也留在 TC VMEM 中,中间要多少周期。一个真实的 kernel 中间的每一步都是这种情形。所以本章的周期数统一这样测:输入和结果都在 TC VMEM,区间从输入就绪开始,到结果就绪结束,区间内没有进出 HBM 的 DMA。 原生 XLA 与 Pallas kernel 放进同一个边界比较。
Pallas 一侧
Pallas kernel 把操作数和结果都声明在 TC VMEM 即可。这样的 kernel 没有 DMA,也没有信号量:
@pl.kernel(
out_type=pltpu.VMEM(x.shape, x.dtype),
mesh=tc_mesh,
name='scale',
compiler_params=pltpu.CompilerParams(
disable_bounds_checks=True,
disable_semaphore_checks=True,
),
)
def kernel(x_vmem: Ref, o_vmem: Ref) -> None:
o_vmem[...] = x_vmem[...] * 2.0
return kernel(pltpu.with_memory_space_constraint(x, pltpu.VMEM))out_type=pltpu.VMEM(...) 让结果留在 TC VMEM,pltpu.with_memory_space_constraint(x, pltpu.VMEM) 要求操作数在进入 kernel 之前已经在 TC VMEM。这也是把一个 kernel 接在另一段片上计算后面时的写法。tpuasm_tools.vmem_kernel(f) 把任意函数 f 包成这样的 kernel:f 的每个参数、每个结果各对应一个 TC VMEM 中的 Ref。
外壳
原生 XLA 的程序没有“操作数在 TC VMEM”这种声明,要在程序里另外表达。被测的函数放进一个外壳,原生 XLA 与 Pallas kernel 用同一个(tpuasm_tools.vmem_shell,这里是只有一个输入时的样子):
from jax._src.state import primitives as state
from jax.experimental.layout import Layout, with_layout_constraint
ROW = Layout(major_to_minor=(0, 1), tiling=((8, 128),))
def pinned(x):
return with_layout_constraint(state.unpin(state.pin(x, to='vmem')), ROW)
def program(x):
result = function(pinned(x)) # 原生 XLA 的运算,或上面的 Pallas kernel
if isinstance(result, jax.Array):
return pinned(result)
return jax.lax.optimization_barrier(tuple(pinned(item) for item in result))外壳做了三件事:
- 内存空间。
state.pin(x, to='vmem')生成一条名为Pin的 custom call,把一个值固定在 TC VMEM;state.unpin把它变回普通的数组,内存空间不变。编译后的 HLO 在每个数组的 layout 末尾标出它所在的内存空间:不写的在 HBM,S(1)是 TC VMEM。 - 布局。
with_layout_constraint(..., ROW)把每个端点固定成第 1 节见过的布局{1,0:T(8,128)},即 Pallas kernel 在 TC VMEM 中的二维数组所用的布局。 - 共同的终点。 有多个结果时,它们共同经过
jax.lax.optimization_barrier,XLA 不能在全部结果就绪之前先处理其中一个。
暂且可以理解为:这三件事缺一样,两边量到的就不是同一件事。第三章第 2 节逐一去掉它们,看各自防住了什么;
pin、unpin是 JAX 的内部接口,那一节也说明为什么公开接口表达不了。
计时区间从输入的 Unpin 结束开始,到最晚的一个结果的 Unpin 结束为止,两端各读一次 LCC。tpuasm_tools.time_in_vmem 完成“放进外壳、编译、找出区间、插入读数、运行并核对输出”这几步:
clock = tpuasm_tools.KernelClock(num_cores=1)
candidates = {'XLA': scale, 'Pallas': tpuasm_tools.vmem_kernel(scale, 'scale')}
timing = tpuasm_tools.time_in_vmem(clock, function, runs, check)runs 是若干份输入,每份运行一次;check 核对每次的输出。返回值中有每次的周期数、区间内的清单,以及三项证据:各个端点在 HLO 中的类型与 layout、区间内有代码的 HLO 指令、区间内的 DMA。脚本开头调用的是 tpu_init.initialize_one_core(),原生 XLA 也只用一个 TensorCore(第 1 节)。
x * 2 的周期数
对 x * 2 测三种大小,每种 8 份不同的输入:
## x * 2,f32[8,128]
XLA:unpin.2 是 f32[8,128]{1,0:T(8,128)S(1)};unpin.3 是 f32[8,128]{1,0:T(8,128)S(1)}
区间是 bundle [521, 547),共 26 个;其中的 HLO 指令:broadcast_multiply_fusion;DMA:无
区间内的向量指令:vld.8x128 1,vmul.8x128.f32 1,vst.8x128 1
8 份输入的周期数:[42, 42, 42, 42, 42, 42, 42, 42]
Pallas:unpin.2.refs_0_ 是 f32[8,128]{1,0:T(8,128)S(1)};unpin.3 是 f32[8,128]{1,0:T(8,128)S(1)}
区间是 bundle [521, 547),共 26 个;其中的 HLO 指令:scale.1;DMA:无
区间内的向量指令:vld.8x128 1,vmul.8x128.f32 1,vst.8x128 1
8 份输入的周期数:[42, 42, 42, 42, 42, 42, 42, 42]
| 输入 | TC VREG 数 | 区间的 bundle 数 | XLA | Pallas |
|---|---|---|---|---|
f32[8,128] |
1 | 26 | 42 | 42 |
f32[64,128] |
8 | 34 | 50 | 50 |
f32[1024,128] |
128 | 154 | 170 | 170 |
两个端点都带 S(1),区间内没有 DMA。三点结论:
- 两边相同。 XLA 的 fusion 与 Pallas kernel 生成的指令相同,周期数也相同。逐元素的乘法没有什么可以手写得更好的地方。
- 读数每次相同。 8 份不同的输入,周期数完全一样。区间内只有指令,周期数是程序的性质,测一次就够,差一个周期也是确定的差别。
- 每多一个 TC VREG 多 1 个周期。 每个 TC VREG 是
vld、vmul、vst三条指令,但它们落在三个不同的槽,不同 TC VREG 的指令可以排进同一个 bundle,所以每个周期完成一个 TC VREG。128 个 TC VREG、384 条指令只要 128 个周期。
f32[64,128] 的清单中间是这样排的:
{ va0: vmul.8x128.f32 v5, 2.0, v4 ; vst: vst.8x128 [vmem:s1], v1 ; vld: vld.8x128 v6, [vmem:s0 + 0x18] }
{ va0: vmul.8x128.f32 v7, 2.0, v6 ; vst: vst.8x128 [vmem:s1 + 0x8], v3 ; vld: vld.8x128 v8, [vmem:s0 + 0x20] }
{ va0: vmul.8x128.f32 v9, 2.0, v8 ; vst: vst.8x128 [vmem:s1 + 0x10], v5 ; vld: vld.8x128 v10, [vmem:s0 + 0x28] }每个 bundle 做三件事:乘上一个 bundle 读进来的 TC VREG,读下一个,写出两个 bundle 之前乘好的那个。vst 比对应的 vmul 晚两个 bundle,正是上面测到的“f32 乘法的结果 2 个周期后可用”:编译器知道这个延迟,把写出排在结果就绪的那个周期,中间不留空等。
区间里固定的 39 个周期
一个 TC VREG 是 42 个周期,而计算只有 3 个 bundle。其余的周期在哪里,可以直接量出来:起点的读数不动,把终点的读数逐个 bundle 向后移,每个位置各运行一次。
for end in range(first + 1, last + 1):
cycles, = clock.interval(compiled, first, end, lambda timed: check(timed(x), x)) 到 bundle 521 为止: 0 { s1: sld s1, [smem:0x3ffe6] }
到 bundle 522 为止: 3 { s0: sadd.s32 s14, s6, s1 ; s1: sld s13, [smem:s6 + s1] }
到 bundle 523 为止: 7 { s0: sadd.s32 s15, 1, s13 }
到 bundle 524 为止: 8 { s1: sst [smem:s14], s15 }
到 bundle 525 为止: 9 { misc: vtrace 0x80000003 } # scale.1 开始
到 bundle 526 为止: 10 { s0: simm.s32 s0, 8 ; s1: simm.s32 s1, 0 }
到 bundle 527 为止: 11 {}
到 bundle 528 为止: 12 { vld: vld.8x128 v0, [vmem:s0] }
到 bundle 529 为止: 13 { va0: vmul.8x128.f32 v1, 2.0, v0 }
到 bundle 530 为止: 14 { vst: vst.8x128 [vmem:s1], v1 }
到 bundle 531 为止: 15 { misc: cfence }
到 bundle 532 为止: 16 { misc: vtrace 0x90000003 }
到 bundle 533 为止: 17 { s0: simm.s32 s17, 0 ; s1: sld s16, [smem:0x3ffe6] }
到 bundle 534 为止: 21 { s0: sadd.s32 s19, s17, s16 ; s1: sld s18, [smem:s17 + s16] }
到 bundle 535 为止: 25 { s0: sadd.s32 s20, 1, s18 }
到 bundle 536 为止: 26 { s1: sst [smem:s19], s20 }
到 bundle 537 为止: 27 { misc: vtrace 0x80000004 } # pin.3 开始
到 bundle 538 为止: 28 { misc: cfence }
到 bundle 539 为止: 29 { misc: vtrace 0x90000004 }
到 bundle 540 为止: 30 { s1: sld s21, [smem:0x3ffe6] }
到 bundle 541 为止: 34 { s0: sadd.s32 s23, s17, s21 ; s1: sld s22, [smem:s17 + s21] }
到 bundle 542 为止: 38 { s0: sadd.s32 s24, 1, s22 }
到 bundle 543 为止: 39 { s1: sst [smem:s23], s24 }
到 bundle 544 为止: 40 { misc: vtrace 0x80000005 } # unpin.3 开始
到 bundle 545 为止: 41 { misc: cfence }
到 bundle 546 为止: 42 { misc: vtrace 0x90000005 }
区间里有三条 HLO 指令:被测的运算、结果的 Pin、结果的 Unpin。每条的结构相同:
- 前面是 runtime 的 4 个 bundle:从 SMEM 读一个指针、读一个计数、加 1、写回。相邻两条
sld至少相隔 4 个周期,用sld结果的指令也要等 4 个周期(本章第 8 节),所以这 4 个 bundle 占 10 个周期(区间开头的那一组因为读法的偏移读作 8)。 - 然后是一对
vtrace标记,中间是这条指令自己的代码,以一条cfence结尾。Pin 和 Unpin 没有代码,只有这三个 bundle。
Pin 与 Unpin 各 13 个周期,被测运算的框架(4 个 runtime bundle、两个标记、cfence、装入地址的 simm 和一个空 bundle)也是 13 个周期,合计 39 个;剩下的 3 个周期才是 vld、vmul、vst。中间的每个 bundle 恰好一个周期。
所以读本章之后的周期数时要记住:区间里有 39 个周期是外壳固定的,对原生 XLA 和 Pallas 相同;超出的部分才是被测的运算。 一个运算如果被 XLA 拆成几条 HLO 指令,每多一条就多一份 13 个周期左右的框架,这是 XLA 那种实现的一部分,计入它的周期数。
这 39 个周期有一处会缩短。Pin 和 Unpin 前面各有 10 个周期花在标量一侧的 sld 上。被测的运算如果自己要在向量一侧等很久(例如等某个单元算出结果),标量一侧不会跟着停,这两组 sld 在等待期间就已经发射完,20 个周期与等待重叠,固定的部分只剩 19 个。后面遇到的运算中,没有等待的固定部分是 39 个周期,有几十个周期等待的是 19 个,等待较短的介于两者之间。
这种“起点不动、逐个移动终点”的办法以后还会用到:它给出区间里每个 bundle 发射的时刻,等待集中在哪一条指令上,一眼就能看出来。注意它每次运行只有两次读数。不能在一次运行里插很多次读数再逐段相减:每次读数带一条 sfence,会打断前后指令的重叠(第三章第 3 节)。
周期数由哪几类开销组成
本节已经出现了几种不同性质的开销,后面各节测到的周期数都由它们组成:
| 开销 | 例子 | 性质 |
|---|---|---|
| 发射一个 bundle | 任何 bundle | 1 个周期;同一个 bundle 中不同槽的指令同时发射 |
| 结果可用的延迟 | f32 加、乘 2 个周期,sld 4 个周期 |
只有下一条指令依赖它时才要等;不依赖它的指令可以填在中间 |
| 提交—取回的往返 | EUP、XLU、MXU(本章第 7、9–12 节) | 几个到几十个周期;期间可以提交别的工作 |
| 单元的发射间隔 | XLU、MXU 每 8 个周期接受一次提交 | 往返很多、互不依赖时,时间由它决定 |
| 等待外部事件 | DMA(本章第 4 节、第二章) | 几百个周期起,每次略有不同 |
设计一个 kernel,就是数清这些资源各被谁占用、占多久,再安排先后。上面 x * 2 的例子是最简单的情形:三条指令落在三个不同的槽,互相错开一个 bundle,每个周期完成一个 TC VREG。后面每认识一个硬件单元,就测出它的延迟和发射间隔;这些数汇总在本章的首页。把它们组合成一个可以从清单推算周期数的模型,是第三章第 4 节的内容;用周期数优化一个 kernel 的完整顺序,见第三章的首页。
LCC 的适用范围
LCC 数的是本芯片时钟的周期,每个 TensorCore 各有一个。由此有两条限制:
- 读数只在同一个 TensorCore 上可比。 两次读数相减才有意义。不同 TensorCore 的 LCC 起点不同,不能直接相减;不同芯片各有自己的时钟,周期的长短还相差百万分之几。
- 周期数不是秒。 换成秒要知道时钟的频率。第三章第 5 节用主机时钟标定出 1050 MHz。
LCC 适合回答“这段代码在这个 TensorCore 上要多少周期”,这正是优化一个 kernel 时要问的问题。要比较不同 TensorCore、不同芯片上的事件,用第三章第 5 节的全局时钟 GTC。
tpu-v4-latency-numbers 用同样的载体与读法测量了第二章引用的全部 DMA 开销。研究报告 48 走了另一条路:为读 LCC 增加一个 Pallas primitive 和它的 lowering,由编译器生成读数指令。两条路得到的读法相同。