"""Видеокарта для плотных стадий конвейера: сетчатка, ламина, кластеризация. Три стадии работают с целым образом 128 × 600 лучей и занимают больше половины кадра: раскладка точек по решётке, ламина (центр минус окружение на трёх масштабах) и кластеризация лучей с учётом глубины. Они и перенесены. Остальное — плоскость пути, ось, оценка движения, треки, решение — остаётся на процессоре: там мелкие массивы и ветвистая логика, и копирование на видеокарту стоило бы дороже самого расчёта. Каждая функция повторяет процессорную версию операция в операцию, в тех же типах. Раскладка по порядку точек, ламина и кластеризация дают результат, совпадающий с процессорным бит в бит (тесты и сверка на записях, EXPERIMENTS п. 21): * сетчатка — только выборки, умножения и корень, всё в float32 без слияния операций; * ламина — равномерные окна считаются через накопленные суммы в float64 и округляются в float32 после каждой оси, ровно как `scipy.ndimage. uniform_filter`: сумма чисел float32 такого диапазона в float64 точна, и результат не зависит от порядка сложения; * кластеризация — целочисленный граф; компоненты нумеруются по наименьшему лучу, как у `scipy.sparse.csgraph.connected_components`. Исключение — раскладка по углам точек (кадры с нарушенным порядком, как в синтетике организаторов): арктангенс видеокарты и процессора может разойтись в последнем знаке, а numpy сортирует ячейки неустойчиво. Расходятся считанные лучи на кадр. Любой сбой видеокарты — нет драйвера, не хватило памяти, ошибка ядра — ловит конвейер: до конца работы он считает на процессоре (`device. notify_cuda_error`), а кадр досчитывается там же, решение не теряется. """ from __future__ import annotations import warnings import numpy as np from .lamina import SCALES, LaminaOutput from .retina import DEG, ORDER_TOL_DEG, RangeImage, ScanLayout class GpuStages: """Плотные стадии на видеокарте; держит таблицы решётки и образ кадра.""" def __init__(self, device: str = "cuda"): import torch self.torch = torch self.dev = torch.device(device) torch.zeros(1, device=self.dev) # поднять контекст сразу, а не на первом кадре self._tab_key = None self._tab: dict = {} self._win: dict = {} self._off: dict = {} self._box: dict = {} self._graphs = True # графы CUDA для ламины; при отказе — по одной операции self._lam_key = None # образ текущего кадра: ламина и кластеризация берут его отсюда, не копируя заново self.r = None self.valid = None self._r_host = None def warmup(self, shape: tuple[int, int] = (128, 600)) -> None: """Прогнать все три стадии на выдуманном кадре. Первый запуск каждого ядра CUDA подгружает его код: без прогрева первый кадр записи шёл 1–1.7 с, очередь подписки переполнялась, и терялись первые 6–8 кадров (замерено в контейнере). Раскладка прогревается на маленькой решётке обоими путями — по порядку точек и по углам. """ from .cdr import PointCloud2 t = self.torch rng = np.random.default_rng(0) n, w, e = 8, 64, 2 lay = ScanLayout(np.linspace(10.0, -20.0, n), -0.1, 3.2, np.zeros(n, np.int64), np.zeros(n), w, e) xyz = lay.dirs.transpose(1, 0, 2)[:, None, :, :] * np.float32(20.0) # столбец·эхо·кольцо xyz = np.broadcast_to(xyz, (w, e, n, 3)).reshape(-1, 3) pts = np.zeros(xyz.shape[0], dtype=[("x", " dict: key = (id(L), L.n_points, L.n_az) if key != self._tab_key: t, d = self.torch, self.dev def up(a, dtype): return t.as_tensor(np.ascontiguousarray(a), dtype=dtype, device=d) self._tab = { "sin_el": up(L._sin_el, t.float32), "el_asc": up(L._el_asc, t.float64), "el_order": up(L._el_order, t.int64), "el_deg": up(L.el_deg, t.float64), "el_step": up(L.el_step_deg, t.float64), "resid": up(L.az_resid_deg, t.float64), } self._win = {} self._tab_key = key return self._tab def _window(self, L: ScanLayout, start: int, stop: int) -> tuple: """Сырые столбцы, нужные сектору, и карта выпрямления внутри них.""" key = (start, stop) if key not in self._win: w = L.n_az lo = start + int(L.col_shift.min()) hi = stop + int(L.col_shift.max()) g = L.gather[:, start:stop] ok = L.gather_ok[:, start:stop] if L.wrap: raw = np.arange(lo, hi) % w g = (g - lo) % w else: lo, hi = max(lo, 0), min(hi, w) raw = np.arange(lo, hi) g = g - lo ok = ok & (g >= 0) & (g < (hi - lo)) g = np.clip(g, 0, hi - lo - 1) # столбцы окна идут подряд, кроме стыка кругового скана cut = np.flatnonzero(np.diff(raw) != 1) + 1 runs = [(int(a[0]), int(a[-1]) + 1) for a in np.split(raw, cut) if a.size] t = self.torch self._win[key] = (runs, t.as_tensor(g, dtype=t.int64, device=self.dev), t.as_tensor(ok, device=self.dev)) return self._win[key] # ------------------------------------------------------------------ загрузка точек def _upload(self, pts: np.ndarray, rows: list[tuple[int, int]]): """Строки структурированного массива → байты на видеокарте, (N, шаг точки).""" t = self.torch step = pts.dtype.itemsize parts = [] for a, b in rows: raw = np.ascontiguousarray(pts[a:b]).view(np.uint8) with warnings.catch_warnings(): # буфер сообщения только для чтения, а тензор мы не пишем warnings.simplefilter("ignore", UserWarning) parts.append(t.from_numpy(raw).to(self.dev, non_blocking=False)) buf = parts[0] if len(parts) == 1 else t.cat(parts) return buf.view(-1, step) def _field(self, buf, pts: np.ndarray, name: str): dt, off = pts.dtype.fields[name][:2] if dt.kind != "f" or dt.itemsize != 4 or dt.byteorder == ">": raise TypeError(f"поле {name}: ожидался float32, а не {dt}") return buf[:, off:off + 4].contiguous().view(self.torch.float32).reshape(-1) # ------------------------------------------------------------------ сетчатка def project(self, L: ScanLayout, pc, cols: slice | None) -> RangeImage: """То же, что `ScanLayout.project`, на видеокарте.""" w = L.n_az start = 0 if cols is None else (cols.start or 0) stop = w if cols is None else (cols.stop if cols.stop is not None else w) self._tables(L) img = None if L.indexed and pc.n_points == L.n_points: img = self._project_indexed(L, pc, start, stop) if img is None: L.n_geometric += 1 img = self._project_geometric(L, pc, start, stop) self._r_host = img.r_near return img def _project_indexed(self, L: ScanLayout, pc, start: int, stop: int): t = self.torch tab = self._tab n, e = L.n_rings, L.n_echo runs, g, ok = self._window(L, start, stop) pts = pc.points per_col = e * n buf = self._upload(pts, [(a * per_col, b * per_col) for a, b in runs]) ncol = sum(b - a for a, b in runs) def cube(name: str): return self._field(buf, pts, name).reshape(ncol, e, n).permute(2, 0, 1) x, y, z = cube("x"), cube("y"), cube("z") good = (x != 0) | (y != 0) | (z != 0) r = t.sqrt(x * x + y * y + z * z) good &= t.isfinite(r) r = t.where(good, r, t.zeros((), dtype=r.dtype, device=self.dev)) dev = t.abs(z - r * tab["sin_el"][:, None, None]) tol = r * np.float32(ORDER_TOL_DEG * DEG) + np.float32(1e-3) if bool(t.any(dev > tol)): return None # яркость — не обязательное поле, без неё нули (как в retina.py) has_i = "intensity" in pts.dtype.names if e == 1: r_near = r[..., 0] r_far = r[..., 0] it = cube("intensity")[..., 0] if has_i else t.zeros_like(r_near) valid = good[..., 0] else: inten = cube("intensity") if has_i else t.zeros_like(r) inf = t.full((), float("inf"), dtype=r.dtype, device=self.dev) near_i = t.argmin(t.where(good, r, inf), dim=-1, keepdim=True) far_i = t.argmax(r, dim=-1, keepdim=True) r_near = t.gather(r, -1, near_i)[..., 0] r_far = t.gather(r, -1, far_i)[..., 0] it = t.gather(inten, -1, near_i)[..., 0] valid = good.any(dim=-1) r_near = t.gather(r_near, 1, g) r_far = t.gather(r_far, 1, g) it = t.gather(it, 1, g) valid = t.gather(valid, 1, g) & ok zero = t.zeros((), dtype=r_near.dtype, device=self.dev) r_near = t.where(valid, r_near, zero) r_far = t.where(valid, r_far, zero) return self._finish(pc.stamp, r_near, r_far, it, valid) def _project_geometric(self, L: ScanLayout, pc, start: int, stop: int): t = self.torch tab = self._tab n, wid = L.n_rings, stop - start pts = pc.points buf = self._upload(pts, [(0, pts.shape[0])]) x, y, z = (self._field(buf, pts, k) for k in ("x", "y", "z")) good = (((x != 0) | (y != 0) | (z != 0)) & t.isfinite(x) & t.isfinite(y) & t.isfinite(z)) idx = t.nonzero(good, as_tuple=True)[0] x, y, z = x[idx], y[idx], z[idx] step = L.az_step_deg az = t.rad2deg(t.atan2(x, -y)) jf = (az - np.float32(L.az0_deg)) / np.float32(step) if L.wrap: jf = t.remainder(jf, L.n_az) margin = float(np.abs(L.az_resid_deg).max()) / abs(step) + 1.0 sel = t.nonzero((jf > start - margin) & (jf < stop - 1 + margin), as_tuple=True)[0] idx, x, y, z, az = idx[sel], x[sel], y[sel], z[sel], az[sel] r = t.sqrt(x * x + y * y + z * z) el = t.rad2deg(t.asin(t.clamp(z / t.clamp_min(r, np.float32(1e-6)), -1.0, 1.0))) el64 = el.double() asc = tab["el_asc"] k = t.clamp(t.searchsorted(asc, el64), 1, n - 1) k = k - ((el64 - asc[k - 1]) < (asc[k] - el64)).long() h = tab["el_order"][k] ok = t.abs(el64 - tab["el_deg"][h]) <= t.clamp_min(tab["el_step"][h], 0.2) j = t.round(((az.double() - tab["resid"][h]) - L.az0_deg) / step).long() if L.wrap: j = t.remainder(j, L.n_az) ok &= (j >= start) & (j < stop) sel = t.nonzero(ok, as_tuple=True)[0] r_near = t.zeros(n * wid, dtype=t.float32, device=self.dev) r_far = t.zeros_like(r_near) it = t.zeros_like(r_near) valid = t.zeros(n * wid, dtype=t.bool, device=self.dev) if sel.numel(): cell = h[sel] * wid + (j[sel] - start) rr = r[sel] mm = t.clamp_max(rr * 1000.0, float((1 << 20) - 1)).long() order = t.argsort(cell * (1 << 20) + mm, stable=True) cs, rs = cell[order], rr[order] new = t.ones(cs.numel(), dtype=t.bool, device=self.dev) new[1:] = cs[1:] != cs[:-1] first = t.nonzero(new, as_tuple=True)[0] last = t.cat([first[1:] - 1, first.new_tensor([cs.numel() - 1])]) r_near[cs[first]] = rs[first] r_far[cs[last]] = rs[last] valid[cs[first]] = True if "intensity" in pts.dtype.names: src = self._field(buf, pts, "intensity")[idx[sel]][order] it[cs[first]] = src[first] return self._finish(pc.stamp, r_near.reshape(n, wid), r_far.reshape(n, wid), it.reshape(n, wid), valid.reshape(n, wid)) def _finish(self, stamp: float, r_near, r_far, it, valid) -> RangeImage: self.r, self.valid = r_near, valid host = [a.cpu().numpy() for a in (r_near, r_far, it, valid)] return RangeImage(stamp, *host) # ------------------------------------------------------------------ ламина def _box_axis(self, X, sizes: list[int], dim: int): """Равномерное окно по одной оси для пачки размеров сразу. `X` — (K, B, H, W) float32, k-й срез фильтруется окном `sizes[k]` вдоль `dim` (2 или 3) с продолжением краевым значением. Суммы окон берутся из накопленной суммы в float64: все частичные суммы чисел float32 такого диапазона в float64 точны, поэтому сумма окна точна при любом порядке сложения. Деление на размер и округление в float32 — как в `scipy.ndimage.uniform_filter1d`. """ t = self.torch Xd = X.double() n = Xd.shape[dim] lo, hi, n_lo, n_hi, size = self._box_tables(tuple(sizes), n, dim) cs = t.cumsum(Xd, dim=dim) cs = t.cat([t.zeros_like(cs.narrow(dim, 0, 1)), cs], dim=dim) def along(a): return a.expand(*Xd.shape[:dim], n, *Xd.shape[dim + 1:]) s = (t.gather(cs, dim, along(hi)) - t.gather(cs, dim, along(lo)) + along(n_lo) * Xd.narrow(dim, 0, 1) + along(n_hi) * Xd.narrow(dim, n - 1, 1)) return (s / size).float() def _box_tables(self, sizes: tuple, n: int, dim: int): """Границы окон по оси: считаются один раз (и до записи графа CUDA).""" key = (sizes, n, dim) if key not in self._box: t = self.torch p = t.as_tensor([s // 2 for s in sizes], dtype=t.int64, device=self.dev) i = t.arange(n, device=self.dev) shape = [len(sizes), 1, 1, 1] shape[dim] = n lo = (i[None, :] - p[:, None]).clamp(min=0).reshape(shape) hi = ((i[None, :] + p[:, None]).clamp(max=n - 1) + 1).reshape(shape) n_lo = (p[:, None] - i[None, :]).clamp(min=0).double().reshape(shape) n_hi = (i[None, :] + p[:, None] - (n - 1)).clamp(min=0).double().reshape(shape) size = t.as_tensor(sizes, dtype=t.float64, device=self.dev).reshape(-1, 1, 1, 1) self._box[key] = (lo, hi, n_lo, n_hi, size) return self._box[key] def _boxes(self, X, sizes: list[tuple[int, int]]): """`scipy.ndimage.uniform_filter(x, size, mode="nearest")` для пачки. `X` — (B, H, W) float32; выход (len(sizes), B, H, W): по осям по очереди, с округлением в float32 после каждой, как в scipy. """ K = len(sizes) Y = self._box_axis(X[None].expand(K, *X.shape), [s[0] for s in sizes], 2) return self._box_axis(Y, [s[1] for s in sizes], 3) def lamina(self, r_max: float = 300.0) -> LaminaOutput: """То же, что `lamina._process_cpu`, по образу текущего кадра. Сама видеокарта считает ламину за доли миллисекунды, а запуск сотни мелких операций из Python стоит в несколько раз дороже. Поэтому ламина записывается в граф CUDA один раз на размер образа и дальше запускается одной командой. Не поддерживает граф драйвер — считаем теми же операциями по одной. """ t = self.torch key = (tuple(self.r.shape), float(r_max)) out = None if self._graphs: try: if self._lam_key != key: self._capture_lamina(key, r_max) self._lam_in[0].copy_(self.r) self._lam_in[1].copy_(self.valid) self._lam_graph.replay() out = self._lam_out except Exception: self._graphs = False self._lam_key = None if out is None: out = self._lamina_core(self.r, self.valid, r_max) f = out[0].cpu().numpy() return LaminaOutput(disp=f[0], on=f[1], off=f[2], on_scale=out[1].cpu().numpy(), surround=f[3], hole=f[4]) def _capture_lamina(self, key: tuple, r_max: float) -> None: t = self.torch self._lam_in = (t.empty_like(self.r), t.empty_like(self.valid)) self._lam_in[0].copy_(self.r) self._lam_in[1].copy_(self.valid) side = t.cuda.Stream(self.dev) side.wait_stream(t.cuda.current_stream(self.dev)) with t.cuda.stream(side): for _ in range(2): # прогрев: таблицы окон и кэш памяти self._lamina_core(*self._lam_in, r_max) t.cuda.current_stream(self.dev).wait_stream(side) g = t.cuda.CUDAGraph() with t.cuda.graph(g): self._lam_out = self._lamina_core(*self._lam_in, r_max) self._lam_graph, self._lam_key = g, key def _lamina_core(self, r, valid, r_max: float): """Ламина целиком на видеокарте: (disp, on, off, surround, hole) и on_scale.""" t = self.torch v = valid.float() zero = t.zeros((), dtype=t.float32, device=self.dev) disp = t.where(valid & (r > 0.05), 1.0 / r, zero) * v # все окна кадра одной пачкой: на каждом масштабе центр и окружение, # для диспаритета и для маски эха, плюс окно «дыры» sizes = [] for r_in, r_out in SCALES: sizes += [(2 * r_in + 1, 2 * r_in + 1), (2 * r_out + 1, 4 * r_out + 1)] sizes.append((5, 15)) box = self._boxes(t.stack([disp, v]), sizes) on = t.zeros_like(disp) off = t.zeros_like(disp) on_scale = t.zeros(disp.shape, dtype=t.int8, device=self.dev) surround_mid = None for k, (r_in, r_out) in enumerate(SCALES): s_in, s_out = sizes[2 * k], sizes[2 * k + 1] n_in = s_in[0] * s_in[1] n_out = s_out[0] * s_out[1] num = box[2 * k + 1, 0] * n_out - box[2 * k, 0] * n_in den = box[2 * k + 1, 1] * n_out - box[2 * k, 1] * n_in sur = t.where(den > 0.5, num / den, zero) enough = den > 8.0 c = t.where(enough, disp - sur, zero) pos = t.clamp_min(c, 0.0) * v neg = t.clamp_min(-(disp - sur), 0.0) * enough better = pos > on on = t.where(better, pos, on) on_scale = t.where(better, t.full((), k, dtype=t.int8, device=self.dev), on_scale) off = t.maximum(off, neg) if k == 1: surround_mid = sur hole = 1.0 - box[-1, 1] on = t.clamp(on, 0.0, 1.0 / max(r_max, 1.0) * 1e4) return t.stack([disp, on, off, surround_mid, hole]), on_scale # ------------------------------------------------------------------ кластеризация def _offsets(self, col_reach: int, row_reach: int): key = (col_reach, row_reach) if key not in self._off: pairs = [(dr, dc) for dr in range(0, row_reach + 1) for dc in range(-col_reach, col_reach + 1) if not (dr == 0 and dc <= 0)] p = self.torch.tensor(pairs, dtype=self.torch.int64, device=self.dev) self._off[key] = (p[:, 0:1], p[:, 1:2]) return self._off[key] def cluster_by_depth(self, mask: np.ndarray, r: np.ndarray, *, rel_tol: float = 0.06, abs_tol: float = 0.35, col_reach: int = 3, row_reach: int = 2): """То же, что `lobula.cluster_by_depth`, с теми же номерами компонент. Рёбра строятся сразу для всех соседей, компоненты — подвешиванием к меньшему номеру со сжатием путей: у каждой компоненты корнем остаётся её наименьший луч, и нумерация по корням совпадает со scipy. """ t = self.torch h, w = mask.shape if r is self._r_host and self.r is not None: rt = self.r.reshape(-1) else: rt = t.as_tensor(np.ascontiguousarray(r), dtype=t.float32, device=self.dev).reshape(-1) m = t.as_tensor(np.ascontiguousarray(mask), device=self.dev).reshape(-1) idx = t.nonzero(m, as_tuple=True)[0] n = idx.numel() if n == 0: return np.zeros(mask.shape, np.int32), 0 lut = t.full((h * w,), -1, dtype=t.int64, device=self.dev) lut[idx] = t.arange(n, device=self.dev) rows = t.div(idx, w, rounding_mode="floor") cols = idx - rows * w ri = rt[idx] dr, dc = self._offsets(col_reach, row_reach) nr = rows[None, :] + dr nc = cols[None, :] + dc ok = (nr < h) & (nc >= 0) & (nc < w) nb = lut[t.where(ok, nr * w + nc, t.zeros_like(nr))] ok &= nb >= 0 rb = ri[nb.clamp(min=0)] ra = ri[None, :].expand_as(rb) ok &= t.abs(ra - rb) <= (abs_tol + rel_tol * t.minimum(ra, rb)) s = t.arange(n, device=self.dev)[None, :].expand_as(nb)[ok] d = nb[ok] parent = t.arange(n, device=self.dev) while True: for _ in range(4): ps, pd = parent[s], parent[d] parent.scatter_reduce_(0, t.maximum(ps, pd), t.minimum(ps, pd), reduce="amin") parent = parent[parent] parent = parent[parent] if not bool((parent[s] != parent[d]).any()): break while True: nxt = parent[parent] if bool(t.equal(nxt, parent)): break parent = nxt roots, comp = t.unique(parent, return_inverse=True) labels = t.zeros(h * w, dtype=t.int32, device=self.dev) labels[idx] = (comp + 1).to(t.int32) return labels.reshape(h, w).cpu().numpy(), int(roots.numel())