[REFACTOR] Collapse custom-allocator setter API to single global entry point

Drop the unused per-type TVMFFISetCustomAllocator(int32_t,...) and rename
TVMFFISetDefaultCustomAllocator to TVMFFISetCustomAllocator. Drop the
type_index parameter from TVMFFIGetCustomAllocator since lookup no
longer needs per-type resolution. The registry collapses to a single
atomic; hot path becomes a single acquire load. Per-type opt-out for
small heap-primitive types remains a clean future addition via a
type-trait dispatched in Handler::New.

Author: cyx-6

include/tvm/ffi/c_api.h

@@ -622,8 +622,7 @@ typedef struct {
 } TVMFFICustomAllocHeader;
 
 /*!
- * \brief Custom allocator entry registered with TVMFFISetCustomAllocator /
- *        TVMFFISetDefaultCustomAllocator.
+ * \brief Custom allocator entry registered with TVMFFISetCustomAllocator.
  */
 typedef struct {
   /*!
@@ -638,36 +637,25 @@ typedef struct {
 } TVMFFICustomAllocator;
 
 /*!
- * \brief Get the custom allocator registered for ``type_index``.
+ * \brief Get the process-wide custom allocator.
  *
- * Resolution order: per-type slot (if non-NULL) -> global default allocator
- * -> NULL (no custom allocator; caller takes the default malloc path).
+ * libtvm_ffi installs a builtin default at registry init, so the result is
+ * never NULL in practice. Frontends can override the global allocator via
+ * TVMFFISetCustomAllocator.
  *
- * \param type_index The type index of the Object being allocated, or -1 to
- *                   query the global default directly.
- * \return The allocator, or NULL when none is registered.
+ * \return The currently registered allocator.
  */
-TVM_FFI_DLL TVMFFICustomAllocator* TVMFFIGetCustomAllocator(int32_t type_index);
+TVM_FFI_DLL TVMFFICustomAllocator* TVMFFIGetCustomAllocator(void);
 
 /*!
- * \brief Register an allocator for a specific type index. Pass
- *        ``allocator==NULL`` to clear the slot and fall back to the global
- *        default.
+ * \brief Register the process-wide custom allocator. Pass ``allocator==NULL``
+ *        to restore the builtin default.
  *
- * \param type_index The type index whose slot is being installed.
- * \param allocator  Pointer to a TVMFFICustomAllocator with process-stable
- *                   lifetime.
+ * \param allocator Pointer to a TVMFFICustomAllocator with process-stable
+ *                  lifetime, or NULL to reset to the builtin default.
  * \return 0 on success, nonzero on failure.
  */
-TVM_FFI_DLL int TVMFFISetCustomAllocator(int32_t type_index, TVMFFICustomAllocator* allocator);
-
-/*!
- * \brief Register a process-wide default allocator used for any type without
- *        a per-type override. Pass ``allocator==NULL`` to clear it.
- *
- * \return 0 on success, nonzero on failure.
- */
-TVM_FFI_DLL int TVMFFISetDefaultCustomAllocator(TVMFFICustomAllocator* allocator);
+TVM_FFI_DLL int TVMFFISetCustomAllocator(TVMFFICustomAllocator* allocator);
 
 /*!
  * \brief Convert type key to type index.

include/tvm/ffi/memory.h

@@ -148,8 +148,8 @@ class ObjAllocatorBase {
     // never null and every Object always has a TVMFFICustomAllocHeader
     // sitting immediately before it. Frontends with richer per-allocation
     // bookkeeping (e.g. Python's PyCustomAllocHeader) override the global
-    // default with their own allocator.
-    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator(T::RuntimeTypeIndex());
+    // allocator with their own.
+    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
     T* ptr = Handler::New(static_cast<Derived*>(this), alloc, std::forward<Args>(args)...);
     TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr);
     ffi_ptr->combined_ref_count = kCombinedRefCountBothOne;
@@ -171,7 +171,7 @@ class ObjAllocatorBase {
     using Handler = typename Derived::template ArrayHandler<ArrayType, ElemType>;
     static_assert(std::is_base_of_v<Object, ArrayType>,
                   "make_inplace_array can only be used to create Object");
-    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator(ArrayType::RuntimeTypeIndex());
+    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
     ArrayType* ptr = Handler::New(static_cast<Derived*>(this), alloc, num_elems,
                                   std::forward<Args>(args)...);
     TVMFFIObject* ffi_ptr = details::ObjectUnsafe::GetHeader(ptr);

python/tvm_ffi/cython/base.pxi

@@ -173,9 +173,8 @@ cdef extern from "tvm/ffi/c_api.h":
     ctypedef struct TVMFFICustomAllocator:
         void* (*allocate)(size_t size, size_t alignment, int32_t type_index)
 
-    TVMFFICustomAllocator* TVMFFIGetCustomAllocator(int32_t type_index) nogil
-    int TVMFFISetCustomAllocator(int32_t type_index, TVMFFICustomAllocator* allocator) nogil
-    int TVMFFISetDefaultCustomAllocator(TVMFFICustomAllocator* allocator) nogil
+    TVMFFICustomAllocator* TVMFFIGetCustomAllocator() nogil
+    int TVMFFISetCustomAllocator(TVMFFICustomAllocator* allocator) nogil
 
     ctypedef struct TVMFFIAny:
         int32_t type_index

python/tvm_ffi/cython/tvm_ffi_python_helpers.h

@@ -151,7 +151,7 @@ TVM_FFI_INLINE PyCustomAllocHeader* TVMFFIPyHeader(void* tptr) {
 inline void TVMFFIPyDeleteSpace(void* tptr);
 
 /*!
- * \brief Allocator entry registered with TVMFFISetDefaultCustomAllocator at
+ * \brief Allocator entry registered with TVMFFISetCustomAllocator at
  *        Cython module init. Allocates ``kPyHeaderOffset + size`` bytes,
  *        zero-inits the prepended PyCustomAllocHeader (in particular
  *        ``py_object = NULL``), wires ``base.delete_space`` to
@@ -201,7 +201,7 @@ inline void TVMFFIPyDeleteSpace(void* tptr) {
  */
 TVM_FFI_INLINE int TVMFFIPyRegisterDefaultAllocator() {
   static TVMFFICustomAllocator allocator{&TVMFFIPyAllocate};
-  return TVMFFISetDefaultCustomAllocator(&allocator);
+  return TVMFFISetCustomAllocator(&allocator);
 }
 
 /*!

src/ffi/custom_allocator.cc

@@ -18,14 +18,14 @@
  */
 /*
  * \file src/ffi/custom_allocator.cc
- * \brief Process-wide registry for custom Object allocators.
+ * \brief Process-wide registry for the custom Object allocator.
  *
  * Every Object allocation goes through the registered allocator. libtvm_ffi
  * installs a builtin default allocator (BuiltinDefaultAllocate / DeleteSpace)
  * at registry construction so the hot path is never null and every Object
  * carries at least the base TVMFFICustomAllocHeader. Frontends (Python via
  * the Cython extension, future Rust binding, etc.) override the global
- * default with TVMFFISetDefaultCustomAllocator when they need richer
+ * allocator with TVMFFISetCustomAllocator when they need richer
  * per-allocation bookkeeping.
  */
 #include <tvm/ffi/base_details.h>
@@ -35,10 +35,7 @@
 
 #include <atomic>
 #include <cstdlib>
-#include <cstring>
-#include <mutex>
 #include <new>
-#include <unordered_map>
 
 namespace tvm {
 namespace ffi {
@@ -114,39 +111,16 @@ class CustomAllocatorRegistry {
     // Install the builtin default at construction so the very first
     // Object allocation in the process (e.g. inside a static initializer
     // that registers a global function) finds a usable allocator.
-    static TVMFFICustomAllocator builtin{&BuiltinDefaultAllocate};
-    default_.store(&builtin, std::memory_order_release);
+    current_.store(BuiltinDefault(), std::memory_order_release);
   }
 
-  TVMFFICustomAllocator* Get(int32_t idx) const {
-    // Hot path: when no per-type override has been registered, skip the
-    // mutex entirely and just read the global default. The size counter is
-    // incremented under mu_ before the first per-type entry is visible, so
-    // an acquire load here pairs with the release store in Set.
-    if (per_type_size_.load(std::memory_order_acquire) == 0) {
-      return default_.load(std::memory_order_acquire);
-    }
-    if (idx >= 0) {
-      std::lock_guard<std::mutex> g(mu_);
-      auto it = per_type_.find(idx);
-      if (it != per_type_.end()) return it->second;
-    }
-    return default_.load(std::memory_order_acquire);
+  TVMFFICustomAllocator* Get() const {
+    return current_.load(std::memory_order_acquire);
   }
 
-  void Set(int32_t idx, TVMFFICustomAllocator* allocator) {
-    if (idx < 0) return;
-    std::lock_guard<std::mutex> g(mu_);
-    if (allocator != nullptr) {
-      per_type_[idx] = allocator;
-    } else {
-      per_type_.erase(idx);
-    }
-    per_type_size_.store(per_type_.size(), std::memory_order_release);
-  }
-
-  void SetDefault(TVMFFICustomAllocator* allocator) {
-    default_.store(allocator, std::memory_order_release);
+  void Set(TVMFFICustomAllocator* allocator) {
+    current_.store(allocator != nullptr ? allocator : BuiltinDefault(),
+                   std::memory_order_release);
   }
 
   static CustomAllocatorRegistry* Global() {
@@ -155,30 +129,26 @@ class CustomAllocatorRegistry {
   }
 
  private:
-  std::atomic<TVMFFICustomAllocator*> default_{nullptr};
-  std::atomic<size_t> per_type_size_{0};
-  mutable std::mutex mu_;
-  std::unordered_map<int32_t, TVMFFICustomAllocator*> per_type_;
+  static TVMFFICustomAllocator* BuiltinDefault() {
+    static TVMFFICustomAllocator builtin{&BuiltinDefaultAllocate};
+    return &builtin;
+  }
+
+  std::atomic<TVMFFICustomAllocator*> current_{nullptr};
 };
 
 }  // namespace
 }  // namespace ffi
 }  // namespace tvm
 
-TVMFFICustomAllocator* TVMFFIGetCustomAllocator(int32_t type_index) {
+TVMFFICustomAllocator* TVMFFIGetCustomAllocator(void) {
   TVM_FFI_LOG_EXCEPTION_CALL_BEGIN();
-  return tvm::ffi::CustomAllocatorRegistry::Global()->Get(type_index);
+  return tvm::ffi::CustomAllocatorRegistry::Global()->Get();
   TVM_FFI_LOG_EXCEPTION_CALL_END(TVMFFIGetCustomAllocator);
 }
 
-int TVMFFISetCustomAllocator(int32_t type_index, TVMFFICustomAllocator* allocator) {
-  TVM_FFI_SAFE_CALL_BEGIN();
-  tvm::ffi::CustomAllocatorRegistry::Global()->Set(type_index, allocator);
-  TVM_FFI_SAFE_CALL_END();
-}
-
-int TVMFFISetDefaultCustomAllocator(TVMFFICustomAllocator* allocator) {
+int TVMFFISetCustomAllocator(TVMFFICustomAllocator* allocator) {
   TVM_FFI_SAFE_CALL_BEGIN();
-  tvm::ffi::CustomAllocatorRegistry::Global()->SetDefault(allocator);
+  tvm::ffi::CustomAllocatorRegistry::Global()->Set(allocator);
   TVM_FFI_SAFE_CALL_END();
 }

src/ffi/extra/dataclass.cc

@@ -2051,7 +2051,7 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index, int32_t total_size) {
     // Both are valid initial states whose destructors / assignment
     // operators handle correctly, so no placement construction is needed.
     size_t alloc_size = static_cast<size_t>(total_size);
-    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator(type_index);
+    TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
     void* obj_ptr = alloc->allocate(alloc_size, alignof(::std::max_align_t), type_index);
     std::memset(obj_ptr, 0, alloc_size);
     TVMFFIObject* ffi_obj = reinterpret_cast<TVMFFIObject*>(obj_ptr);
@@ -2070,7 +2070,7 @@ void PyClassRegisterTypeAttrColumns(int32_t type_index, int32_t total_size) {
       Function::FromTyped([type_index, total_size, type_info](const Object* src) -> ObjectRef {
         // Mirror RegisterFFINew above.
         size_t alloc_size = static_cast<size_t>(total_size);
-        TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator(type_index);
+        TVMFFICustomAllocator* alloc = TVMFFIGetCustomAllocator();
         void* obj_ptr = alloc->allocate(alloc_size, alignof(::std::max_align_t), type_index);
         std::memset(obj_ptr, 0, alloc_size);
         TVMFFIObject* ffi_obj = reinterpret_cast<TVMFFIObject*>(obj_ptr);