Repository navigation
Avoid a full-size copy of broadcast input in NaN reductions - #3108
Open
antonwolfy wants to merge 2 commits into
Open
antonwolfy wants to merge 2 commits into
antonwolfy wants to merge 2 commits into
Conversation
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.
antonwolfy
requested review from
ndgrigorian and
vlad-perevezentsev
as code owners
October 15, 2026 13:34
Contributor
|
View rendered docs @ https://intelpython.github.io/dpnp/pull/3108/index.html |
Contributor
|
Array API standard conformance tests for dpnp=0.21.0dev14=np2py314h8d9cdd5_6 ran successfully. |
Collaborator
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
dpnp.nansum,dpnp.nanprod,dpnp.nancumsumanddpnp.nancumprodreplaceNaNvalues withdpnp.nan_to_numbefore 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 slowTestNansumNanprodHugetests failed withUSMAllocationErroron devices whose maximum allocation size is smaller than that.This PR replaces
NaNvalues 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 throughdpnp.nan_to_numexactly 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-gttcounter of thexedriver, each case in a fresh process:nansum,nanprod,nancumsum)nansum/nanprodof(2048, 1→256, 1024), any axisnansumof(4096, 1→4096), either axisnancumsumof(512, 1→256, 1024),axis=2Time without broadcast axes, median of 101 interleaved runs with the order alternated, ratio this PR / master:
nansumof 1k / 64k elementsnansum/nanprod/nancumsumof 16M elementsnansumof a strided 16M viewThe 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:
nansumof(2048, 1→256, 1024),axis=2nansumof(2048, 1→256, 1024),axis=1nansumof(2048, 1→256, 1024), all axesnanprodof(2048, 1→256, 1024),axis=2nansumof(4096, 1→4096),axis=1/axis=0nancumsumof(512, 1→256, 1024),axis=2New
test_broadcasttests inTestNanProdSumandTestNanCumSumProdcheck the results on inputs broadcast along different axes.dpnp/tests/test_nanfunctions.py, the third-party sum and prod tests and the slowTestNansumNanprodHugetests pass on GPU and CPU; the latter failed before on a GPU with a 1 GiB maximum allocation size.