diff --git a/cuda_core/cuda/core/_memory/_buffer.pxd b/cuda_core/cuda/core/_memory/_buffer.pxd index 98c4b50db31..83dcd4f68c2 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pxd +++ b/cuda_core/cuda/core/_memory/_buffer.pxd @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 from libc.stdint cimport uintptr_t +from libcpp cimport bool as cpp_bool +from libcpp.atomic cimport atomic as std_atomic, memory_order_acquire, memory_order_release from cuda.bindings cimport cydriver from cuda.core._resource_handles cimport DevicePtrHandle @@ -18,13 +20,13 @@ cdef struct _MemAttrs: cdef class Buffer: cdef: - DevicePtrHandle _h_ptr - MemoryResource _memory_resource - object _ipc_data - object _owner - _MemAttrs _mem_attrs - bint _mem_attrs_inited - object __weakref__ + DevicePtrHandle _h_ptr + MemoryResource _memory_resource + object _ipc_data + object _owner + _MemAttrs _mem_attrs + std_atomic[cpp_bool] _mem_attrs_inited + object __weakref__ cdef public: # Python code in _memory/_virtual_memory_resource.py needs to update # this value, though it is technically private. diff --git a/cuda_core/cuda/core/_memory/_buffer.pyx b/cuda_core/cuda/core/_memory/_buffer.pyx index 88f9054385a..7cb52b22ed6 100644 --- a/cuda_core/cuda/core/_memory/_buffer.pyx +++ b/cuda_core/cuda/core/_memory/_buffer.pyx @@ -96,7 +96,7 @@ cdef class Buffer: self._memory_resource = None self._ipc_data = None self._owner = None - self._mem_attrs_inited = False + self._mem_attrs_inited.store(False) def __init__(self, *args, **kwargs) -> None: raise RuntimeError("Buffer objects cannot be instantiated directly. " @@ -126,7 +126,7 @@ cdef class Buffer: self._memory_resource = mr self._ipc_data = IPCDataForBuffer(ipc_descriptor, True) if ipc_descriptor is not None else None self._owner = owner - self._mem_attrs_inited = False + self._mem_attrs_inited.store(False) return self @staticmethod @@ -447,9 +447,9 @@ cdef class Buffer: # ------------------------------ cdef inline void _init_mem_attrs(Buffer self): """Initialize memory attributes by querying the pointer.""" - if not self._mem_attrs_inited: + if not self._mem_attrs_inited.load(memory_order_acquire): _query_memory_attrs(self._mem_attrs, as_cu(self._h_ptr)) - self._mem_attrs_inited = True + self._mem_attrs_inited.store(True, memory_order_release) cdef inline int _query_memory_attrs( @@ -597,7 +597,7 @@ cdef Buffer Buffer_from_deviceptr_handle( buf._memory_resource = mr buf._ipc_data = IPCDataForBuffer(ipc_descriptor, True) if ipc_descriptor is not None else None buf._owner = None - buf._mem_attrs_inited = False + buf._mem_attrs_inited.store(False) return buf