плотные стадии на видеокарте, откат на процессор

README.md

@@ -97,7 +97,7 @@ flyguard/          ядро: стадии обработки, память, сч
   export.py                 3D-рамки, время до столкновения, маркеры RViz
   synth.py                  вставка предметов трассировкой лучей
 tools/             обучение, оценка, разбор, полигон с аугментациями
-tests/             47 тестов, запускаются без данных и без ROS
+tests/             50 тестов, запускаются без данных и без ROS
 docs/              методика и результаты
 artifacts/         обученные модели
 ```

docs/EXPERIMENTS.md

@@ -3223,3 +3223,98 @@ MBON, новизна; у каждого трека ход улики) на си
 | контур габарита в кривой | прямой короб, изогнутой части не видно | оба контура |
 | решения узла | — | без изменений (синтетика: те же находки, фантомы, дальности и число кадров) |
 | тесты | 48 | 50 |
+
+## 21. Плотные стадии на видеокарте с откатом на процессор (26.09)
+
+Решение пользователя: считать на видеокарте, если она есть, и только при её
+отсутствии или сбое — на процессоре. Раньше (п. 7.4) вывод держался на
+процессоре: на видеокарту была перенесена одна ламина, выигрыш 3 мс из 43.
+Перенос одной стадии и правда ничего не даёт — выигрыш даёт перенос всех
+плотных стадий разом, с данными, которые между стадиями не ездят.
+
+### 21.1. Что переносить
+
+Профиль кадра (cProfile, синтетика, 150 кадров; под профилировщиком всё
+медленнее, важны доли): кластеризация с учётом глубины 20 мс, раскладка точек
+по углам 14, ламина 8, оценка движения 8.5, ось пути 6.5, плоскость 3.9 из 74.
+Без профилировщика кадр — 49 мс на синтетике и 36 мс на записи с поезда, из них
+на раскладку, ламину и кластеризацию приходится 22–34 мс. Они работают с целым
+образом 128 × 600 лучей — это и перенесено (`flyguard/gpu.py`). Плоскость, ось,
+движение, треки и решение остались на процессоре: там мелкие массивы и
+ветвистая логика, копирование на видеокарту дороже расчёта.
+
+### 21.2. Те же операции в тех же типах
+
+Цель — не «похожий», а тот же результат, иначе откат на процессор менял бы
+решения посреди поездки.
+
+* **Сетчатка** — выборки, умножения и корень в float32, отдельными ядрами
+  (без слияния умножения со сложением). На видеокарту уходит только окно сырых
+  столбцов, нужное сектору (5 МБ вместо 24).
+* **Ламина.** `scipy.ndimage.uniform_filter` считает окно скользящим средним в
+  double. Суммы чисел float32 такого диапазона (диспаритет 0.003…20, окно до
+  121) в double точны при любом порядке сложения, поэтому окно берётся из
+  накопленной суммы, делится на размер и округляется в float32 после каждой
+  оси — как в scipy. Сверка окон: 0 расхождений на 24.5 млн значений (запись) и
+  20.6 млн (синтетика). Вся ламина записана в граф CUDA: видеокарта считает её
+  за 0.4 мс, а запуск сотни мелких операций из Python стоил 3 мс; с графом —
+  0.6 мс вместе с выгрузкой.
+* **Кластеризация** — рёбра сразу для всех 17 соседей, компоненты
+  подвешиванием к меньшему номеру со сжатием путей: корень компоненты — её
+  наименьший луч, и нумерация совпадает со `scipy.sparse.csgraph`. На 42 тыс.
+  лучей и 581 тыс. рёбер — 8 раундов (медиана).
+
+Исключение — раскладка по углам точек (кадры со сбитым порядком, как в
+синтетике организаторов): арктангенс видеокарты и процессора расходится в
+последнем знаке, а numpy сортирует ячейки неустойчиво.
+
+### 21.3. Сверка
+
+По стадиям, 80 кадров: на `roundT_doubleT` все три — бит в бит; на синтетике
+ламина и кластеризация — бит в бит, раскладка по углам разошлась в 13 кадрах из
+80, 326 значений из 300 тыс.
+
+Конвейер целиком, синтетика и шесть записей, покадрово против процессорного
+прогона той же конфигурации (`v_t33`): на всех шести записях совпали решения,
+объекты, кандидаты и треки в каждом кадре; на синтетике решения, объекты и
+треки — во всех 1499 кадрах, кандидаты разошлись в трёх (раскладка по углам),
+на треки это не повлияло. Синтетика 9 из 10, те же дальности и число кадров.
+
+### 21.4. Откат
+
+Любое исключение видеокарты — нет драйвера, не хватило памяти, ошибка ядра —
+ловит конвейер (`FlyGuard._gpu_off`): до конца работы он считает на
+процессоре, а текущий кадр досчитывается там же. Тест
+`test_pipeline_survives_gpu_failure_mid_frame` роняет видеокарту на 20-м кадре
+и сверяет решения с чисто процессорным прогоном. Узел пишет в журнал, где идёт
+счёт, и предупреждает, если видеокарта отказала на ходу; в диагностике — поле
+`device`.
+
+### 21.5. В контейнере
+
+Образ с PyTorch 2.9.1 + CUDA 12.8, Docker в WSL, видеокарта через `/dev/dxg`
+(на стенде — `--gpus all`). Узел, `ros2 bag play` в том же контейнере:
+
+| запись | счёт | кадр, медиана / p95 | обнаружение |
+|---|---|---|---|
+| `doubleT_obstacle` | видеокарта | 23.5 / 31.3 мс | 189 из 190, 55.8 м |
+| `doubleT_obstacle` | процессор | 32.3 / 36.2 мс | 189 из 190, 55.8 м |
+| синтетика | видеокарта | 26.8 / 43.1 мс | 390 кадров с тревогой |
+| синтетика | процессор | 42.3 / 47.0 мс | 398 кадров с тревогой |
+
+Строка `doubleT_obstacle` с видеокартой — после правки прогрева. До неё кадров
+с видеокартой приходило меньше (195 и 1502 против 201 и 1510, отсюда и 390
+кадров тревоги на синтетике против 398): первый кадр шёл 1.0–1.7 с — при первом
+вызове подгружаются ядра CUDA раскладки, и очередь подписки переполнялась.
+Прогрев (`GpuStages.warmup`) теперь гоняет и раскладку, на маленькой выдуманной
+решётке обоими путями: принято 201 из 201, худший кадр 38 мс.
+
+### 21.6. Образ
+
+PyTorch с библиотеками CUDA — 6.6 ГБ (сами библиотеки NVIDIA 4.3 ГБ, cuDNN и
+cuBLAS из них почти половина); образ — 9.3 ГБ вместо 2.7. Нашим стадиям cuDNN,
+cuBLAS и прочие не нужны, но PyTorch без них не загружается. Ставится одним
+слоем без кэша pip. Колёса — из индекса PyTorch или, где сети нет, с HTTP-адреса
+`TORCH_WHEELS` (`docker/fetch_wheels.py`, `docker/wheels/README.md`). Через
+контекст сборки колёса не передаются: первая попытка так делала, и контекст в
+4.1 ГБ плюс слой с колёсами плюс установка чуть не заняли весь диск C:.

flyguard/gpu.py

@@ -0,0 +1,500 @@
+"""Видеокарта для плотных стадий конвейера: сетчатка, ламина, кластеризация.
+
+Три стадии работают с целым образом 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", "<f4"), ("y", "<f4"), ("z", "<f4"),
+                                             ("intensity", "<f4")])
+        pts["x"], pts["y"], pts["z"] = xyz[:, 0], xyz[:, 1], xyz[:, 2]
+        pts["intensity"] = 10.0
+        cloud = PointCloud2(0.0, "", 1, pts.size, pts.dtype.itemsize, True, pts)
+        self.project(lay, cloud, slice(8, w - 8))
+        mixed = PointCloud2(0.0, "", 1, pts.size, pts.dtype.itemsize, True,
+                            pts[rng.permutation(pts.size)])
+        self.project(lay, mixed, slice(8, w - 8))
+
+        r = (4.0 + 150.0 * rng.random(shape)).astype(np.float32)
+        self.r = t.as_tensor(r, device=self.dev)
+        self.valid = t.ones(shape, dtype=t.bool, device=self.dev)
+        self._r_host = r
+        self.lamina()
+        self.cluster_by_depth(rng.random(shape) < 0.3, r)
+        self.r = self.valid = self._r_host = None
+        t.cuda.synchronize(self.dev)
+
+    # ------------------------------------------------------------------ таблицы
+
+    def _tables(self, L: ScanLayout) -> 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
+
+        if e == 1:
+            r_near = r[..., 0]
+            r_far = r[..., 0]
+            it = cube("intensity")[..., 0]
+            valid = good[..., 0]
+        else:
+            inten = cube("intensity")
+            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())

flyguard/lobula.py

@@ -292,7 +292,7 @@ def find_candidates(tf: TrackFrame, lam: LaminaOutput, corridor: Corridor, *,
                     platform_h_hi: float = 1.25,
                     h_top: float = 0.0, half_width_top: float = 0.0,
                     top_d_max: float = 90.0, top_detect: bool = False,
-                    top_min_rays: int = 3) -> list[Candidate]:
+                    top_min_rays: int = 3, cluster=None) -> list[Candidate]:
     """Выделить кандидатов: связные пятна лучей, попавших в габарит.
 
     Кластеризация ведётся по **расширенной** области, а членство в габарите
@@ -337,7 +337,9 @@ def find_candidates(tf: TrackFrame, lam: LaminaOutput, corridor: Corridor, *,
 
     gap_img = contrast_to_depth_gap(lam.on, tf.r)
 
-    labels, n = cluster_by_depth(context, tf.r, col_reach=dilate_cols + 1)
+    # `cluster` — та же кластеризация на видеокарте (FlyGuard._cluster), с тем
+    # же результатом; здесь самая большая маска кадра, и выигрыш в ней
+    labels, n = (cluster or cluster_by_depth)(context, tf.r, col_reach=dilate_cols + 1)
     if n == 0:
         return []
     if split_gap > 0:

flyguard/pipeline.py

@@ -22,7 +22,7 @@ from .descending import Decision, DescendingNeurons
 from .geometry import (STRAIGHT, Corridor, RailPlane, TrackFrame, fit_corridor,
                        fit_rail_plane)
 from .fan_body import FanBody
-from .lobula import Candidate, find_candidates, gauge_mask
+from .lobula import Candidate, cluster_by_depth, find_candidates, gauge_mask
 from .medulla import EgoMotion, EgoMotionEstimator, EmdBank, looming
 from .mbon_readout import MbonReadout
 from .mushroom_body import Habituation, HabituationConfig, MushroomBody, NoMemory
@@ -221,11 +221,12 @@ class Params:
     # сбрасывать никогда.
     reset_gap_s: float = 2.0
     novelty_gate: float = 0.0      # ниже этой новизны кандидат отбрасывается сразу
-    # Где считать ламину: 'cpu', 'cuda' или 'auto' (видеокарта, если есть).
+    # Где считать плотные стадии — сетчатку, ламину, кластеризацию (gpu.py):
-    # По умолчанию процессор, и это решение, а не осторожность: GPU ускоряет
+    # 'auto' — видеокарта, если PyTorch её видит, иначе процессор; 'cuda';
-    # только ламину (5.8 → 2.9 мс из 43), а на машине проверки требует
+    # 'cpu'. Результат одинаков: на видеокарте те же операции в тех же типах.
-    # другого контейнера и запуска с --gpus all. EXPERIMENTS п. 7.4.
+    # Любой сбой видеокарты переводит конвейер на процессор до конца работы,
-    device: str = "cpu"
+    # кадр досчитывается там же. EXPERIMENTS п. 21.
+    device: str = "auto"
     enable_motion: bool = True
     enable_memory: bool = True
     enable_looming: bool = False   # T4/T5 + LPLC2: нужны для оценки надвигания и для
@@ -271,6 +272,15 @@ class FlyGuard:
         self.p = params or Params()
         from .device import get_device
         self.device = get_device(self.p.device)
+        self.gpu = None                  # плотные стадии на видеокарте (gpu.GpuStages)
+        self.gpu_error: str | None = None
+        if self.device.startswith("cuda"):
+            try:
+                from .gpu import GpuStages
+                self.gpu = GpuStages(self.device)
+                self.gpu.warmup()
+            except Exception as e:      # нет драйвера, не та видеокарта, нет памяти
+                self._gpu_off(e)
         self.memory = memory if (memory is not None and self.p.enable_memory) else NoMemory()
         self.readout = readout if (readout is not None and self.p.enable_mbon) else None
         self.layout_full = layout
@@ -333,6 +343,25 @@ class FlyGuard:
             else DescendingNeurons(warn_evidence=0.0, clear_evidence=0.0,
                                    emergency_evidence=0.0, min_hits=1))
 
+    # ------------------------------------------------------------------ видеокарта
+
+    def _gpu_off(self, exc: Exception) -> None:
+        """Сбой видеокарты: до конца работы считать на процессоре."""
+        from .device import notify_cuda_error
+        notify_cuda_error(exc)
+        self.gpu = None
+        self.device = "cpu"
+        self.gpu_error = f"{type(exc).__name__}: {exc}"
+
+    def _cluster(self, mask: np.ndarray, r: np.ndarray, **kw):
+        """Кластеризация с учётом глубины: видеокарта, при сбое — процессор."""
+        if self.gpu is not None:
+            try:
+                return self.gpu.cluster_by_depth(mask, r, **kw)
+            except Exception as e:
+                self._gpu_off(e)
+        return cluster_by_depth(mask, r, **kw)
+
     # ------------------------------------------------------------------ калибровка
 
     def _ensure_layout(self, pc: PointCloud2) -> bool:
@@ -412,7 +441,14 @@ class FlyGuard:
         self.frames_seen += 1
 
         with t("retina"):
-            img: RangeImage = self.layout_full.project(pc, self.cols)
+            img: RangeImage | None = None
+            if self.gpu is not None:
+                try:
+                    img = self.gpu.project(self.layout_full, pc, self.cols)
+                except Exception as e:
+                    self._gpu_off(e)
+            if img is None:
+                img = self.layout_full.project(pc, self.cols)
 
         with t("stabilize"):
             self.plane = fit_rail_plane(img, self.layout, prev=self.plane)
@@ -423,12 +459,14 @@ class FlyGuard:
                              else STRAIGHT)
 
         with t("lamina"):
-            dev = self.device
+            lam = None
-            if dev.startswith("cuda"):
+            if self.gpu is not None:
-                from .device import is_cuda_available
+                try:
-                if not is_cuda_available():
+                    lam = self.gpu.lamina()      # по образу кадра, уже лежащему на видеокарте
-                    self.device = dev = "cpu"
+                except Exception as e:
-            lam = lamina.process(tf.r, tf.valid, device=dev)
+                    self._gpu_off(e)
+            if lam is None:
+                lam = lamina.process(tf.r, tf.valid, device="cpu")
 
         with t("ego"):
             ego = (self.ego_est.update(tf, pc.stamp) if self.p.enable_motion
@@ -475,7 +513,8 @@ class FlyGuard:
                                     h_top=self.p.h_top,
                                     half_width_top=self.p.half_width_top,
                                     top_d_max=self.p.top_d_max,
-                                    top_detect=self.p.top_detect)
+                                    top_detect=self.p.top_detect,
+                                    cluster=self._cluster)
 
         with t("mushroom"):
             cands = self.memory.annotate(cands)

requirements.txt

@@ -4,6 +4,8 @@ scipy>=1.10
 lightgbm>=4.0
 # графики полигона: tools/plot_benchmark.py
 matplotlib>=3.7
-# только для обучения на GPU (tools/train_mbon.py --device cuda); ядру не нужен
+# видеокарта: плотные стадии конвейера (flyguard/gpu.py) и обучение
+# (tools/train_mbon.py --device cuda). Без него ядро считает на процессоре —
+# с тем же результатом
 # torch>=2.0
 pytest>=7.0

tests/test_pipeline.py

@@ -872,3 +872,110 @@ def test_gauge_outline_bends_with_the_track():
     assert c100 == pytest.approx(-100.0 ** 2 / (2 * R), rel=1e-3)
     assert at100[:, 0].min() == pytest.approx(c100 - 1.2, abs=1e-4)
     assert at100[:, 0].max() == pytest.approx(c100 + 1.2, abs=1e-4)
+
+
+# ------------------------------------------------------------ видеокарта
+
+def _cuda_or_skip():
+    torch = pytest.importorskip("torch")
+    if not torch.cuda.is_available():
+        pytest.skip("нет видеокарты с CUDA")
+
+
+def _scene(seed=0, shape=(128, 600)):
+    """Образ, похожий на тоннель: гладкие стены, ступеньки-предметы, пропуски эха."""
+    rng = np.random.default_rng(seed)
+    h, w = shape
+    r = (np.linspace(3.0, 180.0, w)[None, :] * (1.0 + 0.2 * np.sin(np.arange(h)[:, None] / 9.0))
+         ).astype(np.float32)
+    for _ in range(40):
+        i, j = rng.integers(0, h - 8), rng.integers(0, w - 12)
+        r[i:i + rng.integers(2, 8), j:j + rng.integers(2, 12)] *= np.float32(rng.uniform(0.3, 0.9))
+    r += rng.normal(0.0, 0.02, r.shape).astype(np.float32)
+    valid = rng.random(shape) > 0.1
+    return np.where(valid, r, np.float32(0.0)), valid
+
+
+def test_gpu_lamina_and_clustering_match_cpu_bit_for_bit():
+    """Ламина и кластеризация на видеокарте — те же числа и те же номера компонент."""
+    _cuda_or_skip()
+    import torch
+    from flyguard import lamina
+    from flyguard.gpu import GpuStages
+    from flyguard.lobula import cluster_by_depth
+
+    g = GpuStages("cuda")
+    for seed in range(3):
+        r, valid = _scene(seed)
+        g.r = torch.as_tensor(r, device="cuda")
+        g.valid = torch.as_tensor(valid, device="cuda")
+        a, b = lamina._process_cpu(r, valid), g.lamina()
+        for f in ("disp", "on", "off", "on_scale", "surround", "hole"):
+            assert np.array_equal(getattr(a, f), getattr(b, f)), f
+        mask = valid & (np.random.default_rng(seed).random(r.shape) < 0.5)
+        la, na = cluster_by_depth(mask, r)
+        lb, nb = g.cluster_by_depth(mask, r)
+        assert na == nb and np.array_equal(la, lb)
+
+
+def test_gpu_projection_matches_cpu():
+    """Сетчатка на видеокарте: целый кадр — бит в бит, перемешанный — почти."""
+    _cuda_or_skip()
+    from flyguard.bag import Bag
+    from flyguard.cdr import PointCloud2
+    from flyguard.gpu import GpuStages
+
+    if not DATA.exists():
+        pytest.skip("нет записей")
+    bag = Bag(next(p for p in DATA.iterdir() if p.is_dir()))
+    clouds = [pc for _, pc in bag.frames(start=2, stop=16)]
+    lay = ScanLayout.calibrate(clouds[:12])
+    cols = lay.column_slice(30.0)
+    g = GpuStages("cuda")
+    pc = clouds[-1]
+    a, b = lay.project(pc, cols), g.project(lay, pc, cols)
+    for f in ("r_near", "r_far", "inten", "valid"):
+        assert np.array_equal(getattr(a, f), getattr(b, f)), f
+
+    # порядок точек сбит, как в синтетике организаторов: раскладка по углам
+    rng = np.random.default_rng(0)
+    mixed = PointCloud2(pc.stamp, pc.frame_id, 1, pc.n_points, pc.point_step,
+                        pc.is_dense, pc.points[rng.permutation(pc.n_points)])
+    a, b = lay.project(mixed, cols), g.project(lay, mixed, cols)
+    assert a.valid.sum() > 1000
+    differ = sum(int((getattr(a, f) != getattr(b, f)).sum())
+                 for f in ("r_near", "r_far", "inten", "valid"))
+    assert differ < 1e-3 * a.valid.size     # последний знак арктангенса и порядок равных
+
+
+def test_pipeline_survives_gpu_failure_mid_frame():
+    """Отказ видеокарты посреди кадра: кадр досчитан на процессоре, решение то же."""
+    from flyguard.bag import Bag
+    from flyguard.device import reset_device_cache
+    from flyguard.pipeline import FlyGuard, Params
+
+    if not DATA.exists():
+        pytest.skip("нет записей")
+
+    class Broken:
+        def project(self, *a, **k):
+            raise RuntimeError("CUDA error: device lost (имитация)")
+
+        lamina = cluster_by_depth = project
+
+    frames = [pc for _, pc in Bag(next(p for p in DATA.iterdir() if p.is_dir())).frames(stop=30)]
+    ref = FlyGuard(Params(device="cpu"))
+    fg = FlyGuard(Params(device="cpu"))
+    try:
+        for i, pc in enumerate(frames):
+            if i == 20:                       # видеокарта «была» и отказала
+                fg.gpu, fg.device = Broken(), "cuda"
+            a, b = ref.process(pc), fg.process(pc)
+            assert (a is None) == (b is None)
+            if a is not None:
+                assert a.decision.detected == b.decision.detected
+                assert [c.d for c in a.candidates] == [c.d for c in b.candidates]
+        assert fg.gpu is None and fg.device == "cpu"
+        assert "device lost" in fg.gpu_error
+    finally:
+        reset_device_cache()