jax.experimental.pallas.mosaic_gpu.wait_smem_to_gmem

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 的写入完成。

返回类型：

无