jax.experimental.pallas.mosaic_gpu.wait_load_tmem

jax.experimental.pallas.mosaic_gpu.wait_load_tmem()

等待所有先前由调用线程发起的异步 TMEM 加载完成。

此函数返回后，由调用线程发起的所有 TMEM 加载保证已完成。已读取的 TMEM 区域可以被调用线程安全地覆盖，或被通过 Barrier``s with ``orders_tensor_core=True 信号的任何线程安全地覆盖。