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