From 20449182e562244edf088d608728f6ea41936e2b Mon Sep 17 00:00:00 2001 From: Bohan Hou Date: Sat, 8 Aug 2026 13:16:08 -0400 Subject: [PATCH 1/4] [TIRx][TVMScript] Support uint32 loop and scope-id var dtypes (#53) Loop vars (T.serial/parallel/vectorized/unroll/grid) and scope-id vars (T.thread_id, T.lane_id, T.warp_id, T.cta_id, ...) were pinned to int32. Add an explicit per-site `dtype=` accepting "int32" (the default) or "uint32", so kernels can get u32 index registers, feed .u32 PTX operands directly, and avoid a manual Cast at every use site. - For constructor: promote IntImm bounds across signedness as well as width, with a representability check. - For frames: infer the loop var dtype using C-style signedness promotion instead of taking the code from `min` alone. This also fixes `for i in range()`, which previously tripped the loop_var/extent dtype equality check. - Scope ids: the dtype applies to the def vars only. Extents stay int32, so launch-param IterVars are unchanged and lowering casts at the bind. - Printer: emit `dtype=` for scope ids, whose dtype is not recoverable from the extents. Loops already round-trip via their bound dtypes. T.thread_binding stays int32-only, since the IterVar constructor requires an int-coded extent dtype. --- include/tvm/tirx/script/builder/ir.h | 54 +++- python/tvm/tirx/script/builder/ir.py | 135 +++++++--- src/tirx/ir/stmt.cc | 34 ++- src/tirx/script/builder/ir.cc | 140 +++++++--- src/tirx/script/printer/block.cc | 15 ++ .../python/tirx/codegen/test_codegen_cuda.py | 43 +++ tests/python/tirx/test_parser_printer.py | 246 ++++++++++++++++++ .../transform/test_transform_lower_tirx.py | 27 ++ 8 files changed, 594 insertions(+), 100 deletions(-) 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/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index 7cbae243f650..3c89d61f572e 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 @@ -1149,6 +1175,7 @@ def serial( annotations: dict[str, Any] | None = None, step: Expr | None = None, unroll: bool | None = None, + dtype: str | None = None, ) -> frame.ForFrame: """The serial For statement. @@ -1172,6 +1199,12 @@ def serial( If False, adds ``{"disable_unroll": True}`` annotation. Shorthand for ``annotations={"disable_unroll": True}``. + 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 ------- res : frame.ForFrame @@ -1189,7 +1222,7 @@ def serial( 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 +1231,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 +1249,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 +1264,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 +1273,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 +1291,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 +1306,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 +1315,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 +1333,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 +1348,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 +1399,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 +1409,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 +1420,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/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/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py index 1b817396388a..809d65e49e7d 100644 --- a/tests/python/tirx/codegen/test_codegen_cuda.py +++ b/tests/python/tirx/codegen/test_codegen_cuda.py @@ -1048,5 +1048,48 @@ 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 + + +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/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 51b1fb46ebaa..1b315a6fd0d1 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -2759,5 +2759,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 From 2417cfb047ab2a10ecfc32ca92fab564383f2055 Mon Sep 17 00:00:00 2001 From: Bohan Hou Date: Sat, 8 Aug 2026 14:03:43 -0400 Subject: [PATCH 2/4] [FIX][TIRx] Preserve CUDA sync and unroll directives (#54) --- python/tvm/backend/cuda/intrinsics/header.py | 9 +++----- python/tvm/backend/cuda/intrinsics/sync.py | 8 +++---- python/tvm/backend/cuda/op.py | 2 +- python/tvm/backend/cuda/script.py | 4 ++-- python/tvm/tirx/script/builder/ir.py | 20 ++++++++++++----- src/backend/cuda/codegen/codegen_cuda.cc | 17 ++++++++++++-- src/tirx/script/printer/for_loop.cc | 5 +++-- src/tirx/transform/lower_tirx_opaque.cc | 7 ++++-- .../python/tirx/codegen/test_codegen_cuda.py | 15 ++++++++++++- .../tirx/codegen/test_codegen_hopper.py | 4 +++- tests/python/tirx/test_parser_printer.py | 22 +++++++++++++++++++ 11 files changed, 86 insertions(+), 27 deletions(-) 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/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/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index 3c89d61f572e..bf2005befa7a 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -1174,7 +1174,7 @@ 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. @@ -1193,11 +1193,12 @@ 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 @@ -1212,10 +1213,17 @@ 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): 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/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 809d65e49e7d..69eda965f53e 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.) 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/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 1b315a6fd0d1..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.""" From 0dd67cc3f8095f4390ed60e1fe3baaff2c0ebf38 Mon Sep 17 00:00:00 2001 From: Bohan Hou Date: Sun, 9 Aug 2026 01:58:22 -0400 Subject: [PATCH 3/4] [TIRx][CUDA] Extend the PTX dialect across ISA 9.7.1-9.7.14 (#55) Registers the instructions the table-driven `T.ptx` dialect was missing across nine ISA sections, taking it from 175 to 314 entries and from 162,716 to 198,709 certified variants: 9.7.1 integer arithmetic add/sub/mul/mad (+.wide), mul24, mad24, sad, div, rem, abs, neg, popc, clz, bfind, brev, bfe, bfi, szext, bmsk, dp4a, dp2a 9.7.3 floating point mad, div, abs, sqrt, rsqrt, sin, cos, lg2, tanh, testp, copysign, rcp's f64 approximation 9.7.4 half precision fma (with .relu/.oob), neg, abs, tanh, ex2 9.7.5 mixed precision already covered by the srctype lines; verified cell by cell and test-pinned 9.7.6 comparison/selection set, setp, selp, slct 9.7.7 half comparison set, setp over the half type grid 9.7.8 logic and shift and, or, xor, not, cnot, lop3, shf, shl, shr 9.7.9 data movement scalar mov, prmt, ldu, prefetchu, applypriority, discard, isspacep, getctarank, cvt.pack, shfl.sync, the full cvta grid, st.async, multimem, createpolicy, cp.async.bulk.prefetch, tensormap.replace 9.7.14 parallel sync vote.sync, match.sync, activemask, redux.sync, elect.sync, atom.cas/.exch/ half/vector, red.async, the state-returning mbarrier arrives and non-parity waits, bar.red, tensormap.cp_fenceproxy Four mechanism extensions were needed, each mirroring an existing one: - `OperandSlot.dtype` may be a function of the modifier map, for the operands the ISA types by formula rather than by a written token (mul.wide's double-width result, dp4a's derived accumulator). Same contract as the existing `lanes`/`sinkable` callables. - `OperandSlot.pipe` joins adjacent operands with `|`, for the ISA's paired destinations (setp's p|q, lop3's d|p, shfl.sync, elect.sync). Sits beside the `bracket` mechanism and composes with it. - `b32i`/`b64i` name a bit-size type with the float carrier removed. A handful of instructions refuse a float register in one position while taking either integer signedness; pinning those operands to `.u32` would have dropped the signed spelling ptxas accepts. - The script printer now escapes the family name, so the three keyword mnemonics (and, or, not) round-trip as `T.ptx.and_`. Every registered variant is proven against ptxas rather than read off the spec: each check function was written after probing the cells it rejects, and several encode measured facts where the toolchain and the ISA text disagree -- st.async's release line takes no 8-bit type, `.ftz` on `set` follows the source precision, multimem pairs an ordering semantic with a scope, and bmsk takes no float register. Forms the ISA spells but this toolchain cannot assemble (clmad, multimem.st.async, multimem.red.async, the mbarrier layout facility) are excluded with the reason stated at the entry. --- python/tvm/backend/cuda/ptx/engine.py | 8 +- python/tvm/backend/cuda/ptx/render.py | 20 +- python/tvm/backend/cuda/ptx/table.py | 2672 ++++++++++++++++- python/tvm/script/tirx.pyi | 1106 ++++++- tests/python/tirx/codegen/test_ptx_cvt.py | 6 +- tests/python/tirx/codegen/test_ptx_dialect.py | 1187 +++++++- 6 files changed, 4801 insertions(+), 198 deletions(-) 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/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/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(): From 057c7fe29a18364f1c2fff4be8e9b0daf254aad3 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Sun, 9 Aug 2026 03:25:59 -0400 Subject: [PATCH 4/4] [FIX][TIRx] Skip the uint32 loop-var execution test without CUDA test_uint32_loop_var_runs_correctly launches the kernel on the device, but carried neither of the guards its siblings in this file use, so on a CPU-only node it failed with "Device API cuda is not enabled" instead of skipping. The neighbouring test_uint32_loop_var_and_scope_id_emit_unsigned only inspects the generated source, so it correctly stays unguarded. --- tests/python/tirx/codegen/test_codegen_cuda.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py index 69eda965f53e..daaa9f36d676 100644 --- a/tests/python/tirx/codegen/test_codegen_cuda.py +++ b/tests/python/tirx/codegen/test_codegen_cuda.py @@ -1077,6 +1077,8 @@ def main(A: T.Buffer((128,), "int32")): 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")):