Merge pull request #44 from AMYPAD/pybind11-buffer

pybind11 buffer

cuvec/CMakeLists.txt

@@ -1,6 +1,6 @@
 cmake_minimum_required(VERSION 3.24 FATAL_ERROR)
 if(NOT DEFINED SKBUILD_PROJECT_VERSION)
-  set(SKBUILD_PROJECT_VERSION 4 CACHE STRING "version" FORCE)
+  set(SKBUILD_PROJECT_VERSION 6 CACHE STRING "version" FORCE)
 endif()
 string(REGEX REPLACE [[([0-9]+)\.([0-9]+)\.([0-9]+).*]] [[\1.\2.\3]] SKBUILD_PROJECT_VERSION "${SKBUILD_PROJECT_VERSION}")
 project(cuvec LANGUAGES C CXX VERSION "${SKBUILD_PROJECT_VERSION}")

cuvec/__init__.py

@@ -3,8 +3,8 @@
 
 Python buffered array -> C++11 `std::vector` -> CUDA managed memory.
 """
-__author__ = "Casper O. da Costa-Luis"
-__date__ = "2021"
+__author__ = "Casper da Costa-Luis (https://github.com/casperdcl)"
+__date__ = "2021-2024"
 # version detector. Precedence: installed dist, git, 'UNKNOWN'
 try:
     from ._dist_ver import __version__

cuvec/cpython.py

@@ -11,6 +11,7 @@
 
 __all__ = [
     'CuVec', 'zeros', 'ones', 'zeros_like', 'ones_like', 'copy', 'asarray', 'Shape', 'typecodes']
+__author__, __date__, __version__ = cu.__author__, cu.__date__, cu.__version__
 
 log = logging.getLogger(__name__)
 vec_types = {

cuvec/include/cuvec.cuh

@@ -112,6 +112,13 @@ template <class T> struct NDCuVec {
     if (size != vec.size()) throw std::length_error("reshape: size mismatch");
     this->shape = shape;
   }
+  std::vector<size_t> strides() const {
+    const size_t ndim = this->shape.size();
+    std::vector<size_t> s(ndim);
+    s[ndim - 1] = sizeof(T);
+    for (int i = ndim - 2; i >= 0; i--) s[i] = this->shape[i + 1] * s[i + 1];
+    return s;
+  }
 };
 
 #endif // _CUVEC_H_

cuvec/include/cuvec_pybind11.cuh

@@ -21,12 +21,21 @@ PYBIND11_MAKE_OPAQUE(NDCuVec<long long>);
 PYBIND11_MAKE_OPAQUE(NDCuVec<unsigned long long>);
 #ifdef _CUVEC_HALF
 PYBIND11_MAKE_OPAQUE(NDCuVec<_CUVEC_HALF>);
+template <> struct pybind11::format_descriptor<_CUVEC_HALF> : pybind11::format_descriptor<float> {
+  static std::string format() { return "e"; }
+};
 #endif
 PYBIND11_MAKE_OPAQUE(NDCuVec<float>);
 PYBIND11_MAKE_OPAQUE(NDCuVec<double>);
 
 #define PYBIND11_BIND_NDCUVEC(T, typechar)                                                        \
-  pybind11::class_<NDCuVec<T>>(m, PYBIND11_TOSTRING(NDCuVec_##typechar))                          \
+  pybind11::class_<NDCuVec<T>>(m, PYBIND11_TOSTRING(NDCuVec_##typechar),                          \
+                               pybind11::buffer_protocol())                                       \
+      .def_buffer([](NDCuVec<T> &v) -> pybind11::buffer_info {                                    \
+        return pybind11::buffer_info(v.vec.data(), sizeof(T),                                     \
+                                     pybind11::format_descriptor<T>::format(), v.shape.size(),    \
+                                     v.shape, v.strides());                                       \
+      })                                                                                          \
       .def(pybind11::init<>())                                                                    \
       .def(pybind11::init<std::vector<size_t>>())                                                 \
       .def_property(                                                                              \

cuvec/pybind11.py

@@ -4,79 +4,85 @@
 A pybind11-driven equivalent of the CPython Extension API-driven `cpython.py`
 """
 import logging
-import re
 from collections.abc import Sequence
-from functools import partial
 from textwrap import dedent
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Tuple
 
 import numpy as np
 
 from . import cuvec_pybind11 as cu  # type: ignore [attr-defined] # yapf: disable
-from ._utils import CVector, Shape, _generate_helpers, typecodes
+from ._utils import Shape, _generate_helpers, typecodes
 
 __all__ = [
-    'CuVec', 'zeros', 'ones', 'zeros_like', 'ones_like', 'copy', 'asarray', 'retarray', 'Shape',
-    'typecodes']
+    'CuVec', 'zeros', 'ones', 'zeros_like', 'ones_like', 'copy', 'asarray', 'Shape', 'typecodes']
+__author__, __date__, __version__ = cu.__author__, cu.__date__, cu.__version__
 
 log = logging.getLogger(__name__)
+vec_types = {
+    np.dtype('int8'): cu.NDCuVec_b,
+    np.dtype('uint8'): cu.NDCuVec_B,
+    np.dtype('S1'): cu.NDCuVec_c,
+    np.dtype('int16'): cu.NDCuVec_h,
+    np.dtype('uint16'): cu.NDCuVec_H,
+    np.dtype('int32'): cu.NDCuVec_i,
+    np.dtype('uint32'): cu.NDCuVec_I,
+    np.dtype('int64'): cu.NDCuVec_q,
+    np.dtype('uint64'): cu.NDCuVec_Q,
+    np.dtype('float32'): cu.NDCuVec_f,
+    np.dtype('float64'): cu.NDCuVec_d}
 if hasattr(cu, 'NDCuVec_e'):
     typecodes += 'e'
+    vec_types[np.dtype('float16')] = cu.NDCuVec_e
 
 
-class Pybind11Vector(CVector):
-    RE_CUVEC_TYPE = re.compile(r"<.*NDCuVec_(.) object at 0x\w+>")
-
-    def __init__(self, typechar: str, shape: Shape, cuvec=None):
-        """
-        Args:
-          typechar(char)
-          shape(tuple(int))
-          cuvec(NDCuVec<Type>): if given, `typechar` and `shape` are ignored
-        """
-        if cuvec is None:
-            shape = shape if isinstance(shape, Sequence) else (shape,)
-            cuvec = getattr(cu, f'NDCuVec_{typechar}')(shape)
-        else:
-            typechar = self.is_raw_cuvec(cuvec).group(1)
-        self.cuvec = cuvec
-        super().__init__(typechar)
+def cu_zeros(shape: Shape, dtype="float32"):
+    """
+    Returns a new `<cuvec.cuvec_pybind11.NDCuVec_*>` of the specified shape and data type.
+    """
+    return vec_types[np.dtype(dtype)](shape if isinstance(shape, Sequence) else (shape,))
 
-    @property
-    def shape(self) -> tuple:
-        return tuple(self.cuvec.shape)
 
-    @shape.setter
-    def shape(self, shape: Shape):
-        shape = shape if isinstance(shape, Sequence) else (shape,)
-        self.cuvec.shape = shape
+def cu_copy(arr):
+    """
+    Returns a new `<cuvec.cuvec_pybind11.NDCuVec_*>` with data copied from the specified `arr`.
+    """
+    res = cu_zeros(arr.shape, arr.dtype)
+    np.asarray(res).flat = arr.flat
+    return res
+
+
+_NDCuVec_types = tuple(vec_types.values())
+_NDCuVec_types_s = tuple(map(str, vec_types.values()))
 
-    @property
-    def address(self) -> int:
-        return self.cuvec.address
 
+def is_raw_cuvec(cuvec):
+    """
+    Returns `True` when given the output of
+    pybind11 API functions returning `NDCuVec<T> *` PyObjects.
 
-Pybind11Vector.vec_types = {np.dtype(c): partial(Pybind11Vector, c) for c in typecodes}
+    This is needed since conversely `isinstance(cuvec, CuVec)` may be `False`
+    due to external libraries
+    `#include "cuvec_pybind11.cuh"` making a distinct type object.
+    """
+    return isinstance(cuvec, _NDCuVec_types) or str(type(cuvec)) in _NDCuVec_types_s
 
 
 class CuVec(np.ndarray):
     """
     A `numpy.ndarray` compatible view with a `cuvec` member containing the
-    underlying `Pybind11Vector` object (for use in pybind11 API function calls).
+    underlying `cuvec.cuvec_pybind11.NDCuVec_*` object (for use in pybind11 API function calls).
     """
     def __new__(cls, arr):
-        """arr: `cuvec.pybind11.CuVec`, raw `Pybind11Vector`, or `numpy.ndarray`"""
-        if Pybind11Vector.is_instance(arr):
-            log.debug("wrap pyraw %s", type(arr))
+        """arr: `cuvec.pybind11.CuVec`, raw `cuvec.cuvec_pybind11.NDCuVec_*`, or `numpy.ndarray`"""
+        if is_raw_cuvec(arr):
+            log.debug("wrap raw %s", type(arr))
             obj = np.asarray(arr).view(cls)
-            obj._vec = arr
-            obj.cuvec = arr.cuvec
+            obj.cuvec = arr
             return obj
-        if isinstance(arr, CuVec) and hasattr(arr, '_vec'):
+        if isinstance(arr, CuVec) and hasattr(arr, 'cuvec'):
             log.debug("new view")
             obj = np.asarray(arr).view(cls)
-            obj._vec = arr._vec
-            obj.cuvec = arr._vec.cuvec
+            obj.cuvec = arr.cuvec
             return obj
         if isinstance(arr, np.ndarray):
             log.debug("copy")
@@ -97,12 +103,14 @@ def __cuda_array_interface__(self) -> Dict[str, Any]:
             raise AttributeError(
                 dedent("""\
                 `numpy.ndarray` object has no attribute `cuvec`:
-                try using `cuvec.asarray()` first."""))
-        return self._vec.__cuda_array_interface__
+                try using `cuvec.pybind11.asarray()` first."""))
+        res = self.__array_interface__
+        return {
+            'shape': res['shape'], 'typestr': res['typestr'], 'data': res['data'], 'version': 3}
 
     def resize(self, new_shape: Shape):
         """Change shape (but not size) of array in-place."""
-        self._vec.shape = new_shape
+        self.cuvec.shape = new_shape if isinstance(new_shape, Sequence) else (new_shape,)
         super().resize(new_shape, refcheck=False)
 
     @property
@@ -119,7 +127,7 @@ def zeros(shape: Shape, dtype="float32") -> CuVec:
     Returns a `cuvec.pybind11.CuVec` view of a new `numpy.ndarray`
     of the specified shape and data type (`cuvec` equivalent of `numpy.zeros`).
     """
-    return CuVec(Pybind11Vector.zeros(shape, dtype))
+    return CuVec(cu_zeros(shape, dtype))
 
 
 ones, zeros_like, ones_like = _generate_helpers(zeros, CuVec)
@@ -131,45 +139,16 @@ def copy(arr) -> CuVec:
     with data copied from the specified `arr`
     (`cuvec` equivalent of `numpy.copy`).
     """
-    return CuVec(Pybind11Vector.copy(arr))
+    return CuVec(cu_copy(arr))
 
 
-def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec:
+def asarray(arr, dtype=None, order=None) -> CuVec:
     """
     Returns a `cuvec.pybind11.CuVec` view of `arr`, avoiding memory copies if possible.
     (`cuvec` equivalent of `numpy.asarray`).
-
-    Args:
-      ownership: logging level if `is_raw_cuvec(arr)`.
-        WARNING: `asarray()` should not be used on an existing reference, e.g.:
-        >>> res = asarray(some_pybind11_api_func(..., output=getattr(out, 'cuvec', None)))
-        `res.cuvec` and `out.cuvec` are now the same
-        yet garbage collected separately (dangling ptr).
-        Instead, use the `retarray` helper:
-        >>> raw = some_pybind11_api_func(..., output=getattr(out, 'cuvec', None))
-        >>> res = retarray(raw, out)
-        NB: `asarray()`/`retarray()` are safe if the raw cuvec was created in C++, e.g.:
-        >>> res = retarray(some_pybind11_api_func(..., output=None))
     """
-    if Pybind11Vector.is_raw_cuvec(arr):
-        ownership = ownership.lower()
-        if ownership in {'critical', 'fatal', 'error'}:
-            raise IOError("Can't take ownership of existing cuvec (would create dangling ptr)")
-        getattr(log, ownership)("taking ownership")
-        arr = Pybind11Vector('', (), arr)
-    if not isinstance(arr, np.ndarray) and Pybind11Vector.is_instance(arr):
+    if not isinstance(arr, np.ndarray) and is_raw_cuvec(arr):
         res = CuVec(arr)
         if dtype is None or res.dtype == np.dtype(dtype):
             return CuVec(np.asanyarray(res, order=order))
     return CuVec(np.asanyarray(arr, dtype=dtype, order=order))
-
-
-def retarray(raw, out: Optional[CuVec] = None):
-    """
-    Returns `out if hasattr(out, 'cuvec') else asarray(raw, ownership='debug')`.
-    See `asarray` for explanation.
-    Args:
-      raw: a raw CuVec (returned by C++/pybind11 function).
-      out: preallocated output array.
-    """
-    return out if hasattr(out, 'cuvec') else asarray(raw, ownership='debug')

cuvec/src/cpython.cu

@@ -2,8 +2,6 @@
  * Unifying Python/C++/CUDA memory.
  *
  * Python buffered array -> C++11 `std::vector` -> CUDA managed memory.
- *
- * Copyright (2021) Casper da Costa-Luis
  */
 #include "cuvec_cpython.cuh" // PyCuVec, PyCuVec_tp
 #include <Python.h>
@@ -84,7 +82,7 @@ PyMODINIT_FUNC PyInit_cuvec_cpython(void) {
   if (author == NULL) return NULL;
   PyModule_AddObject(m, "__author__", author);
 
-  PyObject *date = Py_BuildValue("s", "2021");
+  PyObject *date = Py_BuildValue("s", "2021-2024");
   if (date == NULL) return NULL;
   PyModule_AddObject(m, "__date__", date);
 

cuvec/src/cuvec_swig.i

@@ -56,3 +56,12 @@ MKCUVEC(_CUVEC_HALF, e)
 #endif
 MKCUVEC(float, f)
 MKCUVEC(double, d)
+
+%{
+static const char __author__[] = "Casper da Costa-Luis (https://github.com/casperdcl)";
+static const char __date__[] = "2021-2024";
+static const char __version__[] = "4.0.0";
+%}
+static const char __author__[];
+static const char __date__[];
+static const char __version__[];

cuvec/src/example_cpython/example_mod.cu

@@ -1,7 +1,7 @@
 /**
  * Example external extension module using CuVec.
  *
- * Copyright (2021) Casper da Costa-Luis
+ * Copyright (2021-2024) Casper da Costa-Luis
  */
 #include "Python.h"
 #include "cuvec_cpython.cuh" // PyCuVec

cuvec/src/example_pybind11/example_pybind11.cu

@@ -1,7 +1,7 @@
 /**
  * Example external pybind11 extension module using CuVec.
  *
- * Copyright (2021) Casper da Costa-Luis
+ * Copyright (2024) Casper da Costa-Luis
  */
 #include "cuvec.cuh"           // NDCuVec
 #include <pybind11/pybind11.h> // pybind11, PYBIND11_MODULE

cuvec/src/example_swig/example_swig.cu

@@ -1,7 +1,7 @@
 /**
  * Example external SWIG extension module using CuVec.
  *
- * Copyright (2021) Casper da Costa-Luis
+ * Copyright (2021-2024) Casper da Costa-Luis
  */
 #include "cuvec.cuh" // NDCuVec
 #include <stdexcept> // std::length_error

cuvec/src/pybind11.cu

@@ -2,8 +2,6 @@
  * Unifying Python/C++/CUDA memory.
  *
  * pybind11 opaque vector -> C++11 `std::vector` -> CUDA managed memory.
- *
- * Copyright (2024) Casper da Costa-Luis
  */
 #include "cuvec_pybind11.cuh"  // PYBIND11_BIND_NDCUVEC
 #include <pybind11/pybind11.h> // PYBIND11_MODULE
@@ -27,4 +25,7 @@ PYBIND11_MODULE(cuvec_pybind11, m) {
 #endif
   PYBIND11_BIND_NDCUVEC(float, f);
   PYBIND11_BIND_NDCUVEC(double, d);
+  m.attr("__author__") = "Casper da Costa-Luis (https://github.com/casperdcl)";
+  m.attr("__date__") = "2024";
+  m.attr("__version__") = "2.0.0";
 }

cuvec/swig.py

@@ -18,6 +18,7 @@
 __all__ = [
     'CuVec', 'zeros', 'ones', 'zeros_like', 'ones_like', 'copy', 'asarray', 'retarray', 'Shape',
     'typecodes']
+__author__, __date__, __version__ = sw.__author__, sw.__date__, sw.__version__
 
 log = logging.getLogger(__name__)
 if hasattr(sw, 'NDCuVec_e_new'):

docs/index.md

@@ -175,7 +175,7 @@ The following involve no memory copies.
     === "to Python"
         ```py
         import cuvec.pybind11 as cuvec, my_custom_lib
-        arr = cuvec.retarray(my_custom_lib.some_pybind11_api_func())
+        arr = cuvec.asarray(my_custom_lib.some_pybind11_api_func())
         ```
 
     === "to C++"
@@ -245,7 +245,7 @@ Python:
     import cuvec.pybind11 as cuvec, numpy, mymod
     arr = cuvec.zeros((1337, 42, 7), "float32")
     assert all(numpy.mean(arr, axis=(0, 1)) == 0)
-    print(cuvec.retarray(mymod.myfunc(arr.cuvec)).sum())
+    print(cuvec.asarray(mymod.myfunc(arr.cuvec)).sum())
     ```
 
 === "SWIG"

pyproject.toml

@@ -55,7 +55,7 @@ classifiers = [
 dependencies = ['importlib_resources; python_version < "3.9"', "numpy"]
 
 [project.optional-dependencies]
-dev = ["pytest>=6", "pytest-cov", "pytest-timeout", "pytest-xdist"]
+dev = ["pytest>=6", "pytest-cov", "pytest-timeout", "pytest-xdist", "packaging"]
 
 [tool.mypy]
 [[tool.mypy.overrides]]