From fbfd64988d99e95e1242108a1ddc28bbc145cd9a Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Sun, 16 Aug 2026 12:20:37 -0400 Subject: [PATCH 1/2] Grow the surface_distance Dijkstra heap on demand (#3723) _dijkstra and _dijkstra_geodesic sized their lazy-deletion min-heap at height * width. A pixel is re-pushed every time its tentative distance improves, so live occupancy is bounded by the number of improving relaxations, not by the pixel count. On a 48x48 raster with 40% target pixels over 200 m of relief, peak occupancy reached 1.09x the cap and _heap_push wrote past the end of h_keys / h_rows / h_cols, aborting the interpreter with "double free or corruption". The heap now starts at one entry per pixel (which the seeding loop can never exceed) and doubles when full. cost_distance fixed the same defect with a static height * width * (n_neighbors + 1) bound, but that would multiply this module's heap allocation by 9 and force _BYTES_PER_PIXEL from 80 to about 272, tightening the memory guard on every caller. Adds regression tests for the numpy, dask iterative and geodesic paths, plus a SurfaceDistanceDenseTargets asv benchmark. The existing benchmark misses this regime: its integer source raster makes nearly every pixel a target, so almost no relaxation improves a distance. --- .claude/sweep-performance-state.csv | 2 +- benchmarks/benchmarks/surface_distance.py | 35 +++++++ xrspatial/surface_distance.py | 35 ++++++- xrspatial/tests/test_surface_distance.py | 121 ++++++++++++++++++++++ 4 files changed, 190 insertions(+), 3 deletions(-) diff --git a/.claude/sweep-performance-state.csv b/.claude/sweep-performance-state.csv index 77efd2a00..051a49ecb 100644 --- a/.claude/sweep-performance-state.csv +++ b/.claude/sweep-performance-state.csv @@ -43,7 +43,7 @@ resample,2026-04-15T12:00:00Z,SAFE,compute-bound,0,false-positive,Downgraded. GP sieve,2026-04-14T12:00:00Z,WILL OOM,memory-bound,0,false-positive,False positive. Memory guards already in place on both dask paths. CCL is inherently global — documented limitation. CuPy CPU fallback is deliberate and documented. sky_view_factor,2026-03-31T18:00:00Z,SAFE,compute-bound,0,, slope,2026-03-31T18:00:00Z,SAFE,compute-bound,0,, -surface_distance,2026-03-31T18:00:00Z,SAFE,memory-bound,0,1128,Memory guard added to dd_grid allocation. +surface_distance,2026-08-16,RISKY,compute-bound,0,3723,"CRITICAL #3723/PR: _dijkstra and _dijkstra_geodesic sized the lazy-deletion heap height*width; peak occupancy hit 1.09x cap at 40% target density, _heap_push wrote OOB (SIGABRT 'double free or corruption'); fixed by doubling on demand rather than cost_distance's static (n_neighbors+1) bound, which would have raised the 80 B/px guard to ~272. MEDIUM (documented, not fixed): dask iterative path recomputes each source and elev block 6x (2 in _preprocess_tiles_sd, 4 across sweeps, 1 in _assemble_sd) - measured on 40x40/80x80/120x120; MEDIUM: cupy Bellman-Ford relaxation is 6-27x slower than the numpy Dijkstra (512^2: 0.315s vs 0.051s, 373 kernel launches each with a device sync), and DIRECTION mode round-trips to host for _vectorized_calc_direction; MEDIUM: map_overlap depth guard tests pad < max(chunks), so depth may approach chunk size (~9x redundant work) and exceed the smallest chunk when chunking is uneven. Memory guard is correctly chunk-scoped on dask (test_dask_path_bounded_per_chunk) - the morphology #3401 false-MemoryError hazard is NOT present. _sd_relax_kernel has ~8 float64 locals, no register pressure. OOM verdict RISKY: memory stays chunk-bounded but the whole iterative Dijkstra runs eagerly and single-threaded on the client at graph-build time." terrain,2026-03-31T18:00:00Z,RISKY,compute-bound,0,, terrain_metrics,2026-03-31T18:00:00Z,SAFE,memory-bound,0,, viewshed,2026-04-05T12:00:00Z,SAFE,memory-bound,0,fixed-in-tree,Tier B memory estimate tightened from 280 to 368 bytes/pixel (accounts for lexsort double-alloc + computed raster). astype copy=False avoids needless float64 copy. diff --git a/benchmarks/benchmarks/surface_distance.py b/benchmarks/benchmarks/surface_distance.py index 1cf7aad60..bdf1602b3 100644 --- a/benchmarks/benchmarks/surface_distance.py +++ b/benchmarks/benchmarks/surface_distance.py @@ -1,4 +1,5 @@ import numpy as np +import xarray as xr from xrspatial.surface_distance import ( surface_distance, surface_allocation, surface_direction, @@ -24,3 +25,37 @@ def time_surface_allocation(self, nx, type): def time_surface_direction(self, nx, type): surface_direction(self.agg, self.elev) + + +class SurfaceDistanceDenseTargets: + """Dijkstra with a mid-density target raster over rugged relief. + + This is the regime where the lazy-deletion heap holds more than one + live entry per pixel (#3723). The existing SurfaceDistance benchmark + misses it: its integer source raster makes nearly every pixel a + target, so almost no relaxation ever improves a distance. + """ + + params = ([200, 400], [0.05, 0.2, 0.4]) + param_names = ("nx", "target_fraction") + + def setup(self, nx, target_fraction): + ny = nx // 2 + rng = np.random.default_rng(71942) + source = np.zeros((ny, nx), dtype=np.float64) + n_targets = max(1, int(ny * nx * target_fraction)) + source.flat[rng.choice(ny * nx, size=n_targets, replace=False)] = 1.0 + elev = rng.random((ny, nx)) * 200.0 + + coords = dict(y=np.arange(ny, dtype=np.float64), + x=np.arange(nx, dtype=np.float64)) + self.agg = xr.DataArray(source, coords=coords, dims=["y", "x"], + attrs={"res": (1.0, 1.0)}) + self.elev = xr.DataArray(elev, coords=coords, dims=["y", "x"], + attrs={"res": (1.0, 1.0)}) + + def time_surface_distance(self, nx, target_fraction): + surface_distance(self.agg, self.elev) + + def peakmem_surface_distance(self, nx, target_fraction): + surface_distance(self.agg, self.elev) diff --git a/xrspatial/surface_distance.py b/xrspatial/surface_distance.py index c91aeb13e..81b9aa98e 100644 --- a/xrspatial/surface_distance.py +++ b/xrspatial/surface_distance.py @@ -170,6 +170,28 @@ def _check_gpu_memory(rows, cols): # --------------------------------------------------------------------------- +@ngjit +def _heap_grow(keys, rows, cols, size): + """Return copies of the heap arrays with twice the capacity. + + The Dijkstra kernels below use a lazy-deletion min-heap: a pixel is + pushed again every time its tentative distance improves, and stale + entries linger until a pop skips them. Live occupancy is therefore + bounded by the number of improving relaxations, not by the pixel + count, and a fixed ``height * width`` heap can overflow (#3723). + Growing on demand keeps the usual footprint at one entry per pixel + without capping the number of pushes. + """ + cap = 2 * len(keys) + new_keys = np.empty(cap, dtype=np.float64) + new_rows = np.empty(cap, dtype=np.int64) + new_cols = np.empty(cap, dtype=np.int64) + new_keys[:size] = keys[:size] + new_rows[:size] = rows[:size] + new_cols[:size] = cols[:size] + return new_keys, new_rows, new_cols + + @ngjit def _seed_sources(source_data, elev_data, target_values, dist, alloc, src_row, src_col): @@ -211,7 +233,9 @@ def _dijkstra(elev_data, height, width, max_distance, """ n_neighbors = len(dy) - max_heap = height * width + # Starting capacity: one entry per pixel, which the seeding loop below + # can never exceed. Relaxations can, so _heap_grow doubles it there. + max_heap = max(height * width, 1) h_keys = np.empty(max_heap, dtype=np.float64) h_rows = np.empty(max_heap, dtype=np.int64) h_cols = np.empty(max_heap, dtype=np.int64) @@ -260,6 +284,9 @@ def _dijkstra(elev_data, height, width, max_distance, alloc[vr, vc] = alloc[ur, uc] src_row[vr, vc] = src_row[ur, uc] src_col[vr, vc] = src_col[ur, uc] + if h_size == len(h_keys): + h_keys, h_rows, h_cols = _heap_grow( + h_keys, h_rows, h_cols, h_size) h_size = _heap_push(h_keys, h_rows, h_cols, h_size, new_cost, vr, vc) @@ -274,7 +301,8 @@ def _dijkstra_geodesic(elev_data, height, width, max_distance, """ n_neighbors = len(dy) - max_heap = height * width + # See _dijkstra for the heap-capacity rationale. + max_heap = max(height * width, 1) h_keys = np.empty(max_heap, dtype=np.float64) h_rows = np.empty(max_heap, dtype=np.int64) h_cols = np.empty(max_heap, dtype=np.int64) @@ -322,6 +350,9 @@ def _dijkstra_geodesic(elev_data, height, width, max_distance, alloc[vr, vc] = alloc[ur, uc] src_row[vr, vc] = src_row[ur, uc] src_col[vr, vc] = src_col[ur, uc] + if h_size == len(h_keys): + h_keys, h_rows, h_cols = _heap_grow( + h_keys, h_rows, h_cols, h_size) h_size = _heap_push(h_keys, h_rows, h_cols, h_size, new_cost, vr, vc) diff --git a/xrspatial/tests/test_surface_distance.py b/xrspatial/tests/test_surface_distance.py index 74ae62fbd..9bc6d7d35 100644 --- a/xrspatial/tests/test_surface_distance.py +++ b/xrspatial/tests/test_surface_distance.py @@ -792,3 +792,124 @@ def test_error_message_mentions_grid_size(self): surface_distance(raster, elevation) with pytest.raises(MemoryError, match="dask"): surface_distance(raster, elevation) + + +# --------------------------------------------------------------------------- +# Heap capacity regression (#3723) +# --------------------------------------------------------------------------- + + +def _reference_surface_distance(source, elev, connectivity=8, cellsize=1.0): + """Pure-Python multi-source Dijkstra reference for surface distance.""" + import heapq + + h, w = source.shape + diag = np.sqrt(2.0) * cellsize + if connectivity == 8: + nbrs = [(-1, -1, diag), (-1, 0, cellsize), (-1, 1, diag), + (0, -1, cellsize), (0, 1, cellsize), + (1, -1, diag), (1, 0, cellsize), (1, 1, diag)] + else: + nbrs = [(0, -1, cellsize), (-1, 0, cellsize), + (1, 0, cellsize), (0, 1, cellsize)] + + dist = np.full((h, w), np.inf) + heap = [] + for r in range(h): + for c in range(w): + if (source[r, c] != 0 and np.isfinite(source[r, c]) + and np.isfinite(elev[r, c])): + dist[r, c] = 0.0 + heapq.heappush(heap, (0.0, r, c)) + + done = np.zeros((h, w), dtype=bool) + while heap: + d, r, c = heapq.heappop(heap) + if done[r, c]: + continue + done[r, c] = True + for dr, dc, hd in nbrs: + vr, vc = r + dr, c + dc + if not (0 <= vr < h and 0 <= vc < w) or done[vr, vc]: + continue + if not np.isfinite(elev[vr, vc]): + continue + dz = elev[vr, vc] - elev[r, c] + nd = d + np.sqrt(hd * hd + dz * dz) + if nd < dist[vr, vc]: + dist[vr, vc] = nd + heapq.heappush(heap, (nd, vr, vc)) + return dist + + +def _dense_target_scene(n=48, n_targets=921, relief=200.0, seed=1): + """Dense-target scene that overflowed the old height*width heap.""" + rng = np.random.default_rng(seed) + source = np.zeros((n, n), dtype=np.float64) + source.flat[rng.choice(n * n, size=n_targets, replace=False)] = 1.0 + elev = rng.random((n, n)) * relief + return source, elev + + +def test_dense_targets_do_not_overflow_the_heap(): + """A lazy-deletion heap can exceed height*width live entries. + + Before #3723 the heap arrays were sized height*width, so this scene + made _heap_push write past the end of them (SIGABRT without bounds + checking, IndexError with NUMBA_BOUNDSCHECK=1). + """ + source, elev = _dense_target_scene() + raster = _make_raster(source) + elevation = _make_raster(elev) + + result = _compute(surface_distance(raster, elevation, connectivity=8)) + expected = _reference_surface_distance(source, elev, connectivity=8) + + assert np.all(np.isfinite(result)) + np.testing.assert_allclose(result, expected.astype(np.float32), + rtol=1e-5, atol=1e-4) + + +def test_dense_targets_allocation_and_direction_do_not_overflow(): + """The allocation and direction modes share the same Dijkstra kernel.""" + source, elev = _dense_target_scene() + raster = _make_raster(source) + elevation = _make_raster(elev) + + alloc = _compute(surface_allocation(raster, elevation, connectivity=8)) + direction = _compute(surface_direction(raster, elevation, connectivity=8)) + + assert np.all(alloc == 1.0) + assert np.all(np.isfinite(direction)) + + +@pytest.mark.skipif(da is None, reason="dask not installed") +def test_dense_targets_dask_iterative_does_not_overflow(): + """The dask iterative path runs the same kernel per tile.""" + source, elev = _dense_target_scene() + raster_np = _make_raster(source) + elev_np = _make_raster(elev) + raster = _make_raster(source, backend='dask+numpy', chunks=(48, 48)) + elevation = _make_raster(elev, backend='dask+numpy', chunks=(48, 48)) + + np_result = _compute(surface_distance(raster_np, elev_np)) + with pytest.warns(UserWarning, match="iterative"): + dask_result = _compute(surface_distance(raster, elevation)) + + np.testing.assert_allclose(dask_result, np_result, rtol=1e-5, + equal_nan=True) + + +def test_geodesic_dense_targets_do_not_overflow(): + """_dijkstra_geodesic carries the same heap sizing.""" + source, elev = _dense_target_scene(n=32, n_targets=410, relief=200.0) + h, w = source.shape + coords = {'y': np.linspace(10.0, 10.0 + 0.01 * (h - 1), h), + 'x': np.linspace(20.0, 20.0 + 0.01 * (w - 1), w)} + raster = xr.DataArray(source, dims=['y', 'x'], coords=coords, + attrs={'res': (0.01, 0.01)}) + elevation = xr.DataArray(elev, dims=['y', 'x'], coords=coords, + attrs={'res': (0.01, 0.01)}) + + result = _compute(surface_distance(raster, elevation, method='geodesic')) + assert np.all(np.isfinite(result)) From 4afb223dd1ffdc1328aa31bdd15050bdfb10963d Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Sun, 16 Aug 2026 12:24:49 -0400 Subject: [PATCH 2/2] Address review findings on the heap fix (#3723) - Guard the seeding loops in both kernels as well. The old comment argued the loop can push at most one entry per pixel into a heap sized for one entry per pixel, which is true today, but it leaves a memory-safety invariant resting on a comment. Every push site now checks capacity. - Say in the _BYTES_PER_PIXEL comment that the heap terms are a starting capacity, not a bound, so 80 reads as a floor rather than an estimate that quietly went stale. - Move the rationale for the capacity choice from _heap_grow into _dijkstra, where the capacity is picked. - Cover the bounded map_overlap dask branch with a dense-target test. It calls _surface_distance_numpy per padded chunk, so it was already fixed, but only the iterative branch had a test. --- xrspatial/surface_distance.py | 31 +++++++++++++++--------- xrspatial/tests/test_surface_distance.py | 18 ++++++++++++++ 2 files changed, 37 insertions(+), 12 deletions(-) diff --git a/xrspatial/surface_distance.py b/xrspatial/surface_distance.py index 81b9aa98e..e682eaf11 100644 --- a/xrspatial/surface_distance.py +++ b/xrspatial/surface_distance.py @@ -89,6 +89,12 @@ class cupy: # type: ignore[no-redef] # output (float32) 4 # direction-mode temps ~16 # Total ~80 bytes/pixel. A 50000x50000 raster needs ~200 GB. +# +# The heap terms above are the starting capacity of one entry per pixel. +# Dense-target rasters push more than that and _heap_grow doubles the +# three arrays, so the real figure can exceed 80 (48 bytes/pixel of heap +# after one doubling, and ~72 transiently while the copy is in flight). +# This estimate is a floor, not a bound. _BYTES_PER_PIXEL = 80 # CuPy backend skips the explicit binary heap (parallel relaxation instead) @@ -172,16 +178,7 @@ def _check_gpu_memory(rows, cols): @ngjit def _heap_grow(keys, rows, cols, size): - """Return copies of the heap arrays with twice the capacity. - - The Dijkstra kernels below use a lazy-deletion min-heap: a pixel is - pushed again every time its tentative distance improves, and stale - entries linger until a pop skips them. Live occupancy is therefore - bounded by the number of improving relaxations, not by the pixel - count, and a fixed ``height * width`` heap can overflow (#3723). - Growing on demand keeps the usual footprint at one entry per pixel - without capping the number of pushes. - """ + """Return copies of the heap arrays with twice the capacity.""" cap = 2 * len(keys) new_keys = np.empty(cap, dtype=np.float64) new_rows = np.empty(cap, dtype=np.int64) @@ -233,8 +230,12 @@ def _dijkstra(elev_data, height, width, max_distance, """ n_neighbors = len(dy) - # Starting capacity: one entry per pixel, which the seeding loop below - # can never exceed. Relaxations can, so _heap_grow doubles it there. + # This is a lazy-deletion min-heap: a pixel is pushed again every time + # its tentative distance improves, and stale entries linger until a pop + # skips them. Live occupancy therefore tracks the number of improving + # relaxations, not the pixel count, so the old fixed height*width heap + # let _heap_push write past the end of these arrays (#3723). Start at + # one entry per pixel and double whenever a push would overflow. max_heap = max(height * width, 1) h_keys = np.empty(max_heap, dtype=np.float64) h_rows = np.empty(max_heap, dtype=np.int64) @@ -247,6 +248,9 @@ def _dijkstra(elev_data, height, width, max_distance, for r in range(height): for c in range(width): if dist[r, c] < np.inf: + if h_size == len(h_keys): + h_keys, h_rows, h_cols = _heap_grow( + h_keys, h_rows, h_cols, h_size) h_size = _heap_push(h_keys, h_rows, h_cols, h_size, dist[r, c], r, c) @@ -313,6 +317,9 @@ def _dijkstra_geodesic(elev_data, height, width, max_distance, for r in range(height): for c in range(width): if dist[r, c] < np.inf: + if h_size == len(h_keys): + h_keys, h_rows, h_cols = _heap_grow( + h_keys, h_rows, h_cols, h_size) h_size = _heap_push(h_keys, h_rows, h_cols, h_size, dist[r, c], r, c) diff --git a/xrspatial/tests/test_surface_distance.py b/xrspatial/tests/test_surface_distance.py index 9bc6d7d35..ffdf3e439 100644 --- a/xrspatial/tests/test_surface_distance.py +++ b/xrspatial/tests/test_surface_distance.py @@ -900,6 +900,24 @@ def test_dense_targets_dask_iterative_does_not_overflow(): equal_nan=True) +@pytest.mark.skipif(da is None, reason="dask not installed") +def test_dense_targets_dask_bounded_does_not_overflow(): + """The bounded map_overlap path runs the kernel per padded chunk.""" + source, elev = _dense_target_scene() + raster_np = _make_raster(source) + elev_np = _make_raster(elev) + raster = _make_raster(source, backend='dask+numpy', chunks=(16, 16)) + elevation = _make_raster(elev, backend='dask+numpy', chunks=(16, 16)) + + np_result = _compute(surface_distance(raster_np, elev_np, + max_distance=10.0)) + dask_result = _compute(surface_distance(raster, elevation, + max_distance=10.0)) + + np.testing.assert_allclose(dask_result, np_result, rtol=1e-5, + equal_nan=True) + + def test_geodesic_dense_targets_do_not_overflow(): """_dijkstra_geodesic carries the same heap sizing.""" source, elev = _dense_target_scene(n=32, n_targets=410, relief=200.0)