jax.experimental.pallas.mosaic_gpu.Barrier#

class jax.experimental.pallas.mosaic_gpu.Barrier(*, num_arrivals=1, num_barriers=1, orders_tensor_core=False)[source]#

描述一个屏障引用。

参数:

  • num_arrivals (int)

  • num_barriers (int)

  • orders_tensor_core (bool)

num_arrivals#

将被此屏障记录的到达次数。

类型:

int

num_barriers#

将被创建的屏障数量。可以通过索引到屏障 Ref 来访问各个屏障。

类型:

int

orders_tensor_core#

如果为 False，则一个线程成功等待并不能保证其他线程中与 TensorCore 相关的操作已完成。同样，当为 False 时，等待线程中的任何 TensorCore 操作都可以在等待成功之前开始。

类型:

bool

__init__(*, num_arrivals=1, num_barriers=1, orders_tensor_core=False)#
参数:

  • num_arrivals (int)

  • num_barriers (int)

  • orders_tensor_core (bool)

返回类型:

无

方法

__init__(*[, num_arrivals, num_barriers, ...])

get_ref_aval()

属性

num_arrivals

num_barriers

orders_tensor_core