cyx-6 commented on code in PR #593: URL: https://github.com/apache/tvm-ffi/pull/593#discussion_r3499331455
########## python/tvm_ffi/cython/tvm_ffi_python_object.h: ########## @@ -0,0 +1,978 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +/* + * \file tvm_ffi_python_object.h + * \brief PyObject-tying state machine: binds one Python wrapper to one C++ FFI object + * ("chandle") for the object's lifetime so identity is stable (``a.x is a.x``, + * stable ``id()`` across drop+refetch, ``f(x) is x`` for FFI returns). + * + * Split out of tvm_ffi_python_helpers.h. The design overview is the banner comment below. + */ +#ifndef TVM_FFI_PYTHON_OBJECT_H_ +#define TVM_FFI_PYTHON_OBJECT_H_ + +#include <Python.h> +#include <tvm/ffi/c_api.h> +#include <tvm/ffi/memory.h> + +// Define here to avoid dependencies on non-c headers for now +#ifndef TVM_FFI_INLINE +#if defined(_MSC_VER) +#define TVM_FFI_INLINE [[msvc::forceinline]] inline +#else +#define TVM_FFI_INLINE [[gnu::always_inline]] inline +#endif +#endif + +// Managed-dict (`__slots__ = ("__dict__",)` without an explicit dictoffset) +// is a CPython 3.11+ feature. On 3.9/3.10 such types instead use a regular +// ``tp_dictoffset != 0``, which the inactive-eligibility check catches anyway, +// so defining the flag as 0 here yields the correct (no-op) behavior. +#ifndef Py_TPFLAGS_MANAGED_DICT +#define Py_TPFLAGS_MANAGED_DICT 0 +#endif + +#include <atomic> +#include <cassert> +#include <cstring> +#include <utility> + +//================================================================================ +// PyObject-tying state machine. +// +// Ties one Python wrapper to one C++ chandle so that +// - ``a.x is a.x`` while the wrapper is live; +// - ``id(a.x)`` is stable across drop+refetch (when other C++ holders keep +// the chandle alive); +// - ``f(x) is x`` whenever an FFI function returns a chandle that already +// has a canonical wrapper. +// +// Layout +// ------ +// Every Object allocated through the registered Python allocator +// (`TVMFFIPyAllocate`) is preceded by a 16-byte ``PyCustomAllocHeader``: +// +// malloc start +// +-------------------+--------------------------+--------+ +// | tagged_pyobj | TVMFFIObjectAllocHeader | T | +// | (offset 0..8) | delete_space (8..16) | | +// +-------------------+--------------------------+--------+ +// ^ ptr = malloc + 16 +// +// ``tagged_pyobj`` is a tagged pointer to the canonical Python wrapper. The +// wrapper is >= 16-aligned, so the low 4 bits are free; two encode the state +// (see below) without growing the header past its fixed 16 bytes. +// +// States +// ------ +// Bit 0 (Inactive) and bit 1 (InTransit) tag ``tagged_pyobj`` into four states: +// Detached: ``tagged_pyobj == NULL`` -- no wrapper bound to this chandle. +// Active: ``ptr, bits == 00`` -- the live canonical wrapper. +// Inactive: ``ptr | Inactive`` -- dead, untracked allocation cached for +// address-stable revival (settled). +// InTransit: ``ptr | Inactive | InTransit`` -- a transition on this binding is in +// flight (Inactive stays set, so ``TVMFFIPyTagIsInactive`` matches too). +// +// Invariants +// ---------- +// I1. When a PyObject goes out of scope (no Python var refers to it), its +// +1 on chandle is always released (in ``__dealloc__`` -> +// ``TVMFFIPyTpDealloc``). +// I2. When a chandle is destroyed, its cached allocation (if any) is +// reclaimed. +// I3'. ``wrapper.chandle`` is only ever a real C++ object pointer or NULL, +// never a sentinel. A non-NULL chandle owns +1, except inside the +// wrapper's own dealloc window (where it is kept only as a header locator). +// I4. Every ``PyObject*`` the Cython side passes to a helper here is a live +// wrapper (tag bits 0); only this header sets or clears the tag bits. +// I5. InTransit is never a state of its own: it overlays a non-live binding +// mid-transition (Inactive(W) or Detached(NULL)), never the live Active +// wrapper, and a reader that sees it waits the transition out. +// +// The dealloc handshake +// --------------------- +// One allocation can be torn down from two directions, and the handshake stops them +// from racing into a double free or a leak: +// * from Python -- the wrapper's refcount hits 0, so ``tp_dealloc`` -> ``tp_free`` run; +// * from C++ -- the chandle's weak count hits 0, so its Weak deleter fires +// ``TVMFFIPyDeleteSpace``. +// ``tp_dealloc`` cannot read the chandle refcount to decide which side will be last (an +// FFI ``DecRef`` may race it from another thread), so instead of deciding up front it +// pre-tags ``Inactive | InTransit`` and ``DecRef``s unconditionally; the InTransit bit is +// a baton that whichever side settles last clears, and that side performs the single free. +// +// Flow 1 -- wrapper dies, chandle outlives it (cache the allocation): +// tp_dealloc : Active -> ``Inactive | InTransit``, then DecRef (chandle still has +// refs, so no deleter fires). +// tp_free : InTransit still set => clear it, keep ``self`` cached Inactive. +// delete_space : later, when the chandle dies => settled Inactive => reclaim the +// cached wrapper and free the block. +// +// Flow 2 -- wrapper held the last ref (free the allocation now): +// tp_dealloc : Active -> ``Inactive | InTransit``, then DecRef drops the last ref. +// delete_space : fires (same thread, reentrant) => InTransit set => defer the free +// back to tp_free and clear it. +// tp_free : InTransit cleared => free the C++ block here. +// +// Where transitions happen +// ------------------------ +// ``TVMFFIPyMakeRetObject`` (this header), behind ``make_ret_object`` +// (object.pxi) -- owns the whole return-object transition in one frame: +// Detached/Active/Inactive -> Active : fresh / cached / revived-in-place. +// +// ``TVMFFIPyTpDealloc`` (CObject.__dealloc__) -- runs when the wrapper's +// refcount hits 0, before the free: +// Active -> Inactive : eligible; tag Inactive | InTransit, DecRef (the +// handshake; ``tp_free`` / ``TVMFFIPyDeleteSpace`` +// settle it). +// Active -> Detached : type not eligible; detach first, then DecRef. +// +// ``TVMFFIPyArgSetterObjectRValueRef_`` (function.pxi), +// ``__move_handle_from__`` (object.pxi): +// Active -> Detached : detach the binding before a move nulls the +// source chandle. +// +// ``TVMFFIPyDeleteSpace`` (Weak deleter) -- the chandle's weak count hit 0: +// Inactive|InTransit : in-flight dealloc; defer both frees to ``tp_free``. +// Inactive (settled) : reclaim the cached wrapper and free the C++ block. +// +// Slot install +// ------------ +// ``tp_alloc`` / ``tp_free`` are NOT inherited by dynamic subtypes (CPython +// resets them per dynamic subtype), so each registered type needs its own +// install. ``_update_registry`` (object.pxi) -- the choke point every +// registered FFI type funnels through -- calls ``TVMFFIPyInstallTypeSlots`` +// there, once per type. +// +// On free-threaded builds, ``TVMFFIPyWrapDealloc`` additionally replaces each cdef +// carrier's ``tp_dealloc`` with a hand-built slot (per carrier, at its definition site; +// see "Free-threaded builds" below and ``TVMFFIPyWrapDealloc`` for why and how). +// +// Shutdown guard +// -------------- +// ``TVMFFIPyMarkPythonFinalizing`` is wired to atexit from Cython module +// init. After it fires, inactive cached allocations on still-live chandles are +// intentionally leaked (process exiting; OS reclaims) rather than reaching +// for ``PyGILState_Ensure`` on a teardown interpreter. +// +// Free-threaded builds (``Py_GIL_DISABLED``) +// ------------------------------------------ +// Without the GIL the bare ``tagged_pyobj`` reads/writes above race -- the Active-hit +// read is a use-after-free (``make_ret`` reads the wrapper, a concurrent dealloc frees +// it before the IncRef). The tie stays enabled; three FT-only mechanisms close the gap, +// all behind ``#ifdef Py_GIL_DISABLED`` so the GIL build is byte-for-byte unchanged: +// * The word is its own spin-lock (a Locked tag bit, CAS-acquired via ``__atomic_*``), +// so every transition serializes its word edits. Details in the word-access leaves. +// * The Active hit uses ``PyUnstable_TryIncRef`` (inc-if-nonzero), not ``Py_INCREF``, +// so it fails on a wrapper a concurrent dealloc is collecting -- closing the UAF. +// * A hand-built ``tp_dealloc`` slot replaces Cython's thunk, whose resurrection bump would +// otherwise let ``TryIncRef`` revive a dying wrapper. Details at ``TVMFFIPyTpDeallocSlot``. +//================================================================================ + +/*! + * \brief Python-side derived header. ``base.delete_space`` sits at + * ``ptr - sizeof(TVMFFIObjectAllocHeader)`` so the generic C++ + * deleter (which knows nothing about Python) can find it. + */ +struct PyCustomAllocHeader { + PyObject* tagged_pyobj; + TVMFFIObjectAllocHeader base; +}; + +static_assert(sizeof(PyCustomAllocHeader) == 16, + "header must be 16 bytes so T at ptr = malloc + 16 is naturally " + "aligned for alignof(T) up to alignof(max_align_t)"); +static_assert(offsetof(PyCustomAllocHeader, base) == + sizeof(PyCustomAllocHeader) - sizeof(TVMFFIObjectAllocHeader), + "base must sit at ptr - sizeof(TVMFFIObjectAllocHeader) for the " + "C++ deleter to find it"); + +TVM_FFI_INLINE PyCustomAllocHeader* TVMFFIPyHeader(void* ptr) { + return reinterpret_cast<PyCustomAllocHeader*>(static_cast<char*>(ptr) - + sizeof(PyCustomAllocHeader)); +} + +// Low-bit tags on ``tagged_pyobj`` (wrappers are >= 16-aligned -> low 4 bits free); semantics +// are the States/Invariants above. bit 0 Inactive, bit 1 InTransit, bit 2 Locked (free-threaded +// spin-lock only; the GIL build never sets it). ``TVMFFIPyRemoveTag`` masks every defined bit. +constexpr uintptr_t kPyCachedInactiveTagBit = 1; +constexpr uintptr_t kPyInTransitTagBit = 2; +#ifdef Py_GIL_DISABLED +constexpr uintptr_t kPyLockedTagBit = 4; +constexpr uintptr_t kPyTagBitMask = kPyCachedInactiveTagBit | kPyInTransitTagBit | kPyLockedTagBit; +#else +constexpr uintptr_t kPyTagBitMask = kPyCachedInactiveTagBit | kPyInTransitTagBit; +#endif + +TVM_FFI_INLINE bool TVMFFIPyTagIsInactive(PyObject* tagged) { + return (reinterpret_cast<uintptr_t>(tagged) & kPyCachedInactiveTagBit) != 0; +} +TVM_FFI_INLINE bool TVMFFIPyTagInTransit(PyObject* tagged) { + return (reinterpret_cast<uintptr_t>(tagged) & kPyInTransitTagBit) != 0; +} +TVM_FFI_INLINE PyObject* TVMFFIPyRemoveTag(PyObject* tagged) { + return reinterpret_cast<PyObject*>(reinterpret_cast<uintptr_t>(tagged) & ~kPyTagBitMask); +} +// Clear ONLY the InTransit bit (Inactive|InTransit -> Inactive, settled). +TVM_FFI_INLINE PyObject* TVMFFIPyTagClearInTransit(PyObject* tagged) { + return reinterpret_cast<PyObject*>(reinterpret_cast<uintptr_t>(tagged) & ~kPyInTransitTagBit); +} + +//--------------------------------------------------------------- +// Word-access leaves: the ONE place the GIL / free-threaded divergence lives. Every +// transition body below (make_ret, the dealloc family, Rebind) is written once against +// this small vocabulary, so the logic reads identically on both builds and the build +// difference is confined here: +// * lock: ``TVMFFIPyLockWord`` (acquire, return prior state) / ``TVMFFIPyUnlockWord`` +// (release, publish new state) / ``TVMFFIPyUnlockKeep`` (release unchanged). +// * other: ``TVMFFIPyAcquireLoad`` (read without acquiring) / ``TVMFFIPyEnableTryIncRef`` +// (arm a wrapper for a racing reader's TryIncRef before publish) / +// ``TVMFFIPyLockYield`` (GC-safe back-off, free-threaded only). +// +// Free-threaded build: the word is a spin-lock encoded in ``tagged_pyobj`` (the Locked +// bit), CAS-acquired and release-stored via ``__atomic_*``. Held only across short, +// *park-free* sections (no alloc / DecRef / GC op / blocking call), and every wait goes +// through ``TVMFFIPyLockYield`` (detaches the thread state), so the cyclic GC's +// stop-the-world can never freeze a holder nor starve on a waiter. +// +// GIL build: the GIL already serializes every transition, so there is no lock -- each leaf +// collapses to the plain field access the pre-merge code performed (load / store / no-op), +// and the merged bodies emit byte-for-byte unchanged. +//--------------------------------------------------------------- + +#ifdef Py_GIL_DISABLED +TVM_FFI_INLINE bool TVMFFIPyTagIsLocked(PyObject* tagged) { + return (reinterpret_cast<uintptr_t>(tagged) & kPyLockedTagBit) != 0; +} + +/*! \brief GC-safe back-off for any wait on the word. Must run with an attached thread state + * and WITHOUT the word lock held. */ +TVM_FFI_INLINE void TVMFFIPyLockYield() { + PyThreadState* tstate = PyEval_SaveThread(); + PyEval_RestoreThread(tstate); +} + +/*! \brief Acquire the per-word spin-lock (CAS on the Locked bit). Returns the prior binding + * (Locked bit cleared); release it via ``TVMFFIPyUnlockWord`` / ``TVMFFIPyUnlockKeep``. */ +TVM_FFI_INLINE PyObject* TVMFFIPyLockWord(PyCustomAllocHeader* h) { + for (;;) { + PyObject* cur = __atomic_load_n(&h->tagged_pyobj, __ATOMIC_RELAXED); + if (!TVMFFIPyTagIsLocked(cur)) { + PyObject* locked = + reinterpret_cast<PyObject*>(reinterpret_cast<uintptr_t>(cur) | kPyLockedTagBit); + // Acquire on success so the locked section happens-after the matching release. + if (__atomic_compare_exchange_n(&h->tagged_pyobj, &cur, locked, /*weak=*/true, Review Comment: fixed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
