Chunked dask+cupy reproject without full-source eager compute (#1045)

Replaces the eager .compute() approach with a chunked GPU pipeline
that fetches only the needed source window per output chunk. This
handles sources larger than GPU memory while still being 8-20x
faster than the old dask.delayed path.

The key optimizations vs dask.delayed:
- CRS objects and transformer created once (not per chunk)
- CUDA projection + native CUDA resampling per chunk
- Default 2048x2048 GPU chunks (not 512x512)
- Sequential loop avoids dask scheduler overhead

Performance (4096x4096 WGS84 -> UTM, bilinear):
  CuPy single pass:     34ms
  Dask+CuPy (2048):     49ms  (was 958ms)
  Dask+CuPy (512):      71ms
  Dask+CuPy (256):     124ms

All chunk sizes are pixel-exact vs plain CuPy (max_err < 1e-11).

Author: brendancol

xrspatial/reproject/__init__.py

@@ -444,24 +444,12 @@ def reproject(
     tgt_wkt = tgt_crs.to_wkt()
 
     if is_dask and is_cupy:
-        # Dask+CuPy: eagerly compute source to GPU, then single-pass
-        # CuPy reproject.  This avoids per-chunk overhead (pyproj init,
-        # small CUDA kernel launches, dask scheduler) that makes chunked
-        # GPU reproject ~28x slower than a single pass.  The output is
-        # returned as a plain CuPy array; caller can .rechunk() if needed.
-        import cupy as _cp
-        eager_data = raster.data.compute()
-        if not isinstance(eager_data, _cp.ndarray):
-            eager_data = _cp.asarray(eager_data)
-        eager_da = xr.DataArray(
-            eager_data, dims=raster.dims,
-            coords=raster.coords, attrs=raster.attrs,
-        )
-        result_data = _reproject_inmemory_cupy(
-            eager_da, src_bounds, src_shape, y_desc,
+        result_data = _reproject_dask_cupy(
+            raster, src_bounds, src_shape, y_desc,
             src_wkt, tgt_wkt,
             out_bounds, out_shape,
             resampling, nd, transform_precision,
+            chunk_size,
         )
     elif is_dask:
         result_data = _reproject_dask(
@@ -533,14 +521,153 @@ def _reproject_inmemory_cupy(
     )
 
 
+def _reproject_dask_cupy(
+    raster, src_bounds, src_shape, y_desc,
+    src_wkt, tgt_wkt,
+    out_bounds, out_shape,
+    resampling, nodata, precision,
+    chunk_size,
+):
+    """Dask+CuPy backend: process output chunks on GPU sequentially.
+
+    Instead of dask.delayed per chunk (which has ~15ms overhead each from
+    pyproj init + small CUDA launches), we:
+    1. Create CRS/transformer objects once
+    2. Use GPU-sized output chunks (2048x2048 by default)
+    3. For each output chunk, compute CUDA coordinates and fetch only
+       the source window needed from the dask array
+    4. Assemble the result as a CuPy array
+
+    For sources that fit in GPU memory, this is ~22x faster than the
+    dask.delayed path.  For sources that don't fit, each chunk fetches
+    only its required window, so GPU memory usage scales with chunk size,
+    not source size.
+    """
+    import cupy as cp
+
+    from ._crs_utils import _require_pyproj
+
+    pyproj = _require_pyproj()
+    src_crs = pyproj.CRS.from_wkt(src_wkt)
+    tgt_crs = pyproj.CRS.from_wkt(tgt_wkt)
+
+    # Use larger chunks for GPU to amortize kernel launch overhead
+    gpu_chunk = chunk_size or 2048
+    if isinstance(gpu_chunk, int):
+        gpu_chunk = (gpu_chunk, gpu_chunk)
+
+    row_chunks, col_chunks = _compute_chunk_layout(out_shape, gpu_chunk)
+    out_h, out_w = out_shape
+    src_left, src_bottom, src_right, src_top = src_bounds
+    src_h, src_w = src_shape
+    src_res_x = (src_right - src_left) / src_w
+    src_res_y = (src_top - src_bottom) / src_h
+
+    result = cp.full(out_shape, nodata, dtype=cp.float64)
+
+    row_offset = 0
+    for i, rchunk in enumerate(row_chunks):
+        col_offset = 0
+        for j, cchunk in enumerate(col_chunks):
+            cb = _chunk_bounds(
+                out_bounds, out_shape,
+                row_offset, row_offset + rchunk,
+                col_offset, col_offset + cchunk,
+            )
+            chunk_shape = (rchunk, cchunk)
+
+            # CUDA coordinate transform (reuses cached CRS objects)
+            try:
+                from ._projections_cuda import try_cuda_transform
+                cuda_coords = try_cuda_transform(
+                    src_crs, tgt_crs, cb, chunk_shape,
+                )
+            except Exception:
+                cuda_coords = None
+
+            if cuda_coords is not None:
+                src_y, src_x = cuda_coords
+                src_col_px = (src_x - src_left) / src_res_x - 0.5
+                if y_desc:
+                    src_row_px = (src_top - src_y) / src_res_y - 0.5
+                else:
+                    src_row_px = (src_y - src_bottom) / src_res_y - 0.5
+
+                r_min = int(cp.floor(cp.nanmin(src_row_px)).get()) - 2
+                r_max = int(cp.ceil(cp.nanmax(src_row_px)).get()) + 3
+                c_min = int(cp.floor(cp.nanmin(src_col_px)).get()) - 2
+                c_max = int(cp.ceil(cp.nanmax(src_col_px)).get()) + 3
+            else:
+                # CPU fallback for this chunk
+                transformer = pyproj.Transformer.from_crs(
+                    tgt_crs, src_crs, always_xy=True
+                )
+                src_y, src_x = _transform_coords(
+                    transformer, cb, chunk_shape, precision,
+                    src_crs=src_crs, tgt_crs=tgt_crs,
+                )
+                src_col_px = (src_x - src_left) / src_res_x - 0.5
+                if y_desc:
+                    src_row_px = (src_top - src_y) / src_res_y - 0.5
+                else:
+                    src_row_px = (src_y - src_bottom) / src_res_y - 0.5
+                r_min = int(np.floor(np.nanmin(src_row_px))) - 2
+                r_max = int(np.ceil(np.nanmax(src_row_px))) + 3
+                c_min = int(np.floor(np.nanmin(src_col_px))) - 2
+                c_max = int(np.ceil(np.nanmax(src_col_px))) + 3
+
+            # Check overlap
+            if r_min >= src_h or r_max <= 0 or c_min >= src_w or c_max <= 0:
+                col_offset += cchunk
+                continue
+
+            r_min_clip = max(0, r_min)
+            r_max_clip = min(src_h, r_max)
+            c_min_clip = max(0, c_min)
+            c_max_clip = min(src_w, c_max)
+
+            # Fetch only the needed source window from dask
+            window = raster.data[r_min_clip:r_max_clip, c_min_clip:c_max_clip]
+            if hasattr(window, 'compute'):
+                window = window.compute()
+            if not isinstance(window, cp.ndarray):
+                window = cp.asarray(window)
+            window = window.astype(cp.float64)
+
+            if not np.isnan(nodata):
+                window = window.copy()
+                window[window == nodata] = cp.nan
+
+            local_row = src_row_px - r_min_clip
+            local_col = src_col_px - c_min_clip
+
+            if cuda_coords is not None:
+                chunk_data = _resample_cupy_native(
+                    window, local_row, local_col,
+                    resampling=resampling, nodata=nodata,
+                )
+            else:
+                chunk_data = _resample_cupy(
+                    window, local_row, local_col,
+                    resampling=resampling, nodata=nodata,
+                )
+
+            result[row_offset:row_offset + rchunk,
+                   col_offset:col_offset + cchunk] = chunk_data
+            col_offset += cchunk
+        row_offset += rchunk
+
+    return result
+
+
 def _reproject_dask(
     raster, src_bounds, src_shape, y_desc,
     src_wkt, tgt_wkt,
     out_bounds, out_shape,
     resampling, nodata, precision,
     chunk_size, is_cupy,
 ):
-    """Dask backend: build output as ``da.block`` of delayed chunks."""
+    """Dask+NumPy backend: build output as ``da.block`` of delayed chunks."""
     import dask
     import dask.array as da