jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem#

jax.experimental.pallas.mosaic_gpu.copy_gmem_to_smem(src, dst, barrier, *, collective_axes=None)[source]#

将 GMEM 引用异步复制到 SMEM 引用。

参见

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

参数:

  • src (_Ref)

  • dst (_Ref)

  • barrier (_Ref)

  • collective_axes (str | tuple[str, ...] | None | None)

返回类型:

None