jax.experimental.pallas.mosaic_gpu.semaphore_signal_parallel

jax.experimental.pallas.mosaic_gpu.semaphore_signal_parallel#

jax.experimental.pallas.mosaic_gpu.semaphore_signal_parallel(*signals)[源码]#

对多个信号量发出信号，不保证信号到达的顺序。

此原语在很大程度上等同于

for sem in semaphores:
  pl.semaphore_signal(sem, inc, device_id=device_id)

但与上面的循环不同，它不保证信号到达的任何顺序。特别是，目标设备可能在观察到 semaphores[0] 上的信号之前就观察到 semaphores[1] 上的信号。此操作仍然保证在信号到达之前，在信号之前执行的任何副作用都将完全执行并可见。

放宽的要求使得在 GPU 上进行整个操作的成本大大降低，因为可以为所有信号使用一个昂贵的内存围栏（而不是为每个信号使用一个昂贵的围栏）。