Adding Cuda IPC proof-of-concept implementation

.gitignore

@@ -7,6 +7,7 @@
 /cuvis/cuvis_il.py
 /cuvis/_cuvis_pyil.pyd
 /venv
+/__pycache__
 /cuvis/__pycache__
 /cuvis/git-hash.txt
 

CHANGELOG.md

@@ -20,6 +20,13 @@ Pre-releases (`b*`, `rc*`) are not listed.
 - `CI` - `scripts/check_changelog.py` validates this file's structure (header format, allowed section names, descending versions) and the tag/version/changelog agreement at release time.
 - `CONTRIBUTING.md` - documents the branch model, the version scheme, the changelog conventions and the release checklist.
 - `cuvis.BindingInfo` - new frozen dataclass with the fields `built_against: str`, `library_version: str`, `library_path: str` and `missing_symbols: Tuple[str, ...]`, the read-only property `is_complete: bool`, and a `__str__` rendering a report fit for a bug report.
+- `cuvis.CudaImageData` - new class, device-resident image data backed by a shareable CUDA buffer, exposed zero-copy through DLPack or `__cuda_array_interface__` with no host copy.
+- `cuvis.CudaImageData.export_payload` - new method, returns `bytes`.
+- `cuvis.CudaImageData.make_ipc` - new method, returns `bytes`.
+- `cuvis.CudaImageData.to_torch` - new method, returns `torch.Tensor`.
+- `cuvis.Measurement.get_cube` - new method, returns `Union[ImageData, CudaImageData]` depending on whether `cuvis.cuda.enable` was called.
+- `cuvis.Measurement.get_cube_cuda` - new method, returns `CudaImageData`.
+- `cuvis.Measurement.get_cube_cuda_ipc` - new method, returns `CudaImageData`.
 - `cuvis.SdkSettings` - new class, a `MutableMapping` of setting id to value that writes the SDK's `cuvis.settings` file, so the SDK configuration can be built in Python instead of maintained by hand.
   Values are stored as strings: `bool` becomes `true`/`false`, an `Enum` becomes its value, anything else goes through `str()`, and `None` drops the entry.
 - `cuvis.SdkSettings.__enter__`, `cuvis.SdkSettings.__exit__` - new methods; entering the context serializes the settings into a temporary directory and returns its path as `str`, leaving the context removes the directory.
@@ -32,6 +39,25 @@ Pre-releases (`b*`, `rc*`) are not listed.
 - `cuvis.binding.info` - new function, returns `BindingInfo`.
 - `cuvis.binding.missing_symbols` - new function, returns `FrozenSet[str]`.
 - `cuvis.binding.require` - new function, raises `UnavailableSDKFunction` naming whichever of the given functions the installed cuvis library does not provide.
+- `cuvis.binding.unavailable` - new function, returns `Tuple[str, ...]`.
+  A function the binding never exposed is unusable as well, so availability cannot be answered from the reported-missing list alone.
+- `cuvis.cuda` - new module gating the optional CUDA feature surface; CUDA stays off until `cuvis.cuda.enable` is called.
+- `cuvis.cuda.BACKEND_NONE`, `cuvis.cuda.BACKEND_POOL`, `cuvis.cuda.BACKEND_LEGACY`, `cuvis.cuda.BACKEND_VMM` - new constants, the IPC backend codes from `cuvis.h`.
+- `cuvis.cuda.CudaCapabilities` - new `NamedTuple` with the `bool` fields `same_process`, `ipc_pool`, `ipc_legacy`, `ipc_vmm`, `torch` and `cuda_python`, and the read-only property `any_ipc: bool`.
+- `cuvis.cuda.capabilities` - new function, returns `CudaCapabilities`.
+- `cuvis.cuda.disable` - new function.
+- `cuvis.cuda.enable` - new function.
+- `cuvis.cuda.is_enabled` - new function, returns `bool`.
+- `cuvis.cuda.require_device` - new function, raises `UnavailableSDKFunction` unless the installed library provides the same-process device path.
+- `cuvis.cuda.require_ipc` - new function, raises `UnavailableSDKFunction` unless the installed library provides the cross-process export path.
+- `cuvis_ipc` - new top-level module, the consumer side of cross-process CUDA IPC.
+  It sits outside the `cuvis` package because importing `cuvis` requires the SDK and a consumer process does not have one; its only import is `struct`.
+- `cuvis_ipc.ImportedCube` - new class, a mapped IPC buffer in the consumer process, usable as a context manager.
+- `cuvis_ipc.open` - new function, returns `ImportedCube`.
+- `cuvis_ipc.open_descriptor` - new function, returns `ImportedCube`.
+- `cuvis_ipc.pack_payload` - new function, returns `bytes`.
+- `pyproject.toml` - `py-modules` declaring the top-level `cuvis_ipc`.
+- `tests/` - `test_binding.py`, `test_cuda.py` and `test_cuvis_ipc.py`.
 
 ### Changed
 

cuvis/Measurement.py

@@ -14,7 +14,8 @@
     _utc_from_epoch_ms,
 )
 from .cuvis_types import DataFormat, ProcessingMode, ReferenceType
-from .cube_utils import ImageData
+from .cube_utils import ImageData, CudaImageData
+from . import cuda
 
 
 import cuvis.cuvis_types as internal
@@ -37,6 +38,13 @@ class Measurement(object):
     session_info: SessionData  # read-only
     frame_id: int  # read-only
 
+    # When False, refresh() skips fetching image data to the host via
+    # cuvis_measurement_get_data_image. That host fetch moves a GPU-processed cube
+    # into host memory and frees the device copy, which makes the same-process CUDA
+    # path (get_cube_cuda) unavailable. Set False before processing when you intend
+    # to read the cube as CUDA device memory. Default True preserves normal behaviour.
+    _refresh_images = True
+
     def __init__(self, base: Union[int, str, Path]):
         self._handle = None
         self._session = None
@@ -107,6 +115,10 @@ def refresh(self) -> None:
             )
             cdtype = cuvis_il.p_cuvis_data_type_t_value(pType)
             if cdtype == cuvis_il.data_type_image:
+                if not Measurement._refresh_images:
+                    # Skip the host fetch so a GPU-processed cube stays in device
+                    # memory and remains reachable via get_cube_cuda.
+                    continue
                 data = cuvis_il.cuvis_imbuffer_t()
                 cuvis_il.cuvis_measurement_get_data_image(self._handle, key, data)
                 # t0 = datetime.datetime.now()
@@ -275,6 +287,93 @@ def cube(self) -> ImageData:
             "This Measurement does not have a cube saved. Consider reprocessing with a Processing Context."
         )
 
+    def get_cube_cuda(self, key: str = "cube") -> CudaImageData:
+        """Image data as a device-resident CUDA buffer for same-process, zero-copy use.
+
+        .. code-block:: python3
+
+            from cuvis import cuda
+
+            if cuda.capabilities().same_process:
+                cuda.enable()                       # BEFORE loading or processing
+                tensor = mesu.get_cube_cuda().to_torch()
+
+        Returns a :class:`cuvis.CudaImageData` wrapping a CUVIS_CUDA_MEM handle; read it
+        with ``.to_torch()`` (DLPack, which ties the buffer lifetime to the tensor) or
+        through ``__cuda_array_interface__`` (which does not, so keep the CudaImageData
+        alive). No host copy is made.
+
+        The cube must still be on the device, which it is only when `cuda.enable` was
+        called before it was processed; the host fetch in `refresh` otherwise moves it to
+        host memory and frees the device copy.
+
+        :param key: which image entry to read, `"cube"` unless the measurement carries
+            several.
+        :raises cuvis.UnavailableSDKFunction: the installed cuvis library provides no
+            CUDA support; call `cuda.capabilities` first to avoid this.
+        :raises cuvis.cuvis_aux.SDKException: the image data is not device-backed.
+        """
+        cuda.require_device()
+        buf = cuvis_il.cuvis_cuda_imbuffer_t()
+        if cuvis_il.status_ok != cuvis_il.cuvis_measurement_get_data_image_cuda(
+            self._handle, key, buf
+        ):
+            raise SDKException()
+        return CudaImageData(buf)
+
+    def get_cube_cuda_ipc(self, key: str = "cube", backend: int = 0) -> CudaImageData:
+        """Image data as a shareable CUDA buffer for cross-process use.
+
+        Producer side; the consumer opens the payload with :mod:`cuvis_ipc`, which needs
+        no SDK of its own.
+
+        .. code-block:: python3
+
+            cimg = mesu.get_cube_cuda_ipc()         # keep alive until the consumer is done
+            send(cimg.export_payload())            # descriptor + geometry, one blob
+
+            # ... in the consumer process, no cuvis installed ...
+            import cuvis_ipc
+            with cuvis_ipc.open(payload) as cube:
+                tensor = cube.to_torch()
+
+        Fetches the device buffer with `get_cube_cuda` and creates an IPC export on it,
+        filling `.descriptor` with the transportable bytes. The returned object is the
+        in-process pin, since legacy IPC carries no cross-process refcount: drop it and
+        the consumer is reading freed device memory.
+
+        :param key: which image entry to read.
+        :param backend: which mechanism to export with, one of `cuvis.cuda.BACKEND_NONE`
+            (auto), `BACKEND_POOL`, `BACKEND_LEGACY` or `BACKEND_VMM`. `cuda.capabilities`
+            reports which of them this device supports.
+        :raises cuvis.UnavailableSDKFunction: the installed cuvis library provides no
+            CUDA IPC support.
+        :raises cuvis.cuvis_aux.SDKException: the requested backend is unavailable on
+            this device.
+        """
+        cimg = self.get_cube_cuda(key)
+        cimg.make_ipc(backend)
+        return cimg
+
+    def get_cube(self, key: str = "cube") -> Union[ImageData, CudaImageData]:
+        """Cube through whichever mode is active, so one call site serves both.
+
+        .. code-block:: python3
+
+            cube = mesu.get_cube()      # CudaImageData after cuda.enable(), else ImageData
+
+        With CUDA mode on (`cuvis.cuda.enable`) this is `get_cube_cuda` and raises when
+        the device path is unavailable. There is deliberately no silent fallback to the
+        host: a zero-copy path that quietly degrades to two copies is worse than an error,
+        because the cost is invisible.
+
+        :param key: which image entry to read.
+        :return: `CudaImageData` in CUDA mode, otherwise the host `ImageData`.
+        """
+        if cuda.is_enabled():
+            return self.get_cube_cuda(key)
+        return self.cube
+
     @property
     def thumbnail(self):
         thumb = [val for key, val in self.data.items() if "view" in key]

cuvis/__init__.py

@@ -41,7 +41,8 @@
 from .Export import CubeExporter, EnviExporter, TiffExporter, ViewExporter
 from .Calibration import Calibration
 from .AcquisitionContext import AcquisitionContext
-from .cube_utils import ImageData
+from .cube_utils import ImageData, CudaImageData
+from . import cuda
 import os
 import platform
 import sys

cuvis/_dlpack.py

@@ -0,0 +1,104 @@
+"""Minimal DLPack producer over a raw CUDA device pointer.
+
+Builds a DLManagedTensor PyCapsule via ctypes so torch.from_dlpack can consume a
+foreign device buffer zero-copy. The capsule deleter calls a supplied on_delete
+callback when torch releases the tensor, which is how buffer lifetime is tied to
+the tensor (the callback drops the SDK reference that pins the buffer).
+
+Ported from utils_data_cuda/examples/torch_local.py; the only change is that the
+deleter calls an injected callback instead of a hard-coded ctypes SDK.
+"""
+
+import ctypes
+
+_kDLCUDA = 2
+_kDLInt = 0
+_kDLUInt = 1
+_kDLFloat = 2
+
+
+class _DLDevice(ctypes.Structure):
+    _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int)]
+
+
+class _DLDataType(ctypes.Structure):
+    _fields_ = [
+        ("code", ctypes.c_uint8),
+        ("bits", ctypes.c_uint8),
+        ("lanes", ctypes.c_uint16),
+    ]
+
+
+class _DLTensor(ctypes.Structure):
+    _fields_ = [
+        ("data", ctypes.c_void_p),
+        ("device", _DLDevice),
+        ("ndim", ctypes.c_int),
+        ("dtype", _DLDataType),
+        ("shape", ctypes.POINTER(ctypes.c_int64)),
+        ("strides", ctypes.POINTER(ctypes.c_int64)),
+        ("byte_offset", ctypes.c_uint64),
+    ]
+
+
+class _DLManagedTensor(ctypes.Structure):
+    pass
+
+
+_DELETER = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor))
+_DLManagedTensor._fields_ = [
+    ("dl_tensor", _DLTensor),
+    ("manager_ctx", ctypes.c_void_p),
+    ("deleter", _DELETER),
+]
+
+_pycapsule_new = ctypes.pythonapi.PyCapsule_New
+_pycapsule_new.restype = ctypes.py_object
+_pycapsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p]
+
+# Keep ctypes structs / deleters / callbacks alive until torch invokes the deleter.
+_LIVE = {}
+
+
+class CudaDlpack:
+    """DLPack producer over (ptr, nbytes) on a CUDA device. on_delete() runs when
+    torch releases the tensor. Exposes a flat uint8 buffer; reshape/retype in torch."""
+
+    def __init__(self, ptr, nbytes, device, on_delete):
+        self._ptr = int(ptr)
+        self._n = int(nbytes)
+        self._dev = int(device)
+        self._on_delete = on_delete
+
+    def __dlpack_device__(self):
+        return (_kDLCUDA, self._dev)
+
+    def __dlpack__(self, stream=None, max_version=None, dl_device=None, copy=None):
+        shape = (ctypes.c_int64 * 1)(self._n)
+        mt = _DLManagedTensor()
+        mt.dl_tensor.data = ctypes.c_void_p(self._ptr)
+        mt.dl_tensor.device = _DLDevice(_kDLCUDA, self._dev)
+        mt.dl_tensor.ndim = 1
+        mt.dl_tensor.dtype = _DLDataType(_kDLUInt, 8, 1)
+        mt.dl_tensor.shape = shape
+        mt.dl_tensor.strides = None
+        mt.dl_tensor.byte_offset = 0
+
+        on_delete = self._on_delete
+        key = id(mt)
+
+        def _del(_p):
+            try:
+                on_delete()
+            finally:
+                _LIVE.pop(key, None)
+
+        deleter = _DELETER(_del)
+        mt.deleter = deleter
+        _LIVE[key] = (mt, shape, deleter)
+        return _pycapsule_new(ctypes.byref(mt), b"dltensor", None)
+
+
+def make_cuda_dlpack(ptr, nbytes, device, on_delete):
+    """Return an object that torch.from_dlpack consumes into a zero-copy CUDA tensor."""
+    return CudaDlpack(ptr, nbytes, device, on_delete)

cuvis/binding.py

@@ -29,9 +29,11 @@
   derives from :class:`RuntimeError`, so a single ``except RuntimeError`` covers
   both, while ``except SDKException`` still catches it as an ordinary cuvis error.
 
-Against a binding too old to report any of this (an older ``cuvis_il`` wheel), every
-query answers empty: :func:`missing_symbols` is empty, :func:`available` is ``True``
-and :func:`require` never raises. Absence of evidence, not evidence of absence.
+Against a binding too old to report any of this (an older ``cuvis_il`` wheel),
+:func:`missing_symbols` is empty and :func:`info` reports unknown throughout: absence of
+evidence, not evidence of absence. :func:`available` and :func:`require` stay meaningful
+there, because a function such an old binding never exposed is unusable whether or not
+anything reports it missing.
 """
 
 from dataclasses import dataclass, field
@@ -161,25 +163,40 @@ def missing_symbols() -> FrozenSet[str]:
     return frozenset(getattr(cuvis_il, "missing_symbols", ()))
 
 
+def unavailable(*names: str) -> Tuple[str, ...]:
+    """Which of the named functions cannot be called, in the order given.
+
+    A function is unusable for either of two reasons, and callers care about neither:
+    the binding never exposed it, which is what an older ``cuvis_il`` wheel looks like,
+    or the binding exposes it but the loaded library does not export it. Checking only
+    the second would report a function the binding does not even have as available.
+
+    :param names: C function names as they appear in ``cuvis.h``.
+    :return: the subset that is unusable, empty when all of them can be called.
+    """
+    absent = missing_symbols()
+    return tuple(
+        name for name in names if name in absent or not hasattr(cuvis_il, name)
+    )
+
+
 def available(*names: str) -> bool:
-    """Whether every named function is provided by the installed cuvis library.
+    """Whether every named function can actually be called.
 
     .. code-block:: python3
 
         if binding.available("cuvis_measurement_get_data_image_cuda"):
             cube = mesu.get_cube_cuda()
 
     :param names: C function names as they appear in ``cuvis.h``.
-    :return: ``True`` if none of them is reported missing. With a binding too old to
-        report anything this is always ``True``, so treat it as "nothing known to be
-        missing" rather than a guarantee.
+    :return: ``True`` if the binding exposes every one of them and none is reported
+        missing from the loaded library.
     """
-    absent = missing_symbols()
-    return not any(name in absent for name in names)
+    return not unavailable(*names)
 
 
 def require(*names: str) -> None:
-    """Raise unless every named function is provided by the installed cuvis library.
+    """Raise unless every named function can actually be called.
 
     Use it at the start of an operation to fail with a clear explanation, instead of
     letting a call fail deeper in with less context.
@@ -189,20 +206,20 @@ def require(*names: str) -> None:
         binding.require("cuvis_cuda_mem_get_view", "cuvis_cuda_mem_free")
 
     :param names: C function names as they appear in ``cuvis.h``.
-    :raises UnavailableSDKFunction: naming whichever of them are missing; the message
+    :raises UnavailableSDKFunction: naming whichever of them are unusable; the message
         also states the loaded SDK version and the one the binding expects.
     """
-    absent = missing_symbols()
-    unavailable = tuple(name for name in names if name in absent)
-    if unavailable:
-        raise UnavailableSDKFunction(*unavailable)
+    missing = unavailable(*names)
+    if missing:
+        raise UnavailableSDKFunction(*missing)
 
 
 __all__ = [
     "BindingInfo",
     "UnavailableSDKFunction",
     "info",
     "missing_symbols",
+    "unavailable",
     "available",
     "require",
 ]