add PyMemRawAllocator - #6279
Conversation
|
I don't think this needs to be gated behind a feature. Typically we only use features if we can gate some compile time cost behind it, like a dependency or proc macros. |
| extern_libpython! { | ||
| #[cfg_attr(PyPy, link_name = "PyPyMem_RawMalloc")] | ||
| pub fn PyMem_RawMalloc(size: size_t) -> *mut c_void; | ||
| #[cfg_attr(PyPy, link_name = "PyPyMem_RawCalloc")] | ||
| pub fn PyMem_RawCalloc(nelem: size_t, elsize: size_t) -> *mut c_void; | ||
| #[cfg_attr(PyPy, link_name = "PyPyMem_RawRealloc")] | ||
| pub fn PyMem_RawRealloc(ptr: *mut c_void, new_size: size_t) -> *mut c_void; | ||
| #[cfg_attr(PyPy, link_name = "PyPyMem_RawFree")] | ||
| pub fn PyMem_RawFree(ptr: *mut c_void); | ||
|
|
||
| // skipped _PyMem_GetCurrentAllocatorName | ||
| // skipped _PyMem_RawStrdup | ||
| // skipped _PyMem_Strdup | ||
| // skipped _PyMem_RawWcsdup | ||
| } | ||
|
|
There was a problem hiding this comment.
| // skipped PyMem_Del | ||
| // skipped PyMem_DEL | ||
|
|
||
| #[cfg_attr(PyPy, link_name = "PyPyMem_RawMalloc")] |
There was a problem hiding this comment.
Merging this PR will degrade performance by 10.24%
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ❌ | Simulation | bench_fast |
6.1 µs | 6.8 µs | -10.24% |
Tip
Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.
Comparing chirizxc:pymem-raw-alloc (5068b35) with main (92e47e1)
| }; | ||
|
|
||
| /// `GlobalAlloc` implementation backed by CPython's `PyMem_Raw*` functions (`PYMEM_DOMAIN_RAW`). | ||
| /// Safe to use from any thread, attached or not, since the raw domain doesn't require an attached |
There was a problem hiding this comment.
Probably worth adding (if my understanding is correct) that if a user enables tracemalloc, an attached thread sate is required for computing backtraces. It means there is some risk to this for extensions that allocate in detached threads based on whether a user enables it or not (extensions that never detach like protobuf-py don't have this issue)
| /// `GlobalAlloc` implementation backed by CPython's `PyMem_Raw*` functions (`PYMEM_DOMAIN_RAW`). | ||
| /// Safe to use from any thread, attached or not, since the raw domain doesn't require an attached | ||
| /// thread state. | ||
| pub struct PyMemRawAllocator; |
There was a problem hiding this comment.
/// `GlobalAlloc` implementation backed by CPython's `PyMem_Raw*` functions (`PYMEM_DOMAIN_RAW`).
-/// Safe to use from any thread, attached or not, since the raw domain doesn't require an attached
-/// thread state.
+///
+/// The raw domain does not require an attached thread state.
+///
+/// If `tracemalloc` is enabled, CPython replaces the raw allocator with hooks that call
+/// `PyGILState_Ensure()` before recording a traceback, then walk the current thread's Python
+/// frame stack. Calling this allocator (`alloc`/`alloc_zeroed`/`realloc`, not `dealloc`) from a
+/// detached thread will therefore attach a thread state on demand while `tracemalloc` is tracing.
pub struct PyMemRawAllocator;There was a problem hiding this comment.
Thanks this looks good but maybe "attach to the thread" or "attach to Python" are more familiar to pyo3 developers then gilensure.
This is editorial but it means it could be worth pointing out that libraries that allocate frequently on non-Python threads may want to only use this allocator when PY_GIL_DISABLED.
@kumaraditya303 sorry to randomly cc but wonder if you have any insight here. I have been wondering about the implications of tracemalloc on the scientific python changes. I don't think they guard on free-threaded only IIUC
There was a problem hiding this comment.
Looking at Python/tracemalloc.c, the PYMEM_DOMAIN_RAW hooks appear to be unconditional -> no Py_GIL_DISABLED guard:
-
tracemalloc_raw_malloc/tracemalloc_raw_calloc/tracemalloc_raw_reallocall calltracemalloc_alloc/tracemalloc_reallocwithneed_gil=1hardcoded -
Inside
tracemalloc_alloc, theneed_gilflag gates aPyGILState_Ensure()/PyGILState_Release()pair around theTABLES_LOCK()+ADD_TRACEtraceback recording -
These raw hooks are installed unconditionally in
_PyTraceMalloc_StartviaPyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc)
So it looks like raw-domain allocations attach a thread state on demand whenever tracemalloc is tracing, in both GIL and free-threaded builds, I don't see it guarded to free-threaded only 🧐
If scientific-Python extensions allocate frequently from detached/non-Python threads, enabling tracemalloc would introduce this attach cost (and the associated reentrancy/deadlock surface handled by get_reentrant()/set_reentrant()) unconditionally, not just under Py_GIL_DISABLED.
Would appreciate a correction if there's a path I'm missing
There was a problem hiding this comment.
Sorry yeah - the reason I suggest free-threaded vs not is only because attach while costly doesn't block or potentically deadlock in most cases on free-threaded, so the risk of need_gil is much lower. So my understanding is on free-threaded, it can introduce large overhead - but tracemalloc by default introduces large overhead anyways. On GIL, it can have actual threading/scaling issues, code that should be parallel becomes serial and in some cases freeze (I think).
| // CPython documents this alignment as `ALIGNOF_MAX_ALIGN_T` | ||
| const MIN_ALIGN: usize = cfg_select! { | ||
| // Windows: 8 for both `MS_WIN32` / `MS_WIN64`. | ||
| target_os = "windows" => 8, | ||
| // macOS: 16 on Intel (`i386` / `x86_64`). | ||
| all(target_vendor = "apple", any(target_arch = "x86", target_arch = "x86_64")) => 16, | ||
| // macOS: 8 on other archs (e.g. `arm64`). | ||
| all(target_vendor = "apple", not(any(target_arch = "x86", target_arch = "x86_64"))) => 8, | ||
| // Other Unix: autoconf-derived at build time, not checked in, not guaranteed > 8. | ||
| _ => 8, | ||
| }; |
3 participants
Closes #6268
Based on https://github.com/bufbuild/protobuf-py/pull/46/changes#diff-d860a3c1b220287bc41b2a02c7830aab1867404d09509c63d4c7cf07bab2f7a5 (@anuraaga)