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]#

Tracks completion of all preceding tcgen05_mma and async_copy_smem_to_tmem calls.

Parameters:

  • barrier (_Ref) – Barrier Ref for synchronizing with the tensor core. Must have orders_tensor_core set to True.

  • collective_axis (str | None) – The name of the cluster axis along which the operations were performed if it was collective. The cluster axis should have a size of exactly 2, and must be on the minormost cluster axis.

See also