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 3d3343e1da [REFACTOR][TIRx] Remove obsolete AttrStmt infrastructure
(#20573)
3d3343e1da is described below
commit 3d3343e1daa63ef47025893151574dc133a5d115
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 6 15:23:11 2026 -0400
[REFACTOR][TIRx] Remove obsolete AttrStmt infrastructure (#20573)
Remove the generic `AttrStmt` node and `T.attr` construction API now
that scoped operations use `RegionStmt`. Retire the corresponding
printer, FFI, visitor, pass, and backend infrastructure.
Opaque lowering preserves loop annotations directly, including CUDA
unroll hints and inherited automatic-unroll policies, without converting
generic `pragma_*` keys into statement wrappers.
---
include/tvm/s_tir/stmt.h | 2 +-
include/tvm/tirx/script/ir_builder/frame.h | 50 ----------
include/tvm/tirx/script/ir_builder/ir.h | 9 --
include/tvm/tirx/stmt.h | 59 ------------
include/tvm/tirx/stmt_functor.h | 6 --
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/script/ir_builder/frame.py | 4 -
.../tvm/tirx/script/ir_builder/parser_protocol.py | 53 +---------
python/tvm/tirx/stmt.py | 41 --------
src/backend/cuda/transforms/lower_iket.cc | 7 --
src/backend/trn/codegen/codegen_trn.cc | 6 --
src/backend/trn/codegen/codegen_trn.h | 1 -
src/backend/vulkan/codegen/codegen_spirv.cc | 6 --
src/backend/vulkan/codegen/codegen_spirv.h | 1 -
src/s_tir/analysis/calculate_allocated_memory.cc | 6 --
src/s_tir/analysis/estimate_flops.cc | 5 -
.../feature_extractor/per_store_feature.cc | 14 ---
src/s_tir/transform/hoist_expression.cc | 35 -------
src/s_tir/transform/inject_virtual_thread.cc | 18 ----
src/s_tir/transform/lower_opaque_block.cc | 100 +++----------------
src/s_tir/transform/lower_thread_binding.cc | 14 ---
src/script/printer/script_printer.cc | 1 -
src/target/llvm/codegen_cpu.cc | 21 ----
src/target/llvm/codegen_cpu.h | 1 -
src/target/llvm/codegen_llvm.cc | 7 --
src/target/llvm/codegen_llvm.h | 1 -
src/target/source/codegen_c.cc | 6 --
src/target/source/codegen_c.h | 1 -
src/tirx/analysis/verify_tirx_well_formed.cc | 3 -
src/tirx/analysis/verify_well_formed.h | 22 +----
src/tirx/ir/ir_mutator_with_analyzer.cc | 5 -
src/tirx/ir/ir_mutator_with_analyzer.h | 1 -
src/tirx/ir/ir_visitor_with_analyzer.cc | 4 -
src/tirx/ir/ir_visitor_with_analyzer.h | 1 -
src/tirx/ir/stmt.cc | 82 ----------------
src/tirx/ir/stmt_functor.cc | 41 --------
src/tirx/ir/tir_visitor_with_path.cc | 8 --
src/tirx/ir/tir_visitor_with_path.h | 1 -
src/tirx/script/ir_builder/frame.cc | 6 --
src/tirx/script/ir_builder/ir.cc | 10 +-
src/tirx/script/printer/stmt.cc | 59 ------------
src/tirx/transform/common_subexpr_elim.cc | 18 +---
src/tirx/transform/ir_utils.cc | 86 -----------------
src/tirx/transform/ir_utils.h | 2 -
src/tirx/transform/lower_thread_allreduce.h | 2 +-
src/tirx/transform/lower_tirx_opaque.cc | 107 +++------------------
src/tirx/transform/lower_tvm_builtin.cc | 16 +--
src/tirx/transform/remove_no_op.cc | 5 -
src/tirx/transform/storage_rewrite.cc | 42 +-------
src/tirx/transform/tvm_ffi_binder.h | 6 +-
src/tirx/transform/unsupported_dtype_legalize.cc | 25 -----
.../transform/test_s_tir_transform_hoist_if.py | 4 +-
.../test_s_tir_transform_lower_opaque_block.py | 38 --------
tests/python/tirx-base/test_tir_constructor.py | 10 --
tests/python/tirx-base/test_tir_nodes.py | 5 -
.../test_tir_transform_convert_ssa.py | 50 ----------
.../test_tir_transform_split_host_device.py | 2 +-
.../python/tirx/script/test_tirx_script_printer.py | 14 ---
58 files changed, 45 insertions(+), 1107 deletions(-)
diff --git a/include/tvm/s_tir/stmt.h b/include/tvm/s_tir/stmt.h
index 8f45f5dc22..3f94232416 100644
--- a/include/tvm/s_tir/stmt.h
+++ b/include/tvm/s_tir/stmt.h
@@ -91,7 +91,7 @@ class MatchBufferRegion : public ffi::ObjectRef {
* T.where(predicate)
* buffer2 = T.alloc_tensor(shape, dtype)
* buffer3 = Ts.match_buffer(source_buffer[start:end, ...])
- * T.attr({attr_key: attr_value, ...})
+ * Ts.sblock_attr({attr_key: attr_value, ...})
* with T.init():
* // init body
* // body
diff --git a/include/tvm/tirx/script/ir_builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
index 28327c515a..a1621d1cb7 100644
--- a/include/tvm/tirx/script/ir_builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -289,56 +289,6 @@ class RegionFrame : public TIRFrame {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(RegionFrame, TIRFrame,
RegionFrameNode);
};
-/*!
- * \brief A frame that represents attribute node.
- *
- * \sa AttrFrame
- */
-class AttrFrameNode : public TIRFrameNode {
- public:
- explicit AttrFrameNode(ffi::UnsafeInit tag) : value(tag) {}
-
- explicit AttrFrameNode(Expr value) : value(std::move(value)) {}
-
- /*! \brief The node to annotate the attribute. */
- Any node;
- /*! \brief Attribute type key. */
- ffi::String attr_key;
- /*! \brief The value of the attribute. */
- Expr value;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AttrFrameNode>()
- .def_ro("node", &AttrFrameNode::node)
- .def_ro("attr_key", &AttrFrameNode::attr_key)
- .def_ro("value", &AttrFrameNode::value);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AttrFrame",
AttrFrameNode,
- TIRFrameNode);
-
- public:
- /*!
- * \brief The method called when exiting RAII scope.
- * \sa tvm::support::With
- */
- void ExitWithScope() final;
-};
-
-/*!
- * \brief Managed reference to AttrFrameNode.
- *
- * \sa AttrFrameNode
- */
-class AttrFrame : public TIRFrame {
- public:
- explicit AttrFrame(ffi::ObjectPtr<AttrFrameNode> data) :
TIRFrame(ffi::UnsafeInit{}) {
- TVM_FFI_ICHECK(data != nullptr);
- data_ = std::move(data);
- }
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrFrame, TIRFrame,
AttrFrameNode);
-};
-
/*!
* \brief A frame that represents while loop.
*
diff --git a/include/tvm/tirx/script/ir_builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
index 07122fb5dd..2c35dcfda0 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -238,15 +238,6 @@ AssertFrame Assert(PrimExpr condition, ffi::String
error_kind,
Var Bind(Expr value, ffi::Optional<Type> type_annotation = std::nullopt,
ffi::Optional<Var> var = std::nullopt);
-/*!
- * \brief Create an attribute.
- * \param node The node to annotate the attribute.
- * \param attr_key Attribute type key.
- * \param value The value of the attribute.
- * \return The result AttrFrame.
- */
-AttrFrame Attr(ffi::Any node, ffi::String attr_key, Expr value);
-
/*!
* \brief Create a while loop.
* \param condition The termination condition of the loop.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 48fbb8620b..5b414cef47 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -115,56 +115,6 @@ class Bind : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(BindNode);
};
-/*!
- * \brief Define certain auxiliary attribute for the body to be a symbolic
value.
- * This provide auxiliary information for IR passes that transforms body.
- *
- * In terms of effect, this is equivalent to Block(Evaluate(value), body).
- *
- * Examples of possible usage:
- * - Bound of function, variables.
- * - Hint which block corresponds to a parallel region.
- */
-class AttrStmtNode : public StmtNode {
- public:
- explicit AttrStmtNode(ffi::UnsafeInit tag) : value(tag), body(tag) {}
-
- AttrStmtNode(Expr value, Stmt body) : value(std::move(value)),
body(std::move(body)) {}
-
- /*! \brief this is attribute about certain node */
- ffi::Any node;
- /*! \brief the type key of the attribute */
- ffi::String attr_key;
- /*! \brief The attribute value, value is well defined at current scope. */
- Expr value;
- /*! \brief The body statement to be executed */
- Stmt body;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<AttrStmtNode>()
- .def_ro("node", &AttrStmtNode::node)
- .def_ro("attr_key", &AttrStmtNode::attr_key)
- .def_ro("value", &AttrStmtNode::value)
- .def_ro("body", &AttrStmtNode::body);
- }
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.AttrStmt", AttrStmtNode, StmtNode);
-};
-
-/*!
- * \brief Managed reference to AttrStmtNode.
- * \sa AttrStmtNode
- */
-class AttrStmt : public Stmt {
- public:
- TVM_DLL AttrStmt(ffi::Any node, ffi::String attr_key, Expr value, Stmt body,
Span span = Span());
-
- explicit AttrStmt(ffi::ObjectPtr<AttrStmtNode> node) : Stmt(std::move(node))
{}
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AttrStmt, Stmt, AttrStmtNode);
- TVM_DEFINE_OBJECT_REF_COW_METHOD(AttrStmtNode);
-};
-
/*!
* \brief A single-body statement whose semantics are defined by an operator.
* Operands are evaluated in the enclosing scope. Body parameters are
definitions
@@ -876,15 +826,6 @@ constexpr const char* buffer_allocated_addr =
"buffer_allocated_addr";
*/
constexpr const char* kPersistentKernel = "tirx.persistent_kernel";
-/*!
- * \brief Check if attr_key is a pragma key extension
- * \param attr_key The attr key to be compared
- * \return true if it is a pragma key
- */
-inline bool IsPragmaKey(const std::string& attr_key) {
- return attr_key.compare(0, 7, "pragma_") == 0;
-}
-
} // namespace attr
/*!
* \brief Create a type annotation expression
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index b0268f8b8c..6762f53769 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -80,9 +80,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
virtual R Dispatch_(const RegionStmtNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
- virtual R Dispatch_(const AttrStmtNode* node, Args... args) {
- return DispatchDefault_(node, std::forward<Args>(args)...);
- }
virtual R Dispatch_(const IfThenElseNode* node, Args... args) {
return DispatchDefault_(node, std::forward<Args>(args)...);
}
@@ -137,7 +134,6 @@ class StmtFunctor<R(const Stmt&, Args...)> {
return self->DispatchDefault_(node.get(),
std::forward<Args>(args)...);
});
SetDispatch<TSelf, BindNode>(vtable);
- SetDispatch<TSelf, AttrStmtNode>(vtable);
SetDispatch<TSelf, RegionStmtNode>(vtable);
SetDispatch<TSelf, IfThenElseNode>(vtable);
SetDispatch<TSelf, ForNode>(vtable);
@@ -195,7 +191,6 @@ class TVM_DLL StmtExprVisitor : public tvm::ExprVisitor {
using tvm::ExprVisitor::Visit_;
virtual ffi::Optional<VisitInterrupt> Visit_(const BindNode* op);
- virtual ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const IfThenElseNode* op);
virtual ffi::Optional<VisitInterrupt> Visit_(const ForNode* op);
@@ -264,7 +259,6 @@ class TVM_DLL StmtExprMutator : public tvm::ExprMutator {
}
virtual UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode
inplace_mode);
- virtual UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode
inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const IfThenElseNode* op, InplaceMode
inplace_mode);
virtual UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode
inplace_mode);
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 60da49ed28..fbae790743 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -48,7 +48,7 @@ from .stmt import Stmt, Bind, AssertStmt, ForKind, For,
While, Return, Break, Co
# Legacy alias: LetStmt was folded into the body-less Bind statement.
LetStmt = Bind
-from .stmt import BufferStore, AttrStmt, RegionStmt
+from .stmt import BufferStore, RegionStmt
from .stmt import SeqStmt
from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
diff --git a/python/tvm/tirx/script/ir_builder/frame.py
b/python/tvm/tirx/script/ir_builder/frame.py
index cf1065c8de..9ddb9a4d46 100644
--- a/python/tvm/tirx/script/ir_builder/frame.py
+++ b/python/tvm/tirx/script/ir_builder/frame.py
@@ -68,10 +68,6 @@ class ForFrame(TIRFrame):
class AssertFrame(TIRFrame): ...
-@_register_object("script.ir_builder.tirx.AttrFrame")
-class AttrFrame(TIRFrame): ...
-
-
@_register_object("script.ir_builder.tirx.WhileFrame")
class WhileFrame(TIRFrame): ...
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index 34eb8b90e3..a61379f4bc 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -35,7 +35,6 @@ from tvm import ir as _ir
from tvm import tirx as _tir
from tvm.ir import StringImm as _StringImm
from tvm.ir import TensorRegion, Type, is_prim_expr
-from tvm.runtime import convert
from tvm.script.ir_builder import base as _base
from tvm.script.ir_builder.base import AlreadyEmitted
from tvm.script.ir_builder.base import IRBuilder as _IRBuilder
@@ -44,7 +43,7 @@ from tvm.tirx.expr import (
IntImm,
)
-from . import _ffi_api, frame, utils
+from . import _ffi_api, frame
from . import ir as _native
from . import op as _op
from .op import and_ as and_
@@ -438,55 +437,6 @@ def bind( # pylint: disable=invalid-name
return _ffi_api.Bind(value, type_annotation, var) # type:
ignore[attr-defined] # pylint: disable=no-member
-def attr(
- node_or_dict: Any, attr_key: str | None = None, value: Expr | str | None =
None
-) -> frame.AttrFrame | utils._FrameScope:
- """Create an attribute node, or multiple attribute nodes from a dict.
-
- Usage 1 — single attr::
-
- with T.attr(node, key, value):
- ...
-
- Usage 2 — dict sugar (node defaults to ``0``)::
-
- with T.attr({"key1": value1, "key2": value2}):
- ...
-
- Parameters
- ----------
- node_or_dict : Any
- If a dict, each key-value pair becomes an AttrStmt with
- ``node=0``. Otherwise the node to annotate.
-
- attr_key : str, optional
- Attribute type key (required when ``node_or_dict`` is not a dict).
-
- value : Union[Expr, str], optional
- The attribute value (required when ``node_or_dict`` is not a dict).
-
- Returns
- -------
- res : Union[frame.AttrFrame, _FrameScope]
- A single AttrFrame, or a _FrameScope wrapping multiple AttrFrames.
- """
- if isinstance(node_or_dict, dict):
- frames = []
- for k, v in node_or_dict.items():
- if isinstance(v, bool):
- v = IntImm("bool", 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)
- else:
- if attr_key is None or value is None:
- raise ValueError("T.attr(node, attr_key, value) requires all three
arguments")
- node_or_dict = convert(node_or_dict)
- value = convert(value)
- return _ffi_api.Attr(node_or_dict, attr_key, value) # type:
ignore[attr-defined] # pylint: disable=no-member
-
-
def buffer_store(
buffer: Var, # pylint: disable=redefined-outer-name
value: Expr,
@@ -1028,7 +978,6 @@ __all__ = [
"and_",
"arg_",
"assert_",
- "attr",
"bind",
"bind_",
"break_",
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index e91695c811..8c16ca078b 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -319,47 +319,6 @@ class RegionStmt(Stmt):
)
-@tvm_ffi.register_object("tirx.AttrStmt")
-class AttrStmt(Stmt):
- """AttrStmt node.
-
- Parameters
- ----------
- node : Any
- The node to annotate the attribute
-
- attr_key : str
- Attribute type key.
-
- value : Expr
- The value of the attribute
-
- body : Stmt
- The body statement.
-
- span : Optional[Span]
- The location of the stmt in the source code.
- """
-
- node: Any
- attr_key: str
- value: Expr
- body: Stmt
- span: Span | None
-
- def __init__(
- self, node: Any, attr_key: str, value: Expr, body: Stmt, span: Span |
None = None
- ) -> None:
- self.__init_handle_by_constructor__(
- _ffi_api.AttrStmt,
- node,
- attr_key,
- value,
- body,
- span, # type: ignore
- )
-
-
@tvm_ffi.register_object("tirx.SeqStmt")
class SeqStmt(Stmt):
"""Sequence of statements.
diff --git a/src/backend/cuda/transforms/lower_iket.cc
b/src/backend/cuda/transforms/lower_iket.cc
index 64138c9f8e..136b6e2664 100644
--- a/src/backend/cuda/transforms/lower_iket.cc
+++ b/src/backend/cuda/transforms/lower_iket.cc
@@ -638,13 +638,6 @@ class RemoveStrippedIketNoOps : public StmtExprMutator {
return SideEffect(condition) > CallEffectKind::kReadState ?
Evaluate(condition) : Evaluate(0);
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* attr_stmt, InplaceMode
inplace_mode) final {
- auto result = StmtExprMutator::Mutate_(attr_stmt, inplace_mode);
- if (!result.IsUnchanged()) attr_stmt =
ffi::AnyView(result).as<AttrStmtNode>();
- if (IsEvaluateZero(attr_stmt->body)) return attr_stmt->body;
- return result;
- }
-
UnchangedOr<Stmt> Mutate_(const ForNode* loop, InplaceMode inplace_mode)
final {
auto result = StmtExprMutator::Mutate_(loop, inplace_mode);
if (!result.IsUnchanged()) loop = ffi::AnyView(result).as<ForNode>();
diff --git a/src/backend/trn/codegen/codegen_trn.cc
b/src/backend/trn/codegen/codegen_trn.cc
index 6d611a0700..e1bc98c4af 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -277,12 +277,6 @@ void CodeGenTrainium::DispatchAllocTensor(const BindNode*
op, const CallNode* bu
}
}
-void CodeGenTrainium::Dispatch_(const AttrStmtNode* op) {
- TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
- << "Launch attributes are retired; use tirx.launch_thread RegionStmt";
- this->PrintStmt(op->body);
-}
-
void CodeGenTrainium::Dispatch_(const RegionStmtNode* op) {
static const Op tensorized_instruction =
Op::Get("tirx.nki.tensorized_instruction");
if (op->op.same_as(tensorized_instruction)) {
diff --git a/src/backend/trn/codegen/codegen_trn.h
b/src/backend/trn/codegen/codegen_trn.h
index 5f1d16441f..45a1f04ce4 100644
--- a/src/backend/trn/codegen/codegen_trn.h
+++ b/src/backend/trn/codegen/codegen_trn.h
@@ -60,7 +60,6 @@ class CodeGenTrainium final : public CodeGenC {
void PrintType(const PrimType& t, std::ostream& os) final; // NOLINT(*)
void Dispatch_(const BindNode* op) final;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
- void Dispatch_(const AttrStmtNode* op) final; //
NOLINT(*)
void Dispatch_(const RegionStmtNode* op) final; //
NOLINT(*)
void Dispatch_(const ForNode* op) final; //
NOLINT(*)
void Dispatch_(const BufferStoreNode* op) final; //
NOLINT(*)=
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 1ec8d0498b..c9f6d1f5b5 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -1006,12 +1006,6 @@ void CodeGenSPIRV::Dispatch_(const RegionStmtNode* op) {
}
}
-void CodeGenSPIRV::Dispatch_(const AttrStmtNode* op) {
- TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
- << "Launch attributes are retired; use tirx.launch_thread RegionStmt";
- this->Dispatch(op->body);
-}
-
void CodeGenSPIRV::Dispatch_(const AssertStmtNode* op) {
// AssertStmt is a leaf — no body to visit.
}
diff --git a/src/backend/vulkan/codegen/codegen_spirv.h
b/src/backend/vulkan/codegen/codegen_spirv.h
index b7732f88c0..867fac6d5c 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.h
+++ b/src/backend/vulkan/codegen/codegen_spirv.h
@@ -121,7 +121,6 @@ class CodeGenSPIRV : public
tirx::ExprFunctor<spirv::Value(const Expr&)>,
void DispatchDeclTensor(const BindNode* op, const CallNode* buffer_call);
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const RegionStmtNode* op) override;
- void Dispatch_(const AttrStmtNode* op) override;
void Dispatch_(const AssertStmtNode* op) override;
void Dispatch_(const BindNode* op) override;
void Dispatch_(const SeqStmtNode* op) override;
diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc
b/src/s_tir/analysis/calculate_allocated_memory.cc
index 21bfaaa3a6..696a74fe77 100644
--- a/src/s_tir/analysis/calculate_allocated_memory.cc
+++ b/src/s_tir/analysis/calculate_allocated_memory.cc
@@ -107,12 +107,6 @@ class AllocTensorCalculator : public StmtExprVisitor {
_current_size = snapshot;
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) override {
- auto snapshot = _current_size;
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
- _current_size = snapshot;
- return std::nullopt;
- }
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) override {
auto snapshot = _current_size;
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
diff --git a/src/s_tir/analysis/estimate_flops.cc
b/src/s_tir/analysis/estimate_flops.cc
index 49ff6521b2..5ac83f32ec 100644
--- a/src/s_tir/analysis/estimate_flops.cc
+++ b/src/s_tir/analysis/estimate_flops.cc
@@ -166,11 +166,6 @@ class FlopEstimator : private
tirx::ExprFunctor<TResult(const Expr& n)>,
}
TResult Dispatch_(const TensorLoadNode* op) override { return TResult(); }
- TResult Dispatch_(const AttrStmtNode* op) override {
- TResult result = Dispatch(op->body);
- result += Dispatch(op->value);
- return result;
- }
TResult Dispatch_(const RegionStmtNode* op) override {
TResult result = Dispatch(op->body);
for (const Expr& arg : op->args) result += Dispatch(arg);
diff --git a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
index 4f6bbe0351..d50c5c6ec6 100644
--- a/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
+++ b/src/s_tir/meta_schedule/feature_extractor/per_store_feature.cc
@@ -358,20 +358,6 @@ Pass SimplifyForFeatureExtraction(bool
normalize_thread_bindings = false) {
}
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode
inplace_mode) final {
- if (!normalize_thread_bindings_) return StmtExprMutator::Mutate_(op,
inplace_mode);
- auto node = Mutate(op->node, inplace_mode);
- auto value = Mutate(op->value, inplace_mode);
- auto body = Mutate(op->body, inplace_mode);
- if (node.UnchangedOrSameAs(op->node) &&
value.UnchangedOrSameAs(op->value) &&
- body.UnchangedOrSameAs(op->body)) {
- return ffi::Unchanged();
- }
- return AttrStmt(std::move(node).ValueOrUnchanged(op->node), op->attr_key,
- std::move(value).ValueOrUnchanged(op->value),
- std::move(body).ValueOrUnchanged(op->body), op->span);
- }
-
UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode)
final {
if (!normalize_thread_bindings_) return StmtExprMutator::Mutate_(op,
inplace_mode);
auto annotations = Mutate(op->annotations, inplace_mode)
diff --git a/src/s_tir/transform/hoist_expression.cc
b/src/s_tir/transform/hoist_expression.cc
index c641aeeb2a..1a21024639 100644
--- a/src/s_tir/transform/hoist_expression.cc
+++ b/src/s_tir/transform/hoist_expression.cc
@@ -270,30 +270,6 @@ class HoistInfoCollector : public StmtExprVisitor {
return std::nullopt;
}
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) final {
- Var var{ffi::UnsafeInit{}};
- if (const auto* node_iter_var = op->node.as<IterVarNode>()) {
- var = node_iter_var->var;
- } else if (auto opt = op->node.as<Var>()) {
- var = opt.value();
- } else {
- return Parent::Visit_(op);
- }
-
- active_block_vars.insert(var.get());
- active_loop_vars.insert(var.get());
- active_loops.push_back({var, ffi::GetRef<Stmt>(op)});
-
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Parent::Visit_(op));
-
- completed_loops.push_back(active_loops.back());
- active_loops.pop_back();
-
- active_loop_vars.erase(var.get());
- active_block_vars.erase(var.get());
- return std::nullopt;
- }
-
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op) final {
if (op->op.same_as(tirx::builtin::launch_thread())) {
Var var = op->body_params[0].as_or_throw<PrimVar>();
@@ -572,17 +548,6 @@ class ExpressionHoister : public
s_tir::IRMutatorWithAnalyzer {
return WrapHoistedStatements(stmt, it->second);
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
-
- auto it = loop_info_lookup.find(op);
- if (it == loop_info_lookup.end()) {
- return stmt;
- } else {
- return WrapHoistedStatements(stmt, it->second);
- }
- }
-
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
auto it = loop_info_lookup.find(op);
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index a97567cd94..c5b50328a0 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -456,24 +456,6 @@ class VTInjector : public s_tir::IRMutatorWithAnalyzer {
op->span);
}
- // Attribute
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- auto value_result = this->Mutate(op->value, inplace_mode);
- bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
- Expr value = std::move(value_result).ValueOrUnchanged(op->value);
- if (visit_touched_var_ && !vt_loop_injected_) {
- return InjectVTLoop(ffi::GetRef<Stmt>(op), true);
- } else {
- auto body_result = this->Mutate(op->body, inplace_mode);
- bool body_unchanged = body_result.UnchangedOrSameAs(op->body);
- Stmt body = std::move(body_result).ValueOrUnchanged(op->body);
- if (value_unchanged && body_unchanged) {
- return ffi::Unchanged();
- } else {
- return AttrStmt(op->node, op->attr_key, value, body);
- }
- }
- }
// Bind
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
if (const auto* call = op->value.as<CallNode>();
diff --git a/src/s_tir/transform/lower_opaque_block.cc
b/src/s_tir/transform/lower_opaque_block.cc
index 908c3904db..d6f276861f 100644
--- a/src/s_tir/transform/lower_opaque_block.cc
+++ b/src/s_tir/transform/lower_opaque_block.cc
@@ -57,21 +57,6 @@ class OpaqueBlockLower : public StmtExprMutator {
std::optional<ffi::Any> unroll_explicit;
};
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- // Attribute metadata can reference a unit loop variable that is replaced.
- // Rewrite it before descending into body-local definitions.
- auto node = this->Mutate(op->node, inplace_mode);
- auto value = this->Mutate(op->value, inplace_mode);
- auto body = this->Mutate(op->body, inplace_mode);
- if (node.UnchangedOrSameAs(op->node) && value.UnchangedOrSameAs(op->value)
&&
- body.UnchangedOrSameAs(op->body)) {
- return ffi::Unchanged();
- }
- return AttrStmt(std::move(node).ValueOrUnchanged(op->node), op->attr_key,
- std::move(value).ValueOrUnchanged(op->value),
- std::move(body).ValueOrUnchanged(op->body), op->span);
- }
-
UnchangedOr<Stmt> Mutate_(const SBlockRealizeNode* op, InplaceMode
inplace_mode) final {
// We have convert blocks into opaque blocks in previous passes.
TVM_FFI_ICHECK(op->iter_values.empty())
@@ -115,12 +100,6 @@ class OpaqueBlockLower : public StmtExprMutator {
DictAttrs(allocate_annotations))),
std::move(body));
}
- // Step 4. Handle annotations, block annotations are not preserved by
default.
- std::vector<std::pair<std::string, Expr>> pragma_attrs;
- HandleAnnotations(new_block->annotations, &pragma_attrs,
/*is_block=*/true);
- for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) {
- body = AttrStmt(0, it->first, it->second, std::move(body));
- }
return body;
}
@@ -139,7 +118,6 @@ class OpaqueBlockLower : public StmtExprMutator {
// Keep policy on surviving descendants when this owner is lowered away.
auto parent_policy = unroll_policy_.Current();
- std::vector<std::pair<std::string, Expr>> pragma_attrs;
ffi::Map<ffi::String, ffi::Any> new_annotations;
auto annotations = op->annotations;
Stmt body = unroll_policy_.WithNewScope([&]() {
@@ -153,29 +131,18 @@ class OpaqueBlockLower : public StmtExprMutator {
annotations.Set(tirx::attr::unroll_explicit,
policy.unroll_explicit.value());
}
// Rewrite annotations before visiting body-local definitions.
- new_annotations = HandleAnnotations(annotations, &pragma_attrs,
/*is_block=*/false);
+ new_annotations = HandleAnnotations(annotations);
return this->Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
});
VarRemapSet(op->loop_var, previous_remap);
- // Step 4. Keep thread-binding loops until LowerThreadBinding.
+ // Step 2. Keep thread-binding loops until LowerThreadBinding.
if (op->kind != ForKind::kThreadBinding && is_one(extent) &&
op->annotations.empty() &&
!op->annotations.count(s_tir::attr::irregular_loop_mark)) {
return body;
}
- if (op->kind != ForKind::kThreadBinding) {
- body = For(op->loop_var, min, extent, op->kind, std::move(body),
op->thread_binding,
- new_annotations, step);
- }
- // Step 5. Insert nested attrs inside thread-binding scope.
- for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) {
- body = AttrStmt(op->loop_var, it->first, it->second, std::move(body));
- }
- if (op->kind == ForKind::kThreadBinding) {
- body = For(op->loop_var, std::move(min), std::move(extent), op->kind,
std::move(body),
- op->thread_binding, std::move(new_annotations),
std::move(step), op->span);
- }
- return body;
+ return For(op->loop_var, std::move(min), std::move(extent), op->kind,
std::move(body),
+ op->thread_binding, std::move(new_annotations),
std::move(step), op->span);
}
void UpdateUnrollPolicy(const ffi::Map<ffi::String, ffi::Any>& annotations) {
@@ -193,61 +160,22 @@ class OpaqueBlockLower : public StmtExprMutator {
// Effective policy is materialized on each surviving loop, preserving
nested overrides.
ScopeStack<UnrollPolicy> unroll_policy_;
- /*! \brief Convert attr value from annotation map into Expr. */
- Expr ConvertAttrValue(const ffi::String& key, const Any& obj) {
- if (auto expr = obj.try_cast<Expr>()) {
- return expr.value();
- } else if (auto str = obj.try_cast<ffi::String>()) {
- return std::move(StringImm(str.value()));
- } else {
- TVM_FFI_THROW(InternalError) << "Illegal attribute of key " << key << ",
value type "
- << obj.GetTypeKey() << " not supported";
- }
- }
-
- /*!
- * \brief Helper to handle annotation dict.
- * (1) if the attr key is prefixed by `pragma_`, move to ordered kv list.
They
- * are lowered to `AttrStmt` by legacy TE schedule convention.
- * (2) the non-pragma loop annotations are preserved
- * (3) the non-pragma block annotations are dropped
- * \return New annotation dict with preserved keys. Also update pragma attr
pairs ordered by key.
- */
+ // Preserve loop annotations while remapping enclosing bindings. Null
optional
+ // policies do not override inherited values or reach code generation.
ffi::Map<ffi::String, ffi::Any> HandleAnnotations(
- const ffi::Map<ffi::String, ffi::Any>& annotations,
- std::vector<std::pair<std::string, Expr>>* pragma_attrs, bool is_block) {
- ffi::Map<ffi::String, ffi::Any> preserved_annotations;
- pragma_attrs->clear();
- for (const auto& kv : annotations) {
- const ffi::String& key = kv.first;
- if ((key == tirx::attr::auto_unroll_max_step || key ==
tirx::attr::unroll_explicit) &&
- kv.second == nullptr) {
+ const ffi::Map<ffi::String, ffi::Any>& annotations) {
+ ffi::Map<ffi::String, ffi::Any> preserved;
+ for (const auto& [key, value] : annotations) {
+ if ((key == tirx::attr::auto_unroll_max_step || key ==
tirx::attr::unroll_explicit ||
+ key == "pragma_unroll") &&
+ value == nullptr) {
continue;
}
- if (tirx::attr::IsPragmaKey(key)) {
- if (kv.second == nullptr) {
- continue;
- }
-
- auto value = this->Mutate(kv.second,
InplaceMode::kDisallow).ValueOrUnchanged(kv.second);
- pragma_attrs->emplace_back(key, ConvertAttrValue(key, value));
- } else if (!is_block) {
- // Preserve the annotation while remapping enclosing launch bindings,
- // including expressions nested in annotation containers.
- auto value = this->Mutate(kv.second,
InplaceMode::kDisallow).ValueOrUnchanged(kv.second);
- preserved_annotations.Set(key, value);
- }
+ preserved.Set(key, this->Mutate(value,
InplaceMode::kDisallow).ValueOrUnchanged(value));
}
- std::sort(pragma_attrs->begin(), pragma_attrs->end(),
- [](const auto& p1, const auto& p2) { return p1.first < p2.first;
});
- return preserved_annotations;
+ return preserved;
}
- /*! \brief Record the loop_var and loop start value of unit loops, whose
extent is one. */
-
- /*! \brief Attr keys to preserve into loop annotations. */
- std::unordered_set<std::string> preserved_annotations_;
-
/*! \brief The map from buffer var to its storage alignment information. */
std::unordered_map<Var, StorageAlignAnnotation> storage_align_;
};
diff --git a/src/s_tir/transform/lower_thread_binding.cc
b/src/s_tir/transform/lower_thread_binding.cc
index a57fde7973..d6d3dd87a9 100644
--- a/src/s_tir/transform/lower_thread_binding.cc
+++ b/src/s_tir/transform/lower_thread_binding.cc
@@ -34,20 +34,6 @@ class ThreadBindingLowerer : public StmtExprMutator {
using StmtExprMutator::Mutate_;
private:
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- // Attribute metadata can refer to an enclosing loop's lexical binding.
- auto node = Mutate(op->node, inplace_mode);
- auto value = Mutate(op->value, inplace_mode);
- auto body = Mutate(op->body, inplace_mode);
- if (node.UnchangedOrSameAs(op->node) && value.UnchangedOrSameAs(op->value)
&&
- body.UnchangedOrSameAs(op->body)) {
- return ffi::Unchanged();
- }
- return AttrStmt(std::move(node).ValueOrUnchanged(op->node), op->attr_key,
- std::move(value).ValueOrUnchanged(op->value),
- std::move(body).ValueOrUnchanged(op->body), op->span);
- }
-
UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode)
final {
auto annotations = Mutate(op->annotations, inplace_mode)
.as_or_throw<UnchangedOr<ffi::Map<ffi::String,
ffi::Any>>>();
diff --git a/src/script/printer/script_printer.cc
b/src/script/printer/script_printer.cc
index 04bbdcb850..3360ea9fed 100644
--- a/src/script/printer/script_printer.cc
+++ b/src/script/printer/script_printer.cc
@@ -123,7 +123,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
RegisterScriptRepr<te::CommReducerNode>();
RegisterScriptRepr<te::ReduceNode>();
RegisterScriptRepr<tirx::AssertStmtNode>();
- RegisterScriptRepr<tirx::AttrStmtNode>();
RegisterScriptRepr<tirx::RegionStmtNode>();
RegisterScriptRepr<tirx::BindNode>();
RegisterScriptRepr<tirx::BreakNode>();
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 62d0113ac9..35f1260569 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -1141,27 +1141,6 @@ void CodeGenCPU::Dispatch_(const RegionStmtNode* op) {
}
}
-void CodeGenCPU::Dispatch_(const AttrStmtNode* op) {
- EmitDebugLocation(op);
- if (tirx::attr::IsPragmaKey(op->attr_key)) {
- if (op->attr_key == "pragma_parallel_stride_pattern") {
- TVM_FFI_THROW(ValueError)
- << "pragma_parallel_stride_pattern is retired; annotate every
affected parallel For "
- << "with parallel_stride_pattern=True, including later loops in the
same launch "
- << "outside the former attribute body";
- } else if (op->attr_key == "pragma_parallel_barrier_when_finish") {
- TVM_FFI_THROW(ValueError)
- << "pragma_parallel_barrier_when_finish is retired; place
cpu_parallel_barrier() "
- << "after the former attribute body inside the parallel launch";
- } else {
- LOG(WARNING) << "Unknown pragma " << op->attr_key;
- this->Dispatch(op->body);
- }
- } else {
- CodeGenLLVM::Dispatch_(op);
- }
-}
-
void CodeGenCPU::Dispatch_(const ForNode* op) {
EmitDebugLocation(op);
if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) {
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index dba07332d8..e9549a82bd 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -73,7 +73,6 @@ class CodeGenCPU : public CodeGenLLVM {
std::unique_ptr<llvm::Module> Finish() override;
void Dispatch_(const AssertStmtNode* op) override;
void Dispatch_(const RegionStmtNode* op) override;
- void Dispatch_(const AttrStmtNode* op) override;
void Dispatch_(const ForNode* op) override;
llvm::Value* CreateIntrinsic(const CallNode* op) override;
llvm::Value* CreateCallExtern(Type ret_type, ffi::String global_symbol,
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 9d8f6da7b2..dfd8d97dd7 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -2265,13 +2265,6 @@ void CodeGenLLVM::Dispatch_(const RegionStmtNode* op) {
}
}
-void CodeGenLLVM::Dispatch_(const AttrStmtNode* op) {
- TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
- << "Launch attributes are retired; use tirx.launch_thread RegionStmt";
- EmitDebugLocation(op);
- this->Dispatch(op->body);
-}
-
void CodeGenLLVM::Dispatch_(const AssertStmtNode* op) {
EmitDebugLocation(op);
// AssertStmt is a leaf — no body to visit.
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index c077d21563..199953add3 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -246,7 +246,6 @@ class CodeGenLLVM : public
tirx::ExprFunctor<llvm::Value*(const Expr&)>,
void Dispatch_(const IfThenElseNode* op) override;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const RegionStmtNode* op) override;
- void Dispatch_(const AttrStmtNode* op) override;
void Dispatch_(const AssertStmtNode* op) override;
void Dispatch_(const BindNode* op) override;
void Dispatch_(const SeqStmtNode* op) override;
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 68d40b5107..be5bdd5ae0 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -1319,12 +1319,6 @@ void CodeGenC::Dispatch_(const RegionStmtNode* op) {
}
}
-void CodeGenC::Dispatch_(const AttrStmtNode* op) {
- TVM_FFI_CHECK(op->attr_key != "thread_extent" && op->attr_key !=
"virtual_thread", ValueError)
- << "Launch attributes are retired; use tirx.launch_thread RegionStmt";
- this->PrintStmt(op->body);
-}
-
void CodeGenC::PrintEscapedCString(const std::string& str, std::ostream& os) {
os << "\"";
for (unsigned char c : str) {
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index bf96b0a2c5..96af65a167 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -219,7 +219,6 @@ class CodeGenC : public tirx::ExprFunctor<void(const Expr&,
std::ostream&)>,
void Dispatch_(const IfThenElseNode* op) override;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const RegionStmtNode* op) override;
- void Dispatch_(const AttrStmtNode* op) override;
void Dispatch_(const AssertStmtNode* op) override;
void Dispatch_(const EvaluateNode* op) override;
void Dispatch_(const SeqStmtNode* op) override;
diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc
b/src/tirx/analysis/verify_tirx_well_formed.cc
index 1e3acb4bde..8adb59131b 100644
--- a/src/tirx/analysis/verify_tirx_well_formed.cc
+++ b/src/tirx/analysis/verify_tirx_well_formed.cc
@@ -144,9 +144,6 @@ class DeviceFuncVerifier : public
Verifier<DeviceFuncVerifier> {
};
bool VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode, bool
device_func) {
- if (!RetiredLaunchAttrVerifier<TIRVisitorWithPath>::Verify(func,
assert_mode)) {
- return false;
- }
if (!ExecScopeVerifier::Verify(func, assert_mode)) {
return false;
}
diff --git a/src/tirx/analysis/verify_well_formed.h
b/src/tirx/analysis/verify_well_formed.h
index df779c8214..27e92b2615 100644
--- a/src/tirx/analysis/verify_well_formed.h
+++ b/src/tirx/analysis/verify_well_formed.h
@@ -28,25 +28,6 @@ namespace tvm {
namespace tirx {
using AccessPath = ffi::reflection::AccessPath;
-/*! \brief Reject retired launch attributes, including loaded IR. */
-template <typename PathVisitor>
-class RetiredLaunchAttrVerifier
- : public Verifier<RetiredLaunchAttrVerifier<PathVisitor>, PathVisitor> {
- using Verifier = tirx::Verifier<RetiredLaunchAttrVerifier<PathVisitor>,
PathVisitor>;
-
- public:
- using Verifier::Verifier;
- using Verifier::Verify;
-
- private:
- void Dispatch_(const AttrStmtNode* op, AccessPath path) override {
- bool retired = op->attr_key == "thread_extent" || op->attr_key ==
"virtual_thread";
- Verify(!retired) << "ValueError: Retired thread launch attribute at " <<
path
- << "; use RegionStmt launch_thread";
- if (!retired) PathVisitor::Dispatch_(op, path);
- }
-};
-
template <typename PathVisitor>
class UndefinedVarVerifier : public
Verifier<UndefinedVarVerifier<PathVisitor>, PathVisitor> {
using Verifier = tirx::Verifier<UndefinedVarVerifier<PathVisitor>,
PathVisitor>;
@@ -322,8 +303,7 @@ class LoopControlVerifier : public
Verifier<LoopControlVerifier<PathVisitor>, Pa
template <typename PathVisitor, typename NodeRef>
bool VerifyWellFormedCommon(const NodeRef& node, bool assert_mode) {
- return RetiredLaunchAttrVerifier<PathVisitor>::Verify(node, assert_mode) &&
- UndefinedVarVerifier<PathVisitor>::Verify(node, assert_mode) &&
+ return UndefinedVarVerifier<PathVisitor>::Verify(node, assert_mode) &&
UndefinedBufferVerifier<PathVisitor>::Verify(node, assert_mode) &&
TensorLoadTypeVerifier<PathVisitor>::Verify(node, assert_mode) &&
LoopControlVerifier<PathVisitor>::Verify(node, assert_mode);
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.cc
b/src/tirx/ir/ir_mutator_with_analyzer.cc
index 2f430f535d..4b72595516 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -211,11 +211,6 @@ UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const
IfThenElseNode* op,
});
}
-UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const AttrStmtNode* op,
InplaceMode inplace_mode) {
- return constraint_scope_.WithNewScope(
- [&]() { return StmtExprMutator::Mutate_(op, inplace_mode); });
-}
-
UnchangedOr<Stmt> IRMutatorWithAnalyzer::Mutate_(const RegionStmtNode* op,
InplaceMode inplace_mode) {
return constraint_scope_.WithNewScope([&]() -> UnchangedOr<Stmt> {
diff --git a/src/tirx/ir/ir_mutator_with_analyzer.h
b/src/tirx/ir/ir_mutator_with_analyzer.h
index 8cc4e29de1..dd4dd7e1c3 100644
--- a/src/tirx/ir/ir_mutator_with_analyzer.h
+++ b/src/tirx/ir/ir_mutator_with_analyzer.h
@@ -60,7 +60,6 @@ class IRMutatorWithAnalyzer : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode)
override;
UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
override;
UnchangedOr<Stmt> Mutate_(const IfThenElseNode* op, InplaceMode
inplace_mode) override;
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
override;
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) override;
UnchangedOr<Stmt> Mutate_(const AssertStmtNode* op, InplaceMode
inplace_mode) override;
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
override;
diff --git a/src/tirx/ir/ir_visitor_with_analyzer.cc
b/src/tirx/ir/ir_visitor_with_analyzer.cc
index 0094fcfecd..63d8eec00a 100644
--- a/src/tirx/ir/ir_visitor_with_analyzer.cc
+++ b/src/tirx/ir/ir_visitor_with_analyzer.cc
@@ -79,10 +79,6 @@ ffi::Optional<VisitInterrupt>
IRVisitorWithAnalyzer::Visit_(const IfThenElseNode
});
}
-ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const
AttrStmtNode* op) {
- return constraint_scope_.WithNewScope([&]() { return
StmtExprVisitor::Visit_(op); });
-}
-
ffi::Optional<VisitInterrupt> IRVisitorWithAnalyzer::Visit_(const
RegionStmtNode* op) {
return constraint_scope_.WithNewScope([&]() -> ffi::Optional<VisitInterrupt>
{
if (op->op.same_as(tirx::builtin::launch_thread())) {
diff --git a/src/tirx/ir/ir_visitor_with_analyzer.h
b/src/tirx/ir/ir_visitor_with_analyzer.h
index affdde2dc6..1a6a8f3a0b 100644
--- a/src/tirx/ir/ir_visitor_with_analyzer.h
+++ b/src/tirx/ir/ir_visitor_with_analyzer.h
@@ -45,7 +45,6 @@ class IRVisitorWithAnalyzer : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> Visit_(const ForNode* op);
ffi::Optional<VisitInterrupt> Visit_(const BindNode* op);
ffi::Optional<VisitInterrupt> Visit_(const IfThenElseNode* op);
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op);
ffi::Optional<VisitInterrupt> Visit_(const RegionStmtNode* op);
ffi::Optional<VisitInterrupt> Visit_(const AssertStmtNode* op);
ffi::Optional<VisitInterrupt> Visit_(const CallNode* op);
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index 874861b318..7694f664f5 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -149,60 +149,6 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
BindMaybeInplaceMutate(
return ffi::Unchanged();
}
-TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> AttrStmtVisit(
- ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
- // skips: attr_key
- const AttrStmtNode* self =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
AttrStmtNode>(value);
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->node));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->value));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->body));
- return std::nullopt;
-}
-
-TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> AttrStmtMutate(
- ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: attr_key
- const AttrStmtNode* self =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
AttrStmtNode>(value);
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Any>, mapped_node,
- mutator->MutateExpected(self->node));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
- mutator->MutateExpected(self->value));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
- mutator->MutateExpected(self->body));
- if (mapped_node.UnchangedOrSameAs(self->node) &&
mapped_value.UnchangedOrSameAs(self->value) &&
- mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged();
- }
- ffi::ObjectPtr<AttrStmtNode> copy = ffi::make_object<AttrStmtNode>(*self);
- copy->node = std::move(mapped_node).ValueOrUnchanged(std::move(copy->node));
- copy->value =
std::move(mapped_value).ValueOrUnchanged(std::move(copy->value));
- copy->body = std::move(mapped_body).ValueOrUnchanged(std::move(copy->body));
- return ffi::Any(std::move(copy));
-}
-
-TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
AttrStmtMaybeInplaceMutate(
- ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: attr_key
- AttrStmtNode* self = const_cast<AttrStmtNode*>(
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
AttrStmtNode>(value));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Any>, mapped_node,
- mutator->MutateExpected(self->node,
ffi::InplaceMode::kAllow));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Expr>, mapped_value,
- mutator->MutateExpected(self->value,
ffi::InplaceMode::kAllow));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
- mutator->MutateExpected(self->body,
ffi::InplaceMode::kAllow));
- if (mapped_node.UnchangedOrSameAs(self->node) &&
mapped_value.UnchangedOrSameAs(self->value) &&
- mapped_body.UnchangedOrSameAs(self->body)) {
- return ffi::Unchanged();
- }
- if (!mapped_node.IsUnchanged()) self->node =
std::move(mapped_node).ValueUnchecked();
- if (!mapped_value.IsUnchanged()) self->value =
std::move(mapped_value).ValueUnchecked();
- if (!mapped_body.IsUnchanged()) self->body =
std::move(mapped_body).ValueUnchecked();
- return ffi::Unchanged();
-}
-
TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>>
RegionStmtVisit(
ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
const RegionStmtNode* self =
@@ -775,34 +721,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](Var var, Expr value, Span span) { return Bind(var,
value, span); });
}
-// AttrStmt
-AttrStmt::AttrStmt(ffi::Any node, ffi::String attr_key, Expr value, Stmt body,
Span span)
- : Stmt(ffi::UnsafeInit{}) {
- TVM_FFI_CHECK(attr_key != "thread_extent" && attr_key != "virtual_thread",
ValueError)
- << "Thread launch attributes are retired; use launch_thread(tag,
extent)";
- auto n = ffi::make_object<AttrStmtNode>(std::move(value), std::move(body));
- n->node = node;
- n->attr_key = std::move(attr_key);
- n->span = std::move(span);
- data_ = std::move(n);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- AttrStmtNode::RegisterReflection();
- refl::TypeAttrDef<AttrStmtNode>()
- .attr(refl::type_attr::kStructuralVisit,
ffi::FStructuralVisit::FromNative<&AttrStmtVisit>())
- .attr(refl::type_attr::kStructuralMutate,
- ffi::FStructuralMutate::FromNative<&AttrStmtMutate>())
- .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
- ffi::FStructuralMutate::FromNative<&AttrStmtMaybeInplaceMutate>());
-
- refl::GlobalDef().def("tirx.AttrStmt",
- [](Any node, ffi::String attr_key, Expr value, Stmt
body, Span span) {
- return AttrStmt(node, attr_key, value, body, span);
- });
-}
-
// RegionStmt
RegionStmt::RegionStmt(Op op, ffi::Array<Expr> args, ffi::Array<Var>
body_params, DictAttrs attrs,
Stmt body, ffi::Array<Var> result_vars, Span span)
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index d061831506..ae1b53fac8 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -43,7 +43,6 @@ namespace tirx {
void StmtExprVisitor::InitVTable(VTable* vtable) {
tvm::ExprVisitor::InitVTable(vtable);
SetDispatch<StmtExprVisitor, BindNode>(vtable);
- SetDispatch<StmtExprVisitor, AttrStmtNode>(vtable);
SetDispatch<StmtExprVisitor, RegionStmtNode>(vtable);
SetDispatch<StmtExprVisitor, IfThenElseNode>(vtable);
SetDispatch<StmtExprVisitor, ForNode>(vtable);
@@ -134,11 +133,6 @@ ffi::Optional<VisitInterrupt>
StmtExprVisitor::Visit_(const BindNode* op) {
[&]() { return this->Visit(op->var); });
}
-ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const AttrStmtNode* op) {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->value));
- return this->Visit(op->body);
-}
-
ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const RegionStmtNode*
op) {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->attrs));
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->Visit(op->args));
@@ -247,7 +241,6 @@ ffi::Optional<VisitInterrupt> StmtExprVisitor::Visit_(const
TilePrimitiveCallNod
void StmtExprMutator::InitVTable(VTable* vtable) {
tvm::ExprMutator::InitVTable(vtable);
SetDispatch<StmtExprMutator, BindNode>(vtable);
- SetDispatch<StmtExprMutator, AttrStmtNode>(vtable);
SetDispatch<StmtExprMutator, RegionStmtNode>(vtable);
SetDispatch<StmtExprMutator, IfThenElseNode>(vtable);
SetDispatch<StmtExprMutator, ForNode>(vtable);
@@ -293,23 +286,6 @@ UnchangedOr<Stmt> StmtExprMutator::Mutate_(const BindNode*
op, InplaceMode inpla
return Stmt(std::move(copy));
}
-UnchangedOr<Stmt> StmtExprMutator::Mutate_(const AttrStmtNode* op, InplaceMode
inplace_mode) {
- auto value = Mutate(op->value, inplace_mode);
- auto body = Mutate(op->body, inplace_mode);
- if (value.UnchangedOrSameAs(op->value) && body.UnchangedOrSameAs(op->body))
- return ffi::Unchanged();
- if (inplace_mode == InplaceMode::kAllow) {
- auto* writable = const_cast<AttrStmtNode*>(op);
- if (!value.IsUnchanged()) writable->value =
std::move(value).ValueUnchecked();
- if (!body.IsUnchanged()) writable->body = std::move(body).ValueUnchecked();
- return ffi::Unchanged();
- }
- auto copy = ffi::make_object<AttrStmtNode>(*op);
- if (!value.IsUnchanged()) copy->value = std::move(value).ValueUnchecked();
- if (!body.IsUnchanged()) copy->body = std::move(body).ValueUnchecked();
- return Stmt(std::move(copy));
-}
-
UnchangedOr<Stmt> StmtExprMutator::Mutate_(const RegionStmtNode* op,
InplaceMode inplace_mode) {
auto attrs = Mutate(op->attrs,
inplace_mode).as_or_throw<UnchangedOr<DictAttrs>>();
auto args = Mutate(op->args,
inplace_mode).as_or_throw<UnchangedOr<ffi::Array<Expr>>>();
@@ -608,23 +584,6 @@ class IRSubstituteWithDataTypeLegalization : public
DataTypeLegalizer {
}
return load;
}
-
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- Stmt ret = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = ret.as<AttrStmtNode>();
- // remap var node in attr
- if (auto var_node = op->node.as<Var>()) {
- ffi::Any mapped = VarRemapGet(var_node.value());
- if (mapped.type_index() != ffi::TypeIndex::kTVMFFINone) {
- Expr node =
-
std::move(mapped).as_or_throw<UnchangedOr<Expr>>().ValueOrUnchanged(var_node.value());
- if (!node.same_as(var_node.value())) {
- return AttrStmt(node, op->attr_key, op->value, op->body);
- }
- }
- }
- return ret;
- }
};
Stmt SubstituteWithDataTypeLegalization(Stmt stmt,
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index c0595a1168..54cd11dbe8 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -179,14 +179,6 @@ void TIRVisitorWithPath::Dispatch_(const BindNode* op,
AccessPath path) {
bind_scope_.Current().push_back(WithDef(op->var, path->Attr("var")));
}
-void TIRVisitorWithPath::Dispatch_(const AttrStmtNode* op, AccessPath path) {
- Visit(op->value, path->Attr("value"));
- if (auto expr = op->node.as<PrimExpr>()) {
- Visit(expr.value(), path->Attr("node"));
- }
- bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); });
-}
-
void TIRVisitorWithPath::Dispatch_(const RegionStmtNode* op, AccessPath path) {
Visit(op->args, path->Attr("args"));
Visit(ffi::AnyView(op->attrs), path->Attr("attrs"));
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index 72050a9a11..80208d7f7d 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -148,7 +148,6 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
using StmtFunctor::Dispatch;
void Dispatch_(const BindNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const AttrStmtNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const RegionStmtNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const IfThenElseNode* op, ffi::reflection::AccessPath path)
override;
void Dispatch_(const ForNode* op, ffi::reflection::AccessPath path) override;
diff --git a/src/tirx/script/ir_builder/frame.cc
b/src/tirx/script/ir_builder/frame.cc
index 7dba06f37c..5939eaff2d 100644
--- a/src/tirx/script/ir_builder/frame.cc
+++ b/src/tirx/script/ir_builder/frame.cc
@@ -43,7 +43,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ForFrameNode::RegisterReflection();
AssertFrameNode::RegisterReflection();
RegionFrameNode::RegisterReflection();
- AttrFrameNode::RegisterReflection();
WhileFrameNode::RegisterReflection();
IfFrameNode::RegisterReflection();
ThenFrameNode::RegisterReflection();
@@ -220,11 +219,6 @@ void RegionFrameNode::ExitWithScope() {
source_span);
}
-void AttrFrameNode::ExitWithScope() {
- TIRFrameNode::ExitWithScope();
- AddToParent(tvm::tirx::AttrStmt(node, attr_key, value, AsStmt(stmts),
source_span), source_span);
-}
-
void WhileFrameNode::ExitWithScope() {
TIRFrameNode::ExitWithScope();
AddToParent(tvm::tirx::While(condition, AsStmt(stmts), source_span),
source_span);
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 636e079ff3..7d4d468df9 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -188,7 +188,7 @@ ffi::Array<tvm::tirx::Var>
ScopeId(ffi::Optional<ffi::Array<PrimExpr>> extents,
}
// Emit a standalone ScopeIdDefStmt to the current TIRFrame's stmts list.
// The def is visible to all subsequent stmts within the same enclosing
- // scope (PrimFunc body, AttrStmt body, ExecScope body, etc.).
+ // scope (PrimFunc body, RegionStmt body, ExecScope body, etc.).
tvm::tirx::ScopeIdDef def(
scope_ids.Map([](tvm::tirx::Var var) { return
var.as_or_throw<tvm::PrimVar>(); }), extents,
tvm::tirx::StringPairToScopeBinding(parent, cur));
@@ -432,13 +432,6 @@ RegionFrame LaunchThread(ffi::String thread_tag, PrimExpr
extent) {
return Region(tvm::tirx::builtin::launch_thread(), {StringImm(thread_tag),
extent}, {var});
}
-AttrFrame Attr(ffi::Any node, ffi::String attr_key, Expr value) {
- ffi::ObjectPtr<AttrFrameNode> n = ffi::make_object<AttrFrameNode>(value);
- n->node = std::move(node);
- n->attr_key = attr_key;
- return AttrFrame(n);
-}
-
WhileFrame While(PrimExpr condition) {
ffi::ObjectPtr<WhileFrameNode> n =
ffi::make_object<WhileFrameNode>(condition);
return WhileFrame(n);
@@ -705,7 +698,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.def("script.ir_builder.tirx.Grid", Grid)
.def("script.ir_builder.tirx.Assert", Assert)
.def("script.ir_builder.tirx.Bind", Bind)
- .def("script.ir_builder.tirx.Attr", Attr)
.def("script.ir_builder.tirx.While", While)
.def("script.ir_builder.tirx.Return", Return)
.def("script.ir_builder.tirx.Break", Break)
diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc
index 0ffc17c34d..352fc1d69e 100644
--- a/src/tirx/script/printer/stmt.cc
+++ b/src/tirx/script/printer/stmt.cc
@@ -426,65 +426,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
kDocTranslate, FDocTranslate::FromNative<&RegionStmtDocTranslate>());
}
-ffi::Optional<ExprDoc> AttrStmtDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
- const ffi::Object* destination) {
- const auto* stmt =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::AttrStmtNode>(input);
- TVM_FFI_CHECK(destination == nullptr, TypeError)
- << "printer statement-only node cannot fulfill a destination";
- ffi::Optional<ExprDoc> lhs = std::nullopt;
- ExprDoc rhs(ffi::UnsafeInit{});
- tirx::Stmt body = stmt->body;
- TVM_FFI_CHECK(stmt->attr_key != "thread_extent" && stmt->attr_key !=
"virtual_thread", ValueError)
- << "Thread launch attributes are no longer supported; use RegionStmt
launch_thread";
- if (stmt->attr_key == "tirx_hint") {
- if (auto attrs = stmt->node.as<ffi::Map<ffi::String, ffi::Any>>()) {
- ffi::Array<ExprDoc> args;
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc> values;
- for (const auto& [key, value] : attrs.value()) {
- if (key == "message")
- args.push_back(AnyValue(d, value));
- else {
- keys.push_back(key);
- values.push_back(AnyValue(d, value));
- }
- }
- rhs = NamespaceDoc("tirx")->Attr("hint")->Call(args, keys, values);
- }
- }
- if (!rhs.defined()) {
- if (auto zero = stmt->node.as<int64_t>(); zero.has_value() && zero.value()
== 0) {
- ffi::Array<ExprDoc> keys;
- ffi::Array<ExprDoc> values;
- auto current = ffi::GetRef<tirx::AttrStmt>(stmt);
- while (true) {
- keys.push_back(LiteralDoc::Str(current->attr_key, std::nullopt));
- values.push_back(d->Translate(current->value).value());
- auto next = current->body.as<tirx::AttrStmt>();
- if (!next.has_value()) break;
- auto next_zero = next.value()->node.as<int64_t>();
- if (!next_zero.has_value() || next_zero.value() != 0) break;
- current = next.value();
- }
- body = current->body;
- rhs = NamespaceDoc("tirx")->Attr("attr")->Call({DictDoc(keys, values)});
- } else {
- ExprDoc node = AnyValue(d, stmt->node);
- rhs = NamespaceDoc("tirx")->Attr("attr")->Call(
- {node, LiteralDoc::Str(stmt->attr_key, std::nullopt),
d->Translate(stmt->value).value()});
- }
- }
- d->Emit(ScopeDoc(lhs, rhs, Body(body, d), /*allow_concise_scoping=*/true),
- ffi::GetRef<ffi::ObjectRef>(stmt));
- return std::nullopt;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<tirx::AttrStmtNode>().attr(
- kDocTranslate, FDocTranslate::FromNative<&AttrStmtDocTranslate>());
-}
-
} // namespace
} // namespace details
diff --git a/src/tirx/transform/common_subexpr_elim.cc
b/src/tirx/transform/common_subexpr_elim.cc
index d6cb658dab..364492c095 100644
--- a/src/tirx/transform/common_subexpr_elim.cc
+++ b/src/tirx/transform/common_subexpr_elim.cc
@@ -27,7 +27,7 @@
*
* Phase 1 — **CSEPlanner** (analysis, no mutation)
* Walks the TIR tree bottom-up and builds:
- * - A *scope tree* that mirrors the nesting structure of
For/If/While/AttrStmt.
+ * - A *scope tree* that mirrors the nesting structure of
For/If/While/RegionStmt.
* - An *expression DAG* mapping each structurally-unique eligible
expression
* to its occurrence count, LCA scope, first-use location, and direct
* children (which shallower expressions it contains).
@@ -56,7 +56,7 @@
*
* Scope tree
* ----------
- * Each For, IfThenElse (each branch), While, and AttrStmt body creates a new
+ * Each For, IfThenElse (each branch), While, and RegionStmt body creates a new
* scope. The scope tree enables computing the Lowest Common Ancestor (LCA) of
* all scopes where an expression occurs, which determines the correct
insertion
* point — the narrowest scope that dominates all uses.
@@ -122,7 +122,7 @@ using InsertBeforeTable =
* \brief Phase 1 of the two-phase CSE pass.
*
* CSEPlanner is a read-only visitor that scans the TIR tree bottom-up and
builds:
- * 1. A **scope tree** (vector of ScopeEntry) reflecting
For/If/While/AttrStmt nesting.
+ * 1. A **scope tree** (vector of ScopeEntry) reflecting
For/If/While/RegionStmt nesting.
* 2. An **expression DAG** (ExprTable) where each node is an eligible
expression
* with occurrence count, expr_depth, LCA scope, first-use location, and
* direct children (other table entries reachable without passing through
@@ -166,7 +166,7 @@ class CSEPlanner : public StmtExprVisitor {
* \brief One node in the scope tree.
*
* The scope tree mirrors the nesting structure of the TIR program.
- * Each scope-creating statement (For, IfThenElse branch, While, AttrStmt)
+ * Each scope-creating statement (For, IfThenElse branch, While, RegionStmt)
* gets its own ScopeEntry. The root scope (depth 0) represents the function
* body itself.
*/
@@ -625,16 +625,6 @@ class CSEPlanner : public StmtExprVisitor {
return std::nullopt;
}
- /*! \brief AttrStmt: value in parent scope, body in child scope. */
- ffi::Optional<VisitInterrupt> Visit_(const AttrStmtNode* op) override {
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->value));
- int saved = current_scope_;
- current_scope_ = AllocScope(saved, ffi::GetRef<Stmt>(op));
- TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(Visit(op->body));
- current_scope_ = saved;
- return std::nullopt;
- }
-
// ------------------------------------------------------------------
// ComputePlan: convert scan results into the output plan
// ------------------------------------------------------------------
diff --git a/src/tirx/transform/ir_utils.cc b/src/tirx/transform/ir_utils.cc
index caec09540d..ac6167a932 100644
--- a/src/tirx/transform/ir_utils.cc
+++ b/src/tirx/transform/ir_utils.cc
@@ -53,11 +53,6 @@ Stmt MergeNest(const std::vector<Stmt>& nest, Stmt body) {
} else if (const auto* bind = s.as<BindNode>()) {
// Bind has no body -- prepend it before the accumulated body in a
SeqStmt.
body = SeqStmt::Flatten(ffi::GetRef<Stmt>(bind), body);
- } else if (const auto* attr = s.as<AttrStmtNode>()) {
- auto n = ffi::make_object<AttrStmtNode>(*attr);
- TVM_FFI_ICHECK(is_no_op(n->body));
- n->body = body;
- body = Stmt(n);
} else if (const auto* ite = s.as<IfThenElseNode>()) {
auto n = ffi::make_object<IfThenElseNode>(*ite);
TVM_FFI_ICHECK(is_no_op(n->then_case));
@@ -149,7 +144,6 @@ PrimFunc IRConvertSSA::VisitPrimFunc(PrimFunc func) {
// Pop function-scope remaps in reverse order
PopAllRemapsInCurrentScope();
- function_scope_var_remap_.clear();
return func;
}
@@ -197,9 +191,6 @@ Var IRConvertSSA::GetRemappedVar(Var var) {
if (auto it = scoped_var_remap_.find(var.get());
it != scoped_var_remap_.end() && it->second.size()) {
return it->second.back();
- } else if (auto it = function_scope_var_remap_.find(var.get());
- it != function_scope_var_remap_.end()) {
- return it->second;
} else {
return var;
}
@@ -281,83 +272,6 @@ UnchangedOr<Stmt> IRConvertSSA::Mutate_(const
RegionStmtNode* op, InplaceMode in
std::move(results), op->span);
}
-UnchangedOr<Stmt> IRConvertSSA::Mutate_(const AttrStmtNode* op, InplaceMode
inplace_mode) {
- if (const IterVarNode* iter_var = op->node.as<IterVarNode>()) {
- Range dom = iter_var->dom;
- if (dom.defined()) {
- // Retain the original domain while comparing and rebuilding its
replacement.
- auto min = Mutate(dom->min,
InplaceMode::kDisallow).ValueOrUnchanged(dom->min);
- auto extent = Mutate(dom->extent,
InplaceMode::kDisallow).ValueOrUnchanged(dom->extent);
- if (!min.same_as(iter_var->dom->min) ||
!extent.same_as(iter_var->dom->extent)) {
- dom = Range::FromMinExtent(min, extent);
- }
- }
-
- Var var = iter_var->var;
- bool delayed_define = false;
- if (auto it = function_scope_var_remap_.find(var.get());
- it != function_scope_var_remap_.end()) {
- var = it->second;
- } else if (defined_.count(var.get())) {
- Var new_var(var->name, var->ty);
-
- function_scope_var_remap_.insert({var.get(), new_var});
- var = new_var;
- } else {
- // An attribute can refer to a variable defined by a later ForNode.
- // Preserve the same variable definition in the attribute and the loop.
- //
- // Preserve the annotated variable's identity for later definitions
- // and independent functions, without introducing a lexical body binding.
- delayed_define = true;
- }
-
- IterVar new_iter_var;
- if (dom.same_as(iter_var->dom) && var.same_as(iter_var->var)) {
- new_iter_var = ffi::GetRef<IterVar>(iter_var);
- } else {
- new_iter_var = IterVar(dom, var.as_or_throw<PrimVar>(),
iter_var->iter_type,
- iter_var->thread_tag, iter_var->span);
- }
- auto value_result = Mutate(op->value, inplace_mode);
- bool value_unchanged = value_result.UnchangedOrSameAs(op->value);
- auto value = std::move(value_result).ValueOrUnchanged(op->value);
- auto body = scope_.WithNewScope(
- [&]() -> Stmt { return Mutate(op->body,
inplace_mode).ValueOrUnchanged(op->body); });
-
- Stmt output = ffi::GetRef<Stmt>(op);
- if (new_iter_var.get() == iter_var && body.same_as(op->body) &&
value_unchanged) {
- output = ffi::GetRef<Stmt>(op);
- } else {
- output = AttrStmt(new_iter_var, op->attr_key, value, body,
iter_var->span);
- }
-
- if (delayed_define) {
- if (!defined_.count(var.get())) {
- function_scope_var_remap_.insert({var.get(), var});
- defined_.insert(var.get());
- }
- }
-
- return output;
-
- } else if (const VarNode* v = op->node.as<VarNode>()) {
- Stmt stmt = scope_.WithNewScope([&]() -> Stmt {
- return StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- });
- op = stmt.as<AttrStmtNode>();
- if (scoped_var_remap_.count(v) && scoped_var_remap_[v].size() != 0) {
- return AttrStmt(scoped_var_remap_[v].back(), op->attr_key, op->value,
op->body);
- } else {
- return stmt;
- }
- } else {
- return scope_.WithNewScope([&]() -> Stmt {
- return StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- });
- }
-}
-
Var IRConvertSSA::MakeNewVar(const Var& old_var) {
return Var(old_var->name, old_var->ty, old_var->span);
}
diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h
index db29ba0fb5..a3fc4fcb12 100644
--- a/src/tirx/transform/ir_utils.h
+++ b/src/tirx/transform/ir_utils.h
@@ -246,7 +246,6 @@ class IRConvertSSA : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_(const IfThenElseNode* op, InplaceMode
inplace_mode) final;
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final;
UnchangedOr<Stmt> Mutate_(const WhileNode* op, InplaceMode inplace_mode)
final;
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final;
UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final;
UnchangedOr<PrimExpr> Mutate_(const prim::LetNode* op, InplaceMode
inplace_mode) final;
static Var MakeNewVar(const Var& old_var);
@@ -307,7 +306,6 @@ class IRConvertSSA : public StmtExprMutator {
std::unordered_map<const VarNode*, std::vector<Var>> scoped_var_remap_;
std::unordered_set<const VarNode*> defined_;
- std::unordered_map<const VarNode*, Var> function_scope_var_remap_;
ScopeStack<ScopeLevel> scope_;
};
diff --git a/src/tirx/transform/lower_thread_allreduce.h
b/src/tirx/transform/lower_thread_allreduce.h
index 4e03b7d92a..9b8f5b7740 100644
--- a/src/tirx/transform/lower_thread_allreduce.h
+++ b/src/tirx/transform/lower_thread_allreduce.h
@@ -104,7 +104,7 @@ class ThreadAllreduceBuilder final : public DialectMutator {
UnchangedOr<Stmt> MutateAllocTensor(const BindNode* op, InplaceMode
inplace_mode) {
buffer_aliases_.Set(op->var, op->var);
// In flat IR, alloc_remap_ may not yet be populated when this AllocTensor
is visited
- // (the remap is set up by MakeAllreduce which runs during
AttrStmt/Evaluate visit
+ // (the remap is set up by MakeAllreduce which runs during Evaluate visit
// that appears later in the sequence). We record the original data
pointer and
// attempt the remap; if it's not ready, the post-processing pass will
handle it.
const VarNode* orig_data_ptr = op->var.get();
diff --git a/src/tirx/transform/lower_tirx_opaque.cc
b/src/tirx/transform/lower_tirx_opaque.cc
index f30beeaaa4..21aa83bb6d 100644
--- a/src/tirx/transform/lower_tirx_opaque.cc
+++ b/src/tirx/transform/lower_tirx_opaque.cc
@@ -22,10 +22,9 @@
* \brief Lower opaque constructs in TIRX programs. This is the tirx-specific
* counterpart of s_tirx::LowerOpaqueBlock, handling only the non-SBlock
* parts: AllocTensor lowering, For(thread_binding) →
RegionStmt(launch_thread),
- * unit loop elimination, and pragma annotation handling.
+ * unit loop elimination, and loop policy inheritance.
*/
-#include <tvm/ffi/extra/structural_visit.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/runtime/logging.h>
@@ -63,22 +62,6 @@ class TIRxOpaqueLower : public StmtExprMutator {
std::optional<ffi::Any> unroll_explicit;
};
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- // Attribute subjects may contain references to the loop binding being
- // replaced, including expressions and IterVar fields. Rewrite metadata
- // before the body introduces any additional local remaps.
- auto node = Mutate(op->node, inplace_mode);
- auto value = Mutate(op->value, inplace_mode);
- auto body = Mutate(op->body, inplace_mode);
- if (node.UnchangedOrSameAs(op->node) && value.UnchangedOrSameAs(op->value)
&&
- body.UnchangedOrSameAs(op->body)) {
- return ffi::Unchanged();
- }
- return AttrStmt(std::move(node).ValueOrUnchanged(op->node), op->attr_key,
- std::move(value).ValueOrUnchanged(op->value),
- std::move(body).ValueOrUnchanged(op->body), op->span);
- }
-
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
// Step 1. Update unit loop info.
PrimExpr min = this->Mutate(op->min,
inplace_mode).ValueOrUnchanged(op->min);
@@ -117,28 +100,7 @@ class TIRxOpaqueLower : public StmtExprMutator {
});
VarRemapSet(op->loop_var, previous_remap);
- // Step 3. Handle annotations
- std::vector<std::pair<std::string, Expr>> pragma_attrs;
- ffi::Map<ffi::String, ffi::Any> new_annotations =
HandleAnnotations(annotations, &pragma_attrs);
- if (op->kind != ForKind::kThreadBinding) {
- // A pragma that uses this loop's own Var must remain inside its scope.
- // Pragmas using only enclosing bindings retain their wrapper position.
- std::vector<std::pair<std::string, Expr>> outer_pragmas;
- for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) {
- auto uses_loop_var = [&](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
- return var.same_as(op->loop_var) ?
ffi::WalkResult::Interrupt(ffi::VisitInterrupt(var))
- : ffi::WalkResult::Advance();
- };
- if (ffi::StructuralWalk<ffi::WalkOrder::kPreOrder>(it->second,
uses_loop_var).has_value()) {
- body = AttrStmt(op->loop_var, it->first, it->second,
std::move(body));
- } else {
- outer_pragmas.push_back(*it);
- }
- }
- std::reverse(outer_pragmas.begin(), outer_pragmas.end());
- pragma_attrs = std::move(outer_pragmas);
- }
- // Step 4. Create new For loop accordingly
+ // Step 2. Create the lowered loop or launch region.
if (op->kind == ForKind::kThreadBinding) {
// Case 1. Thread binding → RegionStmt(launch_thread)
TVM_FFI_ICHECK(!op->annotations.count("loop_partition_hint") ||
@@ -152,12 +114,7 @@ class TIRxOpaqueLower : public StmtExprMutator {
} else {
// Case 3. An ordinary loop
body = For(op->loop_var, std::move(min), std::move(extent), op->kind,
std::move(body),
- std::nullopt, new_annotations, step);
- }
- // Step 5. Insert nested attrs for pragma annotations
- for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) {
- Var var = op->kind == ForKind::kThreadBinding ? Var(launch_var) :
Var(op->loop_var);
- body = AttrStmt(var, it->first, it->second, std::move(body));
+ std::nullopt, FilterAnnotations(annotations), step);
}
if (op->kind == ForKind::kThreadBinding) {
return RegionStmt(tirx::builtin::launch_thread(),
@@ -182,58 +139,20 @@ class TIRxOpaqueLower : public StmtExprMutator {
// Effective policy is materialized on each surviving loop, preserving
nested overrides.
ScopeStack<UnrollPolicy> unroll_policy_;
- /*! \brief Convert attr value from annotation map into Expr. */
- Expr ConvertAttrValue(const ffi::String& key, const Any& obj) {
- if (auto expr = obj.try_cast<Expr>()) {
- return expr.value();
- } else if (auto str = obj.try_cast<ffi::String>()) {
- return std::move(StringImm(str.value()));
- } else {
- LOG(FATAL) << "Illegal attribute of key " << key << ", value type " <<
obj.GetTypeKey()
- << " not supported";
- }
- }
-
- /*!
- * \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), 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.
- */
- ffi::Map<ffi::String, ffi::Any> HandleAnnotations(
- const ffi::Map<ffi::String, ffi::Any>& annotations,
- std::vector<std::pair<std::string, Expr>>* pragma_attrs) {
- ffi::Map<ffi::String, ffi::Any> preserved_annotations;
- pragma_attrs->clear();
- for (const auto& kv : annotations) {
- const ffi::String& key = kv.first;
- if ((key == tirx::attr::auto_unroll_max_step || key ==
tirx::attr::unroll_explicit) &&
- kv.second == nullptr) {
+ // Null optional policies do not override inherited values or reach code
generation.
+ ffi::Map<ffi::String, ffi::Any> FilterAnnotations(
+ const ffi::Map<ffi::String, ffi::Any>& annotations) {
+ ffi::Map<ffi::String, ffi::Any> preserved;
+ for (const auto& [key, value] : annotations) {
+ if ((key == tirx::attr::auto_unroll_max_step || key ==
tirx::attr::unroll_explicit ||
+ key == "pragma_unroll") &&
+ value == nullptr) {
continue;
}
- if (key == "pragma_unroll") {
- if (kv.second != nullptr) {
- preserved_annotations.Set(key, kv.second);
- }
- } else if (tirx::attr::IsPragmaKey(key)) {
- if (kv.second == nullptr) {
- continue;
- }
-
- pragma_attrs->emplace_back(key, ConvertAttrValue(key, kv.second));
- } else {
- // Loop annotations are always preserved
- preserved_annotations.Set(key, kv.second);
- }
+ preserved.Set(key, value);
}
- std::sort(pragma_attrs->begin(), pragma_attrs->end(),
- [](const auto& p1, const auto& p2) { return p1.first < p2.first;
});
- return preserved_annotations;
+ return preserved;
}
-
- /*! \brief Record the loop_var and loop start value of unit loops, whose
extent is one. */
};
namespace transform {
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index 2046074656..eb5e21ee6e 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -367,20 +367,6 @@ class BuiltinLower : public StmtExprMutator {
});
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- return scope_.WithNewScope([&]() -> Stmt {
- Stmt visited =
- StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- if (!scope_.Current().pending_frees.empty()) {
- const auto* attr = visited.as<AttrStmtNode>();
- if (attr) {
- return AttrStmt(attr->node, attr->attr_key, attr->value,
AppendPendingFrees(attr->body),
- attr->span);
- }
- }
- return visited;
- });
- }
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
auto min_result = this->Mutate(op->min, inplace_mode);
bool min_unchanged = min_result.UnchangedOrSameAs(op->min);
@@ -735,7 +721,7 @@ class BuiltinLower : public StmtExprMutator {
*
* When a Bind allocates via nd_mem_alloc_with_scope, the corresponding
* free_nd stmt is pushed to the current scope's pending_frees. Body-carrying
- * stmts (For, IfThenElse, AttrStmt) create new scopes via
+ * stmts (For, IfThenElse, RegionStmt) create new scopes via
* WithNewScope. On scope exit, pending_frees are appended after the body.
* AllocTensor (flat, no body) pushes its free to the enclosing scope.
*/
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 5d9033d77a..8235e58c26 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -91,11 +91,6 @@ class NoOpRemover : public IRMutatorWithAnalyzer {
: Parent(analyzer), ignore_profiler_call_(ignore_profiler_call) {}
private:
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = stmt.as<AttrStmtNode>();
- return is_no_op(op->body) ? MakeEvaluate(op->value) : stmt;
- }
UnchangedOr<Stmt> Mutate_(const IfThenElseNode* op, InplaceMode
inplace_mode) final {
Stmt stmt = Parent::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
op = stmt.as<IfThenElseNode>();
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 7c5fb6c943..24ac6e62e6 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -346,9 +346,7 @@ class InplaceOpVerifier : public StmtExprVisitor {
dst_ = dst;
src_ = src;
result_ = true;
- if (stmt->IsInstance<AttrStmtNode>()) {
- Visit_(static_cast<const AttrStmtNode*>(stmt));
- } else if (stmt->IsInstance<ForNode>()) {
+ if (stmt->IsInstance<ForNode>()) {
Visit_(static_cast<const ForNode*>(stmt));
} else if (stmt->IsInstance<IfThenElseNode>()) {
Visit_(static_cast<const IfThenElseNode*>(stmt));
@@ -681,23 +679,6 @@ class StoragePlanRewriter : public StmtExprMutator {
return region;
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- if (attr::IsPragmaKey(op->attr_key)) {
- // remake all the allocation at the attach scope.
- if (attach_map_.count(op)) {
- auto& svec = attach_map_[op];
- Stmt stmt =
- StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = stmt.as<AttrStmtNode>();
- return AttrStmt(op->node, op->attr_key, op->value, MakeAttach(svec,
op->body));
- } else {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
- } else {
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
- }
-
UnchangedOr<Stmt> Mutate_(const ForNode* op, InplaceMode inplace_mode) final
{
TVM_FFI_ICHECK(op->kind != ForKind::kVectorized) << "VectorizeLoop before
LiftStorageAlloc";
// remake all the allocation at the attach scope.
@@ -1176,12 +1157,7 @@ class StoragePlanRewriter : public StmtExprMutator {
}
}
// enter/exit new scope
- if (s.stmt->IsInstance<AttrStmtNode>()) {
- const auto* op = static_cast<const AttrStmtNode*>(s.stmt);
- if (attr::IsPragmaKey(op->attr_key)) {
- PlanNewScope(op);
- }
- } else if (s.stmt->IsInstance<RegionStmtNode>()) {
+ if (s.stmt->IsInstance<RegionStmtNode>()) {
PlanNewScope(s.stmt);
} else if (s.stmt->IsInstance<ForNode>()) {
const auto* op = static_cast<const ForNode*>(s.stmt);
@@ -2149,20 +2125,6 @@ class VectorTypeRewriter : public StmtExprMutator {
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode
inplace_mode) final {
- Stmt stmt =
- StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = stmt.as<AttrStmtNode>();
- TVM_FFI_ICHECK(op != nullptr);
- if (auto var = op->node.as<Var>()) {
- auto mapped = VarRemapGet(var.value());
- if (mapped != nullptr) {
- return AttrStmt(mapped.as_or_throw<Var>(), op->attr_key,
op->value, op->body, op->span);
- }
- }
- return stmt;
- }
-
UnchangedOr<Expr> Mutate_(const CallNode* op, InplaceMode inplace_mode)
final {
auto result = StmtExprMutator::Mutate_(op, inplace_mode);
if (!result.IsUnchanged()) {
diff --git a/src/tirx/transform/tvm_ffi_binder.h
b/src/tirx/transform/tvm_ffi_binder.h
index 8fc9ae1200..e012c80c6b 100644
--- a/src/tirx/transform/tvm_ffi_binder.h
+++ b/src/tirx/transform/tvm_ffi_binder.h
@@ -59,7 +59,7 @@ namespace tirx {
* by a later buffer's shape (batch_size). Separating definitions from
* checks guarantees all variables are in scope when assertions reference them.
*
- * - init_nest: Binds, DeclTensors for shape/strides arrays, AttrStmts —
+ * - init_nest: Binds, DeclTensors for shape/strides arrays —
* all value-loading code that defines variables.
* - asserts: AssertStmts — all validation checks.
* - decl_buffers: DeclTensor for buffer-typed parameters — buffer
declarations.
@@ -96,7 +96,7 @@ class TVMFFIABIBuilder {
struct Result {
/*! \brief Var -> VarDefInfo map for defined variables. */
std::unordered_map<const VarNode*, VarDefInfo> var_defs;
- /*! \brief Variable definitions (Binds, shape/strides DeclTensors,
AttrStmts). */
+ /*! \brief Variable definitions (Binds, shape/strides DeclTensors). */
std::vector<Stmt> init_nest;
/*! \brief Validation checks (all AssertStmts). */
std::vector<Stmt> asserts;
@@ -385,7 +385,7 @@ class TVMFFIABIBuilder {
/*! \brief The definition map: VarNode* -> VarDefInfo (value +
first_def_path). */
std::unordered_map<const VarNode*, VarDefInfo> var_defs_;
- /*! \brief Variable definitions: Binds, shape/strides DeclTensors,
AttrStmts. */
+ /*! \brief Variable definitions: Binds, shape/strides DeclTensors. */
std::vector<Stmt> init_nest_;
/*! \brief Validation checks: all AssertStmts. */
std::vector<Stmt> asserts_;
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc
b/src/tirx/transform/unsupported_dtype_legalize.cc
index a629658458..54b4dd1767 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -398,18 +398,6 @@ class ComputeLegalizer : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- Stmt ret = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = ret.as<AttrStmtNode>();
- if (auto var = op->node.as<Var>()) {
- auto mapped = VarRemapGet(var.value());
- if (mapped != nullptr) {
- return AttrStmt(mapped.as_or_throw<Var>(), op->attr_key, op->value,
op->body);
- }
- }
- return ret;
- }
-
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) final {
TensorVar buffer = GetRemappedBuffer(op->source.as_or_throw<TensorVar>());
auto indices = Mutate(op->indices, inplace_mode)
@@ -594,19 +582,6 @@ class StorageLegalizer : public StmtExprMutator {
}
}
- UnchangedOr<Stmt> Mutate_(const AttrStmtNode* op, InplaceMode inplace_mode)
final {
- Stmt ret = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- op = ret.as<AttrStmtNode>();
-
- if (auto var = op->node.as<Var>()) {
- auto mapped = VarRemapGet(var.value());
- if (mapped != nullptr) {
- return AttrStmt(mapped.as_or_throw<Var>(), op->attr_key, op->value,
op->body);
- }
- }
- return ret;
- }
-
UnchangedOr<PrimExpr> Mutate_(const TensorLoadNode* op, InplaceMode
inplace_mode) final {
TensorVar buffer = GetRemappedBuffer(op->source.as_or_throw<TensorVar>());
auto indices = Mutate(op->indices, inplace_mode)
diff --git a/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py
b/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py
index 894689d1b8..4a69d3ec5e 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_hoist_if.py
@@ -50,8 +50,6 @@ def verify_structure(stmt, expected_struct):
elif isinstance(op, tvm.tirx.RegionStmt):
assert op.op.name == "tirx.launch_thread"
val = [(op.body,), ("tirx.RegionStmt", op.op.name,
int(op.args[1]))]
- elif isinstance(op, tvm.tirx.AttrStmt):
- val = [(op.body,), ("tirx.AttrStmt", op.attr_key, int(op.value))]
else:
return
node_dict[key] = val
@@ -161,7 +159,7 @@ def test_no_else():
verify_structure(new_stmt, expected_struct)
-def test_attr_stmt():
+def test_thread_launch():
dshape = (32, 64)
@Ts.prim_func(private=True)
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
index caa0a5fa58..900cdc9593 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
+++ b/tests/python/s_tir/transform/test_s_tir_transform_lower_opaque_block.py
@@ -303,12 +303,6 @@ def transformed_symbolic_strided_buffer_func(A:
T.Tensor((1, n, 10240))):
)
[email protected]_func
-def annotated_loops(A: T.Tensor((16,), "float32")) -> None:
- for i in range(0, 16, annotations={"pragma_1": "str_value", "pragma_2": 1,
"pragma_3": 0.0}):
- A[i] = 0.0
-
-
@Ts.prim_func
def boolean_handling_before(a: T.Tensor(10, "bool"), b: T.Tensor(10, "bool"))
-> None:
for i0 in T.serial(10):
@@ -357,38 +351,6 @@ def test_symbolic_strided_buffer():
_check(compacted_symbolic_strided_buffer_func,
transformed_symbolic_strided_buffer_func)
-def test_annotated_loops():
- mod = tvm.IRModule.from_expr(annotated_loops.with_attr("global_symbol",
"main"))
- mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
- attr1 = mod["main"].body
- attr2 = attr1.body
- attr3 = attr2.body
- assert attr1.attr_key == "pragma_1" and attr1.value == "str_value"
- assert attr2.attr_key == "pragma_2"
- tvm.ir.assert_structural_equal(attr2.value, tvm.tirx.IntImm("int32", 1))
- assert attr3.attr_key == "pragma_3"
- tvm.ir.assert_structural_equal(attr3.value, tvm.tirx.FloatImm("float32",
0.0))
-
-
-def test_annotated_block():
- @Ts.prim_func
- def annotated_block() -> None:
- with Ts.sblock():
- Ts.sblock_attr({"pragma_1": "str_value", "pragma_2": 1,
"pragma_3": 0.0})
- T.evaluate(0)
-
- mod = tvm.IRModule.from_expr(annotated_block.with_attr("global_symbol",
"main"))
- mod = tvm.s_tir.transform.LowerOpaqueBlock()(mod)
- attr1 = mod["main"].body
- attr2 = attr1.body
- attr3 = attr2.body
- assert attr1.attr_key == "pragma_1" and attr1.value == "str_value"
- assert attr2.attr_key == "pragma_2"
- tvm.ir.assert_structural_equal(attr2.value, tvm.tirx.IntImm("int32", 1))
- assert attr3.attr_key == "pragma_3"
- tvm.ir.assert_structural_equal(attr3.value, tvm.tirx.FloatImm("float32",
0.0))
-
-
def test_preserved_annotations():
@Ts.prim_func
def before(A: T.Tensor(8, "float32"), B: T.Tensor(8, "float32")):
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index a65e88eb84..60fc6a9f9a 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -292,10 +292,6 @@ def test_stmt_constructor():
assert x.var == v
assert x.value.value == 1
- x = tvm.tirx.AttrStmt(v == 1, "xx", 1, tvm.tirx.Evaluate(1))
- assert isinstance(x, tvm.tirx.AttrStmt)
- assert x.value.value == 1
-
x = tvm.tirx.AssertStmt(
tvm.tirx.const(1, "bool"),
tvm.ir.StringImm("RuntimeError"),
@@ -339,12 +335,6 @@ def test_stmt_constructor():
assert _is_buffer_binding(x, "tirx.alloc_tensor")
assert x.var == buf
- x = tvm.tirx.AttrStmt(buffer_var, "xyz", 1, nop)
- assert isinstance(x, tvm.tirx.AttrStmt)
- assert x.node == buffer_var
- assert x.attr_key == "xyz"
- assert x.body == nop
-
x = tvm.tirx.IfThenElse(tvm.tirx.const(1, "bool"), tvm.tirx.Evaluate(11),
nop)
assert isinstance(x, tvm.tirx.IfThenElse)
assert x.then_case.value.value == 11
diff --git a/tests/python/tirx-base/test_tir_nodes.py
b/tests/python/tirx-base/test_tir_nodes.py
index f1c03f794a..2f247b37d2 100644
--- a/tests/python/tirx-base/test_tir_nodes.py
+++ b/tests/python/tirx-base/test_tir_nodes.py
@@ -111,11 +111,6 @@ def test_cast():
def test_attr():
- x = tvm.tirx.Var("x", "int32")
- y = tvm.tirx.Var("y", "int32")
- stmt = tvm.tirx.AttrStmt(y, "stride", 10, tvm.tirx.Evaluate(x + 1))
- assert stmt.node == y
-
a = tvm.runtime.convert(1)
assert a == 1
try:
diff --git a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
index 6ef9260492..36dc126ce8 100644
--- a/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
+++ b/tests/python/tirx-transform/test_tir_transform_convert_ssa.py
@@ -453,56 +453,6 @@ def test_thread_idx_reused_within_and_across_functions():
tvm.ir.assert_structural_equal(after, expected)
-def test_track_forward_declarations_in_attr_stmt():
- """T.attr statements may refer to a about-to-be-defined tirx.Var"""
-
- # Generate the PrimFunc, which is already SSA
- #
- # This is constructed directly, rather than using TVMScript.
- # This test case requires a `tirx.AttrStmt` that references a
- # variable, followed by the `tirx.For` defining that variable.
- # This is not expressible in TVMScript, as it only provides the
- # loop iterator within the body of the loop.
- i0_outer_outer = tirx.Var("i0_outer_outer", "int32")
- i0_outer_inner = tirx.Var("i0_outer_inner", "int32")
- i0_inner = tirx.Var("i0_inner", "int32")
-
- A = tirx.decl_tensor(1024, "float32", "A")
- B = tirx.decl_tensor(1024, "float32", "B")
-
- index = i0_outer_outer * 52 + i0_outer_inner * 4 + i0_inner
-
- stmt = tirx.BufferStore(B, tirx.TensorLoad(A, [index]), [index])
- stmt = tirx.IfThenElse(i0_outer_outer * 13 + i0_outer_inner < 256, stmt,
None)
- stmt = tirx.For(i0_inner, 0, 4, tirx.ForKind.VECTORIZED, stmt)
- stmt = tirx.For(i0_outer_inner, 0, 13, tirx.ForKind.PARALLEL, stmt)
- stmt = tirx.AttrStmt(
- T.iter_var(i0_outer_inner, None, "DataPar", ""),
- "pragma_parallal_barrier_when_finish",
- 1,
- stmt,
- )
- stmt = tirx.AttrStmt(
- T.iter_var(i0_outer_inner, None, "DataPar", ""),
- "pragma_parallal_stride_pattern",
- 1,
- stmt,
- )
- stmt = tirx.For(i0_outer_outer, 0, 20, tirx.ForKind.SERIAL, stmt)
- stmt = tirx.AttrStmt(
- T.iter_var(i0_outer_outer, None, "DataPar", ""),
- "pragma_parallal_launch_point",
- 1,
- stmt,
- )
-
- before = tirx.PrimFunc([A, B], stmt)
-
- mod = tvm.IRModule.from_expr(before)
- after = tvm.tirx.transform.ConvertSSA()(mod)
- tvm.ir.assert_structural_equal(after["main"], before)
-
-
def test_shared_shape_var_in_buffer_params_and_alloc_buffer():
"""Shape var shared across buffer params and AllocTensor should not be
renamed.
diff --git
a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
index 2c90bf735e..d71439a79d 100644
--- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py
+++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py
@@ -159,7 +159,7 @@ def test_device_kernel_nonzero_return_is_rejected():
def test_split_host_device_without_func_host_attribute():
"""Like test_split_host_device, but no host specified in the host's target
- The `T.attr` specifying the device still requires splitting out
+ The device scope region still requires splitting out
the kernel.
"""
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index b77418b537..348512e790 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -204,20 +204,6 @@ for i in range(128):
)
-def test_attr_stmt():
- with IRBuilder() as ib:
- with TB.attr("pragma", "unroll", 1):
- TB.evaluate(0)
- obj = ib.get()
- _assert_print(
- obj,
- """
-with T.attr("pragma", "unroll", 1):
- T.evaluate(0)
-""",
- )
-
-
def test_assert_stmt():
with IRBuilder() as ib:
with TB.prim_func():
