tvm_tirx.backend.cuda.op.wait_until

Contents

tvm_tirx.backend.cuda.op.wait_until#

tvm_tirx.backend.cuda.op.wait_until(dst, ptr, predicate, scope='gpu', space='global', ptx_type=None, backoff_ns=None, *, ty=None, loc=UNKNOWN_LOC)#

Read the global word at ptr into dst until predicate holds, and leave the exit value there.

Waiting on an address this way is also what declares it a synchronization word: the checker judges every access to that address against the protocol the wait names, rather than as an ordinary pair of memory accesses.

dst is an initialized thread-local scalar; its current value is tested first, so an already satisfied predicate performs no load. predicate is a trace-time callable taking the current value, or the boolean expression itself. It is re-evaluated on every iteration, so it may test dst against a loop-carried scalar such as a barrier’s phase: the loop body only loads, and nothing it does can move that scalar.

The wait always synchronizes with the contributions that made the predicate hold, so data those threads published elsewhere is visible when it returns. It polls ld.relaxed.<scope> and closes with one ld.acquire.<scope> into a discarded register, which is the cheapest spelling of that edge rather than a separate mode: paying acquire on every poll costs more, and closing with an acq_rel/sc fence instead costs far more, because the fence loses the loop’s fast-path exit.

There is no way to ask for less. A wait whose exit value is the whole message does take an edge it has no use for, and the two such waits in the kernel corpus were measured against this form on three shapes: every ratio landed inside the band that identical code measured against itself, and the two smaller shapes disagreed on the sign. So a relaxed mode would buy nothing here, while what it asks for is a promise made at the call site – that nothing the word guards is read afterwards – which the call site cannot show and the next edit can silently break.

T.nvshmem.wait_until shares this name deliberately: both block until a value satisfies a condition. They differ in what they wait on and how the condition is written – that one names a symmetric object across PEs and takes an enumerated comparison; this one names an address in this device’s global memory and takes a predicate. A protocol that waits within a CTA or cluster belongs on an mbarrier instead, which is why space admits only global.

backoff_ns puts a __nanosleep before each retry, as a contended wait is ordinarily written. It goes before the load, so a predicate that holds on entry still performs no load and no sleep, and a wait whose first poll succeeds pays nothing. A kernel that spells the backoff itself writes ld once and then waits, which is the same instruction sequence.

The backoff is the only thing a wait carries besides its own load, and it stays a scalar for a reason: it runs every iteration, touches no memory, and cannot move what the predicate reads, so it changes nothing the checker concludes. A timeout that has to print and trap is not that, and belongs to a loop the kernel writes itself.

Parameters:

loc (Location)