Exposure Tracked Buffer (first step towards unifying copy-on-write an…

…d spilling) (#13307)

The first step towards unifying copy-on-write and spillable buffers. 

This PR re-implement copy-on-write by introducing a `ExposureTrackedBuffer` and `BufferSlice`. The idea is that when `copy-on-write` (and in a follow-up PR later, when `spill`) is enabled, we use `BufferSlice` throughout cudf. 
`BufferSlice` is a _view_ of a `ExposureTrackedBuffer` that implements copy-on-write semantics by tracking the number of `BufferSlice` that points to the same `ExposureTrackedBuffer`.

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Lawrence Mitchell (https://github.com/wence-)

URL: #13307

docs/cudf/source/developer_guide/library_design.md

@@ -325,30 +325,26 @@ This section describes the internal implementation details of the copy-on-write
 It is recommended that developers familiarize themselves with [the user-facing documentation](copy-on-write-user-doc) of this functionality before reading through the internals
 below.
 
-The core copy-on-write implementation relies on the `CopyOnWriteBuffer` class.
-When the cudf option `"copy_on_write"` is `True`, `as_buffer` will always return a `CopyOnWriteBuffer`.
-This subclass of `cudf.Buffer` contains all the mechanisms to enable copy-on-write behavior.
-The class stores [weak references](https://docs.python.org/3/library/weakref.html) to every existing `CopyOnWriteBuffer` in `CopyOnWriteBuffer._instances`, a mapping from `ptr` keys to `WeakSet`s containing references to `CopyOnWriteBuffer` objects.
-This means that all `CopyOnWriteBuffer`s that point to the same device memory are contained in the same `WeakSet` (corresponding to the same `ptr` key) in `CopyOnWriteBuffer._instances`.
-This data structure is then used to determine whether or not to make a copy when a write operation is performed on a `Column` (see below).
-If multiple buffers point to the same underlying memory, then a copy must be made whenever a modification is attempted.
+The core copy-on-write implementation relies on the factory function `as_exposure_tracked_buffer` and the two classes `ExposureTrackedBuffer` and `BufferSlice`.
+
+An `ExposureTrackedBuffer` is a subclass of the regular `Buffer` that tracks internal and external references to its underlying memory. Internal references are tracked by maintaining [weak references](https://docs.python.org/3/library/weakref.html) to every `BufferSlice` of the underlying memory. External references are tracked through "exposure" status of the underlying memory. A buffer is considered exposed if the device pointer (integer or void*) has been handed out to a library outside of cudf. In this case, we have no way of knowing if the data are being modified by a third party.
+
+`BufferSlice` is a subclass of `ExposureTrackedBuffer` that represents a _slice_ of the memory underlying a exposure tracked buffer.
+
+When the cudf option `"copy_on_write"` is `True`, `as_buffer` calls `as_exposure_tracked_buffer`, which always returns a `BufferSlice`. It is then the slices that determine whether or not to make a copy when a write operation is performed on a `Column` (see below). If multiple slices point to the same underlying memory, then a copy must be made whenever a modification is attempted.
 
 
 ### Eager copies when exposing to third-party libraries
 
-If a `Column`/`CopyOnWriteBuffer` is exposed to a third-party library via `__cuda_array_interface__`, we are no longer able to track whether or not modification of the buffer has occurred. Hence whenever
+If a `Column`/`BufferSlice` is exposed to a third-party library via `__cuda_array_interface__`, we are no longer able to track whether or not modification of the buffer has occurred. Hence whenever
 someone accesses data through the `__cuda_array_interface__`, we eagerly trigger the copy by calling
-`_unlink_shared_buffers` which ensures a true copy of underlying device data is made and
-unlinks the buffer from any shared "weak" references. Any future copy requests must also trigger a true physical copy (since we cannot track the lifetime of the third-party object). To handle this we also mark the `Column`/`CopyOnWriteBuffer` as
-`obj._zero_copied=True` thus indicating that any future shallow-copy requests will trigger a true physical copy
-rather than a copy-on-write shallow copy with weak references.
+`.make_single_owner_inplace` which ensures a true copy of underlying data is made and that the slice is the sole owner. Any future copy requests must also trigger a true physical copy (since we cannot track the lifetime of the third-party object). To handle this we also mark the `Column`/`BufferSlice` as exposed thus indicating that any future shallow-copy requests will trigger a true physical copy rather than a copy-on-write shallow copy.
 
 ### Obtaining a read-only object
 
 A read-only object can be quite useful for operations that will not
-mutate the data. This can be achieved by calling `._get_cuda_array_interface(readonly=True)`, and creating a `SimpleNameSpace` object around it.
-This will not trigger a deep copy even if the `CopyOnWriteBuffer`
-has weak references. This API should only be used when the lifetime of the proxy object is restricted to cudf's internal code execution. Handing this out to external libraries or user-facing APIs will lead to untracked references and undefined copy-on-write behavior. We currently use this API for device to host
+mutate the data. This can be achieved by calling `.get_ptr(mode="read")`, and using `cuda_array_interface_wrapper` to wrap a `__cuda_array_interface__` object around it.
+This will not trigger a deep copy even if multiple `BufferSlice` points to the same `ExposureTrackedBuffer`. This API should only be used when the lifetime of the proxy object is restricted to cudf's internal code execution. Handing this out to external libraries or user-facing APIs will lead to untracked references and undefined copy-on-write behavior. We currently use this API for device to host
 copies like in `ColumnBase.data_array_view(mode="read")` which is used for `Column.values_host`.
 
 

python/cudf/cudf/_lib/column.pyx

@@ -11,7 +11,7 @@ import cudf._lib as libcudf
 from cudf.api.types import is_categorical_dtype, is_datetime64tz_dtype
 from cudf.core.buffer import (
     Buffer,
-    CopyOnWriteBuffer,
+    ExposureTrackedBuffer,
     SpillableBuffer,
     acquire_spill_lock,
     as_buffer,
@@ -339,8 +339,8 @@ cdef class Column:
         if col.base_data is None:
             data = NULL
         else:
-            data = <void*><uintptr_t>(col.base_data.get_ptr(
-                mode="write")
+            data = <void*><uintptr_t>(
+                col.base_data.get_ptr(mode="write")
             )
 
         cdef Column child_column
@@ -534,13 +534,16 @@ cdef class Column:
                     rmm.DeviceBuffer(ptr=data_ptr,
                                      size=(size+offset) * dtype_itemsize)
                 )
-            elif column_owner and isinstance(data_owner, CopyOnWriteBuffer):
-                # TODO: In future, see if we can just pass on the
-                # CopyOnWriteBuffer reference to another column
-                # and still create a weak reference.
-                # With the current design that's not possible.
-                # https://github.com/rapidsai/cudf/issues/12734
-                data = data_owner.copy(deep=False)
+            elif (
+                column_owner and
+                isinstance(data_owner, ExposureTrackedBuffer)
+            ):
+                data = as_buffer(
+                    data=data_ptr,
+                    size=base_nbytes,
+                    owner=data_owner,
+                    exposed=False,
+                )
             elif (
                 # This is an optimization of the most common case where
                 # from_column_view creates a "view" that is identical to
@@ -564,9 +567,9 @@ cdef class Column:
                     owner=data_owner,
                     exposed=True,
                 )
-                if isinstance(data_owner, CopyOnWriteBuffer):
-                    data_owner.get_ptr(mode="write")
-                    # accessing the pointer marks it exposed.
+                if isinstance(data_owner, ExposureTrackedBuffer):
+                    # accessing the pointer marks it exposed permanently.
+                    data_owner.mark_exposed()
                 elif isinstance(data_owner, SpillableBuffer):
                     if data_owner.is_spilled:
                         raise ValueError(

python/cudf/cudf/core/buffer/__init__.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 from cudf.core.buffer.buffer import Buffer, cuda_array_interface_wrapper
-from cudf.core.buffer.cow_buffer import CopyOnWriteBuffer
+from cudf.core.buffer.exposure_tracked_buffer import ExposureTrackedBuffer
 from cudf.core.buffer.spillable_buffer import SpillableBuffer, SpillLock
 from cudf.core.buffer.utils import (
     acquire_spill_lock,

python/cudf/cudf/core/buffer/buffer.py

@@ -5,7 +5,7 @@
 import math
 import pickle
 from types import SimpleNamespace
-from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
+from typing import Any, Dict, Literal, Mapping, Optional, Sequence, Tuple
 
 import numpy
 from typing_extensions import Self
@@ -168,7 +168,7 @@ def _from_host_memory(cls, data: Any) -> Self:
         # Create from device memory
         return cls._from_device_memory(buf)
 
-    def _getitem(self, offset: int, size: int) -> Buffer:
+    def _getitem(self, offset: int, size: int) -> Self:
         """
         Sub-classes can overwrite this to implement __getitem__
         without having to handle non-slice inputs.
@@ -181,7 +181,7 @@ def _getitem(self, offset: int, size: int) -> Buffer:
             )
         )
 
-    def __getitem__(self, key: slice) -> Buffer:
+    def __getitem__(self, key: slice) -> Self:
         """Create a new slice of the buffer."""
         if not isinstance(key, slice):
             raise TypeError(
@@ -193,7 +193,7 @@ def __getitem__(self, key: slice) -> Buffer:
             raise ValueError("slice must be C-contiguous")
         return self._getitem(offset=start, size=stop - start)
 
-    def copy(self, deep: bool = True):
+    def copy(self, deep: bool = True) -> Self:
         """
         Return a copy of Buffer.
 
@@ -233,35 +233,15 @@ def owner(self) -> Any:
     @property
     def __cuda_array_interface__(self) -> Mapping:
         """Implementation of the CUDA Array Interface."""
-        return self._get_cuda_array_interface(readonly=False)
-
-    def _get_cuda_array_interface(self, readonly=False):
-        """Helper function to create a CUDA Array Interface.
-
-        Parameters
-        ----------
-        readonly : bool, default False
-            If True, returns a CUDA Array Interface with
-            readonly flag set to True.
-            If False, returns a CUDA Array Interface with
-            readonly flag set to False.
-
-        Returns
-        -------
-        dict
-        """
         return {
-            "data": (
-                self.get_ptr(mode="read" if readonly else "write"),
-                readonly,
-            ),
+            "data": (self.get_ptr(mode="write"), False),
             "shape": (self.size,),
             "strides": None,
             "typestr": "|u1",
             "version": 0,
         }
 
-    def get_ptr(self, *, mode) -> int:
+    def get_ptr(self, *, mode: Literal["read", "write"]) -> int:
         """Device pointer to the start of the buffer.
 
         Parameters
@@ -274,19 +254,26 @@ def get_ptr(self, *, mode) -> int:
             Failure to fulfill this contract will cause
             incorrect behavior.
 
+        Returns
+        -------
+        int
+            The device pointer as an integer
 
         See Also
         --------
         SpillableBuffer.get_ptr
-        CopyOnWriteBuffer.get_ptr
+        ExposureTrackedBuffer.get_ptr
         """
         return self._ptr
 
-    def memoryview(self) -> memoryview:
+    def memoryview(
+        self, *, offset: int = 0, size: Optional[int] = None
+    ) -> memoryview:
         """Read-only access to the buffer through host memory."""
-        host_buf = host_memory_allocation(self.size)
+        size = self._size if size is None else size
+        host_buf = host_memory_allocation(size)
         rmm._lib.device_buffer.copy_ptr_to_host(
-            self.get_ptr(mode="read"), host_buf
+            self.get_ptr(mode="read") + offset, host_buf
         )
         return memoryview(host_buf).toreadonly()