This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch sequal-upgrade in repository https://gitbox.apache.org/repos/asf/tvm.git
commit bc5c4bc7e53ebbb9090a6fa8959e9c5f9111fc0a Author: tqchen <[email protected]> AuthorDate: Fri Jul 25 10:46:15 2025 -0400 [FFI][REFACTOR] Migrate StructuralEqual/Hash to new reflection This PR migrates the StructuralEqual/Hash to new reflection based approach. The original mechanisms are still kept around and we will phase them out in followup PRs. The new mechanism unifies the structural equal/hash registration with the normal reflection registeration and also brings cleaner implementation for mismatch detection. --- ffi/src/ffi/reflection/structural_equal.cc | 14 ++- ffi/src/ffi/reflection/structural_hash.cc | 11 +- include/tvm/arith/analyzer.h | 2 + include/tvm/arith/int_solver.h | 3 + include/tvm/arith/iter_affine_map.h | 3 + include/tvm/ir/attrs.h | 1 + include/tvm/ir/diagnostic.h | 2 + include/tvm/ir/env_func.h | 6 +- include/tvm/ir/expr.h | 10 +- include/tvm/ir/global_info.h | 2 + include/tvm/ir/module.h | 10 ++ include/tvm/ir/op.h | 1 + include/tvm/ir/source_map.h | 3 + include/tvm/ir/type.h | 8 ++ include/tvm/relax/distributed/struct_info.h | 2 + include/tvm/relax/expr.h | 37 +++++-- include/tvm/relax/struct_info.h | 2 +- include/tvm/target/target.h | 2 +- include/tvm/target/target_kind.h | 11 +- include/tvm/te/tensor.h | 1 + include/tvm/tir/buffer.h | 14 +-- include/tvm/tir/expr.h | 9 +- include/tvm/tir/function.h | 10 +- include/tvm/tir/index_map.h | 7 +- include/tvm/tir/stmt.h | 37 +++---- include/tvm/tir/var.h | 6 +- src/contrib/msc/core/ir/graph.h | 4 + src/contrib/msc/core/ir/plugin.h | 4 + src/ir/module.cc | 54 ++++++++++ src/ir/type.cc | 7 +- src/meta_schedule/module_equality.cc | 33 +++--- src/node/structural_equal.cc | 111 +++++++++++++++++---- src/node/structural_hash.cc | 6 +- src/relax/ir/struct_info.cc | 1 + src/relax/transform/lift_transform_params.cc | 7 +- tests/python/ir/test_node_reflection.py | 9 ++ .../test_tvmscript_printer_structural_equal.py | 10 ++ 37 files changed, 351 insertions(+), 109 deletions(-) diff --git a/ffi/src/ffi/reflection/structural_equal.cc b/ffi/src/ffi/reflection/structural_equal.cc index 03cbdd95be..8ce6dba135 100644 --- a/ffi/src/ffi/reflection/structural_equal.cc +++ b/ffi/src/ffi/reflection/structural_equal.cc @@ -90,12 +90,18 @@ class StructEqualHandler { // NOTE: invariant: lhs and rhs are already the same type const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(lhs->type_index()); if (type_info->metadata == nullptr) { - return lhs.same_as(rhs); + TVM_FFI_THROW(TypeError) << "Type metadata is not set for type `" + << String(type_info->type_key) + << "`, so StructuralHash is not supported for this type"; + } + if (type_info->metadata->structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) { + TVM_FFI_THROW(TypeError) << "_type_s_eq_hash_kind is not set for type `" + << String(type_info->type_key) + << "`, so StructuralHash is not supported for this type"; } - auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind; - if (structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported || - structural_eq_hash_kind == kTVMFFISEqHashKindUniqueInstance) { + auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind; + if (structural_eq_hash_kind == kTVMFFISEqHashKindUniqueInstance) { // use pointer comparison return lhs.same_as(rhs); } diff --git a/ffi/src/ffi/reflection/structural_hash.cc b/ffi/src/ffi/reflection/structural_hash.cc index ba47de5146..80e3afe55c 100644 --- a/ffi/src/ffi/reflection/structural_hash.cc +++ b/ffi/src/ffi/reflection/structural_hash.cc @@ -83,9 +83,16 @@ class StructuralHashHandler { // NOTE: invariant: lhs and rhs are already the same type const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(obj->type_index()); if (type_info->metadata == nullptr) { - // Fallback to pointer hash - return std::hash<const Object*>()(obj.get()); + TVM_FFI_THROW(TypeError) << "Type metadata is not set for type `" + << String(type_info->type_key) + << "`, so StructuralHash is not supported for this type"; } + if (type_info->metadata->structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) { + TVM_FFI_THROW(TypeError) << "_type_s_eq_hash_kind is not set for type `" + << String(type_info->type_key) + << "`, so StructuralHash is not supported for this type"; + } + auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind; if (structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) { // Fallback to pointer hash diff --git a/include/tvm/arith/analyzer.h b/include/tvm/arith/analyzer.h index 78eac07f45..54cbab2586 100644 --- a/include/tvm/arith/analyzer.h +++ b/include/tvm/arith/analyzer.h @@ -106,6 +106,7 @@ class ConstIntBoundNode : public Object { */ static const constexpr int64_t kNegInf = -kPosInf; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "arith.ConstIntBound"; TVM_DECLARE_FINAL_OBJECT_INFO(ConstIntBoundNode, Object); }; @@ -222,6 +223,7 @@ class ModularSetNode : public Object { return equal(coeff, other->coeff) && equal(base, other->base); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "arith.ModularSet"; TVM_DECLARE_FINAL_OBJECT_INFO(ModularSetNode, Object); }; diff --git a/include/tvm/arith/int_solver.h b/include/tvm/arith/int_solver.h index dd9259cf97..e2f384b696 100644 --- a/include/tvm/arith/int_solver.h +++ b/include/tvm/arith/int_solver.h @@ -83,6 +83,7 @@ class IntGroupBoundsNode : public Object { hash_reduce(upper); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const char* _type_key = "arith.IntGroupBounds"; TVM_DECLARE_FINAL_OBJECT_INFO(IntGroupBoundsNode, Object); @@ -173,6 +174,7 @@ class IntConstraintsNode : public Object { hash_reduce(relations); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const char* _type_key = "arith.IntConstraints"; TVM_DECLARE_FINAL_OBJECT_INFO(IntConstraintsNode, Object); @@ -238,6 +240,7 @@ class IntConstraintsTransformNode : public Object { hash_reduce(dst_to_src); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const char* _type_key = "arith.IntConstraintsTransform"; TVM_DECLARE_FINAL_OBJECT_INFO(IntConstraintsTransformNode, Object); diff --git a/include/tvm/arith/iter_affine_map.h b/include/tvm/arith/iter_affine_map.h index b7f0e09e83..3c666b430f 100644 --- a/include/tvm/arith/iter_affine_map.h +++ b/include/tvm/arith/iter_affine_map.h @@ -116,6 +116,7 @@ class IterMarkNode : public Object { hash_reduce(extent); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; static constexpr const char* _type_key = "arith.IterMark"; @@ -176,6 +177,7 @@ class IterSplitExprNode : public IterMapExprNode { hash_reduce(scale); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "arith.IterSplitExpr"; TVM_DECLARE_FINAL_OBJECT_INFO(IterSplitExprNode, IterMapExprNode); }; @@ -239,6 +241,7 @@ class IterSumExprNode : public IterMapExprNode { hash_reduce(base); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "arith.IterSumExpr"; TVM_DECLARE_FINAL_OBJECT_INFO(IterSumExprNode, IterMapExprNode); }; diff --git a/include/tvm/ir/attrs.h b/include/tvm/ir/attrs.h index 952cea2a30..b581854e12 100644 --- a/include/tvm/ir/attrs.h +++ b/include/tvm/ir/attrs.h @@ -122,6 +122,7 @@ class BaseAttrsNode : public Object { TVM_DLL virtual void InitByPackedArgs(const ffi::PackedArgs& kwargs, bool allow_unknown = false) = 0; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; static constexpr const char* _type_key = "ir.Attrs"; diff --git a/include/tvm/ir/diagnostic.h b/include/tvm/ir/diagnostic.h index 8429ac1a62..e1d7abbead 100644 --- a/include/tvm/ir/diagnostic.h +++ b/include/tvm/ir/diagnostic.h @@ -79,6 +79,7 @@ class DiagnosticNode : public Object { equal(this->message, other->message); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "Diagnostic"; TVM_DECLARE_FINAL_OBJECT_INFO(DiagnosticNode, Object); }; @@ -214,6 +215,7 @@ class DiagnosticContextNode : public Object { return equal(module, other->module) && equal(diagnostics, other->diagnostics); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "DiagnosticContext"; TVM_DECLARE_FINAL_OBJECT_INFO(DiagnosticContextNode, Object); }; diff --git a/include/tvm/ir/env_func.h b/include/tvm/ir/env_func.h index 03cf3d625a..c1fdeb6d1c 100644 --- a/include/tvm/ir/env_func.h +++ b/include/tvm/ir/env_func.h @@ -51,7 +51,10 @@ class EnvFuncNode : public Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<EnvFuncNode>().def_ro("name", &EnvFuncNode::name); + // func do not participate in structural equal and hash. + refl::ObjectDef<EnvFuncNode>() + .def_ro("name", &EnvFuncNode::name) + .def_ro("func", &EnvFuncNode::func, refl::AttachFieldFlag::SEqHashIgnore()); } bool SEqualReduce(const EnvFuncNode* other, SEqualReducer equal) const { @@ -64,6 +67,7 @@ class EnvFuncNode : public Object { hash_reduce(name); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "ir.EnvFunc"; static constexpr bool _type_has_method_sequal_reduce = true; static constexpr bool _type_has_method_shash_reduce = true; diff --git a/include/tvm/ir/expr.h b/include/tvm/ir/expr.h index 9a8e290cb9..e2fc996ab4 100644 --- a/include/tvm/ir/expr.h +++ b/include/tvm/ir/expr.h @@ -58,11 +58,14 @@ class BaseExprNode : public Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<BaseExprNode>().def_ro("span", &BaseExprNode::span, refl::DefaultValue(Span())); + // span do not participate in structural equal and hash. + refl::ObjectDef<BaseExprNode>().def_ro("span", &BaseExprNode::span, refl::DefaultValue(Span()), + refl::AttachFieldFlag::SEqHashIgnore()); } static constexpr const char* _type_key = "ir.BaseExpr"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; static constexpr const uint32_t _type_child_slots = 64; @@ -428,7 +431,8 @@ class RelaxExprNode : public BaseExprNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<RelaxExprNode>().def_ro("struct_info_", &RelaxExprNode::struct_info_); + refl::ObjectDef<RelaxExprNode>().def_ro("struct_info_", &RelaxExprNode::struct_info_, + refl::AttachFieldFlag::SEqHashIgnore()); } static constexpr const char* _type_key = "ir.RelaxExpr"; @@ -474,6 +478,7 @@ class GlobalVarNode : public RelaxExprNode { hash_reduce.FreeVarHashImpl(this); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; static constexpr const char* _type_key = "ir.GlobalVar"; TVM_DECLARE_FINAL_OBJECT_INFO(GlobalVarNode, RelaxExprNode); }; @@ -711,6 +716,7 @@ class RangeNode : public Object { } static constexpr const char* _type_key = "ir.Range"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(RangeNode, Object); diff --git a/include/tvm/ir/global_info.h b/include/tvm/ir/global_info.h index 4a0b9ffdae..57eadf2b29 100644 --- a/include/tvm/ir/global_info.h +++ b/include/tvm/ir/global_info.h @@ -43,6 +43,8 @@ using MemoryScope = String; class GlobalInfoNode : public Object { public: static constexpr const char* _type_key = "ir.GlobalInfo"; + + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_BASE_OBJECT_INFO(GlobalInfoNode, Object); diff --git a/include/tvm/ir/module.h b/include/tvm/ir/module.h index a5c2477b8d..bde289b5d0 100644 --- a/include/tvm/ir/module.h +++ b/include/tvm/ir/module.h @@ -138,12 +138,21 @@ class IRModuleNode : public Object { .def_ro("source_map", &IRModuleNode::source_map) .def_ro("attrs", &IRModuleNode::attrs) .def_ro("global_infos", &IRModuleNode::global_infos); + // register custom structural equal and hash. + refl::TypeAttrDef<IRModuleNode>() + .def("__s_equal__", &IRModuleNode::SEqual) + .def("__s_hash__", &IRModuleNode::SHash); } TVM_DLL bool SEqualReduce(const IRModuleNode* other, SEqualReducer equal) const; TVM_DLL void SHashReduce(SHashReducer hash_reduce) const; + TVM_DLL bool SEqual(const IRModuleNode* other, + ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const; + TVM_DLL uint64_t SHash(uint64_t type_key_hash, + ffi::TypedFunction<uint64_t(AnyView, bool)> hash) const; + /*! * \brief Add a function to the global environment. * \param var The var of the global function. @@ -237,6 +246,7 @@ class IRModuleNode : public Object { TVM_OBJECT_ENABLE_SCRIPT_PRINTER(); static constexpr const char* _type_key = "ir.IRModule"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindCustomTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(IRModuleNode, Object); diff --git a/include/tvm/ir/op.h b/include/tvm/ir/op.h index a50f12b167..b7b0a1cd9e 100644 --- a/include/tvm/ir/op.h +++ b/include/tvm/ir/op.h @@ -113,6 +113,7 @@ class OpNode : public RelaxExprNode { hash_reduce(name); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance; static constexpr const char* _type_key = "ir.Op"; TVM_DECLARE_FINAL_OBJECT_INFO(OpNode, RelaxExprNode); diff --git a/include/tvm/ir/source_map.h b/include/tvm/ir/source_map.h index 87b3532786..d53c234690 100644 --- a/include/tvm/ir/source_map.h +++ b/include/tvm/ir/source_map.h @@ -59,6 +59,7 @@ class SourceNameNode : public Object { return equal(name, other->name); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "ir.SourceName"; TVM_DECLARE_FINAL_OBJECT_INFO(SourceNameNode, Object); }; @@ -118,6 +119,7 @@ class SpanNode : public Object { equal(end_column, other->end_column); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "ir.Span"; TVM_DECLARE_BASE_OBJECT_INFO(SpanNode, Object); }; @@ -233,6 +235,7 @@ class SourceMapObj : public Object { return equal(source_map, other->source_map); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "ir.SourceMap"; TVM_DECLARE_FINAL_OBJECT_INFO(SourceMapObj, Object); }; diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h index f879ab5911..a2ab74a3ae 100644 --- a/include/tvm/ir/type.h +++ b/include/tvm/ir/type.h @@ -80,6 +80,14 @@ class TypeNode : public Object { */ mutable Span span; + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + // span do not participate in structural equal and hash. + refl::ObjectDef<TypeNode>().def_ro("span", &TypeNode::span, refl::DefaultValue(Span()), + refl::AttachFieldFlag::SEqHashIgnore()); + } + + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "ir.Type"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; diff --git a/include/tvm/relax/distributed/struct_info.h b/include/tvm/relax/distributed/struct_info.h index 3b4a4a0d1d..7f843a9f2c 100644 --- a/include/tvm/relax/distributed/struct_info.h +++ b/include/tvm/relax/distributed/struct_info.h @@ -61,6 +61,7 @@ class PlacementSpecNode : public Object { } static constexpr const char* _type_key = "relax.distributed.PlacementSpec"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_BASE_OBJECT_INFO(PlacementSpecNode, Object); @@ -119,6 +120,7 @@ class PlacementNode : public Object { static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; static constexpr const char* _type_key = "relax.distributed.Placement"; TVM_DECLARE_FINAL_OBJECT_INFO(PlacementNode, Object); }; diff --git a/include/tvm/relax/expr.h b/include/tvm/relax/expr.h index 34aea7981d..f0660f18ac 100644 --- a/include/tvm/relax/expr.h +++ b/include/tvm/relax/expr.h @@ -57,7 +57,8 @@ class IdNode : public Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<IdNode>().def_ro("name_hint", &IdNode::name_hint); + refl::ObjectDef<IdNode>().def_ro("name_hint", &IdNode::name_hint, + refl::AttachFieldFlag::SEqHashIgnore()); } bool SEqualReduce(const IdNode* other, SEqualReducer equal) const { @@ -66,6 +67,7 @@ class IdNode : public Object { void SHashReduce(SHashReducer hash_reduce) const { hash_reduce.FreeVarHashImpl(this); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; static constexpr const char* _type_key = "relax.Id"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; @@ -120,6 +122,13 @@ class StructInfoNode : public Object { */ mutable Span span; + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef<StructInfoNode>().def_ro("span", &StructInfoNode::span, + refl::AttachFieldFlag::SEqHashIgnore()); + } + + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "ir.StructInfo"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; @@ -409,6 +418,7 @@ class VarNode : public LeafExprNode { hash_reduce(struct_info_); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; static constexpr const char* _type_key = "relax.expr.Var"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; @@ -448,6 +458,7 @@ class DataflowVarNode : public VarNode { hash_reduce(struct_info_); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; static constexpr const char* _type_key = "relax.expr.DataflowVar"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; @@ -655,18 +666,19 @@ class DataTypeImm : public LeafExpr { /*! \brief The base class of a variable binding in Relax. */ class BindingNode : public Object { public: + mutable Span span; /*! \brief The return variable to bound to. */ Var var; - mutable Span span; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BindingNode>() - .def_ro("var", &BindingNode::var) - .def_ro("span", &BindingNode::span); + .def_ro("span", &BindingNode::span, refl::AttachFieldFlag::SEqHashIgnore()) + .def_ro("var", &BindingNode::var, refl::AttachFieldFlag::SEqHashDef()); } static constexpr const char* _type_key = "relax.expr.Binding"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_BASE_OBJECT_INFO(BindingNode, Object); @@ -701,7 +713,6 @@ class MatchCastNode : public BindingNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<MatchCastNode>() - .def_ro("var", &MatchCastNode::var) .def_ro("value", &MatchCastNode::value) .def_ro("struct_info", &MatchCastNode::struct_info); } @@ -734,9 +745,7 @@ class VarBindingNode : public BindingNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<VarBindingNode>() - .def_ro("var", &VarBindingNode::var) - .def_ro("value", &VarBindingNode::value); + refl::ObjectDef<VarBindingNode>().def_ro("value", &VarBindingNode::value); } bool SEqualReduce(const VarBindingNode* other, SEqualReducer equal) const; @@ -757,12 +766,15 @@ class VarBinding : public Binding { class BindingBlockNode : public Object { public: - mutable Span span; Array<Binding> bindings; + mutable Span span; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; - refl::ObjectDef<BindingBlockNode>().def_ro("bindings", &BindingBlockNode::bindings); + refl::ObjectDef<BindingBlockNode>() + .def_ro("bindings", &BindingBlockNode::bindings) + .def_ro("span", &BindingBlockNode::span, refl::AttachFieldFlag::SEqHashIgnore(), + refl::DefaultValue(Span())); } bool SEqualReduce(const BindingBlockNode* other, SEqualReducer equal) const { @@ -771,6 +783,7 @@ class BindingBlockNode : public Object { void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(bindings); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "relax.expr.BindingBlock"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; @@ -906,6 +919,7 @@ class IfNode : public ExprNode { hash_reduce(struct_info_); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; static constexpr const char* _type_key = "relax.expr.If"; TVM_DECLARE_FINAL_OBJECT_INFO(IfNode, ExprNode); }; @@ -960,7 +974,7 @@ class FunctionNode : public BaseFuncNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<FunctionNode>() - .def_ro("params", &FunctionNode::params) + .def_ro("params", &FunctionNode::params, refl::AttachFieldFlag::SEqHashDef()) .def_ro("body", &FunctionNode::body) .def_ro("ret_struct_info", &FunctionNode::ret_struct_info) .def_ro("is_pure", &FunctionNode::is_pure); @@ -983,6 +997,7 @@ class FunctionNode : public BaseFuncNode { hash_reduce(struct_info_); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode; static constexpr const char* _type_key = "relax.expr.Function"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; diff --git a/include/tvm/relax/struct_info.h b/include/tvm/relax/struct_info.h index 25a6b1ef4a..cd9b05ab29 100644 --- a/include/tvm/relax/struct_info.h +++ b/include/tvm/relax/struct_info.h @@ -341,7 +341,7 @@ class FuncStructInfoNode : public StructInfoNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<FuncStructInfoNode>() - .def_ro("params", &FuncStructInfoNode::params) + .def_ro("params", &FuncStructInfoNode::params, refl::AttachFieldFlag::SEqHashDef()) .def_ro("ret", &FuncStructInfoNode::ret) .def_ro("derive_func", &FuncStructInfoNode::derive_func) .def_ro("purity", &FuncStructInfoNode::purity); diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h index 98c984c157..9929791f31 100644 --- a/include/tvm/target/target.h +++ b/include/tvm/target/target.h @@ -177,7 +177,7 @@ class TargetNode : public Object { void SHashReduce(SHashReducer hash_reduce) const; static constexpr const char* _type_key = "target.Target"; - + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(TargetNode, Object); diff --git a/include/tvm/target/target_kind.h b/include/tvm/target/target_kind.h index c9785820b4..15b5f62cd5 100644 --- a/include/tvm/target/target_kind.h +++ b/include/tvm/target/target_kind.h @@ -81,12 +81,14 @@ class TargetKindNode : public Object { namespace refl = tvm::ffi::reflection; refl::ObjectDef<TargetKindNode>() .def_ro("name", &TargetKindNode::name) - .def_ro("default_device_type", &TargetKindNode::default_device_type) - .def_ro("default_keys", &TargetKindNode::default_keys); + .def_ro("default_device_type", &TargetKindNode::default_device_type, + refl::AttachFieldFlag::SEqHashIgnore()) + .def_ro("default_keys", &TargetKindNode::default_keys, + refl::AttachFieldFlag::SEqHashIgnore()); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance; static constexpr const char* _type_key = "target.TargetKind"; - TVM_DECLARE_FINAL_OBJECT_INFO(TargetKindNode, Object); private: @@ -134,10 +136,11 @@ class TargetKind : public ObjectRef { * \return The TargetKind requested */ TVM_DLL static Optional<TargetKind> Get(const String& target_kind_name); - TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TargetKind, ObjectRef, TargetKindNode); /*! \brief Mutable access to the container class */ TargetKindNode* operator->() { return static_cast<TargetKindNode*>(data_.get()); } + TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TargetKind, ObjectRef, TargetKindNode); + private: TVM_DLL static const AttrRegistryMapContainerMap<TargetKind>& GetAttrMapContainer( const String& attr_name); diff --git a/include/tvm/te/tensor.h b/include/tvm/te/tensor.h index 15be5bb069..f45a96df63 100644 --- a/include/tvm/te/tensor.h +++ b/include/tvm/te/tensor.h @@ -88,6 +88,7 @@ class TensorNode : public DataProducerNode { TVM_DLL String GetNameHint() const final; static constexpr const char* _type_key = "te.Tensor"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode; TVM_DECLARE_FINAL_OBJECT_INFO(TensorNode, DataProducerNode); }; diff --git a/include/tvm/tir/buffer.h b/include/tvm/tir/buffer.h index 6b49f619b5..cb16d2912a 100644 --- a/include/tvm/tir/buffer.h +++ b/include/tvm/tir/buffer.h @@ -115,13 +115,14 @@ class BufferNode : public Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BufferNode>() - .def_ro("data", &BufferNode::data) + .def_ro("data", &BufferNode::data, refl::AttachFieldFlag::SEqHashDef()) .def_ro("dtype", &BufferNode::dtype) - .def_ro("shape", &BufferNode::shape) - .def_ro("strides", &BufferNode::strides) - .def_ro("axis_separators", &BufferNode::axis_separators) - .def_ro("elem_offset", &BufferNode::elem_offset) - .def_ro("name", &BufferNode::name) + .def_ro("shape", &BufferNode::shape, refl::AttachFieldFlag::SEqHashDef()) + .def_ro("strides", &BufferNode::strides, refl::AttachFieldFlag::SEqHashDef()) + .def_ro("axis_separators", &BufferNode::axis_separators, + refl::AttachFieldFlag::SEqHashDef()) + .def_ro("elem_offset", &BufferNode::elem_offset, refl::AttachFieldFlag::SEqHashDef()) + .def_ro("name", &BufferNode::name, refl::AttachFieldFlag::SEqHashIgnore()) .def_ro("data_alignment", &BufferNode::data_alignment) .def_ro("offset_factor", &BufferNode::offset_factor) .def_ro("buffer_type", &BufferNode::buffer_type) @@ -163,6 +164,7 @@ class BufferNode : public Object { Array<PrimExpr> ElemOffset(Array<PrimExpr> index) const; static constexpr const char* _type_key = "tir.Buffer"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(BufferNode, Object); diff --git a/include/tvm/tir/expr.h b/include/tvm/tir/expr.h index 9525f88784..3e6a07a6cd 100644 --- a/include/tvm/tir/expr.h +++ b/include/tvm/tir/expr.h @@ -833,7 +833,7 @@ class LetNode : public PrimExprNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<LetNode>() - .def_ro("var", &LetNode::var) + .def_ro("var", &LetNode::var, refl::AttachFieldFlag::SEqHashDef()) .def_ro("value", &LetNode::value) .def_ro("body", &LetNode::body); } @@ -989,11 +989,11 @@ class CommReducerNode : public Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<CommReducerNode>() - .def_ro("lhs", &CommReducerNode::lhs) - .def_ro("rhs", &CommReducerNode::rhs) + .def_ro("lhs", &CommReducerNode::lhs, refl::AttachFieldFlag::SEqHashDef()) + .def_ro("rhs", &CommReducerNode::rhs, refl::AttachFieldFlag::SEqHashDef()) .def_ro("result", &CommReducerNode::result) .def_ro("identity_element", &CommReducerNode::identity_element) - .def_ro("span", &CommReducerNode::span); + .def_ro("span", &CommReducerNode::span, refl::AttachFieldFlag::SEqHashIgnore()); } bool SEqualReduce(const CommReducerNode* other, SEqualReducer equal) const { @@ -1009,6 +1009,7 @@ class CommReducerNode : public Object { } static constexpr const char* _type_key = "tir.CommReducer"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(CommReducerNode, Object); diff --git a/include/tvm/tir/function.h b/include/tvm/tir/function.h index ead03e9676..49df3baff8 100644 --- a/include/tvm/tir/function.h +++ b/include/tvm/tir/function.h @@ -49,8 +49,6 @@ class PrimFuncNode : public BaseFuncNode { public: /*! \brief Function parameters */ Array<tir::Var> params; - /*! \brief The body of the function */ - tir::Stmt body; /*! \brief The return type of the function. */ Type ret_type; /*! @@ -99,14 +97,16 @@ class PrimFuncNode : public BaseFuncNode { * flattened alias of the buffer. */ Map<tir::Var, Buffer> buffer_map; + /*! \brief The body of the function */ + tir::Stmt body; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<PrimFuncNode>() - .def_ro("params", &PrimFuncNode::params) - .def_ro("body", &PrimFuncNode::body) + .def_ro("params", &PrimFuncNode::params, refl::AttachFieldFlag::SEqHashDef()) .def_ro("ret_type", &PrimFuncNode::ret_type) - .def_ro("buffer_map", &PrimFuncNode::buffer_map); + .def_ro("buffer_map", &PrimFuncNode::buffer_map) + .def_ro("body", &PrimFuncNode::body); } bool SEqualReduce(const PrimFuncNode* other, SEqualReducer equal) const { diff --git a/include/tvm/tir/index_map.h b/include/tvm/tir/index_map.h index 1cc6fa950a..55d083834d 100644 --- a/include/tvm/tir/index_map.h +++ b/include/tvm/tir/index_map.h @@ -154,9 +154,11 @@ class IndexMapNode : public Object { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<IndexMapNode>() - .def_ro("initial_indices", &IndexMapNode::initial_indices) + .def_ro("initial_indices", &IndexMapNode::initial_indices, + refl::AttachFieldFlag::SEqHashDef()) .def_ro("final_indices", &IndexMapNode::final_indices) - .def_ro("inverse_index_map", &IndexMapNode::inverse_index_map); + .def_ro("inverse_index_map", &IndexMapNode::inverse_index_map, + refl::AttachFieldFlag::SEqHashIgnore()); } bool SEqualReduce(const IndexMapNode* other, SEqualReducer equal) const { @@ -169,6 +171,7 @@ class IndexMapNode : public Object { hash_reduce(final_indices); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "tir.IndexMap"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h index b89fff0032..9d31d25c39 100644 --- a/include/tvm/tir/stmt.h +++ b/include/tvm/tir/stmt.h @@ -54,6 +54,7 @@ class StmtNode : public Object { TVM_OBJECT_ENABLE_SCRIPT_PRINTER(); static constexpr const char* _type_key = "tir.Stmt"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; static constexpr const uint32_t _type_child_slots = 15; @@ -81,7 +82,7 @@ class LetStmtNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<LetStmtNode>() - .def_ro("var", &LetStmtNode::var) + .def_ro("var", &LetStmtNode::var, refl::AttachFieldFlag::SEqHashDef()) .def_ro("value", &LetStmtNode::value) .def_ro("body", &LetStmtNode::body); } @@ -371,7 +372,7 @@ class AllocateNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<AllocateNode>() - .def_ro("buffer_var", &AllocateNode::buffer_var) + .def_ro("buffer_var", &AllocateNode::buffer_var, refl::AttachFieldFlag::SEqHashDef()) .def_ro("dtype", &AllocateNode::dtype) .def_ro("extents", &AllocateNode::extents) .def_ro("condition", &AllocateNode::condition) @@ -460,7 +461,7 @@ class AllocateConstNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<AllocateConstNode>() - .def_ro("buffer_var", &AllocateConstNode::buffer_var) + .def_ro("buffer_var", &AllocateConstNode::buffer_var, refl::AttachFieldFlag::SEqHashDef()) .def_ro("data", &AllocateConstNode::data) .def_ro("irmod_storage_idx", &AllocateConstNode::irmod_storage_idx) .def_ro("dtype", &AllocateConstNode::dtype) @@ -896,7 +897,7 @@ class ForNode : public StmtNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<ForNode>() - .def_ro("loop_var", &ForNode::loop_var) + .def_ro("loop_var", &ForNode::loop_var, refl::AttachFieldFlag::SEqHashDef()) .def_ro("min", &ForNode::min) .def_ro("extent", &ForNode::extent) .def_ro("kind", &ForNode::kind) @@ -1017,6 +1018,7 @@ class BufferRegionNode : public PrimExprConvertibleNode { TVM_DLL PrimExpr ToPrimExpr() const final; static constexpr const char* _type_key = "tir.BufferRegion"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(BufferRegionNode, PrimExprConvertibleNode); @@ -1082,6 +1084,7 @@ class MatchBufferRegionNode : public Object { } static constexpr const char* _type_key = "tir.MatchBufferRegion"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(MatchBufferRegionNode, Object); @@ -1130,8 +1133,12 @@ class BlockNode : public StmtNode { Array<BufferRegion> writes; /*! \brief The name_hint of the block. */ String name_hint; - /*! \brief The body of the block. */ - Stmt body; + /*! \brief The buffer allocated in the block. */ + Array<Buffer> alloc_buffers; + /*! \brief The match buffer regions. */ + Array<MatchBufferRegion> match_buffers; + /*! \brief The annotation of the block. */ + Map<String, ffi::Any> annotations; /*! * \brief The init statement is executed during the first iteration of reduction loops in a * reduction block. The optional init field allows us to represent initialization and @@ -1140,25 +1147,21 @@ class BlockNode : public StmtNode { * Init field is `std::nullopt` if there is no reduction iter_vars */ Optional<Stmt> init; - /*! \brief The buffer allocated in the block. */ - Array<Buffer> alloc_buffers; - /*! \brief The match buffer regions. */ - Array<MatchBufferRegion> match_buffers; - /*! \brief The annotation of the block. */ - Map<String, ffi::Any> annotations; + /*! \brief The body of the block. */ + Stmt body; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<BlockNode>() - .def_ro("iter_vars", &BlockNode::iter_vars) + .def_ro("iter_vars", &BlockNode::iter_vars, refl::AttachFieldFlag::SEqHashDef()) .def_ro("reads", &BlockNode::reads) .def_ro("writes", &BlockNode::writes) - .def_ro("name_hint", &BlockNode::name_hint) - .def_ro("body", &BlockNode::body) - .def_ro("init", &BlockNode::init) + .def_ro("name_hint", &BlockNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore()) .def_ro("alloc_buffers", &BlockNode::alloc_buffers) .def_ro("match_buffers", &BlockNode::match_buffers) - .def_ro("annotations", &BlockNode::annotations); + .def_ro("annotations", &BlockNode::annotations) + .def_ro("init", &BlockNode::init) + .def_ro("body", &BlockNode::body); } bool SEqualReduce(const BlockNode* other, SEqualReducer equal) const { diff --git a/include/tvm/tir/var.h b/include/tvm/tir/var.h index 24c7e6944d..021b6c301a 100644 --- a/include/tvm/tir/var.h +++ b/include/tvm/tir/var.h @@ -64,7 +64,7 @@ class VarNode : public PrimExprNode { static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef<VarNode>() - .def_ro("name", &VarNode::name_hint) + .def_ro("name", &VarNode::name_hint, refl::AttachFieldFlag::SEqHashIgnore()) .def_ro("type_annotation", &VarNode::type_annotation); } @@ -80,6 +80,7 @@ class VarNode : public PrimExprNode { hash_reduce.FreeVarHashImpl(this); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindFreeVar; static constexpr const char* _type_key = "tir.Var"; static constexpr const uint32_t _type_child_slots = 1; TVM_DECLARE_BASE_OBJECT_INFO(VarNode, PrimExprNode); @@ -290,7 +291,7 @@ class IterVarNode : public PrimExprConvertibleNode { namespace refl = tvm::ffi::reflection; refl::ObjectDef<IterVarNode>() .def_ro("dom", &IterVarNode::dom) - .def_ro("var", &IterVarNode::var) + .def_ro("var", &IterVarNode::var, refl::AttachFieldFlag::SEqHashDef()) .def_ro("iter_type", &IterVarNode::iter_type) .def_ro("thread_tag", &IterVarNode::thread_tag); } @@ -308,6 +309,7 @@ class IterVarNode : public PrimExprConvertibleNode { } static constexpr const char* _type_key = "tir.IterVar"; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; TVM_DECLARE_FINAL_OBJECT_INFO(IterVarNode, PrimExprConvertibleNode); diff --git a/src/contrib/msc/core/ir/graph.h b/src/contrib/msc/core/ir/graph.h index 2ee4aeb05b..1d6b1046d9 100644 --- a/src/contrib/msc/core/ir/graph.h +++ b/src/contrib/msc/core/ir/graph.h @@ -400,6 +400,7 @@ class MSCTensorNode : public Object { hash_reduce(prims); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.MSCTensor"; TVM_DECLARE_FINAL_OBJECT_INFO(MSCTensorNode, Object); }; @@ -514,6 +515,7 @@ class BaseJointNode : public Object { hash_reduce(children); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.BaseJoint"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; @@ -600,6 +602,7 @@ class MSCJointNode : public BaseJointNode { hash_reduce(weights); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.MSCJoint"; TVM_DECLARE_FINAL_OBJECT_INFO(MSCJointNode, BaseJointNode); }; @@ -833,6 +836,7 @@ class BaseGraphNode : public Object { hash_reduce(node_names); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.BaseGraph"; static constexpr const bool _type_has_method_sequal_reduce = true; static constexpr const bool _type_has_method_shash_reduce = true; diff --git a/src/contrib/msc/core/ir/plugin.h b/src/contrib/msc/core/ir/plugin.h index b926edf8e5..291a0e196a 100644 --- a/src/contrib/msc/core/ir/plugin.h +++ b/src/contrib/msc/core/ir/plugin.h @@ -290,6 +290,7 @@ class PluginAttrNode : public Object { hash_reduce(describe); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.PluginAttr"; TVM_DECLARE_FINAL_OBJECT_INFO(PluginAttrNode, Object); }; @@ -371,6 +372,7 @@ class PluginTensorNode : public Object { hash_reduce(describe); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.PluginTensor"; TVM_DECLARE_FINAL_OBJECT_INFO(PluginTensorNode, Object); }; @@ -454,6 +456,7 @@ class PluginExternNode : public Object { hash_reduce(describe); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.PluginExtern"; TVM_DECLARE_FINAL_OBJECT_INFO(PluginExternNode, Object); }; @@ -565,6 +568,7 @@ class PluginNode : public Object { hash_reduce(options); } + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; static constexpr const char* _type_key = "msc.core.Plugin"; TVM_DECLARE_FINAL_OBJECT_INFO(PluginNode, Object); }; diff --git a/src/ir/module.cc b/src/ir/module.cc index 9eedbd5e30..e74f41e698 100644 --- a/src/ir/module.cc +++ b/src/ir/module.cc @@ -96,6 +96,30 @@ bool IRModuleNode::SEqualReduce(const IRModuleNode* other, SEqualReducer equal) return true; } +bool IRModuleNode::SEqual(const IRModuleNode* other, + ffi::TypedFunction<bool(AnyView, AnyView, bool, AnyView)> equal) const { + if (!equal(this->attrs, other->attrs, false, "attrs")) { + return false; + } + if (!equal(this->global_infos, other->global_infos, false, "global_infos")) { + return false; + } + + // Define remaps for GlobalVar and GlobalTypeVar based on their string name. + for (const auto& gv : this->GetGlobalVars()) { + if (other->ContainGlobalVar(gv->name_hint)) { + if (!equal(gv, other->GetGlobalVar(gv->name_hint), true, "functions")) return false; + } + } + + // now check the functions with the GlobalVar remappped + if (!equal(this->functions, other->functions, false, "functions")) { + return false; + } + + return true; +} + void IRModuleNode::SHashReduce(SHashReducer hash_reduce) const { using KV = std::tuple<std::string, ObjectRef, ObjectRef>; // hash the functions. @@ -127,6 +151,36 @@ void IRModuleNode::SHashReduce(SHashReducer hash_reduce) const { hash_reduce(this->global_infos); } +uint64_t IRModuleNode::SHash(uint64_t type_key_hash, + ffi::TypedFunction<uint64_t(AnyView, bool)> hash) const { + uint64_t hash_value = type_key_hash; + hash_value = tvm::ffi::details::StableHashCombine(hash_value, hash(this->attrs, false)); + hash_value = tvm::ffi::details::StableHashCombine(hash_value, hash(this->global_infos, false)); + + // hash the functions. + using KV = std::tuple<std::string, ObjectRef, ObjectRef>; + std::vector<KV> temp; + for (const auto& kv : this->functions) { + temp.emplace_back(kv.first->name_hint, kv.first, kv.second); + } + // sort by the hash key of the keys. + std::sort(temp.begin(), temp.end(), + [](const KV& lhs, const KV& rhs) { return std::get<0>(lhs) < std::get<0>(rhs); }); + hash_value = tvm::ffi::details::StableHashCombine(hash_value, static_cast<uint64_t>(temp.size())); + // first need to define the GlobalVar in the order of the keys + for (size_t i = 0; i < temp.size(); ++i) { + hash_value = tvm::ffi::details::StableHashCombine(hash_value, hash(std::get<1>(temp[i]), true)); + } + // hash the name and content + for (size_t i = 0; i < temp.size(); ++i) { + hash_value = + tvm::ffi::details::StableHashCombine(hash_value, hash(std::get<0>(temp[i]), false)); + hash_value = + tvm::ffi::details::StableHashCombine(hash_value, hash(std::get<2>(temp[i]), false)); + } + return hash_value; +} + bool IRModuleNode::ContainGlobalVar(const String& name) const { return global_var_map_.find(name) != global_var_map_.end(); } diff --git a/src/ir/type.cc b/src/ir/type.cc index 4e580356ff..37b251f1f9 100644 --- a/src/ir/type.cc +++ b/src/ir/type.cc @@ -27,6 +27,7 @@ namespace tvm { TVM_FFI_STATIC_INIT_BLOCK({ + TypeNode::RegisterReflection(); PrimTypeNode::RegisterReflection(); PointerTypeNode::RegisterReflection(); TupleTypeNode::RegisterReflection(); @@ -50,8 +51,12 @@ TVM_FFI_STATIC_INIT_BLOCK({ PointerType::PointerType(Type element_type, String storage_scope) { ObjectPtr<PointerTypeNode> n = make_object<PointerTypeNode>(); + if (storage_scope.empty()) { + n->storage_scope = "global"; + } else { + n->storage_scope = std::move(storage_scope); + } n->element_type = std::move(element_type); - n->storage_scope = std::move(storage_scope); data_ = std::move(n); } diff --git a/src/meta_schedule/module_equality.cc b/src/meta_schedule/module_equality.cc index 986233ca49..8901d5fd8d 100644 --- a/src/meta_schedule/module_equality.cc +++ b/src/meta_schedule/module_equality.cc @@ -18,6 +18,8 @@ */ #include "module_equality.h" +#include <tvm/ffi/reflection/structural_equal.h> +#include <tvm/ffi/reflection/structural_hash.h> #include <tvm/ir/module.h> #include <tvm/node/structural_equal.h> #include <tvm/node/structural_hash.h> @@ -37,28 +39,15 @@ class ModuleEqualityStructural : public ModuleEquality { String GetName() const { return "structural"; } }; -class SEqualHandlerIgnoreNDArray : public SEqualHandlerDefault { - public: - SEqualHandlerIgnoreNDArray() : SEqualHandlerDefault(false, nullptr, false) {} - - protected: - bool DispatchSEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, - const Optional<ObjectPathPair>& current_paths) { - if (auto lhs_ptr = lhs.as<runtime::NDArray::Container>(), - rhs_ptr = rhs.as<runtime::NDArray::Container>(); - lhs_ptr && rhs_ptr) { - SEqualReducer reducer(this, nullptr, map_free_vars); - return NDArrayEqual(lhs_ptr, rhs_ptr, reducer, false); - } - return SEqualHandlerDefault::DispatchSEqualReduce(lhs, rhs, map_free_vars, current_paths); - } -}; - class ModuleEqualityIgnoreNDArray : public ModuleEquality { public: - size_t Hash(IRModule mod) const { return SHashHandlerIgnoreNDArray().Hash(mod, false); } + size_t Hash(IRModule mod) const { + return tvm::ffi::reflection::StructuralHash::Hash(mod, /*map_free_vars=*/false, + /*skip_ndarray_content=*/true); + } bool Equal(IRModule lhs, IRModule rhs) const { - return SEqualHandlerIgnoreNDArray().Equal(lhs, rhs, false); + return tvm::ffi::reflection::StructuralEqual::Equal(lhs, rhs, /*map_free_vars=*/false, + /*skip_ndarray_content=*/true); } String GetName() const { return "ignore-ndarray"; } }; @@ -77,8 +66,10 @@ class ModuleEqualityAnchorBlock : public ModuleEquality { auto anchor_block_lhs = tir::FindAnchorBlock(lhs); auto anchor_block_rhs = tir::FindAnchorBlock(rhs); if (anchor_block_lhs && anchor_block_rhs) { - return SEqualHandlerIgnoreNDArray().Equal(GetRef<tir::Block>(anchor_block_lhs), - GetRef<tir::Block>(anchor_block_rhs), false); + return tvm::ffi::reflection::StructuralEqual::Equal(GetRef<tir::Block>(anchor_block_lhs), + GetRef<tir::Block>(anchor_block_rhs), + /*map_free_vars=*/false, + /*skip_ndarray_content=*/true); } return ModuleEqualityIgnoreNDArray().Equal(lhs, rhs); } diff --git a/src/node/structural_equal.cc b/src/node/structural_equal.cc index 43dee2eb3b..5987692a0f 100644 --- a/src/node/structural_equal.cc +++ b/src/node/structural_equal.cc @@ -20,7 +20,9 @@ * \file src/node/structural_equal.cc */ #include <tvm/ffi/function.h> +#include <tvm/ffi/reflection/access_path.h> #include <tvm/ffi/reflection/registry.h> +#include <tvm/ffi/reflection/structural_equal.h> #include <tvm/ir/module.h> #include <tvm/node/functor.h> #include <tvm/node/node.h> @@ -599,34 +601,107 @@ bool SEqualHandlerDefault::DispatchSEqualReduce(const ObjectRef& lhs, const Obje return impl->DispatchSEqualReduce(lhs, rhs, map_free_vars, current_paths); } +Optional<ObjectPathPair> ObjectPathPairFromAccessPathPair( + Optional<ffi::reflection::AccessPathPair> src) { + if (!src.has_value()) return std::nullopt; + auto translate_path = [](ffi::reflection::AccessPath path) { + ObjectPath result = ObjectPath::Root(); + for (const auto& step : path) { + switch (step->kind) { + case ffi::reflection::AccessKind::kObjectField: { + result = result->Attr(step->key.cast<String>()); + break; + } + case ffi::reflection::AccessKind::kArrayIndex: { + result = result->ArrayIndex(step->key.cast<int64_t>()); + break; + } + case ffi::reflection::AccessKind::kMapKey: { + result = result->MapValue(step->key); + break; + } + case ffi::reflection::AccessKind::kArrayIndexMissing: { + result = result->MissingArrayElement(step->key.cast<int64_t>()); + break; + } + case ffi::reflection::AccessKind::kMapKeyMissing: { + result = result->MissingMapEntry(); + break; + } + default: { + LOG(FATAL) << "Invalid access path kind: " << static_cast<int>(step->kind); + break; + } + } + } + return result; + }; + + return ObjectPathPair(translate_path((*src).get<0>()), translate_path((*src).get<1>())); +} + +bool NodeStructuralEqualAdapter(const Any& lhs, const Any& rhs, bool assert_mode, + bool map_free_vars) { + if (assert_mode) { + auto first_mismatch = ObjectPathPairFromAccessPathPair( + ffi::reflection::StructuralEqual::GetFirstMismatch(lhs, rhs, map_free_vars)); + if (first_mismatch.has_value()) { + std::ostringstream oss; + oss << "StructuralEqual check failed, caused by lhs"; + oss << " at " << (*first_mismatch)->lhs_path; + { + // print lhs + PrinterConfig cfg; + cfg->syntax_sugar = false; + cfg->path_to_underline.push_back((*first_mismatch)->lhs_path); + // The TVMScriptPrinter::Script will fallback to Repr printer, + // if the root node to print is not supported yet, + // e.g. Relax nodes, ArrayObj, MapObj, etc. + oss << ":" << std::endl << TVMScriptPrinter::Script(lhs.cast<ObjectRef>(), cfg); + } + oss << std::endl << "and rhs"; + { + // print rhs + oss << " at " << (*first_mismatch)->rhs_path; + { + PrinterConfig cfg; + cfg->syntax_sugar = false; + cfg->path_to_underline.push_back((*first_mismatch)->rhs_path); + // The TVMScriptPrinter::Script will fallback to Repr printer, + // if the root node to print is not supported yet, + // e.g. Relax nodes, ArrayObj, MapObj, etc. + oss << ":" << std::endl << TVMScriptPrinter::Script(rhs.cast<ObjectRef>(), cfg); + } + } + TVM_FFI_THROW(ValueError) << oss.str(); + } + return true; + } else { + return ffi::reflection::StructuralEqual::Equal(lhs, rhs, map_free_vars); + } +} + TVM_FFI_STATIC_INIT_BLOCK({ namespace refl = tvm::ffi::reflection; refl::GlobalDef() - .def("node.StructuralEqual", - [](const Any& lhs, const Any& rhs, bool assert_mode, bool map_free_vars) { - // If we are asserting on failure, then the `defer_fails` option - // should be enabled, to provide better error messages. For - // example, if the number of bindings in a `relax::BindingBlock` - // differs, highlighting the first difference rather than the - // entire block. - bool defer_fails = assert_mode; - Optional<ObjectPathPair> first_mismatch; - return SEqualHandlerDefault(assert_mode, &first_mismatch, defer_fails) - .Equal(lhs, rhs, map_free_vars); - }) + .def("node.StructuralEqual", NodeStructuralEqualAdapter) .def("node.GetFirstStructuralMismatch", [](const Any& lhs, const Any& rhs, bool map_free_vars) { - Optional<ObjectPathPair> first_mismatch; - bool equal = - SEqualHandlerDefault(false, &first_mismatch, true).Equal(lhs, rhs, map_free_vars); - ICHECK(equal == !first_mismatch.defined()); - return first_mismatch; + /* + Optional<ObjectPathPair> first_mismatch; + bool equal = + SEqualHandlerDefault(false, &first_mismatch, true).Equal(lhs, rhs, map_free_vars); + ICHECK(equal == !first_mismatch.defined()); + return first_mismatch; + */ + return ObjectPathPairFromAccessPathPair( + ffi::reflection::StructuralEqual::GetFirstMismatch(lhs, rhs, map_free_vars)); }); }); bool StructuralEqual::operator()(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_params) const { - return SEqualHandlerDefault(false, nullptr, false).Equal(lhs, rhs, map_free_params); + return ffi::reflection::StructuralEqual::Equal(lhs, rhs, map_free_params); } bool NDArrayEqual(const runtime::NDArray::Container* lhs, const runtime::NDArray::Container* rhs, diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc index 2a60754f7b..6fb8d36784 100644 --- a/src/node/structural_hash.cc +++ b/src/node/structural_hash.cc @@ -22,6 +22,7 @@ #include <dmlc/memory_io.h> #include <tvm/ffi/function.h> #include <tvm/ffi/reflection/registry.h> +#include <tvm/ffi/reflection/structural_hash.h> #include <tvm/node/functor.h> #include <tvm/node/node.h> #include <tvm/node/object_path.h> @@ -296,13 +297,12 @@ TVM_FFI_STATIC_INIT_BLOCK({ namespace refl = tvm::ffi::reflection; refl::GlobalDef().def("node.StructuralHash", [](const Any& object, bool map_free_vars) -> int64_t { - uint64_t hashed_value = SHashHandlerDefault().Hash(object, map_free_vars); - return static_cast<int64_t>(hashed_value); + return ffi::reflection::StructuralHash::Hash(object, map_free_vars); }); }); uint64_t StructuralHash::operator()(const ObjectRef& object) const { - return SHashHandlerDefault().Hash(object, false); + return ffi::reflection::StructuralHash::Hash(object, false); } void SHashHandlerIgnoreNDArray::DispatchSHash(const ObjectRef& object, bool map_free_vars) { diff --git a/src/relax/ir/struct_info.cc b/src/relax/ir/struct_info.cc index 00414fe339..499dd47bb7 100644 --- a/src/relax/ir/struct_info.cc +++ b/src/relax/ir/struct_info.cc @@ -31,6 +31,7 @@ namespace tvm { namespace relax { TVM_FFI_STATIC_INIT_BLOCK({ + StructInfoNode::RegisterReflection(); ObjectStructInfoNode::RegisterReflection(); PrimStructInfoNode::RegisterReflection(); ShapeStructInfoNode::RegisterReflection(); diff --git a/src/relax/transform/lift_transform_params.cc b/src/relax/transform/lift_transform_params.cc index 1a426ec5da..83d978f27d 100644 --- a/src/relax/transform/lift_transform_params.cc +++ b/src/relax/transform/lift_transform_params.cc @@ -23,6 +23,7 @@ */ #include <tvm/ffi/reflection/registry.h> +#include <tvm/ffi/reflection/structural_equal.h> #include <tvm/relax/analysis.h> #include <tvm/relax/expr.h> #include <tvm/relax/expr_functor.h> @@ -540,8 +541,8 @@ class ParamRemapper : private ExprFunctor<void(const Expr&, const Expr&)> { } else { var_remap_.Set(GetRef<Var>(lhs_var), rhs_var); } - CHECK(structural_equal.Equal(lhs_var->struct_info_, rhs_var->struct_info_, - /*map_free_vars=*/true)) + CHECK(tvm::ffi::reflection::StructuralEqual::Equal(lhs_var->struct_info_, rhs_var->struct_info_, + /*map_free_vars=*/true)) << "The struct info of the parameters should be the same for all target functions"; auto lhs_tir_vars = DefinableTIRVarsInStructInfo(GetStructInfo(GetRef<Var>(lhs_var))); auto rhs_tir_vars = DefinableTIRVarsInStructInfo(GetStructInfo(rhs_expr)); @@ -555,8 +556,6 @@ class ParamRemapper : private ExprFunctor<void(const Expr&, const Expr&)> { } } - SEqualHandlerDefault structural_equal{/*assert_mode=*/false, /*first_mismatch=*/nullptr, - /*defer_fail=*/false}; Map<Var, Expr> var_remap_; Map<tir::Var, PrimExpr> tir_var_remap_; }; diff --git a/tests/python/ir/test_node_reflection.py b/tests/python/ir/test_node_reflection.py index 741e61b2eb..be00bc3a47 100644 --- a/tests/python/ir/test_node_reflection.py +++ b/tests/python/ir/test_node_reflection.py @@ -181,6 +181,15 @@ def test_ndarray_dict(): tvm.ir.assert_structural_equal(m1, m2) +def test_free_var_equal(): + x = tvm.tir.Var("x", dtype="int32") + y = tvm.tir.Var("y", dtype="int32") + z = tvm.tir.Var("z", dtype="int32") + v1 = x + y + v1 = y + z + tvm.ir.assert_structural_equal(x, z, map_free_vars=True) + + def test_alloc_const(): dev = tvm.cpu(0) dtype = "float32" diff --git a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py index 58d9402e6f..bbf95801ed 100644 --- a/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py +++ b/tests/python/tvmscript/test_tvmscript_printer_structural_equal.py @@ -45,6 +45,9 @@ def test_prim_func_buffer_map(): A = T.match_buffer(a, (128, 128)) B = T.match_buffer(b, (128, 256)) + func1 = func1.with_attr("global_symbol", "main") + func2 = func2.with_attr("global_symbol", "main") + with pytest.raises(ValueError) as ve: assert_structural_equal(func1, func2) assert _error_message(ve.value) == _expected_result( @@ -109,8 +112,12 @@ def test_allocate(): a_data = T.allocate((256, 128), dtype="float32") a = T.decl_buffer((256, 128), dtype="float32", data=a_data) + func1 = func1.with_attr("global_symbol", "main") + func2 = func2.with_attr("global_symbol", "main") + with pytest.raises(ValueError) as ve: assert_structural_equal(func1, func2) + assert _error_message(ve.value) == _expected_result( func1, func2, @@ -132,6 +139,9 @@ def test_for(): with T.block(): pass + func1 = func1.with_attr("global_symbol", "main") + func2 = func2.with_attr("global_symbol", "main") + with pytest.raises(ValueError) as ve: assert_structural_equal(func1, func2) assert _error_message(ve.value) == _expected_result(
