jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem#

jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem(n, wait_read_only=False)[原始碼]#

等待直到飛行中 (in flight) 的 SMEM->GMEM 複製數量不超過 n 個。

參數:
  • n (int) – 要等待的最大飛行中複製數量。

  • wait_read_only (bool) – 如果為 True,則等待飛行中複製完成從 SMEM 讀取。不等待寫入 GMEM。

返回類型:

None