fb892ea9f2
Comments, docstrings, logs, CLI help, map UI, legends, PDF sheet, scripts, compose files and AGENTS.md are now English. Data keys stay unchanged (relief_oriente, densite_sol, visualisations/, API JSON keys, link params). Wrong comments and help defaults found along the way are corrected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
615 lines
23 KiB
Python
615 lines
23 KiB
Python
"""GPU acceleration helpers for LiDAR pipeline.
|
||
|
||
Auto-selects the best NVIDIA GPU (highest compute capability first).
|
||
Uses CuPy Device API (not CUDA_VISIBLE_DEVICES) so JIT compilation
|
||
works correctly for any architecture (sm_89, sm_120, etc.).
|
||
|
||
In the worker pool, each process is pinned to one GPU (or forced to CPU)
|
||
by gpu_worker_slots / set_active_gpu / force_cpu, bounded by free VRAM.
|
||
"""
|
||
|
||
import logging
|
||
import os
|
||
import numpy as np
|
||
from scipy import ndimage
|
||
|
||
logger = logging.getLogger("lidar")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GPU auto-detection via nvidia-smi (no CUDA context created)
|
||
# ---------------------------------------------------------------------------
|
||
_NUM_GPUS = 0
|
||
HAS_GPU = False
|
||
_gpu_name = None
|
||
_gpu_mem_gb = 0
|
||
_best_gpu_id: int | None = None
|
||
_gpu_reason = None
|
||
|
||
# GPU restriction from -g flag (host-level indices)
|
||
_restricted_gpu_ids: list[int] | None = None
|
||
|
||
# True if CUDA_VISIBLE_DEVICES was written by _init_gpu itself (best-GPU
|
||
# pick). That value MUST NOT be treated as an external restriction: otherwise
|
||
# available_gpu_ids() returns only the chosen GPU and every worker gets the
|
||
# same gpu_id (GPU 1 never used).
|
||
_env_set_by_init = False
|
||
|
||
# Discovered GPU candidates (populated by _pick_gpu)
|
||
_gpu_candidates: list = []
|
||
|
||
|
||
def _pick_gpu() -> list:
|
||
"""List all GPUs from the system, sorted by compute capability (highest first)."""
|
||
try:
|
||
import subprocess
|
||
result = subprocess.run(
|
||
['nvidia-smi', '--query-gpu=index,name,compute_cap,memory.total',
|
||
'--format=csv,noheader,nounits'],
|
||
capture_output=True, text=True, timeout=15,
|
||
)
|
||
if result.returncode != 0:
|
||
return []
|
||
|
||
global _NUM_GPUS
|
||
|
||
gpus = []
|
||
for line in result.stdout.strip().split('\n'):
|
||
parts = [p.strip() for p in line.split(',')]
|
||
if len(parts) < 4:
|
||
continue
|
||
idx = int(parts[0])
|
||
name = parts[1]
|
||
cap_str = parts[2]
|
||
mem_mi = int(parts[3])
|
||
major, minor = (int(x) for x in cap_str.split('.'))
|
||
score = major * 1000 + minor * 100 + mem_mi
|
||
gpus.append((idx, name, cap_str, mem_mi, score, major))
|
||
|
||
_NUM_GPUS = len(gpus)
|
||
gpus.sort(key=lambda g: g[4], reverse=True)
|
||
return gpus
|
||
|
||
except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
|
||
return []
|
||
|
||
|
||
try:
|
||
_gpu_candidates = _pick_gpu() or []
|
||
except Exception:
|
||
pass
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Lazy CuPy initialization — tries each GPU until one works
|
||
# ---------------------------------------------------------------------------
|
||
_cp = None
|
||
_cp_ndimage = None
|
||
_cupy_ndarray = None # type ref that survives disable_gpu() so to_cpu() can still recover orphaned arrays
|
||
_gpu_initialized = False
|
||
|
||
|
||
def _filter_candidates(gpus: list) -> list:
|
||
"""Filter GPU candidates by CUDA_VISIBLE_DEVICES and _restricted_gpu_ids.
|
||
|
||
nvidia-smi lists ALL GPUs even when CUDA_VISIBLE_DEVICES is set
|
||
(driver 580.x behavior), so we must filter manually.
|
||
The value written by _init_gpu (automatic best-GPU pick) is ignored:
|
||
only external restrictions (run.sh -g, compose) count.
|
||
"""
|
||
# Filter by CUDA_VISIBLE_DEVICES if set externally
|
||
cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES')
|
||
if cuda_visible is not None and not _env_set_by_init:
|
||
try:
|
||
visible = {int(i.strip()) for i in cuda_visible.split(',')}
|
||
gpus = [g for g in gpus if g[0] in visible]
|
||
except ValueError:
|
||
pass
|
||
|
||
# Filter by programmatic restriction (-g 0, -g 0,2)
|
||
if _restricted_gpu_ids is not None:
|
||
allowed = set(_restricted_gpu_ids)
|
||
gpus = [g for g in gpus if g[0] in allowed]
|
||
|
||
return gpus
|
||
|
||
|
||
def _runtime_candidates():
|
||
"""Fallback GPU detection through the CuPy runtime (no nvidia-smi).
|
||
|
||
nvidia-smi can exceed its timeout on a loaded system: _pick_gpu() then
|
||
returns [] and workers wrongly fall back to CPU. CuPy indices are
|
||
renumbered according to CUDA_VISIBLE_DEVICES, so they are mapped back to
|
||
host indices to stay compatible with _filter_candidates.
|
||
"""
|
||
try:
|
||
import cupy as _cp_runtime
|
||
n = _cp_runtime.cuda.runtime.getDeviceCount()
|
||
cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES')
|
||
try:
|
||
host_ids = [int(v.strip()) for v in cuda_visible.split(',')]
|
||
except (ValueError, AttributeError):
|
||
host_ids = list(range(n))
|
||
gpus = []
|
||
for i in range(min(n, len(host_ids))):
|
||
props = _cp_runtime.cuda.runtime.getDeviceProperties(i)
|
||
name = props.get('name', b'?')
|
||
if isinstance(name, bytes):
|
||
name = name.decode()
|
||
major = int(props.get('major', 0))
|
||
minor = int(props.get('minor', 0))
|
||
mem_mi = int(props.get('totalGlobalMem', 0)) // (1024 * 1024)
|
||
cap = f"{major}.{minor}"
|
||
score = major * 1000 + minor * 100 + mem_mi
|
||
gpus.append((host_ids[i], name, cap, mem_mi, score, major))
|
||
gpus.sort(key=lambda g: g[4], reverse=True)
|
||
return gpus
|
||
except Exception:
|
||
return []
|
||
|
||
|
||
def _init_gpu():
|
||
"""Lazily initialize CuPy on first GPU use.
|
||
|
||
1. Filters candidates by CUDA_VISIBLE_DEVICES and _restricted_gpu_ids
|
||
2. If CUDA_VISIBLE_DEVICES is already set (e.g. run.sh -g 0):
|
||
import CuPy directly (no subprocess test needed)
|
||
3. Otherwise: test each GPU in a subprocess, pick the first that works
|
||
"""
|
||
global _cp, _cp_ndimage, _cupy_ndarray, _gpu_initialized, HAS_GPU, _best_gpu_id, _gpu_name, _gpu_mem_gb
|
||
if _gpu_initialized:
|
||
return
|
||
_gpu_initialized = True
|
||
|
||
candidates = _filter_candidates(_gpu_candidates)
|
||
if not candidates:
|
||
# Runtime fallback: nvidia-smi failed (timeout on a loaded system)
|
||
candidates = _filter_candidates(_runtime_candidates())
|
||
if not candidates:
|
||
logger.info("No usable GPU — CPU-only mode")
|
||
_cp = None
|
||
_cp_ndimage = None
|
||
HAS_GPU = False
|
||
return
|
||
|
||
cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES')
|
||
|
||
if cuda_visible is not None:
|
||
# CUDA_VISIBLE_DEVICES already set (e.g. by run.sh -g 0).
|
||
# Pick the best visible GPU and import CuPy directly.
|
||
idx, name, cap_str, mem_mi, score, major = candidates[0]
|
||
try:
|
||
import cupy as _real_cupy
|
||
import cupyx.scipy.ndimage as _real_cupy_ndimage
|
||
|
||
# Warm-up kernel to verify GPU works
|
||
x = _real_cupy.array([1.0, 2.0], dtype=_real_cupy.float32)
|
||
s = _real_cupy.sum(x).get()
|
||
if s != 3.0:
|
||
raise RuntimeError("GPU warm-up failed")
|
||
|
||
_best_gpu_id = idx
|
||
_gpu_name = name
|
||
_gpu_mem_gb = mem_mi // 1024
|
||
HAS_GPU = True
|
||
_cp = _real_cupy
|
||
_cp_ndimage = _real_cupy_ndimage
|
||
_cupy_ndarray = _real_cupy.ndarray
|
||
return
|
||
except Exception as e:
|
||
logger.warning(f"GPU unavailable (CUDA_VISIBLE_DEVICES={cuda_visible}): {e}")
|
||
_cp = None
|
||
_cp_ndimage = None
|
||
HAS_GPU = False
|
||
return
|
||
|
||
# No CUDA_VISIBLE_DEVICES set — test each GPU in subprocess
|
||
import subprocess
|
||
_working_gpu = None
|
||
for idx, name, cap_str, mem_mi, score, major in candidates:
|
||
result = subprocess.run(
|
||
['python3', '-c',
|
||
'import cupy; a=cupy.array([1.0,2.0],dtype=cupy.float32); '
|
||
'dev=cupy.cuda.runtime.getDevice(); '
|
||
'print(f"OK:{dev}:{cupy.sum(a).get()}")'],
|
||
capture_output=True, text=True, timeout=120,
|
||
env={**os.environ, 'CUDA_VISIBLE_DEVICES': str(idx)},
|
||
)
|
||
stdout = result.stdout.strip()
|
||
if result.returncode == 0 and stdout.startswith('OK:') and '3.0' in stdout:
|
||
# Verify the device actually used is device 0 (the GPU we targeted)
|
||
parts = stdout.split(':')
|
||
if len(parts) >= 2 and parts[1] == '0':
|
||
_working_gpu = (idx, name, mem_mi)
|
||
break
|
||
else:
|
||
logger.warning(f"GPU {idx} ({name}, sm_{cap_str}): false positive, the test ran on another device")
|
||
logger.warning(f"GPU {idx} ({name}, sm_{cap_str}) not compatible: "
|
||
f"{result.stderr.strip().splitlines()[-1] if result.stderr else 'unknown'}")
|
||
|
||
if _working_gpu is None:
|
||
logger.info("No usable GPU — CPU-only mode")
|
||
_cp = None
|
||
_cp_ndimage = None
|
||
HAS_GPU = False
|
||
return
|
||
|
||
idx, name, mem_mi = _working_gpu
|
||
global _env_set_by_init
|
||
_env_set_by_init = True
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = str(idx)
|
||
|
||
import cupy as _real_cupy
|
||
import cupyx.scipy.ndimage as _real_cupy_ndimage
|
||
|
||
_best_gpu_id = idx
|
||
_gpu_name = name
|
||
_gpu_mem_gb = mem_mi // 1024
|
||
HAS_GPU = True
|
||
_cp = _real_cupy
|
||
_cp_ndimage = _real_cupy_ndimage
|
||
_cupy_ndarray = _real_cupy.ndarray
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public API
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def num_gpus():
|
||
"""Return the number of available GPUs (after restrict_gpus filtering)."""
|
||
return len(_filter_candidates(_gpu_candidates))
|
||
|
||
|
||
def available_gpu_ids():
|
||
"""Return list of host-level GPU indices available for processing.
|
||
|
||
Respects any prior restrict_gpus() call.
|
||
"""
|
||
return [c[0] for c in _filter_candidates(_gpu_candidates)]
|
||
|
||
|
||
def restrict_gpus(gpu_ids: list[int], set_env_var: bool = False):
|
||
"""Restrict GPU selection to specific host-level indices.
|
||
|
||
Stores the restriction to be applied during _init_gpu().
|
||
If set_env_var is True, also sets CUDA_VISIBLE_DEVICES immediately so
|
||
child processes inherit the restriction.
|
||
"""
|
||
global _restricted_gpu_ids
|
||
_restricted_gpu_ids = gpu_ids
|
||
if set_env_var and gpu_ids:
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(i) for i in gpu_ids)
|
||
|
||
|
||
def set_active_gpu(gpu_id):
|
||
"""Restrict to a single GPU by host-level index.
|
||
|
||
Called by spawned workers (_process_file_standalone): CuPy is not yet
|
||
initialized there (the worker does not run log_gpu_status), so
|
||
_init_gpu() must be triggered here, otherwise the worker silently falls
|
||
back to CPU. CUDA_VISIBLE_DEVICES is also narrowed to this single GPU
|
||
before init so that the worker's device 0 is the right one: without it,
|
||
all workers with several visible GPUs share the first of them.
|
||
"""
|
||
global _restricted_gpu_ids
|
||
_restricted_gpu_ids = [gpu_id]
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
|
||
_init_gpu()
|
||
|
||
|
||
def force_cpu():
|
||
"""Forbid the GPU for this process (worker beyond the VRAM budget).
|
||
|
||
No candidate survives the filter, and an empty CUDA_VISIBLE_DEVICES stops
|
||
CuPy from creating a context (~300 MB of VRAM per process otherwise).
|
||
"""
|
||
global _restricted_gpu_ids
|
||
_restricted_gpu_ids = []
|
||
os.environ['CUDA_VISIBLE_DEVICES'] = ''
|
||
_init_gpu()
|
||
|
||
|
||
# Peak VRAM of one worker (MiB): CUDA context (~300) + joint scan-line
|
||
# adjustment on CuPy (~70-100 bytes per ground point in float64, ~15 M points
|
||
# per 0.2 m tile + 100 m buffer) + _fill_nans distance transform (int32
|
||
# indices x2 over ~7000² px ≈ 400 MB); the CuPy pool keeps its blocks between
|
||
# steps. Estimated from the code, tune with LIDAR_GPU_WORKER_MIB.
|
||
GPU_WORKER_MIB = int(os.environ.get("LIDAR_GPU_WORKER_MIB", "2048") or 2048)
|
||
# VRAM left free on each GPU (display, other processes).
|
||
GPU_RESERVE_MIB = int(os.environ.get("LIDAR_GPU_RESERVE_MIB", "512") or 512)
|
||
|
||
|
||
def gpu_free_mib():
|
||
"""Free VRAM per GPU (host index → MiB) via nvidia-smi, without a CUDA context.
|
||
|
||
Empty dict if the query fails: the caller then applies no bound.
|
||
"""
|
||
try:
|
||
import subprocess
|
||
result = subprocess.run(
|
||
['nvidia-smi', '--query-gpu=index,memory.free',
|
||
'--format=csv,noheader,nounits'],
|
||
capture_output=True, text=True, timeout=15,
|
||
)
|
||
if result.returncode != 0:
|
||
return {}
|
||
free = {}
|
||
for line in result.stdout.strip().splitlines():
|
||
parts = [p.strip() for p in line.split(',')]
|
||
if len(parts) >= 2:
|
||
free[int(parts[0])] = int(parts[1])
|
||
return free
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
def gpu_worker_slots(gpu_ids, n_workers, free_mib=None):
|
||
"""Slot of each pool worker: GPU index, -1 (forced CPU) or None.
|
||
|
||
Each GPU gets at most (free VRAM − reserve) / GPU_WORKER_MIB workers (at
|
||
least one), interleaved across GPUs; the surplus runs on CPU instead of
|
||
exhausting VRAM (OOM). No GPU: None everywhere (the worker decides).
|
||
Unknown VRAM: unbounded round-robin (legacy behaviour).
|
||
"""
|
||
if not gpu_ids:
|
||
return [None] * n_workers
|
||
if free_mib is None:
|
||
free_mib = gpu_free_mib()
|
||
if not all(g in free_mib for g in gpu_ids):
|
||
return [gpu_ids[i % len(gpu_ids)] for i in range(n_workers)]
|
||
capacity = {g: max(1, (free_mib[g] - GPU_RESERVE_MIB) // GPU_WORKER_MIB)
|
||
for g in gpu_ids}
|
||
slots = []
|
||
for rank in range(max(capacity.values())):
|
||
slots += [g for g in gpu_ids if capacity[g] > rank]
|
||
slots = slots[:n_workers]
|
||
return slots + [-1] * (n_workers - len(slots))
|
||
|
||
|
||
def _gpu_available():
|
||
"""Check if GPU is usable right now."""
|
||
try:
|
||
_init_gpu()
|
||
return HAS_GPU and _cp is not None
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def log_gpu_status():
|
||
"""Log GPU detection result. Called after logging is configured."""
|
||
if _gpu_available():
|
||
try:
|
||
with _cp.cuda.Device(0):
|
||
props = _cp.cuda.runtime.getDeviceProperties(0)
|
||
name = props['name']
|
||
if isinstance(name, bytes):
|
||
name = name.decode()
|
||
mem_gb = props['totalGlobalMem'] // (1024**3)
|
||
gpu_info = f"GPU: {name} ({mem_gb} GB VRAM) — ID {_best_gpu_id}"
|
||
except Exception:
|
||
gpu_info = f"GPU: {_gpu_name} ({_gpu_mem_gb} GB VRAM)"
|
||
logger.info(gpu_info)
|
||
|
||
# List other GPUs for info
|
||
try:
|
||
import subprocess
|
||
result = subprocess.run(
|
||
['nvidia-smi', '--query-gpu=index,name,compute_cap',
|
||
'--format=csv,noheader,nounits'],
|
||
capture_output=True, text=True, timeout=5,
|
||
)
|
||
if result.returncode == 0:
|
||
for line in result.stdout.strip().split('\n'):
|
||
parts = [p.strip() for p in line.split(',')]
|
||
if len(parts) >= 3:
|
||
idx = int(parts[0])
|
||
if idx != _best_gpu_id:
|
||
cap = parts[2]
|
||
logger.info(f" GPU {idx}: {parts[1]} (sm_{cap}) — available")
|
||
except Exception:
|
||
pass
|
||
else:
|
||
logger.info("No usable GPU — CPU-only mode")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Array transfer
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def to_gpu(arr):
|
||
"""Send array to GPU if available, otherwise return as float32 numpy."""
|
||
if _gpu_available():
|
||
try:
|
||
return _cp.asarray(arr.astype(np.float32))
|
||
except Exception as e:
|
||
# The original error (often OOM) must not be swallowed:
|
||
# without it, a CPU fallback is indistinguishable from a plain bug.
|
||
logger.warning(f"GPU transfer failed ({e}) — falling back to CPU")
|
||
disable_gpu()
|
||
return arr.astype(np.float32)
|
||
|
||
|
||
def to_cpu(arr):
|
||
"""Bring array back to CPU (numpy). No-op if already on CPU.
|
||
|
||
Works even after disable_gpu() thanks to the _cupy_ndarray type reference.
|
||
"""
|
||
if _cupy_ndarray is not None and isinstance(arr, _cupy_ndarray):
|
||
try:
|
||
if _cp is not None:
|
||
return _cp.asnumpy(arr)
|
||
return arr.get()
|
||
except Exception:
|
||
try:
|
||
return np.asarray(arr)
|
||
except Exception:
|
||
pass
|
||
return arr
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Filters — GPU if array is on GPU, CPU otherwise
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def xp_gaussian_filter(arr, sigma):
|
||
if _cp is not None and isinstance(arr, _cp.ndarray):
|
||
try:
|
||
return _cp_ndimage.gaussian_filter(arr, sigma)
|
||
except Exception as e:
|
||
logger.warning(f"GPU Gaussian filter failed ({e}) — falling back to CPU")
|
||
arr = to_cpu(arr)
|
||
return ndimage.gaussian_filter(arr, sigma)
|
||
|
||
|
||
def xp_uniform_filter(arr, size):
|
||
if _cp is not None and isinstance(arr, _cp.ndarray):
|
||
try:
|
||
return _cp_ndimage.uniform_filter(arr, size)
|
||
except Exception as e:
|
||
logger.warning(f"GPU uniform filter failed ({e}) — falling back to CPU")
|
||
arr = to_cpu(arr)
|
||
return ndimage.uniform_filter(arr, size)
|
||
|
||
|
||
def xp_zoom(arr, factor, order=1):
|
||
"""Upsampling with grid_mode (pixel-center alignment over the full image
|
||
extent): each source pixel covers exactly factor×factor output pixels,
|
||
with no half-pixel shift."""
|
||
if _cp is not None and isinstance(arr, _cp.ndarray):
|
||
try:
|
||
return _cp_ndimage.zoom(arr, factor, order=order, mode='nearest', grid_mode=True)
|
||
except Exception as e:
|
||
logger.warning(f"GPU zoom failed ({e}) — falling back to CPU")
|
||
arr = to_cpu(arr)
|
||
return ndimage.zoom(arr, factor, order=order, mode='nearest', grid_mode=True)
|
||
|
||
|
||
def xp_minimum_filter(arr, footprint=None, size=None):
|
||
if _cp is not None and isinstance(arr, _cp.ndarray):
|
||
try:
|
||
return _cp_ndimage.minimum_filter(arr, footprint=footprint, size=size)
|
||
except Exception as e:
|
||
logger.warning(f"GPU minimum filter failed ({e}) — falling back to CPU")
|
||
arr = to_cpu(arr)
|
||
return ndimage.minimum_filter(arr, footprint=footprint, size=size)
|
||
|
||
|
||
def xp_maximum_filter(arr, footprint=None, size=None):
|
||
if _cp is not None and isinstance(arr, _cp.ndarray):
|
||
try:
|
||
return _cp_ndimage.maximum_filter(arr, footprint=footprint, size=size)
|
||
except Exception as e:
|
||
logger.warning(f"GPU maximum filter failed ({e}) — falling back to CPU")
|
||
arr = to_cpu(arr)
|
||
return ndimage.maximum_filter(arr, footprint=footprint, size=size)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# DTM rasterization — mean z per cell (2D bincount)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _bin_mean_core(lib, xs, ys, zs, width, height, x_range, y_range):
|
||
"""Mean z per cell of a regular grid, backend-agnostic.
|
||
|
||
`lib` = numpy or cupy (same primitives). Semantics mirror
|
||
scipy.stats.binned_statistic_2d(statistic='mean', bins=[width, height],
|
||
range=[[xmin, xmax], [ymin, ymax]]): points outside the extent are
|
||
ignored, a value exactly on the right/top edge goes into the last cell.
|
||
Returns (height, width) float64, NaN on empty cells.
|
||
|
||
float64 is required for x/y: at Lambert 93 coordinates (~1e6 m) float32
|
||
resolution is ~6 cm (Y ~7e6 m: ~50 cm), far too coarse for a 0.2 m pixel.
|
||
"""
|
||
(xmin, xmax), (ymin, ymax) = x_range, y_range
|
||
x = lib.asarray(xs, dtype=lib.float64)
|
||
y = lib.asarray(ys, dtype=lib.float64)
|
||
z = lib.asarray(zs, dtype=lib.float64)
|
||
inside = (x >= xmin) & (x <= xmax) & (y >= ymin) & (y <= ymax)
|
||
x, y, z = x[inside], y[inside], z[inside]
|
||
# floor((v - min) / step); interior values are >= 0 so astype truncation
|
||
# == floor. The clip puts the right edge (and the adjacent float rounding)
|
||
# into the last cell, like scipy.
|
||
ix = ((x - xmin) * (float(width) / (xmax - xmin))).astype(lib.int64)
|
||
iy = ((y - ymin) * (float(height) / (ymax - ymin))).astype(lib.int64)
|
||
ix = lib.clip(ix, 0, width - 1)
|
||
iy = lib.clip(iy, 0, height - 1)
|
||
n = int(width) * int(height)
|
||
idx = iy * int(width) + ix
|
||
sums = lib.bincount(idx, weights=z, minlength=n)
|
||
counts = lib.bincount(idx, minlength=n)
|
||
mean = sums / lib.where(counts == 0, 1, counts)
|
||
return lib.where(counts == 0, float("nan"), mean).reshape(int(height), int(width))
|
||
|
||
|
||
def bin_mean_2d(xs, ys, zs, width, height, x_range, y_range):
|
||
""""Mean per cell" rasterization on GPU (caller: dtm.create_dtm_fast).
|
||
|
||
Returns a numpy (height, width) float64 array (NaN = empty cell), or None
|
||
if the GPU is unavailable or fails (usually OOM) — the caller then falls
|
||
back to scipy. A failure here does NOT call disable_gpu(): rasterization
|
||
is a one-off step, and the visualizations that follow must keep their
|
||
accelerator.
|
||
"""
|
||
if not _gpu_available():
|
||
return None
|
||
try:
|
||
result = _bin_mean_core(_cp, xs, ys, zs, width, height,
|
||
x_range, y_range)
|
||
return to_cpu(result)
|
||
except Exception as e:
|
||
logger.warning(f"GPU rasterization failed ({e}) — falling back to scipy")
|
||
gpu_cleanup()
|
||
return None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Misc
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def gpu_cleanup():
|
||
"""Free GPU memory. Call between visualizations to prevent OOM."""
|
||
if _cp is not None:
|
||
try:
|
||
_cp.get_default_memory_pool().free_all_blocks()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def disable_gpu():
|
||
"""Disable GPU acceleration for the rest of this process.
|
||
|
||
Frees the GPU memory pool before dropping the references (avoids VRAM
|
||
leaks). Keeps _cupy_ndarray alive so to_cpu() can recover orphaned arrays.
|
||
"""
|
||
global HAS_GPU, _cp, _cp_ndimage
|
||
if not HAS_GPU:
|
||
return
|
||
logger.warning("GPU disabled — switching to CPU for the rest of this process")
|
||
gpu_cleanup()
|
||
HAS_GPU = False
|
||
_cp = None
|
||
_cp_ndimage = None
|
||
|
||
|
||
def is_gpu_active():
|
||
"""Check if GPU acceleration is currently active."""
|
||
return HAS_GPU
|
||
|
||
|
||
def safe_gpu_call(func, *args, **kwargs):
|
||
"""Call a function with GPU arrays, retrying on CPU if GPU fails."""
|
||
try:
|
||
return func(*args, **kwargs)
|
||
except Exception as e:
|
||
# GPU active: on ANY error (OOM, mixed numpy/cupy types after a
|
||
# transfer failure mid-run, ...) the GPU is disabled and the
|
||
# computation is retried on CPU — a partial failure must not make the
|
||
# whole visualization fail. In CPU mode, the original error is
|
||
# re-raised (already on CPU, nothing to retry).
|
||
if _cp is not None:
|
||
logger.warning(f"GPU error ({e.__class__.__name__}: {e}), "
|
||
f"retrying on CPU...")
|
||
disable_gpu()
|
||
cpu_args = tuple(to_cpu(a) for a in args)
|
||
cpu_kwargs = {k: to_cpu(v) for k, v in kwargs.items()}
|
||
return func(*cpu_args, **cpu_kwargs)
|
||
raise
|