jax.experimental.pallas.mosaic_gpu.copy_smem_to_gmem#

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

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

参数:

  • src (_Ref) – 要复制的源 SMEM 引用。

  • dst (_Ref) – 要复制到的目标 GMEM 引用。

  • predicate (jax.Array | None | None) – 一个布尔值，指示是否应执行复制。如果为 None，则始终执行复制。

  • commit_group (bool) – 如果为 True，则此复制和任何先前未提交的复制都将提交到一个组，并且可以联合通过 jax.experimental.mosaic.gpu.wait_smem_to_gmem() 等待。

  • reduction_op (mgpu.ReductionOp | None | None) – 如果设置，则在存储到 GMEM 时执行指定的归约操作。例如，使用 "add" 在概念上等同于执行 src += dst。

返回类型:

None

另请参阅

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