[REFACTOR][TIRx] Keep AttrStmt node values unboxed (#20030)

AttrStmt.node is an ffi::Any field, so converting POD arguments to
PrimExpr makes its representation depend on the caller rather than the
declared container type.

This change preserves values passed through AttrStmt and T.attr, uses
raw zero sentinel nodes consistently, and updates the printer
canonicalization.

Author: tqchen

python/tvm/tirx/script/builder/ir.py

@@ -1582,7 +1582,7 @@ def attr(
         with T.attr(node, key, value):
             ...
 
-    Usage 2 — dict sugar (node defaults to ``T.int32(0)``)::
+    Usage 2 — dict sugar (node defaults to ``0``)::
 
         with T.attr({"key1": value1, "key2": value2}):
             ...
@@ -1591,7 +1591,7 @@ def attr(
     ----------
     node_or_dict : Any
         If a dict, each key-value pair becomes an AttrStmt with
-        ``node=T.int32(0)``.  Otherwise the node to annotate.
+        ``node=0``.  Otherwise the node to annotate.
 
     attr_key : str, optional
         Attribute type key (required when ``node_or_dict`` is not a dict).
@@ -1609,11 +1609,7 @@ def attr(
         for k, v in node_or_dict.items():
             if isinstance(v, bool):
                 v = IntImm("bool", v)
-            frames.append(
-                _ffi_api.Attr(  # type: ignore[attr-defined]
-                    convert(IntImm("int32", 0)), k, convert(v)
-                )
-            )
+            frames.append(_ffi_api.Attr(0, k, convert(v)))  # type: ignore[attr-defined]
         if len(frames) == 1:
             return frames[0]
         return utils._FrameScope(frames)

python/tvm/tirx/script/builder/utils.py

@@ -108,7 +108,7 @@ def seq_scope():
                     T.evaluate(j)
             result = ib.get()
     """
-    return T.attr(tirx.const(0, "int32"), "pragma_scope", tirx.StringImm("seq"))
+    return T.attr(0, "pragma_scope", tirx.StringImm("seq"))
 
 
 def _unravel_index(index, shape):

python/tvm/tirx/stmt.py

@@ -491,7 +491,7 @@ class AttrStmt(Stmt):
 
     Parameters
     ----------
-    node : Object
+    node : Any
         The node to annotate the attribute
 
     attr_key : str
@@ -507,14 +507,14 @@ class AttrStmt(Stmt):
         The location of the stmt in the source code.
     """
 
-    node: Object
+    node: Any
     attr_key: str
     value: Expr
     body: Stmt
     span: Span | None
 
     def __init__(
-        self, node: Object, attr_key: str, value: Expr, body: Stmt, span: Span | None = None
+        self, node: Any, attr_key: str, value: Expr, body: Stmt, span: Span | None = None
     ) -> None:
         body = _normalize_legacy_stmt(body)
         self.__init_handle_by_constructor__(

src/s_tir/transform/decorate_device_scope.cc

@@ -31,7 +31,7 @@ namespace s_tir {
 using namespace tvm::tirx;
 
 Stmt DecorateDeviceScopeImpl(Stmt&& stmt) {
-  Stmt body = AttrStmt(IntImm::Int32(0), tirx::attr::device_scope, 0, stmt);
+  Stmt body = AttrStmt(0, tirx::attr::device_scope, 0, stmt);
   return body;
 }
 

src/s_tir/transform/inject_software_pipeline.cc

@@ -758,10 +758,9 @@ class PipelineRewriter : public StmtExprMutator {
         auto attach_wait_scope = [&new_blocks](int i, int stage_id, PrimExpr wait_count) {
           auto& block = new_blocks[i].block;
           SBlockNode* n = block.CopyOnWrite();
-          auto zero = IntImm::Int32(0);
           n->body =
-              AttrStmt(zero, s_tir::attr::async_wait_queue_scope, stage_id,
-                       AttrStmt(zero, s_tir::attr::async_wait_inflight_count, wait_count, n->body));
+              AttrStmt(0, s_tir::attr::async_wait_queue_scope, stage_id,
+                       AttrStmt(0, s_tir::attr::async_wait_inflight_count, wait_count, n->body));
         };
 
         if (state.predicate && !ana_normalized->CanProve(state.predicate.value())) {
@@ -805,7 +804,7 @@ class PipelineRewriter : public StmtExprMutator {
 
         for (auto body : group_bodies) {
           auto commit_queue_scope =
-              AttrStmt(IntImm::Int32(0), s_tir::attr::async_commit_queue_scope, stage_id, body);
+              AttrStmt(0, s_tir::attr::async_commit_queue_scope, stage_id, body);
           auto new_block = MakeSBlock(commit_queue_scope, buffer_data_to_buffer_);
           stmts.push_back(SBlockRealize({}, predicate, new_block));
         }
@@ -925,7 +924,7 @@ class PipelineRewriter : public StmtExprMutator {
         }
 
         SBlockNode* n = new_block.CopyOnWrite();
-        n->body = AttrStmt(IntImm::Int32(0), s_tir::attr::async_scope, 1, n->body);
+        n->body = AttrStmt(0, s_tir::attr::async_scope, 1, n->body);
       }
 
       new_blocks.push_back(

src/s_tir/transform/lower_opaque_block.cc

@@ -80,7 +80,7 @@ class OpaqueBlockLower : public StmtExprMutator {
     std::vector<std::pair<std::string, PrimExpr>> pragma_attrs;
     HandleAnnotations(new_block->annotations, &pragma_attrs, /*is_block=*/true);
     for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) {
-      body = AttrStmt(IntImm::Int32(0), it->first, it->second, std::move(body));
+      body = AttrStmt(0, it->first, it->second, std::move(body));
     }
     return body;
   }

src/s_tir/transform/thread_storage_sync.cc

@@ -298,11 +298,9 @@ class ThreadSyncAfterWaitQueueInserter : public StmtExprMutator {
               .as_or_throw<PrimExpr>());
       auto inner = op->body.as<AttrStmtNode>();
       TVM_FFI_ICHECK(inner && inner->attr_key == s_tir::attr::async_wait_inflight_count);
-      auto zero = IntImm::Int32(0);
       auto new_body = SeqStmt({sync, inner->body});
-      return AttrStmt(
-          zero, s_tir::attr::async_wait_queue_scope, op->value,
-          AttrStmt(zero, s_tir::attr::async_wait_inflight_count, inner->value, new_body));
+      return AttrStmt(0, s_tir::attr::async_wait_queue_scope, op->value,
+                      AttrStmt(0, s_tir::attr::async_wait_inflight_count, inner->value, new_body));
     }
     return StmtExprMutator::VisitStmt_(op);
   }

src/tirx/ir/stmt.cc

@@ -91,11 +91,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def("tirx.AttrStmt",
                         [](Any node, ffi::String attr_key, PrimExpr value, Stmt body, Span span) {
-                          // when node is a POD data type like int or bool, first convert to
-                          // primexpr.
-                          if (node.type_index() < ffi::TypeIndex::kTVMFFISmallStr) {
-                            return AttrStmt(node.cast<PrimExpr>(), attr_key, value, body, span);
-                          }
                           return AttrStmt(node, attr_key, value, body, span);
                         });
 }

src/tirx/script/builder/ir.cc

@@ -644,10 +644,6 @@ LaunchThreadFrame LaunchThread(ffi::String thread_tag, PrimExpr extent) {
 }
 
 AttrFrame Attr(ffi::Any node, ffi::String attr_key, PrimExpr value) {
-  // convert POD value to PrimExpr
-  if (node.type_index() < ffi::TypeIndex::kTVMFFISmallStr) {
-    node = node.cast<PrimExpr>();
-  }
   ffi::ObjectPtr<AttrFrameNode> n = ffi::make_object<AttrFrameNode>();
   n->node = std::move(node);
   n->attr_key = attr_key;
@@ -664,8 +660,7 @@ AttrFrame DeviceEntry() {
   // enclosing PrimFuncFrame: ``IRBuilderFrameNode::ExitWithScope`` runs
   // callbacks before popping itself, so the AttrFrame is closed and its
   // emitted ``AttrStmt`` lands in the PrimFunc's body sequence.
-  AttrFrame frame =
-      Attr(IntImm::Int32(0), ffi::String(tvm::tirx::attr::kDeviceEntry), IntImm::Bool(true));
+  AttrFrame frame = Attr(0, ffi::String(tvm::tirx::attr::kDeviceEntry), IntImm::Bool(true));
   IRBuilder builder = IRBuilder::Current();
   ffi::Optional<PrimFuncFrame> pf_frame = builder->FindFrame<PrimFuncFrame>();
   TVM_FFI_ICHECK(pf_frame.has_value())

src/tirx/script/printer/stmt.cc

@@ -796,10 +796,10 @@ ExprDoc DocsifyLaunchThread(const tirx::AttrStmt& attr_stmt, const AccessPath& a
       });
 }
 
-/*! \brief Check whether an AttrStmt has node=IntImm(int32, 0) (the dict-attr pattern). */
+/*! \brief Check whether an AttrStmt has node=0 (the dict-attr pattern). */
 static bool IsDictAttrPattern(const tirx::AttrStmt& stmt) {
-  if (auto int_imm = stmt->node.as<IntImmNode>()) {
-    return int_imm->ty.as_or_throw<PrimType>() == PrimType::Int(32) && int_imm->value == 0;
+  if (auto int_value = stmt->node.as<int64_t>()) {
+    return int_value.value() == 0;
   }
   return false;
 }