diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h index 8346690102c5..fcb6fb4b10a9 100644 --- a/include/tvm/tirx/script/builder/ir.h +++ b/include/tvm/tirx/script/builder/ir.h @@ -131,15 +131,37 @@ SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_sc void TilePrimitiveCall(tvm::tirx::TilePrimitiveCall op_call); -ffi::Array KernelId(ffi::Array extents, ffi::String parent); +/*! + * \brief Define a scope id. Pass `extents=std::nullopt` to defer the extent; it is + * inferred at LowerTIRx from the sibling ScopeIdDef closure. + * \param extents The optional extents of the scope id. + * \param parent The parent scope name. + * \param name The user-facing API name, used in error messages. + * \param cur The current scope name. + * \param dtype The dtype of the introduced scope id vars ("int32" or "uint32"). + * \return The introduced scope id vars. + */ +ffi::Array ScopeId(ffi::Optional> extents, ffi::String parent, + ffi::String name, ffi::String cur, + PrimType dtype = PrimType::Int(32)); + +ffi::Array ClusterId(ffi::Optional> extents, + ffi::String parent, PrimType dtype = PrimType::Int(32)); + +ffi::Array CtaId(ffi::Optional> extents, ffi::String parent, + ffi::Optional> preferred = std::nullopt, + PrimType dtype = PrimType::Int(32)); -ffi::Array CtaId(ffi::Array extents, ffi::String parent); +ffi::Array CtaIdInPair(PrimType dtype = PrimType::Int(32)); -ffi::Array CtaIdInPair(); +ffi::Array WarpgroupId(ffi::Optional> extents, + ffi::String parent, PrimType dtype = PrimType::Int(32)); -ffi::Array WarpId(ffi::Array extents, ffi::String parent); +ffi::Array WarpId(ffi::Optional> extents, ffi::String parent, + PrimType dtype = PrimType::Int(32)); -ffi::Array ThreadId(ffi::Array extents, ffi::String parent); +ffi::Array ThreadId(ffi::Optional> extents, ffi::String parent, + PrimType dtype = PrimType::Int(32)); /*! * \brief The block initialization statement. @@ -249,44 +271,53 @@ ffi::Array Remap(ffi::String kinds, ffi::Array bindings, * \param stop The maximum value of iteration. * \param annotations The optional annotations of the For statement. * \param step The optional step value of iteration. + * \param dtype The optional dtype of the loop var ("int32" or "uint32"). When omitted + * it is inferred from the bounds. * \return The ForFrame. */ ForFrame Serial(PrimExpr start, PrimExpr stop, ffi::Optional> annotations = std::nullopt, - ffi::Optional step = std::nullopt); + ffi::Optional step = std::nullopt, + ffi::Optional dtype = std::nullopt); /*! * \brief The parallel For statement. * \param start The minimum value of iteration. * \param stop The maximum value of iteration. * \param annotations The optional annotations of the For statement. * \param step The optional step value of iteration. + * \param dtype The optional dtype of the loop var ("int32" or "uint32"). * \return The ForFrame. */ ForFrame Parallel(PrimExpr start, PrimExpr stop, ffi::Optional> annotations = std::nullopt, - ffi::Optional step = std::nullopt); + ffi::Optional step = std::nullopt, + ffi::Optional dtype = std::nullopt); /*! * \brief The vectorized For statement. * \param start The minimum value of iteration. * \param stop The maximum value of iteration. * \param annotations The optional annotations of the For statement. * \param step The optional step value of iteration. + * \param dtype The optional dtype of the loop var ("int32" or "uint32"). * \return The ForFrame. */ ForFrame Vectorized(PrimExpr start, PrimExpr stop, ffi::Optional> annotations = std::nullopt, - ffi::Optional step = std::nullopt); + ffi::Optional step = std::nullopt, + ffi::Optional dtype = std::nullopt); /*! * \brief The unrolled For statement. * \param start The minimum value of iteration. * \param stop The maximum value of iteration. * \param annotations The optional annotations of the For statement. * \param step The optional step value of iteration. + * \param dtype The optional dtype of the loop var ("int32" or "uint32"). * \return The ForFrame. */ ForFrame Unroll(PrimExpr start, PrimExpr stop, ffi::Optional> annotations = std::nullopt, - ffi::Optional step = std::nullopt); + ffi::Optional step = std::nullopt, + ffi::Optional dtype = std::nullopt); /*! * \brief The thread-binding For statement. * \param start The minimum value of iteration. @@ -300,9 +331,12 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread, /*! * \brief The grid For statement. * \param extents The extents of the iteration. + * \param dtype The optional dtype of every loop var ("int32" or "uint32"). When omitted + * each loop var takes the dtype of its own extent. * \return The ForFrame. */ -ForFrame Grid(ffi::Array>> extents); +ForFrame Grid(ffi::Array>> extents, + ffi::Optional dtype = std::nullopt); /*! * \brief The assertion statement. diff --git a/python/tvm/backend/cuda/intrinsics/header.py b/python/tvm/backend/cuda/intrinsics/header.py index 2330f75a4f30..b5689352f866 100644 --- a/python/tvm/backend/cuda/intrinsics/header.py +++ b/python/tvm/backend/cuda/intrinsics/header.py @@ -792,16 +792,13 @@ def header_generator(tags): header += R""" __forceinline__ __device__ uint32_t tvm_builtin_elect_one_sync() {{ uint32_t pred = 0; - uint32_t laneid = 0; asm volatile( "{\n" - ".reg .b32 %%rx;\n" ".reg .pred %%px;\n" - " elect.sync %%rx|%%px, %2;\n" - "@%%px mov.s32 %1, 1;\n" - " mov.s32 %0, %%rx;\n" + " elect.sync _|%%px, %1;\n" + "@%%px mov.s32 %0, 1;\n" "}\n" - : "+r"(laneid), "+r"(pred) + : "+r"(pred) : "r"(0xFFFFFFFF)); return pred; }} diff --git a/python/tvm/backend/cuda/intrinsics/sync.py b/python/tvm/backend/cuda/intrinsics/sync.py index f0c084b89ed1..1fcb530d0001 100644 --- a/python/tvm/backend/cuda/intrinsics/sync.py +++ b/python/tvm/backend/cuda/intrinsics/sync.py @@ -97,9 +97,9 @@ def _ptx_barrier_cluster_wait(acquire, aligned): # ============================================================================= -# mbarrier.try_wait.parity.shared::cta.b64 — 1 form. Body wraps the asm in a -# label loop (TIRx convention; the magic ``ticks = 0x989680`` is the timeout -# hint in ns). +# mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 — 1 form. Body wraps the +# asm in a label loop (TIRx convention; the magic ``ticks = 0x989680`` is the +# timeout hint in ns). # ============================================================================= def _mbarrier_wait_parts(barrier, *_rest): """Dispatch on the barrier operand's dtype, as the retired op did. @@ -131,7 +131,7 @@ def _mbarrier_wait_parts(barrier, *_rest): ' "{\\n"\n' ' ".reg .pred P1;\\n"\n' ' "LAB_WAIT:\\n"\n' - ' "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2;\\n"\n' + ' "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\\n"\n' ' "@P1 bra.uni DONE;\\n"\n' ' "bra.uni LAB_WAIT;\\n"\n' ' "DONE:\\n"\n' diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py index e461ebf79dcb..90d100186f11 100644 --- a/python/tvm/backend/cuda/op.py +++ b/python/tvm/backend/cuda/op.py @@ -396,7 +396,7 @@ def _validate_mbarrier_arrive_attrs(sem, scope, space, remote): def cuda_mbarrier_wait(bar, phase): - """TVM intrinsic to call mbarrier.try_wait.parity repeatedly until it returns true + """Retry ``mbarrier.try_wait.parity.acquire.cta`` until it returns true. Parameters ---------- diff --git a/python/tvm/backend/cuda/ptx/engine.py b/python/tvm/backend/cuda/ptx/engine.py index 70a4041a87bc..65eb5c4c6230 100644 --- a/python/tvm/backend/cuda/ptx/engine.py +++ b/python/tvm/backend/cuda/ptx/engine.py @@ -78,7 +78,13 @@ def register_table(table: dict[str, InstructionEntry]) -> None: # The printer name is the *surface* path a user can type, which is the # mnemonic (several `mov_*` entries all answer to `T.ptx.mov`), not the # table key. Reparsing re-dispatches on the operand shape. - family = entry.family + # + # Escaped, because "a user can type it" is the whole requirement and + # three PTX mnemonics are Python keywords: `and`, `or` and `not` (ISA + # 9.7.8) print as `T.ptx.and_` and are read back by `unescape_token` in + # `PTXNamespace.__getattr__`. The escape is the identity for every + # other family, and `gen_stubs` already spells the attribute this way. + family = escape_token(entry.family) register_op_attr(entry.op_name, "TScriptPrinterName", f"ptx.{family}", level=20) register_op_attr(entry.op_name, "TIRxOpCategory", "device_intrin") register_op_attr(entry.op_name, "TDeviceIntrinsicNamespace", "ptx") diff --git a/python/tvm/backend/cuda/ptx/render.py b/python/tvm/backend/cuda/ptx/render.py index 75f127f030b0..6a57fd0fbc22 100644 --- a/python/tvm/backend/cuda/ptx/render.py +++ b/python/tvm/backend/cuda/ptx/render.py @@ -261,9 +261,7 @@ def render_variant( # An immediate lives in the instruction text, never a C parameter. # A literal is table-owned; a choices/open value arrives through # `imms` and is baked here (and into the helper name). - rendered.append( - (slot.bracket, slot.literal if slot.literal is not None else imm_of[slot]) - ) + rendered.append((slot, slot.literal if slot.literal is not None else imm_of[slot])) continue regs = [] n_lanes = lanes_of(slot, mod_map) @@ -355,13 +353,25 @@ def render_variant( asm_post.append(bridge.out_of.format(reg=reg, idx=idx)) regs.append(reg) idx += 1 - rendered.append((slot.bracket, "{" + ", ".join(regs) + "}" if is_group else regs[0])) + rendered.append((slot, "{" + ", ".join(regs) + "}" if is_group else regs[0])) + + # Adjacent slots naming the same `pipe` are one operand written `p|q` + # (setp's two predicate destinations). Merged first, and into a plain text + # entry carrying the group's bracket, so the bracket pass below sees one + # operand where the table declared two. + piped = [] + for key, members in itertools.groupby(rendered, key=lambda pair: pair[0].pipe): + members = list(members) + 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))) # Adjacent slots naming the same `bracket` are one composite memory operand: # `[tensorMap, {c0, c1}]` is a single PTX operand whose members keep their # own registers and constraints. ptx_operands = [] - for key, members in itertools.groupby(rendered, key=lambda pair: pair[0]): + for key, members in itertools.groupby(piped, key=lambda pair: pair[0]): texts = [text for _, text in members] if key is None: ptx_operands.extend(texts) diff --git a/python/tvm/backend/cuda/ptx/table.py b/python/tvm/backend/cuda/ptx/table.py index 4667521f6503..59ca2830aa88 100644 --- a/python/tvm/backend/cuda/ptx/table.py +++ b/python/tvm/backend/cuda/ptx/table.py @@ -74,6 +74,21 @@ "b32": ("uint32", "int32", "float32"), "b64": ("uint64", "int64", "float64"), "b128": ("uint128", "int128"), + # `b32i`/`b64i` are the table's own tokens, not the ISA's: a bit-size type + # with the floating carrier removed. ISA 5.2 says a `.bN` operand accepts + # any fundamental type of that width, and for most instructions it does -- + # but a handful refuse a float register in one position while taking both + # integer signednesses. Measured on ptxas 13.2, at sm_90 unless noted: + # bmsk (all three operands), redux.sync's bitwise line (both), match.sync + # and elect.sync's destination, mbarrier's phase-state token, and + # tensormap.replace's dimension fields (sm_90a) -- each answers + # "Arguments mismatch" to `.f32`/`.f64` and assembles with `.u32`/`.s32` + # (resp. 64-bit). Naming that domain once is what keeps those operands from + # being pinned to `.u32`, which would also throw away the signed spelling + # ptxas accepts. The token never reaches the asm text -- it only chooses + # the dtype axis -- so a name outside the ISA's vocabulary is safe here. + "b32i": ("uint32", "int32"), + "b64i": ("uint64", "int64"), "u8": ("uint8",), "s8": ("int8",), "u16": ("uint16",), @@ -161,6 +176,7 @@ class ModifierSlot: LanesFn = Callable[[dict], int] # modifier map -> registers in this operand's group +DtypeFn = Callable[[dict], str] # modifier map -> this operand's PTX type token @dataclass(frozen=True) @@ -213,6 +229,24 @@ class OperandSlot: ``"rw"``. These are the asm block's sanctioned exceptions to the single-instruction invariant: boundary conversions, never semantics. + Like ``lanes``, ``dtype`` may be a *function* of the modifier map, for the + operands the ISA types by formula rather than by a token in the + instruction text: ``mul.wide``'s destination ("d is twice as wide as a and + b", 9.7.1.3) and ``dp4a``'s accumulator ("c has type .u32 if both .atype + and .btype are .u32 else .s32", 9.7.1.24). Neither type is spellable as a + slot reference, because neither is written in the instruction -- and a + slot could not hold it either: every written token is rendered into the + opcode, so a `.s64` slot would emit `mul.wide.s32.s64`. + + Such a function must be pure in the modifier map, total over every map + ``variants()`` yields (all required slots filled, ``check`` passed), and + return a :data:`PTX_TYPE_DTYPES` key -- the same contract ``lanes`` and + ``sinkable`` carry, and for the same reason: the modifier set is closed, + so every derived token is still enumerable, dispatchable and certifiable. + Write it as a module-level function, never a per-entry lambda: a frozen + dataclass hashes a callable field by identity, so a fresh lambda per entry + would make otherwise-equal slots compare unequal. + ``lanes`` > 1 makes the operand a brace-enclosed register group. PTX writes the group in the operand list (``mov.b64 d, {lo, hi}``), so the group is part of the *shape*, never of the dotted modifier text. @@ -222,7 +256,7 @@ class OperandSlot: rw: str = "r" # "r" | "w" | "rw" -- direction; see the docstring kind: str = "reg" # "reg" | "addr" | "ptr" | "imm" space: str | None = None - dtype: str | None = None + dtype: str | DtypeFn | None = None # 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. @@ -256,6 +290,16 @@ class OperandSlot: # tensorCoords]` is a single PTX operand holding a 64-bit pointer and an # .s32 coordinate vector that have no single-register encoding. bracket: str | None = None + # A pipe-joined operand pair: adjacent slots naming the same `pipe` render + # as one PTX operand written `p|q`, each keeping its own C parameter, + # constraint and bridge. Same idea as `bracket` above, different separator + # and a different reason: `setp.lt.and.s32 p|q, a, b, r;` (ISA 9.7.6.2) + # writes two predicates, the second one from the *complement* of the + # compare result, so `q` is a second destination and not a restatement of + # `p`. The ISA writes the pair in one operand position, so the table does + # too -- a group, like `bracket`, rather than two operands the renderer + # would separate with a comma. + pipe: str | None = None # Whether the ISA lets a lane of this operand be written `_`, the sink # symbol: "this element takes no register". What that means follows the # direction, and it is NOT a destination-only spelling -- the ISA puts it @@ -453,7 +497,15 @@ def operand_type(slot: OperandSlot, mod_map: dict) -> str: that was omitted falls back to ``type``, which is how the mixed-precision forms are typed: ``add.rn.f32.bf16``'s ``a`` is the ``.bf16`` source, while plain ``add.rn.f32``'s ``a`` is just ``.f32``. + + A callable ``dtype`` derives the token from the modifiers instead, for the + operands the ISA types by formula (see :class:`OperandSlot`). This is the + single place any of that is resolved, so every layer downstream -- the + dtype axis, coercion, dispatch, rendering, certification -- sees a derived + token exactly as it sees a written one. """ + if callable(slot.dtype): + return slot.dtype(mod_map) key = slot.dtype or "type" if key in mod_map: return mod_map[key] or mod_map["type"] @@ -856,18 +908,122 @@ def _check_st(m): def _check_rcp(m): - """This entry's rcp.approx is f32-only; .f64 is IEEE-rounded, no .ftz (PTX ISA 9.7.3.13).""" - # Syntax lines: rcp.approx{.ftz}.f32 / rcp.rnd{.ftz}.f32 / rcp.rnd.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. + """ + if m["mode"] == "approx": + if m["type"] == "f64" and not m["ftz"]: + return ( + "there is no rcp.approx.f64: the f64 approximation is spelled " + "rcp.approx.ftz.f64, with .ftz mandatory (PTX ISA 9.7.3.14)" + ) + return None + if m["type"] == "f64" and m["ftz"]: + return "rcp.rnd.f64 takes no .ftz" + return None + + +def _check_sqrt(m): + """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. + """ if m["mode"] == "approx" and m["type"] != "f32": + return "sqrt.approx exists only on the .f32 line; .f64 requires .rn/.rz/.rm/.rp" + if m["type"] == "f64" and m["ftz"]: + return "sqrt.rnd.f64 takes no .ftz" + return None + + +def _check_div_f(m): + """The floating-point divide lines (PTX ISA 9.7.3.8). + + div.approx{.ftz}.f32 d, a, b; div.full{.ftz}.f32 d, a, b; + div.rnd{.ftz}.f32 d, a, b; div.rnd.f64 d, a, b; + + Both approximations (`.approx`, the fast one, and `.full`, the full-range + one) are single-precision only; the ISA spells no rounding modifier on + either, which is why `mode` fuses them with `.rnd` into one slot. + """ + if m["mode"] in ("approx", "full") and m["type"] != "f32": + return f"div.{m['mode']} exists only on the .f32 line; .f64 requires .rn/.rz/.rm/.rp" + if m["type"] == "f64" and m["ftz"]: + return "div.rnd.f64 takes no .ftz" + return None + + +# The .type line shared by most of PTX ISA 9.7.1: mul/mad (9.7.1.3/4), sad +# (9.7.1.8), div (9.7.1.9) and rem (9.7.1.10) all spell exactly this set. +_INT_TYPES = ("u16", "u32", "u64", "s16", "s32", "s64") + +# `.wide`'s result type (ISA 9.7.1.3: "If .wide is specified, then d is twice +# as wide as a and b to receive the full result of the multiplication"; 9.7.1.4 +# says the same of mad's d *and* c). Its keys are also the domain of the .wide +# lines -- "The .wide suffix is supported only for 16- and 32-bit integer +# types" -- so the entry's type slot and this table cannot drift apart. +_WIDE_RESULT = {"u16": "u32", "s16": "s32", "u32": "u64", "s32": "s64"} + + +def _wide_dtype(m): + """`.wide`'s double-width operand type (ISA 9.7.1.3/9.7.1.4).""" + return _WIDE_RESULT[m["type"]] + + +def _dp_acc_dtype(m): + """dp4a/dp2a's accumulator type (ISA 9.7.1.24/9.7.1.25). + + "Operand c has type .u32 if both .atype and .btype are .u32 else operand c + has type .s32" -- and d accumulates into the same type. It is a function of + two modifiers, so no single slot reference can express it. + """ + return "u32" if (m["atype"], m["btype"]) == ("u32", "u32") else "s32" + + +def _check_int_addsub(m): + """Which types the integer add/sub lines take `.sat` on (ISA 9.7.1.1/9.7.1.2). + + add.type1 d, a, b; .type1 = {.u16, .u64, .s16, .s64} + add{.sat}.type2 d, a, b; .type2 = {.u32, .u16x2, .u8x4, .s32, .s16x2, .s8x4} + sub.type1 d, a, b; .type1 = {.u16, .u32, .u64, .s16, .s64} + sub{.sat}.type2 d, a, b; .type2 = {.s32, .u8x4, .s8x4} + + Of the saturating forms only `.s32` is registered: `add.sat` on + .u32/.u16x2/.s16x2 arrived in PTX ISA 9.2 as "sm_120f or higher in the same + family", and the `.u8x4`/`.s8x4` types are excluded outright for the same + reason (see the `_MINMAX_INT_PLAIN` note below). `add.sat.s32` is the PTX + 1.0 form and is supported everywhere. + """ + if m["sat"] and m["type"] != "s32": return ( - "rcp.approx.f64 is a separate syntax line (PTX ISA 9.7.3.14, where .ftz is " - "mandatory) and is not registered here" + f".sat is registered on .s32 only; .sat.{m['type']} is an sm_120f-only form " + f"(PTX ISA 9.2)" ) - if m["type"] == "f64": - if m["mode"] == "approx": - return "rcp.f64 requires an IEEE rounding mode (.rn/.rz/.rm/.rp)" - if m["ftz"]: - return "rcp.rnd.f64 takes no .ftz" + return None + + +def _check_int_mad(m): + """`.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". + """ + if m["sat"] and (m["mode"], m["type"]) != ("hi", "s32"): + return ".sat applies only to the .hi.s32 form" return None @@ -972,19 +1128,446 @@ def _check_half_arith(m): return None -def _check_neg(m): - """`neg{.ftz}.f32 d, a;` and `neg.f64 d, a;` (ISA 9.7.3.10) -- the f64 - line spells no `.ftz`.""" +# The .type line every PTX ISA 9.7.4 instruction spells: a scalar and a packed +# pair in each of the two half formats. `.f16`/`.bf16` ride a 16-bit register +# and the packed pairs a 32-bit one (PTX_TYPE_DTYPES). +_HALF_TYPES = ("f16", "f16x2", "bf16", "bf16x2") + + +def _check_half_fma(m): + """The half-precision fma lines, ISA 9.7.4.4. + + fma.rnd{.ftz}{.sat}.f16 d, a, b, c; fma.rnd{.ftz}.relu.f16 d, a, b, c; + fma.rnd{.ftz}{.sat}.f16x2 d, a, b, c; fma.rnd{.ftz}.relu.f16x2 d, a, b, c; + fma.rnd{.relu}.bf16 d, a, b, c; fma.rnd.oob{.relu}.type d, a, b, c; + fma.rnd{.relu}.bf16x2 d, a, b, c; + + Three rules read off those lines. The bf16 ones take neither `.ftz` nor + `.sat`, as they do for add/sub/mul. `.sat` and `.relu` are two different + clampings on two different syntax lines, so no line offers both. And the + `.oob` line spells only `{.relu}` beside it -- it is written over the whole + `.type` set, but with no `.ftz` and no `.sat`. + """ + if m["type"].startswith("bf16") and (m["ftz"] or m["sat"]): + return f".{m['type']} takes no .ftz or .sat" + if m["oob"] and (m["ftz"] or m["sat"]): + return "the .oob line spells no .ftz or .sat, only {.relu}" + if m["sat"] and m["relu"]: + return ".sat and .relu are separate syntax lines: write one or the other" + return None + + +def _check_half_absneg(m): + """The half-precision abs/neg lines, ISA 9.7.4.6 / 9.7.4.5. + + op{.ftz}.f16 d, a; op{.ftz}.f16x2 d, a; + op.bf16 d, a; op.bf16x2 d, a; + + Same shape as the f32/f64 pair (`_check_absneg`), with the bf16 rule of the + half group in place of the f64 one. It cannot share `_check_half_arith`: + that reads a `.sat` slot these entries do not declare. + """ + if m["type"].startswith("bf16") and m["ftz"]: + return f".{m['type']} takes no .ftz" + return None + + +def _check_half_ex2(m): + """The half-precision ex2 lines, ISA 9.7.4.10. + + ex2.approx.atype d, a; .atype = {.f16, .f16x2} + ex2.approx.ftz.btype d, a; .btype = {.bf16, .bf16x2} + + `.ftz` is not optional here in either direction: the bf16 line spells it + mandatorily and the f16 line does not offer it at all. That is what keeps + these lines out of the `.f32` ex2 entry, whose `.ftz` is optional. + """ + if m["type"].startswith("bf16"): + if not m["ftz"]: + return "the bf16 lines spell .ftz mandatorily: ex2.approx.ftz.bf16{x2}" + return None + if m["ftz"]: + return f".{m['type']} has no .ftz line (only the bf16 lines spell it)" + return None + + +# The comparison operators of set/setp (PTX ISA 9.7.6.1/9.7.6.2), grouped the +# way the Integer Notes and Floating Point Notes group them. +_CMP_ORDERED_OPS = ("eq", "ne", "lt", "le", "gt", "ge") +# "For unsigned values, the comparison operators lo, ls, hi, and hs ... may be +# used instead of lt, le, gt, ge" -- alternates, not additions, and unsigned only. +_CMP_UNSIGNED_OPS = ("lo", "ls", "hi", "hs") +# The unordered counterparts, plus the two NaN predicates. Floating point only. +_CMP_UNORDERED_OPS = ("equ", "neu", "ltu", "leu", "gtu", "geu", "num", "nan") +_CMP_OPS = (*_CMP_ORDERED_OPS, *_CMP_UNSIGNED_OPS, *_CMP_UNORDERED_OPS) +# `.BoolOp`, which combines the compare result with a third predicate operand. +_CMP_BOOL_OPS = ("and", "or", "xor") +# The source type line, shared by set, setp and selp. +_CMP_TYPES = ("b16", "b32", "b64", "u16", "u32", "u64", "s16", "s32", "s64", "f32", "f64") + + +def _check_cmp(m): + """Which comparison operators each source type accepts (ISA 9.7.6.1/9.7.6.2). + + The `.CmpOp` line in the Syntax block is the union over all types; the + per-type rule is in the Notes, and it is three disjoint sets: + + - bit-size: "The untyped, bit-size comparisons are eq and ne." + - signed: the six ordered operators, and nothing else. + - unsigned: those six plus lo/ls/hi/hs, the unsigned-only alternates. + - float: the six ordered ones plus the unordered equ/neu/ltu/leu/gtu/geu + and the NaN predicates num/nan -- never the unsigned alternates. + + `.ftz` is separate and simpler: "Modifier .ftz applies only to .f32 + comparisons." + """ + # `set` names its source type `stype` (its `dtype` is the destination); + # `setp` has a single `type` slot. The comparison is on the source either way. + ty = m.get("stype") or m["type"] + op = m["cmp"] + if m["ftz"] and ty != "f32": + return ".ftz applies only to .f32 comparisons" + if ty.startswith("b"): + if op not in ("eq", "ne"): + return f"the bit-size type .{ty} compares only with eq/ne" + return None + if ty.startswith("s"): + if op not in _CMP_ORDERED_OPS: + return f".{ty} is signed: {'/'.join(_CMP_ORDERED_OPS)} only" + return None + if ty.startswith("u"): + if op in _CMP_UNORDERED_OPS: + return f"{op} is a floating-point comparison, not valid on .{ty}" + return None + if op in _CMP_UNSIGNED_OPS: + return f"{op} is an unsigned-only alternate, not valid on .{ty}" + return None + + +# The `.CmpOp` line of the half-precision comparisons (PTX ISA 9.7.7): the +# ordered six and the unordered six plus num/nan. The unsigned alternates +# lo/ls/hi/hs of 9.7.6 are absent from this section's line entirely. +# +# MEASURED, NOT REGISTERED: ptxas does accept them on an integer source with a +# half destination (`set.lo.f16.u32` assembles at sm_90), but no 9.7.7 syntax +# line spells them, so they stay out -- the table follows the ISA where ptxas +# is the more permissive of the two. +_HALF_CMP_OPS = (*_CMP_ORDERED_OPS, *_CMP_UNORDERED_OPS) +# The `.stype` line of `set.CmpOp{.ftz}.f16.stype` and its bf16 twin. Note what +# is NOT in it: `.bf16` itself. (`set.bf16.bf16` also assembles, and is left out +# for the same reason as above -- the ISA spells no line for it.) +_SET_HALF_STYPES = ( + "b16", "b32", "b64", "u16", "u32", "u64", "s16", "s32", "s64", "f16", "f32", "f64", +) # fmt: skip +# The packed half types, which 9.7.7 treats as one pair of lanes per register. +_HALF_X2 = ("f16x2", "bf16x2") +# The bit-size widths the logic and shift instructions operate on (PTX ISA +# 9.7.8): "fundamentally untyped ... provided the operands are of the same +# size". `_LOGIC_TYPES` adds the predicate line that and/or/xor/not also carry. +_BIT_TYPES = ("b16", "b32", "b64") +_LOGIC_TYPES = ("pred", *_BIT_TYPES) +# scalar mov's type line (PTX ISA 9.7.9.3): "Although only predicate and +# bit-size types are required, we include the arithmetic types for the +# programmer's convenience". +_MOV_TYPES = ("pred", "b16", "b32", "b64", "u16", "u32", "u64", "s16", "s32", "s64", "f32", "f64") +# The state spaces cvta converts between and isspacep queries (ISA 9.7.9.21, +# 9.7.9.20) -- the same eight, spelled the same way, in both instructions. +_CVTA_SPACES = ( + "const", "global", "local", "shared", "shared::cta", "shared::cluster", "param", "param::entry", +) # fmt: skip +# The source types `set` will take `.ftz` with. MEASURED: the ISA writes +# `{.ftz}` across the whole `set.CmpOp{.ftz}.f16.stype` line, but ptxas 13.2 +# answers "Illegal modifier '.ftz' for instruction 'set'" for every source +# outside this set -- probed over all 8 destination x 15 source pairs at +# sm_90. The rule that survives the probe is the one the modifier means: +# `.ftz` flushes subnormal *inputs*, so it attaches to the source's precision, +# and only these three have subnormals to flush at single-or-half width. +_FTZ_SET_STYPES = ("f16", "f32", "f16x2") + + +def _check_half_set(m): + """The six syntax-line groups of the half-precision `set` (ISA 9.7.7.1). + + set.CmpOp{.ftz}.f16.stype / set.CmpOp.bf16.stype .stype = the 12 above + set.CmpOp{.ftz}.dtype.f16 / set.CmpOp.dtype.bf16 .dtype = {u16,s16,u32,s32} + set.CmpOp{.ftz}.dtype.f16x2 .dtype = {f16x2,u32,s32} + set.CmpOp.dtype.bf16x2 .dtype = {bf16x2,u32,s32} + + Three independent rules come out of that, and each was probed against ptxas + before being written here: + + - The (dtype, stype) pairing is a grid, not a product: a half type appears + on exactly one side of every line. The integer destinations reach only + half sources, and the packed destinations only their own packed source. + The wide integer cells this leaves out -- (u32, b32) and friends -- are + not lost: they are the 9.7.6 `set` lines, which own that dtype already. + - `.ftz` follows the *source* precision, not the line it is written on: + only a `.f16`/`.f32`/`.f16x2` source has subnormals to flush at this + width, and no bf16 destination takes the modifier at all. See + `_FTZ_SET_STYPES` -- this one is measured against ptxas rather than read + off the syntax line, which writes `{.ftz}` more broadly than ptxas + accepts it. + - The comparison operator follows the *source*: a bit-size source compares + only for equality, an integer source takes the ordered six, and a + floating-point source takes all fourteen. + """ + dt, st, op = m["dtype"], m["stype"], m["cmp"] + paired = ( + (dt in ("f16", "bf16") and st in _SET_HALF_STYPES) + or (dt in ("u16", "s16", "u32", "s32") and st in ("f16", "bf16")) + or (dt in ("u32", "s32") and st in ("f16x2", "bf16x2")) + or (dt == "f16x2" and st == "f16x2") + or (dt == "bf16x2" and st == "bf16x2") + ) + if not paired: + return f"no half-precision syntax line pairs .{dt} with .{st} (ISA 9.7.7.1)" + if m["ftz"]: + if dt.startswith("bf16"): + return f"a .{dt} destination takes no .ftz" + if st not in _FTZ_SET_STYPES: + return f".ftz flushes subnormal inputs, so a .{st} source does not take it" + if st.startswith("b"): + if op not in ("eq", "ne"): + return f"the bit-size source .{st} compares only with eq/ne" + elif st[0] in "us": + if op not in _CMP_ORDERED_OPS: + return f"the integer source .{st} takes {'/'.join(_CMP_ORDERED_OPS)} only" + return None + + +def _check_half_setp(m): + """`.ftz` on the half-precision `setp` lines (ISA 9.7.7.2). + + Spelled on `.f16`/`.f16x2`, on neither bf16 line -- the same split the half + arithmetic lines make. The other half of this section's shape, single + destination versus `p|q`, is not a check: the ISA gives each type exactly + one of the two, so the type domains of the entries say it instead. + """ + if m["ftz"] and m["type"].startswith("bf16"): + return f".{m['type']} spells no .ftz" + return None + + +# --- PTX ISA 9.7.9.12 (st.async) and 9.7.9.15 (multimem) -------------------- +# The vector lines of st.async: "`.v2` is supported with .b32, .b64, .s32, +# .s64, .u32, .u64, .f32 and .f64 types. `.v4` qualifier is supported with +# .b32, .s32, .u32 and .f32 types." +_ST_ASYNC_TYPES = ("b32", "b64", "b128", "u32", "u64", "s32", "s64", "f32", "f64") +_ST_ASYNC_V2 = ("b32", "b64", "u32", "u64", "s32", "s64", "f32", "f64") +_ST_ASYNC_V4 = ("b32", "u32", "s32", "f32") +# The release line's type list (9.7.9.12, second syntax block), which reaches +# below the 32-bit floor of the mbarrier line. +# MEASURED: the ISA writes `.b8`, `.u8` and `.s8` into that list, but ptxas +# 13.2 rejects all three at sm_100 -- with and without `.mmio`, at either scope +# -- while every 16-bit and wider type assembles. So the line starts at 16 bits +# here, and the byte forms wait for a toolchain that accepts them. +_ST_ASYNC_REL_TYPES = ( + "b16", "b32", "b64", "u16", "u32", "u64", "s16", "s32", "s64", "f32", "f64", +) # fmt: skip + +# createpolicy's `.level::primary_priority` line (ISA 9.7.9.19). Wider than +# the `_L2_EVICT` set the ld/st eviction hints use: this instruction is where +# a priority is *created*, so it spells all four. +_CACHE_PRIORITIES = ( + "L2::evict_last", "L2::evict_normal", "L2::evict_first", "L2::evict_unchanged", +) # fmt: skip + +_MM_INT_TYPES = ("b32", "b64", "u32", "u64", "s32", "s64") +_MM_FLOAT_TYPES = ("f16", "f16x2", "bf16", "bf16x2", "f32", "f64") +# op x type for the integer multimem lines, measured cell by cell against +# ptxas (the ISA states it as a table, and its `.add` row omits .s64). +_MM_INT_OPS = { + "min": ("u32", "u64", "s32", "s64"), + "max": ("u32", "u64", "s32", "s64"), + "add": ("u32", "u64", "s32"), + "and": ("b32", "b64"), + "or": ("b32", "b64"), + "xor": ("b32", "b64"), +} +# .vec x base type for the floating-point multimem lines. The ISA gives this +# as a table ("the size of the specified type along with .vec must equal either +# 32-bits or 64-bits or 128-bits"); every cell below was probed. +_MM_VEC_TYPES = { + "": ("f16x2", "bf16x2", "f32", "f64"), + "v2": ("f16", "f16x2", "bf16", "bf16x2", "f32"), + "v4": ("f16", "f16x2", "bf16", "bf16x2", "f32"), + "v8": ("f16", "bf16"), +} + + +def _check_multimem_sem(m): + """How `.sem` and `.scope` go together on every multimem line (ISA 9.7.9.15). + + The ISA writes two syntax lines per mnemonic -- `{.sem}{.scope}` and a + `.weak` one with no scope position -- which reads as "either may be + omitted". ptxas is stricter, and its rule is the one that makes the two + lines distinct: everywhere but `.weak`, the pair is the same both-or- + neither rule the mbarrier sections state in words, so that check is what + decides it here. ptxas says both halves -- "Modifier '.relaxed' requires + scope with 'multimem.st' instruction" and "Modifier '.cta' requires order + with 'multimem.st' instruction". What multimem adds is `.weak`: the line + with no scope position at all. + """ + if m.get("sem", "") == "weak": + return ".weak is a syntax line of its own and takes no scope" if m["scope"] else None + return _check_mbarrier_sem_scope(m) + + +def _check_multimem_int(m): + """op x type on the integer multimem lines (ISA 9.7.9.15). + + The `.type` line in the Syntax block is the union over ops; the pairing is + the "valid combinations of .op and base type" table, and `.add` is the row + that surprises -- it takes .s32 but not .s64. + + The `.sem`/`.scope` pairing is shared with the floating lines; see + `_check_multimem_sem`. + """ + error = _check_multimem_sem(m) + if error: + return error + op = m.get("op", "") + if op and m["type"] not in _MM_INT_OPS[op]: + return f".{op} takes {' / '.join('.' + t for t in _MM_INT_OPS[op])}" + return None + + +def _check_multimem_f(m): + """`.vec` x type, `.op` x type and `.acc::f32` on the floating multimem lines. + + Three rules, all measured against ptxas over the full grid: + + - The vector width and the element type together have to make 32, 64 or + 128 bits, which is what `_MM_VEC_TYPES` tabulates. Note `.f64` takes no + `.vec` at all, and the scalar (no-`.vec`) line takes no `.f16`/`.bf16` -- + a lone half is not one of the allowed widths. + - `.min`/`.max` are half-only: the ISA's op table gives `.f32`/`.f64` to + `.add` alone. + - `.acc::f32` raises the accumulation precision of a half type, so it is + spelled only where the type is one. + + The `.sem`/`.scope` pairing is shared with the integer lines; see + `_check_multimem_sem`. + """ + error = _check_multimem_sem(m) + if error: + return error + ty, vec, op = m["type"], m.get("vec", ""), m.get("op", "") + if ty not in _MM_VEC_TYPES[vec]: + allowed = " / ".join("." + t for t in _MM_VEC_TYPES[vec]) + return f"{'.' + vec if vec else 'the scalar line'} takes {allowed}" + if op in ("min", "max") and ty in ("f32", "f64"): + return f".{op} is half-precision only; .{ty} is on the .add row" + if m.get("acc"): + if ty in ("f32", "f64"): + return f".acc::f32 raises a half accumulation to .f32; .{ty} is already wider" + if op != "add": + # ptxas: "Illegal reduction operation for instruction + # 'multimem.ld_reduce.acc::f32'". Only a sum accumulates, so only a + # sum has an accumulation precision to raise. + return f".acc::f32 applies to .add; .{op} keeps the operand precision" + return None + + +def _check_st_async(m): + """st.async's mbarrier line (ISA 9.7.9.12), whose `.vec` narrows the types. + + `.v4` is the 32-bit four and `.v2` everything but `.b128`, which the ISA + gives no vector form at all. + """ + vec, ty = m["vec"], m["type"] + allowed = _ST_ASYNC_V2 if vec == "v2" else _ST_ASYNC_V4 + if ty not in allowed: + return f".{vec} takes {' / '.join('.' + t for t in allowed)}" + return None + + +def _check_st_async_rel(m): + """st.async's release line: "If .mmio is specified, .scope must be .sys" + (ISA 9.7.9.12), which ptxas enforces as an illegal-modifier error.""" + if m["mmio"] and m["scope"] != "sys": + return ".mmio requires .sys scope" + return None + + +# --- PTX ISA 9.7.14 warp-level and async reductions ------------------------- +# red.async's four mbarrier lines (9.7.14.7), one op group each. +_RED_ASYNC_OPS = {"inc": ("u32",), "dec": ("u32",), "min": ("u32", "s32"), + "max": ("u32", "s32"), "and": ("b32",), "or": ("b32",), "xor": ("b32",), + "add": ("u32", "s32", "u64")} # fmt: skip +# atom's two vector lines (9.7.14.5). A half element rides `.v2`/`.v4`/`.v8`; +# a packed pair, being twice as wide, stops at `.v4`; and `.f32` has only the +# `.add` line. Every cell probed. +_ATOM_VEC_HALF = ("f16", "bf16") + + +def _check_red_async(m): + """op x type on red.async's mbarrier lines (ISA 9.7.14.7). + + The ISA writes one syntax line per op group rather than a table, so the + grid is read off those four lines: increment/decrement on the unsigned + word, min/max on either signed word, the bitwise ops on the untyped one, + and add reaching to 64 bits. + """ + op, ty = m["op"], m["type"] + if ty not in _RED_ASYNC_OPS[op]: + return f".{op} takes {' / '.join('.' + t for t in _RED_ASYNC_OPS[op])}" + return None + + +def _check_atom_vec(m): + """`.vec` x type on atom's vector lines (ISA 9.7.14.5). + + atom{.sem}{.scope}{.global}.add{.cache}.vec_32_bit.f32 d, [a], b; + atom{...}.op.noftz{.cache}.vec_16_bit.half_word_type d, [a], b; + atom{...}.op.noftz{.cache}.vec_32_bit.packed_type d, [a], b; + + The element width bounds the vector, and one cell of that is all this + check has left to say. The `.f32` line's own entry stops its `.vec` slot at + `.v4`, and the op split falls out of the two entries' op slots; what no + slot can express is that the half line's `.v8` belongs to a *lone* half, + since a packed pair is already 32 bits wide. + """ + if m["vec"] == "v8" and m["type"] in _HALF_X2: + return f".{m['type']} is already a 32-bit pair, so it reaches .v4; .v8 is for a lone half" + return None + + +def _check_slct(m): + """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. + """ + if m["ftz"] and m["ctype"] != "f32": + return ".ftz is spelled only on the .f32 selector line" + return None + + +def _check_absneg(m): + """The one-source sign lines, `op{.ftz}.f32 d, a;` and `op.f64 d, a;`. + + abs (ISA 9.7.3.9) and neg (9.7.3.10) are spelled identically -- the f64 + line of each carries no `.ftz` -- so they share this check. + """ if m["ftz"] and m["type"] != "f32": return ".ftz appears only on the .f32 line" return None def _check_farith(m): - """Which qualifiers each add/sub/mul/fma syntax line allows (PTX ISA 9.7.3.{3,4,5,6}, 9.7.5). + """Which qualifiers each add/sub/mul/fma/mad line allows (PTX ISA 9.7.3.{3,4,5,6,7}, 9.7.5). Same-precision lines: op{.rnd}{.ftz}{.sat}.f32 | op{.rnd}{.ftz}.f32x2 | op{.rnd}.f64 Mixed-precision lines: op{.rnd}{.sat}.f32.atype (.atype = .f16 | .bf16) + + `mad` (9.7.3.7) is the same grid minus `.f32x2` and minus the mixed lines, + both of which its entry excludes by slot domain, so it shares this check: + with no `srctype` slot the mixed branch never fires, and `.ftz`/`.sat` + stay gated to the .f32 line exactly as the syntax spells them. """ ty, src = m["type"], m.get("srctype", "") if src: @@ -1893,6 +2476,18 @@ def _check_tcgen05_mma_block_scale(m): # ------------------------------------------------------------------ # PTX ISA 9.7.1 — Integer Arithmetic Instructions # ------------------------------------------------------------------ + # The section is registered in full, with two exceptions, each stated at + # the entry it belongs to: the packed `.u8x4`/`.s8x4` types (and the + # saturating forms that arrived with them), which the ISA gives only to + # "sm_120f or higher in the same family"; and `clmad` (9.7.1.5), which is + # PTX ISA 9.3 and does not assemble on this toolchain. + # + # Several mnemonics have a floating-point line further down the table. The + # integer lines are separate entries because they share no qualifier with + # the fp ones -- one merged entry would offer `.rnd`/`.ftz` on `.s32` -- and + # dispatch resolves them apart on their type tokens, exactly as it already + # does for `add`/`add_half`. + # # min / max per PTX ISA 9.7.1.13, 9.7.1.14 (integer), 9.7.3.11, 9.7.3.12 # (single/double), 9.7.4.7, 9.7.4.8 (half). # Each mnemonic gets two entries because the ISA gives it two operand @@ -1959,141 +2554,1035 @@ def _check_tcgen05_mma_block_scale(m): ), asm_volatile=False, # legacy fns carried no barrier ), - # ------------------------------------------------------------------ - # PTX ISA 9.7.3 — Floating-Point Instructions - # ------------------------------------------------------------------ - # Floating-point add/sub/mul (PTX ISA 9.7.3.{3,4,5}) together with their - # mixed-precision lines (9.7.5.{1,2}); `mul` has no mixed-precision line. - # add{.rnd}{.ftz}{.sat}.f32 d, a, b; add{.rnd}{.ftz}.f32x2 d, a, b; - # add{.rnd}.f64 d, a, b; add{.rnd}{.sat}.f32.atype d, a, c; - # cert_arch is the family's ceiling, not its floor: .f32/.f64 assemble - # everywhere, but .f32x2 and the mixed lines need sm_100, and certification - # has to run somewhere every legal variant is legal. - *[ - InstructionEntry( - name=name, - cert_arch="sm_100", - slots=( - ModifierSlot("rnd", _FRND, optional=True), - ModifierSlot("ftz", ("ftz",), optional=True), - ModifierSlot("sat", ("sat",), optional=True), - ModifierSlot("type", ("f32", "f64", "f32x2")), - *((ModifierSlot("srctype", ("f16", "bf16"), optional=True),) if mixed else ()), - ), - check=_check_farith, - operands=( - OperandSlot("d", rw="w"), - # On the mixed line `a` is the converted 16-bit source; on every - # other line it is just the instruction type. - OperandSlot("a", dtype="srctype" if mixed else None), - OperandSlot("b"), - ), - ) - # `mul` is the one line with no mixed-precision form (ISA 9.7.5). - for name, mixed in (("add", True), ("sub", True), ("mul", False)) - ], - # NOT REGISTERED: across this whole arithmetic group, the integer lines - # (9.7.1.{1,2,3}), extended-precision add.cc/sub.cc (9.7.2.{1,3}), and the - # half-precision fma line (9.7.4.4), which needs `.relu` and `.oob` slots. - # Those two qualifiers appear on no add/sub/mul half line, which is why - # 9.7.4.{1,2,3} are registered above and 9.7.4.4 is not. + # The rest of 9.7.1, in ISA order. All of it keeps the default + # `asm_volatile`, which is also what every hand-written arithmetic helper + # in the ecosystem does: CUTLASS (`cute/arch/simd_sm100.hpp`'s + # add/mul/fma.f32x2, `fast_math.h`'s tanh, `functional.h`'s rcp/min/max) + # and flashinfer (`math.cuh`'s ex2/lg2/rcp/rsqrt/tanh) write `asm volatile` + # on register-only instructions without exception. The needless-barrier + # worry is about the "memory" clobber, which is derived and stays off here + # (see render.py); `volatile` alone only stops nvcc reordering or commoning + # up the asm, and the whole arithmetic surface is uniform in taking it. # - # fma differs in shape, so it is its own entry: three sources, and .rnd is - # mandatory on every line (PTX ISA 9.7.3.6 / 9.7.5.3). - # fma.rnd{.ftz}{.sat}.f32 d, a, b, c; fma.rnd{.ftz}.f32x2 d, a, b, c; - # fma.rnd.f64 d, a, b, c; fma.rnd{.sat}.f32.abtype d, a, b, c; + # add / sub per PTX ISA 9.7.1.1, 9.7.1.2: + # add.type1 d, a, b; .type1 = {.u16, .u64, .s16, .s64} + # add{.sat}.type2 d, a, b; .type2 = {.u32, .u16x2, .u8x4, .s32, .s16x2, .s8x4} + # sub.type1 d, a, b; .type1 = {.u16, .u32, .u64, .s16, .s64} + # sub{.sat}.type2 d, a, b; .type2 = {.s32, .u8x4, .s8x4} + # The two lines differ only in whether `.sat` is offered, so one entry with + # an optional slot spells both and `_check_int_addsub` draws the boundary. + # NOT REGISTERED: the `.u8x4`/`.s8x4` types on either mnemonic, and `.sat` + # on .u32/.u16x2/.s16x2 -- all sm_120f-only (PTX ISA 9.2), the same + # exclusion the min/max entry above makes for the same reason. InstructionEntry( - name="fma", - cert_arch="sm_100", + name="add_int", + mnemonic="add", slots=( - ModifierSlot("rnd", _FRND), - ModifierSlot("ftz", ("ftz",), optional=True), ModifierSlot("sat", ("sat",), optional=True), - ModifierSlot("type", ("f32", "f64", "f32x2")), - ModifierSlot("srctype", ("f16", "bf16"), optional=True), + ModifierSlot("type", ("u16", "u32", "u64", "u16x2", "s16", "s32", "s64", "s16x2")), ), - check=_check_farith, + check=_check_int_addsub, + # `.u16x2`/`.s16x2` require sm_90 (ISA Target Notes); cert_arch is the + # max floor over the entry's variants. + cert_arch="sm_90", operands=( OperandSlot("d", rw="w"), - # .abtype converts both a and b; c is always the instruction type. - OperandSlot("a", dtype="srctype"), - OperandSlot("b", dtype="srctype"), - OperandSlot("c"), + OperandSlot("a"), + OperandSlot("b"), ), ), - # neg (PTX ISA 9.7.3.10): `neg{.ftz}.f32 d, a;` and `neg.f64 d, a;`. InstructionEntry( - name="neg", + name="sub_int", + mnemonic="sub", slots=( - ModifierSlot("ftz", ("ftz",), optional=True), - ModifierSlot("type", ("f32", "f64")), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", _INT_TYPES), ), - check=_check_neg, + check=_check_int_addsub, operands=( OperandSlot("d", rw="w"), OperandSlot("a"), + OperandSlot("b"), ), ), - # rcp per PTX ISA 9.7.3.13. `rcp.approx.ftz.f64` (a separate syntax line in - # the ISA, with its own sm floor) is excluded until it is needed. + # mul / mad per PTX ISA 9.7.1.3, 9.7.1.4: + # mul.mode.type d, a, b; mad.mode.type d, a, b, c; + # .mode = {.hi, .lo, .wide} mad.hi.sat.s32 d, a, b, c; + # `.wide` is a separate entry per mnemonic, not a third token in `mode`: + # it changes the result *structure* ("d is twice as wide as a and b", and + # for mad so is c), which is the table's rule for where one entry ends. + # Splitting it also lets the slot domains state "supported only for 16- and + # 32-bit integer types" directly, with no check needed. InstructionEntry( - name="rcp", + name="mul_int", + mnemonic="mul", slots=( - ModifierSlot("mode", ("approx", "rn", "rz", "rm", "rp")), - ModifierSlot("ftz", ("ftz",), optional=True), - ModifierSlot("type", ("f32", "f64")), + ModifierSlot("mode", ("hi", "lo")), + ModifierSlot("type", _INT_TYPES), ), - check=_check_rcp, operands=( OperandSlot("d", rw="w"), - OperandSlot("value"), + OperandSlot("a"), + OperandSlot("b"), ), ), - # ex2 per PTX ISA 9.7.3.21 (`ex2.approx{.ftz}.f32`). The half-precision - # forms of 9.7.4.10 (.f16/.f16x2/.bf16/.bf16x2) are deliberately excluded: - # .f16/.bf16 have carriers now, but .f16x2/.bf16x2 need a b32 carrier and - # .ftz is mandatory on the bf16 line while illegal on the f16 line, so they - # cannot share this entry's optional ftz slot. InstructionEntry( - name="ex2", + name="mul_wide", + mnemonic="mul", slots=( - ModifierSlot("mode", ("approx",)), - ModifierSlot("ftz", ("ftz",), optional=True), - ModifierSlot("type", ("f32",)), + ModifierSlot("mode", ("wide",)), + ModifierSlot("type", tuple(_WIDE_RESULT)), + ), + operands=( + OperandSlot("d", rw="w", dtype=_wide_dtype), + OperandSlot("a"), + OperandSlot("b"), + ), + ), + InstructionEntry( + name="mad_int", + mnemonic="mad", + slots=( + ModifierSlot("mode", ("hi", "lo")), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", _INT_TYPES), ), + check=_check_int_mad, operands=( OperandSlot("d", rw="w"), - OperandSlot("value"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), ), ), - # ------------------------------------------------------------------ - # PTX ISA 9.7.4 — Half Precision Floating-Point Instructions - # ------------------------------------------------------------------ - # Half-precision add/sub/mul (PTX ISA 9.7.4.{1,2,3}). Same operand shape as - # the single/double lines above, so they differ only in their slot domains: - # sub{.rnd}{.ftz}{.sat}.f16 / .f16x2 sub{.rnd}.bf16 / .bf16x2 .rnd = {.rn} + InstructionEntry( + name="mad_wide", + mnemonic="mad", + slots=( + ModifierSlot("mode", ("wide",)), + ModifierSlot("type", tuple(_WIDE_RESULT)), + ), + operands=( + OperandSlot("d", rw="w", dtype=_wide_dtype), + OperandSlot("a"), + OperandSlot("b"), + # "If .wide is specified, then d and c are twice as wide as a and b + # to receive the result of the multiplication" (ISA 9.7.1.4). + OperandSlot("c", dtype=_wide_dtype), + ), + ), + # NOT REGISTERED: clmad (PTX ISA 9.7.1.5, `clmad.mode.u64 d, a, b, c;`). + # It was introduced in PTX ISA 9.3; this toolchain assembles 9.2, so ptxas + # rejects every form and certification could not prove a single variant. + # Register it when the toolchain catches up -- it is an ordinary + # mode + fixed-type entry with four .u64 operands. + # + # mul24 / mad24 per PTX ISA 9.7.1.6, 9.7.1.7: a 24x24-bit multiply held in + # 32-bit registers, so there is no `.wide` and the type line is 32-bit only. + # mul24.mode.type d, a, b; mad24.mode.type d, a, b, c; + # .mode = {.hi, .lo} mad24.hi.sat.s32 d, a, b, c; + InstructionEntry( + name="mul24", + slots=( + ModifierSlot("mode", ("hi", "lo")), + ModifierSlot("type", ("u32", "s32")), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + ), + ), + InstructionEntry( + name="mad24", + slots=( + ModifierSlot("mode", ("hi", "lo")), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", ("u32", "s32")), + ), + check=_check_int_mad, + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + ), + ), + # sad per PTX ISA 9.7.1.8: `sad.type d, a, b, c;` -- d = c + |a - b|. + InstructionEntry( + name="sad", + slots=(ModifierSlot("type", _INT_TYPES),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + ), + ), + # div / rem per PTX ISA 9.7.1.9, 9.7.1.10: `div.type d, a, b;`. These are + # the integer lines; the floating-point `div` (9.7.3.8) is the `div_f` + # entry below, sharing this mnemonic. *[ InstructionEntry( - name=f"{name}_half", - mnemonic=name, + name=name, + slots=(ModifierSlot("type", _INT_TYPES),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + ), + ) + for name in ("div", "rem") + ], + # abs / neg per PTX ISA 9.7.1.11, 9.7.1.12 -- signed integers only. Their + # floating-point lines (9.7.3.9, 9.7.3.10) are the `abs_f` and `neg` + # entries below; which of each pair carries a suffix is just which arrived + # second. NOT REGISTERED: `neg.s8x4`, sm_120f-only as above. + InstructionEntry( + name="abs", + slots=(ModifierSlot("type", ("s16", "s32", "s64")),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ), + InstructionEntry( + name="neg_int", + mnemonic="neg", + slots=(ModifierSlot("type", ("s16", "s32", "s64")),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ), + # popc / clz per PTX ISA 9.7.1.15, 9.7.1.16: `popc.type d, a;`. Both count + # bits of a `.b32`/`.b64` source into a `.u32` destination whatever the + # source width -- "destination d has type .u32" -- so d is typed outright + # while a takes the instruction type. + *[ + InstructionEntry( + name=name, + slots=(ModifierSlot("type", ("b32", "b64")),), + operands=( + OperandSlot("d", rw="w", dtype="u32"), + OperandSlot("a"), + ), + ) + for name in ("popc", "clz") + ], + # bfind per PTX ISA 9.7.1.17: `bfind{.shiftamt}.type d, a;`, d again .u32. + InstructionEntry( + name="bfind", + slots=( + ModifierSlot("shiftamt", ("shiftamt",), optional=True), + ModifierSlot("type", ("u32", "u64", "s32", "s64")), + ), + operands=( + OperandSlot("d", rw="w", dtype="u32"), + OperandSlot("a"), + ), + ), + # brev per PTX ISA 9.7.1.19: `brev.type d, a;`. + InstructionEntry( + name="brev", + slots=(ModifierSlot("type", ("b32", "b64")),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ), + # bfe / bfi per PTX ISA 9.7.1.20, 9.7.1.21. Both take their position and + # length operands as `.u32` regardless of the instruction type ("Operands b + # and c are type .u32, but are restricted to the 8-bit value range 0..255"). + # bfe.type d, a, b, c; .type = {.u32, .u64, .s32, .s64} + # bfi.type f, a, b, c, d; .type = {.b32, .b64} + InstructionEntry( + name="bfe", + slots=(ModifierSlot("type", ("u32", "u64", "s32", "s64")),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b", dtype="u32"), # start bit position + OperandSlot("c", dtype="u32"), # field length + ), + ), + InstructionEntry( + name="bfi", + # The ISA's own operand names, kept as they are: the destination is `f` + # and `d` is the *length input*, the one family where `d` is not the + # result. Renaming them would make the helper unreadable against 9.7.1.21. + slots=(ModifierSlot("type", ("b32", "b64")),), + operands=( + OperandSlot("f", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c", dtype="u32"), # start bit position + OperandSlot("d", dtype="u32"), # field length + ), + ), + # szext / bmsk per PTX ISA 9.7.1.22, 9.7.1.23 -- both `.mode = {.clamp, + # .wrap}`, and both take their width operand as an unsigned 32-bit value. + # szext.mode.type d, a, b; .type = {.u32, .s32} + # bmsk.mode.b32 d, a, b; -- a (position) and b (width) are .u32 + InstructionEntry( + name="szext", + slots=( + ModifierSlot("mode", ("clamp", "wrap")), + ModifierSlot("type", ("u32", "s32")), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b", dtype="u32"), # N, the extended-from width + ), + ), + InstructionEntry( + name="bmsk", + slots=( + ModifierSlot("mode", ("clamp", "wrap")), + ModifierSlot("type", ("b32",)), + ), + operands=( + # `b32i`, not `b32`: bmsk is one of the instructions that refuses a + # float register (ptxas: "Arguments mismatch for instruction + # 'bmsk'", in any of the three positions) while taking either + # integer signedness. Probed against popc, clz, brev, bfi and mov, + # which all accept `.f32` -- so this is bmsk's own rule, not the + # bit-size type's. + OperandSlot("d", rw="w", dtype="b32i"), + OperandSlot("a", dtype="b32i"), # start bit position + OperandSlot("b", dtype="b32i"), # mask width + ), + ), + # dp4a / dp2a per PTX ISA 9.7.1.24, 9.7.1.25: + # dp4a.atype.btype d, a, b, c; .atype = .btype = {.u32, .s32} + # dp2a.mode.atype.btype d, a, b, c; .mode = {.lo, .hi} + # The only families in the table with no `type` slot at all: every operand + # is typed individually. a and b name the two written tokens; c and d take + # the accumulator type, which is a function of both (`_dp_acc_dtype`). + InstructionEntry( + name="dp4a", + slots=( + ModifierSlot("atype", ("u32", "s32")), + ModifierSlot("btype", ("u32", "s32")), + ), + operands=( + OperandSlot("d", rw="w", dtype=_dp_acc_dtype), + OperandSlot("a", dtype="atype"), + OperandSlot("b", dtype="btype"), + OperandSlot("c", dtype=_dp_acc_dtype), + ), + ), + InstructionEntry( + name="dp2a", + slots=( + ModifierSlot("mode", ("lo", "hi")), + ModifierSlot("atype", ("u32", "s32")), + ModifierSlot("btype", ("u32", "s32")), + ), + operands=( + OperandSlot("d", rw="w", dtype=_dp_acc_dtype), + OperandSlot("a", dtype="atype"), + OperandSlot("b", dtype="btype"), + OperandSlot("c", dtype=_dp_acc_dtype), + ), + ), + # ------------------------------------------------------------------ + # PTX ISA 9.7.3 — Floating-Point Instructions + # ------------------------------------------------------------------ + # The section is registered in full. Its one exclusion is a form no + # supported architecture assembles: `mad`'s sm_1x line, which omits the + # rounding modifier (stated at the mad entry below). min/max live in the + # 9.7.1 group above, where their integer lines share the entry. + # + # This group also carries the whole of PTX ISA 9.7.5 (Mixed Precision), + # whose three subsections are not separate instructions but a fourth syntax + # line of three that are already here -- `op{.rnd}{.sat}.f32.atype`, which + # converts a 16-bit source to .f32 before operating. They are the optional + # `srctype` slot on `add`/`sub` (9.7.5.{1,2}) and `fma` (9.7.5.3) below; + # `mul` is the one mnemonic of the four the ISA gives no mixed line, which + # is why its entry declares no such slot. + # + # Floating-point add/sub/mul (PTX ISA 9.7.3.{3,4,5}) together with their + # mixed-precision lines (9.7.5.{1,2}); `mul` has no mixed-precision line. + # add{.rnd}{.ftz}{.sat}.f32 d, a, b; add{.rnd}{.ftz}.f32x2 d, a, b; + # add{.rnd}.f64 d, a, b; add{.rnd}{.sat}.f32.atype d, a, c; + # cert_arch is the family's ceiling, not its floor: .f32/.f64 assemble + # everywhere, but .f32x2 and the mixed lines need sm_100, and certification + # has to run somewhere every legal variant is legal. + *[ + InstructionEntry( + name=name, + cert_arch="sm_100", + slots=( + ModifierSlot("rnd", _FRND, optional=True), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", ("f32", "f64", "f32x2")), + *((ModifierSlot("srctype", ("f16", "bf16"), optional=True),) if mixed else ()), + ), + check=_check_farith, + operands=( + OperandSlot("d", rw="w"), + # On the mixed line `a` is the converted 16-bit source; on every + # other line it is just the instruction type. + OperandSlot("a", dtype="srctype" if mixed else None), + OperandSlot("b"), + ), + ) + # `mul` is the one line with no mixed-precision form (ISA 9.7.5). + for name, mixed in (("add", True), ("sub", True), ("mul", False)) + ], + # NOT REGISTERED: across this whole arithmetic group, only the + # extended-precision lines add.cc/addc/sub.cc/subc (9.7.2.{1,2,3,4}), whose + # carry flag is a piece of state no other instruction here has. Everything + # else is registered: the integer lines (9.7.1.{1,2,3}) as + # `add_int`/`sub_int`/`mul_int`/`mul_wide` above, and the half lines + # (9.7.4.{1,2,3,4}) as `add_half`/`sub_half`/`mul_half`/`fma_half` below. + # + # fma differs in shape, so it is its own entry: three sources, and .rnd is + # mandatory on every line (PTX ISA 9.7.3.6 / 9.7.5.3). + # fma.rnd{.ftz}{.sat}.f32 d, a, b, c; fma.rnd{.ftz}.f32x2 d, a, b, c; + # fma.rnd.f64 d, a, b, c; fma.rnd{.sat}.f32.abtype d, a, b, c; + InstructionEntry( + name="fma", + cert_arch="sm_100", + slots=( + ModifierSlot("rnd", _FRND), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", ("f32", "f64", "f32x2")), + ModifierSlot("srctype", ("f16", "bf16"), optional=True), + ), + check=_check_farith, + operands=( + OperandSlot("d", rw="w"), + # .abtype converts both a and b; c is always the instruction type. + OperandSlot("a", dtype="srctype"), + OperandSlot("b", dtype="srctype"), + OperandSlot("c"), + ), + ), + # mad per PTX ISA 9.7.3.7 -- the same grid as fma above minus `.f32x2` and + # minus the mixed-precision lines, so it shares `_check_farith`: + # mad.rnd{.ftz}{.sat}.f32 d, a, b, c; mad.rnd.f64 d, a, b, c; + # The ISA notes mad.{f32,f64} is the same instruction as fma.{f32,f64}; it + # is registered anyway because it is a mnemonic ptxas accepts and PTX + # written elsewhere uses. + # NOT REGISTERED: the sm_1x line `mad{.ftz}{.sat}.f32` with no `.rnd`. The + # ISA's own Errata says ptxas enforces the rounding modifier from PTX ISA + # 3.2 onward, so the no-rnd form assembles on no architecture this dialect + # targets -- which is why `rnd` is a required slot here and optional on the + # add/sub/mul entries above. + InstructionEntry( + name="mad_f", + mnemonic="mad", + slots=( + ModifierSlot("rnd", _FRND), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", ("f32", "f64")), + ), + check=_check_farith, + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + ), + ), + # div per PTX ISA 9.7.3.8 -- the floating-point lines, beside the integer + # `div` in the 9.7.1 group above. `mode` fuses the two approximations with + # the IEEE rounding modes because the ISA requires exactly one of them + # ("one of .approx, .full, or .rnd is required") and spells them in the + # same position. + InstructionEntry( + name="div_f", + mnemonic="div", + slots=( + ModifierSlot("mode", ("approx", "full", "rn", "rz", "rm", "rp")), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", ("f32", "f64")), + ), + check=_check_div_f, + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + ), + ), + # abs / neg per PTX ISA 9.7.3.9, 9.7.3.10 -- one shape, one check: + # op{.ftz}.f32 d, a; op.f64 d, a; + # `abs_f` carries the suffix because the integer `abs` (9.7.1.11) took the + # bare name first; `neg` keeps it because the integer one came second. + *[ + InstructionEntry( + name=name, + mnemonic=mnemonic, + slots=( + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", ("f32", "f64")), + ), + check=_check_absneg, + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ) + for name, mnemonic in (("abs_f", "abs"), ("neg", None)) + ], + # rcp / sqrt per PTX ISA 9.7.3.13 + 9.7.3.14, and 9.7.3.15. Same shape and + # same slots; they differ only in whether an f64 approximation exists, which + # is what separates their checks (see `_check_rcp` on why 9.7.3.14 is a cell + # of this grid rather than an entry of its own). + *[ + InstructionEntry( + name=name, + slots=( + ModifierSlot("mode", ("approx", "rn", "rz", "rm", "rp")), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", ("f32", "f64")), + ), + check=check, + operands=( + OperandSlot("d", rw="w"), + OperandSlot("value"), + ), + ) + for name, check in (("rcp", _check_rcp), ("sqrt", _check_sqrt)) + ], + # rsqrt per PTX ISA 9.7.3.16 + 9.7.3.17. No check: every cell of this 2x2 + # grid is a syntax line the ISA spells -- `rsqrt.approx{.ftz}.f32` and + # `rsqrt.approx.f64` at 9.7.3.16, and `rsqrt.approx.ftz.f64` at 9.7.3.17. + # `.approx` is mandatory on all of them, hence the single-choice slot. + InstructionEntry( + name="rsqrt", + slots=( + ModifierSlot("mode", ("approx",)), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", ("f32", "f64")), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("value"), + ), + ), + # sin / cos / lg2 / ex2 per PTX ISA 9.7.3.18-9.7.3.21: four mnemonics, one + # syntax line each, all of it `op.approx{.ftz}.f32 d, a;`. ex2's + # half-precision forms (9.7.4.10) are the `ex2_half` entry below: they + # cannot share this optional `.ftz` slot, because theirs is mandatory on + # the bf16 line and unspelled on the f16 one. + *[ + InstructionEntry( + name=name, + slots=( + ModifierSlot("mode", ("approx",)), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", ("f32",)), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("value"), + ), + ) + for name in ("ex2", "sin", "cos", "lg2") + ], + # tanh per PTX ISA 9.7.3.22: `tanh.approx.f32 d, a;` -- alone among the + # approximations in spelling no `.ftz`, so it gets no such slot. Requires + # sm_75, below the certification floor. + InstructionEntry( + name="tanh", + slots=( + ModifierSlot("mode", ("approx",)), + ModifierSlot("type", ("f32",)), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("value"), + ), + ), + # testp per PTX ISA 9.7.3.1: `testp.op.type p, a;`, whose result "is .pred". + # The destination is a predicate register like any other in this table -- + # `render.BRIDGE` materializes it through the selp the C boundary needs. + InstructionEntry( + name="testp", + slots=( + ModifierSlot( + "op", ("finite", "infinite", "number", "notanumber", "normal", "subnormal") + ), + ModifierSlot("type", ("f32", "f64")), + ), + operands=( + OperandSlot("p", rw="w", dtype="pred"), + OperandSlot("a"), + ), + ), + # copysign per PTX ISA 9.7.3.2: `copysign.type d, a, b;` -- the sign bit of + # a onto the value of b. + InstructionEntry( + name="copysign", + slots=(ModifierSlot("type", ("f32", "f64")),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + ), + ), + # ------------------------------------------------------------------ + # PTX ISA 9.7.4 — Half Precision Floating-Point Instructions + # ------------------------------------------------------------------ + # The section is registered in full -- every syntax line of all ten + # subsections, with no exclusions. Every entry here carries `_HALF_TYPES`, + # whose four tokens ride two carriers: `.f16`/`.bf16` a 16-bit register and + # the packed pairs a 32-bit one (see PTX_TYPE_DTYPES), which is why none of + # these entries needs to say anything about operand types. + # + # min/max (9.7.4.7, 9.7.4.8) are the exception to the placement, not to the + # coverage: their half lines live in the merged entry up in the 9.7.1 group. + # + # Half-precision add/sub/mul (PTX ISA 9.7.4.{1,2,3}). Same operand shape as + # the single/double lines above, so they differ only in their slot domains: + # sub{.rnd}{.ftz}{.sat}.f16 / .f16x2 sub{.rnd}.bf16 / .bf16x2 .rnd = {.rn} + *[ + InstructionEntry( + name=f"{name}_half", + mnemonic=name, + slots=( + ModifierSlot("rnd", ("rn",), optional=True), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("type", _HALF_TYPES), + ), + check=_check_half_arith, + # The bf16 lines of all three require sm_90 (ISA Target Notes); + # cert_arch is the max floor over the entry's variants. + cert_arch="sm_90", + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + ), + ) + for name in ("add", "sub", "mul") + ], + # Half-precision fma (PTX ISA 9.7.4.4) -- the add/sub/mul grid plus two + # clamping qualifiers the same-precision lines never carry: + # fma.rnd{.ftz}{.sat}.f16{x2} d,a,b,c; fma.rnd{.ftz}.relu.f16{x2} d,a,b,c; + # fma.rnd{.relu}.bf16{x2} d,a,b,c; fma.rnd.oob{.relu}.type d,a,b,c; + # `.relu` clamps negatives to zero and `.oob` forces the result to +0.0 when + # an operand is the out-of-bounds NaN (see Tensors); `_check_half_fma` holds + # which line offers which. `.rnd` is required, as on the f32/f64 fma. + InstructionEntry( + name="fma_half", + mnemonic="fma", + slots=( + ModifierSlot("rnd", ("rn",)), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("sat", ("sat",), optional=True), + ModifierSlot("oob", ("oob",), optional=True), + ModifierSlot("relu", ("relu",), optional=True), + ModifierSlot("type", _HALF_TYPES), + ), + check=_check_half_fma, + cert_arch="sm_90", # the `.oob` floor, the highest over this entry + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + ), + ), + # Half-precision neg / abs (PTX ISA 9.7.4.5, 9.7.4.6): `op{.ftz}.f16{x2}` + # and `op.bf16{x2}`, the same one-source shape as their f32/f64 lines. + *[ + InstructionEntry( + name=f"{name}_half", + mnemonic=name, + slots=( + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", _HALF_TYPES), + ), + check=_check_half_absneg, + cert_arch="sm_90", # the bf16 lines need sm_80; sm_90 clears them + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ) + for name in ("neg", "abs") + ], + # Half-precision tanh (PTX ISA 9.7.4.9): `tanh.approx.type d, a;` over all + # four types, and -- like the .f32 line of 9.7.3.22 -- no `.ftz` anywhere. + InstructionEntry( + name="tanh_half", + mnemonic="tanh", + slots=( + ModifierSlot("mode", ("approx",)), + ModifierSlot("type", _HALF_TYPES), + ), + cert_arch="sm_90", # tanh.approx.bf16{x2} requires sm_90 + operands=( + OperandSlot("d", rw="w"), + OperandSlot("value"), + ), + ), + # Half-precision ex2 (PTX ISA 9.7.4.10): two lines that differ by more than + # their type token -- `ex2.approx.f16{x2}` has no `.ftz` and + # `ex2.approx.ftz.bf16{x2}` requires it. `_check_half_ex2` enforces both + # directions, which is exactly why these cannot join the `.f32` ex2 entry. + InstructionEntry( + name="ex2_half", + mnemonic="ex2", + slots=( + ModifierSlot("mode", ("approx",)), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", _HALF_TYPES), + ), + check=_check_half_ex2, + cert_arch="sm_90", # ex2.approx.ftz.bf16{x2} requires sm_90 + operands=( + OperandSlot("d", rw="w"), + OperandSlot("value"), + ), + ), + # ------------------------------------------------------------------ + # PTX ISA 9.7.6 — Comparison and Selection Instructions + # ------------------------------------------------------------------ + # The section is registered in full. Two spellings are deliberately absent, + # neither of which is a syntax line: + # + # `{!}c`, the negated predicate source of the BoolOp forms. `c` reaches + # this surface as a uint32 evidenced by `T.ptx.pred(x)` and turned into a + # predicate by a setp bridge, so a caller writes the complement into the + # expression it passes -- `!c` spells no program the plain form cannot. + # + # `_` in place of one of setp's two destinations. The ISA allows it + # ("The sink symbol '_' may be used in place of any one of the + # destination operands"), but the same program is written by passing a + # scratch lvalue for the half that is not wanted, and the engine's + # sink guard is per operand rather than per pipe group -- it would reject + # `p|_` as "every lane sunk" without surgery that buys nothing. + # + # set / setp per PTX ISA 9.7.6.1, 9.7.6.2. Both compare a and b and + # optionally fold a third predicate in with a Boolean operator; they differ + # in where the answer goes -- set writes a value of `.dtype` (0xffffffff or + # 1.0f for true), setp writes predicates. + # set.CmpOp{.ftz}.dtype.stype d, a, b; + # set.CmpOp.BoolOp{.ftz}.dtype.stype d, a, b, {!}c; + # setp.CmpOp{.ftz}.type p[|q], a, b; + # setp.CmpOp.BoolOp{.ftz}.type p[|q], a, b, {!}c; + # The `{.BoolOp}` and `[|q]` brackets are each a separate syntax line rather + # than an optional token: the first adds an operand, the second adds a + # destination. That is two independent shape choices, so setp is four + # entries and set is two, all dispatched apart by arity and operand class. + *[ + InstructionEntry( + name="set" if not boolop else "set_bool", + mnemonic="set", + slots=( + ModifierSlot("cmp", _CMP_OPS), + *((ModifierSlot("boolop", _CMP_BOOL_OPS),) if boolop else ()), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("dtype", ("u32", "s32", "f32")), + ModifierSlot("stype", _CMP_TYPES), + ), + check=_check_cmp, + operands=( + # "Operand d has type .dtype; operands a and b have type + # .stype; operand c has type .pred." + OperandSlot("d", rw="w", dtype="dtype"), + OperandSlot("a", dtype="stype"), + OperandSlot("b", dtype="stype"), + *((OperandSlot("c", dtype="pred"),) if boolop else ()), + ), + ) + for boolop in (False, True) + ], + *[ + InstructionEntry( + name="setp" + ("_bool" if boolop else "") + ("_pq" if pq else ""), + mnemonic="setp", + slots=( + ModifierSlot("cmp", _CMP_OPS), + *((ModifierSlot("boolop", _CMP_BOOL_OPS),) if boolop else ()), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", _CMP_TYPES), + ), + check=_check_cmp, + operands=( + # `p|q` is one operand position holding two predicates: q takes + # the Boolean applied to the *complement* of the compare, so it + # is a second result and not a restatement of p. `pipe` renders + # the pair with the ISA's separator (see OperandSlot.pipe). + OperandSlot("p", rw="w", dtype="pred", pipe="pq" if pq else None), + *((OperandSlot("q", rw="w", dtype="pred", pipe="pq"),) if pq else ()), + OperandSlot("a"), + OperandSlot("b"), + *((OperandSlot("c", dtype="pred"),) if boolop else ()), + ), + ) + for boolop in (False, True) + for pq in (False, True) + ], + # selp per PTX ISA 9.7.6.3: `selp.type d, a, b, c;` -- a if the predicate + # is true, b otherwise. The one selection instruction whose selector is a + # predicate rather than a number. + InstructionEntry( + name="selp", + slots=(ModifierSlot("type", _CMP_TYPES),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c", dtype="pred"), + ), + ), + # slct per PTX ISA 9.7.6.4: `slct{.ftz}.dtype.ctype d, a, b, c;` -- a if + # c >= 0, else b. Two instruction types, because the value being selected + # and the number whose sign selects it are independent: "operands d, a and + # b are treated as a bitsize type of the same width as the first + # instruction type; operand c must match the second". + InstructionEntry( + name="slct", + slots=( + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("dtype", _CMP_TYPES), + ModifierSlot("ctype", ("s32", "f32")), + ), + check=_check_slct, + operands=( + OperandSlot("d", rw="w", dtype="dtype"), + OperandSlot("a", dtype="dtype"), + OperandSlot("b", dtype="dtype"), + OperandSlot("c", dtype="ctype"), + ), + ), + # ------------------------------------------------------------------ + # PTX ISA 9.7.7 — Half Precision Comparison Instructions + # ------------------------------------------------------------------ + # The section is registered in full: `set` (9.7.7.1) and `setp` (9.7.7.2) + # over the half types, beside the 9.7.6 entries of the same two mnemonics. + # They are separate entries because the type grids are disjoint -- no 9.7.6 + # slot holds a half token -- and because this section's `.CmpOp` line is a + # different set (no unsigned alternates; see `_HALF_CMP_OPS`). + # + # `{!}c` is left out here for the reason given at the 9.7.6 banner above. + # Three forms ptxas accepts that no syntax line spells are also left out, + # each noted where it belongs: lo/ls/hi/hs on an integer source + # (`_HALF_CMP_OPS`), `set.bf16.bf16` (`_SET_HALF_STYPES`), and a lone `p` + # on the packed setp lines (at the pq entries below). + # + # set per PTX ISA 9.7.7.1. Its two type tokens are a grid rather than a + # product -- `_check_half_set` holds the six line groups. + *[ + InstructionEntry( + name="set_half_bool" if boolop else "set_half", + mnemonic="set", slots=( - ModifierSlot("rnd", ("rn",), optional=True), + ModifierSlot("cmp", _HALF_CMP_OPS), + *((ModifierSlot("boolop", _CMP_BOOL_OPS),) if boolop else ()), ModifierSlot("ftz", ("ftz",), optional=True), - ModifierSlot("sat", ("sat",), optional=True), - ModifierSlot("type", ("f16", "f16x2", "bf16", "bf16x2")), + ModifierSlot("dtype", ("f16", "bf16", "u16", "s16", "u32", "s32", *_HALF_X2)), + ModifierSlot("stype", (*_SET_HALF_STYPES, "bf16", *_HALF_X2)), ), - check=_check_half_arith, + check=_check_half_set, + cert_arch="sm_90", # every bf16 form; the f16 ones need only sm_53 + operands=( + OperandSlot("d", rw="w", dtype="dtype"), + OperandSlot("a", dtype="stype"), + OperandSlot("b", dtype="stype"), + *((OperandSlot("c", dtype="pred"),) if boolop else ()), + ), + ) + for boolop in (False, True) + ], + # setp per PTX ISA 9.7.7.2. Unlike 9.7.6's setp, which offers both + # destination shapes on every type, here the type *decides* the shape: the + # scalar lines spell one predicate, the packed lines spell `p|q`, and there + # q is the second half's comparison rather than the complement of the first + # (`p = BoolOp(t[0], c); q = BoolOp(t[1], c)`). So the two shapes carry + # disjoint type domains instead of a shared one. + # MEASURED, NOT REGISTERED: ptxas also accepts a lone `p` on the packed + # lines, discarding the upper half's result. The ISA spells only `p|q` + # there, and what the discarded form would compute is unstated, so a caller + # who wants one half writes the pair and ignores q. + *[ + InstructionEntry( + name="setp_half" + ("_bool" if boolop else "") + ("_pq" if pq else ""), + mnemonic="setp", + slots=( + ModifierSlot("cmp", _HALF_CMP_OPS), + *((ModifierSlot("boolop", _CMP_BOOL_OPS),) if boolop else ()), + ModifierSlot("ftz", ("ftz",), optional=True), + ModifierSlot("type", _HALF_X2 if pq else ("f16", "bf16")), + ), + check=_check_half_setp, + cert_arch="sm_90", # the bf16 lines; the f16 ones need only sm_53 + operands=( + OperandSlot("p", rw="w", dtype="pred", pipe="pq" if pq else None), + *((OperandSlot("q", rw="w", dtype="pred", pipe="pq"),) if pq else ()), + OperandSlot("a"), + OperandSlot("b"), + *((OperandSlot("c", dtype="pred"),) if boolop else ()), + ), + ) + for boolop in (False, True) + for pq in (False, True) + ], + # ------------------------------------------------------------------ + # PTX ISA 9.7.8 — Logic and Shift Instructions + # ------------------------------------------------------------------ + # The section is registered in full. It is the one group that needs no + # check function anywhere: "fundamentally untyped, performing bit-wise + # operations on operands of any type, provided the operands are of the same + # size", so every entry's slot domain is already exactly its syntax line. + # + # Three of the mnemonics -- `and`, `or`, `not` -- are Python keywords, so + # the surface spells them `T.ptx.and_`, `T.ptx.or_`, `T.ptx.not_` and + # `escape_token`/`unescape_token` carry them across the boundary. The PTX + # text is unaffected; only the attribute a program types is. + # + # and / or / xor (PTX ISA 9.7.8.1-9.7.8.3) and not (9.7.8.4): `op.type d, + # a, b;` over `.pred` and the three bit widths. The predicate line is a + # real one here -- "Allowed types include predicate registers" -- so with + # `.pred` the sources want `T.ptx.pred(x)` (or a bool expression) and the + # destination a uint32 lvalue, and the helper carries the setp/selp bridges + # around the single instruction. + *[ + InstructionEntry( + name=name, + slots=(ModifierSlot("type", _LOGIC_TYPES),), operands=( OperandSlot("d", rw="w"), OperandSlot("a"), OperandSlot("b"), ), ) - for name in ("add", "sub", "mul") + for name in ("and", "or", "xor") + ], + InstructionEntry( + name="not", + slots=(ModifierSlot("type", _LOGIC_TYPES),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ), + # cnot per PTX ISA 9.7.8.5: `d = (a==0) ? 1 : 0`, C's `!` rather than the + # bitwise `not` above. Its type line stops at the bit widths -- ptxas + # answers "Unexpected instruction types specified for 'cnot'" to `.pred`, + # which is why this entry cannot share the domain of the four above. + InstructionEntry( + name="cnot", + slots=(ModifierSlot("type", _BIT_TYPES),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + ), + # lop3 per PTX ISA 9.7.8.6: any of the 256 three-input logical functions, + # named by a look-up table byte rather than by an opcode. + # lop3.b32 d, a, b, c, immLut; + # lop3.BoolOp.b32 d|p, a, b, c, immLut, q; + # `immLut` is an OPEN immediate: it lives in the instruction text, and the + # ISA gives it a range (0..255) rather than a list of meaningful values, so + # enumerating 256 helpers per dtype combination would be certifying an + # arithmetic identity 256 times. Certification samples it, exactly as it + # does for tcgen05's immHalfSplitoff. + InstructionEntry( + name="lop3", + slots=(ModifierSlot("type", ("b32",)),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + OperandSlot("immLut", kind="imm"), + ), + ), + # The BoolOp form folds a predicate in after the fact: `p = (d != 0) + # BoolOp q`, so it writes a bit-size result AND a predicate. `d|p` is the + # pipe pair again (see OperandSlot.pipe), here holding two *different* + # register classes -- the mechanism only groups the text, so each half + # keeps its own constraint and its own bridge. + # NOT REGISTERED: `.xor` (ptxas: "Illegal operation '.xor' for instruction + # 'lop3'" -- the ISA's `.BoolOp` line stops at `.or`/`.and`), and `_` in + # place of `d`, which the ISA allows here but which the engine's per-slot + # sink guard would read as "every lane sunk"; pass a scratch lvalue. + InstructionEntry( + name="lop3_bool", + mnemonic="lop3", + slots=( + ModifierSlot("boolop", ("or", "and")), + ModifierSlot("type", ("b32",)), + ), + operands=( + OperandSlot("d", rw="w", pipe="dp"), + OperandSlot("p", rw="w", dtype="pred", pipe="dp"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + OperandSlot("immLut", kind="imm"), + OperandSlot("q", dtype="pred"), + ), + ), + # shf per PTX ISA 9.7.8.7: the funnel shift. a and b are the low and high + # halves of one 64-bit source ("Operand b holds bits 63:32 and operand a + # holds bits 31:0"), and the direction picks which 32 bits of the shifted + # result land in d. `.mode` bounds the shift amount: 0..32 clamping, or + # 0..31 wrapping. + InstructionEntry( + name="shf", + slots=( + ModifierSlot("dir", ("l", "r")), + ModifierSlot("mode", ("clamp", "wrap")), + ModifierSlot("type", ("b32",)), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c", dtype="u32"), # the shift amount + ), + ), + # shl / shr per PTX ISA 9.7.8.8, 9.7.8.9. They differ in their type line, + # not their shape: shl is untyped because a left shift zero-fills whatever + # the operand means, while shr has to know whether to fill with the sign + # bit, so it carries the signed and unsigned lines too (and keeps the + # bit-size ones "for symmetry with shl"). + # In both, "the b operand must be a 32-bit value, regardless of the + # instruction type" -- hence the fixed .u32 on the shift amount. + *[ + InstructionEntry( + name=name, + slots=(ModifierSlot("type", types),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b", dtype="u32"), # the shift amount + ), + ) + for name, types in ( + ("shl", _BIT_TYPES), + ("shr", (*_BIT_TYPES, "u16", "u32", "u64", "s16", "s32", "s64")), + ) ], # ------------------------------------------------------------------ # PTX ISA 9.7.9 — Data Movement and Conversion Instructions # ------------------------------------------------------------------ + # The largest section of the ISA, and registered across all 27 of its + # subsections. What is left out is listed at the entry it belongs to; + # the exclusions that are about the toolchain or the architecture rather + # than about this table are: + # - `shfl` without `.sync` (9.7.9.5), removed by the ISA for sm_70+; + # - `multimem.st.async` (9.7.9.13), PTX ISA 9.3, which ptxas 13.2 + # (ISA 9.2) cannot assemble; + # - the FP8 multimem types and `.acc::f16` (9.7.9.15), and + # `tensormap.replace.swizzle_atomicity` (9.7.9.27), all scoped to + # architectures above the one their instruction is certified at. + # Everything else absent is an operand the helper ABI cannot carry -- a + # symbol rather than a value -- and says so where it is excluded. + # # mov, vector pack/unpack form (PTX ISA 9.7.9.4) # # mov.type d, a; .type = {.b16, .b32, .b64, .b128} @@ -2390,20 +3879,538 @@ def _check_tcgen05_mma_block_scale(m): ModifierSlot("evict", ("L2::evict_last", "L2::evict_normal"), optional=True), ModifierSlot("tensormap", ("tensormap",), optional=True), ), - check=_check_prefetch, + check=_check_prefetch, + operands=(OperandSlot("addr", kind="addr"),), + ), + # 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 + # release store to global memory. + # st.async{.weak}{.ss}.completion_mechanism{.vec}.type [a], b, [mbar]; + # st.async{.mmio}.sem.scope{.ss}.type [a], b; + # The first is sm_90; the second needs sm_100 ("Feature 'st.async/red.async + # with .global state space' requires .target sm_100"). + *[ + InstructionEntry( + name="st_async_vec" if vec else "st_async", + mnemonic="st.async", + slots=( + ModifierSlot("weak", ("weak",), optional=True), + ModifierSlot("space", ("shared::cluster",), optional=True), + ModifierSlot("completion", ("mbarrier::complete_tx::bytes",)), + *((ModifierSlot("vec", ("v2", "v4")),) if vec else ()), + ModifierSlot("type", _ST_ASYNC_V2 if vec else _ST_ASYNC_TYPES), + ), + check=_check_st_async if vec else None, + cert_arch="sm_90", + operands=( + OperandSlot("addr", kind="addr"), + 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"), + ), + ) + for vec in (False, True) + ], + InstructionEntry( + name="st_async_release", + mnemonic="st.async", + slots=( + ModifierSlot("mmio", ("mmio",), optional=True), + ModifierSlot("sem", ("release",)), + ModifierSlot("scope", ("gpu", "sys")), + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("type", _ST_ASYNC_REL_TYPES), + ), + check=_check_st_async_rel, + cert_arch="sm_100", + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("b"), + ), + ), + # multimem per PTX ISA 9.7.9.15: operations on a multimem address, which + # names one location on each of several GPUs at once. Three mnemonics -- + # ld_reduce reduces across the copies into a register, st writes all of + # them, red reduces into all of them -- and each splits into an integer and + # a floating-point entry, because only the floating lines carry `.vec` and + # `.acc::f32`, and their op x type tables are different. + # NOT REGISTERED: the FP8 types (.e4m3*/.e5m2*) and `.acc::f16`, whose + # Target Notes scope them to the sm_100a/sm_120a family architectures + # rather than to sm_90 where the rest of this subsection lives; and + # `multimem.st.async` (9.7.9.13), which is PTX ISA 9.3 and does not + # assemble on this toolchain. + InstructionEntry( + name="multimem_ld_reduce", + mnemonic="multimem.ld_reduce", + slots=( + ModifierSlot("sem", ("weak", "relaxed", "acquire"), optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("op", tuple(_MM_INT_OPS)), + ModifierSlot("type", _MM_INT_TYPES), + ), + check=_check_multimem_int, + cert_arch="sm_90", + operands=( + OperandSlot("d", rw="w"), + OperandSlot("addr", kind="addr"), + ), + ), + InstructionEntry( + name="multimem_st", + mnemonic="multimem.st", + slots=( + ModifierSlot("sem", ("weak", "relaxed", "release"), optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("type", _MM_INT_TYPES), + ), + check=_check_multimem_int, + cert_arch="sm_90", + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("b"), + ), + ), + InstructionEntry( + name="multimem_red", + mnemonic="multimem.red", + slots=( + # No `.weak` here: ptxas answers "Illegal modifier '.weak' for + # instruction 'multimem.red'", and the ISA gives red a `.redsem` + # line without it -- the same pair the scalar `red` takes. + ModifierSlot("sem", _RED_SEM, optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("op", tuple(_MM_INT_OPS)), + ModifierSlot("type", _MM_INT_TYPES), + ), + check=_check_multimem_int, + cert_arch="sm_90", + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("b"), + ), + ), + # The floating-point multimem lines. Each is two entries rather than one + # with an optional `.vec`, because the qualifier changes the operand's + # *shape*: written, the value is a brace-enclosed register group; omitted, + # it is a plain register. That is the table's rule for where one entry + # ends, and it is also what ptxas insists on ("Vector is not expected for + # argument 1"). + *[ + InstructionEntry( + name=name + ("_vec" if vec else ""), + mnemonic=f"multimem.{mnem}", + slots=( + ModifierSlot("sem", sems, optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", ("global",), optional=True), + *((ModifierSlot("op", ops),) if ops else ()), + *((ModifierSlot("acc", ("acc::f32",), optional=True),) if acc else ()), + *((ModifierSlot("vec", ("v2", "v4", "v8")),) if vec else ()), + ModifierSlot("type", _MM_FLOAT_TYPES), + ), + check=_check_multimem_f, + cert_arch="sm_90", + operands=( + # ld_reduce reads the multimem address into a register; st and + # red send a value the other way. + *( + ( + OperandSlot("d", rw="w", lanes=_vec_lanes if vec else 1), + OperandSlot("addr", kind="addr"), + ) + if mnem == "ld_reduce" + else ( + OperandSlot("addr", kind="addr"), + OperandSlot("b", lanes=_vec_lanes if vec else 1), + ) + ), + ), + ) + # `multimem.st` names no operation (it just writes), and `multimem.red` + # on floats offers only `.add` -- the ISA gives it a `.redop` line of + # its own for exactly that reason. `.acc::f32` is a load-side + # qualifier, so only ld_reduce carries it. + for name, mnem, sems, ops, acc in ( + ( + "multimem_ld_reduce_f", + "ld_reduce", + ("weak", "relaxed", "acquire"), + ("min", "max", "add"), + True, + ), + ("multimem_st_f", "st", ("weak", "relaxed", "release"), (), False), + ("multimem_red_f", "red", ("relaxed", "release"), ("add",), False), + ) + for vec in (False, True) + ], + # createpolicy per PTX ISA 9.7.9.19: build the opaque 64-bit cache-eviction + # policy that the `.level::cache_hint` operand of ld/st/cp takes. Three + # syntax lines, three shapes: + # createpolicy.range{.global}.pri{.sec}.b64 policy, [a], primary, total; + # createpolicy.fractional.pri{.sec}.b64 policy{, fraction}; + # createpolicy.cvt.L2.b64 policy, access-property; + # The fractional line's `{, fraction}` is an optional *operand*, so it is + # two entries split by arity, as elsewhere in this table. + *[ + InstructionEntry( + name=name, + mnemonic="createpolicy", + slots=( + ModifierSlot("kind", (kind,)), + *((ModifierSlot("space", ("global",), optional=True),) if kind == "range" else ()), + ModifierSlot("pri", _CACHE_PRIORITIES), + ModifierSlot("sec", ("L2::evict_first", "L2::evict_unchanged"), optional=True), + ModifierSlot("type", ("b64",)), + ), + cert_arch="sm_90", # the subsection's floor is sm_80 + operands=( + OperandSlot("cache_policy", rw="w", dtype="u64"), + *ops, + ), + ) + for name, kind, ops in ( + ( + "createpolicy_range", + "range", + ( + OperandSlot("addr", kind="addr"), + OperandSlot("primary_size", dtype="u32"), + OperandSlot("total_size", dtype="u32"), + ), + ), + ("createpolicy_fractional", "fractional", ()), + ("createpolicy_fraction", "fractional", (OperandSlot("fraction", dtype="f32"),)), + ) + ], + # The third line converts a CUDA-API access property instead of describing + # one, so it carries neither priority slot. + InstructionEntry( + name="createpolicy_cvt", + mnemonic="createpolicy", + slots=( + ModifierSlot("kind", ("cvt",)), + ModifierSlot("level", ("L2",)), + ModifierSlot("type", ("b64",)), + ), + cert_arch="sm_90", + operands=( + OperandSlot("cache_policy", rw="w", dtype="u64"), + OperandSlot("access_property", dtype="u64"), + ), + ), + # cp.async.bulk.prefetch per PTX ISA 9.7.9.26.4.3 -- the non-tensor bulk + # prefetch, beside the tensor one further down: + # cp.async.bulk.prefetch.L2.src{.level::cache_hint} [srcMem], size{, policy}; + InstructionEntry( + name="cp_async_bulk_prefetch", + mnemonic="cp", + slots=( + ModifierSlot("api", ("async",)), + ModifierSlot("kind", ("bulk",)), + ModifierSlot("op", ("prefetch",)), + ModifierSlot("level", ("L2",)), + ModifierSlot("src", ("global",)), + ModifierSlot("cache", ("L2::cache_hint",), optional=True), + ), + cert_arch="sm_90", + operands=( + OperandSlot("src_mem", kind="addr", space="global"), + OperandSlot("size", dtype="u32"), + OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False), + ), + ), + # tensormap.replace per PTX ISA 9.7.9.27: overwrite one field of a 1024-bit + # tensor-map object in place. Three field groups, and they differ in shape + # rather than in qualifier, so they are three entries: + # .field1 = {global_address, rank} [addr], new_val + # .field2 = {box_dim, global_dim, global_stride, ...} [addr], ord, new_val + # .field3 = {elemtype, interleave_layout, ...} [addr], + # For field3 "the operand new_val must be an immediate" -- ptxas agrees + # ("Integer constant expression expected"), and each field has its own + # closed value table, so they are `imm` operands with per-field choices. + # `ord` is likewise an immediate the ISA bounds to [0..4], which ptxas + # enforces by name. + # NOT REGISTERED: `.swizzle_atomicity` (sm_100a and later only, while the + # rest of this instruction is sm_90a), and the `.elemtype` values 13-15 and + # `.swizzle_mode` value 4, which need newer architectures still. + *[ + InstructionEntry( + name=f"tensormap_replace_{group}", + mnemonic="tensormap.replace", + slots=( + ModifierSlot("mode", ("tile",)), + ModifierSlot("field", fields), + ModifierSlot("space", ("global", "shared::cta"), optional=True), + ModifierSlot("width", ("b1024",)), + ModifierSlot("type", (ty,)), + ), + cert_arch="sm_90a", + operands=( + OperandSlot("addr", kind="addr"), + *ops, + ), + ) + for group, fields, ty, ops in ( + ("address", ("global_address",), "b64", (OperandSlot("new_val", dtype="b64"),)), + ("rank", ("rank",), "b32", (OperandSlot("new_val", dtype="b32"),)), + ( + "stride", + ("global_stride",), + "b64", + ( + OperandSlot("ord", kind="imm", choices=("0", "1", "2", "3", "4")), + OperandSlot("new_val", dtype="b64"), + ), + ), + ( + "dim", + ("box_dim", "global_dim", "element_stride"), + "b32", + ( + OperandSlot("ord", kind="imm", choices=("0", "1", "2", "3", "4")), + # `b32i`: measured at sm_90a over all four 32-bit classes + # -- b32, u32 and s32 assemble here and `.f32` does not, + # while `rank` and the 64-bit fields take all four and so + # stay on the full axis. + OperandSlot("new_val", dtype="b32i"), + ), + ), + ) + ], + # The field3 group, one entry per field because each names its own closed + # set of immediate values (ISA Table 33). + *[ + InstructionEntry( + name=f"tensormap_replace_{field}", + mnemonic="tensormap.replace", + slots=( + ModifierSlot("mode", ("tile",)), + ModifierSlot("field", (field,)), + ModifierSlot("space", ("global", "shared::cta"), optional=True), + ModifierSlot("width", ("b1024",)), + ModifierSlot("type", ("b32",)), + ), + cert_arch="sm_90a", + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("new_val", kind="imm", choices=values), + ), + ) + for field, values in ( + ("elemtype", tuple(str(i) for i in range(13))), # 13-15 need newer archs + ("interleave_layout", ("0", "1", "2")), + ("swizzle_mode", ("0", "1", "2", "3")), # 4 needs sm_103a + ("fill_mode", ("0", "1")), + ) + ], + # mov, scalar form (PTX ISA 9.7.9.3): `mov.type d, a;`. A different + # instruction from the vector pack/unpack lines above that share the + # mnemonic -- this one just copies a register -- and dispatch tells them + # apart by arity, since every pack/unpack shape takes at least three. + # NOT REGISTERED: the other four source forms of the syntax block (`sreg`, + # a variable's address, `avar+imm`, a function name). Each names a *symbol* + # rather than a value, and a symbol cannot cross the helper-function ABI -- + # the same reason `ld` leaves out `.unified` and the `.param` spaces. + InstructionEntry( + name="mov", + slots=(ModifierSlot("type", _MOV_TYPES),), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + ), + asm_volatile=False, # a register copy: let nvcc common it up + ), + # prmt per PTX ISA 9.7.9.7: pick four arbitrary bytes out of the eight in + # {b, a} and reassemble them. Without `.mode`, operand c is four 4-bit + # selectors; with one, its two low bits pick a row of the mode's table. + # The mode comes *after* the type in the text (`prmt.b32{.mode}`), which is + # why the slots are in that order. + InstructionEntry( + name="prmt", + slots=( + ModifierSlot("type", ("b32",)), + ModifierSlot("mode", ("f4e", "b4e", "rc8", "ecl", "ecr", "rc16"), optional=True), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("b"), + OperandSlot("c"), + ), + ), + # ldu per PTX ISA 9.7.9.10: a load whose address is uniform across the + # warp. Two lines, scalar and vector, split the way ld's are. + # Only `.global` or generic addressing: ptxas answers "State space + # incorrect for instruction 'ldu'" to anything else, matching the ISA's + # one-element `.ss` list. + InstructionEntry( + name="ldu", + slots=( + ModifierSlot("space", ("global",), optional=True), # omitted = generic + ModifierSlot("type", _LD_TYPES), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("addr", kind="addr"), + ), + ), + InstructionEntry( + name="ldu_vec", + mnemonic="ldu", + slots=( + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("vec", ("v2", "v4")), + ModifierSlot("type", _LD_VEC_TYPES), + ), + # The same 128-bit ceiling ld's vector lines have (ISA 5.4.2), and + # measured the same way: ptxas answers "Vector type too large, exceeds + # 128 bit limit" to `.v4` with a 64-bit type. ldu has no 256-bit line + # to send those to, so here the check is the whole rule. + check=_check_vec128, + operands=( + OperandSlot("d", rw="w", lanes=_vec_lanes), + OperandSlot("addr", kind="addr"), + ), + ), + # prefetchu per PTX ISA 9.7.9.16, the fourth line of that subsection: the + # uniform cache rather than the data cache, so a different mnemonic and no + # state space (it "requires a generic address"). + InstructionEntry( + name="prefetchu", + slots=(ModifierSlot("level", ("L1",)),), operands=(OperandSlot("addr", kind="addr"),), ), - # cvta per PTX ISA 9.7.9.21. This entry exists to serve the engine's - # shared-address coercion, so it registers exactly the one combination that - # needs: 1 of the 32 legal (direction x space x size) forms. - # NOT REGISTERED: the whole space->generic direction, seven of the eight - # state spaces, and `.u32` -- the last genuinely unusable, since ptxas - # rejects the 32-bit ABI on sm_90 and higher. + # 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 + # other saying it need not be written back at all. + # applypriority{.global}.L2::evict_normal [a], size; + # discard{.global}.L2 [a], size; + # `size` is fixed at 128 by both ISA lines ("The only supported value for + # the size operand is 128"), so it is a table-owned literal: it takes no + # call argument and cannot be got wrong. + *[ + InstructionEntry( + name=name, + slots=( + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("level", (level,)), + ), + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("size", kind="imm", literal="128"), + ), + ) + for name, level in (("applypriority", "L2::evict_normal"), ("discard", "L2")) + ], + # isspacep per PTX ISA 9.7.9.20: does this generic address fall inside the + # window of a given state space? A `.pred` result, so it rides the bridge. + # The address is a plain pointer rather than an `addr` operand: the + # instruction reads the value, it does not dereference it, so there is + # nothing to bracket and no memory to clobber. + InstructionEntry( + name="isspacep", + slots=(ModifierSlot("space", _CVTA_SPACES),), + operands=( + OperandSlot("p", rw="w", dtype="pred"), + OperandSlot("a", kind="ptr"), + ), + asm_volatile=False, # a pure query of an address value + ), + # getctarank per PTX ISA 9.7.9.25: which CTA of the cluster owns this + # address. Two lines by whether the address is a shared one or a generic + # one; `.type` describes the *source*, while "destination d is always a + # 32-bit register". + *[ + InstructionEntry( + name=name, + mnemonic="getctarank", + slots=( + *((ModifierSlot("space", ("shared::cluster",)),) if shared else ()), + ModifierSlot("type", ("u32", "u64")), + ), + cert_arch="sm_90", + operands=( + OperandSlot("d", rw="w", dtype="u32"), + OperandSlot("a"), + ), + asm_volatile=False, + ) + for name, shared in (("getctarank", True), ("getctarank_generic", False)) + ], + # cvt.pack per PTX ISA 9.7.9.23: convert two .s32 values with saturation + # and pack them into one register. Two syntax lines, and the operand count + # is what separates them -- the sub-byte line needs `c` to supply the bits + # the pair does not fill, and the 16-bit line has none left over. + # cvt.pack.sat.convertType.abType d, a, b; 16-bit halves + # cvt.pack.sat.convertType.abType.cType d, a, b, c; sub-byte + *[ + InstructionEntry( + name=name, + mnemonic="cvt.pack", + slots=( + ModifierSlot("sat", ("sat",)), + ModifierSlot("convert", convert), + ModifierSlot("abtype", ("s32",)), + *((ModifierSlot("ctype", ("b32",)),) if has_c else ()), + ), + # `.u4`/`.s4` and `.u2`/`.s2` need sm_75, the rest sm_72. + cert_arch="sm_90", + operands=( + OperandSlot("d", rw="w", dtype="u32"), + OperandSlot("a", dtype="s32"), + OperandSlot("b", dtype="s32"), + *((OperandSlot("c", dtype="b32"),) if has_c else ()), + ), + ) + for name, convert, has_c in ( + ("cvt_pack", ("u16", "s16"), False), + ("cvt_pack_c", ("u2", "s2", "u4", "s4", "u8", "s8"), True), + ) + ], + # shfl.sync per PTX ISA 9.7.9.6: exchange a register between the lanes of a + # warp. `d[|p]` is optional in the syntax and adds a destination, so the + # two shapes are two entries; p reports whether the computed source lane + # was in range. + # NOT REGISTERED: the non-sync `shfl` of 9.7.9.5, which the ISA removed for + # sm_70 and higher in PTX ISA 6.4 -- ptxas agrees ("Instruction 'shfl' + # without '.sync' is not supported on .target sm_70 and higher"), so it + # assembles on no architecture this dialect targets. + *[ + InstructionEntry( + name="shfl_sync_p" if pq else "shfl_sync", + mnemonic="shfl.sync", + slots=( + ModifierSlot("mode", ("up", "down", "bfly", "idx")), + ModifierSlot("type", ("b32",)), + ), + operands=( + OperandSlot("d", rw="w", pipe="dp" if pq else None), + *((OperandSlot("p", rw="w", dtype="pred", pipe="dp"),) if pq else ()), + OperandSlot("a"), + OperandSlot("b", dtype="u32"), # source lane or offset + OperandSlot("c", dtype="u32"), # packed segmask and clamp + OperandSlot("membermask", dtype="u32"), + ), + ) + for pq in (False, True) + ], + # cvta per PTX ISA 9.7.9.21, both directions over all eight state spaces. + # The `.to` direction converts a generic address into a space-specific one + # and the bare direction does the reverse; they are two entries because the + # bare one takes a plain value where `.to` takes a pointer. + # NOT REGISTERED: `.u32`. The 32-bit ABI is unusable at sm_90 and higher -- + # ptxas rejects the program with "Program uses 32-bit address" -- and the + # ISA's own advice for narrowing is to follow cvta with a cvt. InstructionEntry( name="cvta", slots=( ModifierSlot("dir", ("to",)), - ModifierSlot("space", ("shared",)), + ModifierSlot("space", _CVTA_SPACES), ModifierSlot("type", ("u64",)), ), operands=( @@ -2412,6 +4419,21 @@ def _check_tcgen05_mma_block_scale(m): ), asm_volatile=False, # legacy cvta carried no barrier ), + InstructionEntry( + name="cvta_generic", + mnemonic="cvta", + slots=( + ModifierSlot("space", _CVTA_SPACES), + ModifierSlot("type", ("u64",)), + ), + operands=( + OperandSlot("d", rw="w"), + # The source is an address *within* its state space, which is an + # ordinary 64-bit value rather than a pointer the C side owns. + OperandSlot("a", dtype="u64"), + ), + asm_volatile=False, + ), # cvt per PTX ISA 9.7.9.22. # # The generic scalar line first: `cvt{.irnd|.frnd}{.ftz}{.sat}.dtype.atype` @@ -3513,8 +5535,7 @@ def _check_tcgen05_mma_block_scale(m): # The optional thread count is its own syntax line, so `sync` is two # entries sharing a mnemonic, told apart by arity (as `mov`'s shapes are; # `pred` is keyword-only, so arity is unambiguous). - # NOT REGISTERED: the `.red.popc` / `.red.op` lines, which take a `{!}c` - # predicate source and produce a `.pred` result. + # The `.red` lines are registered below, after the plain ones. *[ InstructionEntry( # bar{.cta}.sync a; / barrier{.cta}.sync{.aligned} a; name=f"{mnem}_sync", @@ -3633,6 +5654,174 @@ def _check_tcgen05_mma_block_scale(m): orders_memory=True, operands=(OperandSlot("membermask", dtype="u32"),), ), + # The reduction forms of bar / barrier (PTX ISA 9.7.14.1), which combine a + # predicate across the barrier's threads instead of only waiting: + # bar{.cta}.red.popc.u32 d, a{, b}, {!}c; + # bar{.cta}.red.op.pred p, a{, b}, {!}c; .op = {.and, .or} + # barrier{.cta}.red...{.aligned} likewise + # Four shapes per mnemonic: the count operand `b` is a separate syntax + # line, and `popc` returns how many predicates were true while `.and`/`.or` + # return a predicate. `.xor` is not offered on either -- ptxas: "Illegal + # reduction operation for instruction 'bar.red.pred'". + # NOT REGISTERED: the `{!}c` negation, per the policy stated at `setp`. + *[ + InstructionEntry( + name=f"{mnem}_red_{kind}" + ("_count" if count else ""), + mnemonic=mnem, + slots=( + ModifierSlot("cta", ("cta",), optional=True), + ModifierSlot("action", ("red",)), + ModifierSlot("op", ("popc",) if kind == "popc" else ("and", "or")), + *( + (ModifierSlot("aligned", ("aligned",), optional=True),) + if mnem == "barrier" + else () + ), + ModifierSlot("type", ("u32",) if kind == "popc" else ("pred",)), + ), + orders_memory=True, + operands=( + OperandSlot("d", rw="w", dtype="u32" if kind == "popc" else "pred"), + OperandSlot("a", dtype="u32"), # barrier id + *((OperandSlot("b", dtype="u32"),) if count else ()), # thread count + OperandSlot("c", dtype="pred"), # the predicate being reduced + ), + ) + for mnem in ("bar", "barrier") + for kind in ("popc", "pred") + for count in (False, True) + ], + # vote.sync per PTX ISA 9.7.14.10: reduce a predicate across a warp. The + # `.ballot` line returns one bit per lane instead of a single answer, so it + # is a second entry rather than a fourth `.mode` token. + # NOT REGISTERED: `vote` without `.sync` (9.7.14.9), deprecated in PTX ISA + # 6.0 and rejected outright by ptxas at sm_70 and higher, exactly as the + # non-sync `shfl` is; and the `{!}a` negation, per the usual policy. + *[ + InstructionEntry( + name=name, + mnemonic="vote.sync", + slots=( + ModifierSlot("mode", modes), + ModifierSlot("type", (ty,)), + ), + operands=( + OperandSlot("d", rw="w", dtype="pred" if ty == "pred" else None), + OperandSlot("a", dtype="pred"), + OperandSlot("membermask", dtype="u32"), + ), + ) + for name, modes, ty in ( + ("vote_sync", ("all", "any", "uni"), "pred"), + ("vote_sync_ballot", ("ballot",), "b32"), + ) + ], + # match.sync per PTX ISA 9.7.14.11: which lanes of the warp hold the same + # value. `.all` optionally reports, through a second predicate, whether + # every lane agreed -- `.any` has no such answer to give, and ptxas says so + # ("Predicate output not allowed for instruction 'match.any'"), so the + # pipe form belongs to `.all` alone. + # The destination is a lane mask, so it is `.b32` whatever the type of the + # value being compared. + *[ + InstructionEntry( + name=name, + # The mnemonic is `match`: the ISA spells `match.any.sync.b32`, so + # `.sync` follows the mode rather than preceding it. (ptxas takes + # both orders, but the syntax line is what the table registers.) + mnemonic="match", + slots=( + ModifierSlot("mode", (mode,)), + ModifierSlot("sync", ("sync",)), + ModifierSlot("type", ("b32", "b64")), + ), + cert_arch="sm_90", # the subsection's floor is sm_70 + operands=( + # The lane mask is `b32i`: a float register is rejected here + # (measured), while the value being compared -- `a`, on the + # entry's own type slot -- takes one happily. + OperandSlot("d", rw="w", dtype="b32i", pipe="dp" if pq else None), + *((OperandSlot("p", rw="w", dtype="pred", pipe="dp"),) if pq else ()), + OperandSlot("a"), + OperandSlot("membermask", dtype="u32"), + ), + ) + for name, mode, pq in ( + ("match_any_sync", "any", False), + ("match_all_sync", "all", False), + ("match_all_sync_p", "all", True), + ) + ], + # activemask per PTX ISA 9.7.14.12: the one instruction in the table that + # writes a destination and reads nothing at all. + InstructionEntry( + name="activemask", + slots=(ModifierSlot("type", ("b32",)),), + operands=(OperandSlot("d", rw="w"),), + ), + # elect.sync per PTX ISA 9.7.14.15: pick one leader lane out of the member + # mask. Its `d|p` is not optional the way match's is -- ptxas answers + # "Predicate output expected for instruction 'elect'" to a bare + # destination -- so this family has exactly one shape, and it is the pipe. + InstructionEntry( + name="elect_sync", + mnemonic="elect.sync", + slots=(), + cert_arch="sm_90", + operands=( + # `b32i`: ptxas answers "Arguments mismatch for instruction + # 'elect'" to a float register here, though `.b32` would admit one. + OperandSlot("d", rw="w", dtype="b32i", pipe="dp"), + OperandSlot("p", rw="w", dtype="pred", pipe="dp"), + OperandSlot("membermask", dtype="u32"), + ), + ), + # redux.sync per PTX ISA 9.7.14.13: reduce a value across the warp. + *[ + InstructionEntry( + name=name, + mnemonic="redux.sync", + slots=( + ModifierSlot("op", ops), + ModifierSlot("type", types), + ), + cert_arch="sm_90", # the integer lines need sm_80 + operands=( + OperandSlot("d", rw="w", dtype=dt), + OperandSlot("a", dtype=dt), + OperandSlot("membermask", dtype="u32"), + ), + ) + # The ISA writes the arithmetic and bitwise reductions as two syntax + # lines, and they really are two: the arithmetic one carries the + # signedness in its type, while the bitwise one is untyped. They also + # differ in register class -- ptxas takes no float register on the + # `.b32` line in either position, so that line's operands say `b32i`. + for name, ops, types, dt in ( + ("redux_sync", ("add", "min", "max"), ("u32", "s32"), None), + ("redux_sync_bitwise", ("and", "or", "xor"), ("b32",), "b32i"), + ) + ], + # The `.f32` line of the same subsection, which adds `{.abs}{.NaN}` and is + # scoped to one architecture family: ptxas answers "Instruction + # 'redux.f32' not supported on .target 'sm_100'" and takes it only at + # sm_100a, so it is certified there rather than beside its integer sibling. + InstructionEntry( + name="redux_sync_f32", + mnemonic="redux.sync", + slots=( + ModifierSlot("op", ("min", "max")), + ModifierSlot("abs", ("abs",), optional=True), + ModifierSlot("nan", ("NaN",), optional=True), + ModifierSlot("type", ("f32",)), + ), + cert_arch="sm_100a", + operands=( + OperandSlot("d", rw="w"), + OperandSlot("a"), + OperandSlot("membermask", dtype="u32"), + ), + ), # fence / membar per PTX ISA 9.7.14.4, griddepcontrol per 9.7.14.14. # # These name no address yet constrain everyone else's memory order, so they @@ -3647,7 +5836,7 @@ def _check_tcgen05_mma_block_scale(m): name="fence", slots=( ModifierSlot("sem", ("sc", "acq_rel", "acquire", "release"), optional=True), - ModifierSlot("scope", ("cta", "cluster", "gpu", "sys")), + ModifierSlot("scope", _ATOM_SCOPES), ), orders_memory=True, operands=(), @@ -3682,7 +5871,7 @@ def _check_tcgen05_mma_block_scale(m): ModifierSlot("proxy", ("proxy",)), ModifierSlot("proxykind", ("tensormap::generic",)), ModifierSlot("sem", ("release",)), - ModifierSlot("scope", ("cta", "cluster", "gpu", "sys")), + ModifierSlot("scope", _ATOM_SCOPES), ), orders_memory=True, operands=(), @@ -3694,7 +5883,7 @@ def _check_tcgen05_mma_block_scale(m): ModifierSlot("proxy", ("proxy",)), ModifierSlot("proxykind", ("tensormap::generic",)), ModifierSlot("sem", ("acquire",)), - ModifierSlot("scope", ("cta", "cluster", "gpu", "sys")), + ModifierSlot("scope", _ATOM_SCOPES), ), orders_memory=True, operands=( @@ -3725,12 +5914,116 @@ def _check_tcgen05_mma_block_scale(m): OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False), ), ), + # atom's other syntax lines (PTX ISA 9.7.14.5), each its own entry because + # each differs in shape or in type domain from the `.op` line above. + # + # `.cas` compares and swaps, so it takes a fourth value; it is also the one + # atom line with no `{.level::cache_hint}` (ptxas rejects the qualifier + # there) and the only one that reaches down to `.b16`. + InstructionEntry( + name="atom_cas", + mnemonic="atom", + slots=( + ModifierSlot("sem", _ATOM_SEM, optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", _ATOM_SPACES, optional=True), + ModifierSlot("op", ("cas",)), + ModifierSlot("type", ("b16", "b32", "b64", "b128")), + ), + operands=( + OperandSlot("d", rw="w"), + OperandSlot("addr", kind="addr"), + OperandSlot("compare"), + OperandSlot("value"), + ), + ), + # `.exch` just writes, so it keeps the plain shape; its type line starts at + # 32 bits (ptxas: "Arguments mismatch" for `.b16`) and reaches `.b128`. + InstructionEntry( + name="atom_exch", + mnemonic="atom", + slots=( + ModifierSlot("sem", _ATOM_SEM, optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", _ATOM_SPACES, optional=True), + ModifierSlot("op", ("exch",)), + ModifierSlot("cache", ("L2::cache_hint",), optional=True), + ModifierSlot("type", ("b32", "b64", "b128")), + ), + check=_check_cache_hint, + operands=( + OperandSlot("d", rw="w"), + OperandSlot("addr", kind="addr"), + OperandSlot("value"), + OperandSlot("cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False), + ), + ), + # The half-precision add lines of atom and red. `.noftz` is not optional on + # them -- ptxas: "'.noftz' modifier required for instruction 'atom' with + # type '.f16'" -- so it is a single-choice slot rather than an optional one. + *[ + InstructionEntry( + name=f"{mnem}_half", + mnemonic=mnem, + slots=( + ModifierSlot("sem", _ATOM_SEM if mnem == "atom" else _RED_SEM, optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", _ATOM_SPACES, optional=True), + ModifierSlot("op", ("add",)), + ModifierSlot("noftz", ("noftz",)), + ModifierSlot("cache", ("L2::cache_hint",), optional=True), + ModifierSlot("type", (*_ATOM_VEC_HALF, *_HALF_X2)), + ), + check=_check_cache_hint, + operands=( + *((OperandSlot("d", rw="w"),) if mnem == "atom" else ()), + OperandSlot("addr", kind="addr"), + OperandSlot("value"), + OperandSlot( + "cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False + ), + ), + ) + for mnem in ("atom", "red") + ], + # atom's two vector lines. Both operands are register groups, so this is a + # different shape again; `_check_atom_vec` holds which widths pair with + # which element. `.f32` has no `.noftz` (there is no half to flush), which + # is why the two lines are two entries rather than one. + # The ISA's own examples spell these tokens in a different order than its + # syntax line (`atom.global.v8.f16.max.noftz` against + # `atom{...}.op.noftz{.cache}.vec.type`); ptxas accepts both, and the + # normative syntax line is what is registered here. + *[ + InstructionEntry( + name=name, + mnemonic="atom", + slots=( + ModifierSlot("sem", _ATOM_SEM, optional=True), + ModifierSlot("scope", _ATOM_SCOPES, optional=True), + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("op", ops), + *((ModifierSlot("noftz", ("noftz",)),) if noftz else ()), + ModifierSlot("cache", ("L2::cache_hint",), optional=True), + ModifierSlot("vec", ("v2", "v4", "v8") if noftz else ("v2", "v4")), + ModifierSlot("type", types), + ), + check=_check_atom_vec, + operands=( + OperandSlot("d", rw="w", lanes=_vec_lanes), + OperandSlot("addr", kind="addr"), + OperandSlot("value", lanes=_vec_lanes), + OperandSlot( + "cache_policy", dtype="u64", lanes=_present_lanes("cache"), vector=False + ), + ), + ) + for name, ops, noftz, types in ( + ("atom_vec_half", ("add", "min", "max"), True, (*_ATOM_VEC_HALF, *_HALF_X2)), + ("atom_vec_f32", ("add",), False, ("f32",)), + ) + ], # red / atom scalar `.op` forms per PTX ISA 9.7.14.6 and 9.7.14.5. - # NOT REGISTERED: each is a different syntax shape, i.e. its own entry - # whenever a caller wants it -- - # - the .vec_16_bit/.vec_32_bit vector forms - # - .f16/.bf16/.f16x2/.bf16x2 (add.noftz), which need half carrier types - # - atom's .cas (3 operands) and .exch (own type set): other syntax shapes InstructionEntry( name="red", slots=( @@ -3914,10 +6207,10 @@ def _check_tcgen05_mma_block_scale(m): ), check=_check_mbarrier_sem_scope, operands=( - # Pinned u64, not b64: the bit-type dtype axis would offer an f64 + # `b64i`, not `b64`: the bit-type dtype axis would offer an f64 # carrier, and ptxas rejects an .f64 register as the state operand # ("Arguments mismatch for instruction 'mbarrier.arrive'"). - OperandSlot("state", rw="rw", dtype="u64"), + OperandSlot("state", rw="rw", dtype="b64i"), OperandSlot("addr", kind="addr"), OperandSlot("count", dtype="u32"), ), @@ -3981,6 +6274,135 @@ def _check_tcgen05_mma_block_scale(m): ), operands=(OperandSlot("addr", kind="addr"),), ), + # 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 + # `.shared::cluster` lines offer; these return the phase state instead, and + # the ISA gives them `.shared{::cta}` alone. + *[ + InstructionEntry( + name=f"mbarrier_{act}{suffix}_state", + mnemonic="mbarrier", + slots=( + ModifierSlot("action", (act,)), + *((ModifierSlot("expect_tx", ("expect_tx",)),) if suffix == "_expect_tx" else ()), + ModifierSlot("sem", ("release", "relaxed"), optional=True), + ModifierSlot("scope", ("cta", "cluster"), optional=True), + ModifierSlot("space", ("shared", "shared::cta"), optional=True), + ModifierSlot("type", ("b64",)), + ), + check=_check_mbarrier_sem_scope, + operands=( + # The phase state is an opaque 64-bit token: ptxas takes an + # 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("count", dtype="u32"),) if count else ()), + ), + ) + for act in ("arrive", "arrive_drop") + # Three shapes per action, mirroring the `_` family above: bare, with a + # thread count, and with a transaction count. The noComplete line needs + # no twin -- it has no `_` spelling, so `mbarrier_*_no_complete` above + # is already the state-returning entry. + for suffix, count in (("", False), ("_count", True), ("_expect_tx", True)) + ], + # The non-parity waits (PTX ISA 9.7.14.16.19), which take the state token + # an arrive returned rather than a phase parity. try_wait may also carry a + # sleep hint, which is a separate syntax line and so a separate entry. + # NOT REGISTERED: the `.phase_type::*` qualifiers and the + # `waitComplete|reportPredicate{, reportValue}` report forms -- both are + # PTX ISA 9.3, and ptxas 13.2 (9.2) rejects the token outright. + *[ + InstructionEntry( + name=f"mbarrier_{act}{'_hint' if hint else ''}", + mnemonic="mbarrier", + slots=( + ModifierSlot("action", (act,)), + ModifierSlot("sem", ("acquire", "relaxed"), optional=True), + ModifierSlot("scope", ("cta", "cluster"), optional=True), + ModifierSlot("space", ("shared", "shared::cta"), optional=True), + ModifierSlot("type", ("b64",)), + ), + check=_check_mbarrier_sem_scope, + operands=( + OperandSlot("wait_complete", rw="w", dtype="pred"), + OperandSlot("addr", kind="addr"), + OperandSlot("state", dtype="b64i"), # the token an arrive returned + *((OperandSlot("time_hint", dtype="u32"),) if hint else ()), + ), + ) + for act, hint in (("test_wait", False), ("try_wait", False), ("try_wait", True)) + ], + # tensormap.cp_fenceproxy per PTX ISA 9.7.14.17: copy a tensor-map object + # from shared to global and fence the tensormap proxy over it, in one + # instruction. `size` is the literal 128 -- ptxas both refuses a register + # there ("Arguments mismatch") and names the only legal value ("unexpected + # value '64', expected to be 128"), so it is table-owned. + InstructionEntry( + name="tensormap_cp_fenceproxy", + mnemonic="tensormap.cp_fenceproxy", + slots=( + ModifierSlot("dst", ("global",)), + ModifierSlot("src", ("shared::cta",)), + ModifierSlot("proxy", ("tensormap::generic",)), + ModifierSlot("sem", ("release",)), + ModifierSlot("scope", _ATOM_SCOPES), + ModifierSlot("sync", ("sync",)), + ModifierSlot("aligned", ("aligned",)), + ), + cert_arch="sm_90", + operands=( + OperandSlot("dst_mem", kind="addr", space="global"), + OperandSlot("src_mem", kind="addr", space="shared::cta"), + OperandSlot("size", kind="imm", literal="128"), + ), + ), + # red.async per PTX ISA 9.7.14.7: an asynchronous reduction whose + # completion is signalled through an mbarrier. The ISA writes one syntax + # line per op group; `_check_red_async` is that grid. + InstructionEntry( + name="red_async", + mnemonic="red.async", + slots=( + ModifierSlot("sem", ("relaxed",)), + ModifierSlot("scope", ("cluster",)), + ModifierSlot("space", ("shared::cluster",), optional=True), + ModifierSlot("completion", ("mbarrier::complete_tx::bytes",)), + ModifierSlot("op", tuple(_RED_ASYNC_OPS)), + ModifierSlot("type", ("u32", "s32", "u64", "b32")), + ), + check=_check_red_async, + cert_arch="sm_90", + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("value"), + OperandSlot("mbar", kind="addr"), + ), + ), + # The release line of the same subsection, which reduces straight into + # global memory and needs no mbarrier. Like st.async's release line it is + # sm_100, and `.mmio` is system-scoped only. + # NOT REGISTERED: `multimem.red.async` (9.7.14.8) -- ptxas 13.2 rejects the + # mnemonic at sm_90 and sm_100 alike, as it does `multimem.st.async`. + InstructionEntry( + name="red_async_release", + mnemonic="red.async", + slots=( + ModifierSlot("mmio", ("mmio",), optional=True), + ModifierSlot("sem", ("release",)), + ModifierSlot("scope", ("gpu", "sys")), + ModifierSlot("space", ("global",), optional=True), + ModifierSlot("op", ("add",)), + ModifierSlot("type", ("u32", "s32", "u64", "s64")), + ), + check=_check_st_async_rel, + cert_arch="sm_100", + operands=( + OperandSlot("addr", kind="addr"), + OperandSlot("value"), + ), + ), # cp.async completion tracking: cp.async.mbarrier.arrive per PTX ISA # 9.7.14.16.18, cp.async.commit_group per 9.7.9.26.3.2, cp.async.wait_group # and cp.async.wait_all per 9.7.9.26.3.3, cp.async.bulk.commit_group / @@ -3991,8 +6413,8 @@ def _check_tcgen05_mma_block_scale(m): # set is what makes every one of them certifiable. 0..7 covers every call # site (pipeline depths); widen the tuple if a deeper pipeline appears. # - # NOT REGISTERED: the `cp.async` ca/cg copy lines, which carry an optional - # ignore-src operand. + # (The `cp.async` ca/cg copy lines this note once excluded are registered + # in the 9.7.9 group above, ignore-src operand and all.) InstructionEntry( # cp.async.mbarrier.arrive{.noinc}{.shared{::cta}}.b64 [addr]; name="cp_async_mbarrier_arrive", mnemonic="cp", diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py index 29b110ccce5a..136a89e9dd8a 100644 --- a/python/tvm/backend/cuda/script.py +++ b/python/tvm/backend/cuda/script.py @@ -138,8 +138,8 @@ def __init__(self): self.wgmma = CudaWgmmaNamespace() self.tcgen05 = CudaTcgen05Namespace() self.any_sync = _op_wrapper(_cuda_op.cuda_any_sync) - # elect.sync plus the two movs that materialize its d|p pair: a - # multi-statement asm block, so it belongs here rather than T.ptx. + # elect.sync plus the predicated mov that materializes its predicate: + # a multi-statement asm block, so it belongs here rather than T.ptx. # The warp-specialization passes match this op to build predicates. self.elect_sync: Callable[..., Any] = _op_wrapper(_cuda_op.cuda_elect_sync) # `mov.u32 d, %sreg` -- one PTX instruction, but the special-register diff --git a/python/tvm/script/tirx.pyi b/python/tvm/script/tirx.pyi index ce14826861e9..c991b31aab90 100644 --- a/python/tvm/script/tirx.pyi +++ b/python/tvm/script/tirx.pyi @@ -22,8 +22,31 @@ Regenerate: from typing import Any +class _Chain_abs: + """`abs` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a) + """ + + bf16: _Chain_abs + bf16x2: _Chain_abs + f16: _Chain_abs + f16x2: _Chain_abs + f32: _Chain_abs + f64: _Chain_abs + ftz: _Chain_abs + s16: _Chain_abs + s32: _Chain_abs + s64: _Chain_abs + def __call__(self, *args: Any) -> None: ... + +class _Chain_activemask: + """`activemask` — type∈{b32}""" + + b32: _Chain_activemask + def __call__(self, d: Any, *args: Any) -> None: ... + class _Chain_add: - """`add` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + """`add` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a, b) """ @@ -39,19 +62,36 @@ class _Chain_add: rn: _Chain_add rp: _Chain_add rz: _Chain_add + s16: _Chain_add + s16x2: _Chain_add + s32: _Chain_add + s64: _Chain_add sat: _Chain_add + u16: _Chain_add + u16x2: _Chain_add + u32: _Chain_add + u64: _Chain_add def __call__(self, *args: Any) -> None: ... +class _Chain_and: + """`and` — type∈{pred,b16,b32,b64}""" + + b16: _Chain_and + b32: _Chain_and + b64: _Chain_and + pred: _Chain_and + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_applypriority: + """`applypriority` — space∈{global} (opt); level∈{L2::evict_normal}""" + + L2__evict_normal: _Chain_applypriority + global_: _Chain_applypriority + def __call__(self, addr: Any, *args: Any, pred: Any = None) -> None: ... + class _Chain_atom: - """`atom` — sem∈{relaxed,acquire,release,acq_rel} (opt); scope∈{cta,cluster,gpu,sys} (opt); - space∈{global,shared,shared::cta,shared::cluster} (opt); - op∈{and,or,xor,add,inc,dec,min,max}; cache∈{L2::cache_hint} (opt); - type∈{b32,b64,u32,u64,s32,s64,f32,f64} — op x type pairings for atom/red (PTX ISA - 9.7.14.5 / 9.7.14.6). Normative source: ISA Table 35 (atom) and Table 36 (red), - which give the pairing cell by cell. The `.type = {...}` line in the Syntax block is - only the union across ops, which is why it cannot be transcribed directly. Half- - precision types appear in ptxas' message but are excluded from this entry (they need - .noftz and a half carrier type). + """`atom` — 6 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (*__operands); (d, addr, compare, value) """ L2__cache_hint: _Chain_atom @@ -59,11 +99,19 @@ class _Chain_atom: acquire: _Chain_atom add: _Chain_atom and_: _Chain_atom + b128: _Chain_atom + b16: _Chain_atom b32: _Chain_atom b64: _Chain_atom + bf16: _Chain_atom + bf16x2: _Chain_atom + cas: _Chain_atom cluster: _Chain_atom cta: _Chain_atom dec: _Chain_atom + exch: _Chain_atom + f16: _Chain_atom + f16x2: _Chain_atom f32: _Chain_atom f64: _Chain_atom global_: _Chain_atom @@ -71,6 +119,7 @@ class _Chain_atom: inc: _Chain_atom max: _Chain_atom min: _Chain_atom + noftz: _Chain_atom or_: _Chain_atom relaxed: _Chain_atom release: _Chain_atom @@ -82,35 +131,91 @@ class _Chain_atom: sys: _Chain_atom u32: _Chain_atom u64: _Chain_atom + v2: _Chain_atom + v4: _Chain_atom + v8: _Chain_atom xor: _Chain_atom - def __call__(self, *__operands: Any) -> None: ... + def __call__(self, *args: Any) -> None: ... class _Chain_bar: - """`bar` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list, - so the call selects one. Shapes: (a); (a, b); (membermask) + """`bar` — 8 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (a); (a, b); (membermask); (d, a, c); (d, a, b, c) """ + and_: _Chain_bar arrive: _Chain_bar cta: _Chain_bar + or_: _Chain_bar + popc: _Chain_bar + pred: _Chain_bar + red: _Chain_bar sync: _Chain_bar + u32: _Chain_bar warp: _Chain_bar - def __call__(self, *args: Any, pred: Any = None) -> None: ... + def __call__(self, *args: Any) -> None: ... class _Chain_barrier: - """`barrier` — 5 entries sharing this mnemonic; PTX puts their difference in the operand - list, so the call selects one. Shapes: (a); (a, b); () + """`barrier` — 9 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (a); (a, b); (); (d, a, c); (d, a, b, c) """ acquire: _Chain_barrier aligned: _Chain_barrier + and_: _Chain_barrier arrive: _Chain_barrier cluster: _Chain_barrier cta: _Chain_barrier + or_: _Chain_barrier + popc: _Chain_barrier + pred: _Chain_barrier + red: _Chain_barrier relaxed: _Chain_barrier release: _Chain_barrier sync: _Chain_barrier + u32: _Chain_barrier wait: _Chain_barrier - def __call__(self, *args: Any, pred: Any = None) -> None: ... + def __call__(self, *args: Any) -> None: ... + +class _Chain_bfe: + """`bfe` — type∈{u32,u64,s32,s64}""" + + s32: _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: ... + +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: ... + +class _Chain_bfind: + """`bfind` — shiftamt∈{shiftamt} (opt); type∈{u32,u64,s32,s64}""" + + s32: _Chain_bfind + s64: _Chain_bfind + shiftamt: _Chain_bfind + u32: _Chain_bfind + u64: _Chain_bfind + def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + +class _Chain_bmsk: + """`bmsk` — mode∈{clamp,wrap}; type∈{b32}""" + + b32: _Chain_bmsk + clamp: _Chain_bmsk + wrap: _Chain_bmsk + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_brev: + """`brev` — type∈{b32,b64}""" + + b32: _Chain_brev + b64: _Chain_brev + def __call__(self, d: Any, a: Any, *args: Any) -> None: ... class _Chain_clusterlaunchcontrol: """`clusterlaunchcontrol` — 4 entries sharing this mnemonic; PTX puts their difference in @@ -135,9 +240,39 @@ class _Chain_clusterlaunchcontrol: v4: _Chain_clusterlaunchcontrol def __call__(self, *args: Any) -> None: ... +class _Chain_clz: + """`clz` — type∈{b32,b64}""" + + b32: _Chain_clz + b64: _Chain_clz + def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + +class _Chain_cnot: + """`cnot` — type∈{b16,b32,b64}""" + + b16: _Chain_cnot + b32: _Chain_cnot + b64: _Chain_cnot + def __call__(self, d: Any, a: Any, *args: Any) -> 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: ... + +class _Chain_cos: + """`cos` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" + + approx: _Chain_cos + f32: _Chain_cos + ftz: _Chain_cos + def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + class _Chain_cp: - """`cp` — 21 entries sharing this mnemonic; PTX puts their difference in the operand list, - so the call selects one. Shapes: (group); (*__operands); (); (dst_mem, src_mem, size, + """`cp` — 22 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (*__operands); (group); (); (dst_mem, src_mem, size, mbar); (addr) """ @@ -185,6 +320,24 @@ class _Chain_cp: xor: _Chain_cp def __call__(self, *args: Any, pred: Any = None) -> None: ... +class _Chain_createpolicy: + """`createpolicy` — 4 entries sharing this mnemonic; PTX puts their difference in the + operand list, so the call selects one. Shapes: (cache_policy, addr, primary_size, + total_size); (cache_policy); (cache_policy, fraction); (cache_policy, access_property) + """ + + L2: _Chain_createpolicy + L2__evict_first: _Chain_createpolicy + L2__evict_last: _Chain_createpolicy + L2__evict_normal: _Chain_createpolicy + L2__evict_unchanged: _Chain_createpolicy + b64: _Chain_createpolicy + cvt: _Chain_createpolicy + fractional: _Chain_createpolicy + global_: _Chain_createpolicy + range: _Chain_createpolicy + def __call__(self, *args: Any) -> None: ... + class _Chain_cvt: """`cvt` — 27 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a); (d, a, b); (d, a, b, rbits); (*__operands); (d, @@ -235,21 +388,104 @@ class _Chain_cvt: ue8m0x2: _Chain_cvt def __call__(self, *args: Any) -> None: ... +class _Chain_cvt_pack: + """`cvt_pack` — 2 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (d, a, b); (d, a, b, c) + """ + + b32: _Chain_cvt_pack + s16: _Chain_cvt_pack + s2: _Chain_cvt_pack + s32: _Chain_cvt_pack + s4: _Chain_cvt_pack + s8: _Chain_cvt_pack + sat: _Chain_cvt_pack + u16: _Chain_cvt_pack + u2: _Chain_cvt_pack + u4: _Chain_cvt_pack + u8: _Chain_cvt_pack + def __call__(self, *args: Any) -> None: ... + class _Chain_cvta: - """`cvta` — dir∈{to}; space∈{shared}; type∈{u64}""" + """`cvta` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, ptr); (d, a) + """ + const: _Chain_cvta + global_: _Chain_cvta + local: _Chain_cvta + param: _Chain_cvta + param__entry: _Chain_cvta shared: _Chain_cvta + shared__cluster: _Chain_cvta + shared__cta: _Chain_cvta to: _Chain_cvta u64: _Chain_cvta - def __call__(self, d: Any, ptr: Any, *args: Any) -> None: ... + def __call__(self, *args: Any) -> None: ... + +class _Chain_discard: + """`discard` — space∈{global} (opt); level∈{L2}""" + + L2: _Chain_discard + global_: _Chain_discard + def __call__(self, addr: Any, *args: Any, pred: Any = None) -> None: ... + +class _Chain_div: + """`div` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a, b) + """ + + approx: _Chain_div + f32: _Chain_div + f64: _Chain_div + ftz: _Chain_div + full: _Chain_div + rm: _Chain_div + rn: _Chain_div + rp: _Chain_div + rz: _Chain_div + s16: _Chain_div + s32: _Chain_div + s64: _Chain_div + u16: _Chain_div + u32: _Chain_div + u64: _Chain_div + def __call__(self, *args: Any) -> None: ... + +class _Chain_dp2a: + """`dp2a` — mode∈{lo,hi}; atype∈{u32,s32}; btype∈{u32,s32}""" + + hi: _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: ... + +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: ... + +class _Chain_elect_sync: + """`elect_sync` — (no modifiers)""" + + def __call__(self, d: Any, p: Any, membermask: Any, *args: Any) -> None: ... class _Chain_ex2: - """`ex2` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" + """`ex2` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, value) + """ approx: _Chain_ex2 + bf16: _Chain_ex2 + bf16x2: _Chain_ex2 + f16: _Chain_ex2 + f16x2: _Chain_ex2 f32: _Chain_ex2 ftz: _Chain_ex2 - def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + def __call__(self, *args: Any) -> None: ... class _Chain_fence: """`fence` — 5 entries sharing this mnemonic; PTX puts their difference in the operand @@ -275,25 +511,26 @@ class _Chain_fence: def __call__(self, *args: Any, pred: Any = None) -> None: ... class _Chain_fma: - """`fma` — rnd∈{rn,rz,rm,rp}; ftz∈{ftz} (opt); sat∈{sat} (opt); type∈{f32,f64,f32x2}; - srctype∈{f16,bf16} (opt) — Which qualifiers each add/sub/mul/fma syntax line allows (PTX - ISA 9.7.3.{3,4,5,6}, 9.7.5). Same-precision lines: op{.rnd}{.ftz}{.sat}.f32 | - op{.rnd}{.ftz}.f32x2 | op{.rnd}.f64 Mixed-precision lines: op{.rnd}{.sat}.f32.atype - (.atype = .f16 | .bf16) + """`fma` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a, b, c) """ bf16: _Chain_fma + bf16x2: _Chain_fma f16: _Chain_fma + f16x2: _Chain_fma f32: _Chain_fma f32x2: _Chain_fma f64: _Chain_fma ftz: _Chain_fma + oob: _Chain_fma + relu: _Chain_fma rm: _Chain_fma rn: _Chain_fma rp: _Chain_fma rz: _Chain_fma sat: _Chain_fma - def __call__(self, d: Any, a: Any, b: Any, c: Any, *args: Any) -> None: ... + def __call__(self, *args: Any) -> None: ... class _Chain_fns: """`fns` — type∈{b32}""" @@ -301,6 +538,16 @@ class _Chain_fns: b32: _Chain_fns def __call__(self, d: Any, mask: Any, base: Any, offset: Any, *args: Any) -> None: ... +class _Chain_getctarank: + """`getctarank` — 2 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (d, a) + """ + + shared__cluster: _Chain_getctarank + u32: _Chain_getctarank + u64: _Chain_getctarank + def __call__(self, *args: Any) -> None: ... + class _Chain_griddepcontrol: """`griddepcontrol` — action∈{launch_dependents,wait}""" @@ -308,6 +555,21 @@ class _Chain_griddepcontrol: wait: _Chain_griddepcontrol def __call__(self, *args: Any, pred: Any = None) -> None: ... +class _Chain_isspacep: + """`isspacep` — + space∈{const,global,local,shared,shared::cta,shared::cluster,param,param::entry} + """ + + const: _Chain_isspacep + global_: _Chain_isspacep + local: _Chain_isspacep + param: _Chain_isspacep + param__entry: _Chain_isspacep + shared: _Chain_isspacep + shared__cluster: _Chain_isspacep + shared__cta: _Chain_isspacep + def __call__(self, p: Any, a: Any, *args: Any) -> None: ... + class _Chain_ld: """`ld` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (*__operands) @@ -388,6 +650,88 @@ class _Chain_ldmatrix: x4: _Chain_ldmatrix def __call__(self, *args: Any) -> None: ... +class _Chain_ldu: + """`ldu` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, addr); (*__operands) + """ + + b128: _Chain_ldu + b16: _Chain_ldu + b32: _Chain_ldu + b64: _Chain_ldu + b8: _Chain_ldu + f32: _Chain_ldu + f64: _Chain_ldu + global_: _Chain_ldu + s16: _Chain_ldu + s32: _Chain_ldu + s64: _Chain_ldu + s8: _Chain_ldu + u16: _Chain_ldu + u32: _Chain_ldu + u64: _Chain_ldu + u8: _Chain_ldu + v2: _Chain_ldu + v4: _Chain_ldu + def __call__(self, *args: Any) -> None: ... + +class _Chain_lg2: + """`lg2` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" + + approx: _Chain_lg2 + f32: _Chain_lg2 + ftz: _Chain_lg2 + def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + +class _Chain_lop3: + """`lop3` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a, b, c, immLut); (d, p, a, b, c, immLut, q) + """ + + and_: _Chain_lop3 + b32: _Chain_lop3 + or_: _Chain_lop3 + def __call__(self, *args: Any) -> None: ... + +class _Chain_mad: + """`mad` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a, b, c) + """ + + f32: _Chain_mad + f64: _Chain_mad + ftz: _Chain_mad + hi: _Chain_mad + lo: _Chain_mad + rm: _Chain_mad + rn: _Chain_mad + rp: _Chain_mad + rz: _Chain_mad + s16: _Chain_mad + s32: _Chain_mad + s64: _Chain_mad + sat: _Chain_mad + u16: _Chain_mad + u32: _Chain_mad + u64: _Chain_mad + wide: _Chain_mad + def __call__(self, *args: Any) -> 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". + """ + + hi: _Chain_mad24 + lo: _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: ... + class _Chain_mapa: """`mapa` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a, b) @@ -398,6 +742,18 @@ class _Chain_mapa: u64: _Chain_mapa def __call__(self, *args: Any) -> None: ... +class _Chain_match: + """`match` — 3 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (d, a, membermask); (d, p, a, membermask) + """ + + all: _Chain_match + any: _Chain_match + b32: _Chain_match + b64: _Chain_match + sync: _Chain_match + def __call__(self, *args: Any) -> None: ... + class _Chain_max: """`max` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a, b); (d, a, b, c) @@ -425,9 +781,10 @@ class _Chain_max: def __call__(self, *args: Any) -> None: ... class _Chain_mbarrier: - """`mbarrier` — 15 entries sharing this mnemonic; PTX puts their difference in the operand + """`mbarrier` — 24 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (addr, count); (addr); (addr, tx_count); (state, addr, count); (wait_complete, addr, phase); (wait_complete, addr, phase, time_hint); + (state, addr); (wait_complete, addr, state); (wait_complete, addr, state, time_hint); (count, state) """ @@ -525,18 +882,28 @@ class _Chain_mma: def __call__(self, *args: Any) -> None: ... class _Chain_mov: - """`mov` — 10 entries sharing this mnemonic; PTX puts their difference in the operand list, + """`mov` — 11 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a0, a1); (d0, d1, a); (d, a0, a1, a2, a3); (d0, d1, - d2, d3, a) + d2, d3, a); (d, a) """ b128: _Chain_mov + b16: _Chain_mov b32: _Chain_mov b64: _Chain_mov + f32: _Chain_mov + f64: _Chain_mov + pred: _Chain_mov + s16: _Chain_mov + s32: _Chain_mov + s64: _Chain_mov + u16: _Chain_mov + u32: _Chain_mov + u64: _Chain_mov def __call__(self, *args: Any) -> None: ... class _Chain_mul: - """`mul` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + """`mul` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a, b) """ @@ -548,21 +915,173 @@ class _Chain_mul: f32x2: _Chain_mul f64: _Chain_mul ftz: _Chain_mul + hi: _Chain_mul + lo: _Chain_mul rm: _Chain_mul rn: _Chain_mul rp: _Chain_mul rz: _Chain_mul + s16: _Chain_mul + s32: _Chain_mul + s64: _Chain_mul sat: _Chain_mul + u16: _Chain_mul + u32: _Chain_mul + u64: _Chain_mul + wide: _Chain_mul def __call__(self, *args: Any) -> None: ... +class _Chain_mul24: + """`mul24` — mode∈{hi,lo}; type∈{u32,s32}""" + + hi: _Chain_mul24 + lo: _Chain_mul24 + s32: _Chain_mul24 + u32: _Chain_mul24 + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_multimem_ld_reduce: + """`multimem_ld_reduce` — 3 entries sharing this mnemonic; PTX puts their difference in the + operand list, so the call selects one. Shapes: (d, addr); (*__operands) + """ + + acc__f32: _Chain_multimem_ld_reduce + acquire: _Chain_multimem_ld_reduce + add: _Chain_multimem_ld_reduce + and_: _Chain_multimem_ld_reduce + b32: _Chain_multimem_ld_reduce + b64: _Chain_multimem_ld_reduce + bf16: _Chain_multimem_ld_reduce + bf16x2: _Chain_multimem_ld_reduce + cluster: _Chain_multimem_ld_reduce + cta: _Chain_multimem_ld_reduce + f16: _Chain_multimem_ld_reduce + f16x2: _Chain_multimem_ld_reduce + f32: _Chain_multimem_ld_reduce + f64: _Chain_multimem_ld_reduce + global_: _Chain_multimem_ld_reduce + gpu: _Chain_multimem_ld_reduce + max: _Chain_multimem_ld_reduce + min: _Chain_multimem_ld_reduce + or_: _Chain_multimem_ld_reduce + relaxed: _Chain_multimem_ld_reduce + s32: _Chain_multimem_ld_reduce + s64: _Chain_multimem_ld_reduce + sys: _Chain_multimem_ld_reduce + u32: _Chain_multimem_ld_reduce + u64: _Chain_multimem_ld_reduce + v2: _Chain_multimem_ld_reduce + v4: _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: ... + +class _Chain_multimem_red: + """`multimem_red` — 3 entries sharing this mnemonic; PTX puts their difference in the + operand list, so the call selects one. Shapes: (addr, b); (*__operands) + """ + + add: _Chain_multimem_red + and_: _Chain_multimem_red + b32: _Chain_multimem_red + b64: _Chain_multimem_red + bf16: _Chain_multimem_red + bf16x2: _Chain_multimem_red + cluster: _Chain_multimem_red + cta: _Chain_multimem_red + f16: _Chain_multimem_red + f16x2: _Chain_multimem_red + f32: _Chain_multimem_red + f64: _Chain_multimem_red + global_: _Chain_multimem_red + gpu: _Chain_multimem_red + max: _Chain_multimem_red + min: _Chain_multimem_red + or_: _Chain_multimem_red + relaxed: _Chain_multimem_red + release: _Chain_multimem_red + s32: _Chain_multimem_red + s64: _Chain_multimem_red + sys: _Chain_multimem_red + u32: _Chain_multimem_red + u64: _Chain_multimem_red + v2: _Chain_multimem_red + v4: _Chain_multimem_red + v8: _Chain_multimem_red + xor: _Chain_multimem_red + def __call__(self, *args: Any, pred: Any = None) -> None: ... + +class _Chain_multimem_st: + """`multimem_st` — 3 entries sharing this mnemonic; PTX puts their difference in the + operand list, so the call selects one. Shapes: (addr, b); (*__operands) + """ + + b32: _Chain_multimem_st + b64: _Chain_multimem_st + bf16: _Chain_multimem_st + bf16x2: _Chain_multimem_st + cluster: _Chain_multimem_st + cta: _Chain_multimem_st + f16: _Chain_multimem_st + f16x2: _Chain_multimem_st + f32: _Chain_multimem_st + f64: _Chain_multimem_st + global_: _Chain_multimem_st + gpu: _Chain_multimem_st + relaxed: _Chain_multimem_st + release: _Chain_multimem_st + s32: _Chain_multimem_st + s64: _Chain_multimem_st + sys: _Chain_multimem_st + u32: _Chain_multimem_st + u64: _Chain_multimem_st + v2: _Chain_multimem_st + v4: _Chain_multimem_st + v8: _Chain_multimem_st + weak: _Chain_multimem_st + def __call__(self, *args: Any, pred: Any = None) -> None: ... + class _Chain_neg: - """`neg` — ftz∈{ftz} (opt); type∈{f32,f64} — `neg{.ftz}.f32 d, a;` and `neg.f64 d, a;` (ISA - 9.7.3.10) -- the f64 line spells no `.ftz`. + """`neg` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a) """ + bf16: _Chain_neg + bf16x2: _Chain_neg + f16: _Chain_neg + f16x2: _Chain_neg f32: _Chain_neg f64: _Chain_neg ftz: _Chain_neg + s16: _Chain_neg + s32: _Chain_neg + s64: _Chain_neg + def __call__(self, *args: Any) -> None: ... + +class _Chain_not: + """`not` — type∈{pred,b16,b32,b64}""" + + b16: _Chain_not + b32: _Chain_not + b64: _Chain_not + pred: _Chain_not + def __call__(self, d: Any, a: Any, *args: Any) -> None: ... + +class _Chain_or: + """`or` — type∈{pred,b16,b32,b64}""" + + b16: _Chain_or + b32: _Chain_or + b64: _Chain_or + pred: _Chain_or + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_popc: + """`popc` — type∈{b32,b64}""" + + b32: _Chain_popc + b64: _Chain_popc def __call__(self, d: Any, a: Any, *args: Any) -> None: ... class _Chain_prefetch: @@ -586,9 +1105,34 @@ class _Chain_prefetch: tensormap: _Chain_prefetch def __call__(self, addr: Any, *args: Any, pred: Any = None) -> None: ... +class _Chain_prefetchu: + """`prefetchu` — level∈{L1}""" + + L1: _Chain_prefetchu + def __call__(self, addr: Any, *args: Any, pred: Any = None) -> None: ... + +class _Chain_prmt: + """`prmt` — type∈{b32}; mode∈{f4e,b4e,rc8,ecl,ecr,rc16} (opt)""" + + b32: _Chain_prmt + b4e: _Chain_prmt + ecl: _Chain_prmt + ecr: _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: ... + class _Chain_rcp: - """`rcp` — mode∈{approx,rn,rz,rm,rp}; ftz∈{ftz} (opt); type∈{f32,f64} — This entry's - rcp.approx is f32-only; .f64 is IEEE-rounded, no .ftz (PTX ISA 9.7.3.13). + """`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. """ approx: _Chain_rcp @@ -602,15 +1146,8 @@ class _Chain_rcp: def __call__(self, d: Any, value: Any, *args: Any) -> None: ... class _Chain_red: - """`red` — sem∈{relaxed,release} (opt); scope∈{cta,cluster,gpu,sys} (opt); - space∈{global,shared,shared::cta,shared::cluster} (opt); - op∈{and,or,xor,add,inc,dec,min,max}; cache∈{L2::cache_hint} (opt); - type∈{b32,b64,u32,u64,s32,s64,f32,f64} — op x type pairings for atom/red (PTX ISA - 9.7.14.5 / 9.7.14.6). Normative source: ISA Table 35 (atom) and Table 36 (red), - which give the pairing cell by cell. The `.type = {...}` line in the Syntax block is - only the union across ops, which is why it cannot be transcribed directly. Half- - precision types appear in ptxas' message but are excluded from this entry (they need - .noftz and a half carrier type). + """`red` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (*__operands) """ L2__cache_hint: _Chain_red @@ -618,9 +1155,13 @@ class _Chain_red: and_: _Chain_red b32: _Chain_red b64: _Chain_red + bf16: _Chain_red + bf16x2: _Chain_red cluster: _Chain_red cta: _Chain_red dec: _Chain_red + f16: _Chain_red + f16x2: _Chain_red f32: _Chain_red f64: _Chain_red global_: _Chain_red @@ -628,6 +1169,7 @@ class _Chain_red: inc: _Chain_red max: _Chain_red min: _Chain_red + noftz: _Chain_red or_: _Chain_red relaxed: _Chain_red release: _Chain_red @@ -640,7 +1182,146 @@ class _Chain_red: u32: _Chain_red u64: _Chain_red xor: _Chain_red - def __call__(self, *__operands: Any, pred: Any = None) -> None: ... + def __call__(self, *args: Any, pred: Any = None) -> None: ... + +class _Chain_red_async: + """`red_async` — 2 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (addr, value, mbar); (addr, value) + """ + + add: _Chain_red_async + and_: _Chain_red_async + b32: _Chain_red_async + cluster: _Chain_red_async + dec: _Chain_red_async + global_: _Chain_red_async + gpu: _Chain_red_async + inc: _Chain_red_async + max: _Chain_red_async + mbarrier__complete_tx__bytes: _Chain_red_async + min: _Chain_red_async + mmio: _Chain_red_async + or_: _Chain_red_async + relaxed: _Chain_red_async + release: _Chain_red_async + s32: _Chain_red_async + s64: _Chain_red_async + shared__cluster: _Chain_red_async + sys: _Chain_red_async + u32: _Chain_red_async + u64: _Chain_red_async + xor: _Chain_red_async + def __call__(self, *args: Any, pred: Any = None) -> None: ... + +class _Chain_redux_sync: + """`redux_sync` — 3 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (d, a, membermask) + """ + + NaN: _Chain_redux_sync + abs: _Chain_redux_sync + add: _Chain_redux_sync + and_: _Chain_redux_sync + b32: _Chain_redux_sync + f32: _Chain_redux_sync + max: _Chain_redux_sync + min: _Chain_redux_sync + or_: _Chain_redux_sync + s32: _Chain_redux_sync + u32: _Chain_redux_sync + xor: _Chain_redux_sync + def __call__(self, *args: Any) -> None: ... + +class _Chain_rem: + """`rem` — type∈{u16,u32,u64,s16,s32,s64}""" + + s16: _Chain_rem + s32: _Chain_rem + s64: _Chain_rem + u16: _Chain_rem + u32: _Chain_rem + u64: _Chain_rem + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_rsqrt: + """`rsqrt` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32,f64}""" + + approx: _Chain_rsqrt + f32: _Chain_rsqrt + f64: _Chain_rsqrt + ftz: _Chain_rsqrt + def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + +class _Chain_sad: + """`sad` — type∈{u16,u32,u64,s16,s32,s64}""" + + s16: _Chain_sad + s32: _Chain_sad + s64: _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: ... + +class _Chain_selp: + """`selp` — type∈{b16,b32,b64,u16,u32,u64,s16,s32,s64,f32,f64}""" + + b16: _Chain_selp + b32: _Chain_selp + b64: _Chain_selp + f32: _Chain_selp + f64: _Chain_selp + s16: _Chain_selp + s32: _Chain_selp + s64: _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: ... + +class _Chain_set: + """`set` — 4 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, a, b); (d, a, b, c) + """ + + and_: _Chain_set + b16: _Chain_set + b32: _Chain_set + b64: _Chain_set + bf16: _Chain_set + bf16x2: _Chain_set + eq: _Chain_set + equ: _Chain_set + f16: _Chain_set + f16x2: _Chain_set + f32: _Chain_set + f64: _Chain_set + ftz: _Chain_set + ge: _Chain_set + geu: _Chain_set + gt: _Chain_set + gtu: _Chain_set + hi: _Chain_set + hs: _Chain_set + le: _Chain_set + leu: _Chain_set + lo: _Chain_set + ls: _Chain_set + lt: _Chain_set + ltu: _Chain_set + nan: _Chain_set + ne: _Chain_set + neu: _Chain_set + num: _Chain_set + or_: _Chain_set + s16: _Chain_set + s32: _Chain_set + s64: _Chain_set + u16: _Chain_set + u32: _Chain_set + u64: _Chain_set + xor: _Chain_set + def __call__(self, *args: Any) -> None: ... class _Chain_setmaxnreg: """`setmaxnreg` — action∈{inc,dec}; sync∈{sync}; aligned∈{aligned}; type∈{u32}""" @@ -652,6 +1333,143 @@ class _Chain_setmaxnreg: u32: _Chain_setmaxnreg def __call__(self, nreg: Any, *args: Any, pred: Any = None) -> None: ... +class _Chain_setp: + """`setp` — 8 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (p, a, b); (p, q, a, b); (p, a, b, c); (p, q, a, b, c) + """ + + and_: _Chain_setp + b16: _Chain_setp + b32: _Chain_setp + b64: _Chain_setp + bf16: _Chain_setp + bf16x2: _Chain_setp + eq: _Chain_setp + equ: _Chain_setp + f16: _Chain_setp + f16x2: _Chain_setp + f32: _Chain_setp + f64: _Chain_setp + ftz: _Chain_setp + ge: _Chain_setp + geu: _Chain_setp + gt: _Chain_setp + gtu: _Chain_setp + hi: _Chain_setp + hs: _Chain_setp + le: _Chain_setp + leu: _Chain_setp + lo: _Chain_setp + ls: _Chain_setp + lt: _Chain_setp + ltu: _Chain_setp + nan: _Chain_setp + ne: _Chain_setp + neu: _Chain_setp + num: _Chain_setp + or_: _Chain_setp + s16: _Chain_setp + s32: _Chain_setp + s64: _Chain_setp + u16: _Chain_setp + u32: _Chain_setp + u64: _Chain_setp + xor: _Chain_setp + def __call__(self, *args: Any) -> None: ... + +class _Chain_shf: + """`shf` — dir∈{l,r}; mode∈{clamp,wrap}; type∈{b32}""" + + b32: _Chain_shf + clamp: _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: ... + +class _Chain_shfl_sync: + """`shfl_sync` — 2 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (d, a, b, c, membermask); (d, p, a, b, c, + membermask) + """ + + b32: _Chain_shfl_sync + bfly: _Chain_shfl_sync + down: _Chain_shfl_sync + idx: _Chain_shfl_sync + up: _Chain_shfl_sync + def __call__(self, *args: Any) -> None: ... + +class _Chain_shl: + """`shl` — type∈{b16,b32,b64}""" + + b16: _Chain_shl + b32: _Chain_shl + b64: _Chain_shl + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_shr: + """`shr` — type∈{b16,b32,b64,u16,u32,u64,s16,s32,s64}""" + + b16: _Chain_shr + b32: _Chain_shr + b64: _Chain_shr + s16: _Chain_shr + s32: _Chain_shr + s64: _Chain_shr + u16: _Chain_shr + u32: _Chain_shr + u64: _Chain_shr + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_sin: + """`sin` — mode∈{approx}; ftz∈{ftz} (opt); type∈{f32}""" + + approx: _Chain_sin + f32: _Chain_sin + ftz: _Chain_sin + def __call__(self, d: Any, value: Any, *args: Any) -> 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. + """ + + b16: _Chain_slct + b32: _Chain_slct + b64: _Chain_slct + f32: _Chain_slct + f64: _Chain_slct + ftz: _Chain_slct + s16: _Chain_slct + s32: _Chain_slct + s64: _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: ... + +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. + """ + + approx: _Chain_sqrt + f32: _Chain_sqrt + f64: _Chain_sqrt + ftz: _Chain_sqrt + rm: _Chain_sqrt + rn: _Chain_sqrt + rp: _Chain_sqrt + rz: _Chain_sqrt + def __call__(self, d: Any, value: Any, *args: Any) -> None: ... + class _Chain_st: """`st` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (*__operands) @@ -704,6 +1522,35 @@ class _Chain_st: wt: _Chain_st def __call__(self, *args: Any, pred: Any = None) -> None: ... +class _Chain_st_async: + """`st_async` — 3 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (addr, b, mbar); (*__operands); (addr, b) + """ + + b128: _Chain_st_async + b16: _Chain_st_async + b32: _Chain_st_async + b64: _Chain_st_async + f32: _Chain_st_async + f64: _Chain_st_async + global_: _Chain_st_async + gpu: _Chain_st_async + mbarrier__complete_tx__bytes: _Chain_st_async + mmio: _Chain_st_async + release: _Chain_st_async + s16: _Chain_st_async + s32: _Chain_st_async + s64: _Chain_st_async + shared__cluster: _Chain_st_async + sys: _Chain_st_async + u16: _Chain_st_async + u32: _Chain_st_async + u64: _Chain_st_async + v2: _Chain_st_async + v4: _Chain_st_async + weak: _Chain_st_async + def __call__(self, *args: Any, pred: Any = None) -> None: ... + class _Chain_st_bulk: """`st_bulk` — weak∈{weak} (opt); space∈{shared::cta} (opt)""" @@ -731,7 +1578,7 @@ class _Chain_stmatrix: def __call__(self, *args: Any, pred: Any = None) -> None: ... class _Chain_sub: - """`sub` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + """`sub` — 3 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (d, a, b) """ @@ -747,7 +1594,35 @@ class _Chain_sub: rn: _Chain_sub rp: _Chain_sub rz: _Chain_sub + s16: _Chain_sub + s32: _Chain_sub + s64: _Chain_sub sat: _Chain_sub + u16: _Chain_sub + u32: _Chain_sub + u64: _Chain_sub + def __call__(self, *args: Any) -> None: ... + +class _Chain_szext: + """`szext` — mode∈{clamp,wrap}; type∈{u32,s32}""" + + clamp: _Chain_szext + s32: _Chain_szext + u32: _Chain_szext + wrap: _Chain_szext + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + +class _Chain_tanh: + """`tanh` — 2 entries sharing this mnemonic; PTX puts their difference in the operand list, + so the call selects one. Shapes: (d, value) + """ + + approx: _Chain_tanh + bf16: _Chain_tanh + bf16x2: _Chain_tanh + f16: _Chain_tanh + f16x2: _Chain_tanh + f32: _Chain_tanh def __call__(self, *args: Any) -> None: ... class _Chain_tcgen05: @@ -811,6 +1686,72 @@ class _Chain_tcgen05: x8: _Chain_tcgen05 def __call__(self, *args: Any) -> None: ... +class _Chain_tensormap_cp_fenceproxy: + """`tensormap_cp_fenceproxy` — dst∈{global}; src∈{shared::cta}; proxy∈{tensormap::generic}; + sem∈{release}; scope∈{cta,cluster,gpu,sys}; sync∈{sync}; aligned∈{aligned} + """ + + aligned: _Chain_tensormap_cp_fenceproxy + cluster: _Chain_tensormap_cp_fenceproxy + cta: _Chain_tensormap_cp_fenceproxy + global_: _Chain_tensormap_cp_fenceproxy + gpu: _Chain_tensormap_cp_fenceproxy + release: _Chain_tensormap_cp_fenceproxy + shared__cta: _Chain_tensormap_cp_fenceproxy + sync: _Chain_tensormap_cp_fenceproxy + sys: _Chain_tensormap_cp_fenceproxy + tensormap__generic: _Chain_tensormap_cp_fenceproxy + def __call__(self, dst_mem: Any, src_mem: Any, *args: Any, pred: Any = None) -> None: ... + +class _Chain_tensormap_replace: + """`tensormap_replace` — 8 entries sharing this mnemonic; PTX puts their difference in the + operand list, so the call selects one. Shapes: (addr, new_val); (addr, ord, new_val) + """ + + b1024: _Chain_tensormap_replace + b32: _Chain_tensormap_replace + b64: _Chain_tensormap_replace + box_dim: _Chain_tensormap_replace + element_stride: _Chain_tensormap_replace + elemtype: _Chain_tensormap_replace + fill_mode: _Chain_tensormap_replace + global_: _Chain_tensormap_replace + global_address: _Chain_tensormap_replace + global_dim: _Chain_tensormap_replace + global_stride: _Chain_tensormap_replace + interleave_layout: _Chain_tensormap_replace + rank: _Chain_tensormap_replace + shared__cta: _Chain_tensormap_replace + swizzle_mode: _Chain_tensormap_replace + tile: _Chain_tensormap_replace + def __call__(self, *args: Any, pred: Any = None) -> None: ... + +class _Chain_testp: + """`testp` — op∈{finite,infinite,number,notanumber,normal,subnormal}; type∈{f32,f64}""" + + f32: _Chain_testp + f64: _Chain_testp + finite: _Chain_testp + infinite: _Chain_testp + normal: _Chain_testp + notanumber: _Chain_testp + number: _Chain_testp + subnormal: _Chain_testp + def __call__(self, p: Any, a: Any, *args: Any) -> None: ... + +class _Chain_vote_sync: + """`vote_sync` — 2 entries sharing this mnemonic; PTX puts their difference in the operand + list, so the call selects one. Shapes: (d, a, membermask) + """ + + all: _Chain_vote_sync + any: _Chain_vote_sync + b32: _Chain_vote_sync + ballot: _Chain_vote_sync + pred: _Chain_vote_sync + uni: _Chain_vote_sync + def __call__(self, *args: Any) -> None: ... + class _Chain_wgmma: """`wgmma` — 19 entries sharing this mnemonic; PTX puts their difference in the operand list, so the call selects one. Shapes: (*__operands); (); (group) @@ -951,40 +1892,109 @@ class _Chain_wgmma: wait_group: _Chain_wgmma def __call__(self, *args: Any, pred: Any = None) -> None: ... +class _Chain_xor: + """`xor` — type∈{pred,b16,b32,b64}""" + + b16: _Chain_xor + b32: _Chain_xor + b64: _Chain_xor + pred: _Chain_xor + def __call__(self, d: Any, a: Any, b: Any, *args: Any) -> None: ... + class _PTX: + abs: _Chain_abs + activemask: _Chain_activemask add: _Chain_add + and_: _Chain_and + applypriority: _Chain_applypriority atom: _Chain_atom bar: _Chain_bar barrier: _Chain_barrier + bfe: _Chain_bfe + bfi: _Chain_bfi + bfind: _Chain_bfind + bmsk: _Chain_bmsk + brev: _Chain_brev clusterlaunchcontrol: _Chain_clusterlaunchcontrol + clz: _Chain_clz + cnot: _Chain_cnot + copysign: _Chain_copysign + cos: _Chain_cos cp: _Chain_cp + createpolicy: _Chain_createpolicy cvt: _Chain_cvt + cvt_pack: _Chain_cvt_pack cvta: _Chain_cvta + discard: _Chain_discard + div: _Chain_div + dp2a: _Chain_dp2a + dp4a: _Chain_dp4a + elect_sync: _Chain_elect_sync ex2: _Chain_ex2 fence: _Chain_fence fma: _Chain_fma fns: _Chain_fns + getctarank: _Chain_getctarank griddepcontrol: _Chain_griddepcontrol + isspacep: _Chain_isspacep ld: _Chain_ld ldmatrix: _Chain_ldmatrix + ldu: _Chain_ldu + lg2: _Chain_lg2 + lop3: _Chain_lop3 + mad: _Chain_mad + mad24: _Chain_mad24 mapa: _Chain_mapa + match: _Chain_match max: _Chain_max mbarrier: _Chain_mbarrier min: _Chain_min mma: _Chain_mma mov: _Chain_mov mul: _Chain_mul + mul24: _Chain_mul24 + multimem_ld_reduce: _Chain_multimem_ld_reduce + multimem_red: _Chain_multimem_red + multimem_st: _Chain_multimem_st neg: _Chain_neg + not_: _Chain_not + or_: _Chain_or + popc: _Chain_popc prefetch: _Chain_prefetch + prefetchu: _Chain_prefetchu + prmt: _Chain_prmt rcp: _Chain_rcp red: _Chain_red + red_async: _Chain_red_async + redux_sync: _Chain_redux_sync + rem: _Chain_rem + rsqrt: _Chain_rsqrt + sad: _Chain_sad + selp: _Chain_selp + set: _Chain_set setmaxnreg: _Chain_setmaxnreg + setp: _Chain_setp + shf: _Chain_shf + shfl_sync: _Chain_shfl_sync + shl: _Chain_shl + shr: _Chain_shr + sin: _Chain_sin + slct: _Chain_slct + sqrt: _Chain_sqrt st: _Chain_st + st_async: _Chain_st_async st_bulk: _Chain_st_bulk stmatrix: _Chain_stmatrix sub: _Chain_sub + szext: _Chain_szext + tanh: _Chain_tanh tcgen05: _Chain_tcgen05 + tensormap_cp_fenceproxy: _Chain_tensormap_cp_fenceproxy + tensormap_replace: _Chain_tensormap_replace + testp: _Chain_testp + vote_sync: _Chain_vote_sync wgmma: _Chain_wgmma + xor: _Chain_xor def __getitem__(self, text: str) -> Any: ... ptx: _PTX diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index 7cbae243f650..bf2005befa7a 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -609,94 +609,120 @@ def elected(): ) -def scope_id(extents: list[Expr | int] | None, parent: str, cur: str) -> Var | list[Var]: - ret = _ffi_api.ScopeId(extents, parent, "T.scope_id", cur) # type: ignore[attr-defined] # pylint: disable=no-member +def scope_id( + extents: list[Expr | int] | None, parent: str, cur: str, dtype: str = "int32" +) -> Var | list[Var]: + ret = _ffi_api.ScopeId(extents, parent, "T.scope_id", cur, dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def cluster_id(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def cluster_id(extents: list[Expr | int] | None = None, dtype: str = "int32") -> Var | list[Var]: """Define a kernel→cluster scope id. Pass ``None`` (the default) to defer the - extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.""" - ret = _ffi_api.ClusterId(extents, "kernel") # type: ignore[attr-defined] # pylint: disable=no-member + extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.ClusterId(extents, "kernel", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def cta_id(extents: list[Expr | int] | None = None, preferred=None) -> Var | list[Var]: +def cta_id( + extents: list[Expr | int] | None = None, preferred=None, dtype: str = "int32" +) -> Var | list[Var]: """Define a kernel→cta scope id. Pass ``None`` (the default) to defer the - extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.""" - ret = _ffi_api.CtaId(extents, "kernel", preferred) # type: ignore[attr-defined] # pylint: disable=no-member + extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.CtaId(extents, "kernel", preferred, dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def cta_id_in_cluster(extents: list[Expr | int] | None = None, preferred=None) -> Var | list[Var]: +def cta_id_in_cluster( + extents: list[Expr | int] | None = None, preferred=None, dtype: str = "int32" +) -> Var | list[Var]: """Define a cluster→cta scope id. Pass ``None`` (the default) to defer the - extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure.""" - ret = _ffi_api.CtaId(extents, "cluster", preferred) # type: ignore[attr-defined] # pylint: disable=no-member + extent; it will be inferred at LowerTIRx from sibling ScopeIdDef closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.CtaId(extents, "cluster", preferred, dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def cta_id_in_pair() -> Var: - ret = _ffi_api.CtaIdInPair() # type: ignore[attr-defined] # pylint: disable=no-member +def cta_id_in_pair(dtype: str = "int32") -> Var: + ret = _ffi_api.CtaIdInPair(dtype) # type: ignore[attr-defined] # pylint: disable=no-member return ret[0] -def warpgroup_id(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def warpgroup_id(extents: list[Expr | int] | None = None, dtype: str = "int32") -> Var | list[Var]: """Define a cta→warpgroup scope id. Pass ``None`` (the default) to defer - the extent; it will be inferred at LowerTIRx from sibling closure.""" - ret = _ffi_api.WarpgroupId(extents, "cta") # type: ignore[attr-defined] # pylint: disable=no-member + the extent; it will be inferred at LowerTIRx from sibling closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.WarpgroupId(extents, "cta", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def warp_id(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def warp_id(extents: list[Expr | int] | None = None, dtype: str = "int32") -> Var | list[Var]: """Define a cta→warp scope id. Pass ``None`` (the default) to defer the - extent; it will be inferred at LowerTIRx from sibling closure.""" - ret = _ffi_api.WarpId(extents, "cta") # type: ignore[attr-defined] # pylint: disable=no-member + extent; it will be inferred at LowerTIRx from sibling closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.WarpId(extents, "cta", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def warp_id_in_wg(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def warp_id_in_wg(extents: list[Expr | int] | None = None, dtype: str = "int32") -> Var | list[Var]: """Define a warpgroup→warp scope id. Pass ``None`` (the default) to defer - the extent; it will be inferred at LowerTIRx from sibling closure.""" - ret = _ffi_api.WarpId(extents, "warpgroup") # type: ignore[attr-defined] # pylint: disable=no-member + the extent; it will be inferred at LowerTIRx from sibling closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.WarpId(extents, "warpgroup", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def lane_id(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def lane_id(extents: list[Expr | int] | None = None, dtype: str = "int32") -> Var | list[Var]: """Define a warp→thread scope id. Pass ``None`` (the default) to defer the - extent; it will be inferred at LowerTIRx from sibling closure.""" - ret = _ffi_api.ThreadId(extents, "warp") # type: ignore[attr-defined] # pylint: disable=no-member + extent; it will be inferred at LowerTIRx from sibling closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.ThreadId(extents, "warp", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def thread_id(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def thread_id(extents: list[Expr | int] | None = None, dtype: str = "int32") -> Var | list[Var]: """Define a cta→thread scope id. Pass ``None`` (the default) to defer the - extent; it will be inferred at LowerTIRx from sibling closure.""" - ret = _ffi_api.ThreadId(extents, "cta") # type: ignore[attr-defined] # pylint: disable=no-member + extent; it will be inferred at LowerTIRx from sibling closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.ThreadId(extents, "cta", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret -def thread_id_in_wg(extents: list[Expr | int] | None = None) -> Var | list[Var]: +def thread_id_in_wg( + extents: list[Expr | int] | None = None, dtype: str = "int32" +) -> Var | list[Var]: """Define a warpgroup→thread scope id. Pass ``None`` (the default) to defer - the extent; it will be inferred at LowerTIRx from sibling closure.""" - ret = _ffi_api.ThreadId(extents, "warpgroup") # type: ignore[attr-defined] # pylint: disable=no-member + the extent; it will be inferred at LowerTIRx from sibling closure. + + ``dtype`` selects the dtype of the introduced vars (``"int32"`` or ``"uint32"``).""" + ret = _ffi_api.ThreadId(extents, "warpgroup", dtype) # type: ignore[attr-defined] # pylint: disable=no-member if len(ret) == 1: return ret[0] return ret @@ -1148,7 +1174,8 @@ def serial( *, annotations: dict[str, Any] | None = None, step: Expr | None = None, - unroll: bool | None = None, + unroll: bool | int | None = None, + dtype: str | None = None, ) -> frame.ForFrame: """The serial For statement. @@ -1166,11 +1193,18 @@ def serial( step : Expr The optional step value of iteration. - unroll : bool, optional + unroll : bool or int, optional If True, adds ``{"pragma_unroll": True}`` annotation, which asks CUDA codegen to emit ``#pragma unroll`` while preserving the loop as a C++ ``for``. If False, adds ``{"disable_unroll": True}`` annotation. - Shorthand for ``annotations={"disable_unroll": True}``. + If a positive integer, emits ``#pragma unroll N``. Boolean values are + handled separately from integers, so ``False`` keeps disabling unrolling. + + dtype : str, optional + The dtype of the loop variable, either ``"int32"`` or ``"uint32"``. When + omitted it is inferred from the bounds. Bounds that do not already have this + dtype are converted (literals are retyped, other expressions get a Cast). + Note ``T.thread_binding`` does not support this; its loop var is always int32. Returns ------- @@ -1179,17 +1213,24 @@ def serial( """ if unroll is not None: annotations = dict(annotations) if annotations else {} - if unroll: - annotations["pragma_unroll"] = True + if isinstance(unroll, bool): + if unroll: + annotations["pragma_unroll"] = True + else: + annotations["disable_unroll"] = True + elif isinstance(unroll, int): + if unroll < 1: + raise ValueError("unroll must be a positive integer") + annotations["pragma_unroll"] = unroll else: - annotations["disable_unroll"] = True + raise TypeError("unroll must be a bool, a positive integer, or None") if stop is None: stop = start if is_prim_expr(start): start = IntImm(start.ty, 0) else: start = 0 - return _ffi_api.Serial(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member + return _ffi_api.Serial(start, stop, annotations, step, dtype) # type: ignore[attr-defined] # pylint: disable=no-member def parallel( @@ -1198,6 +1239,7 @@ def parallel( *, annotations: dict[str, Any] | None = None, step: Expr | None = None, + dtype: str | None = None, ) -> frame.ForFrame: """The parallel For statement. @@ -1215,6 +1257,10 @@ def parallel( step : Expr The optional step value of iteration. + dtype : str, optional + The dtype of the loop variable, either ``"int32"`` or ``"uint32"``. When + omitted it is inferred from the bounds. + Returns ------- res : frame.ForFrame @@ -1226,7 +1272,7 @@ def parallel( start = IntImm(start.ty, 0) else: start = 0 - return _ffi_api.Parallel(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member + return _ffi_api.Parallel(start, stop, annotations, step, dtype) # type: ignore[attr-defined] # pylint: disable=no-member def vectorized( @@ -1235,6 +1281,7 @@ def vectorized( *, annotations: dict[str, Any] | None = None, step: Expr | None = None, + dtype: str | None = None, ) -> frame.ForFrame: """The vectorized For statement. @@ -1252,6 +1299,10 @@ def vectorized( step : Expr The optional step value of iteration. + dtype : str, optional + The dtype of the loop variable, either ``"int32"`` or ``"uint32"``. When + omitted it is inferred from the bounds. + Returns ------- res : frame.ForFrame @@ -1263,7 +1314,7 @@ def vectorized( start = IntImm(start.ty, 0) else: start = 0 - return _ffi_api.Vectorized(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member + return _ffi_api.Vectorized(start, stop, annotations, step, dtype) # type: ignore[attr-defined] # pylint: disable=no-member def unroll( @@ -1272,6 +1323,7 @@ def unroll( *, annotations: dict[str, Any] | None = None, step: Expr | None = None, + dtype: str | None = None, ) -> frame.ForFrame: """The unrolled For statement. @@ -1289,6 +1341,10 @@ def unroll( step : Expr The optional step value of iteration. + dtype : str, optional + The dtype of the loop variable, either ``"int32"`` or ``"uint32"``. When + omitted it is inferred from the bounds. + Returns ------- res : frame.ForFrame @@ -1300,7 +1356,7 @@ def unroll( start = IntImm(start.ty, 0) else: start = 0 - return _ffi_api.Unroll(start, stop, annotations, step) # type: ignore[attr-defined] # pylint: disable=no-member + return _ffi_api.Unroll(start, stop, annotations, step, dtype) # type: ignore[attr-defined] # pylint: disable=no-member def thread_binding( @@ -1351,7 +1407,7 @@ def thread_binding( ) -def grid(*extents: tuple[Expr | tuple[Expr, Expr]]) -> frame.ForFrame: +def grid(*extents: tuple[Expr | tuple[Expr, Expr]], dtype: str | None = None) -> frame.ForFrame: """The grid For statement. Parameters @@ -1361,6 +1417,10 @@ def grid(*extents: tuple[Expr | tuple[Expr, Expr]]) -> frame.ForFrame: If a tuple of two Expr is provided, the first is the start of the iteration, and the second is the extent of the iteration. + dtype : str, optional + The dtype of every loop variable, either ``"int32"`` or ``"uint32"``. When + omitted each loop variable takes the dtype of its own extent. + Returns ------- res : frame.ForFrame @@ -1368,17 +1428,20 @@ def grid(*extents: tuple[Expr | tuple[Expr, Expr]]) -> frame.ForFrame: """ # Convert integer extents to IntImm # TODO(@bohan): fix this after FFI refactor + imm_dtype = dtype if dtype is not None else "int32" processed_extents = [] for extent in extents: if isinstance(extent, tuple): start, extent = extent - start = IntImm("int32", start) if isinstance(start, int) else start - extent = IntImm("int32", extent) if isinstance(extent, int) else extent + start = IntImm(imm_dtype, start) if isinstance(start, int) else start + extent = IntImm(imm_dtype, extent) if isinstance(extent, int) else extent processed_extents.append((start, extent)) else: - processed_extents.append(IntImm("int32", extent) if isinstance(extent, int) else extent) + processed_extents.append( + IntImm(imm_dtype, extent) if isinstance(extent, int) else extent + ) extents = tuple(processed_extents) - return _ffi_api.Grid(extents) # type: ignore[attr-defined] # pylint: disable=no-member + return _ffi_api.Grid(extents, dtype) # type: ignore[attr-defined] # pylint: disable=no-member def Assert(condition: Expr, message, error_kind: str = "RuntimeError") -> frame.AssertFrame: # pylint: disable=invalid-name diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index 88b3c99f8379..8b6f2fa2d53c 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/src/backend/cuda/codegen/codegen_cuda.cc @@ -311,9 +311,18 @@ void CodeGenCUDA::VisitStmt_(const tirx::ForNode* op) { if (op->annotations.count("disable_unroll")) { PrintIndent(); stream << "#pragma unroll 1\n"; - } else if (op->kind == tirx::ForKind::kUnrolled || op->annotations.count("pragma_unroll")) { + } else if (op->kind == tirx::ForKind::kUnrolled) { PrintIndent(); stream << "#pragma unroll\n"; + } else if (auto it = op->annotations.find("pragma_unroll"); it != op->annotations.end()) { + PrintIndent(); + stream << "#pragma unroll"; + if (auto count = (*it).second.as()) { + stream << " " << count.value(); + } else if (const auto* count = (*it).second.as()) { + stream << " " << count->value; + } + stream << "\n"; } PrintIndent(); std::string vid = AllocVarID(op->loop_var.get()); @@ -1605,7 +1614,11 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) { return; } else if (op->attr_key == "pragma_unroll") { PrintIndent(); - stream << "#pragma unroll\n"; + stream << "#pragma unroll"; + if (const auto* count = op->value.as(); count && count->value != 1) { + stream << " " << count->value; + } + stream << "\n"; this->VisitStmt(op->body); return; } else if (op->attr_key == tirx::attr::thread_extent) { diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 817da4efedad..a35791c60af3 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -121,6 +121,21 @@ TVM_FFI_STATIC_INIT_BLOCK() { }); } +namespace { +/*! + * \brief Whether an integer literal can be represented exactly by `ty`. + * \note Mirrors the range checks performed by the IntImm constructor. + */ +bool IntImmValueFits(int64_t value, const PrimType& ty) { + int bits = ty.bits(); + if (ty.MatchesCode(DLDataTypeCode::kDLUInt)) { + return value >= 0 && (bits >= 64 || value < (int64_t{1} << bits)); + } + if (bits >= 64) return true; + return value >= -(int64_t{1} << (bits - 1)) && value < (int64_t{1} << (bits - 1)); +} +} // namespace + // For For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt body, ffi::Optional thread_binding, ffi::Map annotations, @@ -141,20 +156,23 @@ For::For(PrimVar loop_var, PrimExpr min, PrimExpr extent, ForKind kind, Stmt bod require_scalar_int_dtype(min, "min"); require_scalar_int_dtype(extent, "extent"); - // When extent, min or step is an IntImm but has narrower dtype than loop_var - // we directly promote them without raising errors. + // When extent, min or step is an IntImm whose dtype differs from loop_var's + // (narrower bits and/or a different signedness code), we directly promote it + // to the loop var's dtype as long as the value stays representable. auto try_promote_imm_dtype = [&](const PrimExpr& e) -> PrimExpr { PrimType e_ty = e.ty(); PrimType loop_var_ty = loop_var.ty(); + if (e_ty == loop_var_ty) return e; + if (const IntImmNode* a = e.as()) { + TVM_FFI_ICHECK(IntImmValueFits(a->value, loop_var_ty)) + << "Literal value " << a->value << " is not representable in the loop variable's dtype (" + << loop_var_ty << ")"; + return IntImm(loop_var_ty, a->value); + } TVM_FFI_ICHECK(e_ty.bits() <= loop_var_ty.bits()) << " Loop variable's dtype (" << loop_var_ty << ") is narrower than that of `min` or `extent` (" << e_ty << ")"; - const IntImmNode* a = e.as(); - if (a && e_ty.bits() < loop_var_ty.bits()) { - return IntImm(loop_var_ty, a->value); - } else { - return e; - } + return e; }; min = try_promote_imm_dtype(min); diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 1538ec6836ed..f4a6a81f8455 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -184,8 +184,19 @@ SBlockFrame Block(ffi::String name, bool no_realize, ffi::String exec_scope) { void TilePrimitiveCall(tvm::tirx::TilePrimitiveCall op_call) { AddToParent(op_call); } +/*! + * \brief Validate a user-requested loop / scope-id var dtype. + * \note Only scalar int32 and uint32 are supported. + */ +void CheckExplicitIndexDtype(const PrimType& dtype) { + TVM_FFI_ICHECK(dtype.IsScalar() && dtype.bits() == 32 && + dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) + << "ValueError: dtype of a loop/scope-id var must be \"int32\" or \"uint32\", got " << dtype; +} + ffi::Array ScopeId(ffi::Optional> extents, ffi::String parent, - ffi::String name, ffi::String cur) { + ffi::String name, ffi::String cur, PrimType dtype) { + CheckExplicitIndexDtype(dtype); // Determine the number of Vars to introduce. Deferred form (extents=None) // is always 1-axis; the verifier closure fills the extent at LowerTIRx. size_t n_vars = extents.has_value() ? extents.value().size() : 1; @@ -195,7 +206,7 @@ ffi::Array ScopeId(ffi::Optional> extents, } ffi::Array scope_ids; for (size_t i = 0; i < n_vars; ++i) { - scope_ids.push_back(tvm::tirx::PrimVar("")); + scope_ids.push_back(tvm::tirx::PrimVar("", dtype)); } // Emit a standalone ScopeIdDefStmt to the current TIRFrame's stmts list. // The def is visible to all subsequent stmts within the same enclosing @@ -208,13 +219,14 @@ ffi::Array ScopeId(ffi::Optional> extents, } ffi::Array ClusterId(ffi::Optional> extents, - ffi::String parent) { - return ScopeId(extents, parent, "T.cluster_id", "cluster"); + ffi::String parent, PrimType dtype) { + return ScopeId(extents, parent, "T.cluster_id", "cluster", dtype); } ffi::Array CtaId(ffi::Optional> extents, ffi::String parent, - ffi::Optional> preferred) { + ffi::Optional> preferred, PrimType dtype) { if (preferred.has_value()) { + CheckExplicitIndexDtype(dtype); TVM_FFI_ICHECK(parent == "cluster") << "ValueError: preferred is only valid when parent=\"cluster\", got parent=\"" << parent << "\""; @@ -222,7 +234,7 @@ ffi::Array CtaId(ffi::Optional> extents, ff << "ValueError: preferred=... requires explicit extents (deferred form is incompatible)"; ffi::Array scope_ids; for (size_t i = 0; i < extents.value().size(); ++i) { - scope_ids.push_back(tvm::tirx::PrimVar("")); + scope_ids.push_back(tvm::tirx::PrimVar("", dtype)); } tvm::tirx::ScopeIdDef def( scope_ids.Map([](tvm::tirx::Var var) { return var.as_or_throw(); }), @@ -230,11 +242,12 @@ ffi::Array CtaId(ffi::Optional> extents, ff AddToParent(tvm::tirx::ScopeIdDefStmt(def)); return scope_ids; } - return ScopeId(extents, parent, "T.cta_id", "cta"); + return ScopeId(extents, parent, "T.cta_id", "cta", dtype); } -ffi::Array CtaIdInPair() { - ffi::Array scope_ids{tvm::tirx::PrimVar("")}; +ffi::Array CtaIdInPair(PrimType dtype) { + CheckExplicitIndexDtype(dtype); + ffi::Array scope_ids{tvm::tirx::PrimVar("", dtype)}; tvm::tirx::ScopeIdDef def( scope_ids.Map([](tvm::tirx::Var var) { return var.as_or_throw(); }), ffi::Array{IntImm::Int32(2)}, tvm::tirx::ScopeBinding::kClusterCtaPair); @@ -243,17 +256,18 @@ ffi::Array CtaIdInPair() { } ffi::Array WarpgroupId(ffi::Optional> extents, - ffi::String parent) { - return ScopeId(extents, parent, "T.warpgroup_id", "warpgroup"); + ffi::String parent, PrimType dtype) { + return ScopeId(extents, parent, "T.warpgroup_id", "warpgroup", dtype); } -ffi::Array WarpId(ffi::Optional> extents, ffi::String parent) { - return ScopeId(extents, parent, "T.warp_id", "warp"); +ffi::Array WarpId(ffi::Optional> extents, ffi::String parent, + PrimType dtype) { + return ScopeId(extents, parent, "T.warp_id", "warp", dtype); } -ffi::Array ThreadId(ffi::Optional> extents, - ffi::String parent) { - return ScopeId(extents, parent, "T.thread_id", "thread"); +ffi::Array ThreadId(ffi::Optional> extents, ffi::String parent, + PrimType dtype) { + return ScopeId(extents, parent, "T.thread_id", "thread", dtype); } BlockInitFrame Init() { return BlockInitFrame(ffi::make_object()); } @@ -476,17 +490,54 @@ ffi::Array Remap(ffi::String kinds, ffi::Array bindings, PrimType } // namespace axis +/*! + * \brief Determine the dtype of a loop var from its bounds, or validate an explicit one. + * + * Without an explicit dtype the bit width is the max of both bounds and the + * signedness follows C-style promotion (unsigned wins), which is exactly what + * `PromoteBinaryOpType` applies to the `stop - start` extent. + */ +PrimType InferLoopVarDtype(const PrimExpr& start, const PrimExpr& stop, + const ffi::Optional& dtype) { + if (dtype.has_value()) { + CheckExplicitIndexDtype(dtype.value()); + return dtype.value(); + } + PrimType start_ty = start.ty(); + PrimType stop_ty = stop.ty(); + bool is_unsigned = + start_ty.MatchesCode(DLDataTypeCode::kDLUInt) || stop_ty.MatchesCode(DLDataTypeCode::kDLUInt); + return PrimType(is_unsigned ? DLDataTypeCode::kDLUInt : DLDataTypeCode::kDLInt, + std::max(start_ty.bits(), stop_ty.bits()), 1); +} + +/*! + * \brief Coerce a loop bound to the loop var's dtype. + * + * Integer literals are re-created in the target dtype; any other mismatched + * expression gets an explicit Cast, so the loop header stays the single place + * where the index dtype has to be spelled out. + */ +PrimExpr ConvertLoopBound(const PrimExpr& e, const PrimType& var_ty) { + if (e.ty() == var_ty) return e; + if (const auto* imm = e.as()) { + return tvm::IntImm(var_ty, imm->value); + } + return tvm::tirx::Cast(var_ty, e); +} + #define TVM_TIRX_IR_BUILDER_FOR_FRAME(Method, Kind) \ ForFrame Method(PrimExpr start, PrimExpr stop, \ ffi::Optional> annotations, \ - ffi::Optional step) { \ - PrimExpr min = start; \ - PrimExpr extent = arith::Analyzer()->Simplify(stop - start); \ + ffi::Optional step, ffi::Optional dtype) { \ + PrimType var_ty = InferLoopVarDtype(start, stop, dtype); \ + PrimExpr min = ConvertLoopBound(start, var_ty); \ + PrimExpr extent = arith::Analyzer()->Simplify(ConvertLoopBound(stop, var_ty) - min); \ + if (step.has_value()) { \ + step = ConvertLoopBound(step.value(), var_ty); \ + } \ ffi::ObjectPtr n = ffi::make_object(); \ - PrimType min_ty = min.ty(); \ - PrimType extent_ty = extent.ty(); \ - int bits = std::max(min_ty.bits(), extent_ty.bits()); \ - n->vars = {Var("v", min_ty.WithBits(bits).WithLanes(1))}; \ + n->vars = {Var("v", var_ty)}; \ n->doms = {Range::FromMinExtent(min, extent)}; \ n->steps = {step}; \ n->f_make_for_loop = [annotations](ffi::Array vars, ffi::Array doms, \ @@ -537,8 +588,12 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread, return ForFrame(n); } -ForFrame Grid(ffi::Array>> extents) { +ForFrame Grid(ffi::Array>> extents, + ffi::Optional dtype) { using namespace tvm::tirx; + if (dtype.has_value()) { + CheckExplicitIndexDtype(dtype.value()); + } ffi::ObjectPtr n = ffi::make_object(); n->vars.reserve(extents.size()); n->doms.reserve(extents.size()); @@ -546,14 +601,15 @@ ForFrame Grid(ffi::Array>> for (const auto& extent : extents) { if (auto prim_expr = extent.as()) { // extent is a single PrimExpr - PrimType dtype = prim_expr.value().ty(); - n->vars.push_back(Var("v", dtype)); - n->doms.push_back(Range(tvm::IntImm(dtype, 0), prim_expr.value())); + PrimType var_ty = dtype.value_or(prim_expr.value().ty()); + n->vars.push_back(Var("v", var_ty)); + n->doms.push_back(Range(tvm::IntImm(var_ty, 0), ConvertLoopBound(prim_expr.value(), var_ty))); } else if (auto tuple = extent.as>()) { // extent is a tuple of two PrimExpr (start, extent) - PrimType dtype = tuple.value().get<0>().ty(); - n->vars.push_back(Var("v", dtype)); - n->doms.push_back(Range::FromMinExtent(tuple.value().get<0>(), tuple.value().get<1>())); + PrimType var_ty = dtype.value_or(tuple.value().get<0>().ty()); + n->vars.push_back(Var("v", var_ty)); + n->doms.push_back(Range::FromMinExtent(ConvertLoopBound(tuple.value().get<0>(), var_ty), + ConvertLoopBound(tuple.value().get<1>(), var_ty))); } else { TVM_FFI_THROW(InternalError) << "TypeError: Invalid type for grid extent"; } @@ -915,30 +971,30 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("script.ir_builder.tirx.Block", Block) .def("script.ir_builder.tirx.TilePrimitiveCall", TilePrimitiveCall) .def("script.ir_builder.tirx.ClusterId", - [](ffi::Optional> extents, ffi::String parent) { - return ClusterId(extents, parent); + [](ffi::Optional> extents, ffi::String parent, PrimType dtype) { + return ClusterId(extents, parent, dtype); }) .def("script.ir_builder.tirx.CtaId", [](ffi::Optional> extents, ffi::String parent, - ffi::Optional> preferred) { - return CtaId(extents, parent, preferred); - }) + ffi::Optional> preferred, + PrimType dtype) { return CtaId(extents, parent, preferred, dtype); }) .def("script.ir_builder.tirx.CtaIdInPair", CtaIdInPair) .def("script.ir_builder.tirx.WarpgroupId", - [](ffi::Optional> extents, ffi::String parent) { - return WarpgroupId(extents, parent); + [](ffi::Optional> extents, ffi::String parent, PrimType dtype) { + return WarpgroupId(extents, parent, dtype); }) .def("script.ir_builder.tirx.WarpId", - [](ffi::Optional> extents, ffi::String parent) { - return WarpId(extents, parent); + [](ffi::Optional> extents, ffi::String parent, PrimType dtype) { + return WarpId(extents, parent, dtype); }) .def("script.ir_builder.tirx.ThreadId", - [](ffi::Optional> extents, ffi::String parent) { - return ThreadId(extents, parent); + [](ffi::Optional> extents, ffi::String parent, PrimType dtype) { + return ThreadId(extents, parent, dtype); }) .def("script.ir_builder.tirx.ScopeId", [](ffi::Optional> extents, ffi::String parent, ffi::String name, - ffi::String cur) { return ScopeId(extents, parent, name, cur); }) + ffi::String cur, + PrimType dtype) { return ScopeId(extents, parent, name, cur, dtype); }) .def("script.ir_builder.tirx.Init", Init) .def("script.ir_builder.tirx.Where", Where) .def("script.ir_builder.tirx.Reads", Reads) diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc index 79271f10629e..93132855d277 100644 --- a/src/tirx/script/printer/block.cc +++ b/src/tirx/script/printer/block.cc @@ -257,6 +257,21 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) kwarg_vals.push_back(d->AsDoc(def->preferred_extents.value(), def_p->Attr("preferred_extents"))); } + // The scope-id dtype is independent of the extents, so it has to be printed + // explicitly whenever it is not the int32 default in order to round-trip. + if (!def->def_ids.empty()) { + PrimType scope_id_ty = def->def_ids[0].ty(); + for (auto scope_id : def->def_ids) { + TVM_FFI_ICHECK(scope_id.ty() == scope_id_ty) + << "mixed scope-id dtypes are unsupported, got " << scope_id_ty << " and " + << scope_id.ty(); + } + if (scope_id_ty != PrimType::Int(32)) { + kwarg_keys.push_back("dtype"); + kwarg_vals.push_back( + LiteralDoc::Str(DType2Str(scope_id_ty->dtype), def_p->Attr("def_ids"))); + } + } ExprDoc rhs = TIR(d, ScopeIdApiName(def->scope))->Call(rhs_args, kwarg_keys, kwarg_vals); return AssignDoc(TupleDoc(lhs), rhs, std::nullopt); }); diff --git a/src/tirx/script/printer/for_loop.cc b/src/tirx/script/printer/for_loop.cc index 456bbfd7249e..363a2d466f0b 100644 --- a/src/tirx/script/printer/for_loop.cc +++ b/src/tirx/script/printer/for_loop.cc @@ -116,7 +116,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) if (annotations.has_value()) { // Check for the special cases: // - annotations == {"disable_unroll": True}: print as unroll=False - // - annotations == {"pragma_unroll": True}: print as unroll=True + // - annotations == {"pragma_unroll": value}: print as unroll=value bool printed_as_unroll = false; if (loop->annotations.size() == 1 && loop->annotations.count("disable_unroll")) { kwargs_keys.push_back("unroll"); @@ -124,7 +124,8 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) printed_as_unroll = true; } else if (loop->annotations.size() == 1 && loop->annotations.count("pragma_unroll")) { kwargs_keys.push_back("unroll"); - kwargs_values.push_back(LiteralDoc::Boolean(true, loop_p->Attr("annotations"))); + kwargs_values.push_back( + d->AsDoc(loop->annotations["pragma_unroll"], loop_p->Attr("annotations"))); printed_as_unroll = true; } if (!printed_as_unroll) { diff --git a/src/tirx/transform/lower_tirx_opaque.cc b/src/tirx/transform/lower_tirx_opaque.cc index 71fed942ec6c..e6ecc9fa7ec4 100644 --- a/src/tirx/transform/lower_tirx_opaque.cc +++ b/src/tirx/transform/lower_tirx_opaque.cc @@ -136,7 +136,8 @@ class TIRxOpaqueLower : public StmtExprMutator { /*! * \brief Handle loop annotation dict. * (1) if the attr key is prefixed by `pragma_`, move to ordered kv list - * (lowered to `AttrStmt` by legacy TE schedule convention). + * (lowered to `AttrStmt` by legacy TE schedule convention), except for + * `pragma_unroll`, whose bool-or-integer value must remain on the loop. * (2) non-pragma loop annotations are preserved. * \return New annotation dict with preserved keys. Also update pragma attr pairs ordered by key. */ @@ -147,7 +148,9 @@ class TIRxOpaqueLower : public StmtExprMutator { pragma_attrs->clear(); for (const auto& kv : annotations) { const ffi::String& key = kv.first; - if (tirx::attr::IsPragmaKey(key)) { + if (key == "pragma_unroll") { + preserved_annotations.Set(key, kv.second); + } else if (tirx::attr::IsPragmaKey(key)) { pragma_attrs->emplace_back(key, ConvertAttrValue(key, kv.second)); } else { // loop annotations are always preserved (no SBlock annotation dropping here) diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py index 1b817396388a..daaa9f36d676 100644 --- a/tests/python/tirx/codegen/test_codegen_cuda.py +++ b/tests/python/tirx/codegen/test_codegen_cuda.py @@ -239,6 +239,19 @@ def main(A: T.Buffer((4,), "int32")): assert "break;" in src +def test_serial_pragma_unroll_count_codegen(): + @T.prim_func + def main(A: T.Buffer((4,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + for i in T.serial(4, unroll=2): + A[i] = A[i] + 1 + + src, _ = _get_source(main) + assert re.search(r"#pragma unroll 2\s*for \(", src) + + def test_serial_disable_unroll_pragma_immediately_precedes_dynamic_for(): @T.prim_func def main(A: T.Buffer((4,), "int32")): @@ -598,7 +611,7 @@ def main(A: T.Buffer((1,), "uint32")): assert "cp.async.mbarrier.arrive.noinc.shared::cta.b64" in src # The spin-wait moved to T.cuda.mbarrier_wait, which takes its timeout as a # parameter rather than baking 10000000 into the asm text. - assert "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2;" in src + assert "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;" in src assert "tvm_builtin_cuda_mbarrier_wait" in src # (No assertion on the timeout local: T.cuda.mbarrier_wait keeps the # `ticks = 0x989680` hint the TIRx spin-wait convention specifies.) @@ -1048,5 +1061,50 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +def test_uint32_loop_var_and_scope_id_emit_unsigned(): + @T.prim_func + def main(A: T.Buffer((128,), "int32")): + T.device_entry() + _ = T.cta_id([1]) + tx = T.thread_id([128], dtype="uint32") + for k in T.serial(4, dtype="uint32"): + A[tx] = A[tx] + T.int32(k) + + src, _ = _get_source(main) + # The loop var is declared unsigned and iterates over unsigned bounds. + assert re.search(r"for \(uint k = \(uint\)0; k < \(uint\)4;", src), src + # The scope id is bound as an unsigned value. + assert re.search(r"uint tx = ", src), src + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_uint32_loop_var_runs_correctly(): + @T.prim_func + def main(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")): + T.device_entry() + _ = T.cta_id([1]) + tx = T.thread_id([128], dtype="uint32") + acc = T.alloc_buffer((1,), "int32", scope="local") + acc[0] = 0 + for k in T.serial(4, dtype="uint32"): + acc[0] = acc[0] + A[tx] + T.int32(k) + B[tx] = acc[0] + + _, mod = _get_source(main) + + A_np = np.arange(128, dtype="int32") + B_ref = A_np * 4 + (0 + 1 + 2 + 3) + + def run_and_check(): + dev = tvm.cuda() + A = tvm.runtime.tensor(A_np, device=dev) + B = tvm.runtime.tensor(np.zeros(128, dtype="int32"), device=dev) + mod(A, B) + np.testing.assert_allclose(B.numpy(), B_ref) + + tvm.testing.run_with_gpu_lock(run_and_check) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py index 0852dfae6734..e25c5a00276f 100644 --- a/tests/python/tirx/codegen/test_codegen_hopper.py +++ b/tests/python/tirx/codegen/test_codegen_hopper.py @@ -367,7 +367,9 @@ def func(A: T.Buffer(1)): src, mod = _get_source(func) print(src) - assert "elect.sync %%rx|%%px, %2;" in src + assert "elect.sync _|%%px, %1;" in src + assert ".reg .b32 %%rx;" not in src + assert "mov.s32 %0, %%rx;" not in src @pytest.mark.gpu diff --git a/tests/python/tirx/codegen/test_ptx_cvt.py b/tests/python/tirx/codegen/test_ptx_cvt.py index 6707c0921dbc..eaf868d5e99a 100644 --- a/tests/python/tirx/codegen/test_ptx_cvt.py +++ b/tests/python/tirx/codegen/test_ptx_cvt.py @@ -235,7 +235,11 @@ # The generic scalar line is one entry named plain "cvt"; the packed lines are # the "cvt_*" family. -_CVT_ENTRIES = {name for name in TABLE if name == "cvt" or name.startswith("cvt_")} +# Every entry of the `cvt` instruction (ISA 9.7.9.22). Keyed off the mnemonic +# rather than the table name: `cvt.pack` (9.7.9.23) is a different instruction +# that happens to sort under the same prefix, and its forms are not points on +# this conversion grid. +_CVT_ENTRIES = {name for name, entry in TABLE.items() if entry.ptx_name == "cvt"} @pytest.mark.parametrize("entry_name,slots,instruction", _FORM_CASES) diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py b/tests/python/tirx/codegen/test_ptx_dialect.py index 36868632b9e9..42d73879158e 100644 --- a/tests/python/tirx/codegen/test_ptx_dialect.py +++ b/tests/python/tirx/codegen/test_ptx_dialect.py @@ -56,13 +56,16 @@ def _assert_ptxas_ok(src: str, rdc: bool = False, arch: str = PTX_ARCH) -> None: def test_ptx_registration(): from tvm.backend.cuda.intrinsics.registry import CODEGEN_REGISTRY - from tvm.backend.cuda.ptx.table import TABLE + from tvm.backend.cuda.ptx.table import TABLE, escape_token assert hasattr(T, "ptx") for entry in TABLE.values(): op = Op.get(entry.op_name) # raises if unregistered assert op.get_attr("TCallEffectKind") is not None, entry.name - family = entry.family # several entries may share a mnemonic + # The printer name is the attribute path a program types, so it is the + # *escaped* family: `and`/`or`/`not` (ISA 9.7.8) are Python keywords and + # print as `T.ptx.and_`. Identity for every other family. + family = escape_token(entry.family) # several entries may share a mnemonic assert op.get_attr("TScriptPrinterName") == f"ptx.{family}", entry.name assert entry.op_name in CODEGEN_REGISTRY, entry.name @@ -486,6 +489,809 @@ def test_ptx_u64_address_handle_is_reinterpreted(): assert addr.args[0].same_as(handle) +def test_ptx_integer_arithmetic_dispatch(): + """The 9.7.1 integer lines share five mnemonics with the floating-point ones. + + Nothing in the call names the entry: `T.ptx.add` reaches three candidates + (integer, single/double, half) and the written type token is what resolves + them, exactly as ptxas resolves them. `.wide` goes further -- it is a + separate entry per mnemonic because its result is "twice as wide as a and + b", a dtype no token in the instruction names. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "int32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + si = T.local_scalar("int32") + sw = T.local_scalar("int64") + fv = T.local_scalar("float32") + cnt = T.local_scalar("uint32") + T.ptx.add.s32(si, A[0], A[1]) # integer line + T.ptx.add.rn.f32(fv, fv, fv) # floating-point line, same mnemonic + T.ptx.mul.lo.s32(si, si, A[2]) # .lo keeps the operand width + T.ptx.mul.wide.s32(sw, si, A[3]) # .wide doubles the destination + T.ptx.mad.wide.s32(sw, si, A[0], sw) # ... and mad's accumulator too + T.ptx.popc.b32(cnt, si) # .b32 source, .u32 destination + T.ptx.dp4a.s32.s32(si, si, si, si) + A[tx % 4] = si + T.int32(sw) + T.int32(cnt) + + src = _cuda_source(kernel) + # Same mnemonic, different entries, each emitting its own ISA line. + assert "add.s32 %0, %1, %2;" in src + assert "add.rn.f32 %0, %1, %2;" in src + assert "mul.lo.s32 %0, %1, %2;" in src + # The derived destination is a 64-bit "l" register while the sources stay + # 32-bit "r" -- the whole reason `.wide` is its own entry. + assert '"=l"(__d) : "r"(__a), "r"(__b)' in src + assert "mul.wide.s32 %0, %1, %2;" in src + assert "mad.wide.s32 %0, %1, %2, %3;" in src + assert "popc.b32 %0, %1;" in src + assert "dp4a.s32.s32 %0, %1, %2, %3;" in src + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # A derived dtype is enforced at trace time like any other: `.wide.s32` + # writes 64 bits, so a 32-bit destination is rejected before codegen. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="must have dtype int64"): + + @T.prim_func + def narrow_wide_dst(out: T.Buffer((1,), "int32")): + T.device_entry() + T.ptx.mul.wide.s32(out[0], T.int32(2), T.int32(3)) + + # `.sat` is a syntax line of its own, not a free qualifier: the check + # rejects it where the ISA does not spell it. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="hi.s32"): + + @T.prim_func + def sat_on_lo(out: T.Buffer((1,), "int32")): + T.device_entry() + T.ptx.mad.lo.sat.s32(out[0], T.int32(2), T.int32(3), T.int32(4)) + + +def test_ptx_floating_point_dispatch(): + """ISA 9.7.3's lines, including three mnemonics shared with the integer group. + + `mad`, `div` and `abs` each name an integer entry and a floating-point one; + as with `add` above, only the written tokens choose between them. The rest + of the section is here for its shapes: a `.pred` destination (testp), the + f64 approximations that exist only with a mandatory `.ftz` (9.7.3.14 and + 9.7.3.17), and the lone approximation with no `.ftz` at all (tanh). + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + f = T.local_scalar("float32") + d = T.local_scalar("float64") + si = T.local_scalar("int32") + isnan = T.local_scalar("uint32") + T.ptx.mad.rn.f32(f, A[0], A[1], A[2]) # fp line + T.ptx.mad.lo.s32(si, si, si, si) # integer line, same mnemonic + T.ptx.div.approx.ftz.f32(f, f, A[3]) # fp divide + T.ptx.div.s32(si, si, si) # integer divide, same mnemonic + T.ptx.abs.ftz.f32(f, f) # fp abs + T.ptx.abs.s32(si, si) # integer abs, same mnemonic + T.ptx.copysign.f32(f, A[0], f) + T.ptx.sqrt.rn.f64(d, d) + T.ptx.rsqrt.approx.ftz.f64(d, d) # ISA 9.7.3.17 + T.ptx.rcp.approx.ftz.f64(d, d) # ISA 9.7.3.14 + T.ptx.sin.approx.f32(f, f) + T.ptx.cos.approx.ftz.f32(f, f) + T.ptx.lg2.approx.f32(f, f) + T.ptx.tanh.approx.f32(f, f) + T.ptx.testp.notanumber.f32(isnan, f) + A[tx % 4] = f + T.float32(d) + T.float32(si) + T.float32(isnan) + + src = _cuda_source(kernel) + for text in ( + "mad.rn.f32 %0, %1, %2, %3;", + "mad.lo.s32 %0, %1, %2, %3;", + "div.approx.ftz.f32 %0, %1, %2;", + "div.s32 %0, %1, %2;", + "abs.ftz.f32 %0, %1;", + "abs.s32 %0, %1;", + "copysign.f32 %0, %1, %2;", + "sqrt.rn.f64 %0, %1;", + "rsqrt.approx.ftz.f64 %0, %1;", + "rcp.approx.ftz.f64 %0, %1;", + "sin.approx.f32 %0, %1;", + "cos.approx.ftz.f32 %0, %1;", + "lg2.approx.f32 %0, %1;", + "tanh.approx.f32 %0, %1;", + ): + assert text in src, text + # The predicate result crosses the C boundary as a uint32 materialized by + # selp, with the real .pred register living inside the asm block. + assert ".reg .pred pd0; testp.notanumber.f32 pd0, %1; selp.b32 %0, 1, 0, pd0;" in src + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # sqrt has no f64 approximation at any spelling, unlike rcp (9.7.3.14). + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="only on the .f32 line"): + + @T.prim_func + def sqrt_approx_f64(out: T.Buffer((1,), "float64")): + T.device_entry() + T.ptx.sqrt.approx.f64(out[0], T.float64(2.0)) + + # ... and rcp's f64 approximation is unreachable without its mandatory .ftz. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="rcp.approx.ftz.f64"): + + @T.prim_func + def rcp_approx_f64(out: T.Buffer((1,), "float64")): + T.device_entry() + T.ptx.rcp.approx.f64(out[0], T.float64(2.0)) + + # One of .approx/.full/.rnd is required: the bare mnemonic names no line. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="missing required modifier"): + + @T.prim_func + def div_without_mode(out: T.Buffer((1,), "float32")): + T.device_entry() + T.ptx.div.f32(out[0], T.float32(1.0), T.float32(2.0)) + + +def test_ptx_half_precision_dispatch(): + """ISA 9.7.4, whose every mnemonic is shared with a wider-precision entry. + + `abs` is the extreme case: three entries answer to it (integer, f32/f64, + half), and only the type token says which. The section also has the two + qualifiers no same-precision line carries -- fma's `.relu` and `.oob` -- + and ex2's split over whether `.ftz` is mandatory or unspellable. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + h = T.local_scalar("uint16") # .f16 / .bf16 carrier + p = T.local_scalar("uint32") # .f16x2 / .bf16x2 carrier + f = T.local_scalar("float32") + si = T.local_scalar("int32") + T.ptx.add.rn.sat.f16(h, h, h) # existing half entry, beside the new ones + T.ptx.fma.rn.relu.f16x2(p, p, p, p) # .relu: no same-precision line has it + T.ptx.fma.rn.oob.bf16x2(p, p, p, p) # .oob likewise + T.ptx.fma.rn.f32(f, f, f, f) # same mnemonic, single precision + T.ptx.abs.ftz.f16(h, h) # three-way family: half ... + T.ptx.abs.ftz.f32(f, f) # ... single ... + T.ptx.abs.s32(si, si) # ... and integer + T.ptx.neg.bf16x2(p, p) + T.ptx.tanh.approx.f16x2(p, p) + T.ptx.ex2.approx.f16(h, h) + T.ptx.ex2.approx.ftz.bf16(h, h) # .ftz mandatory on this line + A[tx % 4] = T.uint32(p) + T.uint32(h) + T.uint32(si) + T.uint32(f) + + src = _cuda_source(kernel) + for text in ( + "add.rn.sat.f16 %0, %1, %2;", + "fma.rn.relu.f16x2 %0, %1, %2, %3;", + "fma.rn.oob.bf16x2 %0, %1, %2, %3;", + "fma.rn.f32 %0, %1, %2, %3;", + "abs.ftz.f16 %0, %1;", + "abs.ftz.f32 %0, %1;", + "abs.s32 %0, %1;", + "neg.bf16x2 %0, %1;", + "tanh.approx.f16x2 %0, %1;", + "ex2.approx.f16 %0, %1;", + "ex2.approx.ftz.bf16 %0, %1;", + ): + assert text in src, text + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # The `.oob` line spells neither .ftz nor .sat beside it. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="oob line spells no"): + + @T.prim_func + def oob_with_ftz(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.fma.rn.oob.ftz.f16(out[0], T.uint16(0), T.uint16(0), T.uint16(0)) + + # .sat and .relu are two clampings on two different syntax lines. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="separate syntax lines"): + + @T.prim_func + def sat_and_relu(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.fma.rn.sat.relu.f16(out[0], T.uint16(0), T.uint16(0), T.uint16(0)) + + # ex2's bf16 line without its mandatory .ftz names no syntax line. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="mandatorily"): + + @T.prim_func + def ex2_bf16_no_ftz(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.ex2.approx.bf16(out[0], T.uint16(0)) + + # ... and the bf16 abs/neg lines spell no .ftz at all. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="takes no .ftz"): + + @T.prim_func + def abs_bf16_ftz(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.abs.ftz.bf16(out[0], T.uint16(0)) + + +def test_ptx_mixed_precision_dispatch(): + """ISA 9.7.5, which adds no instruction of its own. + + Its three subsections are a fourth syntax line of add/sub/fma: + `op{.rnd}{.sat}.f32.atype`, where a second type token names a 16-bit source + converted to .f32 before the operation. That is the `srctype` slot, so + these forms are reached through the same entries as the same-precision + ones, and what selects a line is only whether that second token is written. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + f = T.local_scalar("float32") + h = T.local_scalar("uint16") # the .f16/.bf16 source carrier + T.ptx.fma.rn.f32.f16(f, h, h, A[0]) # both sources converted + T.ptx.add.f32.bf16(f, h, f) # `{.rnd}` omitted: the bare mixed line + T.ptx.sub.rn.sat.f32.f16(f, h, f) + T.ptx.add.rn.f32(f, f, f) # same entry, same-precision line + A[tx % 4] = f + + src = _cuda_source(kernel) + for text in ( + "fma.rn.f32.f16 %0, %1, %2, %3;", + "add.f32.bf16 %0, %1, %2;", + "sub.rn.sat.f32.f16 %0, %1, %2;", + "add.rn.f32 %0, %1, %2;", + ): + assert text in src, text + # Two carriers in one helper: the converted source is 16-bit, the rest f32. + assert '"=f"(__d) : "h"(__a), "f"(__b)' in src + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # The mixed line spells no .ftz, unlike the .f32 line it shares an entry with. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="takes no .ftz"): + + @T.prim_func + def mixed_with_ftz(out: T.Buffer((1,), "float32")): + T.device_entry() + T.ptx.add.rn.ftz.f32.f16(out[0], T.uint16(0), T.float32(0)) + + # A converted source exists only on the .f32 line -- there is no f64 form. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="only exists on the .f32"): + + @T.prim_func + def mixed_f64(out: T.Buffer((1,), "float64")): + T.device_entry() + T.ptx.add.rn.f64.f16(out[0], T.uint16(0), T.float64(0)) + + # And `mul` is the one mnemonic of the four with no mixed line at all, so + # its entry declares no srctype slot and the token does not resolve. + with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"): + + @T.prim_func + def mul_mixed(out: T.Buffer((1,), "float32")): + T.device_entry() + T.ptx.mul.rn.f32.f16(out[0], T.uint16(0), T.float32(0)) + + +def test_ptx_comparison_selection_dispatch(): + """ISA 9.7.6, the section whose results and selectors are predicates. + + Four mnemonics, eight entries: `setp` alone is four, because `{.BoolOp}` + adds an operand and `[|q]` adds a destination, and those are two + independent shape choices rather than optional tokens. Nothing in a call + names the entry -- the arity and the operand classes pick it, exactly as + ptxas picks the syntax line. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "int32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + p = T.local_scalar("uint32") + q = T.local_scalar("uint32") + d = T.local_scalar("uint32") + f = T.local_scalar("float32") + T.ptx.setp.lt.s32(p, A[0], A[1]) # one destination + T.ptx.setp.lt.s32(p, q, A[0], A[1]) # ... two: `p|q`, chosen by arity + T.ptx.setp.lt.and_.s32(p, A[0], A[1], T.ptx.pred(q)) # ... plus a BoolOp + T.ptx.setp.gt.or_.s32(p, q, A[2], A[3], T.ptx.pred(d)) # ... and both + T.ptx.set.lt.u32.f32(d, f, f) # writes a value, not a predicate + T.ptx.selp.b32(d, d, p, T.ptx.pred(q)) # predicate selects + T.ptx.slct.ftz.b32.f32(d, d, p, f) # a sign selects + A[tx % 4] = T.int32(p + q + d) + + src = _cuda_source(kernel) + for text in ( + "setp.lt.s32 pd0, %1, %2;", + "setp.lt.s32 pd0|pd1, %2, %3;", + "setp.lt.and.s32 pd0, %1, %2, ps0;", + "setp.gt.or.s32 pd0|pd1, %2, %3, ps0;", + "set.lt.u32.f32 %0, %1, %2;", + "selp.b32 %0, %1, %2, ps0;", + "slct.ftz.b32.f32 %0, %1, %2, %3;", + ): + assert text in src, text + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # Each source type takes its own operator set, and ptxas agrees with every + # rejection below (probed at sm_90 before they were written into the check). + # lo/ls/hi/hs are unsigned-only alternates for lt/le/gt/ge. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="signed"): + + @T.prim_func + def unsigned_op_on_signed(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.setp.lo.s32(out[0], T.int32(1), T.int32(2)) + + # The unordered comparisons and the NaN predicates are floating point only. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="floating-point comparison"): + + @T.prim_func + def float_op_on_integer(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.setp.nan.u32(out[0], T.uint32(1), T.uint32(2)) + + # A bit-size type has no ordering, only equality. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="only with eq/ne"): + + @T.prim_func + def ordered_on_bitsize(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.setp.lt.b32(out[0], T.uint32(1), T.uint32(2)) + + # .ftz applies only to .f32 comparisons ... + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="only to .f32 comparisons"): + + @T.prim_func + def ftz_off_f32(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.setp.eq.ftz.s32(out[0], T.int32(1), T.int32(2)) + + # ... and on slct, only to the .f32 selector line. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="f32 selector line"): + + @T.prim_func + def ftz_on_s32_selector(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.slct.ftz.b32.s32(out[0], T.uint32(1), T.uint32(2), T.int32(3)) + + +def test_ptx_half_comparison_dispatch(): + """ISA 9.7.7, the half twin of 9.7.6 -- same two mnemonics, different grids. + + Two things separate it from the section it shares names with. `setp`'s + destination shape is decided by the type rather than chosen: the scalar + lines spell one predicate, the packed lines only `p|q`, and there the pair + is the two lanes rather than a result and its complement. And `set` pairs + its two type tokens on a grid, so which (dtype, stype) combinations exist + is a fact of the section, not a product of two slots. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + h = T.local_scalar("uint16") # f16 / bf16 carrier + x2 = T.local_scalar("uint32") # f16x2 / bf16x2 carrier + p = T.local_scalar("uint32") + q = T.local_scalar("uint32") + d32 = T.local_scalar("uint32") + T.ptx.setp.eq.f16(p, h, h) # scalar type: one destination + T.ptx.setp.lt.ftz.f16x2(p, q, x2, x2) # packed type: the pair, per lane + T.ptx.setp.eq.and_.bf16(p, h, h, T.ptx.pred(q)) + T.ptx.setp.gt.or_.bf16x2(p, q, x2, x2, T.ptx.pred(d32)) + T.ptx.set.lt.u32.f16(d32, h, h) # integer answer to a half compare + T.ptx.set.gt.f16.f32(h, T.float32(1.0), T.float32(2.0)) # ... and back + T.ptx.set.eq.f16x2.f16x2(x2, x2, x2) # packed both sides + A[tx % 4] = p + q + d32 + T.uint32(h) + x2 + + src = _cuda_source(kernel) + for text in ( + "setp.eq.f16 pd0, %1, %2;", + "setp.lt.ftz.f16x2 pd0|pd1, %2, %3;", + "setp.eq.and.bf16 pd0, %1, %2, ps0;", + "setp.gt.or.bf16x2 pd0|pd1, %2, %3, ps0;", + "set.lt.u32.f16 %0, %1, %2;", + "set.gt.f16.f32 %0, %1, %2;", + "set.eq.f16x2.f16x2 %0, %1, %2;", + ): + assert text in src, text + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # Every rejection below was probed against ptxas before being written into + # the check. The (dtype, stype) grid: a packed source needs a 32-bit answer. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="no half-precision syntax"): + + @T.prim_func + def bad_pair(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.set.eq.u16.f16x2(out[0], T.uint32(0), T.uint32(0)) + + # `.ftz` follows the source precision: a bf16 source has none to flush. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="does not take it"): + + @T.prim_func + def ftz_bf16_source(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.set.eq.ftz.u32.bf16(out[0], T.uint16(0), T.uint16(0)) + + # ... and no bf16 destination takes it either. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="destination takes no"): + + @T.prim_func + def ftz_bf16_dest(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.set.eq.ftz.bf16.f32(out[0], T.float32(0), T.float32(0)) + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="spells no .ftz"): + + @T.prim_func + def setp_ftz_bf16(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.setp.eq.ftz.bf16(out[0], T.uint16(0), T.uint16(0)) + + # The unordered comparisons need a floating-point source. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="integer source"): + + @T.prim_func + def unordered_on_integer(out: T.Buffer((1,), "uint16")): + T.device_entry() + T.ptx.set.equ.f16.s32(out[0], T.int32(0), T.int32(0)) + + # 9.7.7 spells no unsigned alternates at all, so `lo` never resolves here. + with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"): + + @T.prim_func + def unsigned_alternate(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.setp.lo.f16(out[0], T.uint16(0), T.uint16(0)) + + +def test_ptx_logic_shift_dispatch(): + """ISA 9.7.8, including the three mnemonics that are Python keywords. + + `and`, `or` and `not` cannot be attributes as they stand, so the surface + spells them `and_`/`or_`/`not_` and the escape is carried through both + directions -- typing, and the name the printer emits. The round-trip + assertion below is what proves the second half: a script that printed + `T.ptx.and(...)` would not re-parse at all. + + The section also puts `.pred` in the *type* slot rather than on a single + operand, so a helper can hold three bridges around one instruction. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (4,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + d = T.local_scalar("uint32") + p = T.local_scalar("uint32") + q = T.local_scalar("uint32") + s = T.local_scalar("int32") + T.ptx.and_.b32(d, A[0], A[1]) + T.ptx.or_.b32(d, d, A[2]) + T.ptx.xor.b32(d, d, A[3]) + T.ptx.and_.pred(p, T.ptx.pred(q), T.ptx.pred(d)) # .pred as the type + T.ptx.not_.pred(p, T.ptx.pred(p)) + T.ptx.not_.b32(d, d) + T.ptx.cnot.b32(d, d) + T.ptx.lop3.b32(d, d, d, d, 0x80) # a & b & c, by look-up table + T.ptx.lop3.or_.b32(d, p, d, d, d, 0xFE, T.ptx.pred(q)) # d|p pair + T.ptx.shf.l.clamp.b32(d, d, d, T.uint32(4)) # funnel shift + T.ptx.shl.b32(d, d, T.uint32(2)) + T.ptx.shr.s32(s, s, T.uint32(1)) # signed: fills with the sign bit + A[tx % 4] = d + p + q + T.uint32(s) + + src = _cuda_source(kernel) + for text in ( + "and.b32 %0, %1, %2;", + "or.b32 %0, %1, %2;", + "xor.b32 %0, %1, %2;", + "and.pred pd0, ps0, ps1;", + "not.pred pd0, ps0;", + "not.b32 %0, %1;", + "cnot.b32 %0, %1;", + "lop3.b32 %0, %1, %2, %3, 128;", + "lop3.or.b32 %0|pd0, %2, %3, %4, 254, ps0;", + "shf.l.clamp.b32 %0, %1, %2, %3;", + "shl.b32 %0, %1, %2;", + "shr.s32 %0, %1, %2;", + ): + assert text in src, text + + # The load-bearing one: `and`/`or`/`not` have to survive being printed. + script = kernel.script() + assert "T.ptx.and_(" in script and "T.ptx.not_(" in script + reparsed = tvm.script.from_source(script) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # cnot has no predicate line -- ptxas: "Unexpected instruction types + # specified for 'cnot'" -- so the token is not in its slot at all. + with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"): + + @T.prim_func + def cnot_pred(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.cnot.pred(out[0], T.uint32(0)) + + # lop3's BoolOp line stops at .or/.and; ptxas rejects .xor outright. + with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"): + + @T.prim_func + def lop3_xor(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.lop3.xor.b32(out[0], T.uint32(0), T.uint32(0), T.uint32(0), 0x80) + + # shl is untyped -- a left shift zero-fills whatever the bits mean, so + # there is no signed line (unlike shr, which needs to know how to fill). + with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"): + + @T.prim_func + def shl_signed(out: T.Buffer((1,), "int32")): + 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() + d = T.local_scalar("uint32") + T.ptx.lop3.b32(d, A[0], A[0], A[0], A[0]) + + +def test_ptx_data_movement_dispatch(): + """ISA 9.7.9's newly registered instructions, end to end. + + The section's own difficulty is that its shapes vary more than its + qualifiers: `mov` shares a mnemonic with ten vector pack/unpack entries and + is told apart by arity alone, `shfl.sync` and `multimem` each split on + whether a bracketed part of the syntax line is written, and several + operands are immediates that live in the instruction text. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (8,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + smem = T.alloc_buffer((4,), "uint32", scope="shared") + d = T.local_scalar("uint32") + p = T.local_scalar("uint32") + v = T.local_scalar("uint32") + pol = T.local_scalar("uint64") + gen = T.local_scalar("uint64") + T.ptx.mov.b32(d, A[0]) # scalar mov: two operands + T.ptx.mov.pred(p, T.ptx.pred(d)) # ... and its predicate line + T.ptx.prmt.b32(d, d, d, A[1]) # generic byte permute + T.ptx.prmt.b32.f4e(d, d, d, A[2]) # ... and a mode-driven one + T.ptx.ldu.global_.b32(v, A.ptr_to([3])) + T.ptx.shfl_sync.idx.b32(d, d, T.uint32(0), T.uint32(31), T.uint32(0xFFFFFFFF)) + T.ptx.shfl_sync.up.b32(d, p, d, T.uint32(1), T.uint32(0), T.uint32(0xFFFFFFFF)) + T.ptx.isspacep.global_(p, A.ptr_to([4])) + T.ptx.cvta.to.global_.u64(gen, A.ptr_to([5])) + T.ptx.cvt_pack.sat.u16.s32(d, T.int32(1), T.int32(2)) + T.ptx.cvt_pack.sat.u8.s32.b32(d, T.int32(1), T.int32(2), d) + T.ptx.createpolicy.fractional.L2__evict_last.b64(pol) + T.ptx.createpolicy.cvt.L2.b64(pol, pol) + T.ptx.applypriority.global_.L2__evict_normal(A.ptr_to([6])) + T.ptx.discard.global_.L2(A.ptr_to([7])) + T.ptx.prefetchu.L1(A.ptr_to([0])) + T.ptx.multimem_ld_reduce.add.u32(v, A.ptr_to([0])) + T.ptx.multimem_red.relaxed.gpu.add.u32(A.ptr_to([0]), v) + smem[tx % 4] = d + p + v + A[tx % 8] = smem[tx % 4] + T.uint32(gen) + T.uint32(pol) + + src = _cuda_source(kernel) + for text in ( + "mov.b32 %0, %1;", + "mov.pred pd0, ps0;", + "prmt.b32 %0, %1, %2, %3;", + "prmt.b32.f4e %0, %1, %2, %3;", + "ldu.global.b32 %0, [%1];", + "shfl.sync.idx.b32 %0, %1, %2, %3, %4;", + "shfl.sync.up.b32 %0|pd0, %2, %3, %4, %5;", + "isspacep.global pd0, %1;", + "cvta.to.global.u64 %0, %1;", + "cvt.pack.sat.u16.s32 %0, %1, %2;", + "cvt.pack.sat.u8.s32.b32 %0, %1, %2, %3;", + "createpolicy.fractional.L2::evict_last.b64 %0;", + "createpolicy.cvt.L2.b64 %0, %1;", + "applypriority.global.L2::evict_normal [%0], 128;", + "discard.global.L2 [%0], 128;", + "prefetchu.L1 [%0];", + "multimem.ld_reduce.add.u32 %0, [%1];", + "multimem.red.relaxed.gpu.add.u32 [%0], %1;", + ): + assert text in src, text + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # Every rejection below was probed against ptxas before it was written into + # a check. multimem pairs an ordering semantic with a scope ... + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="go together"): + + @T.prim_func + def sem_without_scope(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint32") + T.device_entry() + v = T.local_scalar("uint32") + T.ptx.multimem_ld_reduce.relaxed.add.u32(v, A.ptr_to([0])) + + # ... and `.weak` is the line that has neither. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="takes no scope"): + + @T.prim_func + def weak_with_scope(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint32") + T.device_entry() + v = T.local_scalar("uint32") + T.ptx.multimem_ld_reduce.weak.gpu.add.u32(v, A.ptr_to([0])) + + # The op x type table: `.add` is the row that takes .s32 but not .s64. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match=r"\.add takes"): + + @T.prim_func + def add_s64(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "int64") + T.device_entry() + v = T.local_scalar("int64") + T.ptx.multimem_ld_reduce.add.s64(v, A.ptr_to([0])) + + # The scalar float line has no lone half: a width has to reach 32 bits. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="the scalar line takes"): + + @T.prim_func + def scalar_f16(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint16") + T.device_entry() + v = T.local_scalar("uint16") + T.ptx.multimem_ld_reduce.add.f16(v, A.ptr_to([0])) + + # `.acc::f32` raises an accumulation, so it needs something to accumulate. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="applies to .add"): + + @T.prim_func + def acc_on_min(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint32") + T.device_entry() + v = T.local_scalar("uint32") + T.ptx.multimem_ld_reduce.min.acc__f32.f16x2(v, A.ptr_to([0])) + + # st.async's mmio line is system-scoped only. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="requires .sys"): + + @T.prim_func + def mmio_gpu(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint32") + T.device_entry() + T.ptx.st_async.mmio.release.gpu.global_.u32(A.ptr_to([0]), T.uint32(1)) + + +def test_ptx_parallel_sync_dispatch(): + """ISA 9.7.14's warp-level primitives and the atom shapes beside `.op`. + + The section is where predicates are most load-bearing: they are sources + (bar.red's `c`), destinations (vote, elect), and both at once. It is also + where one mnemonic carries the most shapes -- `atom` now has the `.op` + line, `.cas`, `.exch`, the half-precision adds and two vector lines, all + resolved by tokens and arity alone. + """ + + @T.prim_func + def kernel(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (8,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + d = T.local_scalar("uint32") + p = T.local_scalar("uint32") + m = T.local_scalar("uint32") + h = T.local_scalar("uint16") + FULL = T.uint32(0xFFFFFFFF) + T.ptx.vote_sync.all.pred(p, T.ptx.pred(d), FULL) + T.ptx.vote_sync.ballot.b32(m, T.ptx.pred(d), FULL) + T.ptx.match.any.sync.b32(m, m, FULL) + T.ptx.match.all.sync.b32(m, p, m, FULL) # the |p twin + T.ptx.activemask.b32(m) + T.ptx.elect_sync(d, p, FULL) + T.ptx.redux_sync.add.u32(d, d, FULL) + T.ptx.redux_sync.and_.b32(d, d, FULL) + T.ptx.atom.global_.add.u32(d, A.ptr_to([0]), T.uint32(1)) # the .op line + T.ptx.atom.global_.cas.b32(d, A.ptr_to([1]), d, d) # ... and .cas + T.ptx.atom.global_.exch.b32(d, A.ptr_to([2]), d) # ... and .exch + T.ptx.atom.global_.add.noftz.f16(h, A.ptr_to([3]), h) + T.ptx.bar.red.popc.u32(d, T.uint32(0), T.ptx.pred(p)) + T.ptx.bar.red.and_.pred(p, T.uint32(0), T.ptx.pred(p)) + A[tx % 8] = d + p + m + T.uint32(h) + + src = _cuda_source(kernel) + for text in ( + "vote.sync.all.pred pd0, ps0, %2;", + "vote.sync.ballot.b32 %0, ps0, %2;", + "match.any.sync.b32 %0, %1, %2;", + "match.all.sync.b32 %0|pd0, %2, %3;", + "activemask.b32 %0;", + "elect.sync %0|pd0, %2;", + "redux.sync.add.u32 %0, %1, %2;", + "redux.sync.and.b32 %0, %1, %2;", + "atom.global.add.u32 %0, [%1], %2;", + "atom.global.cas.b32 %0, [%1], %2, %3;", + "atom.global.exch.b32 %0, [%1], %2;", + "atom.global.add.noftz.f16 %0, [%1], %2;", + "bar.red.popc.u32 %0, %1, ps0;", + "bar.red.and.pred pd0, %1, ps0;", + ): + assert text in src, text + + reparsed = tvm.script.from_source(kernel.script()) + tvm.ir.assert_structural_equal(kernel, reparsed) + + # Every rejection below was probed against ptxas first. redux pairs the + # arithmetic ops with a signed type and the bitwise ones with the untyped + # word -- they are two syntax lines, so the token does not even resolve. + with pytest.raises((AttributeError, tvm.error.DiagnosticError), match="not a valid modifier"): + + @T.prim_func + def redux_add_b32(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.redux_sync.add.b32(out[0], T.uint32(1), T.uint32(0xFFFFFFFF)) + + # atom's vector lines bound the width by the element: a packed pair stops + # at .v4, and only a lone half reaches .v8. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="already a 32-bit pair"): + + @T.prim_func + def packed_v8(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint32") + T.device_entry() + v = T.local_scalar("uint32") + T.ptx.atom.global_.add.noftz.v8.f16x2( + v, v, v, v, v, v, v, v, A.ptr_to([0]), v, v, v, v, v, v, v, v + ) + + # red.async's op groups: `.add` reaches 64 bits, the bitwise ops do not. + with pytest.raises((ValueError, tvm.error.DiagnosticError), match=r"\.and takes"): + + @T.prim_func + def red_async_and_u64(a_ptr: T.handle): + A = T.match_buffer(a_ptr, (1,), "uint64") + T.device_entry() + v = T.local_scalar("uint64") + T.ptx.red_async.relaxed.cluster.shared__cluster.mbarrier__complete_tx__bytes.and_.u64( + A.ptr_to([0]), v, A.ptr_to([0]) + ) + + def test_ptx_parser_roundtrip(): """script() output re-parses to a structurally equal PrimFunc.""" @@ -672,9 +1478,9 @@ def test_ptx_helper_source_golden(): from tvm.backend.cuda.ptx.render import render_variant from tvm.backend.cuda.ptx.table import TABLE, tokens_for - def render(name, predicated=False, dtypes=None, **by_name): + def render(name, predicated=False, dtypes=None, imms=None, **by_name): entry = TABLE[name] - return render_variant(entry, tokens_for(entry, **by_name), predicated, dtypes)[2] + return render_variant(entry, tokens_for(entry, **by_name), predicated, dtypes, imms)[2] # tokens_for is what lets the goldens below name their modifiers. Positional # tuples silently shift when a slot is inserted, which is the one edit this @@ -761,6 +1567,334 @@ def render(name, predicated=False, dtypes=None, **by_name): "}\n" ) + # Integer arithmetic (ISA 9.7.1). The plain shape first: an integer line + # sharing the `add` mnemonic with the floating-point entry, resolved apart + # by its type token alone. + assert render("add_int", sat="sat", type="s32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_add_int_sat_s32" + "(int32_t& __d, int32_t __a, int32_t __b) {\n" + ' asm volatile("add.sat.s32 %0, %1, %2;" : "=r"(__d) : "r"(__a), "r"(__b));\n' + "}\n" + ) + # A dtype the instruction text does not name: `.wide`'s destination is + # "twice as wide as a and b", so one written token (.s16) types the sources + # while the result is the derived .s32 -- 16-bit "h" inputs, a 32-bit "r" + # output. mad.wide derives c the same way, which is why its accumulator is + # the wide type and its multiplicands are not. + assert render("mul_wide", mode="wide", type="s16") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_mul_wide_wide_s16" + "(int32_t& __d, int16_t __a, int16_t __b) {\n" + ' asm volatile("mul.wide.s16 %0, %1, %2;" : "=r"(__d) : "h"(__a), "h"(__b));\n' + "}\n" + ) + assert render("mad_wide", mode="wide", type="u32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_mad_wide_wide_u32" + "(uint64_t& __d, uint32_t __a, uint32_t __b, uint64_t __c) {\n" + ' asm volatile("mad.wide.u32 %0, %1, %2, %3;" : "=l"(__d) : "r"(__a), ' + '"r"(__b), "l"(__c));\n' + "}\n" + ) + # A destination the ISA types outright: popc counts a 64-bit source into a + # .u32 result, so the two operands have unrelated widths. + assert render("popc", type="b64") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_popc_b64" + "(uint32_t& __d, uint64_t __a) {\n" + ' asm volatile("popc.b64 %0, %1;" : "=r"(__d) : "l"(__a));\n' + "}\n" + ) + # Two written type tokens, and an accumulator derived from both: c and d + # are .u32 only when atype and btype both are, so this mixed pair makes + # them .s32 (ISA 9.7.1.24). + assert render("dp4a", atype="u32", btype="s32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_dp4a_u32_s32" + "(int32_t& __d, uint32_t __a, int32_t __b, int32_t __c) {\n" + ' asm volatile("dp4a.u32.s32 %0, %1, %2, %3;" : "=r"(__d) : "r"(__a), ' + '"r"(__b), "r"(__c));\n' + "}\n" + ) + # Five operands under the ISA's own names: the destination is `f`, and `d` + # is the field-length input -- the one family where `d` is not the result. + assert render("bfi", type="b64") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_bfi_b64" + "(uint64_t& __f, uint64_t __a, uint64_t __b, uint32_t __c, uint32_t __d) {\n" + ' asm volatile("bfi.b64 %0, %1, %2, %3, %4;" : "=l"(__f) : "l"(__a), "l"(__b), ' + '"r"(__c), "r"(__d));\n' + "}\n" + ) + # bmsk's destination is written .b32 but declared u32: ptxas takes no float + # register anywhere in this instruction, unlike every other .bN family (see + # the entry's note). So there is no f32 twin of this helper to render. + assert render("bmsk", mode="clamp", type="b32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_bmsk_clamp_b32" + "(uint32_t& __d, uint32_t __a, uint32_t __b) {\n" + ' asm volatile("bmsk.clamp.b32 %0, %1, %2;" : "=r"(__d) : "r"(__a), "r"(__b));\n' + "}\n" + ) + + # Floating point (ISA 9.7.3). A `.pred` destination: the instruction writes + # a predicate register, which no asm constraint can bind, so the helper + # declares one inside the block and materializes 0/1 through selp on the + # way out. Still exactly one instruction -- the rest is the boundary. + assert render("testp", op="notanumber", type="f32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_testp_notanumber_f32" + "(uint32_t& __p, float __a) {\n" + ' asm volatile("{ .reg .pred pd0; testp.notanumber.f32 pd0, %1; ' + 'selp.b32 %0, 1, 0, pd0; }" : "=r"(__p) : "f"(__a));\n' + "}\n" + ) + # `.full` is a divide mode beside the rounding modes, not a qualifier on + # top of one -- the ISA requires exactly one of .approx/.full/.rnd, so the + # entry fuses them into a single mandatory slot. + assert render("div_f", mode="approx", ftz="ftz", type="f32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_div_f_approx_ftz_f32" + "(float& __d, float __a, float __b) {\n" + ' asm volatile("div.approx.ftz.f32 %0, %1, %2;" : "=f"(__d) : "f"(__a), "f"(__b));\n' + "}\n" + ) + # mad shares fma's grid and `_check_farith`, and shares the `mad` mnemonic + # with the two integer entries; `.rn` is what tells them apart. + assert render("mad_f", rnd="rn", sat="sat", type="f32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_mad_f_rn_sat_f32" + "(float& __d, float __a, float __b, float __c) {\n" + ' asm volatile("mad.rn.sat.f32 %0, %1, %2, %3;" : "=f"(__d) : "f"(__a), ' + '"f"(__b), "f"(__c));\n' + "}\n" + ) + # ISA 9.7.3.14 is its own subsection because it is a different computation, + # but syntactically it is one cell of rcp's grid -- reachable only with the + # mandatory .ftz its syntax line spells (see `_check_rcp`). + assert render("rcp", mode="approx", ftz="ftz", type="f64") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_rcp_approx_ftz_f64" + "(double& __d, double __value) {\n" + ' asm volatile("rcp.approx.ftz.f64 %0, %1;" : "=d"(__d) : "d"(__value));\n' + "}\n" + ) + + # Half precision (ISA 9.7.4). fma's two extra clampings, on the packed + # type: `.relu` and `.oob` exist on no same-precision line, which is what + # keeps the half fma out of the f32/f64 entry. A packed pair rides one + # 32-bit register, so every operand is uint32. + assert render("fma_half", rnd="rn", oob="oob", relu="relu", type="f16x2") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_fma_half_rn_oob_relu_f16x2" + "(uint32_t& __d, uint32_t __a, uint32_t __b, uint32_t __c) {\n" + ' asm volatile("fma.rn.oob.relu.f16x2 %0, %1, %2, %3;" : "=r"(__d) : "r"(__a), ' + '"r"(__b), "r"(__c));\n' + "}\n" + ) + # ex2's bf16 line spells `.ftz` mandatorily while its f16 line does not + # offer it -- the reason these are a separate entry from the .f32 ex2. + assert render("ex2_half", mode="approx", ftz="ftz", type="bf16x2") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_ex2_half_approx_ftz_bf16x2" + "(uint32_t& __d, uint32_t __value) {\n" + ' asm volatile("ex2.approx.ftz.bf16x2 %0, %1;" : "=r"(__d) : "r"(__value));\n' + "}\n" + ) + # The scalar half type is the other carrier: a 16-bit register on "h". + assert render("abs_half", ftz="ftz", type="f16") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_abs_half_ftz_f16" + "(uint16_t& __d, uint16_t __a) {\n" + ' asm volatile("abs.ftz.f16 %0, %1;" : "=h"(__d) : "h"(__a));\n' + "}\n" + ) + + # Mixed precision (ISA 9.7.5) is not a separate instruction but a fourth + # syntax line of add/sub/fma: a second type token names a 16-bit source + # that is converted to .f32 before the operation. So one helper carries two + # carriers at once -- "h" for the converted sources, "f" for the rest. + # fma converts both a and b (`.abtype`); add/sub convert only a. + assert render("fma", rnd="rn", sat="sat", type="f32", srctype="bf16") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_fma_rn_sat_f32_bf16" + "(float& __d, uint16_t __a, uint16_t __b, float __c) {\n" + ' asm volatile("fma.rn.sat.f32.bf16 %0, %1, %2, %3;" : "=f"(__d) : "h"(__a), ' + '"h"(__b), "f"(__c));\n' + "}\n" + ) + + # Comparison and selection (ISA 9.7.6). setp's `p|q`: ONE operand position + # holding two predicate destinations, joined by the ISA's own separator + # rather than a comma. Each half is a real register with its own selp on + # the way out -- q carries the Boolean applied to the complement of the + # compare, so it is a second result, not a restatement of p. + assert render("setp_pq", cmp="lt", type="s32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_setp_pq_lt_s32" + "(uint32_t& __p, uint32_t& __q, int32_t __a, int32_t __b) {\n" + ' asm volatile("{ .reg .pred pd0; .reg .pred pd1; setp.lt.s32 pd0|pd1, %2, %3; ' + 'selp.b32 %0, 1, 0, pd0; selp.b32 %1, 1, 0, pd1; }" : "=r"(__p), "=r"(__q) ' + ': "r"(__a), "r"(__b));\n' + "}\n" + ) + # A predicate *source* rides the same carrier in the other direction: setp + # in, then the instruction. selp is the plainest case of it. + assert render("selp", type="f64") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_selp_f64" + "(double& __d, double __a, double __b, uint32_t __c) {\n" + ' asm volatile("{ .reg .pred ps0; setp.ne.b32 ps0, %3, 0; ' + 'selp.f64 %0, %1, %2, ps0; }" : "=d"(__d) : "d"(__a), "d"(__b), "r"(__c));\n' + "}\n" + ) + # `set` writes a value rather than a predicate, so its two type tokens are + # independent: an f32 comparison landing 0xffffffff in a u32 destination. + assert render("set", cmp="lt", dtype="u32", stype="f32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_set_lt_u32_f32" + "(uint32_t& __d, float __a, float __b) {\n" + ' asm volatile("set.lt.u32.f32 %0, %1, %2;" : "=r"(__d) : "f"(__a), "f"(__b));\n' + "}\n" + ) + # slct likewise: the selected value and the number whose sign selects it + # are separately typed. + assert render("slct", ftz="ftz", dtype="b32", ctype="f32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_slct_ftz_b32_f32" + "(uint32_t& __d, uint32_t __a, uint32_t __b, float __c) {\n" + ' asm volatile("slct.ftz.b32.f32 %0, %1, %2, %3;" : "=r"(__d) : "r"(__a), ' + '"r"(__b), "f"(__c));\n' + "}\n" + ) + + # Half-precision comparison (ISA 9.7.7). The packed setp reuses the pipe + # pair, but its two halves mean something else than 9.7.6's: p and q are + # the two lanes' comparisons, not a result and its complement. + assert render("setp_half_pq", cmp="lt", ftz="ftz", type="f16x2") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_setp_half_pq_lt_ftz_f16x2" + "(uint32_t& __p, uint32_t& __q, uint32_t __a, uint32_t __b) {\n" + ' asm volatile("{ .reg .pred pd0; .reg .pred pd1; ' + "setp.lt.ftz.f16x2 pd0|pd1, %2, %3; selp.b32 %0, 1, 0, pd0; " + 'selp.b32 %1, 1, 0, pd1; }" : "=r"(__p), "=r"(__q) : "r"(__a), "r"(__b));\n' + "}\n" + ) + # set's two type tokens run in either direction: a half-valued answer to an + # integer comparison ... + assert render("set_half", cmp="lt", dtype="f16", stype="s32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_set_half_lt_f16_s32" + "(uint16_t& __d, int32_t __a, int32_t __b) {\n" + ' asm volatile("set.lt.f16.s32 %0, %1, %2;" : "=h"(__d) : "r"(__a), "r"(__b));\n' + "}\n" + ) + # ... and an integer-valued answer to a half one, where `.ftz` is legal + # because the *source* is the f16 (see `_FTZ_SET_STYPES`). + assert render("set_half", cmp="gt", ftz="ftz", dtype="u32", stype="f16") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_set_half_gt_ftz_u32_f16" + "(uint32_t& __d, uint16_t __a, uint16_t __b) {\n" + ' asm volatile("set.gt.ftz.u32.f16 %0, %1, %2;" : "=r"(__d) : "h"(__a), "h"(__b));\n' + "}\n" + ) + + # Logic and shift (ISA 9.7.8). `.pred` as an instruction type, rather than + # as one operand's class: three bridges -- two setp in, one selp out -- + # wrapped around a single instruction that never touches the carriers. + assert render("and", type="pred") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_and_pred" + "(uint32_t& __d, uint32_t __a, uint32_t __b) {\n" + ' asm volatile("{ .reg .pred pd0; .reg .pred ps0; .reg .pred ps1; ' + "setp.ne.b32 ps0, %1, 0; setp.ne.b32 ps1, %2, 0; and.pred pd0, ps0, ps1; " + 'selp.b32 %0, 1, 0, pd0; }" : "=r"(__d) : "r"(__a), "r"(__b));\n' + "}\n" + ) + # lop3's BoolOp form: the pipe pair holding two DIFFERENT register classes. + # `d` is an ordinary b32 output bound straight to %0 while `p` rides the + # predicate bridge, so the pair renders `%0|pd0` -- the mechanism groups + # the text and leaves each half its own constraint. The LUT byte is an open + # immediate, baked into the text and into the helper name. + assert render("lop3_bool", boolop="and", type="b32", imms=("128",)) == ( + "__forceinline__ __device__ void tvm_builtin_ptx_lop3_bool_and_b32_128" + "(uint32_t& __d, uint32_t& __p, uint32_t __a, uint32_t __b, uint32_t __c, " + "uint32_t __q) {\n" + ' asm volatile("{ .reg .pred pd0; .reg .pred ps0; setp.ne.b32 ps0, %5, 0; ' + 'lop3.and.b32 %0|pd0, %2, %3, %4, 128, ps0; selp.b32 %1, 1, 0, pd0; }" ' + ': "=r"(__d), "=r"(__p) : "r"(__a), "r"(__b), "r"(__c), "r"(__q));\n' + "}\n" + ) + # The shift amount is a 32-bit value "regardless of the instruction type", + # so a 16-bit shl still takes a uint32 there. + assert render("shl", type="b16") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_shl_b16" + "(uint16_t& __d, uint16_t __a, uint32_t __b) {\n" + ' asm volatile("shl.b16 %0, %1, %2;" : "=h"(__d) : "h"(__a), "r"(__b));\n' + "}\n" + ) + + # Data movement (ISA 9.7.9). shfl.sync's `d|p`: a pipe pair whose halves + # are DIFFERENT register classes -- an ordinary b32 result bound straight + # to %0, and an in-range predicate that rides the bridge. + assert render("shfl_sync_p", mode="up", type="b32") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_shfl_sync_p_up_b32" + "(uint32_t& __d, uint32_t& __p, uint32_t __a, uint32_t __b, uint32_t __c, " + "uint32_t __membermask) {\n" + ' asm volatile("{ .reg .pred pd0; shfl.sync.up.b32 %0|pd0, %2, %3, %4, %5; ' + 'selp.b32 %1, 1, 0, pd0; }" : "=r"(__d), "=r"(__p) : "r"(__a), "r"(__b), ' + '"r"(__c), "r"(__membermask));\n' + "}\n" + ) + # multimem's vector line: `.v4` of `.f16` is one 64-bit access spread over + # four half registers, and `.acc::f32` widens only the accumulation. + assert render("multimem_ld_reduce_f_vec", op="add", acc="acc::f32", vec="v4", type="f16") == ( + "__forceinline__ __device__ void " + "tvm_builtin_ptx_multimem_ld_reduce_f_vec_add_acc__f32_v4_f16" + "(uint16_t& __d0, uint16_t& __d1, uint16_t& __d2, uint16_t& __d3, " + "const void* __addr) {\n" + ' asm volatile("multimem.ld_reduce.add.acc::f32.v4.f16 {%0, %1, %2, %3}, [%4];" ' + ': "=h"(__d0), "=h"(__d1), "=h"(__d2), "=h"(__d3) : "l"(__addr) : "memory");\n' + "}\n" + ) + # A table-owned immediate: the ISA fixes applypriority's size at 128, so it + # is in the text and the helper has no parameter for it at all. + assert render("applypriority", space="global", level="L2::evict_normal") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_applypriority_global_L2__evict_normal" + "(const void* __addr) {\n" + ' asm volatile("applypriority.global.L2::evict_normal [%0], 128;" : : "l"(__addr)' + ' : "memory");\n' + "}\n" + ) + # A caller-chosen immediate: tensormap.replace's field3 new_val must be a + # constant, so each legal value is its own helper, named for it. + assert render( + "tensormap_replace_elemtype", + mode="tile", + field="elemtype", + space="global", + width="b1024", + type="b32", + imms=("7",), + ) == ( + "__forceinline__ __device__ void " + "tvm_builtin_ptx_tensormap_replace_elemtype_tile_elemtype_global_b1024_b32_7" + "(const void* __addr) {\n" + ' asm volatile("tensormap.replace.tile.elemtype.global.b1024.b32 [%0], 7;" : : ' + '"l"(__addr) : "memory");\n' + "}\n" + ) + + # Parallel synchronization (ISA 9.7.14). elect.sync's `d|p` is the one + # pipe pair the ISA makes mandatory -- ptxas rejects a bare destination -- + # and its halves are two different register classes. + assert render("elect_sync") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_elect_sync" + "(uint32_t& __d, uint32_t& __p, uint32_t __membermask) {\n" + ' asm volatile("{ .reg .pred pd0; elect.sync %0|pd0, %2; ' + 'selp.b32 %1, 1, 0, pd0; }" : "=r"(__d), "=r"(__p) : "r"(__membermask));\n' + "}\n" + ) + # atom.cas is the four-value shape, and the only atom line reaching 128 + # bits in every position -- the "q" constraint the C boundary needs. + assert render("atom_cas", space="global", op="cas", type="b128") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_atom_cas_global_cas_b128" + "(__uint128_t& __d, const void* __addr, __uint128_t __compare, " + "__uint128_t __value) {\n" + ' asm volatile("atom.global.cas.b128 %0, [%1], %2, %3;" : "=q"(__d) : ' + '"l"(__addr), "q"(__compare), "q"(__value) : "memory");\n' + "}\n" + ) + # bar.red reduces a predicate across a barrier, so a predicate crosses the + # boundary in BOTH directions around one instruction: setp on the way in, + # selp on the way out. + assert render("bar_red_pred", action="red", op="and", type="pred") == ( + "__forceinline__ __device__ void tvm_builtin_ptx_bar_red_pred_red_and_pred" + "(uint32_t& __d, uint32_t __a, uint32_t __c) {\n" + ' asm volatile("{ .reg .pred pd0; .reg .pred ps0; setp.ne.b32 ps0, %2, 0; ' + 'bar.red.and.pred pd0, %1, ps0; selp.b32 %0, 1, 0, pd0; }" : "=r"(__d) : ' + '"r"(__a), "r"(__c) : "memory");\n' + "}\n" + ) + # Register group: `{%1, %2}` is ONE PTX operand spanning two registers, so # the group is braces in the asm text and flat C parameters around it. An # unpack turns that into two "=" outputs. @@ -857,19 +1991,31 @@ def test_ptx_coercion_ir_forms(): _ASM_RE = re.compile(r'asm(?: volatile)?\("(.*?)"\s*:', re.S) _BLOCK_RE = re.compile(r"^\{ (?P.*) \}$") -# The sanctioned in-block boundary conversions, and nothing else: predicate -# register declarations, setp conversions in (@p's own guard and pred_src -# operands), selp materializations out (pred_dst operands). # The asm block's sanctioned non-instructions: `render.BRIDGE`'s register # declarations and the conversions that move a value between the class the ISA -# names and the carrier inline asm can bind. Never semantics -- see BRIDGE. -_BOUNDARY_PREFIXES = ( - ".reg .pred ", - ".reg .b8 ", - "setp.ne.b32 ", - "selp.b32 ", - "cvt.u8.u16 ", - "cvt.u16.u8 ", +# names and the carrier inline asm can bind, plus `@p`'s own guard. Never +# semantics -- see BRIDGE. +# +# Matched in FULL, not by opcode prefix, because an opcode is no longer a +# discriminator: `setp` and `selp` are registered instructions (ISA 9.7.6), so +# a helper's one real instruction can carry the same mnemonic as the bridge +# statements around it. What separates them is the shape the bridge always +# has -- it names a bridge-local register (`p`, `ps`, `pd`, `raw_`) +# and compares or selects against the literals the conversion is made of, while +# a registered instruction's operands are `%` throughout. Matching the whole +# statement also makes this stricter than the prefix form it replaces: a stray +# conversion in some other shape is now a failure instead of being peeled. +_BOUNDARY_RE = re.compile( + r"|".join( + ( + r"\.reg \.pred (?:p|ps\d+|pd\d+);", # @p guard / pred bridge declarations + r"\.reg \.b8 raw_\w+;", # b8 bridge declaration + r"setp\.ne\.b32 (?:p|ps\d+), %\d+, 0;", # @p guard, pred_src conversion in + r"selp\.b32 %\d+, 1, 0, pd\d+;", # pred_dst materialization out + r"cvt\.u8\.u16 raw_\w+, %\d+;", # b8 conversion in + r"cvt\.u16\.u8 %\d+, raw_\w+;", # b8 conversion out + ) + ) ) @@ -890,7 +2036,7 @@ def _sole_instruction(asm_text): m = _BLOCK_RE.match(asm_text) if m: stmts = [f"{part.strip()};" for part in m.group("body").split(";") if part.strip()] - core = [st for st in stmts if not st.startswith(_BOUNDARY_PREFIXES)] + core = [st for st in stmts if not _BOUNDARY_RE.fullmatch(st)] if len(core) != 1: return None body = core[0].removeprefix("@p ") @@ -954,7 +2100,12 @@ def test_ptx_single_instruction_invariant(): assert instr.startswith(opcode + " ") or instr == opcode + ";", ( f"{opcode}: emitted instruction does not match the opcode: {instr!r}" ) - assert "cvta" not in source or entry.name == "cvta", ( + # The address coercion `cvta` is a separate IR node, so no OTHER + # instruction's helper may smuggle one in. The cvta entries + # themselves are the exception, being that instruction: the guard + # keys off the mnemonic rather than the table name so that every + # syntax line of it (`cvta.to.space` and `cvta.space`) is covered. + assert "cvta" not in source or entry.ptx_name == "cvta", ( f"{opcode}: cvta must be a separate IR node, never inside a helper" ) # Every helper is void: a PTX destination is an operand, never a C @@ -1022,7 +2173,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 == 162716 # update when the table grows + assert total == 198709 # update when the table grows def test_ptx_no_instruction_registered_twice(): diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 51b1fb46ebaa..0e523c29cd9e 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -2297,6 +2297,28 @@ def test(A_ptr: T.handle) -> None: assert_structural_equal(test, from_source(code)) +def test_roundtrip_serial_unroll_count(): + """T.serial(N, unroll=2) should preserve the requested unroll count.""" + + # fmt: off + @T.prim_func + def test(A_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (128,), "float32", scope="global") + T.device_entry() + cta_id = T.cta_id([1]) + warp_id = T.warp_id([1]) + lane_id = T.lane_id([32]) + for _ in T.serial(10, unroll=2): + Tx.cta.fill(A[0:32], T.float32(0)) + # fmt: on + + code = test.script() + assert "unroll=2" in code, f"printer should emit unroll=2, got:\n{code}" + assert "annotations" not in code, "printer should NOT emit annotations dict" + assert from_source(code).script() == code + assert_structural_equal(test, from_source(code)) + + def test_roundtrip_serial_unroll_false_with_other_annotations(): """When other annotations exist alongside disable_unroll, fall back to full dict.""" @@ -2759,5 +2781,251 @@ def func(A_ptr: T.handle): assert_structural_equal(func, from_source(code)) +def _assert_roundtrip(func): + code = func.script() + assert from_source(code).script() == code + assert_structural_equal(func, from_source(code)) + + +def test_loop_var_dtype_uint32(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (128,), "float32") + for i in T.serial(128, dtype="uint32"): + A[i] = T.float32(1) + # fmt: on + + loop = func.body + assert loop.loop_var.ty == PrimType("uint32") + assert loop.min.ty == PrimType("uint32") + assert loop.extent.ty == PrimType("uint32") + _assert_roundtrip(func) + + +def test_loop_var_dtype_uint32_with_step(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (128,), "float32") + for i in T.serial(4, 128, step=2, dtype="uint32"): + A[i] = T.float32(1) + # fmt: on + + loop = func.body + assert loop.loop_var.ty == PrimType("uint32") + assert loop.min.ty == PrimType("uint32") + assert loop.extent.ty == PrimType("uint32") + assert loop.step.ty == PrimType("uint32") + _assert_roundtrip(func) + + +@pytest.mark.parametrize("for_kind", ["serial", "parallel", "vectorized", "unroll"]) +def test_loop_var_dtype_uint32_all_for_kinds(for_kind): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (4,), "float32") + for i in getattr(T, for_kind)(4, dtype="uint32"): + A[i] = T.float32(1) + # fmt: on + + assert func.body.loop_var.ty == PrimType("uint32") + _assert_roundtrip(func) + + +def test_grid_loop_var_dtype_uint32(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (8, 16), "float32") + for i, j in T.grid(8, 16, dtype="uint32"): + A[i, j] = T.float32(1) + # fmt: on + + outer = func.body + assert outer.loop_var.ty == PrimType("uint32") + assert outer.body.loop_var.ty == PrimType("uint32") + _assert_roundtrip(func) + + +def test_loop_var_dtype_defaults_to_int32(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (128,), "float32") + for i in range(128): + A[i] = T.float32(1) + # fmt: on + + assert func.body.loop_var.ty == PrimType("int32") + _assert_roundtrip(func) + + +def test_loop_var_dtype_inferred_from_unsigned_extent(): + """A uint32 extent makes the loop var uint32 without an explicit dtype.""" + + # fmt: off + @T.prim_func + def func(A_ptr: T.handle, n: T.uint32): + A = T.match_buffer(A_ptr, (128,), "float32") + for i in range(n): + A[i] = T.float32(1) + # fmt: on + + assert func.body.loop_var.ty == PrimType("uint32") + _assert_roundtrip(func) + + +def test_loop_var_dtype_casts_mismatched_bound(): + """A non-literal bound of another dtype is cast to the requested loop dtype.""" + + # fmt: off + @T.prim_func + def func(A_ptr: T.handle, n: T.int32): + A = T.match_buffer(A_ptr, (128,), "float32") + for i in T.serial(n, dtype="uint32"): + A[i] = T.float32(1) + # fmt: on + + loop = func.body + assert loop.loop_var.ty == PrimType("uint32") + assert loop.extent.ty == PrimType("uint32") + _assert_roundtrip(func) + + +@pytest.mark.parametrize("dtype", ["int64", "uint64", "int16", "float32"]) +def test_loop_var_dtype_rejects_unsupported(dtype): + with pytest.raises(Exception, match='must be "int32" or "uint32"'): + T.serial(4, dtype=dtype) + + +def test_thread_binding_has_no_dtype_parameter(): + with pytest.raises(TypeError): + T.thread_binding(0, 128, "threadIdx.x", dtype="uint32") + + +def test_hand_built_for_promotes_int_literal_bounds_to_uint32(): + """The For constructor retypes literal bounds to the loop var's dtype.""" + loop_var = tvm.tirx.Var("i", "uint32") + loop = tvm.tirx.For(loop_var, 0, 128, tvm.tirx.ForKind.SERIAL, tvm.tirx.Evaluate(0)) + assert loop.min.ty == PrimType("uint32") + assert loop.extent.ty == PrimType("uint32") + + +def test_hand_built_for_rejects_negative_literal_for_uint32(): + loop_var = tvm.tirx.Var("i", "uint32") + with pytest.raises(Exception, match="not representable"): + tvm.tirx.For(loop_var, -1, 128, tvm.tirx.ForKind.SERIAL, tvm.tirx.Evaluate(0)) + + +def test_scope_id_dtype_uint32(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (128,), "float32") + T.device_entry() + bx = T.cta_id([1]) + tx = T.thread_id([128], dtype="uint32") + A[tx] = T.float32(bx) + # fmt: on + + scope_defs = [] + tvm.tirx.stmt_functor.post_order_visit( + func.body, + lambda s: ( + scope_defs.append(getattr(s, "def")) if isinstance(s, tvm.tirx.ScopeIdDefStmt) else None + ), + ) + dtypes = {str(d.def_ids[0].ty) for d in scope_defs} + assert dtypes == {"int32", "uint32"} + # The extents stay int32 regardless of the def var dtype. + for d in scope_defs: + assert d.extents[0].ty == PrimType("int32") + + code = func.script() + assert 'T.thread_id([128], dtype="uint32")' in code + assert "T.cta_id([1])" in code + _assert_roundtrip(func) + + +def test_scope_id_dtype_uint32_lane_and_warp(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (32,), "float32") + T.device_entry() + _ = T.cta_id([1]) + warp = T.warp_id([4], dtype="uint32") + lane = T.lane_id([32], dtype="uint32") + A[lane] = T.float32(warp) + # fmt: on + + code = func.script() + assert 'T.warp_id([4], dtype="uint32")' in code + assert 'T.lane_id([32], dtype="uint32")' in code + _assert_roundtrip(func) + + +def test_scope_id_dtype_uint32_with_preferred(): + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (4,), "float32") + T.device_entry() + _ = T.cluster_id([2]) + cx, cy = T.cta_id_in_cluster([2, 2], preferred=[2, 2], dtype="uint32") + tx = T.thread_id([32]) + if tx == 0: + A[cx + cy] = T.float32(1) + # fmt: on + + code = func.script() + assert 'dtype="uint32"' in code + _assert_roundtrip(func) + + +def test_scope_id_dtype_uint32_deferred_extent(): + """The deferred (extent=None) form carries the dtype too.""" + + # fmt: off + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (32,), "float32") + T.device_entry() + _ = T.cta_id([1]) + lane = T.lane_id(dtype="uint32") + warp = T.warp_id([4]) + A[lane] = T.float32(warp) + # fmt: on + + scope_defs = [] + tvm.tirx.stmt_functor.post_order_visit( + func.body, + lambda s: ( + scope_defs.append(getattr(s, "def")) if isinstance(s, tvm.tirx.ScopeIdDefStmt) else None + ), + ) + deferred = [d for d in scope_defs if d.extents is None] + assert len(deferred) == 1 + assert deferred[0].def_ids[0].ty == PrimType("uint32") + _assert_roundtrip(func) + + +@pytest.mark.parametrize("dtype", ["int64", "float32"]) +def test_scope_id_dtype_rejects_unsupported(dtype): + # fmt: off + with pytest.raises(Exception, match='must be "int32" or "uint32"'): + + @T.prim_func + def func(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (128,), "float32") + T.device_entry() + _ = T.cta_id([1]) + tx = T.thread_id([128], dtype=dtype) + A[tx] = T.float32(1) + # fmt: on + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py index d0ab45d3289e..86f6fb0b29dc 100644 --- a/tests/python/tirx/transform/test_transform_lower_tirx.py +++ b/tests/python/tirx/transform/test_transform_lower_tirx.py @@ -644,6 +644,33 @@ def after(): compare(before, after, LowerTIRx) +def test_lower_uint32_scope_id_casts_at_bind(): + """uint32 scope ids get a Cast at the bind; launch params stay int32.""" + + @T.prim_func(private=True) + def before(): + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([128], dtype="uint32") + for k in T.serial(4, dtype="uint32"): + T.evaluate(tx + k) + + @T.prim_func(private=True) + def after(): + blockIdx_x = T.launch_thread("blockIdx.x", 1) + threadIdx_x = T.launch_thread("threadIdx.x", 128) + warp_id_in_cta: T.let[T.int32] = T.tvm_warp_shuffle( + T.uint32(4294967295), threadIdx_x // 32, 0, 32, 32 + ) + v: T.let[T.int32] = blockIdx_x + tx: T.let[T.uint32] = T.Cast("uint32", threadIdx_x) + T.evaluate(v) + for k in T.serial(T.uint32(4)): + T.evaluate(tx + k) + + compare(before, after, LowerTIRx) + + def test_lower_exec_context_infers_plain_predicate_for_dispatch(): import tvm.tirx.operator.tile_primitive as _ # noqa: F401 from tvm.tirx.operator.tile_primitive.dispatcher import register_dispatch