From ef5249c27d78f7c17062fd17084fea8927b6acee Mon Sep 17 00:00:00 2001 From: Charlie Lin Date: Thu, 23 Jul 2026 18:23:01 -0400 Subject: [PATCH] Make param_packer.feed() free-threading-safe via import-time init Restructure param_packer.h so the ctypes type pointers and the feeder table are built once in a new init_param_packer(), called at module import (from utils.pxi) while single-threaded. feed() becomes a read-only lookup on a never-mutated std::map, so concurrent kernel launches from multiple host threads no longer race the map or the ctypes lazy-init -- races that were live in the cp315t (Py_MOD_GIL_NOT_USED) wheels. init_param_packer is declared 'except +' so a failed 'import ctypes' surfaces as a Python exception at import instead of std::terminate across the Cython boundary, and feed() is now non-throwing. The ctypes module strong ref is kept deliberately (it keeps the borrowed type pointers valid). Harden resource_handles.cpp: make mr_dealloc_cb std::atomic with acquire/release ordering, and document the verified single-init of initialize_deferred_cleanup() and the mutation-only lazy path for GraphBox::slot_table (both already safe; no behavior change). --- .../cuda/bindings/_lib/param_packer.h | 51 ++++++++++++++++--- .../cuda/bindings/_lib/param_packer.pxd | 6 +++ cuda_bindings/cuda/bindings/_lib/utils.pxi | 5 ++ cuda_core/cuda/core/_cpp/resource_handles.cpp | 29 +++++++++-- 4 files changed, 79 insertions(+), 12 deletions(-) diff --git a/cuda_bindings/cuda/bindings/_lib/param_packer.h b/cuda_bindings/cuda/bindings/_lib/param_packer.h index 160ef5f7c92..744c58f49cd 100644 --- a/cuda_bindings/cuda/bindings/_lib/param_packer.h +++ b/cuda_bindings/cuda/bindings/_lib/param_packer.h @@ -8,6 +8,11 @@ #include #include +// Strong reference to the ctypes module, acquired once at import in +// init_param_packer() and intentionally never released. This ref is +// load-bearing: the ctypes_c_* pointers below are *borrowed* references into +// the ctypes module dict, and stay valid only while this strong ref keeps the +// module (and therefore its type objects) alive. Do not Py_DECREF it. static PyObject* ctypes_module = nullptr; static PyTypeObject* ctypes_c_char = nullptr; @@ -28,6 +33,10 @@ static PyTypeObject* ctypes_c_float = nullptr; static PyTypeObject* ctypes_c_double = nullptr; static PyTypeObject* ctypes_c_void_p = nullptr; +// Import ctypes and cache pointers to its scalar type objects. May throw +// std::runtime_error (translated to a Python exception via the `except +` +// declaration on init_param_packer in param_packer.pxd). Called exactly once, +// from init_param_packer() at module import while single-threaded. static void fetch_ctypes() { ctypes_module = PyImport_ImportModule("ctypes"); @@ -127,17 +136,45 @@ static void populate_feeders(PyTypeObject* target_t, PyTypeObject* source_t) } } +// Initialize all shared state once, at module import, while single-threaded +// (called as a bare module-level statement from utils.pxi -- mirroring the +// _resource_handles.pyx pattern of calling initialize_deferred_cleanup() at +// import). This fetches the ctypes type pointers and pre-builds the *entire* +// feeder table, so the hot feed() path below is afterwards a pure read of +// never-mutated global state. That is what makes feed() safe to call +// concurrently from multiple threads under free-threading (Py_MOD_GIL_NOT_USED) +// without any lock: doing the work lazily inside feed() would race the map and +// the ctypes lazy-init across threads launching distinct kernels. +// +// May throw (e.g. if `import ctypes` fails); declared `except +` in the pxd so +// the failure surfaces as a Python exception during module import rather than +// std::terminate. +static void init_param_packer() +{ + if (ctypes_module != nullptr) + return; // defensive: module import already runs exactly once + fetch_ctypes(); + // Pre-build every feeder the old lazy path could ever have created. This is + // a fixed, finite set of (target, source) type pairs, so a table built here + // is identical to the one feed() used to build on demand; any unmatched + // pair still falls through to feed() returning 0 -> the ctype() fallback in + // utils.pxi. After this, m_feeders is never mutated again. + populate_feeders(ctypes_c_int, &PyLong_Type); + populate_feeders(ctypes_c_bool, &PyBool_Type); + populate_feeders(ctypes_c_byte, &PyLong_Type); + populate_feeders(ctypes_c_double, &PyFloat_Type); + populate_feeders(ctypes_c_float, &PyFloat_Type); + populate_feeders(ctypes_c_longlong, &PyLong_Type); +} + +// Hot path. Read-only lookup on the import-time-populated feeder table; no +// mutation, no allocation, no throwing C-API translation -> effectively +// noexcept and safe under concurrent, GIL-free calls. Returns 0 for any +// unhandled (target, source) pair so the caller applies its ctype() fallback. static int feed(void* ptr, PyObject* value, PyObject* type) { PyTypeObject* pto = (PyTypeObject*)type; - if (ctypes_c_int == nullptr) - fetch_ctypes(); auto found = m_feeders.find({pto,value->ob_type}); - if (found == m_feeders.end()) - { - populate_feeders(pto, value->ob_type); - found = m_feeders.find({pto,value->ob_type}); - } if (found != m_feeders.end()) { return found->second(ptr, value); diff --git a/cuda_bindings/cuda/bindings/_lib/param_packer.pxd b/cuda_bindings/cuda/bindings/_lib/param_packer.pxd index 1c0ad690be4..2da96fcf19f 100644 --- a/cuda_bindings/cuda/bindings/_lib/param_packer.pxd +++ b/cuda_bindings/cuda/bindings/_lib/param_packer.pxd @@ -4,4 +4,10 @@ # Include "param_packer.h" so its contents get compiled into every # Cython extension module that depends on param_packer.pxd. cdef extern from "param_packer.h": + # Populates ctypes pointers + the feeder table once, at import. `except +` + # translates a C++ throw (e.g. failed `import ctypes`) into a Python + # exception during module import, preserving any already-set Python error. + void init_param_packer() except + + # Hot path: read-only lookup, does not throw -> intentionally left as an + # implicit-noexcept extern (no `except +` needed). int feed(void* ptr, object o, object ct) diff --git a/cuda_bindings/cuda/bindings/_lib/utils.pxi b/cuda_bindings/cuda/bindings/_lib/utils.pxi index 2dd4d5c1a27..4b302938685 100644 --- a/cuda_bindings/cuda/bindings/_lib/utils.pxi +++ b/cuda_bindings/cuda/bindings/_lib/utils.pxi @@ -11,6 +11,11 @@ import ctypes as _ctypes cimport cuda.bindings.cydriver as cydriver cimport cuda.bindings._lib.param_packer as param_packer +# Initialize the param_packer feeder table + ctypes pointers once, at import, +# while single-threaded, so param_packer.feed() is a lock-free read-only lookup +# on the kernel-launch hot path under free-threading. See param_packer.h. +param_packer.init_param_packer() + cdef void* _callocWrapper(length, size): cdef void* out = calloc(length, size) if out is NULL: diff --git a/cuda_core/cuda/core/_cpp/resource_handles.cpp b/cuda_core/cuda/core/_cpp/resource_handles.cpp index 56e20b9572c..45ef5049b63 100644 --- a/cuda_core/cuda/core/_cpp/resource_handles.cpp +++ b/cuda_core/cuda/core/_cpp/resource_handles.cpp @@ -331,6 +331,13 @@ void enqueue_cleanup(void* item) noexcept { } // namespace +// Invoked exactly once, at import, from _resource_handles.pyx as a bare +// module-level statement (verified single call site) -- so this runs +// single-threaded before any concurrent access is possible. The check-then-act +// below is therefore not a race: it is a cheap idempotency guard, not a +// concurrent-init CAS. If this were ever reachable concurrently it would need a +// compare-exchange (with the loser deleting its queue and skipping Py_AtExit); +// keep it single-init instead. void initialize_deferred_cleanup() { if (deferred_cleanup_queue.load(std::memory_order_acquire)) { return; @@ -1030,10 +1037,16 @@ DevicePtrHandle deviceptr_create_mapped_graphics( // MemoryResource-owned Device Pointer Handles // ============================================================================ -static MRDeallocCallback mr_dealloc_cb = nullptr; +// Registered exactly once at import from _buffer.pyx (module-level statement, +// runs single-threaded before any MR-backed device pointer exists), so a plain +// pointer would already be safe. Made atomic as free-threading hardening: it is +// read/called from shared_ptr deleters on arbitrary threads, and the +// acquire/release pair gives a well-formed happens-before between the import +// registration and every later deleter read. +static std::atomic mr_dealloc_cb{nullptr}; void register_mr_dealloc_callback(MRDeallocCallback cb) { - mr_dealloc_cb = cb; + mr_dealloc_cb.store(cb, std::memory_order_release); } DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* mr) { @@ -1051,8 +1064,9 @@ DevicePtrHandle deviceptr_create_with_mr(CUdeviceptr ptr, size_t size, PyObject* [mr, size](DevicePtrBox* b) { GILAcquireGuard gil; if (gil.acquired()) { - if (mr_dealloc_cb) { - mr_dealloc_cb(mr, b->resource, size, b->h_stream); + if (MRDeallocCallback cb = + mr_dealloc_cb.load(std::memory_order_acquire)) { + cb(mr, b->resource, size, b->h_stream); } Py_DECREF(mr); } @@ -1310,7 +1324,12 @@ void cleanup_graph_slot_table(void* table) noexcept { struct GraphBox { CUgraph resource; GraphHandle h_parent; // Keeps parent alive for child/branch graphs - mutable GraphSlotTable* slot_table = nullptr; // Lazily created; owned by the graph's user object + // Lazily created; owned by the graph's user object. The lazy creation is + // reached only from graph_set_slot() (a graph *mutation*), never from a + // read-only accessor, so two concurrent *reads* of the same graph cannot + // race on it -- concurrent mutation of one graph is the user's + // responsibility per docs/source/concurrency.rst. No guard needed. + mutable GraphSlotTable* slot_table = nullptr; }; const GraphBox* get_box(const GraphHandle& h) {