This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 93098fbbbe [TIRx][CUDA] Replace source helpers with typed PTX forms
(#20140)
93098fbbbe is described below
commit 93098fbbbe95aea6663d64fd038f1d566582b2b1
Author: Hongyi Jin <[email protected]>
AuthorDate: Mon Aug 17 15:40:40 2026 -0400
[TIRx][CUDA] Replace source helpers with typed PTX forms (#20140)
## Motivation
TIRx low-level kernels need to reject `tirx.cuda.func_call`: it embeds
arbitrary CUDA source outside the typed IR and PTX dialect, so the
compiler cannot validate operand types, destination liveness,
instruction legality, or target placement.
The downstream audit in mlc-ai/TIRx-kernels#62 found four missing
general capabilities: inactive-path liveness for predicated
destinations, the MXF4 block32 MMA spelling, typed SMEM descriptor
operations, and correct target binding for private PrimFunc helpers
called from a `T.device_entry` region. This PR adds those capabilities
without kernel-specific hooks.
## Change-by-change rationale
### 1. One `preserve_dst` flag for predicated destinations
Files: `python/tvm/backend/cuda/ptx/engine.py`, `render.py`, and
`table.py`.
- Change: PTX calls use one Boolean choice, `preserve_dst=False` by
default and `preserve_dst=True` when the old destination must remain
live.
- Why one flag is sufficient: a false PTX predicate always suppresses
the instruction write. The only compiler-level choice is whether the
inline-asm destination also needs its previous C++ value as an input.
`False` means it does not; `True` means it does.
- `preserve_dst=False`: use the normal write-only `"="` constraint. The
inactive value is undefined to the caller and must be merged or guarded
before use.
- `preserve_dst=True`: use a read-write `"+"` constraint and seed any
carrier bridge from the old destination. The inactive path therefore
retains that value.
- Why not two flags: `undefined_dst` was exactly the complement of
`preserve_dst`, so exposing both created contradictory and invalid
Boolean combinations. The renderer now derives the default undefined
policy from `entry.has_dst` plus `preserve_dst=False`.
- Validation: `preserve_dst=True` requires both `pred=...` and an
ordinary written destination. Accumulators already have `rw="rw"` and
need no extra flag; stores and other destination-free instructions
reject it.
- IR/codegen identity: the marker is now only `pred` or `pred,keep`. The
renderer still emits distinct `_pred_undef` and `_pred_keep` helper
names, preserving helper identity and the exact `"="` versus `"+"` ABI.
Hardware behavior is unchanged in both cases: predicate false means the
PTX instruction does not execute. The flag only states whether the
surrounding compiler must keep the previous destination value live.
The downstream GDN path needs both choices in one expression chain: its
predicated shared loads use the default undefined destination until a
`selp`, while the later predicated `ex2` uses `preserve_dst=True` to
retain the zero selected for inactive lanes.
### 2. Explicit MXF4 block32 `tcgen05.mma` form
File: `python/tvm/backend/cuda/ptx/table.py`.
- Change: add SS and TS entries for
`tcgen05.mma.*.block_scale.block16/block32` with a dedicated
`block_size` modifier.
- Why: the paged MQA kernel needs
`tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32`. The table
previously covered the scale-vector spelling but not the explicit
scale-block-size spelling, forcing the kernel to hide a valid
instruction in a CUDA source helper.
- ISA constraint: a validator encodes the legal matrix instead of
accepting every token combination: `mxf8f6f4` and `mxf4` allow only
block32; `mxf4nvf4` allows block16 and block32. Keeping this as a
separate table form preserves exact instruction spelling and rejects
invalid `mxf4.block16` at IR construction time.
### 3. Typed SMEM descriptor address updates
File: `python/tvm/backend/cuda/tile_primitive/common.py`.
- Change: introduce `smem_desc_replace_lo` and implement
`smem_desc_add_16B_offset` by reinterpreting the descriptor as
`uint32x2`, updating only lane 0, and repacking it.
- Why: the low 32 bits are the address lane and the high 32 bits contain
descriptor control fields. A generic 64-bit add could carry into those
control fields. The previous implementation used a CUDA union inside
`func_call`; the typed implementation makes the no-carry uint32 behavior
explicit and gives descriptor address replacement one reusable
authority.
### 4. Typed descriptor warp uniformization
Files: `python/tvm/backend/cuda/lang/smem_desc.py` and
`python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py`.
- Change: replace `smem_desc_make_lo_uniform` CUDA helpers with
`mov.b64` unpacking, `shfl.sync.idx.b32` of the low lane, and `mov.b64`
repacking while preserving the high lane.
- Why: only the address lane needs to be broadcast from lane 0; the
descriptor metadata must remain unchanged. Expressing the sequence in
typed PTX removes an arbitrary source boundary and lets normal
operand/type validation see the operation.
- Lowering detail: the tile-primitive implementation emits explicit
local buffers and a `SeqStmt` because it constructs IR nodes below the
script layer. The script-level `SmemDescriptor` method emits the same
typed operation sequence at its abstraction layer.
### 5. `T.device_entry` as the device-side target boundary
File: `src/tirx/transform/bind_target.cc`.
- Change: both the function classifier and call substitutor treat
`attr::kDeviceEntry` as entering GPU scope, alongside thread extents and
virtual-thread attributes.
- Why: a TIRx PrimFunc may contain host preparation followed by a device
body marked by `T.device_entry`. A private typed helper called from that
device body must receive the CUDA target. Without recognizing the
marker, BindTarget can classify or clone the helper as host-side,
producing a cross-target call instead of a device helper.
- Scope: this does not globally classify the whole function as device
code; it changes classification only while visiting the body dominated
by the existing canonical device-entry marker.
This enables the downstream weighted-ReLU reduction and acquire polling
loop to remain typed private PrimFuncs rather than CUDA source helpers.
### 6. Generated PTX typing surface
Files: `python/tvm/backend/cuda/ptx/gen_stubs.py` and
`python/tvm/script/tirx.pyi`.
- Change: generated PTX chain signatures expose `pred` and, for families
with an ordinary written destination, the single `preserve_dst: bool =
False` flag.
- Why: the static typing surface must match the runtime dialect API. The
stub remains generated from the instruction table rather than becoming a
second hand-maintained source of truth.
- Generator fix: line wrapping now checks the length of the final
indented stub line directly. Removing the redundant flag exposed an
off-by-one boundary in the previous indirect threshold.
### 7. Regression tests
Files: `tests/python/tirx/codegen/test_ptx_dialect.py`,
`tests/python/tirx-transform/test_tir_transform_helpers.py`, and
`tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py`.
- Predicated destinations: verify default write-only versus opted-in
read-write constraints, `_pred_undef` versus `_pred_keep` helper
identity, IR markers and round trips, invalid flag placement, and ptxas
acceptance.
- MXF4 block size: verify the exact block32 opcode reaches generated
CUDA/PTX and ptxas, and verify invalid `mxf4.block16` is rejected.
- BindTarget: use one private helper from both host code and the region
after `T.device_entry`, then verify separate LLVM and CUDA-targeted
functions are produced.
- Descriptor lowering: verify all descriptor modes still build and that
warp-scoped uniformization appears as typed `T.ptx.shfl_sync` before the
kernel replacement point instead of checking for source-helper names.
## Downstream evidence
The downstream migration is mlc-ai/TIRx-kernels#62.
- GDN predication produced byte-identical SASS; fixed-config
baseline/candidate ratios were 0.9998x, 1.0001x, and 1.0098x. After
simplifying to one flag, all three configs passed again against the
frozen FlashInfer oracle.
- The typed private weighted-ReLU helper retained 168 registers, zero
stack/local memory, and the intended packed arithmetic; ratios were
1.0011x and 1.0116x.
- The typed acquire polling helper retained the acquire load/control
structure and completed TP1 and TP4 correctness; ratios were about
1.0003x and 0.9955x.
- The affected downstream correctness matrix passed, and the low-level
IR registry audit no longer needs source helpers for any of these cases.
## Validation
- `python -m pytest -q tests/python/tirx/codegen/test_ptx_dialect.py`:
47 passed, 32 skipped
- BindTarget regression selection: 1 passed
- descriptor lowering regression selection: 6 passed
- downstream GDN frozen-oracle correctness: 3 configs passed
- downstream low-level IR checker: 2 passed
- changed-file pre-commit: passed in both repositories
- generated `tirx.pyi`: byte-identical to `python -m
tvm.backend.cuda.ptx.gen_stubs` output
## Related
Downstream kernel migration:
https://github.com/mlc-ai/TIRx-kernels/pull/62
---
python/tvm/backend/cuda/lang/smem_desc.py | 19 +-
python/tvm/backend/cuda/ptx/engine.py | 49 ++-
python/tvm/backend/cuda/ptx/gen_stubs.py | 10 +-
python/tvm/backend/cuda/ptx/render.py | 30 +-
python/tvm/backend/cuda/ptx/table.py | 59 ++-
python/tvm/backend/cuda/tile_primitive/common.py | 26 +-
.../cuda/tile_primitive/gemm_async/tcgen05.py | 32 +-
python/tvm/script/tirx.pyi | 487 ++++++++++++++++-----
src/tirx/transform/bind_target.cc | 6 +-
.../tirx-transform/test_tir_transform_helpers.py | 55 +++
tests/python/tirx/codegen/test_ptx_dialect.py | 104 ++++-
.../cuda/gemm_async/test_gemm_async.py | 17 +-
12 files changed, 699 insertions(+), 195 deletions(-)
diff --git a/python/tvm/backend/cuda/lang/smem_desc.py
b/python/tvm/backend/cuda/lang/smem_desc.py
index b1d850abd2..d6244fd73a 100644
--- a/python/tvm/backend/cuda/lang/smem_desc.py
+++ b/python/tvm/backend/cuda/lang/smem_desc.py
@@ -43,13 +43,14 @@ class SmemDescriptor:
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])
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 43a9d4d7a7..3e29523af2 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -118,11 +118,13 @@ def _make_codegen(entry: InstructionEntry):
# 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 _make_codegen(entry: InstructionEntry):
# 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 @@ class _InstrChain:
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 @@ class _InstrChain:
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 @@ class _InstrChain:
# 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 9d31cb723b..2d8c22222c 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 6a57fd0fbc..96bdab0a91 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 71c299a919..e55c4b095b 100644
--- a/python/tvm/backend/cuda/ptx/table.py
+++ b/python/tvm/backend/cuda/ptx/table.py
@@ -420,13 +420,12 @@ class InstructionEntry:
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 @@ _ENTRIES = [
# 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 @@ _ENTRIES = [
)
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 157753d2fe..d229809631 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<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):
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 619385c732..8ea5f29166 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 gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
_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<SmemDescriptor*>(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 gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
)
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 5bf4865f2b..ac9e092c12 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 a03b8a252d..e64670b451 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 1932098c55..982c13a513 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 test_bind_target_with_device_host_call_same_func():
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 76449d571a..f08d6ea419 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -201,6 +201,49 @@ def test_ptx_predication_codegen():
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 test_ptx_destination_errors():
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 test_ptx_pred_operand_roundtrip():
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 ad5942aa90..02d972ce58 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"
)
