[FFI] Structural equal and hash based on reflection (#18156)

This PR add initial support for structural equal and hash
via the new reflection mechanism.

It will helps us to streamline the structural equality/hash
with broader support and clean error reports via AccessPath.
It also gives us ability to unify all struct equal/hash
registration into the extra meta-data in reflection registration.

Author: tqchen

ffi/CMakeLists.txt

@@ -57,6 +57,9 @@ add_library(tvm_ffi_objs OBJECT
   "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/dtype.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/testing.cc"
   "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/container.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/reflection/access_path.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/reflection/structural_equal.cc"
+  "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/reflection/structural_hash.cc"
 )
 set_target_properties(
   tvm_ffi_objs PROPERTIES

ffi/include/tvm/ffi/c_api.h

@@ -56,6 +56,27 @@
 #define TVM_FFI_DLL_EXPORT __attribute__((visibility("default")))
 #endif
 
+/*!
+ * \brief Marks the API as extra c++ api that is defined in cc files.
+ *
+ * These APIs are extra features that depend on, but are not required to
+ * support essential core functionality, such as function calling and object
+ * access.
+ *
+ * They are implemented in cc files to reduce compile-time overhead.
+ * The input/output only uses POD/Any/ObjectRef for ABI stability.
+ * However, these extra APIs may have an issue across MSVC/Itanium ABI,
+ *
+ * Related features are also available through reflection based function
+ * that is fully based on C API
+ *
+ * The project aims to minimize the number of extra C++ APIs and only
+ * restrict the use to non-core functionalities.
+ */
+#ifndef TVM_FFI_EXTRA_CXX_API
+#define TVM_FFI_EXTRA_CXX_API TVM_FFI_DLL
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -326,12 +347,89 @@ typedef enum {
   kTVMFFIFieldFlagBitMaskHasDefault = 1 << 1,
   /*! \brief The field is a static method. */
   kTVMFFIFieldFlagBitMaskIsStaticMethod = 1 << 2,
+  /*!
+   * \brief The field should be ignored when performing structural eq/hash
+   *
+   * This is an optional meta-data for structural eq/hash.
+   */
+  kTVMFFIFieldFlagBitMaskSEqHashIgnore = 1 << 3,
+  /*!
+   * \brief The field enters a def region where var can be defined/matched.
+   *
+   * This is an optional meta-data for structural eq/hash.
+   */
+  kTVMFFIFieldFlagBitMaskSEqHashDef = 1 << 4,
 #ifdef __cplusplus
 };
 #else
 } TVMFFIFieldFlagBitMask;
 #endif
 
+/*!
+ * \brief Optional meta-data for structural eq/hash.
+ *
+ * This meta-data is only useful when we want to leverage the information
+ * to perform richer semantics aware structural comparison and hash.
+ * It can be safely ignored if such information is not needed.
+ *
+ * The meta-data record comparison method in tree node and DAG node.
+ *
+ * \code
+ * x = VarNode()
+ * v0 = AddNode(x, 1)
+ * v1 = AddNode(x, 1)
+ * v2 = AddNode(v0, v0)
+ * v3 = AddNode(v1, v0)
+ * \endcode
+ *
+ * Consider the construct sequence of AddNode below,
+ * if AddNode is treated as a tree node, then v2 and v3
+ * structural equals to each other, but if AddNode is
+ * treated as a DAG node, then v2 and v3 does not
+ * structural equals to each other.
+ */
+#ifdef __cplusplus
+enum TVMFFISEqHashKind : int32_t {
+#else
+typedef enum {
+#endif
+  /*! \brief Do not support structural eq/hash. */
+  kTVMFFISEqHashKindUnsupported = 0,
+  /*!
+   * \brief The object be compared as a tree node.
+   */
+  kTVMFFISEqHashKindTreeNode = 1,
+  /*!
+   * \brief The object is treated as a free variable that can be mapped
+   *        to another free variable in the definition region.
+   */
+  kTVMFFISEqHashKindFreeVar = 2,
+  /*!
+   * \brief The field should be compared as a DAG node.
+   */
+  kTVMFFISEqHashKindDAGNode = 3,
+  /*!
+   * \brief The object is treated as a constant tree node.
+   *
+   * Same as tree node, but the object does not contain free var
+   * as any of its nested children.
+   *
+   * That means we can use pointer equality for equality.
+   */
+  kTVMFFISEqHashKindConstTreeNode = 4,
+  /*!
+   * \brief One can simply use pointer equality for equality.
+   *
+   * This is useful for "singleton"-style object that can
+   * is only an unique copy of each value.
+   */
+  kTVMFFISEqHashKindUniqueInstance = 5,
+#ifdef __cplusplus
+};
+#else
+} TVMFFISEqHashKind;
+#endif
+
 /*!
  * \brief Information support for optional object reflection.
  */
@@ -431,7 +529,11 @@ typedef struct {
    *
    * This field is set optional and set to 0 if not registered.
    */
-  int64_t total_size;
+  int32_t total_size;
+  /*!
+   * \brief Optional meta-data for structural eq/hash.
+   */
+  TVMFFISEqHashKind structural_eq_hash_kind;
 } TVMFFITypeExtraInfo;
 
 /*!

ffi/include/tvm/ffi/object.h

@@ -212,6 +212,8 @@ class Object {
   static constexpr int32_t _type_index = TypeIndex::kTVMFFIObject;
   // the static type depth of the class
   static constexpr int32_t _type_depth = 0;
+  // the structural equality and hash kind of the type
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUnsupported;
   // extra fields used by plug-ins for attribute visiting
   // and structural information
   static constexpr const bool _type_has_method_sequal_reduce = false;

ffi/include/tvm/ffi/reflection/access_path.h

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/ffi/reflection/registry.h
+ * \brief Registry of reflection metadata.
+ */
+#ifndef TVM_FFI_REFLECTION_ACCESS_PATH_H_
+#define TVM_FFI_REFLECTION_ACCESS_PATH_H_
+
+#include <tvm/ffi/any.h>
+#include <tvm/ffi/c_api.h>
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/reflection/registry.h>
+
+namespace tvm {
+namespace ffi {
+namespace reflection {
+
+enum class AccessKind : int32_t {
+  kObjectField = 0,
+  kArrayIndex = 1,
+  kMapKey = 2,
+  // the following two are used for error reporting when
+  // the supposed access field is not available
+  kArrayIndexMissing = 3,
+  kMapKeyMissing = 4,
+};
+
+/*!
+ * \brief Represent a single step in object field, map key, array index access.
+ */
+class AccessStepObj : public Object {
+ public:
+  /*!
+   * \brief The kind of the access pattern.
+   */
+  AccessKind kind;
+  /*!
+   * \brief The access key
+   * \note for array access, it will always be integer
+   *       for field access, it will be string
+   */
+  Any key;
+
+  AccessStepObj(AccessKind kind, Any key) : kind(kind), key(key) {}
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<AccessStepObj>()
+        .def_ro("kind", &AccessStepObj::kind)
+        .def_ro("key", &AccessStepObj::key);
+  }
+
+  static constexpr const char* _type_key = "tvm.ffi.reflection.AccessStep";
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindConstTreeNode;
+  TVM_FFI_DECLARE_FINAL_OBJECT_INFO(AccessStepObj, Object);
+};
+
+/*!
+ * \brief ObjectRef class of AccessStepObj.
+ *
+ * \sa AccessStepObj
+ */
+class AccessStep : public ObjectRef {
+ public:
+  AccessStep(AccessKind kind, Any key) : ObjectRef(make_object<AccessStepObj>(kind, key)) {}
+
+  static AccessStep ObjectField(String field_name) {
+    return AccessStep(AccessKind::kObjectField, field_name);
+  }
+
+  static AccessStep ArrayIndex(int64_t index) { return AccessStep(AccessKind::kArrayIndex, index); }
+
+  static AccessStep ArrayIndexMissing(int64_t index) {
+    return AccessStep(AccessKind::kArrayIndexMissing, index);
+  }
+
+  static AccessStep MapKey(Any key) { return AccessStep(AccessKind::kMapKey, key); }
+
+  static AccessStep MapKeyMissing(Any key) { return AccessStep(AccessKind::kMapKeyMissing, key); }
+
+  TVM_FFI_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AccessStep, ObjectRef, AccessStepObj);
+};
+
+using AccessPath = Array<AccessStep>;
+using AccessPathPair = Tuple<AccessPath, AccessPath>;
+
+}  // namespace reflection
+}  // namespace ffi
+}  // namespace tvm
+#endif  // TVM_FFI_REFLECTION_ACCESS_PATH_H_

ffi/include/tvm/ffi/reflection/registry.h

@@ -55,6 +55,39 @@ class DefaultValue : public FieldInfoTrait {
   Any value_;
 };
 
+/*
+ * \brief Trait that can be used to attach field flag
+ */
+class AttachFieldFlag : public FieldInfoTrait {
+ public:
+  /*!
+   * \brief Attach a field flag to the field
+   *
+   * \param flag The flag to be set
+   *
+   * \return The trait object.
+   */
+  explicit AttachFieldFlag(int32_t flag) : flag_(flag) {}
+
+  /*!
+   * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashDef
+   */
+  TVM_FFI_INLINE static AttachFieldFlag SEqHashDef() {
+    return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashDef);
+  }
+  /*!
+   * \brief Attach kTVMFFIFieldFlagBitMaskSEqHashIgnore
+   */
+  TVM_FFI_INLINE static AttachFieldFlag SEqHashIgnore() {
+    return AttachFieldFlag(kTVMFFIFieldFlagBitMaskSEqHashIgnore);
+  }
+
+  TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const { info->flags |= flag_; }
+
+ private:
+  int32_t flag_;
+};
+
 /*!
  * \brief Get the byte offset of a class member field.
  *
@@ -83,7 +116,11 @@ class ReflectionDefBase {
   template <typename T>
   static int FieldSetter(void* field, const TVMFFIAny* value) {
     TVM_FFI_SAFE_CALL_BEGIN();
-    *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value).cast<T>();
+    if constexpr (std::is_same_v<T, Any>) {
+      *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value);
+    } else {
+      *reinterpret_cast<T*>(field) = AnyView::CopyFromTVMFFIAny(*value).cast<T>();
+    }
     TVM_FFI_SAFE_CALL_END();
   }
 
@@ -346,6 +383,7 @@ class ObjectDef : public ReflectionDefBase {
   void RegisterExtraInfo(ExtraArgs&&... extra_args) {
     TVMFFITypeExtraInfo info;
     info.total_size = sizeof(Class);
+    info.structural_eq_hash_kind = Class::_type_s_eq_hash_kind;
     info.creator = nullptr;
     info.doc = TVMFFIByteArray{nullptr, 0};
     if constexpr (std::is_default_constructible_v<Class>) {