structural_equal.h

Go to the documentation of this file.

/*
 * 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_NODE_STRUCTURAL_EQUAL_H_
#define TVM_NODE_STRUCTURAL_EQUAL_H_

#include <tvm/node/functor.h>
#include <tvm/node/object_path.h>
#include <tvm/runtime/container/array.h>
#include <tvm/runtime/data_type.h>

#include <string>

namespace tvm {

class BaseValueEqual {
 public:
  bool operator()(const double& lhs, const double& rhs) const {
    // fuzzy float pt comparison
    constexpr double atol = 1e-9;
    if (lhs == rhs) return true;
    double diff = lhs - rhs;
    return diff > -atol && diff < atol;
  }

  bool operator()(const int64_t& lhs, const int64_t& rhs) const { return lhs == rhs; }
  bool operator()(const uint64_t& lhs, const uint64_t& rhs) const { return lhs == rhs; }
  bool operator()(const int& lhs, const int& rhs) const { return lhs == rhs; }
  bool operator()(const bool& lhs, const bool& rhs) const { return lhs == rhs; }
  bool operator()(const std::string& lhs, const std::string& rhs) const { return lhs == rhs; }
  bool operator()(const DataType& lhs, const DataType& rhs) const { return lhs == rhs; }
  template <typename ENum, typename = typename std::enable_if<std::is_enum<ENum>::value>::type>
  bool operator()(const ENum& lhs, const ENum& rhs) const {
    return lhs == rhs;
  }
};

class ObjectPathPairNode : public Object {
 public:
  ObjectPath lhs_path;
  ObjectPath rhs_path;

  ObjectPathPairNode(ObjectPath lhs_path, ObjectPath rhs_path);

  static constexpr const char* _type_key = "ObjectPathPair";
  TVM_DECLARE_FINAL_OBJECT_INFO(ObjectPathPairNode, Object);
};

class ObjectPathPair : public ObjectRef {
 public:
  ObjectPathPair(ObjectPath lhs_path, ObjectPath rhs_path);

  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ObjectPathPair, ObjectRef, ObjectPathPairNode);
};

class StructuralEqual : public BaseValueEqual {
 public:
  // inheritate operator()
  using BaseValueEqual::operator();
  TVM_DLL bool operator()(const ObjectRef& lhs, const ObjectRef& rhs) const;
};

class SEqualReducer {
 private:
  struct PathTracingData;

 public:
  class Handler {
   public:
    virtual bool SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars,
                              const Optional<ObjectPathPair>& current_paths) = 0;

    virtual void DeferFail(const ObjectPathPair& mismatch_paths) = 0;

    virtual bool IsFailDeferralEnabled() = 0;

    virtual ObjectRef MapLhsToRhs(const ObjectRef& lhs) = 0;
    virtual void MarkGraphNode() = 0;

   protected:
    using PathTracingData = SEqualReducer::PathTracingData;
  };

  SEqualReducer() = default;
  explicit SEqualReducer(Handler* handler, const PathTracingData* tracing_data, bool map_free_vars)
      : handler_(handler), tracing_data_(tracing_data), map_free_vars_(map_free_vars) {}

  bool operator()(const double& lhs, const double& rhs,
                  Optional<ObjectPathPair> paths = NullOpt) const;
  bool operator()(const int64_t& lhs, const int64_t& rhs,
                  Optional<ObjectPathPair> paths = NullOpt) const;
  bool operator()(const uint64_t& lhs, const uint64_t& rhs,
                  Optional<ObjectPathPair> paths = NullOpt) const;
  bool operator()(const int& lhs, const int& rhs, Optional<ObjectPathPair> paths = NullOpt) const;
  bool operator()(const bool& lhs, const bool& rhs, Optional<ObjectPathPair> paths = NullOpt) const;
  bool operator()(const std::string& lhs, const std::string& rhs,
                  Optional<ObjectPathPair> paths = NullOpt) const;
  bool operator()(const DataType& lhs, const DataType& rhs,
                  Optional<ObjectPathPair> paths = NullOpt) const;

  template <typename ENum, typename = typename std::enable_if<std::is_enum<ENum>::value>::type>
  bool operator()(const ENum& lhs, const ENum& rhs,
                  Optional<ObjectPathPair> paths = NullOpt) const {
    using Underlying = typename std::underlying_type<ENum>::type;
    static_assert(std::is_same<Underlying, int>::value,
                  "Enum must have `int` as the underlying type");
    return EnumAttrsEqual(static_cast<int>(lhs), static_cast<int>(rhs), &lhs, &rhs, paths);
  }

  template <typename T, typename Callable,
            typename = std::enable_if_t<
                std::is_same_v<std::invoke_result_t<Callable, const ObjectPath&>, ObjectPath>>>
  bool operator()(const T& lhs, const T& rhs, const Callable& callable) {
    if (IsPathTracingEnabled()) {
      ObjectPathPair current_paths = GetCurrentObjectPaths();
      ObjectPathPair new_paths = {callable(current_paths->lhs_path),
                                  callable(current_paths->rhs_path)};
      return (*this)(lhs, rhs, new_paths);
    } else {
      return (*this)(lhs, rhs);
    }
  }

  bool operator()(const ObjectRef& lhs, const ObjectRef& rhs) const;

  bool operator()(const ObjectRef& lhs, const ObjectRef& rhs, const ObjectPathPair& paths) const {
    ICHECK(IsPathTracingEnabled()) << "Path tracing must be enabled when calling this function";
    return ObjectAttrsEqual(lhs, rhs, map_free_vars_, &paths);
  }

  bool DefEqual(const ObjectRef& lhs, const ObjectRef& rhs);

  template <typename T>
  bool operator()(const Array<T>& lhs, const Array<T>& rhs) const {
    if (tracing_data_ == nullptr) {
      // quick specialization for Array to reduce amount of recursion
      // depth as array comparison is pretty common.
      if (lhs.size() != rhs.size()) return false;
      for (size_t i = 0; i < lhs.size(); ++i) {
        if (!(operator()(lhs[i], rhs[i]))) return false;
      }
      return true;
    }

    // If tracing is enabled, fall back to the regular path
    const ObjectRef& lhs_obj = lhs;
    const ObjectRef& rhs_obj = rhs;
    return (*this)(lhs_obj, rhs_obj);
  }
  bool FreeVarEqualImpl(const runtime::Object* lhs, const runtime::Object* rhs) const {
    // var need to be remapped, so it belongs to graph node.
    handler_->MarkGraphNode();
    // We only map free vars if they corresponds to the same address
    // or map free_var option is set to be true.
    return lhs == rhs || map_free_vars_;
  }

  Handler* operator->() const { return handler_; }

  bool IsPathTracingEnabled() const { return tracing_data_ != nullptr; }

  const ObjectPathPair& GetCurrentObjectPaths() const;

  void RecordMismatchPaths(const ObjectPathPair& paths) const;

 private:
  bool EnumAttrsEqual(int lhs, int rhs, const void* lhs_address, const void* rhs_address,
                      Optional<ObjectPathPair> paths = NullOpt) const;

  bool ObjectAttrsEqual(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars,
                        const ObjectPathPair* paths) const;

  static void GetPathsFromAttrAddressesAndStoreMismatch(const void* lhs_address,
                                                        const void* rhs_address,
                                                        const PathTracingData* tracing_data);

  template <typename T>
  static bool CompareAttributeValues(const T& lhs, const T& rhs,
                                     const PathTracingData* tracing_data,
                                     Optional<ObjectPathPair> paths = NullOpt);

  Handler* handler_ = nullptr;
  const PathTracingData* tracing_data_ = nullptr;
  bool map_free_vars_ = false;
};

class SEqualHandlerDefault : public SEqualReducer::Handler {
 public:
  SEqualHandlerDefault(bool assert_mode, Optional<ObjectPathPair>* first_mismatch,
                       bool defer_fails);
  virtual ~SEqualHandlerDefault();

  bool SEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars,
                    const Optional<ObjectPathPair>& current_paths) override;
  void DeferFail(const ObjectPathPair& mismatch_paths) override;
  bool IsFailDeferralEnabled() override;
  ObjectRef MapLhsToRhs(const ObjectRef& lhs) override;
  void MarkGraphNode() override;

  virtual bool Equal(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars);

 protected:
  virtual bool DispatchSEqualReduce(const ObjectRef& lhs, const ObjectRef& rhs, bool map_free_vars,
                                    const Optional<ObjectPathPair>& current_paths);

 private:
  class Impl;
  Impl* impl;
};

}  // namespace tvm
#endif  // TVM_NODE_STRUCTURAL_EQUAL_H_