From 320aa735534ba8c36336073c4ecb2e5d0292af5e Mon Sep 17 00:00:00 2001 From: tqchen Date: Thu, 18 Jun 2026 17:08:14 +0000 Subject: [PATCH] [OBJECT] Optimize dynamic RuntimeTypeIndex loads Cache dynamically allocated type indices in a static inline const _type_index and make RuntimeTypeIndex return the cached value directly. This moves _GetOrAllocRuntimeTypeIndex to cold/startup code and reduces the hot accessor to a plain load. A side codegen probe shows ReadDynTypeIndex shrinking from 176 B with guard/allocation checks to 14 B, while total probe object size drops from 769 B to 634 B despite the added _type_index guard/init storage. --- include/tvm/ffi/enum.h | 8 +++++--- include/tvm/ffi/object.h | 13 +++++++------ include/tvm/ffi/reflection/enum_def.h | 9 +++++---- include/tvm/ffi/reflection/registry.h | 2 +- 4 files changed, 18 insertions(+), 14 deletions(-) diff --git a/include/tvm/ffi/enum.h b/include/tvm/ffi/enum.h index 0cd97d993..a42a18226 100644 --- a/include/tvm/ffi/enum.h +++ b/include/tvm/ffi/enum.h @@ -131,16 +131,18 @@ inline Enum EnumObj::Get(const String& name) { static_assert(std::is_base_of_v, "EnumObj::Get requires T to be a subclass of EnumObj"); const TVMFFITypeAttrColumn* column = GetEnumEntriesColumn(); - int32_t type_index = EnumClsObj::RuntimeTypeIndex(); + int32_t type_index = EnumClsObj::_GetOrAllocRuntimeTypeIndex(); if (column != nullptr) { int32_t offset = type_index - column->begin_index; if (offset >= 0 && offset < column->size) { const TVMFFIAny* stored = &column->data[offset]; if (stored->type_index != kTVMFFINone) { - Dict entries = AnyView::CopyFromTVMFFIAny(*stored).cast>(); + Dict entries = + AnyView::CopyFromTVMFFIAny(*stored).cast>(); auto it = entries.find(name); if (it != entries.end()) { - return (*it).second; + return details::ObjectUnsafe::ObjectRefFromObjectPtr( + details::ObjectUnsafe::ObjectPtrFromObjectRef((*it).second)); } } } diff --git a/include/tvm/ffi/object.h b/include/tvm/ffi/object.h index 262c5708f..71c69f16f 100644 --- a/include/tvm/ffi/object.h +++ b/include/tvm/ffi/object.h @@ -307,11 +307,11 @@ class Object { * \brief Get the runtime allocated type index of the type * \note Getting this information may need dynamic calls into a global table. */ - static int32_t RuntimeTypeIndex() { return TypeIndex::kTVMFFIObject; } + TVM_FFI_INLINE static int32_t RuntimeTypeIndex() noexcept { return TypeIndex::kTVMFFIObject; } /*! * \brief Internal function to get or allocate a runtime index. */ - static int32_t _GetOrAllocRuntimeTypeIndex() { // NOLINT(bugprone-reserved-identifier) + TVM_FFI_COLD_CODE static int32_t _GetOrAllocRuntimeTypeIndex() { // NOLINT(*) return TypeIndex::kTVMFFIObject; } @@ -936,7 +936,7 @@ struct ObjectPtrEqual { */ #define TVM_FFI_DECLARE_OBJECT_INFO_STATIC(TypeKey, TypeName, ParentType) \ static constexpr int32_t _type_depth = ParentType::_type_depth + 1; \ - static int32_t _GetOrAllocRuntimeTypeIndex() { \ + TVM_FFI_COLD_CODE static int32_t _GetOrAllocRuntimeTypeIndex() { \ static_assert(!ParentType::_type_final, "ParentType marked as final"); \ static_assert(TypeName::_type_child_slots == 0 || ParentType::_type_child_slots == 0 || \ TypeName::_type_child_slots < ParentType::_type_child_slots, \ @@ -948,7 +948,7 @@ struct ObjectPtrEqual { TypeName::_type_child_slots_can_overflow, ParentType::_GetOrAllocRuntimeTypeIndex()); \ return TypeName::_type_index; \ } \ - static int32_t RuntimeTypeIndex() { return TypeName::_type_index; } \ + TVM_FFI_INLINE static int32_t RuntimeTypeIndex() noexcept { return TypeName::_type_index; } \ static constexpr const char* _type_key = TypeKey /*! @@ -959,7 +959,7 @@ struct ObjectPtrEqual { */ #define TVM_FFI_DECLARE_OBJECT_INFO_PREDEFINED_TYPE_KEY(TypeName, ParentType) \ static constexpr int32_t _type_depth = ParentType::_type_depth + 1; \ - static int32_t _GetOrAllocRuntimeTypeIndex() { \ + TVM_FFI_COLD_CODE static int32_t _GetOrAllocRuntimeTypeIndex() { \ static_assert(!ParentType::_type_final, "ParentType marked as final"); \ static_assert(TypeName::_type_child_slots == 0 || ParentType::_type_child_slots == 0 || \ TypeName::_type_child_slots < ParentType::_type_child_slots, \ @@ -971,7 +971,8 @@ struct ObjectPtrEqual { TypeName::_type_child_slots_can_overflow, ParentType::_GetOrAllocRuntimeTypeIndex()); \ return tindex; \ } \ - static int32_t RuntimeTypeIndex() { return _GetOrAllocRuntimeTypeIndex(); } + static inline const int32_t _type_index = TypeName::_GetOrAllocRuntimeTypeIndex(); \ + TVM_FFI_INLINE static int32_t RuntimeTypeIndex() noexcept { return TypeName::_type_index; } /*! * \brief Helper macro to declare object information with dynamic type index. diff --git a/include/tvm/ffi/reflection/enum_def.h b/include/tvm/ffi/reflection/enum_def.h index 6e1e91c2e..1970d24b6 100644 --- a/include/tvm/ffi/reflection/enum_def.h +++ b/include/tvm/ffi/reflection/enum_def.h @@ -74,8 +74,8 @@ class EnumDef : public ReflectionDefBase { * \param instance_name The instance's string name (e.g., ``"Add"``). */ explicit EnumDef(const char* instance_name) - : type_index_(EnumClsObj::RuntimeTypeIndex()), name_(instance_name) { - Dict entries = EnsureEntriesDict(); + : type_index_(EnumClsObj::_GetOrAllocRuntimeTypeIndex()), name_(instance_name) { + Dict entries = EnsureEntriesDict(); String name_str(name_); if (entries.count(name_str) != 0) { TVM_FFI_THROW(RuntimeError) << "Duplicate enum entry `" << name_ << "` for type `" @@ -83,6 +83,7 @@ class EnumDef : public ReflectionDefBase { } ordinal_ = static_cast(entries.size()); ObjectPtr obj = make_object(); + ::tvm::ffi::details::ObjectUnsafe::GetHeader(obj.get())->type_index = type_index_; obj->_value = ordinal_; obj->_name = name_str; instance_ = Enum(ObjectPtr(std::move(obj))); @@ -134,8 +135,8 @@ class EnumDef : public ReflectionDefBase { int64_t ordinal() const { return ordinal_; } private: - Dict EnsureEntriesDict() { - return EnsureDict>(type_attr::kEnumEntries); + Dict EnsureEntriesDict() { + return EnsureDict>(type_attr::kEnumEntries); } Dict> EnsureAttrsDict() { diff --git a/include/tvm/ffi/reflection/registry.h b/include/tvm/ffi/reflection/registry.h index f810e5ff0..2acdce754 100644 --- a/include/tvm/ffi/reflection/registry.h +++ b/include/tvm/ffi/reflection/registry.h @@ -992,7 +992,7 @@ class TypeAttrDef : public ReflectionDefBase { */ template explicit TypeAttrDef(ExtraArgs&&... extra_args) - : type_index_(Class::RuntimeTypeIndex()), type_key_(Class::_type_key) {} + : type_index_(Class::_GetOrAllocRuntimeTypeIndex()), type_key_(Class::_type_key) {} /*! * \brief Define a function-valued type attribute.