hydro: twi_d8 dask+cupy round-trips chunks through host; watershed init loops run in pure Python

[brendancol]

Describe the bug

Two performance problems in the hydro subpackage, found during a performance audit on 2026-07-23.

1. twi_d8 dask+cupy path round-trips every chunk through the host

_twi_dask_cupy (xrspatial/hydro/twi_d8.py:116-130) pulls each chunk to the host with b.get(), runs the elementwise math (tan, where, log) on numpy, then copies the result back to the device with cp.asarray. TWI has no neighborhood or graph structure, so the same expression runs natively on cupy-backed dask arrays via _twi_dask. The host detour looks copied from the D8 iterative drivers, which actually need it (their kernels are numba CPU code). TWI does not.

Measured on this host (8192x8192 float64, 2048-square chunks, warm runs, best of 5):

scheduler	current round-trip	direct dask+cupy	ratio
threads	0.672 s	0.038 s	17.7x
synchronous	2.170 s	0.035 s	62x

Max abs difference between the two paths: 1.8e-15.

2. Watershed family initializes labels/state with plain-Python full-grid loops

The numpy paths of watershed_d8, watershed_dinf, and watershed_mfd (and the cupy fallback paths of the dinf/mfd variants) build the labels/state arrays with nested Python loops over every cell before handing off to the jitted trace kernel:

• xrspatial/hydro/watershed_d8.py:1048-1056 (numpy path)
• xrspatial/hydro/watershed_dinf.py:704-712 (numpy) and 224-232 (cupy)
• xrspatial/hydro/watershed_mfd.py:706-714 (numpy) and 232-240 (cupy)

On a 2000x2000 grid the init loop takes 0.95 s while the jitted trace kernel takes 0.069 s, so roughly 90% of the numpy-backend runtime is Python loop overhead. A vectorized np.where init produces identical labels/state arrays (checked with np.array_equal, NaN positions included) and runs 17x faster, which makes the whole call about 6x faster.

Expected behavior

twi(..., routing='d8') on dask+cupy input keeps chunks on the device. Watershed init is vectorized so the trace kernel dominates the runtime.

Additional context

Found by /sweep-performance (hydro pass). The fix should extend the asv TWI benchmark with a dask+cupy case so a reintroduced round-trip shows up in benchmarks. The existing Watershed numpy benchmark already covers the init path.

Environment: Linux (WSL2), CUDA available, cupy installed.