jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive

jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive#

jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive(barrier, collective_axis=None)[source]#

跟踪先前 tcgen05_mma 调用的完成情况。

参数:

barrier (_Ref) – 用于与张量核心同步的 Barrier Ref。必须将 orders_tensor_core 设置为 True。
collective_axis (str | None) – 如果 MMA 是集体操作，则为执行 MMA 的集群轴的名称。集群轴的大小必须正好为 2，并且必须位于最次的集群轴上。

另请参阅