jax.experimental.pallas.mosaic_gpu.copy_smem_to_gmem#

jax.experimental.pallas.mosaic_gpu.copy_smem_to_gmem(src, dst, predicate=None)[source]#

Asynchronously copies a SMEM reference to a GMEM reference.

Parameters:

  • src (pallas_core.AbstractMemoryRef) – The SMEM reference to copy from.

  • dst (pallas_core.AbstractMemoryRef) – The GMEM reference to copy to.

  • predicate (jax.Array | None | None) – A boolean indicating whether the copy should be performed. If None, the copy is always performed.

Return type:

None

See also

jax.experimental.mosaic.gpu.wait_smem_to_gmem() jax.experimental.mosaic.gpu.commit_smem()