Skip to content

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

Open
antonwolfy wants to merge 2 commits into
masterfrom
nansum-no-full-copy
Open

antonwolfy wants to merge 2 commits into
masterfrom
nansum-no-full-copy

Conversation

@antonwolfy

Copy link
Copy Markdown
Contributor

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.

  • Have you provided a meaningful PR description?
  • Have you added a test, reproducer or referred to an issue with a reproducer?
  • Have you tested your changes locally for CPU and GPU devices?
  • Have you made sure that new changes do not introduce compiler warnings?
  • Have you checked performance impact of proposed changes?
  • Have you added documentation for your changes, if necessary?
  • Have you added your changes to the changelog?

nansum, nanprod, nancumsum and nancumprod replaced NaN values with
nan_to_num, which copied a broadcast input to its full shape. Replace
NaN values only in the distinct elements of such an input and broadcast
the result back as a view. Inputs without broadcast axes are handled as
before.
@github-actions

Copy link
Copy Markdown
Contributor

View rendered docs @ https://intelpython.github.io/dpnp/pull/3108/index.html

@github-actions

Copy link
Copy Markdown
Contributor

Array API standard conformance tests for dpnp=0.21.0dev14=np2py314h8d9cdd5_6 ran successfully.
Passed: 1376
Failed: 0
Skipped: 6

@coveralls

Copy link
Copy Markdown
Collaborator

Coverage Status

No base build to compare — nansum-no-full-copy into master

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants