[ABI] Add begin_index to TypeAttrColumn (#471)

This PR adds a begin_index field to TypeAttrColumn. The begin_index
enables the type attributes to store narrowly a range of type indices
which can be useful when type attribute is narrowed to specific subscope
where objects are allocated continuously so we can optimize for space
and locality.

As of now the accessor of the TypeAttrColumn is limited to extra/cc so
impact is limited. To be careful, we begin_index is set to 0 for next
few versions and will migrate to nonzero size in 1.0 (so i64 platform
size is compatible)

Author: tqchen

include/tvm/ffi/any.h

@@ -626,12 +626,14 @@ struct AnyHash {
     } else {
       if (src.data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) {
         static const TVMFFITypeAttrColumn* custom_hash_column = GetAnyHashTypeAttrColumn();
-        if (custom_hash_column != nullptr &&
-            static_cast<size_t>(src.data_.type_index) < custom_hash_column->size) {
-          const TVMFFIAny& custom_any_hash = custom_hash_column->data[src.data_.type_index];
-          if (custom_any_hash.type_index != TypeIndex::kTVMFFINone) {
-            return details::StableHashCombine(src.data_.type_index,
-                                              CallCustomAnyHash(custom_any_hash, src));
+        if (custom_hash_column != nullptr) {
+          int32_t offset = src.data_.type_index - custom_hash_column->begin_index;
+          if (offset >= 0 && offset < custom_hash_column->size) {
+            const TVMFFIAny& custom_any_hash = custom_hash_column->data[offset];
+            if (custom_any_hash.type_index != TypeIndex::kTVMFFINone) {
+              return details::StableHashCombine(src.data_.type_index,
+                                                CallCustomAnyHash(custom_any_hash, src));
+            }
           }
         }
       }
@@ -724,11 +726,13 @@ struct AnyEqual {
       }
       if (lhs.data_.type_index >= TypeIndex::kTVMFFIStaticObjectBegin) {
         static const TVMFFITypeAttrColumn* custom_equal_column = GetAnyEqualTypeAttrColumn();
-        if (custom_equal_column != nullptr &&
-            static_cast<size_t>(lhs.data_.type_index) < custom_equal_column->size) {
-          const TVMFFIAny& custom_any_equal = custom_equal_column->data[lhs.data_.type_index];
-          if (custom_any_equal.type_index != TypeIndex::kTVMFFINone) {
-            return CallCustomAnyEqual(custom_any_equal, lhs, rhs);
+        if (custom_equal_column != nullptr) {
+          int32_t offset = lhs.data_.type_index - custom_equal_column->begin_index;
+          if (offset >= 0 && offset < custom_equal_column->size) {
+            const TVMFFIAny& custom_any_equal = custom_equal_column->data[offset];
+            if (custom_any_equal.type_index != TypeIndex::kTVMFFINone) {
+              return CallCustomAnyEqual(custom_any_equal, lhs, rhs);
+            }
           }
         }
       }

include/tvm/ffi/c_api.h

@@ -1070,20 +1070,42 @@ typedef struct {
 } TVMFFITypeMetadata;
 
 /*!
- * \brief Column array that stores extra attributes about types
+ * \brief One column of a type–attribute table: extra attributes keyed by runtime type index.
  *
- * The attributes stored in a column array that can be looked up by type index.
- * Note that the TypeAttr behaves like type_traits so column[T] so not contain
- * attributes from base classes.
+ * TypeAttr is designed to support an open set of possible attributes that can be
+ * registered and queried across different downstream use cases.
+ *
+ * Conceptually, TypeAttr is a dynamic variant of TypeTraits as seen in C++/Rust.
+ * It behaves like c++ type_traits, so column[T] does not contain attributes from base classes.
+ *
+ * Typical use cases for TypeAttr include:
+ * - Storing extra magic trait functions that can customize behaviors like hash/eq/repr.
+ * - Storing extra metadata about type data structures (e.g., mutability) and properties of
+ *   op/operator structures that can be used by compiler transformations (e.g., passes).
+ *
+ * This column covers type indices in the range [begin_index, begin_index + size).
+ * For a given type_index, the corresponding entry is data[type_index - begin_index]
+ * when begin_index <= type_index < begin_index + size; otherwise the entry is
+ * not present (treat as None/null).
  *
- * \note
  * \sa TVMFFIRegisterTypeAttr
  */
 typedef struct {
-  /*! \brief The data of the column. */
+  /*! \brief The data of the column, indexed by (type_index - begin_index). */
   const TVMFFIAny* data;
-  /*! \brief The size of the column. */
-  size_t size;
+  /*!
+   * \brief The number of elements in the data array.
+   *
+   * The column covers type indices in the range [begin_index, begin_index + size).
+   * For a given type_index, the corresponding entry is data[type_index - begin_index]
+   * when begin_index <= type_index < begin_index + size.
+   */
+  int32_t size;
+  /*!
+   * \brief The starting type index of the column data.
+   * A value of 0 means the data array starts at type_index 0.
+   */
+  int32_t begin_index;
 } TVMFFITypeAttrColumn;
 
 /*!

include/tvm/ffi/reflection/accessor.h

@@ -154,12 +154,12 @@ class TypeAttrColumn {
    * \return The type attribute column.
    */
   AnyView operator[](int32_t type_index) const {
-    size_t tindex = static_cast<size_t>(type_index);
-    if (tindex >= column_->size) {
+    int32_t offset = type_index - column_->begin_index;
+    if (offset < 0 || offset >= column_->size) {
       return AnyView();
     }
     const AnyView* any_view_data = reinterpret_cast<const AnyView*>(column_->data);
-    return any_view_data[tindex];
+    return any_view_data[offset];
   }
 
  private:

python/tvm_ffi/cython/base.pxi

@@ -249,7 +249,8 @@ cdef extern from "tvm/ffi/c_api.h":
 
     ctypedef struct TVMFFITypeAttrColumn:
         const TVMFFIAny* data
-        size_t size
+        int32_t size
+        int32_t begin_index
 
     int TVMFFIObjectDecRef(TVMFFIObjectHandle obj) nogil
     int TVMFFIObjectIncRef(TVMFFIObjectHandle obj) nogil

python/tvm_ffi/cython/object.pxi

@@ -535,9 +535,13 @@ def _lookup_type_attr(type_index: int32_t, attr_key: str) -> Any:
     cdef ByteArrayArg attr_key_bytes = ByteArrayArg(c_str(attr_key))
     cdef const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&attr_key_bytes.cdata)
     cdef TVMFFIAny data
-    if column == NULL or column.size <= type_index:
+    cdef int32_t offset
+    if column == NULL:
         return None
-    return make_ret(column.data[type_index])
+    offset = type_index - column.begin_index
+    if offset < 0 or offset >= column.size:
+        return None
+    return make_ret(column.data[offset])
 
 
 def _type_cls_to_type_info(type_cls: type) -> TypeInfo | None:

rust/tvm-ffi-sys/src/c_api.rs

@@ -339,10 +339,15 @@ pub struct TVMFFITypeMetadata {
 /// attributes from base classes.
 #[repr(C)]
 pub struct TVMFFITypeAttrColumn {
-    /// The data of the column
+    /// The data of the column, indexed by (type_index - begin_index).
     pub data: *const TVMFFIAny,
-    /// The size of the column
-    pub size: usize,
+    /// The number of elements in the data array.
+    /// The column covers type indices [begin_index, begin_index + size).
+    pub size: i32,
+    /// The starting type index of the column data.
+    /// Lookup: if begin_index <= type_index < begin_index + size,
+    /// the entry is data[(type_index - begin_index) as usize].
+    pub begin_index: i32,
 }
 
 /// Runtime type information for object type checking

src/ffi/object.cc

@@ -263,17 +263,29 @@ class TypeTable {
       column_index = (*it).second;
     }
     TypeAttrColumnData* column = type_attr_columns_[column_index].get();
-    if (column->data_.size() < static_cast<size_t>(type_index) + 1) {
+    if (type_index == kTVMFFINone) {
+      // Sentinel: just ensure the column exists without registering a value.
+      if (column->data_.empty()) {
+        column->data = reinterpret_cast<const TVMFFIAny*>(column->data_.data());
+        column->size = 0;
+        column->begin_index = 0;
+      }
+      return;
+    }
+    // TODO(1.0): set begin_index to first registered type_index for sparse column storage
+    // For now, begin_index is always 0 (resize from index 0).
+    if (static_cast<size_t>(type_index) >= column->data_.size()) {
+      // Extend back from index 0.
       column->data_.resize(static_cast<size_t>(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) {
+    column->data = reinterpret_cast<const TVMFFIAny*>(column->data_.data());
+    column->size = static_cast<int32_t>(column->data_.size());
+    column->begin_index = 0;
+    if (column->data_[type_index - column->begin_index] != nullptr) {
       TVM_FFI_THROW(RuntimeError) << "Type attribute `" << name_str << "` is already set for type `"
                                   << TypeIndexToTypeKey(type_index) << "`";
     }
-    column->data_[type_index] = value_view;
+    column->data_[type_index - column->begin_index] = value_view;
   }
   const TVMFFITypeAttrColumn* GetTypeAttrColumn(const TVMFFIByteArray* name) {
     String name_str(*name);

tests/cpp/test_reflection.cc

@@ -179,6 +179,24 @@ TEST(Reflection, TypeAttrColumn) {
   EXPECT_EQ(size_attr[TIntObj::RuntimeTypeIndex()].cast<int>(), sizeof(TIntObj));
 }
 
+TEST(Reflection, TypeAttrColumnBeginIndex) {
+  // Get the column and verify begin_index
+  TVMFFIByteArray attr_name = {"test.size", std::char_traits<char>::length("test.size")};
+  const TVMFFITypeAttrColumn* column = TVMFFIGetTypeAttrColumn(&attr_name);
+  ASSERT_NE(column, nullptr);
+  // begin_index should be >= 0
+  EXPECT_GE(column->begin_index, 0);
+  // size should cover the range from begin_index
+  EXPECT_GT(column->size, 0);
+  // verify that lookup of a type_index below begin_index returns None
+  reflection::TypeAttrColumn size_attr("test.size");
+  AnyView result = size_attr[0];  // index 0 is kTVMFFINone, unlikely to have this attr
+  (void)result;  // suppress unused variable warning; we only verify no crash occurs
+  // The result may or may not be None depending on begin_index; the key is no crash.
+  // verify the known registered entry still works
+  EXPECT_EQ(size_attr[TIntObj::RuntimeTypeIndex()].cast<int>(), sizeof(TIntObj));
+}
+
 TVM_FFI_STATIC_INIT_BLOCK() {
   namespace refl = tvm::ffi::reflection;
   refl::GlobalDef().def_method("testing.Int_GetValue", &TIntObj::GetValue);