diff --git a/python/tvm/backend/cuda/lang/smem_desc.py b/python/tvm/backend/cuda/lang/smem_desc.py index b1d850abd216..d6244fd73aa8 100644 --- a/python/tvm/backend/cuda/lang/smem_desc.py +++ b/python/tvm/backend/cuda/lang/smem_desc.py @@ -43,13 +43,14 @@ def add_16B_offset(self, offset): def make_lo_uniform(self): """Broadcast the lower 32 bits to all warp lanes via ``__shfl_sync``.""" - func_name = "smem_desc_make_lo_uniform" - source_code = f""" -__forceinline__ __device__ void {func_name}(uint64_t* desc) {{ - SmemDescriptor* d = reinterpret_cast(desc); - d->lo = __shfl_sync(0xffffffff, d->lo, 0); -}} -""" - return T.cuda.func_call( - func_name, T.address_of(self._buf[0]), source_code=source_code, return_type="void" + desc_lo = T.alloc_local([1], "uint32") + desc_hi = T.alloc_local([1], "uint32") + T.ptx.mov.b64(desc_lo[0], desc_hi[0], self._buf[0]) + T.ptx.shfl_sync.idx.b32( + desc_lo[0], + desc_lo[0], + T.uint32(0), + T.uint32(0x1F), + T.uint32(0xFFFFFFFF), ) + T.ptx.mov.b64(self._buf[0], desc_lo[0], desc_hi[0]) diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py index 43a9d4d7a776..3e29523af2fe 100644 --- a/python/tvm/backend/cuda/ptx/engine.py +++ b/python/tvm/backend/cuda/ptx/engine.py @@ -118,11 +118,13 @@ def codegen(*args): # operand layout -- which may depend on them, a register group's length # being a function of the modifiers -- is looked up from them (memoized # per token combination in the table). - # The marker is a comma-joined flag set ("pred" and/or "p"); codegen - # only cares about @p, since the register classes are already in the - # table. See the arg-layout note in `_emit`. + # The marker is a comma-joined flag set ("pred", destination policy, + # and/or "p"); codegen only cares about @p and the destination + # binding policy, since the register classes are already in the table. + # See the arg-layout note in `_emit`. flags = parse_str(args[-1]).split(",") predicated = "pred" in flags + preserve_dst = "keep" in flags tokens = [parse_str(a) for a in args[len(args) - n_slots - 1 : -1]] rest = args[: len(args) - n_slots - 1] # operands, plus pred when present mod_map = mods(entry, tokens) @@ -167,7 +169,15 @@ def _slot_dtype(slot, i, n): # the helper (which has no parameter for them). imm_at = {i for slot, i, _ in layout if slot.kind == "imm"} imms = tuple(str(int(rest[at[i]])) for i in sorted(imm_at)) - _, helper, source = render_variant(entry, tokens, predicated, dtypes, imms, sinks) + _, helper, source = render_variant( + entry, + tokens, + predicated, + dtypes, + imms, + sinks, + preserve_dst=preserve_dst, + ) # Every helper is void; a destination is an ordinary argument, printed # by the C codegen as the lvalue it binds the reference parameter to. # A predicate rides after the operands, so everything past n_operands @@ -458,7 +468,7 @@ def _coerce_pred(entry, pred): raise ValueError(f"{entry.name}: pred must be a bool/uint32/int32 expression") -def _emit(entry, filled, operands, pred=None): +def _emit(entry, filled, operands, pred=None, preserve_dst=False): # Modifiers resolve before the operands are looked at: the attribute chain # is parsed before the call happens, and a register group's length may be a # function of the modifiers, so the expected arity needs the modifier map. @@ -477,14 +487,12 @@ def _emit(entry, filled, operands, pred=None): if len(operands) != n_args: names = ", ".join(f"{slot.name}[{n}]" if n > 1 else slot.name for slot, _, n in layout) raise ValueError(f"{entry.name} expects {n_args} operand(s) ({names}), got {len(operands)}") + if preserve_dst and not entry.has_dst: + raise ValueError(f"{entry.name}: preserve_dst=True requires a written destination") if pred is not None: - if entry.has_dst: - raise ValueError( - f"{entry.name}: @p predication is only supported on instructions without a " - f"destination (a false predicate leaves the destination unwritten, and the " - f'"=" output constraint would discard its prior value)' - ) pred = _coerce_pred(entry, pred) + elif preserve_dst: + raise ValueError(f"{entry.name}: preserve_dst=True requires pred=...") # The sink symbol `_`: a lane the caller discards. It is checked here # rather than in `_coerce_operand` because it is not a value at all -- the # instruction names the symbol, so the lane leaves no Call argument, and @@ -528,6 +536,8 @@ def _emit(entry, filled, operands, pred=None): # indistinguishable from the integer that shares its carrier. The marker # makes the round trip exact; a call with neither prints "" as before. flags = ["pred"] if pred is not None else [] + if preserve_dst: + flags.append("keep") flags += [ f"p{i}" for slot, i, lanes in layout @@ -587,7 +597,7 @@ def __getattr__(self, name): raise AttributeError(name) return _InstrChain(_narrow(self._cands, unescape_token(name))) - def __call__(self, *args, pred=None): + def __call__(self, *args, pred=None, preserve_dst=False): # Also accepts the printed round-trip form: trailing modifier-token # strings in slot order ("" = omitted slot) followed by the pred # marker, and, when the marker says "pred", the predicate as the last @@ -603,6 +613,8 @@ def __call__(self, *args, pred=None): flags = marker.split(",") if marker else [] if "pred" in flags and pred is None and args: args, pred = args[:-1], args[-1] + if "keep" in flags: + preserve_dst = True # Put back what the printed text cannot carry: the sunk lanes # (which left no argument at all, so they are re-inserted first, in # ascending order, to restore the operand positions) and then the @@ -631,7 +643,18 @@ def __call__(self, *args, pred=None): # what lets optional trailing operands dispatch by arity. for entry, filled in cands: try: - hits.append((entry, _emit(entry, filled, operands, pred=pred))) + hits.append( + ( + entry, + _emit( + entry, + filled, + operands, + pred=pred, + preserve_dst=preserve_dst, + ), + ) + ) except ValueError as err: # Keep the exception; a lone candidate re-raises it untouched, # and only the aggregate view needs entry names in front. diff --git a/python/tvm/backend/cuda/ptx/gen_stubs.py b/python/tvm/backend/cuda/ptx/gen_stubs.py index 9d31cb723b17..2d8c22222c97 100644 --- a/python/tvm/backend/cuda/ptx/gen_stubs.py +++ b/python/tvm/backend/cuda/ptx/gen_stubs.py @@ -78,7 +78,8 @@ def _chain_class(family: str, entries: list[InstructionEntry]) -> str: f"{s.name}∈{{{','.join(s.choices)}}}{' (opt)' if s.optional else ''}" for s in entry.slots ) if entry.check is not None and entry.check.__doc__: - doc = f"{doc} — {entry.check.__doc__.strip()}" if doc else entry.check.__doc__.strip() + check_doc = " ".join(entry.check.__doc__.split()) + doc = f"{doc} — {check_doc}" if doc else check_doc if len(entries) > 1: shapes = dict.fromkeys( "(" + ", ".join(_operand_params(e)).replace(": Any", "") + ")" for e in entries @@ -98,8 +99,9 @@ def _chain_class(family: str, entries: list[InstructionEntry]) -> str: # spells its operands as the catch-all; a second one would not even # parse ("Only one '*' parameter allowed"). params.append("*args: Any") - if not any(e.has_dst for e in entries): - params.append("pred: Any = None") + params.append("pred: Any = None") + if any(e.has_dst for e in entries): + params.append("preserve_dst: bool = False") signature = f"def __call__({', '.join(params)}) -> None" # Emit the shape ruff format would produce, so the generated text needs no # formatter to be canonical: a docstring that fits on one line closes on @@ -121,7 +123,7 @@ def _chain_class(family: str, entries: list[InstructionEntry]) -> str: # `T.ptx["tcgen05.ld.sync.aligned.16x64b.x4.b32"](...)`. continue lines.append(f" {attr}: {cls}") - if len(signature) > 92: # keep the generated stub within the repo line limit + if len(f" {signature}: ...") > 100: joined = ",\n ".join(params) ret = signature[signature.rindex(")") + 1 :] signature = f"def __call__(\n {joined},\n ){ret}" diff --git a/python/tvm/backend/cuda/ptx/render.py b/python/tvm/backend/cuda/ptx/render.py index 6a57fd0fbc22..96bdab0a91a4 100644 --- a/python/tvm/backend/cuda/ptx/render.py +++ b/python/tvm/backend/cuda/ptx/render.py @@ -196,7 +196,13 @@ def _helper_name( def render_variant( - entry: InstructionEntry, tokens, predicated=False, dtypes=None, imms=None, sinks=frozenset() + entry: InstructionEntry, + tokens, + predicated=False, + dtypes=None, + imms=None, + sinks=frozenset(), + preserve_dst=False, ): """Render one variant: ``(opcode, helper_name, helper_source)``. @@ -218,8 +224,11 @@ def render_variant( ``predicated`` is a framework-level axis (never in the table): the helper gains a trailing ``uint32_t __pred`` operand, and the instruction is - guarded with ``.reg .pred p; setp.ne.b32 p, %N, 0; @p ...``. Only valid - for instructions without a destination — see ``InstructionEntry.has_dst``. + guarded with ``.reg .pred p; setp.ne.b32 p, %N, 0; @p ...``. A written + destination uses the normal write-only binding by default, which states + that the caller will not consume the inactive value before selecting or + guarding it. ``preserve_dst`` instead binds read-write and retains the old + value on the inactive path. """ mod_map = mods(entry, tokens) written = [tok for tok in tokens if tok] @@ -242,9 +251,14 @@ def render_variant( # non-canonical (atom's d and b do exactly that). imm_of = dict(zip(imm_slots(entry), imms or (), strict=True)) helper = _helper_name(entry, written, imms, dtypes, canonical, mod_map, sinks) + assert not preserve_dst or entry.has_dst, "preserve_dst requires a written destination" if predicated: - assert not entry.has_dst, "@p is only supported on instructions without a destination" - helper += "_pred" + if entry.has_dst: + helper += "_pred_keep" if preserve_dst else "_pred_undef" + else: + helper += "_pred" + else: + assert not preserve_dst, "preserve_dst requires predication" params, inputs, outputs, rendered = [], [], [], [] pre, post = [], [] # C-side carrier declarations / bit puns @@ -300,7 +314,7 @@ def render_variant( elif slot.kind == "ptr": params.append(f"const void* {lname}") inputs.append(f'"l"({lname})') - elif slot.rw == "rw": + elif slot.rw == "rw" or (preserve_dst and slot.rw == "w"): # A register the instruction both reads and writes -- an # in-place accumulator. "+" tells the compiler the prior value # is live, which "=" would declare dead; it is also what makes @@ -311,7 +325,7 @@ def render_variant( # which has no coherent read-modify-write story; the dtypes # that need one are refused rather than half-supported. raise ValueError( - f"{entry.name}: accumulator '{slot.name}' cannot take dtype " + f"{entry.name}: read-write destination '{slot.name}' cannot take dtype " f"{dtype_of[slot]} (it binds through a carrier)" ) params.append(f"{cb.c_type}& {lname}") @@ -347,7 +361,7 @@ def render_variant( reg = template.format(slot=slot.name, n=bridge_counts[template]) bridge_counts[template] += 1 bridge_decls.append(f".reg {bridge.reg_class} {reg};") - if slot.rw in ("r", "rw"): + if slot.rw in ("r", "rw") or (preserve_dst and slot.rw == "w"): asm_pre.append(bridge.into.format(reg=reg, idx=idx)) if slot.rw in ("w", "rw"): asm_post.append(bridge.out_of.format(reg=reg, idx=idx)) diff --git a/python/tvm/backend/cuda/ptx/table.py b/python/tvm/backend/cuda/ptx/table.py index 71c299a919ec..e55c4b095bdc 100644 --- a/python/tvm/backend/cuda/ptx/table.py +++ b/python/tvm/backend/cuda/ptx/table.py @@ -420,13 +420,12 @@ def typed_operands(self) -> tuple[OperandSlot, ...]: def has_dst(self) -> bool: """Whether the instruction writes a destination operand. - Gates ``@p``: a false predicate leaves destinations unwritten, and the - ``"="`` output constraint tells nvcc the prior value is dead, so a - predicated destination silently loses it. An accumulator (``rw="rw"``) - binds "+" instead, which keeps the old value live -- so it does not - count here and @p remains available on it. A ``.pred`` result is a - ``rw="w"`` register like any other, written through "=" the same way, - so it counts without needing a case of its own. + A false predicate leaves destinations unwritten. The default ``"="`` + output constraint means the inactive value is undefined to the caller; + ``preserve_dst=True`` explicitly requests a read-write binding instead. + An accumulator (``rw="rw"``) already binds "+", so it does not count + here. A ``.pred`` result is a ``rw="w"`` register like any other and + counts without needing a case of its own. """ return any(s.kind == "reg" and s.rw == "w" for s in self.operands) @@ -2482,6 +2481,19 @@ def _check_tcgen05_mma_block_scale(m): return None +def _check_tcgen05_mma_block_scale_block(m): + """Valid block sizes per kind: mxf8f6f4/mxf4 use block32, while + mxf4nvf4 supports block16 and block32.""" + valid = { + "kind::mxf8f6f4": ("block32",), + "kind::mxf4": ("block32",), + "kind::mxf4nvf4": ("block16", "block32"), + }[m["kind"]] + if m["block_size"] not in valid: + return f"{m['kind']} supports {'/'.join(valid)}" + return None + + # `{, byteMask}` exists exactly when `.cp_mask` is written. _cp_mask_lanes = _present_lanes("cp_mask") @@ -7647,8 +7659,8 @@ def _check_tcgen05_mma_block_scale(m): # only differ by those qualifiers: no call sites. # - .ws without the zero-column-mask-desc operand: every caller passes # the mask (as literal zero). - # - block_scale's .block16/.block32 vector sizes and its - # scale_vec-omitted spelling: the library always writes .scale_vec::NX. + # - block_scale's scale_vec-omitted spelling without a .block16/.block32 + # size: no call site uses that form. *[ InstructionEntry( name=f"tcgen05_mma_{form}", @@ -7707,6 +7719,35 @@ def _check_tcgen05_mma_block_scale(m): ) for form in ("ss", "ts") ], + *[ + InstructionEntry( # block-scaled with an explicit scale block size + name=f"tcgen05_mma_block_scale_block_{form}", + mnemonic="tcgen05", + slots=( + ModifierSlot("action", ("mma",)), + ModifierSlot("cta_group", ("cta_group::1", "cta_group::2")), + ModifierSlot("kind", ("kind::mxf8f6f4", "kind::mxf4", "kind::mxf4nvf4")), + ModifierSlot("block_scale", ("block_scale",)), + ModifierSlot("block_size", ("block16", "block32")), + ), + cert_arch="sm_100a", + check=_check_tcgen05_mma_block_scale_block, + operands=( + OperandSlot("d_tmem", kind="addr", space="tmem"), + *( + (OperandSlot("a_desc", dtype="u64"),) + if form == "ss" + else (OperandSlot("a_tmem", kind="addr", space="tmem"),) + ), + OperandSlot("b_desc", dtype="u64"), + OperandSlot("idesc", dtype="u32"), + OperandSlot("sfa_tmem", kind="addr", space="tmem"), + OperandSlot("sfb_tmem", kind="addr", space="tmem"), + OperandSlot("enable_input_d", dtype="pred"), + ), + ) + for form in ("ss", "ts") + ], *[ InstructionEntry( # weight-stationary: no mask vector, a zero-column desc name=f"tcgen05_mma_ws_{form}", diff --git a/python/tvm/backend/cuda/tile_primitive/common.py b/python/tvm/backend/cuda/tile_primitive/common.py index 157753d2feb3..d2298096316a 100644 --- a/python/tvm/backend/cuda/tile_primitive/common.py +++ b/python/tvm/backend/cuda/tile_primitive/common.py @@ -55,24 +55,22 @@ def get_indices(nth, start, extent): return [r + s for r, s in zip(reversed(relative), start)] +def smem_desc_replace_lo(desc_val, desc_lo): + """Replace the lower address lane of a 64-bit SMEM descriptor.""" + desc_halves = T.reinterpret("uint32x2", desc_val) + desc_hi = T.Shuffle([desc_halves], [1]) + return T.reinterpret("uint64", T.Shuffle([T.cast(desc_lo, "uint32"), desc_hi], [0, 1])) + + def smem_desc_add_16B_offset(desc_val, offset): """Add a 16B-aligned byte offset to the lower 32 bits of a SMEM descriptor. - Uses the SmemDescriptor union defined in the CUDA header (header.py). - All callers must share a single implementation to avoid codegen conflicts. + The address lane wraps as uint32 without carrying into the descriptor's + upper control bits. """ - func_name = "tvm_builtin_smem_desc_add_16B_offset" - source_code = f""" -__forceinline__ __device__ uint64_t {func_name}(uint64_t desc_base, int32_t offset) {{ - SmemDescriptor desc; - desc.desc_ = desc_base; - desc.lo += static_cast(offset); - return desc.desc_; -}} -""" - return T.cuda.func_call( - func_name, desc_val, offset, source_code=source_code, return_type="uint64" - ) + desc_halves = T.reinterpret("uint32x2", desc_val) + desc_lo = T.Shuffle([desc_halves], [0]) + T.cast(offset, "uint32") + return smem_desc_replace_lo(desc_val, desc_lo) class CopyInstType(Enum): diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py index 619385c732dd..8ea5f2916643 100644 --- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py +++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py @@ -1074,16 +1074,26 @@ def _atom_off(dim): _SWIZZLE_TO_LAYOUT = {0: 0, 1: 6, 2: 4, 3: 2, 4: 1} _krp = Evaluate(tirx_op.tvm_kernel_replace_point()) - def _make_lo_uniform(desc): - func_name = "smem_desc_make_lo_uniform_" - source_code = f""" - __forceinline__ __device__ void {func_name}(uint64_t* desc) {{ - SmemDescriptor* d = reinterpret_cast(desc); - d->lo = __shfl_sync(0xffffffff, d->lo, 0); - }} - """ - return T.cuda.func_call( - func_name, T.address_of(desc), source_code=source_code, return_type="void" + def _make_lo_uniform(desc_buf): + desc_lo = tvm.tirx.decl_buffer((1,), "uint32", name=f"{desc_buf.name}_lo", scope="local") + desc_hi = tvm.tirx.decl_buffer((1,), "uint32", name=f"{desc_buf.name}_hi", scope="local") + unpack = T.ptx.mov.b64(desc_lo[0], desc_hi[0], desc_buf[0]) + shuffle = T.ptx.shfl_sync.idx.b32( + desc_lo[0], + desc_lo[0], + T.uint32(0), + T.uint32(0x1F), + T.uint32(0xFFFFFFFF), + ) + pack = T.ptx.mov.b64(desc_buf[0], desc_lo[0], desc_hi[0]) + return SeqStmt( + [ + AllocBuffer(desc_lo), + AllocBuffer(desc_hi), + Evaluate(unpack), + Evaluate(shuffle), + Evaluate(pack), + ] ) def _make_desc(smem_buf, ldo, sdo, swizzle_val, name): @@ -1103,7 +1113,7 @@ def _make_desc(smem_buf, ldo, sdo, swizzle_val, name): ) wrap_stmts = [AllocBuffer(desc_buf), Evaluate(encode_call)] if warp_scope: - wrap_stmts.append(Evaluate(_make_lo_uniform(desc_buf[0]))) + wrap_stmts.append(_make_lo_uniform(desc_buf)) wrap_stmts.append(_krp) wrap = SeqStmt(wrap_stmts) sctx.add_post_buffer_def_stmt(smem_buf, wrap) diff --git a/python/tvm/script/tirx.pyi b/python/tvm/script/tirx.pyi index 5bf4865f2ba6..ac9e092c1208 100644 --- a/python/tvm/script/tirx.pyi +++ b/python/tvm/script/tirx.pyi @@ -37,13 +37,19 @@ class _Chain_abs: s16: _Chain_abs s32: _Chain_abs s64: _Chain_abs - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_activemask: """`activemask` — type∈{b32}""" b32: _Chain_activemask - def __call__(self, d: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_add: """`add` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -71,7 +77,7 @@ class _Chain_add: u16x2: _Chain_add u32: _Chain_add u64: _Chain_add - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_and: """`and` — type∈{pred,b16,b32,b64}""" @@ -80,7 +86,15 @@ class _Chain_and: b32: _Chain_and b64: _Chain_and pred: _Chain_and - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_applypriority: """`applypriority` — space∈{global} (opt); level∈{L2::evict_normal}""" @@ -135,7 +149,7 @@ class _Chain_atom: v4: _Chain_atom v8: _Chain_atom xor: _Chain_atom - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_bar: """`bar` — 8 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -152,7 +166,7 @@ class _Chain_bar: sync: _Chain_bar u32: _Chain_bar warp: _Chain_bar - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_barrier: """`barrier` — 9 entries sharing this mnemonic; PTX puts their difference in the operand @@ -174,7 +188,7 @@ class _Chain_barrier: sync: _Chain_barrier u32: _Chain_barrier wait: _Chain_barrier - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_bfe: """`bfe` — type∈{u32,u64,s32,s64}""" @@ -183,14 +197,33 @@ class _Chain_bfe: s64: _Chain_bfe u32: _Chain_bfe u64: _Chain_bfe - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_bfi: """`bfi` — type∈{b32,b64}""" b32: _Chain_bfi b64: _Chain_bfi - def __call__(self, f: Any, a: Any, b: Any, c: Any, d: Any, *args: Any) -> None: ... + def __call__( + self, + f: Any, + a: Any, + b: Any, + c: Any, + d: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_bfind: """`bfind` — shiftamt∈{shiftamt} (opt); type∈{u32,u64,s32,s64}""" @@ -200,7 +233,14 @@ class _Chain_bfind: shiftamt: _Chain_bfind u32: _Chain_bfind u64: _Chain_bfind - def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_bmsk: """`bmsk` — mode∈{clamp,wrap}; type∈{b32}""" @@ -208,14 +248,29 @@ class _Chain_bmsk: b32: _Chain_bmsk clamp: _Chain_bmsk wrap: _Chain_bmsk - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_brev: """`brev` — type∈{b32,b64}""" b32: _Chain_brev b64: _Chain_brev - def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_clusterlaunchcontrol: """`clusterlaunchcontrol` — 4 entries sharing this mnemonic; PTX puts their difference in @@ -238,14 +293,21 @@ class _Chain_clusterlaunchcontrol: shared__cta: _Chain_clusterlaunchcontrol try_cancel: _Chain_clusterlaunchcontrol v4: _Chain_clusterlaunchcontrol - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_clz: """`clz` — type∈{b32,b64}""" b32: _Chain_clz b64: _Chain_clz - def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_cnot: """`cnot` — type∈{b16,b32,b64}""" @@ -253,14 +315,29 @@ class _Chain_cnot: b16: _Chain_cnot b32: _Chain_cnot b64: _Chain_cnot - def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_copysign: """`copysign` — type∈{f32,f64}""" f32: _Chain_copysign f64: _Chain_copysign - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_cos: """`cos` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" @@ -268,7 +345,14 @@ class _Chain_cos: approx: _Chain_cos f32: _Chain_cos ftz: _Chain_cos - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + value: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_cp: """`cp` — 22 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -336,7 +420,7 @@ class _Chain_createpolicy: fractional: _Chain_createpolicy global_: _Chain_createpolicy range: _Chain_createpolicy - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_cvt: """`cvt` — 27 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -386,7 +470,7 @@ class _Chain_cvt: u64: _Chain_cvt u8: _Chain_cvt ue8m0x2: _Chain_cvt - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_cvt_pack: """`cvt_pack` — 2 entries sharing this mnemonic; PTX puts their difference in the operand @@ -404,7 +488,7 @@ class _Chain_cvt_pack: u2: _Chain_cvt_pack u4: _Chain_cvt_pack u8: _Chain_cvt_pack - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_cvta: """`cvta` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -421,7 +505,7 @@ class _Chain_cvta: shared__cta: _Chain_cvta to: _Chain_cvta u64: _Chain_cvta - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_discard: """`discard` — space∈{global} (opt); level∈{L2}""" @@ -450,7 +534,7 @@ class _Chain_div: u16: _Chain_div u32: _Chain_div u64: _Chain_div - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_dp2a: """`dp2a` — mode∈{lo,hi}; atype∈{u32,s32}; btype∈{u32,s32}""" @@ -459,19 +543,45 @@ class _Chain_dp2a: lo: _Chain_dp2a s32: _Chain_dp2a u32: _Chain_dp2a - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_dp4a: """`dp4a` — atype∈{u32,s32}; btype∈{u32,s32}""" s32: _Chain_dp4a u32: _Chain_dp4a - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_elect_sync: """`elect_sync` — (no modifiers)""" - def __call__(self, d: Any, p: Any, membermask: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + p: Any, + membermask: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_ex2: """`ex2` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -485,7 +595,7 @@ class _Chain_ex2: f16x2: _Chain_ex2 f32: _Chain_ex2 ftz: _Chain_ex2 - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_fence: """`fence` — 5 entries sharing this mnemonic; PTX puts their difference in the operand @@ -530,13 +640,22 @@ class _Chain_fma: rp: _Chain_fma rz: _Chain_fma sat: _Chain_fma - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_fns: """`fns` — type∈{b32}""" b32: _Chain_fns - def __call__(self, d: Any, mask: Any, base: Any, offset: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + mask: Any, + base: Any, + offset: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_getctarank: """`getctarank` — 2 entries sharing this mnemonic; PTX puts their difference in the operand @@ -546,7 +665,7 @@ class _Chain_getctarank: shared__cluster: _Chain_getctarank u32: _Chain_getctarank u64: _Chain_getctarank - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_griddepcontrol: """`griddepcontrol` — action∈{launch_dependents,wait}""" @@ -568,7 +687,14 @@ class _Chain_isspacep: shared: _Chain_isspacep shared__cluster: _Chain_isspacep shared__cta: _Chain_isspacep - def __call__(self, p: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + p: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_ld: """`ld` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -625,7 +751,7 @@ class _Chain_ld: v8: _Chain_ld volatile: _Chain_ld weak: _Chain_ld - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_ldmatrix: """`ldmatrix` — 3 entries sharing this mnemonic; PTX puts their difference in the operand @@ -648,7 +774,7 @@ class _Chain_ldmatrix: x1: _Chain_ldmatrix x2: _Chain_ldmatrix x4: _Chain_ldmatrix - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_ldu: """`ldu` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -673,7 +799,7 @@ class _Chain_ldu: u8: _Chain_ldu v2: _Chain_ldu v4: _Chain_ldu - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_lg2: """`lg2` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" @@ -681,7 +807,14 @@ class _Chain_lg2: approx: _Chain_lg2 f32: _Chain_lg2 ftz: _Chain_lg2 - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + value: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_lop3: """`lop3` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -691,7 +824,7 @@ class _Chain_lop3: and_: _Chain_lop3 b32: _Chain_lop3 or_: _Chain_lop3 - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mad: """`mad` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -715,14 +848,13 @@ class _Chain_mad: u32: _Chain_mad u64: _Chain_mad wide: _Chain_mad - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mad24: """`mad24` — mode∈{hi,lo}; sat∈{sat} (opt); type∈{u32,s32} — `.sat` on the multiply-add - lines: `.hi` mode, `.s32` type, nothing else. Both lines spell it as a syntax line - of its own -- `mad.hi.sat.s32 d, a, b, c;` (ISA 9.7.1.4) and `mad24.hi.sat.s32 d, a, - b, c;` (9.7.1.7) -- with the Notes repeating "Applies only to .s32 type in .hi - mode". + lines: `.hi` mode, `.s32` type, nothing else. Both lines spell it as a syntax line of + its own -- `mad.hi.sat.s32 d, a, b, c;` (ISA 9.7.1.4) and `mad24.hi.sat.s32 d, a, b, c;` + (9.7.1.7) -- with the Notes repeating "Applies only to .s32 type in .hi mode". """ hi: _Chain_mad24 @@ -730,7 +862,16 @@ class _Chain_mad24: s32: _Chain_mad24 sat: _Chain_mad24 u32: _Chain_mad24 - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_mapa: """`mapa` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -740,7 +881,7 @@ class _Chain_mapa: shared__cluster: _Chain_mapa u32: _Chain_mapa u64: _Chain_mapa - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_match: """`match` — 3 entries sharing this mnemonic; PTX puts their difference in the operand @@ -752,7 +893,7 @@ class _Chain_match: b32: _Chain_match b64: _Chain_match sync: _Chain_match - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_max: """`max` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -778,7 +919,7 @@ class _Chain_max: u32: _Chain_max u64: _Chain_max xorsign: _Chain_max - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mbarrier: """`mbarrier` — 25 entries sharing this mnemonic; PTX puts their difference in the operand @@ -808,7 +949,7 @@ class _Chain_mbarrier: shared__cta: _Chain_mbarrier test_wait: _Chain_mbarrier try_wait: _Chain_mbarrier - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_min: """`min` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -834,7 +975,7 @@ class _Chain_min: u32: _Chain_min u64: _Chain_min xorsign: _Chain_min - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mma: """`mma` — 12 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -879,7 +1020,7 @@ class _Chain_mma: u4: _Chain_mma u8: _Chain_mma xor: _Chain_mma - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mov: """`mov` — 11 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -900,7 +1041,7 @@ class _Chain_mov: u16: _Chain_mov u32: _Chain_mov u64: _Chain_mov - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mul: """`mul` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -929,7 +1070,7 @@ class _Chain_mul: u32: _Chain_mul u64: _Chain_mul wide: _Chain_mul - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_mul24: """`mul24` — mode∈{hi,lo}; type∈{u32,s32}""" @@ -938,7 +1079,15 @@ class _Chain_mul24: lo: _Chain_mul24 s32: _Chain_mul24 u32: _Chain_mul24 - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_multimem_ld_reduce: """`multimem_ld_reduce` — 3 entries sharing this mnemonic; PTX puts their difference in the @@ -975,7 +1124,7 @@ class _Chain_multimem_ld_reduce: v8: _Chain_multimem_ld_reduce weak: _Chain_multimem_ld_reduce xor: _Chain_multimem_ld_reduce - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_multimem_red: """`multimem_red` — 3 entries sharing this mnemonic; PTX puts their difference in the @@ -1057,7 +1206,7 @@ class _Chain_neg: s16: _Chain_neg s32: _Chain_neg s64: _Chain_neg - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_not: """`not` — type∈{pred,b16,b32,b64}""" @@ -1066,7 +1215,14 @@ class _Chain_not: b32: _Chain_not b64: _Chain_not pred: _Chain_not - def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_or: """`or` — type∈{pred,b16,b32,b64}""" @@ -1075,23 +1231,38 @@ class _Chain_or: b32: _Chain_or b64: _Chain_or pred: _Chain_or - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_popc: """`popc` — type∈{b32,b64}""" b32: _Chain_popc b64: _Chain_popc - def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_prefetch: """`prefetch` — space∈{global,local,const,param} (opt); level∈{L1,L2} (opt); evict∈{L2::evict_last,L2::evict_normal} (opt); tensormap∈{tensormap} (opt) — Each prefetch syntax line names exactly one target (PTX ISA 9.7.9.16). - `.level::eviction_priority` stays bound to `.global` on purpose: its syntax line is - `prefetch.global.level::eviction_priority`, with `.global` written in rather than - the `{.ss}` that the `ld` lines carry. Generic addressing is not offered there, so - neither is it here. + `.level::eviction_priority` stays bound to `.global` on purpose: its syntax line is + `prefetch.global.level::eviction_priority`, with `.global` written in rather than the + `{.ss}` that the `ld` lines carry. Generic addressing is not offered there, so neither + is it here. """ L1: _Chain_prefetch @@ -1121,18 +1292,26 @@ class _Chain_prmt: f4e: _Chain_prmt rc16: _Chain_prmt rc8: _Chain_prmt - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_rcp: """`rcp` — mode∈{approx,rn,rz,rm,rp}; ftz∈{ftz} (opt); type∈{f32,f64} — rcp's four syntax - lines, across two ISA subsections. rcp.approx{.ftz}.f32 d, a; - rcp.rnd{.ftz}.f32 d, a; (9.7.3.13) rcp.rnd.f64 d, a; - rcp.approx.ftz.f64 d, a; (9.7.3.14) The ISA gives - the last one a subsection of its own because it is a different *computation* -- a - gross approximation off the top 20 mantissa bits, with its own corner-case table -- - but its syntax is one more cell of this grid, and the shape (`d, a`) is unchanged. - So it lives here, with the mandatory `.ftz` of its syntax line enforced below rather - than by a second entry that would render identically. + lines, across two ISA subsections. rcp.approx{.ftz}.f32 d, a; rcp.rnd{.ftz}.f32 d, a; + (9.7.3.13) rcp.rnd.f64 d, a; rcp.approx.ftz.f64 d, a; (9.7.3.14) The ISA gives the last + one a subsection of its own because it is a different *computation* -- a gross + approximation off the top 20 mantissa bits, with its own corner-case table -- but its + syntax is one more cell of this grid, and the shape (`d, a`) is unchanged. So it lives + here, with the mandatory `.ftz` of its syntax line enforced below rather than by a + second entry that would render identically. """ approx: _Chain_rcp @@ -1143,7 +1322,14 @@ class _Chain_rcp: rn: _Chain_rcp rp: _Chain_rcp rz: _Chain_rcp - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + value: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_red: """`red` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -1230,7 +1416,7 @@ class _Chain_redux_sync: s32: _Chain_redux_sync u32: _Chain_redux_sync xor: _Chain_redux_sync - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_rem: """`rem` — type∈{u16,u32,u64,s16,s32,s64}""" @@ -1241,7 +1427,15 @@ class _Chain_rem: u16: _Chain_rem u32: _Chain_rem u64: _Chain_rem - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_rsqrt: """`rsqrt` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32,f64}""" @@ -1250,7 +1444,14 @@ class _Chain_rsqrt: f32: _Chain_rsqrt f64: _Chain_rsqrt ftz: _Chain_rsqrt - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + value: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_sad: """`sad` — type∈{u16,u32,u64,s16,s32,s64}""" @@ -1261,7 +1462,16 @@ class _Chain_sad: u16: _Chain_sad u32: _Chain_sad u64: _Chain_sad - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_selp: """`selp` — type∈{b16,b32,b64,u16,u32,u64,s16,s32,s64,f32,f64}""" @@ -1277,7 +1487,16 @@ class _Chain_selp: u16: _Chain_selp u32: _Chain_selp u64: _Chain_selp - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_set: """`set` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -1321,7 +1540,7 @@ class _Chain_set: u32: _Chain_set u64: _Chain_set xor: _Chain_set - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_setmaxnreg: """`setmaxnreg` — action∈{inc,dec}; sync∈{sync}; aligned∈{aligned}; type∈{u32}""" @@ -1375,7 +1594,7 @@ class _Chain_setp: u32: _Chain_setp u64: _Chain_setp xor: _Chain_setp - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_shf: """`shf` — dir∈{l,r}; mode∈{clamp,wrap}; type∈{b32}""" @@ -1385,7 +1604,16 @@ class _Chain_shf: l: _Chain_shf r: _Chain_shf wrap: _Chain_shf - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_shfl_sync: """`shfl_sync` — 2 entries sharing this mnemonic; PTX puts their difference in the operand @@ -1398,7 +1626,7 @@ class _Chain_shfl_sync: down: _Chain_shfl_sync idx: _Chain_shfl_sync up: _Chain_shfl_sync - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_shl: """`shl` — type∈{b16,b32,b64}""" @@ -1406,7 +1634,15 @@ class _Chain_shl: b16: _Chain_shl b32: _Chain_shl b64: _Chain_shl - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_shr: """`shr` — type∈{b16,b32,b64,u16,u32,u64,s16,s32,s64}""" @@ -1420,7 +1656,15 @@ class _Chain_shr: u16: _Chain_shr u32: _Chain_shr u64: _Chain_shr - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_sin: """`sin` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" @@ -1428,14 +1672,21 @@ class _Chain_sin: approx: _Chain_sin f32: _Chain_sin ftz: _Chain_sin - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + value: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_slct: """`slct` — ftz∈{ftz} (opt); dtype∈{b16,b32,b64,u16,u32,u64,s16,s32,s64,f32,f64}; ctype∈{s32,f32} — slct's two lines (ISA 9.7.6.4), which differ only in the selector - type. slct.dtype.s32 d, a, b, c; slct{.ftz}.dtype.f32 d, a, b, - c; `.ftz` is spelled on the .f32 selector line alone -- there is nothing to - flush when the sign being tested is an integer's. + type. slct.dtype.s32 d, a, b, c; slct{.ftz}.dtype.f32 d, a, b, c; `.ftz` is spelled on + the .f32 selector line alone -- there is nothing to flush when the sign being tested is + an integer's. """ b16: _Chain_slct @@ -1450,14 +1701,22 @@ class _Chain_slct: u16: _Chain_slct u32: _Chain_slct u64: _Chain_slct - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + c: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_sqrt: """`sqrt` — mode∈{approx,rn,rz,rm,rp}; ftz∈{ftz} (opt); type∈{f32,f64} — sqrt's three lines - (PTX ISA 9.7.3.15). sqrt.approx{.ftz}.f32 d, a; sqrt.rnd{.ftz}.f32 d, a; - sqrt.rnd.f64 d, a; Unlike rcp, there is no f64 approximation at any - spelling -- 9.7.3.15 is the whole of sqrt, and it offers `.approx` on the .f32 line - only. + (PTX ISA 9.7.3.15). sqrt.approx{.ftz}.f32 d, a; sqrt.rnd{.ftz}.f32 d, a; sqrt.rnd.f64 d, + a; Unlike rcp, there is no f64 approximation at any spelling -- 9.7.3.15 is the whole of + sqrt, and it offers `.approx` on the .f32 line only. """ approx: _Chain_sqrt @@ -1468,7 +1727,14 @@ class _Chain_sqrt: rn: _Chain_sqrt rp: _Chain_sqrt rz: _Chain_sqrt - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + value: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_st: """`st` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -1601,7 +1867,7 @@ class _Chain_sub: u16: _Chain_sub u32: _Chain_sub u64: _Chain_sub - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_szext: """`szext` — mode∈{clamp,wrap}; type∈{u32,s32}""" @@ -1610,7 +1876,15 @@ class _Chain_szext: s32: _Chain_szext u32: _Chain_szext wrap: _Chain_szext - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_tanh: """`tanh` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, @@ -1623,10 +1897,10 @@ class _Chain_tanh: f16: _Chain_tanh f16x2: _Chain_tanh f32: _Chain_tanh - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_tcgen05: - """`tcgen05` — 18 entries sharing this mnemonic; PTX puts their difference in the operand + """`tcgen05` — 20 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (dst, ncols); (taddr, ncols); (); (*__operands); (taddr, s_desc); (d_tmem, a_desc, b_desc, idesc, sfa_tmem, sfb_tmem, enable_input_d); (d_tmem, a_tmem, b_desc, idesc, sfa_tmem, sfb_tmem, enable_input_d); (d_tmem, a_desc, @@ -1641,6 +1915,8 @@ class _Chain_tcgen05: b64: _Chain_tcgen05 b6x16_p32: _Chain_tcgen05 b8x16: _Chain_tcgen05 + block16: _Chain_tcgen05 + block32: _Chain_tcgen05 block_scale: _Chain_tcgen05 commit: _Chain_tcgen05 cp: _Chain_tcgen05 @@ -1684,7 +1960,7 @@ class _Chain_tcgen05: x4: _Chain_tcgen05 x64: _Chain_tcgen05 x8: _Chain_tcgen05 - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_tensormap_cp_fenceproxy: """`tensormap_cp_fenceproxy` — dst∈{global}; src∈{shared::cta}; proxy∈{tensormap::generic}; @@ -1737,7 +2013,14 @@ class _Chain_testp: notanumber: _Chain_testp number: _Chain_testp subnormal: _Chain_testp - def __call__(self, p: Any, a: Any, *args: Any) -> None: ... + def __call__( + self, + p: Any, + a: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _Chain_vote_sync: """`vote_sync` — 2 entries sharing this mnemonic; PTX puts their difference in the operand @@ -1750,7 +2033,7 @@ class _Chain_vote_sync: ballot: _Chain_vote_sync pred: _Chain_vote_sync uni: _Chain_vote_sync - def __call__(self, *args: Any) -> None: ... + def __call__(self, *args: Any, pred: Any = None, preserve_dst: bool = False) -> None: ... class _Chain_wgmma: """`wgmma` — 19 entries sharing this mnemonic; PTX puts their difference in the operand @@ -1899,7 +2182,15 @@ class _Chain_xor: b32: _Chain_xor b64: _Chain_xor pred: _Chain_xor - def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + def __call__( + self, + d: Any, + a: Any, + b: Any, + *args: Any, + pred: Any = None, + preserve_dst: bool = False, + ) -> None: ... class _PTX: abs: _Chain_abs diff --git a/src/tirx/transform/bind_target.cc b/src/tirx/transform/bind_target.cc index a03b8a252d2d..e64670b451b7 100644 --- a/src/tirx/transform/bind_target.cc +++ b/src/tirx/transform/bind_target.cc @@ -114,7 +114,8 @@ class FunctionClassifierVisitor : public StmtExprVisitor { } void VisitStmt_(const AttrStmtNode* op) final { - if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) { + if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread || + op->attr_key == attr::kDeviceEntry) { // Enter GPU scope for thread extent and virtual thread attributes bool last_is_under_gpu_scope = is_under_gpu_scope_; is_under_gpu_scope_ = true; @@ -199,7 +200,8 @@ class CallSubstitutor : public StmtExprMutator { } Stmt VisitStmt_(const AttrStmtNode* op) final { - if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread) { + if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread || + op->attr_key == attr::kDeviceEntry) { // Enter GPU scope for thread extent and virtual thread attributes bool last_is_under_gpu_scope = is_under_gpu_scope_; is_under_gpu_scope_ = true; diff --git a/tests/python/tirx-transform/test_tir_transform_helpers.py b/tests/python/tirx-transform/test_tir_transform_helpers.py index 1932098c5574..982c13a51325 100644 --- a/tests/python/tirx-transform/test_tir_transform_helpers.py +++ b/tests/python/tirx-transform/test_tir_transform_helpers.py @@ -284,6 +284,61 @@ def main( tvm.ir.assert_structural_equal(After, Expected) +def test_bind_target_with_tirx_device_entry(): + """BindTarget classifies calls after T.device_entry as device-side.""" + + @I.ir_module + class Before: + @T.prim_func(private=True) + def add(a: T.int32, b: T.int32) -> T.int32: + return a + b + + @T.prim_func + def main(A: T.Buffer((1,), "int32")): + T.func_attr({"global_symbol": "main"}) + host_value: T.let[T.int32] = Before.add(1, 2) + T.device_entry() + tx = T.thread_id([1]) + A[tx] = Before.add(host_value, 3) + + @I.ir_module + class Expected: + @T.prim_func(private=True) + def add(a: T.int32, b: T.int32) -> T.int32: + T.func_attr({"target": T.target({"arch": "sm_100a", "kind": "cuda"})}) + return a + b + + @T.prim_func(private=True) + def add_host(a: T.int32, b: T.int32) -> T.int32: + T.func_attr({"target": T.target({"kind": "llvm", "opt-level": 0})}) + return a + b + + @T.prim_func + def main(A: T.Buffer((1,), "int32")): + T.func_attr( + { + "global_symbol": "main", + "target": T.target( + { + "arch": "sm_100a", + "host": {"kind": "llvm", "opt-level": 0}, + "kind": "cuda", + } + ), + } + ) + host_value: T.let[T.int32] = Expected.add_host(1, 2) + T.device_entry() + tx = T.thread_id([1]) + A[tx] = Expected.add(host_value, 3) + + target = tvm.target.Target( + {"kind": "cuda", "arch": "sm_100a"}, host={"kind": "llvm", "opt-level": 0} + ) + After = tvm.tirx.transform.BindTarget(target)(Before) + tvm.ir.assert_structural_equal(After, Expected) + + def test_filter_primfunc(): mod = MockModule assert mod diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py index 76449d571af8..f08d6ea419eb 100644 --- a/tests/python/tirx/codegen/test_ptx_dialect.py +++ b/tests/python/tirx/codegen/test_ptx_dialect.py @@ -201,6 +201,49 @@ def kernel(a_ptr: T.handle): assert "setp.ne.b32 p, %2, 0; @p red.relaxed.gpu.global.add.u32 [%0], %1;" in src +@requires_nvcc +def test_ptx_predicated_destination_preserves_old_value(): + @T.prim_func + def kernel(a_ptr: T.handle, out_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "float32") + Out = T.match_buffer(out_ptr, (32,), "float32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + value: T.float32 = T.float32(0) + pred: T.uint32 = T.cast(tx == 0, "uint32") + T.ptx.ld.global_.f32(value, A.ptr_to([0]), pred=pred, preserve_dst=True) + T.ptx.ex2.approx.ftz.f32(value, value, pred=pred, preserve_dst=True) + Out[tx] = value + + src = _cuda_source(kernel) + assert "tvm_builtin_ptx_ld_global_f32_pred_keep" in src + assert "tvm_builtin_ptx_ex2_approx_ftz_f32_pred_keep" in src + assert '"+f"(__d)' in src + _assert_ptxas_ok(src) + + +@requires_nvcc +def test_ptx_predicated_destination_is_undefined_by_default(): + @T.prim_func + def kernel(a_ptr: T.handle, out_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "float32") + Out = T.match_buffer(out_ptr, (32,), "float32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + value = T.local_scalar("float32") + pred: T.uint32 = T.cast(tx == 0, "uint32") + T.ptx.ld.global_.f32(value, A.ptr_to([0]), pred=pred) + Out[tx] = T.if_then_else(pred != 0, value, T.float32(0)) + + src = _cuda_source(kernel) + assert "tvm_builtin_ptx_ld_global_f32_pred_undef" in src + assert '"=f"(__d)' in src + assert '"+f"(__d)' not in src + _assert_ptxas_ok(src) + + def test_ptx_string_form_matches_chain(): chain_call = T.ptx.ld.global_.acquire.gpu.b32 string_call = T.ptx["ld.acquire.gpu.global.b32"] @@ -308,15 +351,6 @@ def rvalue_destination(a_ptr: T.handle): T.device_entry() T.ptx.ld.global_.b32(T.uint32(0), A.ptr_to([0])) - # @p is rejected on any instruction that writes a destination. - with pytest.raises((ValueError, tvm.error.DiagnosticError), match="without a destination"): - - @T.prim_func - def predicated_destination(out: T.Buffer((1,), "uint32"), a_ptr: T.handle): - A = T.match_buffer(a_ptr, (1,), "uint32") - T.device_entry() - T.ptx.ld.global_.b32(out[0], A.ptr_to([0]), pred=T.uint32(1)) - def test_ptx_register_group_codegen(): """A `.lanes > 1` operand renders as braces in the asm, flat params in C. @@ -1377,6 +1411,42 @@ def kernel(a_ptr: T.handle): tvm.ir.assert_structural_equal(kernel, reparsed) +@requires_nvcc +def test_ptx_tcgen05_mma_block_size_form(): + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (32,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + if tx == 0: + tmem = T.local_scalar("uint32") + desc = T.local_scalar("uint64") + idesc = T.local_scalar("uint32") + flag = T.local_scalar("uint32") + T.ptx["tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32"]( + tmem, desc, desc, idesc, tmem, tmem, T.ptx.pred(flag) + ) + A[tx] = A[tx] + + src = _cuda_source(kernel) + assert "tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32" in src + _assert_ptxas_ok(src, arch="sm_100a") + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="mxf4.*block32"): + + @T.prim_func + def invalid_mxf4_block16(): + T.device_entry() + tmem = T.local_scalar("uint32") + desc = T.local_scalar("uint64") + idesc = T.local_scalar("uint32") + flag = T.local_scalar("uint32") + T.ptx["tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block16"]( + tmem, desc, desc, idesc, tmem, tmem, T.ptx.pred(flag) + ) + + def test_ptx_pred_operand_rejects_untagged_integer(): """An untagged integer at a `.pred` position is refused, by name. @@ -2007,11 +2077,15 @@ def test_ptx_coercion_ir_forms(): call = T.ptx.st.release.gpu.global_.b32(global_ptr, val, pred=flag) assert call.args[2].same_as(flag) assert len(call.args) == 2 + 1 + 8 + 1 # operands + pred + slot tokens + marker - # @p on an instruction with a destination is rejected: a false predicate - # leaves it unwritten while "=" tells nvcc its prior value is dead. - dst = tvm.tirx.Var("d", "uint32") - with pytest.raises(ValueError, match="without a destination"): - T.ptx.ld.global_.b32(dst, global_ptr, pred=flag) + out = tvm.tirx.decl_buffer((1,), "uint32", name="out", scope="local") + call = T.ptx.ld.global_.b32(out[0], global_ptr, pred=flag) + assert str(call.args[-1]).strip('"') == "pred" + call = T.ptx.ld.global_.b32(out[0], global_ptr, pred=flag, preserve_dst=True) + assert str(call.args[-1]).strip('"') == "pred,keep" + with pytest.raises(ValueError, match="requires pred"): + T.ptx.ld.global_.b32(out[0], global_ptr, preserve_dst=True) + with pytest.raises(ValueError, match="requires a written destination"): + T.ptx.st.release.gpu.global_.b32(global_ptr, val, pred=flag, preserve_dst=True) # fp16/bf16 dtypes bring in __half / __nv_bfloat16 and their bit-cast helpers. @@ -2201,7 +2275,7 @@ def test_ptx_all_variants_render_unique(): or f"; {opcode};" in source ) total += not predicated # a @p twin is not a separate variant - assert total == 200002 # update when the table grows + assert total == 200018 # update when the table grows def test_ptx_no_instruction_registered_twice(): diff --git a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py index ad5942aa9026..02d972ce582d 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py @@ -3057,10 +3057,9 @@ def test_gemm_smem_desc_modes_codegen(smem_desc): base address, selected by the ``smem_desc`` config. - ``hoist`` (default): allocate + encode one descriptor per operand - (``descA`` / ``descB``), then add the per-MMA 16B offset via - ``smem_desc_add_16B_offset``. This builder uses a single-thread scope, - where the encoding thread is also the consumer and no warp shuffle is - valid or needed. + (``descA`` / ``descB``), then update the descriptor address lane for each + MMA. This builder uses a single-thread scope, where the encoding thread + is also the consumer and no warp shuffle is valid or needed. - ``recompute``: build the full descriptor inline per MMA (``_uniform_desc``) with no allocated/encoded descriptor cell — trades a few ALU ops for one fewer live register on the hot path. @@ -3084,18 +3083,12 @@ def test_gemm_smem_desc_modes_codegen(smem_desc): if smem_desc == "hoist": assert "encode_matrix_descriptor" in src, "hoist mode must encode a descriptor" - assert "smem_desc_make_lo_uniform" not in src, "hoist mode must not warp-shuffle" - assert "smem_desc_add_16B_offset" in src, "hoist mode must add the per-MMA 16B offset" elif smem_desc == "local_hoist": assert "descA_local" in src and "descB_local" in src - assert "smem_desc_add_16B_offset" in src assert "encode_matrix_descriptor" in src elif smem_desc == "encode": assert "encode_matrix_descriptor" in src - assert "smem_desc_add_16B_offset" not in src else: - assert "smem_desc_make_lo_uniform" not in src, "recompute mode must not hoist a descriptor" - assert "smem_desc_add_16B_offset" not in src, "recompute mode must not add a 16B offset" assert "encode_matrix_descriptor" not in src, "recompute mode must not encode a descriptor" @@ -3736,10 +3729,10 @@ def test_gemm_tcgen05_hoisted_descriptor_uniformization(scope_kind, expect_unifo callback_text = str(sctx.callbacks["post_buffer_def_stmt"]) assert "encode_matrix_descriptor" in callback_text - assert ("smem_desc_make_lo_uniform" in callback_text) == expect_uniform + assert ("T.ptx.shfl_sync" in callback_text) == expect_uniform assert ("elect_sync" in impl.script()) == expect_uniform if expect_uniform: - assert callback_text.index("smem_desc_make_lo_uniform") < callback_text.index( + assert callback_text.index("T.ptx.shfl_sync") < callback_text.index( "tvm_kernel_replace_point" )