You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by ju...@apache.org on 2026/04/03 22:05:11 UTC
(tvm) branch junrus/2025-04-02/always-use-bundled-tvm-ffi updated: tmp
This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch junrus/2025-04-02/always-use-bundled-tvm-ffi
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/junrus/2025-04-02/always-use-bundled-tvm-ffi by this push:
new fd47c1a470 tmp
fd47c1a470 is described below
commit fd47c1a4708bd492fae8cc8ffd3c15022ead851a
Author: Junru Shao <ju...@nvidia.com>
AuthorDate: Fri Apr 3 15:05:00 2026 -0700
tmp
---
.gitignore | 6 +
CMakeLists.txt | 1 +
python/tvm/runtime/script_printer.py | 13 +
python/tvm/script/ir_builder/tirx/ir.py | 2 +
python/tvm/tirx/op.py | 13 +
src/ir/attrs.cc | 11 +
src/ir/expr.cc | 81 ++++
src/ir/global_info.cc | 50 ++
src/ir/op.cc | 14 +-
src/ir/type.cc | 28 ++
src/relax/ir/expr.cc | 130 +++++
src/relax/ir/struct_info.cc | 34 ++
src/relax/ir/type.cc | 30 ++
src/script_v2/printer/ir/module.cc | 160 ++++++
src/script_v2/printer/ir/primitives.cc | 137 ++++++
src/script_v2/printer/relax/expr.cc | 829 ++++++++++++++++++++++++++++++++
src/script_v2/printer/tirx/block.cc | 194 ++++++++
src/script_v2/printer/tirx/buffer.cc | 222 +++++++++
src/script_v2/printer/tirx/expr.cc | 184 +++++++
src/script_v2/printer/tirx/for_loop.cc | 155 ++++++
src/script_v2/printer/tirx/function.cc | 336 +++++++++++++
src/script_v2/printer/tirx/stmt.cc | 103 ++++
src/script_v2/printer/tirx/utils.h | 304 ++++++++++++
src/script_v2/printer/utils.h | 136 ++++++
src/tirx/ir/buffer.cc | 22 +-
src/tirx/ir/expr.cc | 385 +++++++++++++++
src/tirx/ir/function.cc | 27 ++
src/tirx/ir/stmt.cc | 204 ++++++++
28 files changed, 3809 insertions(+), 2 deletions(-)
diff --git a/.gitignore b/.gitignore
index 93f5841047..eef0b3f781 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,9 @@
+plan-changed.md
+plan.md
+traits-printer-parser-v4.md
+testsuite/
+ir_testsuite.jsonl
+
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0950db7b0b..dc4d904991 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -277,6 +277,7 @@ tvm_file_glob(GLOB_RECURSE COMPILER_SRCS
src/driver/*.cc
src/support/*.cc
src/script/*.cc
+ src/script_v2/*.cc
src/relax/ir/*.cc
src/relax/op/*.cc
src/relax/analysis/*.cc
diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py
index 6ba9abb032..e324cbf8ba 100644
--- a/python/tvm/runtime/script_printer.py
+++ b/python/tvm/runtime/script_printer.py
@@ -224,6 +224,19 @@ class Scriptable:
),
)
+ def script_v2(self, *, indent_spaces: int = 4) -> str:
+ """Print using the traits-based V2 printer.
+
+ Returns
+ -------
+ script : str
+ The V2 TVM Script of the given TVM IR
+ """
+ from tvm_ffi.ir.text import to_python, PrinterConfig
+
+ cfg = PrinterConfig(indent_spaces=indent_spaces)
+ return to_python(self, cfg)
+
def _relax_script(
self,
*,
diff --git a/python/tvm/script/ir_builder/tirx/ir.py b/python/tvm/script/ir_builder/tirx/ir.py
index ce88c56315..4113d4c07e 100644
--- a/python/tvm/script/ir_builder/tirx/ir.py
+++ b/python/tvm/script/ir_builder/tirx/ir.py
@@ -1928,6 +1928,7 @@ truncdiv = _op_wrapper(_tir_op.truncdiv)
truncmod = _op_wrapper(_tir_op.truncmod)
tvm_access_ptr = _op_wrapper(_tir_op.tvm_access_ptr)
tvm_throw_last_error = _op_wrapper(_tir_op.tvm_throw_last_error)
+tvm_static_handle = _op_wrapper(_tir_op.tvm_static_handle)
tvm_stack_alloca = _op_wrapper(_tir_op.tvm_stack_alloca)
tvm_stack_make_shape = _op_wrapper(_tir_op.tvm_stack_make_shape)
tvm_stack_make_array = _op_wrapper(_tir_op.tvm_stack_make_array)
@@ -2208,6 +2209,7 @@ __all__ = float_types + [
"truncmod",
"tvm_access_ptr",
"tvm_throw_last_error",
+ "tvm_static_handle",
"tvm_stack_alloca",
"tvm_stack_make_shape",
"tvm_stack_make_array",
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 6b4a636f30..b8d7244dfa 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -801,6 +801,19 @@ def tvm_throw_last_error():
return call_intrin("handle", "tirx.tvm_throw_last_error")
+def tvm_static_handle():
+ """Create a function-local static handle initialized to nullptr.
+
+ Can be used to cache function-local static resources.
+
+ Returns
+ -------
+ call : PrimExpr
+ The call expression.
+ """
+ return call_intrin("handle", "tirx.tvm_static_handle")
+
+
def tvm_load_matrix_sync(fragment, m, n, k, index, buffer_ptr, stride, layout):
"""TVM intrinsic for tensor core load operators
diff --git a/src/ir/attrs.cc b/src/ir/attrs.cc
index 748f4bf5c9..062285224f 100644
--- a/src/ir/attrs.cc
+++ b/src/ir/attrs.cc
@@ -21,6 +21,7 @@
* \file attrs.cc
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
@@ -31,6 +32,16 @@ namespace tvm {
TVM_FFI_STATIC_INIT_BLOCK() {
AttrFieldInfoNode::RegisterReflection();
DictAttrsNode::RegisterReflection();
+
+ // DictAttrs → Literal (transparent wrapper, prints inner dict)
+ namespace tr = ::tvm::ffi::ir::traits;
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ auto trait = tr::Trait(ffi::make_object<tr::LiteralObj>(
+ ffi::String("$field:__dict__"), ffi::Optional<ffi::String>()));
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(DictAttrsNode::RuntimeTypeIndex(), &name_arr, raw);
}
DictAttrs WithAttrs(DictAttrs attrs, ffi::Map<ffi::String, ffi::Any> new_attrs) {
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index f9d6e0fc60..8b4aac30c3 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -23,6 +23,9 @@
*/
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/text/ast.h>
+#include <tvm/ffi/ir/text/printer.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/function.h>
@@ -42,6 +45,55 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IntImmNode::RegisterReflection();
FloatImmNode::RegisterReflection();
RangeNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for base IR expr nodes ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+
+ // IntImm → Literal
+ register_trait(IntImmNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:value"), OS("int"))));
+ // FloatImm → Literal
+ register_trait(FloatImmNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:value"), OS("float"))));
+ // GlobalVar → Call("I.GlobalVar", "$field:name_hint")
+ // NOTE: The trait is only used when __ffi_text_print__ doesn't apply (see below).
+ register_trait(GlobalVarNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ S("I.GlobalVar"), S("$field:name_hint"), ffi::Optional<ffi::String>())));
+ // Range → Call("I.Range", "$method:range_args")
+ register_trait(RangeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ S("I.Range"), S("$method:range_args"), ffi::Optional<ffi::String>())));
+}
+
+// GlobalVar: __ffi_text_print__ override — check printer var table first.
+// When a GlobalVar is registered in the printer (e.g., as Module.func_name),
+// return the stored reference. Otherwise, fall through to I.GlobalVar("name").
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace text = ::tvm::ffi::ir::text;
+ namespace refl = ::tvm::ffi::reflection;
+ refl::TypeAttrDef<GlobalVarNode>().def(
+ "__ffi_text_print__",
+ [](GlobalVar gv, text::IRPrinter printer, refl::AccessPath path) -> text::NodeAST {
+ // Check if this GlobalVar is bound in the printer's variable table
+ if (auto doc = printer->VarGet(gv)) {
+ return doc.value();
+ }
+ // Fallback: print as I.GlobalVar("name_hint")
+ text::ExprAST callee = text::ExprAttr(text::IdAST("I"), "GlobalVar");
+ ffi::List<text::ExprAST> args;
+ args.push_back(text::LiteralAST::Str(gv->name_hint));
+ return text::CallAST(callee, std::move(args), {}, {});
+ });
}
PrimExpr::PrimExpr(int32_t value) : PrimExpr(IntImm(DataType::Int(32), value)) {}
@@ -231,4 +283,33 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
+// --- Register computed-property methods for $method: trait references ---
+namespace {
+template <typename Func>
+void RegisterComputedMethod(int32_t type_index, const char* type_key, const char* method_name,
+ Func&& func) {
+ using namespace ::tvm::ffi;
+ TVMFFIMethodInfo info;
+ info.name = TVMFFIByteArray{method_name, std::char_traits<char>::length(method_name)};
+ info.doc = TVMFFIByteArray{nullptr, 0};
+ info.flags = 0;
+ std::string full_name = std::string(type_key) + "." + method_name;
+ Function method = Function::FromTyped(std::forward<Func>(func), std::move(full_name));
+ info.method = AnyView(method).CopyToTVMFFIAny();
+ static const char kEmptyMetadata[] = "{}";
+ info.metadata = TVMFFIByteArray{kEmptyMetadata, 2};
+ TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index, &info));
+}
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // Range: args = [min, min + extent]
+ RegisterComputedMethod(
+ RangeNode::RuntimeTypeIndex(), RangeNode::_type_key, "range_args",
+ [](Range node) -> ffi::Array<PrimExpr> {
+ PrimExpr max_val = node->min + node->extent;
+ return {node->min, max_val};
+ });
+}
+
} // namespace tvm
diff --git a/src/ir/global_info.cc b/src/ir/global_info.cc
index 151387d3c2..d1888f25bd 100644
--- a/src/ir/global_info.cc
+++ b/src/ir/global_info.cc
@@ -22,13 +22,32 @@
* \brief Module global info.
*/
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/global_info.h>
+#include <tvm/tirx/expr.h>
+
namespace tvm {
TVM_FFI_STATIC_INIT_BLOCK() {
VDeviceNode::RegisterReflection();
DummyGlobalInfoNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for VDevice ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ // VDevice → Call("I.vdevice", "$method:vdevice_args")
+ register_trait(VDeviceNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ S("I.vdevice"), S("$method:vdevice_args"), ffi::Optional<ffi::String>())));
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -53,4 +72,35 @@ TVM_FFI_STATIC_INIT_BLOCK() {
return VDevice(tgt, dev_id, mem_scope);
});
}
+// --- Register computed-property methods for $method: trait references ---
+namespace {
+template <typename Func>
+void RegisterComputedMethod(int32_t type_index, const char* type_key, const char* method_name,
+ Func&& func) {
+ using namespace ::tvm::ffi;
+ TVMFFIMethodInfo info;
+ info.name = TVMFFIByteArray{method_name, std::char_traits<char>::length(method_name)};
+ info.doc = TVMFFIByteArray{nullptr, 0};
+ info.flags = 0;
+ std::string full_name = std::string(type_key) + "." + method_name;
+ Function method = Function::FromTyped(std::forward<Func>(func), std::move(full_name));
+ info.method = AnyView(method).CopyToTVMFFIAny();
+ static const char kEmptyMetadata[] = "{}";
+ info.metadata = TVMFFIByteArray{kEmptyMetadata, 2};
+ TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index, &info));
+}
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // VDevice: args = [target.ToConfig(), IntImm(vdevice_id), StringImm(memory_scope)]
+ RegisterComputedMethod(
+ VDeviceNode::RuntimeTypeIndex(), VDeviceNode::_type_key, "vdevice_args",
+ [](VDevice node) -> ffi::Array<ObjectRef> {
+ ffi::Map<ffi::String, ffi::Any> config = node->target->ToConfig();
+ IntImm dev_id(DataType::Int(64), static_cast<int64_t>(node->vdevice_id));
+ tirx::StringImm mem_scope(node->memory_scope);
+ return {config, dev_id, mem_scope};
+ });
+}
+
} // namespace tvm
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 0f8ff2ea7d..43de0905ce 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -22,6 +22,7 @@
* \brief Primitive operators and intrinsics.
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/ir/type.h>
@@ -34,7 +35,18 @@
namespace tvm {
-TVM_FFI_STATIC_INIT_BLOCK() { OpNode::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ OpNode::RegisterReflection();
+
+ namespace tr = ::tvm::ffi::ir::traits;
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ auto trait = tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("I.Op"), ffi::String("$field:name"), ffi::Optional<ffi::String>()));
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(OpNode::RuntimeTypeIndex(), &name_arr, raw);
+}
using ffi::Any;
using ffi::Function;
diff --git a/src/ir/type.cc b/src/ir/type.cc
index b28e20a78f..1d577dfa8f 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -22,6 +22,7 @@
* \brief Common type system AST nodes throughout the IR.
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/type.h>
namespace tvm {
@@ -33,6 +34,33 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TupleTypeNode::RegisterReflection();
FuncTypeNode::RegisterReflection();
TensorMapTypeNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for base type nodes ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+ auto NS = []() { return ffi::Optional<ffi::String>(); };
+
+ // PrimType → PrimTy
+ register_trait(PrimTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::PrimTyObj>(S("$field:dtype"))));
+ // PointerType → PointerTy
+ register_trait(PointerTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::PointerTyObj>(
+ S("$field:element_type"), OS("$field:storage_scope"))));
+ // TupleType → TupleTy
+ register_trait(TupleTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::TupleTyObj>(S("$field:fields"))));
+ // FuncType → FuncTy
+ register_trait(FuncTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::FuncTyObj>(OS("$field:arg_types"), OS("$field:ret_type"))));
}
PrimType::PrimType(runtime::DataType dtype, Span span) {
diff --git a/src/relax/ir/expr.cc b/src/relax/ir/expr.cc
index 2fbd573a5f..438077d1bb 100644
--- a/src/relax/ir/expr.cc
+++ b/src/relax/ir/expr.cc
@@ -16,9 +16,12 @@
* specific language governing permissions and limitations
* under the License.
*/
+#include <tvm/ffi/ir/text/ast.h>
+#include <tvm/ffi/ir/text/printer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/relax/struct_info.h>
#include <tvm/relax/type.h>
@@ -50,6 +53,133 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IfNode::RegisterReflection();
FunctionNode::RegisterReflection();
ExternFuncNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for ALL Relax nodes ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+ auto NS = []() { return ffi::Optional<ffi::String>(); };
+ auto NR = []() { return ffi::Optional<tr::Region>(); };
+
+ // Var → Value (name from vid.name_hint)
+ register_trait(VarNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::ValueObj>(S("$field:vid"), NS(), NS())));
+ // DataflowVar → Value
+ register_trait(DataflowVarNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::ValueObj>(S("$field:vid"), NS(), NS())));
+ // VarBinding → Assign
+ register_trait(VarBindingNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssignObj>(S("$field:var"), S("$field:value"))));
+ // MatchCast → Assign
+ register_trait(MatchCastNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssignObj>(S("$field:var"), S("$field:value"))));
+ // If → If
+ register_trait(IfNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::IfObj>(
+ S("$field:cond"),
+ tr::Region(S("$field:true_branch")),
+ ffi::Optional<tr::Region>(tr::Region(S("$field:false_branch"))))));
+ // Function → Func (symbol is a literal "func" since name comes from GlobalVar binding)
+ register_trait(FunctionNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::FuncObj>(
+ S("func"),
+ tr::Region(S("$field:body"), OS("$field:params")),
+ NS(), NS(), NS())));
+ // ExternFunc → Literal
+ register_trait(ExternFuncNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:global_symbol"),
+ OS("extern_func"))));
+ // Call → Call
+ register_trait(CallNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ S("$field:op"), S("$field:args"), OS("$field:attrs"))));
+ // Tuple → Call-like (R.tuple(fields...))
+ register_trait(TupleNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(S("Tuple"), S("$field:fields"), NS())));
+ // TupleGetItem → Load-like (index into tuple)
+ register_trait(TupleGetItemNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LoadObj>(S("$field:tuple_value"),
+ OS("$field:index"))));
+ // ShapeExpr → Literal (shape value)
+ register_trait(ShapeExprNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:values"), OS("shape"))));
+ // Constant → Literal
+ register_trait(ConstantNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:data"), OS("constant"))));
+ // PrimValue → Literal
+ register_trait(PrimValueNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:value"), OS("prim"))));
+ // StringImm → Literal
+ register_trait(StringImmNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:value"), OS("string"))));
+ // DataTypeImm → Literal
+ register_trait(DataTypeImmNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(S("$field:value"), OS("dtype"))));
+ // BindingBlock → Seq
+ register_trait(BindingBlockNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::SeqObj>(S("$field:bindings"))));
+ // DataflowBlock → With (scoped dataflow region with R.output epilogue)
+ register_trait(DataflowBlockNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::WithObj>(
+ tr::Region(S("$field:bindings")),
+ NS(), NS(), OS("R.dataflow"), NS(), OS("$method:print_outputs"))));
+ // SeqExpr → Seq with ret
+ register_trait(SeqExprNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::SeqObj>(S("$field:blocks"), OS("$field:body"))));
+}
+
+// --- Register computed-property methods for $method: trait references ---
+namespace {
+template <typename Func>
+void RegisterComputedMethod(int32_t type_index, const char* type_key, const char* method_name,
+ Func&& func) {
+ TVMFFIMethodInfo info;
+ info.name = TVMFFIByteArray{method_name, std::char_traits<char>::length(method_name)};
+ info.doc = TVMFFIByteArray{nullptr, 0};
+ info.flags = 0;
+ std::string full_name = std::string(type_key) + "." + method_name;
+ ::tvm::ffi::Function method =
+ ::tvm::ffi::Function::FromTyped(std::forward<Func>(func), std::move(full_name));
+ info.method = ::tvm::ffi::AnyView(method).CopyToTVMFFIAny();
+ static const char kEmptyMetadata[] = "{}";
+ info.metadata = TVMFFIByteArray{kEmptyMetadata, 2};
+ TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index, &info));
+}
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace text = ::tvm::ffi::ir::text;
+ // DataflowBlock.print_outputs(): returns a Function that appends R.output(...)
+ // for non-DataflowVar bindings as an epilogue statement.
+ RegisterComputedMethod(
+ DataflowBlockNode::RuntimeTypeIndex(), DataflowBlockNode::_type_key, "print_outputs",
+ [](DataflowBlock block) -> ::tvm::ffi::Function {
+ return ::tvm::ffi::Function::FromTyped(
+ [](::tvm::ffi::ObjectRef obj, text::IRPrinter printer, text::DefaultFrame frame) {
+ DataflowBlock blk = ::tvm::Downcast<DataflowBlock>(obj);
+ ffi::List<text::ExprAST> outputs;
+ for (const auto& b : blk->bindings) {
+ if (!b->var->IsInstance<DataflowVarNode>()) {
+ ffi::Optional<text::ExprAST> var_expr = printer->VarGet(b->var);
+ if (var_expr.has_value()) {
+ outputs.push_back(var_expr.value());
+ }
+ }
+ }
+ if (!outputs.empty()) {
+ text::ExprAST callee = text::ExprAttr(text::IdAST("R"), "output");
+ frame->stmts.push_back(
+ text::ExprStmtAST(text::ExprCall(callee, std::move(outputs))));
+ }
+ });
+ });
}
Id::Id(ffi::String name_hint) {
diff --git a/src/relax/ir/struct_info.cc b/src/relax/ir/struct_info.cc
index 434917bd0f..3fb6a18be3 100644
--- a/src/relax/ir/struct_info.cc
+++ b/src/relax/ir/struct_info.cc
@@ -22,6 +22,7 @@
* \brief Relax struct info.
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/struct_info.h>
@@ -38,6 +39,39 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TensorStructInfoNode::RegisterReflection();
TupleStructInfoNode::RegisterReflection();
FuncStructInfoNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for StructInfo types ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+ auto NS = []() { return ffi::Optional<ffi::String>(); };
+
+ // TensorStructInfo → TensorTy
+ register_trait(TensorStructInfoNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::TensorTyObj>(OS("$field:shape"), OS("$field:dtype"),
+ OS("$field:vdevice"))));
+ // PrimStructInfo → PrimTy
+ register_trait(PrimStructInfoNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::PrimTyObj>(S("$field:dtype"))));
+ // ShapeStructInfo → ShapeTy
+ register_trait(ShapeStructInfoNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::ShapeTyObj>(OS("$field:values"), OS("$field:ndim"))));
+ // TupleStructInfo → TupleTy
+ register_trait(TupleStructInfoNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::TupleTyObj>(S("$field:fields"))));
+ // FuncStructInfo → FuncTy
+ register_trait(FuncStructInfoNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::FuncTyObj>(OS("$field:params"), OS("$field:ret"))));
+ // ObjectStructInfo → Ty (base marker)
+ register_trait(ObjectStructInfoNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::TyObj>()));
}
ObjectStructInfo::ObjectStructInfo(Span span) {
diff --git a/src/relax/ir/type.cc b/src/relax/ir/type.cc
index faa0814f4c..5e522fdafb 100644
--- a/src/relax/ir/type.cc
+++ b/src/relax/ir/type.cc
@@ -22,6 +22,7 @@
* \brief Relax type system.
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/type.h>
@@ -33,6 +34,35 @@ TVM_FFI_STATIC_INIT_BLOCK() {
TensorTypeNode::RegisterReflection();
ObjectTypeNode::RegisterReflection();
PackedFuncTypeNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for Relax type nodes ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+
+ // ShapeType → ShapeTy
+ register_trait(ShapeTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::ShapeTyObj>(
+ ffi::Optional<ffi::String>(),
+ ffi::Optional<ffi::String>(ffi::String("$field:ndim")))));
+ // ObjectType → Ty (base marker)
+ register_trait(ObjectTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::TyObj>()));
+ // PackedFuncType → FuncTy
+ register_trait(PackedFuncTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::FuncTyObj>(
+ ffi::Optional<ffi::String>(), ffi::Optional<ffi::String>())));
+ // TensorType (DynTensorType) → TensorTy
+ register_trait(TensorTypeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::TensorTyObj>(
+ ffi::Optional<ffi::String>(),
+ ffi::Optional<ffi::String>(ffi::String("$field:dtype")),
+ ffi::Optional<ffi::String>())));
}
ShapeType::ShapeType(int ndim, Span span) {
diff --git a/src/script_v2/printer/ir/module.cc b/src/script_v2/printer/ir/module.cc
new file mode 100644
index 0000000000..573e86e79d
--- /dev/null
+++ b/src/script_v2/printer/ir/module.cc
@@ -0,0 +1,160 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/ir/module.cc
+ * \brief __ffi_text_print__ override for IRModule.
+ *
+ * Produces:
+ * # from tvm.script import ir as I
+ * # from tvm.script import tirx as T
+ * # from tvm.script import relax as R
+ * #
+ * @I.ir_module
+ * class Module:
+ * I.module_attrs(...)
+ * I.module_global_infos(...)
+ * @T.prim_func / @R.function
+ * def func_name(...): ...
+ */
+#include <tvm/ir/module.h>
+
+#include "../utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+namespace {
+
+struct SortableFunction {
+ int priority;
+ GlobalVar gv;
+ BaseFunc func;
+
+ explicit SortableFunction(const std::pair<GlobalVar, BaseFunc>& obj)
+ : priority(0), gv(obj.first), func(obj.second) {
+ if (gv->name_hint == "main") {
+ priority = 1000;
+ } else if (obj.second->GetTypeKey() == "tirx.PrimFunc") {
+ priority = 1;
+ } else if (obj.second->GetTypeKey() == "relax.expr.ExternFunc") {
+ priority = 2;
+ } else if (obj.second->GetTypeKey() == "relax.expr.Function") {
+ priority = 3;
+ } else {
+ priority = 0;
+ }
+ }
+
+ bool operator<(const SortableFunction& other) const {
+ if (this->priority != other.priority) {
+ return this->priority < other.priority;
+ }
+ return this->gv->name_hint < other.gv->name_hint;
+ }
+};
+
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ refl::TypeAttrDef<IRModuleNode>().def(
+ "__ffi_text_print__",
+ [](IRModule mod, IRPrinter printer, AccessPath path) -> NodeAST {
+ // Sort functions by priority (PrimFunc < ExternFunc < Function < main)
+ std::vector<SortableFunction> functions;
+ for (const auto& kv : mod->functions) {
+ functions.push_back(SortableFunction(kv));
+ }
+ std::sort(functions.begin(), functions.end());
+
+ IdAST module_doc = IdAST("Module");
+ List<ExprAST> decorators;
+ decorators.push_back(IR("ir_module"));
+
+ // Push frame for prologue statements
+ DefaultFrame frame;
+ printer->FramePush(frame);
+
+ // Define GlobalVars: bind each to Module.name_hint so inner functions
+ // can reference them as "Module.func_name(args...)" (matching V1 behavior).
+ for (const auto& entry : functions) {
+ const GlobalVar& gv = entry.gv;
+ String name = gv->name_hint;
+ ffi::Function creator = ffi::Function::FromTyped(
+ [name]() -> ExprAST { return ExprAttr(IdAST("Module"), std::string(name)); });
+ printer->VarDefNoName(creator, gv, Optional<ObjectRef>(frame));
+ }
+
+ // Print attrs prologue
+ if (mod->attrs.defined() && !mod->attrs->dict.empty()) {
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(IR("module_attrs"),
+ {Print(printer, mod->attrs, path->Attr("attrs"))})));
+ }
+
+ // Print global_infos prologue
+ if (mod->global_infos.defined() && !mod->global_infos.empty()) {
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(IR("module_global_infos"),
+ {Print(printer, mod->global_infos,
+ path->Attr("global_infos"))})));
+ }
+
+ // Print each function
+ List<StmtAST> body;
+ for (const auto& s : frame->stmts) body.push_back(s);
+
+ for (const auto& entry : functions) {
+ const GlobalVar& gv = entry.gv;
+ const BaseFunc& base_func = entry.func;
+ AccessPath func_path = path->Attr("functions")->MapItem(gv);
+ // Set the binding name so the function printer can retrieve it
+ printer->VarDef(gv->name_hint, base_func, frame);
+ NodeAST doc = printer->operator()(Any(base_func), func_path).cast<NodeAST>();
+ if (auto* block = doc.as<StmtBlockASTObj>()) {
+ body.push_back(block->stmts.back());
+ } else if (doc.get()->IsInstance<StmtASTObj>()) {
+ body.push_back(GetRef<StmtAST>(static_cast<const StmtASTObj*>(doc.get())));
+ } else if (doc.get()->IsInstance<ExprASTObj>()) {
+ ExprAST lhs = IdAST(gv->name_hint);
+ body.push_back(AssignAST(lhs,
+ GetRef<ExprAST>(static_cast<const ExprASTObj*>(doc.get())),
+ Optional<ExprAST>()));
+ }
+ }
+
+ printer->FramePop();
+
+ ClassAST class_ast(module_doc, {}, decorators, body);
+
+ // Add import header comments
+ List<StmtAST> result;
+ result.push_back(CommentAST(
+ Optional<String>(String("from tvm.script import ir as I"))));
+ result.push_back(CommentAST(
+ Optional<String>(String("from tvm.script import tirx as T"))));
+ result.push_back(CommentAST(
+ Optional<String>(String("from tvm.script import relax as R"))));
+ result.push_back(CommentAST(Optional<String>()));
+ result.push_back(class_ast);
+ return StmtBlockAST(result);
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/ir/primitives.cc b/src/script_v2/printer/ir/primitives.cc
new file mode 100644
index 0000000000..9cdb96225b
--- /dev/null
+++ b/src/script_v2/printer/ir/primitives.cc
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file script_v2/printer/ir/primitives.cc
+ * \brief __ffi_text_print__ overrides for base IR types.
+ *
+ * Nodes now handled purely by traits (no override needed):
+ * - IntImm, FloatImm → Literal trait (dtype-aware PrintLiteral)
+ * - PrimType → PrimTy trait (PrintPrimTy)
+ * - TupleType → TupleTy trait (PrintTupleTy)
+ * - FuncType → FuncTy trait (PrintFuncTy)
+ * - GlobalVar → Call("I.GlobalVar", "$field:name_hint")
+ * - Op → Call("I.Op", "$field:name")
+ * - DictAttrs → Literal("$field:dict")
+ * - PointerType → PointerTy trait (PrintPointerTy)
+ * - Range → Call("I.Range", "$method:range_args")
+ * - VDevice → Call("I.vdevice", "$method:vdevice_args")
+ */
+
+#include <tvm/ffi/container/map_base.h>
+#include <tvm/ir/global_info.h>
+#include <tvm/ir/type.h>
+#include <tvm/target/target.h>
+
+#include "../utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // Target → T.target({config_dict})
+ refl::TypeAttrDef<TargetNode>().def(
+ "__ffi_text_print__",
+ [](Target t, IRPrinter printer, AccessPath path) -> NodeAST {
+ Map<String, Any> config = t->ToConfig();
+ return ExprCall(TIR("target"), {Print(printer, config, path)});
+ });
+
+ // String → "str" literal (prevents ffi.String objects from falling to DefaultPrint)
+ refl::TypeAttrDef<ffi::details::StringObj>().def(
+ "__ffi_text_print__",
+ [](ffi::String s, IRPrinter printer, AccessPath path) -> NodeAST {
+ return LiteralAST::Str(s);
+ });
+
+ // Map → {key: val, ...} dict literal
+ // Uses raw MapBaseObj iterator to handle maps with non-object values
+ // (e.g., raw bools/ints in Target config maps)
+ refl::TypeAttrDef<ffi::MapObj>().def(
+ "__ffi_text_print__",
+ [](ObjectRef obj, IRPrinter printer, AccessPath path) -> NodeAST {
+ const ffi::MapBaseObj* map_obj = static_cast<const ffi::MapBaseObj*>(obj.get());
+ List<ExprAST> keys;
+ List<ExprAST> vals;
+ for (auto it = map_obj->begin(); it != map_obj->end(); ++it) {
+ auto kv = *it;
+ keys.push_back(printer->operator()(ffi::Any(kv.first), path->Attr("key"))
+ .cast<ExprAST>());
+ vals.push_back(printer->operator()(ffi::Any(kv.second), path->Attr("value"))
+ .cast<ExprAST>());
+ }
+ return DictAST(std::move(keys), std::move(vals));
+ });
+
+ // Array → [e1, e2, ...] list literal
+ // Uses ObjectRef arg + raw ArrayObj access to handle arrays with non-object elements
+ // (e.g., raw integers in annotation arrays like buffer_dim_align)
+ refl::TypeAttrDef<ffi::ArrayObj>().def(
+ "__ffi_text_print__",
+ [](ObjectRef obj, IRPrinter printer, AccessPath path) -> NodeAST {
+ const ffi::ArrayObj* arr = static_cast<const ffi::ArrayObj*>(obj.get());
+ int64_t n = arr->size();
+ List<ExprAST> elts;
+ for (int64_t i = 0; i < n; ++i) {
+ ffi::AnyView elem = arr->at(i);
+ ffi::Any elem_any(elem);
+ elts.push_back(printer->operator()(std::move(elem_any), path->ArrayItem(i))
+ .cast<ExprAST>());
+ }
+ return ListAST({}, std::move(elts));
+ });
+
+ // PointerType → T.handle("dtype", "scope")
+ // Override trait-based PointerTy to match V1 behavior:
+ // PrimType element_type → string dtype literal (e.g. "int32", "void")
+ // Non-PrimType element_type → fall through to printer dispatch
+ refl::TypeAttrDef<PointerTypeNode>().def(
+ "__ffi_text_print__",
+ [](PointerType ty, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST elem_expr(ffi::UnsafeInit{});
+ if (const auto* prim_type = ty->element_type.as<PrimTypeNode>()) {
+ std::string ds = DType2Str(prim_type->dtype);
+ elem_expr = LiteralAST::Str(ds);
+ } else {
+ elem_expr = Print(printer, ty->element_type, path->Attr("element_type"));
+ }
+ List<ExprAST> args;
+ args.push_back(elem_expr);
+ if (!ty->storage_scope.empty()) {
+ args.push_back(LiteralAST::Str(std::string(ty->storage_scope)));
+ }
+ return ExprCall(TIR("handle"), std::move(args));
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // VDevice → I.vdevice({config}, vdevice_id, "memory_scope")
+ refl::TypeAttrDef<VDeviceNode>().def(
+ "__ffi_text_print__",
+ [](VDevice vdev, IRPrinter printer, AccessPath path) -> NodeAST {
+ Map<String, Any> config = vdev->target->ToConfig();
+ ExprAST config_doc = Print(printer, config, path->Attr("target"));
+ ExprAST dev_id_doc = LiteralAST::Int(vdev->vdevice_id);
+ ExprAST mem_scope_doc = LiteralAST::Str(std::string(vdev->memory_scope));
+ return ExprCall(IR("vdevice"), {config_doc, dev_id_doc, mem_scope_doc});
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/relax/expr.cc b/src/script_v2/printer/relax/expr.cc
new file mode 100644
index 0000000000..7cdad009dd
--- /dev/null
+++ b/src/script_v2/printer/relax/expr.cc
@@ -0,0 +1,829 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/relax/expr.cc
+ * \brief __ffi_text_print__ overrides for Relax nodes.
+ *
+ * Nodes now handled purely by traits:
+ * - VarBinding → Assign trait (PrintAssign handles stmt-AST RHS)
+ * - BindingBlock → Seq trait (PrintSeq handles flattening)
+ * - SeqExpr → Seq trait with ret field
+ * - DataflowBlock → With trait (R.dataflow) + text_printer_post hook
+ */
+#include <tvm/ir/expr.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/ffi/reflection/accessor.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <cmath>
+#include <limits>
+#include <unordered_set>
+
+#include "../utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace relax;
+
+namespace {
+
+/*!
+ * \brief Determine the function name from context.
+ *
+ * Priority: (1) GlobalVar from AccessPath MapItem key (module context),
+ * (2) global_symbol attribute,
+ * (3) fallback "main".
+ */
+String FindFuncName(const relax::Function& func, const IRPrinter& printer,
+ const AccessPath& path) {
+ // Priority 1: binding name from module context (set via VarDefNoName in IRModule prologue)
+ if (auto binding_expr = printer->VarGet(func)) {
+ if (auto* id_node = binding_expr.value().as<text::IdASTObj>()) {
+ return id_node->name;
+ }
+ }
+ // Priority 2: GlobalVar from AccessPath MapItem key
+ if (path->step.defined()) {
+ const auto& step = path->step.value();
+ if (step->kind == ffi::reflection::AccessKind::kMapItem) {
+ if (const auto* gv = step->key.as<GlobalVarNode>()) {
+ return gv->name_hint;
+ }
+ }
+ }
+ // Priority 3: global_symbol attribute
+ if (func->attrs.defined()) {
+ auto it = func->attrs->dict.find("global_symbol");
+ if (it != func->attrs->dict.end()) {
+ return (*it).second.cast<String>();
+ }
+ }
+ return "main";
+}
+
+/*!
+ * \brief Check if this function is at top level in a module.
+ *
+ * A function is at the top level if it is reached via
+ * Root->Attr("functions")->MapItem(gv).
+ */
+bool AtTopLevelInModule(const AccessPath& path) {
+ // path should be MapItem(gv)
+ if (!path->step.defined()) return false;
+ if (path->step.value()->kind != ffi::reflection::AccessKind::kMapItem) return false;
+ // parent should be Attr("functions")
+ auto parent_opt = path->GetParent();
+ if (!parent_opt.has_value()) return false;
+ AccessPath parent = parent_opt.value();
+ if (!parent->step.defined()) return false;
+ if (parent->step.value()->kind != ffi::reflection::AccessKind::kAttr) return false;
+ return true;
+}
+
+/*!
+ * \brief Collect all tirx::Var references from a PrimExpr (e.g. shape dims).
+ */
+void CollectTIRVarsFromPrimExpr(const PrimExpr& expr,
+ std::vector<tirx::Var>* out,
+ std::unordered_set<const Object*>* seen) {
+ tirx::PostOrderVisit(expr, [&](const ObjectRef& obj) {
+ if (const auto* tv = obj.as<tirx::VarNode>()) {
+ if (seen->insert(tv).second) {
+ out->push_back(GetRef<tirx::Var>(tv));
+ }
+ }
+ });
+}
+
+/*!
+ * \brief Collect all tirx::Var references from a relax param's struct_info.
+ *
+ * Walks TensorStructInfo shapes, ShapeStructInfo values, and PrimStructInfo
+ * values to find symbolic dimension variables.
+ */
+void CollectTIRVarsFromStructInfo(const StructInfo& sinfo,
+ std::vector<tirx::Var>* out,
+ std::unordered_set<const Object*>* seen) {
+ if (const auto* tsi = sinfo.as<TensorStructInfoNode>()) {
+ if (tsi->shape.defined()) {
+ if (const auto* se = tsi->shape.value().as<ShapeExprNode>()) {
+ for (const auto& dim : se->values) {
+ CollectTIRVarsFromPrimExpr(dim, out, seen);
+ }
+ }
+ }
+ } else if (const auto* ssi = sinfo.as<ShapeStructInfoNode>()) {
+ if (ssi->values.defined()) {
+ for (const auto& val : ssi->values.value()) {
+ CollectTIRVarsFromPrimExpr(val, out, seen);
+ }
+ }
+ } else if (const auto* psi = sinfo.as<PrimStructInfoNode>()) {
+ if (psi->value.defined()) {
+ CollectTIRVarsFromPrimExpr(psi->value.value(), out, seen);
+ }
+ } else if (const auto* tsi = sinfo.as<TupleStructInfoNode>()) {
+ for (const auto& field : tsi->fields) {
+ CollectTIRVarsFromStructInfo(field, out, seen);
+ }
+ } else if (const auto* fsi = sinfo.as<FuncStructInfoNode>()) {
+ if (fsi->params.defined()) {
+ for (const auto& param : fsi->params.value()) {
+ CollectTIRVarsFromStructInfo(param, out, seen);
+ }
+ }
+ CollectTIRVarsFromStructInfo(fsi->ret, out, seen);
+ }
+}
+
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // DataflowBlock: with R.dataflow(): bindings + R.output(...)
+ // Must use __ffi_text_print__ (not trait) to preserve output vars in parent scope.
+ // The With trait's FramePop removes all vars, causing the SeqExpr body return to fail.
+ refl::TypeAttrDef<DataflowBlockNode>().def(
+ "__ffi_text_print__",
+ [](DataflowBlock node, IRPrinter printer, AccessPath path) -> NodeAST {
+ DefaultFrame frame;
+ printer->FramePush(frame);
+ List<StmtAST> body;
+ for (int i = 0; i < static_cast<int>(node->bindings.size()); ++i) {
+ NodeAST s = printer->operator()(Any(node->bindings[i]),
+ path->Attr("bindings")->ArrayItem(i))
+ .cast<NodeAST>();
+ if (auto* block = s.as<StmtBlockASTObj>()) {
+ for (const auto& st : block->stmts) body.push_back(st);
+ } else if (s.get()->IsInstance<StmtASTObj>()) {
+ body.push_back(GetRef<StmtAST>(static_cast<const StmtASTObj*>(s.get())));
+ } else if (s.get()->IsInstance<ExprASTObj>()) {
+ body.push_back(ExprStmtAST(GetRef<ExprAST>(static_cast<const ExprASTObj*>(s.get()))));
+ }
+ }
+ // R.output() for non-DataflowVar bindings
+ List<ExprAST> outputs;
+ for (const auto& b : node->bindings) {
+ if (!b->var->IsInstance<DataflowVarNode>()) {
+ if (auto var_doc = printer->VarGet(b->var)) {
+ outputs.push_back(var_doc.value());
+ }
+ }
+ }
+ if (!outputs.empty()) {
+ body.push_back(ExprStmtAST(ExprCall(Relax("output"), std::move(outputs))));
+ }
+ // Pop frame BUT re-register output vars in the PARENT frame so they're
+ // accessible to the enclosing SeqExpr's return statement.
+ List<ffi::Any> output_vars;
+ for (const auto& b : node->bindings) {
+ if (!b->var->IsInstance<DataflowVarNode>()) {
+ output_vars.push_back(b->var);
+ }
+ }
+ // Save the var info before popping
+ ffi::Map<ffi::Any, text::VarInfo> saved_info;
+ for (const auto& ov : output_vars) {
+ auto it = printer->obj2info.find(ov);
+ if (it != printer->obj2info.end()) {
+ saved_info.Set(ov, (*it).second);
+ }
+ }
+ printer->FramePop();
+ // Re-register output vars in parent frame
+ if (!printer->frames.empty()) {
+ ObjectRef parent_frame = printer->frames.back().cast<ObjectRef>();
+ for (const auto& kv : saved_info) {
+ printer->obj2info.Set(kv.first, kv.second);
+ // Add to parent frame's var list
+ auto frame_vars = printer->frame_vars[parent_frame].cast<List<ffi::Any>>();
+ frame_vars.push_back(kv.first);
+ }
+ }
+
+ ExprAST ctx = ExprCall(Relax("dataflow"), {});
+ return WithAST(Optional<ExprAST>(), ctx, body);
+ });
+
+ // relax.Function: @R.function with struct_info annotations and prologue
+ refl::TypeAttrDef<FunctionNode>().def(
+ "__ffi_text_print__",
+ [](relax::Function func, IRPrinter printer, AccessPath path) -> NodeAST {
+ // Step 1. Determine function name
+ String func_name = FindFuncName(func, printer, path);
+ IdAST name = IdAST(func_name);
+ bool at_top_level = AtTopLevelInModule(path);
+ bool has_global_symbol = func->attrs.defined() &&
+ func->attrs->dict.count("global_symbol");
+
+ // Step 2. Build decorator: @R.function with optional kwargs
+ List<String> dec_keys;
+ List<ExprAST> dec_values;
+ // pure=False when function is impure
+ if (!func->is_pure) {
+ dec_keys.push_back(String("pure"));
+ dec_values.push_back(LiteralAST::Bool(false));
+ }
+ // private=True when at top level (or standalone) without global_symbol
+ // V1 marks standalone functions and top-level module functions as private
+ // when they lack a global_symbol attribute.
+ bool is_standalone = !path->step.defined() || path->depth <= 1;
+ if ((at_top_level || is_standalone) && !has_global_symbol) {
+ dec_keys.push_back(String("private"));
+ dec_values.push_back(LiteralAST::Bool(true));
+ }
+ List<ExprAST> decorators;
+ if (dec_keys.size() > 0) {
+ decorators.push_back(
+ ExprCallKw(Relax("function"), {}, std::move(dec_keys), std::move(dec_values)));
+ } else {
+ decorators.push_back(Relax("function"));
+ }
+
+ // Step 3. Push frame, define symbolic TIR vars, then print params.
+ DefaultFrame frame;
+ printer->FramePush(frame);
+ int n = func->params.size();
+
+ // Step 3a. Collect and define symbolic TIR vars from param struct_info
+ // BEFORE printing params. This ensures that when struct_info shape
+ // annotations (e.g. `R.Tensor((n, m), dtype="float32")`) are printed,
+ // the tirx::Var `n` and `m` are already defined and VarGet returns
+ // the correct IdAST instead of auto-defining a free var named "v".
+ // The `var = T.int64()` prologue statements are emitted into the frame.
+ {
+ std::vector<tirx::Var> tir_vars;
+ std::unordered_set<const Object*> seen;
+ for (int i = 0; i < n; ++i) {
+ Var var = func->params[i];
+ if (var->struct_info_.defined()) {
+ StructInfo si = Downcast<StructInfo>(var->struct_info_.value());
+ CollectTIRVarsFromStructInfo(si, &tir_vars, &seen);
+ }
+ }
+ for (const auto& tir_var : tir_vars) {
+ if (!printer->VarGet(tir_var).has_value()) {
+ printer->VarDef(tir_var->name_hint, tir_var, frame);
+ ExprAST var_id = printer->VarGet(tir_var).value();
+ std::string dtype_str = DType2Str(tir_var->dtype);
+ frame->stmts.push_back(
+ AssignAST(var_id, ExprCall(TIR(dtype_str), {}), Optional<ExprAST>()));
+ }
+ }
+ }
+
+ // Step 3b. Print params (tirx::Vars in struct_info shapes are now defined)
+ List<AssignAST> params;
+ for (int i = 0; i < n; ++i) {
+ Var var = func->params[i];
+ AccessPath var_p = path->Attr("params")->ArrayItem(i);
+ printer->VarDef(var->vid->name_hint, var, frame);
+ ExprAST var_id = printer->VarGet(var).value();
+ Optional<ExprAST> annotation;
+ if (var->struct_info_.defined()) {
+ annotation = Print(printer, var->struct_info_.value(),
+ var_p->Attr("struct_info_"));
+ }
+ params.push_back(AssignAST(var_id, Optional<ExprAST>(), annotation));
+ }
+
+ // Step 4. Print attributes (filter global_symbol when it matches func name)
+ // V1 filters global_symbol for top-level functions (both in-module and standalone)
+ // when the symbol matches the function name being used.
+ bool should_filter_global_symbol = has_global_symbol &&
+ func->attrs->dict.at("global_symbol").cast<ffi::String>() == func_name &&
+ (at_top_level || is_standalone);
+ if (func->attrs.defined() && !func->attrs->dict.empty()) {
+ if (should_filter_global_symbol) {
+ // global_symbol matches func name: filter it out
+ ffi::Map<ffi::String, ffi::Any> filtered;
+ for (const auto& kv : func->attrs->dict) {
+ if (kv.first != "global_symbol") {
+ filtered.Set(kv.first, kv.second);
+ }
+ }
+ if (!filtered.empty()) {
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(Relax("func_attr"),
+ {Print(printer, DictAttrs(filtered),
+ path->Attr("attrs"))})));
+ }
+ } else {
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(Relax("func_attr"),
+ {Print(printer, func->attrs, path->Attr("attrs"))})));
+ }
+ }
+
+ // Step 5. Print body: inline SeqExpr handling.
+ // We must NOT delegate to the SeqExpr Seq trait because the trait's
+ // PrintBody/PrintSeq pushes its own frame for the blocks. When a
+ // DataflowBlock's With trait FramePop removes vars defined inside
+ // the block, the subsequent `ret` resolution can no longer find the
+ // body Var, producing "Undefined variable: v".
+ // Instead, we print each block inline and resolve the return var
+ // while DataflowBlock vars are still in scope.
+ List<StmtAST> body;
+ relax::SeqExpr seq = Downcast<relax::SeqExpr>(func->body);
+ for (int i = 0; i < static_cast<int>(seq->blocks.size()); ++i) {
+ NodeAST block_ast = printer->operator()(
+ Any(seq->blocks[i]),
+ path->Attr("body")->Attr("blocks")->ArrayItem(i)).cast<NodeAST>();
+ if (auto* sb = block_ast.as<StmtBlockASTObj>()) {
+ for (const auto& s : sb->stmts) body.push_back(s);
+ } else if (block_ast.get()->IsInstance<StmtASTObj>()) {
+ body.push_back(
+ GetRef<StmtAST>(static_cast<const StmtASTObj*>(block_ast.get())));
+ }
+ }
+ // Resolve the return var BEFORE FramePop, while all block vars
+ // (including DataflowBlock vars) are still accessible.
+ ExprAST ret_expr = Print(printer, seq->body, path->Attr("body")->Attr("body"));
+ body.push_back(ReturnAST(ret_expr));
+
+ List<StmtAST> all_body;
+ for (const auto& s : frame->stmts) all_body.push_back(s);
+ for (const auto& s : body) all_body.push_back(s);
+ printer->FramePop();
+
+ // Step 6. Print return type from FuncStructInfo->ret (matching V1)
+ Optional<ExprAST> ret_type;
+ if (const auto* func_sinfo = func->struct_info_.as<FuncStructInfoNode>()) {
+ ret_type = Print(printer, func_sinfo->ret, path->Attr("struct_info_")->Attr("ret"));
+ }
+ return FunctionAST(name, params, decorators, ret_type, all_body);
+ });
+}
+
+// ---- Shape value printing helpers ----
+
+namespace {
+
+/*!
+ * \brief Print a PrimExpr for use in struct_info shape contexts.
+ *
+ * In V1, the "relax" dispatch for IntImm/FloatImm prints them as plain
+ * integer/float literals (not T.int64(10)). This matches that behavior.
+ */
+ExprAST PrintShapeValue(const PrimExpr& e, const AccessPath& e_p, const IRPrinter& printer,
+ bool stringify_vars = false) {
+ if (const auto* int_imm = e.as<IntImmNode>()) {
+ if (int_imm->dtype.is_bool()) {
+ return LiteralAST::Bool(int_imm->value != 0);
+ }
+ return LiteralAST::Int(int_imm->value);
+ }
+ if (const auto* float_imm = e.as<FloatImmNode>()) {
+ return LiteralAST::Float(float_imm->value);
+ }
+ // For PrimExpr containing symbolic TIR Vars, stringify them (matching V1 PrintShapeVar).
+ // Only do this in param/return annotation contexts (stringify_vars=true).
+ if (stringify_vars) {
+ bool has_tir_var = false;
+ tirx::PostOrderVisit(e, [&](const ObjectRef& obj) {
+ if (obj->IsInstance<tirx::VarNode>()) has_tir_var = true;
+ });
+ if (has_tir_var) {
+ // Simple Var: just use name_hint as string
+ if (const auto* v = e.as<tirx::VarNode>()) {
+ return LiteralAST::Str(v->name_hint);
+ }
+ // Compound expressions (n * 2, etc.): build string from parts
+ // Use a simple recursive stringifier
+ std::function<std::string(const PrimExpr&)> stringify;
+ stringify = [&](const PrimExpr& expr) -> std::string {
+ if (const auto* v = expr.as<tirx::VarNode>()) return std::string(v->name_hint);
+ if (const auto* imm = expr.as<IntImmNode>()) return std::to_string(imm->value);
+ if (const auto* add = expr.as<tirx::AddNode>()) return stringify(add->a) + " + " + stringify(add->b);
+ if (const auto* sub = expr.as<tirx::SubNode>()) return stringify(sub->a) + " - " + stringify(sub->b);
+ if (const auto* mul = expr.as<tirx::MulNode>()) return stringify(mul->a) + " * " + stringify(mul->b);
+ if (const auto* div = expr.as<tirx::FloorDivNode>()) return stringify(div->a) + " // " + stringify(div->b);
+ if (const auto* mod = expr.as<tirx::FloorModNode>()) return stringify(mod->a) + " % " + stringify(mod->b);
+ return "?";
+ };
+ return LiteralAST::Str(stringify(e));
+ }
+ }
+ // For other PrimExpr types without TIR Vars, use the general printer
+ return Print(printer, e, e_p);
+}
+
+/*!
+ * \brief Extract a scalar value from a 0-d CPU tensor as an ExprAST literal.
+ *
+ * Returns std::nullopt for non-scalar or non-CPU tensors, or unsupported dtypes.
+ * Matches the V1 SpecialScalar logic.
+ */
+ffi::Optional<ExprAST> SpecialScalar(const runtime::Tensor& tensor, const AccessPath& p) {
+ DataType dtype(tensor->dtype);
+ const void* data = tensor->data;
+ if (tensor->ndim != 0 || tensor->device.device_type != kDLCPU) {
+ return std::nullopt;
+ }
+ if (dtype == DataType::Int(8)) {
+ return LiteralAST::Int(*reinterpret_cast<const int8_t*>(data));
+ } else if (dtype == DataType::Int(16)) {
+ return LiteralAST::Int(*reinterpret_cast<const int16_t*>(data));
+ } else if (dtype == DataType::Int(32)) {
+ return LiteralAST::Int(*reinterpret_cast<const int32_t*>(data));
+ } else if (dtype == DataType::Int(64)) {
+ return LiteralAST::Int(*reinterpret_cast<const int64_t*>(data));
+ } else if (dtype == DataType::Float(16)) {
+ uint16_t bits = *reinterpret_cast<const uint16_t*>(data);
+ uint16_t sign_bit = (bits & 0x8000) >> 15;
+ uint16_t exponent = (bits & 0x7C00) >> 10;
+ uint16_t fraction = (bits & 0x03FF) >> 0;
+ double value;
+ if (exponent == 0x1F && fraction == 0) {
+ value = std::numeric_limits<double>::infinity();
+ } else if (exponent == 0x1F) {
+ value = std::numeric_limits<double>::quiet_NaN();
+ } else if (exponent == 0 && fraction == 0) {
+ value = 0.0;
+ } else if (exponent == 0) {
+ value = ::std::pow(2.0, -24) * static_cast<double>(fraction);
+ } else {
+ value = ::std::pow(2.0, static_cast<double>(exponent) - 25) *
+ static_cast<double>(fraction | (1 << 10));
+ }
+ if (sign_bit) {
+ value *= -1.0;
+ }
+ return LiteralAST::Float(value);
+ } else if (dtype == DataType::Float(32)) {
+ return LiteralAST::Float(*reinterpret_cast<const float*>(data));
+ } else if (dtype == DataType::Float(64)) {
+ return LiteralAST::Float(*reinterpret_cast<const double*>(data));
+ } else if (dtype == DataType::Bool()) {
+ return LiteralAST::Bool(*reinterpret_cast<const uint8_t*>(data) != 0);
+ } else {
+ return std::nullopt;
+ }
+}
+
+} // namespace
+
+// ---- Relax leaf-expression printers ----
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // PrimValue → R.prim_value(value)
+ refl::TypeAttrDef<PrimValueNode>().def(
+ "__ffi_text_print__",
+ [](PrimValue node, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST val = PrintShapeValue(node->value, path->Attr("value"), printer, false);
+ return ExprCall(Relax("prim_value"), {val});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // StringImm → R.str("value")
+ refl::TypeAttrDef<relax::StringImmNode>().def(
+ "__ffi_text_print__",
+ [](relax::StringImm node, IRPrinter printer, AccessPath path) -> NodeAST {
+ return ExprCall(Relax("str"), {LiteralAST::Str(node->value)});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // DataTypeImm → R.dtype("dtype_string")
+ refl::TypeAttrDef<relax::DataTypeImmNode>().def(
+ "__ffi_text_print__",
+ [](relax::DataTypeImm node, IRPrinter printer, AccessPath path) -> NodeAST {
+ std::string ds = DType2Str(node->value);
+ return ExprCall(Relax("dtype"), {LiteralAST::Str(ds)});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // Tuple → (a, b, ...) or R.tuple() for empty
+ refl::TypeAttrDef<relax::TupleNode>().def(
+ "__ffi_text_print__",
+ [](relax::Tuple node, IRPrinter printer, AccessPath path) -> NodeAST {
+ if (node->fields.empty()) {
+ return ExprCall(Relax("tuple"), {});
+ }
+ List<ExprAST> elts;
+ for (int i = 0; i < static_cast<int>(node->fields.size()); ++i) {
+ elts.push_back(Print(printer, node->fields[i], path->Attr("fields")->ArrayItem(i)));
+ }
+ return TupleAST({}, std::move(elts));
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // TupleGetItem → tuple[index]
+ refl::TypeAttrDef<relax::TupleGetItemNode>().def(
+ "__ffi_text_print__",
+ [](relax::TupleGetItem node, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST tuple_doc = Print(printer, node->tuple, path->Attr("tuple_value"));
+ return ExprIndex(tuple_doc, {LiteralAST::Int(node->index)});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // ShapeExpr → R.shape([dims...])
+ refl::TypeAttrDef<relax::ShapeExprNode>().def(
+ "__ffi_text_print__",
+ [](relax::ShapeExpr node, IRPrinter printer, AccessPath path) -> NodeAST {
+ List<ExprAST> dims;
+ for (int i = 0; i < static_cast<int>(node->values.size()); ++i) {
+ dims.push_back(PrintShapeValue(node->values[i], path->Attr("values")->ArrayItem(i),
+ printer));
+ }
+ return ExprCall(Relax("shape"), {ListAST({}, std::move(dims))});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // Constant → R.const(value, dtype) for scalars, metadata[...] placeholder for tensors
+ refl::TypeAttrDef<relax::ConstantNode>().def(
+ "__ffi_text_print__",
+ [](relax::Constant node, IRPrinter printer, AccessPath path) -> NodeAST {
+ if (ffi::Optional<ExprAST> s = SpecialScalar(node->data, path->Attr("data"))) {
+ return ExprCall(Relax("const"),
+ {s.value(), LiteralAST::Str(DType2Str(DataType(node->data->dtype)))});
+ }
+ // Non-scalar: emit R.const(0, dtype) as a lossy placeholder.
+ // V2 does not have a metadata registry like V1, so we cannot
+ // faithfully round-trip non-scalar constants yet.
+ return ExprCall(Relax("const"),
+ {LiteralAST::Int(0),
+ LiteralAST::Str(DType2Str(DataType(node->data->dtype)))});
+ });
+}
+
+// ---- StructInfo printers matching V1 output ----
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // ObjectStructInfo → R.Object
+ refl::TypeAttrDef<ObjectStructInfoNode>().def(
+ "__ffi_text_print__",
+ [](relax::ObjectStructInfo, IRPrinter, AccessPath) -> NodeAST {
+ return Relax("Object");
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // PrimStructInfo → R.Prim(dtype) or R.Prim(value=...)
+ refl::TypeAttrDef<PrimStructInfoNode>().def(
+ "__ffi_text_print__",
+ [](relax::PrimStructInfo n, IRPrinter printer, AccessPath path) -> NodeAST {
+ List<ExprAST> args;
+ List<String> kwargs_keys;
+ List<ExprAST> kwargs_values;
+ if (n->value.defined()) {
+ kwargs_keys.push_back(String("value"));
+ kwargs_values.push_back(PrintShapeValue(n->value.value(), path->Attr("value"), printer, true));
+ } else {
+ args.push_back(LiteralAST::Str(DType2Str(n->dtype)));
+ }
+ return ExprCallKw(Relax("Prim"), std::move(args),
+ std::move(kwargs_keys), std::move(kwargs_values));
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // ShapeStructInfo → R.Shape([dims]) or R.Shape(ndim=N)
+ refl::TypeAttrDef<ShapeStructInfoNode>().def(
+ "__ffi_text_print__",
+ [](relax::ShapeStructInfo n, IRPrinter printer, AccessPath path) -> NodeAST {
+ if (n->values.defined()) {
+ ffi::Array<PrimExpr> shape = n->values.value();
+ AccessPath shape_p = path->Attr("values");
+ List<ExprAST> shape_docs;
+ for (int i = 0, ndim = shape.size(); i < ndim; ++i) {
+ shape_docs.push_back(PrintShapeValue(shape[i], shape_p->ArrayItem(i), printer, true));
+ }
+ return ExprCall(Relax("Shape"), {ListAST({}, std::move(shape_docs))});
+ }
+ return ExprCallKw(Relax("Shape"), {},
+ {String("ndim")}, {LiteralAST::Int(n->ndim)});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // TensorStructInfo → R.Tensor((shape,), dtype=...) etc.
+ refl::TypeAttrDef<TensorStructInfoNode>().def(
+ "__ffi_text_print__",
+ [](relax::TensorStructInfo n, IRPrinter printer, AccessPath path) -> NodeAST {
+ List<ExprAST> args;
+ List<String> kwargs_keys;
+ List<ExprAST> kwargs_values;
+ if (n->shape.defined()) {
+ // Dig into ShapeExpr to get individual dims
+ if (const auto* shape = n->shape.value().as<relax::ShapeExprNode>()) {
+ auto shape_expr = ffi::GetRef<relax::ShapeExpr>(shape);
+ AccessPath shape_p = path->Attr("shape")->Attr("values");
+ List<ExprAST> shape_docs;
+ for (int i = 0, ndim = shape_expr->values.size(); i < ndim; ++i) {
+ shape_docs.push_back(PrintShapeValue(shape_expr->values[i],
+ shape_p->ArrayItem(i), printer,
+ /*stringify_vars=*/true));
+ }
+ args.push_back(TupleAST({}, std::move(shape_docs)));
+ } else {
+ args.push_back(Print(printer, n->shape.value(), path->Attr("shape")));
+ }
+ }
+ if (!n->IsUnknownDtype()) {
+ kwargs_keys.push_back(String("dtype"));
+ kwargs_values.push_back(LiteralAST::Str(DType2Str(n->dtype)));
+ }
+ if (!n->shape.defined() && !n->IsUnknownNdim()) {
+ kwargs_keys.push_back(String("ndim"));
+ kwargs_values.push_back(LiteralAST::Int(n->ndim));
+ }
+ if (args.empty() && kwargs_keys.empty()) {
+ return Relax("Tensor");
+ }
+ return ExprCallKw(Relax("Tensor"), std::move(args),
+ std::move(kwargs_keys), std::move(kwargs_values));
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // TupleStructInfo → R.Tuple(fields...)
+ refl::TypeAttrDef<TupleStructInfoNode>().def(
+ "__ffi_text_print__",
+ [](relax::TupleStructInfo n, IRPrinter printer, AccessPath path) -> NodeAST {
+ if (n->fields.empty()) {
+ return Relax("Tuple");
+ }
+ List<ExprAST> fields_doc;
+ AccessPath fields_p = path->Attr("fields");
+ for (int i = 0, l = n->fields.size(); i < l; ++i) {
+ fields_doc.push_back(Print(printer, n->fields[i], fields_p->ArrayItem(i)));
+ }
+ return ExprCall(Relax("Tuple"), std::move(fields_doc));
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // FuncStructInfo → R.Callable(...) matching V1 format
+ refl::TypeAttrDef<FuncStructInfoNode>().def(
+ "__ffi_text_print__",
+ [](relax::FuncStructInfo n, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST ret_doc = Print(printer, n->ret, path->Attr("ret"));
+ ExprAST purity_doc = LiteralAST::Bool(n->purity);
+
+ if (n->IsOpaque()) {
+ List<String> keys;
+ List<ExprAST> values;
+ if (!n->ret->IsInstance<relax::ObjectStructInfoNode>()) {
+ keys.push_back(String("ret"));
+ values.push_back(ret_doc);
+ }
+ if (n->purity) {
+ keys.push_back(String("purity"));
+ values.push_back(purity_doc);
+ }
+ if (keys.size() > 0) {
+ return ExprCallKw(Relax("Callable"), {}, std::move(keys), std::move(values));
+ } else {
+ return Relax("Callable");
+ }
+ }
+ // Non-opaque: R.Callable((params...), ret, purity)
+ List<ExprAST> params_doc;
+ ffi::Array<relax::StructInfo> params = n->params.value();
+ AccessPath params_p = path->Attr("params");
+ for (int i = 0, n_params = params.size(); i < n_params; ++i) {
+ params_doc.push_back(Print(printer, params[i], params_p->ArrayItem(i)));
+ }
+ return ExprCall(Relax("Callable"),
+ {TupleAST({}, std::move(params_doc)), ret_doc, purity_doc});
+ });
+}
+
+// ---- relax.Call: __ffi_text_print__ override ----
+// Maps Op-based calls to R.name(args...) syntax matching V1.
+// Handles ExternFunc, Op, Var, and GlobalVar ops.
+TVM_FFI_STATIC_INIT_BLOCK() {
+ refl::TypeAttrDef<relax::CallNode>().def(
+ "__ffi_text_print__",
+ [](relax::Call call, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST prefix(ffi::UnsafeInit{});
+ List<ExprAST> args;
+ List<String> kw_keys;
+ List<ExprAST> kw_vals;
+
+ // Step 1. Determine callee prefix
+ if (const auto* ef = call->op.as<relax::ExternFuncNode>()) {
+ prefix = Relax("call_packed");
+ args.push_back(LiteralAST::Str(ef->global_symbol));
+ } else if (const auto* op_node = call->op.as<OpNode>()) {
+ std::string name = op_node->name;
+ if (name.rfind("relax.", 0) == 0) {
+ prefix = Relax(name.substr(6));
+ } else {
+ prefix = IdAST(name);
+ }
+ } else if (call->op->IsInstance<relax::VarNode>()) {
+ prefix = Print(printer, call->op, path->Attr("op"));
+ } else if (call->op->IsInstance<GlobalVarNode>()) {
+ // Check if the GlobalVar has been registered (e.g. from module prologue)
+ if (auto bound = printer->VarGet(call->op)) {
+ prefix = bound.value();
+ } else {
+ prefix = Print(printer, call->op, path->Attr("op"));
+ }
+ } else {
+ prefix = Print(printer, call->op, path->Attr("op"));
+ }
+
+ // Step 2. Print args
+ for (int i = 0, n = call->args.size(); i < n; ++i) {
+ args.push_back(Print(printer, call->args[i], path->Attr("args")->ArrayItem(i)));
+ }
+
+ // Step 3. Print attrs as kwargs
+ if (call->attrs.defined()) {
+ if (call->op->IsInstance<relax::ExternFuncNode>()) {
+ kw_keys.push_back(String("attrs_type_key"));
+ kw_vals.push_back(LiteralAST::Str(call->attrs->GetTypeKey()));
+ }
+ if (const auto* dict_attrs = call->attrs.as<DictAttrsNode>()) {
+ // Sort attrs by key for deterministic output
+ std::vector<std::pair<ffi::String, ffi::Any>> sorted;
+ for (const auto& kv : dict_attrs->dict) {
+ sorted.push_back(kv);
+ }
+ std::sort(sorted.begin(), sorted.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ for (const auto& kv : sorted) {
+ kw_keys.push_back(kv.first);
+ kw_vals.push_back(
+ Print(printer, kv.second, path->Attr("attrs")->Attr(kv.first)));
+ }
+ } else if (call->attrs.defined()) {
+ // Non-DictAttrs: use reflection to iterate fields
+ const TVMFFITypeInfo* info = TVMFFIGetTypeInfo(call->attrs->type_index());
+ ffi::reflection::ForEachFieldInfo(info, [&](const TVMFFIFieldInfo* fi) {
+ ffi::String fname(fi->name.data, fi->name.size);
+ if (fname == "span") return; // skip span field
+ Any field_val = ffi::reflection::FieldGetter(fi)(call->attrs);
+ kw_keys.push_back(fname);
+ kw_vals.push_back(
+ Print(printer, std::move(field_val), path->Attr("attrs")->Attr(fname)));
+ });
+ }
+ }
+
+ // Step 4. Print sinfo_args / out_sinfo
+ if (call->sinfo_args.size() > 0) {
+ // For call_tir and related ops, use "out_sinfo" and unwrap single-element tuples.
+ bool use_out_sinfo = false;
+ if (const auto* op_node = call->op.as<OpNode>()) {
+ std::string name = op_node->name;
+ use_out_sinfo = (name == "relax.call_tir" ||
+ name == "relax.call_tir_inplace" ||
+ name == "relax.call_dps_packed" ||
+ name == "relax.call_tir_with_grad" ||
+ name == "relax.dist.call_tir_local_view");
+ }
+ AccessPath sinfo_p = path->Attr("sinfo_args");
+ List<ExprAST> sinfo_docs;
+ for (int i = 0, n = call->sinfo_args.size(); i < n; ++i) {
+ sinfo_docs.push_back(Print(printer, call->sinfo_args[i], sinfo_p->ArrayItem(i)));
+ }
+ if (use_out_sinfo) {
+ kw_keys.push_back(String("out_sinfo"));
+ // Unwrap single-element tuple
+ if (sinfo_docs.size() == 1) {
+ kw_vals.push_back(sinfo_docs[0]);
+ } else {
+ kw_vals.push_back(TupleAST({}, std::move(sinfo_docs)));
+ }
+ } else {
+ kw_keys.push_back(String("sinfo_args"));
+ kw_vals.push_back(TupleAST({}, std::move(sinfo_docs)));
+ }
+ }
+
+ if (!kw_keys.empty()) {
+ return CallAST(prefix, std::move(args), std::move(kw_keys), std::move(kw_vals));
+ }
+ return CallAST(prefix, std::move(args), {}, {});
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/block.cc b/src/script_v2/printer/tirx/block.cc
new file mode 100644
index 0000000000..a3c3342b07
--- /dev/null
+++ b/src/script_v2/printer/tirx/block.cc
@@ -0,0 +1,194 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/tirx/block.cc
+ * \brief V2 text printers for SBlock and SBlockRealize.
+ */
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt.h>
+
+#include "utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+
+static NodeAST PrintSBlockRealize(SBlockRealize realize, IRPrinter printer, AccessPath path) {
+ SBlock block = realize->block;
+ AccessPath block_p = path->Attr("block");
+
+ DefaultFrame frame;
+ printer->FramePush(frame);
+
+ // Build context expr: T.sblock("name")
+ ExprAST ctx = ExprCall(TIR("sblock"), {LiteralAST::Str(block->name_hint)});
+
+ // Define iter_vars and build as_var
+ // Each iter_var: vi = T.axis.spatial(dom) or vi = T.axis.reduce(dom)
+ for (int i = 0; i < static_cast<int>(block->iter_vars.size()); ++i) {
+ IterVar iv = block->iter_vars[i];
+ AccessPath iv_p = block_p->Attr("iter_vars")->ArrayItem(i);
+
+ // Define the variable
+ ExprAST var_id = DefineVar(iv->var, printer, iv_p->Attr("var"));
+
+ // Build the axis call
+ std::string axis_type;
+ switch (static_cast<int>(iv->iter_type)) {
+ case kDataPar: axis_type = "spatial"; break;
+ case kCommReduce: axis_type = "reduce"; break;
+ case kOrdered: axis_type = "scan"; break;
+ default: axis_type = "opaque"; break;
+ }
+
+ // Domain
+ ExprAST dom(ffi::UnsafeInit{});
+ if (iv->dom.defined()) {
+ if (is_zero(iv->dom->min)) {
+ dom = Print(printer, iv->dom->extent, iv_p->Attr("dom")->Attr("extent"));
+ } else {
+ dom = TupleAST({}, {Print(printer, iv->dom->min, iv_p->Attr("dom")->Attr("min")),
+ Print(printer, iv->dom->min + iv->dom->extent,
+ iv_p->Attr("dom")->Attr("extent"))});
+ }
+ } else {
+ dom = LiteralAST::Null();
+ }
+
+ // Value (from iter_values)
+ ExprAST val(ffi::UnsafeInit{});
+ if (i < static_cast<int>(realize->iter_values.size())) {
+ val = Print(printer, realize->iter_values[i], path->Attr("iter_values")->ArrayItem(i));
+ } else {
+ val = LiteralAST::Null();
+ }
+
+ // vi = T.axis.spatial(dom, val)
+ ExprAST rhs = ExprCall(ExprAttr(ExprAttr(IdAST("T"), "axis"), axis_type), {dom, val});
+ frame->stmts.push_back(
+ AssignAST(var_id, rhs, Optional<ExprAST>()));
+ }
+
+ // Predicate
+ if (!is_one(realize->predicate)) {
+ ExprAST pred = Print(printer, realize->predicate, path->Attr("predicate"));
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(TIR("where"), {pred})));
+ }
+
+ // Reads — always emit T.reads() even when empty (required for opaque blocks
+ // where the parser can't auto-detect access regions)
+ {
+ List<ExprAST> reads;
+ for (int i = 0; i < static_cast<int>(block->reads.size()); ++i) {
+ reads.push_back(
+ Print(printer, block->reads[i], block_p->Attr("reads")->ArrayItem(i)));
+ }
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(TIR("reads"), std::move(reads))));
+ }
+
+ // Writes — always emit T.writes() even when empty (same reason as reads)
+ {
+ List<ExprAST> writes;
+ for (int i = 0; i < static_cast<int>(block->writes.size()); ++i) {
+ writes.push_back(
+ Print(printer, block->writes[i], block_p->Attr("writes")->ArrayItem(i)));
+ }
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(TIR("writes"), std::move(writes))));
+ }
+
+ // Annotations
+ if (!block->annotations.empty()) {
+ ExprAST annot = Print(printer, block->annotations, block_p->Attr("annotations"));
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(TIR("sblock_attr"), {annot})));
+ }
+
+ // Alloc buffers (use PrintBufferDecl for V1-compatible kwarg elision)
+ for (int i = 0; i < static_cast<int>(block->alloc_buffers.size()); ++i) {
+ Buffer buf = block->alloc_buffers[i];
+ AccessPath buffer_p = block_p->Attr("alloc_buffers")->ArrayItem(i);
+ printer->VarDef(buf->name, buf, frame);
+ DefineBufferDataVar(buf, printer);
+ ExprAST buf_id = printer->VarGet(buf).value();
+ List<ExprAST> no_extra;
+ ExprAST rhs = PrintBufferDecl(buf, "sblock_alloc_buffer", std::move(no_extra),
+ printer, buffer_p);
+ frame->stmts.push_back(
+ AssignAST(buf_id, rhs, Optional<ExprAST>()));
+ }
+
+ // Match buffers
+ for (int i = 0; i < static_cast<int>(block->match_buffers.size()); ++i) {
+ NodeAST s = printer->operator()(Any(block->match_buffers[i]),
+ block_p->Attr("match_buffers")->ArrayItem(i))
+ .cast<NodeAST>();
+ if (s.get()->IsInstance<StmtASTObj>()) {
+ frame->stmts.push_back(GetRef<StmtAST>(static_cast<const StmtASTObj*>(s.get())));
+ }
+ }
+
+ // Init
+ if (block->init.defined()) {
+ List<StmtAST> init_body = PrintBodyStmts(block->init.value(), printer,
+ block_p->Attr("init"));
+ ExprAST init_ctx = ExprCall(TIR("init"), {});
+ frame->stmts.push_back(
+ WithAST(Optional<ExprAST>(), init_ctx, init_body));
+ }
+
+ // Body
+ List<StmtAST> body = PrintBodyStmts(block->body, printer, block_p->Attr("body"));
+
+ // Merge
+ List<StmtAST> all_body;
+ for (const auto& s : frame->stmts) all_body.push_back(s);
+ for (const auto& s : body) all_body.push_back(s);
+
+ printer->FramePop();
+ return WithAST(Optional<ExprAST>(), ctx, all_body);
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // SBlockRealize is the main entry point
+ refl::TypeAttrDef<SBlockRealizeNode>().def(
+ "__ffi_text_print__",
+ [](SBlockRealize node, IRPrinter printer, AccessPath path) -> NodeAST {
+ return PrintSBlockRealize(node, printer, path);
+ });
+
+ // SBlock (standalone, without realize) → same but no iter_values/predicate
+ refl::TypeAttrDef<SBlockNode>().def(
+ "__ffi_text_print__",
+ [](SBlock block, IRPrinter printer, AccessPath path) -> NodeAST {
+ // Create a dummy realize
+ Array<PrimExpr> iter_values;
+ for (const auto& iv : block->iter_vars) {
+ iter_values.push_back(iv->var);
+ }
+ SBlockRealize realize(iter_values, Bool(true), block);
+ return PrintSBlockRealize(realize, printer, path);
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/buffer.cc b/src/script_v2/printer/tirx/buffer.cc
new file mode 100644
index 0000000000..cd396bec09
--- /dev/null
+++ b/src/script_v2/printer/tirx/buffer.cc
@@ -0,0 +1,222 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/tirx/buffer.cc
+ * \brief Buffer __ffi_text_print__ kept to prevent infinite recursion
+ * when PrintBufferTy resolves buffer fields that reference other
+ * buffers. AllocBuffer/DeclBuffer/MatchBufferRegion/BufferRegion
+ * are now handled by AssignCall/Load traits.
+ */
+#include <tvm/tirx/buffer.h>
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt.h>
+
+#include "../utils.h"
+#include "utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // BufferRegion: Range→slice conversion
+ refl::TypeAttrDef<BufferRegionNode>().def(
+ "__ffi_text_print__",
+ [](BufferRegion node, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST buf = Print(printer, node->buffer, path->Attr("buffer"));
+ List<ExprAST> idx;
+ for (int i = 0; i < static_cast<int>(node->region.size()); ++i) {
+ Range r = node->region[i];
+ AccessPath rp = path->Attr("region")->ArrayItem(i);
+ // If extent is 1, use plain index instead of slice
+ if (is_one(r->extent)) {
+ idx.push_back(Print(printer, r->min, rp->Attr("min")));
+ } else {
+ ExprAST start = Print(printer, r->min, rp->Attr("min"));
+ PrimExpr end_expr = r->min + r->extent;
+ ExprAST end = Print(printer, end_expr, rp->Attr("extent"));
+ idx.push_back(SliceAST({}, start, end, Optional<ExprAST>()));
+ }
+ }
+ return ExprIndex(buf, std::move(idx));
+ });
+
+ // AllocBuffer: buf = T.alloc_buffer(shape, dtype, scope=..., ...)
+ // Uses PrintBufferDecl for V1-compatible positional dtype and kwarg elision.
+ refl::TypeAttrDef<AllocBufferNode>().def(
+ "__ffi_text_print__",
+ [](AllocBuffer node, IRPrinter printer, AccessPath path) -> NodeAST {
+ Buffer buf = node->buffer;
+ DefaultFrame frame = printer->frames.back().cast<DefaultFrame>();
+ printer->VarDef(buf->name, buf, frame);
+ DefineBufferDataVar(buf, printer);
+ ExprAST buf_id = printer->VarGet(buf).value();
+ AccessPath buffer_p = path->Attr("buffer");
+ List<ExprAST> no_extra;
+ // Use PrintBufferDecl for core buffer kwargs
+ ExprAST rhs = PrintBufferDecl(buf, "alloc_buffer", std::move(no_extra),
+ printer, buffer_p);
+ return AssignAST(buf_id, rhs, Optional<ExprAST>());
+ });
+
+ // DeclBuffer: buf = T.decl_buffer(shape, dtype, kwargs...)
+ // Uses PrintBufferDecl for full kwargs (data=, strides=, etc.)
+ refl::TypeAttrDef<DeclBufferNode>().def(
+ "__ffi_text_print__",
+ [](DeclBuffer node, IRPrinter printer, AccessPath path) -> NodeAST {
+ Buffer buf = node->buffer;
+ DefaultFrame frame = printer->frames.back().cast<DefaultFrame>();
+ DefineBufferVars(buf, printer, frame);
+ printer->VarDef(buf->name, buf, frame);
+ DefineBufferDataVar(buf, printer);
+ ExprAST buf_id = printer->VarGet(buf).value();
+ List<ExprAST> no_extra;
+ ExprAST rhs = PrintBufferDecl(buf, "decl_buffer", std::move(no_extra),
+ printer, path->Attr("buffer"));
+ return AssignAST(buf_id, rhs, Optional<ExprAST>());
+ });
+
+ // MatchBufferRegion: buf = T.match_buffer(source, shape, dtype, kwargs...)
+ refl::TypeAttrDef<MatchBufferRegionNode>().def(
+ "__ffi_text_print__",
+ [](MatchBufferRegion node, IRPrinter printer, AccessPath path) -> NodeAST {
+ Buffer buf = node->buffer;
+ DefaultFrame frame = printer->frames.back().cast<DefaultFrame>();
+ // Define any new shape/stride/elem_offset vars before printing
+ DefineBufferVars(buf, printer, frame);
+ printer->VarDef(buf->name, buf, frame);
+ DefineBufferDataVar(buf, printer);
+ ExprAST buf_id = printer->VarGet(buf).value();
+ ExprAST source = Print(printer, node->source, path->Attr("source"));
+ List<ExprAST> extra_args;
+ extra_args.push_back(source);
+ ExprAST rhs = PrintBufferDecl(buf, "match_buffer", std::move(extra_args),
+ printer, path->Attr("buffer"));
+ return AssignAST(buf_id, rhs, Optional<ExprAST>());
+ });
+
+ // BufferLoad: convert Ramp indices to slices; fall back to vload() for predicate
+ refl::TypeAttrDef<BufferLoadNode>().def(
+ "__ffi_text_print__",
+ [](BufferLoad node, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST buf = Print(printer, node->buffer, path->Attr("buffer"));
+ // Fall back to buf.vload(...) when predicate is present
+ if (node->predicate.defined()) {
+ ExprAST indices = Print(printer, node->indices, path->Attr("indices"));
+ ExprAST predicate = Print(printer, node->predicate.value(), path->Attr("predicate"));
+ return ExprCallKw(ExprAttr(buf, "vload"), {indices},
+ {String("predicate")}, {predicate});
+ }
+ // Convert Ramp indices to slices
+ List<ExprAST> idx;
+ for (int i = 0; i < static_cast<int>(node->indices.size()); ++i) {
+ AccessPath idx_p = path->Attr("indices")->ArrayItem(i);
+ if (const auto* ramp = node->indices[i].as<RampNode>()) {
+ if (ramp->stride.as<IntImmNode>()) {
+ ExprAST start = Print(printer, ramp->base, idx_p->Attr("base"));
+ PrimExpr stop_val = ramp->base + ramp->lanes * ramp->stride;
+ ExprAST stop = Print(printer, stop_val, idx_p->Attr("lanes"));
+ Optional<ExprAST> step;
+ if (!is_one(ramp->stride)) {
+ step = Print(printer, ramp->stride, idx_p->Attr("stride"));
+ }
+ idx.push_back(SliceAST({}, start, stop, step));
+ continue;
+ }
+ }
+ idx.push_back(Print(printer, node->indices[i], idx_p));
+ }
+ return ExprIndex(buf, std::move(idx));
+ });
+
+ // BufferStore: convert Ramp indices to slices; fall back to vstore() for predicate
+ refl::TypeAttrDef<BufferStoreNode>().def(
+ "__ffi_text_print__",
+ [](BufferStore node, IRPrinter printer, AccessPath path) -> NodeAST {
+ ExprAST buf = Print(printer, node->buffer, path->Attr("buffer"));
+ ExprAST value = Print(printer, node->value, path->Attr("value"));
+ // Fall back to buf.vstore(...) when predicate is present
+ if (node->predicate.defined()) {
+ ExprAST indices = Print(printer, node->indices, path->Attr("indices"));
+ ExprAST predicate = Print(printer, node->predicate.value(), path->Attr("predicate"));
+ return ExprStmtAST(ExprCallKw(ExprAttr(buf, "vstore"), {indices, value},
+ {String("predicate")}, {predicate}));
+ }
+ // Convert Ramp indices to slices
+ List<ExprAST> idx;
+ for (int i = 0; i < static_cast<int>(node->indices.size()); ++i) {
+ AccessPath idx_p = path->Attr("indices")->ArrayItem(i);
+ if (const auto* ramp = node->indices[i].as<RampNode>()) {
+ if (ramp->stride.as<IntImmNode>()) {
+ ExprAST start = Print(printer, ramp->base, idx_p->Attr("base"));
+ PrimExpr stop_val = ramp->base + ramp->lanes * ramp->stride;
+ ExprAST stop = Print(printer, stop_val, idx_p->Attr("lanes"));
+ Optional<ExprAST> step;
+ if (!is_one(ramp->stride)) {
+ step = Print(printer, ramp->stride, idx_p->Attr("stride"));
+ }
+ idx.push_back(SliceAST({}, start, stop, step));
+ continue;
+ }
+ }
+ idx.push_back(Print(printer, node->indices[i], idx_p));
+ }
+ return AssignAST(ExprIndex(buf, std::move(idx)), value, Optional<ExprAST>());
+ });
+
+ // Buffer: variable check + kwarg elision
+ refl::TypeAttrDef<BufferNode>().def(
+ "__ffi_text_print__",
+ [](Buffer buf, IRPrinter printer, AccessPath path) -> NodeAST {
+ if (printer->VarGet(buf).has_value()) {
+ return printer->VarGet(buf).value();
+ }
+ List<ExprAST> args;
+ args.push_back(PrintTuple(printer, buf->shape, path->Attr("shape")));
+ args.push_back(LiteralAST::Str(DType2Str(buf->dtype)));
+ List<String> kw_keys;
+ List<ExprAST> kw_vals;
+ if (!buf->strides.empty()) {
+ kw_keys.push_back(String("strides"));
+ kw_vals.push_back(PrintList(printer, buf->strides, path->Attr("strides")));
+ }
+ if (const auto* offset = buf->elem_offset.as<IntImmNode>()) {
+ if (offset->value != 0) {
+ kw_keys.push_back(String("elem_offset"));
+ kw_vals.push_back(Print(printer, buf->elem_offset, path->Attr("elem_offset")));
+ }
+ } else {
+ kw_keys.push_back(String("elem_offset"));
+ kw_vals.push_back(Print(printer, buf->elem_offset, path->Attr("elem_offset")));
+ }
+ if (buf.scope() != "global") {
+ kw_keys.push_back(String("scope"));
+ kw_vals.push_back(LiteralAST::Str(buf.scope()));
+ }
+ if (!kw_keys.empty()) {
+ return ExprCallKw(TIR("Buffer"), std::move(args), std::move(kw_keys),
+ std::move(kw_vals));
+ }
+ return ExprCall(TIR("Buffer"), std::move(args));
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/expr.cc b/src/script_v2/printer/tirx/expr.cc
new file mode 100644
index 0000000000..64acfc6c05
--- /dev/null
+++ b/src/script_v2/printer/tirx/expr.cc
@@ -0,0 +1,184 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/tirx/expr.cc
+ * \brief __ffi_text_print__ overrides for genuinely irregular TIRx expressions.
+ *
+ * Nodes handled purely by traits: Add..Or, Div/Mod/Min/Max (BinOp), Not (UnaryOp),
+ * Var/SizeVar (Value), StringImm (Literal), BufferLoad/ProducerLoad (Load),
+ * Cast/Select/Ramp/Broadcast/Shuffle/IterVar (Call+$method), tirx.Call (Call+$method),
+ * Let (Call+kwargs via $method), Reduce (Call+kwargs via $method).
+ *
+ * Only CommReducer remains here — its lambda construction is genuinely irreducible.
+ */
+#include <tvm/ir/expr.h>
+#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/op_attr_types.h>
+
+#include "../utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+
+static ExprAST TIRCall(const std::string& name, List<ExprAST> args) {
+ return ExprCall(TIR(name), std::move(args));
+}
+
+// tirx.Call: __ffi_text_print__ override
+// Only intercepts when op is a GlobalVar (needs printer context for "Module.func_name").
+// All other tirx::Call nodes (Op-based) are handled by the trait+computed-method path,
+// which already uses TScriptPrinterName and TScriptDtypePrintLocation.
+TVM_FFI_STATIC_INIT_BLOCK() {
+ refl::TypeAttrDef<CallNode>().def(
+ "__ffi_text_print__",
+ [](tirx::Call call, IRPrinter printer, AccessPath path) -> NodeAST {
+ // Only handle GlobalVar ops here; for Op-based calls, return the result of
+ // the generic trait dispatch (which uses call_callee + call_args computed methods).
+ if (!call->op->IsInstance<GlobalVarNode>()) {
+ // Delegate to trait-based printing by invoking the generic dispatch.
+ // We return the trait-printed result by calling the printer on the object
+ // with trait dispatch. We need to trigger trait-based printing manually.
+ // The simplest approach: compute callee and args inline.
+ //
+ // Callee: use TScriptPrinterName if available, else strip prefix
+ std::string callee_name;
+ if (auto* op = call->op.as<OpNode>()) {
+ static const OpAttrMap<ffi::String> op_names =
+ Op::GetAttrMap<ffi::String>("TScriptPrinterName");
+ Op op_ref = ffi::GetRef<Op>(op);
+ if (op_names.count(op_ref)) {
+ callee_name = std::string(op_names[op_ref]);
+ } else {
+ std::string full(op->name);
+ auto pos = full.rfind('.');
+ callee_name = (pos != std::string::npos) ? full.substr(pos + 1) : full;
+ }
+ } else {
+ callee_name = "call";
+ }
+ ExprAST callee = TIR(callee_name);
+ // Args: compute using TScriptDtypePrintLocation
+ List<ExprAST> arg_docs;
+ int print_location = 0; // kNone
+ if (auto* op = call->op.as<OpNode>()) {
+ static const OpAttrMap<TScriptDtypePrintLocation> dtype_locations =
+ Op::GetAttrMap<TScriptDtypePrintLocation>("TScriptDtypePrintLocation");
+ Op op_ref = ffi::GetRef<Op>(op);
+ if (dtype_locations.count(op_ref)) {
+ print_location = dtype_locations[op_ref].IntValue();
+ }
+ }
+ std::string dtype_str = call->dtype.is_void()
+ ? "void"
+ : ffi::DLDataTypeToString(call->dtype);
+ // Check for LLVM intrinsic (needs name lookup for first arg)
+ bool is_llvm_intrin = false;
+ if (auto* op = call->op.as<OpNode>()) {
+ static const OpAttrMap<ffi::String> op_names2 =
+ Op::GetAttrMap<ffi::String>("TScriptPrinterName");
+ Op op_ref = ffi::GetRef<Op>(op);
+ if (op_names2.count(op_ref)) {
+ std::string n(op_names2[op_ref]);
+ is_llvm_intrin = (n == "call_llvm_pure_intrin" || n == "call_llvm_intrin");
+ }
+ }
+ if (print_location == 1) { // kFirst
+ arg_docs.push_back(LiteralAST::Str(dtype_str));
+ }
+ for (int i = 0, n = static_cast<int>(call->args.size()); i < n; ++i) {
+ if (i == 0 && is_llvm_intrin) {
+ auto f_lookup = ffi::Function::GetGlobal("target.llvm_get_intrinsic_name");
+ if (f_lookup.has_value() && call->args[0].as<IntImmNode>()) {
+ int64_t id = call->args[0].as<IntImmNode>()->value;
+ ffi::Any ret;
+ ffi::AnyView args_view[1] = {ffi::AnyView(id)};
+ f_lookup.value().CallPacked(args_view, 1, &ret);
+ ffi::String name = ret.cast<ffi::String>();
+ arg_docs.push_back(LiteralAST::Str(std::string(name)));
+ continue;
+ }
+ }
+ arg_docs.push_back(
+ Print(printer, call->args[i], path->Attr("args")->ArrayItem(i)));
+ }
+ if (print_location == 2) { // kLast
+ arg_docs.push_back(LiteralAST::Str(dtype_str));
+ }
+ return CallAST(callee, std::move(arg_docs), {}, {});
+ }
+ // GlobalVar op: look up the printer's variable table for "Module.func_name"
+ ExprAST callee(ffi::UnsafeInit{});
+ if (auto gv_doc = printer->VarGet(call->op)) {
+ callee = gv_doc.value();
+ } else {
+ // Fallback: the GlobalVar in call->op may be a different object instance
+ // than the one registered by the module printer. Search by name_hint.
+ GlobalVar op_gv = Downcast<GlobalVar>(call->op);
+ bool found = false;
+ for (const auto& kv : printer->obj2info) {
+ if (const auto* gv_node = kv.first.as<GlobalVarNode>()) {
+ if (gv_node->name_hint == op_gv->name_hint) {
+ callee = kv.second->creator().cast<ExprAST>();
+ found = true;
+ break;
+ }
+ }
+ }
+ if (!found) {
+ // Last resort: print the GlobalVar through dispatch
+ callee = Print(printer, call->op, path->Attr("op"));
+ }
+ }
+ List<ExprAST> args;
+ for (int i = 0, n = call->args.size(); i < n; ++i) {
+ args.push_back(Print(printer, call->args[i], path->Attr("args")->ArrayItem(i)));
+ }
+ return CallAST(callee, std::move(args), {}, {});
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // CommReducer: lambda construction — genuinely irreducible
+ refl::TypeAttrDef<CommReducerNode>().def(
+ "__ffi_text_print__",
+ [](CommReducer node, IRPrinter printer, AccessPath path) -> NodeAST {
+ List<ExprAST> lhs_vars, rhs_vars, results;
+ for (int i = 0; i < static_cast<int>(node->lhs.size()); ++i)
+ lhs_vars.push_back(Print(printer, node->lhs[i], path->Attr("lhs")->ArrayItem(i)));
+ for (int i = 0; i < static_cast<int>(node->rhs.size()); ++i)
+ rhs_vars.push_back(Print(printer, node->rhs[i], path->Attr("rhs")->ArrayItem(i)));
+ for (int i = 0; i < static_cast<int>(node->result.size()); ++i)
+ results.push_back(Print(printer, node->result[i], path->Attr("result")->ArrayItem(i)));
+ List<ExprAST> params;
+ params.insert(params.end(), lhs_vars.begin(), lhs_vars.end());
+ params.insert(params.end(), rhs_vars.begin(), rhs_vars.end());
+ ExprAST lambda_body = (results.size() == 1) ? results[0] : TupleAST({}, results);
+ LambdaAST lambda_ast({}, params, lambda_body);
+ return TIRCall("comm_reducer",
+ {lambda_ast, PrintList(printer, node->identity_element,
+ path->Attr("identity_element"))});
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/for_loop.cc b/src/script_v2/printer/tirx/for_loop.cc
new file mode 100644
index 0000000000..b4cf63b0f4
--- /dev/null
+++ b/src/script_v2/printer/tirx/for_loop.cc
@@ -0,0 +1,155 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/tirx/for_loop.cc
+ * \brief For loop __ffi_text_print__ kept due to cross-field logic
+ * (kind + thread_binding + annotations + range sugar + T.grid).
+ */
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt.h>
+
+#include <unordered_set>
+#include <vector>
+
+#include "utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ refl::TypeAttrDef<ForNode>().def(
+ "__ffi_text_print__",
+ [](tirx::For loop, IRPrinter printer, AccessPath path) -> NodeAST {
+ // Step 1. Check syntactic sugar: T.grid
+ // Collect consecutive serial for-loops with min=0, no annotations,
+ // trivial step, and whose extents don't depend on prior grid loop vars.
+ std::vector<const ForNode*> grid;
+ std::unordered_set<const tirx::VarNode*> grid_loop_vars;
+ {
+ for (const ForNode* l = loop.get(); l != nullptr;
+ l = l->body.as<ForNode>()) {
+ if (l->kind != ForKind::kSerial ||
+ !is_zero(l->min) ||
+ !l->annotations.empty() ||
+ !l->HasTrivialStep() ||
+ tirx::UsesVar(l->extent, [&grid_loop_vars](const tirx::VarNode* v) {
+ return grid_loop_vars.count(v) > 0;
+ })) {
+ break;
+ }
+ grid.push_back(l);
+ grid_loop_vars.insert(l->loop_var.get());
+ }
+ }
+
+ // Step 2. If grid.size() > 1, print as T.grid
+ if (grid.size() > 1) {
+ DefaultFrame frame;
+ printer->FramePush(frame);
+ int n = grid.size();
+ List<ExprAST> lhs_vars;
+ List<ExprAST> extents;
+ AccessPath cur_p = path;
+ for (int i = 0; i < n; ++i) {
+ const ForNode* g = grid[i];
+ lhs_vars.push_back(
+ DefineVar(GetRef<Var>(static_cast<const VarNode*>(g->loop_var.get())),
+ printer, cur_p->Attr("loop_var")));
+ extents.push_back(Print(printer, g->extent, cur_p->Attr("extent")));
+ cur_p = cur_p->Attr("body");
+ }
+ ExprAST lhs = TupleAST({}, std::move(lhs_vars));
+ ExprAST rhs = ExprCall(TIR("grid"), std::move(extents));
+ List<StmtAST> body = PrintBodyStmts(
+ GetRef<Stmt>(static_cast<const StmtNode*>(grid.back()->body.get())),
+ printer, cur_p);
+ List<StmtAST> all_body;
+ for (const auto& s : frame->stmts) all_body.push_back(s);
+ for (const auto& s : body) all_body.push_back(s);
+ printer->FramePop();
+ return ForAST(lhs, rhs, all_body);
+ }
+
+ // Step 3. Single for loop (no grid sugar)
+ DefaultFrame frame;
+ printer->FramePush(frame);
+ ExprAST lhs = DefineVar(loop->loop_var, printer, path->Attr("loop_var"));
+
+ ExprAST rhs(ffi::UnsafeInit{});
+ bool is_zero_min = is_zero(loop->min);
+ bool has_trivial_step = loop->HasTrivialStep();
+
+ if (loop->kind == ForKind::kSerial && loop->annotations.empty()) {
+ if (is_zero_min && has_trivial_step) {
+ rhs = ExprCall(IdAST("range"),
+ {Print(printer, loop->extent, path->Attr("extent"))});
+ } else {
+ PrimExpr end = loop->min + loop->extent;
+ rhs = ExprCall(IdAST("range"),
+ {Print(printer, loop->min, path->Attr("min")),
+ Print(printer, end, path->Attr("extent"))});
+ }
+ } else {
+ std::string prefix;
+ switch (loop->kind) {
+ case ForKind::kSerial: prefix = "serial"; break;
+ case ForKind::kParallel: prefix = "parallel"; break;
+ case ForKind::kUnrolled: prefix = "unroll"; break;
+ case ForKind::kVectorized: prefix = "vectorized"; break;
+ case ForKind::kThreadBinding: prefix = "thread_binding"; break;
+ default: prefix = "serial"; break;
+ }
+ List<ExprAST> args;
+ if (is_zero_min && has_trivial_step) {
+ args.push_back(Print(printer, loop->extent, path->Attr("extent")));
+ } else {
+ PrimExpr end = loop->min + loop->extent;
+ args.push_back(Print(printer, loop->min, path->Attr("min")));
+ args.push_back(Print(printer, end, path->Attr("extent")));
+ }
+ List<String> kw_keys;
+ List<ExprAST> kw_vals;
+ if (!loop->annotations.empty()) {
+ kw_keys.push_back(String("annotations"));
+ kw_vals.push_back(Print(printer, loop->annotations, path->Attr("annotations")));
+ }
+ if (loop->kind == ForKind::kThreadBinding && loop->thread_binding.defined()) {
+ kw_keys.push_back(String("thread"));
+ kw_vals.push_back(
+ LiteralAST::Str(loop->thread_binding.value()->thread_tag,
+ {path->Attr("thread_binding")}));
+ }
+ rhs = !kw_keys.empty()
+ ? ExprCallKw(TIR(prefix), std::move(args), std::move(kw_keys), std::move(kw_vals))
+ : ExprCall(TIR(prefix), std::move(args));
+ }
+
+ List<StmtAST> body = PrintBodyStmts(loop->body, printer, path->Attr("body"));
+ List<StmtAST> all_body;
+ for (const auto& s : frame->stmts) all_body.push_back(s);
+ for (const auto& s : body) all_body.push_back(s);
+ printer->FramePop();
+ return ForAST(lhs, rhs, all_body);
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/function.cc b/src/script_v2/printer/tirx/function.cc
new file mode 100644
index 0000000000..3be84cfec6
--- /dev/null
+++ b/src/script_v2/printer/tirx/function.cc
@@ -0,0 +1,336 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/tirx/function.cc
+ * \brief V2 text printer for tirx.PrimFunc.
+ */
+#include <tvm/ir/attrs.h>
+#include <tvm/tirx/function.h>
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <unordered_map>
+#include <unordered_set>
+
+#include "utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // PrimFunc → @T.prim_func \n def name(params...): body
+ refl::TypeAttrDef<PrimFuncNode>().def(
+ "__ffi_text_print__",
+ [](PrimFunc func, IRPrinter printer, AccessPath path) -> NodeAST {
+ // Determine function name: prefer binding name from module context,
+ // then attrs["global_symbol"], then "main" (matching V1 FindFunctionName).
+ String func_name = "main";
+ bool in_module = printer->VarIsDefined(func);
+ if (in_module) {
+ if (auto binding_expr = printer->VarGet(func)) {
+ if (auto* id_node = binding_expr.value().as<::tvm::ffi::ir::text::IdASTObj>()) {
+ func_name = id_node->name;
+ }
+ }
+ } else if (func->attrs.defined()) {
+ auto it = func->attrs->dict.find("global_symbol");
+ if (it != func->attrs->dict.end()) {
+ func_name = (*it).second.cast<String>();
+ }
+ }
+ IdAST name = IdAST(func_name);
+
+ // Build decorator: @T.prim_func or @T.prim_func(private=True)
+ List<ExprAST> decorators;
+ bool has_global_symbol = func->attrs.defined() &&
+ func->attrs->dict.count("global_symbol");
+ if (!has_global_symbol) {
+ decorators.push_back(
+ ExprCallKw(TIR("prim_func"), {}, {String("private")}, {LiteralAST::Bool(true)}));
+ } else {
+ decorators.push_back(TIR("prim_func"));
+ }
+
+ // Push frame
+ DefaultFrame frame;
+ printer->FramePush(frame);
+
+ // Pre-compute buffer_data_counter: how many buffers in buffer_map share each data var
+ int n_args = func->params.size();
+ std::unordered_map<const tirx::VarNode*, int> buffer_data_counter;
+ for (const auto& pair : func->buffer_map) {
+ const tirx::VarNode* data_var = pair.second->data.get();
+ if (!buffer_data_counter.count(data_var)) {
+ buffer_data_counter.insert({data_var, 0});
+ }
+ ++buffer_data_counter.at(data_var);
+ }
+
+ // Step 1. Handle params with buffer inlining (matching V1 logic)
+ List<AssignAST> params;
+ std::unordered_set<const BufferNode*> buffer_inlined;
+
+ for (int i = 0; i < n_args; ++i) {
+ Var var = func->params[i];
+ AccessPath var_p = path->Attr("params")->ArrayItem(i);
+
+ // V1 inlining: syntax_sugar && CountVarOccurrence==2 && buffer_map has var
+ // && IsSimpleBuffer && buffer_data_counter==1
+ if (CountVarOccurrence(func, var) == 2 && func->buffer_map.count(var)) {
+ Buffer buffer = func->buffer_map[var];
+ if (IsSimpleBuffer(buffer) && buffer_data_counter.at(buffer->data.get()) == 1) {
+ AccessPath buffer_p = path->Attr("buffer_map")->MapItem(var);
+ // Define buffer variable (use buffer name, not param var name)
+ printer->VarDef(buffer->name, buffer, frame);
+ DefineBufferDataVar(buffer, printer);
+ ExprAST buf_id = printer->VarGet(buffer).value();
+ ExprAST annotation = PrintBufferAnnotation(buffer, printer, buffer_p);
+ params.push_back(AssignAST(buf_id, Optional<ExprAST>(),
+ Optional<ExprAST>(annotation)));
+ buffer_inlined.insert(buffer.get());
+ continue;
+ }
+ }
+
+ // Regular param with type annotation
+ ExprAST var_id = DefineVar(var, printer, var_p);
+ Optional<ExprAST> annotation;
+ if (var->type_annotation.defined()) {
+ annotation = Print(printer, var->type_annotation, var_p->Attr("type_annotation"));
+ }
+ params.push_back(
+ AssignAST(var_id, Optional<ExprAST>(), annotation));
+ }
+
+ // Step 2. Handle func->attrs (filter out global_symbol if it matches func name)
+ // Helper: build a DictAST from a Map<String, Any>
+ auto PrintAttrDict = [&](const Map<String, Any>& dict,
+ const AccessPath& dict_p) -> ExprAST {
+ List<ExprAST> keys;
+ List<ExprAST> vals;
+ for (const auto& kv : dict) {
+ keys.push_back(LiteralAST::Str(kv.first));
+ vals.push_back(Print(printer, kv.second, dict_p));
+ }
+ return DictAST(std::move(keys), std::move(vals));
+ };
+
+ if (func->attrs.defined() && !func->attrs->dict.empty()) {
+ if (func->attrs->dict.count("global_symbol") &&
+ func->attrs->dict.at("global_symbol").cast<String>() == func_name) {
+ // global_symbol matches func name: filter it out
+ Map<String, Any> new_attrs;
+ for (const auto& kv : func->attrs->dict) {
+ if (kv.first != "global_symbol") {
+ new_attrs.Set(kv.first, kv.second);
+ }
+ }
+ if (!new_attrs.empty()) {
+ ExprAST attr_dict = PrintAttrDict(new_attrs, path->Attr("attrs"));
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(TIR("func_attr"), {attr_dict})));
+ }
+ } else {
+ // No global_symbol or it doesn't match: print all attrs as dict
+ ExprAST attr_dict = PrintAttrDict(func->attrs->dict, path->Attr("attrs"));
+ frame->stmts.push_back(
+ ExprStmtAST(ExprCall(TIR("func_attr"), {attr_dict})));
+ }
+ }
+
+ // Step 3. Handle buffer_map: non-inlined entries → match_buffer with full kwargs
+ // Process each buffer sequentially: emit var defs then match_buffer for each
+ // (matching V1's interleaved behavior where BufferAttrs defines vars inline)
+ for (int i = 0; i < n_args; ++i) {
+ Var param = func->params[i];
+ if (func->buffer_map.count(param)) {
+ Buffer buffer = func->buffer_map[param];
+ if (buffer_inlined.count(buffer.get())) continue;
+ // Define any new shape vars for this buffer
+ DefineBufferVars(buffer, printer, frame);
+ // Emit match_buffer statement
+ AccessPath buffer_p = path->Attr("buffer_map")->MapItem(param);
+ printer->VarDef(buffer->name, buffer, frame);
+ DefineBufferDataVar(buffer, printer);
+ ExprAST buf_id = printer->VarGet(buffer).value();
+ ExprAST param_doc = params[i]->lhs;
+ List<ExprAST> extra_args;
+ extra_args.push_back(param_doc);
+ ExprAST rhs = PrintBufferDecl(buffer, "match_buffer", std::move(extra_args),
+ printer, buffer_p);
+ frame->stmts.push_back(AssignAST(buf_id, rhs, Optional<ExprAST>()));
+ }
+ }
+
+ // Step 3b. Emit declarations for undefined Vars
+ // V1's SetCommonPrefix + PrintVar detects Vars used before definition
+ // and emits `var = T.dtype()`. Replicate this: collect all vars already
+ // defined (params, buffer_map vars), then find undefined vars in the body.
+ // Skip vars that will be defined by launch_thread AttrStmts (thread_extent
+ // or virtual_thread with IterVar nodes), since those get their definition
+ // from the assignment form `var = T.launch_thread(...)`.
+ {
+ // Collect thread vars that launch_thread will define
+ std::unordered_set<const tirx::VarNode*> thread_vars;
+ {
+ class ThreadVarCollector : public tirx::StmtVisitor {
+ public:
+ std::unordered_set<const tirx::VarNode*>* result;
+ void VisitStmt_(const tirx::AttrStmtNode* op) final {
+ if ((op->attr_key == "thread_extent" || op->attr_key == "virtual_thread") &&
+ (op->node.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) &&
+ op->node.cast<ObjectRef>()->IsInstance<IterVarNode>()) {
+ IterVar iv = op->node.cast<IterVar>();
+ result->insert(iv->var.get());
+ }
+ tirx::StmtVisitor::VisitStmt_(op);
+ }
+ };
+ ThreadVarCollector collector;
+ collector.result = &thread_vars;
+ collector(func->body);
+ }
+
+ ffi::Array<Var> defined_vars;
+ // Params
+ for (const auto& param : func->params) {
+ defined_vars.push_back(param);
+ }
+ // Buffer data + shape vars from buffer_map
+ for (const auto& pair : func->buffer_map) {
+ Buffer buf = pair.second;
+ defined_vars.push_back(buf->data);
+ for (const PrimExpr& s : buf->shape) {
+ if (const auto* v = s.as<tirx::VarNode>()) {
+ defined_vars.push_back(GetRef<Var>(v));
+ }
+ }
+ for (const PrimExpr& s : buf->strides) {
+ if (const auto* v = s.as<tirx::VarNode>()) {
+ defined_vars.push_back(GetRef<Var>(v));
+ }
+ }
+ if (const auto* v = buf->elem_offset.as<tirx::VarNode>()) {
+ defined_vars.push_back(GetRef<Var>(v));
+ }
+ }
+ // Find body to scan (handle implicit root block)
+ Stmt body_to_scan = func->body;
+ ffi::Array<Var> undef = tirx::UndefinedVars(body_to_scan, defined_vars);
+ // Deduplicate (UndefinedVars may return duplicates) while preserving order
+ std::unordered_set<const tirx::VarNode*> seen;
+ for (const Var& v : undef) {
+ if (seen.count(v.get())) continue;
+ seen.insert(v.get());
+ // Skip vars that will be defined by launch_thread
+ if (thread_vars.count(v.get())) continue;
+ // Emit: var_name = T.dtype()
+ if (!printer->VarGet(v).has_value()) {
+ DefineNewTIRVar(v, printer, frame);
+ }
+ }
+ }
+
+ // Step 4. Handle func->body with implicit root block detection
+ // V1 logic: detect single SBlockRealize with name "root", no iter_values,
+ // predicate==1, no annotations/match_buffers/reads/writes/init.
+ // If it has alloc_buffers, or body contains further SBlockRealizes with iter_vars,
+ // then treat it as implicit root.
+ Optional<SBlock> implicit_root_block;
+ {
+ const SBlockRealizeNode* root_block_realize =
+ func->body.as<SBlockRealizeNode>();
+ if (root_block_realize && !root_block_realize->iter_values.size() &&
+ is_one(root_block_realize->predicate)) {
+ SBlock root_block = root_block_realize->block;
+ if (!root_block->annotations.size() && !root_block->match_buffers.size() &&
+ !root_block->reads.size() && !root_block->writes.size() &&
+ !root_block->init.defined()) {
+ const SBlockRealizeNode* block_realize =
+ root_block->body.as<SBlockRealizeNode>();
+ if (root_block->alloc_buffers.size() ||
+ (block_realize && block_realize->block->iter_vars.size()) ||
+ (!block_realize &&
+ ContainsNode<SBlockRealizeNode>(root_block->body))) {
+ implicit_root_block = root_block;
+ }
+ }
+ }
+ }
+
+ List<StmtAST> body_stmts;
+ if (implicit_root_block.defined()) {
+ SBlock root_block = implicit_root_block.value();
+ AccessPath root_block_p = path->Attr("body")->Attr("block");
+ // Comment out the root block
+ frame->stmts.push_back(CommentAST(Optional<String>(String("with T.sblock(\"root\"):"))));
+ // Handle root block alloc_buffers
+ for (int i = 0, n = root_block->alloc_buffers.size(); i < n; ++i) {
+ Buffer buffer = root_block->alloc_buffers[i];
+ AccessPath buffer_p = root_block_p->Attr("alloc_buffers")->ArrayItem(i);
+ printer->VarDef(buffer->name, buffer, frame);
+ DefineBufferDataVar(buffer, printer);
+ ExprAST buf_id = printer->VarGet(buffer).value();
+ List<ExprAST> no_extra;
+ ExprAST rhs = PrintBufferDecl(buffer, "sblock_alloc_buffer", std::move(no_extra),
+ printer, buffer_p);
+ frame->stmts.push_back(AssignAST(buf_id, rhs, Optional<ExprAST>()));
+ }
+ // Print root block body
+ body_stmts = PrintBodyStmts(root_block->body, printer, root_block_p->Attr("body"));
+ } else {
+ body_stmts = PrintBodyStmts(func->body, printer, path->Attr("body"));
+ }
+
+ // Merge frame stmts + body
+ List<StmtAST> all_body;
+ for (const auto& s : frame->stmts) all_body.push_back(s);
+ for (const auto& s : body_stmts) all_body.push_back(s);
+
+ printer->FramePop();
+
+ // Return type annotation (suppress empty TupleType, matching V1)
+ Optional<ExprAST> ret_type;
+ if (func->ret_type.defined()) {
+ const auto* as_tuple = func->ret_type.as<TupleTypeNode>();
+ if (!as_tuple || as_tuple->fields.size()) {
+ ret_type = Print(printer, func->ret_type, path->Attr("ret_type"));
+ }
+ }
+
+ FunctionAST func_ast(name, params, decorators, ret_type, all_body);
+
+ // When printing standalone (not in module context), add header comments
+ // matching V1's HeaderWrapper behavior.
+ if (!in_module) {
+ List<StmtAST> result;
+ result.push_back(CommentAST(
+ Optional<String>(String("from tvm.script import tirx as T"))));
+ result.push_back(CommentAST(Optional<String>()));
+ result.push_back(func_ast);
+ return StmtBlockAST(result);
+ }
+ return func_ast;
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/stmt.cc b/src/script_v2/printer/tirx/stmt.cc
new file mode 100644
index 0000000000..ec4b22cd4f
--- /dev/null
+++ b/src/script_v2/printer/tirx/stmt.cc
@@ -0,0 +1,103 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file script_v2/printer/tirx/stmt.cc
+ * \brief Evaluate override (return sugar) and AttrStmt override
+ * (launch_thread vs T.attr).
+ */
+#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt.h>
+
+#include "../utils.h"
+#include "utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // Evaluate: detect Evaluate(Call(tirx.ret, [value])) and print as `return value`
+ refl::TypeAttrDef<EvaluateNode>().def(
+ "__ffi_text_print__",
+ [](Evaluate node, IRPrinter printer, AccessPath path) -> NodeAST {
+ if (auto* call = node->value.as<tirx::CallNode>()) {
+ if (call->op.same_as(tirx::builtin::ret()) && call->args.size() == 1) {
+ ExprAST val = Print(printer, call->args[0],
+ path->Attr("value")->Attr("args")->ArrayItem(0));
+ return ReturnAST(val);
+ }
+ }
+ // Default: if value is a Call, print as bare expression statement;
+ // otherwise wrap in T.evaluate(value) (matching V1 behavior).
+ ExprAST val = Print(printer, node->value, path->Attr("value"));
+ if (node->value->IsInstance<tirx::CallNode>()) {
+ return ExprStmtAST(val);
+ }
+ return ExprStmtAST(ExprCall(TIR("evaluate"), {val}));
+ });
+
+ // AttrStmt: conditional launch_thread vs T.attr()
+ refl::TypeAttrDef<AttrStmtNode>().def(
+ "__ffi_text_print__",
+ [](AttrStmt node, IRPrinter printer, AccessPath path) -> NodeAST {
+ bool is_iter_var = (node->node.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) &&
+ node->node.cast<ObjectRef>()->IsInstance<IterVarNode>();
+ if ((node->attr_key == "thread_extent" || node->attr_key == "virtual_thread") &&
+ is_iter_var) {
+ IterVar iv = node->node.cast<IterVar>();
+ // Check BEFORE printing the body, since body printing may auto-define the var
+ // via the def_free_var trait.
+ bool was_defined = printer->VarGet(iv->var).has_value();
+ if (!was_defined) {
+ // Pre-define the var so body printing uses the correct name.
+ DefineVar(iv->var, printer, path->Attr("node")->Attr("var"));
+ }
+ List<StmtAST> body = PrintBodyStmts(node->body, printer, path->Attr("body"));
+ if (was_defined) {
+ // Already defined by a prior launch_thread — use with-block form
+ ExprAST var = printer->VarGet(iv->var).value();
+ ExprAST ctx = ExprCall(TIR("launch_thread"),
+ {var, Print(printer, node->value, path->Attr("value"))});
+ return WithAST(Optional<ExprAST>(), ctx, body);
+ } else {
+ // First definition — use concise assignment form
+ ExprAST var = printer->VarGet(iv->var).value();
+ ExprAST ctx = ExprCall(TIR("launch_thread"),
+ {LiteralAST::Str(iv->thread_tag),
+ Print(printer, node->value, path->Attr("value"))});
+ List<StmtAST> result;
+ result.push_back(AssignAST(var, ctx));
+ for (const auto& s : body) result.push_back(s);
+ return StmtBlockAST(std::move(result));
+ }
+ }
+ List<StmtAST> body = PrintBodyStmts(node->body, printer, path->Attr("body"));
+ ExprAST ctx = ExprCall(TIR("attr"),
+ {Print(printer, node->node, path->Attr("node")),
+ LiteralAST::Str(node->attr_key, {path->Attr("attr_key")}),
+ Print(printer, node->value, path->Attr("value"))});
+ return WithAST(Optional<ExprAST>(), ctx, body);
+ });
+}
+
+} // namespace script_v2
+} // namespace tvm
diff --git a/src/script_v2/printer/tirx/utils.h b/src/script_v2/printer/tirx/utils.h
new file mode 100644
index 0000000000..9c25d6037c
--- /dev/null
+++ b/src/script_v2/printer/tirx/utils.h
@@ -0,0 +1,304 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_V2_PRINTER_TIRX_UTILS_H_
+#define TVM_SCRIPT_V2_PRINTER_TIRX_UTILS_H_
+
+#include <tvm/runtime/device_api.h>
+#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/buffer.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/function.h>
+#include <tvm/tirx/stmt.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <unordered_map>
+
+#include "../utils.h"
+
+namespace tvm {
+namespace script_v2 {
+
+using namespace tirx;
+using text::AssignASTObj;
+
+/*! \brief Define a TIR variable in the printer and return its IdAST. */
+inline ExprAST DefineVar(const tirx::Var& var, const IRPrinter& printer,
+ const AccessPath& path) {
+ if (printer->VarGet(var).has_value()) {
+ return printer->VarGet(var).value();
+ }
+ DefaultFrame frame = printer->frames.back().cast<DefaultFrame>();
+ printer->VarDef(var->name_hint, var, frame);
+ return printer->VarGet(var).value();
+}
+
+/*! \brief Define a TIR variable and return AssignAST with annotation if typed. */
+inline StmtAST DefineVarAssign(const tirx::Var& var, ExprAST rhs,
+ const IRPrinter& printer, const AccessPath& path) {
+ ExprAST id = DefineVar(var, printer, path);
+ Optional<ExprAST> annotation;
+ if (var->type_annotation.defined()) {
+ annotation = Print(printer, var->type_annotation, path->Attr("type_annotation"));
+ }
+ return AssignAST(id, rhs, annotation);
+}
+
+/*! \brief Print a statement body, flattening SeqStmt. Returns a List<StmtAST>. */
+inline List<StmtAST> PrintBodyStmts(const Stmt& stmt, const IRPrinter& printer,
+ const AccessPath& path) {
+ List<StmtAST> result;
+ NodeAST ast = printer->operator()(Any(stmt), path).cast<NodeAST>();
+ if (auto* block = ast.as<StmtBlockASTObj>()) {
+ for (const auto& s : block->stmts) {
+ result.push_back(s);
+ }
+ } else if (ast.get()->IsInstance<StmtASTObj>()) {
+ result.push_back(GetRef<StmtAST>(static_cast<const StmtASTObj*>(ast.get())));
+ } else if (ast.get()->IsInstance<ExprASTObj>()) {
+ result.push_back(ExprStmtAST(
+ GetRef<ExprAST>(static_cast<const ExprASTObj*>(ast.get()))));
+ }
+ return result;
+}
+
+/*! \brief A Var occurrence counter visitor (matches V1 OccurrenceCounter). */
+class OccurrenceCounter : public tirx::StmtExprVisitor {
+ public:
+ int count = 0;
+ const tirx::VarNode* v = nullptr;
+
+ void VisitExpr_(const tirx::VarNode* op) final {
+ if (op == v) ++count;
+ tirx::StmtExprVisitor::VisitExpr_(op);
+ }
+ void VisitStmt_(const tirx::BufferStoreNode* op) final {
+ VisitBuffer(op->buffer.get());
+ tirx::StmtExprVisitor::VisitStmt_(op);
+ }
+ void VisitExpr_(const tirx::BufferLoadNode* op) final {
+ VisitBuffer(op->buffer.get());
+ tirx::StmtExprVisitor::VisitExpr_(op);
+ }
+ void VisitStmt_(const tirx::AllocBufferNode* op) final {
+ VisitBuffer(op->buffer.get());
+ tirx::StmtExprVisitor::VisitStmt_(op);
+ }
+ void VisitStmt_(const tirx::DeclBufferNode* op) final {
+ VisitBuffer(op->buffer.get());
+ tirx::StmtExprVisitor::VisitStmt_(op);
+ }
+ void VisitBuffer(const tirx::BufferNode* buffer) {
+ VisitExpr(buffer->data);
+ for (const PrimExpr& shape_i : buffer->shape) VisitExpr(shape_i);
+ for (const PrimExpr& stride_i : buffer->strides) VisitExpr(stride_i);
+ VisitExpr(buffer->elem_offset);
+ }
+ explicit OccurrenceCounter(const tirx::VarNode* var) { v = var; }
+};
+
+/*! \brief Count how many times a Var occurs in a PrimFunc (params + buffer_map + body). */
+inline int CountVarOccurrence(const tirx::PrimFunc& f, const tirx::Var& v) {
+ OccurrenceCounter counter(v.get());
+ counter(f->body);
+ for (const tirx::Var& param : f->params) {
+ counter(param);
+ }
+ for (const auto& pair : f->buffer_map) {
+ counter(pair.first);
+ counter.VisitBuffer(pair.second.get());
+ }
+ return counter.count;
+}
+
+/*! \brief Check if a buffer is "simple" (can be inlined as param annotation).
+ * Matches V1 IsSimpleBuffer logic exactly. */
+inline bool IsSimpleBuffer(const Buffer& buf) {
+ if (!buf->strides.empty()) return false;
+ for (const PrimExpr& shp_i : buf->shape) {
+ if (!tirx::UndefinedVars(shp_i).empty()) return false;
+ }
+ if (!tirx::UndefinedVars(buf->elem_offset).empty()) {
+ return false;
+ } else if (buf->elem_offset->IsInstance<IntImmNode>()) {
+ IntImm elem_offset = Downcast<IntImm>(buf->elem_offset);
+ if (elem_offset->value != 0) return false;
+ }
+ return buf.scope() == "global" &&
+ buf->data_alignment == runtime::kAllocAlignment &&
+ buf->offset_factor == 1 &&
+ buf->buffer_type == tirx::BufferType::kDefault &&
+ buf->axis_separators.empty();
+}
+
+/*! \brief Print buffer as T.Buffer(shape, dtype) annotation (simple buffer). */
+inline ExprAST PrintBufferAnnotation(const Buffer& buf, const IRPrinter& printer,
+ const AccessPath& path) {
+ List<ExprAST> args;
+ args.push_back(PrintTuple(printer, buf->shape, path->Attr("shape")));
+ args.push_back(LiteralAST::Str(DType2Str(buf->dtype)));
+ return ExprCall(TIR("Buffer"), std::move(args));
+}
+
+/*!
+ * \brief Check if a PrimExpr is a new (undefined) Var in the printer.
+ */
+inline bool IsNewVar(const PrimExpr& e, const IRPrinter& printer) {
+ return e->IsInstance<tirx::VarNode>() && !printer->VarGet(e).has_value();
+}
+
+/*!
+ * \brief Define a new TIR Var if not already defined, and emit
+ * `var_name = T.<dtype>()` into the frame.
+ * Returns the IdAST for the newly defined var.
+ */
+inline ExprAST DefineNewTIRVar(const tirx::Var& var, const IRPrinter& printer,
+ DefaultFrame& frame) {
+ ExprAST var_id = DefineVar(var, printer, AccessPath::Root());
+ std::string dtype_str = DType2Str(var->dtype);
+ ExprAST rhs = ExprCall(TIR(dtype_str), {});
+ frame->stmts.push_back(AssignAST(var_id, rhs, Optional<ExprAST>()));
+ return var_id;
+}
+
+/*!
+ * \brief Define any new vars in a buffer's shape, strides, and elem_offset.
+ * Emits `var = T.int32()` etc. into the frame for each undefined var.
+ * Must be called BEFORE PrintBufferDecl so that all buffer vars are available.
+ */
+inline void DefineBufferVars(const Buffer& buf, const IRPrinter& printer,
+ DefaultFrame& frame) {
+ for (const PrimExpr& e : buf->shape) {
+ if (IsNewVar(e, printer)) {
+ DefineNewTIRVar(Downcast<tirx::Var>(e), printer, frame);
+ }
+ }
+ for (const PrimExpr& e : buf->strides) {
+ if (IsNewVar(e, printer)) {
+ DefineNewTIRVar(Downcast<tirx::Var>(e), printer, frame);
+ }
+ }
+ if (IsNewVar(buf->elem_offset, printer)) {
+ DefineNewTIRVar(Downcast<tirx::Var>(buf->elem_offset), printer, frame);
+ }
+}
+
+/*!
+ * \brief Define the buffer's data variable as `buf_name.data` using VarDefNoName.
+ * This allows references to buf->data to render as `A.data` instead of `A_1`.
+ * Must be called AFTER the buffer itself has been defined via VarDef.
+ */
+inline void DefineBufferDataVar(const Buffer& buf, const IRPrinter& printer) {
+ if (!printer->VarGet(buf->data).has_value()) {
+ ExprAST buf_expr = printer->VarGet(buf).value();
+ // Capture buf_expr in a Function that returns buf_name.data
+ Function creator = Function::FromTyped([buf_expr]() -> ExprAST {
+ return ExprAttr(buf_expr, "data");
+ });
+ printer->VarDefNoName(creator, buf->data,
+ Optional<ObjectRef>(printer->frames.back().cast<ObjectRef>()));
+ }
+}
+
+/*!
+ * \brief Print a buffer declaration call: T.<method>(extra_args..., shape, dtype, kwargs...)
+ *
+ * Matches V1's BufferDecl/BufferCall: positional args are (extra_args..., shape, dtype),
+ * then kwargs for non-default: data, strides, elem_offset, scope, align, offset_factor,
+ * buffer_type, axis_separators.
+ *
+ * NOTE: Call DefineBufferVars() first to define any new shape/stride vars.
+ */
+inline ExprAST PrintBufferDecl(const Buffer& buf, const std::string& method,
+ List<ExprAST> extra_args,
+ const IRPrinter& printer, const AccessPath& path) {
+ // Positional: extra_args, shape, dtype
+ List<ExprAST> args;
+ for (const auto& a : extra_args) args.push_back(a);
+ args.push_back(PrintTuple(printer, buf->shape, path->Attr("shape")));
+ if (DType2Str(buf->dtype) != "float32") {
+ args.push_back(LiteralAST::Str(DType2Str(buf->dtype)));
+ }
+ // Kwargs for non-default fields
+ List<String> kw_keys;
+ List<ExprAST> kw_vals;
+ // strides (skip for alloc_buffer/sblock_alloc_buffer — V1 parser doesn't accept it)
+ if (!buf->strides.empty() &&
+ method != "alloc_buffer" && method != "sblock_alloc_buffer") {
+ kw_keys.push_back(String("strides"));
+ List<ExprAST> stride_elts;
+ for (int i = 0; i < static_cast<int>(buf->strides.size()); ++i) {
+ stride_elts.push_back(Print(printer, buf->strides[i], path->Attr("strides")->ArrayItem(i)));
+ }
+ kw_vals.push_back(TupleAST({}, std::move(stride_elts)));
+ }
+ // elem_offset
+ // V1 logic: if IntImm, print only if non-zero.
+ // If new var, DON'T print elem_offset kwarg (but set needs_print_factor).
+ // If existing var, print elem_offset kwarg.
+ bool needs_print_factor_for_elem_offset = false;
+ if (const auto* int_imm = buf->elem_offset.as<IntImmNode>()) {
+ if (int_imm->value != 0) {
+ kw_keys.push_back(String("elem_offset"));
+ kw_vals.push_back(Print(printer, buf->elem_offset, path->Attr("elem_offset")));
+ }
+ } else if (IsNewVar(buf->elem_offset, printer)) {
+ // New var: don't print elem_offset kwarg, but force offset_factor printing
+ needs_print_factor_for_elem_offset = true;
+ } else {
+ // Existing var: print elem_offset kwarg
+ kw_keys.push_back(String("elem_offset"));
+ kw_vals.push_back(Print(printer, buf->elem_offset, path->Attr("elem_offset")));
+ }
+ // scope
+ if (buf.scope() != "global") {
+ kw_keys.push_back(String("scope"));
+ kw_vals.push_back(LiteralAST::Str(buf.scope()));
+ }
+ // align (data_alignment)
+ if (buf->data_alignment != runtime::kAllocAlignment) {
+ kw_keys.push_back(String("align"));
+ kw_vals.push_back(LiteralAST::Int(buf->data_alignment));
+ }
+ // offset_factor
+ if (needs_print_factor_for_elem_offset || buf->offset_factor != 1) {
+ kw_keys.push_back(String("offset_factor"));
+ kw_vals.push_back(LiteralAST::Int(buf->offset_factor));
+ }
+ // buffer_type
+ if (buf->buffer_type != tirx::BufferType::kDefault) {
+ kw_keys.push_back(String("buffer_type"));
+ kw_vals.push_back(LiteralAST::Str("auto"));
+ }
+ // axis_separators (skip for alloc_buffer/sblock_alloc_buffer — V1 parser doesn't accept it)
+ if (!buf->axis_separators.empty() &&
+ method != "alloc_buffer" && method != "sblock_alloc_buffer") {
+ kw_keys.push_back(String("axis_separators"));
+ kw_vals.push_back(PrintList(printer, buf->axis_separators, path->Attr("axis_separators")));
+ }
+
+ if (!kw_keys.empty()) {
+ return ExprCallKw(TIR(method), std::move(args), std::move(kw_keys), std::move(kw_vals));
+ }
+ return ExprCall(TIR(method), std::move(args));
+}
+
+} // namespace script_v2
+} // namespace tvm
+
+#endif // TVM_SCRIPT_V2_PRINTER_TIRX_UTILS_H_
diff --git a/src/script_v2/printer/utils.h b/src/script_v2/printer/utils.h
new file mode 100644
index 0000000000..7b5954b982
--- /dev/null
+++ b/src/script_v2/printer/utils.h
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_SCRIPT_V2_PRINTER_UTILS_H_
+#define TVM_SCRIPT_V2_PRINTER_UTILS_H_
+
+#include <tvm/ffi/ir/text/ast.h>
+#include <tvm/ffi/ir/text/printer.h>
+#include <tvm/ffi/ir/traits.h>
+#include <tvm/ffi/reflection/access_path.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/dtype.h>
+
+namespace tvm {
+namespace script_v2 {
+
+namespace text = ::tvm::ffi::ir::text;
+namespace traits = ::tvm::ffi::ir::traits;
+namespace refl = ::tvm::ffi::reflection;
+
+using text::AccessPath;
+using text::AssertAST;
+using text::AssignAST;
+using text::CallAST;
+using text::ClassAST;
+using text::CommentAST;
+using text::DefaultFrame;
+using text::DictAST;
+using text::ExprAST;
+using text::ExprAttr;
+using text::ExprCall;
+using text::ExprCallKw;
+using text::ExprIndex;
+using text::ExprStmtAST;
+using text::ForAST;
+using text::FunctionAST;
+using text::IdAST;
+using text::IfAST;
+using text::IRPrinter;
+using text::LambdaAST;
+using text::LiteralAST;
+using text::ListAST;
+using text::NodeAST;
+using text::OperationAST;
+using text::OperationASTObj;
+using text::PrinterConfig;
+using text::ReturnAST;
+using text::SliceAST;
+using text::StmtAST;
+using text::StmtBlockAST;
+using text::TupleAST;
+using text::WhileAST;
+using text::WithAST;
+using text::YieldAST;
+
+using ffi::Any;
+using ffi::AnyView;
+using ffi::Array;
+using ffi::Function;
+using ffi::GetRef;
+using ffi::List;
+using ffi::Map;
+using ffi::ObjectRef;
+using ffi::Optional;
+using ffi::String;
+
+using text::AssignASTObj;
+using text::ExprASTObj;
+using text::FunctionASTObj;
+using text::IfASTObj;
+using text::StmtASTObj;
+using text::StmtBlockASTObj;
+
+/*! \brief Convert a DataType to its string representation. */
+inline std::string DType2Str(const runtime::DataType& dtype) {
+ return dtype.is_void() ? "void" : runtime::DLDataTypeToString(dtype);
+}
+
+/*! \brief Build `prefix.attr` as an IdAST with dot notation. */
+inline ExprAST PrefixedId(const std::string& prefix, const std::string& attr) {
+ return ExprAttr(IdAST(prefix), attr);
+}
+
+/*! \brief Build `T.attr` */
+inline ExprAST TIR(const std::string& attr) { return PrefixedId("T", attr); }
+
+/*! \brief Build `R.attr` */
+inline ExprAST Relax(const std::string& attr) { return PrefixedId("R", attr); }
+
+/*! \brief Build `I.attr` */
+inline ExprAST IR(const std::string& attr) { return PrefixedId("I", attr); }
+
+/*! \brief Print an object through the printer dispatch. */
+inline ExprAST Print(const IRPrinter& printer, Any obj, AccessPath path) {
+ return printer->operator()(std::move(obj), std::move(path)).cast<ExprAST>();
+}
+
+/*! \brief Print an Array as a ListAST. */
+template <typename T>
+ListAST PrintList(const IRPrinter& printer, const Array<T>& arr, const AccessPath& path) {
+ List<ExprAST> elts;
+ for (int i = 0; i < static_cast<int>(arr.size()); ++i) {
+ elts.push_back(Print(printer, arr[i], path->ArrayItem(i)));
+ }
+ return ListAST({}, std::move(elts));
+}
+
+/*! \brief Print an Array as a TupleAST. */
+template <typename T>
+TupleAST PrintTuple(const IRPrinter& printer, const Array<T>& arr, const AccessPath& path) {
+ List<ExprAST> elts;
+ for (int i = 0; i < static_cast<int>(arr.size()); ++i) {
+ elts.push_back(Print(printer, arr[i], path->ArrayItem(i)));
+ }
+ return TupleAST({}, std::move(elts));
+}
+
+} // namespace script_v2
+} // namespace tvm
+
+#endif // TVM_SCRIPT_V2_PRINTER_UTILS_H_
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 8a8f81068c..538c51bac3 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -22,6 +22,7 @@
*/
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
#include <tvm/tirx/analysis.h>
@@ -39,7 +40,26 @@
namespace tvm {
namespace tirx {
-TVM_FFI_STATIC_INIT_BLOCK() { BufferNode::RegisterReflection(); }
+TVM_FFI_STATIC_INIT_BLOCK() {
+ BufferNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+ auto NS = []() { return ffi::Optional<ffi::String>(); };
+ {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ // Buffer → BufferTy
+ auto trait = tr::Trait(ffi::make_object<tr::BufferTyObj>(
+ S("$field:shape"), S("$field:dtype"),
+ OS("$field:strides"), OS("$field:elem_offset"), OS("$field:name")));
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(BufferNode::RuntimeTypeIndex(), &name_arr, raw);
+ }
+}
using IndexMod = tirx::FloorModNode;
using IndexDiv = tirx::FloorDivNode;
diff --git a/src/tirx/ir/expr.cc b/src/tirx/ir/expr.cc
index f4130e70d6..27bda1effb 100644
--- a/src/tirx/ir/expr.cc
+++ b/src/tirx/ir/expr.cc
@@ -21,10 +21,12 @@
* \file expr.cc
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
+#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt_functor.h>
#include <optional>
@@ -70,6 +72,389 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ShuffleNode::RegisterReflection();
CommReducerNode::RegisterReflection();
ReduceNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for expression nodes ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ // Use Any to properly manage ownership
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto binop = [&](int32_t ti, const char* lhs, const char* rhs, const char* op) {
+ register_trait(ti, tr::Trait(ffi::make_object<tr::BinOpObj>(
+ ffi::String(lhs), ffi::String(rhs), ffi::String(op))));
+ };
+ auto val = [&](int32_t ti, const char* name, const char* ty) {
+ register_trait(ti, tr::Trait(ffi::make_object<tr::ValueObj>(
+ ffi::String(name), ffi::Optional<ffi::String>(ffi::String(ty)),
+ ffi::Optional<ffi::String>())));
+ };
+
+ // --- BinOp: all binary operations ---
+ binop(AddNode::RuntimeTypeIndex(), "$field:a", "$field:b", "+");
+ binop(SubNode::RuntimeTypeIndex(), "$field:a", "$field:b", "-");
+ binop(MulNode::RuntimeTypeIndex(), "$field:a", "$field:b", "*");
+ binop(DivNode::RuntimeTypeIndex(), "$field:a", "$field:b", "Div");
+ binop(ModNode::RuntimeTypeIndex(), "$field:a", "$field:b", "truncmod");
+ binop(FloorDivNode::RuntimeTypeIndex(), "$field:a", "$field:b", "//");
+ binop(FloorModNode::RuntimeTypeIndex(), "$field:a", "$field:b", "%");
+ binop(MinNode::RuntimeTypeIndex(), "$field:a", "$field:b", "min");
+ binop(MaxNode::RuntimeTypeIndex(), "$field:a", "$field:b", "max");
+ binop(EQNode::RuntimeTypeIndex(), "$field:a", "$field:b", "==");
+ binop(NENode::RuntimeTypeIndex(), "$field:a", "$field:b", "!=");
+ binop(LTNode::RuntimeTypeIndex(), "$field:a", "$field:b", "<");
+ binop(LENode::RuntimeTypeIndex(), "$field:a", "$field:b", "<=");
+ binop(GTNode::RuntimeTypeIndex(), "$field:a", "$field:b", ">");
+ binop(GENode::RuntimeTypeIndex(), "$field:a", "$field:b", ">=");
+ binop(AndNode::RuntimeTypeIndex(), "$field:a", "$field:b", "and");
+ binop(OrNode::RuntimeTypeIndex(), "$field:a", "$field:b", "or");
+
+ // --- UnaryOp ---
+ register_trait(NotNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::UnaryOpObj>(ffi::String("$field:a"), ffi::String("~"))));
+
+ // --- Value ---
+ val(VarNode::RuntimeTypeIndex(), "$field:name", "$field:type_annotation");
+ val(SizeVarNode::RuntimeTypeIndex(), "$field:name", "$field:type_annotation");
+
+ // --- Literal ---
+ register_trait(StringImmNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LiteralObj>(ffi::String("$field:value"),
+ ffi::Optional<ffi::String>(ffi::String("string")))));
+
+ // --- Load ---
+ register_trait(BufferLoadNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LoadObj>(ffi::String("$field:buffer"),
+ ffi::Optional<ffi::String>(ffi::String("$field:indices")))));
+ register_trait(ProducerLoadNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LoadObj>(ffi::String("$field:producer"),
+ ffi::Optional<ffi::String>(ffi::String("$field:indices")))));
+
+ // --- Call: T.op_name(dtype?, args...) via $method ---
+ register_trait(CallNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(ffi::String("$method:call_callee"),
+ ffi::String("$method:call_args"),
+ ffi::Optional<ffi::String>())));
+
+ // --- Select: T.Select(cond, t, f) via $method ---
+ register_trait(SelectNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(ffi::String("T.Select"),
+ ffi::String("$method:select_args"),
+ ffi::Optional<ffi::String>())));
+
+ // --- Cast: T.Cast(dtype_str, value) via $method ---
+ register_trait(CastNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(ffi::String("T.Cast"),
+ ffi::String("$method:cast_args"),
+ ffi::Optional<ffi::String>())));
+
+ // --- Ramp/Broadcast/Shuffle: Call-like via $method ---
+ register_trait(RampNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("T.Ramp"), ffi::String("$method:ramp_args"),
+ ffi::Optional<ffi::String>())));
+ register_trait(BroadcastNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("T.Broadcast"), ffi::String("$method:broadcast_args"),
+ ffi::Optional<ffi::String>())));
+ register_trait(ShuffleNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("T.Shuffle"), ffi::String("$method:shuffle_args"),
+ ffi::Optional<ffi::String>())));
+
+ // --- IterVar: T.iter_var(var, dom, type_str [, thread_tag]) via $method ---
+ register_trait(IterVarNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("T.iter_var"), ffi::String("$method:iter_var_args"),
+ ffi::Optional<ffi::String>())));
+
+ // --- Let: T.Let(body, where={var: value}) via Call with kwargs ---
+ register_trait(LetNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("T.Let"), ffi::String("$method:let_positional"),
+ ffi::Optional<ffi::String>(),
+ ffi::Optional<ffi::String>(ffi::String("$method:let_kwargs")))));
+
+ // --- CommReducer/Reduce: Call-like ---
+ register_trait(CommReducerNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("comm_reducer"), ffi::String("$field:lhs"),
+ ffi::Optional<ffi::String>())));
+ // --- Reduce: T.reduce(combiner, source=..., init=..., axis=..., condition=..., value_index=...) ---
+ register_trait(ReduceNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ ffi::String("T.reduce"), ffi::String("$method:reduce_positional"),
+ ffi::Optional<ffi::String>(),
+ ffi::Optional<ffi::String>(ffi::String("$method:reduce_kwargs")))));
+}
+
+// --- Register computed-property methods for $method: trait references ---
+// These methods are called by the trait system's $method: resolver when
+// printing expression nodes. We register them directly via
+// TVMFFITypeRegisterMethod to avoid creating a second ObjectDef<T> (which
+// would fail because RegisterReflection was already called above).
+namespace {
+template <typename Func>
+void RegisterComputedMethod(int32_t type_index, const char* type_key, const char* method_name,
+ Func&& func) {
+ using namespace ::tvm::ffi;
+ TVMFFIMethodInfo info;
+ info.name = TVMFFIByteArray{method_name, std::char_traits<char>::length(method_name)};
+ info.doc = TVMFFIByteArray{nullptr, 0};
+ info.flags = 0;
+ std::string full_name = std::string(type_key) + "." + method_name;
+ Function method = Function::FromTyped(std::forward<Func>(func), std::move(full_name));
+ info.method = AnyView(method).CopyToTVMFFIAny();
+ static const char kEmptyMetadata[] = "{}";
+ info.metadata = TVMFFIByteArray{kEmptyMetadata, 2};
+ TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index, &info));
+}
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // Cast: args = [dtype_string, value]
+ RegisterComputedMethod(
+ CastNode::RuntimeTypeIndex(), CastNode::_type_key, "cast_args",
+ [](Cast node) -> ffi::Array<ObjectRef> {
+ StringImm dtype_str(ffi::DLDataTypeToString(node->dtype));
+ return {dtype_str, node->value};
+ });
+
+ // Select: args = [condition, true_value, false_value]
+ RegisterComputedMethod(
+ SelectNode::RuntimeTypeIndex(), SelectNode::_type_key, "select_args",
+ [](Select node) -> ffi::Array<PrimExpr> {
+ return {node->condition, node->true_value, node->false_value};
+ });
+
+ // Ramp: args = [base, stride, lanes]
+ RegisterComputedMethod(
+ RampNode::RuntimeTypeIndex(), RampNode::_type_key, "ramp_args",
+ [](Ramp node) -> ffi::Array<PrimExpr> {
+ return {node->base, node->stride, node->lanes};
+ });
+
+ // Broadcast: args = [value, lanes]
+ RegisterComputedMethod(
+ BroadcastNode::RuntimeTypeIndex(), BroadcastNode::_type_key, "broadcast_args",
+ [](Broadcast node) -> ffi::Array<PrimExpr> {
+ return {node->value, node->lanes};
+ });
+
+ // Shuffle: args = [vectors_array, indices_array]
+ RegisterComputedMethod(
+ ShuffleNode::RuntimeTypeIndex(), ShuffleNode::_type_key, "shuffle_args",
+ [](Shuffle node) -> ffi::Array<ObjectRef> {
+ ffi::Array<ObjectRef> result;
+ result.push_back(node->vectors);
+ result.push_back(node->indices);
+ return result;
+ });
+
+ // IterVar: args = [var, dom, type_str] or [var, dom, type_str, thread_tag]
+ RegisterComputedMethod(
+ IterVarNode::RuntimeTypeIndex(), IterVarNode::_type_key, "iter_var_args",
+ [](IterVar node) -> ffi::Array<ObjectRef> {
+ const char* type_str;
+ switch (static_cast<int>(node->iter_type)) {
+ case kDataPar: type_str = "DataPar"; break;
+ case kThreadIndex: type_str = "ThreadIndex"; break;
+ case kCommReduce: type_str = "CommReduce"; break;
+ case kOrdered: type_str = "Ordered"; break;
+ case kOpaque: type_str = "DimInfo"; break;
+ default: type_str = "Unrolled"; break;
+ }
+ ffi::Array<ObjectRef> result;
+ result.push_back(node->var);
+ result.push_back(node->dom);
+ result.push_back(StringImm(ffi::String(type_str)));
+ if (!node->thread_tag.empty()) {
+ result.push_back(StringImm(ffi::String(std::string(node->thread_tag))));
+ }
+ return result;
+ });
+
+ // tirx.Call: callee = "T." + printer_name
+ // Uses TScriptPrinterName when available, otherwise falls back to short op name.
+ // Note: GlobalVar ops are handled in __ffi_text_print__ override (script_v2/printer/tirx/expr.cc).
+ RegisterComputedMethod(
+ tirx::CallNode::RuntimeTypeIndex(), tirx::CallNode::_type_key, "call_callee",
+ [](tirx::Call node) -> ffi::String {
+ if (auto* op = node->op.as<OpNode>()) {
+ // Check TScriptPrinterName first (matches V1 behavior)
+ static const OpAttrMap<ffi::String> op_names =
+ Op::GetAttrMap<ffi::String>("TScriptPrinterName");
+ Op op_ref = ffi::GetRef<Op>(op);
+ if (op_names.count(op_ref)) {
+ return ffi::String("T." + std::string(op_names[op_ref]));
+ }
+ // Fallback: strip prefix
+ std::string full(op->name);
+ auto pos = full.rfind('.');
+ return ffi::String("T." + ((pos != std::string::npos) ? full.substr(pos + 1) : full));
+ }
+ return ffi::String("T.call");
+ });
+
+ // tirx.Call: args = [dtype_str?, ...args]
+ // Respects TScriptDtypePrintLocation: kNone(0)=no dtype, kFirst(1)=prepend, kLast(2)=append
+ RegisterComputedMethod(
+ tirx::CallNode::RuntimeTypeIndex(), tirx::CallNode::_type_key, "call_args",
+ [](tirx::Call node) -> ffi::Array<ObjectRef> {
+ ffi::Array<ObjectRef> result;
+ // Determine dtype print location from Op attribute (default: kNone)
+ int print_location = static_cast<int>(ScriptDtypePrintLocation::kNone);
+ if (auto* op = node->op.as<OpNode>()) {
+ static const OpAttrMap<TScriptDtypePrintLocation> dtype_locations =
+ Op::GetAttrMap<TScriptDtypePrintLocation>("TScriptDtypePrintLocation");
+ Op op_ref = ffi::GetRef<Op>(op);
+ if (dtype_locations.count(op_ref)) {
+ print_location = dtype_locations[op_ref].IntValue();
+ }
+ }
+ std::string dtype_str = node->dtype.is_void() ? "void"
+ : ffi::DLDataTypeToString(node->dtype);
+ // Check if this is an LLVM intrinsic call (needs name lookup for first arg)
+ bool is_llvm_intrin = false;
+ if (auto* op = node->op.as<OpNode>()) {
+ static const OpAttrMap<ffi::String> op_names =
+ Op::GetAttrMap<ffi::String>("TScriptPrinterName");
+ Op op_ref = ffi::GetRef<Op>(op);
+ if (op_names.count(op_ref)) {
+ std::string name(op_names[op_ref]);
+ is_llvm_intrin = (name == "call_llvm_pure_intrin" || name == "call_llvm_intrin");
+ }
+ }
+ // Prepend dtype if kFirst (V1 always prints dtype for kFirst ops)
+ if (print_location == static_cast<int>(ScriptDtypePrintLocation::kFirst)) {
+ result.push_back(StringImm(ffi::String(dtype_str)));
+ }
+ for (int i = 0; i < static_cast<int>(node->args.size()); ++i) {
+ if (i == 0 && is_llvm_intrin) {
+ // Replace first arg (integer intrinsic ID) with name string
+ auto f_lookup = ffi::Function::GetGlobal("target.llvm_get_intrinsic_name");
+ if (f_lookup.has_value() && node->args[0].as<IntImmNode>()) {
+ int64_t id = node->args[0].as<IntImmNode>()->value;
+ ffi::Any ret;
+ ffi::AnyView args_view[1] = {ffi::AnyView(id)};
+ f_lookup.value().CallPacked(args_view, 1, &ret);
+ ffi::String name = ret.cast<ffi::String>();
+ result.push_back(StringImm(name));
+ } else {
+ result.push_back(node->args[i]);
+ }
+ } else {
+ result.push_back(node->args[i]);
+ }
+ }
+ // Append dtype if kLast (V1 always prints dtype for kLast ops)
+ if (print_location == static_cast<int>(ScriptDtypePrintLocation::kLast)) {
+ result.push_back(StringImm(ffi::String(dtype_str)));
+ }
+ return result;
+ });
+
+ // For: computed_end = min + extent
+ RegisterComputedMethod(
+ ForNode::RuntimeTypeIndex(), ForNode::_type_key, "computed_end",
+ [](tirx::For node) -> PrimExpr { return node->min + node->extent; });
+
+ // For: kind_prefix — returns "T.serial" / "T.parallel" / etc., or "" for simple range
+ RegisterComputedMethod(
+ ForNode::RuntimeTypeIndex(), ForNode::_type_key, "kind_prefix",
+ [](tirx::For node) -> ffi::String {
+ if (node->kind == ForKind::kSerial && node->annotations.empty()) {
+ return ffi::String(""); // empty → PrintFor uses range()
+ }
+ switch (node->kind) {
+ case ForKind::kSerial: return ffi::String("T.serial");
+ case ForKind::kParallel: return ffi::String("T.parallel");
+ case ForKind::kUnrolled: return ffi::String("T.unroll");
+ case ForKind::kVectorized: return ffi::String("T.vectorized");
+ case ForKind::kThreadBinding: return ffi::String("T.thread_binding");
+ default: return ffi::String("T.serial");
+ }
+ });
+
+ // For: extra kwargs (thread="tag" for thread_binding, annotations for non-empty)
+ RegisterComputedMethod(
+ ForNode::RuntimeTypeIndex(), ForNode::_type_key, "for_kwargs",
+ [](tirx::For node) -> ffi::Map<ffi::String, ObjectRef> {
+ ffi::Map<ffi::String, ObjectRef> result;
+ if (!node->annotations.empty()) {
+ result.Set(ffi::String("annotations"), node->annotations);
+ }
+ if (node->kind == ForKind::kThreadBinding && node->thread_binding.defined()) {
+ result.Set(ffi::String("thread"),
+ StringImm(ffi::String(std::string(node->thread_binding.value()->thread_tag))));
+ }
+ return result;
+ });
+
+ // AssertStmt: structured_msg → (error_kind, [message_parts...])
+ RegisterComputedMethod(
+ AssertStmtNode::RuntimeTypeIndex(), AssertStmtNode::_type_key, "structured_msg",
+ [](AssertStmt node) -> ffi::Array<ObjectRef> {
+ ffi::Array<ObjectRef> result;
+ result.push_back(node->error_kind);
+ result.push_back(node->message_parts);
+ return result;
+ });
+
+ // BufferRegion: printed_indices → array of [min, max] pairs for each Range
+ RegisterComputedMethod(
+ BufferRegionNode::RuntimeTypeIndex(), BufferRegionNode::_type_key, "printed_indices",
+ [](BufferRegion node) -> ffi::Array<ObjectRef> {
+ ffi::Array<ObjectRef> result;
+ for (const auto& r : node->region) {
+ ffi::Array<PrimExpr> pair;
+ pair.push_back(r->min);
+ pair.push_back(r->min + r->extent);
+ result.push_back(pair);
+ }
+ return result;
+ });
+
+ // Let: positional args = [body]
+ RegisterComputedMethod(
+ LetNode::RuntimeTypeIndex(), LetNode::_type_key, "let_positional",
+ [](Let node) -> ffi::Array<ObjectRef> {
+ return {node->body};
+ });
+
+ // Let: kwargs = {"where": {var_obj: value_obj}}
+ RegisterComputedMethod(
+ LetNode::RuntimeTypeIndex(), LetNode::_type_key, "let_kwargs",
+ [](Let node) -> ffi::Map<ffi::String, ObjectRef> {
+ ffi::Map<ObjectRef, ObjectRef> where_dict;
+ where_dict.Set(ObjectRef(node->var), ObjectRef(node->value));
+ ffi::Map<ffi::String, ObjectRef> result;
+ result.Set(ffi::String("where"), where_dict);
+ return result;
+ });
+
+ // Reduce: positional args = [combiner]
+ RegisterComputedMethod(
+ ReduceNode::RuntimeTypeIndex(), ReduceNode::_type_key, "reduce_positional",
+ [](Reduce node) -> ffi::Array<ObjectRef> {
+ return {node->combiner};
+ });
+
+ // Reduce: kwargs = {source, init, axis, condition, value_index}
+ RegisterComputedMethod(
+ ReduceNode::RuntimeTypeIndex(), ReduceNode::_type_key, "reduce_kwargs",
+ [](Reduce node) -> ffi::Map<ffi::String, ObjectRef> {
+ ffi::Map<ffi::String, ObjectRef> result;
+ result.Set(ffi::String("source"), node->source);
+ result.Set(ffi::String("init"), node->init);
+ result.Set(ffi::String("axis"), node->axis);
+ result.Set(ffi::String("condition"), node->condition);
+ result.Set(ffi::String("value_index"), IntImm(DataType::Int(32), node->value_index));
+ return result;
+ });
}
/* \brief Convert an object to a PrimExpr
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index b817ccfd73..b9955fed2d 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -22,6 +22,7 @@
* \brief The function data structure.
*/
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/struct_info.h>
#include <tvm/s_tir/analysis.h>
@@ -34,6 +35,32 @@ namespace tirx {
TVM_FFI_STATIC_INIT_BLOCK() {
PrimFuncNode::RegisterReflection();
TensorIntrinNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+ auto NS = []() { return ffi::Optional<ffi::String>(); };
+
+ // PrimFunc → Func (symbol is literal; actual name comes from GlobalVar binding or attrs)
+ register_trait(PrimFuncNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::FuncObj>(
+ S("main"),
+ tr::Region(S("$field:body"), OS("$field:params")),
+ OS("$field:attrs"),
+ NS(), // text_printer_kind
+ NS()))); // text_printer_pre
+ // TensorIntrin → Call-like
+ register_trait(TensorIntrinNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::CallObj>(
+ S("TensorIntrin"), S("$field:desc"), NS())));
}
namespace {
diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc
index fb5b06954c..5cba339881 100644
--- a/src/tirx/ir/stmt.cc
+++ b/src/tirx/ir/stmt.cc
@@ -22,6 +22,7 @@
*/
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/ir/traits.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/op_attr_types.h>
@@ -50,6 +51,88 @@ TVM_FFI_STATIC_INIT_BLOCK() {
MatchBufferRegionNode::RegisterReflection();
SBlockNode::RegisterReflection();
SBlockRealizeNode::RegisterReflection();
+
+ // --- V2 printer: register __ffi_traits__ for statement nodes ---
+ namespace tr = ::tvm::ffi::ir::traits;
+ auto register_trait = [](int32_t type_index, ffi::ObjectRef trait) {
+ static const char kName[] = "__ffi_traits__";
+ TVMFFIByteArray name_arr = {kName, sizeof(kName) - 1};
+ ffi::Any val(trait);
+ TVMFFIAny* raw = reinterpret_cast<TVMFFIAny*>(&val);
+ TVMFFITypeRegisterAttr(type_index, &name_arr, raw);
+ };
+ auto S = [](const char* s) { return ffi::String(s); };
+ auto OS = [](const char* s) { return ffi::Optional<ffi::String>(ffi::String(s)); };
+ auto NS = []() { return ffi::Optional<ffi::String>(); };
+ auto NR = []() { return ffi::Optional<tr::Region>(); };
+
+ // Bind → Assign
+ register_trait(BindNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssignObj>(S("$field:var"), S("$field:value"))));
+ // IfThenElse → If
+ register_trait(IfThenElseNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::IfObj>(
+ S("$field:condition"),
+ tr::Region(S("$field:then_case")),
+ ffi::Optional<tr::Region>(tr::Region(S("$field:else_case"))))));
+ // While → While
+ register_trait(WhileNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::WhileObj>(
+ S("$field:condition"), tr::Region(S("$field:body")))));
+ // BufferStore → Store
+ register_trait(BufferStoreNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::StoreObj>(
+ S("$field:buffer"), S("$field:value"), OS("$field:indices"))));
+ // AssertStmt → Assert with $method:structured_msg → (error_kind, [parts...])
+ register_trait(AssertStmtNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssertObj>(
+ S("$field:condition"), OS("$method:structured_msg"))));
+ // Evaluate → ExprStmt with T.evaluate() wrapper
+ register_trait(EvaluateNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::ExprStmtObj>(S("$field:value"), OS("T.evaluate"))));
+ // SeqStmt → Seq
+ register_trait(SeqStmtNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::SeqObj>(S("$field:seq"))));
+ // AttrStmt → With with $method dispatch for launch_thread vs T.attr
+ register_trait(AttrStmtNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::WithObj>(
+ tr::Region(S("$field:body"), OS("$method:defined_var"), OS("$method:context_args")),
+ NS(), NS(), OS("$method:context_call"), NS(), NS())));
+ // For → For (semantically; __ffi_text_print__ handles printing due to
+ // cross-field logic: kind_prefix + thread_binding + annotations)
+ register_trait(ForNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::ForObj>(
+ tr::Region(S("$field:body"), OS("$field:loop_var")),
+ OS("$field:min"), OS("$method:computed_end"),
+ OS("$field:step"), NS(), NS(),
+ OS("$field:annotations"), OS("$method:kind_prefix"))));
+ // DeclBuffer → AssignCall: buf = T.decl_buffer(shape, dtype)
+ register_trait(DeclBufferNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssignCallObj>(
+ S("$field:buffer"), S("T.decl_buffer"), S("$method:decl_args"))));
+ // AllocBuffer → AssignCall: buf = T.alloc_buffer(shape, dtype=..., scope=...)
+ register_trait(AllocBufferNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssignCallObj>(
+ S("$field:buffer"), S("T.alloc_buffer"), S("$method:alloc_args"),
+ OS("$method:alloc_kwargs"))));
+ // BufferRegion → Load with $method:printed_indices for Range→slice conversion
+ register_trait(BufferRegionNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::LoadObj>(S("$field:buffer"),
+ OS("$method:printed_indices"))));
+ // MatchBufferRegion → AssignCall: buf = T.match_buffer(source, shape, dtype)
+ register_trait(MatchBufferRegionNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::AssignCallObj>(
+ S("$field:buffer"), S("T.match_buffer"), S("$method:match_args"))));
+ // SBlock → With (block scope)
+ register_trait(SBlockNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::WithObj>(
+ tr::Region(S("$field:body"), OS("$field:iter_vars")),
+ NS(), NS(), OS("$field:name_hint"), NS(), NS())));
+ // SBlockRealize → With (realized block scope)
+ register_trait(SBlockRealizeNode::RuntimeTypeIndex(),
+ tr::Trait(ffi::make_object<tr::WithObj>(
+ tr::Region(S("$field:block")),
+ NS(), NS(), NS(), NS(), NS())));
}
// Bind
@@ -608,5 +691,126 @@ TVM_TIR_REGISTER_OP("type_annotation")
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));
+// --- Register computed-property methods for $method: trait references ---
+namespace {
+template <typename Func>
+void RegisterComputedMethod(int32_t type_index, const char* type_key, const char* method_name,
+ Func&& func) {
+ using namespace ::tvm::ffi;
+ TVMFFIMethodInfo info;
+ info.name = TVMFFIByteArray{method_name, std::char_traits<char>::length(method_name)};
+ info.doc = TVMFFIByteArray{nullptr, 0};
+ info.flags = 0;
+ std::string full_name = std::string(type_key) + "." + method_name;
+ Function method = Function::FromTyped(std::forward<Func>(func), std::move(full_name));
+ info.method = AnyView(method).CopyToTVMFFIAny();
+ static const char kEmptyMetadata[] = "{}";
+ info.metadata = TVMFFIByteArray{kEmptyMetadata, 2};
+ TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index, &info));
+}
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ // AssertStmt: structured_msg = [error_kind, [part0, part1, ...]]
+ // PrintAssert renders outer Array as tuple, inner Array as list:
+ // assert cond, ("ValueError", ["msg part 0", "msg part 1"])
+ RegisterComputedMethod(
+ AssertStmtNode::RuntimeTypeIndex(), AssertStmtNode::_type_key, "structured_msg",
+ [](AssertStmt node) -> ffi::Array<ObjectRef> {
+ ffi::Array<ObjectRef> parts_arr;
+ for (const StringImm& part : node->message_parts) {
+ parts_arr.push_back(part);
+ }
+ return {node->error_kind, parts_arr};
+ });
+
+ // AllocBuffer: alloc_args = [shape]
+ RegisterComputedMethod(
+ AllocBufferNode::RuntimeTypeIndex(), AllocBufferNode::_type_key, "alloc_args",
+ [](AllocBuffer node) -> ffi::Array<ObjectRef> {
+ return {node->buffer->shape};
+ });
+
+ // AllocBuffer: alloc_kwargs = {dtype: ..., scope: ...} (elide defaults)
+ RegisterComputedMethod(
+ AllocBufferNode::RuntimeTypeIndex(), AllocBufferNode::_type_key, "alloc_kwargs",
+ [](AllocBuffer node) -> ffi::Map<ffi::String, ObjectRef> {
+ ffi::Map<ffi::String, ObjectRef> result;
+ std::string dtype_str = ffi::DLDataTypeToString(node->buffer->dtype);
+ if (dtype_str != "float32") {
+ result.Set(ffi::String("dtype"), StringImm(ffi::String(dtype_str)));
+ }
+ std::string scope_str(node->buffer.scope());
+ if (scope_str != "global") {
+ result.Set(ffi::String("scope"), StringImm(ffi::String(scope_str)));
+ }
+ return result;
+ });
+
+ // DeclBuffer: decl_args = [shape, dtype_str]
+ RegisterComputedMethod(
+ DeclBufferNode::RuntimeTypeIndex(), DeclBufferNode::_type_key, "decl_args",
+ [](DeclBuffer node) -> ffi::Array<ObjectRef> {
+ return {node->buffer->shape,
+ StringImm(ffi::String(ffi::DLDataTypeToString(node->buffer->dtype)))};
+ });
+
+ // MatchBufferRegion: match_args = [source, shape, dtype_str]
+ RegisterComputedMethod(
+ MatchBufferRegionNode::RuntimeTypeIndex(), MatchBufferRegionNode::_type_key, "match_args",
+ [](MatchBufferRegion node) -> ffi::Array<ObjectRef> {
+ return {node->source, node->buffer->shape,
+ StringImm(ffi::String(ffi::DLDataTypeToString(node->buffer->dtype)))};
+ });
+
+ // AttrStmt: context_call → "T.launch_thread" or "T.attr"
+ RegisterComputedMethod(
+ AttrStmtNode::RuntimeTypeIndex(), AttrStmtNode::_type_key, "context_call",
+ [](AttrStmt node) -> ffi::String {
+ bool is_iter_var =
+ (node->node.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) &&
+ node->node.cast<ObjectRef>()->IsInstance<IterVarNode>();
+ if ((node->attr_key == "thread_extent" || node->attr_key == "virtual_thread") &&
+ is_iter_var) {
+ return ffi::String("T.launch_thread");
+ }
+ return ffi::String("T.attr");
+ });
+
+ // AttrStmt: context_args → args for the context call
+ RegisterComputedMethod(
+ AttrStmtNode::RuntimeTypeIndex(), AttrStmtNode::_type_key, "context_args",
+ [](AttrStmt node) -> ffi::Array<ObjectRef> {
+ bool is_iter_var =
+ (node->node.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) &&
+ node->node.cast<ObjectRef>()->IsInstance<IterVarNode>();
+ if ((node->attr_key == "thread_extent" || node->attr_key == "virtual_thread") &&
+ is_iter_var) {
+ IterVar iv = node->node.cast<IterVar>();
+ return {StringImm(ffi::String(std::string(iv->thread_tag))), node->value};
+ }
+ ffi::Array<ObjectRef> result;
+ result.push_back(node->node.cast<ObjectRef>());
+ result.push_back(StringImm(ffi::String(std::string(node->attr_key))));
+ result.push_back(node->value);
+ return result;
+ });
+
+ // AttrStmt: defined_var → IterVar's var for launch_thread, or None
+ RegisterComputedMethod(
+ AttrStmtNode::RuntimeTypeIndex(), AttrStmtNode::_type_key, "defined_var",
+ [](AttrStmt node) -> ffi::Optional<ObjectRef> {
+ bool is_iter_var =
+ (node->node.type_index() >= ffi::TypeIndex::kTVMFFIStaticObjectBegin) &&
+ node->node.cast<ObjectRef>()->IsInstance<IterVarNode>();
+ if ((node->attr_key == "thread_extent" || node->attr_key == "virtual_thread") &&
+ is_iter_var) {
+ IterVar iv = node->node.cast<IterVar>();
+ return iv->var;
+ }
+ return ffi::Optional<ObjectRef>();
+ });
+}
+
} // namespace tirx
} // namespace tvm