Skip to content

Drop the per-cell NaN early-out in the planar CPU slope kernel - #3742

Open
brendancol wants to merge 2 commits into
mainfrom
issue-3739
Open

Drop the per-cell NaN early-out in the planar CPU slope kernel#3742
brendancol wants to merge 2 commits into
mainfrom
issue-3739

Conversation

@brendancol

Copy link
Copy Markdown
Contributor

Closes #3739

  • Drop the if np.isnan(data[y, x]): continue early-out at the top of the planar CPU slope kernel _cpu. Neighbour NaN already propagates through the Horn stencil, and out is 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.
  • Add a SlopeNaN asv benchmark (numpy and dask, same nx grid as Slope) so the nodata path gets timed from now on.
  • Add tests that pin the NaN footprint on a speckled raster: NaN at every centre-NaN cell and its 8-neighbours, finite everywhere else in the interior, on numpy and dask+numpy.

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 arctan the skip avoids.

Timings

2000x4000 float32 DEM (Gaussian bump plus default_rng(71942).normal(0, 2) noise), _cpu called 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.

NaN pattern before after ratio
none 102 to 107 ms 92 to 101 ms 0.90x to 0.98x
30% random NaN 90 to 94 ms 53 to 67 ms 0.58x to 0.72x
left half NaN 79 to 81 ms 70 to 77 ms 0.86x to 0.97x

Identical results in every regime: np.array_equal(old, new, equal_nan=True) holds against the _cpu from origin/main for 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) in benchmarks/benchmarks/common.py sets 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.setup calls it with include_nan=True and then masks 30% of cells at random (fixed seed) with DataArray.where, which keeps the dask array lazy for the dask case. Verified the resulting NaN fraction is 0.301 on both backends. common.py is 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 passed
  • Direct A/B against origin/main's _cpu on the nine inputs listed above, all bitwise identical
  • SlopeNaN.setup and time_slope_nan run for numpy and dask at nx=1000
  • flake8 on the three touched files reports only the pre-existing aligned rows in the QGIS fixture (same 30 hits before and after)

_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.
@github-actions github-actions Bot added the performance PR touches performance-sensitive code label Sep 4, 2026

@brendancol brendancol left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_nan forces .compute() for dask, but the existing Slope.time_slope goes through Benchmarking.time, which does not. So for type="dask" the two classes time different things (graph build vs actual kernel work) and a side-by-side read of Slope vs SlopeNaN is misleading. The compute is the right call for the new class. Either add a one-line comment saying the dask numbers are not comparable with Slope, or leave Slope as is and accept it; changing Slope is out of scope here.
  • xrspatial/tests/test_slope.py:642-648: _expected_nan_footprint builds 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)) at nx=10000 allocates 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 _cpu still 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 p NaN and arctan(NaN) is NaN; the only case that needed help was a finite window with a NaN centre, and ctr == ctr covers it. inf behaves the same on both sides too (inf - inf is NaN in the stencil, an inf centre still computes).
  • out is 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 _cpu so both are covered by one comparison.
  • The benchmark reads include_nan correctly: common.py sets exactly z[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_overlap depth or boundary
  • No premature materialization or unnecessary copies: benchmark compute is deliberate
  • Benchmark exists: SlopeNaN added
  • 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 brendancol left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

  • SlopeNaN vs Slope dask comparability: fixed, comment added at benchmarks/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.float32 at benchmarks/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.py still 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.3 selects 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

@brendancol

Copy link
Copy Markdown
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.

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

Labels

performance PR touches performance-sensitive code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

slope: drop the per-cell NaN early-out in the planar CPU kernel

1 participant