From 0866334a0556df2f4771e6cd9c16e8367f4dd4e5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 21:55:30 +0200 Subject: [PATCH 1/9] fix: segment_sum on the fake CuPy sums on the host buffers --- CHANGELOG.md | 3 +++ src/cunumpy/_algorithms.py | 13 ++++++++++++- tests/unit/test_segment_plans.py | 31 ++++++++++++++++++++++++++++++- 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 97bd2cd..7dd14f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 warnings as errors. Docstring style: `docs/source/contributing.md`. ### Fixed +- `xp.algorithms.segment_sum` and `SegmentPlan.sum` work on the fake CuPy: it + cannot compile the CUDA kernel, so the sum runs on the host buffers of the + fake arrays and the CuPy path of a program runs without a GPU. - The ordering guide identifies 0.6.3 as the first release preserving F order in CuPy-to-host conversions. - `CudaStruct.from_signature` accepts quoted pyccel annotations diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index b5cd42c..11dd85f 100644 --- a/src/cunumpy/_algorithms.py +++ b/src/cunumpy/_algorithms.py @@ -8,6 +8,7 @@ import array_api_compat.numpy as np +from cunumpy import _fake_cupy from cunumpy.xp import assert_same_backend, get_array_backend, get_array_module @@ -207,7 +208,9 @@ def sum(self, values: Any, *, out: Any = None) -> Any: Floating-point and complex dtypes are kept; integer and bool values give float64. CUDA supports float16/32/64 and complex64/128 (float16 accumulates in float32) and sums all components in one launch with atomics, so the - floating-point order is not deterministic. + floating-point order is not deterministic. On the fake CuPy, which cannot + compile CUDA kernels, the sum runs on the host buffers of the fake arrays, + so the CuPy path of a program works without a GPU. Parameters ---------- @@ -275,6 +278,14 @@ def sum(self, values: Any, *, out: Any = None) -> Any: self._keys[self._valid], values[self._valid].astype(dtype, copy=False), ) + elif values.size and self.n_segments and _fake_cupy.is_active(): + keys = _fake_cupy.host_buffer(self._keys) + valid = keys >= 0 + np.add.at( + _fake_cupy.host_buffer(out), + keys[valid], + _fake_cupy.host_buffer(values)[valid].astype(dtype, copy=False), + ) elif values.size and self.n_segments: half = np.dtype(dtype) == np.dtype(np.float16) packed = xpm.ascontiguousarray(values, dtype=np.float32 if half else dtype) diff --git a/tests/unit/test_segment_plans.py b/tests/unit/test_segment_plans.py index 6301d5c..9c1c0e2 100644 --- a/tests/unit/test_segment_plans.py +++ b/tests/unit/test_segment_plans.py @@ -1,5 +1,7 @@ """Grouping and reusable reductions on CPU and (when available) CUDA.""" +from pathlib import Path + import numpy as np import pytest @@ -10,7 +12,9 @@ segment_boundaries, segment_sum, ) -from cunumpy.kernel_testing import BACKENDS +from cunumpy.kernel_testing import BACKENDS, run_in_fake_cupy_subprocess + +SRC = str(Path(__file__).resolve().parents[2] / "src") @pytest.mark.parametrize("backend", BACKENDS) @@ -160,3 +164,28 @@ def test_fused_sum_kernel_matches_reference_in_cpu_emulation(dtype): expected = np.zeros_like(out) np.add.at(expected, keys[keys >= 0], values[keys >= 0]) np.testing.assert_array_equal(out, expected) + + +def test_segment_sum_runs_on_the_fake_cupy(): + """The fake CuPy cannot compile the CUDA kernel; the sum uses its host buffers.""" + code = """ +import numpy as np +import cunumpy as xp +xp.set_backend("cupy") +keys = xp.asarray([0, 2, 0, -1, 2]) +values = xp.asarray(np.arange(15.0).reshape(5, 3)) +total = xp.algorithms.segment_sum(values, keys, 3) +assert xp.is_gpu(total) +expected = np.array([[6.0, 8.0, 10.0], [0.0, 0.0, 0.0], [15.0, 17.0, 19.0]]) +np.testing.assert_array_equal(xp.to_numpy(total), expected) +plan = xp.algorithms.SegmentPlan(keys, 3) +out = xp.empty((3, 3)) +assert plan.sum(values, out=out) is out +np.testing.assert_array_equal(xp.to_numpy(out), expected) +counts = xp.algorithms.segment_sum(xp.asarray([True, True, False, True, True]), keys, 3) +np.testing.assert_array_equal(xp.to_numpy(counts), [1.0, 0.0, 2.0]) +complex_sum = xp.algorithms.segment_sum(xp.asarray([1 + 2j, 3j]), xp.asarray([1, 1]), 2) +np.testing.assert_array_equal(xp.to_numpy(complex_sum), [0, 1 + 5j]) +print("ok") +""" + assert "ok" in run_in_fake_cupy_subprocess(code, env={"PYTHONPATH": SRC}).stdout From 26d9177a8e364cce3a1c225422562963dd0a5051 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 22:08:57 +0200 Subject: [PATCH 2/9] add xp.scipy.linalg.solve_circulant on the device --- CHANGELOG.md | 7 +- docs/source/guides/solvers.md | 10 +- src/cunumpy/LLM_GUIDE.md | 1 + src/cunumpy/_fake_cupy.py | 2 +- src/cunumpy/_fake_cupyx.py | 178 ++++++++++++++++++++++++------- src/cunumpy/_linalg.py | 103 ++++++++++++++++++ src/cunumpy/_scipy_backend.py | 26 ++++- tests/unit/test_fake_cupyx.py | 56 ++++++++++ tests/unit/test_linalg.py | 44 ++++++++ tests/unit/test_scipy_backend.py | 16 +++ 10 files changed, 394 insertions(+), 49 deletions(-) create mode 100644 src/cunumpy/_linalg.py create mode 100644 tests/unit/test_linalg.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7dd14f5..a401cdb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,8 +14,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 copies, fixed shapes, one synchronization per iteration). - The fake CuPy has `cupyx.scipy`: SciPy behind the CuPy rules (host arrays rejected, fake device arrays and proxy objects returned), so `xp.scipy` code - runs its CuPy path without a GPU. `cupyx.scipy.interpolate` has only the - names CuPy provides. `cupyx.empty_pinned`/`zeros_pinned` return NumPy arrays. + runs its CuPy path without a GPU. Every `cupyx.scipy` subpackage has only + the names CuPy v14.2.0 provides, so e.g. `special.jv` fails as on a GPU. + `cupyx.empty_pinned`/`zeros_pinned` return NumPy arrays. +- `xp.scipy.linalg.solve_circulant` on the CuPy backend, which `cupyx.scipy` + lacks: SciPy's FFT solve and arguments, on the device. - Solvers guide: splines on the device with `xp.scipy.interpolate`, and batched root finding. - Compiled Pyccel regression coverage for F-ordered arrays and column-block diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index f4b874c..29ead61 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -44,6 +44,12 @@ def periodic_poisson(rho, length): * `cupyx.scipy` has only part of SciPy. A missing name raises `AttributeError` saying which backend lacks it; check with `xp.scipy.special.available("erfcx")` where a fallback is possible. + Commonly missing: `special.jv` (only `j0`, `j1`, `yn`), `special.erfi`, + `linalg.solve`/`inv` (use `xp.linalg`), `RectBivariateSpline`, and all of + `integrate` and `optimize`. +* cunumpy fills in `xp.scipy.linalg.solve_circulant` on the CuPy backend: an + FFT solve on the device, with SciPy's arguments (a circulant preconditioner + per iteration stays on the GPU). * Keyword arguments can differ between the two: SciPy's `cg` takes `rtol` (since SciPy 1.12), CuPy's takes `tol`. Pass tolerances only after checking both, or through a small wrapper. @@ -54,8 +60,8 @@ def periodic_poisson(rho, length): (`erf`, `erfc`, Bessel functions `i0`, `i1`, `k0`, ...) are in `xp.scipy.special` on both backends. * On the fake CuPy (tests without a GPU), `xp.scipy` runs SciPy with the CuPy - rules: host arrays are rejected, results are fake device arrays, and - `interpolate` has only the names CuPy has. + rules: host arrays are rejected, results are fake device arrays, and each + subpackage has only the names CuPy has. ## Splines and interpolation diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 83c02bf..d0aa12e 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -100,6 +100,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | CUDA struct from a Pyccel argument class | `xp.arguments.CudaStruct.from_signature(Cls.__init__, "Name")`, `xp.arguments.write_cuda_header(...)` | | SciPy (sparse, sparse.linalg, fft, special, ndimage, ...) on either backend | `xp.scipy..` (SciPy or `cupyx.scipy`); `xp.scipy.special.available(name)` | | splines on the device | `xp.scipy.interpolate.UnivariateSpline`, `make_interp_spline`, `NdBSpline` (no `RectBivariateSpline`: `make_interp_spline` per axis + `NdBSpline`) | +| circulant solve (FFT preconditioner) on the device | `xp.scipy.linalg.solve_circulant(c, b)` (cunumpy fill-in on CuPy, SciPy's arguments) | | many independent scalar root problems (one per ray, per marker) | `xp.optimize.newton(func, x0_array, fprime=None)` (SciPy's array Newton/secant, on the device) | | chain of elementwise operations as one GPU kernel | `@xp.kernels.fuse` (`cupy.fuse` for CuPy arrays, plain call otherwise) | | PETSc solve on device arrays without copies | `xp.petsc.petsc_vec(array)` (CUDA/HIP petsc4py for CuPy arrays); `xp.synchronize()` around PETSc calls | diff --git a/src/cunumpy/_fake_cupy.py b/src/cunumpy/_fake_cupy.py index 35a1cc7..a96b1be 100644 --- a/src/cunumpy/_fake_cupy.py +++ b/src/cunumpy/_fake_cupy.py @@ -16,7 +16,7 @@ :class:`~cunumpy.arguments.CudaStruct` packing and the argument checks of :class:`~cunumpy.kernels.CudaKernel` work; * ``cupyx.scipy`` forwards to SciPy under the same rules, with the names of - ``cupyx.scipy.interpolate`` only (see :mod:`cunumpy._fake_cupyx`); + ``cupyx.scipy`` (see :mod:`cunumpy._fake_cupyx`); * CUDA kernels cannot run: ``RawKernel`` and friends raise ``NotImplementedError`` when called, and :func:`cunumpy.kernel_testing.requires_cupy` skips tests while the fake is active. Inside diff --git a/src/cunumpy/_fake_cupyx.py b/src/cunumpy/_fake_cupyx.py index 58b17a0..900b0ef 100644 --- a/src/cunumpy/_fake_cupyx.py +++ b/src/cunumpy/_fake_cupyx.py @@ -9,10 +9,10 @@ in CuPy; * array results are fake device arrays, and objects (a fitted spline, a factorization) are proxies whose methods follow the same rules; -* ``cupyx.scipy.interpolate`` has only the names CuPy provides - (:data:`INTERPOLATE`), so code using a SciPy name that CuPy lacks (e.g. - ``RectBivariateSpline``) fails here as it would on a GPU. The other - subpackages forward every SciPy name. +* each subpackage has only the names CuPy provides (:data:`NAMES`), so code + using a SciPy name that CuPy lacks (e.g. ``RectBivariateSpline``, + ``special.jv``, ``linalg.solve_circulant``) fails here as it would on a + GPU. A CuPy name that the installed SciPy lacks is missing too. ``cupyx.empty_pinned`` and ``cupyx.zeros_pinned`` return NumPy arrays, like CuPy's (which are backed by pinned memory). @@ -43,41 +43,139 @@ "stats", ) -#: The names of ``cupyx.scipy.interpolate`` (CuPy reference, stable). -INTERPOLATE = frozenset( - { - "Akima1DInterpolator", - "BPoly", - "BSpline", - "BarycentricInterpolator", - "CloughTocher2DInterpolator", - "CubicHermiteSpline", - "CubicSpline", - "InterpolatedUnivariateSpline", - "KroghInterpolator", - "LSQUnivariateSpline", - "LinearNDInterpolator", - "NdBSpline", - "NdPPoly", - "NearestNDInterpolator", - "PPoly", - "PchipInterpolator", - "RBFInterpolator", - "RegularGridInterpolator", - "UnivariateSpline", - "barycentric_interpolate", - "interp1d", - "interpn", - "krogh_interpolate", - "make_interp_spline", - "make_lsq_spline", - "pchip_interpolate", - "splantider", - "splder", - }, -) -_ALLOWED = {"interpolate": INTERPOLATE} +def _names(text: str) -> frozenset[str]: + return frozenset(text.split()) + + +#: The public names of each ``cupyx.scipy`` subpackage (CuPy v14.2.0). +NAMES = { + "fft": _names( + """ + dct dctn dst dstn fft fft2 fftfreq fftn fftshift fht get_fft_plan hfft + hfft2 hfftn idct idctn idst idstn ifft ifft2 ifftn ifftshift ifht ihfft + ihfft2 ihfftn irfft irfft2 irfftn next_fast_len rfft rfft2 rfftfreq + rfftn + """, + ), + "fftpack": _names( + """ + fft fft2 fftn get_fft_plan ifft ifft2 ifftn irfft rfft + """, + ), + "interpolate": _names( + """ + Akima1DInterpolator BPoly BSpline BarycentricInterpolator + CloughTocher2DInterpolator CubicHermiteSpline CubicSpline + InterpolatedUnivariateSpline KroghInterpolator LSQUnivariateSpline + LinearNDInterpolator NdBSpline NdPPoly NearestNDInterpolator PPoly + PchipInterpolator RBFInterpolator RegularGridInterpolator + UnivariateSpline barycentric_interpolate interp1d interpn + krogh_interpolate make_interp_spline make_lsq_spline pchip + pchip_interpolate splantider splder + """, + ), + "linalg": _names( + """ + bandwidth block_diag circulant companion convolution_matrix dft expm + fiedler fiedler_companion hadamard hankel helmert hilbert khatri_rao + kron leslie lu lu_factor lu_solve solve_triangular toeplitz + """, + ), + "ndimage": _names( + """ + affine_transform binary_closing binary_dilation binary_erosion + binary_fill_holes binary_hit_or_miss binary_opening binary_propagation + black_tophat center_of_mass convolve convolve1d correlate correlate1d + distance_transform_edt extrema find_objects fourier_ellipsoid + fourier_gaussian fourier_shift fourier_uniform gaussian_filter + gaussian_filter1d gaussian_gradient_magnitude gaussian_laplace + generate_binary_structure generic_filter generic_filter1d + generic_gradient_magnitude generic_laplace grey_closing grey_dilation + grey_erosion grey_opening histogram iterate_structure label + labeled_comprehension laplace map_coordinates maximum maximum_filter + maximum_filter1d maximum_position mean median median_filter minimum + minimum_filter minimum_filter1d minimum_position morphological_gradient + morphological_laplace percentile_filter prewitt rank_filter rotate shift + sobel spline_filter spline_filter1d standard_deviation sum sum_labels + uniform_filter uniform_filter1d value_indices variance white_tophat zoom + """, + ), + "signal": _names( + """ + BadCoefficients CZT StateSpace TransferFunction ZerosPolesGain ZoomFFT + abcd_normalize argrelextrema argrelmax argrelmin band_stop_obj bilinear + bilinear_zpk bode buttap butter buttord cheb1ap cheb1ord cheb2ap + cheb2ord cheby1 cheby2 check_COLA check_NOLA chirp choose_conv_method + coherence cont2discrete convolve convolve2d correlate correlate2d + correlation_lags csd cspline1d cspline1d_eval cspline2d cwt czt + czt_points dbode decimate deconvolve detrend dfreqresp dimpulse dlsim + dlti dstep ellip ellipap ellipord fftconvolve filtfilt find_peaks + findfreqs firls firwin firwin2 freqresp freqs freqs_zpk freqz freqz_sos + freqz_zpk gammatone gauss_spline gausspulse get_window group_delay + hilbert hilbert2 iircomb iirdesign iirfilter iirnotch iirpeak impulse + invres invresz istft kaiser_atten kaiser_beta kaiserord lfilter + lfilter_zi lfiltic lombscargle lp2bp lp2bp_zpk lp2bs lp2bs_zpk lp2hp + lp2hp_zpk lp2lp lp2lp_zpk lsim lti max_len_seq medfilt medfilt2d + minimum_phase morlet morlet2 normalize oaconvolve order_filter + peak_prominences peak_widths periodogram place_poles qmf qspline1d + qspline1d_eval qspline2d resample resample_poly residue residuez ricker + savgol_coeffs savgol_filter sawtooth sepfir2d sos2tf sos2zpk sosfilt + sosfilt_zi sosfiltfilt sosfreqz spectrogram spline_filter square ss2tf + ss2zpk step stft sweep_poly symiirorder1 symiirorder2 tf2sos tf2ss + tf2zpk unique_roots unit_impulse upfirdn vectorstrength welch wiener + zoom_fft zpk2sos zpk2ss zpk2tf + """, + ), + "sparse": _names( + """ + SparseEfficiencyWarning SparseWarning block_array block_diag bmat + coo_array coo_matrix csc_array csc_matrix csr_array csr_matrix dia_array + dia_matrix diags diags_array eye eye_array find get_index_dtype hstack + identity issparse isspmatrix isspmatrix_coo isspmatrix_csc + isspmatrix_csr isspmatrix_dia kron kronsum matrix_transpose permute_dims + rand random random_array safely_cast_index_arrays sparray spdiags + spmatrix swapaxes tril triu vstack + """, + ), + "sparse.csgraph": _names( + """ + connected_components + """, + ), + "sparse.linalg": _names( + """ + LinearOperator SuperLU aslinearoperator bicgstab cg cgs eigsh factorized + gmres lobpcg lsmr lsqr matrix_power minres norm spilu splu spsolve + spsolve_triangular svds + """, + ), + "spatial": _names( + """ + Delaunay KDTree distance_matrix + """, + ), + "special": _names( + """ + bdtr bdtrc bdtri beta betainc betaincinv betaln binom boxcox boxcox1p + btdtr btdtri cbrt chdtr chdtrc chdtri cosdg cosm1 cotdg digamma + ellipeinc ellipj ellipk ellipkinc ellipkm1 entr erf erfc erfcinv erfcx + erfinv exp1 exp10 exp2 expi expit expm1 expn exprel fdtr fdtrc fdtri + gamma gammainc gammaincc gammainccinv gammaincinv gammaln gammasgn gdtr + gdtrc huber i0 i0e i1 i1e inv_boxcox inv_boxcox1p j0 j1 k0 k0e k1 k1e + kl_div lambertw log1p log_expit log_ndtr log_softmax loggamma logit + logsumexp lpmv multigammaln nbdtr nbdtrc nbdtri ndtr ndtri pdtr pdtrc + pdtri poch polygamma pseudo_huber psi radian rel_entr rgamma round + shichi sici sinc sindg softmax sph_harm sph_harm_y spherical_yn tandg + wright_bessel xlog1py xlogy y0 y1 yn zeta zetac + """, + ), + "stats": _names( + """ + boxcox_llf entropy trim_mean zmap zscore + """, + ), +} class _Proxy: @@ -168,10 +266,10 @@ def call(*args: Any, **kwargs: Any) -> Any: def _subpackage(path: str, cupy: Any) -> types.ModuleType: module = types.ModuleType(f"cupyx.scipy.{path}") - allowed = _ALLOWED.get(path) + allowed = NAMES[path] def __getattr__(name: str) -> Any: - if name.startswith("__") or (allowed is not None and name not in allowed): + if name.startswith("__") or name not in allowed: raise AttributeError( f"module 'cupyx.scipy.{path}' has no attribute {name!r}" ) diff --git a/src/cunumpy/_linalg.py b/src/cunumpy/_linalg.py new file mode 100644 index 0000000..41e8ffb --- /dev/null +++ b/src/cunumpy/_linalg.py @@ -0,0 +1,103 @@ +"""SciPy linear algebra that ``cupyx.scipy.linalg`` lacks, on the array's backend.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from cunumpy.xp import get_array_module + + +def _moveaxis(xpm: Any, a: Any, source: int, destination: int) -> Any: + if source % a.ndim == destination % a.ndim: + return a + return xpm.moveaxis(a, source, destination) + + +def _as_inexact(xpm: Any, a: Any) -> Any: + a = xpm.asarray(a) + if a.ndim == 0: + a = xpm.reshape(a, (1,)) + if not xpm.isdtype(a.dtype, ("real floating", "complex floating")): + a = xpm.astype(a, xpm.float64) + return a + + +def solve_circulant( + c: Any, + b: Any, + singular: str = "raise", + tol: float | None = None, + caxis: int = -1, + baxis: int = 0, + outaxis: int = 0, +) -> Any: + """Solve ``C x = b`` for a circulant matrix ``C`` with first column `c`. + + The steps and arguments of :func:`scipy.linalg.solve_circulant`, with FFTs + on the backend of `c`: on the CuPy backend nothing leaves the device, and + the singularity check synchronizes once. Available as + ``xp.scipy.linalg.solve_circulant`` on both backends. + + Parameters + ---------- + c : array + First column of ``C``; a batch of columns along the other axes. + b : array + Right-hand side(s), of length ``len(c)`` along `baxis`. + singular : {"raise", "lstsq"}, optional + On a (near) singular ``C``: raise, or return the least-squares + solution. + tol : float, optional + Eigenvalues of ``C`` with magnitude at most `tol` count as zero. + Default ``max(abs(fft(c))) * len(c) * eps``. + caxis, baxis, outaxis : int, optional + The axis of `c`, `b` and the result along which the vectors lie. + + Returns + ------- + array + The solution ``x``, real unless `c` or `b` is complex. + + Raises + ------ + numpy.linalg.LinAlgError + If ``C`` is near singular and `singular` is ``"raise"``. + ValueError + If the lengths of `c` and `b` differ. + + Examples + -------- + >>> c = xp.asarray([4.0, 1.0, 0.0, 1.0]) + >>> xp.scipy.linalg.solve_circulant(c, xp.asarray([6.0, 6.0, 6.0, 6.0])) + array([1., 1., 1., 1.]) + """ + xpm = get_array_module(c) + c = _as_inexact(xpm, c) + b = _as_inexact(xpm, b) + n = c.shape[caxis] + if b.shape[baxis] != n: + raise ValueError(f"Shapes of c {c.shape} and b {b.shape} are incompatible") + fc = xpm.fft.fft(_moveaxis(xpm, c, caxis, -1), axis=-1) + abs_fc = xpm.abs(fc) + if tol is None: + # the tolerance of numpy.linalg.matrix_rank, as in SciPy + tol = xpm.max(abs_fc, axis=-1, keepdims=True) * n * np.finfo(np.float64).eps + near_zeros = abs_fc <= tol + is_near_singular = bool(xpm.any(near_zeros)) + if is_near_singular: + if singular == "raise": + raise np.linalg.LinAlgError("near singular circulant matrix.") + fc = xpm.where(near_zeros, xpm.ones_like(fc), fc) + fb = xpm.fft.fft(_moveaxis(xpm, b, baxis, -1), axis=-1) + q = fb / fc + if is_near_singular: + q = xpm.where(near_zeros, xpm.zeros_like(q), q) + x = xpm.fft.ifft(q, axis=-1) + if not ( + xpm.isdtype(c.dtype, "complex floating") + or xpm.isdtype(b.dtype, "complex floating") + ): + x = xpm.real(x) + return _moveaxis(xpm, x, -1, outaxis) diff --git a/src/cunumpy/_scipy_backend.py b/src/cunumpy/_scipy_backend.py index 3e00026..cc9f527 100644 --- a/src/cunumpy/_scipy_backend.py +++ b/src/cunumpy/_scipy_backend.py @@ -11,7 +11,9 @@ Names are resolved at every access, so a backend switch takes effect at once; nothing is imported until a name is used, and SciPy is not a dependency of cunumpy. ``cupyx.scipy`` covers only part of SciPy, and keyword arguments can -differ (SciPy's ``cg`` takes ``rtol``, CuPy's ``tol``). A SciPy sparse matrix +differ (SciPy's ``cg`` takes ``rtol``, CuPy's ``tol``). A few missing +functions are filled in by cunumpy on the device (``FILLS``, e.g. +``xp.scipy.linalg.solve_circulant``). A SciPy sparse matrix moves to the device once with ``xp.scipy.sparse.csr_matrix(host_matrix)`` on the CuPy backend, and back with ``matrix.get()``. See :doc:`/guides/solvers`. """ @@ -22,9 +24,10 @@ from types import ModuleType from typing import Any +from cunumpy import _linalg from cunumpy.xp import get_backend -__all__ = ["SUBMODULES", "ScipyNamespace", "scipy"] +__all__ = ["FILLS", "SUBMODULES", "ScipyNamespace", "scipy"] #: The SciPy subpackages forwarded by ``xp.scipy``: those that ``cupyx.scipy`` #: provides as well. @@ -43,6 +46,10 @@ "stats", ) +#: Functions ``cupyx.scipy`` lacks, implemented by cunumpy on the device; used +#: on the CuPy backend only, where the ``cupyx`` module has no such name. +FILLS = {"linalg": {"solve_circulant": _linalg.solve_circulant}} + _ROOTS = {"numpy": "scipy", "cupy": "cupyx.scipy"} _INSTALL = { "numpy": "SciPy is not installed (pip install scipy)", @@ -54,7 +61,8 @@ class ScipyNamespace: """A SciPy (sub)package of the active backend; see :data:`cunumpy.scipy`. Attribute access forwards to the module of the active backend. A name the - backend's module lacks raises ``AttributeError`` naming the backend; a + backend's module lacks raises ``AttributeError`` naming the backend, unless + cunumpy fills it in (``FILLS``, CuPy backend only); a missing SciPy (NumPy backend) or CuPy raises ``ImportError``. Parameters @@ -124,6 +132,9 @@ def __getattr__(self, name: str) -> Any: try: return getattr(module, name) except AttributeError: + fill = self._fill(name) + if fill is not None: + return fill other = "cupy" if get_backend() == "numpy" else "numpy" raise AttributeError( f"{module.__name__} has no attribute {name!r}: it is not available " @@ -131,6 +142,11 @@ def __getattr__(self, name: str) -> Any: f"{self._name(_ROOTS[other])})", ) from None + def _fill(self, name: str) -> Any: + if get_backend() != "cupy": + return None + return FILLS.get(self._path, {}).get(name) + def available(self, name: str) -> bool: """Return whether `name` exists in this namespace on the active backend. @@ -157,7 +173,7 @@ def available(self, name: str) -> bool: return False return True try: - return hasattr(self.resolve(), name) + return hasattr(self.resolve(), name) or self._fill(name) is not None except ImportError: return False @@ -172,6 +188,8 @@ def __dir__(self) -> list[str]: names = dir(self.resolve()) except ImportError: names = [] + if get_backend() == "cupy": + names = [*names, *FILLS.get(self._path, {})] return sorted(set(names) | set(children) | {"available", "resolve"}) diff --git a/tests/unit/test_fake_cupyx.py b/tests/unit/test_fake_cupyx.py index 3fa26cf..f392ad2 100644 --- a/tests/unit/test_fake_cupyx.py +++ b/tests/unit/test_fake_cupyx.py @@ -18,6 +18,7 @@ def run(code): def test_same_subpackages_as_xp_scipy(): assert _fake_cupyx.SUBPACKAGES == SUBMODULES + assert tuple(_fake_cupyx.NAMES) == SUBMODULES def test_splines_fit_and_evaluate_on_the_device(): @@ -86,6 +87,61 @@ def test_interpolate_has_only_the_names_cupy_provides(): assert "ok" in run(code) +def test_every_subpackage_has_only_the_names_cupy_provides(): + code = """ +import cunumpy as xp +import cupyx.scipy.linalg +xp.set_backend("cupy") +missing = { + "special": ("jv", "erfi"), + "linalg": ("solve_circulant", "inv", "eigh"), + "sparse.linalg": ("inv",), +} +present = { + "special": ("yn", "erf", "erfinv", "ndtri"), + "linalg": ("circulant", "lu_factor", "lu_solve"), + "sparse": ("csr_matrix", "kron", "bmat", "identity"), + "sparse.linalg": ("splu", "spsolve", "cg"), + "signal": ("argrelextrema",), + "fft": ("rfft", "irfft"), +} +ns = lambda path: eval("xp.scipy." + path) +for path, names in missing.items(): + for name in names: + assert not hasattr(ns(path).resolve(), name), (path, name) +for path, names in present.items(): + for name in names: + assert ns(path).available(name), (path, name) +# cunumpy fills in solve_circulant, which raw cupyx lacks +assert not hasattr(cupyx.scipy.linalg, "solve_circulant") +assert xp.scipy.linalg.available("solve_circulant") +print("ok") +""" + assert "ok" in run(code) + + +def test_solve_circulant_stays_on_the_device(): + code = """ +import numpy as np +import scipy.linalg +import cunumpy as xp +xp.set_backend("cupy") +rng = np.random.default_rng(0) +c_h, b_h = rng.random(6) + 2.0, rng.random((6, 3)) +x = xp.scipy.linalg.solve_circulant(xp.asarray(c_h), xp.asarray(b_h)) +assert xp.is_gpu(x) +assert np.allclose(xp.to_numpy(x), scipy.linalg.solve_circulant(c_h, b_h)) +try: + xp.scipy.linalg.solve_circulant(xp.asarray([1.0, 1.0, 0.0, 0.0]), xp.ones(4)) +except np.linalg.LinAlgError: + pass +else: + raise AssertionError("singular matrix accepted") +print("ok") +""" + assert "ok" in run(code) + + def test_other_subpackages_and_pinned_memory(): code = """ import numpy as np diff --git a/tests/unit/test_linalg.py b/tests/unit/test_linalg.py new file mode 100644 index 0000000..e17df92 --- /dev/null +++ b/tests/unit/test_linalg.py @@ -0,0 +1,44 @@ +"""Tests for `solve_circulant`, the cunumpy fill-in for `cupyx.scipy.linalg`.""" + +import numpy as np +import pytest + +from cunumpy._linalg import solve_circulant + +scipy_linalg = pytest.importorskip("scipy.linalg") + +RNG = np.random.default_rng(0) + + +@pytest.mark.parametrize( + ("c", "b", "kwargs"), + [ + (RNG.random(5), RNG.random(5), {}), + (RNG.random(5), RNG.random((5, 3)), {}), + (RNG.random((2, 5)), RNG.random((5, 2)), {}), + (RNG.random(6) + 1j * RNG.random(6), RNG.random(6), {}), + ( + RNG.random((5, 3)), + RNG.random((3, 5)), + {"caxis": 0, "baxis": 1, "outaxis": -1}, + ), + (np.array([1.0, 1.0, 0.0, 0.0]), np.ones(4), {"singular": "lstsq"}), + (np.array([4, 1, 0, 1]), np.array([6, 6, 6, 6]), {}), + (np.float32([4, 1, 0, 1]), np.float32([1, 2, 3, 4]), {}), + ], +) +def test_matches_scipy(c, b, kwargs): + x = solve_circulant(c, b, **kwargs) + ref = scipy_linalg.solve_circulant(c, b, **kwargs) + assert x.dtype == ref.dtype + np.testing.assert_allclose(x, ref, rtol=1e-6, atol=1e-12) + + +def test_singular_raises(): + with pytest.raises(np.linalg.LinAlgError, match="near singular"): + solve_circulant(np.array([1.0, 1.0, 0.0, 0.0]), np.ones(4)) + + +def test_incompatible_shapes(): + with pytest.raises(ValueError, match="incompatible"): + solve_circulant(np.ones(4), np.ones(5)) diff --git a/tests/unit/test_scipy_backend.py b/tests/unit/test_scipy_backend.py index 01f9777..ed928b7 100644 --- a/tests/unit/test_scipy_backend.py +++ b/tests/unit/test_scipy_backend.py @@ -7,6 +7,7 @@ import pytest import cunumpy as xp +from cunumpy import _linalg from cunumpy import _scipy_backend as scipy_backend from cunumpy._scipy_backend import SUBMODULES, ScipyNamespace @@ -85,6 +86,21 @@ def test_cupy_backend_forwards_to_cupyx(fake_cupyx): xp.scipy.ndimage.resolve() +def test_cupy_backend_fills_in_missing_functions(fake_cupyx): + modules = {"cupyx.scipy.linalg": types.ModuleType("cupyx.scipy.linalg")} + fake_cupyx["cupyx.scipy"].linalg = modules["cupyx.scipy.linalg"] + sys.modules.update(modules) + try: + assert xp.scipy.linalg.solve_circulant is _linalg.solve_circulant + assert xp.scipy.linalg.available("solve_circulant") + assert "solve_circulant" in dir(xp.scipy.linalg) + # a cupyx version would win over the fill + modules["cupyx.scipy.linalg"].solve_circulant = "device solve_circulant" + assert xp.scipy.linalg.solve_circulant == "device solve_circulant" + finally: + del sys.modules["cupyx.scipy.linalg"] + + def test_backend_switch_takes_effect_immediately(fake_cupyx, monkeypatch): pytest.importorskip("scipy") import scipy.special From 702f263321cbd78b801f6e1aaaaf9d5eeae26099 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 22:19:07 +0200 Subject: [PATCH 3/9] Add xp.optimize.fsolve/root/minimize and xp.integrate.quad/odeint with device callbacks --- CHANGELOG.md | 4 + docs/source/api.md | 3 +- docs/source/guides/solvers.md | 28 ++- docs/source/reference/index.md | 1 + docs/source/reference/integrate.md | 8 + docs/source/reference/optimize.md | 5 +- src/cunumpy/LLM_GUIDE.md | 1 + src/cunumpy/__init__.py | 2 + src/cunumpy/__init__.pyi | 1 + src/cunumpy/_scipy_host.py | 370 +++++++++++++++++++++++++++++ src/cunumpy/integrate.py | 10 + src/cunumpy/optimize.py | 11 +- tests/unit/test_scipy_host.py | 136 +++++++++++ 13 files changed, 572 insertions(+), 8 deletions(-) create mode 100644 docs/source/reference/integrate.md create mode 100644 src/cunumpy/_scipy_host.py create mode 100644 src/cunumpy/integrate.py create mode 100644 tests/unit/test_scipy_host.py diff --git a/CHANGELOG.md b/CHANGELOG.md index a401cdb..f0a56ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 runs its CuPy path without a GPU. Every `cupyx.scipy` subpackage has only the names CuPy v14.2.0 provides, so e.g. `special.jv` fails as on a GPU. `cupyx.empty_pinned`/`zeros_pinned` return NumPy arrays. +- `xp.optimize.fsolve`, `root`, `minimize` and `xp.integrate.quad`, `odeint` + (`cupyx.scipy` has neither): SciPy on the host with SciPy's arguments, the + callbacks get their arrays on the backend of `x0`/`y0`, and results come + back on that backend. - `xp.scipy.linalg.solve_circulant` on the CuPy backend, which `cupyx.scipy` lacks: SciPy's FFT solve and arguments, on the device. - Solvers guide: splines on the device with `xp.scipy.interpolate`, and diff --git a/docs/source/api.md b/docs/source/api.md index 3771573..fdb68db 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -49,7 +49,8 @@ The submodules are named so that they do not hide a NumPy name (`rng`, not | [`cunumpy.cuda`](reference/cuda.md) | CUDA only | device selection and memory, `stream`, streams/events, `pin_memory`, debug mode, CUDA headers and source tools (`cuda_include_dir`, `parse_cuda_signature`) | | [`cunumpy.rng`](reference/rng.md) | both | `random_streams`, `get_rng`, `philox_*` | | [`cunumpy.algorithms`](reference/algorithms.md) | both | `morton_*`, `sort_by_key`, `segment_sum`, `compact_by_mask` | -| [`cunumpy.optimize`](reference/optimize.md) | both | `newton` (batched Newton/secant), `NewtonResult` | +| [`cunumpy.optimize`](reference/optimize.md) | both | `newton` (batched Newton/secant), `NewtonResult`; `fsolve`, `root`, `minimize` (SciPy on the host, callbacks on the device) | +| [`cunumpy.integrate`](reference/integrate.md) | both | `quad`, `odeint` (SciPy on the host, callbacks on the device) | | [`cunumpy.mpi`](reference/mpi.md) | both | `mpi_buffer`, `MPIStaging`, CUDA-aware MPI detection, `synchronize_for_mpi`, and the MPI stand-ins of maybempi (`get_mpi`, `local_rank`, ...) | | [`cunumpy.profiling`](reference/profiling.md) | both | `timed_region`, `nvtx_range`, `count_transfers`, `assert_no_transfers`, `TransferBudget` | | [`cunumpy.memory`](reference/memory.md) | both | `HostStaging`, `HostCopy`, `DeviceMirror` | diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index 29ead61..ee2d948 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -46,7 +46,7 @@ def periodic_poisson(rho, length): `xp.scipy.special.available("erfcx")` where a fallback is possible. Commonly missing: `special.jv` (only `j0`, `j1`, `yn`), `special.erfi`, `linalg.solve`/`inv` (use `xp.linalg`), `RectBivariateSpline`, and all of - `integrate` and `optimize`. + `integrate` and `optimize` (see `xp.optimize` and `xp.integrate` below). * cunumpy fills in `xp.scipy.linalg.solve_circulant` on the CuPy backend: an FFT solve on the device, with SciPy's arguments (a circulant preconditioner per iteration stays on the GPU). @@ -93,7 +93,7 @@ transfers. SciPy's optimizers run on the host, and `cupyx.scipy` has no `optimize`. Sequential scalar problems (one `fsolve`, a `quad` integral, an ODE) belong on -the host at setup (see `xp.setup_on_host`). Many independent scalar equations +the host at setup (see the next section and `xp.setup_on_host`). Many independent scalar equations (a flux surface on every ray, an inverse mapping at every marker) are one array problem: `xp.optimize.newton` runs Newton's method (with `fprime`) or the secant method on all of them at once, with the steps of @@ -110,6 +110,30 @@ r = xp.optimize.newton(residual, xp.full(levels.shape, 0.3)) Each iteration synchronizes once to test convergence. `full_output=True` returns `root`, `converged` and `zero_der` arrays. +## SciPy solvers with device callbacks + +For a single problem, `xp.optimize.fsolve`, `root` and `minimize` and +`xp.integrate.quad` and `odeint` run SciPy on the host with SciPy's +arguments, so they work on both backends. With `x0` (`y0` for `odeint`) on +the device, the callbacks get `x` on the device and the result comes back on +the device, so a residual written with `xp` that uses device data needs no +host copy of that data: + +```python +B = equilibrium.b_field_on_grid() # device array + +def residual(x): + return interpolate(B, x) - b_target + +x = xp.optimize.fsolve(residual, xp.asarray([0.5, 0.0])) # device array +``` + +Every evaluation copies `x` to the device and the result to the host, so use +them at setup, not in a time loop. `quad` calls its integrand with a Python +float; a 0-d device array returned is copied to the host. Bounds and +constraint matrices of `minimize` are host data; constraint functions get +device arrays like `fun`. + ## Fused elementwise updates: `xp.kernels.fuse` On the GPU, `(gamma - 1) * (E - 0.5 * rho * u**2)` runs as five kernels, each diff --git a/docs/source/reference/index.md b/docs/source/reference/index.md index 5b9c93f..2c74f22 100644 --- a/docs/source/reference/index.md +++ b/docs/source/reference/index.md @@ -15,6 +15,7 @@ cuda rng algorithms optimize +integrate mpi profiling memory diff --git a/docs/source/reference/integrate.md b/docs/source/reference/integrate.md new file mode 100644 index 0000000..94a36ce --- /dev/null +++ b/docs/source/reference/integrate.md @@ -0,0 +1,8 @@ +# Integration: `cunumpy.integrate` + +SciPy's `quad` and `odeint` with callbacks on the device. + +```{eval-rst} +.. automodule:: cunumpy.integrate + :members: +``` diff --git a/docs/source/reference/optimize.md b/docs/source/reference/optimize.md index abf497f..e1a2f7e 100644 --- a/docs/source/reference/optimize.md +++ b/docs/source/reference/optimize.md @@ -1,6 +1,7 @@ -# Root finding: `cunumpy.optimize` +# Optimization: `cunumpy.optimize` -Batched Newton and secant iterations on either backend. +Batched Newton and secant iterations on the device, and SciPy's `fsolve`, +`root` and `minimize` with callbacks on the device. ```{eval-rst} .. automodule:: cunumpy.optimize diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index d0aa12e..0fea615 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -101,6 +101,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | SciPy (sparse, sparse.linalg, fft, special, ndimage, ...) on either backend | `xp.scipy..` (SciPy or `cupyx.scipy`); `xp.scipy.special.available(name)` | | splines on the device | `xp.scipy.interpolate.UnivariateSpline`, `make_interp_spline`, `NdBSpline` (no `RectBivariateSpline`: `make_interp_spline` per axis + `NdBSpline`) | | circulant solve (FFT preconditioner) on the device | `xp.scipy.linalg.solve_circulant(c, b)` (cunumpy fill-in on CuPy, SciPy's arguments) | +| one root/minimization/integral/ODE at setup with an `xp` callback using device data | `xp.optimize.fsolve`/`root`/`minimize`, `xp.integrate.quad`/`odeint` (SciPy on the host, callback args and results on `x0`'s backend) | | many independent scalar root problems (one per ray, per marker) | `xp.optimize.newton(func, x0_array, fprime=None)` (SciPy's array Newton/secant, on the device) | | chain of elementwise operations as one GPU kernel | `@xp.kernels.fuse` (`cupy.fuse` for CuPy arrays, plain call otherwise) | | PETSc solve on device arrays without copies | `xp.petsc.petsc_vec(array)` (CUDA/HIP petsc4py for CuPy arrays); `xp.synchronize()` around PETSc calls | diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 48106b2..1d5247b 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -19,6 +19,7 @@ algorithms, arguments, cuda, + integrate, kernels, memory, mpi, @@ -110,6 +111,7 @@ def require_version(minimum: str) -> None: "get_array_module", "get_backend", "host_call", + "integrate", "is_cpu", "is_gpu", "kernels", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 8625124..9e23036 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -11,6 +11,7 @@ from numpy import * from cunumpy import algorithms as algorithms from cunumpy import arguments as arguments from cunumpy import cuda as cuda +from cunumpy import integrate as integrate from cunumpy import kernels as kernels from cunumpy import memory as memory from cunumpy import mpi as mpi diff --git a/src/cunumpy/_scipy_host.py b/src/cunumpy/_scipy_host.py new file mode 100644 index 0000000..0a62dff --- /dev/null +++ b/src/cunumpy/_scipy_host.py @@ -0,0 +1,370 @@ +"""SciPy solvers that ``cupyx.scipy`` lacks, run on the host with callbacks on the device.""" + +from __future__ import annotations + +import copy +import functools +import importlib +from collections.abc import Callable +from typing import Any + +from cunumpy import xp +from cunumpy._host import _to_device + + +def _scipy(name: str, function: str) -> Any: + try: + return importlib.import_module(f"scipy.{name}") + except ImportError as error: + raise ImportError( + f"cunumpy.{name}.{function} needs SciPy (pip install scipy)", + ) from error + + +def _host(value: Any) -> Any: + """Copy device arrays in `value` (or a tuple/list of them) to the host.""" + if isinstance(value, (tuple, list)): + return type(value)(_host(v) for v in value) + return xp.to_numpy(value) if xp.is_gpu(value) else value + + +def _device(value: Any) -> Any: + """Copy NumPy arrays in `value` (a tuple, list or dict of them) to the device.""" + if isinstance(value, dict): # OptimizeResult, infodict: a copy of the same type + out = copy.copy(value) + for key, item in value.items(): + out[key] = _device(item) + return out + if isinstance(value, (tuple, list)): + return type(value)(_device(v) for v in value) + return _to_device(value) + + +class _Bridge: + """Callbacks get arrays on the backend of `x0`; SciPy sees host arrays.""" + + def __init__(self, x0: Any) -> None: + self.device = xp.is_gpu(x0) + + def callback(self, fun: Any, n_arrays: int = 1) -> Any: + # the first `n_arrays` positional and all keyword arguments of `fun` + # arrive on the device + if not (self.device and callable(fun)): + return fun + + # wraps: SciPy reads the signature (a callback named `intermediate_result`) + @functools.wraps(fun) + def call(*args: Any, **kwargs: Any) -> Any: + args = (*_device(list(args[:n_arrays])), *args[n_arrays:]) + kwargs = {k: _device(v) for k, v in kwargs.items()} # intermediate_result + return _host(fun(*args, **kwargs)) + + return call + + def result(self, value: Any) -> Any: + return _device(value) if self.device else value + + def constraints(self, constraints: Any) -> Any: + if isinstance(constraints, (tuple, list)): + return type(constraints)(self.constraints(c) for c in constraints) + if isinstance(constraints, dict): # {"type": "ineq", "fun": ..., "jac": ...} + out = dict(constraints) + for key in ("fun", "jac"): + out[key] = self.callback(out.get(key)) + return {k: v for k, v in out.items() if v is not None} + out = copy.copy(constraints) # NonlinearConstraint, LinearConstraint + for name, n_arrays in (("fun", 1), ("jac", 1), ("hess", 2)): + if callable(getattr(out, name, None)): + setattr(out, name, self.callback(getattr(out, name), n_arrays)) + return out + + +def fsolve( + func: Callable[..., Any], + x0: Any, + args: tuple = (), + fprime: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Find a root of a system of equations with :func:`scipy.optimize.fsolve`. + + SciPy (MINPACK) runs on the host. With `x0` on the device, `func` and + `fprime` get ``x`` on the device, their results are copied to the host, + and the root comes back on the device: two transfers per evaluation, so + use it for setup, not in a time loop. + + Parameters + ---------- + func : callable + ``func(x, *args)``, the residual, an array of the shape of `x0`. + x0 : array + Initial guess; its backend is the backend of the callbacks and result. + args : tuple, optional + Extra arguments of `func` and `fprime`, passed unchanged. + fprime : callable, optional + ``fprime(x, *args)``, the Jacobian. + **kwargs + The other arguments of :func:`scipy.optimize.fsolve`. + + Returns + ------- + array or tuple + The root, on the backend of `x0`; with ``full_output=True`` SciPy's + tuple, with its arrays on that backend. + + See Also + -------- + root : The same with a choice of method. + newton : Many independent scalar equations at once, on the device. + + Examples + -------- + >>> xp.optimize.fsolve(lambda x: x**2 - xp.asarray([4.0, 9.0]), xp.ones(2)) + array([2., 3.]) + """ + optimize = _scipy("optimize", "fsolve") + bridge = _Bridge(x0) + out = optimize.fsolve( + bridge.callback(func), + _host(x0), + args=args, + fprime=bridge.callback(fprime), + **kwargs, + ) + return bridge.result(out) + + +def root( + fun: Callable[..., Any], + x0: Any, + args: tuple = (), + jac: Any = None, + callback: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Find a root of a system of equations with :func:`scipy.optimize.root`. + + SciPy runs on the host; with `x0` on the device, `fun`, `jac` and + `callback` get ``x`` on the device and the arrays of the result are + copied back to the device. See :func:`fsolve`. + + Parameters + ---------- + fun : callable + ``fun(x, *args)``, the residual (or ``(residual, jacobian)`` with + ``jac=True``). + x0 : array + Initial guess; its backend is the backend of the callbacks and result. + args : tuple, optional + Extra arguments of `fun` and `jac`, passed unchanged. + jac : callable or bool, optional + ``jac(x, *args)``, the Jacobian; ``True`` if `fun` returns it too. + callback : callable, optional + ``callback(x, f)``, called after each iteration (some methods only). + **kwargs + The other arguments of :func:`scipy.optimize.root` (``method``, + ``tol``, ``options``). + + Returns + ------- + OptimizeResult + SciPy's result; ``x``, ``fun`` and the other arrays on the backend of + `x0`. + + Examples + -------- + >>> res = xp.optimize.root(lambda x: x**3 - 8.0, xp.ones(1)) + >>> res.success, res.x + (True, array([2.])) + """ + optimize = _scipy("optimize", "root") + bridge = _Bridge(x0) + out = optimize.root( + bridge.callback(fun), + _host(x0), + args=args, + jac=bridge.callback(jac), + callback=bridge.callback(callback, 2), + **kwargs, + ) + return bridge.result(out) + + +def minimize( + fun: Callable[..., Any], + x0: Any, + args: tuple = (), + jac: Any = None, + hess: Any = None, + hessp: Callable[..., Any] | None = None, + bounds: Any = None, + constraints: Any = (), + callback: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Minimize a scalar function with :func:`scipy.optimize.minimize`. + + SciPy runs on the host; with `x0` on the device, every callback (`fun`, + `jac`, `hess`, `hessp`, `callback` and the functions of `constraints`) + gets its arrays on the device, and the arrays of the result are copied + back to the device. See :func:`fsolve`. + + Parameters + ---------- + fun : callable + ``fun(x, *args)``, a scalar (or ``(value, gradient)`` with + ``jac=True``). + x0 : array + Initial guess; its backend is the backend of the callbacks and result. + args : tuple, optional + Extra arguments of the callbacks, passed unchanged. + jac, hess : callable, str or bool, optional + Gradient and Hessian, as in SciPy (a callable, a finite-difference + scheme such as ``"2-point"``, or a ``HessianUpdateStrategy``). + hessp : callable, optional + ``hessp(x, p, *args)``, the Hessian times ``p``. + bounds : sequence or Bounds, optional + Bounds on ``x``, as in SciPy (host data). + constraints : dict, constraint object or sequence of them, optional + As in SciPy; their functions get device arrays like `fun`, their + bounds and matrices are host data. + callback : callable, optional + Called after each iteration, as in SciPy (``callback(intermediate_result)`` + or ``callback(xk)``). + **kwargs + The other arguments of :func:`scipy.optimize.minimize` (``method``, + ``tol``, ``options``). + + Returns + ------- + OptimizeResult + SciPy's result; ``x``, ``jac`` and the other arrays on the backend of + `x0`. + + Examples + -------- + >>> res = xp.optimize.minimize(lambda x: xp.sum((x - 1.0) ** 2), xp.zeros(2)) + >>> bool(xp.allclose(res.x, 1.0, atol=1e-6)) + True + """ + optimize = _scipy("optimize", "minimize") + bridge = _Bridge(x0) + if bridge.device: + constraints = bridge.constraints(constraints) + out = optimize.minimize( + bridge.callback(fun), + _host(x0), + args=args, + jac=bridge.callback(jac), + hess=bridge.callback(hess), + hessp=bridge.callback(hessp, 2), + bounds=bounds, + constraints=constraints, + callback=bridge.callback(callback, 2), # trust-constr: (xk, state) + **kwargs, + ) + return bridge.result(out) + + +def quad( + func: Callable[..., Any], + a: float, + b: float, + args: tuple = (), + **kwargs: Any, +) -> Any: + """Integrate a scalar function with :func:`scipy.integrate.quad`. + + SciPy (QUADPACK) runs on the host and calls `func` with a Python float. + `func` may compute on the device and return a 0-d device array: it is + copied to the host, one synchronization per evaluation. For many + integrals at once, evaluate the integrand on a fixed grid instead. + + Parameters + ---------- + func : callable + ``func(x, *args)``, a scalar. + a, b : float + Integration limits (``numpy.inf`` allowed). + args : tuple, optional + Extra arguments of `func`, passed unchanged. + **kwargs + The other arguments of :func:`scipy.integrate.quad`. + + Returns + ------- + tuple + SciPy's result, ``(value, abserr)`` plus the outputs that + ``full_output`` adds, on the host. + + Examples + -------- + >>> value, error = xp.integrate.quad(lambda x: x**2, 0.0, 3.0) + >>> round(value, 12) + 9.0 + """ + integrate = _scipy("integrate", "quad") + return integrate.quad( + lambda x, *a: _host(func(x, *a)), + _host(a), + _host(b), + args=args, + **kwargs, + ) + + +def odeint( + func: Callable[..., Any], + y0: Any, + t: Any, + args: tuple = (), + Dfun: Callable[..., Any] | None = None, + **kwargs: Any, +) -> Any: + """Integrate a system of ODEs with :func:`scipy.integrate.odeint`. + + SciPy (LSODA) runs on the host; with `y0` on the device, `func` and + `Dfun` get ``y`` on the device, and the solution comes back on the + device. See :func:`cunumpy.optimize.fsolve` for the cost. + + Parameters + ---------- + func : callable + ``func(y, t, *args)`` (``func(t, y, *args)`` with ``tfirst=True``), + the derivative of ``y``. + y0 : array + Initial condition; its backend is the backend of the callbacks and + result. + t : array + Times at which to return ``y``, starting with the time of `y0`; either + backend. + args : tuple, optional + Extra arguments of `func` and `Dfun`, passed unchanged. + Dfun : callable, optional + The Jacobian of `func`, with the same arguments. + **kwargs + The other arguments of :func:`scipy.integrate.odeint`. + + Returns + ------- + array or tuple + ``y`` at the times `t`, shape ``(len(t), len(y0))``, on the backend of + `y0`; with ``full_output=True`` also SciPy's info dict. + + Examples + -------- + >>> y = xp.integrate.odeint(lambda y, t: -y, xp.ones(1), xp.asarray([0.0, 1.0])) + >>> bool(xp.allclose(y[-1], np.exp(-1.0), rtol=1e-6)) + True + """ + integrate = _scipy("integrate", "odeint") + bridge = _Bridge(y0) + out = integrate.odeint( + bridge.callback(func, 2), + _host(y0), + _host(t), + args=args, + Dfun=bridge.callback(Dfun, 2), + **kwargs, + ) + return bridge.result(out) diff --git a/src/cunumpy/integrate.py b/src/cunumpy/integrate.py new file mode 100644 index 0000000..fd9b5a2 --- /dev/null +++ b/src/cunumpy/integrate.py @@ -0,0 +1,10 @@ +"""Integration on either backend. + +``cupyx.scipy`` has no ``integrate``. :func:`quad` and :func:`odeint` run +SciPy on the host, with callbacks that may compute on the device and results +on the backend of the input. See :doc:`/guides/solvers`. +""" + +from cunumpy._scipy_host import odeint, quad + +__all__ = ["odeint", "quad"] diff --git a/src/cunumpy/optimize.py b/src/cunumpy/optimize.py index 3014989..86b612d 100644 --- a/src/cunumpy/optimize.py +++ b/src/cunumpy/optimize.py @@ -1,14 +1,19 @@ """Root finding on either backend. SciPy's optimizers run on the host only, and ``cupyx.scipy`` has no -``optimize``. This module has the parts that profit from the device: many -independent problems solved at once, each step one array operation:: +``optimize``. :func:`newton` runs on the device: many independent problems +solved at once, each step one array operation:: r = xp.optimize.newton(lambda r: psi(r) - levels, r0) +:func:`fsolve`, :func:`root` and :func:`minimize` run SciPy on the host, with +callbacks that get their arrays on the backend of ``x0`` (so code written with +``xp`` that uses device data works unchanged) and results on that backend. + See :doc:`/guides/solvers`. """ from cunumpy._optimize import NewtonResult, newton +from cunumpy._scipy_host import fsolve, minimize, root -__all__ = ["NewtonResult", "newton"] +__all__ = ["NewtonResult", "fsolve", "minimize", "newton", "root"] diff --git a/tests/unit/test_scipy_host.py b/tests/unit/test_scipy_host.py new file mode 100644 index 0000000..5a20d3d --- /dev/null +++ b/tests/unit/test_scipy_host.py @@ -0,0 +1,136 @@ +"""SciPy solvers cupyx lacks: `xp.optimize.fsolve/root/minimize`, `xp.integrate`.""" + +from pathlib import Path + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy.kernel_testing import run_in_fake_cupy_subprocess + +optimize = pytest.importorskip("scipy.optimize") +integrate = pytest.importorskip("scipy.integrate") + +SRC = str(Path(__file__).resolve().parents[2] / "src") + + +def residual(x): + return np.array( + [x[0] + 0.5 * (x[0] - x[1]) ** 3 - 1.0, 0.5 * (x[1] - x[0]) ** 3 + x[1]] + ) + + +def test_exported(): + assert set(xp.optimize.__all__) >= {"fsolve", "root", "minimize", "newton"} + assert xp.integrate.__all__ == ["odeint", "quad"] + assert "integrate" in xp.__all__ + + +def test_host_arrays_give_scipy_results(): + x0 = np.zeros(2) + np.testing.assert_allclose( + xp.optimize.fsolve(residual, x0), optimize.fsolve(residual, x0) + ) + _, info, ier, _ = xp.optimize.fsolve(residual, x0, full_output=True) + assert ier == 1 and info["nfev"] > 0 + res = xp.optimize.root(residual, x0, method="lm") + np.testing.assert_allclose(res.x, optimize.root(residual, x0, method="lm").x) + rosen = optimize.rosen + res = xp.optimize.minimize( + rosen, np.zeros(3), jac=optimize.rosen_der, method="BFGS" + ) + np.testing.assert_allclose(res.x, 1.0, atol=1e-5) + value, _ = xp.integrate.quad(np.exp, 0.0, 1.0) + assert value == pytest.approx(np.e - 1.0) + t = np.linspace(0.0, 1.0, 5) + np.testing.assert_allclose( + xp.integrate.odeint(lambda y, t: -2.0 * y, np.ones(2), t), + integrate.odeint(lambda y, t: -2.0 * y, np.ones(2), t), + ) + + +def test_missing_scipy(monkeypatch): + import sys + + monkeypatch.setitem(sys.modules, "scipy.optimize", None) + with pytest.raises(ImportError, match=r"cunumpy.optimize.fsolve needs SciPy"): + xp.optimize.fsolve(residual, np.zeros(2)) + + +def test_callbacks_and_results_on_the_device(): + code = """ +import numpy as np +import scipy.optimize as so +import cunumpy as xp +xp.set_backend("cupy") + +target = xp.asarray([4.0, 9.0]) # device data captured by the callbacks + +def on_device(*arrays): + for a in arrays: + assert xp.is_gpu(a), type(a) + +def f(x): + on_device(x) + return x**2 - target + +def jac(x): + on_device(x) + return xp.diag(2 * x) + +x = xp.optimize.fsolve(f, xp.ones(2), fprime=jac) +assert xp.is_gpu(x) and np.allclose(xp.to_numpy(x), [2.0, 3.0]) +x, info, ier, msg = xp.optimize.fsolve(f, xp.ones(2), full_output=True) +assert xp.is_gpu(x) and xp.is_gpu(info["fvec"]) + +seen = [] +res = xp.optimize.root(f, xp.ones(2), jac=jac, method="hybr") +assert res.success and xp.is_gpu(res.x) and xp.is_gpu(res.fun) +res = xp.optimize.root(lambda x: (f(x), jac(x)), xp.ones(2), jac=True) +assert np.allclose(xp.to_numpy(res.x), [2.0, 3.0]) + +def fun(x): + on_device(x) + return xp.sum((x - target) ** 2) + +def grad(x): + on_device(x) + return 2 * (x - target) + +def hessp(x, p): + on_device(x, p) + return 2 * p + +def callback(intermediate_result): + on_device(intermediate_result.x) + seen.append(1) + +res = xp.optimize.minimize(fun, xp.zeros(2), jac=grad, hessp=hessp, method="trust-ncg", callback=callback) +assert seen and xp.is_gpu(res.x) and np.allclose(xp.to_numpy(res.x), [4.0, 9.0]) +cons = [ + {"type": "ineq", "fun": lambda x: (on_device(x), 5.0 - x[1])[1]}, + so.NonlinearConstraint(lambda x: (on_device(x), x[0])[1], -np.inf, 3.0), + so.LinearConstraint(np.asarray([[1.0, 1.0]]), -np.inf, 7.0), +] +bounds = so.Bounds(np.zeros(2), np.full(2, 10.0)) +res = xp.optimize.minimize(fun, xp.ones(2), method="SLSQP", constraints=cons, bounds=bounds) +assert res.success and np.allclose(xp.to_numpy(res.x), [2.0, 5.0], atol=1e-6) +res = xp.optimize.minimize(fun, xp.ones(2), bounds=[(0.0, 3.0), (0.0, 3.0)], method="L-BFGS-B") +assert np.allclose(xp.to_numpy(res.x), [3.0, 3.0]) + +rate = xp.asarray([1.0, 2.0]) +def rhs(y, t): + on_device(y) + return -rate * y +t = xp.linspace(0.0, 1.0, 5) +y = xp.integrate.odeint(rhs, xp.ones(2), t) +assert xp.is_gpu(y) and y.shape == (5, 2) +assert np.allclose(xp.to_numpy(y[-1]), np.exp([-1.0, -2.0]), rtol=1e-6) +y = xp.integrate.odeint(lambda t, y: rhs(y, t), xp.ones(2), t, tfirst=True) +assert np.allclose(xp.to_numpy(y[-1]), np.exp([-1.0, -2.0]), rtol=1e-6) + +value, error = xp.integrate.quad(lambda s: xp.sum(rate) * s, 0.0, 1.0) +assert abs(value - 1.5) < 1e-12 +print("ok") +""" + assert "ok" in run_in_fake_cupy_subprocess(code, env={"PYTHONPATH": SRC}).stdout From 617c16cd79f7f44273da01bc467e297be785d22c Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 22:27:36 +0200 Subject: [PATCH 4/9] Added MIT license --- CHANGELOG.md | 1 + LICENSE.txt | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+) create mode 100644 LICENSE.txt diff --git a/CHANGELOG.md b/CHANGELOG.md index f0a56ac..652f241 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- MIT license (`LICENSE.txt`, referenced by `pyproject.toml`). - `xp.optimize.newton`: Newton's method or the secant method on many independent scalar equations at once, with the steps of `scipy.optimize.newton` for an array `x0`, on the array's backend (no host diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000..4b4eeec --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Max Lindqvist + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. From 82b7160fd080e1e4f0787ee6818eaba68ee0013a Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 22:50:11 +0200 Subject: [PATCH 5/9] Support CuPy >= 14.0: fake cupyx.scipy names of 14.0.0, checked against CuPy in the GPU CI --- .gitlab-ci.yml | 16 ++- CHANGELOG.md | 9 +- docs/source/guides/solvers.md | 4 +- docs/source/installation.md | 7 + docs/source/reference/top-level.md | 1 + src/cunumpy/__init__.py | 4 + src/cunumpy/__init__.pyi | 4 + src/cunumpy/_fake_cupyx.py | 206 +++++++++++++++++++---------- src/cunumpy/xp.py | 53 +++++++- testing/cupyx_names.py | 75 +++++++++++ tests/unit/test_cupyx_names.py | 84 ++++++++++++ 11 files changed, 386 insertions(+), 77 deletions(-) create mode 100644 testing/cupyx_names.py create mode 100644 tests/unit/test_cupyx_names.py diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 86f7a54..1e91db3 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -3,6 +3,8 @@ variables: PIP_DISABLE_PIP_VERSION_CHECK: "1" CUNUMPY_BACKEND: "cupy" CUNUMPY_REQUIRE_CUDA: "1" + # empty: the newest CuPy; gpu_tests_min_cupy pins the oldest supported one + CUPY_PIN: "" stages: - test @@ -19,7 +21,7 @@ stages: - module load fftw-serial/3.3.10 - nvidia-smi # Keep NVRTC, CUDA libraries and headers on the same toolkit generation. - - python3 -m pip install --user "cupy-cuda13x[ctk]" "cuda-toolkit==13.2.*" + - python3 -m pip install --user "cupy-cuda13x[ctk]${CUPY_PIN}" "cuda-toolkit==13.2.*" - python3 -m pip install --user -e ".[test]" - export PATH="$HOME/.local/bin:$PATH" - python3 -c "import cupy; cupy.show_config()" @@ -28,6 +30,18 @@ gpu_tests: extends: .gpu_job script: - python3 -m pytest -q tests + - python3 testing/cupyx_names.py + +# The oldest supported CuPy (cunumpy.MIN_CUPY_VERSION): the whole suite, and the +# fake cupyx.scipy must list exactly its names (_fake_cupyx.NAMES). +gpu_tests_min_cupy: + extends: .gpu_job + variables: + CUPY_PIN: "==14.0.0" + script: + - python3 -c "import cupy, cunumpy; assert cupy.__version__ == cunumpy.MIN_CUPY_VERSION, cupy.__version__" + - python3 -m pytest -q tests + - python3 testing/cupyx_names.py --exact gpu_sanitizers: extends: .gpu_job diff --git a/CHANGELOG.md b/CHANGELOG.md index 652f241..bff2a62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- Supported CuPy versions: 14.0 or newer (`xp.MIN_CUPY_VERSION`, + `xp.cupy_version_supported`); an older CuPy warns when the CuPy backend is + first checked, and `backend_info()` reports the minimum. The GPU CI runs the + whole suite on CuPy 14.0.0 too (`gpu_tests_min_cupy`). - MIT license (`LICENSE.txt`, referenced by `pyproject.toml`). - `xp.optimize.newton`: Newton's method or the secant method on many independent scalar equations at once, with the steps of @@ -16,7 +20,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - The fake CuPy has `cupyx.scipy`: SciPy behind the CuPy rules (host arrays rejected, fake device arrays and proxy objects returned), so `xp.scipy` code runs its CuPy path without a GPU. Every `cupyx.scipy` subpackage has only - the names CuPy v14.2.0 provides, so e.g. `special.jv` fails as on a GPU. + the names the oldest supported CuPy (14.0.0) provides, so e.g. `special.jv` + or `sparse.csr_array` fails as on a GPU. `testing/cupyx_names.py` and + `tests/unit/test_cupyx_names.py` check these names against the installed + CuPy in the GPU CI. `cupyx.empty_pinned`/`zeros_pinned` return NumPy arrays. - `xp.optimize.fsolve`, `root`, `minimize` and `xp.integrate.quad`, `odeint` (`cupyx.scipy` has neither): SciPy on the host with SciPy's arguments, the diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index ee2d948..ee013fe 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -61,7 +61,9 @@ def periodic_poisson(rho, length): `xp.scipy.special` on both backends. * On the fake CuPy (tests without a GPU), `xp.scipy` runs SciPy with the CuPy rules: host arrays are rejected, results are fake device arrays, and each - subpackage has only the names CuPy has. + subpackage has only the names of the oldest supported CuPy + (`xp.MIN_CUPY_VERSION`), so code that passes there runs on every supported + CuPy. ## Splines and interpolation diff --git a/docs/source/installation.md b/docs/source/installation.md index 6eb5b94..fdb0a7a 100644 --- a/docs/source/installation.md +++ b/docs/source/installation.md @@ -20,8 +20,15 @@ example: ```bash python -m pip install cupy-cuda12x # CUDA 12.x +python -m pip install cupy-cuda13x # CUDA 13.x ``` +CuNumpy supports CuPy 14.0 or newer (`xp.MIN_CUPY_VERSION`); the GPU CI tests +14.0 and the newest release. CuPy 13 lacks parts of `cupyx.scipy` that +CuNumpy relies on (`interpolate.UnivariateSpline`, `NdBSpline`, +`CubicSpline`, ...): selecting the CuPy backend with an older CuPy warns. +`xp.backend_info()` shows the installed and the minimum version. + On HPC clusters, load the site's CUDA module first and check the CuPy installation guide for the matching package name. Then verify that CuNumpy can use it: diff --git a/docs/source/reference/top-level.md b/docs/source/reference/top-level.md index 52b6248..2df44ea 100644 --- a/docs/source/reference/top-level.md +++ b/docs/source/reference/top-level.md @@ -12,4 +12,5 @@ The NumPy/CuPy namespace (see the [API overview](../api.md)) plus backend select ```{eval-rst} .. autodata:: cunumpy.scipy +.. autodata:: cunumpy.MIN_CUPY_VERSION ``` diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 1d5247b..cfa7e67 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -33,10 +33,12 @@ from cunumpy._scipy_backend import scipy from cunumpy._staging import to_host_async from cunumpy.xp import ( + MIN_CUPY_VERSION, as_device_array, assert_same_backend, backend_info, cupy_available, + cupy_version_supported, default_float_dtype, get_array_backend, get_array_module, @@ -96,6 +98,7 @@ def require_version(minimum: str) -> None: __all__ = [ + "MIN_CUPY_VERSION", "__version__", "algorithms", "arguments", @@ -105,6 +108,7 @@ def require_version(minimum: str) -> None: "cuda", "cupy_available", "cupy_backend", + "cupy_version_supported", "default_float_dtype", "evaluate_on_host", "get_array_backend", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 9e23036..e239928 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -28,6 +28,10 @@ def to_cupy(array: Any) -> Any: ... def to_cunumpy(array: Any) -> Any: ... def to_host_async(array: Any) -> HostCopy: ... def cupy_available() -> bool: ... +def cupy_version_supported(cupy_version: str) -> bool: ... + +MIN_CUPY_VERSION: str + def backend_info() -> dict[str, Any]: ... def get_array_module(array: Any) -> Any: ... def get_backend() -> str: ... diff --git a/src/cunumpy/_fake_cupyx.py b/src/cunumpy/_fake_cupyx.py index 900b0ef..0a85e11 100644 --- a/src/cunumpy/_fake_cupyx.py +++ b/src/cunumpy/_fake_cupyx.py @@ -9,10 +9,11 @@ in CuPy; * array results are fake device arrays, and objects (a fitted spline, a factorization) are proxies whose methods follow the same rules; -* each subpackage has only the names CuPy provides (:data:`NAMES`), so code - using a SciPy name that CuPy lacks (e.g. ``RectBivariateSpline``, - ``special.jv``, ``linalg.solve_circulant``) fails here as it would on a - GPU. A CuPy name that the installed SciPy lacks is missing too. +* each subpackage has only the names the oldest supported CuPy provides + (:data:`NAMES`, CuPy :data:`NAMES_VERSION`), so code using a SciPy name + that CuPy lacks (e.g. ``RectBivariateSpline``, ``special.jv``, + ``linalg.solve_circulant``) fails here as it would on a GPU. A CuPy name + that the installed SciPy lacks is missing too. ``cupyx.empty_pinned`` and ``cupyx.zeros_pinned`` return NumPy arrays, like CuPy's (which are backed by pinned memory). @@ -48,14 +49,22 @@ def _names(text: str) -> frozenset[str]: return frozenset(text.split()) -#: The public names of each ``cupyx.scipy`` subpackage (CuPy v14.2.0). +#: The CuPy release whose ``cupyx.scipy`` names :data:`NAMES` lists: the oldest +#: supported one (``cunumpy.xp.MIN_CUPY_VERSION``), so code that runs on the +#: fake runs on every supported CuPy. Newer CuPy releases have more names. +NAMES_VERSION = "14.0.0" + +#: The public names of each ``cupyx.scipy`` subpackage in CuPy +#: :data:`NAMES_VERSION`. Check them against an installed CuPy with +#: ``python testing/cupyx_names.py`` (the GPU CI does, in +#: ``tests/unit/test_cupyx_names.py``). NAMES = { "fft": _names( """ - dct dctn dst dstn fft fft2 fftfreq fftn fftshift fht get_fft_plan hfft - hfft2 hfftn idct idctn idst idstn ifft ifft2 ifftn ifftshift ifht ihfft - ihfft2 ihfftn irfft irfft2 irfftn next_fast_len rfft rfft2 rfftfreq - rfftn + dct dctn dst dstn fft fft2 fftfreq fftn fftshift fht + get_fft_plan hfft hfft2 hfftn idct idctn idst idstn ifft ifft2 + ifftn ifftshift ifht ihfft ihfft2 ihfftn irfft irfft2 irfftn + next_fast_len rfft rfft2 rfftfreq rfftn """, ), "fftpack": _names( @@ -67,75 +76,80 @@ def _names(text: str) -> frozenset[str]: """ Akima1DInterpolator BPoly BSpline BarycentricInterpolator CloughTocher2DInterpolator CubicHermiteSpline CubicSpline - InterpolatedUnivariateSpline KroghInterpolator LSQUnivariateSpline - LinearNDInterpolator NdBSpline NdPPoly NearestNDInterpolator PPoly - PchipInterpolator RBFInterpolator RegularGridInterpolator - UnivariateSpline barycentric_interpolate interp1d interpn - krogh_interpolate make_interp_spline make_lsq_spline pchip - pchip_interpolate splantider splder + InterpolatedUnivariateSpline KroghInterpolator + LSQUnivariateSpline LinearNDInterpolator NdBSpline NdPPoly + NearestNDInterpolator PPoly PchipInterpolator RBFInterpolator + RegularGridInterpolator UnivariateSpline barycentric_interpolate + interp1d interpn krogh_interpolate make_interp_spline + make_lsq_spline pchip pchip_interpolate splantider splder """, ), "linalg": _names( """ - bandwidth block_diag circulant companion convolution_matrix dft expm - fiedler fiedler_companion hadamard hankel helmert hilbert khatri_rao - kron leslie lu lu_factor lu_solve solve_triangular toeplitz + bandwidth block_diag circulant companion convolution_matrix dft + expm fiedler fiedler_companion hadamard hankel helmert hilbert + khatri_rao kron leslie lu lu_factor lu_solve solve_triangular + toeplitz """, ), "ndimage": _names( """ affine_transform binary_closing binary_dilation binary_erosion - binary_fill_holes binary_hit_or_miss binary_opening binary_propagation - black_tophat center_of_mass convolve convolve1d correlate correlate1d - distance_transform_edt extrema find_objects fourier_ellipsoid - fourier_gaussian fourier_shift fourier_uniform gaussian_filter - gaussian_filter1d gaussian_gradient_magnitude gaussian_laplace + binary_fill_holes binary_hit_or_miss binary_opening + binary_propagation black_tophat center_of_mass convolve + convolve1d correlate correlate1d distance_transform_edt extrema + find_objects fourier_ellipsoid fourier_gaussian fourier_shift + fourier_uniform gaussian_filter gaussian_filter1d + gaussian_gradient_magnitude gaussian_laplace generate_binary_structure generic_filter generic_filter1d - generic_gradient_magnitude generic_laplace grey_closing grey_dilation - grey_erosion grey_opening histogram iterate_structure label - labeled_comprehension laplace map_coordinates maximum maximum_filter - maximum_filter1d maximum_position mean median median_filter minimum - minimum_filter minimum_filter1d minimum_position morphological_gradient - morphological_laplace percentile_filter prewitt rank_filter rotate shift - sobel spline_filter spline_filter1d standard_deviation sum sum_labels - uniform_filter uniform_filter1d value_indices variance white_tophat zoom + generic_gradient_magnitude generic_laplace grey_closing + grey_dilation grey_erosion grey_opening histogram + iterate_structure label labeled_comprehension laplace + map_coordinates maximum maximum_filter maximum_filter1d + maximum_position mean median median_filter minimum + minimum_filter minimum_filter1d minimum_position + morphological_gradient morphological_laplace percentile_filter + prewitt rank_filter rotate shift sobel spline_filter + spline_filter1d standard_deviation sum sum_labels uniform_filter + uniform_filter1d value_indices variance white_tophat zoom """, ), "signal": _names( """ - BadCoefficients CZT StateSpace TransferFunction ZerosPolesGain ZoomFFT - abcd_normalize argrelextrema argrelmax argrelmin band_stop_obj bilinear - bilinear_zpk bode buttap butter buttord cheb1ap cheb1ord cheb2ap - cheb2ord cheby1 cheby2 check_COLA check_NOLA chirp choose_conv_method - coherence cont2discrete convolve convolve2d correlate correlate2d - correlation_lags csd cspline1d cspline1d_eval cspline2d cwt czt - czt_points dbode decimate deconvolve detrend dfreqresp dimpulse dlsim - dlti dstep ellip ellipap ellipord fftconvolve filtfilt find_peaks - findfreqs firls firwin firwin2 freqresp freqs freqs_zpk freqz freqz_sos - freqz_zpk gammatone gauss_spline gausspulse get_window group_delay - hilbert hilbert2 iircomb iirdesign iirfilter iirnotch iirpeak impulse - invres invresz istft kaiser_atten kaiser_beta kaiserord lfilter - lfilter_zi lfiltic lombscargle lp2bp lp2bp_zpk lp2bs lp2bs_zpk lp2hp - lp2hp_zpk lp2lp lp2lp_zpk lsim lti max_len_seq medfilt medfilt2d - minimum_phase morlet morlet2 normalize oaconvolve order_filter - peak_prominences peak_widths periodogram place_poles qmf qspline1d - qspline1d_eval qspline2d resample resample_poly residue residuez ricker - savgol_coeffs savgol_filter sawtooth sepfir2d sos2tf sos2zpk sosfilt - sosfilt_zi sosfiltfilt sosfreqz spectrogram spline_filter square ss2tf - ss2zpk step stft sweep_poly symiirorder1 symiirorder2 tf2sos tf2ss - tf2zpk unique_roots unit_impulse upfirdn vectorstrength welch wiener - zoom_fft zpk2sos zpk2ss zpk2tf + BadCoefficients CZT StateSpace TransferFunction ZerosPolesGain + ZoomFFT abcd_normalize argrelextrema argrelmax argrelmin + band_stop_obj bilinear bilinear_zpk bode buttap butter buttord + cheb1ap cheb1ord cheb2ap cheb2ord cheby1 cheby2 check_COLA + check_NOLA chirp choose_conv_method coherence cont2discrete + convolve convolve2d correlate correlate2d correlation_lags csd + cspline1d cspline1d_eval cspline2d cwt czt czt_points dbode + decimate deconvolve detrend dfreqresp dimpulse dlsim dlti dstep + ellip ellipap ellipord fftconvolve filtfilt find_peaks findfreqs + firls firwin firwin2 freqresp freqs freqs_zpk freqz freqz_sos + freqz_zpk gammatone gauss_spline gausspulse get_window + group_delay hilbert hilbert2 iircomb iirdesign iirfilter + iirnotch iirpeak impulse invres invresz istft kaiser_atten + kaiser_beta kaiserord lfilter lfilter_zi lfiltic lombscargle + lp2bp lp2bp_zpk lp2bs lp2bs_zpk lp2hp lp2hp_zpk lp2lp lp2lp_zpk + lsim lti max_len_seq medfilt medfilt2d minimum_phase morlet + morlet2 normalize oaconvolve order_filter peak_prominences + peak_widths periodogram place_poles qmf qspline1d qspline1d_eval + qspline2d resample resample_poly residue residuez ricker + savgol_coeffs savgol_filter sawtooth sepfir2d sos2tf sos2zpk + sosfilt sosfilt_zi sosfiltfilt sosfreqz spectrogram + spline_filter square ss2tf ss2zpk step stft sweep_poly + symiirorder1 symiirorder2 tf2sos tf2ss tf2zpk unique_roots + unit_impulse upfirdn vectorstrength welch wiener zoom_fft + zpk2sos zpk2ss zpk2tf """, ), "sparse": _names( """ - SparseEfficiencyWarning SparseWarning block_array block_diag bmat - coo_array coo_matrix csc_array csc_matrix csr_array csr_matrix dia_array - dia_matrix diags diags_array eye eye_array find get_index_dtype hstack - identity issparse isspmatrix isspmatrix_coo isspmatrix_csc - isspmatrix_csr isspmatrix_dia kron kronsum matrix_transpose permute_dims - rand random random_array safely_cast_index_arrays sparray spdiags - spmatrix swapaxes tril triu vstack + SparseEfficiencyWarning SparseWarning bmat coo_matrix csc_matrix + csr_matrix dia_matrix diags eye find hstack identity issparse + isspmatrix isspmatrix_coo isspmatrix_csc isspmatrix_csr + isspmatrix_dia kron kronsum rand random spdiags spmatrix tril + triu vstack """, ), "sparse.csgraph": _names( @@ -145,8 +159,8 @@ def _names(text: str) -> frozenset[str]: ), "sparse.linalg": _names( """ - LinearOperator SuperLU aslinearoperator bicgstab cg cgs eigsh factorized - gmres lobpcg lsmr lsqr matrix_power minres norm spilu splu spsolve + LinearOperator SuperLU aslinearoperator cg cgs eigsh factorized + gmres lobpcg lsmr lsqr minres norm spilu splu spsolve spsolve_triangular svds """, ), @@ -157,17 +171,19 @@ def _names(text: str) -> frozenset[str]: ), "special": _names( """ - bdtr bdtrc bdtri beta betainc betaincinv betaln binom boxcox boxcox1p - btdtr btdtri cbrt chdtr chdtrc chdtri cosdg cosm1 cotdg digamma - ellipeinc ellipj ellipk ellipkinc ellipkm1 entr erf erfc erfcinv erfcx - erfinv exp1 exp10 exp2 expi expit expm1 expn exprel fdtr fdtrc fdtri - gamma gammainc gammaincc gammainccinv gammaincinv gammaln gammasgn gdtr - gdtrc huber i0 i0e i1 i1e inv_boxcox inv_boxcox1p j0 j1 k0 k0e k1 k1e - kl_div lambertw log1p log_expit log_ndtr log_softmax loggamma logit - logsumexp lpmv multigammaln nbdtr nbdtrc nbdtri ndtr ndtri pdtr pdtrc - pdtri poch polygamma pseudo_huber psi radian rel_entr rgamma round - shichi sici sinc sindg softmax sph_harm sph_harm_y spherical_yn tandg - wright_bessel xlog1py xlogy y0 y1 yn zeta zetac + bdtr bdtrc bdtri beta betainc betaincinv betaln binom boxcox + boxcox1p btdtr btdtri cbrt chdtr chdtrc chdtri cosdg cosm1 cotdg + digamma ellipeinc ellipj ellipk ellipkinc ellipkm1 entr erf erfc + erfcinv erfcx erfinv exp1 exp10 exp2 expi expit expm1 expn + exprel fdtr fdtrc fdtri gamma gammainc gammaincc gammainccinv + gammaincinv gammaln gammasgn gdtr gdtrc huber i0 i0e i1 i1e + inv_boxcox inv_boxcox1p j0 j1 k0 k0e k1 k1e kl_div lambertw + log1p log_expit log_ndtr log_softmax loggamma logit logsumexp + lpmv multigammaln nbdtr nbdtrc nbdtri ndtr ndtri pdtr pdtrc + pdtri poch polygamma pseudo_huber psi radian rel_entr rgamma + round shichi sici sinc sindg softmax sph_harm sph_harm_y + spherical_yn tandg wright_bessel xlog1py xlogy y0 y1 yn zeta + zetac """, ), "stats": _names( @@ -178,6 +194,50 @@ def _names(text: str) -> frozenset[str]: } +def public_names(module: types.ModuleType) -> frozenset[str]: + """The public names of a (real) ``cupyx.scipy`` subpackage, as in :data:`NAMES`. + + Names without a leading underscore, except submodules and the + ``annotations`` of ``from __future__ import annotations``. + """ + return frozenset( + name + for name in dir(module) + if not name.startswith("_") + and name != "annotations" + and not isinstance(getattr(module, name, None), types.ModuleType) + ) + + +def installed_names() -> dict[str, frozenset[str]]: + """:func:`public_names` of each subpackage of the installed (real) ``cupyx.scipy``.""" + return { + path: public_names(importlib.import_module(f"cupyx.scipy.{path}")) + for path in SUBPACKAGES + } + + +def compare_names( + real: dict[str, frozenset[str]], +) -> tuple[dict[str, list[str]], dict[str, list[str]]]: + """Compare :data:`NAMES` with the names of a real ``cupyx.scipy``. + + Returns ``(missing, added)``: per subpackage (only those with a + difference), the names of :data:`NAMES` that `real` lacks, and the names + `real` has beyond :data:`NAMES`. On CuPy :data:`NAMES_VERSION` both are + empty; a newer CuPy may add names, but a name it lacks means the fake + accepts code that fails on that CuPy. + """ + missing, added = {}, {} + for path in SUBPACKAGES: + have = real.get(path, frozenset()) + if NAMES[path] - have: + missing[path] = sorted(NAMES[path] - have) + if have - NAMES[path]: + added[path] = sorted(have - NAMES[path]) + return missing, added + + class _Proxy: """A SciPy object (spline, interpolator, ...) seen as a CuPy object.""" diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index bb22e9a..4ae131b 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -10,6 +10,7 @@ import logging import os +import re import sys import warnings from collections.abc import Callable, Generator @@ -46,6 +47,44 @@ _CUPY_AVAILABLE_CACHE = None _CUPY_UNAVAILABLE_REASON: str | None = None +#: The oldest supported CuPy release. Older ones lack parts of ``cupyx.scipy`` +#: that cunumpy documents (``interpolate.UnivariateSpline``, ``NdBSpline``, +#: ``CubicSpline``, ...); the fake ``cupyx.scipy`` lists the names of this +#: release. The GPU CI tests this release and the newest one. +MIN_CUPY_VERSION = "14.0.0" + + +def cupy_version_supported(cupy_version: str) -> bool: + """Return whether a CuPy version is :data:`MIN_CUPY_VERSION` or newer. + + Pre-releases of the minimum itself (``14.0.0rc1``) are older than it. + + Parameters + ---------- + cupy_version : str + A version string, e.g. ``cupy.__version__``. + + Returns + ------- + bool + False also for a string that is not a version. + + Examples + -------- + >>> xp.cupy_version_supported("13.6.0") + False + >>> xp.cupy_version_supported("14.2.0") + True + """ + match = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?(.*)", cupy_version) + if match is None: + return False + release = tuple(int(part or 0) for part in match.groups()[:3]) + minimum = tuple(int(part) for part in MIN_CUPY_VERSION.split(".")) + if release != minimum: + return release > minimum + return re.match(r"\.?(a|b|rc|dev)", match.group(4)) is None + def cupy_available() -> bool: """Return whether CuPy can be imported and reports a usable CUDA device. @@ -71,6 +110,16 @@ def cupy_available() -> bool: try: import cupy as cp + cupy_version = getattr(cp, "__version__", "") + fake = getattr(cp, "__cunumpy_fake__", False) + if cupy_version and not fake and not cupy_version_supported(cupy_version): + warnings.warn( + f"CuPy {cupy_version} is older than the oldest version cunumpy " + f"supports ({MIN_CUPY_VERSION}); parts of xp.scipy (e.g. " + "interpolate.UnivariateSpline) are missing on the GPU", + RuntimeWarning, + stacklevel=2, + ) # Check if a GPU is available _CUPY_AVAILABLE_CACHE = cp.is_available() _CUPY_UNAVAILABLE_REASON = ( @@ -321,7 +370,8 @@ def backend_info() -> dict[str, Any]: ------- dict Keys ``backend``, ``cupy_available``, ``cuda_unavailable_reason``, - ``versions``, ``device`` (None without CUDA) and + ``versions``, ``min_cupy_version`` (:data:`MIN_CUPY_VERSION`), + ``device`` (None without CUDA) and ``cuda_visible_devices``, plus ``cuda_inspection_error`` on failure. Examples @@ -342,6 +392,7 @@ def backend_info() -> dict[str, Any]: "cupy_available": available, "cuda_unavailable_reason": None if available else _CUPY_UNAVAILABLE_REASON, "versions": versions, + "min_cupy_version": MIN_CUPY_VERSION, "device": None, "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), } diff --git a/testing/cupyx_names.py b/testing/cupyx_names.py new file mode 100644 index 0000000..632e797 --- /dev/null +++ b/testing/cupyx_names.py @@ -0,0 +1,75 @@ +"""Compare the fake CuPy's ``cupyx.scipy`` names with the installed CuPy. + +The fake ``cupyx.scipy`` (``cunumpy._fake_cupyx.NAMES``) lists the names of +the oldest supported CuPy (``NAMES_VERSION``). Run this where the real CuPy +is installed:: + + python testing/cupyx_names.py # report; exit 1 if the fake lists a name CuPy lacks + python testing/cupyx_names.py --exact # also exit 1 if CuPy has names the fake lacks + python testing/cupyx_names.py --print # print NAMES for the installed CuPy + +``--print`` gives the new ``NAMES`` when raising the minimum CuPy version: run +it with that version installed and paste the output into ``_fake_cupyx.py``. +The GPU CI runs the same check in ``tests/unit/test_cupyx_names.py``. +""" + +import argparse +import sys +import textwrap + + +def main() -> int: + """Run the comparison; return the exit code.""" + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument( + "--exact", action="store_true", help="fail on names the fake lacks too" + ) + parser.add_argument( + "--print", action="store_true", help="print NAMES for the installed CuPy" + ) + args = parser.parse_args() + + import cupy + + if getattr(cupy, "__cunumpy_fake__", False): + sys.exit("the fake CuPy is active: unset CUNUMPY_FAKE_CUPY") + + from cunumpy._fake_cupyx import ( + NAMES_VERSION, + SUBPACKAGES, + compare_names, + installed_names, + ) + + real = installed_names() + if args.print: + print( + f"#: The public names of each ``cupyx.scipy`` subpackage (CuPy {cupy.__version__})." + ) + print("NAMES = {") + for path in SUBPACKAGES: + body = textwrap.fill( + " ".join(sorted(real[path])), + width=72, + initial_indent=" " * 8, + subsequent_indent=" " * 8, + break_on_hyphens=False, + break_long_words=False, + ) + print(f' "{path}": _names(\n """\n{body}\n """,\n ),') + print("}") + return 0 + + missing, added = compare_names(real) + print(f"installed CuPy {cupy.__version__}; the fake lists CuPy {NAMES_VERSION}") + for path, names in missing.items(): + print(f"cupyx.scipy.{path} lacks names the fake has: {' '.join(names)}") + for path, names in added.items(): + print(f"cupyx.scipy.{path} has names the fake lacks: {' '.join(names)}") + if not missing and not added: + print("the names agree") + return 1 if missing or (args.exact and added) else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unit/test_cupyx_names.py b/tests/unit/test_cupyx_names.py new file mode 100644 index 0000000..f529a40 --- /dev/null +++ b/tests/unit/test_cupyx_names.py @@ -0,0 +1,84 @@ +"""The fake cupyx.scipy names (_fake_cupyx.NAMES) against the oldest supported and the installed CuPy.""" + +import types + +import pytest + +import cunumpy as xp +from cunumpy import _fake_cupyx +from cunumpy._fake_cupyx import NAMES, NAMES_VERSION, compare_names, public_names +from cunumpy.kernel_testing import requires_cupy + + +def test_names_are_those_of_the_oldest_supported_cupy(): + assert NAMES_VERSION == xp.MIN_CUPY_VERSION + assert tuple(NAMES) == _fake_cupyx.SUBPACKAGES + # the reason for the minimum: CuPy 13 has no smoothing splines + assert {"UnivariateSpline", "NdBSpline"} <= NAMES["interpolate"] + + +def test_compare_names(): + real = {path: set(names) for path, names in NAMES.items()} + assert compare_names(real) == ({}, {}) + real["special"] = (real["special"] - {"erf"}) | {"jv"} + missing, added = compare_names(real) + assert missing == {"special": ["erf"]} + assert added == {"special": ["jv"]} + + +def test_public_names_skip_private_names_and_submodules(): + module = types.ModuleType("cupyx.scipy.example") + module.solve = lambda: None + module._helper = lambda: None + module.annotations = object() # from __future__ import annotations + module.linalg = types.ModuleType("cupyx.scipy.example.linalg") + assert public_names(module) == {"solve"} + + +@requires_cupy +def test_installed_cupy_has_every_name_of_the_fake(): + """On a GPU: the fake has no name the installed CuPy lacks; on the oldest supported CuPy, the same names.""" + import cupy + + missing, added = compare_names(_fake_cupyx.installed_names()) + assert not missing, f"CuPy {cupy.__version__} lacks names of the fake: {missing}" + if cupy.__version__ == NAMES_VERSION: + assert not added, f"CuPy {NAMES_VERSION} has names the fake lacks: {added}" + + +@requires_cupy +def test_installed_cupy_is_supported(): + import cupy + + assert xp.cupy_version_supported(cupy.__version__) + + +@pytest.mark.parametrize( + ("version", "supported"), + [ + ("13.6.0", False), + ("14.0.0", True), + ("14.0.0rc1", False), + ("14.2.0", True), + ("15.0.0a1", True), + ], +) +def test_cupy_version_supported(version, supported): + assert xp.cupy_version_supported(version) is supported + + +def test_old_cupy_warns_once(monkeypatch): + import sys + + from cunumpy import xp as backend + + old = types.ModuleType("cupy") + old.__version__ = "13.6.0" + old.is_available = lambda: False + monkeypatch.setitem(sys.modules, "cupy", old) + monkeypatch.setattr(backend, "_CUPY_AVAILABLE_CACHE", None) + with pytest.warns( + RuntimeWarning, match=r"CuPy 13\.6\.0 is older than .* \(14\.0\.0\)" + ): + assert not backend.cupy_available() + assert not backend.cupy_available() # cached: no second warning From 384ea8191f4bb15c64c52b2bf1c411496d15748b Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 23:32:42 +0200 Subject: [PATCH 6/9] Hide the real cupyx in the fake_cupyx test fixture --- tests/unit/test_scipy_backend.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/test_scipy_backend.py b/tests/unit/test_scipy_backend.py index ed928b7..dcc0f82 100644 --- a/tests/unit/test_scipy_backend.py +++ b/tests/unit/test_scipy_backend.py @@ -15,6 +15,10 @@ @pytest.fixture def fake_cupyx(monkeypatch): """A fake `cupyx.scipy` with `special.erf` and `sparse.linalg.cg`, as backend.""" + # hide the real cupyx (imported by earlier tests on a GPU): its other + # subpackages, e.g. cupyx.scipy.ndimage, would stay importable + for name in [n for n in sys.modules if n.split(".")[0] == "cupyx"]: + monkeypatch.delitem(sys.modules, name) modules = {} for name in ("cupyx", "cupyx.scipy", "cupyx.scipy.special", "cupyx.scipy.sparse"): modules[name] = types.ModuleType(name) From 747066867415eba9d716c6de3700d6728172aa45 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 17:08:37 +0200 Subject: [PATCH 7/9] more xp scipy --- src/cunumpy/_dispatch.py | 8 ++++++ src/cunumpy/_kernel.py | 58 +++++++++++++++++++++++++++++++--------- src/cunumpy/xp.py | 11 ++++++++ 3 files changed, 64 insertions(+), 13 deletions(-) diff --git a/src/cunumpy/_dispatch.py b/src/cunumpy/_dispatch.py index d7bb198..cebb9ca 100644 --- a/src/cunumpy/_dispatch.py +++ b/src/cunumpy/_dispatch.py @@ -29,6 +29,7 @@ from typing import Any import array_api_compat +import numpy as np from cunumpy._cuda_kernel import CudaKernel, _compile_in_threads from cunumpy._kernel import ( @@ -51,8 +52,15 @@ _is_device_array = array_api_compat.is_cupy_array +#: Argument types that are never on the device: checked first, since a kernel +#: call tests every argument and most are host arrays or scalars. +_HOST_TYPES = frozenset({np.ndarray, float, int, bool, np.float64, np.int64, np.bool_}) + + def _on_device(arg: Any) -> bool: """Whether `arg` is a CuPy array or a CUDA argument object (``__cuda_args__``).""" + if type(arg) in _HOST_TYPES: + return False return _is_device_array(arg) or callable(getattr(type(arg), "__cuda_args__", None)) diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index 16de9c4..c4a4022 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -19,9 +19,9 @@ import inspect import os import warnings -from collections.abc import Callable, Generator, Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager -from types import ModuleType +from types import ModuleType, TracebackType from typing import Any import array_api_compat @@ -1000,6 +1000,17 @@ def as_kernel_array( >>> xp.kernels.as_kernel_array([1, 2], like=grid, dtype=float) array([1., 2.]) """ + # fast path: a C-contiguous NumPy array of the dtype for a host kernel, the + # common case, without the conversions below (subclasses take the full path) + if ( + type(value) is np.ndarray + and (type(like) is np.ndarray or not is_gpu(like)) + and (dtype is None or value.dtype == dtype) + and value.flags.c_contiguous + and value.ndim > 0 + and (value.ndim == 1 or 1 not in value.shape) + ): + return value if is_gpu(like): import cupy @@ -1041,10 +1052,39 @@ def _same_memory(buffer: Any, out: Any) -> bool: ) -@contextmanager +class _KernelOutput: + """The context manager of :func:`kernel_output` (a class: cheaper than a generator).""" + + __slots__ = ("_buffer", "_out") + + def __init__(self, out: Any, buffer: Any) -> None: + self._out = out + self._buffer = buffer + + def __enter__(self) -> Any: + return self._buffer + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + if exc_type is not None: + return + buffer, out = self._buffer, self._out + if not _same_memory(buffer, out): + if is_gpu(out) and not is_gpu(buffer): + out[...] = to_cupy(buffer) + elif not is_gpu(out) and is_gpu(buffer): + out[...] = to_numpy(buffer) + else: + out[...] = buffer + + def kernel_output( out: Any, like: Any, dtype: Any = None, *, strided: bool = False -) -> Generator[Any]: +) -> _KernelOutput: """Provide the output buffer for the kernel chosen for `like`, copied into `out`. The buffer is `out` itself if :func:`~cunumpy.kernels.as_kernel_array` @@ -1075,15 +1115,7 @@ def kernel_output( >>> out array([1.5, 1.5], dtype=float32) """ - buffer = as_kernel_array(out, like, dtype, strided=strided) - yield buffer - if not _same_memory(buffer, out): - if is_gpu(out) and not is_gpu(buffer): - out[...] = to_cupy(buffer) - elif not is_gpu(out) and is_gpu(buffer): - out[...] = to_numpy(buffer) - else: - out[...] = buffer + return _KernelOutput(out, as_kernel_array(out, like, dtype, strided=strided)) _SCALAR_ANNOTATIONS = {"int", "float", "bool", "complex", "str"} diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 4ae131b..5a16a20 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -21,6 +21,11 @@ import array_api_compat import array_api_compat.numpy as np +import numpy + +#: The exact type of a host array; a fast path checks ``type(a) is _NDARRAY`` before +#: the slower ``array_api_compat`` tests (subclasses take the full path). +_NDARRAY = numpy.ndarray from cunumpy._transfers import _ACTIVE as _COUNTERS from cunumpy._transfers import ( @@ -493,6 +498,8 @@ def synchronize() -> None: def _to_numpy(array: Any) -> np.ndarray: """`to_numpy` without transfer counting, for internal use.""" + if type(array) is _NDARRAY: + return array if get_array_backend(array) == "cupy": return array.get(order="A") @@ -533,6 +540,8 @@ def to_numpy(array: Any) -> np.ndarray: >>> xp.to_numpy([1, 2]) array([1, 2]) """ + if type(array) is _NDARRAY: + return array result = _to_numpy(array) if _COUNTERS and get_array_backend(array) == "cupy": _record("to_host", f"to_numpy({_describe(array)})", nbytes=_nbytes(result)) @@ -719,6 +728,8 @@ def get_array_backend(array: Any) -> BackendType: >>> xp.get_array_backend(np.zeros(3)) 'numpy' """ + if type(array) is _NDARRAY: + return "numpy" return "cupy" if array_api_compat.is_cupy_array(array) else "numpy" From e13608004578a8575b2d88dff10a4202da01d1d8 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 17:25:35 +0200 Subject: [PATCH 8/9] formatting --- docs/source/conf.py | 5 ++++- docs/source/guides/solvers.md | 7 +++++-- docs/source/kernels/cuda-headers.md | 8 ++++++-- 3 files changed, 15 insertions(+), 5 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 5f966c0..96f089f 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -42,7 +42,10 @@ def _defining_comment(name, obj): from sphinx.pycode import ModuleAnalyzer for modname, module in list(sys.modules.items()): - if not modname.startswith("cunumpy._") or getattr(module, name, None) is not obj: + if ( + not modname.startswith("cunumpy._") + or getattr(module, name, None) is not obj + ): continue try: docs = ModuleAnalyzer.for_module(modname).find_attr_docs() diff --git a/docs/source/guides/solvers.md b/docs/source/guides/solvers.md index ee013fe..c13277e 100644 --- a/docs/source/guides/solvers.md +++ b/docs/source/guides/solvers.md @@ -77,8 +77,8 @@ least-squares variants), `BSpline`, `make_interp_spline`, `make_lsq_spline`, ```python interp = xp.scipy.interpolate -bx = interp.make_interp_spline(x, values, k=3) # fit along x (axis 0) -by = interp.make_interp_spline(y, bx.c.T, k=3) # then along y +bx = interp.make_interp_spline(x, values, k=3) # fit along x (axis 0) +by = interp.make_interp_spline(y, bx.c.T, k=3) # then along y spline = interp.NdBSpline((bx.t, by.t), by.c.T, (3, 3)) # RectBivariateSpline clamps points to the grid; NdBSpline extrapolates points = xp.stack([xp.clip(R, x[0], x[-1]), xp.clip(Z, y[0], y[-1])], axis=-1) @@ -106,6 +106,7 @@ the secant method on all of them at once, with the steps of def residual(r): return psi(R0 + r * cos_theta, Z0 + r * sin_theta) - levels + r = xp.optimize.newton(residual, xp.full(levels.shape, 0.3)) ``` @@ -124,9 +125,11 @@ host copy of that data: ```python B = equilibrium.b_field_on_grid() # device array + def residual(x): return interpolate(B, x) - b_target + x = xp.optimize.fsolve(residual, xp.asarray([0.5, 0.0])) # device array ``` diff --git a/docs/source/kernels/cuda-headers.md b/docs/source/kernels/cuda-headers.md index 6c2115f..1d2eea3 100644 --- a/docs/source/kernels/cuda-headers.md +++ b/docs/source/kernels/cuda-headers.md @@ -185,12 +185,16 @@ same box form a contiguous range, the starting point of tree builds on the GPU. On the host: ```python -keys = xp.algorithms.morton_keys(positions, lower, upper, levels) # (n, 2|3) -> (n,) uint64 +keys = xp.algorithms.morton_keys( + positions, lower, upper, levels +) # (n, 2|3) -> (n,) uint64 keys, order, positions = xp.algorithms.sort_by_key(keys, positions) node = keys >> np.uint64(ndim * (levels - level)) # node index at `level` cells = xp.algorithms.morton_decode(node, ndim) # its integer coordinates key = xp.algorithms.morton_encode(ix, iy) # from integer cells -scales = xp.algorithms.morton_scales(lower, upper, levels) # 2**levels / (upper - lower) +scales = xp.algorithms.morton_scales( + lower, upper, levels +) # 2**levels / (upper - lower) ``` In a kernel: From 6205b67d61eccca8f2e76e6cc99789242383b1db Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 17:29:08 +0200 Subject: [PATCH 9/9] numpydoc fix --- src/cunumpy/_kernel.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index c4a4022..44f7339 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -1102,10 +1102,10 @@ def kernel_output( strided : bool, optional As for :func:`~cunumpy.kernels.as_kernel_array`. Default False. - Yields - ------ - array - The buffer to pass to the kernel. + Returns + ------- + context manager + Yields the buffer to pass to the kernel. Examples --------