IntelPython / dpnp / 39480275836

83%

push

github

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 devi... (continued)

1604 of 2994 branches covered (53.57%)

Branch coverage included in aggregate %.

7 of 7 new or added lines in 1 file covered. (100.0%)

29399 of 34533 relevant lines covered (85.13%)

19629.49 hits per line