"""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