[do not merge] tvmscript printer v2

.gitignore

@@ -1,3 +1,10 @@
+.claude/
+ir_testsuite.jsonl
+ir_testsuite_filtered.jsonl
+main.py
+main2.py
+testsuite/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

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

include/tvm/ir/attrs.h

@@ -32,6 +32,7 @@
 #include <tvm/ffi/extra/structural_equal.h>
 #include <tvm/ffi/extra/structural_hash.h>
 #include <tvm/ffi/function.h>
+#include <tvm/ffi/extra/ir_traits.h>
 #include <tvm/ffi/reflection/accessor.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/expr.h>
@@ -145,7 +146,10 @@ class DictAttrsNode : public BaseAttrsNode {
 
   static void RegisterReflection() {
     namespace rfl = ffi::reflection;
-    rfl::ObjectDef<DictAttrsNode>().def_ro("__dict__", &DictAttrsNode::dict);
+    namespace tr = tvm::ffi::ir_traits;
+    rfl::ObjectDef<DictAttrsNode>()
+        .def_ro("__dict__", &DictAttrsNode::dict)
+        .def_ir_traits<tr::LiteralTraitsObj>("$field:__dict__");
   }
 
   void InitByPackedArgs(const ffi::PackedArgs& args, bool allow_unknown) final;

include/tvm/ir/expr.h

@@ -24,6 +24,7 @@
 #ifndef TVM_IR_EXPR_H_
 #define TVM_IR_EXPR_H_
 
+#include <tvm/ffi/extra/ir_traits.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ffi/string.h>
 #include <tvm/ir/source_map.h>
@@ -459,7 +460,10 @@ class GlobalVarNode : public RelaxExprNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<GlobalVarNode>().def_ro("name_hint", &GlobalVarNode::name_hint);
+    namespace tr = tvm::ffi::ir_traits;
+    refl::ObjectDef<GlobalVarNode>()
+        .def_ro("name_hint", &GlobalVarNode::name_hint)
+        .def_ir_traits<tr::CallTraitsObj>("I.GlobalVar", "$field:name_hint");
   }
 
   bool SEqual(const GlobalVarNode* other,
@@ -498,7 +502,10 @@ class IntImmNode : public PrimExprNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<IntImmNode>().def_ro("value", &IntImmNode::value);
+    namespace tr = tvm::ffi::ir_traits;
+    refl::ObjectDef<IntImmNode>()
+        .def_ro("value", &IntImmNode::value)
+        .def_ir_traits<tr::LiteralTraitsObj>("$field:value", "int");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.IntImm", IntImmNode, PrimExprNode);
 };
@@ -533,7 +540,10 @@ class FloatImmNode : public PrimExprNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<FloatImmNode>().def_ro("value", &FloatImmNode::value);
+    namespace tr = tvm::ffi::ir_traits;
+    refl::ObjectDef<FloatImmNode>()
+        .def_ro("value", &FloatImmNode::value)
+        .def_ir_traits<tr::LiteralTraitsObj>("$field:value", "float");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.FloatImm", FloatImmNode, PrimExprNode);
 };
@@ -675,10 +685,12 @@ class RangeNode : public Object {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
+    namespace tr = tvm::ffi::ir_traits;
     refl::ObjectDef<RangeNode>()
         .def_ro("min", &RangeNode::min)
         .def_ro("extent", &RangeNode::extent)
-        .def_ro("span", &RangeNode::span, refl::AttachFieldFlag::SEqHashIgnore());
+        .def_ro("span", &RangeNode::span, refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ir_traits<tr::CallTraitsObj>("I.Range", "$global:ir._range_args");
   }
 
   static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;

include/tvm/ir/global_info.h

@@ -25,6 +25,7 @@
 #ifndef TVM_IR_GLOBAL_INFO_H_
 #define TVM_IR_GLOBAL_INFO_H_
 
+#include <tvm/ffi/extra/ir_traits.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/expr.h>
 #include <tvm/target/target.h>
@@ -71,10 +72,12 @@ class VDeviceNode : public GlobalInfoNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
+    namespace tr = tvm::ffi::ir_traits;
     refl::ObjectDef<VDeviceNode>()
         .def_ro("target", &VDeviceNode::target)
         .def_ro("vdevice_id", &VDeviceNode::vdevice_id)
-        .def_ro("memory_scope", &VDeviceNode::memory_scope);
+        .def_ro("memory_scope", &VDeviceNode::memory_scope)
+        .def_ir_traits<tr::CallTraitsObj>("I.vdevice", "$global:ir._vdevice_args");
   }
 
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.VDevice", VDeviceNode, GlobalInfoNode);
@@ -97,7 +100,9 @@ class DummyGlobalInfoNode : public GlobalInfoNode {
  public:
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<DummyGlobalInfoNode>();
+    namespace tr = tvm::ffi::ir_traits;
+    refl::ObjectDef<DummyGlobalInfoNode>()
+        .def_ir_traits<tr::CallTraitsObj>("I.dummy_global_info", "$global:ir._dummy_args");
   }
 
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.DummyGlobalInfo", DummyGlobalInfoNode, GlobalInfoNode);

include/tvm/ir/op.h

@@ -26,6 +26,7 @@
 #define TVM_IR_OP_H_
 
 #include <tvm/ffi/function.h>
+#include <tvm/ffi/extra/ir_traits.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/attrs.h>
 #include <tvm/ir/env_func.h>
@@ -93,14 +94,16 @@ class OpNode : public RelaxExprNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
+    namespace tr = tvm::ffi::ir_traits;
     refl::ObjectDef<OpNode>()
         .def_ro("name", &OpNode::name)
         .def_ro("op_type", &OpNode::op_type, refl::AttachFieldFlag::SEqHashIgnore())
         .def_ro("description", &OpNode::description, refl::AttachFieldFlag::SEqHashIgnore())
         .def_ro("arguments", &OpNode::arguments, refl::AttachFieldFlag::SEqHashIgnore())
         .def_ro("attrs_type_key", &OpNode::attrs_type_key, refl::AttachFieldFlag::SEqHashIgnore())
         .def_ro("num_inputs", &OpNode::num_inputs, refl::AttachFieldFlag::SEqHashIgnore())
-        .def_ro("support_level", &OpNode::support_level, refl::AttachFieldFlag::SEqHashIgnore());
+        .def_ro("support_level", &OpNode::support_level, refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ir_traits<tr::CallTraitsObj>("I.Op", "$field:name");
   }
 
   static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;

include/tvm/ir/type.h

@@ -50,6 +50,7 @@
 #define TVM_IR_TYPE_H_
 
 #include <tvm/ffi/container/array.h>
+#include <tvm/ffi/extra/ir_traits.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/source_map.h>
 #include <tvm/runtime/data_type.h>
@@ -118,7 +119,10 @@ class PrimTypeNode : public TypeNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<PrimTypeNode>().def_ro("dtype", &PrimTypeNode::dtype);
+    namespace tr = tvm::ffi::ir_traits;
+    refl::ObjectDef<PrimTypeNode>()
+        .def_ro("dtype", &PrimTypeNode::dtype)
+        .def_ir_traits<tr::PrimTyTraitsObj>("$field:dtype");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PrimType", PrimTypeNode, TypeNode);
 };
@@ -162,9 +166,11 @@ class PointerTypeNode : public TypeNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
+    namespace tr = tvm::ffi::ir_traits;
     refl::ObjectDef<PointerTypeNode>()
         .def_ro("element_type", &PointerTypeNode::element_type)
-        .def_ro("storage_scope", &PointerTypeNode::storage_scope);
+        .def_ro("storage_scope", &PointerTypeNode::storage_scope)
+        .def_ir_traits<tr::CallTraitsObj>("T.handle", "$global:ir._handle_args");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PointerType", PointerTypeNode, TypeNode);
 };
@@ -198,9 +204,11 @@ class TupleTypeNode : public TypeNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
+    namespace tr = tvm::ffi::ir_traits;
     refl::ObjectDef<TupleTypeNode>()
         .def_ro("fields", &TupleTypeNode::fields)
-        .def_ro("span", &TupleTypeNode::span);
+        .def_ro("span", &TupleTypeNode::span)
+        .def_ir_traits<tr::TupleTyTraitsObj>("$field:fields");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.TupleType", TupleTypeNode, TypeNode);
 };
@@ -258,10 +266,12 @@ class FuncTypeNode : public TypeNode {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
+    namespace tr = tvm::ffi::ir_traits;
     refl::ObjectDef<FuncTypeNode>()
         .def_ro("arg_types", &FuncTypeNode::arg_types)
         .def_ro("ret_type", &FuncTypeNode::ret_type)
-        .def_ro("span", &FuncTypeNode::span);
+        .def_ro("span", &FuncTypeNode::span)
+        .def_ir_traits<tr::FuncTyTraitsObj>("$field:arg_types", "$field:ret_type");
   }
   TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.FuncType", FuncTypeNode, TypeNode);
 };

include/tvm/relax/distributed/struct_info.h

@@ -25,6 +25,7 @@
 #ifndef TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
 #define TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
 
+#include <tvm/ffi/extra/ir_traits.h>
 #include <tvm/relax/distributed/global_info.h>
 #include <tvm/relax/struct_info.h>
 namespace tvm {
@@ -91,7 +92,10 @@ class PlacementNode : public Object {
 
   static void RegisterReflection() {
     namespace refl = tvm::ffi::reflection;
-    refl::ObjectDef<PlacementNode>().def_ro("dim_specs", &PlacementNode::dim_specs);
+    namespace tr = tvm::ffi::ir_traits;
+    refl::ObjectDef<PlacementNode>()
+        .def_ro("dim_specs", &PlacementNode::dim_specs)
+        .def_ir_traits<tr::LiteralTraitsObj>("$global:relax._placement_str");
   }
 
   static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;