From 297c549a06eaa120a863336c309b2b782232eaa8 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Thu, 23 Jul 2026 14:08:53 -0400 Subject: [PATCH 1/3] Drop twi dask+cupy host round-trip; vectorize watershed init loops (#3692) --- benchmarks/benchmarks/twi.py | 13 ++++++-- xrspatial/hydro/tests/test_twi_d8.py | 36 ++++++++++++++++++++++ xrspatial/hydro/tests/test_watershed_d8.py | 18 +++++++++++ xrspatial/hydro/twi_d8.py | 22 +++---------- xrspatial/hydro/watershed_d8.py | 16 +++------- xrspatial/hydro/watershed_dinf.py | 32 ++++++------------- xrspatial/hydro/watershed_mfd.py | 32 ++++++------------- 7 files changed, 94 insertions(+), 75 deletions(-) diff --git a/benchmarks/benchmarks/twi.py b/benchmarks/benchmarks/twi.py index c74773081..d34d4aec7 100644 --- a/benchmarks/benchmarks/twi.py +++ b/benchmarks/benchmarks/twi.py @@ -1,20 +1,29 @@ from xrspatial import flow_direction, flow_accumulation, slope, twi +from xrspatial.utils import has_cuda_and_cupy from .common import get_xr_dataarray class TWI: - params = ([100, 300, 1000], ["numpy", "dask"]) + params = ([100, 300, 1000], ["numpy", "dask", "cupy", "dask+cupy"]) param_names = ("nx", "type") def setup(self, nx, type): + if type in ("cupy", "dask+cupy") and not has_cuda_and_cupy(): + raise NotImplementedError() + ny = nx // 2 elev = get_xr_dataarray((ny, nx), "numpy") flow_dir = flow_direction(elev) self.flow_accum = flow_accumulation(flow_dir) self.slope_agg = slope(elev) - if type == "dask": + if type in ("cupy", "dask+cupy"): + import cupy + self.flow_accum.data = cupy.asarray(self.flow_accum.data) + self.slope_agg.data = cupy.asarray(self.slope_agg.data) + + if type in ("dask", "dask+cupy"): import dask.array as da chunks = (max(1, ny // 2), max(1, nx // 2)) self.flow_accum.data = da.from_array( diff --git a/xrspatial/hydro/tests/test_twi_d8.py b/xrspatial/hydro/tests/test_twi_d8.py index 69f787d93..5bda47aa5 100644 --- a/xrspatial/hydro/tests/test_twi_d8.py +++ b/xrspatial/hydro/tests/test_twi_d8.py @@ -180,3 +180,39 @@ def test_numpy_equals_dask_cupy(self): np.testing.assert_allclose( result_np.data, result_dc.data.compute().get(), equal_nan=True, rtol=1e-10) + + def test_dask_cupy_stays_on_device(self): + """The dask+cupy path must not round-trip chunks through the host. + + Regression test for the removed ``_twi_dask_cupy`` wrapper, which + pulled every chunk to numpy and pushed it back with cp.asarray. + The lazy result must keep a cupy meta, and twi must add the same + graph layers on dask+cupy input as on plain dask input (the + round-trip added a host-transfer map_blocks layer per input plus + a device-transfer layer on the output). + """ + import cupy + + np.random.seed(42) + fa_data = np.random.uniform(1, 1000, (6, 6)).astype(np.float64) + sl_data = np.random.uniform(0.1, 60, (6, 6)).astype(np.float64) + + fa_da, sl_da = _make_twi_rasters(fa_data, sl_data, backend='dask', + chunks=(3, 3)) + fa_dc, sl_dc = _make_twi_rasters(fa_data, sl_data, + backend='dask+cupy', chunks=(3, 3)) + + result_da = twi(fa_da, sl_da) + result_dc = twi(fa_dc, sl_dc) + + assert isinstance(result_dc.data._meta, cupy.ndarray) + + def added_layers(result, *inputs): + input_layers = set() + for arr in inputs: + input_layers |= set(arr.data.__dask_graph__().layers) + return len(set(result.data.__dask_graph__().layers) + - input_layers) + + assert (added_layers(result_dc, fa_dc, sl_dc) + == added_layers(result_da, fa_da, sl_da)) diff --git a/xrspatial/hydro/tests/test_watershed_d8.py b/xrspatial/hydro/tests/test_watershed_d8.py index 13f4cb799..8e653a963 100644 --- a/xrspatial/hydro/tests/test_watershed_d8.py +++ b/xrspatial/hydro/tests/test_watershed_d8.py @@ -95,6 +95,24 @@ def test_nan_handling(): assert np.isnan(result.data[1, 1]) +def test_pour_point_on_nan_flow_dir(): + """A pour point placed on a NaN flow_dir cell is nodata, not a label.""" + flow_dir = np.array([ + [1.0, 0.0], + [np.nan, 1.0], + ], dtype=np.float64) + pour_points = np.full((2, 2), np.nan, dtype=np.float64) + pour_points[1, 0] = 9.0 # on the NaN flow_dir cell + pour_points[0, 1] = 7.0 + + fd_da = create_test_raster(flow_dir) + pp_da = create_test_raster(pour_points) + result = watershed(fd_da, pp_da) + + assert np.isnan(result.data[1, 0]) + assert result.data[0, 1] == 7.0 + + def test_linear_chain(): """Row of cells flowing east to a pour point at the end.""" N = 6 diff --git a/xrspatial/hydro/twi_d8.py b/xrspatial/hydro/twi_d8.py index a07e27abf..4a667c27f 100644 --- a/xrspatial/hydro/twi_d8.py +++ b/xrspatial/hydro/twi_d8.py @@ -65,7 +65,9 @@ def twi_d8(flow_accum: xr.DataArray, sl_data = slope_agg.data if has_cuda_and_cupy() and is_dask_cupy(flow_accum): - out = _twi_dask_cupy(fa_data, sl_data, cellsize) + # TWI is purely elementwise, so cupy-backed dask arrays run the + # same expression natively; no host round-trip needed. + out = _twi_dask(fa_data, sl_data, cellsize) elif has_cuda_and_cupy() and is_cupy_array(fa_data): out = _twi_cupy(fa_data, sl_data, cellsize) @@ -106,25 +108,9 @@ def _twi_cupy(fa, sl, cellsize): def _twi_dask(fa, sl, cellsize): + """Elementwise TWI on dask arrays (numpy- or cupy-backed chunks).""" import dask.array as _da sca = fa * cellsize tan_slope = _da.tan(_da.radians(sl)) tan_slope = _da.where(tan_slope < _TAN_MIN, _TAN_MIN, tan_slope) return _da.log(sca / tan_slope) - - -def _twi_dask_cupy(fa, sl, cellsize): - import cupy as cp - fa_np = fa.map_blocks( - lambda b: b.get(), dtype=fa.dtype, - meta=np.array((), dtype=fa.dtype), - ) - sl_np = sl.map_blocks( - lambda b: b.get(), dtype=sl.dtype, - meta=np.array((), dtype=sl.dtype), - ) - result = _twi_dask(fa_np, sl_np, cellsize) - return result.map_blocks( - cp.asarray, dtype=result.dtype, - meta=cp.array((), dtype=result.dtype), - ) diff --git a/xrspatial/hydro/watershed_d8.py b/xrspatial/hydro/watershed_d8.py index 892a40028..70667e589 100644 --- a/xrspatial/hydro/watershed_d8.py +++ b/xrspatial/hydro/watershed_d8.py @@ -1043,17 +1043,11 @@ def watershed_d8(flow_dir: xr.DataArray, h, w = fd.shape # Init labels and state: pour points → resolved (state 3), # NaN flow_dir → nodata (state 0), others → unresolved (state 1) - labels = np.full((h, w), np.nan, dtype=np.float64) - state = np.zeros((h, w), dtype=np.int8) - for r in range(h): - for c in range(w): - if fd[r, c] != fd[r, c]: # NaN - pass # state 0, label NaN - elif pp[r, c] == pp[r, c]: # not NaN → pour point - labels[r, c] = pp[r, c] - state[r, c] = 3 - else: - state[r, c] = 1 # unresolved + fd_valid = ~np.isnan(fd) + pp_valid = ~np.isnan(pp) + labels = np.where(fd_valid & pp_valid, pp, np.nan) + state = np.where(fd_valid, + np.where(pp_valid, 3, 1), 0).astype(np.int8) out = _watershed_cpu(fd, labels, state, h, w) elif has_cuda_and_cupy() and is_cupy_array(data): diff --git a/xrspatial/hydro/watershed_dinf.py b/xrspatial/hydro/watershed_dinf.py index 17221b6b8..90ea81388 100644 --- a/xrspatial/hydro/watershed_dinf.py +++ b/xrspatial/hydro/watershed_dinf.py @@ -219,17 +219,11 @@ def _watershed_dinf_cupy(flow_dir_data, pour_points_data): fd_np = _to_numpy_f64(flow_dir_data) pp_np = _to_numpy_f64(pour_points_data) h, w = fd_np.shape - labels = np.full((h, w), np.nan, dtype=np.float64) - state = np.zeros((h, w), dtype=np.int8) - for r in range(h): - for c in range(w): - if fd_np[r, c] != fd_np[r, c]: - pass - elif pp_np[r, c] == pp_np[r, c]: - labels[r, c] = pp_np[r, c] - state[r, c] = 3 - else: - state[r, c] = 1 + fd_valid = ~np.isnan(fd_np) + pp_valid = ~np.isnan(pp_np) + labels = np.where(fd_valid & pp_valid, pp_np, np.nan) + state = np.where(fd_valid, + np.where(pp_valid, 3, 1), 0).astype(np.int8) out = _watershed_dinf_cpu(fd_np, labels, state, h, w) return cp.asarray(out) @@ -699,17 +693,11 @@ def watershed_dinf(flow_dir_dinf: xr.DataArray, fd = data.astype(np.float64) pp = np.asarray(pp_data, dtype=np.float64) h, w = fd.shape - labels = np.full((h, w), np.nan, dtype=np.float64) - state = np.zeros((h, w), dtype=np.int8) - for r in range(h): - for c in range(w): - if fd[r, c] != fd[r, c]: - pass - elif pp[r, c] == pp[r, c]: - labels[r, c] = pp[r, c] - state[r, c] = 3 - else: - state[r, c] = 1 + fd_valid = ~np.isnan(fd) + pp_valid = ~np.isnan(pp) + labels = np.where(fd_valid & pp_valid, pp, np.nan) + state = np.where(fd_valid, + np.where(pp_valid, 3, 1), 0).astype(np.int8) out = _watershed_dinf_cpu(fd, labels, state, h, w) elif has_cuda_and_cupy() and is_cupy_array(data): diff --git a/xrspatial/hydro/watershed_mfd.py b/xrspatial/hydro/watershed_mfd.py index 6c51c423e..68397fc1b 100644 --- a/xrspatial/hydro/watershed_mfd.py +++ b/xrspatial/hydro/watershed_mfd.py @@ -227,17 +227,11 @@ def _watershed_mfd_cupy(fractions_data, pour_points_data): fr_np = _to_numpy_f64(fractions_data) pp_np = _to_numpy_f64(pour_points_data) _, h, w = fr_np.shape - labels = np.full((h, w), np.nan, dtype=np.float64) - state = np.zeros((h, w), dtype=np.int8) - for r in range(h): - for c in range(w): - if fr_np[0, r, c] != fr_np[0, r, c]: - pass - elif pp_np[r, c] == pp_np[r, c]: - labels[r, c] = pp_np[r, c] - state[r, c] = 3 - else: - state[r, c] = 1 + fr_valid = ~np.isnan(fr_np[0]) + pp_valid = ~np.isnan(pp_np) + labels = np.where(fr_valid & pp_valid, pp_np, np.nan) + state = np.where(fr_valid, + np.where(pp_valid, 3, 1), 0).astype(np.int8) out = _watershed_mfd_cpu(fr_np, labels, state, h, w) return cp.asarray(out) @@ -701,17 +695,11 @@ def watershed_mfd(flow_dir_mfd: xr.DataArray, _check_memory(H, W) fr = data.astype(np.float64) pp = np.asarray(pp_data, dtype=np.float64) - labels = np.full((H, W), np.nan, dtype=np.float64) - state = np.zeros((H, W), dtype=np.int8) - for r in range(H): - for c in range(W): - if fr[0, r, c] != fr[0, r, c]: - pass - elif pp[r, c] == pp[r, c]: - labels[r, c] = pp[r, c] - state[r, c] = 3 - else: - state[r, c] = 1 + fr_valid = ~np.isnan(fr[0]) + pp_valid = ~np.isnan(pp) + labels = np.where(fr_valid & pp_valid, pp, np.nan) + state = np.where(fr_valid, + np.where(pp_valid, 3, 1), 0).astype(np.int8) out = _watershed_mfd_cpu(fr, labels, state, H, W) elif has_cuda_and_cupy() and is_cupy_array(data): From fd3f5ee2b3c483ad6d4fd193d00bf510ecc005b5 Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Thu, 23 Jul 2026 14:11:12 -0400 Subject: [PATCH 2/3] Update sweep-performance state for hydro pass (#3692) --- .claude/sweep-performance-state.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.claude/sweep-performance-state.csv b/.claude/sweep-performance-state.csv index 9c0e8fb7e..6909955b7 100644 --- a/.claude/sweep-performance-state.csv +++ b/.claude/sweep-performance-state.csv @@ -21,7 +21,7 @@ geodesic,2026-03-31T18:00:00Z,N/A,compute-bound,0,, geotiff,2026-07-05,SAFE,IO-bound,0,,"Pass 18 (2026-07-05): re-audit, 0 CRITICAL/0 HIGH/0 MEDIUM. Only 2 commits since Pass 17: #3604 (f0275b61, O(1) ndim gate in _writers/eager.py before dispatch -- no perf surface) and #3605 (6b0887fb, tests only). Category scans across the whole subpackage came back matching Pass 17's known-clean state: Cat1 .compute()/.get() sites are the documented streaming writes (_writer.py row bands, _writers/gpu.py per-band stream #3166/#3117) and batched D2H; Cat2 50k-task cap intact at _backends/dask.py:580, #3597 StreamingStats single-pass intact, _dask_finite_stats fuses 5 reductions into one dask.compute; Cat3 all Device().synchronize() sites batch-end/error-recovery (#2212/#2107 fixes intact); Cat5 _overview_kernels.py @ngjit row-major float-stable. Probes on-device this pass (PYTHONPATH=worktree, verified xrspatial.__file__): dask read 2560x2560 deflate chunks=256 = 400 tasks/100 chunks (4/chunk); eager numpy parity 0.0; eager GPU read returns cupy, parity 0.0, 400ms incl warmup, no host round trip; dask+cupy lazy (401 tasks), corner chunk is cupy, parity 0.0. _gpu_decode.py untouched since Pass 17 so recorded register values stand (inflate 67 / LZW 29; introspection still unavailable via numba _Kernel API). Prior deferred LOWs unchanged (redundant .copy() in _writer predictor encode, identity pack arithmetic, twice-built IFD byte length, per-level device sync in _block_reduce_2d_gpu, _nvcomp_batch_compress adler32 staging, _CloudSource.read_range per-call open). No HIGH/MEDIUM => no /rockout, no new benchmark required. SAFE/IO-bound holds. | Pass 17 (2026-07-02): re-audit, 0 CRITICAL/0 HIGH/0 MEDIUM. Prior row logged as needing GPU re-validation; CUDA IS available on this host and all GPU decode/read paths were re-run on-device this pass (no cuda-unavailable token). Probes: dask read 2560x2560 chunks=256 = 400 tasks / 100 chunks (4/chunk, 50k cap intact at _backends/dask.py:580); eager numpy parity 0.0; eager GPU read of deflate-tiled file returns cupy, parity 0.0, 165ms, no host round trip; dask+cupy read lazy (401 tasks) producing cupy corner parity 0.0; GPU deflate+fp-predictor(predictor=3) decode parity 0.0 (routes through nvcomp GPU path on this host, numba _inflate_tiles fallback not specialized). Cat1: no .values/np.asarray/np.array on dask or cupy arrays in _backends/; the _writer.py:1221-1394 .compute() calls are the intentional streaming row-band writes (single-pass verified in #3117/#3235/#3597, not per-chunk re-execution). Cat3: all Device().synchronize() in _gpu_decode.py (1008/1029/1264/1758/2531/2824) are batch-end or error-recovery, not per-tile in loops -- #2212/#2107 fixes intact. Cat5: _overview_kernels.py are @ngjit CPU kernels (row-major inner loop, float-stable accumulators), not CUDA. Reviewed the 6 commits since last inspection (#3595/#3599 PAM overwrite cleanup, #3592/#3598 + #3593/#3596 docstrings, #3594 non-finite RAT guard, #3600 color_ramp streaming accumulation = Pass 16's own fix, #3588/#3589 isort): all correctness/docs, no new performance surface, no regressions. Register introspection unavailable in this numba's _Kernel API; kernels unchanged since Pass 7 (inflate 67 regs / LZW 29 regs). Prior deferred LOWs unchanged (redundant .copy() in _writer predictor encode, identity pack arithmetic, twice-built IFD byte length). No HIGH/MEDIUM => no /rockout, no new benchmark required. SAFE/IO-bound holds. | Pass 16 (2026-07-01): 1 MEDIUM found and fixed, 0 HIGH. to_geotiff(dask, color_ramp=...) executed the source graph twice: streaming write computed every chunk, then _write_sidecars -> _finite_stats ran dask.compute over the same source for the PAM/QML statistics (measured 32 chunk executions for a 16-chunk source; color_ramp_range escape hatch stays at 16). Filed #3597, fixed by threading a chunk_observer through _write_streaming's three materialisation sites (row band, wide-raster segment, strip band) into a new StreamingStats accumulator (_symbology.py; Chan mean/M2 combine, float64 accumulators, ddof=0, nodata/finite exclusion matches _finite_stats; all-valid buffers skip the boolean-index copy and reductions use dtype=float64 on the original buffer to avoid an astype copy). Post-fix executions 16/16; round-trip bench 8192x8192 f32 deflate write+stats 1.53s -> 1.44s and total source reads halve (the real win is expensive-to-recompute sources: HTTP COGs, cold storage, long pipelines). GPU writer and VRT paths keep the documented second pass (GPU writer fully materialises anyway); docstrings updated. 13 new tests in test_color_ramp_single_pass_3597.py incl. execution counters (row-band, strip, segmented wide path with full-width chunks), nodata/int/all-NaN/multiband gating, accumulator-vs-_finite_stats parity, and a dask+cupy gpu=False streaming leg (run on-device). Audited all 26 geotiff commits since 2026-06-11: pack range guards #3272/#3277 correctly defer per-chunk scans into the write's single compute (no #3235 regression); #3374 fixed Pass 14's deferred LOW (chunked GPU read now parses header/IFDs once); xarray engine #3375/#3377/#3380 is a thin wrapper (no eager compute); PAM sidecar read on open_geotiff is a local os.path.exists probe (no HTTP cost); _CloudSource cat_file change is neutral-to-better. Probes on-device this pass: dask read 400 tasks/100 chunks (4/chunk, cap intact), eager GPU read cupy parity 0.0 (222ms), dask+GPU lazy with cupy meta, symbology _finite_stats GPU/CPU parity 0.0. LOWs unchanged from prior passes. SAFE/IO-bound holds. | Pass 15 (2026-06-11): 1 MEDIUM found and fixed. _pack (_attrs.py:~1795) guarded the no-sentinel integer restore with an eager bool(out.isnull().any()), which executed the whole upstream dask graph at to_geotiff(pack=True) call time; the streaming writer then executed it again, so every source chunk computed twice (measured 32 decode-task executions for 16 chunks on a 512x512 int16 SCALE/OFFSET no-GDAL_NODATA source; 71->33 total task starts post-fix). Filed #3235, fixed by mapping a per-chunk NaN guard (_pack_guard_no_nan) into the graph for dask-backed data (raises from the write's single compute; numpy keeps the eager call-time check; meta= preserves cupy backing). 9 new tests in test_pack_lazy_nan_guard_3235.py incl. fusion-proof execution counter and cupy-chunk guard unit test (dask+cupy e2e still blocked upstream by #3112). Scrutinised all 16 commits since 2026-06-08 (pack/unpack series #3065/#3075/#3079/#3129/#3174/#3175, VRT placement #3135, compression_level gate #3176, streaming banding #3136, dask+cupy writer order fix #3171): no other regressions; #3171's get-then-asarray order is intentional D2H for gpu=False. GPU validated on-device this pass: eager GPU unpack returns cupy with exact parity (387ms incl warmup, only 0-d scalar .get()s -- no bulk host round trip), dask+GPU unpack lazy (112 tasks/16 chunks, cupy meta, compute returns cupy, parity 0.0), GDS fast path intact without unpack (4 tasks/chunk); unpack disqualifying GDS is documented intentional. Dask CPU probe 4 tasks/chunk, 50k-task cap intact. Note: #1714 (_write_vrt_tiled synchronous scheduler) is now FIXED+CLOSED (scheduler='threads' at _writers/eager.py:1517) -- drop from the open-issue list. LOW noted (no fix): _pack does identity (data-0.0)/1.0 arithmetic allocating two full-array temporaries when scale==1/offset==0 (masked_nodata-only pack); prior deferred LOWs unchanged. SAFE/IO-bound holds. | Pass 14 (2026-06-09): MEDIUM found and fixed -- _write_streaming ran one dask .compute() per 256-row tile-row/strip, so a source chunk taller than the band re-executed once per band it overlapped (measured 2x at chunks=512, 4x at chunks=1024, whole upstream graph re-runs for computed pipelines). Filed #3117, fixed via _stream_row_bands: consecutive tile-rows/strips group into row bands sized by the source chunk-row span (one-chunk halo, #3007 accounting) under streaming_buffer_bytes; each band computes once and tiles/strips are carved from the materialised band. Wide rasters needing column segmentation keep the per-tile-row path. Post-fix per-chunk executions == 1 on the default read->write round trip. 5 new tests (TestRowBandRecompute3117 + _stream_row_bands unit); write/integration/parity suites pass (2195). LOW deferred (no fix): _read_geotiff_gpu_chunked parses header+all IFDs twice at graph build (_backends/gpu.py ~1367-1419, cap check then GDS probe; build-time only). GPU paths validated on-device this pass: eager gpu read returns cupy with parity, dask+GPU chunked read lazy (17 tasks/4 chunks) with parity; GPU writer full materialisation is documented intentional (streaming_buffer_bytes no-op). Read path keeps 50k-task graph cap; dask read probe 4 tasks/chunk. SAFE/IO-bound holds. | Pass 13 (2026-05-20): 1 MEDIUM found and fixed. _nvjpeg_batch_encode (_gpu_decode.py:~L1560) and _nvjpeg2k_batch_encode (~L2958) called cupy.cuda.Device().synchronize() inside the per-tile encode loops, a whole-device fence that blocked every CUDA stream and serialised concurrent work (e.g. predictor encodes on other streams). The decode-side counterpart _try_nvjpeg_batch_decode already used cupy.cuda.Stream.null.synchronize() at L1442; the encoder side was inconsistent. Filed #2212 and fixed both encoders to use Stream.null.synchronize(), scoping the per-tile sync to the default stream the encode/retrieve calls were issued on. nvJPEG / nvJPEG2000 encoders maintain a single shared state per encoder so encodes within a batch are inherently serial; the fix removes the device-wide blocker without changing the API ordering contract. 5 new tests in test_nvjpeg_encode_stream_sync_2212.py (AST checks that neither encoder contains Device().synchronize() inside a for-loop, that both call Stream.null.synchronize() in the loop, and that the decoder reference pattern stays pinned). All 5 new tests + 19 existing related encode/decode tests pass. nvjpeg/nvjpeg2k shared libs not present on this host so end-to-end encode verification is gated; add cuda-unavailable-libs note to re-validate on a host with the RAPIDS conda env. SAFE/IO-bound verdict holds; no change in dask graph cost. Dask probe: 2560x2560 deflate-tiled file via read_geotiff_dask(chunks=256) yields 400 tasks for 100 chunks (4 tasks/chunk), well under the 50K cap. LOW deferred (no fix in this PR): _build_ifd called twice per IFD level in _assemble_standard_layout (_writer.py:1531+1543), _assemble_cog_layout (1582+1625), and the COG overview path (2519+2546+2740) -- the first call's bytes are discarded; only the overflow byte length is used to compute pixel_data_offset. Cost is bounded by IFD count (typically 1-5 overview levels) so absolute impact is minor. Pre-existing pattern. | Pass 12 (2026-05-18): 1 MEDIUM found and fixed. _try_nvjpeg2k_batch_decode at _gpu_decode.py:~L2725-2778 allocated per-tile per-component cupy.empty buffers (N*S round-trips through the cupy memory pool) and called cupy.cuda.Device().synchronize() once per tile, forcing default-stream serialisation that defeats nvJPEG2000's internal pipelining. Filed #2107 and fixed: pre-allocate a single d_comp_pool sized n_tiles*samples*tile_height*pitch under a _check_gpu_memory guard, derive per-tile/per-component views as slab offsets, and replace the per-tile sync with a single batch-end sync. Same pattern as #1659 (_try_nvcomp_from_device_bufs), #1688 (_try_kvikio_read_tiles), #1712 (_nvcomp_batch_compress). 7 new tests in test_nvjpeg2k_single_alloc_2107.py: AST-level structural assertions confirm no cupy.empty inside the for-loop and no Device().synchronize() inside the loop, plus pool/per_tile_comp_bytes presence and _check_gpu_memory guard checks; lib-absent short-circuit; unsupported-dtype cleanup contract; cupy-only pool slab-non-overlap test (gpu-marked). libnvjpeg2k.so not present on this host so the end-to-end nvJPEG2000 decode is gated -- note added to re-validate on a host with the RAPIDS conda env. All 30 jpeg2000/compression tests + 7 new tests pass. SAFE/IO-bound verdict holds (no change in dask graph cost). Dask probe: 4096x4096 deflate-tiled file via read_geotiff_dask(chunks=512) yields 256 tasks for 64 chunks (4 tasks/chunk), well under the 50K cap. | Pass 11 (2026-05-18): 1 MEDIUM found and fixed. _read_strips (_reader.py:~L1972) and _fetch_decode_cog_http_strips (_reader.py:~L2670) decoded strips sequentially in a Python for-loop while the tile counterparts (_read_tiles L2146, _fetch_decode_cog_http_tiles L2898) gated parallel decode on _PARALLEL_DECODE_PIXEL_THRESHOLD via ThreadPoolExecutor. Filed #2100 and fixed: both strip paths now collect jobs, parallel-decode when n_strips > 1 and strip_pixels >= 64K, then place sequentially. Measured (uint16, 4-core): 4096x4096 deflate 130ms->34ms (3.82x), 8192x8192 deflate 531ms->146ms (3.63x), 8192x8192 zstd 211ms->85ms (2.48x), uncompressed 25ms->22ms (1.14x). 5 new tests in test_parallel_strip_decode_2100.py (parallel/serial parity, pool-engaged on multi-strip, serial-path for single-strip, windowed cross-strip read, HTTP COG strip parity). 3998 tests pass; 8 pre-existing failures predating this change (predictor2 BE + size_param_validation_gpu_vrt reference now-private read_to_array attr). SAFE/IO-bound verdict holds. | Pass 10 (2026-05-15): 1 new MEDIUM found and fixed; 2 LOW noted. MEDIUM (_reader.py:2737): _fetch_decode_cog_http_tiles decoded tiles sequentially in a Python for-loop after the concurrent fetch landed (issue #1480). Local _read_tiles parallelises decode whenever tile_pixels >= 64K via ThreadPoolExecutor (_reader.py:2017); the HTTP path was structurally similar but never picked up the same gate, so wide windowed reads of multi-tile COGs left deflate/zstd decode single-threaded. Mirrored the local-path threshold + pool. 5 new tests in test_cog_http_parallel_decode_2026_05_15.py (parallel + serial round-trip correctness, pool-instantiation branch selection above the threshold, single-tile path skips the pool, structural _decode_strip_or_tile call count == n_tiles). All 262 COG/HTTP tests pass; 3162 of 3164 selected geotiff tests pass overall (2 pre-existing failures predating Pass 9 per prior notes -- test_predictor2_big_endian_gpu_1517 references the now-private read_to_array attr, and the test_size_param_validation_gpu_vrt_1776 tile_size=4 validator failure). LOW deferred (no fix in this PR): (1) _block_reduce_2d_gpu (_gpu_decode.py:3142/3163/3189) does bool(mask.any().item()) per overview level when nodata is set, paying one device sync per level; the alternative (unconditional cupy.putmask) always pays the work cost and the short-circuit is correct under the current API. (2) _nvcomp_batch_compress adler32 staging (_gpu_decode.py:2543-2546) issues n_tiles slice-assign kernels into a fresh contig buffer despite all callers passing slices of a single underlying d_tile_buf; an API refactor to accept the source buffer directly would skip the rebuild. SAFE/IO-bound verdict holds. Dask probe: 2560x2560 chunks=256 yields 400 tasks (4 per chunk), well under the 50000 cap. GPU probe: 1024x1024 float32 zstd read returns CuPy-backed in 236 ms with no host round-trip. | Rockout 2026-05-15: LOW filed #1934 -- _apply_nodata_mask_gpu used cupy.where (allocating); switched to cupy.putmask on the already-owned buffer (float path) and on the post-astype float64 buffer (int path). Saves one chunk-sized device allocation per call. 7 new tests in test_apply_nodata_mask_gpu_inplace_1934.py; 52 related nodata tests pass. | Pass 8 (2026-05-12): 1 new MEDIUM found and fixed. _assemble_standard_layout/_assemble_cog_layout returned bytes(bytearray), doubling peak memory transiently during eager writes. Filed #1756, fixed by returning the bytearray directly. Measured: 95 MB uint8 raster peak drops 202 MB -> 107 MB. _write_bytes / parse_header already accepted the buffer protocol so the change is transparent to callers. 6 new tests in test_assemble_layout_no_bytes_copy_1756.py. 2123 existing geotiff tests pass; the 10 unrelated failures (test_no_georef_windowed_coords_1710, test_predictor2_big_endian_gpu_1517) reference the now-private read_to_array attribute (commit 8adb749, issue #1708) and predate this change. SAFE/IO-bound verdict holds. | Pass 7 (2026-05-12): re-audit identified 4 MEDIUM findings, all real, all backed by microbenches. (1) unpack_bits sub-byte loops for bps=2/4/12 in _compression.py:836-878 were 100-200x slower than vectorised numpy (filed #1713, fixed in this branch: bps=4 2M pixels drops from 165ms to 3ms = 55x; bps=2/12 similar). (2) _write_vrt_tiled at __init__.py:1708 uses scheduler='synchronous' on independent tile writes; measured 33% slowdown on 256-tile zstd write vs threads scheduler (filed #1714, no fix yet). (3) _nvcomp_batch_compress at _gpu_decode.py:2522-2526 still does per-tile cupy.get().tobytes() despite #1552 / #1659 fixing the same pattern elsewhere; measured 45% reduction with concat+single get on n=1024 (filed #1712, no fix yet). (4) _nvcomp_batch_compress at _gpu_decode.py:2457 uses per-tile cupy.empty allocations; 1024 tiles 16KB drops from 4.7ms to 1.0ms with single contiguous + views (bundled into #1712). Cat 6 OOM verdict: SAFE/IO-bound holds -- read_geotiff_dask caps task count at _MAX_DASK_CHUNKS=50_000 and per-chunk memory is bounded by chunk size. _inflate_tiles_kernel resource usage on Ampere: 67 regs/thread, 2896B local/thread, 8192B shared/block (LZW kernel: 29 regs, 24576B shared) -- register pressure under control; high local memory in inflate is unavoidable (LZ77 state) but only thread 0 in each block uses it. | Pass 4 (2026-05-10): re-audit after #1559 (centralise attrs across all read backends). New _populate_attrs_from_geo_info helper at __init__.py:301 runs once per read, not per-chunk -- no perf impact. Probe: 2560x2560 deflate-tiled file opened via read_geotiff_dask yields 400 tasks (4 tasks/chunk for 100 chunks), well under 1M cap. read_geotiff_gpu(1024x1024) returns cupy.ndarray end-to-end with no host round-trip (226ms incl. write+decode). No new HIGH/MEDIUM findings. SAFE/IO-bound holds. | Pass 3 (2026-05-10): SAFE/IO-bound. Audited 4 perf commits: #1558 (in-place NaN writes on uniquely-owned buffers correct), #1556 (fp-predictor ngjit ~297us/tile for 256x256 float32), #1552 (single cupy.concatenate + one .get() for batched D2H at _gpu_decode.py:870-913), #1551 (parallel decode threshold >=65536px engages 256x256 default at _reader.py:1121). Bench: 8192x8192 f32 deflate+pred2 256-tile write 782ms; 4096x4096 f32 deflate read 83ms with parallel decode. Deferred LOW (none filed, all <10% MEDIUM threshold): _writer.py:459/1109 redundant .copy() before predictor encode (~1% per tile), _compression.py:280 lzw_decompress dst[:n].copy() (~2% per LZW tile decode), _writer.py:1419 seg_np.copy() before in-place NaN substitution (negligible, conditional path), _CloudSource.read_range opens fresh fsspec handle per range (pre-existing, predates audit scope). nvCOMP per-tile D2H batching break-even confirmed (variable sizes need staging buffer, no win). | Pass 3 (2026-05-10): audited f157746,39322c3,f23ec8f,1aac3b7. All 5 commits correct. Redundant .copy() in _writer.py:459,1109 and _compression.py:280 (1-2% overhead, LOW). _CloudSource.read_range() per-call open is pre-existing arch issue. No HIGH/MEDIUM regressions. SAFE. | re-audit 2026-05-02: 6 commits since 2026-04-16 (predictor=3 CPU encode/decode, GPU predictor stride fix, validate_tile_layout, BigTIFF LONG8 offsets, AREA_OR_POINT VRT, per-tile alloc guard). 1M dask chunk cap intact at __init__.py:948; adler32 batch transfer intact at _gpu_decode.py:1825. New code is metadata validation and dispatcher logic with no extra materialization or per-tile sync points. No HIGH/MEDIUM regressions. | Pass 5 (2026-05-12): re-audit identified MEDIUM in _gpu_decode.py:1577 _try_nvcomp_from_device_bufs: per-tile cupy.empty + trailing cupy.concatenate doubled peak VRAM and added serial concat. Filed #1659 and fixed to single-buffer + pointer offsets (matches LZW/deflate/host-buffer patterns at L1847/L1878/L1114). Microbench (alloc+concat overhead only, not full nvCOMP latency): n=256 tile_bytes=65536 drops 3.66ms->0.69ms, n=256 tile_bytes=262144 drops 8.18ms->0.13ms. Tests: 5 new tests in test_nvcomp_from_device_bufs_single_alloc_1659.py (codec short-circuit, no-lib short-circuit, memory-guard contract, real ZSTD round-trip via nvCOMP, structural single-buffer check). 1458 existing geotiff tests pass, 3 unrelated matplotlib/py3.14 failures pre-existing. SAFE/IO-bound verdict holds. | Pass 6 (2026-05-12): re-audit on top of #1659. New HIGH in _try_kvikio_read_tiles at _gpu_decode.py:941: per-tile cupy.empty() + blocking IOFuture.get() inside loop serialised GDS reads to ~1 outstanding pread, missed parallelism the kvikio worker pool was designed for, paid per-tile cupy.empty setup (matches #1659 anti-pattern in nvCOMP path), and lacked _check_gpu_memory guard. Filed #1688 and fixed to single contiguous buffer + batched submit + guard. Microbench with 8-worker pool simulation: 256 tiles@1ms latency drops 256ms->38.7ms (~6.6x); single-thread simulation 256ms->28.5ms (9x). Tests: 9 new tests in test_kvikio_batched_pread_1688.py (kvikio-absent path, single-buffer pointer arithmetic, submit-before-get ordering, memory guard, partial-read fallback, round-trip data, zero-size/all-sparse tiles). All 1577 geotiff tests pass except pre-existing matplotlib/py3.14 failures." glcm,2026-03-31T18:00:00Z,SAFE,compute-bound,0,,"Downgraded to MEDIUM. da.stack without rechunk is scheduling overhead, not OOM risk." hillshade,2026-04-16T12:00:00Z,SAFE,compute-bound,0,,"Re-audit after Horn's method rewrite (PR 1175): clean stencil, map_overlap depth=(1,1), no materialization. Zero findings." -hydro,2026-05-01,RISKY,memory-bound,0,1416,"Fixed-in-tree 2026-05-01: hand_mfd._hand_mfd_dask now assembles via da.map_blocks instead of eager da.block of pre-computed tiles (matches hand_dinf pattern). Remaining MEDIUM: sink_d8 CCL fully materializes labels (inherently global), flow_accumulation_mfd frac_bdry held in driver dict instead of memmap-backed BoundaryStore. D8 iterative paths (flow_accum/fill/watershed/basin/stream_*) use serial-tile sweep with memmap-backed boundary store -- per-tile RAM bounded but driver iterates O(diameter) times. flow_direction_*, flow_path/snap_pour_point/twi/hand_d8/hand_dinf are SAFE." +hydro,2026-07-23,RISKY,memory-bound,1,3692,"Fixed-in-tree 2026-07-23 (issue 3692): twi_d8 dask+cupy dropped per-chunk GPU->CPU->GPU round-trip (HIGH, 17.7x threaded / 62x sync on 8192^2; now runs _twi_dask natively, asv TWI benchmark gained cupy+dask+cupy params); watershed d8/dinf/mfd numpy+cupy init loops vectorized (MEDIUM, loop was ~90% of numpy runtime, 6-7x end-to-end). Known remaining MEDIUMs unchanged: sink_d8 CCL fully materializes labels (inherently global), flow_accumulation_mfd frac_bdry held in driver dict instead of memmap BoundaryStore, D8 iterative paths use serial-tile sweep with memmap boundary store (per-tile RAM bounded, driver iterates O(diameter)). hand_mfd lazy map_blocks assembly from 2026-05-01 confirmed intact. flow_direction_*, flow_path/snap_pour_point/hand_d8/hand_dinf remain SAFE; graph probes: flow_direction_d8 ~20 tasks/chunk (map_overlap), twi 8/chunk. _detect_flow_type per-block compute loop is test-only, dead in production." interpolate,2026-06-12,SAFE,compute-bound,0,3298,"3 MEDIUM fixed via #3298: kriging no-variance path now uses dual form k0 @ (K_inv @ z_aug) (1.6x, drops (P,N+1) w temp); dask variance computed in one map_blocks pass (was 2.06x); dask+cupy chunk-invariant uploads (idw x/y/z, kriging x/y/z/K_inv) hoisted and cKDTree built once for dask k-nearest. 1 LOW documented, not fixed: _experimental_variogram bins pairs with per-lag boolean masks, O(nlags*N^2) passes where np.bincount would do one. Dask graphs are plain map_blocks, 2 tasks/chunk, no fan-in; memory guards cover host allocations. GPU paths executed on this host (CUDA available)." interpolate-kriging,2026-06-04,SAFE,graph-bound,0,2923,"MEDIUM: memory guard used full-grid k0 term on dask templates -> spurious MemoryError (issue #2923, fixed). LOW: _experimental_variogram nlags python loop vectorizable via bincount (~1.2x, pair-array materialization dominates) - doc only. Dask graph clean (2 tasks/chunk); cupy returns device arrays; no .values/.compute/.data.get materialization." interpolate_spline,2026-06-04,SAFE,compute-bound,0,,"scope=spline-only. Audited _spline.py + _validation.py only (not _idw/_kriging). 1 MEDIUM (Cat3 GPU transfer): _spline_dask_cupy/_spline_cupy re-uploaded invariant x_pts/y_pts/weights host->device once per chunk. Fixed in PR #2929: added _tps_evaluate_gpu taking on-device point/weight arrays + only per-chunk grid slices; dask+cupy uploads invariants once at graph build (verified 48->3 on 16 chunks, scales with chunk count). numpy/cupy/dask+cupy parity ~1e-14. Added cupy+dask+cupy parity tests and an upload-count regression test (red without fix: 48!=3). _tps_cuda_kernel 30 regs/thread, 6 scalar locals -- no register pressure. CPU/dask+numpy eval @ngjit, row-major, no materialization. Dask graph probe 2560x2560/256 chunks = 200 tasks (2/chunk), no fan-in. Memory guard _check_spline_memory bounds N^2 solve. No issue filed -- gh issue create denied by auto-mode classifier; finding surfaced directly by sweep. GitHub issue field left empty." From 38bf988b3f9c3bc45b2d80bd3d8b5f230c1eab2c Mon Sep 17 00:00:00 2001 From: Brendan Collins Date: Thu, 23 Jul 2026 14:14:43 -0400 Subject: [PATCH 3/3] Address review: collapse redundant twi dask+cupy branch, note init mask overhead (#3692) --- xrspatial/hydro/twi_d8.py | 12 +++++------- xrspatial/hydro/watershed_d8.py | 4 +++- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/xrspatial/hydro/twi_d8.py b/xrspatial/hydro/twi_d8.py index 4a667c27f..3d9b1433b 100644 --- a/xrspatial/hydro/twi_d8.py +++ b/xrspatial/hydro/twi_d8.py @@ -18,7 +18,7 @@ da = None from xrspatial.utils import (_validate_raster, get_dataarray_resolution, has_cuda_and_cupy, - is_cupy_array, is_dask_cupy) + is_cupy_array) # Minimum tan(slope) clamp: tan(0.001°) _TAN_MIN = np.tan(np.radians(0.001)) @@ -64,15 +64,13 @@ def twi_d8(flow_accum: xr.DataArray, fa_data = flow_accum.data sl_data = slope_agg.data - if has_cuda_and_cupy() and is_dask_cupy(flow_accum): - # TWI is purely elementwise, so cupy-backed dask arrays run the - # same expression natively; no host round-trip needed. - out = _twi_dask(fa_data, sl_data, cellsize) - - elif has_cuda_and_cupy() and is_cupy_array(fa_data): + if has_cuda_and_cupy() and is_cupy_array(fa_data): out = _twi_cupy(fa_data, sl_data, cellsize) elif da is not None and isinstance(fa_data, da.Array): + # Covers numpy- and cupy-backed dask arrays alike: TWI is purely + # elementwise, so cupy chunks run the same expression natively + # with no host round-trip. out = _twi_dask(fa_data, sl_data, cellsize) elif isinstance(fa_data, np.ndarray): diff --git a/xrspatial/hydro/watershed_d8.py b/xrspatial/hydro/watershed_d8.py index 70667e589..3a051791b 100644 --- a/xrspatial/hydro/watershed_d8.py +++ b/xrspatial/hydro/watershed_d8.py @@ -54,7 +54,9 @@ def _to_numpy_f64(arr): # state (int8) -> 1 # path_r (int64) -> 8 # path_c (int64) -> 8 -# Total ~33 bytes/pixel. The caller's ``flow_dir`` and ``pour_points`` +# Total ~33 bytes/pixel. The vectorized init also holds two boolean +# masks plus their conjunction (~3 B/px) transiently; that fits in the +# 50% guard headroom. The caller's ``flow_dir`` and ``pour_points`` # arrays already live in RAM before dispatch and are not double-counted. _BYTES_PER_PIXEL = 33