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，则始终执行复制。

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

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

返回类型:

无

另请参阅

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