This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new a6c7377bae [REFACTOR][TIRx] Keep AttrStmt node values unboxed (#20030)
a6c7377bae is described below
commit a6c7377bae6d721880a23afce9642f3676292ab4
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Jul 20 06:21:51 2026 +0800
[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.
---
python/tvm/tirx/script/builder/ir.py | 10 +++-------
python/tvm/tirx/script/builder/utils.py | 2 +-
python/tvm/tirx/stmt.py | 6 +++---
src/s_tir/transform/decorate_device_scope.cc | 2 +-
src/s_tir/transform/inject_software_pipeline.cc | 9 ++++-----
src/s_tir/transform/lower_opaque_block.cc | 2 +-
src/s_tir/transform/thread_storage_sync.cc | 6 ++----
src/tirx/ir/stmt.cc | 5 -----
src/tirx/script/builder/ir.cc | 7 +------
src/tirx/script/printer/stmt.cc | 6 +++---
src/tirx/transform/make_packed_api.cc | 2 +-
src/tirx/transform/tile_primitive_dispatch.cc | 6 ++----
12 files changed, 22 insertions(+), 41 deletions(-)
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index 560a3b152b..43a383a1d2 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/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)
diff --git a/python/tvm/tirx/script/builder/utils.py
b/python/tvm/tirx/script/builder/utils.py
index 04f8373452..6649d83a04 100644
--- a/python/tvm/tirx/script/builder/utils.py
+++ b/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):
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index 1c46f4e011..7270e04c8b 100644
--- a/python/tvm/tirx/stmt.py
+++ b/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__(
diff --git a/src/s_tir/transform/decorate_device_scope.cc
b/src/s_tir/transform/decorate_device_scope.cc
index 7af0d8eb72..d787a1bfcb 100644
--- a/src/s_tir/transform/decorate_device_scope.cc
+++ b/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;
}
diff --git a/src/s_tir/transform/inject_software_pipeline.cc
b/src/s_tir/transform/inject_software_pipeline.cc
index 21470d38b5..2c837e9ed3 100644
--- a/src/s_tir/transform/inject_software_pipeline.cc
+++ b/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(
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index a48dd8208b..88b366c9ef 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/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;
}
diff --git a/src/s_tir/transform/thread_storage_sync.cc
b/src/s_tir/transform/thread_storage_sync.cc
index 6253b64802..c38b2e2687 100644
--- a/src/s_tir/transform/thread_storage_sync.cc
+++ b/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);
}
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 2a5907fb6f..c768e450a6 100644
--- a/src/tirx/ir/stmt.cc
+++ b/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);
});
}
diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc
index c1f37c7f7b..1f198b4238 100644
--- a/src/tirx/script/builder/ir.cc
+++ b/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())
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index ffafa41494..8601ecb62a 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/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;
}
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index 162ae78bab..8beb6d918c 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -249,7 +249,7 @@ PrimFunc MakePackedAPI(PrimFunc func) {
ffi::symbol::tvm_ffi_symbol_prefix +
global_symbol.value()}});
Stmt body = ReturnRewriter(v_result)(func_ptr->body);
- body = AttrStmt(IntImm::Int32(0), attr::compute_scope, StringImm(name_hint +
"_compute_"), body);
+ body = AttrStmt(0, attr::compute_scope, StringImm(name_hint + "_compute_"),
body);
// Set device context
if (need_set_device) {
ffi::Any node = ffi::String("default");
diff --git a/src/tirx/transform/tile_primitive_dispatch.cc
b/src/tirx/transform/tile_primitive_dispatch.cc
index de80b9cfd3..5322c45b37 100644
--- a/src/tirx/transform/tile_primitive_dispatch.cc
+++ b/src/tirx/transform/tile_primitive_dispatch.cc
@@ -592,8 +592,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator {
// dispatched impls too.
void PrepareLaunchParams(const AttrStmtNode* entry_node, Stmt body,
std::vector<std::pair<Var, PrimExpr>>* scope_binds)
{
- Stmt gather_target =
- AttrStmt(IntImm::Int32(0), tvm::tirx::attr::kDeviceEntry,
IntImm::Bool(true), body);
+ Stmt gather_target = AttrStmt(0, tvm::tirx::attr::kDeviceEntry,
IntImm::Bool(true), body);
std::vector<ScopeIdDefWithSource> gathered =
ScopeIdDefGather::Gather(gather_target);
Array<ScopeIdDef> defs;
defs.reserve(gathered.size());
@@ -624,8 +623,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator {
std::vector<std::pair<Var, const StmtNode*>>*
implicit_scope_id_evals) {
// Gather from a temporary stmt synthesized as the device-entry marker
// so direct ScopeIdDefStmt children are attributed back to entry_node.
- Stmt gather_target =
- AttrStmt(IntImm::Int32(0), tvm::tirx::attr::kDeviceEntry,
IntImm::Bool(true), body);
+ Stmt gather_target = AttrStmt(0, tvm::tirx::attr::kDeviceEntry,
IntImm::Bool(true), body);
std::vector<ScopeIdDefWithSource> gathered =
ScopeIdDefGather::Gather(gather_target);
// Remap the synthetic source pointer back to the real entry_node so the
// injector matches against the actual node present in the post-processed
