Use atomic flag for buffer memory-attribute initialization (NVIDIA#2216)

Store and read Buffer._mem_attrs_inited via std::atomic using
acquire/release semantics to avoid races during lazy pointer memory-
attribute initialization under free-threaded execution.

Signed-off-by: Sebastian Berg <sebastianb@nvidia.com>

Author: seberg

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.

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -103,7 +103,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. "
@@ -133,7 +133,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
@@ -462,9 +462,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(
@@ -612,7 +612,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