Avoid a full-size copy of broadcast input in NaN reductions (#3108)

`dpnp.nansum`, `dpnp.nanprod`, `dpnp.nancumsum` and `dpnp.nancumprod`
replace `NaN` values with `dpnp.nan_to_num` before reducing, which
copies the whole input. For a broadcast view this materializes every
repeated element: a `(2048, 1, 1024)` float32 array broadcast to `(2048,
256, 1024)` becomes a 2 GiB copy, so the slow `TestNansumNanprodHuge`
tests failed with `USMAllocationError` on devices whose maximum
allocation size is smaller than that.

This PR replaces `NaN` values only in the distinct elements of a
broadcast input, by slicing each zero-stride axis to length 1, and
broadcasts the result back to the input shape as a view; the result is
only read by the reduction that follows. Inputs without broadcast axes
go through `dpnp.nan_to_num` exactly as before: the only extra work is a
check of the shape and strides, which takes 0.5-0.8 µs.

Measured against master on Intel Arc Graphics (Lunar Lake, Level Zero)
and on its CPU (OpenCL).

**Peak GPU memory**, from the per-process `drm-total-gtt` counter of the
`xe` driver, each case in a fresh process:

| Case (float32) | master | this PR |
|---|---|---|
| 11 cases without broadcast axes (contiguous, strided, transposed, with
NaNs, complex64, int32; `nansum`, `nanprod`, `nancumsum`) | same | same
(within 0.1 MiB) |
| `nansum`/`nanprod` of `(2048, 1→256, 1024)`, any axis | 2055-2063 MiB
| 14-22 MiB |
| `nansum` of `(4096, 1→4096)`, either axis | 71 MiB | 6 MiB |
| `nancumsum` of `(512, 1→256, 1024)`, `axis=2` | 1031 MiB | 520 MiB
(the output itself is 512 MiB) |

**Time without broadcast axes**, median of 101 interleaved runs with the
order alternated, ratio this PR / master:

| Case (float32) | GPU | CPU |
|---|---|---|
| `nansum` of 1k / 64k elements | 1.03 / 1.05 | 1.06 / 1.02 |
| `nansum` / `nanprod` / `nancumsum` of 16M elements | 0.98 / 0.77 /
1.01 | 0.98 / 1.01 / 1.00 |
| `nansum` of a strided 16M view | 1.00 | 1.00 |

The deviations from 1.00 are within the run-to-run noise; the added
check costs under 1 µs.

**Time with broadcast axes**, median of 31 interleaved runs:

| Case (float32) | GPU master → this PR | CPU master → this PR |
|---|---|---|
| `nansum` of `(2048, 1→256, 1024)`, `axis=2` | 296 → 33 ms | 1091 → 91
ms |
| `nansum` of `(2048, 1→256, 1024)`, `axis=1` | 377 → 78 ms | 939 → 118
ms |
| `nansum` of `(2048, 1→256, 1024)`, all axes | 358 → 101 ms | 921 → 409
ms |
| `nanprod` of `(2048, 1→256, 1024)`, `axis=2` | 429 → 60 ms | 1097 → 92
ms |
| `nansum` of `(4096, 1→4096)`, `axis=1` / `axis=0` | 12.5 / 10.7 → 1.2
/ 1.2 ms | 29.7 / 54.7 → 3.6 / 3.9 ms |
| `nancumsum` of `(512, 1→256, 1024)`, `axis=2` | 136 → 74 ms | 533 →
303 ms |

New `test_broadcast` tests in `TestNanProdSum` and `TestNanCumSumProd`
check the results on inputs broadcast along different axes.
`dpnp/tests/test_nanfunctions.py`, the third-party sum and prod tests
and the slow `TestNansumNanprodHuge` tests pass on GPU and CPU; the
latter failed before on a GPU with a 1 GiB maximum allocation size.

Author: antonwolfy

CHANGELOG.md

@@ -59,6 +59,7 @@ This release is compatible with NumPy 2.5.
 * Changed `dpnp.sort`, `dpnp.argsort`, and their `dpnp.ndarray`/`dpnp.tensor` counterparts to place `NaN` values last instead of first when sorting in descending order [#3066](https://github.com/IntelPython/dpnp/pull/3066)
 * Documented and made explicit support for `copy=None` ("copy only if needed") in `dpnp.nan_to_num` [#3083](https://github.com/IntelPython/dpnp/pull/3083)
 * Changed `dpnp.linalg.pinv` to compute the default tolerance from the machine epsilon of the data type used for the computation, and to return a floating-point result for an empty integer or boolean array [#3090](https://github.com/IntelPython/dpnp/pull/3090)
+* Avoided a full-size copy of a broadcast input in `dpnp.nansum`, `dpnp.nanprod`, `dpnp.nancumsum` and `dpnp.nancumprod` by replacing `NaN` values only in the distinct elements of the input [#3108](https://github.com/IntelPython/dpnp/pull/3108)
 
 ### Deprecated
 

dpnp/dpnp_iface_nanfunctions.py

@@ -53,6 +53,10 @@ def _replace_nan_no_mask(a, val):
     the `val` value, and return the copy. If `a` is not of inexact type,
     do nothing and return `a`.
 
+    If `a` is a broadcast view, only its distinct elements are copied and
+    the copy is broadcast back to the shape of `a`, so the result must only
+    be read.
+
     Parameters
     ----------
     a : {dpnp.ndarray, usm_ndarray}
@@ -69,10 +73,17 @@ def _replace_nan_no_mask(a, val):
     """
 
     dpnp.check_supported_arrays_type(a)
-    if dpnp.issubdtype(a.dtype, dpnp.inexact):
+    if not dpnp.issubdtype(a.dtype, dpnp.inexact):
+        return a
+
+    repeated = [st == 0 and sh > 1 for sh, st in zip(a.shape, a.strides)]
+    if not any(repeated):
         return dpnp.nan_to_num(a, nan=val, posinf=dpnp.inf, neginf=-dpnp.inf)
 
-    return a
+    # a zero-stride axis repeats the same elements, keep one of them
+    base = a[tuple(slice(0, 1) if r else slice(None) for r in repeated)]
+    res = dpnp.nan_to_num(base, nan=val, posinf=dpnp.inf, neginf=-dpnp.inf)
+    return dpnp.broadcast_to(res, a.shape)
 
 
 def _replace_nan(a, val):

dpnp/tests/test_nanfunctions.py

@@ -114,6 +114,23 @@ def test_allnans(self, func, dtype, array):
         # Array API indicates that the behavior is unspecified
         assert_dtype_allclose(result, expected, check_shape=False)
 
+    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
+    @pytest.mark.parametrize("axis", [None, 0, 1, 2])
+    @pytest.mark.parametrize(
+        "shape", [(1, 3, 4), (2, 1, 4), (2, 3, 1), (1, 3, 1)]
+    )
+    def test_broadcast(self, func, shape, axis, dtype):
+        # positive values avoid cancellation in the cumulative sums
+        a = generate_random_numpy_array(shape, dtype=dtype, low=1)
+        a.flat[::3] = numpy.nan
+        ia = dpnp.array(a)
+
+        a = numpy.broadcast_to(a, (2, 3, 4))
+        ia = dpnp.broadcast_to(ia, (2, 3, 4))
+        result = getattr(dpnp, func)(ia, axis=axis)
+        expected = getattr(numpy, func)(a, axis=axis)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("axis", [None, 0, 1])
     def test_empty(self, func, axis):
         a = numpy.zeros((0, 3))
@@ -540,6 +557,23 @@ def test_basic(self, func, axis, keepdims, dtype):
 
         assert_allclose(result, expected, rtol=1e-5)
 
+    @pytest.mark.parametrize("dtype", get_float_complex_dtypes())
+    @pytest.mark.parametrize("axis", [None, 0, 1, 2, (0, 2)])
+    @pytest.mark.parametrize(
+        "shape", [(1, 3, 4), (2, 1, 4), (2, 3, 1), (1, 3, 1)]
+    )
+    def test_broadcast(self, func, shape, axis, dtype):
+        # positive values avoid cancellation in the sums
+        a = generate_random_numpy_array(shape, dtype=dtype, low=1)
+        a.flat[::3] = numpy.nan
+        ia = dpnp.array(a)
+
+        a = numpy.broadcast_to(a, (2, 3, 4))
+        ia = dpnp.broadcast_to(ia, (2, 3, 4))
+        result = getattr(dpnp, func)(ia, axis=axis)
+        expected = getattr(numpy, func)(a, axis=axis)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.usefixtures("suppress_complex_warning")
     @pytest.mark.usefixtures("suppress_invalid_numpy_warnings")
     @pytest.mark.parametrize("in_dt", get_float_complex_dtypes())