statgpu 0.1.0__py3-none-any.whl

statgpu/backends/_cupy.py

@@ -0,0 +1,453 @@
+"""
+CuPy GPU backend.
+"""
+
+import numpy as np
+
+from statgpu.backends._base import BackendBase
+from statgpu.backends._utils import _torch_to_cupy_dlpack
+
+
+class CuPyBackend(BackendBase):
+    """
+    GPU backend powered by CuPy.
+
+    Requires ``cupy`` (install via ``pip install statgpu[gpu11]`` for CUDA 11
+    or ``pip install statgpu[gpu12]`` for CUDA 12).
+    """
+
+    name = "cupy"
+
+    @property
+    def xp(self):
+        import cupy as cp  # deferred so import doesn't fail without cupy
+        return cp
+
+    def asarray(self, x, dtype=None):
+        import cupy as cp
+        if hasattr(x, "cpu"):
+            arr = _torch_to_cupy_dlpack(x)
+            if arr is not None:
+                return arr.astype(dtype, copy=False) if dtype is not None else arr
+            # PyTorch tensors expose a .cpu() method that moves the tensor to
+            # CPU memory before converting to NumPy.  Duck-typing avoids a
+            # mandatory torch import.
+            x = x.detach().cpu().numpy()
+        return cp.asarray(x, dtype=dtype)
+
+    def to_numpy(self, x) -> np.ndarray:
+        import cupy as cp
+        if isinstance(x, cp.ndarray):
+            return cp.asnumpy(x)
+        # Fallback for numpy or other array-likes
+        if hasattr(x, "get"):
+            return x.get()
+        return np.asarray(x)
+
+    def is_available(self) -> bool:
+        try:
+            import cupy as cp
+            cp.cuda.Device(0).use()
+            return True
+        except Exception:
+            return False
+
+    def lstsq(self, A, b, rcond=None):
+        import cupy as cp
+        # CuPy's lstsq signature matches NumPy's
+        return cp.linalg.lstsq(A, b, rcond=rcond)
+
+    def solve_triangular(self, A, b, lower=False, trans=False, unit_triangular=False):
+        """
+        Solve the triangular system Ax = b.
+
+        Parameters
+        ----------
+        A : cupy.ndarray
+            Triangular matrix (n, n).
+        b : cupy.ndarray
+            Right-hand side (n,) or (n, k).
+        lower : bool, default=False
+            Whether to use the lower triangle of A.
+        trans : bool, default=False
+            Whether to transpose A.
+        unit_triangular : bool, default=False
+            Whether to assume the diagonal of A is all ones.
+
+        Returns
+        -------
+        x : cupy.ndarray
+            Solution to the system.
+        """
+        import cupy as cp
+        # Use cupyx.scipy.linalg.solve_triangular for proper triangular solve
+        # This is much faster than generic solve for triangular systems
+        try:
+            from cupyx.scipy.linalg import solve_triangular
+            return solve_triangular(A, b, lower=lower, trans=trans, unit_diagonal=unit_triangular)
+        except ImportError:
+            # Fallback to generic solve if cupyx.scipy not available
+            return cp.linalg.solve(A, b)
+
+    # ------------------------------------------------------------------
+    # Helper methods for array operations
+    # ------------------------------------------------------------------
+
+    def sum(self, x, axis=None, keepdims=False):
+        """Sum over specified axis/axes."""
+        import cupy as cp
+        return cp.sum(x, axis=axis, keepdims=keepdims)
+
+    def mean(self, x, axis=None, keepdims=False):
+        """Mean over specified axis/axes."""
+        import cupy as cp
+        return cp.mean(x, axis=axis, keepdims=keepdims)
+
+    def sqrt(self, x):
+        """Element-wise square root."""
+        import cupy as cp
+        return cp.sqrt(x)
+
+    def abs(self, x):
+        """Element-wise absolute value."""
+        import cupy as cp
+        return cp.abs(x)
+
+    def max(self, x, axis=None, keepdims=False):
+        """Maximum value along axis."""
+        import cupy as cp
+        return cp.max(x, axis=axis, keepdims=keepdims)
+
+    def outer(self, a, b):
+        """Outer product."""
+        import cupy as cp
+        return cp.outer(a.flatten(), b.flatten())
+
+    def stack(self, arrays, axis=0):
+        """Stack arrays along a new axis."""
+        import cupy as cp
+        return cp.stack(arrays, axis=axis)
+
+    def zeros(self, shape, dtype=None):
+        """Create array of zeros."""
+        import cupy as cp
+        return cp.zeros(shape, dtype=dtype)
+
+    def arange(self, start, stop=None, step=1, dtype=None):
+        """Create range array."""
+        import cupy as cp
+        if stop is None:
+            result = cp.arange(start, step=step)
+        else:
+            result = cp.arange(start, stop, step=step)
+        if dtype is not None:
+            result = result.astype(dtype)
+        return result
+
+    def array(self, val, dtype=None):
+        """Create a scalar or array from a value."""
+        import cupy as cp
+        return cp.array(val, dtype=dtype)
+
+    def atleast_1d(self, x):
+        """Ensure array is at least 1D."""
+        import cupy as cp
+        return cp.atleast_1d(x)
+
+    @property
+    def newaxis(self):
+        """Alias for None, used in indexing."""
+        import cupy as cp
+        return cp.newaxis
+
+    @property
+    def float64(self):
+        """float64 dtype."""
+        import cupy as cp
+        return cp.float64
+
+    @property
+    def float32(self):
+        """float32 dtype."""
+        import cupy as cp
+        return cp.float32
+
+    @property
+    def int64(self):
+        """int64 dtype."""
+        import cupy as cp
+        return cp.int64
+
+    @property
+    def int32(self):
+        """int32 dtype."""
+        import cupy as cp
+        return cp.int32
+
+    def clip(self, x, min_val, max_val):
+        """Clip values to [min_val, max_val]."""
+        import cupy as cp
+        return cp.clip(x, min_val, max_val)
+
+    def minimum(self, x, y):
+        """Element-wise minimum of two arrays."""
+        import cupy as cp
+        return cp.minimum(x, y)
+
+    def maximum(self, x, y):
+        """Element-wise maximum of two arrays."""
+        import cupy as cp
+        return cp.maximum(x, y)
+
+    def matmul(self, a, b):
+        """Matrix multiplication."""
+        import cupy as cp
+        return cp.matmul(a, b)
+
+    def min(self, x, axis=None, keepdims=False):
+        """Minimum value along axis."""
+        import cupy as cp
+        return cp.min(x, axis=axis, keepdims=keepdims)
+
+    def expand_dims(self, x, axis):
+        """Expand array dimensions."""
+        import cupy as cp
+        return cp.expand_dims(x, axis)
+
+    def eigh(self, a):
+        """Eigenvalue decomposition for symmetric/Hermitian matrices."""
+        import cupy as cp
+        return cp.linalg.eigh(a)
+
+    def argmin(self, x, axis=None):
+        """Indices of minimum values along axis."""
+        import cupy as cp
+        return cp.argmin(x, axis=axis)
+
+    def argmax(self, x, axis=None):
+        """Indices of maximum values along axis."""
+        import cupy as cp
+        return cp.argmax(x, axis=axis)
+
+    def argsort(self, x, axis=-1):
+        """Indices that would sort the array."""
+        import cupy as cp
+        return cp.argsort(x, axis=axis)
+
+    def where(self, condition, x, y):
+        """Element-wise conditional selection."""
+        import cupy as cp
+        return cp.where(condition, x, y)
+
+    def flip(self, x, axis=None):
+        """Reverse array order along axis."""
+        import cupy as cp
+        return cp.flip(x, axis=axis)
+
+    def exp(self, x):
+        """Element-wise exponential."""
+        import cupy as cp
+        return cp.exp(x)
+
+    def log(self, x):
+        """Element-wise natural logarithm."""
+        import cupy as cp
+        return cp.log(x)
+
+    def copy(self, x):
+        """Return a copy of x."""
+        import cupy as cp
+        return x.copy()
+
+    def ones(self, shape, dtype=None):
+        """Create array of ones."""
+        import cupy as cp
+        return cp.ones(shape, dtype=dtype)
+
+    def full(self, shape, fill_value, dtype=None):
+        """Create array filled with a constant value."""
+        import cupy as cp
+        return cp.full(shape, fill_value, dtype=dtype)
+
+    def diag(self, x, k=0):
+        """Extract diagonal or create diagonal matrix."""
+        import cupy as cp
+        return cp.diag(x, k=k)
+
+    def transpose(self, x, axes=None):
+        """Transpose array."""
+        import cupy as cp
+        return cp.transpose(x, axes)
+
+    def eye(self, n, m=None, dtype=None):
+        """Create identity matrix."""
+        import cupy as cp
+        if m is None:
+            m = n
+        return cp.eye(n, m, dtype=dtype)
+
+    def cummin(self, arr, axis=0):
+        """Cumulative minimum along *axis* (GPU-native for small arrays)."""
+        import cupy as cp
+        if arr.size == 0 or arr.shape[axis] == 0:
+            return arr.copy()
+        if str(arr.dtype) not in _CUPY_CUMOP_DTYPES:
+            return cp.minimum.accumulate(arr, axis=axis)
+        if arr.ndim == 1:
+            return self._cumop_1d(arr, cp.minimum)
+        # Multi-dim: transpose target axis to last, scan, transpose back
+        if axis != arr.ndim - 1:
+            axes = list(range(arr.ndim))
+            axes[axis], axes[-1] = axes[-1], axes[axis]
+            arr = cp.transpose(arr, axes)
+            return cp.transpose(self._cumop_last_axis(arr, cp.minimum), axes)
+        return self._cumop_last_axis(arr, cp.minimum)
+
+    def cummax(self, arr, axis=0):
+        """Cumulative maximum along *axis* (GPU-native for small arrays)."""
+        import cupy as cp
+        if arr.size == 0 or arr.shape[axis] == 0:
+            return arr.copy()
+        if str(arr.dtype) not in _CUPY_CUMOP_DTYPES:
+            return cp.maximum.accumulate(arr, axis=axis)
+        if arr.ndim == 1:
+            return self._cumop_1d(arr, cp.maximum)
+        if axis != arr.ndim - 1:
+            axes = list(range(arr.ndim))
+            axes[axis], axes[-1] = axes[-1], axes[axis]
+            arr = cp.transpose(arr, axes)
+            return cp.transpose(self._cumop_last_axis(arr, cp.maximum), axes)
+        return self._cumop_last_axis(arr, cp.maximum)
+
+    @staticmethod
+    def _cumop_1d(arr, op):
+        """1D cumulative op using sequential write."""
+        import cupy as cp
+        # Ensure contiguous for CUDA kernel compatibility
+        if not arr.flags.c_contiguous:
+            arr = cp.ascontiguousarray(arr)
+        n = len(arr)
+        if n == 0:
+            return cp.empty_like(arr)
+        result = cp.empty_like(arr)
+        result[0] = arr[0]
+        if n > 1:
+            _launch_cumop_1d(arr, result, n, op is cp.minimum)
+        return result
+
+    @staticmethod
+    def _cumop_last_axis(arr, op):
+        """Cumulative op along last axis for N-D arrays."""
+        import cupy as cp
+        # Ensure contiguous for CUDA kernel compatibility
+        if not arr.flags.c_contiguous:
+            arr = cp.ascontiguousarray(arr)
+        shape = arr.shape
+        K = shape[-1]
+        if K == 0:
+            return cp.empty_like(arr)
+        flat = arr.reshape(-1, K)
+        N = flat.shape[0]
+        if N == 0:
+            return cp.empty_like(arr)
+        result = cp.empty_like(flat)
+        result[:, 0] = flat[:, 0]
+        if K > 1:
+            _launch_cumop_2d(flat, result, N, K, op is cp.minimum)
+        return result.reshape(shape)
+
+
+# ── Raw CUDA kernels for cumulative scan ──
+_cumop_1d_template = r'''
+extern "C" __global__
+void {name}(const {dtype}* __restrict__ x,
+            {dtype}* __restrict__ out, int n) {{
+    {dtype} cur = x[0];
+    out[0] = cur;
+    for (int j = 1; j < n; j++) {{
+        if ({cmp}) cur = x[j];
+        out[j] = cur;
+    }}
+}}
+'''
+
+_cumop_2d_template = r'''
+extern "C" __global__
+void {name}(const {dtype}* __restrict__ x,
+            {dtype}* __restrict__ out, int N, int K) {{
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= N) return;
+    const {dtype}* row = x + tid * K;
+    {dtype}* orow = out + tid * K;
+    {dtype} cur = row[0];
+    orow[0] = cur;
+    for (int j = 1; j < K; j++) {{
+        if ({cmp}) cur = row[j];
+        orow[j] = cur;
+    }}
+}}
+'''
+_CUPY_CUMOP_DTYPES = {
+    "float64": "double",
+    "float32": "float",
+    "int64": "long long",
+    "int32": "int",
+}
+_cumop_kernels = {}
+
+
+def _get_cumop_kernels(dtype):
+    dtype = str(dtype)
+    if dtype not in _CUPY_CUMOP_DTYPES:
+        raise TypeError(f"Unsupported dtype for CuPy cumop kernels: {dtype}")
+    if dtype in _cumop_kernels:
+        return _cumop_kernels[dtype]
+    import cupy as cp
+    ctype = _CUPY_CUMOP_DTYPES[dtype]
+
+    kmin1_mod = cp.RawModule(code=_cumop_1d_template.format(name="cummin_1d", dtype=ctype, cmp="x[j] < cur"))
+    kmax1_mod = cp.RawModule(code=_cumop_1d_template.format(name="cummax_1d", dtype=ctype, cmp="x[j] > cur"))
+    kmin2_mod = cp.RawModule(code=_cumop_2d_template.format(name="cummin_2d", dtype=ctype, cmp="row[j] < cur"))
+    kmax2_mod = cp.RawModule(code=_cumop_2d_template.format(name="cummax_2d", dtype=ctype, cmp="row[j] > cur"))
+
+    kernels = (
+        kmin1_mod.get_function('cummin_1d'),
+        kmax1_mod.get_function('cummax_1d'),
+        kmin2_mod.get_function('cummin_2d'),
+        kmax2_mod.get_function('cummax_2d'),
+    )
+    _cumop_kernels[dtype] = kernels
+    return kernels
+
+
+def _cumop_kernels_available(dtype=None):
+    """Check if CuPy cumop kernels can be compiled (lazy, caches on first call)."""
+    try:
+        _get_cumop_kernels(dtype or "float64")
+        return True
+    except Exception:
+        return False
+
+
+def _launch_cumop_1d(arr, result, n, is_min):
+    if arr is None or result is None:
+        raise RuntimeError(
+            "CuPy cumop kernels failed to compile or unavailable. "
+            "Cannot run cummin/cummax on this device."
+        )
+    kmin1, kmax1, _, _ = _get_cumop_kernels(arr.dtype)
+    kernel = kmin1 if is_min else kmax1
+    kernel((1,), (1,), (arr, result, n))
+
+
+def _launch_cumop_2d(arr, result, N, K, is_min):
+    if arr is None or result is None:
+        raise RuntimeError(
+            "CuPy cumop kernels failed to compile or unavailable. "
+            "Cannot run cummin/cummax on this device."
+        )
+    _, _, kmin2, kmax2 = _get_cumop_kernels(arr.dtype)
+    kernel = kmin2 if is_min else kmax2
+    block = min(N, 256)
+    grid = (N + block - 1) // block
+    kernel((grid,), (block,), (arr, result, N, K))

statgpu/backends/_factory.py

@@ -0,0 +1,65 @@
+"""
+Backend factory: select the appropriate compute backend automatically or
+explicitly by name.
+"""
+
+from statgpu.backends._base import BackendBase
+from statgpu.backends._numpy import NumpyBackend
+from statgpu.backends._cupy import CuPyBackend
+from statgpu.backends._torch import TorchBackend
+
+# Module-level singletons (one instance per library, shared across calls).
+_numpy_backend = NumpyBackend()
+_cupy_backend = CuPyBackend()
+_torch_backend = TorchBackend()
+
+
+def get_backend(backend: str = "auto", device: str = "auto") -> BackendBase:
+    """
+    Return a compute backend instance.
+
+    Parameters
+    ----------
+    backend : {'auto', 'numpy', 'cupy', 'torch'}, default='auto'
+        Which array library to use.
+
+        * ``'numpy'`` – always use NumPy (CPU).
+        * ``'cupy'``  – use CuPy (requires a CUDA GPU and the ``cupy`` package).
+        * ``'torch'`` – use PyTorch (requires the ``torch`` package; defaults
+          to CUDA if available, else CPU).
+        * ``'auto'``  – pick automatically: CuPy if available, else PyTorch
+          CUDA if available, else NumPy.
+
+    device : {'auto', 'cpu', 'cuda'}, default='auto'
+        Hint about the target device.  Ignored when *backend* is explicitly
+        set to a non-``'auto'`` value.  When ``'cpu'``, always returns the
+        NumPy backend regardless of GPU availability.
+
+    Returns
+    -------
+    BackendBase
+        A backend instance that can be used to create/convert arrays.
+
+    Examples
+    --------
+    >>> from statgpu.backends import get_backend
+    >>> xp = get_backend().xp      # numpy, cupy, or torch depending on hw
+    >>> arr = xp.zeros((3, 3))
+    """
+    if backend == "numpy":
+        return _numpy_backend
+    if backend == "cupy":
+        return _cupy_backend
+    if backend == "torch":
+        return _torch_backend
+
+    # --- auto-selection ---
+    if device == "cpu":
+        return _numpy_backend
+
+    # Prefer CuPy → PyTorch CUDA → NumPy
+    if _cupy_backend.is_available():
+        return _cupy_backend
+    if _torch_backend.is_available():
+        return _torch_backend
+    return _numpy_backend