jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem#

jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem(n, wait_read_only=False)[源代码]#

等待,直到调用线程发出的 SMEM->GMEM 副本不再超过最新的 n 个副本在进行中。

参数:
  • n (int) – 等待中的副本的最大数量。

  • wait_read_only (bool) – 如果为 True,则等待正在进行的副本从 SMEM 读取完成。不等待写入 GMEM。

返回类型: