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 f20fa692d5 [TIRx][CUDA] Add PTX address expressions with immediate
byte offsets (#20153)
f20fa692d5 is described below
commit f20fa692d5dd71d875a9e310eae3e754169888fb
Author: Bohan Hou <[email protected]>
AuthorDate: Tue Aug 18 12:34:09 2026 -0400
[TIRx][CUDA] Add PTX address expressions with immediate byte offsets
(#20153)
This PR adds `T.ptx.addr(base, byte_offset)` to the TIRx PTX dialect: a
pure expression that folds a compile-time signed byte displacement into
the PTX address operand, rendering as `[%N+imm]` instead of requiring a
separate address computation before the instruction.
- `tirx.ptx.addr` is an expression only in the outer PTX call's IR: the
helper still receives the coerced base register, while the displacement
becomes renderer metadata baked into the instruction text and the helper
name (`_addr<slot>_p<imm>` / `_m<imm>`).
- Table-level `allow_imm_offset` classification of address slots, with
validation that rejects immediate offsets on operand classes that cannot
take them (e.g. `tmem` addresses).
- Immediate operands are now validated to be compile-time `IntImm` at
CUDA codegen, with an actionable error pointing at explicitly-unrolled
loops.
- Displacements are range-checked to int32; zero offsets normalize to
the bare form so existing helper names are untouched.
Also includes a small test fix:
`test_tirx_kernels_registry_correctness.py` accepts both the old and new
MegaMoE kernel registry names (`deepgemm_fp8_fp4_mega_moe` /
`sm100_fp8_fp4_mega_moe`), so the test works against tirx-kernels
checkouts from either side of the rename.
Tested with `tests/python/tirx/codegen/test_ptx_addr.py` (new, 12
cases), plus the full `tests/python/tirx/` suite on sm100.
---
python/tvm/backend/cuda/ptx/__init__.py | 5 +-
python/tvm/backend/cuda/ptx/engine.py | 176 ++++++++++-
python/tvm/backend/cuda/ptx/gen_stubs.py | 1 +
python/tvm/backend/cuda/ptx/render.py | 61 +++-
python/tvm/backend/cuda/ptx/table.py | 188 ++++++-----
python/tvm/script/tirx.pyi | 1 +
tests/python/tirx/codegen/test_ptx_addr.py | 351 +++++++++++++++++++++
tests/python/tirx/codegen/test_ptx_dialect.py | 87 ++++-
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 19 ++
.../tirx/test_tirx_kernels_registry_correctness.py | 9 +-
10 files changed, 793 insertions(+), 105 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/__init__.py
b/python/tvm/backend/cuda/ptx/__init__.py
index ec1b48643b..e30175828a 100644
--- a/python/tvm/backend/cuda/ptx/__init__.py
+++ b/python/tvm/backend/cuda/ptx/__init__.py
@@ -23,9 +23,10 @@ registers every table entry as a TVM Op with a generic
codegen; the
``register_backend()`` via ``script_namespaces()``.
"""
-from .engine import PTXNamespace, register_table
+from .engine import PTXNamespace, register_addr, register_table
from .table import TABLE
+register_addr()
register_table(TABLE)
-__all__ = ["TABLE", "PTXNamespace", "register_table"]
+__all__ = ["TABLE", "PTXNamespace", "register_addr", "register_table"]
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 3e29523af2..2b5069491f 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -40,6 +40,7 @@ per-instruction generated or hand-written code:
from tvm.backend.cuda.codegen.registry import register_codegen
from tvm.backend.cuda.codegen.utils import parse_str
from tvm.backend.cuda.op import cuda_cvta_generic_to_shared, cuda_func_call
+from tvm.ir import Call
from tvm.ir.op import register_op_attr
from tvm.ir.type import PointerType, PrimType
from tvm.runtime import const
@@ -63,6 +64,22 @@ from .table import (
# It is also the honest answer: "do not touch my instruction" is exactly the
# contract a hand-written PTX call wants.
_EFFECT_OPAQUE = CallEffectKind.Opaque.value
+_EFFECT_PURE = CallEffectKind.Pure.value
+_ADDR_OP_NAME = "tirx.ptx.addr"
+_INT32_MIN = -(1 << 31)
+_INT32_MAX = (1 << 31) - 1
+_INTEGER_DTYPES = frozenset(
+ {
+ "int8",
+ "int16",
+ "int32",
+ "int64",
+ "uint8",
+ "uint16",
+ "uint32",
+ "uint64",
+ }
+)
# ---------------------------------------------------------------------------
# Registration (import time)
@@ -91,6 +108,22 @@ def register_table(table: dict[str, InstructionEntry]) ->
None:
register_codegen(f"ptx.{entry.name}")(_make_codegen(entry))
+def register_addr() -> None:
+ """Register the pure address-expression op consumed by PTX instructions."""
+ register_op_attr(_ADDR_OP_NAME, "TCallEffectKind", _EFFECT_PURE)
+ register_op_attr(_ADDR_OP_NAME, "TScriptPrinterName", "ptx.addr", level=20)
+ register_op_attr(_ADDR_OP_NAME, "TIRxOpCategory", "device_intrin")
+ register_op_attr(_ADDR_OP_NAME, "TDeviceIntrinsicNamespace", "ptx")
+ register_codegen("ptx.addr")(_unconsumed_addr_codegen)
+
+
+def _unconsumed_addr_codegen(*_args):
+ raise ValueError(
+ "T.ptx.addr(...) must be consumed by a PTX address operand that
supports "
+ "immediate byte offsets"
+ )
+
+
# ---------------------------------------------------------------------------
# Codegen (compile time): table -> asm volatile helper
# ---------------------------------------------------------------------------
@@ -109,6 +142,32 @@ def arg_dtype(value) -> str:
return ty.dtype if isinstance(ty, PrimType) else type(value).__name__
+def _is_addr_call(value) -> bool:
+ return isinstance(value, Call) and getattr(value.op, "name", None) ==
_ADDR_OP_NAME
+
+
+def _codegen_addr_offset(entry, slot, value) -> tuple[object, int]:
+ """Unpack one nested ``tirx.ptx.addr`` call at CUDA codegen time."""
+ if not slot.allow_imm_offset:
+ raise ValueError(f"{entry.name}: operand '{slot.name}' does not
support T.ptx.addr(...)")
+ if len(value.args) != 2:
+ raise ValueError("malformed tirx.ptx.addr call: expected base and
byte_offset")
+ base, offset = value.args
+ if _is_addr_call(base):
+ raise ValueError("T.ptx.addr(...) cannot be nested")
+ if not isinstance(offset, IntImm) or arg_dtype(offset) not in
_INTEGER_DTYPES:
+ raise ValueError(
+ f"{entry.name}: T.ptx.addr byte_offset must become a compile-time "
+ "signed int32 constant before CUDA codegen (use an
explicitly-unrolled loop)"
+ )
+ offset = int(offset)
+ if not _INT32_MIN <= offset <= _INT32_MAX:
+ raise ValueError(
+ f"{entry.name}: T.ptx.addr byte_offset {offset} is outside signed
int32 range"
+ )
+ return base, offset
+
+
def _make_codegen(entry: InstructionEntry):
n_slots = len(entry.slots)
@@ -126,7 +185,7 @@ def _make_codegen(entry: InstructionEntry):
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
+ rest = list(args[: len(args) - n_slots - 1]) # operands, plus pred
when present
mod_map = mods(entry, tokens)
layout = operand_layout(entry, mod_map)
n_operands = sum(n for _, _, n in layout)
@@ -164,11 +223,37 @@ def _make_codegen(entry: InstructionEntry):
return operand_dtypes(slot, mod_map)[0]
dtypes = tuple(_slot_dtype(slot, i, n) for slot, i, n in layout if
slot.kind == "reg")
- # Caller-chosen immediates ride the Call as IntImm args but are baked
- # into the instruction text, so they are read here and NOT forwarded to
- # the helper (which has no parameter for them).
+ # Caller-chosen immediates ride the Call until device codegen so an
+ # explicitly-unrolled loop may specialize them. They must be IntImm by
+ # this point, then are baked into the instruction text and NOT
+ # forwarded to 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))
+ imm_values = [rest[at[i]] for i in sorted(imm_at)]
+ if any(not isinstance(value, IntImm) for value in imm_values):
+ raise ValueError(
+ f"{entry.name}: immediate operands must become compile-time
constants "
+ "before CUDA codegen (use an explicitly-unrolled loop)"
+ )
+ imms = tuple(str(int(value)) for value in imm_values)
+ # ``tirx.ptx.addr`` is an expression only in the outer PTX call's IR.
+ # The helper still receives the coerced base, while the signed byte
+ # displacement becomes renderer metadata baked into ``[%N+imm]``.
+ addr_offsets = []
+ logical_addr_slot = 0
+ for slot, i, lanes in layout:
+ if slot.kind != "addr":
+ continue
+ if lanes != 1:
+ raise AssertionError(
+ f"{entry.name}: address operand '{slot.name}' must occupy
one register"
+ )
+ value = rest[at[i]]
+ if _is_addr_call(value):
+ base, offset = _codegen_addr_offset(entry, slot, value)
+ rest[at[i]] = base
+ if offset:
+ addr_offsets.append((logical_addr_slot, offset))
+ logical_addr_slot += 1
_, helper, source = render_variant(
entry,
tokens,
@@ -177,6 +262,7 @@ def _make_codegen(entry: InstructionEntry):
imms,
sinks,
preserve_dst=preserve_dst,
+ addr_offsets=tuple(addr_offsets),
)
# Every helper is void; a destination is an ordinary argument, printed
# by the C codegen as the lvalue it binds the reference parameter to.
@@ -217,6 +303,48 @@ class _Sink:
SINK = _Sink()
+class AddrArg:
+ """Temporary trace-time wrapper for ``T.ptx.addr(base, byte_offset)``.
+
+ It deliberately is not a TIR expression. An eligible outer PTX address
+ operand supplies the state space, coerces ``base``, and only then creates
+ the nested pure ``tirx.ptx.addr`` Call.
+ """
+
+ __slots__ = ("base", "byte_offset")
+
+ def __init__(self, base, byte_offset):
+ if isinstance(base, AddrArg):
+ raise ValueError("T.ptx.addr(...) cannot be nested")
+ self.base = base
+ self.byte_offset = _coerce_addr_offset(byte_offset)
+
+ def __repr__(self):
+ return f"T.ptx.addr({self.base!r}, {self.byte_offset!r})"
+
+
+def _coerce_addr_offset(value):
+ """Validate the byte displacement while preserving unrollable
expressions."""
+ if isinstance(value, bool):
+ raise ValueError("T.ptx.addr byte_offset must be a signed int32
integer, not bool")
+ if isinstance(value, IntImm):
+ if arg_dtype(value) not in _INTEGER_DTYPES:
+ raise ValueError(
+ f"T.ptx.addr byte_offset must be a signed int32 integer, got
{arg_dtype(value)}"
+ )
+ value = int(value)
+ if isinstance(value, int):
+ if not _INT32_MIN <= value <= _INT32_MAX:
+ raise ValueError(f"T.ptx.addr byte_offset {value} is outside
signed int32 range")
+ return const(value, "int32")
+ dtype = arg_dtype(value)
+ if dtype in _INTEGER_DTYPES:
+ # An explicitly-unrolled loop may specialize this expression later.
+ # CUDA codegen performs the final IntImm and range checks.
+ return value
+ raise ValueError(f"T.ptx.addr byte_offset must be a scalar integer
expression, got {dtype}")
+
+
class PredArg:
"""``T.ptx.pred(x)`` -- "this operand is a ``.pred`` register".
@@ -264,6 +392,17 @@ def _coerce_operand(entry, slot, values, mod_map):
# unwrapped here: whether the tag is present is the discriminator, so each
# branch below has to be able to see it.
values = [v if isinstance(v, PredArg) else getattr(v, "scalar", v) for v
in values]
+ addr_args = [v for v in values if isinstance(v, AddrArg)]
+ if addr_args:
+ if slot.kind != "addr" or not slot.allow_imm_offset:
+ raise ValueError(
+ f"{entry.name}: operand '{slot.name}' does not support
T.ptx.addr(...)"
+ )
+ if len(addr_args) != len(values):
+ raise ValueError(
+ f"{entry.name}: operand '{slot.name}' cannot mix offset and
plain addresses"
+ )
+ return [_coerce_addr_arg(entry, slot, value, mod_map) for value in
values]
is_pred = slot.kind == "reg" and operand_type(slot, mod_map) == "pred"
tagged = [v for v in values if isinstance(v, PredArg)]
if tagged and not is_pred:
@@ -280,6 +419,12 @@ def _coerce_operand(entry, slot, values, mod_map):
return _coerce_typed(entry, slot, values, mod_map)
+def _coerce_addr_arg(entry, slot, value, mod_map):
+ base = getattr(value.base, "scalar", value.base)
+ base = _coerce_address(entry, slot, base, mod_map)
+ return call_intrin(base.ty, _ADDR_OP_NAME, base, value.byte_offset)
+
+
def _coerce_pred_operand(entry, slot, values):
"""Coerce a ``.pred`` operand -- the one register class the C boundary
cannot bind.
@@ -445,6 +590,20 @@ def _coerce_imm(entry, slot, value):
if isinstance(value, IntImm):
value = value.value
if not isinstance(value, int) or isinstance(value, bool):
+ # Open immediates may be produced by an explicitly-unrolled TIR loop.
+ # Keep the integer expression in the Call; the unroll/simplify pipeline
+ # must turn it into IntImm before the codegen hook bakes it into text.
+ if slot.choices is None and arg_dtype(value) in (
+ "int8",
+ "int16",
+ "int32",
+ "int64",
+ "uint8",
+ "uint16",
+ "uint32",
+ "uint64",
+ ):
+ return value
raise ValueError(
f"{entry.name}: operand '{slot.name}' is an immediate in the
instruction "
f"text; it needs a compile-time integer constant, got
{type(value).__name__}"
@@ -735,6 +894,11 @@ class PTXNamespace:
"""Tag an operand as a ``.pred`` register -- see :class:`PredArg`."""
return PredArg(value)
+ @staticmethod
+ def addr(base, byte_offset):
+ """Form ``[base+byte_offset]`` for an eligible PTX address operand."""
+ return AddrArg(base, byte_offset)
+
def _family(self, token):
cands = self._by_family.get(token)
return _InstrChain(list(cands)) if cands else None
@@ -771,7 +935,7 @@ class PTXNamespace:
def __dir__(self):
"""Family names — drives tab completion."""
- return sorted(self._family_names() | set(super().__dir__()))
+ return sorted(self._family_names() | {"addr"} | set(super().__dir__()))
def __repr__(self):
return f"<T.ptx: {len(self._family_names())} instruction families>"
diff --git a/python/tvm/backend/cuda/ptx/gen_stubs.py
b/python/tvm/backend/cuda/ptx/gen_stubs.py
index 2d8c22222c..85b2526033 100644
--- a/python/tvm/backend/cuda/ptx/gen_stubs.py
+++ b/python/tvm/backend/cuda/ptx/gen_stubs.py
@@ -172,6 +172,7 @@ def generate() -> str:
out.append("class _PTX:")
for family in sorted(families):
out.append(f" {escape_token(family)}: _Chain_{family}")
+ out.append(" def addr(self, base: Any, byte_offset: Any) -> Any: ...")
out.append(" def __getitem__(self, text: str) -> Any: ...")
out.append("")
out.append("ptx: _PTX")
diff --git a/python/tvm/backend/cuda/ptx/render.py
b/python/tvm/backend/cuda/ptx/render.py
index 96bdab0a91..b2b2f70846 100644
--- a/python/tvm/backend/cuda/ptx/render.py
+++ b/python/tvm/backend/cuda/ptx/render.py
@@ -147,8 +147,40 @@ BRIDGE = {
}
+def _normalize_addr_offsets(entry: InstructionEntry, addr_offsets) ->
tuple[tuple[int, int], ...]:
+ """Validate and canonicalize ``(logical_address_slot, byte_offset)``
metadata."""
+ address_slots = tuple(slot for slot in entry.operands if slot.kind ==
"addr")
+ normalized = {}
+ for logical_slot, offset in addr_offsets or ():
+ if isinstance(logical_slot, bool) or not isinstance(logical_slot, int):
+ raise ValueError(f"{entry.name}: address slot index must be an
integer")
+ if not 0 <= logical_slot < len(address_slots):
+ raise ValueError(f"{entry.name}: no address slot {logical_slot}")
+ slot = address_slots[logical_slot]
+ if not slot.allow_imm_offset:
+ raise ValueError(
+ f"{entry.name}: operand '{slot.name}' does not support an
immediate offset"
+ )
+ if isinstance(offset, bool) or not isinstance(offset, int):
+ raise ValueError(f"{entry.name}: address byte offset must be an
integer")
+ if not -(1 << 31) <= offset <= (1 << 31) - 1:
+ raise ValueError(f"{entry.name}: address byte offset {offset} is
outside int32 range")
+ if logical_slot in normalized:
+ raise ValueError(f"{entry.name}: duplicate address slot
{logical_slot}")
+ if offset:
+ normalized[logical_slot] = offset
+ return tuple(sorted(normalized.items()))
+
+
def _helper_name(
- entry: InstructionEntry, written, imms, dtypes, canonical, mod_map,
sinks=()
+ entry: InstructionEntry,
+ written,
+ imms,
+ dtypes,
+ canonical,
+ mod_map,
+ sinks=(),
+ addr_offsets=(),
) -> str:
"""The helper's C identifier: the instruction's ISA identity, plus a
signature discriminator only when it is no longer enough.
@@ -184,7 +216,11 @@ def _helper_name(
f"sink_{name}" + "".join(str(lane) for _, lane in sorted(group))
for name, group in itertools.groupby(sorted(sinks), key=lambda pair:
pair[0])
]
- isa_name = [entry.name, *written, *(imms or ()), *sunk]
+ offset_suffixes = [
+ f"addr{logical_slot}_{'p' if offset > 0 else 'm'}{abs(offset)}"
+ for logical_slot, offset in addr_offsets
+ ]
+ isa_name = [entry.name, *written, *(imms or ()), *offset_suffixes, *sunk]
discriminator = (
[]
if all(dtype == canon for dtype, canon in present)
@@ -203,6 +239,7 @@ def render_variant(
imms=None,
sinks=frozenset(),
preserve_dst=False,
+ addr_offsets=(),
):
"""Render one variant: ``(opcode, helper_name, helper_source)``.
@@ -231,6 +268,8 @@ def render_variant(
value on the inactive path.
"""
mod_map = mods(entry, tokens)
+ addr_offsets = _normalize_addr_offsets(entry, addr_offsets)
+ addr_offset_of = dict(addr_offsets)
written = [tok for tok in tokens if tok]
opcode = ".".join([entry.ptx_name, *written])
canonical = canonical_dtypes(entry, tokens)
@@ -241,7 +280,7 @@ def render_variant(
# derived the same way so dispatch, stubs and certification do not
# need to know the difference.
assert not predicated, f"{opcode}: raw entries have no @p twin"
- helper = _helper_name(entry, written, imms, dtypes, canonical,
mod_map, sinks)
+ helper = _helper_name(entry, written, imms, dtypes, canonical,
mod_map, sinks, addr_offsets)
return opcode, helper, entry.raw_render(entry, opcode, helper, tokens,
tuple(dtypes))
# A helper name is the instruction's ISA identity plus, only when it is no
# longer enough, a signature discriminator. The opcode alone stopped being
@@ -250,7 +289,7 @@ def render_variant(
# ones that changed collides whenever two operands swap which of them is
# 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)
+ helper = _helper_name(entry, written, imms, dtypes, canonical, mod_map,
sinks, addr_offsets)
assert not preserve_dst or entry.has_dst, "preserve_dst requires a written
destination"
if predicated:
if entry.has_dst:
@@ -269,6 +308,7 @@ def render_variant(
bridge_counts: dict[str, int] = collections.defaultdict(int)
dtype_of = dict(zip(entry.typed_operands, dtypes, strict=True))
idx = 0
+ logical_addr_slot = 0
for slot in entry.operands:
pname = f"__{slot.name}"
if slot.kind == "imm":
@@ -350,9 +390,11 @@ def render_variant(
inputs.append(f'"{cb.constraint}"({cb.to_carrier.format(lname)})')
bridge = BRIDGE.get(operand_type(slot, mod_map)) if slot.kind ==
"reg" else None
if bridge is None:
- regs.append(
- f"[%{idx}]" if slot.kind == "addr" and slot.bracket is
None else f"%{idx}"
- )
+ if slot.kind == "addr" and slot.bracket is None:
+ offset = addr_offset_of.get(logical_addr_slot)
+ regs.append(f"[%{idx}{f'+{offset}' if offset is not None
else ''}]")
+ else:
+ regs.append(f"%{idx}")
else:
# The C side above bound the carrier; the instruction names a
# block-local register of the class the ISA actually asks for,
@@ -368,6 +410,8 @@ def render_variant(
regs.append(reg)
idx += 1
rendered.append((slot, "{" + ", ".join(regs) + "}" if is_group else
regs[0]))
+ if slot.kind == "addr":
+ logical_addr_slot += 1
# Adjacent slots naming the same `pipe` are one operand written `p|q`
# (setp's two predicate destinations). Merged first, and into a plain text
@@ -379,7 +423,8 @@ def render_variant(
if key is None:
piped.extend((slot.bracket, text) for slot, text in members)
else:
- piped.append((members[0][0].bracket, "|".join(text for _, text in
members)))
+ slot = members[0][0]
+ piped.append((slot.bracket, "|".join(text for _, text in members)))
# Adjacent slots naming the same `bracket` are one composite memory
operand:
# `[tensorMap, {c0, c1}]` is a single PTX operand whose members keep their
diff --git a/python/tvm/backend/cuda/ptx/table.py
b/python/tvm/backend/cuda/ptx/table.py
index e55c4b095b..209651c59a 100644
--- a/python/tvm/backend/cuda/ptx/table.py
+++ b/python/tvm/backend/cuda/ptx/table.py
@@ -265,6 +265,10 @@ class OperandSlot:
space: str | None = None
dtype: str | DtypeFn | None = None
dtypes: tuple[str, ...] | DtypesFn | None = None
+ # Whether this independent byte-address operand accepts
+ # ``T.ptx.addr(base, byte_offset)``. Composite address members and tmem
+ # addresses are different PTX operand classes and must leave this false.
+ allow_imm_offset: bool = False
# kind="imm" is a value in the instruction *text* (never a C parameter),
# in one of three states, by who owns the value:
# literal set -- the ISA fixed it; invisible to programs.
@@ -3722,7 +3726,7 @@ _ENTRIES = [
check=_check_ld,
operands=(
OperandSlot("d", rw="w", dtypes=_ld_dst_dtypes),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
),
@@ -3771,7 +3775,7 @@ _ENTRIES = [
check=_check_ld_vec,
operands=(
OperandSlot("d", rw="w", lanes=_vec_lanes),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
),
@@ -3798,7 +3802,7 @@ _ENTRIES = [
operands=(
# `_` means this element is not read from memory.
OperandSlot("d", rw="w", lanes=_vec_lanes, sinkable=_sink256),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
),
# Complete scalar `st` per PTX ISA 9.7.9.11, at parity with `ld`.
@@ -3823,7 +3827,7 @@ _ENTRIES = [
),
check=_check_st,
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
@@ -3847,7 +3851,7 @@ _ENTRIES = [
),
check=_check_st_vec,
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value", lanes=_vec_lanes),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
@@ -3868,7 +3872,7 @@ _ENTRIES = [
cert_arch="sm_100",
check=_check_st_vec256,
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
# ISA 9.7.9.11 puts the sink in "vector expression b" -- the data
# being stored -- so here `_` means this element is not written to
# memory. Sink is not a destination-only spelling.
@@ -3892,7 +3896,7 @@ _ENTRIES = [
ModifierSlot("space", ("shared::cta",), optional=True),
),
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("size", dtype="u64"),
OperandSlot("initval", kind="imm", literal="0"),
),
@@ -3911,7 +3915,7 @@ _ENTRIES = [
ModifierSlot("tensormap", ("tensormap",), optional=True),
),
check=_check_prefetch,
- operands=(OperandSlot("addr", kind="addr"),),
+ operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
# st.async per PTX ISA 9.7.9.12. Two syntax blocks that share nothing but
# the mnemonic: one signals completion through an mbarrier, the other is a
@@ -3934,12 +3938,12 @@ _ENTRIES = [
check=_check_st_async if vec else None,
cert_arch="sm_90",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("b", lanes=_vec_lanes if vec else 1),
# The mbarrier lives in the same state space as the destination
# ("`.ss` specifies the state space of the destination operand
# a and the mbarrier operand mbar").
- OperandSlot("mbar", kind="addr"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True),
),
)
for vec in (False, True)
@@ -3957,7 +3961,7 @@ _ENTRIES = [
check=_check_st_async_rel,
cert_arch="sm_100",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("b"),
),
),
@@ -3986,7 +3990,7 @@ _ENTRIES = [
cert_arch="sm_90",
operands=(
OperandSlot("d", rw="w"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
),
InstructionEntry(
@@ -4001,7 +4005,7 @@ _ENTRIES = [
check=_check_multimem_int,
cert_arch="sm_90",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("b"),
),
),
@@ -4021,7 +4025,7 @@ _ENTRIES = [
check=_check_multimem_int,
cert_arch="sm_90",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("b"),
),
),
@@ -4052,11 +4056,11 @@ _ENTRIES = [
*(
(
OperandSlot("d", rw="w", lanes=_vec_lanes if vec else
1),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr",
allow_imm_offset=True),
)
if mnem == "ld_reduce"
else (
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr",
allow_imm_offset=True),
OperandSlot("b", lanes=_vec_lanes if vec else 1),
)
),
@@ -4109,7 +4113,7 @@ _ENTRIES = [
"createpolicy_range",
"range",
(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("primary_size", dtype="u32"),
OperandSlot("total_size", dtype="u32"),
),
@@ -4150,7 +4154,7 @@ _ENTRIES = [
),
cert_arch="sm_90",
operands=(
- OperandSlot("src_mem", kind="addr", space="global"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="global"),
OperandSlot("size", dtype="u32"),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
@@ -4182,7 +4186,7 @@ _ENTRIES = [
),
cert_arch="sm_90a",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
*ops,
),
)
@@ -4228,7 +4232,7 @@ _ENTRIES = [
),
cert_arch="sm_90a",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("new_val", kind="imm", choices=values),
),
)
@@ -4287,7 +4291,7 @@ _ENTRIES = [
),
operands=(
OperandSlot("d", rw="w"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
),
InstructionEntry(
@@ -4305,7 +4309,7 @@ _ENTRIES = [
check=_check_vec128,
operands=(
OperandSlot("d", rw="w", lanes=_vec_lanes),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
),
# prefetchu per PTX ISA 9.7.9.16, the fourth line of that subsection: the
@@ -4314,7 +4318,7 @@ _ENTRIES = [
InstructionEntry(
name="prefetchu",
slots=(ModifierSlot("level", ("L1",)),),
- operands=(OperandSlot("addr", kind="addr"),),
+ operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
# applypriority / discard per PTX ISA 9.7.9.17, 9.7.9.18. Same shape: an
# address range and a cache level, one hinting how to evict it and the
@@ -4332,7 +4336,7 @@ _ENTRIES = [
ModifierSlot("level", (level,)),
),
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("size", kind="imm", literal="128"),
),
)
@@ -5201,8 +5205,8 @@ _ENTRIES = [
),
cert_arch="sm_90",
operands=(
- OperandSlot("dst_mem", kind="addr", space="shared"),
- OperandSlot("src_mem", kind="addr", space="global"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="shared"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="global"),
OperandSlot(
"cp_size", kind="imm", choices=("4", "8", "16") if cop ==
"ca" else ("16",)
),
@@ -5305,8 +5309,8 @@ _ENTRIES = [
),
cert_arch="sm_90",
operands=(
- OperandSlot("dst_mem", kind="addr", space="shared::cta"),
- OperandSlot("src_mem", kind="addr", space="global"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="global"),
OperandSlot("size", dtype="u32"),
OperandSlot(
"ignore_bytes_left",
@@ -5320,7 +5324,7 @@ _ENTRIES = [
lanes=_ignore_oob_lanes,
vector=False,
),
- OperandSlot("mbar", kind="addr", space="shared"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared"),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
),
),
@@ -5338,10 +5342,10 @@ _ENTRIES = [
),
cert_arch="sm_90",
operands=(
- OperandSlot("dst_mem", kind="addr", space="shared::cluster"),
- OperandSlot("src_mem", kind="addr", space="global"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="shared::cluster"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="global"),
OperandSlot("size", dtype="u32"),
- OperandSlot("mbar", kind="addr", space="shared"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared"),
OperandSlot("cta_mask", dtype="u16", lanes=_tma_mask_lanes,
vector=False),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
),
@@ -5358,10 +5362,10 @@ _ENTRIES = [
),
cert_arch="sm_90",
operands=(
- OperandSlot("dst_mem", kind="addr", space="shared::cluster"),
- OperandSlot("src_mem", kind="addr", space="shared::cta"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="shared::cluster"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("size", dtype="u32"),
- OperandSlot("mbar", kind="addr", space="shared"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared"),
),
),
InstructionEntry( # shared::cta -> global
@@ -5396,8 +5400,8 @@ _ENTRIES = [
),
cert_arch="sm_100a",
operands=(
- OperandSlot("dst_mem", kind="addr", space="global"),
- OperandSlot("src_mem", kind="addr", space="shared::cta"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="global"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("size", dtype="u32"),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
# "the 16-bit wide byteMask operand" -- the legacy helper bound it
@@ -5441,10 +5445,10 @@ _ENTRIES = [
cert_arch="sm_100a",
check=_check_tma_gather4,
operands=(
- OperandSlot("dst_mem", kind="addr", space="shared::cluster"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="shared::cluster"),
OperandSlot("tmap", kind="addr", space="global", bracket="src"),
OperandSlot("coords", dtype="s32", lanes=_tma_coords_lanes,
bracket="src"),
- OperandSlot("mbar", kind="addr", space="shared"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared"),
OperandSlot("cta_mask", dtype="u16", lanes=_tma_mask_lanes,
vector=False),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
),
@@ -5467,10 +5471,10 @@ _ENTRIES = [
cert_arch="sm_100a",
check=_check_tma_gather4,
operands=(
- OperandSlot("dst_mem", kind="addr", space="shared::cta"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("tmap", kind="addr", space="global", bracket="src"),
OperandSlot("coords", dtype="s32", lanes=_tma_coords_lanes,
bracket="src"),
- OperandSlot("mbar", kind="addr", space="shared"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared"),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
),
),
@@ -5493,7 +5497,7 @@ _ENTRIES = [
operands=(
OperandSlot("tmap", kind="addr", space="global", bracket="dst"),
OperandSlot("coords", dtype="s32", lanes=_tma_coords_lanes,
bracket="dst"),
- OperandSlot("src_mem", kind="addr", space="shared::cta"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
),
),
@@ -5517,7 +5521,7 @@ _ENTRIES = [
operands=(
OperandSlot("tmap", kind="addr", space="global", bracket="dst"),
OperandSlot("coords", dtype="s32", lanes=_tma_coords_lanes,
bracket="dst"),
- OperandSlot("src_mem", kind="addr", space="shared::cta"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("cache_policy", dtype="u64", lanes=_tma_cache_lanes,
vector=False),
),
),
@@ -5918,7 +5922,7 @@ _ENTRIES = [
),
orders_memory=True,
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
# "The only supported value for the size operand is 128, which must
# be a constant integer literal" -- ISA 9.7.14.4.
OperandSlot("size", kind="imm", literal="128"),
@@ -5940,7 +5944,7 @@ _ENTRIES = [
check=_check_atomic,
operands=(
OperandSlot("d", rw="w"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
@@ -5963,7 +5967,7 @@ _ENTRIES = [
),
operands=(
OperandSlot("d", rw="w"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("compare"),
OperandSlot("value"),
),
@@ -5984,7 +5988,7 @@ _ENTRIES = [
check=_check_cache_hint,
operands=(
OperandSlot("d", rw="w"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
@@ -6008,7 +6012,7 @@ _ENTRIES = [
check=_check_cache_hint,
operands=(
*((OperandSlot("d", rw="w"),) if mnem == "atom" else ()),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
OperandSlot(
"cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False
@@ -6042,7 +6046,7 @@ _ENTRIES = [
check=_check_atom_vec,
operands=(
OperandSlot("d", rw="w", lanes=_vec_lanes),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value", lanes=_vec_lanes),
OperandSlot(
"cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False
@@ -6067,7 +6071,7 @@ _ENTRIES = [
),
check=_check_atomic,
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
OperandSlot("cache_policy", dtype="u64",
lanes=_present_lanes("cache"), vector=False),
),
@@ -6134,7 +6138,7 @@ _ENTRIES = [
ModifierSlot("type", ("b64",)),
),
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("count", dtype="u32"),
),
),
@@ -6172,7 +6176,7 @@ _ENTRIES = [
check=_check_mbarrier_sem_scope,
operands=(
OperandSlot("state", kind="imm", literal="_"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
),
)
for act in ("arrive", "arrive_drop")
@@ -6191,7 +6195,7 @@ _ENTRIES = [
check=_check_mbarrier_sem_scope,
operands=(
OperandSlot("state", kind="imm", literal="_"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("count", dtype="u32"),
),
)
@@ -6212,7 +6216,7 @@ _ENTRIES = [
check=_check_mbarrier_sem_scope,
operands=(
OperandSlot("state", kind="imm", literal="_"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("tx_count", dtype="u32"),
),
)
@@ -6242,7 +6246,7 @@ _ENTRIES = [
# carrier, and ptxas rejects an .f64 register as the state
operand
# ("Arguments mismatch for instruction 'mbarrier.arrive'").
OperandSlot("state", rw="rw", dtype="b64i"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("count", dtype="u32"),
),
)
@@ -6268,7 +6272,7 @@ _ENTRIES = [
check=_check_mbarrier_sem_scope,
operands=(
OperandSlot("wait_complete", rw="w", dtype="pred"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("phase", dtype="u32"),
*((OperandSlot("time_hint", dtype="u32"),) if act ==
"try_wait" else ()),
),
@@ -6289,7 +6293,7 @@ _ENTRIES = [
check=_check_mbarrier_sem_scope,
operands=(
OperandSlot("wait_complete", rw="w", dtype="pred"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("phase", dtype="u32"),
),
),
@@ -6306,7 +6310,7 @@ _ENTRIES = [
),
check=_check_mbarrier_sem_scope,
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("tx_count", dtype="u32"),
),
)
@@ -6320,7 +6324,7 @@ _ENTRIES = [
ModifierSlot("space", ("shared", "shared::cta"), optional=True),
ModifierSlot("type", ("b64",)),
),
- operands=(OperandSlot("addr", kind="addr"),),
+ operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
# The state-returning arrive lines (PTX ISA 9.7.14.16.16 / .17). The
# entries above bake `_` into the text, which is the only spelling the
@@ -6344,7 +6348,7 @@ _ENTRIES = [
# integer register of either signedness and rejects a float one
# ("Arguments mismatch"), which is what `b64i` names.
OperandSlot("state", rw="w", dtype="b64i"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
*((OperandSlot("count", dtype="u32"),) if count else ()),
),
)
@@ -6375,7 +6379,7 @@ _ENTRIES = [
check=_check_mbarrier_sem_scope,
operands=(
OperandSlot("wait_complete", rw="w", dtype="pred"),
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("state", dtype="b64i"), # the token an arrive
returned
*((OperandSlot("time_hint", dtype="u32"),) if hint else ()),
),
@@ -6401,8 +6405,8 @@ _ENTRIES = [
),
cert_arch="sm_90",
operands=(
- OperandSlot("dst_mem", kind="addr", space="global"),
- OperandSlot("src_mem", kind="addr", space="shared::cta"),
+ OperandSlot("dst_mem", kind="addr", allow_imm_offset=True,
space="global"),
+ OperandSlot("src_mem", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("size", kind="imm", literal="128"),
),
),
@@ -6423,9 +6427,9 @@ _ENTRIES = [
check=_check_red_async,
cert_arch="sm_90",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
- OperandSlot("mbar", kind="addr"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True),
),
),
# The release line of the same subsection, which reduces straight into
@@ -6447,7 +6451,7 @@ _ENTRIES = [
check=_check_st_async_rel,
cert_arch="sm_100",
operands=(
- OperandSlot("addr", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
OperandSlot("value"),
),
),
@@ -6480,7 +6484,7 @@ _ENTRIES = [
# sm_90+, so pinning the operand to shared would bind a 32-bit register
# under the space-omitted spelling. `operand_space` reads the entry's
# `space` slot instead, so the carrier follows the spelling.
- operands=(OperandSlot("addr", kind="addr"),),
+ operands=(OperandSlot("addr", kind="addr", allow_imm_offset=True),),
),
InstructionEntry( # mbarrier.pending_count.b64 count, state;
# The reader of the `state` result the two `.noComplete` entries above
@@ -6534,8 +6538,8 @@ _ENTRIES = [
# `operand_space` read the entry's `space` slot gives each variant the
# carrier its own spelling promises -- the mbarrier-family rule.
operands=(
- OperandSlot("addr", kind="addr"),
- OperandSlot("mbar", kind="addr"),
+ OperandSlot("addr", kind="addr", allow_imm_offset=True),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True),
),
),
# clusterlaunchcontrol.query_cancel per PTX ISA 9.7.14.19: decode the
@@ -6942,7 +6946,7 @@ _ENTRIES = [
),
operands=(
OperandSlot("r", rw="w", dtype="b32", lanes=_ldmatrix_lanes),
- OperandSlot("p", kind="addr"),
+ OperandSlot("p", kind="addr", allow_imm_offset=True),
),
),
InstructionEntry( # line 1, .m16n16.b8: "only .x1 and .x2 are valid"
@@ -6964,7 +6968,7 @@ _ENTRIES = [
cert_arch="sm_100a",
operands=(
OperandSlot("r", rw="w", dtype="b32", lanes=_ldmatrix_lanes),
- OperandSlot("p", kind="addr"),
+ OperandSlot("p", kind="addr", allow_imm_offset=True),
),
),
InstructionEntry( # lines 2+3: the 6/4-bit decompression loads
@@ -6984,7 +6988,7 @@ _ENTRIES = [
check=_check_ldmatrix_b8fmt,
operands=(
OperandSlot("r", rw="w", dtype="b32", lanes=_ldmatrix_lanes),
- OperandSlot("p", kind="addr"),
+ OperandSlot("p", kind="addr", allow_imm_offset=True),
),
),
# stmatrix per PTX ISA 9.7.15.5.16 -- the store mirror of ldmatrix. One
@@ -7010,7 +7014,7 @@ _ENTRIES = [
),
cert_arch="sm_90", # ISA: "Requires sm_90 or higher."
operands=(
- OperandSlot("p", kind="addr"),
+ OperandSlot("p", kind="addr", allow_imm_offset=True),
OperandSlot("r", dtype="b32", lanes=_matrix_num_lanes),
),
),
@@ -7029,7 +7033,7 @@ _ENTRIES = [
),
cert_arch="sm_100a",
operands=(
- OperandSlot("p", kind="addr"),
+ OperandSlot("p", kind="addr", allow_imm_offset=True),
OperandSlot("r", dtype="b32", lanes=_matrix_num_lanes),
),
),
@@ -7465,7 +7469,7 @@ _ENTRIES = [
cert_arch="sm_100a",
orders_memory=True,
operands=(
- OperandSlot("dst", kind="addr", space="shared::cta"),
+ OperandSlot("dst", kind="addr", allow_imm_offset=True,
space="shared::cta"),
OperandSlot("ncols", dtype="u32"),
),
),
@@ -7795,7 +7799,9 @@ _ENTRIES = [
),
cert_arch="sm_100a",
orders_memory=True,
- operands=(OperandSlot("mbar", kind="addr", space="shared::cluster"),),
+ operands=(
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared::cluster"),
+ ),
),
# tcgen05.commit...{.shared::cluster}.multicast::cluster.b64 [mbar],
ctaMask;
# The multicast form: `pred` is keyword-only, so the trailing mask
@@ -7815,7 +7821,7 @@ _ENTRIES = [
cert_arch="sm_100a",
orders_memory=True,
operands=(
- OperandSlot("mbar", kind="addr", space="shared::cluster"),
+ OperandSlot("mbar", kind="addr", allow_imm_offset=True,
space="shared::cluster"),
OperandSlot("mask", dtype="u16"),
),
),
@@ -7840,6 +7846,34 @@ _ENTRIES = [
),
]
+
+def _validate_imm_offset_slots(entries) -> None:
+ """Reject capability bits on operand classes that cannot spell
``[addr+imm]``."""
+ errors = []
+ for entry in entries:
+ for slot in entry.operands:
+ if not slot.allow_imm_offset:
+ continue
+ reasons = []
+ if slot.kind != "addr":
+ reasons.append(f"kind={slot.kind!r}, expected 'addr'")
+ if slot.bracket is not None:
+ reasons.append("is a composite bracket member")
+ if slot.space == "tmem" or (
+ slot.space is None
+ and any(
+ modifier.name == "space" and "tmem" in modifier.choices
+ for modifier in entry.slots
+ )
+ ):
+ reasons.append("is a tmem address")
+ if reasons:
+ errors.append(f"{entry.name}.{slot.name}: " + ";
".join(reasons))
+ if errors:
+ raise ValueError("invalid allow_imm_offset slots:\n " + "\n
".join(errors))
+
+
+_validate_imm_offset_slots(_ENTRIES)
TABLE: dict[str, InstructionEntry] = {e.name: e for e in _ENTRIES}
# Keying by name silently drops a duplicate, and a dropped entry is an ISA line
# that stops being reachable. Two entries never legitimately share a name.
diff --git a/python/tvm/script/tirx.pyi b/python/tvm/script/tirx.pyi
index ac9e092c12..89c3d9fd11 100644
--- a/python/tvm/script/tirx.pyi
+++ b/python/tvm/script/tirx.pyi
@@ -2286,6 +2286,7 @@ class _PTX:
vote_sync: _Chain_vote_sync
wgmma: _Chain_wgmma
xor: _Chain_xor
+ def addr(self, base: Any, byte_offset: Any) -> Any: ...
def __getitem__(self, text: str) -> Any: ...
ptx: _PTX
diff --git a/tests/python/tirx/codegen/test_ptx_addr.py
b/tests/python/tirx/codegen/test_ptx_addr.py
new file mode 100644
index 0000000000..fa4dac30b6
--- /dev/null
+++ b/tests/python/tirx/codegen/test_ptx_addr.py
@@ -0,0 +1,351 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for ``T.ptx.addr(base, byte_offset)``."""
+
+import pytest
+
+import tvm
+from tvm.ir import Call, Op
+from tvm.runtime import const
+from tvm.script import tirx as T
+from tvm.tirx.expr import Broadcast, CallEffectKind
+
+TARGET = tvm.target.Target("cuda")
+
+
+def _cuda_source(func) -> str:
+ with TARGET:
+ mod = tvm.compile(tvm.IRModule({"main": func}), target=TARGET,
tir_pipeline="tirx")
+ return mod.mod.imports[0].inspect_source("cuda")
+
+
+def _calls(func, op_name):
+ calls = []
+
+ def visit(node):
+ if isinstance(node, Call) and getattr(node.op, "name", None) ==
op_name:
+ calls.append(node)
+
+ tvm.tirx.stmt_functor.post_order_visit(func.body, visit)
+ return calls
+
+
+def test_ptx_addr_registration_and_table_capabilities():
+ from tvm.backend.cuda.codegen.registry import CODEGEN_REGISTRY
+ from tvm.backend.cuda.ptx.table import TABLE
+
+ op = Op.get("tirx.ptx.addr")
+ assert int(op.get_attr("TCallEffectKind")) == CallEffectKind.Pure.value
+ assert op.get_attr("TScriptPrinterName") == "ptx.addr"
+ assert "tirx.ptx.addr" in CODEGEN_REGISTRY
+
+ addresses = [slot for entry in TABLE.values() for slot in entry.operands
if slot.kind == "addr"]
+ assert len(addresses) == 145
+ assert sum(slot.allow_imm_offset for slot in addresses) == 115
+ assert sum(slot.bracket is not None and not slot.allow_imm_offset for slot
in addresses) == 5
+ assert sum(slot.space == "tmem" and not slot.allow_imm_offset for slot in
addresses) == 25
+
+
+def test_ptx_addr_table_validation_rejects_wrong_operand_classes():
+ from tvm.backend.cuda.ptx.table import (
+ InstructionEntry,
+ OperandSlot,
+ _validate_imm_offset_slots,
+ )
+
+ bad_entries = (
+ InstructionEntry("reg", (OperandSlot("x", allow_imm_offset=True),)),
+ InstructionEntry("ptr", (OperandSlot("x", kind="ptr",
allow_imm_offset=True),)),
+ InstructionEntry(
+ "composite",
+ (OperandSlot("x", kind="addr", bracket="pair",
allow_imm_offset=True),),
+ ),
+ InstructionEntry(
+ "tmem",
+ (OperandSlot("x", kind="addr", space="tmem",
allow_imm_offset=True),),
+ ),
+ )
+ with pytest.raises(ValueError, match="reg.x.*kind='reg'"):
+ _validate_imm_offset_slots(bad_entries)
+
+
+def test_ptx_addr_coercion_ir_order_and_shared_codegen():
+ @T.prim_func
+ def kernel(global_buf: T.Buffer((8,), "uint64"), raw_shared: T.uint32,
raw_global: T.uint64):
+ T.device_entry()
+ tx = T.thread_id([32])
+ shared_buf = T.alloc_buffer((8,), "uint64", scope="shared")
+ value = T.local_scalar("uint64")
+ if tx == 0:
+ T.ptx.ld.shared.b64(value, T.ptx.addr(shared_buf.data, 4))
+ T.ptx.ld.shared.b64(value, T.ptx.addr(raw_shared, 8))
+ T.ptx.ld.global_.b64(value, T.ptx.addr(global_buf.data, 12))
+ T.ptx.ld.global_.b64(value, T.ptx.addr(raw_global, 16))
+
+ calls = _calls(kernel, "tirx.ptx.addr")
+ assert len(calls) == 4
+ assert getattr(calls[0].args[0].op, "name", None) ==
"tirx.cuda.cvta_generic_to_shared"
+ assert not isinstance(calls[1].args[0], Call)
+ assert getattr(calls[2].args[0].op, "name", None) == "tirx.buffer_data"
+ assert getattr(calls[3].args[0].op, "name", None) == "tirx.reinterpret"
+ assert [int(call.args[1]) for call in calls] == [4, 8, 12, 16]
+ assert all(call.ty == call.args[0].ty for call in calls)
+
+ source = _cuda_source(kernel)
+ assert "ld.shared.b64 %0, [%1+4];" in source
+ assert "ld.shared.b64 %0, [%1+8];" in source
+ # The shared pointer needs cvta; the raw uint32 shared-window address does
not.
+ assert source.count("__cvta_generic_to_shared") == 1
+
+
+def test_ptx_addr_scalar_vector_cache_predicate_and_multi_address_codegen():
+ @T.prim_func
+ def kernel(
+ src: T.Buffer((64,), "uint32"),
+ dst: T.Buffer((64,), "uint32"),
+ policy: T.Buffer((1,), "uint64"),
+ ):
+ T.device_entry()
+ tx = T.thread_id([32])
+ shared_buf = T.alloc_buffer((64,), "uint32", scope="shared")
+ barrier = T.alloc_buffer((1,), "uint64", scope="shared")
+ values = T.alloc_local((2,), "uint32")
+ T.ptx.ld.global_.b32(values[0], T.ptx.addr(src.data, 16))
+ T.ptx.ld.global_.L2__cache_hint.b32(values[1], T.ptx.addr(src.data,
-16), policy[0])
+ T.ptx.ld.shared.v2.b32(values[0], values[1],
T.ptx.addr(shared_buf.data, 8))
+ T.ptx.st.global_.b32(T.ptx.addr(dst.data, 0), values[0])
+ T.ptx.st.global_.v2.b32(T.ptx.addr(dst.data, 32), values[0],
values[1], pred=tx == 0)
+ T.ptx["cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes"](
+ T.ptx.addr(shared_buf.data, 16),
+ T.ptx.addr(src.data, -16),
+ T.uint32(16),
+ T.ptx.addr(barrier.data, 8),
+ )
+
+ source = _cuda_source(kernel)
+ assert "ld.global.b32 %0, [%1+16];" in source
+ assert "ld.global.L2::cache_hint.b32 %0, [%1+-16], %2;" in source
+ assert "ld.shared.v2.b32 {%0, %1}, [%2+8];" in source
+ assert "st.global.b32 [%0], %1;" in source
+ assert "@p st.global.v2.b32 [%0+32], {%1, %2};" in source
+ assert (
+ "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes "
+ "[%0+16], [%1+-16], %2, [%3+8];"
+ ) in source
+
+
+def test_ptx_addr_zero_sign_boundaries_and_helper_names():
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for
+
+ entry = TABLE["ld"]
+ tokens = tokens_for(entry, space="global", type="b32")
+ bare = render_variant(entry, tokens)
+ zero = render_variant(entry, tokens, addr_offsets=((0, 0),))
+ positive = render_variant(entry, tokens, addr_offsets=((0, 16),))
+ negative = render_variant(entry, tokens, addr_offsets=((0, -16),))
+ low = render_variant(entry, tokens, addr_offsets=((0, -(1 << 31)),))
+ high = render_variant(entry, tokens, addr_offsets=((0, (1 << 31) - 1),))
+
+ assert zero == bare
+ assert bare[1] == "tvm_builtin_ptx_ld_global_b32"
+ assert positive[1].endswith("_addr0_p16")
+ assert negative[1].endswith("_addr0_m16")
+ assert "[%1+16]" in positive[2]
+ assert "[%1+-16]" in negative[2]
+ assert "[%1+-2147483648]" in low[2]
+ assert "[%1+2147483647]" in high[2]
+
+ from tvm.backend.cuda.ptx.table import renderings
+
+ cp_entry = TABLE["cp_async_ca"]
+ cp_tokens, cp_dtypes, cp_predicated, cp_imms, cp_sinks =
next(iter(renderings(cp_entry)))
+ _, helper, source = render_variant(
+ cp_entry,
+ cp_tokens,
+ cp_predicated,
+ cp_dtypes,
+ cp_imms,
+ cp_sinks,
+ addr_offsets=((0, 16), (1, -16)),
+ )
+ assert helper.endswith("_addr0_p16_addr1_m16")
+ assert "[%0+16], [%1+-16]" in source
+
+ for value in (-(1 << 31) - 1, 1 << 31):
+ with pytest.raises(ValueError, match="outside int32 range"):
+ render_variant(entry, tokens, addr_offsets=((0, value),))
+
+
+def test_ptx_addr_unrolled_expression_and_dynamic_rejection():
+ @T.prim_func
+ def unrolled(src: T.Buffer((16,), "uint32")):
+ T.device_entry()
+ tx = T.thread_id([32])
+ value = T.local_scalar("uint32")
+ if tx == 0:
+ for i in T.unroll(3):
+ T.ptx.ld.global_.b32(value, T.ptx.addr(src.data, i * 16))
+
+ source = _cuda_source(unrolled)
+ assert "ld.global.b32 %0, [%1];" in source
+ assert "ld.global.b32 %0, [%1+16];" in source
+ assert "ld.global.b32 %0, [%1+32];" in source
+
+ @T.prim_func
+ def thread_dynamic(src: T.Buffer((16,), "uint32")):
+ T.device_entry()
+ tx = T.thread_id([32])
+ value = T.local_scalar("uint32")
+ T.ptx.ld.global_.b32(value, T.ptx.addr(src.data, tx * 4))
+
+ @T.prim_func
+ def loop_dynamic(src: T.Buffer((16,), "uint32")):
+ T.device_entry()
+ tx = T.thread_id([32])
+ value = T.local_scalar("uint32")
+ if tx == 0:
+ for i in T.serial(2):
+ T.ptx.ld.global_.b32(value, T.ptx.addr(src.data, i * 4))
+
+ for func in (thread_dynamic, loop_dynamic):
+ with pytest.raises(
+ (ValueError, tvm.error.InternalError), match="must become a
compile-time"
+ ):
+ _cuda_source(func)
+
+
+def test_ptx_addr_offset_type_and_range_rejections():
+ for value in (True, 1.5, Broadcast(const(1, "int32"), 4)):
+ with pytest.raises(ValueError, match="byte_offset"):
+ T.ptx.addr(None, value)
+ for value in (-(1 << 31) - 1, 1 << 31):
+ with pytest.raises(ValueError, match="outside signed int32 range"):
+ T.ptx.addr(None, value)
+ with pytest.raises(ValueError, match="cannot be nested"):
+ T.ptx.addr(T.ptx.addr(None, 0), 4)
+
+
+def test_ptx_addr_pointer_and_raw_address_validation():
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match="uint32 address
requires shared"
+ ):
+
+ @T.prim_func
+ def global_u32(raw: T.uint32):
+ T.device_entry()
+ value = T.local_scalar("uint32")
+ T.ptx.ld.global_.b32(value, T.ptx.addr(raw, 4))
+
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match="does not support
T.ptx.addr"
+ ):
+
+ @T.prim_func
+ def ptr_operand(src: T.Buffer((8,), "uint32")):
+ T.device_entry()
+ result = T.local_scalar("uint32")
+ T.ptx.isspacep.global_(result, T.ptx.addr(src.data, 4))
+
+
+def test_ptx_addr_tma_tmem_and_independent_immediate_rejections():
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match="does not support
T.ptx.addr"
+ ):
+
+ @T.prim_func
+ def tma(tmap: T.Buffer((8,), "uint64")):
+ T.device_entry()
+ shared_buf = T.alloc_buffer((16,), "uint32", scope="shared")
+ barrier = T.alloc_buffer((1,), "uint64", scope="shared")
+
T.ptx["cp.async.bulk.tensor.1d.shared::cta.global.mbarrier::complete_tx::bytes"](
+ shared_buf.data, T.ptx.addr(tmap.data, 16), T.int32(0),
barrier.data
+ )
+
+ with pytest.raises(
+ (ValueError, tvm.error.DiagnosticError), match="does not support
T.ptx.addr"
+ ):
+
+ @T.prim_func
+ def tmem(raw: T.uint32):
+ T.device_entry()
+ value = T.local_scalar("uint32")
+ T.ptx["tcgen05.ld.sync.aligned.32x32b.x1.b32"](value,
T.ptx.addr(raw, 16))
+
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, variants
+
+ entry = TABLE["tcgen05_ld_split"]
+ tokens = variants(entry)[0]
+ _, helper, source = render_variant(entry, tokens, imms=("16",))
+ assert helper.endswith("_16")
+ assert "[%" in source and "], 16;" in source
+ with pytest.raises(ValueError, match="does not support an immediate
offset"):
+ render_variant(entry, tokens, imms=("16",), addr_offsets=((0, 16),))
+
+
+def test_ptx_addr_printer_script_and_json_roundtrip():
+ @T.prim_func
+ def kernel(src: T.Buffer((8,), "uint32"), dst: T.Buffer((8,), "uint32")):
+ T.device_entry()
+ value = T.local_scalar("uint32")
+ T.ptx.ld.global_.b32(value, T.ptx.addr(src.data, -16))
+ T.ptx.st.global_.b32(T.ptx.addr(dst.data, 16), value)
+
+ script = kernel.script()
+ assert script.count("T.ptx.addr(") == 2
+ tvm.ir.assert_structural_equal(kernel, tvm.script.from_source(script))
+ tvm.ir.assert_structural_equal(kernel,
tvm.ir.load_json(tvm.ir.save_json(kernel)))
+
+
+def test_ptx_addr_legacy_positional_offsets_rejected():
+ with pytest.raises((ValueError, tvm.error.DiagnosticError)):
+
+ @T.prim_func
+ def scalar_load(src: T.Buffer((8,), "uint32")):
+ T.device_entry()
+ value = T.local_scalar("uint32")
+ T.ptx.ld.global_.b32(value, src.data, 16)
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError)):
+
+ @T.prim_func
+ def vector_load(src: T.Buffer((8,), "uint32")):
+ T.device_entry()
+ values = T.alloc_local((2,), "uint32")
+ T.ptx.ld.global_.v2.b32(values[0], values[1], src.data, 16)
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError)):
+
+ @T.prim_func
+ def scalar_store(dst: T.Buffer((8,), "uint32")):
+ T.device_entry()
+ T.ptx.st.global_.b32(dst.data, 16, T.uint32(0))
+
+ with pytest.raises((ValueError, tvm.error.DiagnosticError)):
+
+ @T.prim_func
+ def vector_store(dst: T.Buffer((8,), "uint32")):
+ T.device_entry()
+ T.ptx.st.global_.v2.b32(dst.data, 16, T.uint32(0), T.uint32(0))
+
+
+def test_ptx_addr_unconsumed_codegen_diagnostic():
+ from tvm.backend.cuda.codegen.registry import CODEGEN_REGISTRY
+
+ with pytest.raises(ValueError, match="must be consumed by a PTX address
operand"):
+ CODEGEN_REGISTRY["tirx.ptx.addr"]([const(0, "uint64"), const(16,
"int32")])
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py
b/tests/python/tirx/codegen/test_ptx_dialect.py
index f08d6ea419..fe1ffaebe0 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -16,6 +16,7 @@
# under the License.
"""Tests for the table-driven PTX dialect (``T.ptx``)."""
+import itertools
import os
import re
import shutil
@@ -1112,16 +1113,35 @@ def test_ptx_logic_shift_dispatch():
T.device_entry()
T.ptx.shl.s32(out[0], T.int32(1), T.uint32(2))
- # The LUT byte lives in the instruction text, so it has to be a constant.
- with pytest.raises((ValueError, tvm.error.DiagnosticError),
match="compile-time integer"):
-
- @T.prim_func
- def lut_runtime(a_ptr: T.handle):
- A = T.match_buffer(a_ptr, (1,), "uint32")
- T.device_entry()
+ # Open immediates may survive tracing so explicitly-unrolled expressions
+ # can specialize, but a runtime LUT byte still has no register form and is
+ # rejected at CUDA codegen.
+ @T.prim_func
+ def lut_runtime(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ tx = T.thread_id([32])
+ if tx == 0:
d = T.local_scalar("uint32")
T.ptx.lop3.b32(d, A[0], A[0], A[0], A[0])
+ with pytest.raises((ValueError, tvm.error.InternalError),
match="compile-time constants"):
+ _cuda_source(lut_runtime)
+
+ @T.prim_func
+ def lut_unrolled(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (1,), "uint32")
+ T.device_entry()
+ tx = T.thread_id([32])
+ if tx == 0:
+ d = T.local_scalar("uint32")
+ for i in T.unroll(2):
+ T.ptx.lop3.b32(d, A[0], A[0], A[0], i * 128)
+
+ unrolled_src = _cuda_source(lut_unrolled)
+ assert "lop3.b32 %0, %1, %2, %3, 0;" in unrolled_src
+ assert "lop3.b32 %0, %1, %2, %3, 128;" in unrolled_src
+
def test_ptx_data_movement_dispatch():
"""ISA 9.7.9's newly registered instructions, end to end.
@@ -2128,6 +2148,36 @@ def _as_render_args(rendering):
return tokens, predicated, dtypes, imms, sinks
+def _addr_offset_samples(entry):
+ """Small certification axis for address immediates, separate from modifier
products."""
+ from tvm.backend.cuda.ptx.table import renderings
+
+ enabled = [
+ logical_slot
+ for logical_slot, slot in enumerate(s for s in entry.operands if
s.kind == "addr")
+ if slot.allow_imm_offset
+ ]
+ if not enabled:
+ return ()
+ representative = _as_render_args(next(iter(renderings(entry))))
+ samples = [
+ (representative, ((logical_slot, offset),))
+ for logical_slot in enabled
+ for offset in (16, -16)
+ ]
+ if len(enabled) > 1:
+ samples.append(
+ (
+ representative,
+ tuple(
+ (logical_slot, 16 if index % 2 == 0 else -16)
+ for index, logical_slot in enumerate(enabled)
+ ),
+ )
+ )
+ return tuple(samples)
+
+
def _sole_instruction(asm_text):
"""The single PTX statement in ``asm_text``, or None if it is not exactly
one.
@@ -2275,6 +2325,10 @@ def test_ptx_all_variants_render_unique():
or f"; {opcode};" in source
)
total += not predicated # a @p twin is not a separate variant
+ for args, addr_offsets in _addr_offset_samples(entry):
+ _, helper, _ = render_variant(entry, *args,
addr_offsets=addr_offsets)
+ assert helper not in names, f"address-offset helper name
collision: {helper}"
+ names.add(helper)
assert total == 200018 # update when the table grows
@@ -2434,6 +2488,9 @@ def test_ptx_sampled_helpers_assemble():
for i in rng.sample(range(len(rendered)), min(48, len(rendered))):
_, _, source = render_variant(entry, *_as_render_args(rendered[i]))
by_arch.setdefault(arch,
[]).append(source.replace("__forceinline__ ", ""))
+ for args, addr_offsets in _addr_offset_samples(entry):
+ _, _, source = render_variant(entry, *args,
addr_offsets=addr_offsets)
+ by_arch.setdefault(arch,
[]).append(source.replace("__forceinline__ ", ""))
for arch, sources in by_arch.items():
_assert_ptxas_ok("\n".join([_CERT_PRELUDE, *sources]), rdc=True,
arch=arch)
@@ -2473,13 +2530,21 @@ def test_ptx_all_helpers_certify(shard):
by_arch = {}
covered = 0
- for index, (entry, rendering) in enumerate(
- (TABLE[name], r) for name in sorted(TABLE) for r in
renderings(TABLE[name])
- ):
+ baseline = (
+ (TABLE[name], _as_render_args(rendering), ())
+ for name in sorted(TABLE)
+ for rendering in renderings(TABLE[name])
+ )
+ address_samples = (
+ (TABLE[name], args, addr_offsets)
+ for name in sorted(TABLE)
+ for args, addr_offsets in _addr_offset_samples(TABLE[name])
+ )
+ for index, (entry, args, addr_offsets) in
enumerate(itertools.chain(baseline, address_samples)):
if index % _CERT_SHARDS == shard:
covered += 1
arch = entry.cert_arch or PTX_ARCH
- _, _, src = render_variant(entry, *_as_render_args(rendering))
+ _, _, src = render_variant(entry, *args, addr_offsets=addr_offsets)
by_arch.setdefault(arch, []).append(src.replace("__forceinline__
", ""))
assert covered, "empty shard: lower _CERT_SHARDS"
for arch, sources in by_arch.items():
diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
index 88d8fc818e..53bbdec724 100644
--- a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
+++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
@@ -162,6 +162,25 @@ def test_ptx_ld_st_raw_shared_address_codegen():
assert '"q"(__value)' in src
+def test_ptx_ld_st_immediate_offset_codegen():
+ """An immediate displacement must stay inside the PTX memory operand."""
+
+ @T.prim_func
+ def main(src: T.Buffer((4,), "uint64"), out: T.Buffer((4,), "uint64")):
+ T.device_entry()
+ tx = T.thread_id([32])
+ values = T.alloc_local((2,), "uint64")
+ if tx == 0:
+ T.ptx.ld.global_.v2.b64(values[0], values[1], T.ptx.addr(src.data,
16))
+ T.ptx.st.global_.v2.b64(T.ptx.addr(out.data, 16), values[0],
values[1])
+
+ with TARGET:
+ mod = tvm.compile(tvm.IRModule({"main": main}), target=TARGET,
tir_pipeline="tirx")
+ src = mod.mod.imports[0].inspect_source("cuda")
+ assert "ld.global.v2.b64 {%0, %1}, [%2+16];" in src
+ assert "st.global.v2.b64 [%0+16], {%1, %2};" in src
+
+
def test_ptx_ld_global_nc_v8_codegen():
"""FlashMLA index loads need ``ld.global.nc`` with a 256B prefetch."""
diff --git a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
index 5485477188..43b8a71b98 100644
--- a/tests/python/tirx/test_tirx_kernels_registry_correctness.py
+++ b/tests/python/tirx/test_tirx_kernels_registry_correctness.py
@@ -51,7 +51,14 @@ _KERNELS = {
for kernel_name in sorted({workload["kernel"] for workload in _WORKLOADS})
}
_DISTRIBUTED_KERNELS = frozenset(
- {"allgather_gemm", "deepgemm_fp8_fp4_mega_moe", "gemm_reduce_scatter"}
+ # Both MegaMoE names are listed so the test works across tirx-kernels
+ # checkouts from before and after the sm100_fp8_fp4_mega_moe rename.
+ {
+ "allgather_gemm",
+ "deepgemm_fp8_fp4_mega_moe",
+ "gemm_reduce_scatter",
+ "sm100_fp8_fp4_mega_moe",
+ }
)
_MEGA_MOE_KERNELS = frozenset({"deepgemm_fp8_fp4_mega_moe",
"sm100_fp8_fp4_mega_moe"})
_XDIST_CUDA_DEVICE = None
