[TIRx][CUDA] Replace source helpers with typed PTX forms

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<SmemDescriptor*>(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])

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<i>"); 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<i>"); 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.

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}"

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))

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}",

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<uint32_t>(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):