[FFI][REFACTOR] Introduce TypeAttr in reflection

This PR introduces TypeAttr to reflection to bring extra
optional attribute registration that can be used to extend behaviors such
as structural equality. Also renames TypeExtraInfo to TypeMetadata for better clarity.

Author: tqchen

ffi/include/tvm/ffi/c_api.h

@@ -534,7 +534,22 @@ typedef struct {
    * \brief Optional meta-data for structural eq/hash.
    */
   TVMFFISEqHashKind structural_eq_hash_kind;
-} TVMFFITypeExtraInfo;
+} TVMFFITypeMetadata;
+
+/*
+ * \brief Column array that stores extra attributes about types
+ *
+ * The attributes stored in column arrays that can be looked up by type index.
+ *
+ * \note
+ * \sa TVMFFIRegisterTypeAttr
+ */
+typedef struct {
+  /*! \brief The data of the column. */
+  const TVMFFIAny* data;
+  /*! \brief The size of the column. */
+  size_t size;
+} TVMFFITypeAttrColumn;
 
 /*!
  * \brief Runtime type information for object type checking.
@@ -567,7 +582,7 @@ typedef struct TVMFFITypeInfo {
   /*! \brief The reflection method. */
   const TVMFFIMethodInfo* methods;
   /*! \brief The extra information of the type. */
-  const TVMFFITypeExtraInfo* extra_info;
+  const TVMFFITypeMetadata* metadata;
 } TVMFFITypeInfo;
 
 //------------------------------------------------------------
@@ -738,11 +753,27 @@ TVM_FFI_DLL int TVMFFITypeRegisterMethod(int32_t type_index, const TVMFFIMethodI
 /*!
  * \brief Register type creator information for runtime reflection.
  * \param type_index The type index
- * \param extra_info The extra information to be registered.
+ * \param metadata The extra information to be registered.
  * \return 0 when success, nonzero when failure happens
  */
-TVM_FFI_DLL int TVMFFITypeRegisterExtraInfo(int32_t type_index,
-                                            const TVMFFITypeExtraInfo* extra_info);
+TVM_FFI_DLL int TVMFFITypeRegisterMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata);
+
+/*!
+ * \brief Register extra type attributes that can be looked up during runtime.
+ * \param type_index The type index
+ * \param attr_value The attribute value to be registered.
+ * \return 0 when success, nonzero when failure happens
+ */
+TVM_FFI_DLL int TVMFFITypeRegisterAttr(int32_t type_index, const TVMFFIByteArray* attr_name,
+                                       const TVMFFIAny* attr_value);
+
+/*!
+ * \brief Get the type attribute column by name.
+ * \param attr_name The name of the attribute.
+ * \return The pointer to the type attribute column.
+ * \return NULL if the attribute was not registered in the system
+ */
+TVM_FFI_DLL const TVMFFITypeAttrColumn* TVMFFIGetTypeAttrColumn(const TVMFFIByteArray* attr_name);
 
 //------------------------------------------------------------
 // Section: DLPack support APIs

ffi/include/tvm/ffi/reflection/accessor.h

@@ -104,6 +104,29 @@ class FieldSetter {
   const TVMFFIFieldInfo* field_info_;
 };
 
+class TypeAttrColumn {
+ public:
+  TypeAttrColumn(std::string_view attr_name) {
+    TVMFFIByteArray attr_name_array = {attr_name.data(), attr_name.size()};
+    column_ = TVMFFIGetTypeAttrColumn(&attr_name_array);
+    if (column_ == nullptr) {
+      TVM_FFI_THROW(RuntimeError) << "Cannot find type attribute " << attr_name;
+    }
+  }
+
+  AnyView operator[](int32_t type_index) const {
+    size_t tindex = static_cast<size_t>(type_index);
+    if (tindex >= column_->size) {
+      return AnyView();
+    }
+    const AnyView* any_view_data = reinterpret_cast<const AnyView*>(column_->data);
+    return any_view_data[tindex];
+  }
+
+ private:
+  const TVMFFITypeAttrColumn* column_;
+};
+
 /*!
  * \brief helper function to get reflection method info by type key and method name
  *

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

@@ -150,7 +150,7 @@ class ReflectionDefBase {
   }
 
   template <typename T>
-  TVM_FFI_INLINE static void ApplyExtraInfoTrait(TVMFFITypeExtraInfo* info, const T& value) {
+  TVM_FFI_INLINE static void ApplyExtraInfoTrait(TVMFFITypeMetadata* info, const T& value) {
     if constexpr (std::is_same_v<std::decay_t<T>, char*>) {
       info->doc = TVMFFIByteArray{value, std::char_traits<char>::length(value)};
     }
@@ -381,7 +381,7 @@ class ObjectDef : public ReflectionDefBase {
  private:
   template <typename... ExtraArgs>
   void RegisterExtraInfo(ExtraArgs&&... extra_args) {
-    TVMFFITypeExtraInfo info;
+    TVMFFITypeMetadata info;
     info.total_size = sizeof(Class);
     info.structural_eq_hash_kind = Class::_type_s_eq_hash_kind;
     info.creator = nullptr;
@@ -391,7 +391,7 @@ class ObjectDef : public ReflectionDefBase {
     }
     // apply extra info traits
     ((ApplyExtraInfoTrait(&info, std::forward<ExtraArgs>(extra_args)), ...));
-    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterExtraInfo(type_index_, &info));
+    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMetadata(type_index_, &info));
   }
 
   template <typename T, typename BaseClass, typename... ExtraArgs>
@@ -446,6 +446,68 @@ class ObjectDef : public ReflectionDefBase {
   const char* type_key_;
 };
 
+template <typename Class, typename = std::enable_if_t<std::is_base_of_v<Object, Class>>>
+class TypeAttrDef : public ReflectionDefBase {
+ public:
+  template <typename... ExtraArgs>
+  explicit TypeAttrDef(ExtraArgs&&... extra_args)
+      : type_index_(Class::RuntimeTypeIndex()), type_key_(Class::_type_key) {}
+
+  /*
+   * \brief Define a function-valued type attribute.
+   *
+   * \tparam Func The function type.
+   *
+   * \param name The name of the function.
+   * \param func The function to be registered.
+   *
+   * \return The TypeAttrDef object.
+   */
+  template <typename Func>
+  TypeAttrDef& def(const char* name, Func&& func) {
+    TVMFFIByteArray name_array = {name, std::char_traits<char>::length(name)};
+    ffi::Function ffi_func =
+        GetMethod<Class>(std::string(type_key_) + "." + name, std::forward<Func>(func));
+    TVMFFIAny value_any = AnyView(ffi_func).CopyToTVMFFIAny();
+    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any));
+    return *this;
+  }
+
+  /*
+   * \brief Define a constant-valued type attribute.
+   *
+   * \tparam T The type of the value.
+   *
+   * \param name The name of the attribute.
+   * \param value The value of the attribute.
+   *
+   * \return The TypeAttrDef object.
+   */
+  template <typename T>
+  TypeAttrDef& attr(const char* name, T value) {
+    TVMFFIByteArray name_array = {name, std::char_traits<char>::length(name)};
+    TVMFFIAny value_any = AnyView(value).CopyToTVMFFIAny();
+    TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &name_array, &value_any));
+    return *this;
+  }
+
+ private:
+  int32_t type_index_;
+  const char* type_key_;
+};
+
+/*!
+ * \brief Ensure the type attribute column is presented in the system.
+ *
+ * \param name The name of the type attribute.
+ */
+inline void EnsureTypeAttrColumn(std::string_view name) {
+  TVMFFIByteArray name_array = {name.data(), name.size()};
+  AnyView any_view(nullptr);
+  TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(kTVMFFINone, &name_array,
+                                                 reinterpret_cast<const TVMFFIAny*>(&any_view)));
+}
+
 }  // namespace reflection
 }  // namespace ffi
 }  // namespace tvm

ffi/src/ffi/object.cc

@@ -59,7 +59,7 @@ class TypeTable {
     /*! \brief type methods informaton */
     std::vector<TVMFFIMethodInfo> type_methods_data;
     /*! \brief extra information */
-    TVMFFITypeExtraInfo extra_info_data;
+    TVMFFITypeMetadata metadata_data;
     // NOTE: the indices in [index, index + num_reserved_slots) are
     // reserved for the child-class of this type.
     /*! \brief Total number of slots reserved for the type and its children. */
@@ -100,10 +100,14 @@ class TypeTable {
       this->num_methods = 0;
       this->fields = nullptr;
       this->methods = nullptr;
-      this->extra_info = nullptr;
+      this->metadata = nullptr;
     }
   };
 
+  struct TypeAttrColumnData : public TVMFFITypeAttrColumn {
+    std::vector<Any> data_;
+  };
+
   int32_t GetOrAllocTypeIndex(String type_key, int32_t static_type_index, int32_t type_depth,
                               int32_t num_child_slots, bool child_slots_can_overflow,
                               int32_t parent_type_index) {
@@ -219,19 +223,49 @@ class TypeTable {
     entry->num_methods = static_cast<int32_t>(entry->type_methods_data.size());
   }
 
-  void RegisterTypeExtraInfo(int32_t type_index, const TVMFFITypeExtraInfo* extra_info) {
+  void RegisterTypeMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata) {
     Entry* entry = GetTypeEntry(type_index);
-    if (entry->extra_info != nullptr) {
+    if (entry->metadata != nullptr) {
       TVM_FFI_LOG_AND_THROW(RuntimeError)
           << "Overriding " << ToStringView(entry->type_key) << ", possible causes:\n"
           << "- two ObjectDef<T>() calls for the same T \n"
           << "- when we forget to assign _type_key to ObjectRef<Y> that inherits from T\n"
           << "- another type with the same key is already registered\n"
           << "Cross check the reflection registration.";
     }
-    entry->extra_info_data = *extra_info;
-    entry->extra_info_data.doc = this->CopyString(extra_info->doc);
-    entry->extra_info = &(entry->extra_info_data);
+    entry->metadata_data = *metadata;
+    entry->metadata_data.doc = this->CopyString(metadata->doc);
+    entry->metadata = &(entry->metadata_data);
+  }
+
+  void RegisterTypeAttr(int32_t type_index, const TVMFFIByteArray* name, const TVMFFIAny* value) {
+    AnyView value_view = AnyView::CopyFromTVMFFIAny(*value);
+    String name_str(*name);
+    size_t column_index = 0;
+    auto it = type_attr_name_to_column_index_.find(name_str);
+    if (it == type_attr_name_to_column_index_.end()) {
+      column_index = type_attr_columns_.size();
+      type_attr_columns_.emplace_back(std::make_unique<TypeAttrColumnData>());
+      type_attr_name_to_column_index_.Set(name_str, column_index);
+    }
+    TypeAttrColumnData* column = type_attr_columns_[column_index].get();
+    if (column->data_.size() < static_cast<size_t>(type_index + 1)) {
+      column->data_.resize(type_index + 1, Any(nullptr));
+      column->data = reinterpret_cast<const TVMFFIAny*>(column->data_.data());
+      column->size = column->data_.size();
+    }
+    if (type_index == kTVMFFINone) return;
+    if (column->data_[type_index] != nullptr) {
+      TVM_FFI_THROW(RuntimeError) << "Type attribute `" << name_str << "` is already set for type `"
+                                  << TypeIndexToTypeKey(type_index) << "`";
+    }
+    column->data_[type_index] = value_view;
+  }
+  const TVMFFITypeAttrColumn* GetTypeAttrColumn(const TVMFFIByteArray* name) {
+    String name_str(*name);
+    auto it = type_attr_name_to_column_index_.find(name_str);
+    if (it == type_attr_name_to_column_index_.end()) return nullptr;
+    return type_attr_columns_[(*it).second].get();
   }
 
   void Dump(int min_children_count) {
@@ -284,11 +318,11 @@ class TypeTable {
     this->GetOrAllocTypeIndex(Object::_type_key, Object::_type_index, Object::_type_depth,
                               Object::_type_child_slots, Object::_type_child_slots_can_overflow,
                               -1);
-    TVMFFITypeExtraInfo info;
+    TVMFFITypeMetadata info;
     info.total_size = sizeof(Object);
     info.creator = nullptr;
     info.doc = TVMFFIByteArray{nullptr, 0};
-    RegisterTypeExtraInfo(Object::_type_index, &info);
+    RegisterTypeMetadata(Object::_type_index, &info);
     // reserve the static types
     ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFINone, TypeIndex::kTVMFFINone);
     ReserveBuiltinTypeIndex(StaticTypeKey::kTVMFFIInt, TypeIndex::kTVMFFIInt);
@@ -328,6 +362,9 @@ class TypeTable {
   std::vector<std::unique_ptr<Entry>> type_table_;
   Map<String, int64_t> type_key2index_;
   std::vector<Any> any_pool_;
+  // type attribute columns
+  std::vector<std::unique_ptr<TypeAttrColumnData>> type_attr_columns_;
+  Map<String, int64_t> type_attr_name_to_column_index_;
 };
 
 void MakeObjectFromPackedArgs(ffi::PackedArgs args, Any* ret) {
@@ -343,12 +380,12 @@ void MakeObjectFromPackedArgs(ffi::PackedArgs args, Any* ret) {
   TVM_FFI_ICHECK(args.size() % 2 == 1);
   const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
 
-  if (type_info->extra_info == nullptr || type_info->extra_info->creator == nullptr) {
+  if (type_info->metadata == nullptr || type_info->metadata->creator == nullptr) {
     TVM_FFI_THROW(RuntimeError) << "Type `" << TypeIndexToTypeKey(type_index)
                                 << "` does not support reflection creation";
   }
   TVMFFIObjectHandle handle;
-  TVM_FFI_CHECK_SAFE_CALL(type_info->extra_info->creator(&handle));
+  TVM_FFI_CHECK_SAFE_CALL(type_info->metadata->creator(&handle));
   ObjectPtr<Object> ptr =
       details::ObjectUnsafe::ObjectPtrFromOwned<Object>(static_cast<TVMFFIObject*>(handle));
 
@@ -437,12 +474,25 @@ int TVMFFITypeRegisterMethod(int32_t type_index, const TVMFFIMethodInfo* info) {
   TVM_FFI_SAFE_CALL_END();
 }
 
-int TVMFFITypeRegisterExtraInfo(int32_t type_index, const TVMFFITypeExtraInfo* extra_info) {
+int TVMFFITypeRegisterMetadata(int32_t type_index, const TVMFFITypeMetadata* metadata) {
   TVM_FFI_SAFE_CALL_BEGIN();
-  tvm::ffi::TypeTable::Global()->RegisterTypeExtraInfo(type_index, extra_info);
+  tvm::ffi::TypeTable::Global()->RegisterTypeMetadata(type_index, metadata);
   TVM_FFI_SAFE_CALL_END();
 }
 
+int TVMFFITypeRegisterAttr(int32_t type_index, const TVMFFIByteArray* name,
+                           const TVMFFIAny* value) {
+  TVM_FFI_SAFE_CALL_BEGIN();
+  tvm::ffi::TypeTable::Global()->RegisterTypeAttr(type_index, name, value);
+  TVM_FFI_SAFE_CALL_END();
+}
+
+const TVMFFITypeAttrColumn* TVMFFIGetTypeAttrColumn(const TVMFFIByteArray* name) {
+  TVM_FFI_LOG_EXCEPTION_CALL_BEGIN();
+  return tvm::ffi::TypeTable::Global()->GetTypeAttrColumn(name);
+  TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIGetTypeAttrColumn);
+}
+
 int32_t TVMFFITypeGetOrAllocIndex(const TVMFFIByteArray* type_key, int32_t static_type_index,
                                   int32_t type_depth, int32_t num_child_slots,
                                   int32_t child_slots_can_overflow, int32_t parent_type_index) {

ffi/src/ffi/reflection/structural_equal.cc

@@ -89,10 +89,10 @@ class StructEqualHandler {
   bool CompareObject(ObjectRef lhs, ObjectRef rhs) {
     // NOTE: invariant: lhs and rhs are already the same type
     const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(lhs->type_index());
-    if (type_info->extra_info == nullptr) {
+    if (type_info->metadata == nullptr) {
       return lhs.same_as(rhs);
     }
-    auto structural_eq_hash_kind = type_info->extra_info->structural_eq_hash_kind;
+    auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind;
 
     if (structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported ||
         structural_eq_hash_kind == kTVMFFISEqHashKindUniqueInstance) {

ffi/src/ffi/reflection/structural_hash.cc

@@ -82,11 +82,11 @@ class StructuralHashHandler {
   uint64_t HashObject(ObjectRef obj) {
     // NOTE: invariant: lhs and rhs are already the same type
     const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(obj->type_index());
-    if (type_info->extra_info == nullptr) {
+    if (type_info->metadata == nullptr) {
       // Fallback to pointer hash
       return std::hash<const Object*>()(obj.get());
     }
-    auto structural_eq_hash_kind = type_info->extra_info->structural_eq_hash_kind;
+    auto structural_eq_hash_kind = type_info->metadata->structural_eq_hash_kind;
     if (structural_eq_hash_kind == kTVMFFISEqHashKindUnsupported) {
       // Fallback to pointer hash
       return std::hash<const Object*>()(obj.get());

ffi/tests/cpp/test_reflection_accessor.cc

@@ -143,6 +143,11 @@ TEST(Reflection, ForEachFieldInfo) {
   EXPECT_EQ(field_name_to_offset["z"], 16 + sizeof(TVMFFIObject));
 }
 
+TEST(Reflection, TypeAttrColumn) {
+  reflection::TypeAttrColumn size_attr("test.size");
+  EXPECT_EQ(size_attr[TIntObj::_type_index].cast<int>(), sizeof(TIntObj));
+}
+
 TVM_FFI_STATIC_INIT_BLOCK({
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def_method("testing.Int_GetValue", &TIntObj::GetValue);

ffi/tests/cpp/testing_object.h

@@ -83,6 +83,10 @@ inline void TIntObj::RegisterReflection() {
   refl::ObjectDef<TIntObj>()
       .def_ro("value", &TIntObj::value)
       .def_static("static_add", &TInt::StaticAdd, "static add method");
+  // define extra type attributes
+  refl::TypeAttrDef<TIntObj>()
+      .def("test.GetValue", &TIntObj::GetValue)
+      .attr("test.size", sizeof(TIntObj));
 }
 
 class TFloatObj : public TNumberObj {