This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 615abe2a56 [TIRx][CUDA] Add the declared synchronization-word wait
(#20353)
615abe2a56 is described below
commit 615abe2a568893f109bc0f90e26626aa03a0934c
Author: Guanjie Wang <[email protected]>
AuthorDate: Tue Sep 15 20:47:24 2026 -0400
[TIRx][CUDA] Add the declared synchronization-word wait (#20353)
## Problem
This exists for `tirx-kernels`, where every cross-thread protocol is a
hand-rolled spin on a global word -- `ld; while (!done) { ld; }`,
sometimes with a `__nanosleep` between polls and a fence after the loop.
`sm100_fp8_fp4_mega_moe` alone has eight. To a race checker that loop is
N ordinary reads racing another CTA's write, and it has to report every
one: nothing in the source separates a read that retries until a
condition holds from a read that took whatever it got.
`wait_until` is that shape, spelled once, and it can replace all of them
-- the 18 such loops in the corpus are an acquiring poll, a relaxed
poll, or a volatile poll with a backoff, and all three are this
operation. Adopting it costs nothing: the table below measures the
emitted form against the hand-written loop it replaces, on every kernel
that owns one. What it buys is that the word becomes declared -- the
address the wait names is the protocol's, and accesses to it are judged
against the protocol instead of as an ordinary pair.
## The operation
T.cuda.wait_until(dst, ptr, predicate, scope=, space=, ptx_type=,
backoff_ns=)
Spin on one global word until `predicate` holds; the exit value stays in
`dst`. Publishers stay raw PTX: `st.release`, `red`, `atom`. The name is
shared with `T.nvshmem.wait_until` deliberately -- that one names a
symmetric object and an enumerated comparison, this one a global address
and a predicate.
## Generated code
ld_relaxed(dst, ptr); // peeled
if (!predicate) {
#pragma unroll 1
do { ld_relaxed(dst, ptr); } while (!predicate);
}
{ T edge; ld_acquire(edge, ptr); (void)edge; } // the edge
`ld.relaxed.<scope>.<space>.<ptx_type>` polls, one `ld.acquire.<scope>`
closes. The closing read is discarded on purpose: it may observe a value
later than the predicate accepted, and these predicates are not all
monotone -- `mega_moe`'s grid barrier tests a sign-bit flip, its ring
waits test equality -- so reaching `dst` would hand the caller a value
its own predicate rejects. The edge survives: such words are published
by release RMWs, so acquiring any contribution synchronizes with all
earlier ones.
## Why this shape
Three spellings of the same acquire. The baseline is the loop these
kernels wrote by hand -- `ld.acquire` on every poll, no closing read --
and both other columns are wall time relative to it, negative meaning
faster. Interleaved, every benchmarkable kernel that owns a spin wait,
first round dropped as warm-up.
| kernel | sites | relaxed poll + acquire exit (emitted) | relaxed poll
+ `fence.acq_rel` exit |
|---|---|---|---|
| `sm100_fp8_fp4_mega_moe` | 8 | **-1.04%** | +49.9~51.8% |
| `radix_topk_multi_cta` | 1 | **-0.50%** | +5.89~6.10% |
| `agent_evolved_moe_fp8_blockscale_dsv3` | 3 | +0.03% | +1.53~2.15% |
| `agent_evolved_kda_backward_packed` | 4 | -0.08% | -0.10~+0.03% |
| `cudnn_sm100_flex_attention_backward` | 1 | -0.25% | -0.06~+0.96% |
The fence exit is never faster than the baseline on any of the five.
`fence.sc` costs what `fence.acq_rel` costs, and keeping the acquiring
poll *and* adding a fence is just as bad, which puts the cost in the
fence rather than in the loop body.
`ld.volatile` polls within 0.1% of `ld.relaxed` everywhere, which is ISA
8.4.2 measured; `.relaxed.<scope>` is emitted because it names the scope
instead of resting on `volatile` meaning `.sys`.
Peeling the first load gives ptxas one loop body to schedule, and the
SASS then matches the hand-written loop this replaces. That axis was not
swept separately.
## Other parameters
- `scope`, `ptx_type` spell the access. `ptx_type` respells at the same
width (a counter read as raw bits) and is required when the address is
untyped. 128 bits is refused: that exit value is not a scalar a
predicate can test.
- `space` admits `global` only. A wait within a CTA or cluster belongs
on an `mbarrier`; `shared` is refused with a message that says so.
- `backoff_ns` puts `__nanosleep` between polls. Nothing helps: monotone
degradation from 0, +1.35% at 20ns, +30.08% at 2us. Measured on short
waits only, so it stays opt-in.
---
python/tvm/backend/cuda/cpp/asm.py | 236 +++++++++++++++
python/tvm/backend/cuda/op.py | 142 +++++++++
python/tvm/backend/cuda/script.py | 4 +
src/backend/cuda/op/target_builtin.cc | 1 +
tests/python/tirx/codegen/test_cuda_wait_until.py | 353 ++++++++++++++++++++++
5 files changed, 736 insertions(+)
diff --git a/python/tvm/backend/cuda/cpp/asm.py
b/python/tvm/backend/cuda/cpp/asm.py
index 40171de139..b2c94b47f5 100644
--- a/python/tvm/backend/cuda/cpp/asm.py
+++ b/python/tvm/backend/cuda/cpp/asm.py
@@ -25,6 +25,9 @@ the PTX text, or the offset-scaling ``cp.async`` form the
legacy
``InjectPTXAsyncCopy`` pass emits.
"""
+import re
+
+from tvm import ir
from tvm.backend.cuda.op import cuda_func_call
from ..codegen.registry import CODEGEN_REGISTRY, register_codegen
@@ -33,6 +36,239 @@ from ..codegen.types import PTXDataType
from ..codegen.utils import parse_str
+# =============================================================================
+# Declared synchronization words. The four direct forms emit exactly what their
+# raw PTX spellings do; only the operation's identity differs, which is what
+# lets a checker tell a protocol's own accesses from a stray one. The wait is
+# the one that generates something new: a loop, which is why it lives here and
+# not in the instruction table.
+# =============================================================================
+_WAIT_UNTIL_SCALARS = {
+ "int32": "s32",
+ "uint32": "u32",
+ "int64": "s64",
+ "uint64": "u64",
+ "uint128": "b128",
+ "int128": "b128",
+}
+
+
+_WAIT_UNTIL_PTX_WIDTH = {
+ "b32": 32,
+ "s32": 32,
+ "u32": 32,
+ "b64": 64,
+ "s64": 64,
+ "u64": 64,
+ "b128": 128,
+}
+
+
+def _wait_until_scalar_suffix(ty, requested=""):
+ """The PTX type this access is spelled with.
+
+ Defaults to the word's own signedness. A caller may ask for another
+ spelling of the same width -- `b32` for a load or a store that only moves
+ the value, which is how kernels ordinarily write one -- but not for another
+ width, which would access a different number of bytes. Whether the
+ instruction has the requested form at all is settled by the instruction
+ table: `add` has no bit-typed form and rejects one there.
+ """
+ suffix = _WAIT_UNTIL_SCALARS.get(str(ty))
+ if suffix is None:
+ raise TypeError(f"sync word must be a 32/64/128-bit integer scalar,
got {ty}")
+ if not requested:
+ return suffix
+ width = _WAIT_UNTIL_PTX_WIDTH.get(requested)
+ if width is None:
+ raise TypeError(f"invalid ptx_type {requested!r}")
+ if width != _WAIT_UNTIL_PTX_WIDTH[suffix]:
+ raise TypeError(
+ f"ptx_type {requested!r} is {width} bits but the sync word is "
+ f"{_WAIT_UNTIL_PTX_WIDTH[suffix]} bits"
+ )
+ return requested
+
+
+def _wait_until_pointee(ptr, what):
+ if not isinstance(ptr.ty, ir.PointerType):
+ raise TypeError(f"{what} ptr must be a pointer to the synchronization
word")
+ return ptr.ty.element_type
+
+
+def _wait_until_word_suffix(ptr, requested, what):
+ """The PTX type the word is accessed as.
+
+ Normally the word's own type decides it and `ptx_type` may respell it at
+ the same width. A kernel that addresses its workspace by byte offset has
+ no pointee type to read -- the address is an untyped handle -- and there
+ `ptx_type` is not a respelling but the statement of how wide the word is,
+ so it is required.
+ """
+ pointee = str(_wait_until_pointee(ptr, what))
+ # A 128-bit word spans two 64-bit elements, so a pointer into the pair is
+ # how a kernel names it; asking for `.b128` there is not a respelling of
+ # the pointee but the statement that the word is the wider one.
+ if requested == "b128" and _WAIT_UNTIL_PTX_WIDTH.get(
+ _WAIT_UNTIL_SCALARS.get(pointee, "")
+ ) in (32, 64):
+ return requested
+ if pointee in _WAIT_UNTIL_SCALARS:
+ return _wait_until_scalar_suffix(pointee, requested)
+ if not requested:
+ raise TypeError(
+ f"{what} ptr is an untyped address; pass ptx_type to say how wide
the word is"
+ )
+ if requested not in _WAIT_UNTIL_PTX_WIDTH:
+ raise TypeError(f"invalid ptx_type {requested!r}")
+ return requested
+
+
+def _wait_until_same_width(dst_ty, suffix, what):
+ """The destination register must be as wide as the word.
+
+ Not the same type: a bit-typed access moves 32 or 64 bits into whatever
+ register of that width the caller named, which is how kernels ordinarily
+ read a counter they treat as signed out of an unsigned word.
+ """
+ dst = _wait_until_scalar_suffix(dst_ty)
+ if _WAIT_UNTIL_PTX_WIDTH[dst] != _WAIT_UNTIL_PTX_WIDTH[suffix]:
+ raise TypeError(
+ f"{what} destination is {_WAIT_UNTIL_PTX_WIDTH[dst]} bits but the "
+ f"sync word is {_WAIT_UNTIL_PTX_WIDTH[suffix]} bits"
+ )
+
+
+def _wait_until_thread_local_scalar(dst, what):
+ if not isinstance(dst, ir.TensorLoad) or dst.source.scope() not in {
+ "local",
+ "local_scalar",
+ "register",
+ "reg",
+ }:
+ raise TypeError(f"{what} dst must be a writable thread-local scalar")
+ return dst.ty
+
+
+def _wait_until_forward(spelling, *args):
+ """Emit one instruction-table PTX operation and return its codegen
result."""
+ from ..ptx import PTXNamespace # pylint: disable=import-outside-toplevel
+
+ call = PTXNamespace()[spelling](*args)
+ return CODEGEN_REGISTRY[call.op.name](call.args)
+
+
+@register_codegen("cuda_wait_until")
+def cuda_wait_until(dst, ptr, condition, scope, space, ptx_type, backoff_ns):
+ """Lower a declared wait to a pre-tested loop around one scoped load."""
+ scope, space, ptx_type = (parse_str(x) for x in (scope, space, ptx_type))
+ dtype = _wait_until_thread_local_scalar(dst, "wait_until")
+ suffix = _wait_until_word_suffix(ptr, ptx_type, "wait_until")
+ _wait_until_same_width(dtype, suffix, "wait_until")
+
+ # The wait polls relaxed and closes with a single `ld.acquire`, rather than
+ # paying acquire semantics on every poll. Measured against the acquiring
+ # poll over every benchmarkable kernel that owns a spin wait, interleaved,
+ # round 1 dropped:
+ #
+ # sm100_fp8_fp4_mega_moe 8 sites -1.04%
+ # radix_topk_multi_cta 1 site -0.50%
+ # agent_evolved_moe_fp8_blockscale_dsv3 3 sites +0.03%
+ # agent_evolved_kda_backward_packed 4 sites -0.08%
+ # cudnn_sm100_flex_attention_backward 1 site -0.25%
+ #
+ # The closing read is what takes the edge, and it may observe a value later
+ # than the one that satisfied the predicate. That is still the edge the
+ # protocol means: these words are published by `red`/`atom` release RMWs,
+ # so every contribution sits in one release sequence and an acquire reading
+ # any of them synchronizes with all the earlier ones. It reads into a
+ # discarded temporary precisely so the later value cannot reach `dst` --
+ # the wait's exit value stays the value the predicate accepted, which
+ # matters because predicates here are not all monotone (`mega_moe`'s grid
+ # barrier tests a sign-bit flip, and the ring waits test equality).
+ #
+ # This is the wait's only lowering. `ld.volatile` polls the same way -- PTX
+ # ISA 8.4.2 puts it and `.relaxed` in one class, and measured they are
+ # within 0.1% everywhere -- but `.relaxed.<scope>` says the scope out loud
+ # instead of resting on `volatile` meaning `.sys`.
+ load_call, tags = _wait_until_forward(
+ f"ld.relaxed.{scope}.{space}.{suffix}", dst, ptr
+ )
+ # The helper is named after the poll, which is the load that repeats; the
+ # closing acquire hangs off that name with an `_acquire` suffix.
+ load_name = parse_str(load_call.args[0])
+ name = load_name.replace("ptx_ld_", "cuda_wait_until_")
+ source = load_call.args[-1].value.replace(load_name, name + "_load")
+ acquire_call, _ = _wait_until_forward(
+ f"ld.acquire.{scope}.{space}.{suffix}", dst, ptr
+ )
+ acquire_name = parse_str(acquire_call.args[0])
+ acquire_source = acquire_call.args[-1].value.replace(acquire_name, name +
"_acquire")
+ source += "\n" + acquire_source
+ # NVRTC has no `__typeof__`, and `decltype` on the destination yields a
+ # reference that cannot be declared uninitialized, so the scratch takes
+ # the C type the generated helper already spells in its signature.
+ scratch_type = re.search(
+ rf"void\s+{re.escape(name)}_acquire\(\s*([\w:]+)\s*&", acquire_source
+ )
+ if scratch_type is None: # pragma: no cover - the helper shape is fixed
+ raise RuntimeError(f"cannot read the destination type of
{name}_acquire")
+ closing = (
+ f" {{ {scratch_type.group(1)} __tirx_wait_edge; "
+ f"{name}_acquire(__tirx_wait_edge, (ptr)); (void)__tirx_wait_edge; }}"
+ )
+ # The predicate stays at the call site: an ordinary bool argument would be
+ # evaluated once, before the load ever updates the destination.
+ # Load first, test after, which is how every spin loop in this repository
+ # is written: `ld; while (!done) { ld; }` reads once before it can decide
+ # anything, so a pre-tested loop would need the caller to seed the
+ # destination -- and the only way to seed it honestly is another load of
+ # the same word, which is one more unguaranteed read for a checker to
+ # judge. A `do`/`while` needs no seed.
+ #
+ # `unroll 1` is what a hand-written spin carries, and what the loop this
+ # replaces emitted. Without it nothing stops ptxas from duplicating the
+ # load into an unrolled body, which changes how often a waiter polls even
+ # though the instruction sequence is the same one.
+ #
+ # The first load and test are peeled out of the loop. The executed sequence
+ # is the same either way, but a single `do`/`while` gives ptxas one body to
+ # schedule and it stops emitting the early-exit branch a hand-written
+ # `ld; while (!done) { ld; }` gets -- so a waiter whose predicate already
+ # holds pays a second load and a poll it did not pay before. Measured on
+ # `sm100_fp8_fp4_mega_moe`'s grid barrier at +2.2% and on
+ # `radix_topk_multi_cta` at +1.25%, reproduced on two idle B200s with the
+ # states interleaved. Peeling costs nothing when the wait does spin.
+ #
+ # A backoff of zero is no backoff: the macro, and so the emitted code, is
+ # the one a wait without one produces, down to the argument list.
+ backoff = 0 if not hasattr(backoff_ns, "value") else int(backoff_ns.value)
+ if backoff == 0:
+ source += (
+ f"\n#define {name}(dst, ptr, predicate) "
+ f"do {{ {name}_load((dst), (ptr)); if (!(predicate)) {{ "
+ f"_Pragma(\"unroll 1\") "
+ f"do {{ {name}_load((dst), (ptr)); }} while (!(predicate)); }}"
+ f"{closing} }} while (0)\n"
+ )
+ operands = (condition,)
+ else:
+ # Between polls, never before the first one and never after the last:
+ # the shape `allgather_gemm` and `gemm_reduce_scatter` write by hand.
+ name = f"{name}_backoff"
+ source = source.replace(f"{name[: -len('_backoff')]}_load",
f"{name}_load")
+ source += (
+ f"\n#define {name}(dst, ptr, predicate, backoff_ns) "
+ f"do {{ {name}_load((dst), (ptr)); if (!(predicate)) {{ "
+ f"_Pragma(\"unroll 1\") "
+ f"while (1) {{ __nanosleep(backoff_ns); {name}_load((dst), (ptr));
"
+ f"if (predicate) break; }} }}"
+ f"{closing} }} while (0)\n"
+ )
+ operands = (condition, backoff_ns)
+ return cuda_func_call(name, *load_call.args[1:-1], *operands,
source_code=source), tags
+
+
# =============================================================================
# mbarrier waits — ``mbarrier.try_wait`` only polls once, so the body wraps it
# in a branch loop that retries until the parity flips. The magic
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index ca8598aec0..0d6ea8806c 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -382,6 +382,148 @@ def cuda_float8tohalf8(src_addr, dst_addr):
return call_intrin("", "tirx.cuda.float8tohalf8", src_addr, dst_addr)
+_WAIT_UNTIL_SCOPE = ("cta", "cluster", "gpu", "sys")
+# Global only, and that is the whole surface a declared word needs. A protocol
+# that waits inside a CTA or a cluster has `mbarrier`, which is the hardware's
+# own primitive for it and which the checker already models by generation; a
+# polled flag in shared memory would be a worse spelling of the same thing. The
+# shared-memory reads that look like a protocol in practice -- fetching a TMEM
+# allocation address out of a mailbox -- are ordinary reads after a barrier,
+# with no spin and no predicate, so they are not protocols at all.
+_WAIT_UNTIL_SPACE = ("global",)
+_WAIT_UNTIL_DTYPE = ("int32", "uint32", "int64", "uint64")
+# Widths, so a requested PTX type can be checked against the word it accesses.
+# Which spellings an instruction actually takes is the instruction table's
+# business: `add` has no `.b32` form, so a bit-typed arrival is rejected there
+# rather than listed as illegal here.
+_WAIT_UNTIL_PTX_TYPES = {
+ "b32": 32,
+ "s32": 32,
+ "u32": 32,
+ "b64": 64,
+ "s64": 64,
+ "u64": 64,
+ # A 128-bit word is a 16-byte naturally aligned span, moved by the `.b128`
+ # forms of `ld`/`st` (PTX ISA 8.3, sm_70). It carries a value no scalar
+ # predicate can test, so `wait` rejects it; `store` and `load` take it.
+ "b128": 128,
+}
+
+
+def _validate_wait_until_attrs(scope, space, ptx_type=None):
+ """The attributes a declared synchronization-word wait still carries."""
+ if scope not in _WAIT_UNTIL_SCOPE:
+ raise ValueError(f"invalid scope={scope!r}; expected one of
{_WAIT_UNTIL_SCOPE}")
+ if space not in _WAIT_UNTIL_SPACE:
+ detail = (
+ "; a protocol that waits within a CTA or cluster belongs on an "
+ "`mbarrier`, which the checker models by generation"
+ if space == "shared"
+ else ""
+ )
+ raise ValueError(
+ f"invalid space={space!r}; expected one of
{_WAIT_UNTIL_SPACE}{detail}"
+ )
+ if ptx_type is not None and ptx_type not in _WAIT_UNTIL_PTX_TYPES:
+ raise ValueError(
+ f"invalid ptx_type={ptx_type!r}; expected one of "
+ f"{tuple(sorted(_WAIT_UNTIL_PTX_TYPES))}"
+ )
+
+
+def _reject_wide_word_for_predicate(ptx_type, what):
+ """A predicate tests one scalar, so a 128-bit word cannot be waited on.
+
+ The exit value of a 16-byte word is not a number the loop can compare, and
+ the checker's record of a declared word's writes holds one scalar per
+ write. A kernel that has to read one spells the `ld` itself.
+ """
+ if ptx_type is not None and _WAIT_UNTIL_PTX_TYPES[ptx_type] > 64:
+ raise ValueError(
+ f"{what} does not take a {_WAIT_UNTIL_PTX_TYPES[ptx_type]}-bit
word: "
+ "its exit value is not a scalar a predicate can test; read it with
"
+ "a plain `ld` instead"
+ )
+
+
+def cuda_wait_until(
+ dst,
+ ptr,
+ predicate,
+ scope="gpu",
+ space="global",
+ ptx_type=None,
+ backoff_ns=None,
+):
+ """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.
+ """
+ _validate_wait_until_attrs(scope, space, ptx_type)
+ _reject_wide_word_for_predicate(ptx_type, "wait_until")
+ if tirx.is_buffer_var(dst):
+ dst = dst[0]
+ condition = tirx.convert(predicate(dst) if callable(predicate) else
predicate)
+ return call_intrin(
+ "",
+ "tirx.cuda.wait_until",
+ dst,
+ ptr,
+ condition,
+ scope,
+ space,
+ ptx_type or "",
+ tirx.convert(0 if backoff_ns is None else backoff_ns),
+ )
+
+
def _validate_mbarrier_arrive_attrs(sem, scope, space, remote):
if (sem == "") != (scope == ""):
raise ValueError("mbarrier.arrive sem and scope must be specified
together")
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index 136a89e9dd..e5c1c846af 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -148,6 +148,10 @@ class CUDANamespace:
self.mov_sreg: Callable[..., Any] = _op_wrapper(_cuda_op.cuda_mov_sreg)
# Spin-until-ready mbarrier waits: label-loop asm blocks, not single
# PTX instructions -- which is why they live here and not in T.ptx.
+ # One declared synchronization word: every access a protocol makes to
+ # it goes through these, so a checker can separate them from a stray
+ # access and read the word's write history off the declaration.
+ self.wait_until = _op_wrapper(_cuda_op.cuda_wait_until)
self.mbarrier_wait = _op_wrapper(_cuda_op.cuda_mbarrier_wait)
self.mbarrier_wait_acquire_cluster = _op_wrapper(
_cuda_op.cuda_mbarrier_wait_acquire_cluster
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index 19526f319f..5619bcc9aa 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -167,6 +167,7 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = {
TIRX_DEVICE_INTRIN_ALIAS(cuda_any_sync, cuda, kPure),
TIRX_DEVICE_INTRIN_ALIAS(cuda_atomic_add, cuda, kOpaque),
TIRX_DEVICE_INTRIN_ALIAS(cuda_atomic_cas, cuda, kOpaque),
+ TIRX_DEVICE_INTRIN_ALIAS(cuda_wait_until, cuda, kOpaque),
TIRX_DEVICE_INTRIN_ALIAS(cuda_ballot_sync, cuda, kOpaque),
TIRX_DEVICE_INTRIN_ALIAS(cuda_bfloat1622float2, cuda, kOpaque),
TIRX_DEVICE_INTRIN_ALIAS(cuda_bfloat162float, cuda, kOpaque),
diff --git a/tests/python/tirx/codegen/test_cuda_wait_until.py
b/tests/python/tirx/codegen/test_cuda_wait_until.py
new file mode 100644
index 0000000000..26fa522626
--- /dev/null
+++ b/tests/python/tirx/codegen/test_cuda_wait_until.py
@@ -0,0 +1,353 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Codegen contract for the declared wait.
+
+`wait_until` is the only operation of the set that survives. The four
+direct forms were removed: they emitted exactly what their raw spellings do, so
+they bought nothing at codegen, and a protocol's word is now recognized by the
+address the wait names rather than by each access carrying an identity. A
+publisher is spelled `red`/`atom`/`st` directly, as these tests do.
+
+What the wait itself generates is a loop, and the shape of that loop is the
+contract asserted here: it polls relaxed, closes an acquiring wait with one
+`ld.acquire`, peels the first poll, and never lets the closing read reach the
+caller's destination. The protocols exercised are the shapes real kernels use.
+"""
+
+import pytest
+
+import tvm
+from tvm.script import tirx as T
+
+
+def build(func):
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"})
+ with target:
+ mod = tvm.compile(tvm.IRModule({"main": func}), target=target,
tir_pipeline="tirx")
+ return getattr(mod, "mod", mod).imports[0].inspect_source("")
+
+
+def rendezvous(backoff_ns=None, ptx_type=None):
+ """An N-way barrier on a monotone counter, as radix_topk_multi_cta writes
it."""
+
+ @T.prim_func
+ def kernel(state: T.Buffer((1,), "int32"), participants: T.int32):
+ T.device_entry()
+ T.cta_id([2])
+ lane = T.thread_id([32])
+ spin = T.alloc_local((1,), "int32")
+ phase = T.alloc_local((1,), "int32")
+ if lane == 0:
+ phase[0] = T.int32(0)
+ spin[0] = T.int32(0)
+ T.ptx.red.release.gpu.global_.add.s32(state.ptr_to([0]),
T.int32(1))
+ T.cuda.wait_until(
+ spin[0],
+ state.ptr_to([0]),
+ predicate=lambda v: v >= (phase[0] + T.int32(1)) *
participants,
+ scope="gpu",
+ ptx_type=ptx_type,
+ backoff_ns=backoff_ns,
+ )
+
+ return kernel
+
+
+def packed_contribution():
+ """Counter in the high half, payload in the low half: DeepEP's notify
slot."""
+
+ @T.prim_func
+ def kernel(slot: T.Buffer((1,), "uint64"), out: T.Buffer((1,), "uint64"),
n: T.int32):
+ T.device_entry()
+ T.cta_id([2])
+ lane = T.thread_id([32])
+ observed = T.alloc_local((1,), "uint64")
+ if lane == 0:
+ T.ptx.red.relaxed.gpu.global_.add.u64(
+ slot.ptr_to([0]),
+ T.bitwise_or(T.shift_left(T.uint64(1), T.uint64(32)),
T.uint64(7)),
+ )
+ observed[0] = T.uint64(0)
+ T.cuda.wait_until(
+ observed[0],
+ slot.ptr_to([0]),
+ predicate=lambda v: T.shift_right(v, T.uint64(32)) ==
T.Cast("uint64", n),
+ scope="gpu",
+ )
+ out[0] = T.bitwise_and(observed[0], T.uint64(0xFFFFFFFF))
+
+ return kernel
+
+
+def _wait_macro(source):
+ return next(
+ line
+ for line in source.splitlines()
+ if line.startswith("#define") and "wait_until" in line
+ )
+
+
+# =============================================================================
+# The emitted loop
+# =============================================================================
+
+
+def test_the_wait_is_a_loop_beside_a_raw_publisher():
+ """The publisher is ordinary PTX; only the wait generates something new."""
+
+ source = build(rendezvous())
+ assert "red.release.gpu.global.add.s32" in source
+ assert "while (!(" in source
+
+
+def test_a_wait_loads_before_it_tests():
+ """`do { ld } while (!done)`, the way every spin here is written by hand.
+
+ A pre-tested loop would need the caller to seed the destination, and the
+ only honest seed is another load of the same word -- one more unguaranteed
+ read, in the kernel, for a checker to have an opinion about. Loading first
+ needs none.
+ """
+
+ macro = _wait_macro(build(rendezvous()))
+ assert "do { tvm_builtin_cuda_wait_until" in macro
+ assert macro.index("_load(") < macro.index("while (!(")
+
+
+def test_an_acquiring_wait_polls_relaxed_and_closes_with_one_acquire():
+ """The poll carries no ordering; a single closing `ld.acquire` takes the
edge.
+
+ Measured against the acquiring poll over every benchmarkable kernel that
+ owns a spin wait, interleaved, round 1 dropped: -1.04% on
+ `sm100_fp8_fp4_mega_moe` (8 wait sites), -0.50% on `radix_topk_multi_cta`,
+ and within 0.3% on the other three. Paying acquire semantics on every poll
+ buys nothing -- only the last read decides what the waiter goes on to
+ observe.
+ """
+
+ source = build(rendezvous())
+ macro = _wait_macro(source)
+ assert "ld.relaxed.gpu.global.s32" in source
+ assert "ld.acquire.gpu.global.s32" in source
+ # the loop polls relaxed, and the acquire is reached once, after it
+ assert macro.count("_load(") == 2
+ assert macro.count("_acquire(") == 1
+ assert macro.index("while (!(") < macro.index("_acquire(")
+
+
+def test_the_closing_acquire_cannot_overwrite_the_waited_value():
+ """The edge read goes to a scratch, never to the caller's destination.
+
+ It may observe a value later than the one the predicate accepted, and the
+ predicates here are not all monotone -- `sm100_fp8_fp4_mega_moe`'s grid
+ barrier tests a sign-bit flip and its ring waits test equality -- so
letting
+ it reach `dst` would hand the caller a value its own predicate rejects.
+
+ The edge survives that: these words are published by `red`/`atom` release
+ RMWs, so every contribution sits in one release sequence and an acquire
+ reading any of them synchronizes with all the earlier ones.
+ """
+
+ macro = _wait_macro(build(rendezvous()))
+ scratch = "__tirx_wait_edge"
+ assert f"{scratch};" in macro
+ assert f"_acquire({scratch}, (ptr))" in macro
+ # every write to the caller's destination comes from a polling load
+ assert "_acquire((dst)" not in macro
+
+
+def test_every_wait_takes_the_edge():
+ """There is one lowering, so a caller cannot ask for a wait worth less.
+
+ A wait that consumes only its own exit value takes an edge it does not
+ need. Both such waits in the kernel corpus were measured against this form
+ and neither could tell the difference, so the mode they would have needed
+ is not worth the thing it costs: a promise, made at the call site, that
+ nothing the word guards is read afterwards -- which is a promise the call
+ site cannot show.
+ """
+
+ macro = _wait_macro(build(rendezvous()))
+ assert "_acquire(__tirx_wait_edge, (ptr))" in macro
+ assert "ld.volatile" not in build(rendezvous())
+
+
+def test_a_bit_typed_word_keeps_its_handwritten_ptx():
+ """A kernel picks the PTX type per operation, not per word.
+
+ `red.add` has no bit-typed form, while a load that only moves the value is
+ ordinarily spelled `.b32`. Reproducing that split is what lets a migration
+ leave the emitted PTX unchanged.
+ """
+
+ source = build(rendezvous(ptx_type="b32"))
+ assert "red.release.gpu.global.add.s32" in source
+ assert "ld.relaxed.gpu.global.b32" in source
+ assert "ld.acquire.gpu.global.b32" in source
+ # the default spelling must be gone, or the override did nothing
+ assert "ld.acquire.gpu.global.s32" not in source
+
+
+def test_a_packed_word_round_trips():
+ source = build(packed_contribution())
+ assert "red.relaxed.gpu.global.add.u64" in source
+ assert "ld.relaxed.gpu.global.u64" in source
+
+
+def test_a_backoff_sleeps_between_polls_and_never_around_them():
+ """Sleep between polls, as the contended waits are written by hand: never
+ before the first poll, never after the last.
+
+ The first poll is peeled out of the loop, so the shape is
+ `ld; if (!done) { while (1) { sleep; ld; if (done) break; } }`. That runs
+ the same sequence a hand-written `while (1) { ld; if (done) break; sleep;
}`
+ runs -- load, test, sleep, load, test -- and the peel is what keeps the
+ early exit a hand-written spin gets."""
+
+ source = build(rendezvous(backoff_ns=40))
+ macro = _wait_macro(source)
+ assert "__nanosleep(backoff_ns)" in macro
+ # The peeled poll runs before any sleep: nothing sleeps before polling
once.
+ assert macro.index("_load(") < macro.index("__nanosleep")
+ # The peeled poll is guarded, so an already satisfied predicate never
+ # reaches the loop at all.
+ assert macro.index("_load(") < macro.index("if (!(predicate))")
+ # Inside the loop the sleep precedes the poll it separates, and the break
+ # follows that poll, so nothing sleeps after the last one.
+ loop = macro[macro.index("while (1)") :]
+ assert loop.index("__nanosleep") < loop.index("_load(")
+ assert loop.index("_load(") < loop.index("if (predicate) break;")
+ # The call site is what fixes the 40.
+ assert "_backoff(" in source and ", 40)" in source
+
+ # No backoff is the spelling a wait had before the field existed.
+ plain = build(rendezvous())
+ assert "__nanosleep" not in plain
+ assert "_backoff" not in plain
+
+
+def test_a_backoff_still_closes_the_acquire_after_the_last_poll():
+ """The edge read belongs after the loop, never inside it."""
+
+ macro = _wait_macro(build(rendezvous(backoff_ns=40)))
+ assert macro.index("if (predicate) break;") < macro.index("_acquire(")
+
+
+def test_a_wait_may_test_against_a_thread_local_scalar():
+ """A barrier tests the counter against its loop-carried phase.
+
+ The macro re-evaluates the predicate each iteration, and the loop body only
+ loads, so the phase cannot move under the wait even though it is re-read.
+ """
+
+ source = build(rendezvous())
+ assert "while (!(" in source
+ assert "phase" in source
+
+
+# =============================================================================
+# What the wait refuses
+# =============================================================================
+
+
[email protected](
+ "kwargs, error, message",
+ [
+ ({"order": "acquire"}, TypeError, "order"),
+ ({"impl": "volatile"}, TypeError, "impl"),
+ ({"scope": "warp"}, ValueError, "scope"),
+ ],
+)
+def test_a_wait_rejects_attributes_it_cannot_mean(kwargs, error, message):
+ """`order` and `impl` are gone, so a kernel still passing one is told.
+
+ They named the two axes this wait was measured on. Both settled, and a
+ silently accepted keyword would let a kernel written against the old
+ spelling keep compiling while meaning something else.
+ """
+
+ from tvm.backend.cuda.op import cuda_wait_until
+
+ with pytest.raises(error, match=message):
+ cuda_wait_until(None, None, None, **kwargs)
+
+
+def test_a_ptx_type_of_another_width_is_refused():
+ with pytest.raises(Exception, match="64 bits but the sync word is 32
bits"):
+ build(rendezvous(ptx_type="b64"))
+
+
+def test_a_wide_word_cannot_be_waited_on():
+ """A predicate tests one scalar, and a 16-byte exit value is not one.
+
+ The wait refuses the width rather than testing half of it. The refusal
lands
+ while the function is traced, so the definition is what is guarded here.
+ """
+
+ with pytest.raises(Exception, match="does not take a 128-bit word"):
+
+ @T.prim_func
+ def kernel(response: T.Buffer((2,), "uint64")):
+ T.device_entry()
+ T.cta_id([1])
+ lane = T.thread_id([32])
+ seen = T.local_scalar("uint64")
+ if lane == 0:
+ seen = T.uint64(0)
+ T.cuda.wait_until(
+ seen,
+ response.ptr_to([0]),
+ seen != T.uint64(0),
+ scope="gpu",
+ ptx_type="b128",
+ )
+
+
+def test_a_declared_word_is_global_and_says_where_a_shared_wait_belongs():
+ """Shared memory is out of scope by design, not by omission.
+
+ A protocol that waits within a CTA or a cluster has `mbarrier`, which is
+ the hardware's primitive for it and which the checker models by generation.
+ A polled flag in shared memory would be a worse spelling of the same thing,
+ so the refusal points at the primitive that does belong there.
+ """
+
+ with pytest.raises(ValueError, match="mbarrier"):
+ T.cuda.wait_until(
+ None, None, None, scope="cta", space="shared", ptx_type="b32"
+ )
+
+
+def test_the_four_direct_forms_are_gone():
+ """Only the wait remains: a publisher is spelled in raw PTX.
+
+ The removed forms emitted exactly what their raw spellings do, so they
never
+ changed codegen; the protocol's word is recognized by the address the wait
+ names instead.
+ """
+
+ retired = (
+ "atomic_ref_store",
+ "atomic_ref_add",
+ "atomic_ref_fetch_add",
+ "atomic_ref_load",
+ # The surviving wait kept the family's name until it was the only
+ # member left; `atomic` claimed an atomicity its scoped loads never
+ # had, and `ref` grouped a family of one.
+ "atomic_ref_wait",
+ )
+ assert [name for name in retired if hasattr(T.cuda, name)] == []
