jax.experimental.pallas.mosaic_gpu.copy_smem_to_gmem

jax.experimental.pallas.mosaic_gpu.copy_smem_to_gmem#

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

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

参数:

src (_Ref) – 要从中复制的 SMEM 引用。
dst (_Ref) – 要复制到的 GMEM 引用。
predicate (jax.Array | None | None) – 一个布尔值，指示是否应执行复制。如果为 None，则始终执行复制。

返回类型:

None

另请参阅

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