tvm_tirx.backend.cuda.lang.MBarrier

tvm_tirx.backend.cuda.lang.MBarrier#

class tvm_tirx.backend.cuda.lang.MBarrier(*args, **kwargs)#

Bases: object

Mbarrier wrapper with regular mbarrier.arrive.

Parameters:

  • pool (SMEMPool) – Shared memory pool allocator.

  • depth (int) – Number of barrier slots (one per pipeline stage).

  • phase_offset (int) – XORed into the phase bit on every wait / arrive.

  • leader (Expr, optional) – Boolean predicate selecting the single thread that runs mbarrier.init. Defaults to T.cuda.thread_rank() == 0 – thread 0 of the enclosing CTA, which always picks exactly one thread regardless of which index variables the caller binds. Override only when you want a different CTA-local thread to do the init.

init(count)#
wait(stage, phase)#
arrive(stage, remote=None, pred=None, count=None)#
ptr_to(idx)#
remote_view(rank)#

Create a view of this barrier mapped to another CTA’s shared memory.

The returned view retains the local barrier and target CTA so arrive emits the cluster form. Its mapped buffer remains available through ptr_to for operations that consume a remote shared-memory pointer. init and wait are local-only and reject remote views.