Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

pybind11 buffer #44

Merged
merged 7 commits into from
Mar 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cuvec/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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}")
Expand Down
4 changes: 2 additions & 2 deletions cuvec/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__
Expand Down
1 change: 1 addition & 0 deletions cuvec/cpython.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
7 changes: 7 additions & 0 deletions cuvec/include/cuvec.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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_
11 changes: 10 additions & 1 deletion cuvec/include/cuvec_pybind11.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -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( \
Expand Down
137 changes: 58 additions & 79 deletions cuvec/pybind11.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Check warning on line 35 in cuvec/pybind11.py

View check run for this annotation

Codecov / codecov/patch

cuvec/pybind11.py#L35

Added line #L35 was not covered by tests


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")
Expand All @@ -97,12 +103,14 @@
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 {

Check warning on line 108 in cuvec/pybind11.py

View check run for this annotation

Codecov / codecov/patch

cuvec/pybind11.py#L107-L108

Added lines #L107 - L108 were not covered by tests
'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
Expand All @@ -119,7 +127,7 @@
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)
Expand All @@ -131,45 +139,16 @@
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')
4 changes: 1 addition & 3 deletions cuvec/src/cpython.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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);

Expand Down
9 changes: 9 additions & 0 deletions cuvec/src/cuvec_swig.i
Original file line number Diff line number Diff line change
Expand Up @@ -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__[];
2 changes: 1 addition & 1 deletion cuvec/src/example_cpython/example_mod.cu
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion cuvec/src/example_pybind11/example_pybind11.cu
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion cuvec/src/example_swig/example_swig.cu
Original file line number Diff line number Diff line change
@@ -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
Expand Down
5 changes: 3 additions & 2 deletions cuvec/src/pybind11.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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";
}
1 change: 1 addition & 0 deletions cuvec/swig.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'):
Expand Down
4 changes: 2 additions & 2 deletions docs/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -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++"
Expand Down Expand Up @@ -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"
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Expand Down
Loading