[FFI] Structural equal and hash based on reflection

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>) {