Segment_sum on the fake CuPy sums on the host buffers

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

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)

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