Centralize DLPack device classification into classify_dl_device()

Extract the duplicated device-type mapping logic from
Buffer.__dlpack_device__(), setup_dl_tensor_device(), and
_smv_get_dl_device() into a single classify_dl_device() function
in _dlpack.pyx. All three call sites now delegate to it.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: rparolin

cuda_core/cuda/core/_dlpack.pyx

@@ -88,21 +88,28 @@ cdef inline int setup_dl_tensor_layout(DLTensor* dl_tensor, object buf) except -
     return 0
 
 
+def classify_dl_device(buf) -> tuple[int, int]:
+    """Classify a buffer into a DLPack (device_type, device_id) pair.
+
+    ``buf`` must expose ``is_device_accessible``, ``is_host_accessible``,
+    ``is_managed``, and ``device_id`` attributes.
+    """
+    cdef bint d = buf.is_device_accessible
+    cdef bint h = buf.is_host_accessible
+    if d and not h:
+        return (_kDLCUDA, buf.device_id)
+    if d and h:
+        return (_kDLCUDAManaged if buf.is_managed else _kDLCUDAHost, 0)
+    if not d and h:
+        return (_kDLCPU, 0)
+    raise BufferError("buffer is neither device-accessible nor host-accessible")
+
+
 cdef inline int setup_dl_tensor_device(DLTensor* dl_tensor, object buf) except -1:
     cdef DLDevice* device = &dl_tensor.device
-    # buf should be a Buffer instance
-    if buf.is_device_accessible and not buf.is_host_accessible:
-        device.device_type = _kDLCUDA
-        device.device_id = buf.device_id
-    elif buf.is_device_accessible and buf.is_host_accessible:
-        # Keep in sync with Buffer.__dlpack_device__() and _smv_get_dl_device().
-        device.device_type = _kDLCUDAManaged if buf.is_managed else _kDLCUDAHost
-        device.device_id = 0
-    elif not buf.is_device_accessible and buf.is_host_accessible:
-        device.device_type = _kDLCPU
-        device.device_id = 0
-    else:  # not buf.is_device_accessible and not buf.is_host_accessible
-        raise BufferError("invalid buffer")
+    dev_type, dev_id = classify_dl_device(buf)
+    device.device_type = <_DLDeviceType>dev_type
+    device.device_id = <int32_t>dev_id
     return 0
 
 

cuda_core/cuda/core/_memory/_buffer.pyx

@@ -34,7 +34,7 @@ if sys.version_info >= (3, 12):
 else:
     BufferProtocol = object
 
-from cuda.core._dlpack import DLDeviceType, make_py_capsule
+from cuda.core._dlpack import DLDeviceType, classify_dl_device, make_py_capsule
 from cuda.core._utils.cuda_utils import driver
 from cuda.core._device import Device
 
@@ -323,18 +323,7 @@ cdef class Buffer:
         return capsule
 
     def __dlpack_device__(self) -> tuple[int, int]:
-        cdef bint d = self.is_device_accessible
-        cdef bint h = self.is_host_accessible
-        if d and (not h):
-            return (DLDeviceType.kDLCUDA, self.device_id)
-        if d and h:
-            # Keep in sync with setup_dl_tensor_device() and _smv_get_dl_device().
-            if self.is_managed:
-                return (DLDeviceType.kDLCUDAManaged, 0)
-            return (DLDeviceType.kDLCUDAHost, 0)
-        if (not d) and h:
-            return (DLDeviceType.kDLCPU, 0)
-        raise BufferError("buffer is neither device-accessible nor host-accessible")
+        return classify_dl_device(self)
 
     def __buffer__(self, flags: int, /) -> memoryview:
         # Support for Python-level buffer protocol as per PEP 688.

cuda_core/cuda/core/_memoryview.pyx

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from ._dlpack cimport *
+from ._dlpack import classify_dl_device
 from libc.stdint cimport intptr_t
 from cuda.core._layout cimport _StridedLayout, get_strides_ptr
 from cuda.core._stream import Stream
@@ -590,8 +591,6 @@ cdef inline int _smv_get_dl_device(
     cdef _DLDeviceType device_type
     cdef int32_t device_id
     cdef object buf
-    cdef bint d
-    cdef bint h
     if view.dl_tensor != NULL:
         device_type = view.dl_tensor.device.device_type
         if device_type == _kDLCUDA:
@@ -601,23 +600,9 @@ cdef inline int _smv_get_dl_device(
             device_id = 0
     elif view.is_device_accessible:
         buf = view.get_buffer()
-        d = buf.is_device_accessible
-        h = buf.is_host_accessible
-        if d and (not h):
-            device_type = _kDLCUDA
-            device_id = buf.device_id
-        elif d and h:
-            # Keep in sync with Buffer.__dlpack_device__() and setup_dl_tensor_device().
-            if buf.is_managed:
-                device_type = _kDLCUDAManaged
-            else:
-                device_type = _kDLCUDAHost
-            device_id = 0
-        elif (not d) and h:
-            device_type = _kDLCPU
-            device_id = 0
-        else:
-            raise BufferError("buffer is neither device-accessible nor host-accessible")
+        dev_type, dev_id = classify_dl_device(buf)
+        device_type = <_DLDeviceType>dev_type
+        device_id = <int32_t>dev_id
     else:
         device_type = _kDLCPU
         device_id = 0