[TIRx][CUDA] uint32 index dtypes, directive fixes, and PTX ISA coverage

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<tvm::tirx::Var> KernelId(ffi::Array<PrimExpr> 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<tvm::tirx::Var> ScopeId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
+                                   ffi::String name, ffi::String cur,
+                                   PrimType dtype = PrimType::Int(32));
+
+ffi::Array<tvm::tirx::Var> ClusterId(ffi::Optional<ffi::Array<PrimExpr>> extents,
+                                     ffi::String parent, PrimType dtype = PrimType::Int(32));
+
+ffi::Array<tvm::tirx::Var> CtaId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
+                                 ffi::Optional<ffi::Array<PrimExpr>> preferred = std::nullopt,
+                                 PrimType dtype = PrimType::Int(32));
 
-ffi::Array<tvm::tirx::Var> CtaId(ffi::Array<PrimExpr> extents, ffi::String parent);
+ffi::Array<tvm::tirx::Var> CtaIdInPair(PrimType dtype = PrimType::Int(32));
 
-ffi::Array<tvm::tirx::Var> CtaIdInPair();
+ffi::Array<tvm::tirx::Var> WarpgroupId(ffi::Optional<ffi::Array<PrimExpr>> extents,
+                                       ffi::String parent, PrimType dtype = PrimType::Int(32));
 
-ffi::Array<tvm::tirx::Var> WarpId(ffi::Array<PrimExpr> extents, ffi::String parent);
+ffi::Array<tvm::tirx::Var> WarpId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
+                                  PrimType dtype = PrimType::Int(32));
 
-ffi::Array<tvm::tirx::Var> ThreadId(ffi::Array<PrimExpr> extents, ffi::String parent);
+ffi::Array<tvm::tirx::Var> ThreadId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
+                                    PrimType dtype = PrimType::Int(32));
 
 /*!
  * \brief The block initialization statement.
@@ -249,44 +271,53 @@ ffi::Array<Var> Remap(ffi::String kinds, ffi::Array<PrimExpr> 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<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
-                ffi::Optional<PrimExpr> step = std::nullopt);
+                ffi::Optional<PrimExpr> step = std::nullopt,
+                ffi::Optional<PrimType> 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<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
-                  ffi::Optional<PrimExpr> step = std::nullopt);
+                  ffi::Optional<PrimExpr> step = std::nullopt,
+                  ffi::Optional<PrimType> 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<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
-                    ffi::Optional<PrimExpr> step = std::nullopt);
+                    ffi::Optional<PrimExpr> step = std::nullopt,
+                    ffi::Optional<PrimType> 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<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
-                ffi::Optional<PrimExpr> step = std::nullopt);
+                ffi::Optional<PrimExpr> step = std::nullopt,
+                ffi::Optional<PrimType> 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<Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>> extents);
+ForFrame Grid(ffi::Array<Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>> extents,
+              ffi::Optional<PrimType> dtype = std::nullopt);
 
 /*!
  * \brief The assertion statement.

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

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'

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

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

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)