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()