fix(py_class): remove __ffi_init_inplace__ to fix memory leak (#551)

## Summary

- **Fix memory leak** in `@py_class`: the custom `__new__` override
allocated a spurious C++ object via `__ffi_new__` whenever a py_class
instance was returned from C++, which was then orphaned when
`make_ret_object` overwrote `chandle` with the actual return value.
- **Remove `__ffi_init_inplace__`** entirely — both the C++
`MakeInitInplace` function, its registration, and the `kInitInplace`
constant.
- **Unify `c_class` and `py_class`** auto-generated `__init__` to both
use `self.__init_handle_by_constructor__(ffi_init, *args)`, differing
only in guard logic: py_class uses a `chandle != 0` guard (no-op when
already allocated), c_class uses a `type_info` identity guard
(`TypeError` from subclass).
- **Remove the custom `__new__` override** — no more spurious allocation
on the `make_ret` path.
- **Wrap user-defined `__init__`** on py_class with a `chandle == 0`
guard that pre-allocates via `__ffi_new__` before user code runs; use
`functools.wraps` to preserve metadata.

### Files changed

| File | Change |
|------|--------|
| `python/tvm_ffi/_dunder.py` | Core redesign: removed `inplace` param,
added `py_class_mode`, removed `__new__` override, added user-init
wrapping with `functools.wraps` |
| `python/tvm_ffi/registry.py` | Removed `inplace=False` kwarg from
`_make_init` call |
| `include/tvm/ffi/reflection/accessor.h` | Removed `kInitInplace`
constant |
| `include/tvm/ffi/reflection/registry.h` | Updated comments (removed
`__ffi_init_inplace__` references) |
| `src/ffi/extra/dataclass.cc` | Removed `MakeInitInplace()`, its
registration, `EnsureTypeAttrColumn` call |
| `python/tvm_ffi/dataclasses/py_class.py` | Updated comment |
| `tests/python/test_dataclass_init.py` | Removed `TestInitInplace`
class (5 tests) |
| `tests/python/test_dataclass_py_class.py` | Added 8 new tests in
`TestPyClassNoLeak` |

## Test plan

- [x] All 27 pre-commit linters pass
- [x] All 2103 Python tests pass (2095 existing + 8 new)
- [x] All 349 C++ tests pass
- [x] Memory leak validation: 20k construct+deepcopy cycles with no leak

Author: junrushao

include/tvm/ffi/reflection/accessor.h

@@ -340,19 +340,6 @@ inline constexpr const char* kNew = "__ffi_new__";
  * Keyword arguments are packed as ``[KWARGS, key0, val0, key1, val1, ...]``.
  */
 inline constexpr const char* kInit = "__ffi_init__";
-/*!
- * \brief In-place init on a pre-allocated object (no allocation).
- *
- * Used by ``@py_class`` auto-generated ``__init__``: ``__new__`` has already
- * allocated the object via ``kNew``, so this function only sets fields.
- * The first argument is ``self`` (the pre-allocated object).
- *
- * Signature: ``(self: TSelf, *args, **kwargs) -> void``, where ``TSelf`` is a subclass of
- * ObjectRef.
- *
- * Keyword arguments are packed as ``[KWARGS, key0, val0, key1, val1, ...]``.
- */
-inline constexpr const char* kInitInplace = "__ffi_init_inplace__";
 /*!
  * \brief Convert ``AnyView`` to a specific reflected ``TSelf`` type.
  *

include/tvm/ffi/reflection/registry.h

@@ -726,7 +726,7 @@ class ObjectDef : public ReflectionDefBase {
       TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &fn_any));
     }
     // Step 2. Register `__ffi_new__` <== info->metadata->creator
-    // Also, `__ffi_init__` and `__ffi_init_inplace__` if no explicit init is defined.
+    // Also, `__ffi_init__` if no explicit init is defined.
     if (info->metadata != nullptr && info->metadata->creator != nullptr) {
       Function fn = Function::FromTyped(
           [creator = info->metadata->creator]() -> ObjectRef {
@@ -739,7 +739,7 @@ class ObjectDef : public ReflectionDefBase {
       TVMFFIByteArray attr = AsByteArray(type_attr::kNew);
       TVMFFIAny fn_any = AnyView(fn).CopyToTVMFFIAny();
       TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index_, &attr, &fn_any));
-      // Step 3. Register `__ffi_init__` and `__ffi_init_inplace__` if no explicit init is defined.
+      // Step 3. Register `__ffi_init__` if no explicit init is defined.
       // Use Function::GetGlobal to look up the registration function, which lives in
       // dataclass.cc and may not be loaded yet for early (builtin) types.
       if (!has_explicit_init_) {

python/tvm_ffi/_dunder.py

@@ -31,42 +31,49 @@
 def _make_init(
     type_cls: type,
     type_info: TypeInfo,
-    inplace: bool,
     ffi_init: Function,
+    py_class_mode: bool = False,
 ) -> Callable[..., None]:
-    """Build ``__init__`` that delegates to a C++ init function.
+    """Build ``__init__`` that delegates to ``__ffi_init__``.
+
+    Both ``@c_class`` and ``@py_class`` use the same constructor-call path
+    (``self.__init_handle_by_constructor__(ffi_init, *args)``).  The only
+    difference is how ``super().__init__()`` from a subclass is handled:
+
+    * **c_class** — raises ``TypeError`` (subclass must define its own init).
+    * **py_class** — silently skips when the C++ handle is already set, so
+      ``super().__init__()`` is a harmless no-op.
 
     Parameters
     ----------
     type_cls
-        The class to build an __init__ for.  Used for signature and error messages.
+        The class to build an __init__ for.
     type_info
-        The TypeInfo for *type_cls*, used to build the signature and for type checks.
-    inplace : bool
-        If True (py_class), *ffi_init* is ``__ffi_init_inplace__`` — called as
-        ``ffi_init(self, *args)`` on a pre-allocated object.
-        If False (c_class), *ffi_init* is ``__ffi_init__`` — called via
-        ``self.__init_handle_by_constructor__(ffi_init, *args)``.
+        The TypeInfo for *type_cls*.
     ffi_init
-        The C++ initialiser resolved from TypeAttrColumn at install time.
+        The C++ ``__ffi_init__`` resolved at install time.
+    py_class_mode
+        If True, use a ``chandle`` guard instead of a ``TypeError`` guard.
 
     """
     sig = _make_init_signature(type_info)
     kwargs_obj = core.KWARGS
     missing = core.MISSING
     has_post_init = hasattr(type_cls, "__post_init__")
 
-    if inplace:
+    if py_class_mode:
 
         def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
-            ffi_args: list[Any] = [self, *args]
+            if self.__chandle__() != 0:
+                return
+            ffi_args: list[Any] = list(args)
             if kwargs:
                 ffi_args.append(kwargs_obj)
                 for key, val in kwargs.items():
                     if val is not missing:
                         ffi_args.append(key)
                         ffi_args.append(val)
-            ffi_init(*ffi_args)
+            self.__init_handle_by_constructor__(ffi_init, *ffi_args)
             if has_post_init:
                 self.__post_init__()
 
@@ -244,16 +251,15 @@ def _install_dataclass_dunders(  # noqa: PLR0912, PLR0915
     unsafe_hash
         If True, install ``__hash__`` using ``RecursiveHash``.
     py_class_mode
-        If True, use C++ ``__ffi_init_inplace__`` for ``__init__``
-        (object already allocated by ``__new__``).
+        If True, use a ``chandle`` guard for ``__init__`` so that
+        ``super().__init__()`` is a no-op, and wrap user-defined
+        ``__init__`` to allocate via ``__ffi_init__`` before user code.
 
     """
     from . import _ffi_api  # noqa: PLC0415
 
     type_info: TypeInfo = cls.__tvm_ffi_type_info__  # type: ignore[assignment]
     type_index: int = type_info.type_index
-    ffi_new: Function | None = core._lookup_type_attr(type_index, "__ffi_new__")
-    ffi_init_inplace: Function | None = core._lookup_type_attr(type_index, "__ffi_init_inplace__")
     # Look up __ffi_init__ from TypeMethod (preferred) or TypeAttrColumn (fallback).
     ffi_init: Function | None = None
     for method in type_info.methods:
@@ -262,45 +268,25 @@ def _install_dataclass_dunders(  # noqa: PLR0912, PLR0915
             break
     if ffi_init is None:
         ffi_init = core._lookup_type_attr(type_index, "__ffi_init__")
+    ffi_new: Function | None = core._lookup_type_attr(type_index, "__ffi_new__")
     ffi_shallow_copy: Function | None = core._lookup_type_attr(type_index, "__ffi_shallow_copy__")
-    pyobject_new = core.Object.__new__
-
-    # __new__ (py_class only: allocate via __ffi_new__)
-    if py_class_mode and "__new__" not in cls.__dict__:
-        if ffi_new is not None:
-
-            def __new__(klass: type, *args: Any, **kwargs: Any) -> Any:
-                obj = pyobject_new(klass)
-                obj.__init_handle_by_constructor__(ffi_new)
-                return obj
-
-            cls.__new__ = __new__  # type: ignore[attr-defined]
 
     # __init__
     # ┌─────────────────────┬──────────────────────┬──────────────────────┐
     # │                     │ user __init__        │ no user __init__     │
     # ├─────────────────────┼──────────────────────┼──────────────────────┤
     # │ c_class, init=True  │ keep as-is           │ _make_init           │
     # │ c_class, init=False │ keep as-is           │ TypeError guard      │
-    # │ py_class, init=True │ keep as-is           │ _make_init(inplace)  │
-    # │ py_class, init=False│ keep as-is           │ TypeError guard      │
+    # │ py_class, init=True │ wrap chandle guard   │ _make_init(py_class) │
+    # │ py_class, init=False│ wrap chandle guard   │ TypeError guard      │
     # └─────────────────────┴──────────────────────┴──────────────────────┘
     if "__init__" not in cls.__dict__:
-        if init and py_class_mode and ffi_init_inplace is not None:
-            # py_class init=True: delegate to C++ __ffi_init_inplace__
-            cls.__init__ = _make_init(  # type: ignore[attr-defined]
-                cls,
-                type_info,
-                ffi_init=ffi_init_inplace,
-                inplace=True,
-            )
-        elif init and not py_class_mode and ffi_init is not None:
-            # c_class init=True: delegate to __ffi_init__ via TypeAttrColumn
+        if init and ffi_init is not None:
             cls.__init__ = _make_init(  # type: ignore[attr-defined]
                 cls,
                 type_info,
                 ffi_init=ffi_init,
-                inplace=False,
+                py_class_mode=py_class_mode,
             )
         elif not init:
             # init=False, no user __init__: TypeError guard
@@ -315,6 +301,24 @@ def __init___(self: Any, *args: Any, **kwargs: Any) -> None:
             __init___.__qualname__ = f"{cls.__qualname__}.__init__"
             __init___.__module__ = cls.__module__
             cls.__init__ = __init___  # type: ignore[attr-defined]
+    elif py_class_mode and ffi_new is not None:
+        # User-defined __init__: wrap with chandle guard so the C++ object
+        # is allocated (via __ffi_new__) before the user's code runs.
+        # We use __ffi_new__ (zero-arg allocator) rather than __ffi_init__
+        # because the user's __init__ signature may not match the field
+        # layout.  super().__init__() from within the user's code becomes
+        # a no-op because chandle is already set.
+        import functools  # noqa: PLC0415
+
+        user_init = cls.__dict__["__init__"]
+
+        @functools.wraps(user_init)
+        def __init__(self: Any, *args: Any, **kwargs: Any) -> None:
+            if self.__chandle__() == 0:
+                self.__init_handle_by_constructor__(ffi_new)
+            user_init(self, *args, **kwargs)
+
+        cls.__init__ = __init__  # type: ignore[attr-defined]
 
     # __repr__
     if repr and "__repr__" not in cls.__dict__:

python/tvm_ffi/dataclasses/py_class.py

@@ -422,7 +422,7 @@ def _install_deferred_init(
 #: TypeAttrColumn entries; all other names are registered as TypeMethod.
 #:
 #: System-managed names (``__ffi_new__``, ``__ffi_init__``,
-#: ``__ffi_init_inplace__``, ``__ffi_shallow_copy__``) are intentionally
+#: ``__ffi_shallow_copy__``) are intentionally
 #: absent because the C++ runtime generates them.
 _FFI_RECOGNIZED_METHODS: frozenset[str] = _FFI_TYPE_ATTR_NAMES
 

python/tvm_ffi/registry.py

@@ -381,7 +381,6 @@ def _install_init(cls: type, type_info: TypeInfo) -> None:
             cls,
             type_info,
             ffi_init=ffi_init,
-            inplace=False,
         )
     elif issubclass(cls, core.Object):
         type_name = cls.__name__

src/ffi/extra/dataclass.cc

@@ -1827,7 +1827,7 @@ Function MakeFieldSetter(int32_t field_type_index, int64_t type_converter_int,
 }
 
 // ============================================================================
-// Shared helpers for MakeInit / MakeInitInplace
+// Shared helpers for MakeInit
 // ============================================================================
 
 /*!
@@ -1992,39 +1992,22 @@ Function MakeInit(int32_t type_index) {
   });
 }
 
-Function MakeInitInplace(int32_t type_index) {
-  auto info = BuildAutoInitInfo(type_index);
-  return Function::FromPacked([info](PackedArgs args, Any* rv) {
-    TVM_FFI_ICHECK(args.size() >= 1)
-        << "__ffi_init_inplace__ requires at least one argument (self)";
-    ObjectRef self = args[0].cast<ObjectRef>();
-    Object* obj = const_cast<Object*>(self.get());
-    const auto raw_args = reinterpret_cast<const TVMFFIAny*>(args.data());
-    BindFieldArgs(obj, *info, raw_args + 1, args.size() - 1);
-  });
-}
-
 void RegisterFFIInit(int32_t type_index) {
   namespace refl = ::tvm::ffi::reflection;
   Function auto_init_fn = MakeInit(type_index);
   TVMFFIByteArray attr_name = refl::AsByteArray(refl::type_attr::kInit);
   TVMFFIAny attr_value = AnyView(auto_init_fn).CopyToTVMFFIAny();
   TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, &attr_name, &attr_value));
-
-  Function init_inplace_fn = MakeInitInplace(type_index);
-  TVMFFIByteArray ip_attr_name = refl::AsByteArray(refl::type_attr::kInitInplace);
-  TVMFFIAny ip_attr_value = AnyView(init_inplace_fn).CopyToTVMFFIAny();
-  TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterAttr(type_index, &ip_attr_name, &ip_attr_value));
 }
 
 /*!
  * \brief Combined registration for Python-defined types:
- * ``__ffi_init__``, ``__ffi_init_inplace__``, ``__ffi_new__``, ``__ffi_shallow_copy__``
+ * ``__ffi_init__``, ``__ffi_new__``, ``__ffi_shallow_copy__``
  */
 void PyClassRegisterTypeAttrColumns(int32_t type_index, int32_t total_size) {
   namespace refl = ::tvm::ffi::reflection;
   const TVMFFITypeInfo* type_info = TVMFFIGetTypeInfo(type_index);
-  // Step 1. Register `__ffi_init__` and `__ffi_init_inplace__`
+  // Step 1. Register `__ffi_init__`
   RegisterFFIInit(type_index);
   // Step 2. Register `__ffi_new__`
   Function new_fn = Function::FromTyped([type_index, total_size]() -> ObjectRef {
@@ -2148,7 +2131,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
   refl::EnsureTypeAttrColumn(refl::type_attr::kCompare);
   refl::EnsureTypeAttrColumn(refl::type_attr::kNew);
   refl::EnsureTypeAttrColumn(refl::type_attr::kInit);
-  refl::EnsureTypeAttrColumn(refl::type_attr::kInitInplace);
   refl::GlobalDef()
       .def("ffi._RegisterFFIInit", RegisterFFIInit)
       .def("ffi.MakeFieldSetter", MakeFieldSetter)

tests/python/test_dataclass_init.py

@@ -944,65 +944,6 @@ def test_low_level_kw_only_error_says_keyword_only(self) -> None:
             _ffi_init(obj, 1)
 
 
-# ###########################################################################
-#  __ffi_init_inplace__ protocol tests
-# ###########################################################################
-
-
-def _ffi_init_inplace(obj: Any, *args: Any) -> None:
-    """Look up __ffi_init_inplace__ from TypeAttrColumn and call it directly."""
-    type_index = type(obj).__tvm_ffi_type_info__.type_index
-    fn = core._lookup_type_attr(type_index, "__ffi_init_inplace__")
-    assert fn is not None, f"No __ffi_init_inplace__ for {type(obj)}"
-    fn(obj, *args)
-
-
-class TestInitInplace:
-    """Test the __ffi_init_inplace__ TypeAttrColumn."""
-
-    def test_inplace_exists(self) -> None:
-        type_index = getattr(_TestCxxAutoInit, "__tvm_ffi_type_info__").type_index
-        fn = core._lookup_type_attr(type_index, "__ffi_init_inplace__")
-        assert fn is not None
-
-    def test_inplace_positional(self) -> None:
-        obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit)
-        obj.__init_handle_by_constructor__(
-            core._lookup_type_attr(type(obj).__tvm_ffi_type_info__.type_index, "__ffi_new__")
-        )
-        _ffi_init_inplace(obj, 10, core.KWARGS, "c", 30)
-        assert obj.a == 10
-        assert obj.b == 42  # default
-        assert obj.c == 30
-        assert obj.d == 99  # default
-
-    def test_inplace_all_kwargs(self) -> None:
-        obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit)
-        obj.__init_handle_by_constructor__(
-            core._lookup_type_attr(type(obj).__tvm_ffi_type_info__.type_index, "__ffi_new__")
-        )
-        _ffi_init_inplace(obj, core.KWARGS, "a", 1, "c", 2, "d", 3)
-        assert obj.a == 1
-        assert obj.c == 2
-        assert obj.d == 3
-
-    def test_inplace_kw_only_rejects_positional(self) -> None:
-        obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit)
-        obj.__init_handle_by_constructor__(
-            core._lookup_type_attr(type(obj).__tvm_ffi_type_info__.type_index, "__ffi_new__")
-        )
-        with pytest.raises(TypeError):
-            _ffi_init_inplace(obj, 1, 2)
-
-    def test_inplace_missing_kw_only_error(self) -> None:
-        obj = _TestCxxAutoInit.__new__(_TestCxxAutoInit)
-        obj.__init_handle_by_constructor__(
-            core._lookup_type_attr(type(obj).__tvm_ffi_type_info__.type_index, "__ffi_new__")
-        )
-        with pytest.raises(TypeError, match="keyword-only"):
-            _ffi_init_inplace(obj, 1)
-
-
 # ###########################################################################
 #  __ffi_init__ TypeMethod registration regression tests
 #