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 的操作完成。

返回类型:

None