This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new c85fd42d [ABI] Add begin_index to TypeAttrColumn (#471)
c85fd42d is described below
commit c85fd42df6eae4ae0ec1aaa4ebb67ac859758cf5
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Feb 22 11:15:39 2026 -0500
[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)
---
include/tvm/ffi/any.h | 26 ++++++++++++++----------
include/tvm/ffi/c_api.h | 38 +++++++++++++++++++++++++++--------
include/tvm/ffi/reflection/accessor.h | 6 +++---
python/tvm_ffi/cython/base.pxi | 3 ++-
python/tvm_ffi/cython/object.pxi | 8 ++++++--
rust/tvm-ffi-sys/src/c_api.rs | 11 +++++++---
src/ffi/object.cc | 24 ++++++++++++++++------
tests/cpp/test_reflection.cc | 18 +++++++++++++++++
8 files changed, 100 insertions(+), 34 deletions(-)
diff --git a/include/tvm/ffi/any.h b/include/tvm/ffi/any.h
index 530d9aaa..47a44729 100644
--- a/include/tvm/ffi/any.h
+++ b/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);
+ }
}
}
}
diff --git a/include/tvm/ffi/c_api.h b/include/tvm/ffi/c_api.h
index c195ad30..125fb315 100644
--- a/include/tvm/ffi/c_api.h
+++ b/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;
/*!
diff --git a/include/tvm/ffi/reflection/accessor.h
b/include/tvm/ffi/reflection/accessor.h
index 68c8b0d5..700f7b8c 100644
--- a/include/tvm/ffi/reflection/accessor.h
+++ b/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:
diff --git a/python/tvm_ffi/cython/base.pxi b/python/tvm_ffi/cython/base.pxi
index e80a4802..2b201eee 100644
--- a/python/tvm_ffi/cython/base.pxi
+++ b/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
diff --git a/python/tvm_ffi/cython/object.pxi b/python/tvm_ffi/cython/object.pxi
index c20fe1bf..15978387 100644
--- a/python/tvm_ffi/cython/object.pxi
+++ b/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:
diff --git a/rust/tvm-ffi-sys/src/c_api.rs b/rust/tvm-ffi-sys/src/c_api.rs
index ee555e2a..2035eda9 100644
--- a/rust/tvm-ffi-sys/src/c_api.rs
+++ b/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
diff --git a/src/ffi/object.cc b/src/ffi/object.cc
index cb43adf8..2c3ee1fc 100644
--- a/src/ffi/object.cc
+++ b/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);
diff --git a/tests/cpp/test_reflection.cc b/tests/cpp/test_reflection.cc
index 6570857c..26c7920a 100644
--- a/tests/cpp/test_reflection.cc
+++ b/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);
