diff --git a/cuvec/_common.py b/cuvec/_common.py index 42e901d..c5b63b1 100644 --- a/cuvec/_common.py +++ b/cuvec/_common.py @@ -1,11 +1,83 @@ """Common helpers for cuvec.{cpython,pybind11,swig} modules.""" import array +import re +from abc import ABC, abstractmethod +from typing import Any, Dict from typing import Sequence as Seq from typing import Union +import numpy as np + Shape = Union[Seq[int], int] # u: non-standard np.dype('S2'); l/L: inconsistent between `array` and `numpy` typecodes = ''.join(i for i in array.typecodes if i not in "ulL") +C_TYPES = { + "signed char": 'b', + "unsigned char": 'B', + "char": 'c', + "short": 'h', + "unsigned short": 'H', + "int": 'i', + "unsigned int": 'I', + "long long": 'q', + "unsigned long long": 'Q', + "float": 'f', + "double": 'd'} # yapf: disable + + +class CVector(ABC): + """Thin wrapper around `CuVec`. Always takes ownership.""" + vec_types: Dict[np.dtype, Any] + RE_CUVEC_TYPE: re.Pattern + + def __init__(self, typechar: str): + self.typechar = typechar + + @property + @abstractmethod + def shape(self) -> tuple: + pass + + @property + @abstractmethod + def address(self) -> int: + pass + + @property + def __array_interface__(self) -> Dict[str, Any]: + return { + 'shape': self.shape, 'typestr': np.dtype(self.typechar).str, + 'data': (self.address, False), 'version': 3} + + @property + def __cuda_array_interface__(self) -> Dict[str, Any]: + return self.__array_interface__ + + def __repr__(self) -> str: + return f"{type(self).__name__}('{self.typechar}', {self.shape})" + + def __str__(self) -> str: + return f"{np.dtype(self.typechar)}{self.shape} at 0x{self.address:x}" + + @classmethod + def zeros(cls, shape: Shape, dtype="float32"): + """Returns a new Vector of the specified shape and data type.""" + return cls.vec_types[np.dtype(dtype)](shape) + + @classmethod + def copy(cls, arr): + """Returns a new Vector with data copied from the specified `arr`.""" + res = cls.zeros(arr.shape, arr.dtype) + np.asarray(res).flat = arr.flat + return res + + @classmethod + def is_instance(cls, arr): + return isinstance(arr, cls) or type(arr).__name__ == cls.__name__ + + @classmethod + def is_raw_cuvec(cls, arr): + return cls.RE_CUVEC_TYPE.match(str(arr)) def _generate_helpers(zeros, CuVec): diff --git a/cuvec/pybind11.py b/cuvec/pybind11.py index 6d94122..a6092e1 100644 --- a/cuvec/pybind11.py +++ b/cuvec/pybind11.py @@ -1,4 +1,8 @@ -"""Thin wrappers around `cuvec_pybind11` C++/CUDA module""" +""" +Thin wrappers around `cuvec_pybind11` C++/CUDA module + +A pybind11-driven equivalent of the CPython Extension API-driven `cpython.py` +""" import logging import re from collections.abc import Sequence @@ -9,45 +13,32 @@ import numpy as np from . import cuvec_pybind11 as cu # type: ignore [attr-defined] # yapf: disable -from ._common import Shape, _generate_helpers, typecodes +from ._common import C_TYPES, CVector, Shape, _generate_helpers, typecodes log = logging.getLogger(__name__) -RE_NDCUVEC_TYPE = r"<.*NDCuVec_(.) object at 0x\w+>" -NDCUVEC_TYPES = { - "signed char": 'b', - "unsigned char": 'B', - "char": 'c', - "short": 'h', - "unsigned short": 'H', - "int": 'i', - "unsigned int": 'I', - "long long": 'q', - "unsigned long long": 'Q', - "float": 'f', - "double": 'd'} # yapf: disable +NDCUVEC_TYPES = dict(C_TYPES) if hasattr(cu, 'NDCuVec_e'): typecodes += 'e' NDCUVEC_TYPES["__half"] = 'e' -class Pybind11Vector: +class Pybind11Vector(CVector): + RE_CUVEC_TYPE = re.compile(r"<.*NDCuVec_(.) object at 0x\w+>") + def __init__(self, typechar: str, shape: Shape, cuvec=None): """ - Thin wrapper around `NDCuVec`. Always takes ownership. Args: - typechar(str) + typechar(char) shape(tuple(int)) cuvec(NDCuVec): if given, `typechar` and `shape` are ignored """ - if cuvec is not None: - assert is_raw_cuvec(cuvec) - self.typechar = re.match(RE_NDCUVEC_TYPE, str(cuvec)).group(1) # type: ignore - self.cuvec = cuvec - return - - self.typechar = typechar - shape = cu.Shape(shape if isinstance(shape, Sequence) else (shape,)) - self.cuvec = getattr(cu, f'NDCuVec_{typechar}')(shape) + if cuvec is None: + shape = cu.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) @property def shape(self) -> tuple: @@ -57,48 +48,8 @@ def shape(self) -> tuple: def address(self) -> int: return self.cuvec.address() - @property - def __array_interface__(self) -> Dict[str, Any]: - return { - 'shape': self.shape, 'typestr': np.dtype(self.typechar).str, - 'data': (self.address, False), 'version': 3} - - @property - def __cuda_array_interface__(self) -> Dict[str, Any]: - return self.__array_interface__ - - def __repr__(self) -> str: - return f"{type(self).__name__}('{self.typechar}', {self.shape})" - - def __str__(self) -> str: - return f"{np.dtype(self.typechar)}{self.shape} at 0x{self.address:x}" - - -vec_types = {np.dtype(c): partial(Pybind11Vector, c) for c in typecodes} - - -def cu_zeros(shape: Shape, dtype="float32"): - """ - Returns a new `Pybind11Vector` of the specified shape and data type. - """ - return vec_types[np.dtype(dtype)](shape) - - -def cu_copy(arr): - """ - Returns a new `Pybind11Vector` with data copied from the specified `arr`. - """ - res = cu_zeros(arr.shape, arr.dtype) - np.asarray(res).flat = arr.flat - return res - - -def is_raw_cuvec(arr): - return re.match(RE_NDCUVEC_TYPE, str(arr)) - -def is_raw_pyvec(arr): - return isinstance(arr, Pybind11Vector) or type(arr).__name__ == "Pybind11Vector" +Pybind11Vector.vec_types = {np.dtype(c): partial(Pybind11Vector, c) for c in typecodes} class CuVec(np.ndarray): @@ -108,7 +59,7 @@ class CuVec(np.ndarray): """ def __new__(cls, arr): """arr: `cuvec.pybind11.CuVec`, raw `Pybind11Vector`, or `numpy.ndarray`""" - if is_raw_pyvec(arr): + if Pybind11Vector.is_instance(arr): log.debug("wrap pyraw %s", type(arr)) obj = np.asarray(arr).view(cls) obj.pyvec = arr @@ -146,7 +97,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(cu_zeros(shape, dtype)) + return CuVec(Pybind11Vector.zeros(shape, dtype)) ones, zeros_like, ones_like = _generate_helpers(zeros, CuVec) @@ -158,7 +109,7 @@ def copy(arr) -> CuVec: with data copied from the specified `arr` (`cuvec` equivalent of `numpy.copy`). """ - return CuVec(cu_copy(arr)) + return CuVec(Pybind11Vector.copy(arr)) def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec: @@ -178,13 +129,13 @@ def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec: NB: `asarray()`/`retarray()` are safe if the raw cuvec was created in C++, e.g.: >>> res = retarray(some_pybind11_api_func(..., output=None)) """ - if is_raw_cuvec(arr): + 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 is_raw_pyvec(arr): + if not isinstance(arr, np.ndarray) and Pybind11Vector.is_instance(arr): res = CuVec(arr) if dtype is None or res.dtype == np.dtype(dtype): return CuVec(np.asanyarray(res, order=order)) diff --git a/cuvec/swig.py b/cuvec/swig.py index c27c9d3..fd2ef8f 100644 --- a/cuvec/swig.py +++ b/cuvec/swig.py @@ -13,47 +13,35 @@ import numpy as np from . import swvec as sw # type: ignore [attr-defined] # yapf: disable -from ._common import Shape, _generate_helpers, typecodes +from ._common import C_TYPES, CVector, Shape, _generate_helpers, typecodes log = logging.getLogger(__name__) -RE_SWIG_TYPE = ("<.*SwigCuVec_(.); proxy of \s*\*' at 0x\w+>") -SWIG_TYPES = { - "signed char": 'b', - "unsigned char": 'B', - "char": 'c', - "short": 'h', - "unsigned short": 'H', - "int": 'i', - "unsigned int": 'I', - "long long": 'q', - "unsigned long long": 'Q', - "float": 'f', - "double": 'd'} # yapf: disable +SWIG_TYPES = dict(C_TYPES) if hasattr(sw, 'SwigCuVec_e_new'): typecodes += 'e' SWIG_TYPES["__half"] = 'e' -class SWIGVector: +class SWIGVector(CVector): + RE_CUVEC_TYPE = re.compile("<.*SwigCuVec_(.); proxy of \s*\*' at 0x\w+>") + def __init__(self, typechar: str, shape: Shape, cuvec=None): """ - Thin wrapper around `SwigPyObject>`. Always takes ownership. Args: typechar(char) shape(tuple(int)) cuvec(SwigPyObject>): if given, `typechar` and `shape` are ignored """ - if cuvec is not None: - assert is_raw_cuvec(cuvec) - self.typechar = re.match(RE_SWIG_TYPE, str(cuvec)).group(1) # type: ignore - self.cuvec = cuvec - return - - self.typechar = typechar # type: ignore - self.cuvec = getattr( - sw, f'SwigCuVec_{typechar}_new')(shape if isinstance(shape, Sequence) else (shape,)) + if cuvec is None: + cuvec = getattr( + sw, + f'SwigCuVec_{typechar}_new')(shape if isinstance(shape, Sequence) else (shape,)) + else: + typechar = self.is_raw_cuvec(cuvec).group(1) + self.cuvec = cuvec + super().__init__(typechar) def __del__(self): getattr(sw, f'SwigCuVec_{self.typechar}_del')(self.cuvec) @@ -66,48 +54,8 @@ def shape(self) -> tuple: def address(self) -> int: return getattr(sw, f'SwigCuVec_{self.typechar}_address')(self.cuvec) - @property - def __array_interface__(self) -> Dict[str, Any]: - return { - 'shape': self.shape, 'typestr': np.dtype(self.typechar).str, - 'data': (self.address, False), 'version': 3} - - @property - def __cuda_array_interface__(self) -> Dict[str, Any]: - return self.__array_interface__ - - def __repr__(self) -> str: - return f"{type(self).__name__}('{self.typechar}', {self.shape})" - - def __str__(self) -> str: - return f"{np.dtype(self.typechar)}{self.shape} at 0x{self.address:x}" - - -vec_types = {np.dtype(c): partial(SWIGVector, c) for c in typecodes} - - -def cu_zeros(shape: Shape, dtype="float32"): - """ - Returns a new `SWIGVector` of the specified shape and data type. - """ - return vec_types[np.dtype(dtype)](shape) - - -def cu_copy(arr): - """ - Returns a new `SWIGVector` with data copied from the specified `arr`. - """ - res = cu_zeros(arr.shape, arr.dtype) - np.asarray(res).flat = arr.flat - return res - - -def is_raw_cuvec(arr): - return re.match(RE_SWIG_TYPE, str(arr)) - -def is_raw_swvec(arr): - return isinstance(arr, SWIGVector) or type(arr).__name__ == "SWIGVector" +SWIGVector.vec_types = {np.dtype(c): partial(SWIGVector, c) for c in typecodes} class CuVec(np.ndarray): @@ -117,7 +65,7 @@ class CuVec(np.ndarray): """ def __new__(cls, arr): """arr: `cuvec.swig.CuVec`, raw `SWIGVector`, or `numpy.ndarray`""" - if is_raw_swvec(arr): + if SWIGVector.is_instance(arr): log.debug("wrap swraw %s", type(arr)) obj = np.asarray(arr).view(cls) obj.swvec = arr @@ -155,7 +103,7 @@ def zeros(shape: Shape, dtype="float32") -> CuVec: Returns a `cuvec.swig.CuVec` view of a new `numpy.ndarray` of the specified shape and data type (`cuvec` equivalent of `numpy.zeros`). """ - return CuVec(cu_zeros(shape, dtype)) + return CuVec(SWIGVector.zeros(shape, dtype)) ones, zeros_like, ones_like = _generate_helpers(zeros, CuVec) @@ -167,7 +115,7 @@ def copy(arr) -> CuVec: with data copied from the specified `arr` (`cuvec` equivalent of `numpy.copy`). """ - return CuVec(cu_copy(arr)) + return CuVec(SWIGVector.copy(arr)) def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec: @@ -187,13 +135,13 @@ def asarray(arr, dtype=None, order=None, ownership: str = 'warning') -> CuVec: NB: `asarray()`/`retarray()` are safe if the raw cuvec was created in C++/SWIG, e.g.: >>> res = retarray(some_swig_api_func(..., output=None)) """ - if is_raw_cuvec(arr): + if SWIGVector.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 = SWIGVector('', (), arr) - if not isinstance(arr, np.ndarray) and is_raw_swvec(arr): + if not isinstance(arr, np.ndarray) and SWIGVector.is_instance(arr): res = CuVec(arr) if dtype is None or res.dtype == np.dtype(dtype): return CuVec(np.asanyarray(res, order=order))