Drop the per-cell NaN early-out in the planar CPU slope kernel - #3742
Open
brendancol wants to merge 2 commits into
Open
Drop the per-cell NaN early-out in the planar CPU slope kernel#3742brendancol wants to merge 2 commits into
brendancol wants to merge 2 commits into
Conversation
_cpu started every interior cell with `if np.isnan(data[y, x]): continue`. On DEMs with scattered nodata that branch is data-dependent and mispredicts more than the few float32 ops and one arctan it skips. Remove the branch and fold the centre back in with a select after the arithmetic. Neighbour NaN already propagates through the Horn stencil, so the select only has to cover a NaN centre with valid neighbours. Results are bitwise identical to the old kernel on 2000x4000 float32 DEMs with no NaN, 30% random NaN and a NaN left half, and on 3x3, 1xN and int16 inputs. On the 30% random NaN raster the kernel goes from ~92 ms to ~55 ms; the NaN-free and contiguous-NaN cases are within noise. Add a SlopeNaN asv benchmark with 30% random NaN over the interior so the nodata path is timed from now on. get_xr_dataarray(include_nan=True) only sets the [0, 0] corner to NaN, which the kernel never visits, so the benchmark adds its own speckle. Add tests pinning the NaN footprint on a speckled raster: NaN at every centre-NaN cell and its 8-neighbours, finite elsewhere, on numpy and dask+numpy.
brendancol
commented
Sep 4, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
PR Review: Drop the per-cell NaN early-out in the planar CPU slope kernel
Blockers (must fix before merge)
None.
Suggestions (should fix, not blocking)
-
benchmarks/benchmarks/slope.py:35-38:time_slope_nanforces.compute()for dask, but the existingSlope.time_slopegoes throughBenchmarking.time, which does not. So fortype="dask"the two classes time different things (graph build vs actual kernel work) and a side-by-side read ofSlopevsSlopeNaNis misleading. The compute is the right call for the new class. Either add a one-line comment saying the dask numbers are not comparable withSlope, or leaveSlopeas is and accept it; changingSlopeis out of scope here. -
xrspatial/tests/test_slope.py:642-648:_expected_nan_footprintbuilds the expected mask by hand. Worth a sentence in the comment block that these tests also pass on the pre-PR kernel; they pin behaviour rather than reproduce a regression, so nobody reverts the kernel expecting them to go red. The bitwise A/B in the PR body is what proves equivalence.
Nits (optional improvements)
-
benchmarks/benchmarks/slope.py:31:rng.random((ny, nx))atnx=10000allocates a 50M-element float64 array (400 MB) just to derive a bool mask.rng.random((ny, nx), dtype=np.float32)halves that. Setup time is not timed, so this only matters for peak memory on the asv runner. -
xrspatial/slope.py:46: the section banner above_cpustill reads "Planar backend functions (unchanged)". It predates this PR and is now wrong twice over. Not this PR's doing, but the diff sits directly under it.
What looks good
- The select semantics are exactly what the old branch gave. Any NaN in the 3x3 window makes
pNaN andarctan(NaN)is NaN; the only case that needed help was a finite window with a NaN centre, andctr == ctrcovers it.infbehaves the same on both sides too (inf - infis NaN in the stencil, aninfcentre still computes). outis pre-filled with NaN so the border rows and columns stay NaN without any change to the loop bounds.- The A/B in the PR body covers the degenerate shapes (3x3, 1x50, 50x1, 2x50) and the int16 cast, and the numpy and dask+numpy paths share
_cpuso both are covered by one comparison. - The benchmark reads
include_nancorrectly:common.pysets exactlyz[0, 0], a border cell, so the class adds its own 30% speckle. Verified the mask keeps the dask array lazy and chunked. - GPU and geodesic kernels untouched, as intended.
Checklist
- Algorithm matches reference: Horn 3x3 stencil unchanged, same coefficients and load order
- All implemented backends produce consistent results: numpy and dask+numpy share the kernel; cupy paths not modified
- NaN handling is correct: select on the centre, stencil propagates neighbours
- Edge cases are covered by tests: speckled footprint on numpy and dask, existing centre-NaN and 1xN/2xN tests still pass
- Dask chunk boundaries handled correctly: no change to
map_overlapdepth or boundary - No premature materialization or unnecessary copies: benchmark compute is deliberate
- Benchmark exists:
SlopeNaNadded - README feature matrix: not applicable, no new function or backend
- Docstrings: no public signature change
Note in SlopeNaN that its dask timing forces the compute while Slope.time_slope does not, so the two classes' dask numbers are not comparable. Build the speckle mask from a float32 draw to halve the setup allocation at nx=10000. State in the footprint test comment that the tests pass on the old kernel too, since they pin behaviour rather than reproduce a regression. Drop the stale "(unchanged)" from the planar section banner.
brendancol
commented
Sep 4, 2026
brendancol
left a comment
Contributor
Author
There was a problem hiding this comment.
PR Review: Drop the per-cell NaN early-out in the planar CPU slope kernel (follow-up)
Second pass over b2c5a22, which responds to the first review.
Blockers (must fix before merge)
None.
Suggestions (should fix, not blocking)
None.
Nits (optional improvements)
None.
Disposition of the first-pass findings
SlopeNaNvsSlopedask comparability: fixed, comment added atbenchmarks/benchmarks/slope.py:36-38.- Footprint tests pass on the old kernel too: fixed, stated in the comment block at
xrspatial/tests/test_slope.py:634-636. - float64 draw for the speckle mask: fixed,
dtype=np.float32atbenchmarks/benchmarks/slope.py:32. - Stale "(unchanged)" banner: fixed at
xrspatial/slope.py:46.
What looks good
- The kernel itself is unchanged since the first pass;
pytest xrspatial/tests/test_slope.pystill reports 137 passed and the benchmark class still runs on numpy and dask. - The float32 draw does not change the mask:
rng.random(dtype=np.float32) < 0.3selects with the same seed, so the NaN fraction stays at 0.30.
Checklist
- Algorithm matches reference
- All implemented backends produce consistent results
- NaN handling is correct
- Edge cases are covered by tests
- Dask chunk boundaries handled correctly
- No premature materialization or unnecessary copies
- Benchmark exists
- README feature matrix: not applicable
- Docstrings: no public signature change
Contributor
Author
|
Correction to the follow-up review: the float32 draw does not reproduce the float64 mask cell for cell (the generator consumes different bits per value), so the speckle pattern changed with that commit. The NaN fraction is unchanged at 0.30, which is what the benchmark depends on. The kernel is not affected. |
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.
Closes #3739
if np.isnan(data[y, x]): continueearly-out at the top of the planar CPU slope kernel_cpu. Neighbour NaN already propagates through the Horn stencil, andoutis pre-filled with NaN, so the only case the branch covered was a NaN centre with valid neighbours. That is now handled by a select after the arithmetic (out[y, x] = r if ctr == ctr else np.nan). No fastmath, no dtype change, neighbour loads in the same order.SlopeNaNasv benchmark (numpy and dask, samenxgrid asSlope) so the nodata path gets timed from now on.Backends: numpy and dask+numpy share
_cpu, so both get the change. The cupy and dask+cupy kernels and the geodesic kernels are untouched.Why
On DEMs with scattered nodata the early-out is a data-dependent branch the CPU cannot predict, and the mispredicts cost more than the few float32 ops and one
arctanthe skip avoids.Timings
2000x4000 float32 DEM (Gaussian bump plus
default_rng(71942).normal(0, 2)noise),_cpucalled directly,time.perf_counter, median of 9 after warmup. Three separate runs on a box that had other test suites running at the same time, so the absolute numbers wobble; the ratios hold up.Identical results in every regime:
np.array_equal(old, new, equal_nan=True)holds against the_cpufromorigin/mainfor the three rasters above plus 3x3, 1x50, 50x1, 2x50, a 3x3 with a NaN centre, and an int16 200x300 input.The NaN-free case is unchanged within noise. A contiguous NaN block is only a small win because the predictor learns it. The gain is specific to speckled nodata such as masked water or cloud holes, where the kernel is roughly 1.5x to 2x faster. The spike that motivated this measured 0.47x on a quiet machine; I could not get the box that quiet with the sibling runs going.
On the benchmark's NaN pattern
get_xr_dataarray(include_nan=True)inbenchmarks/benchmarks/common.pysets exactly one cell,z[0, 0], to NaN. That is a border cell the kernel never visits, so on its own it times the NaN-free path again.SlopeNaN.setupcalls it withinclude_nan=Trueand then masks 30% of cells at random (fixed seed) withDataArray.where, which keeps the dask array lazy for the dask case. Verified the resulting NaN fraction is 0.301 on both backends.common.pyis not modified. The dask case forces.compute()so asv times the kernel rather than graph construction, following the twi and convolution benchmarks.Test plan
pytest xrspatial/tests/test_slope.py -x -q: 137 passedorigin/main's_cpuon the nine inputs listed above, all bitwise identicalSlopeNaN.setupandtime_slope_nanrun for numpy and dask at nx=1000