jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem

jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem(src, dst, barrier)

Asynchronously copies a GMEM reference to a SMEM reference.

See also

jax.experimental.mosaic.gpu.barrier_arrive() jax.experimental.mosaic.gpu.barrier_wait()

Parameters:

• src (pallas_core.AbstractMemoryRef)

• dst (pallas_core.AbstractMemoryRef)

• barrier (pallas_core.AbstractMemoryRef)

Return type:

None