From bb751ba76293a49cd5101c75769077df6b53e02f Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Thu, 6 Oct 2022 18:24:23 +0900 Subject: [PATCH 1/3] [Node] Expose StructuralEqual/Hash handler implemenation to header --- include/tvm/node/structural_equal.h | 42 +++++++++++ include/tvm/node/structural_hash.h | 38 ++++++++++ src/node/structural_equal.cc | 109 +++++++++++++++++++--------- src/node/structural_hash.cc | 65 ++++++++++++----- 4 files changed, 201 insertions(+), 53 deletions(-) diff --git a/include/tvm/node/structural_equal.h b/include/tvm/node/structural_equal.h index b51021fe4076..1260f2ecf9e3 100644 --- a/include/tvm/node/structural_equal.h +++ b/include/tvm/node/structural_equal.h @@ -324,5 +324,47 @@ class SEqualReducer { bool map_free_vars_ = false; }; +/*! \brief The default handler for equality testing. + * + * Users can derive from this class and override the DispatchSEqualReduce method, + * to customize equality testing. + */ +class SEqualHandlerDefault : public SEqualReducer::Handler { + public: + SEqualHandlerDefault(bool assert_mode, Optional* first_mismatch); + virtual ~SEqualHandlerDefault(); + + virtual bool SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, + const Optional& current_paths) override; + virtual void DeferFail(const ObjectPathPair& mismatch_paths) override; + virtual ObjectRef MapLhsToRhs(const ObjectRef& lhs) override; + virtual void MarkGraphNode() override; + + /*! + * \brief The entry point for equality testing + * \param lhs The left operand. + * \param rhs The right operand. + * \param map_free_vars Whether or not to remap variables if possible. + * \return The equality result. + */ + virtual bool Equal(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars); + + protected: + /*! + * \brief The dispatcher for equality testing of intermediate objects + * \param lhs The left operand. + * \param rhs The right operand. + * \param map_free_vars Whether or not to remap variables if possible. + * \param current_paths Optional paths to `lhs` and `rhs` objects, for error traceability. + * \return The equality result. + */ + virtual bool DispatchSEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, + const Optional& current_paths); + + private: + class Impl; + Impl* impl; +}; + } // namespace tvm #endif // TVM_NODE_STRUCTURAL_EQUAL_H_ diff --git a/include/tvm/node/structural_hash.h b/include/tvm/node/structural_hash.h index a30a2c59d0d1..591804395025 100644 --- a/include/tvm/node/structural_hash.h +++ b/include/tvm/node/structural_hash.h @@ -200,6 +200,44 @@ class SHashReducer { bool map_free_vars_; }; +/*! \brief The default handler for hash key computation + * + * Users can derive from this class and override the DispatchSHash method, + * to customize hashing. + */ +class SHashHandlerDefault : public SHashReducer::Handler { + public: + SHashHandlerDefault(); + virtual ~SHashHandlerDefault(); + + virtual void SHashReduceHashedValue(size_t hashed_value) override; + virtual void SHashReduce(const ObjectRef& key, bool map_free_vars) override; + virtual void SHashReduceFreeVar(const runtime::Object* var, bool map_free_vars) override; + virtual bool LookupHashedValue(const ObjectRef& key, size_t* hashed_value) override; + virtual void MarkGraphNode() override; + + /*! + * \brief The entry point for hashing + * \param object The object to be hashed. + * \param map_free_vars Whether or not to remap variables if possible. + * \return The hash result. + */ + virtual size_t Hash(const ObjectRef& object, bool map_free_vars); + + protected: + /*! + * \brief The dispatcher for hashing of intermediate objects + * \param object An intermediate object to be hashed. + * \param map_free_vars Whether or not to remap variables if possible. + * \return The hash result. + */ + virtual void DispatchSHash(const ObjectRef& object, bool map_free_vars); + + private: + class Impl; + Impl* impl; +}; + class SEqualReducer; struct NDArrayContainerTrait { static constexpr const std::nullptr_t VisitAttrs = nullptr; diff --git a/src/node/structural_equal.cc b/src/node/structural_equal.cc index 01874c0536ae..2f49d9ef5629 100644 --- a/src/node/structural_equal.cc +++ b/src/node/structural_equal.cc @@ -198,13 +198,13 @@ bool SEqualReducer::ObjectAttrsEqual(const ObjectRef& lhs, const ObjectRef& rhs, * The order of SEqual being called is the same as the order as if we * eagerly do recursive calls in SEqualReduce. */ -class RemapVarSEqualHandler : public SEqualReducer::Handler { +class SEqualHandlerDefault::Impl { public: - explicit RemapVarSEqualHandler(bool assert_mode, Optional* first_mismatch) - : assert_mode_(assert_mode), first_mismatch_(first_mismatch) {} + Impl(SEqualHandlerDefault* parent, bool assert_mode, Optional* first_mismatch) + : parent_(parent), assert_mode_(assert_mode), first_mismatch_(first_mismatch) {} bool SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, - const Optional& current_paths) final { + const Optional& current_paths) { // We cannot use check lhs.same_as(rhs) to check equality. // if we choose to enable var remapping. // @@ -239,17 +239,17 @@ class RemapVarSEqualHandler : public SEqualReducer::Handler { return CheckResult(run(), lhs, rhs, current_paths); } - void DeferFail(const ObjectPathPair& mismatch_paths) final { + void DeferFail(const ObjectPathPair& mismatch_paths) { pending_tasks_.emplace_back(Task::ForceFailTag{}, mismatch_paths); } - void MarkGraphNode() final { + void MarkGraphNode() { // need to push to pending tasks in this case ICHECK(!allow_push_to_stack_ && !task_stack_.empty()); task_stack_.back().graph_equal = true; } - ObjectRef MapLhsToRhs(const ObjectRef& lhs) final { + ObjectRef MapLhsToRhs(const ObjectRef& lhs) { auto it = equal_map_lhs_.find(lhs); if (it != equal_map_lhs_.end()) return it->second; return ObjectRef(nullptr); @@ -279,7 +279,35 @@ class RemapVarSEqualHandler : public SEqualReducer::Handler { return RunTasks(); } + // The default equal as registered in the structural equal vtable. + bool DispatchSEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, + const Optional& current_paths) { + auto compute = [=]() { + ICHECK(lhs.defined() && rhs.defined() && lhs->type_index() == rhs->type_index()); + // skip entries that already have equality maps. + auto it = equal_map_lhs_.find(lhs); + if (it != equal_map_lhs_.end()) { + return it->second.same_as(rhs); + } + if (equal_map_rhs_.count(rhs)) return false; + + SEqualReducer reducer = GetReducer(lhs, rhs, map_free_vars, current_paths); + return vtable_->SEqualReduce(lhs.get(), rhs.get(), reducer); + }; + return CheckResult(compute(), lhs, rhs, current_paths); + } + protected: + SEqualReducer GetReducer(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, + const Optional& current_paths) { + if (!IsPathTracingEnabled()) { + return SEqualReducer(parent_, nullptr, map_free_vars); + } else { + PathTracingData tracing_data = {current_paths.value(), lhs, rhs, first_mismatch_}; + return SEqualReducer(parent_, &tracing_data, map_free_vars); + } + } + // Check the result. bool CheckResult(bool result, const ObjectRef& lhs, const ObjectRef& rhs, const Optional& current_paths) { @@ -335,7 +363,8 @@ class RemapVarSEqualHandler : public SEqualReducer::Handler { // which populates the pending tasks. ICHECK_EQ(pending_tasks_.size(), 0U); allow_push_to_stack_ = false; - if (!DispatchSEqualReduce(entry.lhs, entry.rhs, entry.map_free_vars, entry.current_paths)) + if (!parent_->DispatchSEqualReduce(entry.lhs, entry.rhs, entry.map_free_vars, + entry.current_paths)) return false; allow_push_to_stack_ = true; // Push pending tasks in reverse order, so earlier tasks get to @@ -349,31 +378,6 @@ class RemapVarSEqualHandler : public SEqualReducer::Handler { return true; } - // The default equal as registered in the structural equal vtable. - bool DispatchSEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, - const Optional& current_paths) { - auto compute = [=]() { - ICHECK(lhs.defined() && rhs.defined() && lhs->type_index() == rhs->type_index()); - // skip entries that already have equality maps. - auto it = equal_map_lhs_.find(lhs); - if (it != equal_map_lhs_.end()) { - return it->second.same_as(rhs); - } - if (equal_map_rhs_.count(rhs)) return false; - - // Run reduce check for free nodes. - if (!IsPathTracingEnabled()) { - return vtable_->SEqualReduce(lhs.get(), rhs.get(), - SEqualReducer(this, nullptr, map_free_vars)); - } else { - PathTracingData tracing_data = {current_paths.value(), lhs, rhs, first_mismatch_}; - return vtable_->SEqualReduce(lhs.get(), rhs.get(), - SEqualReducer(this, &tracing_data, map_free_vars)); - } - }; - return CheckResult(compute(), lhs, rhs, current_paths); - } - private: /*! \brief Pending reduce tasks. */ struct Task { @@ -407,6 +411,8 @@ class RemapVarSEqualHandler : public SEqualReducer::Handler { bool IsPathTracingEnabled() const { return first_mismatch_ != nullptr; } + // The owner of this impl + SEqualHandlerDefault* parent_; // list of pending tasks to be pushed to the stack. std::vector pending_tasks_; // Internal task stack to executed the task. @@ -425,22 +431,53 @@ class RemapVarSEqualHandler : public SEqualReducer::Handler { std::unordered_map equal_map_rhs_; }; +SEqualHandlerDefault::SEqualHandlerDefault(bool assert_mode, + Optional* first_mismatch) { + impl = new Impl(this, assert_mode, first_mismatch); +} + +SEqualHandlerDefault::~SEqualHandlerDefault() { delete impl; } + +bool SEqualHandlerDefault::SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, + bool map_free_vars, + const Optional& current_paths) { + return impl->SEqualReduce(lhs, rhs, map_free_vars, current_paths); +} + +void SEqualHandlerDefault::DeferFail(const ObjectPathPair& mismatch_paths) { + impl->DeferFail(mismatch_paths); +} + +ObjectRef SEqualHandlerDefault::MapLhsToRhs(const ObjectRef& lhs) { return impl->MapLhsToRhs(lhs); } + +void SEqualHandlerDefault::MarkGraphNode() { impl->MarkGraphNode(); } + +bool SEqualHandlerDefault::Equal(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars) { + return impl->Equal(lhs, rhs, map_free_vars); +} + +bool SEqualHandlerDefault::DispatchSEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, + bool map_free_vars, + const Optional& current_paths) { + return impl->DispatchSEqualReduce(lhs, rhs, map_free_vars, current_paths); +} + TVM_REGISTER_GLOBAL("node.StructuralEqual") .set_body_typed([](const ObjectRef& lhs, const ObjectRef& rhs, bool assert_mode, bool map_free_vars) { - return RemapVarSEqualHandler(assert_mode, nullptr).Equal(lhs, rhs, map_free_vars); + return SEqualHandlerDefault(assert_mode, nullptr).Equal(lhs, rhs, map_free_vars); }); TVM_REGISTER_GLOBAL("node.GetFirstStructuralMismatch") .set_body_typed([](const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars) { Optional first_mismatch; - bool equal = RemapVarSEqualHandler(false, &first_mismatch).Equal(lhs, rhs, map_free_vars); + bool equal = SEqualHandlerDefault(false, &first_mismatch).Equal(lhs, rhs, map_free_vars); ICHECK(equal == !first_mismatch.defined()); return first_mismatch; }); bool StructuralEqual::operator()(const ObjectRef& lhs, const ObjectRef& rhs) const { - return RemapVarSEqualHandler(false, nullptr).Equal(lhs, rhs, false); + return SEqualHandlerDefault(false, nullptr).Equal(lhs, rhs, false); } } // namespace tvm diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc index b40b1751fb78..030533333a41 100644 --- a/src/node/structural_hash.cc +++ b/src/node/structural_hash.cc @@ -55,8 +55,11 @@ void ReflectionVTable::SHashReduce(const Object* self, SHashReducer reducer) con // In particular, when we traverse unordered_map, we should first sort // the entries by keys(or hash of keys) before traversing. -class VarCountingSHashHandler : public SHashReducer::Handler { +class SHashHandlerDefault::Impl { public: + explicit Impl(SHashHandlerDefault* parent) : parent_(parent) {} + virtual ~Impl() = default; + /*! \brief Pending reduce tasks. */ struct Task { /*! @@ -81,15 +84,13 @@ class VarCountingSHashHandler : public SHashReducer::Handler { : object(object), reduced_hash(reduced_hash), map_free_vars(map_free_vars) {} }; - VarCountingSHashHandler() {} - - void MarkGraphNode() final { + void MarkGraphNode() { // need to push to pending tasks in this case ICHECK(!allow_push_to_stack_ && !task_stack_.empty()); task_stack_.back().graph_node_hash = true; } - bool LookupHashedValue(const ObjectRef& key, size_t* hash_value) final { + bool LookupHashedValue(const ObjectRef& key, size_t* hash_value) { auto it = hash_memo_.find(key); if (it != hash_memo_.end()) { hash_value[0] = it->second; @@ -98,11 +99,11 @@ class VarCountingSHashHandler : public SHashReducer::Handler { return false; } - void SHashReduceHashedValue(size_t hashed_value) final { + void SHashReduceHashedValue(size_t hashed_value) { pending_tasks_.emplace_back(Task(ObjectRef(nullptr), hashed_value, false)); } - void SHashReduceFreeVar(const runtime::Object* var, bool map_free_vars) final { + void SHashReduceFreeVar(const runtime::Object* var, bool map_free_vars) { ICHECK(!hash_memo_.count(GetRef(var))); if (map_free_vars) { // use counter value. @@ -115,7 +116,7 @@ class VarCountingSHashHandler : public SHashReducer::Handler { } } - void SHashReduce(const ObjectRef& object, bool map_free_vars) final { + void SHashReduce(const ObjectRef& object, bool map_free_vars) { // Directly push the result // Note: it is still important to push the result to pendng tasks // so that the reduction order of hash values stays the same. @@ -151,6 +152,11 @@ class VarCountingSHashHandler : public SHashReducer::Handler { return ret; } + void DispatchSHash(const ObjectRef& object, bool map_free_vars) { + ICHECK(object.defined()); + vtable_->SHashReduce(object.get(), SHashReducer(parent_, map_free_vars)); + } + protected: /*! * \brief Pop the top entry of the task stack and push the hash into the result stack. @@ -219,7 +225,7 @@ class VarCountingSHashHandler : public SHashReducer::Handler { ICHECK_EQ(pending_tasks_.size(), 0U); allow_push_to_stack_ = false; // dispatch hash, reduce to the current slot. - this->DispatchSHash(entry.object, entry.map_free_vars); + parent_->DispatchSHash(entry.object, entry.map_free_vars); allow_push_to_stack_ = true; // Move pending tasks to the stack until the marked point. while (pending_tasks_.size() != 0) { @@ -231,13 +237,9 @@ class VarCountingSHashHandler : public SHashReducer::Handler { } } - // The default equal as registered in the structural equal vtable. - void DispatchSHash(const ObjectRef& object, bool map_free_vars) { - ICHECK(object.defined()); - vtable_->SHashReduce(object.get(), SHashReducer(this, map_free_vars)); - } - private: + // The owner of this impl + SHashHandlerDefault* parent_; // free var counter. size_t free_var_counter_{0}; // graph node counter. @@ -256,14 +258,43 @@ class VarCountingSHashHandler : public SHashReducer::Handler { std::unordered_map hash_memo_; }; +SHashHandlerDefault::SHashHandlerDefault() { impl = new Impl(this); } +SHashHandlerDefault::~SHashHandlerDefault() { delete impl; } + +void SHashHandlerDefault::SHashReduceHashedValue(size_t hashed_value) { + return impl->SHashReduceHashedValue(hashed_value); +} + +void SHashHandlerDefault::SHashReduce(const ObjectRef& key, bool map_free_vars) { + impl->SHashReduce(key, map_free_vars); +} + +void SHashHandlerDefault::SHashReduceFreeVar(const runtime::Object* var, bool map_free_vars) { + impl->SHashReduceFreeVar(var, map_free_vars); +} + +bool SHashHandlerDefault::LookupHashedValue(const ObjectRef& key, size_t* hashed_value) { + return impl->LookupHashedValue(key, hashed_value); +} + +void SHashHandlerDefault::MarkGraphNode() { impl->MarkGraphNode(); } + +size_t SHashHandlerDefault::Hash(const ObjectRef& object, bool map_free_vars) { + return impl->Hash(object, map_free_vars); +} + +void SHashHandlerDefault::DispatchSHash(const ObjectRef& key, bool map_free_vars) { + impl->DispatchSHash(key, map_free_vars); +} + TVM_REGISTER_GLOBAL("node.StructuralHash") .set_body_typed([](const ObjectRef& object, bool map_free_vars) -> int64_t { - size_t hashed_value = VarCountingSHashHandler().Hash(object, map_free_vars); + size_t hashed_value = SHashHandlerDefault().Hash(object, map_free_vars); return static_cast(hashed_value); }); size_t StructuralHash::operator()(const ObjectRef& object) const { - return VarCountingSHashHandler().Hash(object, false); + return SHashHandlerDefault().Hash(object, false); } // SEQualReduce traits for runtime containers. From 29e42774049e5a7867013afd8b25790ff226412f Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 7 Oct 2022 05:22:29 +0900 Subject: [PATCH 2/3] minor fix --- src/node/structural_hash.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/src/node/structural_hash.cc b/src/node/structural_hash.cc index 030533333a41..a355e44028b6 100644 --- a/src/node/structural_hash.cc +++ b/src/node/structural_hash.cc @@ -58,7 +58,6 @@ void ReflectionVTable::SHashReduce(const Object* self, SHashReducer reducer) con class SHashHandlerDefault::Impl { public: explicit Impl(SHashHandlerDefault* parent) : parent_(parent) {} - virtual ~Impl() = default; /*! \brief Pending reduce tasks. */ struct Task { From c21cae5315d1154e79047a562778de5de5d4525b Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 7 Oct 2022 08:08:51 +0900 Subject: [PATCH 3/3] lint fix --- include/tvm/node/structural_equal.h | 10 +++++----- include/tvm/node/structural_hash.h | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/include/tvm/node/structural_equal.h b/include/tvm/node/structural_equal.h index 1260f2ecf9e3..371b8f9c7bd9 100644 --- a/include/tvm/node/structural_equal.h +++ b/include/tvm/node/structural_equal.h @@ -334,11 +334,11 @@ class SEqualHandlerDefault : public SEqualReducer::Handler { SEqualHandlerDefault(bool assert_mode, Optional* first_mismatch); virtual ~SEqualHandlerDefault(); - virtual bool SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, - const Optional& current_paths) override; - virtual void DeferFail(const ObjectPathPair& mismatch_paths) override; - virtual ObjectRef MapLhsToRhs(const ObjectRef& lhs) override; - virtual void MarkGraphNode() override; + bool SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars, + const Optional& current_paths) override; + void DeferFail(const ObjectPathPair& mismatch_paths) override; + ObjectRef MapLhsToRhs(const ObjectRef& lhs) override; + void MarkGraphNode() override; /*! * \brief The entry point for equality testing diff --git a/include/tvm/node/structural_hash.h b/include/tvm/node/structural_hash.h index 591804395025..8b8a403326c4 100644 --- a/include/tvm/node/structural_hash.h +++ b/include/tvm/node/structural_hash.h @@ -210,11 +210,11 @@ class SHashHandlerDefault : public SHashReducer::Handler { SHashHandlerDefault(); virtual ~SHashHandlerDefault(); - virtual void SHashReduceHashedValue(size_t hashed_value) override; - virtual void SHashReduce(const ObjectRef& key, bool map_free_vars) override; - virtual void SHashReduceFreeVar(const runtime::Object* var, bool map_free_vars) override; - virtual bool LookupHashedValue(const ObjectRef& key, size_t* hashed_value) override; - virtual void MarkGraphNode() override; + void SHashReduceHashedValue(size_t hashed_value) override; + void SHashReduce(const ObjectRef& key, bool map_free_vars) override; + void SHashReduceFreeVar(const runtime::Object* var, bool map_free_vars) override; + bool LookupHashedValue(const ObjectRef& key, size_t* hashed_value) override; + void MarkGraphNode() override; /*! * \brief The entry point for hashing