Blackwell Tensor Core 的同步模型
动机
tcgen05 的 PTX 手册:同一线程中的非 pipelined 指令给出两个例子:
1 | |
推断:这个例子中,在tcgen05.st完成前,tcgen05::wait会一直阻塞住程序。无法发射tcgen05.ld。因此在blackwell的tensor pipeline execution中,确保了这两个tcgen05操作的execution order,自然不需要fence
1 | |
mbarrier.try_wait.relaxed.cluster语义解释
mbarrier:操作的是 memory barrier 对象。try_wait:不是无条件卡死等待,而是“试着等一下 / 检查一下”。完成返回 true,没完成返回 false,所以常放在循环里反复检查。(后面重点验证)relaxed:只负责”等待这个barriere完成”,不保证普通内存读写的先后和可见性cluster:这个同步/内存语义的 scope 是 thread block cluster。
疑问
如果mbarrier.try_wait.relaxed.cluster的语义是:在前面的tcgen05操作,即这个例子中的tcgen05.mma完成前,一直阻塞住程序,不让后面的指令被发射,则也保证了tcgen05.mma和tcgen05.ld的有序性,但是ptx文档在此处却加了fence,无法理解😕。
所以这里需要弄清楚mbarrier.try_wait.relaxed.cluster是否具备**”在前面的tcgen05操作,即这个例子中的tcgen05.mma完成前,一直阻塞住程序,不让后面的指令被发射”**,如果是,则和ptx官方文档矛盾了。
本文会先介绍tcgen05的异步指令的同步方式。再做实验验证
Blackwell:tcgen05 指令之间怎样同步
tcgen05.mma 由谁发出,结果在哪里
Blackwell 的 tcgen05.mma 由单线程发出语义,一个线程发出指令就能启动整次矩阵乘法;不像 mma.sync 或 wgmma.mma_async 需要所有参与线程共同发出。它仍是异步计算:发出指令的线程继续执行,不等于矩阵乘法已经完成。PTX:tcgen05.mma 的发出语义
Blackwell 增加了 TMEM(Tensor Memory),用于存放 Tensor Core 相关的数据和累加结果。tcgen05.mma 的累加结果留在 TMEM;普通线程想处理结果,得另外发出 tcgen05.ld 把它读回寄存器。这与 Hopper 的 WGMMA 累加结果留在参与线程寄存器中不同。
1 | |
tcgen05 ptx表
tcgen05.* 指令 |
类别 | 作用 |
|---|---|---|
.mma |
异步 | 发出矩阵乘加,结果写 TMEM |
.cp |
异步 | 从 shared memory 向 TMEM 搬数据 |
.shift |
异步 | 对 TMEM 中的数据执行 shift |
.ld |
异步 | 从 TMEM 读到普通寄存器 |
.st |
异步 | 从普通寄存器写入 TMEM |
.commit |
同步 | 让此前由同一发出线程提交的 mma/cp/shift 由 barrier 跟踪 |
.wait::ld、.wait::st |
同步 | 等该线程此前相应的 ld/st 完成 |
.fence::* |
同步 | 规定前后 tcgen05 操作之间的顺序 |
.alloc、.dealloc、.relinquish_alloc_permit |
同步 | 管理 TMEM 的分配与使用许可 |
哪些异步指令可以直接接着发出
tcgen05的异步指令交由tensorcore pipeline处理,其execution order和程序顺序不一定一致,因此需要保证异步指令之间/异步指令和同步指令之间的同步。
Blackwell 有一批明确列出的由tensorcore pipeline自动保证execution order的异步指令对。称为pipelined instructions。它们访问相关 TMEM 位置的顺序由手册保证,不要求前一条先全部完成。例如,同一个 warp 发出 tcgen05.cp 往 TMEM 写输入,随后发出从相同位置读输入的 tcgen05.mma;或者前一次 tcgen05.mma 写累加结果,后一次 tcgen05.mma 接着读取这个结果,这些tcgen05异步指令对不需要程序显示保证execution order。参考PTX:pipelined 指令组合的完整清单
下表是部分示例pipelined instructions pair。
| 编号 | 先发出的异步操作 | 后发出的异步操作 | 额外条件 |
|---|---|---|---|
| 1 | mma 读 A/元数据 |
mma 写 D |
相同 N |
| 2 | mma 读 A/元数据 |
cp 写相关位置 |
相同 N |
| 3 | mma 读 A/元数据 |
shift 写相关位置 |
相同 N |
| 4 | mma 读 C |
mma 写 D |
相同 N、累加器地址和形状 |
| 5 | mma 写 D |
mma 读 C |
相同 N、累加器地址和形状 |
| 6 | mma 写 D |
mma 写 D |
相同 N、累加器地址和形状 |
| 7 | 非 .4x256b 的 cp 写入 |
mma 读 A/元数据 |
相同 N 和地址 |
| 8 | .4x256b 的 cp 写入 |
mma 读 A/元数据 |
相同 N |
| 9 | shift 写入 |
mma 读 A/元数据 |
相同 N |
| 10 | shift 读写 |
mma.cp.4x256b 写入 |
相同 N |
| 11 | st 写入 |
ld 读取 |
同一 warp |
| 12 | st 写入 |
下一条 st 写入 |
同一 warp |
| 13 | ld 读取 |
st 写入 |
同一 warp |
异步指令如何同步
tcgen05.mma、tcgen05.cp、tcgen05.shift 可以用 tcgen05.commit 关联到一个 mbarrier,以后通过等待该 barrier 的相应阶段,确认这些操作完成。tcgen05.ld 和 tcgen05.st 则用 tcgen05.wait::ld、tcgen05.wait::st 等待。
手册的两个例子:为什么后一段需要 fence
直接等待先前的 tcgen05.st
例一:
1 | |
tcgen05.wait::st 对应此前同一线程发出的 tcgen05.st。这条 wait 会阻塞执行线程,直到 st 完成;不需要另加 tcgen05.fence::after_thread_sync。
经由 commit 和 barrier 等 MMA 完成
例二:
1 | |
疑问
如果mbarrier.try_wait.relaxed.cluster的语义是:在前面的tcgen05操作,即这个例子中的tcgen05.mma完成前,一直阻塞住程序,不让后面的指令被发射,则也保证了tcgen05.mma和tcgen05.ld的有序性,但是ptx文档在此处却加了fence。
接下来需要做实验弄清楚mbarrier.try_wait.relaxed.cluster是否具备**”在前面的tcgen05操作,即这个例子中的tcgen05.mma完成前,一直阻塞住程序,不让后面的指令被发射”**。
B200 实验
待验证问题
我们想知道:try_wait是否会在barrier完成前,一直阻塞后面程序的指令发射
1 | |
测试程序流程
启动一个 CTA,共 64 个线程,分成两个 warp:
| 线程 | 做什么 |
|---|---|
| 0–31 | 等待方:在循环中反复执行 try_wait,每次紧跟着读一次全局计时器 |
| 32 | 生产者:延迟一段时间,然后执行唯一一次 mbarrier.arrive |
| 33–63 | 不执行 arrive;只参加 CTA 同步 |
线程 0 初始化 barrier,让它还需要 一次 arrive 才完成。 等待方通过waiter_armed通知生产者“现在可以开始”的变量
1 | |
等待方先确认 barrier 尚未完成,再设置 waiter_armed。线程 32 看到它以后,故意等约一百万个 SM 时钟周期。整个流程是:
1 | |
生产者的执行逻辑
生产者的代码如下:
1 | |
第一行用 clock64() 控制延迟多久。
第二行用 %globaltimer 记录时间 T_producer_before_arrive。
第三行执行唯一一次 arrive。我们比较的是第二行的 %globaltimer 读数。
需要观察的数据
只要等待方的某次 %globaltimer 读数比 T_producer_before_arrive 还早,那次读钟就一定早于生产者的 arrive。这就代表了try_wait不会一直阻塞指令发射
等待方的执行逻辑
等待方先做预检查,确认 barrier 尚未完成;线程 0 再设置 waiter_armed。真正取得实验时间戳的是下面这段代码:
1 | |
try_wait_parity_generic_with_clock做三件事:
1 | |
- 第一行检查 barrier,把这一轮的答案放进
complete_pred。尚未完成时,try_wait可能暂停执行线程,也可能稍后返回false;调用一次不保证成功。 - 第二行读
%globaltimer,把读数写到函数的输出参数。它没有使用complete_pred。 - 第三行才使用
complete_pred,转成 C++ 的complete;循环根据这个值决定重试还是退出。
这里要区分两条后续路径:读钟不依赖等待结果;“是否退出循环”依赖等待结果。不能因为读钟取值早,就说循环已经提前退出。
保存的是最后一轮,不是前一轮的旧读数
第一次调用把读数写到 independent_after_wait_clock,再赋给 latest_independent_clock。如果返回 false,每次重试都把本轮新的读数直接写进 latest_independent_clock;返回 true 的那一次也一样。循环退出后,保存的就是这个变量,而不是只保存第一次失败时的读数。
两个时间戳具体说明了什么
从同一次运行里取两个 %globaltimer 读数:
1 | |
这两个读数能够确认的顺序是:
1 | |
我测试了256 次,这个差值始终为负,范围是 −2816 到 −2304 。也就是说,即使这次调用最终让循环得到 true,它后面那条不依赖谓词的读钟也在生产者 arrive 前取到了值。这就推翻了“try_wait是否会在barrier完成前,一直阻塞后面程序的指令发射”
结论
try_wait不会在barrier完成前,一直阻塞后面程序的指令发射,因此不能用barrier commit + try_wait保证tcgen05异步指令的execution order
ptx文档关于mbarrier.try_wait的描述如下:
mbarrier.try_wait is a potentially blocking instruction which tests for the completion of the phase. If the phase is not complete, the executing thread may be suspended. Suspended thread resumes execution when the specified phase completes OR before the phase completes following a system-dependent time limit. The optional 32-bit unsigned integer operand suspendTimeHint specifies the time limit, in nanoseconds, that may be used for the time limit instead of the system-dependent limit.
在mbarrier跟踪的事件完成前,try_wait可能阻塞程序。