diff --git a/.gitignore b/.gitignore index 32214a7..27b1227 100644 --- a/.gitignore +++ b/.gitignore @@ -4,5 +4,12 @@ viz/ # Rendered video outputs (only the GIFs in assets/ are committed, via Git LFS) *.mp4 +.cache/ + +build/ +build_tests/ + +*.pyc + # Design notes are kept locally, not in the repository docs/design/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 15ec8ca..a7b0527 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -215,6 +215,7 @@ if(BUILD_TESTING) tests/factor_manifold_facade_test.cpp tests/dense_linear_solver_test.cpp tests/dense_cholesky_solver_test.cpp + tests/block_sparse_cholesky_solver_test.cpp tests/dense_qr_solver_test.cpp tests/gauss_newton_test.cpp tests/so2_manifold_test.cpp @@ -235,6 +236,9 @@ if(BUILD_TESTING) tests/decreasing_scale_minimizer_test.cpp tests/synthetic_pgo_test.cpp tests/synthetic_sba_test.cpp + tests/schur_system_test.cpp + tests/stereo_loop_ba_test.cpp + tests/triple_smooth_factor_batch.cu tests/dense_matrix_ops_test.cpp tests/utils_test.cpp tests/device_vector_test.cpp diff --git a/README.md b/README.md index 142e3fd..de3b79a 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ See [`python/examples`](python/examples/README.md) to run them. | **Built-in factors** | Reprojection, PnP, between (SO(2)/SO(3)/SE(2)/SE(3)/Sim(2)/Sim(3)/SL(4)/vector), point-to-point, point-to-plane, symmetric point-to-plane, prior, constant-velocity/constant-acceleration motion priors (SO(2)/SO(3)/SE(2)/SE(3)) | | **Custom factors and states** | CuPy / NVIDIA Warp kernels in Python, or user-defined CUDA kernels via `SizedFactorBatch` / `SizedStateBatch` in C++; the same types work with every minimizer, including RANSAC — see [Custom factors and states](docs/sphinx/custom_factors_and_states.rst) | | **Numeric Jacobians** | Finite-difference Jacobians for any factor batch (manifold-aware, reuses each state's `Plus` retraction), selectable globally (`MinimizerOptions.jacobian_mode`) or per factor group (the `jacobian_mode_override` of `Problem.add_factor_batch`) — write a factor with only a residual and let cuNLS differentiate it; see [Numeric Jacobians](docs/sphinx/numeric_jacobians.rst) | -| **Linear solver** | Block-sparse PCG (variable block-Jacobi preconditioner, default), NVIDIA cuDSS (optional, loaded via `dlopen()` at runtime — see [C++ Installation](docs/sphinx/installation.rst)), dense LDLT, dense Cholesky (cuSOLVER), dense QR (cuSOLVER) | +| **Linear solver** | Supernodal block-sparse Cholesky (built in, default: exact steps; also solves the Schur complement of bundle adjustment), block-sparse PCG (variable block-Jacobi preconditioner), NVIDIA cuDSS (optional, loaded via `dlopen()` at runtime — see [C++ Installation](docs/sphinx/installation.rst)), block-tridiagonal Cholesky, dense LDLT, dense Cholesky (cuSOLVER), dense QR (cuSOLVER). Guide: [Linear systems and solvers](docs/sphinx/linear_solvers.rst) | | **Safety checks** | Optional runtime validation (linear-solver diagnostics and more), off by default for low-latency solves — enable with `MinimizerOptions::disable_safety_checks = false` when debugging | | **Execution model** | Fully asynchronous via CUDA streams | diff --git a/cunls/common/types.h b/cunls/common/types.h index 6833af3..d638b24 100644 --- a/cunls/common/types.h +++ b/cunls/common/types.h @@ -151,32 +151,6 @@ using pvector = PinnedVector; * - col_ids: column indices of non-zero entries (size = num_nonzeros). * - values: non-zero values (size = num_nonzeros). */ -/** - * @brief Cached sparse matrix dimensions. - * - * Stores num_rows, num_cols, and num_nonzeros so that repeated GPU scans - * (e.g. thrust::max_element + cudaStreamSynchronize) can be avoided in - * hot loops. Populated once during initialization via ExtractMatrixMetadata - * and reused thereafter. - */ -struct CSRMatrixDimensions { - int num_rows = -1; - int num_cols = -1; - int num_nonzeros = -1; - - void Set(int rows, int cols, int nnz) { - num_rows = rows; - num_cols = cols; - num_nonzeros = nnz; - } - - void Invalidate() { - num_rows = -1; - num_cols = -1; - num_nonzeros = -1; - } -}; - /** * @brief Compressed Sparse Row (CSR) matrix stored in GPU memory. * diff --git a/cunls/linear_solver/CMakeLists.txt b/cunls/linear_solver/CMakeLists.txt index f97ad20..deaed51 100644 --- a/cunls/linear_solver/CMakeLists.txt +++ b/cunls/linear_solver/CMakeLists.txt @@ -1,4 +1,6 @@ add_library(cunls_linear_solver OBJECT + block_sparse_cholesky_analysis.cpp + block_sparse_cholesky_solver.cu block_sparse_pcg_solver.cu block_tridiagonal_solver.cu dense_cholesky_solver.cu diff --git a/cunls/linear_solver/block_sparse_cholesky_analysis.cpp b/cunls/linear_solver/block_sparse_cholesky_analysis.cpp new file mode 100644 index 0000000..bf3f3f3 --- /dev/null +++ b/cunls/linear_solver/block_sparse_cholesky_analysis.cpp @@ -0,0 +1,703 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cunls/linear_solver/block_sparse_cholesky_analysis.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace cunls { +namespace internal { +namespace { + +/** Supervariable graph: consecutive columns with identical patterns, as one vertex. */ +struct Graph { + std::vector start; ///< Per vertex: first column; size + 1 entries. + std::vector weight; ///< Per vertex: columns. + std::vector> adj; ///< Per vertex: neighbors (no self). + int Size() const { return static_cast(weight.size()); } +}; + +/** A 64-bit mix (splitmix64 finalizer) for order-independent set hashes. */ +uint64_t Mix(uint64_t x) { + x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull; + x = (x ^ (x >> 27)) * 0x94D049BB133111EBull; + return x ^ (x >> 31); +} + +/** + * Approximate minimum degree ordering of the subgraph induced by `nodes`, on + * the quotient graph: an eliminated vertex becomes an element standing for + * the clique of its neighbors, so the fill is never formed. Each variable + * keeps its adjacent elements and the variables it still reaches directly; + * its external degree (in columns) is bounded from above by its direct + * neighbors, the new element and, for every other adjacent element e, the + * part of e outside the new element. Elements inside the new element are + * absorbed, and variables with identical adjacency are merged + * (indistinguishable: they are eliminated together). `local` is scratch of + * size g.Size() filled with -1, and left so. + */ +std::vector MinimumDegree(const Graph &g, const std::vector &nodes, + std::vector &local) { + const int m = static_cast(nodes.size()); + for (int i = 0; i < m; ++i) local[nodes[i]] = i; + std::vector> vars(m), elems(m); // Per variable: adjacency. + std::vector> members(m); // Per element: its variables (dead ones skipped). + std::vector nv(m), degree(m, 0), element_weight(m, 0), outside(m, 0); + int64_t remaining = 0; // Columns of the variables not eliminated yet. + size_t edges = 0; + for (int i = 0; i < m; ++i) { + nv[i] = g.weight[nodes[i]]; + remaining += nv[i]; + for (int u : g.adj[nodes[i]]) { + if (local[u] >= 0) { + vars[i].push_back(local[u]); + degree[i] += g.weight[u]; + } + } + edges += vars[i].size(); + } + // Variables whose lists grow far beyond the average degree are updated lazily (below). + const size_t long_lists = std::max(64, 16 * edges / std::max(m, 1)); + for (int i = 0; i < m; ++i) local[nodes[i]] = -1; + // State per index: variable (principal), merged into another variable, or element. + enum : char { kVariable, kMerged, kElement }; + std::vector state(m, kVariable); + std::vector merged_next(m, -1); // Variables merged into a principal one (a list). + // Degree buckets (doubly linked), degrees capped at the remaining columns. + std::vector head(remaining + 1, -1), next(m, -1), prev(m, -1); + int64_t min_degree = 0; + auto insert = [&](int i) { + const int64_t d = degree[i]; + next[i] = head[d]; + prev[i] = -1; + if (head[d] >= 0) prev[head[d]] = i; + head[d] = i; + min_degree = std::min(min_degree, d); + }; + auto remove = [&](int i) { + if (prev[i] >= 0) { + next[prev[i]] = next[i]; + } else { + head[degree[i]] = next[i]; + } + if (next[i] >= 0) prev[next[i]] = prev[i]; + }; + for (int i = 0; i < m; ++i) { + degree[i] = std::min(degree[i], remaining - nv[i]); + insert(i); + } + std::vector mark(m, -1), stamp(m, -1); + std::vector order; + order.reserve(m); + std::vector lp; + std::vector> hashes; + std::vector heavy; + std::vector pruned(m, 0); // Per variable: list lengths after its last pruning. + for (int pass = 0; static_cast(order.size()) < m; ++pass) { + while (head[min_degree] < 0) ++min_degree; + const int p = head[min_degree]; + remove(p); + // The new element: variables of p's elements (absorbed) and its own variables. + lp.clear(); + mark[p] = pass; + int64_t lp_weight = 0; + auto take = [&](int v) { + if (state[v] == kVariable && mark[v] != pass) { + mark[v] = pass; + lp.push_back(v); + lp_weight += nv[v]; + } + }; + for (int e : elems[p]) { + if (state[e] != kElement) continue; // absorbed already (stale entry) + for (int v : members[e]) take(v); + members[e].clear(); + members[e].shrink_to_fit(); + state[e] = kMerged; // absorbed into p + } + for (int v : vars[p]) take(v); + vars[p].clear(); + vars[p].shrink_to_fit(); + elems[p].clear(); + elems[p].shrink_to_fit(); + state[p] = kElement; + members[p] = lp; + element_weight[p] = lp_weight; + remaining -= nv[p]; + for (int v = p; v >= 0; v = merged_next[v]) order.push_back(nodes[v]); + // Prune the adjacency of the element's variables and take them out of the + // buckets. A variable with lists far longer than the average degree (a + // pose of bundle adjustment collects an element per eliminated landmark) + // is pruned only when its lists doubled, keeps the cheap degree bound and + // is not merged: scanning its lists at every elimination would be + // quadratic. + heavy.assign(lp.size(), 0); + for (size_t k = 0; k < lp.size(); ++k) { + const int i = lp[k]; + remove(i); + auto &e = elems[i]; + auto &a = vars[i]; + heavy[k] = e.size() + a.size() > long_lists; + if (!heavy[k] || e.size() + a.size() >= 2 * pruned[i]) { + e.erase(std::remove_if(e.begin(), e.end(), [&](int x) { return state[x] != kElement; }), + e.end()); + a.erase(std::remove_if(a.begin(), a.end(), + [&](int x) { return state[x] != kVariable || mark[x] == pass; }), + a.end()); + pruned[i] = e.size() + a.size() + 1; + } + e.push_back(p); + } + // |e \ Lp| for every other element adjacent to the new one (without the + // heavy variables: an overestimate, which only bounds degrees from above + // and skips absorptions). + for (size_t k = 0; k < lp.size(); ++k) { + if (heavy[k]) continue; + const int i = lp[k]; + for (int e : elems[i]) { + if (e == p) continue; + if (stamp[e] != pass) { + stamp[e] = pass; + outside[e] = element_weight[e]; + } + outside[e] -= nv[i]; + } + } + // Approximate external degrees; elements inside Lp are absorbed. + for (size_t k = 0; k < lp.size(); ++k) { + const int i = lp[k]; + const int64_t bound = std::min(degree[i] + lp_weight - nv[i], remaining - nv[i]); + if (heavy[k]) { + degree[i] = bound; + continue; + } + int64_t d = lp_weight - nv[i]; + auto &e = elems[i]; + size_t keep = 0; + for (int x : e) { + if (x != p && (state[x] != kElement || outside[x] <= 0)) { + state[x] = kMerged; // aggressive absorption: x ⊂ Lp + continue; + } + if (x != p) d += outside[x]; + e[keep++] = x; + } + e.resize(keep); + for (int v : vars[i]) d += nv[v]; + degree[i] = std::min(d, bound); + } + // Indistinguishable variables: identical elements and variables (an + // order-independent hash, then an exact comparison). + hashes.clear(); + for (size_t k = 0; k < lp.size(); ++k) { + if (heavy[k]) continue; + const int i = lp[k]; + uint64_t h = elems[i].size() * 0x9E3779B97F4A7C15ull + vars[i].size(); + for (int x : elems[i]) h += Mix(static_cast(x)); + for (int x : vars[i]) h += Mix(static_cast(x) + (1ull << 40)); + hashes.push_back({h, i}); + } + std::sort(hashes.begin(), hashes.end()); + for (size_t a = 0; a < hashes.size();) { + size_t b = a + 1; + while (b < hashes.size() && hashes[b].first == hashes[a].first) ++b; + if (b - a > 1) { + for (size_t x = a; x < b; ++x) { + const int i = hashes[x].second; + std::sort(elems[i].begin(), elems[i].end()); + std::sort(vars[i].begin(), vars[i].end()); + } + } + for (size_t x = a; x < b; ++x) { + const int i = hashes[x].second; + if (state[i] != kVariable) continue; + for (size_t y = x + 1; y < b; ++y) { + const int j = hashes[y].second; + if (state[j] != kVariable || elems[i] != elems[j] || vars[i] != vars[j]) continue; + // Merge j into i: j is eliminated right after i. + int tail = i; + while (merged_next[tail] >= 0) tail = merged_next[tail]; + merged_next[tail] = j; + nv[i] += nv[j]; + degree[i] -= nv[j]; + nv[j] = 0; + state[j] = kMerged; + elems[j].clear(); + vars[j].clear(); + } + } + a = b; + } + for (int i : lp) { + if (state[i] != kVariable) continue; + degree[i] = std::max(0, std::min(degree[i], remaining - nv[i])); + insert(i); + } + // Drop the merged variables from the element. + auto &mp = members[p]; + mp.erase(std::remove_if(mp.begin(), mp.end(), [&](int x) { return state[x] != kVariable; }), + mp.end()); + } + return order; +} + +/** Nested dissection with breadth-first level-set separators and minimum-degree leaves. */ +class Dissection { + public: + Dissection(const Graph &g, int leaf_weight) + : g_(g), leaf_(leaf_weight), set_(g.Size(), -1), level_(g.Size(), -1), local_(g.Size(), -1) {} + + std::vector Order() { + std::vector all(g_.Size()); + for (int v = 0; v < g_.Size(); ++v) all[v] = v; + Recurse(std::move(all)); + return std::move(order_); + } + + private: + int64_t Weight(const std::vector &nodes) const { + int64_t w = 0; + for (int v : nodes) w += g_.weight[v]; + return w; + } + + /** Breadth-first search from `root` within the set `id`; levels in level_, returns the visit + * order. */ + std::vector Bfs(int root, int id, int visit) { + std::vector visited{root}; + level_[root] = 0; + set_[root] = visit; + for (size_t head = 0; head < visited.size(); ++head) { + const int u = visited[head]; + for (int v : g_.adj[u]) { + if (set_[v] == id) { + set_[v] = visit; + level_[v] = level_[u] + 1; + visited.push_back(v); + } + } + } + for (int v : visited) set_[v] = id; + return visited; + } + + /** + * Orders a part that is not split further: minimum degree, unless the part + * is nearly complete (any order then gives a dense factor, and minimum + * degree would cost the cube of its size). + */ + void Leaf(const std::vector &nodes) { + const int m = static_cast(nodes.size()); + for (int v : nodes) local_[v] = 0; + int64_t edges = 0; + for (int v : nodes) { + for (int u : g_.adj[v]) edges += local_[u] == 0; + } + for (int v : nodes) local_[v] = -1; + if (edges >= static_cast(m) * (m - 1) * 3 / 4) { + std::vector o = nodes; + std::sort(o.begin(), o.end()); + order_.insert(order_.end(), o.begin(), o.end()); + return; + } + std::vector o = MinimumDegree(g_, nodes, local_); + order_.insert(order_.end(), o.begin(), o.end()); + } + + void Recurse(std::vector nodes) { + if (nodes.empty()) return; + const int64_t total = Weight(nodes); + if (total <= leaf_) { + Leaf(nodes); + return; + } + const int id = next_id_++; + for (int v : nodes) set_[v] = id; + // Connected components are ordered independently. + std::vector first = Bfs(nodes[0], id, next_id_++); + if (first.size() < nodes.size()) { + const int done = next_id_++; + std::vector> components; + for (int v : nodes) { + if (set_[v] != id) continue; + components.push_back(Bfs(v, id, next_id_++)); + for (int u : components.back()) set_[u] = done; + } + for (auto &c : components) Recurse(std::move(c)); + return; + } + // Pseudo-peripheral root: restart from the farthest vertex while the depth grows. + std::vector visit = std::move(first); + for (int attempt = 0; attempt < 4; ++attempt) { + const int depth = level_[visit.back()]; + std::vector again = Bfs(visit.back(), id, next_id_++); + if (level_[again.back()] <= depth) { + visit = Bfs(visit.front(), id, next_id_++); // restore the deeper levels + break; + } + visit = std::move(again); + } + const int depth = level_[visit.back()]; + if (depth < 2) { + Leaf(nodes); + return; + } + std::vector level_weight(depth + 1, 0); + for (int v : visit) level_weight[level_[v]] += g_.weight[v]; + // The separator level minimizing its weight relative to the product of + // the halves (a ratio cut): small separators near the root would peel + // the part off a few vertices at a time. + int64_t before = 0; + double best = std::numeric_limits::max(); + int cut = -1; + for (int l = 1; l < depth; ++l) { + before += level_weight[l - 1]; + const int64_t sep = level_weight[l], after = total - before - sep; + if (after <= 0) break; + double score = + static_cast(sep) / (static_cast(before) * static_cast(after)); + if (8 * std::min(before, after) < total) score *= 1e6; // unbalanced: a last resort + if (score < best) { + best = score; + cut = l; + } + } + if (2 * level_weight[cut] > total) { + Leaf(nodes); + return; + } + std::vector a, b, sep; + for (int v : visit) { + const int l = level_[v]; + (l < cut ? a : l > cut ? b : sep).push_back(v); + } + Recurse(std::move(a)); + Recurse(std::move(b)); + Leaf(sep); + } + + const Graph &g_; + const int leaf_; + std::vector set_, level_, local_; + std::vector order_; + int next_id_ = 0; +}; + +/** Block (supervariable) structure of L for an ordering. */ +struct Symbolic { + std::vector order; ///< Position -> vertex. + std::vector> rows; ///< Per position: positions below the diagonal, ascending. + std::vector parent; ///< Per position: elimination-tree parent, or -1. + std::vector> supernodes; ///< [first, last] positions. + double flops = 0.0; + int chain = 0; ///< Longest etree path counted in panels (sequential steps). +}; + +Symbolic Factorize(const Graph &g, std::vector order) { + const int nb = g.Size(); + Symbolic s; + std::vector pos(nb); + for (int k = 0; k < nb; ++k) pos[order[k]] = k; + s.rows.resize(nb); + s.parent.assign(nb, -1); + std::vector> children(nb); + std::vector mark(nb, -1); + for (int k = 0; k < nb; ++k) { + std::vector &r = s.rows[k]; + mark[k] = k; + for (int u : g.adj[order[k]]) { + const int p = pos[u]; + if (p > k && mark[p] != k) { + mark[p] = k; + r.push_back(p); + } + } + for (int c : children[k]) { + for (int p : s.rows[c]) { + if (mark[p] != k) { + mark[p] = k; + r.push_back(p); + } + } + } + std::sort(r.begin(), r.end()); + if (!r.empty()) { + s.parent[k] = r.front(); + children[r.front()].push_back(k); + } + } + // Fundamental supernodes, then relaxed amalgamation: a supernode merges + // into its parent supernode (the next one, when it is its parent) if that + // saves a panel step and the explicit zeros stay under kRelaxedZeros of the + // merged entries. Latency per step dominates on the GPU, flops are cheap. + constexpr double kRelaxedZeros = 0.5; + std::vector wo(nb); + for (int k = 0; k < nb; ++k) wo[k] = g.weight[order[k]]; + auto below_of = [&](int k) { + double b = 0; + for (int p : s.rows[k]) b += wo[p]; + return b; + }; + auto panels = [](double width) { + return static_cast((width + kCholeskyPanelWidth - 1) / kCholeskyPanelWidth); + }; + struct Node { + int first, last; + double width, below, zeros; + }; + std::vector nodes; + for (int k = 0; k < nb; ++k) { + const bool extends = + k > 0 && s.parent[k - 1] == k && s.rows[k - 1].size() == s.rows[k].size() + 1; + if (extends) { + nodes.back().last = k; + nodes.back().width += wo[k]; + nodes.back().below = below_of(k); + } else { + nodes.push_back({k, k, static_cast(wo[k]), below_of(k), 0.0}); + } + } + std::vector merged; + for (const Node &node : nodes) { + if (!merged.empty()) { + Node &a = merged.back(); + if (s.parent[a.last] == node.first) { + const double width = a.width + node.width; + const double zeros = a.zeros + node.zeros + a.width * (node.width + node.below - a.below); + const double entries = width * (width + 1) / 2 + width * node.below; + if (panels(a.width) + panels(node.width) > panels(width) && + zeros <= kRelaxedZeros * entries) { + a = {a.first, node.last, width, node.below, zeros}; + continue; + } + } + } + merged.push_back(node); + } + for (const Node &node : merged) { + s.supernodes.push_back({node.first, node.last}); + for (int i = 0; i < node.width; ++i) { + const double c = node.below + node.width - i; + s.flops += c * c; + } + } + std::vector supernode_of(nb); + for (size_t i = 0; i < s.supernodes.size(); ++i) { + for (int k = s.supernodes[i].first; k <= s.supernodes[i].second; ++k) supernode_of[k] = i; + } + std::vector critical(s.supernodes.size(), 0); + for (size_t i = 0; i < s.supernodes.size(); ++i) { + critical[i] += panels(merged[i].width); + const int p = s.parent[s.supernodes[i].second]; + if (p >= 0) { + int &c = critical[supernode_of[p]]; + c = std::max(c, critical[i]); + } + s.chain = std::max(s.chain, critical[i]); + } + s.order = std::move(order); + return s; +} + +/** Estimated factorization time [us]: a fixed latency per sequential panel step plus the flops. */ +double EstimatedMicroseconds(const Symbolic &s) { + constexpr double kStepLatency = 40.0; // the dependent kernels of a level (measured) + constexpr double kFlopsPerMicrosecond = 5e6; + return s.chain * kStepLatency + s.flops / kFlopsPerMicrosecond; +} + +} // namespace + +namespace { + +/** The supervariable graph from its CSR adjacency. */ +Graph MakeGraph(const std::vector &supervariable_start, + const std::vector &adjacency_offsets, const std::vector &adjacency) { + Graph g; + g.start = supervariable_start; + const int nb = static_cast(g.start.size()) - 1; + g.weight.resize(nb); + g.adj.resize(nb); + for (int v = 0; v < nb; ++v) { + g.weight[v] = g.start[v + 1] - g.start[v]; + g.adj[v].assign(adjacency.begin() + adjacency_offsets[v], + adjacency.begin() + adjacency_offsets[v + 1]); + } + return g; +} + +/** + * The symbolic factorization of the cheaper of two orderings by + * EstimatedMicroseconds(): nested dissection, and minimum degree unless the + * graph is nearly complete (any ordering then gives a dense factor). Sets the + * ordering's name in `name`. + */ +Symbolic ChooseOrdering(const Graph &g, int64_t num_edges, const char **name) { + const int nb = g.Size(); + Symbolic best = Factorize(g, Dissection(g, kCholeskyPanelWidth).Order()); + *name = "nested dissection"; + if (num_edges < static_cast(nb) * nb / 4) { + std::vector all(nb), local(nb, -1); + for (int v = 0; v < nb; ++v) all[v] = v; + Symbolic md = Factorize(g, MinimumDegree(g, all, local)); + if (EstimatedMicroseconds(md) < EstimatedMicroseconds(best)) { + best = std::move(md); + *name = "minimum degree"; + } + } + return best; +} + +/** + * The scalar permutation (supervariables in order, their columns consecutive) + * and the panels: supernodes cut into kCholeskyPanelWidth-column pieces. + * Returns the panel of each (permuted) column. + */ +std::vector BuildPanels(const Graph &g, const Symbolic &s, CholeskyPlan &plan) { + const int nb = g.Size(); + std::vector first(nb + 1, 0); // per position: first permuted column + for (int k = 0; k < nb; ++k) first[k + 1] = first[k] + g.weight[s.order[k]]; + plan.perm.resize(plan.n); + for (int k = 0; k < nb; ++k) { + const int v = s.order[k]; + for (int t = 0; t < g.weight[v]; ++t) plan.perm[first[k] + t] = g.start[v] + t; + } + std::vector col_panel(plan.n); + for (const auto &[k0, k1] : s.supernodes) { + const int c0 = first[k0], c1 = first[k1 + 1]; + std::vector below; // the rows below the supernode + for (int p : s.rows[k1]) { + for (int c = first[p]; c < first[p + 1]; ++c) below.push_back(c); + } + for (int c = c0; c < c1; c += kCholeskyPanelWidth) { + CholeskyPlan::Panel panel; + panel.col = c; + panel.width = std::min(kCholeskyPanelWidth, c1 - c); + panel.rows_offset = static_cast(plan.below_rows.size()); + for (int r = c + panel.width; r < c1; ++r) plan.below_rows.push_back(r); + plan.below_rows.insert(plan.below_rows.end(), below.begin(), below.end()); + panel.height = panel.width + static_cast(plan.below_rows.size()) - panel.rows_offset; + panel.offset = plan.factor_size; + panel.diagonal_offset = plan.diagonal_size; + plan.factor_size += static_cast(panel.height) * panel.width; + plan.diagonal_size += static_cast(panel.width) * panel.width; + std::fill(col_panel.begin() + c, col_panel.begin() + c + panel.width, + static_cast(plan.panels.size())); + plan.panels.push_back(panel); + } + } + if (plan.factor_size > INT_MAX) { + throw std::runtime_error("BlockSparseCholeskySolver: the factor exceeds 2^31 entries"); + } + return col_panel; +} + +/** + * The level schedule: a panel's level is one more than the highest level of + * the panels that update it. Then per level the factor items, the runs of + * each panel's below rows by target panel with their update items, and + * whether no two panels of the level update the same target. + */ +void Schedule(const std::vector &col_panel, CholeskyPlan &plan) { + const int np = static_cast(plan.panels.size()); + auto below_of = [&](int p) { + const auto &panel = plan.panels[p]; + return std::make_pair(plan.below_rows.data() + panel.rows_offset, panel.height - panel.width); + }; + std::vector level(np, 0); + for (int p = 0; p < np; ++p) { + const auto [rows, count] = below_of(p); + if (count > 0) { + int &l = level[col_panel[rows[0]]]; + l = std::max(l, level[p] + 1); + } + } + plan.num_levels = np > 0 ? *std::max_element(level.begin(), level.end()) + 1 : 0; + + std::vector> factor_by_level(plan.num_levels); + std::vector> update_by_level(plan.num_levels); + std::vector> targets; // (level, target, source) + int64_t relmap_size = 0; + for (int p = 0; p < np; ++p) { + const int l = level[p]; + const auto [rows, count] = below_of(p); + if (count == 0) { + factor_by_level[l].push_back({p, 0}); + continue; + } + for (int r = 0; r < count; r += kCholeskyTileRows) factor_by_level[l].push_back({p, r}); + for (int a = 0; a < count;) { + const int q = col_panel[rows[a]]; + int b = a + 1; + while (b < count && col_panel[rows[b]] == q) ++b; + const int run = static_cast(plan.runs.size()); + if (relmap_size + (count - a) > INT_MAX) { + throw std::runtime_error( + "BlockSparseCholeskySolver: the relative map exceeds 2^31 entries"); + } + plan.runs.push_back({p, q, a, b, static_cast(relmap_size)}); + relmap_size += count - a; + for (int r = a; r < count; r += kCholeskyTileRows) update_by_level[l].push_back({p, run, r}); + targets.emplace_back(l, q, p); + a = b; + } + } + plan.relmap_size = static_cast(relmap_size); + std::sort(targets.begin(), targets.end()); + plan.update_exclusive.assign(plan.num_levels, 1); + for (size_t i = 1; i < targets.size(); ++i) { + const auto &[l0, q0, p0] = targets[i - 1]; + const auto &[l1, q1, p1] = targets[i]; + if (l0 == l1 && q0 == q1 && p0 != p1) plan.update_exclusive[l1] = 0; + } + for (int l = 0; l < plan.num_levels; ++l) { + plan.factor_items.insert(plan.factor_items.end(), factor_by_level[l].begin(), + factor_by_level[l].end()); + plan.factor_offsets.push_back(static_cast(plan.factor_items.size())); + plan.update_items.insert(plan.update_items.end(), update_by_level[l].begin(), + update_by_level[l].end()); + plan.update_offsets.push_back(static_cast(plan.update_items.size())); + } +} + +} // namespace + +CholeskyPlan AnalyzeCholesky(int n, const std::vector &supervariable_start, + const std::vector &adjacency_offsets, + const std::vector &adjacency) { + CholeskyPlan plan; + plan.n = n; + plan.factor_offsets = {0}; + plan.update_offsets = {0}; + if (n == 0) return plan; + const Graph g = MakeGraph(supervariable_start, adjacency_offsets, adjacency); + plan.num_supervariables = g.Size(); + const Symbolic s = ChooseOrdering(g, static_cast(adjacency.size()), &plan.ordering); + plan.flops = s.flops; + plan.num_supernodes = static_cast(s.supernodes.size()); + Schedule(BuildPanels(g, s, plan), plan); + return plan; +} + +} // namespace internal +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_cholesky_analysis.h b/cunls/linear_solver/block_sparse_cholesky_analysis.h new file mode 100644 index 0000000..bd6ef73 --- /dev/null +++ b/cunls/linear_solver/block_sparse_cholesky_analysis.h @@ -0,0 +1,139 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +namespace cunls { +namespace internal { + +/** @brief Widest panel (columns) of the supernodal factor; one CTA factors a panel's diagonal. */ +constexpr int kCholeskyPanelWidth = 64; + +/** @brief Rows of a panel processed by one CTA (factor-solve and update tiles). */ +constexpr int kCholeskyTileRows = 64; + +/** + * @brief The symbolic analysis of a sparse SPD matrix for BlockSparseCholeskySolver: + * the fill-reducing ordering, the supernodal structure of the factor L, and + * the level schedule and index maps of the numeric phase. + * + * L is stored as panels: groups of at most kCholeskyPanelWidth consecutive + * columns (in the permuted order) with identical structure below their + * diagonal block. Panel p is a dense column-major height_p x width_p block + * whose first width_p rows are its own columns (the diagonal block, lower + * triangle used) and whose remaining rows are its below rows (`below_rows` + * from `rows_offset`, ascending, all after the panel's columns). The numeric + * phase processes the panels level by level: a panel's level is one more than + * the highest level among the panels that update it, so the panels of a level + * are independent. + */ +struct CholeskyPlan { + /** @brief A panel of L. */ + struct Panel { + int col; ///< First column (permuted order). + int width; ///< Columns, at most kCholeskyPanelWidth. + int height; ///< Rows: width plus the rows below the diagonal block. + int rows_offset; ///< Offset of its below rows in below_rows. + int64_t offset; ///< Offset of its column-major values in the factor storage. + /// Offset of its factored diagonal block L_D (row-major width x width, the + /// reciprocals of its diagonal on the diagonal). + int64_t diagonal_offset; + }; + + /** + * @brief A run of a panel's below rows that are the columns of one target + * panel: the update L_b L_bᵀ of the source panel restricted to those + * columns and to the rows from the run on goes to the target. + */ + struct Run { + int source; ///< Source panel. + int target; ///< Target panel. + int begin; ///< First below-row position (of the source) in the run. + int end; ///< One past the last. + int relmap_offset; ///< Offset in the relative map of the target positions of below rows + ///< begin... of the source (the map is built on the device). + }; + + /** @brief One CTA of the factor kernel: a panel and the first of its <= 64 below rows. */ + struct FactorItem { + int panel; ///< The panel. + int row_begin; ///< First of its below rows (position) handled by the CTA. + }; + + /** @brief One CTA of the update kernel: a run and the first of its <= 64 rows. */ + struct UpdateItem { + int panel; ///< The source panel. + int run; ///< Its run (index into runs): the target and the columns. + int row_begin; ///< First below row (position, at or after the run) handled by the CTA. + }; + + int n = 0; ///< Matrix size. + std::vector perm; ///< Permuted column -> original column. + std::vector panels; ///< Panels in column order. + std::vector below_rows; ///< Below rows (permuted indices, ascending) of all panels. + int64_t factor_size = 0; ///< Floats of the factor storage. + int64_t diagonal_size = 0; ///< Floats of the factored diagonal blocks. + std::vector runs; ///< Update runs of all panels. + int relmap_size = 0; ///< Entries of the relative map (see Run). + + int num_levels = 0; ///< Levels of the schedule. + std::vector factor_items; ///< Factor CTAs by level. + std::vector factor_offsets; ///< num_levels + 1 offsets into factor_items. + std::vector update_items; ///< Update CTAs by level. + std::vector update_offsets; ///< num_levels + 1 offsets into update_items. + /** @brief Per level: whether no two panels of the level update the same target (no atomics). */ + std::vector update_exclusive; + + // Statistics of the analysis (for logs and tests). + int num_supervariables = 0; ///< Groups of consecutive columns with identical structure. + int num_supernodes = 0; ///< Supernodes before splitting into panels. + double flops = 0.0; ///< Factorization flops (multiply-add counted as 2). + const char *ordering = ""; ///< Name of the chosen ordering. +}; + +/** + * @brief Analyzes the pattern of a symmetric matrix, given as its + * supervariables (groups of consecutive columns with identical patterns: a + * state's tangent block, or several states) and their graph. + * + * Two fill-reducing orderings of the supervariable graph are built and the + * one with the lower estimated factorization time kept (a fixed latency per + * sequential panel step plus the flops): + * - approximate minimum degree on the quotient graph (least fill on + * irregular graphs: covisibility, bundle adjustment, random loop + * closures); + * - nested dissection with balanced breadth-first level-set separators and + * minimum-degree leaves (short level schedules and less fill on chains + * and grids: trajectories, visual-inertial windows, pose graphs with + * local loop closures). + * Fundamental supernodes are amalgamated with their parents where that saves + * a panel step at no more than 50% explicit zeros, then cut into panels. + * + * @param n Matrix size. + * @param supervariable_start First column of every supervariable, then n. + * @param adjacency_offsets Per supervariable: offsets into @p adjacency (CSR). + * @param adjacency Neighbor supervariables (without the supervariable itself). + */ +CholeskyPlan AnalyzeCholesky(int n, const std::vector &supervariable_start, + const std::vector &adjacency_offsets, + const std::vector &adjacency); + +} // namespace internal +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_cholesky_solver.cu b/cunls/linear_solver/block_sparse_cholesky_solver.cu new file mode 100644 index 0000000..a4dd39f --- /dev/null +++ b/cunls/linear_solver/block_sparse_cholesky_solver.cu @@ -0,0 +1,817 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "cunls/common/helper.h" +#include "cunls/common/log.h" +#include "cunls/linear_solver/block_sparse_cholesky_solver.h" + +namespace cunls { +namespace { + +constexpr int kThreads = 256; +constexpr int kW = internal::kCholeskyPanelWidth; +constexpr int kR = internal::kCholeskyTileRows; +constexpr int kPad = kW + 1; // shared-memory row stride: conflict-free column access + +/** + * A pivot within this fraction of its original diagonal entry (either sign) + * is a rank deficiency, not a failure: about 100 float epsilons, the + * round-off floor of a pivot whose diagonal entry is cancelled by the + * updates (free gauges land near 1e-6). Its unknown is fixed at zero (a zero + * column with a zero reciprocal), which solves a consistent semidefinite + * system along a free gauge (Gauss-Newton without an anchored state). + * Regular pivots can be far below their diagonal entries (strongly coupled + * states, e.g. motion priors: 1e-4), so the threshold stays tight. + */ +constexpr float kZeroPivot = 1e-5f; + +using Panel = internal::CholeskyPlan::Panel; +using Run = internal::CholeskyPlan::Run; +using FactorItem = internal::CholeskyPlan::FactorItem; +using UpdateItem = internal::CholeskyPlan::UpdateItem; + +unsigned Blocks(size_t n) { return static_cast((n + kThreads - 1) / kThreads); } + +/** factor[dst[e]] = values[e] for the nonzeros with dst[e] >= 0 (the lower triangle of PAPᵀ). */ +__global__ void AssembleKernel(int nnz, const int *dst, const float *values, float *factor) { + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e < nnz) { + const int d = dst[e]; + if (d >= 0) factor[d] = values[e]; + } +} + +/** One warp per row: flag[i] = 1 if row i starts a supervariable (its pattern differs from row i - + * 1). */ +__global__ void SupervariableKernel(int n, const int *row_offsets, const int *col_ids, int *flag) { + const int i = (blockIdx.x * blockDim.x + threadIdx.x) / 32, lane = threadIdx.x % 32; + if (i >= n) return; + bool differs = i == 0; + if (i > 0) { + const int a = row_offsets[i - 1], b = row_offsets[i], c = row_offsets[i + 1]; + if (c - b != b - a) { + differs = true; + } else { + bool d = false; + for (int k = lane; k < b - a; k += 32) d = d || col_ids[a + k] != col_ids[b + k]; + differs = __any_sync(0xffffffffu, d); + } + } + if (lane == 0) flag[i] = differs; +} + +/** + * One warp per supervariable v: its neighbors, from the pattern of its first + * row (column indices ascending, so a neighbor's columns are consecutive). + * Counts them when `adjacency` is null, else writes them at offsets[v]. + */ +__global__ void AdjacencyKernel(int nb, const int *first_row, const int *row_offsets, + const int *col_ids, const int *supervariable, const int *offsets, + int *count, int *adjacency) { + const int v = (blockIdx.x * blockDim.x + threadIdx.x) / 32, lane = threadIdx.x % 32; + if (v >= nb) return; + const int r = first_row[v], begin = row_offsets[r], end = row_offsets[r + 1]; + int written = 0; + for (int base = begin; base < end; base += 32) { + const int e = base + lane; + bool keep = false; + int u = -1; + if (e < end) { + u = supervariable[col_ids[e]]; + keep = u != v && (e == begin || u != supervariable[col_ids[e - 1]]); + } + const unsigned mask = __ballot_sync(0xffffffffu, keep); + if (adjacency != nullptr && keep) { + adjacency[offsets[v] + written + __popc(mask & ((1u << lane) - 1))] = u; + } + written += __popc(mask); + } + if (adjacency == nullptr && lane == 0) count[v] = written; +} + +/** + * One warp per row i of A: the factor offset of every nonzero (i, j) whose + * permuted position (i', j') is in the lower triangle, else -1. The row i' + * is found in the panel of column j' by arithmetic (its own columns) or by + * binary search in its below rows. diagonal_index[i'] = the nonzero (i, i). + */ +__global__ void AssemblyMapKernel(int n, const int *row_offsets, const int *col_ids, + const int *iperm, const int *col_panel, const Panel *panels, + const int *below_rows, int *dst, int *diagonal_index) { + const int i = (blockIdx.x * blockDim.x + threadIdx.x) / 32, lane = threadIdx.x % 32; + if (i >= n) return; + const int ip = iperm[i]; + for (int e = row_offsets[i] + lane; e < row_offsets[i + 1]; e += 32) { + const int jp = iperm[col_ids[e]]; + int d = -1; + if (ip >= jp) { + const int p = col_panel[jp]; + const Panel panel = panels[p]; + int pos; + if (ip < panel.col + panel.width) { + pos = ip - panel.col; + } else { + const int *rows = below_rows + panel.rows_offset; + int lo = 0, hi = panel.height - panel.width; + while (lo < hi) { + const int mid = (lo + hi) / 2; + if (rows[mid] < ip) { + lo = mid + 1; + } else { + hi = mid; + } + } + pos = panel.width + lo; + } + d = static_cast(panel.offset + static_cast(jp - panel.col) * panel.height + + pos); + if (ip == jp) diagonal_index[ip] = e; + } + dst[e] = d; + } +} + +/** + * The relative map: for run r (one warp), the positions in its target panel + * of the source's below rows from the run on (the target's own columns by + * arithmetic, its below rows by binary search). + */ +__global__ void RelmapKernel(int num_runs, const Run *runs, const Panel *panels, + const int *below_rows, int *relmap) { + const int r = (blockIdx.x * blockDim.x + threadIdx.x) / 32, lane = threadIdx.x % 32; + if (r >= num_runs) return; + const Run run = runs[r]; + const Panel source = panels[run.source], target = panels[run.target]; + const int *rows = below_rows + source.rows_offset, *target_rows = below_rows + target.rows_offset; + const int count = source.height - source.width, target_count = target.height - target.width; + for (int i = run.begin + lane; i < count; i += 32) { + const int row = rows[i]; + int pos; + if (row < target.col + target.width) { + pos = row - target.col; + } else { + int lo = 0, hi = target_count; + while (lo < hi) { + const int mid = (lo + hi) / 2; + if (target_rows[mid] < row) { + lo = mid + 1; + } else { + hi = mid; + } + } + pos = target.width + lo; + } + relmap[run.relmap_offset + i - run.begin] = pos; + } +} + +/** + * out[i] = in[perm[i]], and original[i] = the diagonal entry of A at permuted + * column i (values[diagonal_index[i]], 0 without one). + */ +__global__ void PermuteKernel(int n, const int *perm, const float *in, float *out, + const int *diagonal_index, const float *values, float *original) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + out[i] = in[perm[i]]; + const int e = diagonal_index[i]; + original[i] = e >= 0 ? values[e] : 0.f; + } +} + +/** out[perm[i]] = in[i]. */ +__global__ void UnpermuteKernel(int n, const int *perm, const float *in, float *out) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[perm[i]] = in[i]; +} + +/** + * One CTA per (panel, 64 below rows) of a level: factors the stacked panel + * [D; B; y_pᵀ] (D the diagonal block, B the CTA's rows below it, y_p the + * panel's part of the forward-substitution vector): L_D = chol(D), B L_D^{-T} + * and, as the last row, z_pᵀ = y_pᵀ L_D^{-T} (the forward substitution of the + * panel comes out of the row solve). Then y[rows of B] -= (B L_D^{-T}) z_p. + * Every CTA of a panel factors D redundantly; the one with row_begin 0 + * stores L_D (row-major, in `diagonal`, reciprocals on its diagonal) and z_p, + * each stores its rows of B L_D^{-T} in place. y_p is final: the panels + * below finished in earlier levels. + * + * Blocked by 16 columns: one warp factors the 16 x 16 diagonal block in + * groups of 4 columns (a lane per row), a thread per row solves the rows below it, and + * the trailing columns take the rank-16 update, 3 x 7 outputs per thread + * with the 16 multipliers of its columns in registers. The upper triangle of + * D takes garbage that is never read. + */ +__global__ void __launch_bounds__(kThreads) + FactorKernel(const FactorItem *items, const Panel *panels, const int *below_rows, + const float *original, float *factor, float *diagonal, float *y, float *z, + int *failure) { + constexpr int kRows = kW + kR; // stacked rows, plus y_pᵀ + constexpr int kB = 16; // column block + __shared__ float tile[kRows + 1][kPad]; + __shared__ float reciprocal[kW]; // 1 / L_D(k, k) + __shared__ __align__(16) float corner[4][4]; // the 4 x 4 diagonal block of a column group + __shared__ __align__(16) float group[4][kB]; // the 4 columns of a column group + __shared__ __align__(16) float limit[kW]; // zero-pivot limits (kZeroPivot) + __shared__ float partial[kThreads / kR][kR]; + const FactorItem item = items[blockIdx.x]; + const Panel panel_info = panels[item.panel]; + const int col = panel_info.col, w = panel_info.width, h = panel_info.height; + const int below = max(0, min(kR, h - w - item.row_begin)); + const int rows = w + below + 1; // y_pᵀ last + float *panel = factor + panel_info.offset; + const int t = threadIdx.x, lane = t % 32; + { + // Lanes along rows (coalesced), all loads in flight before the stores. + const int i = t % kRows, j0 = t / kRows; + const float *src = panel + (i < w ? i : item.row_begin + i); + float v[kW / 2]; +#pragma unroll + for (int m = 0; m < kW / 2; ++m) { + const int j = j0 + 2 * m; + v[m] = i < rows - 1 && j < w ? src[j * h] : 0.f; + } + if (i < rows - 1) { +#pragma unroll + for (int m = 0; m < kW / 2; ++m) tile[i][j0 + 2 * m] = v[m]; + } + } + if (t < w) tile[rows - 1][t] = y[col + t]; + if (t < kW) limit[t] = t < w ? kZeroPivot * original[col + t] : 0.f; + __syncthreads(); + for (int kb = 0; kb < w; kb += kB) { + const int bw = min(kB, w - kb); + if (t < 32) { + // 16 x 16 diagonal block: lane i holds row kb + i. + float r[kB]; +#pragma unroll + for (int c = 0; c < kB; ++c) r[c] = lane < bw && c <= lane ? tile[kb + lane][kb + c] : 0.f; + float lim[kB]; +#pragma unroll + for (int q = 0; q < kB / 4; ++q) { + const float4 v = reinterpret_cast(limit + kb)[q]; + lim[4 * q] = v.x; + lim[4 * q + 1] = v.y; + lim[4 * q + 2] = v.z; + lim[4 * q + 3] = v.w; + } + // Groups of 4 columns, two warp round trips each: the owner lanes + // publish the group's 4 x 4 diagonal block, every lane factors it in + // registers and solves its own row against it, then the 4 columns are + // published and every row takes the rank-4 update. Lane c keeps the + // reciprocal and the failure of column c (written after the loop). + float own_scale = 0.f; + bool own_failure = false; +#pragma unroll + for (int c0 = 0; c0 < kB; c0 += 4) { + if (c0 < bw) { + if (lane >= c0 && lane < c0 + 4) { +#pragma unroll + for (int b = 0; b < 4; ++b) corner[lane - c0][b] = r[c0 + b]; + } + __syncwarp(); + float a[4][4]; +#pragma unroll + for (int x = 0; x < 4; ++x) { + const float4 v = reinterpret_cast(corner[x])[0]; + a[x][0] = v.x; + a[x][1] = v.y; + a[x][2] = v.z; + a[x][3] = v.w; + } + // The group's block: L (lower), scales 1 / L(b, b); a vanishing + // pivot (rank deficiency) fixes its unknown: zero scale. + float lg[4][4], scale[4]; +#pragma unroll + for (int b = 0; b < 4; ++b) { + float pivot = a[b][b]; +#pragma unroll + for (int q = 0; q < b; ++q) pivot -= lg[b][q] * lg[b][q]; + const bool zero = fabsf(pivot) <= lim[c0 + b]; + asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(scale[b]) : "f"(pivot)); + if (zero) scale[b] = 0.f; + const bool own = lane == c0 + b; + own_scale = own ? scale[b] : own_scale; + own_failure |= own && !zero && (!(pivot > 0.f) || !isfinite(pivot)); +#pragma unroll + for (int x = b + 1; x < 4; ++x) { + float v = a[x][b]; +#pragma unroll + for (int q = 0; q < b; ++q) v -= lg[x][q] * lg[b][q]; + lg[x][b] = v * scale[b]; + } + } + // The lane's own row of the group's columns (zero above the diagonal). + float xv[4]; +#pragma unroll + for (int b = 0; b < 4; ++b) { + float v = r[c0 + b]; +#pragma unroll + for (int q = 0; q < b; ++q) v -= xv[q] * lg[b][q]; + xv[b] = lane >= c0 + b ? v * scale[b] : 0.f; + } + if (lane < kB) { +#pragma unroll + for (int b = 0; b < 4; ++b) group[b][lane] = xv[b]; + } + __syncwarp(); +#pragma unroll + for (int b = 0; b < 4; ++b) { + float lc[kB]; +#pragma unroll + for (int q = c0 / 4 + 1; q < kB / 4; ++q) { + const float4 v = reinterpret_cast(group[b])[q]; + lc[4 * q] = v.x; + lc[4 * q + 1] = v.y; + lc[4 * q + 2] = v.z; + lc[4 * q + 3] = v.w; + } +#pragma unroll + for (int j = c0 + 4; j < kB; ++j) r[j] -= xv[b] * lc[j]; + if (lane >= c0 + b) r[c0 + b] = xv[b]; + } + } + } + if (lane < bw) reciprocal[kb + lane] = own_scale; + if (own_failure) *failure = 1; + if (lane < bw) { +#pragma unroll + for (int c = 0; c < kB; ++c) { + if (c <= lane) tile[kb + lane][kb + c] = r[c]; + } + } + } + __syncthreads(); + // Rows below the block: x L_blockᵀ = row, one thread per row. + const int first = kb + bw; + if (t < rows - first) { + const int i = first + t; + float x[kB]; +#pragma unroll + for (int c = 0; c < kB; ++c) { + if (c < bw) { + float acc0 = tile[i][kb + c], acc1 = 0.f; // two chains +#pragma unroll + for (int j = 0; j + 1 < c; j += 2) { + acc0 -= x[j] * tile[kb + c][kb + j]; + acc1 -= x[j + 1] * tile[kb + c][kb + j + 1]; + } + if (c % 2 == 1) acc0 -= x[c - 1] * tile[kb + c][kb + c - 1]; + x[c] = (acc0 + acc1) * reciprocal[kb + c]; + tile[i][kb + c] = x[c]; + } + } + } + __syncthreads(); + // Trailing update of columns first..w-1, rows first..rows-1. + if (first < w) { + const int jj = t % 16, ii = t / 16; + float lj[3][kB]; +#pragma unroll + for (int q = 0; q < 3; ++q) { + const int j = first + jj + 16 * q; +#pragma unroll + for (int c = 0; c < kB; ++c) lj[q][c] = j < w && c < bw ? tile[j][kb + c] : 0.f; + } +#pragma unroll + for (int m = 0; m < 8; ++m) { + const int i = first + ii + 16 * m; + if (i < rows) { + float li[kB]; +#pragma unroll + for (int c = 0; c < kB; ++c) li[c] = tile[i][kb + c]; +#pragma unroll + for (int q = 0; q < 3; ++q) { + const int j = first + jj + 16 * q; + float acc = 0.f; +#pragma unroll + for (int c = 0; c < kB; ++c) acc += li[c] * lj[q][c]; + if (j < w) tile[i][j] -= acc; + } + } + } + } + __syncthreads(); + } + const float *zp = tile[rows - 1]; + if (item.row_begin == 0) { + if (t < w) z[col + t] = zp[t]; + float *d = diagonal + panel_info.diagonal_offset; + const int c = t % kW; +#pragma unroll 4 + for (int i = t / kW; i < w; i += kThreads / kW) { + if (c < w) d[i * w + c] = c < i ? tile[i][c] : c == i ? reciprocal[c] : 0.f; + } + } + if (below > 0) { + const int r = t % kR, q = t / kR; + float *dst = panel + w + item.row_begin + r; + float acc = 0.f; +#pragma unroll 4 + for (int c = q; c < w; c += kThreads / kR) { + if (r < below) { + const float v = tile[w + r][c]; + dst[c * h] = v; + acc += v * zp[c]; + } + } + partial[q][r] = acc; + __syncthreads(); + if (t < below) { + float sum = 0.f; +#pragma unroll + for (int s = 0; s < kThreads / kR; ++s) sum += partial[s][t]; + atomicAdd(y + below_rows[panel_info.rows_offset + item.row_begin + t], -sum); + } + } +} + +/** Loads the row-major w x w L_D of a panel into shared memory (whole CTA). */ +__device__ void LoadDiagonal(const float *d, int w, float (*ld)[kPad]) { + const int c = threadIdx.x % kW, i0 = threadIdx.x / kW; + float v[kW * kW / kThreads]; +#pragma unroll + for (int m = 0; m < kW * kW / kThreads; ++m) { + const int i = i0 + m * (kThreads / kW); + v[m] = i < w && c < w ? d[i * w + c] : 0.f; + } +#pragma unroll + for (int m = 0; m < kW * kW / kThreads; ++m) ld[i0 + m * (kThreads / kW)][c] = v[m]; +} + +/** + * One warp: v := L_D^{-1} v (forward) or L_D^{-T} v (backward), v in shared + * memory; ld holds L_D with the reciprocals of its diagonal on the diagonal. + */ +template +__device__ void WarpTriangularSolve(const float (*ld)[kPad], int w, float *v) { + const int lane = threadIdx.x % 32; + float lo = lane < w ? v[lane] : 0.f, hi = lane + 32 < w ? v[lane + 32] : 0.f; + for (int s = 0; s < w; ++s) { + const int k = kTranspose ? w - 1 - s : s; + const bool owner_hi = k >= 32; + if (lane == k % 32) { + if (owner_hi) { + hi *= ld[k][k]; + } else { + lo *= ld[k][k]; + } + } + const float xk = __shfl_sync(0xffffffffu, owner_hi ? hi : lo, k % 32); + if (kTranspose) { + if (lane < k) lo -= ld[k][lane] * xk; + if (lane + 32 < k) hi -= ld[k][lane + 32] * xk; + } else { + if (lane > k && lane < w) lo -= ld[lane][k] * xk; + if (lane + 32 > k && lane + 32 < w) hi -= ld[lane + 32][k] * xk; + } + } + if (lane < w) v[lane] = lo; + if (lane + 32 < w) v[lane + 32] = hi; +} + +/** + * One CTA per (run, 64 rows): the update -L_b[rows] L_b[run]ᵀ of a panel + * into the target panel of the run (rows at or below the run, columns of the + * run), lower triangle of the target's diagonal block only. Each thread + * computes 4 x 4 outputs from float4 shared loads (rows 4 tx.., columns + * 4 ty..); the tile then goes through shared memory so the read-modify-write + * of the target runs along its rows (coalesced). + */ +template +__global__ void __launch_bounds__(kThreads) + UpdateKernel(const UpdateItem *items, const Run *runs, const int *relmap, const Panel *panels, + float *factor) { + __shared__ __align__(16) float as[kW][kR]; + __shared__ __align__(16) float bs[kW][kR]; // 32 KB in all: 3 CTAs per SM on sm_86 + const UpdateItem item = items[blockIdx.x]; + const Panel source = panels[item.panel]; + const int w = source.width, h = source.height; + const Run run = runs[item.run]; + const int rows = min(kR, h - w - item.row_begin), cols = run.end - run.begin; + const float *src = factor + source.offset + w; + const int t = threadIdx.x; + { + const int i = t % kR, k0 = t / kR; + float va[kW * kR / kThreads], vb[kW * kR / kThreads]; +#pragma unroll + for (int m = 0; m < kW * kR / kThreads; ++m) { + const int k = k0 + m * (kThreads / kR); + va[m] = i < rows && k < w ? src[k * h + item.row_begin + i] : 0.f; + vb[m] = i < cols && k < w ? src[k * h + run.begin + i] : 0.f; + } +#pragma unroll + for (int m = 0; m < kW * kR / kThreads; ++m) { + as[k0 + m * (kThreads / kR)][i] = va[m]; + bs[k0 + m * (kThreads / kR)][i] = vb[m]; + } + } + __syncthreads(); + const int tx = t % 16, ty = t / 16; + float acc[4][4] = {}; + for (int k = 0; k < w; ++k) { + const float4 a = reinterpret_cast(as[k])[tx]; + const float4 b = reinterpret_cast(bs[k])[ty]; + const float av[4] = {a.x, a.y, a.z, a.w}, bv[4] = {b.x, b.y, b.z, b.w}; +#pragma unroll + for (int m = 0; m < 4; ++m) { +#pragma unroll + for (int n = 0; n < 4; ++n) acc[m][n] += av[m] * bv[n]; + } + } + __syncthreads(); + // as := the tile, column-major (as[j][i]). +#pragma unroll + for (int n = 0; n < 4; ++n) { + reinterpret_cast(as[4 * ty + n])[tx] = + make_float4(acc[0][n], acc[1][n], acc[2][n], acc[3][n]); + } + __syncthreads(); + const Panel target = panels[run.target]; + float *dst = factor + target.offset; + const int lane = t % 32, warp = t / 32; + const int *rel = relmap + run.relmap_offset + item.row_begin - run.begin; + int row_pos[kR / 32]; +#pragma unroll + for (int half = 0; half < kR / 32; ++half) { + const int i = lane + 32 * half; + row_pos[half] = i < rows ? rel[i] : -1; + } + for (int j = warp; j < cols; j += kThreads / 32) { + const int col_pos = relmap[run.relmap_offset + j]; + float *column = dst + static_cast(col_pos) * target.height; +#pragma unroll + for (int half = 0; half < kR / 32; ++half) { + if (row_pos[half] >= col_pos) { + const float v = as[j][lane + 32 * half]; + if (kAtomic) { + atomicAdd(column + row_pos[half], -v); + } else { + column[row_pos[half]] -= v; + } + } + } + } +} + +/** + * Backward substitution, one CTA per (panel, 64 below rows) of a level + * (levels in reverse): each CTA adds its rows' L_b[rows]ᵀ x[rows] to + * accumulator[cols]; the last CTA of the panel then solves + * x_p = L_D^{-T} (z_p - accumulator) and resets the accumulator and counter. + */ +__global__ void __launch_bounds__(kThreads) + BackwardKernel(const FactorItem *items, const Panel *panels, const int *below_rows, + const float *factor, const float *diagonal, float *x, float *accumulator, + int *counter) { + __shared__ float ld[kW][kPad]; + __shared__ float v[kW]; + __shared__ bool last; + const FactorItem item = items[blockIdx.x]; + const Panel panel = panels[item.panel]; + const int col = panel.col, w = panel.width, h = panel.height; + const int t = threadIdx.x, lane = t % 32; + const int below = min(kR, h - w - item.row_begin); + const int tiles = max(1, (h - w + kR - 1) / kR); + if (below > 0) { + const float *l = factor + panel.offset + w + item.row_begin; + const int r = t % kR, q = t / kR; + const float xr = r < below ? x[below_rows[panel.rows_offset + item.row_begin + r]] : 0.f; + constexpr int kPerThread = kW / (kThreads / kR); + float v[kPerThread]; +#pragma unroll + for (int m = 0; m < kPerThread; ++m) { + const int k = q + m * (kThreads / kR); + v[m] = r < below && k < w ? l[k * h + r] : 0.f; + } +#pragma unroll + for (int m = 0; m < kPerThread; ++m) { + const int k = q + m * (kThreads / kR); + float acc = v[m] * xr; +#pragma unroll + for (int s = 16; s > 0; s /= 2) acc += __shfl_xor_sync(0xffffffffu, acc, s); + if (lane == 0 && k < w) atomicAdd(accumulator + col + k, acc); + } + } + if (tiles > 1) { + __threadfence(); + __syncthreads(); + if (t == 0) last = atomicAdd(counter + item.panel, 1) == tiles - 1; + __syncthreads(); + if (!last) return; + __threadfence(); + } else { + __syncthreads(); + } + LoadDiagonal(diagonal + panel.diagonal_offset, w, ld); + if (t < w) { + v[t] = x[col + t] - __ldcg(accumulator + col + t); + accumulator[col + t] = 0.f; + } + if (t == 0) counter[item.panel] = 0; + __syncthreads(); + if (t < 32) WarpTriangularSolve(ld, w, v); + __syncthreads(); + if (t < w) x[col + t] = v[t]; +} + +/** The supervariables of a symmetric pattern and their adjacency (no self loops). */ +struct SupervariableGraph { + std::vector start{0}; ///< Per supervariable: first row; plus n. + std::vector adjacency_offsets{0}; ///< Per supervariable: its neighbors [offset, next). + std::vector adjacency; ///< Neighbor supervariables. +}; + +/** + * Finds the supervariables (consecutive rows with identical patterns) and + * their graph on the device; only the graph comes to the host. + */ +SupervariableGraph FindSupervariableGraph(cudaStream_t stream, const CSRSparseMatrix &matrix) { + SupervariableGraph graph; + const size_t n = matrix.NumRows(); + if (n == 0) return graph; + const int *ro = matrix.row_offsets.data(), *ci = matrix.col_ids.data(); + dvector flag(n); + SupervariableKernel<<>>(static_cast(n), ro, ci, + flag.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + std::vector h_flag(n), h_supervariable(n); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + flag.CopyToHost(h_flag.data(), n); + for (size_t i = 0; i < n; ++i) { + if (i > 0 && h_flag[i]) graph.start.push_back(static_cast(i)); + h_supervariable[i] = static_cast(graph.start.size()) - 1; + } + const int nb = static_cast(graph.start.size()); + graph.start.push_back(static_cast(n)); + dvector supervariable(h_supervariable); + dvector first_row(std::vector(graph.start.begin(), graph.start.end() - 1)), count(nb); + AdjacencyKernel<<(nb)), kThreads, 0, stream>>>( + nb, first_row.data(), ro, ci, supervariable.data(), nullptr, count.data(), nullptr); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + std::vector h_count(nb); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + count.CopyToHost(h_count.data(), nb); + for (int v = 0; v < nb; ++v) { + graph.adjacency_offsets.push_back(graph.adjacency_offsets.back() + h_count[v]); + } + graph.adjacency.resize(graph.adjacency_offsets.back()); + if (!graph.adjacency.empty()) { + dvector offsets(graph.adjacency_offsets), d_adjacency(graph.adjacency.size()); + AdjacencyKernel<<(nb)), kThreads, 0, stream>>>( + nb, first_row.data(), ro, ci, supervariable.data(), offsets.data(), nullptr, + d_adjacency.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + d_adjacency.CopyToHost(graph.adjacency.data(), graph.adjacency.size()); + } + return graph; +} + +} // namespace + +BlockSparseCholeskySolver::BlockSparseCholeskySolver() : failure_(1), failure_host_(1) {} + +bool BlockSparseCholeskySolver::Initialize(cudaStream_t stream, const Problem & /*problem*/, + const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) { + const size_t n = spd_matrix.NumRows(); + if (rhs.size() != n || result.size() != n) { + LogError("BlockSparseCholeskySolver: matrix size {} does not match rhs {} / result {}", n, + rhs.size(), result.size()); + return false; + } + const SupervariableGraph graph = FindSupervariableGraph(stream, spd_matrix); + plan_ = internal::AnalyzeCholesky(static_cast(n), graph.start, graph.adjacency_offsets, + graph.adjacency); + const auto &plan = plan_; + const size_t np = plan.panels.size(); + perm_ = dvector(plan.perm); + panels_ = dvector(plan.panels); + below_rows_ = dvector(plan.below_rows); + runs_ = dvector(plan.runs); + factor_items_ = dvector(plan.factor_items); + update_items_ = dvector(plan.update_items); + + // Where each nonzero of A goes in the factor, and the target positions of the updates. + assembly_dst_.resize(spd_matrix.col_ids.size()); + diagonal_index_.resize(n); + if (n > 0) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(diagonal_index_.data(), 0xFF, n * sizeof(int), stream)); + std::vector iperm(n), col_panel(n); + for (size_t c = 0; c < n; ++c) iperm[plan.perm[c]] = static_cast(c); + for (size_t p = 0; p < np; ++p) { + const auto &panel = plan.panels[p]; + std::fill(col_panel.begin() + panel.col, col_panel.begin() + panel.col + panel.width, + static_cast(p)); + } + dvector d_iperm(iperm), d_col_panel(col_panel); + AssemblyMapKernel<<>>( + static_cast(n), spd_matrix.row_offsets.data(), spd_matrix.col_ids.data(), + d_iperm.data(), d_col_panel.data(), panels_.data(), below_rows_.data(), + assembly_dst_.data(), diagonal_index_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); // before the temporaries go + } + relmap_.resize(plan.relmap_size); + if (!plan.runs.empty()) { + RelmapKernel<<>>( + static_cast(plan.runs.size()), runs_.data(), panels_.data(), below_rows_.data(), + relmap_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + factor_.resize(plan.factor_size); + diagonal_.resize(plan.diagonal_size); + original_.resize(n); + work_.resize(n); + solution_.resize(n); + accumulator_.resize(n); + counter_.resize(np); + if (n > 0) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(accumulator_.data(), 0, n * sizeof(float), stream)); + } + if (np > 0) THROW_ON_CUDA_ERROR(cudaMemsetAsync(counter_.data(), 0, np * sizeof(int), stream)); + LogMessage( + "BlockSparseCholesky: n {}, {} supervariables, {} ordering, {} supernodes, {} panels, {} " + "levels, factor {:.2f} M entries, {:.3f} GFLOP", + n, plan.num_supervariables, plan.ordering, plan.num_supernodes, np, plan.num_levels, + plan.factor_size / 1e6, plan.flops / 1e9); + return true; +} + +bool BlockSparseCholeskySolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) { + const auto &plan = plan_; + const int n = plan.n; + if (n == 0) return true; + if (safety_checks_) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(failure_.data(), 0, sizeof(int), stream)); + } + // Assembly of PAPᵀ into the panels, and the permuted right-hand side. + THROW_ON_CUDA_ERROR(cudaMemsetAsync(factor_.data(), 0, factor_.size() * sizeof(float), stream)); + const int nnz = static_cast(assembly_dst_.size()); + AssembleKernel<<>>(nnz, assembly_dst_.data(), + spd_matrix.values.data(), factor_.data()); + PermuteKernel<<>>(n, perm_.data(), rhs.data(), work_.data(), + diagonal_index_.data(), + spd_matrix.values.data(), original_.data()); + // Factorization and forward substitution, level by level. + for (int l = 0; l < plan.num_levels; ++l) { + const int f0 = plan.factor_offsets[l], f1 = plan.factor_offsets[l + 1]; + FactorKernel<<>>( + factor_items_.data() + f0, panels_.data(), below_rows_.data(), original_.data(), + factor_.data(), diagonal_.data(), work_.data(), solution_.data(), failure_.data()); + const int u0 = plan.update_offsets[l], u1 = plan.update_offsets[l + 1]; + if (u1 > u0) { + if (plan.update_exclusive[l]) { + UpdateKernel<<>>(update_items_.data() + u0, + runs_.data(), relmap_.data(), + panels_.data(), factor_.data()); + } else { + UpdateKernel<<>>(update_items_.data() + u0, + runs_.data(), relmap_.data(), + panels_.data(), factor_.data()); + } + } + } + // Backward substitution, levels in reverse. + for (int l = plan.num_levels - 1; l >= 0; --l) { + const int f0 = plan.factor_offsets[l], f1 = plan.factor_offsets[l + 1]; + BackwardKernel<<>>( + factor_items_.data() + f0, panels_.data(), below_rows_.data(), factor_.data(), + diagonal_.data(), solution_.data(), accumulator_.data(), counter_.data()); + } + UnpermuteKernel<<>>(n, perm_.data(), solution_.data(), + result.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + if (safety_checks_) { + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(failure_host_.data(), failure_.data(), sizeof(int), + cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + if (failure_host_[0] != 0) { + LogError("BlockSparseCholeskySolver: the matrix is not positive definite"); + return false; + } + } + return true; +} + +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_cholesky_solver.h b/cunls/linear_solver/block_sparse_cholesky_solver.h new file mode 100644 index 0000000..ab6eb7e --- /dev/null +++ b/cunls/linear_solver/block_sparse_cholesky_solver.h @@ -0,0 +1,110 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "cunls/common/types.h" +#include "cunls/linear_solver/block_sparse_cholesky_analysis.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" + +namespace cunls { + +/** + * @brief Supernodal sparse Cholesky factorization on the GPU, without + * external dependencies (no cuDSS). + * + * Initialize() finds the supervariables and their graph on the device, then + * analyzes the graph on the host (internal::AnalyzeCholesky): a fill-reducing + * ordering, supernodes cut into panels of at most 64 columns, and a level + * schedule; the map of A's nonzeros into the factor is built on the device. + * Every Solve() then factors A = L Lᵀ and solves with it, all on the device, + * without host synchronization unless safety checks are on: + * - assembly: the lower triangle of A scattered into the panels; + * - per level, one kernel factors every panel of the level, one CTA per + * 64 rows below its diagonal block (blocked by 16 columns in shared + * memory; the diagonal block redundantly per CTA), and takes the panel's + * step of the forward substitution along (the right-hand side as one + * more row); a second kernel applies each panel's update L_b L_bᵀ to the + * panels it targets, one CTA per 64 x 64 tile (atomics only where two + * panels of the level update the same target); + * - backward substitution, level by level in reverse, one CTA per 64 rows + * of a panel; the last CTA of a panel solves its diagonal block. + * + * Meant for the reduced systems of bundle adjustment (Schur complements), + * where it replaces a general sparse direct solver: dense covisibility (a + * few thousand columns, about 20% block fill), banded visual-inertial chains + * (tens of thousands of columns) and small dense systems. + * + * A pivot within 1e-5 of zero relative to its original diagonal entry + * (about 100 float epsilons) is a rank deficiency (a free gauge of Gauss-Newton): its + * unknown is fixed at zero, which solves a consistent semidefinite system. + * A clearly negative or non-finite pivot sets a device flag; with safety + * checks (the default) Solve() reads it back and returns false. Without them + * the failure shows up as non-finite values in the result. + */ +class BlockSparseCholeskySolver : public SparseLinearSolver { + public: + BlockSparseCholeskySolver(); + + using SparseLinearSolver::Initialize; // the BSR overloads: block storage is not supported + using SparseLinearSolver::Solve; + + /** @brief Analyzes the pattern of @p spd_matrix (host) and allocates the factor. */ + bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) override; + + /** @brief Factors @p spd_matrix (the pattern of Initialize()) and solves for @p rhs. */ + bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) override; + + /** @brief Skips reading back the pivot-failure flag after each factorization. */ + void DisableSafetyChecks() override { safety_checks_ = false; } + + /** @brief The analysis of the last Initialize() (host side). */ + const internal::CholeskyPlan &Plan() const { return plan_; } + + private: + internal::CholeskyPlan plan_; ///< Host analysis. + bool safety_checks_ = true; ///< Read back the failure flag after each Solve(). + + // Device copies of the plan. + dvector perm_; ///< Permuted column -> original column. + dvector panels_; ///< The panels. + dvector below_rows_; ///< Below rows of all panels. + dvector runs_; ///< Update runs of all panels. + dvector relmap_; ///< Target row positions of the runs (built on the device). + dvector factor_items_; ///< Factor CTAs by level. + dvector update_items_; ///< Update CTAs by level. + dvector assembly_dst_; ///< Per nonzero of A: its factor offset, or -1 (upper). + dvector diagonal_index_; ///< Per permuted column: its diagonal nonzero of A, or -1. + + // Numeric storage. + dvector factor_; ///< Panels of L (column-major). + dvector diagonal_; ///< Per panel: L_D, row-major. + dvector original_; ///< Per permuted column: the diagonal entry of A (zero pivots). + dvector work_; ///< Permuted right-hand side, reduced by the forward substitution. + dvector solution_; ///< Permuted L^{-1} b, then the solution. + /// Backward substitution: L_bᵀ x per column (zero between solves). + dvector accumulator_; + dvector counter_; ///< Backward substitution: finished CTAs per panel (zero between). + dvector failure_; ///< Set by a non-positive or non-finite pivot. + pvector failure_host_; ///< Pinned read-back of failure_. +}; + +} // namespace cunls diff --git a/cunls/linear_solver/llms.txt b/cunls/linear_solver/llms.txt index 9ade221..3e1a4e1 100644 --- a/cunls/linear_solver/llms.txt +++ b/cunls/linear_solver/llms.txt @@ -26,7 +26,7 @@ Sparse linear system abstraction and implementation used by minimizers. - `dense_qr_solver.h`: `DenseQRSolver` — cuSOLVER geqrf/ormqr + cuBLAS trsm. - `cudss_sparse_linear_solver.h`: - `cuDSSLinearSolverOptions` - - `cuDSSLinearSolver` — the only backend that declines block storage. + - `cuDSSLinearSolver` — declines block storage (CSR only). - `block_sparse_pcg_solver.h`: - `BlockSparsePCGOptions` — `block_size` / `block_layout`, `max_iterations`, `relative_tolerance`, `absolute_tolerance`, @@ -48,6 +48,19 @@ Sparse linear system abstraction and implementation used by minimizers. recursion, one warp per subproblem, registers + shuffles, cp.async prefetch of the next stage, kernel templated on the stage block size M <= 32 (dispatch table). No host synchronization (graph-capturable). +- `block_sparse_cholesky_solver.h` / `block_sparse_cholesky_analysis.h`: + `BlockSparseCholeskySolver` — supernodal sparse Cholesky, no external + dependency; the default backend (`MinimizerOptions::sparse_linear_solver_type`) + and the reduced-system solver of `LinearSystemType::SchurComplement`. CSR + only (declines block storage). `Initialize` finds supervariables and their + graph on the GPU, picks approximate minimum degree (quotient graph) or + nested dissection (balanced BFS level-set separators) by a cost model + (`AnalyzeCholesky`, host), + amalgamates supernodes and cuts them into panels of <= 64 columns, and + builds the assembly map and the relative map of the updates on the GPU. + `Solve`: per level a factor kernel (one + CTA per panel and 64 rows; forward substitution folded in) and an update + kernel (64 x 64 tiles), then a level-scheduled backward substitution. ## Expected system form diff --git a/cunls/linear_solver/sparse_linear_solver.cpp b/cunls/linear_solver/sparse_linear_solver.cpp index 6b1e552..6e03405 100644 --- a/cunls/linear_solver/sparse_linear_solver.cpp +++ b/cunls/linear_solver/sparse_linear_solver.cpp @@ -50,6 +50,8 @@ SparseLinearSolverPtr CreateSparseLinearSolver(SparseLinearSolverType type, return std::make_unique(config.block_sparse_pcg_options); case SparseLinearSolverType::BlockTridiagonal: return std::make_unique(); + case SparseLinearSolverType::BlockSparseCholesky: + return std::make_unique(); default: throw std::invalid_argument("Invalid sparse linear solver type"); } diff --git a/cunls/linear_solver/sparse_linear_solver.h b/cunls/linear_solver/sparse_linear_solver.h index 6cb5e3e..477ea49 100644 --- a/cunls/linear_solver/sparse_linear_solver.h +++ b/cunls/linear_solver/sparse_linear_solver.h @@ -21,6 +21,7 @@ #include +#include "cunls/linear_solver/block_sparse_cholesky_solver.h" #include "cunls/linear_solver/block_sparse_pcg_solver.h" #include "cunls/linear_solver/block_tridiagonal_solver.h" #include "cunls/linear_solver/cudss_sparse_linear_solver.h" @@ -35,8 +36,8 @@ namespace cunls { * @brief Selects the linear solver backend for the Gauss-Newton system. */ enum class SparseLinearSolverType { - cuDSS, ///< Sparse direct solver using NVIDIA's cuDSS library. The - ///< only backend that cannot consume block storage. + cuDSS, ///< Sparse direct solver using NVIDIA's cuDSS library + ///< (optional dependency, loaded at run time). DenseLDLT, ///< Densifies the coefficient matrix and solves with a custom ///< CUDA pivoted LDLT kernel. DenseCholesky, ///< Densifies the coefficient matrix and solves with cuSOLVER @@ -54,6 +55,14 @@ enum class SparseLinearSolverType { ///< in time (trajectories, MPC): needs the states' ///< stages (Problem::SetStateStages). One warp per ///< subproblem, no host synchronization. + /** + * Supernodal sparse Cholesky on the GPU without external dependencies + * (BlockSparseCholeskySolver): fill-reducing ordering, dense panels of at + * most 64 columns, level-scheduled kernels. Requires an SPD matrix. The + * default backend, and the solver of the reduced system of + * LinearSystemType::SchurComplement. + */ + BlockSparseCholesky, }; /** diff --git a/cunls/llms.txt b/cunls/llms.txt index a087195..63b5137 100644 --- a/cunls/llms.txt +++ b/cunls/llms.txt @@ -14,8 +14,9 @@ This directory contains all core library modules and the public umbrella header - `common/llms.txt`: CUDA handles, device vectors, sparse/dense type aliases. - `factor/llms.txt`: factor interfaces and built-in factor implementations. -- `linear_solver/llms.txt`: sparse solver interface, BlockSparsePCG - iterative solver (default), cuDSS direct solver, and dense fallbacks. +- `linear_solver/llms.txt`: sparse solver interface, BlockSparseCholesky + direct solver (default), BlockSparsePCG iterative solver, cuDSS direct + solver, and dense fallbacks. - `math/llms.txt`: Lie and matrix math kernels/utilities. - `minimizer/llms.txt`: `Problem`, residual wrappers, GN/LM solvers, RANSAC minimizers. - `robustifier/llms.txt`: robust loss batches. diff --git a/cunls/minimizer/CMakeLists.txt b/cunls/minimizer/CMakeLists.txt index 17bcb00..b8d8f2d 100644 --- a/cunls/minimizer/CMakeLists.txt +++ b/cunls/minimizer/CMakeLists.txt @@ -6,6 +6,7 @@ add_library(cunls_minimizer OBJECT device_reduction.cu gauss_newton_minimizer.cu levenberg_marquardt_minimizer.cu + linear_system.cu minimizer.cu minimizer_state.cu normal_equations.cu @@ -23,5 +24,6 @@ add_library(cunls_minimizer OBJECT ransac/slot_kernels.cu ransac/slot_set.cpp residual_batch.cu + schur_system.cu sparse_matrix.cu ) diff --git a/cunls/minimizer/augmented_lagrangian_minimizer.cu b/cunls/minimizer/augmented_lagrangian_minimizer.cu index a3cf5d4..e705919 100644 --- a/cunls/minimizer/augmented_lagrangian_minimizer.cu +++ b/cunls/minimizer/augmented_lagrangian_minimizer.cu @@ -319,7 +319,7 @@ void AugmentedLagrangianMinimizer::SetOptions(const AugmentedLagrangianMinimizer bool AugmentedLagrangianMinimizer::HoldActiveBounds(cudaStream_t stream, const MinimizerState &state, - NormalEquations &system, dvector &rhs, + LinearSystem &system, dvector &rhs, const dvector *step) { if (rhs.empty()) return false; // mask[column] = 0 for the components at a bound that `direction` pushes outward. @@ -342,7 +342,7 @@ bool AugmentedLagrangianMinimizer::HoldActiveBounds(cudaStream_t stream, // Before the solve: hold where the steepest-descent direction points outward. refinements_ = 0; mark(rhs.data(), bound_mask_); - system.ZeroMaskedLhsRowsColumns(stream, bound_mask_); + system.SetMask(&bound_mask_); ElementwiseMultiplyInPlace(stream, rhs.data(), bound_mask_.data(), rhs.size()); return false; } @@ -369,8 +369,10 @@ bool AugmentedLagrangianMinimizer::HoldActiveBounds(cudaStream_t stream, if (count == 0) return false; LogMessage("Bounds: {} more components held", count); } - system.ZeroMaskedLhsRowsColumns(stream, bound_extra_); + // hold_outward_kernel added the newly held components to bound_mask_: the + // next solve holds them too. ElementwiseMultiplyInPlace(stream, rhs.data(), bound_extra_.data(), rhs.size()); + system.SetMask(&bound_mask_); return true; } @@ -440,9 +442,8 @@ AugmentedLagrangianMinimizerSummary AugmentedLagrangianMinimizer::Minimize(cudaS THROW_ON_CUDA_ERROR(cudaGetLastError()); } } - auto restrict_system = [this](cudaStream_t s, const MinimizerState &state, - NormalEquations &system, dvector &rhs, - const dvector *step) { + auto restrict_system = [this](cudaStream_t s, const MinimizerState &state, LinearSystem &system, + dvector &rhs, const dvector *step) { return HoldActiveBounds(s, state, system, rhs, step); }; signature_ = signature; diff --git a/cunls/minimizer/augmented_lagrangian_minimizer.h b/cunls/minimizer/augmented_lagrangian_minimizer.h index 11c3eb2..239c97b 100644 --- a/cunls/minimizer/augmented_lagrangian_minimizer.h +++ b/cunls/minimizer/augmented_lagrangian_minimizer.h @@ -244,18 +244,19 @@ class AugmentedLagrangianMinimizer { * @brief Projected Gauss-Newton for the box-bounded states, run by every * inner solve as its InnerSolve::restrict_system, at the current `state`. * - * Before the solve (`step` null), once the normal equations are built (and, + * Before the solve (`step` null), once the linear system is built (and, * for Levenberg-Marquardt, damped): the components that sit at a bound * with the steepest-descent direction (the right-hand side) pointing - * outward leave the system (rows, columns and rhs entries zeroed, diagonal - * kept), so their step is exactly 0. After a solve (`step` given): a free - * component the step still pushes through its bound is held too and true - * asks for another solve (active-set refinement), at most + * outward are masked out of the system (LinearSystem::SetMask with + * bound_mask_; their rhs entries zeroed here), so their step is exactly 0. + * After a solve (`step` given): a free component the step still pushes + * through its bound is held too and true asks for another solve (active-set + * refinement), at most * max_bound_refinements times; false keeps the step. Together with * VectorStateBatch's bounded Plus, which clamps the moved components into * the box, the iterates stay feasible. */ - bool HoldActiveBounds(cudaStream_t stream, const MinimizerState &state, NormalEquations &system, + bool HoldActiveBounds(cudaStream_t stream, const MinimizerState &state, LinearSystem &system, dvector &rhs, const dvector *step); // Implementation state; the behaviour lives in augmented_lagrangian_minimizer.cu. @@ -282,8 +283,9 @@ class AugmentedLagrangianMinimizer { std::vector bounded_; ///< Box-bounded state batches of the problem. std::vector> - column_offsets_; ///< Per state batch: column of each state (-1: constant). - dvector bound_mask_; ///< Per column: 0 where held at a bound. + column_offsets_; ///< Per state batch: column of each state (-1: constant). + /// Per column: 0 where held at a bound; the LinearSystem's declared mask (by pointer). + dvector bound_mask_; dvector bound_step_mask_; ///< Per column: 0 where the step leaves the box. dvector bound_extra_; ///< Per column: 0 where newly held by a refinement. dvector bound_count_; ///< Components newly held by the last refinement. diff --git a/cunls/minimizer/bsr_matrix.cu b/cunls/minimizer/bsr_matrix.cu index 3bccf4e..8c54b6c 100644 --- a/cunls/minimizer/bsr_matrix.cu +++ b/cunls/minimizer/bsr_matrix.cu @@ -421,6 +421,13 @@ void ScaleSymmetric(cudaStream_t stream, BSRSparseMatrix &matrix, const dvector< THROW_ON_CUDA_ERROR(cudaGetLastError()); } +void CopyBSRValues(cudaStream_t stream, const BSRSparseMatrix &input, BSRSparseMatrix &output) { + if (&input == &output || input.values.empty()) return; + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(output.values.data(), input.values.data(), + input.values.size() * sizeof(float), cudaMemcpyDeviceToDevice, + stream)); +} + void CopyBSRSparseMatrix(cudaStream_t stream, const BSRSparseMatrix &input, BSRSparseMatrix &output) { if (&input == &output) { diff --git a/cunls/minimizer/bsr_matrix.h b/cunls/minimizer/bsr_matrix.h index 5cc8207..4a522c1 100644 --- a/cunls/minimizer/bsr_matrix.h +++ b/cunls/minimizer/bsr_matrix.h @@ -92,6 +92,9 @@ void ScaleSymmetric(cudaStream_t stream, BSRSparseMatrix &matrix, const dvector< void CopyBSRSparseMatrix(cudaStream_t stream, const BSRSparseMatrix &input, BSRSparseMatrix &output); +/** @brief Copies only the values of a BSR matrix with the same structure. */ +void CopyBSRValues(cudaStream_t stream, const BSRSparseMatrix &input, BSRSparseMatrix &output); + /** @brief y = A * x for block storage (y is resized to the row count). */ void MultiplyBSRByDenseVector(cudaStream_t stream, const BSRSparseMatrix &matrix, const dvector &x, dvector &y); diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index 735b4e1..a843005 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -40,7 +40,7 @@ __global__ void classify_kernel(size_t num_problems, const float *cost, const fl GaussNewtonMinimizer::GaussNewtonMinimizer(const MinimizerOptions &options) : Minimizer(options) {} -void GaussNewtonMinimizer::ClassifySteps(cudaStream_t stream, const NormalEquations &, +void GaussNewtonMinimizer::ClassifySteps(cudaStream_t stream, const LinearSystem &, const dvector &, ProblemPartition &partition) { const size_t n = partition.NumProblems(); classify_kernel<<((n + kThreads - 1) / kThreads), kThreads, 0, stream>>>( diff --git a/cunls/minimizer/gauss_newton_minimizer.h b/cunls/minimizer/gauss_newton_minimizer.h index ba99f59..cdece91 100644 --- a/cunls/minimizer/gauss_newton_minimizer.h +++ b/cunls/minimizer/gauss_newton_minimizer.h @@ -47,10 +47,10 @@ class GaussNewtonMinimizer : public Minimizer { private: /** @brief Leaves the Gauss-Newton system unchanged. */ - void UpdateSystem(cudaStream_t, size_t, NormalEquations &, const ProblemPartition &) override {} + void UpdateSystem(cudaStream_t, size_t, LinearSystem &, const ProblemPartition &) override {} /** @brief Rejects steps that do not lower the cost; see the class comment. */ - void ClassifySteps(cudaStream_t stream, const NormalEquations &system, const dvector &step, + void ClassifySteps(cudaStream_t stream, const LinearSystem &system, const dvector &step, ProblemPartition &partition) override; }; diff --git a/cunls/minimizer/hessian_structure.cu b/cunls/minimizer/hessian_structure.cu index 9489dc7..b4774ec 100644 --- a/cunls/minimizer/hessian_structure.cu +++ b/cunls/minimizer/hessian_structure.cu @@ -332,7 +332,7 @@ void HessianStructureBuilder::BuildLayout(const Problem &problem) { } void HessianStructureBuilder::ResolveFactorColumns(cudaStream_t stream, const Problem &problem, - int num_cols, dvector &tangent_at_col) { + int num_cols, dvector *tangent_at_col) { const auto &residual_batches = problem.GetResidualBatches(); if (factor_cols_.empty()) { @@ -341,10 +341,12 @@ void HessianStructureBuilder::ResolveFactorColumns(cudaStream_t stream, const Pr THROW_ON_CUDA_ERROR( cudaMemsetAsync(factor_cols_.data(), 0xFF, factor_cols_.size() * sizeof(int), stream)); - tangent_at_col.resize(static_cast(num_cols)); - if (num_cols > 0) { - THROW_ON_CUDA_ERROR( - cudaMemsetAsync(tangent_at_col.data(), 0, tangent_at_col.size() * sizeof(int), stream)); + if (tangent_at_col != nullptr) { + tangent_at_col->resize(static_cast(num_cols)); + if (num_cols > 0) { + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(tangent_at_col->data(), 0, tangent_at_col->size() * sizeof(int), stream)); + } } // The state pointers are read from the problem's device tables @@ -365,10 +367,10 @@ void HessianStructureBuilder::ResolveFactorColumns(cudaStream_t stream, const Pr ComputeStateColumnOffsets(stream, last_col, state_batch, block_col_map); const size_t tangent_entries = state_batch->NumActiveStates() * state_batch->TangentSize(); - if (tangent_entries > 0) { + if (tangent_at_col != nullptr && tangent_entries > 0) { FillTangentAtColKernel<<>>( static_cast(tangent_entries), static_cast(state_batch->TangentSize()), - block_col_map.data(), tangent_at_col.data()); + block_col_map.data(), tangent_at_col->data()); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -396,18 +398,23 @@ void HessianStructureBuilder::ResolveFactorColumns(cudaStream_t stream, const Pr THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); } +void HessianStructureBuilder::ResolveColumns(cudaStream_t stream, const Problem &problem, + int num_cols, dvector *tangent_at_col) { + BuildLayout(problem); + ResolveFactorColumns(stream, problem, num_cols, tangent_at_col); +} + int HessianStructureBuilder::DiscoverBlockPairs(cudaStream_t stream, const Problem &problem, int num_cols, dvector &pair_fields, dvector &pair_order, dvector &pair_group, size_t &num_valid) { auto stream_policy = thrust::cuda::par_nosync.on(stream); - BuildLayout(problem); // Every buffer below is scratch for this call only. Holding them as members // would keep ~20 bytes per candidate pair alive for the whole solve, which on // a large SBA problem is hundreds of MiB against no reuse benefit. dvector tangent_at_col; - ResolveFactorColumns(stream, problem, num_cols, tangent_at_col); + ResolveColumns(stream, problem, num_cols, &tangent_at_col); num_valid = 0; if (total_pairs_ == 0 || num_cols == 0) { diff --git a/cunls/minimizer/hessian_structure.h b/cunls/minimizer/hessian_structure.h index 53e3976..d8366c1 100644 --- a/cunls/minimizer/hessian_structure.h +++ b/cunls/minimizer/hessian_structure.h @@ -68,7 +68,7 @@ class HessianStructureBuilder { * * @param stream CUDA stream for GPU operations. * @param problem The optimization problem. - * @param num_cols Free tangent dimensions in the reduced system. + * @param num_cols Free (non-constant) tangent columns. * @param[out] output CSR matrix; row_offsets/col_ids filled, values sized. * @param want_scatter_maps When false, WriteOffsets() is left empty and the * final scatter pass is skipped (the J^T J path does not need it). @@ -96,6 +96,20 @@ class HessianStructureBuilder { void Build(cudaStream_t stream, const Problem &problem, int num_cols, int block_size, BSRSparseMatrix &output, bool want_scatter_maps); + /** + * @brief Only the layout and the column of every (factor, state slot) + * (Layout(), FactorCols(), JacobianValuesSize()), without the pattern; for + * linear systems that build their own (SchurSystem). + * + * @param stream CUDA stream for GPU operations (synchronized on return). + * @param problem The optimization problem. + * @param num_cols Free (non-constant) tangent columns. + * @param[out] tangent_at_col If not null: the tangent size of the state + * owning each column. + */ + void ResolveColumns(cudaStream_t stream, const Problem &problem, int num_cols, + dvector *tangent_at_col = nullptr); + /** @brief Per-residual-batch geometry, indexed as the problem's batches. */ const std::vector &Layout() const { return layout_; } @@ -130,10 +144,11 @@ class HessianStructureBuilder { /** * @brief Resolves every (factor, state slot) to a global column. * - * @param[out] tangent_at_col Tangent size of the state owning each column. + * @param[out] tangent_at_col If not null: the tangent size of the state + * owning each column. */ void ResolveFactorColumns(cudaStream_t stream, const Problem &problem, int num_cols, - dvector &tangent_at_col); + dvector *tangent_at_col); std::vector layout_; dvector factor_cols_; diff --git a/cunls/minimizer/levenberg_marquardt_minimizer.cu b/cunls/minimizer/levenberg_marquardt_minimizer.cu index 7b50493..6c297c5 100644 --- a/cunls/minimizer/levenberg_marquardt_minimizer.cu +++ b/cunls/minimizer/levenberg_marquardt_minimizer.cu @@ -64,6 +64,9 @@ MinimizerOptions BaseOptions(const LevenbergMarquardtMinimizerOptions &o) { if (!(o.lambda_upscale > 1.f)) { throw std::invalid_argument("LevenbergMarquardtMinimizer: lambda_upscale must be > 1"); } + if (!(o.min_diagonal >= 0.f)) { + throw std::invalid_argument("LevenbergMarquardtMinimizer: min_diagonal must be >= 0"); + } MinimizerOptions base = o.base_options; if (base.max_consecutive_rejected_steps > 0) { base.max_consecutive_rejected_steps += RejectionsToMaxDamping(o); @@ -74,7 +77,7 @@ MinimizerOptions BaseOptions(const LevenbergMarquardtMinimizerOptions &o) { /** The damping options the λ update reads (kernel arguments must be trivially copyable). */ struct DampingRule { float initial_lambda, lambda_min, lambda_max, lambda_upscale, lambda_downscale, - lambda_downscale_threshold; + lambda_downscale_threshold, min_diagonal; }; /** λ of a subproblem after its previous step's outcome (see the class comment). */ @@ -100,7 +103,7 @@ __device__ float UpdatedLambda(const DampingRule &o, bool first, float lambda, i /** * One thread per row and per subproblem (whichever is more): thread p < P * writes subproblem p's updated λ to lambda_out; thread r < rows writes - * λ'_{problem(r)} diag(H)_r to damping (recomputing λ', so no thread reads + * λ'_{problem(r)} max(diag(H)_r, min_diagonal) to damping (recomputing λ', so no thread reads * what another writes). */ __global__ void damping_kernel(size_t num_rows, size_t num_problems, const int *row_problem, @@ -116,7 +119,7 @@ __global__ void damping_kernel(size_t num_rows, size_t num_problems, const int * const int p = row_problem != nullptr ? row_problem[i] : 0; const float lambda = UpdatedLambda(o, first, first ? 0.f : lambda_in[p], first ? 0 : outcome[p], first ? 0 : rejected[p], first ? 0.f : quality[p]); - damping[i] = lambda * diagonal[i]; + damping[i] = lambda * fmaxf(diagonal[i], o.min_diagonal); } } @@ -148,7 +151,7 @@ LevenbergMarquardtMinimizer::LevenbergMarquardtMinimizer( : Minimizer(BaseOptions(options)), options_(options) {} void LevenbergMarquardtMinimizer::UpdateSystem(cudaStream_t stream, size_t iteration, - NormalEquations &system, + LinearSystem &system, const ProblemPartition &partition) { const size_t num_problems = partition.NumProblems(); const bool first = iteration == 0; @@ -157,23 +160,24 @@ void LevenbergMarquardtMinimizer::UpdateSystem(cudaStream_t stream, size_t itera next_lambda_.resize(num_problems); quality_.resize(num_problems); } - // (H + λ_p diag(H)) on the rows of subproblem p, with λ updated in the same pass. - system.ExtractLhsDiagonal(stream, diagonal_); - damping_.resize(diagonal_.size()); - const size_t n = std::max(diagonal_.size(), num_problems); + // The damping λ_p max(diag(S H S), min_diagonal) of the rows of subproblem + // p, with λ updated in the same pass. + system.ExtractScaledHessianDiagonal(stream, scaled_diagonal_); + damping_.resize(scaled_diagonal_.size()); + const size_t n = std::max(scaled_diagonal_.size(), num_problems); damping_kernel<<>>( - diagonal_.size(), num_problems, partition.RowProblem(), + scaled_diagonal_.size(), num_problems, partition.RowProblem(), DampingRule{options_.initial_lambda, options_.lambda_min, options_.lambda_max, options_.lambda_upscale, options_.lambda_downscale, - options_.lambda_downscale_threshold}, + options_.lambda_downscale_threshold, options_.min_diagonal}, first, partition.Outcome(), partition.Rejected(), quality_.data(), lambda_.data(), - next_lambda_.data(), diagonal_.data(), damping_.data()); + next_lambda_.data(), scaled_diagonal_.data(), damping_.data()); THROW_ON_CUDA_ERROR(cudaGetLastError()); std::swap(lambda_, next_lambda_); - system.AddScaledDiagonalToLhs(stream, 1.f, damping_); + system.SetDamping(&damping_); } -void LevenbergMarquardtMinimizer::ClassifySteps(cudaStream_t stream, const NormalEquations &system, +void LevenbergMarquardtMinimizer::ClassifySteps(cudaStream_t stream, const LinearSystem &system, const dvector &step, ProblemPartition &partition) { const size_t n = partition.NumProblems(); @@ -182,16 +186,16 @@ void LevenbergMarquardtMinimizer::ClassifySteps(cudaStream_t stream, const Norma // The predicted reduction's terms δᵀDδ and δᵀHδ, per subproblem. The step // is in physical coordinates (dx = S z, see ColumnScaling), so the damping // term λ zᵀ diag(S H S) z is λ dxᵀ diag(H) dx with the unscaled Hessian; - // without column scaling that is the diagonal UpdateSystem extracted. - const dvector *hessian_diagonal = &diagonal_; + // without column scaling that is the diagonal UpdateSystem extracted. The + // min_diagonal floor is left out: it only matters for unobservable unknowns, + // whose step is zero. + const dvector *hessian_diagonal = &scaled_diagonal_; if (Options().column_scaling != ColumnScaling::None) { - system.ExtractHessianDiagonal(stream, hessian_diagonal_); - hessian_diagonal = &hessian_diagonal_; + system.ExtractHessianDiagonal(stream, unscaled_diagonal_); + hessian_diagonal = &unscaled_diagonal_; } partition.SumRows(stream, step.data(), nullptr, hessian_diagonal->data(), diag_weight_.data()); - hessian_step_.resize(step.size()); - system.MultiplyHessian(stream, cusparse_handle_.GetHandle(stream), step, hessian_step_, buffer_); - partition.SumRows(stream, step.data(), hessian_step_.data(), nullptr, matrix_weight_.data()); + system.StepCurvature(stream, step, partition, matrix_weight_.data()); const MinimizerOptions &base = Options(); classify_kernel<<>>( n, partition.Cost(), partition.NewCost(), partition.StepSquared(), diag_weight_.data(), diff --git a/cunls/minimizer/levenberg_marquardt_minimizer.h b/cunls/minimizer/levenberg_marquardt_minimizer.h index 3118238..ecf107c 100644 --- a/cunls/minimizer/levenberg_marquardt_minimizer.h +++ b/cunls/minimizer/levenberg_marquardt_minimizer.h @@ -19,7 +19,6 @@ #include -#include "cunls/common/cusparse_helper.h" #include "cunls/minimizer/minimizer.h" namespace cunls { @@ -93,6 +92,19 @@ struct LevenbergMarquardtMinimizerOptions { */ float lambda_min = 1e-6; + /** + * @brief Floor of the diagonal the damping scales: λ max(diag(H)_i, + * min_diagonal) (the diagonal of the column-scaled system when column + * scaling is on). + * + * An unobservable unknown (a zero Jacobian column, e.g. a pose whose + * observations are all invalid) has a zero diagonal, which no λ diag(H) + * damping fixes: the system stays singular and every step fails. With the + * floor it gets a positive pivot and, its right-hand side being zero, a + * zero step. Ceres' min_lm_diagonal. Default: 1e-6 + */ + float min_diagonal = 1e-6f; + /** * @brief Threshold for step acceptance based on rho (step quality). * @@ -115,10 +127,11 @@ struct LevenbergMarquardtMinimizerOptions { /** * @brief Levenberg-Marquardt: Gauss-Newton with adaptive damping. * - * Solves (H + λ diag(H)) δ = -g with one λ per subproblem. A large λ gives - * short, gradient-like steps; a small λ gives Gauss-Newton steps. More robust - * than GaussNewtonMinimizer far from the solution and when H is close to - * singular. + * Solves (H + λ D) δ = -g with one λ per subproblem, D = max(diag(H), + * min_diagonal) (of the column-scaled system with column scaling). A large λ + * gives short, gradient-like steps; a small λ gives Gauss-Newton steps. More + * robust than GaussNewtonMinimizer far from the solution and when H is close + * to singular. * * Per subproblem, ρ = actual cost reduction / reduction predicted by the * linear model, ½ δᵀHδ + λ δᵀ diag(H) δ: @@ -144,7 +157,8 @@ class LevenbergMarquardtMinimizer : public Minimizer { * lambda_max. * * @throws std::invalid_argument unless 0 < lambda_min <= lambda_max, - * initial_lambda <= lambda_max and lambda_upscale > 1. + * initial_lambda <= lambda_max, lambda_upscale > 1 and + * min_diagonal >= 0. */ explicit LevenbergMarquardtMinimizer( const LevenbergMarquardtMinimizerOptions &options = LevenbergMarquardtMinimizerOptions()); @@ -153,28 +167,33 @@ class LevenbergMarquardtMinimizer : public Minimizer { /** * @brief Sets every subproblem's λ to initial_lambda on iteration 0, and * otherwise updates it from the outcome of the previous step (see the class - * comment); then adds λ_p diag(H) to the rows of subproblem p. + * comment); then declares the damping λ_p max(diag(S H S), min_diagonal) + * of the rows of subproblem p (LinearSystem::SetDamping; S the column + * scale). */ - void UpdateSystem(cudaStream_t stream, size_t iteration, NormalEquations &system, + void UpdateSystem(cudaStream_t stream, size_t iteration, LinearSystem &system, const ProblemPartition &partition) override; + /** @brief A failed solve is a rejected step: λ grows until the system factors. */ + bool RejectsFailedSolves() const override { return true; } + /** @brief Computes ρ per subproblem and applies the rules of the class comment. */ - void ClassifySteps(cudaStream_t stream, const NormalEquations &system, const dvector &step, + void ClassifySteps(cudaStream_t stream, const LinearSystem &system, const dvector &step, ProblemPartition &partition) override; const LevenbergMarquardtMinimizerOptions options_; ///< As constructed (not widened). - dvector lambda_; ///< Per subproblem: current damping λ. - dvector next_lambda_; ///< Per subproblem: λ being updated (swapped with lambda_). - dvector quality_; ///< Per subproblem: ρ of the last step (for the λ update). - dvector diag_weight_; ///< Per subproblem: δᵀ diag(H) δ. - dvector matrix_weight_; ///< Per subproblem: δᵀ H δ. - dvector diagonal_; ///< Per row: diag of the working left-hand side (S H S). - dvector hessian_diagonal_; ///< Per row: diag(H) unscaled (with column scaling). - dvector damping_; ///< Per row: λ diag(H). - dvector hessian_step_; ///< Per row: H δ. - cuSPARSEHandle cusparse_handle_; ///< For H δ in CSR storage. - dvector buffer_; ///< cuSPARSE scratch of H δ. + dvector lambda_; ///< Per subproblem: current damping λ. + dvector next_lambda_; ///< Per subproblem: λ being updated (swapped with lambda_). + dvector quality_; ///< Per subproblem: ρ of the last step (for the λ update). + dvector diag_weight_; ///< Per subproblem: δᵀ diag(H) δ. + dvector matrix_weight_; ///< Per subproblem: δᵀ H δ. + /// Per row: diag(S H S) of the current linearization (UpdateSystem()). + dvector scaled_diagonal_; + /// Per row: diag(H), unscaled; extracted by ClassifySteps() with column scaling only. + dvector unscaled_diagonal_; + /// Per row: λ max(diag(S H S), min_diagonal), declared to the system by pointer. + dvector damping_; }; } // namespace cunls diff --git a/cunls/minimizer/linear_system.cu b/cunls/minimizer/linear_system.cu new file mode 100644 index 0000000..77754fa --- /dev/null +++ b/cunls/minimizer/linear_system.cu @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cunls/minimizer/linear_system.h" +#include "cunls/minimizer/sparse_matrix.h" + +namespace cunls { + +void LinearSystem::ExtractScaledHessianDiagonal(cudaStream_t stream, + dvector &diagonal) const { + ExtractHessianDiagonal(stream, diagonal); + if (column_scale_ == nullptr || diagonal.empty()) return; + ElementwiseMultiplyInPlace(stream, diagonal.data(), column_scale_->data(), diagonal.size()); + ElementwiseMultiplyInPlace(stream, diagonal.data(), column_scale_->data(), diagonal.size()); +} + +} // namespace cunls diff --git a/cunls/minimizer/linear_system.h b/cunls/minimizer/linear_system.h new file mode 100644 index 0000000..ebd2b99 --- /dev/null +++ b/cunls/minimizer/linear_system.h @@ -0,0 +1,203 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include + +#include "cunls/common/types.h" + +namespace cunls { + +class Problem; +class ProblemPartition; + +/** + * @brief How each iteration's linear system is formed and solved + * (MinimizerOptions::linear_system). + */ +enum class LinearSystemType { + /** + * Assemble the full normal equations H = JᵀJ and solve them with + * MinimizerOptions::sparse_linear_solver_type. Works for every problem. + */ + Assembled, + /** + * Eliminate the state batches of Problem::SetEliminatedStateBatches (the + * landmarks of a bundle adjustment) with a Schur complement, solve the + * reduced system of the remaining states with BlockSparseCholeskySolver + * (sparse_linear_solver_type is not used), and recover the eliminated + * states per landmark. + */ + SchurComplement, +}; + +/** + * @brief The linear system of one minimizer iteration. + * + * Created by the minimizer (MinimizerOptions::linear_system); subclasses of + * Minimizer and InnerSolve::restrict_system receive it to declare their + * modifications. Not meant to be implemented outside the library. + * + * The Gauss-Newton system H δ = b at the current states, with H = JᵀJ and + * b = -Jᵀr from the per-factor Jacobians and residuals, as modified by the + * minimizer: a column scaling s (solve S H S z = S b), a damping a added to + * the diagonal of S H S, and a mask m whose zero entries hold their unknowns + * at 0. The modifications are declared with SetColumnScale(), SetDamping() + * and SetMask() (by pointer: the vectors stay owned by the caller and must + * outlive the next Solve()) and applied by Solve(); H itself is never + * modified, so the damping can be changed and the system solved again + * without a new Assemble(). Each new linearization starts with + * ClearModifications(). + * + * Implementations: NormalEquations (assembles H and hands it to a + * SparseLinearSolver) and SchurSystem (eliminates the landmark states and + * solves the reduced system). The minimizer owns one, chosen by + * MinimizerOptions::linear_system. + */ +class LinearSystem { + public: + virtual ~LinearSystem() = default; + + /** + * @brief Symbolic setup for `problem`: sparsity patterns and maps. Called + * whenever the problem's structure changes. + * + * @param stream CUDA stream. + * @param problem The problem; its free tangent columns are laid out by state + * batch (see ComputeStateColumnOffsets). + * @param num_cols Number of free tangent columns. + * @param reorderable Per residual batch: whether its Jacobians can be + * evaluated in EvaluationOrder() (analytic Jacobians; numeric + * differentiation keeps the natural order). + */ + virtual void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + const std::vector &reorderable) = 0; + + /** @brief Floats of the per-factor Jacobian buffer Assemble() reads. */ + virtual size_t JacobianValuesSize() const = 0; + + /** + * @brief The order in which the factors of residual batch `batch` are + * evaluated for Assemble() (device array of NumActiveFactors() factor + * indices: item t evaluates factor order[t]; FactorBatch::Evaluate's + * factor_ids), or nullptr for the natural order. Only for batches marked + * reorderable in Initialize(). + */ + virtual const int *EvaluationOrder(size_t batch) const { return nullptr; } + + /** + * @brief Builds H and b = -Jᵀr at the current states. + * + * @param stream CUDA stream. + * @param problem The problem of Initialize(). + * @param jacobians Per-item row-major Jacobian blocks (items in + * EvaluationOrder()), residual batches concatenated in problem order. + * Kept by pointer: must stay unchanged until the next Assemble(). + * @param residuals Residuals, laid out the same way; kept like `jacobians`. + * @param[out] rhs b, one entry per free tangent column (unscaled, unmasked). + */ + virtual void Assemble(cudaStream_t stream, const Problem &problem, const float *jacobians, + const float *residuals, dvector &rhs) = 0; + + /** @brief diag(H), unscaled and undamped, one entry per column. */ + virtual void ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const = 0; + + /** + * @brief Prepares the solver for this structure (symbolic factorization), + * with the system as currently modified. Called once per structure, after + * the first Assemble() and before the first Solve(). + * + * @param stream CUDA stream. + * @param problem The problem of Initialize(). + * @param rhs Right-hand side vector (sizes the solver's setup). + * @param step Solution vector (sizes the solver's setup; not written). + * @return false if the solver rejects the system. + */ + virtual bool PrepareSolver(cudaStream_t stream, const Problem &problem, const dvector &rhs, + dvector &step) = 0; + + /** + * @brief Solves the modified system S H S z + diag(a) z = rhs (masked). + * + * @param stream CUDA stream. + * @param rhs Right-hand side, already scaled and masked by the caller. + * @param[out] step z, in scaled coordinates (the caller maps it back). + * @return false on a solver failure. + */ + virtual bool Solve(cudaStream_t stream, const dvector &rhs, dvector &step) = 0; + + /** + * @brief out[p] = δᵀ H δ over the rows of subproblem p, with the undamped, + * unscaled H of the last Assemble() (the curvature term of the + * Levenberg-Marquardt model). + * + * @param stream CUDA stream. + * @param step δ, one entry per column, in physical coordinates. + * @param partition The subproblems. + * @param[out] out Device array of partition.NumProblems() floats. + */ + virtual void StepCurvature(cudaStream_t stream, const dvector &step, + ProblemPartition &partition, float *out) const = 0; + + /** + * @brief Column scale s, or nullptr for none. Not copied: the vector must + * stay alive, and is read by every later Solve(). + */ + void SetColumnScale(const dvector *scale) { column_scale_ = scale; } + + /** + * @brief Damping a added to the diagonal of the scaled system, or nullptr + * for none. Not copied, as SetColumnScale(). + */ + void SetDamping(const dvector *damping) { damping_ = damping; } + + /** + * @brief Mask m, or nullptr for none: the unknowns with m = 0 get a zero row + * and column and a unit diagonal, so their step is 0 (the caller zeroes + * their right-hand side). Not copied, as SetColumnScale(). + */ + void SetMask(const dvector *mask) { mask_ = mask; } + + /** @brief Clears the column scale, damping and mask (a new linearization). */ + void ClearModifications() { + column_scale_ = nullptr; + damping_ = nullptr; + mask_ = nullptr; + } + + /** @brief diag(S H S): the diagonal of the scaled, undamped system. */ + void ExtractScaledHessianDiagonal(cudaStream_t stream, dvector &diagonal) const; + + protected: + /** @brief The declared column scale; nullptr for none. */ + const dvector *ColumnScale() const { return column_scale_; } + /** @brief The declared damping; nullptr for none. */ + const dvector *Damping() const { return damping_; } + /** @brief The declared mask; nullptr for none. */ + const dvector *Mask() const { return mask_; } + + private: + const dvector *column_scale_ = nullptr; ///< See SetColumnScale(). + const dvector *damping_ = nullptr; ///< See SetDamping(). + const dvector *mask_ = nullptr; ///< See SetMask(). +}; + +} // namespace cunls diff --git a/cunls/minimizer/minimizer.cu b/cunls/minimizer/minimizer.cu index c9b115f..0f91442 100644 --- a/cunls/minimizer/minimizer.cu +++ b/cunls/minimizer/minimizer.cu @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -28,7 +29,9 @@ #include "cunls/factor/constraint_factor_batch.h" #include "cunls/minimizer/device_reduction.h" #include "cunls/minimizer/minimizer.h" +#include "cunls/minimizer/normal_equations.h" #include "cunls/minimizer/residual_batch.h" +#include "cunls/minimizer/schur_system.h" #include "cunls/minimizer/sparse_matrix.h" namespace cunls { @@ -106,6 +109,32 @@ float ComputeCost(cudaStream_t stream, const Problem &problem, const MinimizerSt return scratch.h_scalars[0]; } +/** dst[t * B + b] = src[order[t] * B + b]: the state pointers of the items of `order`. */ +__global__ void GatherStatePointersKernel(size_t num_items, int num_slots, const int *order, + float *const *src, float **dst) { + const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= num_items * num_slots) return; + const size_t t = i / num_slots; + dst[i] = src[static_cast(order[t]) * num_slots + (i - t * num_slots)]; +} + +/** + * The state pointers of residual batch items evaluated in `order` (item t + * evaluates factor order[t]), gathered into `buffer`. + */ +float const *const *GatherItemStatePointers(cudaStream_t stream, const int *order, + float *const *state_pointers, size_t num_factors, + int num_slots, dvector &buffer) { + constexpr int kThreads = 256; + const size_t count = num_factors * num_slots; + if (buffer.size() < count) buffer.resize(count); + GatherStatePointersKernel<<((count + kThreads - 1) / kThreads), kThreads, 0, + stream>>>(num_factors, num_slots, order, state_pointers, + buffer.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + return buffer.data(); +} + /** * Residuals and per-factor Jacobian blocks at `state`. Per residual batch, * either the FactorBatch's analytic Jacobian (ResidualBatch::Evaluate, which @@ -135,7 +164,18 @@ void ComputeResidualAndJacobian(cudaStream_t stream, const Problem &problem, auto ptrs = state_pointers[i].data(); const auto &factor_batch = rb.GetFactorBatch(); if (problem.JacobianModeFor(i, options.jacobian_mode) == JacobianMode::kAnalytic) { - CheckEvaluate(rb.Evaluate(stream, workspace_ptr, residuals_ptr, ptrs, nullptr, jacobian_ptr), + // Items in the linear system's order (residuals and Jacobians alike); + // the state pointers are per item (FactorBatch::Evaluate). + const int *order = system.linear_system->EvaluationOrder(i); + const size_t num_factors = factor_batch->NumActiveFactors(); + float const *const *item_ptrs = ptrs; + if (order != nullptr && num_factors > 0) { + item_ptrs = GatherItemStatePointers(stream, order, ptrs, num_factors, + static_cast(factor_batch->StateSizes().size()), + scratch.item_state_pointers); + } + CheckEvaluate(rb.Evaluate(stream, workspace_ptr, residuals_ptr, item_ptrs, nullptr, + jacobian_ptr, order, order != nullptr ? num_factors : 0), i); } else { CheckEvaluate(factor_batch->Evaluate(residuals_ptr, nullptr, ptrs, stream), i); @@ -159,18 +199,27 @@ void ComputeResidualAndJacobian(cudaStream_t stream, const Problem &problem, * column scaling (dx = S z). */ void SolveStep(cudaStream_t stream, const MinimizerOptions &options, const MinimizerState &state, - const internal::InnerSolve &inner, MinimizerSystem &system) { + const internal::InnerSolve &inner, bool reject_failed_solve, + MinimizerSystem &system) { + // A failed solve either throws or, when the minimizer can recover (by more + // damping), becomes a NaN step: its trial cost is non-finite, so step + // control rejects it like any other bad step. auto solve = [&]() { - if (!system.normal_equations.Solve(stream, *system.solver, system.rhs, system.step)) { + if (system.linear_system->Solve(stream, system.rhs, system.step)) return true; + if (!reject_failed_solve) { std::string str = "Failed to solve linear system"; LogError(str); throw std::runtime_error(str); } + LogWarning("Failed to solve the linear system; the step is rejected"); + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(system.step.data(), 0xff, system.step.size() * sizeof(float), stream)); + return false; }; - solve(); + if (!solve()) return; while (inner.restrict_system && - inner.restrict_system(stream, state, system.normal_equations, system.rhs, &system.step)) { - solve(); + inner.restrict_system(stream, state, *system.linear_system, system.rhs, &system.step)) { + if (!solve()) return; } if (options.column_scaling != ColumnScaling::None && !system.step.empty()) { ElementwiseMultiplyInPlace(stream, system.step.data(), system.column_scale.data(), @@ -225,27 +274,57 @@ void SetUpStructure(cudaStream_t stream, Problem &problem, const MinimizerOption state_ops.Preprocess(stream, problem.GetStateBatches()); // The Hessian pattern comes straight from the factor graph; no global // Jacobian is involved. - system.normal_equations.Initialize(stream, problem, - static_cast(state_ops.NumReducedStates()), - system.solver->SupportsBlockStorage()); - const size_t jacobian_floats = system.normal_equations.JacobianValuesSize(); + std::vector analytic(problem.GetResidualBatches().size()); + for (size_t i = 0; i < analytic.size(); i++) { + analytic[i] = problem.JacobianModeFor(i, options.jacobian_mode) == JacobianMode::kAnalytic; + } + system.linear_system->Initialize(stream, problem, static_cast(state_ops.NumReducedStates()), + analytic); + const size_t jacobian_floats = system.linear_system->JacobianValuesSize(); if (system.factor_jacobians.size() != jacobian_floats) { system.factor_jacobians.resize(jacobian_floats); } } +std::unique_ptr CreateLinearSystem(const MinimizerOptions &options) { + const auto create_solver = [&](SparseLinearSolverType type) { + SparseLinearSolverPtr solver = + CreateSparseLinearSolver(type, options.sparse_linear_solver_config); + if (options.disable_safety_checks) solver->DisableSafetyChecks(); + return solver; + }; + switch (options.linear_system) { + case LinearSystemType::Assembled: + return std::make_unique(create_solver(options.sparse_linear_solver_type)); + case LinearSystemType::SchurComplement: + // The reduced system has a sparse block pattern (dense covisibility + // blocks, banded for visual-inertial chains) and must be solved exactly. + if (options.sparse_linear_solver_type != SparseLinearSolverType::BlockSparseCholesky) { + LogWarning( + "LinearSystemType::SchurComplement solves the reduced system with " + "BlockSparseCholesky; MinimizerOptions::sparse_linear_solver_type is not used"); + } + return std::make_unique( + create_solver(SparseLinearSolverType::BlockSparseCholesky), + !options.disable_safety_checks); + } + throw std::invalid_argument("MinimizerOptions::linear_system: unknown LinearSystemType"); +} + void BuildSystem(cudaStream_t stream, const Problem &problem, const MinimizerState &state, const MinimizerOptions &options, MinimizerSystem &system, MinimizerScratch &scratch) { ComputeResidualAndJacobian(stream, problem, state, options, system, scratch); - system.normal_equations.Assemble(stream, problem, system.factor_jacobians.data(), - system.residuals.data(), system.rhs); + LinearSystem &linear_system = *system.linear_system; + linear_system.Assemble(stream, problem, system.factor_jacobians.data(), system.residuals.data(), + system.rhs); + linear_system.ClearModifications(); if (options.column_scaling != ColumnScaling::None) { // H_jj = ||J_{:,j}||^2: scaling by the Hessian diagonal is scaling by the // Jacobian column norms. - system.normal_equations.ExtractHessianDiagonal(stream, system.column_scale); + linear_system.ExtractHessianDiagonal(stream, system.column_scale); InvertSqrtWithFloorInPlace(stream, system.column_scale); - system.normal_equations.ScaleLhsSymmetric(stream, system.column_scale); + linear_system.SetColumnScale(&system.column_scale); ElementwiseMultiplyInPlace(stream, system.rhs.data(), system.column_scale.data(), system.rhs.size()); } @@ -254,9 +333,7 @@ void BuildSystem(cudaStream_t stream, const Problem &problem, const MinimizerSta } // namespace internal Minimizer::Minimizer(const MinimizerOptions &options) : options_(options) { - system_.solver = CreateSparseLinearSolver(options_.sparse_linear_solver_type, - options_.sparse_linear_solver_config); - if (options_.disable_safety_checks) system_.solver->DisableSafetyChecks(); + system_.linear_system = internal::CreateLinearSystem(options_); } MinimizerSummary Minimizer::Minimize(cudaStream_t stream, Problem &problem) { @@ -367,20 +444,24 @@ MinimizerSummary internal::InnerMinimize(Minimizer &m, cudaStream_t stream, Prob return summary; } + // The subclass's update (damping) and the caller's restriction, on the + // system built at the current states. + auto update = [&](size_t iteration) { + m.UpdateSystem(stream, iteration, *system.linear_system, partition); + if (inner.restrict_system) { + inner.restrict_system(stream, current, *system.linear_system, system.rhs, nullptr); + } + }; auto build = [&](size_t iteration) { auto build_range = m.profiler_domain_.CreateDomainRange("BuildSystem"); BuildSystem(stream, problem, current, options, system, scratch); - m.UpdateSystem(stream, iteration, system.normal_equations, partition); - if (inner.restrict_system) { - inner.restrict_system(stream, current, system.normal_equations, system.rhs, nullptr); - } + update(iteration); }; build(0); system.step.resize(system.rhs.size()); if (!structure.solver_ready) { auto sa_range = m.profiler_domain_.CreateDomainRange("PerformSymbolicAnalysis"); - if (!system.normal_equations.InitializeSolver(stream, *system.solver, problem, system.rhs, - system.step)) { + if (!system.linear_system->PrepareSolver(stream, problem, system.rhs, system.step)) { std::string str = "Failed to initialize linear solver"; LogError(str); throw std::runtime_error(str); @@ -398,7 +479,7 @@ MinimizerSummary internal::InnerMinimize(Minimizer &m, cudaStream_t stream, Prob summary.iteration_costs.push_back(summary.final_cost); { auto solve_range = m.profiler_domain_.CreateDomainRange("LinearSolve"); - SolveStep(stream, options, current, inner, system); + SolveStep(stream, options, current, inner, m.RejectsFailedSolves(), system); } UpdateStates(stream, m.state_ops_, current, system.step, updated); if (params.line_search) partition.ResetLineSearch(stream); @@ -421,22 +502,34 @@ MinimizerSummary internal::InnerMinimize(Minimizer &m, cudaStream_t stream, Prob evaluate(updated, partition.NewCost()); } partition.SumRows(stream, system.step.data(), nullptr, nullptr, partition.StepSquared()); - m.ClassifySteps(stream, system.normal_equations, system.step, partition); + m.ClassifySteps(stream, *system.linear_system, system.step, partition); partition.StepControl(stream, params, d_out); // One subproblem, with a read-back anyway: a taken step swaps the trial // states in instead of copying them. const bool swap_taken = !fixed && partition.NumProblems() == 1; if (!swap_taken) partition.CopyAccepted(stream, problem, updated, current); summary.num_iterations++; + // A rejected step leaves the states, so the linearization stands: only + // the update (new damping) is redone. Not with restrict_system, which + // zeroed right-hand side entries in place: the system is rebuilt. + bool relinearize = true; if (!fixed) { std::tie(cost, active) = read_back(swap_taken ? 3 : 2); - if (swap_taken && scratch.h_scalars[2] != 0.f) std::swap(current, updated); + const bool step_taken = swap_taken && scratch.h_scalars[2] != 0.f; + if (step_taken) std::swap(current, updated); + relinearize = !swap_taken || step_taken || inner.restrict_system; summary.final_cost = cost; LogMessage("Iteration #{}: cost = {}, {} subproblems active", summary.num_iterations, cost, active); if (active == 0) break; } - if (summary.num_iterations < inner.max_num_iterations) build(summary.num_iterations); + if (summary.num_iterations < inner.max_num_iterations) { + if (relinearize) { + build(summary.num_iterations); + } else { + update(summary.num_iterations); + } + } } if (inner.problem_at_cap != nullptr) { diff --git a/cunls/minimizer/minimizer.h b/cunls/minimizer/minimizer.h index 9ad0794..179d904 100644 --- a/cunls/minimizer/minimizer.h +++ b/cunls/minimizer/minimizer.h @@ -20,6 +20,7 @@ #include #include +#include #include #include "cunls/common/pinned_vector.h" @@ -27,8 +28,8 @@ #include "cunls/common/types.h" #include "cunls/linear_solver/sparse_linear_solver.h" #include "cunls/minimizer/jacobian_mode.h" +#include "cunls/minimizer/linear_system.h" #include "cunls/minimizer/minimizer_state.h" -#include "cunls/minimizer/normal_equations.h" #include "cunls/minimizer/numeric_diff_jacobian.h" #include "cunls/minimizer/problem.h" #include "cunls/minimizer/problem_partition.h" @@ -120,22 +121,30 @@ struct MinimizerOptions { * @brief Type of sparse linear solver to use. * * Supported options: - * - BlockSparsePCG (default): block-Jacobi preconditioned CG; derives - * the per-state-batch block layout automatically. Recommended for - * most SBA and PGO workloads — see `profile/PCG_RESULTS.md` and - * `profile/benchmark_a6000.png`. - * - cuDSS: NVIDIA cuDSS sparse direct solver. Preferred when many - * near-identical small Hessians are solved back-to-back and the - * per-iteration kernel-launch overhead of PCG dominates. + * - BlockSparseCholesky (default): supernodal sparse Cholesky on the GPU + * (no external dependencies); exact steps. The fastest choice for pose + * graphs, bundle adjustment and visual-inertial problems, whose graphs + * have local structure (small separators). Requires an SPD system: + * Gauss-Newton needs every gauge fixed (Levenberg-Marquardt damps it). + * - BlockSparsePCG: block-Jacobi preconditioned CG; derives the + * per-state-batch block layout automatically. Inexact steps (stops at + * a relative residual). Preferred when the graph has no small + * separators, so that a direct factor fills in (many long-range + * random edges), and for very large well-conditioned systems of small + * blocks, where a few CG iterations beat a factorization. + * - cuDSS: NVIDIA cuDSS sparse direct solver (optional dependency). * - DenseLDLT: converts CSR to dense, solves with a custom pivoted * LDLT kernel. Suitable for small dense Hessians. * - DenseCholesky / DenseQR: cuSOLVER-backed dense solves. * - BlockTridiagonal: block-tridiagonal Cholesky for stage-ordered * problems (Problem::SetStateStages). * - * Default: BlockSparsePCG. + * Not used with LinearSystemType::SchurComplement (always + * BlockSparseCholesky). + * + * Default: BlockSparseCholesky. */ - SparseLinearSolverType sparse_linear_solver_type = SparseLinearSolverType::BlockSparsePCG; + SparseLinearSolverType sparse_linear_solver_type = SparseLinearSolverType::BlockSparseCholesky; /** * @brief Configuration for the sparse linear solver. @@ -156,6 +165,12 @@ struct MinimizerOptions { */ SparseLinearSolverConfig sparse_linear_solver_config = {}; + /** + * @brief How each iteration's linear system is formed; see LinearSystemType. + * Default: Assembled. + */ + LinearSystemType linear_system = LinearSystemType::Assembled; + /** * @brief Optional diagonal scaling of the normal equations. * @@ -169,6 +184,10 @@ struct MinimizerOptions { * When false, the minimizer enables all optional runtime validation * that can catch numerical problems early. Currently this covers * post-factorization checks in the linear solver: + * - BlockSparseCholesky: the read-back of the pivot-failure flag (a + * clearly negative or non-finite pivot). + * - SchurComplement: the read-back of the landmark-block failure flag (a + * landmark block that cannot be factored even with a diagonal shift). * - Cholesky: cuSOLVER devInfo after potrf and potrs (non-SPD or * invalid-parameter detection). * - QR: diagonal-of-R inspection for rank deficiency. @@ -176,8 +195,9 @@ struct MinimizerOptions { * and solve. * * On failure the solver's Solve() returns false and a diagnostic is - * emitted via LogError(). The minimizer treats a false return as a - * fatal error and throws std::runtime_error. + * emitted via LogError(). Levenberg-Marquardt treats a false return as + * a rejected step and raises its damping; Gauss-Newton throws + * std::runtime_error (see Minimizer::RejectsFailedSolves()). * * When true, every check listed above is skipped: no device-to-host * memcpy, no stream synchronization, and no in-kernel validation. @@ -262,14 +282,14 @@ struct InnerSolve { /// and box-bounded states are kept in their bounds by its restrict_system /// (otherwise a problem with either is rejected). bool constraints_managed = false; - /// Removes unknowns from the linear system: zeroes their rows, columns and - /// right-hand side entries (keeping the diagonal), so their step is exactly - /// 0. Called every iteration once the normal equations are built (and + /// Removes unknowns from the linear system: declares the mask that holds + /// them (LinearSystem::SetMask) and zeroes their right-hand side entries, + /// so their step is exactly 0. Called every iteration once the system is built (and /// updated by the subclass) with `step` null, before the solve; then after /// each solve with the solved `step`: returning true asks for another solve /// (it removed more unknowns), false keeps the step. `state` holds the /// current states. Empty: nothing is removed. - std::function &rhs, const dvector *step)> restrict_system; /// Device array of Problem::NumProblems() ints, written: nonzero where a @@ -294,22 +314,22 @@ struct MinimizerStructure { /** @brief The linearized problem and its solution. */ struct MinimizerSystem { - SparseLinearSolverPtr solver; ///< Linear solver of the normal equations. - NumericDiffJacobianBuilder numeric_diff; ///< Jacobians of kNumeric residual batches. - dvector residuals; ///< Residuals of every active factor. - PerFactorJacobians factor_jacobians; ///< Dense Jacobian block of every factor. - NormalEquations normal_equations; ///< H = JᵀJ and the working left-hand side. - dvector rhs; ///< -Jᵀr (scaled and masked like the lhs). - dvector column_scale; ///< S of ColumnScaling::HessianDiagonal. - dvector step; ///< The step δ, one entry per reduced row. + std::unique_ptr linear_system; ///< H = JᵀJ, its modifications and solver. + NumericDiffJacobianBuilder numeric_diff; ///< Jacobians of kNumeric residual batches. + dvector residuals; ///< Residuals of every active factor. + PerFactorJacobians factor_jacobians; ///< Dense Jacobian block of every factor. + dvector rhs; ///< -Jᵀr (scaled and masked like the system). + dvector column_scale; ///< S of ColumnScaling::HessianDiagonal. + dvector step; ///< The step δ, one entry per free column. }; /** @brief Scratch buffers and read-back slots. */ struct MinimizerScratch { - dvector buffer; ///< Per-factor costs, residuals and workspaces. - dvector d_scalars; ///< Device scalars: totals and read-back slots. - PinnedVector h_scalars; ///< Host mirror of d_scalars. - dvector d_partials; ///< Reduction partials. + dvector buffer; ///< Per-factor costs, residuals and workspaces. + dvector d_scalars; ///< Device scalars: totals and read-back slots. + PinnedVector h_scalars; ///< Host mirror of d_scalars. + dvector d_partials; ///< Reduction partials. + dvector item_state_pointers; ///< State pointers in LinearSystem::EvaluationOrder(). }; /** @@ -321,6 +341,17 @@ struct MinimizerScratch { void SetUpStructure(cudaStream_t stream, Problem &problem, const MinimizerOptions &options, StateBatchOps &state_ops, MinimizerSystem &system); +/** + * @brief The linear system MinimizerOptions::linear_system selects: the + * assembled normal equations with the solver of + * MinimizerOptions::sparse_linear_solver_type, or the Schur complement with + * BlockSparseCholesky (sparse_linear_solver_type ignored, with a warning when + * set to another solver). + * + * @throws std::invalid_argument for an unknown LinearSystemType. + */ +std::unique_ptr CreateLinearSystem(const MinimizerOptions &options); + /** * @brief Step 1 of Minimize(): the Gauss-Newton system at `state`, H = JᵀJ and * rhs = -Jᵀr, then S H S z = S b with column scaling. Requires SetUpStructure @@ -337,8 +368,8 @@ void BuildSystem(cudaStream_t stream, const Problem &problem, const MinimizerSta * @brief Common base of GaussNewtonMinimizer and LevenbergMarquardtMinimizer. * * Minimizes the problem's total cost, ½ Σ ρ(‖r‖²), by repeating: - * 1. build the normal equations H δ = -g at the current states (with - * column scaling), + * 1. build the linear system H δ = -g at the current states (LinearSystem; + * with column scaling), * 2. let the subclass update them (UpdateSystem), * 3. solve for the step δ, * 4. evaluate the cost at the trial states, shortening δ by line search @@ -347,6 +378,11 @@ void BuildSystem(cudaStream_t stream, const Problem &problem, const MinimizerSta * 6. take or reject each step, and stop each subproblem that converged or * hit the rejection cap (MinimizerOptions::max_consecutive_rejected_steps). * + * When the only subproblem rejects its step, the states are unchanged and + * the linearization is kept: the next iteration redoes step 2 only (new + * damping), unless InnerSolve::restrict_system is set. With + * RejectsFailedSolves() a failed solve in step 3 counts as a rejected step. + * * A problem made of independent subproblems (Problem::SetProblemPartition) * is solved in one linear system, but every decision in steps 4 to 6 is * taken per subproblem. A problem without a partition is one subproblem. @@ -381,7 +417,8 @@ class Minimizer { * AugmentedLagrangianMinimizer), or its * sizes, connectivity or partition are invalid (see * Problem::CheckSizes). - * @throws std::runtime_error if the linear solver fails. + * @throws std::runtime_error if the linear solver fails and the method + * cannot recover from it (see RejectsFailedSolves()). */ MinimizerSummary Minimize(cudaStream_t stream, Problem &problem); @@ -394,26 +431,26 @@ class Minimizer { protected: /** - * @brief For subclasses: stores the options and creates the linear solver. + * @brief For subclasses: stores the options and creates the linear system. * @param options Options common to all minimizers. */ explicit Minimizer(const MinimizerOptions &options); private: /** - * @brief Subclass step 2: modifies the normal equations before the solve. + * @brief Subclass step 2: modifies the linear system before the solve. * * Called once per iteration, after the Gauss-Newton system for the current * states has been assembled into `system`. Gauss-Newton leaves it as is; - * Levenberg-Marquardt adds its damping. + * Levenberg-Marquardt declares its damping (LinearSystem::SetDamping). * * @param stream CUDA stream; enqueue only, do not synchronize. * @param iteration 0 on the first iteration of this Minimize() call. - * @param system Normal equations; the working left-hand side may be modified. + * @param system The linear system; its damping may be declared. * @param partition Subproblems, with the outcome of the previous * iteration's steps (ProblemPartition::Outcome); read only. */ - virtual void UpdateSystem(cudaStream_t stream, size_t iteration, NormalEquations &system, + virtual void UpdateSystem(cudaStream_t stream, size_t iteration, LinearSystem &system, const ProblemPartition &partition) = 0; /** @@ -431,16 +468,23 @@ class Minimizer { * non-finite cost is always rejected. * * @param stream CUDA stream; enqueue only, do not synchronize. - * @param system Normal equations the step was solved from (the assembled - * Hessian and the working left-hand side). - * @param step The (possibly shortened) step δ, one entry per reduced row. + * @param system The linear system the step was solved from. + * @param step The (possibly shortened) step δ, one entry per free column. * @param partition Subproblems with current cost (Cost), trial cost * (NewCost) and ‖δ‖² (StepSquared); write Reject() and Converged() * only. */ - virtual void ClassifySteps(cudaStream_t stream, const NormalEquations &system, + virtual void ClassifySteps(cudaStream_t stream, const LinearSystem &system, const dvector &step, ProblemPartition &partition) = 0; + /** + * @brief Whether a failed linear solve (LinearSystem::Solve() returning + * false, e.g. a factorization that found the system not positive definite) + * counts as a rejected step, which this method answers with more damping. + * Otherwise Minimize() throws std::runtime_error. Default: false. + */ + virtual bool RejectsFailedSolves() const { return false; } + friend MinimizerSummary internal::InnerMinimize(Minimizer &minimizer, cudaStream_t stream, Problem &problem, const internal::InnerSolve &inner); @@ -448,7 +492,7 @@ class Minimizer { // Implementation state; the behaviour lives in minimizer.cu. const MinimizerOptions options_; internal::MinimizerStructure structure_; ///< Cached setup of the last problem. - internal::MinimizerSystem system_; ///< Linearization, normal equations, step. + internal::MinimizerSystem system_; ///< Linearization, linear system, step. internal::MinimizerScratch scratch_; ///< Buffers and read-back slots. StateBatchOps state_ops_; ///< Manifold plus and tangent layout. ProblemPartition partition_; ///< Subproblems and their step bookkeeping. diff --git a/cunls/minimizer/normal_equations.cu b/cunls/minimizer/normal_equations.cu index 55a9294..5cd5c9e 100644 --- a/cunls/minimizer/normal_equations.cu +++ b/cunls/minimizer/normal_equations.cu @@ -15,26 +15,30 @@ * limitations under the License. */ +#include + #include "cunls/common/log.h" #include "cunls/minimizer/normal_equations.h" #include "cunls/minimizer/problem.h" +#include "cunls/minimizer/problem_partition.h" #include "cunls/minimizer/sparse_matrix.h" namespace cunls { -void NormalEquations::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, - bool solver_supports_block_storage) { - csr_dims_.Invalidate(); +NormalEquations::NormalEquations(SparseLinearSolverPtr solver) : solver_(std::move(solver)) {} +void NormalEquations::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + const std::vector &) { // Two independent conditions have to hold. The tangent dimensions must share // a factor, or there are no tiles to form; and the solver must be able to read // tiles, since a CSR-only backend is better served by a natively assembled CSR // than by one expanded from blocks. const int problem_block_size = ChooseHessianBlockSize(problem); - block_size_ = solver_supports_block_storage ? problem_block_size : 1; + block_size_ = solver_->SupportsBlockStorage() ? problem_block_size : 1; if (UsesBlockStorage()) { assembler_.Initialize(stream, problem, num_cols, block_size_, bsr_hessian_); + CopyBSRSparseMatrix(stream, bsr_hessian_, bsr_lhs_); // the pattern, once LogMessage("Hessian storage: BSR with block size {}", block_size_); return; } @@ -52,28 +56,18 @@ void NormalEquations::Initialize(cudaStream_t stream, const Problem &problem, in } assembler_.Initialize(stream, problem, num_cols, csr_hessian_); - int num_rows = 0, num_matrix_cols = 0, num_nonzeros = 0; - ExtractMatrixMetadata(stream, csr_hessian_, num_rows, num_matrix_cols, num_nonzeros); - csr_dims_.Set(num_rows, num_matrix_cols, num_nonzeros); - // The pattern is fixed from here on: locate the diagonal once. + // The pattern is fixed from here on: copy it and locate the diagonal once. + CopyCSRSparseMatrix(stream, csr_hessian_, csr_lhs_); FindDiagonalPositions(stream, csr_hessian_, csr_diagonal_); - csr_lhs_pattern_ready_ = false; } void NormalEquations::Assemble(cudaStream_t stream, const Problem &problem, const float *jacobians, const float *residuals, dvector &rhs) { if (UsesBlockStorage()) { assembler_.Assemble(stream, problem, jacobians, residuals, bsr_hessian_, rhs); - CopyBSRSparseMatrix(stream, bsr_hessian_, bsr_lhs_); return; } assembler_.Assemble(stream, problem, jacobians, residuals, csr_hessian_, rhs); - if (csr_lhs_pattern_ready_) { - CopyCSRValues(stream, csr_hessian_, csr_lhs_); - } else { - CopyCSRSparseMatrix(stream, csr_hessian_, csr_lhs_); - csr_lhs_pattern_ready_ = true; - } } void NormalEquations::ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const { @@ -84,65 +78,55 @@ void NormalEquations::ExtractHessianDiagonal(cudaStream_t stream, dvector ExtractDiagonalAt(stream, csr_hessian_, csr_diagonal_, diagonal); } -void NormalEquations::ExtractLhsDiagonal(cudaStream_t stream, dvector &diagonal) const { - if (UsesBlockStorage()) { - ExtractDiagonal(stream, bsr_lhs_, diagonal); - return; - } - ExtractDiagonalAt(stream, csr_lhs_, csr_diagonal_, diagonal); -} - -void NormalEquations::AddScaledDiagonalToLhs(cudaStream_t stream, float scale, - const dvector &diagonal) { +void NormalEquations::BuildWorkingLhs(cudaStream_t stream) { if (UsesBlockStorage()) { - AddScaledDiagonal(stream, scale, diagonal, bsr_lhs_, bsr_lhs_); + CopyBSRValues(stream, bsr_hessian_, bsr_lhs_); + if (ColumnScale() != nullptr) { + ScaleSymmetric(stream, bsr_lhs_, *ColumnScale(), tile_row_scratch_); + } + if (Damping() != nullptr) AddScaledDiagonal(stream, 1.f, *Damping(), bsr_lhs_, bsr_lhs_); + if (Mask() != nullptr) ZeroMaskedRowsColumns(stream, bsr_lhs_, *Mask()); return; } - AddScaledDiagonalAt(stream, scale, diagonal, csr_diagonal_, csr_lhs_); + CopyCSRValues(stream, csr_hessian_, csr_lhs_); + if (ColumnScale() != nullptr) ScaleSymmetricCSR(stream, csr_lhs_, *ColumnScale()); + if (Damping() != nullptr) AddScaledDiagonalAt(stream, 1.f, *Damping(), csr_diagonal_, csr_lhs_); + if (Mask() != nullptr) ZeroMaskedRowsColumns(stream, csr_lhs_, *Mask()); } -void NormalEquations::ScaleLhsSymmetric(cudaStream_t stream, const dvector &scale) { +void NormalEquations::MultiplyHessian(cudaStream_t stream, const dvector &x, + dvector &y) const { if (UsesBlockStorage()) { - ScaleSymmetric(stream, bsr_lhs_, scale, tile_row_scratch_); - return; - } - ScaleSymmetricCSR(stream, csr_lhs_, scale); -} - -void NormalEquations::ZeroMaskedLhsRowsColumns(cudaStream_t stream, const dvector &mask) { - if (UsesBlockStorage()) { - ZeroMaskedRowsColumns(stream, bsr_lhs_, mask); + MultiplyBSRByDenseVector(stream, bsr_hessian_, x, y); return; } - ZeroMaskedRowsColumns(stream, csr_lhs_, mask); + const int n = static_cast(csr_hessian_.NumRows()); + MultiplyCSRByDenseVector(stream, cusparse_handle_.GetHandle(stream), csr_hessian_, n, n, + static_cast(csr_hessian_.NumNonZeros()), x, y, spmv_buffer_); } -void NormalEquations::MultiplyHessian(cudaStream_t stream, void *cusparse_handle, - const dvector &x, dvector &y, - dvector &buffer) const { - if (UsesBlockStorage()) { - MultiplyBSRByDenseVector(stream, bsr_hessian_, x, y); - return; - } - MultiplyCSRByDenseVector(stream, cusparse_handle, csr_hessian_, csr_dims_.num_rows, - csr_dims_.num_cols, csr_dims_.num_nonzeros, x, y, buffer); +void NormalEquations::StepCurvature(cudaStream_t stream, const dvector &step, + ProblemPartition &partition, float *out) const { + hessian_step_.resize(step.size()); + MultiplyHessian(stream, step, hessian_step_); + partition.SumRows(stream, step.data(), hessian_step_.data(), nullptr, out); } -bool NormalEquations::InitializeSolver(cudaStream_t stream, SparseLinearSolver &solver, - const Problem &problem, const dvector &rhs, - dvector &step) { +bool NormalEquations::PrepareSolver(cudaStream_t stream, const Problem &problem, + const dvector &rhs, dvector &step) { + BuildWorkingLhs(stream); if (UsesBlockStorage()) { - return solver.Initialize(stream, problem, bsr_lhs_, rhs, step); + return solver_->Initialize(stream, problem, bsr_lhs_, rhs, step); } - return solver.Initialize(stream, problem, csr_lhs_, rhs, step); + return solver_->Initialize(stream, problem, csr_lhs_, rhs, step); } -bool NormalEquations::Solve(cudaStream_t stream, SparseLinearSolver &solver, - const dvector &rhs, dvector &step) { +bool NormalEquations::Solve(cudaStream_t stream, const dvector &rhs, dvector &step) { + BuildWorkingLhs(stream); if (UsesBlockStorage()) { - return solver.Solve(stream, bsr_lhs_, rhs, step); + return solver_->Solve(stream, bsr_lhs_, rhs, step); } - return solver.Solve(stream, csr_lhs_, rhs, step); + return solver_->Solve(stream, csr_lhs_, rhs, step); } } // namespace cunls diff --git a/cunls/minimizer/normal_equations.h b/cunls/minimizer/normal_equations.h index 73c64a3..6065961 100644 --- a/cunls/minimizer/normal_equations.h +++ b/cunls/minimizer/normal_equations.h @@ -19,30 +19,34 @@ #include +#include "cunls/common/cusparse_helper.h" #include "cunls/common/types.h" -#include "cunls/linear_solver/sparse_linear_solver_base.h" +#include "cunls/linear_solver/sparse_linear_solver.h" #include "cunls/minimizer/block_hessian_assembler.h" #include "cunls/minimizer/bsr_matrix.h" +#include "cunls/minimizer/linear_system.h" namespace cunls { class Problem; +class ProblemPartition; /** - * @brief The assembled normal equations, and the storage they live in. + * @brief The assembled normal equations (LinearSystemType::Assembled). * - * Holds two matrices: the Hessian `H` as assembled, and a working copy the - * minimizer is free to damp and scale before handing it to the solver. Both - * are stored in one of two layouts, chosen once at Initialize(): + * Holds two matrices: the Hessian `H` as assembled, and the working + * left-hand side handed to the solver, rebuilt from `H` before every solver + * call (PrepareSolver(), Solve()) with the declared column scale, damping and + * mask. Both are stored in one of two + * layouts, chosen once at Initialize(): * * - **Block (BSR)** when the state tangent dimensions share a common factor and * the solver can read tiles. One column index per dense tile instead of one * per scalar entry, which is bandwidth the solver's SpMV no longer moves. * - **Scalar (CSR)** otherwise. * - * Which layout is live is an implementation detail. Callers work in terms of - * "the Hessian" and "the working left-hand side"; every operation dispatches - * internally, so no branch on storage leaks into the minimizers. + * Which layout is live is an implementation detail; every operation + * dispatches internally. * * A BSR matrix with `block_size == 1` is byte-for-byte the same layout as CSR, * so the scalar case could in principle be expressed as block storage too. It @@ -52,68 +56,56 @@ class Problem; * one-tile-per-entry BSR kernel cannot beat. Block storage is a win only when * the tiles are real. */ -class NormalEquations { +class NormalEquations : public LinearSystem { public: + /** @brief Takes ownership of the solver of the working left-hand side (not null). */ + explicit NormalEquations(SparseLinearSolverPtr solver); + /** - * @brief Derives the sparsity pattern and picks the storage layout. - * - * @param stream CUDA stream for GPU operations. - * @param problem The optimization problem. - * @param num_cols Number of free tangent dimensions in the reduced system. - * @param solver_supports_block_storage Whether the active solver can consume - * block storage; see SparseLinearSolver::SupportsBlockStorage. + * @brief Derives the sparsity pattern of H and picks the storage layout: + * block storage when the tangent sizes share a factor and the solver + * accepts it (SparseLinearSolver::SupportsBlockStorage()), scalar CSR + * otherwise. Every batch is evaluated in its natural order (`reorderable` + * is not used). */ void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, - bool solver_supports_block_storage); + const std::vector &reorderable) override; - /** @brief Floats needed for the per-factor Jacobian buffer Assemble() reads. */ - size_t JacobianValuesSize() const { return assembler_.JacobianValuesSize(); } + /** @copydoc LinearSystem::JacobianValuesSize */ + size_t JacobianValuesSize() const override { return assembler_.JacobianValuesSize(); } - /** - * @brief Assembles `H` and `rhs`, then refreshes the working left-hand side. - * - * @param stream CUDA stream for GPU operations. - * @param problem The optimization problem (must match Initialize). - * @param jacobians Per-factor dense Jacobian blocks. - * @param residuals Residual vector. - * @param[out] rhs Right-hand side `-J^T r`. - */ + /** @brief Assembles `H` and `rhs` = -Jᵀr. */ void Assemble(cudaStream_t stream, const Problem &problem, const float *jacobians, - const float *residuals, dvector &rhs); - - /** @brief Diagonal of the assembled (undamped) Hessian. */ - void ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const; + const float *residuals, dvector &rhs) override; - /** @brief Diagonal of the working left-hand side. */ - void ExtractLhsDiagonal(cudaStream_t stream, dvector &diagonal) const; + /** @copydoc LinearSystem::ExtractHessianDiagonal */ + void ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const override; - /** @brief Adds `scale * diag(diagonal)` to the working left-hand side. */ - void AddScaledDiagonalToLhs(cudaStream_t stream, float scale, const dvector &diagonal); + /** @brief Builds the working left-hand side and runs the solver's symbolic setup. */ + bool PrepareSolver(cudaStream_t stream, const Problem &problem, const dvector &rhs, + dvector &step) override; - /** @brief Applies `A_ij *= scale[i] * scale[j]` to the working left-hand side. */ - void ScaleLhsSymmetric(cudaStream_t stream, const dvector &scale); - - /** - * @brief Zeroes the rows and columns i with mask[i] == 0 of the working - * left-hand side, keeping a positive diagonal (see ZeroMaskedRowsColumns). - */ - void ZeroMaskedLhsRowsColumns(cudaStream_t stream, const dvector &mask); + /** @brief Builds the working left-hand side and solves with it. */ + bool Solve(cudaStream_t stream, const dvector &rhs, dvector &step) override; - /** @brief y = (assembled, undamped Hessian) * x, in either storage layout. */ - void MultiplyHessian(cudaStream_t stream, void *cusparse_handle, const dvector &x, - dvector &y, dvector &buffer) const; + /** @brief δᵀHδ per subproblem, from H δ (SpMV) in either layout. */ + void StepCurvature(cudaStream_t stream, const dvector &step, ProblemPartition &partition, + float *out) const override; - /** @brief Hands the working left-hand side to the solver for symbolic setup. */ - bool InitializeSolver(cudaStream_t stream, SparseLinearSolver &solver, const Problem &problem, - const dvector &rhs, dvector &step); + // Test hooks: the matrices and products the layouts produce. - /** @brief Solves with the working left-hand side. */ - bool Solve(cudaStream_t stream, SparseLinearSolver &solver, const dvector &rhs, - dvector &step); + /** @brief y = H x with the assembled, undamped Hessian, in either layout. */ + void MultiplyHessian(cudaStream_t stream, const dvector &x, dvector &y) const; /** @brief True when the block layout is live for the current problem. */ bool UsesBlockStorage() const { return block_size_ > 1; } + /** + * @brief Rebuilds the working left-hand side from H with the declared scale, + * damping and mask (PrepareSolver() and Solve() do this first). + */ + void BuildWorkingLhs(cudaStream_t stream); + /** @brief Working left-hand side in scalar storage; empty under block storage. */ const CSRSparseMatrix &LhsCSR() const { return csr_lhs_; } @@ -121,20 +113,22 @@ class NormalEquations { const BSRSparseMatrix &LhsBSR() const { return bsr_lhs_; } private: - BlockHessianAssembler assembler_; - - // Exactly one pair is populated, decided by block_size_. - CSRSparseMatrix csr_hessian_; - CSRSparseMatrix csr_lhs_; - BSRSparseMatrix bsr_hessian_; - BSRSparseMatrix bsr_lhs_; - - CSRMatrixDimensions csr_dims_; ///< Cached dims for the scalar SpMV. - dvector tile_row_scratch_; ///< Tile-to-block-row map for scaling. - dvector csr_diagonal_; ///< Scalar layout: position of each diagonal entry. - bool csr_lhs_pattern_ready_ = false; ///< csr_lhs_ holds the pattern; copy values only. - - int block_size_ = 1; + SparseLinearSolverPtr solver_; ///< Solver of the working left-hand side. + BlockHessianAssembler assembler_; ///< Pattern and assembly of H from the factors. + int block_size_ = 1; ///< Tile size of the block layout; 1: scalar layout. + + // Exactly one layout is populated, decided by block_size_. + CSRSparseMatrix csr_hessian_; ///< Scalar layout: H as assembled. + CSRSparseMatrix csr_lhs_; ///< Scalar layout: the working left-hand side (H's pattern). + dvector csr_diagonal_; ///< Scalar layout: position of each diagonal entry. + BSRSparseMatrix bsr_hessian_; ///< Block layout: H as assembled. + BSRSparseMatrix bsr_lhs_; ///< Block layout: the working left-hand side. + dvector tile_row_scratch_; ///< Block layout: tile -> block row (column scaling). + + // Scratch of MultiplyHessian() / StepCurvature(). + mutable cuSPARSEHandle cusparse_handle_; ///< Scalar-layout SpMV. + mutable dvector spmv_buffer_; ///< cuSPARSE workspace of the SpMV. + mutable dvector hessian_step_; ///< H δ. }; } // namespace cunls diff --git a/cunls/minimizer/problem.cpp b/cunls/minimizer/problem.cpp index 72a35d0..667da5e 100644 --- a/cunls/minimizer/problem.cpp +++ b/cunls/minimizer/problem.cpp @@ -353,6 +353,24 @@ void Problem::SetStateStages(const std::vector &state_stages) { state_stages_ = state_stages; } +void Problem::SetEliminatedStateBatches(const std::vector &batches) { + for (size_t i = 0; i < batches.size(); ++i) { + const StateBatch *batch = batches[i]; + if (batch == nullptr) { + throw std::invalid_argument("Problem::SetEliminatedStateBatches: null state batch"); + } + if (std::find(state_batches_.begin(), state_batches_.end(), batch) == state_batches_.end()) { + throw std::invalid_argument( + "Problem::SetEliminatedStateBatches: the state batch is not part of this problem"); + } + if (std::find(batches.begin(), batches.begin() + i, batch) != batches.begin() + i) { + throw std::invalid_argument( + "Problem::SetEliminatedStateBatches: a state batch is listed twice"); + } + } + eliminated_state_batches_ = batches; +} + void Problem::CheckSizes() const { const std::string error = SizeError(); if (!error.empty()) { diff --git a/cunls/minimizer/problem.h b/cunls/minimizer/problem.h index 5da3f9e..f7a2e33 100644 --- a/cunls/minimizer/problem.h +++ b/cunls/minimizer/problem.h @@ -333,6 +333,31 @@ class Problem { /** @brief Per state batch: device stage ids (empty when not set). */ const std::vector &StateStages() const { return state_stages_; } + /** + * @brief Declares the state batches LinearSystemType::SchurComplement + * eliminates (the landmarks of a bundle adjustment). + * + * Every factor may touch at most one non-constant state of these batches; + * the batches must share one tangent size of 1 to 3, and the states that + * share factors with them one tangent size of 3, 6 or 9. Constant states of + * a listed batch are not eliminated (their factors act on the other states + * only). The requirements are checked when the minimizer sets up the + * structure (std::invalid_argument). Without the declaration + * SchurComplement throws; with LinearSystemType::Assembled the declaration + * has no effect. + * + * @param batches State batches already added to this problem; an empty + * vector clears the declaration. + * @throws std::invalid_argument if a batch is null, not part of this + * problem, or listed twice. + */ + void SetEliminatedStateBatches(const std::vector &batches); + + /** @brief The batches of SetEliminatedStateBatches() (empty when not set). */ + const std::vector &EliminatedStateBatches() const { + return eliminated_state_batches_; + } + private: /** * @brief Validates that all inputs are non-null and sizes are consistent. @@ -380,9 +405,10 @@ class Problem { /// Host pointer lists returned by GetStatePointers(): the user's lists for /// host-list batches, downloads of the device tables otherwise. mutable std::vector> state_pointers_; - size_t num_problems_ = 0; ///< SetProblemPartition(). - std::vector state_problem_ids_; ///< SetProblemPartition(). - std::vector state_stages_; ///< SetStateStages(). + size_t num_problems_ = 0; ///< SetProblemPartition(). + std::vector state_problem_ids_; ///< SetProblemPartition(). + std::vector state_stages_; ///< SetStateStages(). + std::vector eliminated_state_batches_; ///< SetEliminatedStateBatches(). std::vector> jacobian_mode_overrides_; ///< Per-residual-batch JacobianMode ///< override, index-aligned with diff --git a/cunls/minimizer/problem_partition.cu b/cunls/minimizer/problem_partition.cu index a233228..b110cd0 100644 --- a/cunls/minimizer/problem_partition.cu +++ b/cunls/minimizer/problem_partition.cu @@ -208,19 +208,22 @@ __global__ void line_search_kernel(size_t num_problems, const float *cost, const * minimizer's classification (reject, converged); see ProblemPartition::StepControl. */ __global__ void step_control_kernel(StepControlParams params, size_t num_problems, float *cost, - const float *new_cost, const int *reject, const int *converged, - const int *shortened, int *active, int *rejected, int *accept, - int *outcome, float *out) { + const float *new_cost, const float *step_squared, + const int *reject, const int *converged, const int *shortened, + int *active, int *rejected, int *accept, int *outcome, + float *out) { const size_t p = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; if (p >= num_problems) return; int take = 0; int result = kStepInactive; if (active[p]) { const float current = cost[p], updated = new_cost[p]; - // A step to a non-finite cost is rejected, never taken. - const bool finite = isfinite(updated); + // A non-finite step (a failed linear solve) or a step to a non-finite + // cost is rejected, never taken. The step is checked on its own: factors + // may give an invalid state (e.g. a NaN point) a finite, even zero, cost. + const bool finite = isfinite(updated) && isfinite(step_squared[p]); const bool is_rejected = reject[p] != 0 || !finite; - if (params.line_search && updated < current && (shortened[p] || is_rejected)) { + if (params.line_search && finite && updated < current && (shortened[p] || is_rejected)) { // Line search: a (possibly shortened) step that decreases the cost is // taken (see MinimizerOptions::max_line_search_steps). take = 1; @@ -465,8 +468,9 @@ void ProblemPartition::StepControl(cudaStream_t stream, const StepControlParams float *d_out) { if (num_problems_ > 1) THROW_ON_CUDA_ERROR(cudaMemsetAsync(d_out, 0, 2 * sizeof(float), stream)); step_control_kernel<<>>( - params, num_problems_, cost_.data(), new_cost_.data(), reject_.data(), converged_.data(), - shortened_.data(), active_.data(), rejected_.data(), accept_.data(), outcome_.data(), d_out); + params, num_problems_, cost_.data(), new_cost_.data(), step_squared_.data(), reject_.data(), + converged_.data(), shortened_.data(), active_.data(), rejected_.data(), accept_.data(), + outcome_.data(), d_out); THROW_ON_CUDA_ERROR(cudaGetLastError()); } diff --git a/cunls/minimizer/problem_partition.h b/cunls/minimizer/problem_partition.h index d34cdce..07d7f2f 100644 --- a/cunls/minimizer/problem_partition.h +++ b/cunls/minimizer/problem_partition.h @@ -122,7 +122,8 @@ class ProblemPartition { /** * @brief Takes or rejects the step of every active subproblem from - * NewCost(), Reject() and Converged() (see StepOutcome): a step to a + * NewCost(), Reject() and Converged() (see StepOutcome): a non-finite step + * (StepSquared(), e.g. after a failed linear solve) or a step to a * non-finite cost is rejected; with line search, a shortened or rejected * step that lowers the cost is taken; a converged subproblem stops, taking * its step if the cost did not rise; a rejection counts towards the cap. diff --git a/cunls/minimizer/residual_batch.cu b/cunls/minimizer/residual_batch.cu index 32ccaed..39bd06a 100644 --- a/cunls/minimizer/residual_batch.cu +++ b/cunls/minimizer/residual_batch.cu @@ -25,7 +25,6 @@ namespace cunls { constexpr int kBlockSize = 256; -constexpr int kWarpSize = 32; namespace { @@ -127,54 +126,34 @@ __global__ void scale_residuals_thread_kernel(float *__restrict__ residuals, } // ============================================================================ -// Kernel 4: Scale Jacobians (optimized) +// Kernel 4: Scale Jacobians // ============================================================================ // S*J_col = sqrt_rho1*(J_col - alpha*(r'*J_col)*r) // -// Old kernel: 32x32 block per factor (1024 threads). For 2x9 Jacobian only -// 18 values to process, so 99.8% of thread-cycles were wasted. -// -// New design: one warp per factor. Lane 0 computes alpha and sqrt_rho1. -// Lanes cooperatively process columns in a stride pattern. For residual_dim -// <= 32 (covers all practical NLS problems), the dot product r'*J_col is -// computed via warp shuffle reduction, eliminating shared memory entirely. - -__global__ void scale_jacobians_warp_kernel(const float *__restrict__ residuals, - float *__restrict__ jacobians, - const float *__restrict__ squared_error, - const float3 *__restrict__ rho_coeffs, - int num_residuals, int residual_dim, int num_cols) { - const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / kWarpSize; - const int lane = threadIdx.x % kWarpSize; - if (warp_id >= num_residuals) return; - - const int row_offset = warp_id * residual_dim; - const float *res_base = residuals + row_offset; - - float sqrt_rho1, alpha; - { - float sq = squared_error[warp_id]; - float3 rh = rho_coeffs[warp_id]; - sqrt_rho1 = sqrtf(rh.y); - alpha = jacobian_scaling_alpha(sq, rh); - } - - // For each Jacobian column, compute dot = r'*J_col, then apply scaling. - // Each lane handles different columns in a stride pattern. - for (int col = lane; col < num_cols; col += kWarpSize) { - // Compute dot product r' * J_col - float dot = 0.0f; - for (int row = 0; row < residual_dim; ++row) { - dot += res_base[row] * jacobians[num_cols * (row_offset + row) + col]; - } - - float scale = sqrt_rho1 * alpha * dot; - - // Apply: J_col = sqrt_rho1 * J_col - scale * r - for (int row = 0; row < residual_dim; ++row) { - int idx = num_cols * (row_offset + row) + col; - jacobians[idx] = sqrt_rho1 * jacobians[idx] - scale * res_base[row]; - } +// One thread per (factor, Jacobian column). The columns of a factor are +// consecutive threads, so each Jacobian row is read and written coalesced and +// every lane is busy (a warp per factor left 23 of 32 lanes idle on a 2x9 +// reprojection Jacobian). The residual and the loss coefficients of a factor +// are shared by its columns' threads (cache hits). +__global__ void scale_jacobians_kernel(const float *__restrict__ residuals, + float *__restrict__ jacobians, + const float *__restrict__ squared_error, + const float3 *__restrict__ rho_coeffs, int num_residuals, + int residual_dim, int num_cols) { + const long long i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= static_cast(num_residuals) * num_cols) return; + const int factor = static_cast(i / num_cols); + const int col = static_cast(i - static_cast(factor) * num_cols); + const float3 rh = rho_coeffs[factor]; + const float sqrt_rho1 = sqrtf(rh.y); + const float alpha = jacobian_scaling_alpha(squared_error[factor], rh); + const float *r = residuals + static_cast(factor) * residual_dim; + float *J = jacobians + static_cast(factor) * residual_dim * num_cols + col; + float dot = 0.0f; + for (int row = 0; row < residual_dim; ++row) dot += r[row] * J[row * num_cols]; + const float scale = sqrt_rho1 * alpha * dot; + for (int row = 0; row < residual_dim; ++row) { + J[row * num_cols] = sqrt_rho1 * J[row * num_cols] - scale * r[row]; } } @@ -262,12 +241,13 @@ bool ResidualBatch::ApplyLossToItems(cudaStream_t stream, float *workspace, floa int num_cols = 0; for (const auto &d : factor_batch_->StateSizes()) num_cols += d; - int warps_per_block = kBlockSize / kWarpSize; - int jac_blocks = (num_residuals + warps_per_block - 1) / warps_per_block; - - scale_jacobians_warp_kernel<<>>( - residuals, jacobians, sq_err_ptr, rho_ptr, num_residuals, residual_dim, num_cols); - THROW_ON_CUDA_ERROR(cudaGetLastError()); + const size_t num_threads = static_cast(num_residuals) * num_cols; + if (num_threads > 0) { + scale_jacobians_kernel<<((num_threads + kBlockSize - 1) / kBlockSize), + kBlockSize, 0, stream>>>(residuals, jacobians, sq_err_ptr, rho_ptr, + num_residuals, residual_dim, num_cols); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } } // Scale residuals diff --git a/cunls/minimizer/schur_system.cu b/cunls/minimizer/schur_system.cu new file mode 100644 index 0000000..f2f32c8 --- /dev/null +++ b/cunls/minimizer/schur_system.cu @@ -0,0 +1,1993 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "cunls/common/helper.h" +#include "cunls/common/log.h" +#include "cunls/minimizer/problem.h" +#include "cunls/minimizer/problem_partition.h" +#include "cunls/minimizer/schur_system.h" +#include "cunls/state/state_batch_ops.h" + +namespace cunls { +namespace { + +using internal::schur::BatchGeometry; +using internal::schur::ColumnRange; + +constexpr int kThreads = 256; +constexpr unsigned kFullMask = 0xffffffffu; +/// Pose-major pairs per CTA of the off-diagonal pass (long rows are split). +constexpr int kChunkPairs = 256; + +unsigned Blocks(size_t n) { return static_cast((n + kThreads - 1) / kThreads); } + +/** Bits needed for the values [0, n). */ +int BitsFor(long long n) { + int bits = 1; + while ((1LL << bits) < n) ++bits; + return bits; +} + +/** Index of (j, k), j <= k, in the packed upper triangle of a d x d matrix. */ +__host__ __device__ constexpr int Tri(int j, int k, int d) { + return j * d - j * (j - 1) / 2 + (k - j); +} + +/** The residual batch of global factor g: the last one starting at or before g. */ +__device__ __forceinline__ int FindBatch(const BatchGeometry *batches, int num_batches, int g) { + int lo = 0, hi = num_batches - 1; + while (lo < hi) { + const int mid = (lo + hi + 1) >> 1; + if (batches[mid].factor_begin <= g) { + lo = mid; + } else { + hi = mid - 1; + } + } + return lo; +} + +/** The column range holding free column `col`. */ +__device__ __forceinline__ const ColumnRange &FindRange(const ColumnRange *ranges, int num_ranges, + int col) { + int r = 0; + while (r + 1 < num_ranges && col >= ranges[r].end_col) ++r; + return ranges[r]; +} + +/** Landmark or retained block id of free column `col` in `range`. */ +__device__ __forceinline__ int IdOf(const ColumnRange &range, int col) { + return range.first_id + (col - range.first_col) / range.tangent; +} + +// --------------------------------------------------------------------------- +// CUB wrappers (setup only) +// --------------------------------------------------------------------------- + +template +void SortPairs(cudaStream_t stream, dvector &temp, const K *keys_in, K *keys_out, + const V *values_in, V *values_out, int n, int end_bit) { + if (n == 0) return; + size_t bytes = 0; + THROW_ON_CUDA_ERROR(cub::DeviceRadixSort::SortPairs(nullptr, bytes, keys_in, keys_out, values_in, + values_out, n, 0, end_bit, stream)); + if (temp.size() < bytes) temp.resize(bytes); + THROW_ON_CUDA_ERROR(cub::DeviceRadixSort::SortPairs( + temp.data(), bytes, keys_in, keys_out, values_in, values_out, n, 0, end_bit, stream)); +} + +template +void ExclusiveSum(cudaStream_t stream, dvector &temp, const T *in, T *out, int n) { + if (n == 0) return; + size_t bytes = 0; + THROW_ON_CUDA_ERROR(cub::DeviceScan::ExclusiveSum(nullptr, bytes, in, out, n, stream)); + if (temp.size() < bytes) temp.resize(bytes); + THROW_ON_CUDA_ERROR(cub::DeviceScan::ExclusiveSum(temp.data(), bytes, in, out, n, stream)); +} + +template +void InclusiveSum(cudaStream_t stream, dvector &temp, const T *in, T *out, int n) { + if (n == 0) return; + size_t bytes = 0; + THROW_ON_CUDA_ERROR(cub::DeviceScan::InclusiveSum(nullptr, bytes, in, out, n, stream)); + if (temp.size() < bytes) temp.resize(bytes); + THROW_ON_CUDA_ERROR(cub::DeviceScan::InclusiveSum(temp.data(), bytes, in, out, n, stream)); +} + +/** Reads one device value back (synchronizes the stream). */ +template +T ReadBack(cudaStream_t stream, const T *device_value) { + T value{}; + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&value, device_value, sizeof(T), cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + return value; +} + +/** Fills `out[i] = i`. */ +__global__ void IotaKernel(int n, int *out) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) out[i] = i; +} + +void Iota(cudaStream_t stream, dvector &out, int n) { + out.resize(n); + if (n == 0) return; + IotaKernel<<>>(n, out.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +// --------------------------------------------------------------------------- +// Symbolic phase +// --------------------------------------------------------------------------- + +/** Error flags of the setup checks (SchurSystem::scratch.error[0]). */ +enum SetupError : int { kTwoLandmarks = 1 }; + +/** + * Per factor: its landmark (the single free slot in an eliminated batch), + * the number of (landmark, retained slot) entries and of C^R items + * (retained slot pairs: a < b for landmark factors, a <= b otherwise). + * error[0] |= kTwoLandmarks for a factor with two free landmark slots; + * error[1], error[2]: min and max tangent size of retained slots of landmark + * factors. + */ +__global__ void ClassifyFactorsKernel(int num_factors, const BatchGeometry *batches, + int num_batches, const int *slot_sizes, + const int *factor_cols, const ColumnRange *ranges, + int num_ranges, int *factor_landmark, int *factor_eslot, + int *entry_count, int *rr_count, int *batch_has_landmarks, + int *error) { + const int g = blockIdx.x * blockDim.x + threadIdx.x; + if (g >= num_factors) return; + const int batch_index = FindBatch(batches, num_batches, g); + const BatchGeometry &batch = batches[batch_index]; + const int *cols = factor_cols + batch.col_offset + (g - batch.factor_begin) * batch.num_slots; + int num_landmarks = 0, eslot = -1, landmark = -1, num_retained = 0; + for (int s = 0; s < batch.num_slots; ++s) { + const int col = cols[s]; + if (col < 0) continue; + const ColumnRange &range = FindRange(ranges, num_ranges, col); + if (range.eliminated) { + ++num_landmarks; + eslot = s; + landmark = IdOf(range, col); + } else { + ++num_retained; + } + } + if (num_landmarks > 1) atomicOr(error, kTwoLandmarks); + const bool landmark_factor = num_landmarks == 1; + factor_landmark[g] = landmark_factor ? landmark : -1; + factor_eslot[g] = landmark_factor ? eslot : -1; + entry_count[g] = landmark_factor ? num_retained : 0; + rr_count[g] = landmark_factor ? num_retained * (num_retained - 1) / 2 + : num_retained * (num_retained + 1) / 2; + if (landmark_factor) { + if (batch_has_landmarks[batch_index] == 0) atomicExch(batch_has_landmarks + batch_index, 1); + for (int s = 0; s < batch.num_slots; ++s) { + if (s == eslot || cols[s] < 0) continue; + const int size = slot_sizes[batch.slot_begin + s]; + atomicMin(error + 1, size); + atomicMax(error + 2, size); + } + } +} + +/** + * Writes each factor's entries (key: landmark << block_bits | block; value: + * factor, slot) and C^R items (factor, packed slots, key of the S block). + */ +__global__ void EmitItemsKernel(int num_factors, const BatchGeometry *batches, int num_batches, + const int *factor_cols, const ColumnRange *ranges, int num_ranges, + const int *factor_landmark, const int *factor_eslot, + const int *entry_offsets, const int *rr_offsets, int block_bits, + unsigned long long *entry_keys, int2 *entry_values, int *rr_factor, + int *rr_slots, unsigned long long *rr_keys) { + const int g = blockIdx.x * blockDim.x + threadIdx.x; + if (g >= num_factors) return; + const BatchGeometry &batch = batches[FindBatch(batches, num_batches, g)]; + const int *cols = factor_cols + batch.col_offset + (g - batch.factor_begin) * batch.num_slots; + const int landmark = factor_landmark[g]; + const int eslot = factor_eslot[g]; + int entry = entry_offsets[g]; + int rr = rr_offsets[g]; + for (int a = 0; a < batch.num_slots; ++a) { + const int col_a = cols[a]; + if (col_a < 0 || a == eslot) continue; + const int block_a = IdOf(FindRange(ranges, num_ranges, col_a), col_a); + if (landmark >= 0) { + entry_keys[entry] = (static_cast(landmark) << block_bits) | + static_cast(block_a); + entry_values[entry] = make_int2(g, a); + ++entry; + } + for (int b = landmark >= 0 ? a + 1 : a; b < batch.num_slots; ++b) { + const int col_b = cols[b]; + if (col_b < 0 || b == eslot) continue; + const int block_b = IdOf(FindRange(ranges, num_ranges, col_b), col_b); + const int lo = min(block_a, block_b), hi = max(block_a, block_b); + rr_factor[rr] = g; + rr_slots[rr] = a | (b << 8) | ((block_a > block_b) << 16) | ((block_a == block_b) << 17); + rr_keys[rr] = (static_cast(lo) << block_bits) | static_cast(hi); + ++rr; + } + } +} + +/** Unpacks the sorted entries and flags the first entry of every pair. */ +__global__ void UnpackEntriesKernel(int num_entries, const unsigned long long *keys, + const int2 *values, int *entry_factor, int *entry_slot, + int *head) { + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= num_entries) return; + entry_factor[e] = values[e].x; + entry_slot[e] = values[e].y; + head[e] = (e == 0 || keys[e] != keys[e - 1]) ? 1 : 0; +} + +/** Writes the pairs from the heads of the sorted entries. */ +__global__ void BuildPairsKernel(int num_entries, const unsigned long long *keys, const int *head, + const int *pair_id_plus_one, int block_bits, int *pair_landmark, + int *pair_block, int *pair_entry_offsets) { + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= num_entries) return; + if (head[e]) { + const int pair = pair_id_plus_one[e] - 1; + pair_landmark[pair] = static_cast(keys[e] >> block_bits); + pair_block[pair] = static_cast(keys[e] & ((1ull << block_bits) - 1)); + pair_entry_offsets[pair] = e; + } + if (e == num_entries - 1) pair_entry_offsets[pair_id_plus_one[e]] = num_entries; +} + +/** landmark_pair_offsets from the landmark-sorted pairs (landmarks without pairs get empty ranges). + */ +__global__ void LandmarkOffsetsKernel(int num_pairs, int num_landmarks, const int *pair_landmark, + int *landmark_pair_offsets) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k > num_pairs) return; + const int prev = k == 0 ? -1 : pair_landmark[k - 1]; + const int cur = k == num_pairs ? num_landmarks : pair_landmark[k]; + for (int l = prev + 1; l <= cur; ++l) landmark_pair_offsets[l] = k; +} + +/** Scatters the inverse permutation. */ +__global__ void InvertPermutationKernel(int n, const int *perm, int *inverse) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) inverse[perm[i]] = i; +} + +/** Off-diagonal clique products k (k - 1) / 2 of each landmark (index L: 0). */ +__global__ void CliqueCountsKernel(int num_landmarks, const int *landmark_pair_offsets, + long long *counts) { + const int l = blockIdx.x * blockDim.x + threadIdx.x; + if (l > num_landmarks) return; + long long k = l < num_landmarks ? landmark_pair_offsets[l + 1] - landmark_pair_offsets[l] : 0; + counts[l] = k * (k - 1) / 2; +} + +/** Keys of the S blocks of every clique product (a < b) of the pair's landmark, pair a. */ +__global__ void CliqueKeysKernel(int num_pairs, const int *pair_landmark, const int *pair_block, + const int *landmark_pair_offsets, const long long *clique_offsets, + int block_bits, unsigned long long *keys, int *values) { + const int pair = blockIdx.x * blockDim.x + threadIdx.x; + if (pair >= num_pairs) return; + const int l = pair_landmark[pair]; + const int begin = landmark_pair_offsets[l]; + const long long k = landmark_pair_offsets[l + 1] - begin; + const long long a = pair - begin; + const unsigned long long p = static_cast(pair_block[pair]); + long long index = clique_offsets[l] + a * k - a * (a + 1) / 2; + for (long long b = a + 1; b < k; ++b, ++index) { + keys[index] = (p << block_bits) | static_cast(pair_block[begin + b]); + values[index] = static_cast(index); + } +} + +/** Keys of the diagonal blocks and of the C^R items, after the clique keys. */ +__global__ void OtherKeysKernel(int num_blocks, int num_rr, long long num_cliques, + const unsigned long long *rr_keys, int block_bits, + unsigned long long *keys, int *values) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_blocks + num_rr) return; + const long long index = num_cliques + i; + keys[index] = i < num_blocks + ? (static_cast(i) << block_bits) | static_cast(i) + : rr_keys[i - num_blocks]; + values[index] = static_cast(index); +} + +/** Head flags of the sorted keys (first of each distinct S block). */ +__global__ void HeadFlagsKernel(long long n, const unsigned long long *keys, int *head) { + const long long i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + head[i] = (i == 0 || keys[i] != keys[i - 1]) ? 1 : 0; +} + +/** + * Writes each distinct S block (row, col) and routes the block id of every + * diagonal and C^R key back to its source. + */ +__global__ void AssignBlocksKernel(long long n, const unsigned long long *keys, const int *values, + const int *head, const int *block_id_plus_one, int block_bits, + long long num_cliques, int num_blocks, int *s_row, int *s_col, + int *diag_block, int *rr_block) { + const long long i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (i >= n) return; + const int block = block_id_plus_one[i] - 1; + if (head[i]) { + s_row[block] = static_cast(keys[i] >> block_bits); + s_col[block] = static_cast(keys[i] & ((1ull << block_bits) - 1)); + } + const long long v = values[i]; + if (v < num_cliques) { + // Clique products find their block by search at run time. + } else if (v < num_cliques + num_blocks) { + diag_block[v - num_cliques] = block; + } else { + rr_block[v - num_cliques - num_blocks] = block; + } +} + +/** Value count of each S block, and the per-row upper and lower block counts. */ +__global__ void BlockSizesKernel(int num_s_blocks, const int *s_row, const int *s_col, + const int *block_dim, int *value_counts, int *upper_count, + int *lower_count) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= num_s_blocks) return; + const int p = s_row[j], q = s_col[j]; + value_counts[j] = block_dim[p] * block_dim[q]; + atomicAdd(upper_count + p, 1); + if (p != q) atomicAdd(lower_count + q, 1); +} + +/** Keys (col << bits | row) of the strictly upper blocks, to order the mirrored lower blocks. */ +__global__ void LowerKeysKernel(int num_s_blocks, const int *s_row, const int *s_col, + int block_bits, unsigned long long *keys, int *values, + int *num_lower) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= num_s_blocks) return; + const int p = s_row[j], q = s_col[j]; + if (p == q) return; + const int slot = atomicAdd(num_lower, 1); + keys[slot] = (static_cast(q) << block_bits) | static_cast(p); + values[slot] = j; +} + +/** Full block row offsets: the upper and the mirrored lower blocks of each row. */ +__global__ void FullRowOffsetsKernel(int num_blocks, const int *upper_row_offsets, + const int *lower_row_offsets, int *full_row_offsets) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p <= num_blocks) full_row_offsets[p] = upper_row_offsets[p] + lower_row_offsets[p]; +} + +/** + * Places the upper blocks and the mirrored lower blocks (sorted by row, then + * column) into the full block rows, lower part first, so columns ascend: + * upper block i of row p lands at i + lower_row_offsets[p + 1], lower block t + * of row q at t + upper_row_offsets[q]. + */ +__global__ void PlaceFullBlocksKernel(int num_s_blocks, int num_lower, const int *s_row, + const int *s_col, const unsigned long long *lower_keys, + const int *lower_values, const int *upper_row_offsets, + const int *lower_row_offsets, int block_bits, int *full_row, + int *full_col, int *full_src) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < num_s_blocks) { + const int p = s_row[i]; + const int pos = i + lower_row_offsets[p + 1]; + full_row[pos] = p; + full_col[pos] = s_col[i]; + full_src[pos] = i * 2; + } else if (i < num_s_blocks + num_lower) { + const int t = i - num_s_blocks; + const int q = static_cast(lower_keys[t] >> block_bits); + const int p = static_cast(lower_keys[t] & ((1ull << block_bits) - 1)); + const int pos = t + upper_row_offsets[q]; + full_row[pos] = q; + full_col[pos] = p; + full_src[pos] = lower_values[t] * 2 + 1; + } +} + +/** Width (columns) of each full block. */ +__global__ void FullWidthsKernel(int num_full, const int *full_col, const int *block_dim, + int *widths) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < num_full) widths[i] = block_dim[full_col[i]]; +} + +/** Scalar row length of each block row, and its nonzero count (rows x length). */ +__global__ void RowLengthsKernel(int num_blocks, const int *full_row_offsets, const int *width_pos, + const int *block_dim, int *row_length, int *row_nonzeros) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= num_blocks) return; + const int length = width_pos[full_row_offsets[p + 1]] - width_pos[full_row_offsets[p]]; + row_length[p] = length; + row_nonzeros[p] = length * block_dim[p]; +} + +/** CSR row offsets of S: rows of block p start at row_base[p] + r * row_length[p]. */ +__global__ void CsrRowOffsetsKernel(int num_blocks, const int *block_dim, const int *block_col, + const int *row_base, const int *row_length, int num_nonzeros, + int num_cols, int *row_offsets) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= num_blocks) return; + for (int r = 0; r < block_dim[p]; ++r) { + row_offsets[block_col[p] + r] = row_base[p] + r * row_length[p]; + } + if (p == num_blocks - 1) row_offsets[num_cols] = num_nonzeros; +} + +/** Column ids of S and the position in the upper block storage each nonzero copies. */ +__global__ void CsrFillKernel(int num_full, const int *full_row, const int *full_col, + const int *full_src, const int *full_row_offsets, + const int *width_pos, const int *block_dim, const int *block_col, + const int *row_base, const int *row_length, + const int *s_value_offsets, int *col_ids, int *csr_src) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_full) return; + const int p = full_row[i], q = full_col[i]; + const int dp = block_dim[p], dq = block_dim[q]; + const int col_in_row = width_pos[i] - width_pos[full_row_offsets[p]]; + const int src_block = full_src[i] >> 1; + const bool mirrored = full_src[i] & 1; + const int value0 = s_value_offsets[src_block]; + for (int r = 0; r < dp; ++r) { + const int row_start = row_base[p] + r * row_length[p] + col_in_row; + for (int c = 0; c < dq; ++c) { + col_ids[row_start + c] = block_col[q] + c; + int src; + if (p == q) { + src = value0 + min(r, c) * dp + max(r, c); // diagonal blocks: upper triangle only + } else if (mirrored) { + src = value0 + c * dp + r; // stored as block (q, p), dq x dp + } else { + src = value0 + r * dq + c; + } + csr_src[row_start + c] = src; + } + } +} + +/** Retained blocks of one state batch: dims, columns of S and of the full system. */ +__global__ void RetainedBlocksKernel(int num_states, int tangent, int first_block, + int first_natural_col, int first_reduced_col, int *block_dim, + int *block_col, int *block_natural_col) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_states) return; + block_dim[first_block + i] = tangent; + block_col[first_block + i] = first_reduced_col + i * tangent; + block_natural_col[first_block + i] = first_natural_col + i * tangent; +} + +/** Landmarks of one eliminated state batch: their first free column. */ +__global__ void LandmarkColumnsKernel(int num_states, int tangent, int first_landmark, + int first_natural_col, int *landmark_natural_col) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < num_states) landmark_natural_col[first_landmark + i] = first_natural_col + i * tangent; +} + +/** Per column of S: its free column in the full system and its diagonal in the S values. */ +__global__ void ColumnMapsKernel(int num_blocks, const int *block_dim, const int *block_col, + const int *block_natural_col, const int *diag_block, + const int *s_value_offsets, int *col_natural, + int *col_diag_value) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p >= num_blocks) return; + const int d = block_dim[p]; + const int value0 = s_value_offsets[diag_block[p]]; + for (int r = 0; r < d; ++r) { + col_natural[block_col[p] + r] = block_natural_col[p] + r; + col_diag_value[block_col[p] + r] = value0 + r * d + r; + } +} + +/** Per pose-major position: the pair's entry count. */ +__global__ void PoseMajorPairsKernel(int num_pairs, const int *pm_pair, + const int *pair_entry_offsets, int *pm_entry_count) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= num_pairs) return; + const int pair = pm_pair[t]; + pm_entry_count[t] = pair_entry_offsets[pair + 1] - pair_entry_offsets[pair]; +} + +/** + * Copies the entries of each pair to pose-major order, and gives each + * landmark factor the first pose-major position of its pairs (its key in the + * evaluation order). + */ +__global__ void PoseMajorEntriesKernel(int num_pairs, const int *pm_pair, + const int *pair_entry_offsets, const int *pm_entry_offsets, + const int *entry_factor, const int *entry_slot, + int *pm_entry_factor, int *pm_entry_slot, int *factor_key) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= num_pairs) return; + const int pair = pm_pair[t]; + int out = pm_entry_offsets[t]; + for (int e = pair_entry_offsets[pair]; e < pair_entry_offsets[pair + 1]; ++e, ++out) { + pm_entry_factor[out] = entry_factor[e]; + pm_entry_slot[out] = entry_slot[e]; + atomicMin(factor_key + entry_factor[e], t); + } +} + +/** + * Sort keys of the evaluation order: batch-major, then the factor's key + * (landmark factors by their first pose-major pair, the others after them), + * or the natural order for a batch that is not reordered. + */ +__global__ void OrderKeysKernel(int num_factors, const BatchGeometry *batches, int num_batches, + const int *reordered, const int *factor_key, + unsigned long long *keys, int *values) { + const int g = blockIdx.x * blockDim.x + threadIdx.x; + if (g >= num_factors) return; + const int b = FindBatch(batches, num_batches, g); + const unsigned key = reordered[b] ? static_cast(factor_key[g]) + : static_cast(g - batches[b].factor_begin); + keys[g] = (static_cast(b) << 32) | key; + values[g] = g; +} + +/** + * The evaluation order (batch-local factor of each item), each factor's item, + * and the slot columns of each item (FactorCols() in item order). + */ +__global__ void ItemsKernel(int num_factors, const BatchGeometry *batches, int num_batches, + const int *sorted_factors, const int *factor_cols, int *factor_item, + int *order, int *item_cols) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_factors) return; + const BatchGeometry &batch = batches[FindBatch(batches, num_batches, i)]; + const int g = sorted_factors[i]; + const int item = i - batch.factor_begin, local = g - batch.factor_begin; + factor_item[g] = item; + order[i] = local; + for (int s = 0; s < batch.num_slots; ++s) { + item_cols[batch.col_offset + item * batch.num_slots + s] = + factor_cols[batch.col_offset + local * batch.num_slots + s]; + } +} + +/** + * Where the pose-major entries' Jacobians and residuals are: the factor's + * item start, and packed (retained slot column, landmark slot column, m, n). + */ +__global__ void EntryLocationsKernel(int num_entries, const int *pm_entry_factor, + const int *pm_entry_slot, const BatchGeometry *batches, + int num_batches, const int *slot_offsets, + const int *factor_eslot, const int *factor_item, + long long *jacobian_base, long long *residual_base, + int *meta) { + const int e = blockIdx.x * blockDim.x + threadIdx.x; + if (e >= num_entries) return; + const int g = pm_entry_factor[e]; + const BatchGeometry &batch = batches[FindBatch(batches, num_batches, g)]; + const long long item = factor_item[g]; + const int m = batch.num_residuals, n = batch.tangent_dim; + jacobian_base[e] = batch.jacobian_offset + item * m * n; + residual_base[e] = batch.residual_offset + item * m; + meta[e] = slot_offsets[batch.slot_begin + pm_entry_slot[e]] | + (slot_offsets[batch.slot_begin + factor_eslot[g]] << 8) | (m << 16) | (n << 24); +} + +// --------------------------------------------------------------------------- +// Numeric phase +// --------------------------------------------------------------------------- + +/** Column scale, mask and damping (null: identity / none) of the free columns. */ +struct Modifiers { + const float *scale; + const float *mask; + const float *damping; + __device__ __forceinline__ float Scale(int col) const { return scale ? scale[col] : 1.f; } + __device__ __forceinline__ float Mask(int col) const { return mask ? mask[col] : 1.f; } + __device__ __forceinline__ float Damping(int col) const { return damping ? damping[col] : 0.f; } +}; + +/** + * Segments of equal keys over the lanes of a warp (consecutive lanes with the + * same key; give inactive lanes distinct negative keys). All 32 lanes must + * call it and Sum() together. + */ +struct WarpSegments { + int start; ///< First lane of this lane's segment. + bool tail; ///< This lane is the last of its segment (holds the sum). + + __device__ __forceinline__ explicit WarpSegments(int key) { + const int lane = threadIdx.x & 31; + const int prev = __shfl_up_sync(kFullMask, key, 1); + const int next = __shfl_down_sync(kFullMask, key, 1); + const unsigned heads = __ballot_sync(kFullMask, lane == 0 || prev != key); + start = 31 - __clz(heads & (lane == 31 ? kFullMask : ((2u << lane) - 1u))); + tail = lane == 31 || next != key; + } + + /** Inclusive segmented sum: the tail lane gets the sum of its segment. */ + __device__ __forceinline__ float Sum(float v) const { + const int lane = threadIdx.x & 31; +#pragma unroll + for (int d = 1; d < 32; d <<= 1) { + const float u = __shfl_up_sync(kFullMask, v, d); + if (lane - d >= start) v += u; + } + return v; + } +}; + +/** + * One thread per item (evaluation order), for the terms without a retained-state pair: the + * landmark terms of a landmark factor (b_e, diag, D_i = J_eᵀJ_e, AoS upper + * triangle; few factors per landmark, so plain atomics), and every term of a + * factor between retained states only. The retained terms of landmark + * factors are AssemblePairsKernel's. + */ +__global__ void AssembleFactorsKernel(int num_factors, const BatchGeometry *batches, + int num_batches, const int *slot_sizes, + const int *slot_offsets, const int *item_cols, + const int *factor_landmark, const int *factor_eslot, + const int *evaluation_order, int landmark_dim, + const float *__restrict__ jacobians, + const float *__restrict__ residuals, float *rhs, + float *diagonal, float *landmark_hessian) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; // item, so J and r stream + if (i >= num_factors) return; + const BatchGeometry &batch = batches[FindBatch(batches, num_batches, i)]; + const long long item = i - batch.factor_begin; + const int local = evaluation_order[i]; + const int g = batch.factor_begin + local; + const int m = batch.num_residuals, n = batch.tangent_dim; + const float *J = jacobians + batch.jacobian_offset + item * m * n; + const float *r = residuals + batch.residual_offset + item * m; + const int *cols = item_cols + batch.col_offset + item * batch.num_slots; + const int landmark = factor_landmark[g]; + for (int s = 0; s < batch.num_slots; ++s) { + const int col = cols[s]; + if (col < 0 || (landmark >= 0 && s != factor_eslot[g])) continue; + const int offset = slot_offsets[batch.slot_begin + s]; + for (int c = 0; c < slot_sizes[batch.slot_begin + s]; ++c) { + float gradient = 0.f, norm = 0.f; + for (int row = 0; row < m; ++row) { + const float j = J[row * n + offset + c]; + gradient += j * r[row]; + norm += j * j; + } + atomicAdd(rhs + col + c, -gradient); + atomicAdd(diagonal + col + c, norm); + } + } + if (landmark < 0) return; + const int offset = slot_offsets[batch.slot_begin + factor_eslot[g]]; + const int tri = landmark_dim * (landmark_dim + 1) / 2; + for (int j = 0; j < landmark_dim; ++j) { + for (int k = j; k < landmark_dim; ++k) { + float sum = 0.f; + for (int row = 0; row < m; ++row) sum += J[row * n + offset + j] * J[row * n + offset + k]; + atomicAdd(landmark_hessian + landmark * tri + Tri(j, k, landmark_dim), sum); + } + } +} + +/** + * One thread per pair, pose-major: the retained terms of the pair's landmark + * factors (b_p = -Σ J_pᵀ r, diag J_pᵀJ_p), summed over each retained state's + * run of pairs with WarpSegments, one atomic per state and warp. + */ +template +__global__ void __launch_bounds__(kThreads) + AssemblePairsKernel(int num_pairs, const int *__restrict__ pm_entry_offsets, + const long long *__restrict__ entry_jacobian, + const long long *__restrict__ entry_residual, + const int *__restrict__ entry_meta, const int *__restrict__ pm_pair, + const int *__restrict__ pair_block, const float *__restrict__ jacobians, + const float *__restrict__ residuals, + const int *__restrict__ block_natural_col, float *rhs, float *diagonal) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + float gradient[DP], norm[DP]; +#pragma unroll + for (int c = 0; c < DP; ++c) gradient[c] = norm[c] = 0.f; + int p = -1 - static_cast(threadIdx.x & 31); + if (t < num_pairs) { + p = pair_block[pm_pair[t]]; + for (int e = pm_entry_offsets[t]; e < pm_entry_offsets[t + 1]; ++e) { + const int meta = entry_meta[e]; + const int m = (meta >> 16) & 0xff, n = meta >> 24; + const float *Jp = jacobians + entry_jacobian[e] + (meta & 0xff); + const float *r = residuals + entry_residual[e]; + for (int row = 0; row < m; ++row) { + const float rr = r[row]; +#pragma unroll + for (int c = 0; c < DP; ++c) { + const float j = Jp[row * n + c]; + gradient[c] += j * rr; + norm[c] += j * j; + } + } + } + } + const WarpSegments segments(p); +#pragma unroll + for (int c = 0; c < DP; ++c) { + gradient[c] = segments.Sum(gradient[c]); + norm[c] = segments.Sum(norm[c]); + } + if (segments.tail && p >= 0) { + const int col = block_natural_col[p]; +#pragma unroll + for (int c = 0; c < DP; ++c) { + atomicAdd(rhs + col + c, -gradient[c]); + atomicAdd(diagonal + col + c, norm[c]); + } + } +} + +/** g = the retained part of the (scaled, masked) right-hand side. */ +__global__ void GatherRhsKernel(int num_cols, const int *col_natural, const float *rhs, + float *reduced_rhs) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c < num_cols) reduced_rhs[c] = rhs[col_natural[c]]; +} + +/** + * One thread per C^R item: J_aᵀ J_b of its factor, scaled and masked, into + * its block of S (transposed if the block is stored the other way; for two + * slots on one state, as the symmetric sum; diagonal blocks keep only their + * upper triangle). + */ +__global__ void RetainedProductsKernel(int num_rr, const int *rr_factor, const int *rr_slots, + const int *rr_block, const BatchGeometry *batches, + int num_batches, const int *slot_sizes, + const int *slot_offsets, const int *factor_cols, + const int *factor_item, const float *__restrict__ jacobians, + const int *s_value_offsets, Modifiers mod, float *s_values) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_rr) return; + const int g = rr_factor[i]; + const BatchGeometry &batch = batches[FindBatch(batches, num_batches, g)]; + const int local = g - batch.factor_begin; + const int m = batch.num_residuals, n = batch.tangent_dim; + const float *J = + jacobians + batch.jacobian_offset + static_cast(factor_item[g]) * m * n; + const int packed = rr_slots[i]; + const int a = packed & 0xff, b = (packed >> 8) & 0xff; + const bool transposed = (packed >> 16) & 1, same_state = (packed >> 17) & 1; + const int da = slot_sizes[batch.slot_begin + a], db = slot_sizes[batch.slot_begin + b]; + const int oa = slot_offsets[batch.slot_begin + a], ob = slot_offsets[batch.slot_begin + b]; + const int *cols = factor_cols + batch.col_offset + local * batch.num_slots; + const int col_a = cols[a], col_b = cols[b]; + float *block = s_values + s_value_offsets[rr_block[i]]; + for (int r = 0; r < da; ++r) { + const float wr = mod.Scale(col_a + r) * mod.Mask(col_a + r); + for (int c = 0; c < db; ++c) { + if (a == b && c < r) continue; + float sum = 0.f; + for (int row = 0; row < m; ++row) sum += J[row * n + oa + r] * J[row * n + ob + c]; + const float v = sum * wr * mod.Scale(col_b + c) * mod.Mask(col_b + c); + if (a == b) { + atomicAdd(block + r * da + c, v); + } else if (same_state) { + atomicAdd(block + min(r, c) * da + max(r, c), r == c ? 2.f * v : v); + } else if (transposed) { + atomicAdd(block + c * da + r, v); + } else { + atomicAdd(block + r * db + c, v); + } + } + } +} + +/** + * One stored row of V (DE floats), padded for one vector load or store: a + * landmark of size 3 takes a float4. + */ +template +struct VRow; +template <> +struct VRow<1> { + static constexpr int kFloats = 1; + __device__ static void Load(const float *p, float (&v)[1]) { v[0] = *p; } + __device__ static void Store(float *p, const float (&v)[1]) { *p = v[0]; } +}; +template <> +struct VRow<2> { + static constexpr int kFloats = 2; + __device__ static void Load(const float *p, float (&v)[2]) { + const float2 x = *reinterpret_cast(p); + v[0] = x.x; + v[1] = x.y; + } + __device__ static void Store(float *p, const float (&v)[2]) { + *reinterpret_cast(p) = make_float2(v[0], v[1]); + } +}; +template <> +struct VRow<3> { + static constexpr int kFloats = 4; + __device__ static void Load(const float *p, float (&v)[3]) { + const float4 x = *reinterpret_cast(p); + v[0] = x.x; + v[1] = x.y; + v[2] = x.z; + } + __device__ static void Store(float *p, const float (&v)[3]) { + *reinterpret_cast(p) = make_float4(v[0], v[1], v[2], 0.f); + } +}; + +/** In-place Cholesky of the DE x DE matrix A (lower triangle used); false if not positive. */ +template +__device__ __forceinline__ bool Cholesky(float (&A)[DE][DE]) { + bool ok = true; +#pragma unroll + for (int j = 0; j < DE; ++j) { + float d = A[j][j]; +#pragma unroll + for (int k = 0; k < j; ++k) d -= A[j][k] * A[j][k]; + ok = ok && d > 0.f; + const float l = sqrtf(fmaxf(d, 1e-30f)); + A[j][j] = l; + const float inv = 1.f / l; +#pragma unroll + for (int i = j + 1; i < DE; ++i) { + float v = A[i][j]; +#pragma unroll + for (int k = 0; k < j; ++k) v -= A[i][k] * A[j][k]; + A[i][j] = v * inv; + } + } + return ok; +} + +/** Solves L x = b in place (L lower, from Cholesky). */ +template +__device__ __forceinline__ void ForwardSolve(const float (&L)[DE][DE], float (&x)[DE]) { +#pragma unroll + for (int j = 0; j < DE; ++j) { + float v = x[j]; +#pragma unroll + for (int k = 0; k < j; ++k) v -= L[j][k] * x[k]; + x[j] = v / L[j][j]; + } +} + +/** Solves Lᵀ x = b in place. */ +template +__device__ __forceinline__ void BackwardSolve(const float (&L)[DE][DE], float (&x)[DE]) { +#pragma unroll + for (int j = DE - 1; j >= 0; --j) { + float v = x[j]; +#pragma unroll + for (int k = j + 1; k < DE; ++k) v -= L[k][j] * x[k]; + x[j] = v / L[j][j]; + } +} + +/** + * The working landmark block of landmark l: s_e D_i s_e, masked (unit + * diagonal), plus the damping, factored in place. Its free columns start at + * `col`. A block that is not positive definite (a landmark seen along one + * ray, or with a zero Jacobian, e.g. behind every camera, which no damping of + * the form λ diag(D_i) fixes) is factored again with a small diagonal shift, + * relative to its largest diagonal entry: in the spirit of the pseudo-inverse + * Ceres uses for rank-deficient blocks, the unobservable directions then get + * (almost) no step instead of failing the solve. Returns false only if that + * fails too (non-finite values). + */ +template +__device__ __forceinline__ bool LandmarkFactor(const float *__restrict__ landmark_hessian, + int num_landmarks, int l, int col, + const Modifiers &mod, float (&L)[DE][DE]) { + constexpr float kRelativeShift = 1e-4f; + constexpr float kMinShift = 1e-20f; + float D[DE][DE], w[DE]; +#pragma unroll + for (int j = 0; j < DE; ++j) w[j] = mod.Scale(col + j) * mod.Mask(col + j); + float largest = 0.f; +#pragma unroll + for (int j = 0; j < DE; ++j) { +#pragma unroll + for (int k = j; k < DE; ++k) { + const float v = landmark_hessian[l * (DE * (DE + 1) / 2) + Tri(j, k, DE)] * w[j] * w[k]; + D[j][k] = v; + D[k][j] = v; + } + D[j][j] = mod.Mask(col + j) == 0.f ? 1.f : D[j][j] + mod.Damping(col + j); + largest = fmaxf(largest, D[j][j]); + } +#pragma unroll + for (int j = 0; j < DE; ++j) { +#pragma unroll + for (int k = 0; k < DE; ++k) L[j][k] = D[j][k]; + } + if (Cholesky(L)) return true; + const float shift = fmaxf(kRelativeShift * largest, kMinShift); +#pragma unroll + for (int j = 0; j < DE; ++j) { +#pragma unroll + for (int k = 0; k < DE; ++k) L[j][k] = D[j][k] + (j == k ? shift : 0.f); + } + return Cholesky(L); +} + +/** + * One thread per pair, in pose-major order: contracts the pair's factors into + * W = J_pᵀJ_e and A = J_pᵀJ_p (scaled, masked), eliminates the landmark + * (V = W L⁻ᵀ, h = L⁻¹ b_e), stores V, and adds the diagonal bracket + * A - V Vᵀ and -V h of all pairs of one retained state with a + * warp-segmented sum, one atomic per state and warp. + */ +template +__global__ void __launch_bounds__(kThreads) + EliminatePairsKernel(int num_pairs, const int *__restrict__ pm_entry_offsets, + const long long *__restrict__ entry_jacobian, + const int *__restrict__ entry_meta, const int *__restrict__ pm_pair, + const int *__restrict__ pair_landmark, const int *__restrict__ pair_block, + const float *__restrict__ jacobians, + const float *__restrict__ landmark_hessian, int num_landmarks, + const int *__restrict__ landmark_natural_col, + const int *__restrict__ block_natural_col, + const int *__restrict__ block_col, const int *__restrict__ diag_block, + const int *__restrict__ s_value_offsets, const float *__restrict__ rhs, + Modifiers mod, float *__restrict__ pair_v, float *s_values, + float *reduced_rhs, int *failure) { + constexpr int kTri = DP * (DP + 1) / 2; + const int t = blockIdx.x * blockDim.x + threadIdx.x; + const int lane = threadIdx.x & 31; + float bracket[kTri]; + float gp[DP]; +#pragma unroll + for (int i = 0; i < kTri; ++i) bracket[i] = 0.f; +#pragma unroll + for (int i = 0; i < DP; ++i) gp[i] = 0.f; + int p = -1 - lane; // inactive lanes: distinct, negative keys (WarpSegments) + if (t < num_pairs) { + const int pair = pm_pair[t]; + p = pair_block[pair]; + const int l = pair_landmark[pair]; + float W[DP][DE]; +#pragma unroll + for (int r = 0; r < DP; ++r) { +#pragma unroll + for (int j = 0; j < DE; ++j) W[r][j] = 0.f; + } + for (int e = pm_entry_offsets[t]; e < pm_entry_offsets[t + 1]; ++e) { + const int meta = entry_meta[e]; + const int m = (meta >> 16) & 0xff, n = meta >> 24; + const float *Jp = jacobians + entry_jacobian[e] + (meta & 0xff); + const float *Je = jacobians + entry_jacobian[e] + ((meta >> 8) & 0xff); + for (int row = 0; row < m; ++row) { + float jp[DP], je[DE]; +#pragma unroll + for (int c = 0; c < DP; ++c) jp[c] = Jp[row * n + c]; +#pragma unroll + for (int j = 0; j < DE; ++j) je[j] = Je[row * n + j]; +#pragma unroll + for (int r = 0; r < DP; ++r) { +#pragma unroll + for (int j = 0; j < DE; ++j) W[r][j] += jp[r] * je[j]; +#pragma unroll + for (int c = r; c < DP; ++c) bracket[Tri(r, c, DP)] += jp[r] * jp[c]; + } + } + } + const int pcol = block_natural_col[p]; + const int ecol = landmark_natural_col[l]; + float wp[DP], we[DE]; +#pragma unroll + for (int r = 0; r < DP; ++r) wp[r] = mod.Scale(pcol + r) * mod.Mask(pcol + r); +#pragma unroll + for (int j = 0; j < DE; ++j) we[j] = mod.Scale(ecol + j) * mod.Mask(ecol + j); + float L[DE][DE]; + const bool factored = LandmarkFactor(landmark_hessian, num_landmarks, l, ecol, mod, L); + float h[DE]; +#pragma unroll + for (int j = 0; j < DE; ++j) h[j] = rhs[ecol + j]; + ForwardSolve(L, h); + // A landmark block that cannot be factored even with the shift of + // LandmarkFactor (non-finite values) poisons g, so the whole step is NaN + // and rejected. + if (!factored) { + h[0] = CUDART_NAN_F; + *failure = 1; + } + // V = W L⁻ᵀ, row by row (L v = w for each row w of W). + float V[DP][DE]; +#pragma unroll + for (int r = 0; r < DP; ++r) { +#pragma unroll + for (int j = 0; j < DE; ++j) V[r][j] = W[r][j] * wp[r] * we[j]; + ForwardSolve(L, V[r]); + VRow::Store(pair_v + (static_cast(t) * DP + r) * VRow::kFloats, V[r]); + } +#pragma unroll + for (int r = 0; r < DP; ++r) { +#pragma unroll + for (int c = r; c < DP; ++c) { + float vv = 0.f; +#pragma unroll + for (int j = 0; j < DE; ++j) vv += V[r][j] * V[c][j]; + bracket[Tri(r, c, DP)] = bracket[Tri(r, c, DP)] * wp[r] * wp[c] - vv; + } + float vh = 0.f; +#pragma unroll + for (int j = 0; j < DE; ++j) vh += V[r][j] * h[j]; + gp[r] = -vh; + } + } + // Sum over runs of equal p (pose-major: a run is one retained state). + const WarpSegments segments(p); +#pragma unroll + for (int i = 0; i < kTri; ++i) bracket[i] = segments.Sum(bracket[i]); +#pragma unroll + for (int i = 0; i < DP; ++i) gp[i] = segments.Sum(gp[i]); + if (segments.tail && p >= 0) { + float *block = s_values + s_value_offsets[diag_block[p]]; +#pragma unroll + for (int r = 0; r < DP; ++r) { +#pragma unroll + for (int c = r; c < DP; ++c) atomicAdd(block + r * DP + c, bracket[Tri(r, c, DP)]); + atomicAdd(reduced_rhs + block_col[p] + r, gp[r]); + } + } +} + +/** + * The off-diagonal blocks of S, by block row: one CTA per chunk of the + * pose-major pairs (p, i) of retained state p, one warp per pair at a time. + * The warps of a CTA share row p's blocks and V_pi's landmarks in cache. + * For each later pair (q, i) of landmark i (q > p: the landmark's pairs are + * sorted by state), -V_pi V_qiᵀ is added to block (p, q). The products are + * taken 32 at a time: each lane finds one product's block (binary search in + * the row) and V_qi's position; then the lanes sweep the products' entries, + * so each warp-wide atomic (a native global reduction) covers contiguous + * values of one block. An entry costs two vector loads (rows of V_pi in shared + * memory and of V_qi) and one reduction. Chunks bound the work per CTA however + * long the tracks are. + */ +template +__global__ void __launch_bounds__(kThreads) + OffDiagonalRowsKernel(const int3 *__restrict__ chunks, + const int *__restrict__ upper_row_offsets, const int *__restrict__ s_col, + const int *__restrict__ s_value_offsets, const int *__restrict__ pm_pair, + const int *__restrict__ pair_landmark, const int *__restrict__ pair_block, + const int *__restrict__ pair_pm, + const int *__restrict__ landmark_pair_offsets, + const float *__restrict__ pair_v, float *s_values) { + using Row = VRow; + constexpr int kVFloats = DP * Row::kFloats; + constexpr int kBlock = DP * DP; + constexpr int kWarps = kThreads / 32; + // Per product: first value of its block of S, pose-major position of (q, i). + __shared__ int2 product[kWarps][32]; + __shared__ __align__(16) float vp[kWarps][kVFloats]; // V_pi + const int3 chunk = chunks[blockIdx.x]; // retained state, pose-major pairs [begin, end) + const int p = chunk.x; + const int first_block = upper_row_offsets[p], end_block = upper_row_offsets[p + 1]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + for (int t = chunk.y + warp; t < chunk.z; t += kWarps) { + const int pair = pm_pair[t]; + const int later = landmark_pair_offsets[pair_landmark[pair] + 1] - pair - 1; + if (later <= 0) continue; + __syncwarp(); + for (int k = lane; k < kVFloats; k += 32) + vp[warp][k] = pair_v[static_cast(t) * kVFloats + k]; + for (int base = 0; base < later; base += 32) { + const int count = min(32, later - base); + __syncwarp(); + if (lane < count) { + const int other = pair + 1 + base + lane; + const int q = pair_block[other]; + int lo = first_block + 1, hi = end_block - 1; // the row's first block is (p, p) + while (lo < hi) { + const int mid = (lo + hi) >> 1; + if (s_col[mid] < q) { + lo = mid + 1; + } else { + hi = mid; + } + } + product[warp][lane] = make_int2(s_value_offsets[lo], pair_pm[other]); + } + __syncwarp(); + // kUnroll entries per lane per step, their loads issued together: the + // V_qi gathers are the latency this loop waits on. + constexpr int kUnroll = 4; + const int total = count * kBlock; + for (int idx0 = lane; idx0 < total; idx0 += 32 * kUnroll) { + float a[kUnroll][DE], b[kUnroll][DE]; + int target[kUnroll]; +#pragma unroll + for (int u = 0; u < kUnroll; ++u) { + const int idx = min(idx0 + 32 * u, total - 1); + const int k = idx / kBlock; + const int e = idx - k * kBlock; + const int r = e / DP, c = e - r * DP; + const int2 info = product[warp][k]; + Row::Load(&vp[warp][r * Row::kFloats], a[u]); + Row::Load(pair_v + (static_cast(info.y) * DP + c) * Row::kFloats, b[u]); + target[u] = info.x + e; + } +#pragma unroll + for (int u = 0; u < kUnroll; ++u) { + if (idx0 + 32 * u >= total) break; + float v = 0.f; +#pragma unroll + for (int j = 0; j < DE; ++j) v += a[u][j] * b[u][j]; + atomicAdd(s_values + target[u], -v); + } + } + } + } +} + +/** Damping (or a unit diagonal for a masked column) on the diagonal of S. */ +__global__ void FinalizeDiagonalKernel(int num_cols, const int *col_natural, + const int *col_diag_value, Modifiers mod, float *s_values) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c >= num_cols) return; + const int col = col_natural[c]; + float &d = s_values[col_diag_value[c]]; + d = mod.Mask(col) == 0.f ? 1.f : d + mod.Damping(col); +} + +/** Expands the upper block storage of S to its full scalar CSR. */ +__global__ void ExpandKernel(int num_nonzeros, const int *__restrict__ csr_src, + const float *__restrict__ s_values, float *__restrict__ values) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < num_nonzeros) values[i] = s_values[csr_src[i]]; +} + +/** step at the retained columns = x_r. */ +__global__ void ScatterRetainedKernel(int num_cols, const int *col_natural, + const float *reduced_step, float *step) { + const int c = blockIdx.x * blockDim.x + threadIdx.x; + if (c < num_cols) step[col_natural[c]] = reduced_step[c]; +} + +/** One thread per landmark: x_e = L⁻ᵀ (L⁻¹ b_e - Σ_p V_piᵀ x_p). */ +template +__global__ void BackSubstituteKernel( + int num_landmarks, int num_pairs, const int *__restrict__ landmark_pair_offsets, + const int *__restrict__ pair_block, const int *__restrict__ pair_pm, + const float *__restrict__ landmark_hessian, const int *__restrict__ landmark_natural_col, + const int *__restrict__ block_col, const float *__restrict__ pair_v, + const float *__restrict__ reduced_step, const float *__restrict__ rhs, Modifiers mod, + float *__restrict__ step, int *failure) { + const int l = blockIdx.x * blockDim.x + threadIdx.x; + if (l >= num_landmarks) return; + const int ecol = landmark_natural_col[l]; + float L[DE][DE]; + const bool factored = LandmarkFactor(landmark_hessian, num_landmarks, l, ecol, mod, L); + float u[DE]; +#pragma unroll + for (int j = 0; j < DE; ++j) u[j] = factored ? rhs[ecol + j] : CUDART_NAN_F; + if (!factored) *failure = 1; + ForwardSolve(L, u); + for (int pair = landmark_pair_offsets[l]; pair < landmark_pair_offsets[l + 1]; ++pair) { + const long long pm = pair_pm[pair]; + const float *x = reduced_step + block_col[pair_block[pair]]; +#pragma unroll + for (int r = 0; r < DP; ++r) { + const float xr = x[r]; + float v[DE]; + VRow::Load(pair_v + (pm * DP + r) * VRow::kFloats, v); +#pragma unroll + for (int j = 0; j < DE; ++j) u[j] -= v[j] * xr; + } + } + BackwardSolve(L, u); +#pragma unroll + for (int j = 0; j < DE; ++j) step[ecol + j] = u[j]; +} + +/** One thread per item: ‖J_f δ_f‖² of its factor f, written at f. */ +__global__ void FactorCurvatureKernel(int num_factors, const BatchGeometry *batches, + int num_batches, const int *slot_sizes, + const int *slot_offsets, const int *item_cols, + const int *evaluation_order, + const float *__restrict__ jacobians, + const float *__restrict__ step, float *curvature) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; // item, so J streams + if (i >= num_factors) return; + const BatchGeometry &batch = batches[FindBatch(batches, num_batches, i)]; + const int item = i - batch.factor_begin; + const int g = batch.factor_begin + evaluation_order[i]; + const int m = batch.num_residuals, n = batch.tangent_dim; + const float *J = jacobians + batch.jacobian_offset + static_cast(item) * m * n; + const int *cols = item_cols + batch.col_offset + item * batch.num_slots; + float sum = 0.f; + for (int row = 0; row < m; ++row) { + float v = 0.f; + for (int s = 0; s < batch.num_slots; ++s) { + const int col = cols[s]; + if (col < 0) continue; + const int offset = slot_offsets[batch.slot_begin + s]; + for (int c = 0; c < slot_sizes[batch.slot_begin + s]; ++c) { + v += J[row * n + offset + c] * step[col + c]; + } + } + sum += v * v; + } + curvature[g] = sum; +} + +/** + * Runs the statement(s) with the constants DE and DP bound to the supported (landmark, + * retained) tangent sizes (see SupportedDims). + */ +#define SCHUR_DISPATCH_DIMS(de, dp, ...) \ + do { \ + if ((de) == 3 && (dp) == 6) { \ + constexpr int DE = 3, DP = 6; \ + static_cast(DE); \ + static_cast(DP); \ + __VA_ARGS__; \ + } else if ((de) == 1 && (dp) == 6) { \ + constexpr int DE = 1, DP = 6; \ + static_cast(DE); \ + static_cast(DP); \ + __VA_ARGS__; \ + } else if ((de) == 2 && (dp) == 6) { \ + constexpr int DE = 2, DP = 6; \ + static_cast(DE); \ + static_cast(DP); \ + __VA_ARGS__; \ + } else if ((de) == 3 && (dp) == 9) { \ + constexpr int DE = 3, DP = 9; \ + static_cast(DE); \ + static_cast(DP); \ + __VA_ARGS__; \ + } else if ((de) == 3 && (dp) == 3) { \ + constexpr int DE = 3, DP = 3; \ + static_cast(DE); \ + static_cast(DP); \ + __VA_ARGS__; \ + } \ + } while (0) + +/** Whether (de, dp) is a supported combination (see DispatchDims). */ +bool SupportedDims(int de, int dp) { + return (dp == 6 && de >= 1 && de <= 3) || (de == 3 && (dp == 9 || dp == 3)); +} + +} // namespace + +SchurSystem::SchurSystem(SparseLinearSolverPtr solver, bool check_failures) + : solver_(std::move(solver)), check_failures_(check_failures), failure_(0, 1) {} + +/** Setup temporaries passed between the phases of SchurSystem::Initialize(). */ +struct SchurSystem::SetupScratch { + dvector temp; ///< CUB temporary storage. + int block_bits = 0; ///< Key bits of a retained state. + int landmark_bits = 0; ///< Key bits of a landmark. + dvector entry_offsets; ///< Per factor: its first entry. + dvector rr_offsets; ///< Per factor: its first C^R item. + dvector batch_has_landmarks; ///< Per residual batch: has landmark factors. + dvector rr_keys; ///< Per C^R item: the key of its block of S. + dvector entry_factor; ///< Per entry (landmark-major): its factor. + dvector entry_slot; ///< Per entry (landmark-major): its retained slot. + dvector pair_entry_offsets; ///< Per pair: its entries [offset, next). + // The full (both triangles) block rows of S. + dvector full_row_offsets; ///< Per retained state: its full blocks [offset, next). + dvector full_row; ///< Per full block: its block row. + dvector full_col; ///< Per full block: its block column. + dvector full_src; ///< Per full block: its upper block (transposed if lower). + dvector width_pos; ///< Per full block: first column within its row. + int num_s_blocks = 0; ///< Blocks of S (upper triangle). + int num_full = 0; ///< Full blocks. + int num_values = 0; ///< Values of S (upper block triangle). + int num_nonzeros = 0; ///< Nonzeros of the scalar CSR of S. + // Geometry and counts. + dvector ranges; ///< Per state batch with free columns. + dvector block_dim; ///< Per retained state: tangent size. + dvector error; ///< Setup error flags (see ClassifyFactorsKernel). + int num_entries = 0; ///< (landmark factor, retained slot) entries. + int num_rr = 0; ///< (factor, retained slot, retained slot) items of C^R. + dvector clique_offsets; ///< Per landmark: start of its clique products. + long long num_cliques = 0; ///< Σ_i k_i (k_i - 1) / 2 off-diagonal clique products. +}; + +void SchurSystem::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + const std::vector &reorderable) { + const auto &eliminated = problem.EliminatedStateBatches(); + if (eliminated.empty()) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: declare the state batches to eliminate with " + "Problem::SetEliminatedStateBatches"); + } + num_cols_ = num_cols; + structure_.ResolveColumns(stream, problem, num_cols); + SetupScratch scratch; + SetUpGeometry(stream, problem, scratch); + ClassifyFactors(stream, scratch); + BuildPairs(stream, scratch); + BuildEvaluationOrder(stream, reorderable, scratch); + BuildReducedPattern(stream, scratch); + BuildReducedCsr(stream, scratch); + AllocateNumericBuffers(stream, scratch); +} + +/** Residual batch geometry and slot tables; the landmarks and the retained states, numbered in + * column order. */ +void SchurSystem::SetUpGeometry(cudaStream_t stream, const Problem &problem, + SetupScratch &scratch) { + const auto &eliminated = problem.EliminatedStateBatches(); + const auto &residual_batches = problem.GetResidualBatches(); + const auto &layout = structure_.Layout(); + std::vector batches; + std::vector slot_sizes, slot_offsets; + factors_.batch_offsets.clear(); + int factor_cursor = 0; + for (size_t b = 0; b < residual_batches.size(); ++b) { + const HessianBatchLayout &lb = layout[b]; + BatchGeometry geometry{}; + geometry.jacobian_offset = static_cast(lb.jacobian_offset); + geometry.residual_offset = static_cast(lb.residual_offset); + geometry.factor_begin = factor_cursor; + geometry.num_residuals = lb.residual_dim; + geometry.tangent_dim = lb.tangent_dim; + geometry.num_slots = lb.num_blocks; + geometry.slot_begin = static_cast(slot_sizes.size()); + geometry.col_offset = static_cast(lb.col_offset); + int offset = 0; + for (size_t size : residual_batches[b].GetFactorBatch()->StateSizes()) { + slot_sizes.push_back(static_cast(size)); + slot_offsets.push_back(offset); + offset += static_cast(size); + } + batches.push_back(geometry); + factors_.batch_offsets.push_back(factor_cursor); + factor_cursor += lb.num_factors; + } + factors_.landmark.resize(factor_cursor); + factors_.landmark_slot.resize(factor_cursor); + factors_.batches = dvector(batches); + factors_.slot_sizes = dvector(slot_sizes); + factors_.slot_offsets = dvector(slot_offsets); + + // Column ranges: landmarks and retained blocks, numbered in column order. + std::vector ranges; + int col = 0, landmarks = 0, blocks = 0, reduced_col = 0; + landmark_dim_ = 0; + for (const StateBatch *batch : problem.GetStateBatches()) { + const int num_free = static_cast(batch->NumActiveStates() - batch->NumConstStates()); + const int tangent = static_cast(batch->TangentSize()); + if (num_free == 0 || tangent == 0) continue; + const bool is_eliminated = + std::find(eliminated.begin(), eliminated.end(), batch) != eliminated.end(); + if (is_eliminated) { + if (landmark_dim_ != 0 && landmark_dim_ != tangent) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: the eliminated state batches have different " + "tangent sizes"); + } + landmark_dim_ = tangent; + } + ranges.push_back({col, col + num_free * tangent, tangent, is_eliminated ? landmarks : blocks, + is_eliminated ? 1 : 0}); + col += num_free * tangent; + if (is_eliminated) { + landmarks += num_free; + } else { + blocks += num_free; + } + } + if (landmark_dim_ > 3) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: eliminated states must have tangent size 1 to 3"); + } + + scratch.ranges = dvector(ranges); + scratch.block_dim.resize(blocks); + pattern_.state_col.resize(blocks); + pattern_.state_natural_col.resize(blocks); + pattern_.landmark_natural_col.resize(landmarks); + for (const ColumnRange &range : ranges) { + const int num_states = (range.end_col - range.first_col) / range.tangent; + if (range.eliminated) { + LandmarkColumnsKernel<<>>( + num_states, range.tangent, range.first_id, range.first_col, + pattern_.landmark_natural_col.data()); + } else { + RetainedBlocksKernel<<>>( + num_states, range.tangent, range.first_id, range.first_col, reduced_col, + scratch.block_dim.data(), pattern_.state_col.data(), pattern_.state_natural_col.data()); + reduced_col += num_states * range.tangent; + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + num_retained_cols_ = reduced_col; + if (NumLandmarks() == 0) { + LogWarning( + "LinearSystemType::SchurComplement: no free state to eliminate; solving the full " + "system densely through the reduced system"); + } +} + +/** Classifies the factors (landmark, entries, C^R items) and checks the requirements of the class + * comment. */ +void SchurSystem::ClassifyFactors(cudaStream_t stream, SetupScratch &scratch) { + dvector &temp = scratch.temp; + const int F = NumFactors(); + const int *factor_cols = structure_.FactorCols().data(); + std::vector error_init = {0, INT_MAX, 0}; + scratch.error = dvector(error_init); + dvector entry_count(F + 1), rr_count(F + 1); + dvector &entry_offsets = scratch.entry_offsets, &rr_offsets = scratch.rr_offsets; + entry_offsets.resize(F + 1); + rr_offsets.resize(F + 1); + dvector &batch_has_landmarks = scratch.batch_has_landmarks; + batch_has_landmarks = dvector(0, std::max(NumBatches(), 1)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(entry_count.data(), 0, (F + 1) * sizeof(int), stream)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(rr_count.data(), 0, (F + 1) * sizeof(int), stream)); + if (F > 0) { + ClassifyFactorsKernel<<>>( + F, factors_.batches.data(), NumBatches(), factors_.slot_sizes.data(), factor_cols, + scratch.ranges.data(), static_cast(scratch.ranges.size()), factors_.landmark.data(), + factors_.landmark_slot.data(), entry_count.data(), rr_count.data(), + batch_has_landmarks.data(), scratch.error.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + ExclusiveSum(stream, temp, entry_count.data(), entry_offsets.data(), F + 1); + ExclusiveSum(stream, temp, rr_count.data(), rr_offsets.data(), F + 1); + std::vector errors(3); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(errors.data(), scratch.error.data(), 3 * sizeof(int), + cudaMemcpyDeviceToHost, stream)); + scratch.num_entries = ReadBack(stream, entry_offsets.data() + F); + scratch.num_rr = ReadBack(stream, rr_offsets.data() + F); + if (errors[0] & kTwoLandmarks) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: a factor touches two states of the eliminated state " + "batches; keep such states retained"); + } + pair_block_dim_ = scratch.num_entries > 0 ? errors[2] : 6; + if (scratch.num_entries > 0 && errors[1] != errors[2]) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: the states that share factors with landmarks must all " + "have one tangent size"); + } + if (NumLandmarks() > 0 && !SupportedDims(landmark_dim_, pair_block_dim_)) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: unsupported tangent sizes (landmarks " + + std::to_string(landmark_dim_) + ", states observing them " + + std::to_string(pair_block_dim_) + "); supported: 1-3 with 6, and 3 with 3 or 9"); + } + scratch.block_bits = BitsFor(std::max(NumRetainedStates(), 2)); + scratch.landmark_bits = BitsFor(std::max(NumLandmarks(), 2)); + if (2 * scratch.block_bits > 64 || scratch.landmark_bits + scratch.block_bits > 64) { + throw std::invalid_argument("LinearSystemType::SchurComplement: problem too large"); + } +} + +/** The (retained state, landmark) pairs, landmark-major, and their pose-major order cut into chunks + * for the off-diagonal pass. */ +void SchurSystem::BuildPairs(cudaStream_t stream, SetupScratch &scratch) { + // Entries, sorted by (landmark, retained state); stable, so by factor within. + dvector &temp = scratch.temp; + const int F = NumFactors(); + const int *factor_cols = structure_.FactorCols().data(); + const int block_bits = scratch.block_bits, landmark_bits = scratch.landmark_bits; + const dvector &entry_offsets = scratch.entry_offsets, &rr_offsets = scratch.rr_offsets; + const int E = scratch.num_entries, R = scratch.num_rr; + dvector entry_keys(E), entry_keys_sorted(E); + dvector entry_values(E), entry_values_sorted(E); + pattern_.rr_factor.resize(R); + pattern_.rr_slots.resize(R); + dvector &rr_keys = scratch.rr_keys; + rr_keys.resize(R); + if (F > 0) { + EmitItemsKernel<<>>( + F, factors_.batches.data(), NumBatches(), factor_cols, scratch.ranges.data(), + static_cast(scratch.ranges.size()), factors_.landmark.data(), + factors_.landmark_slot.data(), entry_offsets.data(), rr_offsets.data(), block_bits, + entry_keys.data(), entry_values.data(), pattern_.rr_factor.data(), pattern_.rr_slots.data(), + rr_keys.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + SortPairs(stream, temp, entry_keys.data(), entry_keys_sorted.data(), entry_values.data(), + entry_values_sorted.data(), E, landmark_bits + block_bits); + + // Pairs: runs of equal keys. + dvector &entry_factor = scratch.entry_factor, &entry_slot = scratch.entry_slot; + entry_factor.resize(E); + entry_slot.resize(E); + dvector head(E), pair_id(E); + int P = 0; + if (E > 0) { + UnpackEntriesKernel<<>>( + E, entry_keys_sorted.data(), entry_values_sorted.data(), entry_factor.data(), + entry_slot.data(), head.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + InclusiveSum(stream, temp, head.data(), pair_id.data(), E); + P = ReadBack(stream, pair_id.data() + E - 1); + } + pairs_.landmark.resize(P); + pairs_.state.resize(P); + dvector &pair_entry_offsets = scratch.pair_entry_offsets; + pair_entry_offsets.resize(P + 1); + pairs_.landmark_offsets.resize(NumLandmarks() + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(pairs_.landmark_offsets.data(), 0, + (NumLandmarks() + 1) * sizeof(int), stream)); + if (E > 0) { + BuildPairsKernel<<>>( + E, entry_keys_sorted.data(), head.data(), pair_id.data(), block_bits, + pairs_.landmark.data(), pairs_.state.data(), pair_entry_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + LandmarkOffsetsKernel<<>>( + P, NumLandmarks(), pairs_.landmark.data(), pairs_.landmark_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + // Pose-major order of the pairs (stable: landmark-ascending within a state). + dvector pair_iota; + Iota(stream, pair_iota, P); + dvector pm_keys(P); + pairs_.pose_major.resize(P); + pairs_.pose_major_position.resize(P); + SortPairs(stream, temp, reinterpret_cast(pairs_.state.data()), pm_keys.data(), + pair_iota.data(), pairs_.pose_major.data(), P, block_bits); + if (P > 0) { + InvertPermutationKernel<<>>(P, pairs_.pose_major.data(), + pairs_.pose_major_position.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + // Pose-major runs of the retained states, cut into chunks for the + // off-diagonal pass. + { + dvector pm_row_offsets(NumRetainedStates() + 1); + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(pm_row_offsets.data(), 0, (NumRetainedStates() + 1) * sizeof(int), stream)); + if (P > 0) { + LandmarkOffsetsKernel<<>>( + P, NumRetainedStates(), reinterpret_cast(pm_keys.data()), + pm_row_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + std::vector offsets(NumRetainedStates() + 1); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(offsets.data(), pm_row_offsets.data(), + offsets.size() * sizeof(int), cudaMemcpyDeviceToHost, + stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + std::vector chunks; + for (int b = 0; b < NumRetainedStates(); ++b) { + for (int t = offsets[b]; t < offsets[b + 1]; t += kChunkPairs) { + chunks.push_back(make_int3(b, t, std::min(t + kChunkPairs, offsets[b + 1]))); + } + } + pairs_.chunks = dvector(chunks); + } +} + +/** The pose-major entries and the evaluation order that lays their Jacobians out in that order. */ +void SchurSystem::BuildEvaluationOrder(cudaStream_t stream, const std::vector &reorderable, + SetupScratch &scratch) { + dvector &temp = scratch.temp; + const int F = NumFactors(), E = scratch.num_entries, P = NumPairs(); + const int *factor_cols = structure_.FactorCols().data(); + const dvector &pair_entry_offsets = scratch.pair_entry_offsets; + const dvector &entry_factor = scratch.entry_factor, &entry_slot = scratch.entry_slot; + const dvector &batch_has_landmarks = scratch.batch_has_landmarks; + pairs_.entry_offsets.resize(P + 1); + { + dvector pm_entry_count(P + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(pm_entry_count.data(), 0, (P + 1) * sizeof(int), stream)); + dvector pm_entry_factor(E), pm_entry_slot(E); + dvector factor_key(F); + if (F > 0) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(factor_key.data(), 0x7f, F * sizeof(int), stream)); + } + if (P > 0) { + PoseMajorPairsKernel<<>>( + P, pairs_.pose_major.data(), pair_entry_offsets.data(), pm_entry_count.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + ExclusiveSum(stream, temp, pm_entry_count.data(), pairs_.entry_offsets.data(), P + 1); + if (P > 0) { + PoseMajorEntriesKernel<<>>( + P, pairs_.pose_major.data(), pair_entry_offsets.data(), pairs_.entry_offsets.data(), + entry_factor.data(), entry_slot.data(), pm_entry_factor.data(), pm_entry_slot.data(), + factor_key.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + std::vector has_landmarks(NumBatches()); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(has_landmarks.data(), batch_has_landmarks.data(), + NumBatches() * sizeof(int), cudaMemcpyDeviceToHost, + stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + factors_.reordered.assign(NumBatches(), false); + std::vector reordered(NumBatches()); + for (int b = 0; b < NumBatches(); ++b) { + factors_.reordered[b] = + has_landmarks[b] != 0 && b < static_cast(reorderable.size()) && reorderable[b]; + reordered[b] = factors_.reordered[b] ? 1 : 0; + } + dvector reordered_device(reordered); + dvector order_keys(F), order_keys_sorted(F); + dvector order_values(F), order_values_sorted(F); + factors_.item.resize(F); + factors_.evaluation_order.resize(F); + factors_.item_cols.resize(structure_.FactorCols().size()); + if (F > 0) { + OrderKeysKernel<<>>( + F, factors_.batches.data(), NumBatches(), reordered_device.data(), factor_key.data(), + order_keys.data(), order_values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + SortPairs(stream, temp, order_keys.data(), order_keys_sorted.data(), order_values.data(), + order_values_sorted.data(), F, 32 + BitsFor(std::max(NumBatches(), 2))); + ItemsKernel<<>>( + F, factors_.batches.data(), NumBatches(), order_values_sorted.data(), factor_cols, + factors_.item.data(), factors_.evaluation_order.data(), factors_.item_cols.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + pairs_.entry_jacobian.resize(E); + pairs_.entry_residual.resize(E); + pairs_.entry_meta.resize(E); + if (E > 0) { + EntryLocationsKernel<<>>( + E, pm_entry_factor.data(), pm_entry_slot.data(), factors_.batches.data(), NumBatches(), + factors_.slot_offsets.data(), factors_.landmark_slot.data(), factors_.item.data(), + pairs_.entry_jacobian.data(), pairs_.entry_residual.data(), pairs_.entry_meta.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + } +} + +/** The block pattern of S: clique products, C^R items and the diagonal; value offsets and the full + * block rows. */ +void SchurSystem::BuildReducedPattern(cudaStream_t stream, SetupScratch &scratch) { + dvector &temp = scratch.temp; + const int P = NumPairs(), R = scratch.num_rr, block_bits = scratch.block_bits; + const dvector &rr_keys = scratch.rr_keys; + // Clique products and the block pattern of S (upper triangle). + const int L = NumLandmarks(); + dvector clique_counts(L + 1); + scratch.clique_offsets.resize(L + 1); + CliqueCountsKernel<<>>(L, pairs_.landmark_offsets.data(), + clique_counts.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + ExclusiveSum(stream, temp, clique_counts.data(), scratch.clique_offsets.data(), L + 1); + scratch.num_cliques = ReadBack(stream, scratch.clique_offsets.data() + L); + const long long K = scratch.num_cliques + NumRetainedStates() + R; + if (K >= INT_MAX) { + throw std::invalid_argument( + "LinearSystemType::SchurComplement: too many landmark observation pairs"); + } + const int Ki = static_cast(K); + dvector keys(Ki), keys_sorted(Ki); + dvector values(Ki), values_sorted(Ki); + if (P > 0) { + CliqueKeysKernel<<>>( + P, pairs_.landmark.data(), pairs_.state.data(), pairs_.landmark_offsets.data(), + scratch.clique_offsets.data(), block_bits, keys.data(), values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + if (NumRetainedStates() + R > 0) { + OtherKeysKernel<<>>( + NumRetainedStates(), R, scratch.num_cliques, rr_keys.data(), block_bits, keys.data(), + values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + SortPairs(stream, temp, keys.data(), keys_sorted.data(), values.data(), values_sorted.data(), Ki, + 2 * block_bits); + dvector key_head(Ki), block_id(Ki); + int &num_s_blocks = scratch.num_s_blocks; + if (Ki > 0) { + HeadFlagsKernel<<>>(Ki, keys_sorted.data(), key_head.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + InclusiveSum(stream, temp, key_head.data(), block_id.data(), Ki); + num_s_blocks = ReadBack(stream, block_id.data() + Ki - 1); + } + dvector s_row(num_s_blocks); + pattern_.upper_col.resize(num_s_blocks); + dvector &s_col = pattern_.upper_col; + pattern_.state_diagonal.resize(NumRetainedStates()); + pattern_.rr_block.resize(R); + if (Ki > 0) { + AssignBlocksKernel<<>>( + Ki, keys_sorted.data(), values_sorted.data(), key_head.data(), block_id.data(), block_bits, + scratch.num_cliques, NumRetainedStates(), s_row.data(), s_col.data(), + pattern_.state_diagonal.data(), pattern_.rr_block.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + // The key arrays are the largest setup buffers; release them early. + keys = {}; + keys_sorted = {}; + values = {}; + values_sorted = {}; + key_head = {}; + block_id = {}; + + // Value offsets of the S blocks; full (both triangles) block rows. + const int NB = NumRetainedStates(); + dvector value_counts(num_s_blocks + 1), upper_count(NB + 1), lower_count(NB + 1); + pattern_.value_offsets.resize(num_s_blocks + 1); + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(value_counts.data(), 0, (num_s_blocks + 1) * sizeof(int), stream)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(upper_count.data(), 0, (NB + 1) * sizeof(int), stream)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(lower_count.data(), 0, (NB + 1) * sizeof(int), stream)); + dvector num_lower_device(0, 1); + dvector lower_keys(num_s_blocks), lower_keys_sorted(num_s_blocks); + dvector lower_values(num_s_blocks), lower_values_sorted(num_s_blocks); + if (num_s_blocks > 0) { + BlockSizesKernel<<>>( + num_s_blocks, s_row.data(), s_col.data(), scratch.block_dim.data(), value_counts.data(), + upper_count.data(), lower_count.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + LowerKeysKernel<<>>( + num_s_blocks, s_row.data(), s_col.data(), block_bits, lower_keys.data(), + lower_values.data(), num_lower_device.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + ExclusiveSum(stream, temp, value_counts.data(), pattern_.value_offsets.data(), num_s_blocks + 1); + const int num_lower = ReadBack(stream, num_lower_device.data()); + scratch.num_values = ReadBack(stream, pattern_.value_offsets.data() + num_s_blocks); + SortPairs(stream, temp, lower_keys.data(), lower_keys_sorted.data(), lower_values.data(), + lower_values_sorted.data(), num_lower, 2 * block_bits); + pattern_.upper_row_offsets.resize(NB + 1); + dvector &upper_row_offsets = pattern_.upper_row_offsets; + dvector lower_row_offsets(NB + 1); + dvector &full_row_offsets = scratch.full_row_offsets; + full_row_offsets.resize(NB + 1); + ExclusiveSum(stream, temp, upper_count.data(), upper_row_offsets.data(), NB + 1); + ExclusiveSum(stream, temp, lower_count.data(), lower_row_offsets.data(), NB + 1); + FullRowOffsetsKernel<<>>( + NB, upper_row_offsets.data(), lower_row_offsets.data(), full_row_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + const int num_full = scratch.num_full = num_s_blocks + num_lower; + dvector &full_row = scratch.full_row, &full_col = scratch.full_col; + dvector &full_src = scratch.full_src, &width_pos = scratch.width_pos; + full_row.resize(num_full); + full_col.resize(num_full); + full_src.resize(num_full); + width_pos.resize(num_full + 1); + dvector widths(num_full + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(widths.data(), 0, (num_full + 1) * sizeof(int), stream)); + if (num_full > 0) { + PlaceFullBlocksKernel<<>>( + num_s_blocks, num_lower, s_row.data(), s_col.data(), lower_keys_sorted.data(), + lower_values_sorted.data(), upper_row_offsets.data(), lower_row_offsets.data(), block_bits, + full_row.data(), full_col.data(), full_src.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + FullWidthsKernel<<>>( + num_full, full_col.data(), scratch.block_dim.data(), widths.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + ExclusiveSum(stream, temp, widths.data(), width_pos.data(), num_full + 1); +} + +/** The scalar CSR of S (both triangles) and the column maps. */ +void SchurSystem::BuildReducedCsr(cudaStream_t stream, SetupScratch &scratch) { + dvector &temp = scratch.temp; + const int NB = NumRetainedStates(), num_full = scratch.num_full; + const dvector &full_row_offsets = scratch.full_row_offsets; + const dvector &full_row = scratch.full_row, &full_col = scratch.full_col; + const dvector &full_src = scratch.full_src, &width_pos = scratch.width_pos; + dvector row_length(NB), row_nonzeros(NB + 1), row_base(NB + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(row_nonzeros.data(), 0, (NB + 1) * sizeof(int), stream)); + if (NB > 0) { + RowLengthsKernel<<>>( + NB, full_row_offsets.data(), width_pos.data(), scratch.block_dim.data(), row_length.data(), + row_nonzeros.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + ExclusiveSum(stream, temp, row_nonzeros.data(), row_base.data(), NB + 1); + const int num_nonzeros = scratch.num_nonzeros = ReadBack(stream, row_base.data() + NB); + reduced_.row_offsets.resize(num_retained_cols_ + 1); + reduced_.col_ids.resize(num_nonzeros); + reduced_.values.resize(num_nonzeros); + pattern_.csr_source.resize(num_nonzeros); + if (NB > 0) { + CsrRowOffsetsKernel<<>>( + NB, scratch.block_dim.data(), pattern_.state_col.data(), row_base.data(), row_length.data(), + num_nonzeros, num_retained_cols_, reduced_.row_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + if (num_full > 0) { + CsrFillKernel<<>>( + num_full, full_row.data(), full_col.data(), full_src.data(), full_row_offsets.data(), + width_pos.data(), scratch.block_dim.data(), pattern_.state_col.data(), row_base.data(), + row_length.data(), pattern_.value_offsets.data(), reduced_.col_ids.data(), + pattern_.csr_source.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + pattern_.col_natural.resize(num_retained_cols_); + pattern_.col_diagonal.resize(num_retained_cols_); + if (NB > 0) { + ColumnMapsKernel<<>>( + NB, scratch.block_dim.data(), pattern_.state_col.data(), pattern_.state_natural_col.data(), + pattern_.state_diagonal.data(), pattern_.value_offsets.data(), pattern_.col_natural.data(), + pattern_.col_diagonal.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } +} + +/** The per-solve buffers. */ +void SchurSystem::AllocateNumericBuffers(cudaStream_t stream, const SetupScratch &scratch) { + const int L = NumLandmarks(), P = NumPairs(), F = NumFactors(); + const int tri = landmark_dim_ * (landmark_dim_ + 1) / 2; + hessian_diagonal_.resize(num_cols_); + landmark_hessian_.resize(static_cast(tri) * L); + s_values_.resize(scratch.num_values); + pair_v_.resize(static_cast(P) * pair_block_dim_ * + (landmark_dim_ == 3 ? 4 : landmark_dim_)); + reduced_rhs_.resize(num_retained_cols_); + reduced_step_.resize(num_retained_cols_); + factor_curvature_.resize(F); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + LogMessage( + "Schur complement: {} landmarks (dim {}), {} retained states ({} columns, dim {} with " + "landmarks), {} pairs, {} clique products, S: {} blocks, {} nonzeros", + L, landmark_dim_, NumRetainedStates(), num_retained_cols_, pair_block_dim_, P, + scratch.num_cliques, scratch.num_s_blocks, scratch.num_nonzeros); +} + +void SchurSystem::Assemble(cudaStream_t stream, const Problem &, const float *jacobians, + const float *residuals, dvector &rhs) { + jacobians_ = jacobians; + rhs.resize(num_cols_); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(rhs.data(), 0, rhs.size() * sizeof(float), stream)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(hessian_diagonal_.data(), 0, + hessian_diagonal_.size() * sizeof(float), stream)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(landmark_hessian_.data(), 0, + landmark_hessian_.size() * sizeof(float), stream)); + if (NumFactors() == 0) return; + AssembleFactorsKernel<<>>( + NumFactors(), factors_.batches.data(), NumBatches(), factors_.slot_sizes.data(), + factors_.slot_offsets.data(), factors_.item_cols.data(), factors_.landmark.data(), + factors_.landmark_slot.data(), factors_.evaluation_order.data(), landmark_dim_, jacobians, + residuals, rhs.data(), hessian_diagonal_.data(), landmark_hessian_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + if (NumPairs() == 0) return; + SCHUR_DISPATCH_DIMS(landmark_dim_, pair_block_dim_, { + AssemblePairsKernel<<>>( + NumPairs(), pairs_.entry_offsets.data(), pairs_.entry_jacobian.data(), + pairs_.entry_residual.data(), pairs_.entry_meta.data(), pairs_.pose_major.data(), + pairs_.state.data(), jacobians, residuals, pattern_.state_natural_col.data(), rhs.data(), + hessian_diagonal_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + }); +} + +void SchurSystem::ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const { + diagonal.resize(hessian_diagonal_.size()); + if (diagonal.empty()) return; + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(diagonal.data(), hessian_diagonal_.data(), + diagonal.size() * sizeof(float), cudaMemcpyDeviceToDevice, + stream)); +} + +void SchurSystem::FormReducedSystem(cudaStream_t stream, const dvector &rhs) { + const Modifiers mod{ColumnScale() ? ColumnScale()->data() : nullptr, + Mask() ? Mask()->data() : nullptr, Damping() ? Damping()->data() : nullptr}; + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(s_values_.data(), 0, s_values_.size() * sizeof(float), stream)); + if (num_retained_cols_ > 0) { + GatherRhsKernel<<>>( + num_retained_cols_, pattern_.col_natural.data(), rhs.data(), reduced_rhs_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + const int *factor_cols = structure_.FactorCols().data(); + const int num_rr = static_cast(pattern_.rr_factor.size()); + if (num_rr > 0) { + RetainedProductsKernel<<>>( + num_rr, pattern_.rr_factor.data(), pattern_.rr_slots.data(), pattern_.rr_block.data(), + factors_.batches.data(), NumBatches(), factors_.slot_sizes.data(), + factors_.slot_offsets.data(), factor_cols, factors_.item.data(), jacobians_, + pattern_.value_offsets.data(), mod, s_values_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + if (NumPairs() > 0) { + SCHUR_DISPATCH_DIMS(landmark_dim_, pair_block_dim_, { + EliminatePairsKernel<<>>( + NumPairs(), pairs_.entry_offsets.data(), pairs_.entry_jacobian.data(), + pairs_.entry_meta.data(), pairs_.pose_major.data(), pairs_.landmark.data(), + pairs_.state.data(), jacobians_, landmark_hessian_.data(), NumLandmarks(), + pattern_.landmark_natural_col.data(), pattern_.state_natural_col.data(), + pattern_.state_col.data(), pattern_.state_diagonal.data(), pattern_.value_offsets.data(), + rhs.data(), mod, pair_v_.data(), s_values_.data(), reduced_rhs_.data(), failure_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + OffDiagonalRowsKernel + <<(pairs_.chunks.size()), kThreads, 0, stream>>>( + pairs_.chunks.data(), pattern_.upper_row_offsets.data(), pattern_.upper_col.data(), + pattern_.value_offsets.data(), pairs_.pose_major.data(), pairs_.landmark.data(), + pairs_.state.data(), pairs_.pose_major_position.data(), + pairs_.landmark_offsets.data(), pair_v_.data(), s_values_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + }); + } + if (num_retained_cols_ > 0) { + FinalizeDiagonalKernel<<>>( + num_retained_cols_, pattern_.col_natural.data(), pattern_.col_diagonal.data(), mod, + s_values_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + const int nnz = static_cast(pattern_.csr_source.size()); + if (nnz > 0) { + ExpandKernel<<>>(nnz, pattern_.csr_source.data(), + s_values_.data(), reduced_.values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } +} + +bool SchurSystem::PrepareSolver(cudaStream_t stream, const Problem &problem, + const dvector &rhs, dvector &) { + FormReducedSystem(stream, rhs); + return solver_->Initialize(stream, problem, reduced_, reduced_rhs_, reduced_step_); +} + +bool SchurSystem::Solve(cudaStream_t stream, const dvector &rhs, dvector &step) { + if (check_failures_) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(failure_.data(), 0, sizeof(int), stream)); + } + FormReducedSystem(stream, rhs); + if (!solver_->Solve(stream, reduced_, reduced_rhs_, reduced_step_)) return false; + step.resize(num_cols_); + if (num_retained_cols_ > 0) { + ScatterRetainedKernel<<>>( + num_retained_cols_, pattern_.col_natural.data(), reduced_step_.data(), step.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + if (NumLandmarks() > 0) { + const Modifiers mod{ColumnScale() ? ColumnScale()->data() : nullptr, + Mask() ? Mask()->data() : nullptr, Damping() ? Damping()->data() : nullptr}; + SCHUR_DISPATCH_DIMS(landmark_dim_, pair_block_dim_, { + BackSubstituteKernel<<>>( + NumLandmarks(), NumPairs(), pairs_.landmark_offsets.data(), pairs_.state.data(), + pairs_.pose_major_position.data(), landmark_hessian_.data(), + pattern_.landmark_natural_col.data(), pattern_.state_col.data(), pair_v_.data(), + reduced_step_.data(), rhs.data(), mod, step.data(), failure_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + }); + } + if (check_failures_ && ReadBack(stream, failure_.data()) != 0) { + LogError("Schur complement: a landmark block could not be factored (non-finite values)"); + return false; + } + return true; +} + +void SchurSystem::StepCurvature(cudaStream_t stream, const dvector &step, + ProblemPartition &partition, float *out) const { + if (NumFactors() > 0) { + FactorCurvatureKernel<<>>( + NumFactors(), factors_.batches.data(), NumBatches(), factors_.slot_sizes.data(), + factors_.slot_offsets.data(), factors_.item_cols.data(), factors_.evaluation_order.data(), + jacobians_, step.data(), factor_curvature_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + partition.SumFactorCosts(stream, factor_curvature_.data(), out); +} + +const int *SchurSystem::EvaluationOrder(size_t batch) const { + if (batch >= factors_.reordered.size() || !factors_.reordered[batch]) return nullptr; + return factors_.evaluation_order.data() + factors_.batch_offsets[batch]; +} + +} // namespace cunls diff --git a/cunls/minimizer/schur_system.h b/cunls/minimizer/schur_system.h new file mode 100644 index 0000000..46db6cc --- /dev/null +++ b/cunls/minimizer/schur_system.h @@ -0,0 +1,265 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include + +#include "cunls/common/types.h" +#include "cunls/linear_solver/sparse_linear_solver.h" +#include "cunls/minimizer/hessian_structure.h" +#include "cunls/minimizer/linear_system.h" + +namespace cunls { + +namespace internal { +namespace schur { + +/** @brief Per residual batch geometry, for kernels that walk all factors (internal). */ +struct BatchGeometry { + long long jacobian_offset; ///< Start of the batch in the Jacobian buffer. + long long residual_offset; ///< Start of the batch in the residual buffer. + int factor_begin; ///< Global index of the batch's first factor. + int num_residuals; ///< m: residuals per factor. + int tangent_dim; ///< n: Jacobian columns per factor (sum of slot sizes). + int num_slots; ///< State slots per factor. + int slot_begin; ///< Start of the batch's slots in the slot tables. + int col_offset; ///< Start of the batch in HessianStructureBuilder::FactorCols(). +}; + +/** @brief A state batch's range of free columns and its role in the elimination (internal). */ +struct ColumnRange { + int first_col; ///< First free column of the batch. + int end_col; ///< One past its last free column. + int tangent; ///< Tangent size of its states. + int first_id; ///< Landmark id (eliminated) or retained block id of its first free state. + int eliminated; ///< Nonzero for a batch of Problem::EliminatedStateBatches(). +}; + +} // namespace schur +} // namespace internal + +/** + * @brief The linear system with the landmark states eliminated + * (LinearSystemType::SchurComplement). + * + * (Here S is the Schur complement; the column scale of LinearSystem is + * applied to the full system before the elimination.) + * + * Partitions the free columns into the eliminated landmarks e (the batches of + * Problem::SetEliminatedStateBatches) and the retained states r. No factor + * touches two landmarks, so the landmark block D of the Hessian is block + * diagonal and each Solve() computes, for the modified system (column scale, + * damping and mask applied before elimination): + * + * S = C - W D⁻¹ Wᵀ, g = b_r - W D⁻¹ b_e, S x_r = g, + * x_e,i = D_i⁻¹ (b_e,i - Σ_p W_piᵀ x_r,p), + * + * without ever forming the full Hessian. Per (retained state p, landmark i) + * pair, with L_i the Cholesky factor of D_i, V_pi = W_pi L_i⁻ᵀ and + * h_i = L_i⁻¹ b_e,i: + * + * S_pp = C^R_pp + Σ_i (A_pi - V_pi V_piᵀ) (A_pi: the pair's own J_pᵀJ_p; + * no large-minus-large cancellation) + * S_pq = C^R_pq - Σ_i V_pi V_qiᵀ, g_p = b_r,p - Σ_i V_pi h_i. + * + * S (upper block triangle, variable block sizes) is expanded to the full + * scalar CSR and solved with BlockSparseCholeskySolver (CreateLinearSystem; + * S is sparse, and banded for visual-inertial chains, so a dense + * factorization does not scale and an iterative one is not exact enough). + * + * Requirements (checked in Initialize()): every factor touches at most one + * free landmark; landmarks have tangent size 1 to 3; the retained states that + * share factors with landmarks all have one tangent size, 3, 6 or 9. Other + * factors (between retained states only) may have any shape. + * + * Work per Solve() (P pairs, L landmarks, Σ k_i² clique products): + * - pose-major pass over the pairs: contract the pair's Jacobians, eliminate + * the landmark (3x3 Cholesky in registers), write V, and reduce the + * diagonal brackets and g with warp-segmented sums (one atomic per + * retained state and warp); + * - landmark-major pass: one warp per landmark forms -V_pi V_qiᵀ for every + * pair of its observations from shared memory, with coalesced atomics; + * - the backend solve of S, then one thread per landmark back-substitutes. + */ +class SchurSystem : public LinearSystem { + public: + /** + * @brief Takes ownership of the solver of the reduced system. + * + * @param solver Solver of the reduced system S. + * @param check_failures Read back, after every Solve(), whether a landmark + * block could not be factored even with a diagonal shift (non-finite + * values), and return false if so (one synchronization; + * MinimizerOptions::disable_safety_checks off). Without it, such a + * landmark makes the step NaN, which the minimizer rejects. + */ + SchurSystem(SparseLinearSolverPtr solver, bool check_failures); + + /** + * @brief Builds the elimination plan on the GPU: landmark and retained + * numbering, the (state, landmark) pairs, the pattern of S and the maps + * that write into it, and the expansion of S to scalar CSR. + * + * @throws std::invalid_argument if the problem declares no eliminated state + * batches, or violates the requirements of the class comment. + */ + void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + const std::vector &reorderable) override; + + size_t JacobianValuesSize() const override { return structure_.JacobianValuesSize(); } + + /** + * @brief For a reorderable batch with landmark factors: its landmark + * factors in the pose-major order of their pairs (then its other factors), + * so the pair passes stream the Jacobians. + */ + const int *EvaluationOrder(size_t batch) const override; + + /** @brief rhs = -Jᵀr, diag(H) and the landmark blocks D, one pass over the factors. */ + void Assemble(cudaStream_t stream, const Problem &problem, const float *jacobians, + const float *residuals, dvector &rhs) override; + + void ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const override; + + /** @brief Forms S and runs the backend's symbolic setup on it. */ + bool PrepareSolver(cudaStream_t stream, const Problem &problem, const dvector &rhs, + dvector &step) override; + + /** @brief Eliminates, solves S x_r = g and back-substitutes the landmarks. */ + bool Solve(cudaStream_t stream, const dvector &rhs, dvector &step) override; + + /** @brief δᵀHδ per subproblem as Σ over factors of ‖J_f δ_f‖². */ + void StepCurvature(cudaStream_t stream, const dvector &step, ProblemPartition &partition, + float *out) const override; + + /** @brief The reduced system S of the last Solve() (full scalar CSR, for tests). */ + const CSRSparseMatrix &ReducedMatrix() const { return reduced_; } + + /** @brief Retained (reduced) dimension. */ + int NumRetainedColumns() const { return num_retained_cols_; } + + private: + /// Temporaries shared by the setup phases (defined in the implementation). + struct SetupScratch; + + // Setup phases of Initialize(), in order. + /** @brief Batch geometry, slot tables, and the landmark and retained numbering. */ + void SetUpGeometry(cudaStream_t stream, const Problem &problem, SetupScratch &scratch); + /** @brief Classifies the factors and checks the requirements of the class comment. */ + void ClassifyFactors(cudaStream_t stream, SetupScratch &scratch); + /** @brief The (retained state, landmark) pairs and their pose-major order. */ + void BuildPairs(cudaStream_t stream, SetupScratch &scratch); + /** @brief The pose-major entries and the Jacobian evaluation order. */ + void BuildEvaluationOrder(cudaStream_t stream, const std::vector &reorderable, + SetupScratch &scratch); + /** @brief The block pattern of S and the maps that write into it. */ + void BuildReducedPattern(cudaStream_t stream, SetupScratch &scratch); + /** @brief The scalar CSR of S handed to the solver. */ + void BuildReducedCsr(cudaStream_t stream, SetupScratch &scratch); + /** @brief The per-solve buffers. */ + void AllocateNumericBuffers(cudaStream_t stream, const SetupScratch &scratch); + + /** @brief Builds S and g from the current Jacobians and declared modifications. */ + void FormReducedSystem(cudaStream_t stream, const dvector &rhs); + + /** @brief Per residual batch and per factor: where its Jacobians, slots and landmark are. */ + struct FactorTables { + dvector batches; ///< Per residual batch. + std::vector batch_offsets; ///< Per residual batch: its first factor (global index). + std::vector reordered; ///< Per residual batch: evaluated in evaluation_order. + dvector slot_sizes; ///< Per (batch, slot): tangent size. + dvector slot_offsets; ///< Per (batch, slot): first Jacobian column of the slot. + dvector landmark; ///< Per factor: its landmark, or -1. + dvector landmark_slot; ///< Per factor: its landmark's slot, or -1. + dvector item; ///< Per factor: its item (position in its batch's evaluation). + dvector evaluation_order; ///< Per item, batches concatenated: its batch-local factor. + dvector item_cols; ///< HessianStructureBuilder::FactorCols() in item order. + }; + + /** + * @brief The (retained state, landmark) pairs that share factors, numbered + * landmark-major, and their entries (one per landmark factor and retained + * slot) in pose-major order, the order of V and of the Jacobians. + */ + struct PairTables { + dvector landmark; ///< Per pair: its landmark. + dvector state; ///< Per pair: its retained state (block row of S). + dvector landmark_offsets; ///< Per landmark: its pairs [offset, next). + dvector pose_major; ///< Pose-major position -> pair. + dvector pose_major_position; ///< Pair -> pose-major position. + dvector entry_offsets; ///< Per pose-major pair: its entries [offset, next). + dvector entry_jacobian; ///< Per entry: its factor's first Jacobian value. + dvector entry_residual; ///< Per entry: its factor's first residual. + /// Per entry: retained slot column | landmark slot column << 8 | m << 16 | n << 24. + dvector entry_meta; + /// Off-diagonal pass, one CTA each: retained state, pose-major pairs [begin, end). + dvector chunks; + }; + + /** @brief The pattern of S and the maps that write into it and expand it. */ + struct ReducedPattern { + dvector state_col; ///< Per retained state: first column of S. + dvector state_natural_col; ///< Per retained state: first column of the full system. + dvector state_diagonal; ///< Per retained state: its diagonal block of S. + dvector landmark_natural_col; ///< Per landmark: first column of the full system. + dvector upper_row_offsets; ///< Per retained state: its upper blocks of S [offset, next). + dvector upper_col; ///< Per block of S (upper): its column block. + dvector value_offsets; ///< Per block of S (upper): its first value (row-major). + dvector rr_factor; ///< Per C^R item: its factor. + dvector rr_slots; ///< Per C^R item: slot a | slot b << 8 | transposed << 16. + dvector rr_block; ///< Per C^R item: its block of S. + dvector col_natural; ///< Per column of S: column of the full system. + dvector col_diagonal; ///< Per column of S: position of its diagonal value. + dvector csr_source; ///< Per nonzero of the scalar CSR: position in the values. + }; + + int NumFactors() const { return static_cast(factors_.landmark.size()); } + int NumBatches() const { return static_cast(factors_.batches.size()); } + int NumLandmarks() const { return static_cast(pattern_.landmark_natural_col.size()); } + int NumRetainedStates() const { return static_cast(pattern_.state_col.size()); } + int NumPairs() const { return static_cast(pairs_.landmark.size()); } + + SparseLinearSolverPtr solver_; ///< Solver of the reduced system. + bool check_failures_ = false; ///< See the constructor. + HessianStructureBuilder structure_; ///< Layout and per-(factor, slot) columns. + + int num_cols_ = 0; ///< Free columns of the full system. + int num_retained_cols_ = 0; ///< Columns of S. + int landmark_dim_ = 0; ///< DE: tangent size of the landmarks. + int pair_block_dim_ = 0; ///< DP: tangent size of the retained states paired with landmarks. + FactorTables factors_; + PairTables pairs_; + ReducedPattern pattern_; + + // Numeric state. + const float *jacobians_ = nullptr; ///< Assemble(): per-factor Jacobians. + dvector hessian_diagonal_; ///< diag(H), unscaled. + dvector landmark_hessian_; ///< Per landmark: upper triangle of D_i, unscaled (SoA). + dvector s_values_; ///< S, upper block triangle, row-major blocks. + dvector pair_v_; ///< V per pair, pose-major SoA. + CSRSparseMatrix reduced_; ///< S, full scalar CSR, for the solver. + dvector reduced_rhs_; ///< g. + dvector reduced_step_; ///< x_r. + mutable dvector factor_curvature_; ///< Per factor ‖J_f δ_f‖² (StepCurvature()). + dvector failure_; ///< Nonzero: a landmark block failed to factor. +}; + +} // namespace cunls diff --git a/docs/sphinx/api/linear_solver.rst b/docs/sphinx/api/linear_solver.rst index 7060ad4..aef20ca 100644 --- a/docs/sphinx/api/linear_solver.rst +++ b/docs/sphinx/api/linear_solver.rst @@ -2,11 +2,13 @@ Linear Solver API ################################################################################ -`cunls/linear_solver` hosts linear-system abstractions (block-Jacobi PCG, -cuDSS integration, dense pivoted LDLT, dense Cholesky, and dense QR -solvers) behind a common interface. Every backend accepts scalar CSR; all -but cuDSS also accept block BSR directly, which is what -:cpp:func:`SupportsBlockStorage` reports. +`cunls/linear_solver` hosts the sparse linear solvers (supernodal sparse +Cholesky, block-Jacobi PCG, cuDSS integration, block-tridiagonal Cholesky, +dense pivoted LDLT, dense Cholesky and dense QR) behind a common interface. +Every backend accepts scalar CSR; BlockSparsePCG and the dense backends also +accept block BSR directly, which is what :cpp:func:`SupportsBlockStorage` +reports. How to choose a backend, and how BlockSparseCholesky and the Schur +complement work: :doc:`../linear_solvers`. This page documents the C++ API only: ``pycunls`` does not expose the linear solver classes. From Python, the backend is selected through the @@ -19,13 +21,17 @@ SparseLinearSolverType Enum in `cunls/linear_solver/sparse_linear_solver.h`: +- `BlockSparseCholesky` — supernodal sparse Cholesky on the GPU without + external dependencies. Default backend for Gauss-Newton and + Levenberg-Marquardt, and the solver of the reduced system of + ``LinearSystemType::SchurComplement``: exact steps, faster than + ``BlockSparsePCG`` and ``cuDSS`` on pose graphs, bundle adjustment and + visual-inertial problems. Requires an SPD system. - `BlockSparsePCG` — iterative block-Jacobi preconditioned conjugate - gradient solver. Default backend for Gauss-Newton and - Levenberg-Marquardt: outperforms ``cuDSS`` on most SBA / PGO - workloads (see ``profile/PCG_RESULTS.md``). -- `cuDSS` — NVIDIA cuDSS sparse direct solver. Pick when each Solve - sees a tiny system and PCG's per-iter kernel-launch overhead - dominates. + gradient solver (inexact steps). Prefer it when the graph has no small + separators (many long-range random edges, where a direct factor fills + in) and for very large well-conditioned systems of small blocks. +- `cuDSS` — NVIDIA cuDSS sparse direct solver (optional dependency). - `DenseLDLT` — densifies the coefficient matrix and solves with a custom CUDA pivoted LDLT kernel. - `DenseCholesky` — densifies the coefficient matrix and solves with @@ -193,6 +199,49 @@ residual norm is copied to the host, and only once every :returns: Number of PCG iterations consumed by the most recent ``Solve`` call. Useful for convergence diagnostics. +BlockSparseCholeskySolver +------------------------- + +Concrete implementation in `block_sparse_cholesky_solver.h`, the default +backend (``SparseLinearSolverType::BlockSparseCholesky``) and the solver of the +reduced system of ``LinearSystemType::SchurComplement``. A supernodal sparse +Cholesky factorization ``P A Pᵀ = L Lᵀ`` on the GPU without external +dependencies: ``Initialize`` analyzes the pattern once (fill-reducing +ordering, supernodes cut into panels of at most 64 columns, a level +schedule), every ``Solve`` factors the current values and solves, with no +host synchronization unless safety checks are on. The design is described in +:doc:`../linear_solvers`. + +A pivot within ``1e-5`` of its original diagonal entry (relative) is a rank +deficiency, such as a free gauge under Gauss-Newton: its unknown is fixed at +zero, which solves a consistent semidefinite system. A clearly negative or +non-finite pivot is a failure. + +.. cpp:function:: BlockSparseCholeskySolver::BlockSparseCholeskySolver() + +.. cpp:function:: bool BlockSparseCholeskySolver::Initialize(cudaStream_t stream, const Problem& problem, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) + + Symbolic analysis of the pattern of ``spd_matrix`` (both triangles; the + values are not read) and allocation of the factor. ``problem`` is not used. + + :returns: ``false`` on a dimension mismatch. + :throws std::runtime_error: the factor or its update maps would exceed + :math:`2^{31}` entries. + +.. cpp:function:: bool BlockSparseCholeskySolver::Solve(cudaStream_t stream, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) + + Factors ``spd_matrix`` (the pattern of ``Initialize``) and solves for + ``rhs``. + + :returns: ``false`` with safety checks on (the default) when a pivot is + clearly negative or non-finite; without them the result is non-finite + instead. + +.. cpp:function:: void BlockSparseCholeskySolver::DisableSafetyChecks() + + Skips reading back the pivot-failure flag after each factorization (one + stream synchronization). + cuDSSLinearSolverMode --------------------- diff --git a/docs/sphinx/api/minimizer.rst b/docs/sphinx/api/minimizer.rst index f40f845..b3a5629 100644 --- a/docs/sphinx/api/minimizer.rst +++ b/docs/sphinx/api/minimizer.rst @@ -156,14 +156,29 @@ then override individual fields. needs to escalate from ``lambda_min`` to ``lambda_max``, so the cap counts the rejections at full damping. Set to ``0`` to disable this criterion. - **sparse_linear_solver_type** (``SparseLinearSolverType``, default - ``BlockSparsePCG``) — selects the linear-system backend. - ``BlockSparsePCG`` runs block-Jacobi preconditioned conjugate gradient - with the block layout derived automatically from the problem's state - batches. ``cuDSS`` uses NVIDIA's sparse direct solver; ``DenseLDLT`` + ``BlockSparseCholesky``) — selects the linear-system backend. + ``BlockSparseCholesky`` factors the system with the built-in supernodal + sparse Cholesky (exact steps, no external dependencies; requires an SPD + system, so Gauss-Newton needs every gauge fixed). ``BlockSparsePCG`` runs + block-Jacobi preconditioned conjugate gradient with the block layout + derived automatically from the problem's state batches (inexact steps); + prefer it for graphs without small separators (many long-range random + edges) and for very large well-conditioned systems of small blocks. + ``cuDSS`` uses NVIDIA's sparse direct solver; ``DenseLDLT`` converts to dense and factorizes with a custom pivoted LDLT kernel; ``DenseCholesky`` converts to dense and uses cuSOLVER Cholesky (requires SPD); ``DenseQR`` converts to dense and uses cuSOLVER QR - factorization (works for any non-singular matrix). + factorization (works for any non-singular matrix). Not used with + ``LinearSystemType.SchurComplement``. How to choose: :doc:`../linear_solvers`. +- **linear_system** (``LinearSystemType``, default ``Assembled``) — how each + iteration's linear system is formed. ``Assembled`` assembles the full normal + equations and solves them with ``sparse_linear_solver_type``. + ``SchurComplement`` eliminates the state batches declared with + ``Problem.set_eliminated_state_batches`` (the landmarks of a bundle + adjustment), factors the reduced system of the other states with + ``BlockSparseCholesky`` and recovers the eliminated states: the same step, + much cheaper when the eliminated states outnumber the others. See + :doc:`../linear_solvers`. - **reuse_structure** (``bool``, default ``False``) — the problem's structure is unchanged since this minimizer's previous ``minimize`` on it (same batches, connectivity, active and constant counts, partition; only state @@ -285,6 +300,11 @@ parameters for Levenberg-Marquardt. quality above which :math:`\lambda` is decreased. Steps with :math:`\rho \ge \text{threshold}` are considered "very successful" and the solver becomes more Gauss-Newton-like. +- **min_diagonal** (``float``, default ``1e-6``) — floor of the diagonal the + damping scales: :math:`\lambda \max(\operatorname{diag}(H)_i, + \text{min\_diagonal})` (of the column-scaled system with column scaling). + An unobservable unknown (a zero Jacobian column) then gets a positive pivot + and a zero step instead of a singular system. **Example** @@ -562,20 +582,45 @@ Creates an empty problem with no states or factors. batch in mind: ``DenseCholesky`` runs without failure checks by default (``disable_safety_checks``), and a factorization that breaks down on one ill-conditioned subproblem's block corrupts the step of every subproblem; - ``DenseLDLT`` / ``DenseQR`` (pivoted) and ``cuDSS`` keep the blocks - independent. ``BlockSparsePCG`` stops on the residual of the whole batch, so + ``DenseLDLT`` / ``DenseQR`` (pivoted), ``BlockSparseCholesky`` and ``cuDSS`` + keep the blocks independent. ``BlockSparsePCG`` stops on the residual of the whole batch, so a subproblem with a large residual sets the accuracy the others get. +- ``set_eliminated_state_batches(state_batches) -> None`` — declares the + state batches that ``LinearSystemType.SchurComplement`` eliminates (the + landmarks of a bundle adjustment). Every factor may touch at most one + non-constant state of these batches, and they must share one tangent size + of 1 to 3; the states that share factors with them must have one tangent + size of 3, 6 or 9 (checked at ``minimize``; ``ValueError`` otherwise). + Constant states of a listed batch are not eliminated. An empty list clears + the declaration. See :doc:`../linear_solvers`. + .. _py-enums-label: +-------------------------------------------------------------------------------- +``pycunls.LinearSystemType`` +-------------------------------------------------------------------------------- + +Enum used by ``MinimizerOptions.linear_system``: + +- ``LinearSystemType.Assembled`` (default) — the full normal equations, + solved with ``sparse_linear_solver_type``. Works for every problem. +- ``LinearSystemType.SchurComplement`` — eliminates the state batches of + ``Problem.set_eliminated_state_batches`` with a Schur complement and factors + the reduced system with ``BlockSparseCholesky`` (``sparse_linear_solver_type`` + is not used). + -------------------------------------------------------------------------------- ``pycunls.SparseLinearSolverType`` -------------------------------------------------------------------------------- Integer enum selecting the linear-system backend. -- ``SparseLinearSolverType.BlockSparsePCG`` (default) — block-Jacobi - preconditioned conjugate gradient solver. The block layout is derived +- ``SparseLinearSolverType.BlockSparseCholesky`` (default) — supernodal + sparse Cholesky on the GPU, without external dependencies (exact steps; + requires an SPD system). +- ``SparseLinearSolverType.BlockSparsePCG`` — block-Jacobi + preconditioned conjugate gradient solver (inexact steps). The block layout is derived automatically from the problem's state batches at minimizer-initialize time. - ``SparseLinearSolverType.cuDSS`` — sparse direct solver via NVIDIA cuDSS. @@ -666,14 +711,22 @@ criterion is applied per subproblem. (Levenberg-Marquardt: counted at full damping, after the rejections that escalate :math:`\lambda` to ``lambda_max``). Set to 0 to disable. Default: 5. - **sparse_linear_solver_type** [in]: Linear backend; options are - ``BlockSparsePCG`` (block-Jacobi preconditioned conjugate gradient; - layout auto-derived from the problem's state batches), ``cuDSS`` (sparse + ``BlockSparseCholesky`` (built-in supernodal sparse Cholesky; requires an + SPD system), ``BlockSparsePCG`` (block-Jacobi preconditioned conjugate + gradient; layout auto-derived from the problem's state batches), ``cuDSS`` (sparse direct solver via NVIDIA's cuDSS library), ``DenseLDLT`` (converts CSR to dense and solves with a custom CUDA pivoted LDLT factorization), ``DenseCholesky`` (converts CSR to dense and solves via cuSOLVER Cholesky; requires SPD matrix), and ``DenseQR`` (converts CSR to dense and solves via cuSOLVER QR factorization; works for any non-singular matrix). - Default: ``BlockSparsePCG``. + Default: ``BlockSparseCholesky``. Not used with + ``LinearSystemType::SchurComplement``. See :doc:`../linear_solvers`. +- **linear_system** [in]: How each iteration's linear system is formed: + ``LinearSystemType::Assembled`` (the full normal equations, solved with + ``sparse_linear_solver_type``) or ``LinearSystemType::SchurComplement`` (the + states of :cpp:func:`Problem::SetEliminatedStateBatches` eliminated, the + reduced system factored with ``BlockSparseCholesky``). Default: + ``Assembled``. - **sparse_linear_solver_config** [in]: Backend-specific options. For ``BlockSparsePCG`` contains :code:`block_sparse_pcg_options` (``block_size`` / ``block_layout``, ``max_iterations``, @@ -771,6 +824,9 @@ constructing a :code:`LevenbergMarquardtMinimizer`. cost reduction) to accept a step. Default: 0.25. - **lambda_downscale_threshold** [in]: Step quality above which :math:`\lambda` is decreased. Default: 0.75. +- **min_diagonal** [in]: Floor of the diagonal the damping scales, + :math:`\lambda \max(\operatorname{diag}(H)_i, \text{min\_diagonal})`. + Default: 1e-6. .. _ransac-structures-label: @@ -1200,6 +1256,25 @@ minimizers accept, damp and stop each one on its own (on the device; one :param ``state_problem_ids``: [in] One device array per state batch (not owned, read at every solve): the subproblem of each state. +:code:`Problem::SetEliminatedStateBatches` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Purpose:** Declares the state batches that +``LinearSystemType::SchurComplement`` eliminates (the landmarks of a bundle +adjustment). See the Python ``set_eliminated_state_batches`` above for the +requirements, and :doc:`../linear_solvers` for how the elimination works. + +.. cpp:function:: void SetEliminatedStateBatches(const std::vector &batches) + + :param ``batches``: [in] State batches already added to the problem; an + empty vector clears the declaration. + :throws std::invalid_argument: a batch is null, not part of the problem, or + listed twice. + +.. cpp:function:: const std::vector &EliminatedStateBatches() const + + :returns: [out] The declared batches (empty when not set). + ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ :code:`Problem::SetStatePointers` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/sphinx/index.rst b/docs/sphinx/index.rst index 25789ae..c1a2fe5 100644 --- a/docs/sphinx/index.rst +++ b/docs/sphinx/index.rst @@ -48,6 +48,7 @@ Guides constraints dynamics imu + linear_solvers mpc custom_factors_and_states numeric_jacobians diff --git a/docs/sphinx/installation.rst b/docs/sphinx/installation.rst index 6638368..76330fb 100644 --- a/docs/sphinx/installation.rst +++ b/docs/sphinx/installation.rst @@ -128,8 +128,9 @@ Notes - Neither ``libcunls.so``/``libcunls.a`` nor the ``pycunls`` wheel bundle or depend on cuDSS at load time — every other solver (``DenseLDLT``, - ``DenseCholesky``, ``DenseQR``, ``BlockSparsePCG``, ``BlockTridiagonal``) works with no cuDSS - installed at all. + ``DenseCholesky``, ``DenseQR``, ``BlockSparsePCG``, ``BlockTridiagonal``, + ``BlockSparseCholesky``, and the Schur complement linear system) works with + no cuDSS installed at all. - To use the cuDSS solver, download a cuDSS release matching your CUDA major version from `NVIDIA's cuDSS redistributables `_ diff --git a/docs/sphinx/linear_solvers.rst b/docs/sphinx/linear_solvers.rst new file mode 100644 index 0000000..ed49a48 --- /dev/null +++ b/docs/sphinx/linear_solvers.rst @@ -0,0 +1,329 @@ +############################################################################### +Linear systems and solvers +############################################################################### + +Every Gauss-Newton or Levenberg-Marquardt iteration solves the normal +equations of the linearized problem for the step :math:`\delta`: + +.. math:: + + (H + \lambda D)\,\delta = b, \qquad H = J^T J, \quad b = -J^T f, + +with the damping :math:`\lambda D` of Levenberg-Marquardt (zero for +Gauss-Newton), optionally column-scaled (see :doc:`api/minimizer`). Two options +of ``MinimizerOptions`` decide how this system is formed and solved: + +- ``linear_system`` (``LinearSystemType``): ``Assembled`` (default) assembles + the full system; ``SchurComplement`` eliminates the landmark states of a + bundle adjustment first and solves a much smaller reduced system. +- ``sparse_linear_solver_type`` (``SparseLinearSolverType``): the solver of the + assembled system. The default, ``BlockSparseCholesky``, is cuNLS's own + supernodal sparse Cholesky factorization; the Schur complement always uses + it for its reduced system. + +This page explains how to choose, then how the two main pieces work. + +=============================================================================== +Choosing +=============================================================================== + +.. list-table:: + :header-rows: 1 + :widths: 40 22 38 + + * - Problem + - ``linear_system`` + - ``sparse_linear_solver_type`` + * - Bundle adjustment: poses and landmarks, each landmark observed by + several poses + - ``SchurComplement`` + - (not used: ``BlockSparseCholesky``) + * - Visual-inertial bundle adjustment (keyframe poses, velocities, biases + and landmarks) + - ``SchurComplement`` + - (not used) + * - Pose graphs, trajectories, SLAM back ends: graphs whose edges are + mostly local (odometry, loop closures between nearby places) + - ``Assembled`` + - ``BlockSparseCholesky`` (default) + * - Graphs with many long-range random edges (no small separators), and + very large well-conditioned systems of small blocks + - ``Assembled`` + - ``BlockSparsePCG`` + * - Stage-ordered problems (MPC, trajectories with + ``Problem::SetStateStages``) + - ``Assembled`` + - ``BlockTridiagonal`` + * - A single small dense problem (one pose: PnP, registration) + - ``Assembled`` + - ``DenseLDLT`` or ``DenseCholesky`` + +``BlockSparseCholesky`` and the Schur complement compute the exact step (up to +float32 rounding). ``BlockSparsePCG`` stops at a relative residual (``1e-3`` by +default): its steps are inexact, which costs Levenberg-Marquardt iterations on +ill-conditioned problems and can stop it short (bundle adjustment with long +tracks, visual-inertial windows). Its strength is graphs on which any direct +factorization fills in: when many edges connect distant states at random, the +factor of :math:`H` becomes nearly dense, while conjugate gradient only +multiplies by :math:`H`. ``cuDSS`` (NVIDIA's sparse direct solver) is an +optional dependency; ``BlockSparseCholesky`` is faster on the problems above. + +Measured end to end (Levenberg-Marquardt to convergence, RTX A6000, every call +including the symbolic setup): + +.. list-table:: + :header-rows: 1 + :widths: 46 18 18 18 + + * - Problem + - BlockSparseCholesky + - BlockSparsePCG + - cuDSS + * - SE(3) pose graph, 2 500 poses, grid-like loop closures + - **15 ms** + - 75 ms + - 50 ms + * - SE(3) pose graph, 100 000 poses, grid-like loop closures + - **541 ms** + - 592 ms + - 2 124 ms + * - SE(3) pose graph, 5 000 poses, 2 000 random loop closures + - 106 ms + - **52 ms** + - 329 ms + * - 256 independent SE(2) graphs of 12 poses (one solve) + - **4.6 ms** + - 49 ms + - 8.1 ms + * - Bundle adjustment, full system (119 real problems) + - **281 ms** + - 661 ms + - 345 ms + * - Visual-inertial BA, 501 keyframes, full system + - **23 ms** + - 19 ms, 100x the final cost + - 46 ms + * - Visual-inertial BA, 501 keyframes, ``SchurComplement`` + - **14 ms** + - + - + +=============================================================================== +BlockSparseCholesky +=============================================================================== + +``BlockSparseCholeskySolver`` factors :math:`P A P^T = L L^T` (:math:`P` a +fill-reducing permutation) and solves with :math:`L`, entirely on the GPU and +without external dependencies. The matrices of nonlinear least squares are +block structured (one dense block per pair of connected states), which the +solver exploits throughout: it works on supervariables (the columns of one +state) and on dense panels of :math:`L`. + +------------------------------------------------------------------------------- +Setup: symbolic analysis (``Initialize``) +------------------------------------------------------------------------------- + +Run once per sparsity pattern (every ``minimize`` call, unless +``MinimizerOptions::reuse_structure`` vouches for an unchanged structure). +It reads the pattern only: + +1. **Supervariables.** Consecutive columns with identical patterns (a state's + tangent block, or several states always seen together) become one vertex. + They and their graph are found on the GPU; only the graph comes to the + host. +2. **Ordering.** Two fill-reducing orderings of the supervariable graph are + computed and the one with the lower estimated factorization time kept (a + fixed latency per sequential panel step plus the flops): + + - *approximate minimum degree* on the quotient graph (an eliminated vertex + becomes an element standing for the clique of its neighbors, so the fill + is never formed explicitly): the least fill on irregular graphs, such as + covisibility, bundle adjustment and random loop closures; + - *nested dissection* with balanced breadth-first level-set separators and + minimum-degree leaves: short schedules and little fill on chains and + grids, such as trajectories, visual-inertial windows and pose graphs with + local loop closures. + +3. **Supernodes and panels.** The elimination tree gives the fundamental + supernodes (columns with nested structure). A supernode is merged into its + parent when that saves a sequential step and keeps the explicit zeros at + no more than 50%. Supernodes are cut into **panels** of at most 64 + columns: dense column-major blocks holding the panel's diagonal block and + all its rows below. +4. **Schedule.** A panel's level is one more than the highest level among the + panels that update it, so the panels of a level are independent. The work + items of every level, and the maps from the nonzeros of :math:`A` into the + panels and from each panel's rows into the panels it updates, are built + (the maps on the GPU). + +Setup takes a few milliseconds for thousands of states (2-4 ms for a +camera system with 700 poses, 10 ms for a full bundle adjustment with 10 000 +landmarks) and grows to about 60 ms for 40 000 poses. + +------------------------------------------------------------------------------- +Numeric phase (``Solve``) +------------------------------------------------------------------------------- + +Every solve factors the current values and solves; nothing is synchronized +with the host unless safety checks are on: + +1. **Assembly.** The lower triangle of :math:`PAP^T` is scattered into the + panels and the right-hand side is permuted. +2. **Factorization and forward substitution, level by level.** One kernel + launch per level factors all its panels, one CTA per panel and 64 rows + below its diagonal block. The CTA stacks the diagonal block :math:`D`, its + rows :math:`B` and the panel's part of the right-hand side :math:`y^T` in + shared memory and factors them in blocks of 16 columns: one warp factors + the 16 x 16 diagonal block (in groups of 4 columns, each group with two + warp round trips), one thread per row solves the rows below it, and all + threads apply the rank-16 update to the remaining columns. This yields + :math:`L_D`, :math:`B L_D^{-T}` and, from the last row, the panel's step of + the forward substitution :math:`z = L^{-1} b` at no extra cost. A second + launch applies each panel's update :math:`L_B L_B^T` to the panels it + targets, one CTA per 64 x 64 tile; atomics are used only on levels where two + panels update the same target. +3. **Backward substitution**, levels in reverse: each CTA multiplies its rows' + :math:`L_B^T` by the solution below, and the last CTA of a panel solves the + diagonal block. + +The cost of a level is mostly latency for small problems (each level is two +dependent launches) and throughput for large ones, so the number of levels +matters as much as the flops; this is what the cost model of the ordering +weighs. + +------------------------------------------------------------------------------- +Rank deficiency and failures +------------------------------------------------------------------------------- + +A pivot within :math:`10^{-5}` of its original diagonal entry (relative, either +sign) is treated as a rank deficiency rather than a failure: its unknown is +fixed at zero. This solves consistent semidefinite systems, such as +Gauss-Newton on a pose graph without an anchored pose (a free gauge), the +way conjugate gradient does. A clearly negative or non-finite pivot is a +failure: with safety checks on (``disable_safety_checks = false``) the solve +returns an error, otherwise the step becomes non-finite and is rejected. +Levenberg-Marquardt's damping keeps the system positive definite. + +------------------------------------------------------------------------------- +Limits +------------------------------------------------------------------------------- + +- The factor and its update maps are limited to :math:`2^{31}` entries each + (``std::runtime_error`` at setup). Graphs with many long-range random edges + reach it first: their factors are nearly dense. Use ``BlockSparsePCG`` there. +- Scalar CSR input only (no block storage): the assembly of the normal + equations is scalar for this backend, at no measurable cost. +- Small, latency-bound problems (a few hundred poses) solve in about 0.1 ms + per factorization; when the structure is the same at every call, set + ``reuse_structure`` to skip the setup. + +=============================================================================== +Schur complement +=============================================================================== + +In a bundle adjustment every factor connects one landmark with one or a few +poses, and landmarks outnumber poses by orders of magnitude. Ordering the +unknowns as retained states :math:`r` (poses, and any other states) and +eliminated landmarks :math:`e`, the normal equations read + +.. math:: + + \begin{bmatrix} C & W \\ W^T & D \end{bmatrix} + \begin{bmatrix} \delta_r \\ \delta_e \end{bmatrix} = + \begin{bmatrix} b_r \\ b_e \end{bmatrix}, + +where :math:`D` is **block diagonal**: no factor touches two landmarks, so each +landmark has its own small block :math:`D_i` (3 x 3 for a point). Eliminating +the landmarks gives the **reduced system** of the retained states + +.. math:: + + S\,\delta_r = g, \qquad S = C - W D^{-1} W^T, \qquad g = b_r - W D^{-1} b_e, + +and each landmark's step follows from its own block: + +.. math:: + + \delta_{e,i} = D_i^{-1} \big(b_{e,i} - W_i^T \delta_r\big). + +The step is the exact solution of the full system. :math:`S` has one block per +pair of poses that observe a common landmark: dense when every camera sees +most of the scene, banded for a visual-inertial trajectory, and in any case +far smaller than the full system. + +------------------------------------------------------------------------------- +Usage +------------------------------------------------------------------------------- + +Declare the landmark state batches as eliminated, and select the Schur +complement: + +.. code-block:: cpp + + problem.SetEliminatedStateBatches({&point_states}); + cunls::LevenbergMarquardtMinimizerOptions options; + options.base_options.linear_system = cunls::LinearSystemType::SchurComplement; + cunls::LevenbergMarquardtMinimizer minimizer(options); + minimizer.Minimize(stream, problem); + +.. code-block:: python + + problem.set_eliminated_state_batches([point_states]) + options = pycunls.LevenbergMarquardtMinimizerOptions() + options.base_options.linear_system = pycunls.LinearSystemType.SchurComplement + pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) + +Complete programs: ``examples/sparse_bundle_adjustment`` (C++) and +``python/examples/sparse_bundle_adjustment.py``; visual-inertial bundle +adjustment: ``python/examples/imu_bundle_adjustment.py``. + +Requirements, checked when the minimizer sets up the structure +(``std::invalid_argument`` / ``ValueError`` otherwise): + +- every factor touches at most one non-constant state of the eliminated + batches (constant states are not eliminated); +- the eliminated states share one tangent size of 1 to 3; +- the retained states that share factors with them share one tangent size of + 6 (poses), 3 or 9. Other factors, between retained states only (IMU, + priors, odometry), may have any shape. + +------------------------------------------------------------------------------- +How it works +------------------------------------------------------------------------------- + +The full Hessian is never formed. The setup (once per structure, on the GPU) +numbers the landmarks and the retained states, finds every (retained state, +landmark) pair that shares a factor, orders the pairs by retained state +("pose-major"), lays the landmark factors' Jacobians out in that order, and +builds the block pattern of :math:`S` and the maps that write into it. + +Each iteration, with the column scaling, damping and mask of the minimizer +applied to the full system before the elimination: + +1. **Assembly** of :math:`b`, :math:`\operatorname{diag}(H)` and the landmark + blocks :math:`D_i`, in one pass over the factors. +2. **Pair pass** (one thread per pair, pose-major): the pair's Jacobians are + contracted into :math:`W_{pi}` and :math:`A_{pi} = J_p^T J_p`, the + landmark block is factored, :math:`D_i = L_i L_i^T`, and + :math:`V_{pi} = W_{pi} L_i^{-T}` and :math:`h_i = L_i^{-1} b_{e,i}` are + formed. The diagonal blocks of :math:`S` receive + :math:`A_{pi} - V_{pi} V_{pi}^T`, written per pair so no large term + cancels another, and :math:`g` receives :math:`-V_{pi} h_i`. Warp-segmented + sums leave one atomic per retained state and warp. +3. **Off-diagonal pass** (one warp per landmark track): the blocks + :math:`-V_{pi} V_{qi}^T` for every pair of poses observing landmark + :math:`i`. +4. **Reduced solve**: :math:`S` (expanded to scalar CSR) is factored by + ``BlockSparseCholesky``. +5. **Back-substitution**, one thread per landmark: + :math:`\delta_{e,i} = L_i^{-T}\big(h_i - \sum_p V_{pi}^T \delta_{r,p}\big)`. + +A landmark block that is not positive definite (for example a point seen +along a single ray) is regularized with a small diagonal shift. With safety +checks on, a block that still cannot be factored makes the solve fail; +without them its landmark's step is non-finite and the minimizer rejects the +step. + +Per solve of the reduced system (RTX A6000): 1.6 ms for a 723-pose cuSFM +bundle adjustment (dense covisibility, :math:`S` of size 4 338), 0.6 ms for a +2 001-keyframe visual-inertial window (:math:`S` of size 30 009, banded). diff --git a/examples/pose_graph_optimization/README.md b/examples/pose_graph_optimization/README.md index 94be041..ec7b02c 100644 --- a/examples/pose_graph_optimization/README.md +++ b/examples/pose_graph_optimization/README.md @@ -62,6 +62,17 @@ active count equals its capacity. `LevenbergMarquardtMinimizer` allocates its working buffers during initialization before the first linearized solve. +### Linear solver + +Each LM step solves the normal equations of the pose graph (6 unknowns per +free pose). The example sets `sparse_linear_solver_type` to +`BlockSparseCholesky`, the default: cuNLS's supernodal sparse Cholesky, which +orders the graph to keep the factor sparse (a chain with local loop closures +stays almost banded) and gives the exact step. `BlockSparsePCG` (block-Jacobi +preconditioned conjugate gradient) is the iterative alternative, preferable +when many loop closures connect distant poses at random, which makes any +factor dense. + ## Validation After optimization, the example prints: diff --git a/examples/pose_graph_optimization/main.cpp b/examples/pose_graph_optimization/main.cpp index 2f7968d..80467db 100644 --- a/examples/pose_graph_optimization/main.cpp +++ b/examples/pose_graph_optimization/main.cpp @@ -92,10 +92,15 @@ int main() { } // 6. Solve with Levenberg-Marquardt; the result is written into poses. + // Each step factors the pose graph's normal equations with + // BlockSparseCholesky (the default; exact steps). BlockSparsePCG is the + // iterative alternative for graphs with many long-range loop closures. cunls::LevenbergMarquardtMinimizerOptions options; options.base_options.max_num_iterations = 60; options.base_options.state_tolerance = 1e-8f; options.base_options.cost_tolerance = 1e-8f; + options.base_options.sparse_linear_solver_type = + cunls::SparseLinearSolverType::BlockSparseCholesky; options.initial_lambda = 1e-3f; cunls::LevenbergMarquardtMinimizer minimizer(options); cunls::CudaStream stream; diff --git a/examples/sparse_bundle_adjustment/README.md b/examples/sparse_bundle_adjustment/README.md index 2a837f2..a5a2866 100644 --- a/examples/sparse_bundle_adjustment/README.md +++ b/examples/sparse_bundle_adjustment/README.md @@ -4,7 +4,8 @@ This example demonstrates synthetic sparse bundle adjustment with: - `ReprojectionFactorBatch` for reprojection residuals - `SE3StateBatch` for camera poses (first two fixed, rest optimized) - `VectorStateBatch<3>` for 3D points (all optimized) -- `LevenbergMarquardtMinimizer` for nonlinear optimization +- `LevenbergMarquardtMinimizer` for nonlinear optimization, with the points + eliminated by a Schur complement (`LinearSystemType::SchurComplement`) The setup is intentionally compact but complete: 1. Generate synthetic camera poses in SE(3). @@ -68,6 +69,20 @@ active count equals its capacity. `LevenbergMarquardtMinimizer::Minimize` runs `Initialize`, which allocates minimizer working buffers and prepares state/factor batches for the solve. +### Linear system: Schur complement + +Every LM iteration solves a linear system for the step of all poses and +points. The points are declared eliminated with +`problem.SetEliminatedStateBatches({&point_states})` and the minimizer is set +to `LinearSystemType::SchurComplement`: each point only connects to poses, +so its 3x3 block is eliminated in closed form, the reduced camera system (6 +unknowns per free pose) is solved with cuNLS's supernodal sparse Cholesky +(`BlockSparseCholesky`; `sparse_linear_solver_type` is not used), and the +points' steps are recovered from it. The step is the +exact solution of the full system, at a fraction of the cost. Without the +declaration (`LinearSystemType::Assembled`, the default) the full system is +assembled and solved by `sparse_linear_solver_type` directly. + ### Optimization and checks The example runs LM and prints: diff --git a/examples/sparse_bundle_adjustment/main.cpp b/examples/sparse_bundle_adjustment/main.cpp index 614119c..79e60eb 100644 --- a/examples/sparse_bundle_adjustment/main.cpp +++ b/examples/sparse_bundle_adjustment/main.cpp @@ -104,10 +104,17 @@ int main() { } // 6. Solve with Levenberg-Marquardt; the result is written into poses / points. + // The points are eliminated with a Schur complement: each point only + // connects to poses, so its 3x3 block is eliminated in closed form, the + // reduced camera system (6 unknowns per free pose) is factored by + // BlockSparseCholesky, and every point is recovered from it. The step is + // the exact solution of the full system, much cheaper to compute. + problem.SetEliminatedStateBatches({&point_states}); cunls::LevenbergMarquardtMinimizerOptions options; options.base_options.max_num_iterations = 80; options.base_options.state_tolerance = 1e-8f; options.base_options.cost_tolerance = 1e-8f; + options.base_options.linear_system = cunls::LinearSystemType::SchurComplement; options.initial_lambda = 1e-3f; cunls::LevenbergMarquardtMinimizer minimizer(options); cunls::CudaStream stream; diff --git a/llms.txt b/llms.txt index 141ac67..b1a8be2 100644 --- a/llms.txt +++ b/llms.txt @@ -21,8 +21,9 @@ cuNLS organizes optimization into: - `pycunls.mpc` for model predictive control (dynamics factors, costs, constraints, receding horizon); guide: `docs/sphinx/mpc.rst`, example: `python/examples/supermarket_drones.py` (batched fleet with moving obstacles). -- Sparse backends (`BlockSparsePCG` default, plus `cuDSS` direct and - dense LDLT/Cholesky/QR fallbacks) and robust losses. +- Sparse backends (`BlockSparseCholesky` default, plus `BlockSparsePCG` + iterative, optional `cuDSS` direct and dense LDLT/Cholesky/QR fallbacks) + and robust losses. ## Entry points diff --git a/python/examples/README.md b/python/examples/README.md index 9430a66..9391897 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -6,12 +6,12 @@ visualization). | Example | What it shows | |---|---| -| `sparse_bundle_adjustment.py` | Joint camera-pose and landmark optimization with CuPy | +| `sparse_bundle_adjustment.py` | Joint camera-pose and landmark optimization with CuPy, landmarks eliminated by a Schur complement | | `pose_graph_optimization.py` | SE(3) pose-graph optimization with CuPy | | `custom_warp_factor.py` | A custom factor kernel written with NVIDIA Warp | | `custom_warp_state.py` | A custom state batch (positive-scalar manifold) written with NVIDIA Warp | | `ransac_pnp.py` | Robust PnP with 50% outliers using `RansacLevenbergMarquardtMinimizer` | -| `imu_bundle_adjustment.py` | Visual-inertial bundle adjustment with `ImuFactorBatch` | +| `imu_bundle_adjustment.py` | Visual-inertial bundle adjustment with `ImuFactorBatch`, landmarks eliminated by a Schur complement | | [`tartan_vio.py`](#visual-inertial-odometry-with-ransac-on-tartanground) | RGB-D inertial odometry on TartanGround with injected outliers: IMU factor + RANSAC | | [`supermarket_drones.py`](#a-drone-fleet-in-a-supermarket-batched-mpc) | A fleet of quadrotors in a real store with walking shoppers: batched MPC | diff --git a/python/examples/imu_bundle_adjustment.py b/python/examples/imu_bundle_adjustment.py index cf73779..ee03e34 100644 --- a/python/examples/imu_bundle_adjustment.py +++ b/python/examples/imu_bundle_adjustment.py @@ -20,7 +20,9 @@ landmarks through ReprojectionFactorBatch. Both factor types read the same world_from_rig pose states (the rig's pose in the world). Velocities and IMU biases start at zero and are recovered from the solve; poses and landmarks start perturbed. The data are -noiseless; the solve recovers the trajectory to about a millimeter. +noiseless; the solve recovers the trajectory to about a millimeter. The +landmarks are eliminated with a Schur complement, whose reduced system is +solved with the built-in supernodal Cholesky. Usage:: @@ -182,8 +184,14 @@ def main(): problem.add_factor_batch(imu, imu_ptrs) problem.add_factor_batch(reproj, reproj_ptrs) + # The landmarks are eliminated with a Schur complement: each iteration + # factors only the reduced system of the keyframe states (pose, velocity and + # bias, 15 unknowns per keyframe; banded along the trajectory) with + # BlockSparseCholesky and recovers every landmark from it. + problem.set_eliminated_state_batches([points]) options = pycunls.LevenbergMarquardtMinimizerOptions() options.base_options.max_num_iterations = 30 + options.base_options.linear_system = pycunls.LinearSystemType.SchurComplement stream = pycunls.CudaStream() summary = pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) diff --git a/python/examples/pose_graph_optimization.py b/python/examples/pose_graph_optimization.py index 4a7a3c9..55a17a3 100644 --- a/python/examples/pose_graph_optimization.py +++ b/python/examples/pose_graph_optimization.py @@ -68,11 +68,16 @@ def main(): problem.add_factor_batch(between, state_pointers) assert problem.check_consistency(), "Problem consistency check failed" - # 5. Solve with Levenberg-Marquardt. + # 5. Solve with Levenberg-Marquardt. Each step factors the pose graph's + # normal equations with BlockSparseCholesky (the default; exact steps). + # BlockSparsePCG is the iterative alternative for graphs with many + # long-range loop closures. options = pycunls.LevenbergMarquardtMinimizerOptions() options.base_options.max_num_iterations = 60 options.base_options.state_tolerance = 1e-8 options.base_options.cost_tolerance = 1e-8 + options.base_options.sparse_linear_solver_type = ( + pycunls.SparseLinearSolverType.BlockSparseCholesky) options.initial_lambda = 1e-3 stream = pycunls.CudaStream() summary = pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) diff --git a/python/examples/sparse_bundle_adjustment.py b/python/examples/sparse_bundle_adjustment.py index fcf995e..bb17763 100644 --- a/python/examples/sparse_bundle_adjustment.py +++ b/python/examples/sparse_bundle_adjustment.py @@ -20,6 +20,9 @@ factors (normalized image coordinates). Poses are the cameras' world poses (world_from_camera, the pose convention of cuNLS). Poses 0 and 1 are held constant to fix the gauge: the global rigid-body transform and the scale. + +The points are eliminated with a Schur complement, whose reduced system is +solved with the built-in supernodal Cholesky. """ import cupy as cp @@ -80,11 +83,17 @@ def main(): problem.add_factor_batch(reprojection, state_pointers) assert problem.check_consistency(), "Problem consistency check failed" - # 6. Solve with Levenberg-Marquardt. + # 6. Solve with Levenberg-Marquardt. The points are eliminated with a Schur + # complement: each point only connects to poses, so its 3x3 block is + # eliminated in closed form, the reduced camera system (6 unknowns per free + # pose) is factored by BlockSparseCholesky, and every point is recovered + # from it. The step is the exact solution of the full system, much cheaper. + problem.set_eliminated_state_batches([point_states]) options = pycunls.LevenbergMarquardtMinimizerOptions() options.base_options.max_num_iterations = 80 options.base_options.state_tolerance = 1e-8 options.base_options.cost_tolerance = 1e-8 + options.base_options.linear_system = pycunls.LinearSystemType.SchurComplement options.initial_lambda = 1e-3 stream = pycunls.CudaStream() summary = pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) diff --git a/python/pycunls/__init__.py b/python/pycunls/__init__.py index 19be086..f02f77e 100644 --- a/python/pycunls/__init__.py +++ b/python/pycunls/__init__.py @@ -56,6 +56,7 @@ Verbosity, set_log_verbosity, # --- Enumerations --- + LinearSystemType, SparseLinearSolverType, ColumnScaling, JacobianMode, @@ -204,6 +205,7 @@ "CudaStream", "Verbosity", "set_log_verbosity", + "LinearSystemType", "SparseLinearSolverType", "ColumnScaling", "JacobianMode", diff --git a/python/pycunls/_pycunls_core.pyi b/python/pycunls/_pycunls_core.pyi index 45bd50b..d4cc340 100644 --- a/python/pycunls/_pycunls_core.pyi +++ b/python/pycunls/_pycunls_core.pyi @@ -60,13 +60,26 @@ def set_log_verbosity(verbosity: Verbosity) -> None: # Enumerations # =================================================================== +class LinearSystemType(enum.Enum): + """How each iteration's linear system is formed (MinimizerOptions.linear_system).""" + + Assembled = ... + """The full normal equations, solved with sparse_linear_solver_type.""" + SchurComplement = ... + """Eliminate the state batches of Problem.set_eliminated_state_batches (the landmarks of a + bundle adjustment) and solve the reduced system with BlockSparseCholesky.""" + class SparseLinearSolverType(enum.Enum): + """Solver of the assembled system (MinimizerOptions.sparse_linear_solver_type; not used with + LinearSystemType.SchurComplement). Default: BlockSparseCholesky.""" + cuDSS = ... DenseLDLT = ... DenseCholesky = ... DenseQR = ... BlockSparsePCG = ... BlockTridiagonal = ... + BlockSparseCholesky = ... class ColumnScaling(enum.Enum): """Diagonal scaling mode for the GN/LM normal equations.""" @@ -111,6 +124,7 @@ class MinimizerOptions: max_line_search_steps: int reuse_structure: bool sparse_linear_solver_type: SparseLinearSolverType + linear_system: LinearSystemType column_scaling: ColumnScaling jacobian_mode: JacobianMode numeric_diff_options: NumericDiffOptions @@ -141,6 +155,7 @@ class LevenbergMarquardtMinimizerOptions: lambda_downscale: float lambda_max: float lambda_min: float + min_diagonal: float step_accept_threshold: float lambda_downscale_threshold: float @@ -1689,6 +1704,12 @@ class Problem: states of the same or adjacent stages within a subproblem. An empty list clears the stages.""" ... + def set_eliminated_state_batches(self, state_batches: list[StateBatch]) -> None: + """Declare the state batches LinearSystemType.SchurComplement eliminates + (the landmarks of a bundle adjustment). Every factor may touch at most + one non-constant state of these batches; they must share one tangent + size of 1 to 3. An empty list clears the declaration.""" + ... @property def num_problems(self) -> int: ... diff --git a/python/src/bind_problem.cpp b/python/src/bind_problem.cpp index 7661d54..b2a5b3f 100644 --- a/python/src/bind_problem.cpp +++ b/python/src/bind_problem.cpp @@ -187,5 +187,16 @@ void bind_problem(nb::module_ &m) { "s of state batch b belongs to stage state_stages[b][s] (device int32 arrays, one per " "state batch). Factors may only connect states of the same or adjacent stages within " "a subproblem. An empty list clears the stages.") + .def( + "set_eliminated_state_batches", + [](cunls::Problem &self, const std::vector &batches) { + self.SetEliminatedStateBatches( + std::vector(batches.begin(), batches.end())); + }, + nb::arg("state_batches"), + "Declare the state batches LinearSystemType.SchurComplement eliminates (the " + "landmarks of a bundle adjustment). Every factor may touch at most one non-constant " + "state of these batches; they must share one tangent size of 1 to 3. An empty list " + "clears the declaration.") .def_prop_ro("num_problems", &cunls::Problem::NumProblems); } diff --git a/python/src/bind_types.cpp b/python/src/bind_types.cpp index 0abdd43..6400482 100644 --- a/python/src/bind_types.cpp +++ b/python/src/bind_types.cpp @@ -88,13 +88,36 @@ void bind_types(nb::module_ &m) { // --- Enumerations for solver/multiplier strategy selection --- - nb::enum_(m, "SparseLinearSolverType") - .value("cuDSS", cunls::SparseLinearSolverType::cuDSS) - .value("DenseLDLT", cunls::SparseLinearSolverType::DenseLDLT) - .value("DenseCholesky", cunls::SparseLinearSolverType::DenseCholesky) - .value("DenseQR", cunls::SparseLinearSolverType::DenseQR) - .value("BlockSparsePCG", cunls::SparseLinearSolverType::BlockSparsePCG) - .value("BlockTridiagonal", cunls::SparseLinearSolverType::BlockTridiagonal); + nb::enum_( + m, "SparseLinearSolverType", + "Solver of the assembled system (MinimizerOptions.sparse_linear_solver_type; not used " + "with LinearSystemType.SchurComplement). Default: BlockSparseCholesky.") + .value("cuDSS", cunls::SparseLinearSolverType::cuDSS, + "NVIDIA cuDSS sparse direct solver (optional dependency, loaded at run time).") + .value("DenseLDLT", cunls::SparseLinearSolverType::DenseLDLT, + "Dense pivoted LDLT (small systems).") + .value("DenseCholesky", cunls::SparseLinearSolverType::DenseCholesky, + "Dense cuSOLVER Cholesky (small SPD systems).") + .value("DenseQR", cunls::SparseLinearSolverType::DenseQR, + "Dense cuSOLVER QR (small non-singular systems).") + .value("BlockSparsePCG", cunls::SparseLinearSolverType::BlockSparsePCG, + "Block-Jacobi preconditioned conjugate gradient (inexact steps); for graphs with " + "many long-range random edges.") + .value("BlockTridiagonal", cunls::SparseLinearSolverType::BlockTridiagonal, + "Block-tridiagonal Cholesky of stage-ordered problems (Problem.set_state_stages).") + .value("BlockSparseCholesky", cunls::SparseLinearSolverType::BlockSparseCholesky, + "Supernodal sparse Cholesky on the GPU (default; exact steps, no dependencies)."); + + nb::enum_(m, "LinearSystemType", + "How each iteration's linear system is formed " + "(MinimizerOptions.linear_system).") + .value("Assembled", cunls::LinearSystemType::Assembled, + "The full normal equations, solved with sparse_linear_solver_type.") + .value("SchurComplement", cunls::LinearSystemType::SchurComplement, + "Eliminate the state batches of Problem.set_eliminated_state_batches (the " + "landmarks of a bundle adjustment) and solve the reduced system with the " + "built-in supernodal Cholesky (BlockSparseCholesky; sparse_linear_solver_type " + "is not used)."); nb::enum_(m, "ColumnScaling") .value("none", cunls::ColumnScaling::None) @@ -151,6 +174,9 @@ void bind_types(nb::module_ &m) { "calls skip the structure setup. A size change falls back to the full setup. " "Default: False.") .def_rw("sparse_linear_solver_type", &cunls::MinimizerOptions::sparse_linear_solver_type) + .def_rw("linear_system", &cunls::MinimizerOptions::linear_system, + "How each iteration's linear system is formed: LinearSystemType.Assembled " + "(default) or SchurComplement (needs Problem.set_eliminated_state_batches).") .def_rw("column_scaling", &cunls::MinimizerOptions::column_scaling) .def_rw("jacobian_mode", &cunls::MinimizerOptions::jacobian_mode, "Global default JacobianMode for every factor batch, unless " @@ -196,6 +222,10 @@ void bind_types(nb::module_ &m) { .def_rw("lambda_downscale", &cunls::LevenbergMarquardtMinimizerOptions::lambda_downscale) .def_rw("lambda_max", &cunls::LevenbergMarquardtMinimizerOptions::lambda_max) .def_rw("lambda_min", &cunls::LevenbergMarquardtMinimizerOptions::lambda_min) + .def_rw("min_diagonal", &cunls::LevenbergMarquardtMinimizerOptions::min_diagonal, + "Floor of the diagonal the damping scales, lambda * max(diag(H), min_diagonal), " + "so unobservable unknowns (zero Jacobian columns) get a zero step instead of a " + "singular system. Default: 1e-6.") .def_rw("step_accept_threshold", &cunls::LevenbergMarquardtMinimizerOptions::step_accept_threshold) .def_rw("lambda_downscale_threshold", diff --git a/python/tests/test_minimizer.py b/python/tests/test_minimizer.py index 223dc03..9126adc 100644 --- a/python/tests/test_minimizer.py +++ b/python/tests/test_minimizer.py @@ -170,7 +170,7 @@ def test_defaults(self): assert opts.state_tolerance == pytest.approx(1e-6) assert opts.cost_tolerance == pytest.approx(1e-6) assert (opts.sparse_linear_solver_type - == pycunls.SparseLinearSolverType.BlockSparsePCG) + == pycunls.SparseLinearSolverType.BlockSparseCholesky) assert opts.column_scaling == pycunls.ColumnScaling.none assert opts.jacobian_mode == pycunls.JacobianMode.analytic assert opts.numeric_diff_options.method == pycunls.NumericDiffMethod.central @@ -240,3 +240,79 @@ def test_per_group_override_converges(self, stream): assert summary.final_cost < 1e-3 result = cp.asnumpy(states_gpu) np.testing.assert_allclose(result, target, atol=1e-2) + + +# --------------------------------------------------------------------------- +# Schur complement linear system +# --------------------------------------------------------------------------- + +def _bundle(num_poses=6, num_points=200, seed=0): + """Cameras along x looking down +z (world_from_camera, 4x4 row-major) at a + point cloud, one reprojection factor per (pose, point); poses 0 and 1 are + constant (gauge). Returns the problem, its point batch and the buffers.""" + rng = np.random.default_rng(seed) + poses = np.tile(np.eye(4, dtype=np.float32), (num_poses, 1, 1)) + poses[:, 0, 3] = np.linspace(-1.0, 1.0, num_poses) + points = np.column_stack([rng.uniform(-2, 2, num_points), rng.uniform(-2, 2, num_points), + rng.uniform(4, 8, num_points)]).astype(np.float32) + observations = [] + for k in range(num_poses): + cam = points - poses[k, :3, 3] + observations.append(cam[:, :2] / cam[:, 2:3]) + observations = np.concatenate(observations).astype(np.float32) + initial_poses = poses.copy() + initial_poses[2:, :3, 3] += rng.normal(0, 0.05, (num_poses - 2, 3)).astype(np.float32) + initial_points = points + rng.normal(0, 0.1, points.shape).astype(np.float32) + + buffers = { + "poses": cp.asarray(initial_poses.reshape(-1)), + "points": cp.asarray(initial_points.reshape(-1)), + "observations": cp.asarray(observations.reshape(-1)), + "const_ids": cp.array([0, 1], dtype=cp.int32), + } + pose_states = pycunls.SE3StateBatch(buffers["poses"], num_poses, buffers["const_ids"], 2) + pose_states.set_num_active_states(num_poses, 2) + point_states = pycunls.VectorStateBatch3(buffers["points"], num_points) + point_states.set_num_active_states(num_points) + reprojection = pycunls.ReprojectionFactorBatch(buffers["observations"], + num_poses * num_points, 1e-3) + reprojection.set_num_active_factors(num_poses * num_points) + pointers = [] + for k in range(num_poses): + for i in range(num_points): + pointers += [pose_states.state_device_ptr(k), point_states.state_device_ptr(i)] + problem = pycunls.Problem() + problem.add_state_batch(pose_states) + problem.add_state_batch(point_states) + problem.add_factor_batch(reprojection, pointers) + buffers["batches"] = (pose_states, point_states, reprojection) + return problem, point_states, buffers + + +class TestSchurComplement: + """LinearSystemType.SchurComplement reaches the assembled system's minimum.""" + + def test_matches_assembled(self, stream): + costs = [] + for system in (pycunls.LinearSystemType.Assembled, + pycunls.LinearSystemType.SchurComplement): + problem, points, _buffers = _bundle() + problem.set_eliminated_state_batches([points]) + options = pycunls.LevenbergMarquardtMinimizerOptions() + options.base_options.max_num_iterations = 30 + options.base_options.linear_system = system + # The exact reference for Assembled; SchurComplement has its own solver. + options.base_options.sparse_linear_solver_type = ( + pycunls.SparseLinearSolverType.DenseCholesky) + summary = pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) + cp.cuda.runtime.streamSynchronize(stream.get_stream()) + assert summary.final_cost < 1e-3 * summary.initial_cost + costs.append(summary.final_cost) + assert costs[1] == pytest.approx(costs[0], rel=1e-2, abs=1e-6) + + def test_requires_eliminated_batches(self, stream): + problem, _points, _buffers = _bundle() + options = pycunls.LevenbergMarquardtMinimizerOptions() + options.base_options.linear_system = pycunls.LinearSystemType.SchurComplement + with pytest.raises(ValueError, match="SetEliminatedStateBatches"): + pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) diff --git a/tests/block_hessian_assembler_test.cpp b/tests/block_hessian_assembler_test.cpp index 7c5e4ae..615883c 100644 --- a/tests/block_hessian_assembler_test.cpp +++ b/tests/block_hessian_assembler_test.cpp @@ -130,15 +130,18 @@ class SystemBuilder { internal::SetUpStructure(stream, problem, options_, state_ops_, system_); state_.Recreate(stream, problem); internal::BuildSystem(stream, problem, state_, options_, system_, scratch_); + Equations().BuildWorkingLhs(stream); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); } + /** @brief The working left-hand side (scaled, damped as declared), as CSR. */ const CSRSparseMatrix &HessianAsCSR(cudaStream_t stream) { - if (!system_.normal_equations.UsesBlockStorage()) { - return system_.normal_equations.LhsCSR(); + Equations().BuildWorkingLhs(stream); + if (!Equations().UsesBlockStorage()) { + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + return Equations().LhsCSR(); } - test_utils::ExpandBSRToCSR(stream, system_.normal_equations.LhsBSR(), csr_mirror_, - expand_scratch_); + test_utils::ExpandBSRToCSR(stream, Equations().LhsBSR(), csr_mirror_, expand_scratch_); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); return csr_mirror_; } @@ -153,8 +156,7 @@ class SystemBuilder { */ float WeightedSquaredStep(cudaStream_t stream, const dvector &step) { dvector h_step; - system_.normal_equations.MultiplyHessian(stream, cusparse_handle_.GetHandle(stream), step, - h_step, buffer_); + Equations().MultiplyHessian(stream, step, h_step); std::vector a(step.size()), b(step.size()); step.CopyToHost(a.data(), a.size()); h_step.CopyToHost(b.data(), b.size()); @@ -162,12 +164,14 @@ class SystemBuilder { for (size_t i = 0; i < a.size(); ++i) out += static_cast(a[i]) * b[i]; return static_cast(out); } - NormalEquations &Equations() { return system_.normal_equations; } - bool UsesBlockStorage() const { return system_.normal_equations.UsesBlockStorage(); } + NormalEquations &Equations() { return static_cast(*system_.linear_system); } + bool UsesBlockStorage() const { + return static_cast(*system_.linear_system).UsesBlockStorage(); + } private: SystemBuilder(const MinimizerOptions &options, SparseLinearSolverPtr solver) : options_(options) { - system_.solver = std::move(solver); + system_.linear_system = std::make_unique(std::move(solver)); } static MinimizerOptions MakeOptions(SparseLinearSolverType solver) { @@ -181,8 +185,6 @@ class SystemBuilder { internal::MinimizerSystem system_; internal::MinimizerScratch scratch_; MinimizerState state_; - cuSPARSEHandle cusparse_handle_; - dvector buffer_; CSRSparseMatrix csr_mirror_; dvector expand_scratch_; }; @@ -956,8 +958,8 @@ TEST(HessianStorageTest, BlockDiagonalOpsMatchScalar) { // 1. Diagonal extraction. dvector csr_diag; dvector bsr_diag; - scalar.Equations().ExtractLhsDiagonal(s, csr_diag); - block.Equations().ExtractLhsDiagonal(s, bsr_diag); + scalar.Equations().ExtractScaledHessianDiagonal(s, csr_diag); + block.Equations().ExtractScaledHessianDiagonal(s, bsr_diag); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); std::vector expected_diag(csr_diag.size()); @@ -970,11 +972,13 @@ TEST(HessianStorageTest, BlockDiagonalOpsMatchScalar) { ASSERT_NEAR(expected_diag[i], actual_diag[i], diag_tol) << "diagonal mismatch at " << i; } - // 2. Damped update, in place, exactly as LM performs it. + // 2. Damped system, exactly as LM declares it. constexpr float kLambda = 0.017f; - scalar.Equations().AddScaledDiagonalToLhs(s, kLambda, csr_diag); - block.Equations().AddScaledDiagonalToLhs(s, kLambda, bsr_diag); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + for (float &v : expected_diag) v *= kLambda; + dvector damping(expected_diag.size()); + damping.CopyFromHost(expected_diag.data(), expected_diag.size()); + scalar.Equations().SetDamping(&damping); + block.Equations().SetDamping(&damping); SystemSnapshot expected = Snapshot(scalar, s); SystemSnapshot actual = Snapshot(block, s); @@ -1278,8 +1282,9 @@ void ExpectDenseBackendAgreesAcrossLayouts(SparseLinearSolver &solver, Problem & // and Cholesky rightly refuses it. Damping is applied to the block LHS // before the scalar copy is expanded from it, so both layouts hold the same // matrix and the comparison stays about the scatter. - dvector ones(1.f, n); - builder.Equations().AddScaledDiagonalToLhs(s, 1e-2f, ones); + dvector damping(1e-2f, n); + builder.Equations().SetDamping(&damping); + builder.Equations().BuildWorkingLhs(s); const BSRSparseMatrix &bsr = builder.Equations().LhsBSR(); CSRSparseMatrix csr; @@ -1440,6 +1445,7 @@ TEST(BlockHessianAssemblerTest, MatchesReferenceWithColumnScaling) { CudaStream stream; MinimizerOptions block_options; block_options.column_scaling = ColumnScaling::HessianDiagonal; + block_options.sparse_linear_solver_type = SparseLinearSolverType::BlockSparsePCG; // BSR backend // Column scaling reads the Hessian diagonal and rescales the LHS in place, and // both of those are implemented once per layout. The reference therefore has diff --git a/tests/block_sparse_cholesky_solver_test.cpp b/tests/block_sparse_cholesky_solver_test.cpp new file mode 100644 index 0000000..06d18ef --- /dev/null +++ b/tests/block_sparse_cholesky_solver_test.cpp @@ -0,0 +1,433 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * @file block_sparse_cholesky_solver_test.cpp + * @brief BlockSparseCholeskySolver on block-sparse SPD matrices shaped like + * the reduced systems of bundle adjustment: banded chains (visual-inertial), + * random covisibility, dense blocks wider than a panel, disconnected parts, + * tiny and empty systems; refactorization with new values; non-SPD input. + * Checked by the double-precision residual. + */ + +#include "cunls/linear_solver/block_sparse_cholesky_solver.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include "cunls/common/cuda_stream.h" +#include "cunls/common/helper.h" +#include "cunls/common/types.h" +#include "cunls/linear_solver/sparse_linear_solver.h" +#include "cunls/minimizer/problem.h" + +namespace cunls { +namespace { + +/** A symmetric block pattern: block sizes and the coupled block pairs (i < j). */ +struct BlockPattern { + std::vector sizes; + std::set> pairs; +}; + +/** A host SPD matrix (full, double) on a block pattern, diagonally dominant. */ +struct HostSystem { + int n = 0; + std::vector a; // row-major n x n + std::vector row_offsets, col_ids; +}; + +HostSystem MakeSystem(const BlockPattern &pattern, unsigned seed) { + HostSystem s; + std::vector start{0}; + for (int size : pattern.sizes) start.push_back(start.back() + size); + s.n = start.back(); + const int n = s.n; + s.a.assign(static_cast(n) * n, 0.0); + std::mt19937 rng(seed); + std::uniform_real_distribution u(-1.0, 1.0); + auto fill = [&](int bi, int bj) { + for (int i = start[bi]; i < start[bi + 1]; ++i) { + for (int j = start[bj]; j < start[bj + 1]; ++j) { + if (bi == bj && j > i) continue; + const double v = u(rng); + s.a[static_cast(i) * n + j] = v; + s.a[static_cast(j) * n + i] = v; + } + } + }; + for (size_t b = 0; b < pattern.sizes.size(); ++b) fill(static_cast(b), static_cast(b)); + for (const auto &[bi, bj] : pattern.pairs) fill(bi, bj); + for (int i = 0; i < n; ++i) { + double off = 0.0; + for (int j = 0; j < n; ++j) { + if (j != i) off += std::fabs(s.a[static_cast(i) * n + j]); + } + s.a[static_cast(i) * n + i] = off + 1.0 + std::fabs(u(rng)); + } + // Pattern CSR: the blocks, not the values (explicit zeros allowed). + std::vector> coupled(pattern.sizes.size(), + std::vector(pattern.sizes.size(), 0)); + for (size_t b = 0; b < pattern.sizes.size(); ++b) coupled[b][b] = 1; + for (const auto &[bi, bj] : pattern.pairs) coupled[bi][bj] = coupled[bj][bi] = 1; + s.row_offsets.push_back(0); + for (size_t bi = 0; bi < pattern.sizes.size(); ++bi) { + for (int i = start[bi]; i < start[bi + 1]; ++i) { + for (size_t bj = 0; bj < pattern.sizes.size(); ++bj) { + if (!coupled[bi][bj]) continue; + for (int j = start[bj]; j < start[bj + 1]; ++j) s.col_ids.push_back(j); + } + s.row_offsets.push_back(static_cast(s.col_ids.size())); + } + } + return s; +} + +std::vector CsrValues(const HostSystem &s) { + std::vector values; + for (int i = 0; i < s.n; ++i) { + for (int e = s.row_offsets[i]; e < s.row_offsets[i + 1]; ++e) { + values.push_back(static_cast(s.a[static_cast(i) * s.n + s.col_ids[e]])); + } + } + return values; +} + +/** Device copy of a host system plus a right-hand side. */ +struct DeviceSystem { + explicit DeviceSystem(const HostSystem &s) : rhs(s.n), result(s.n) { + matrix.row_offsets = dvector(s.row_offsets); + if (!s.col_ids.empty()) matrix.col_ids = dvector(s.col_ids); + SetValues(s); + } + void SetValues(const HostSystem &s) { + const std::vector values = CsrValues(s); + matrix.values.resize(values.size()); + if (!values.empty()) matrix.values.CopyFromHost(values.data(), values.size()); + } + CSRSparseMatrix matrix; + dvector rhs, result; +}; + +/** + * Relative residual ‖A x - b‖ / ‖b‖ (double) of the GPU solution; the test + * matrices are diagonally dominant, so it bounds the relative error too. + */ +double SolveError(BlockSparseCholeskySolver &solver, cudaStream_t stream, const HostSystem &s, + DeviceSystem &d, unsigned seed, bool *ok) { + std::mt19937 rng(seed); + std::uniform_real_distribution u(-1.0, 1.0); + std::vector b(s.n); + std::vector bf(s.n); + for (int i = 0; i < s.n; ++i) bf[i] = static_cast(b[i] = u(rng)); + if (s.n > 0) d.rhs.CopyFromHost(bf.data(), bf.size()); + *ok = solver.Solve(stream, d.matrix, d.rhs, d.result); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + if (s.n == 0) return 0.0; + std::vector x(s.n); + d.result.CopyToHost(x.data(), x.size()); + double num = 0.0, den = 0.0; + for (int i = 0; i < s.n; ++i) { + double r = -b[i]; + for (int e = s.row_offsets[i]; e < s.row_offsets[i + 1]; ++e) { + r += s.a[static_cast(i) * s.n + s.col_ids[e]] * x[s.col_ids[e]]; + } + num += r * r; + den += b[i] * b[i]; + } + return std::sqrt(num / den); +} + +/** Keyframes of state blocks {6, 3, 6} (pose, velocity, bias), each coupled to the next `band`. */ +BlockPattern Chain(int keyframes, int band) { + BlockPattern p; + for (int k = 0; k < keyframes; ++k) { + for (int size : {6, 3, 6}) p.sizes.push_back(size); + } + for (int k = 0; k < keyframes; ++k) { + for (int a = 0; a < 3; ++a) { + for (int b = a + 1; b < 3; ++b) p.pairs.insert({3 * k + a, 3 * k + b}); + } + for (int d = 1; d <= band && k + d < keyframes; ++d) { + for (int a = 0; a < 3; ++a) { + for (int b = 0; b < 3; ++b) p.pairs.insert({3 * k + a, 3 * (k + d) + b}); + } + } + } + return p; +} + +/** Poses (6) along a trajectory, coupled to neighbors and to random others with probability `fill`. + */ +BlockPattern Covisibility(int poses, double fill, unsigned seed) { + BlockPattern p; + p.sizes.assign(poses, 6); + std::mt19937 rng(seed); + std::uniform_real_distribution u(0.0, 1.0); + for (int i = 0; i < poses; ++i) { + for (int j = i + 1; j < poses; ++j) { + if (j - i <= 3 || u(rng) < fill) p.pairs.insert({i, j}); + } + } + return p; +} + +BlockPattern Dense(int blocks, int size) { + BlockPattern p; + p.sizes.assign(blocks, size); + for (int i = 0; i < blocks; ++i) { + for (int j = i + 1; j < blocks; ++j) p.pairs.insert({i, j}); + } + return p; +} + +/** A w x h grid of blocks of the given size, each coupled to its right and lower neighbors. */ +BlockPattern Grid(int w, int h, int size) { + BlockPattern p; + p.sizes.assign(w * h, size); + for (int i = 0; i < w * h; ++i) { + if (i % w + 1 < w) p.pairs.insert({i, i + 1}); + if (i + w < w * h) p.pairs.insert({i, i + w}); + } + return p; +} + +/** + * The full system of bundle adjustment: a chain of poses (6) and landmarks + * (3) each seen by 3 to 5 random poses, so that every pose reaches far more + * vertices than the average (the lazily updated vertices of minimum degree). + */ +BlockPattern BundleAdjustment(int poses, int landmarks, unsigned seed) { + BlockPattern p; + p.sizes.assign(poses, 6); + p.sizes.insert(p.sizes.end(), landmarks, 3); + for (int i = 0; i + 1 < poses; ++i) p.pairs.insert({i, i + 1}); + std::mt19937 rng(seed); + std::uniform_int_distribution pose(0, poses - 1), views(3, 5); + for (int l = 0; l < landmarks; ++l) { + for (int v = views(rng); v > 0; --v) p.pairs.insert({pose(rng), poses + l}); + } + return p; +} + +struct Case { + const char *name; + BlockPattern pattern; +}; + +std::vector Cases() { + std::vector cases; + cases.push_back({"Scalar", {{1}, {}}}); + cases.push_back({"TwoBlocks", {{3, 6}, {{0, 1}}}}); + cases.push_back({"Chain", Chain(120, 6)}); + cases.push_back({"LongChain", Chain(400, 3)}); + cases.push_back({"Covisibility", Covisibility(250, 0.15, 7)}); + cases.push_back({"SparseCovisibility", Covisibility(300, 0.02, 11)}); + cases.push_back({"Dense", Dense(50, 6)}); // 300 columns: several panels per supernode + cases.push_back({"WideBlock", {{150, 6, 70}, {{0, 1}, {1, 2}}}}); + cases.push_back({"Grid", Grid(30, 30, 3)}); + cases.push_back({"ScalarGrid", Grid(40, 40, 1)}); + cases.push_back({"BundleAdjustment", BundleAdjustment(12, 500, 5)}); + BlockPattern two = Chain(60, 4); + const BlockPattern other = Covisibility(40, 0.1, 3); + const int offset = static_cast(two.sizes.size()); + two.sizes.insert(two.sizes.end(), other.sizes.begin(), other.sizes.end()); + for (const auto &[i, j] : other.pairs) two.pairs.insert({i + offset, j + offset}); + cases.push_back({"Disconnected", two}); + return cases; +} + +class BlockSparseCholeskyTest : public ::testing::TestWithParam {}; + +TEST_P(BlockSparseCholeskyTest, SolvesAccurately) { + const HostSystem s = MakeSystem(GetParam().pattern, 1); + DeviceSystem d(s); + CudaStream stream; + BlockSparseCholeskySolver solver; + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + bool ok = false; + EXPECT_LT(SolveError(solver, stream.GetStream(), s, d, 2, &ok), 1e-5); + EXPECT_TRUE(ok); + // A second right-hand side, and new values on the same pattern. + EXPECT_LT(SolveError(solver, stream.GetStream(), s, d, 3, &ok), 1e-5); + const HostSystem s2 = MakeSystem(GetParam().pattern, 4); + d.SetValues(s2); + EXPECT_LT(SolveError(solver, stream.GetStream(), s2, d, 5, &ok), 1e-5); + EXPECT_TRUE(ok); +} + +INSTANTIATE_TEST_SUITE_P(Structures, BlockSparseCholeskyTest, ::testing::ValuesIn(Cases()), + [](const ::testing::TestParamInfo &info) { + return std::string(info.param.name); + }); + +/** Panels never exceed the width limit, and every column belongs to exactly one panel. */ +TEST(BlockSparseCholeskyPlanTest, PanelsCoverTheColumns) { + const HostSystem s = MakeSystem(Dense(50, 6), 1); + DeviceSystem d(s); + CudaStream stream; + BlockSparseCholeskySolver solver; + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + const internal::CholeskyPlan &plan = solver.Plan(); + int next = 0; + for (const auto &panel : plan.panels) { + EXPECT_LE(panel.width, internal::kCholeskyPanelWidth); + EXPECT_GE(panel.col, next); + next = panel.col + panel.width; + } + EXPECT_EQ(next, s.n); + EXPECT_GT(plan.num_levels, 1); +} + +TEST(BlockSparseCholeskySolverTest, EmptySystem) { + HostSystem s; + s.row_offsets = {0}; + DeviceSystem d(s); + CudaStream stream; + BlockSparseCholeskySolver solver; + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + EXPECT_TRUE(solver.Solve(stream.GetStream(), d.matrix, d.rhs, d.result)); +} + +/** A negative pivot is reported with safety checks; without them the result is non-finite. */ +TEST(BlockSparseCholeskySolverTest, ReportsNonPositiveDefinite) { + HostSystem s = MakeSystem(Chain(30, 3), 1); + s.a[static_cast(200) * s.n + 200] = -1.0; + for (bool checks : {true, false}) { + DeviceSystem d(s); + CudaStream stream; + BlockSparseCholeskySolver solver; + if (!checks) solver.DisableSafetyChecks(); + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + bool ok = true; + SolveError(solver, stream.GetStream(), s, d, 2, &ok); + std::vector x(s.n); + d.result.CopyToHost(x.data(), x.size()); + bool finite = true; + for (float v : x) finite = finite && std::isfinite(v); + if (checks) { + EXPECT_FALSE(ok); + } else { + EXPECT_TRUE(ok); + EXPECT_FALSE(finite); + } + } +} + +/** + * A consistent semidefinite system, the normal equations of relative + * measurements between poses without an anchor (Gauss-Newton with a free + * gauge): the vanishing pivots fix their unknowns, and the solution + * satisfies the system. + */ +TEST(BlockSparseCholeskySolverTest, SolvesConsistentSemidefiniteSystems) { + const BlockPattern pattern = Covisibility(40, 0.1, 3); + HostSystem s = MakeSystem(pattern, 1); + const int n = s.n; + std::fill(s.a.begin(), s.a.end(), 0.0); + std::mt19937 rng(9); + std::uniform_real_distribution u(-1.0, 1.0); + for (const auto &[bi, bj] : pattern.pairs) { + // Residual R (x_i - x_j): blocks RᵀR on the diagonal, -RᵀR off it. + double r[6][6], rtr[6][6] = {}; + for (auto &row : r) { + for (double &v : row) v = u(rng); + } + for (int i = 0; i < 6; ++i) { + for (int j = 0; j < 6; ++j) { + for (int k = 0; k < 6; ++k) rtr[i][j] += r[k][i] * r[k][j]; + } + } + for (int i = 0; i < 6; ++i) { + for (int j = 0; j < 6; ++j) { + s.a[static_cast(6 * bi + i) * n + 6 * bi + j] += rtr[i][j]; + s.a[static_cast(6 * bj + i) * n + 6 * bj + j] += rtr[i][j]; + s.a[static_cast(6 * bi + i) * n + 6 * bj + j] -= rtr[i][j]; + s.a[static_cast(6 * bj + i) * n + 6 * bi + j] -= rtr[i][j]; + } + } + } + // b = A y: consistent (orthogonal to the null space, a common 6-vector). + std::vector y(n), b(n, 0.0); + for (double &v : y) v = u(rng); + for (int i = 0; i < n; ++i) { + for (int j = 0; j < n; ++j) b[i] += s.a[static_cast(i) * n + j] * y[j]; + } + std::vector bf(b.begin(), b.end()); + DeviceSystem d(s); + d.rhs.CopyFromHost(bf.data(), n); + CudaStream stream; + BlockSparseCholeskySolver solver; + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + EXPECT_TRUE(solver.Solve(stream.GetStream(), d.matrix, d.rhs, d.result)); + std::vector x(n); + d.result.CopyToHost(x.data(), n); + double num = 0.0, den = 0.0; + for (int i = 0; i < n; ++i) { + double r = -b[i]; + for (int j = 0; j < n; ++j) r += s.a[static_cast(i) * n + j] * x[j]; + num += r * r; + den += b[i] * b[i]; + } + EXPECT_LT(std::sqrt(num / den), 1e-4); +} + +/** + * Strongly coupled unknowns (pivot 1e-4 of its diagonal entry, as with stiff + * motion priors) are regular, not a rank deficiency: solved accurately. + */ +TEST(BlockSparseCholeskySolverTest, KeepsSmallRegularPivots) { + HostSystem s = MakeSystem(Dense(2, 1), 1); + const double c = 1.0 - 5e-5; // second pivot 1 - c^2 = 1e-4 + s.a = {1.0, c, c, 1.0}; + DeviceSystem d(s); + const std::vector b = {1.f, -1.f}; + d.rhs.CopyFromHost(b.data(), 2); + CudaStream stream; + BlockSparseCholeskySolver solver; + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + EXPECT_TRUE(solver.Solve(stream.GetStream(), d.matrix, d.rhs, d.result)); + std::vector x(2); + d.result.CopyToHost(x.data(), 2); + // x = (1, -1) / (1 - c) = ±2e4; float32 keeps about 4 digits at this conditioning. + EXPECT_NEAR(x[0], 1.0 / (1.0 - c), 2e2); + EXPECT_NEAR(x[1], -1.0 / (1.0 - c), 2e2); +} + +TEST(BlockSparseCholeskySolverTest, CreatedByTheFactory) { + const HostSystem s = MakeSystem(Covisibility(40, 0.2, 5), 1); + DeviceSystem d(s); + CudaStream stream; + SparseLinearSolverPtr solver = + CreateSparseLinearSolver(SparseLinearSolverType::BlockSparseCholesky, {}); + ASSERT_NE(dynamic_cast(solver.get()), nullptr); + ASSERT_TRUE(solver->Initialize(stream.GetStream(), Problem(), d.matrix, d.rhs, d.result)); + bool ok = false; + EXPECT_LT(SolveError(*dynamic_cast(solver.get()), stream.GetStream(), + s, d, 2, &ok), + 1e-5); + EXPECT_TRUE(ok); +} + +} // namespace +} // namespace cunls diff --git a/tests/gauss_newton_test.cpp b/tests/gauss_newton_test.cpp index 5bc0415..dca542f 100644 --- a/tests/gauss_newton_test.cpp +++ b/tests/gauss_newton_test.cpp @@ -143,6 +143,7 @@ struct TestParam { }; constexpr int kPCG = test_utils::kBlockSparsePCGSolverId; +constexpr int kChol = test_utils::kBlockSparseCholeskySolverId; /** @brief Test types: 1D, 2D, 3D, and 4D vectors for each available solver. */ typedef ::testing::Types< @@ -150,7 +151,8 @@ typedef ::testing::Types< TestParam<1, 0>, TestParam<2, 0>, TestParam<3, 0>, TestParam<4, 0>, TestParam<1, 1>, TestParam<2, 1>, TestParam<3, 1>, TestParam<4, 1>, #endif - TestParam<1, kPCG>, TestParam<2, kPCG>, TestParam<3, kPCG>, TestParam<4, kPCG>> + TestParam<1, kPCG>, TestParam<2, kPCG>, TestParam<3, kPCG>, TestParam<4, kPCG>, + TestParam<1, kChol>, TestParam<2, kChol>, TestParam<3, kChol>, TestParam<4, kChol>> TestParams; TYPED_TEST_CASE(GaussNewtonMinimizerTest, TestParams); diff --git a/tests/imu_factor_batch_test.cpp b/tests/imu_factor_batch_test.cpp index 99767b6..6637497 100644 --- a/tests/imu_factor_batch_test.cpp +++ b/tests/imu_factor_batch_test.cpp @@ -40,6 +40,7 @@ #include "cunls/state/se3_state_batch.h" #include "cunls/state/vector_state_batch.h" #include "dynamics_test_support.h" +#include "tests/utils.h" namespace cunls { namespace { @@ -892,14 +893,16 @@ int EnvInt(const char *name, int fallback) { * read the same world_from_rig pose states (camera = rig = IMU); truth is * consistent, the initial values are perturbed. Sizes: CUNLS_IMU_PROFILE_KEYFRAMES (101), * CUNLS_IMU_PROFILE_SAMPLES (200), CUNLS_IMU_PROFILE_LANDMARKS (10 K); - * CUNLS_IMU_PROFILE_SOLVER=cudss selects - * cuDSS instead of the default PCG. Run with --gtest_also_run_disabled_tests. + * CUNLS_IMU_PROFILE_SOLVER=cudss or pcg selects cuDSS or BlockSparsePCG + * instead of the default BlockSparseCholesky. Run with + * --gtest_also_run_disabled_tests. */ TEST(ImuFactorBatch, DISABLED_VioProfile) { const int K = EnvInt("CUNLS_IMU_PROFILE_KEYFRAMES", 101); const int N = EnvInt("CUNLS_IMU_PROFILE_SAMPLES", 200); const char *solver_env = std::getenv("CUNLS_IMU_PROFILE_SOLVER"); const bool cudss = solver_env != nullptr && std::string(solver_env) == "cudss"; + const bool pcg = solver_env != nullptr && std::string(solver_env) == "pcg"; const int L = EnvInt("CUNLS_IMU_PROFILE_LANDMARKS", 10 * K); // Moderate noise so that the float32 solve converges (see ImuParameters); // the kernels' cost does not depend on the values. @@ -1008,6 +1011,13 @@ TEST(ImuFactorBatch, DISABLED_VioProfile) { LevenbergMarquardtMinimizerOptions options; options.base_options.max_num_iterations = 20; if (cudss) options.base_options.sparse_linear_solver_type = SparseLinearSolverType::cuDSS; + if (pcg) options.base_options.sparse_linear_solver_type = SparseLinearSolverType::BlockSparsePCG; + options.base_options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(200); + options.base_options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = + test_utils::PCGTolFromEnv(1e-3f); + const bool schur = test_utils::UseSchurComplementFromEnv(options.base_options); + if (schur) problem.SetEliminatedStateBatches({&point_states}); LevenbergMarquardtMinimizer minimizer(options); CudaStream stream; // Solve twice from the same initial values: the first call includes the @@ -1025,11 +1035,12 @@ TEST(ImuFactorBatch, DISABLED_VioProfile) { ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); std::printf(" %s solve: %.2f ms\n", run == 0 ? "cold" : "warm", ms); } + const char *label = schur ? "Schur" : cudss ? "cuDSS" : pcg ? "PCG" : "BlockSparseCholesky"; std::printf( "VIO profile: %d keyframes, %d samples per pair, %d landmarks, %d observations, %s\n" " %zu iterations, cost %g -> %g, %.2f ms (%.3f ms per iteration)\n", - K, N, L, num_obs, cudss ? "cuDSS" : "PCG", summary.num_iterations, summary.initial_cost, - summary.final_cost, ms, ms / std::max(1, summary.num_iterations)); + K, N, L, num_obs, label, summary.num_iterations, summary.initial_cost, summary.final_cost, ms, + ms / std::max(1, summary.num_iterations)); } // Throughput: run with --gtest_also_run_disabled_tests. diff --git a/tests/minimizer_benchmark_test.cpp b/tests/minimizer_benchmark_test.cpp index 96e5494..9c6c0a2 100644 --- a/tests/minimizer_benchmark_test.cpp +++ b/tests/minimizer_benchmark_test.cpp @@ -529,5 +529,264 @@ TEST(MinimizerBenchmark, DISABLED_RealTimeMpc) { } } +// ============================================================================= +// Linear solver comparison: the same problem with each backend in BENCH_SOLVERS +// (comma separated, default "pcg,chol,cudss"), Levenberg-Marquardt with the +// default options apart from the solver. Every Minimize() repeats the solver's +// setup (no reuse_structure); cold_ms is the first call of the process. +// BENCH LM/ median_ms=... cold_ms=... iterations=... final_cost=... +// ============================================================================= + +/** The backends named in BENCH_SOLVERS. */ +std::vector> BenchSolvers() { + const char *env = std::getenv("BENCH_SOLVERS"); + const std::string list = env != nullptr ? env : "pcg,chol,cudss"; + std::vector> out; + if (list.find("pcg") != std::string::npos) { + out.push_back({"pcg", SparseLinearSolverType::BlockSparsePCG}); + } + if (list.find("chol") != std::string::npos) { + out.push_back({"chol", SparseLinearSolverType::BlockSparseCholesky}); + } +#ifdef CUNLS_ENABLE_CUDSS + if (list.find("cudss") != std::string::npos) { + out.push_back({"cudss", SparseLinearSolverType::cuDSS}); + } +#endif + return out; +} + +/** Runs LM on the problem with every backend of BENCH_SOLVERS. */ +void CompareSolvers(const std::string &name, Problem &problem, int reps, + const std::function &reset) { + CudaStream stream; + for (const auto &[label, solver] : BenchSolvers()) { + LevenbergMarquardtMinimizerOptions lm; + lm.base_options.sparse_linear_solver_type = solver; + LevenbergMarquardtMinimizer minimizer(lm); + auto solve = [&] { return minimizer.Minimize(stream.GetStream(), problem); }; + reset(); + THROW_ON_CUDA_ERROR(cudaDeviceSynchronize()); + const auto t0 = std::chrono::steady_clock::now(); + solve(); + THROW_ON_CUDA_ERROR(cudaDeviceSynchronize()); + const double cold = + std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + const Timing t = Time(reps, reset, solve); + std::printf( + "BENCH %-30s LM/%-5s median_ms=%9.3f cold_ms=%9.3f iterations=%3zu final_cost=%.6g\n", + name.c_str(), label.c_str(), t.median_ms, cold, t.iterations, t.final_cost); + std::fflush(stdout); + } +} + +/** Edges i - i+1 along a chain of n poses plus i - i+ring for every pose (a ring x n/ring grid). */ +std::vector> GridEdges(int n, int ring) { + std::vector> edges; + for (int i = 0; i + 1 < n; ++i) edges.push_back({i, i + 1}); + for (int i = 0; i + ring < n; ++i) edges.push_back({i, i + ring}); + return edges; +} + +/** A chain of n poses plus `closures` loop closures between random pose pairs. */ +std::vector> RandomEdges(int n, int closures, std::mt19937 &rng) { + std::vector> edges; + for (int i = 0; i + 1 < n; ++i) edges.push_back({i, i + 1}); + std::uniform_int_distribution pick(0, n - 1); + while (static_cast(edges.size()) < n - 1 + closures) { + const int a = pick(rng), b = pick(rng); + if (std::abs(a - b) > 1) edges.push_back({a, b}); + } + return edges; +} + +/** + * SE(3) pose graph over a random walk of n poses with the given edges: noisy + * relative-pose measurements (Δ = T_b⁻¹ T_a for the edge (a, b), as + * SE3BetweenFactorBatch expects), every pose perturbed, the first one constant. + */ +void SE3Graph(const std::string &name, int n, const std::vector> &edges, + int reps) { + std::mt19937 rng(7); + std::uniform_real_distribution r(-0.1, 0.1), t(-0.4, 0.4); + std::normal_distribution normal(0.0, 1.0); + auto noise = [&](double s) { + return Pose3({s * normal(rng), s * normal(rng), s * normal(rng)}, + {s * normal(rng), s * normal(rng), s * normal(rng)}); + }; + std::vector gt(n), init(n); + gt[0] = Pose3({0, 0, 0}, {0, 0, 0}); + for (int i = 1; i < n; ++i) { + gt[i] = Mul<16, 4>(gt[i - 1], Pose3({r(rng), r(rng), r(rng)}, {t(rng), t(rng), t(rng)})); + } + init[0] = gt[0]; + for (int i = 1; i < n; ++i) init[i] = Mul<16, 4>(gt[i], noise(0.05)); + std::vector deltas; + for (auto [a, b] : edges) { + deltas.push_back(Mul<16, 4>(Mul<16, 4>(Inv<16, 4>(gt[b]), gt[a]), noise(0.01))); + } + const std::vector init_f = ToFloat(init); + dvector init_d(init_f), poses(init_f), deltas_d(ToFloat(deltas)); + dvector const_ids(std::vector{0}); + SE3StateBatch states(poses.data(), n, const_ids.data(), 1); + states.SetNumActiveStates(n, 1); + SE3BetweenFactorBatch between(reinterpret_cast(deltas_d.data()), + edges.size()); + between.SetNumActiveFactors(edges.size()); + std::vector ptrs; + for (auto [a, b] : edges) { + ptrs.push_back(states.StateDevicePtr(a)); + ptrs.push_back(states.StateDevicePtr(b)); + } + Problem problem; + problem.AddStateBatch(&states); + problem.AddFactorBatch(&between, ptrs); + CompareSolvers(name, problem, reps, [&] { + THROW_ON_CUDA_ERROR(cudaMemcpy(poses.data(), init_d.data(), init_f.size() * sizeof(float), + cudaMemcpyDeviceToDevice)); + }); +} + +/** SE(2) counterpart of SE3Graph, with a prior on the first pose instead of a constant one. */ +void SE2Graph(const std::string &name, int n, const std::vector> &edges, + int reps) { + std::mt19937 rng(3); + std::normal_distribution normal(0.0, 1.0); + std::vector gt, init, deltas; + Mat3 t = Pose2(0, 0, 0); + for (int i = 0; i < n; ++i) { + gt.push_back(t); + t = Mul<9, 3>(t, Pose2(1.0 + 0.3 * normal(rng), 0.3 * normal(rng), 0.2 * normal(rng))); + } + for (int i = 0; i < n; ++i) { + init.push_back( + Mul<9, 3>(gt[i], Pose2(0.1 * normal(rng), 0.1 * normal(rng), 0.1 * normal(rng)))); + } + for (auto [a, b] : edges) { + const Mat3 noise = Pose2(0.02 * normal(rng), 0.02 * normal(rng), 0.02 * normal(rng)); + deltas.push_back(Mul<9, 3>(Mul<9, 3>(Inv<9, 3>(gt[b]), gt[a]), noise)); + } + const std::vector init_f = ToFloat(init); + dvector init_d(init_f), poses(init_f); + dvector deltas_d(ToFloat(deltas)), prior_d(ToFloat(std::vector{gt[0]})); + SE2StateBatch states(poses.data(), n); + states.SetNumActiveStates(n); + SE2PriorFactorBatch prior(reinterpret_cast(prior_d.data()), 1); + prior.SetNumActiveFactors(1); + SE2BetweenFactorBatch between(reinterpret_cast(deltas_d.data()), + edges.size()); + between.SetNumActiveFactors(edges.size()); + std::vector between_ptrs; + for (auto [a, b] : edges) { + between_ptrs.push_back(states.StateDevicePtr(a)); + between_ptrs.push_back(states.StateDevicePtr(b)); + } + Problem problem; + problem.AddStateBatch(&states); + problem.AddFactorBatch(&prior, std::vector{states.StateDevicePtr(0)}); + problem.AddFactorBatch(&between, between_ptrs); + CompareSolvers(name, problem, reps, [&] { + THROW_ON_CUDA_ERROR(cudaMemcpy(poses.data(), init_d.data(), init_f.size() * sizeof(float), + cudaMemcpyDeviceToDevice)); + }); +} + +/** + * Dim-dimensional vectors on a w x w grid (a Gaussian Markov random field): + * a weak prior (weight 0.1) to random targets on every vector and a + * difference factor to random offsets between grid neighbours. Linear, so a + * direct solve converges in one step. + */ +template +void VectorGrid(const std::string &name, int w, int reps) { + const int n = w * w; + std::mt19937 rng(11); + std::normal_distribution normal(0.f, 1.f); + std::vector targets(n * Dim), x0(n * Dim, 0.f), offsets; + for (float &v : targets) v = normal(rng); + std::vector prior_ptrs, edge_ptrs; + dvector x(x0), init(x0); + VectorStateBatch states(x.data(), n); + states.SetNumActiveStates(n); + for (int i = 0; i < n; ++i) prior_ptrs.push_back(states.StateDevicePtr(i)); + for (int r = 0; r < w; ++r) { + for (int c = 0; c < w; ++c) { + const int i = r * w + c; + for (int j : {c + 1 < w ? i + 1 : -1, r + 1 < w ? i + w : -1}) { + if (j < 0) continue; + edge_ptrs.push_back(states.StateDevicePtr(i)); + edge_ptrs.push_back(states.StateDevicePtr(j)); + for (int k = 0; k < Dim; ++k) offsets.push_back(0.1f * normal(rng)); + } + } + } + const size_t edges = edge_ptrs.size() / 2; + dvector d_targets(targets), d_offsets(offsets); + WeightedFactorBatch> prior( + 0.1f, reinterpret_cast *>(d_targets.data()), static_cast(n)); + prior.SetNumActiveFactors(n); + VectorBetweenFactorBatch smooth(reinterpret_cast *>(d_offsets.data()), + edges); + smooth.SetNumActiveFactors(edges); + Problem problem; + problem.AddStateBatch(&states); + problem.AddFactorBatch(&prior, prior_ptrs); + problem.AddFactorBatch(&smooth, edge_ptrs); + CompareSolvers(name, problem, reps, [&] { + THROW_ON_CUDA_ERROR( + cudaMemcpy(x.data(), init.data(), x0.size() * sizeof(float), cudaMemcpyDeviceToDevice)); + }); +} + +TEST(SolverBenchmark, DISABLED_SE3PoseGraphs) { + // Grid-like loop closures (revisiting the previous row): sphere2500-like topology. + for (auto [rings, per_ring] : + std::vector>{{10, 10}, {50, 50}, {100, 100}, {200, 200}, {400, 250}}) { + const int n = rings * per_ring; + SE3Graph("se3_grid_" + std::to_string(n), n, GridEdges(n, per_ring), n > 20000 ? 5 : 10); + } + // Long-range loop closures between random poses: the worst case for fill-in. + for (auto [n, closures] : + std::vector>{{1000, 100}, {5000, 2000}, {20000, 1000}}) { + std::mt19937 rng(5); + SE3Graph("se3_random_" + std::to_string(n) + "_lc" + std::to_string(closures), n, + RandomEdges(n, closures, rng), 10); + } + // A plain odometry chain. + std::mt19937 rng(1); + SE3Graph("se3_chain_10000", 10000, RandomEdges(10000, 0, rng), 10); +} + +TEST(SolverBenchmark, DISABLED_SE2PoseGraphs) { + for (auto [rows, per_row] : std::vector>{{35, 100}, {100, 100}, {300, 300}}) { + const int n = rows * per_row; + SE2Graph("se2_grid_" + std::to_string(n), n, GridEdges(n, per_row), 10); + } + for (auto [n, closures] : std::vector>{{5000, 1000}, {20000, 1000}}) { + std::mt19937 rng(5); + SE2Graph("se2_random_" + std::to_string(n) + "_lc" + std::to_string(closures), n, + RandomEdges(n, closures, rng), 10); + } +} + +TEST(SolverBenchmark, DISABLED_VectorGrids) { + VectorGrid<1>("vec1_grid_300x300", 300, 10); + VectorGrid<2>("vec2_grid_200x200", 200, 10); + VectorGrid<3>("vec3_grid_100x100", 100, 10); + VectorGrid<3>("vec3_grid_300x300", 300, 10); + VectorGrid<4>("vec4_grid_100x100", 100, 10); + VectorGrid<6>("vec6_grid_100x100", 100, 10); + VectorGrid<9>("vec9_grid_100x100", 100, 10); + VectorGrid<12>("vec12_grid_100x100", 100, 10); + VectorGrid<16>("vec16_grid_64x64", 64, 10); +} + +TEST(SolverBenchmark, DISABLED_BatchedGraphs) { + for (const auto &[label, solver] : BenchSolvers()) { + SE2Graphs("se2_graph_256x12/" + label, 256, 12, 2, solver, 50); + SE2Graphs("se2_graph_16x500/" + label, 16, 500, 100, solver, 20); + } +} + } // namespace } // namespace cunls diff --git a/tests/numeric_diff_perf_test.cpp b/tests/numeric_diff_perf_test.cpp index 9ae83a2..1022c2a 100644 --- a/tests/numeric_diff_perf_test.cpp +++ b/tests/numeric_diff_perf_test.cpp @@ -108,8 +108,7 @@ const char *ToString(JacobianMode m) { class SystemBuilder { public: explicit SystemBuilder(const MinimizerOptions &options) : options_(options) { - system_.solver = CreateSparseLinearSolver(options.sparse_linear_solver_type, - options.sparse_linear_solver_config); + system_.linear_system = internal::CreateLinearSystem(options); } void Prepare(cudaStream_t stream, Problem &problem) { diff --git a/tests/pnp_factor_batch_test.cpp b/tests/pnp_factor_batch_test.cpp index d6ed78a..69c16c7 100644 --- a/tests/pnp_factor_batch_test.cpp +++ b/tests/pnp_factor_batch_test.cpp @@ -401,7 +401,8 @@ INSTANTIATE_TEST_SUITE_P(AllSolvers, PnPSolverTest, ::testing::ValuesIn(test_utils::AvailableSolverTypes( {SparseLinearSolverType::cuDSS, SparseLinearSolverType::DenseLDLT, SparseLinearSolverType::DenseCholesky, - SparseLinearSolverType::DenseQR})), + SparseLinearSolverType::DenseQR, + SparseLinearSolverType::BlockSparseCholesky})), [](const ::testing::TestParamInfo &info) { switch (info.param) { case SparseLinearSolverType::cuDSS: @@ -412,6 +413,8 @@ INSTANTIATE_TEST_SUITE_P(AllSolvers, PnPSolverTest, return std::string("DenseCholesky"); case SparseLinearSolverType::DenseQR: return std::string("DenseQR"); + case SparseLinearSolverType::BlockSparseCholesky: + return std::string("BlockSparseCholesky"); default: return std::string("Unknown"); } diff --git a/tests/sba_minimizer_test.cpp b/tests/sba_minimizer_test.cpp index d562a05..d8d1e35 100644 --- a/tests/sba_minimizer_test.cpp +++ b/tests/sba_minimizer_test.cpp @@ -36,7 +36,9 @@ #include +#include #include +#include #include #include #include @@ -367,6 +369,7 @@ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { test_utils::PCGMaxIterFromEnv(400); options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); + const bool schur = test_utils::UseSchurComplementFromEnv(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -394,12 +397,22 @@ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { Problem problem; BuildProblemFromHost(host, &problem); ASSERT_TRUE(problem.CheckConsistency()); + if (schur) problem.SetEliminatedStateBatches({problem.GetStateBatches()[1]}); // points MinimizerSummary summary; { auto minimize_range = profiler_domain_.CreateDomainRange("Minimize"); + const auto start = std::chrono::steady_clock::now(); summary = minimizer.Minimize(stream.GetStream(), problem); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + const double ms = + std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); + std::printf( + "[SBA problem %d] %d poses, %d points, %d observations: %zu iterations, " + "cost %.6g -> %.6g, %.1f ms\n", + current_index, host.Nposes, host.Npoints, host.N_obs, summary.num_iterations, + summary.initial_cost, summary.final_cost, ms); } EXPECT_TRUE(std::isfinite(summary.initial_cost)) << "Initial cost must be finite (no NaNs)"; diff --git a/tests/schur_system_test.cpp b/tests/schur_system_test.cpp new file mode 100644 index 0000000..4a45818 --- /dev/null +++ b/tests/schur_system_test.cpp @@ -0,0 +1,577 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * @file schur_system_test.cpp + * @brief LinearSystemType::SchurComplement against the assembled normal + * equations on a small bundle adjustment: the same right-hand side, Hessian + * diagonal, step (with column scaling, damping and masks) and step + * curvature, the same minimization, and the setup checks. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cunls/common/cuda_stream.h" +#include "cunls/common/helper.h" +#include "cunls/common/types.h" +#include "cunls/factor/between/se3_between_factor_batch.h" +#include "cunls/factor/between/vector_between_factor_batch.h" +#include "cunls/factor/reprojection_factor_batch.h" +#include "cunls/minimizer/gauss_newton_minimizer.h" +#include "cunls/minimizer/levenberg_marquardt_minimizer.h" +#include "cunls/minimizer/minimizer.h" +#include "cunls/minimizer/problem.h" +#include "cunls/minimizer/problem_partition.h" +#include "cunls/robustifier/cauchy_loss_function_batch.h" +#include "cunls/state/se3_state_batch.h" +#include "cunls/state/vector_state_batch.h" + +namespace cunls { +namespace { + +/** Row-major 3x4 [R | t] with R a rotation about z by `yaw`. */ +SE3Transform Pose(float yaw, float x, float y, float z) { + const float c = std::cos(yaw), s = std::sin(yaw); + return {c, -s, 0.f, x, s, c, 0.f, y, 0.f, 0.f, 1.f, z}; +} + +/** a * b of two rigid transforms (row-major 3x4). */ +SE3Transform Compose(const SE3Transform &a, const SE3Transform &b) { + SE3Transform out{}; + for (int r = 0; r < 3; ++r) { + for (int c = 0; c < 4; ++c) { + float v = c == 3 ? a[r * 4 + 3] : 0.f; + for (int k = 0; k < 3; ++k) v += a[r * 4 + k] * b[k * 4 + c]; + out[r * 4 + c] = v; + } + } + return out; +} + +/** The inverse of a rigid transform. */ +SE3Transform Inverse(const SE3Transform &t) { + SE3Transform out{}; + for (int r = 0; r < 3; ++r) { + for (int c = 0; c < 3; ++c) out[r * 4 + c] = t[c * 4 + r]; + out[r * 4 + 3] = -(t[r] * t[3] + t[4 + r] * t[7] + t[8 + r] * t[11]); + } + return out; +} + +/** world_from_rig of ring position k (rig z axis toward the center, y down). */ +SE3Transform RingPose(int k, int num_poses, float radius) { + const float a = 2.f * static_cast(M_PI) * k / num_poses; + return {-std::sin(a), 0.f, -std::cos(a), radius * std::cos(a), // + std::cos(a), 0.f, -std::sin(a), radius * std::sin(a), // + 0.f, -1.f, 0.f, 0.f}; +} + +/** Variants of SmallBundle. */ +struct BundleOptions { + /// Independent rings, declared as subproblems (Problem::SetProblemPartition). + int rings = 1; + /// Points per ring in a second, retained point batch (reprojection factors + /// between retained states only: blocks of S of sizes 6 and 3). + int retained_points = 0; + /// Adds a landmark seen once, on the optical axis of an axis-aligned camera + /// of a constant pose: its Jacobian has an exactly zero column, so its + /// block of the Hessian is singular however it is damped. + bool singular_landmark = false; +}; + +/** + * Rings of cameras (world_from_rig, rig z axis toward the ring center, y + * down) around point clouds, seen through two cameras of a stereo rig (two + * reprojection batches, one with a Cauchy loss), plus SE3 between factors + * along each ring (factors between retained states only). The first pose and + * the first point of each ring are constant, so some reprojection factors + * touch no free pose. + */ +class SmallBundle { + public: + static constexpr int kNumPoses = 12; ///< Per ring. + static constexpr int kNumPoints = 300; ///< Per ring. + static constexpr float kRadius = 6.f; + static constexpr float kRingSpacing = 40.f; + + explicit SmallBundle(const BundleOptions &options = BundleOptions()) { + std::mt19937 rng(7); + std::normal_distribution noise(0.f, 1.f); + std::uniform_real_distribution cube(-2.f, 2.f); + const int rings = options.rings; + std::vector poses, ring; + std::vector> points, truth, retained, retained_truth; + std::vector pose_problem, point_problem, retained_problem, pose_const, point_const; + for (int g = 0; g < rings; ++g) { + for (int k = 0; k < kNumPoses; ++k) { + ring.push_back(RingPose(k, kNumPoses, kRadius)); + ring.back()[3] += kRingSpacing * g; + poses.push_back(ring.back()); + if (k > 0) { + for (int r = 0; r < 3; ++r) poses.back()[r * 4 + 3] += 0.05f * noise(rng); + } + pose_problem.push_back(g); + } + pose_const.push_back(g * kNumPoses); + point_const.push_back(g * kNumPoints); + for (int i = 0; i < kNumPoints; ++i) { + truth.push_back({kRingSpacing * g + cube(rng), cube(rng), cube(rng)}); + points.push_back({truth.back()[0] + 0.05f * noise(rng), + truth.back()[1] + 0.05f * noise(rng), + truth.back()[2] + 0.05f * noise(rng)}); + point_problem.push_back(g); + } + for (int i = 0; i < options.retained_points; ++i) { + retained_truth.push_back({kRingSpacing * g + cube(rng), cube(rng), cube(rng)}); + retained.push_back({retained_truth.back()[0] + 0.05f * noise(rng), + retained_truth.back()[1] + 0.05f * noise(rng), + retained_truth.back()[2] + 0.05f * noise(rng)}); + retained_problem.push_back(g); + } + } + if (options.singular_landmark) { + // Ring pose 0 has an axis-aligned rotation, and camera 0 sits at rig + // x = 0.1 (camera_from_rig translation -0.1): the rig point (0.1, 0, 3) + // lies exactly on its optical axis. + const SE3Transform &t = ring[0]; + const float pr[3] = {0.1f, 0.f, 3.f}; + Vector<3> p{}; + for (int r = 0; r < 3; ++r) { + p[r] = t[r * 4 + 3]; + for (int c = 0; c < 3; ++c) p[r] += t[r * 4 + c] * pr[c]; + } + truth.push_back(p); + points.push_back(p); + point_problem.push_back(0); + } + const int num_poses = static_cast(poses.size()); + const int num_points = static_cast(points.size()); + poses_ = dvector(poses); + points_ = dvector>(points); + pose_const_ = dvector(pose_const); + point_const_ = dvector(point_const); + pose_batch_ = std::make_unique(reinterpret_cast(poses_.data()), + num_poses, pose_const_.data(), rings); + pose_batch_->SetNumActiveStates(num_poses, rings); + point_batch_ = std::make_unique>( + reinterpret_cast(points_.data()), num_points, point_const_.data(), rings); + point_batch_->SetNumActiveStates(num_points, rings); + problem_.AddStateBatch(pose_batch_.get()); + problem_.AddStateBatch(point_batch_.get()); + if (!retained.empty()) { + retained_ = dvector>(retained); + retained_batch_ = std::make_unique>( + reinterpret_cast(retained_.data()), retained.size()); + retained_batch_->SetNumActiveStates(retained.size()); + problem_.AddStateBatch(retained_batch_.get()); + } + + // Two cameras 0.2 m apart on the rig; observations of the true points from + // the true rings (the initial poses are slightly off). + auto observe = [&](const SE3Transform &camera_from_world, const Vector<3> &x, bool exact) { + float pc[3]; + for (int r = 0; r < 3; ++r) { + pc[r] = camera_from_world[r * 4 + 3]; + for (int c = 0; c < 3; ++c) pc[r] += camera_from_world[r * 4 + c] * x[c]; + } + const float sigma = exact ? 0.f : 1e-3f; + return Vector<2>{pc[0] / pc[2] + sigma * noise(rng), pc[1] / pc[2] + sigma * noise(rng)}; + }; + for (int camera = 0; camera < 3; ++camera) { + // Batch 2: camera 0 seeing the retained points. + if (camera == 2 && retained.empty()) break; + const int lens = camera == 2 ? 0 : camera; + const SE3Transform camera_from_rig = Pose(0.f, lens == 0 ? -0.1f : 0.1f, 0.f, 0.f); + std::vector> observations; + std::vector extrinsics; + std::vector states; + for (int k = 0; k < num_poses; ++k) { + const int g = k / kNumPoses; + const SE3Transform camera_from_world = Compose(camera_from_rig, Inverse(ring[k])); + const int count = camera == 2 ? options.retained_points : kNumPoints; + for (int j = 0; j < count; ++j) { + if ((j + k + camera) % 3 == 0) continue; // not every camera sees every point + if (camera == 2) { + const int i = g * options.retained_points + j; + observations.push_back(observe(camera_from_world, retained_truth[i], false)); + states.push_back(retained_batch_->StateDevicePtr(i)); + } else { + const int i = g * kNumPoints + j; + observations.push_back(observe(camera_from_world, truth[i], false)); + states.push_back(point_batch_->StateDevicePtr(i)); + } + extrinsics.push_back(camera_from_rig); + states.insert(states.end() - 1, pose_batch_->StateDevicePtr(k)); + } + if (camera == 0 && k == 0 && options.singular_landmark) { + observations.push_back(observe(camera_from_world, truth.back(), true)); + extrinsics.push_back(camera_from_rig); + states.push_back(pose_batch_->StateDevicePtr(0)); + states.push_back(point_batch_->StateDevicePtr(num_points - 1)); + } + } + const size_t n = observations.size(); + observations_[camera] = dvector>(observations); + extrinsics_[camera] = dvector(extrinsics); + reprojection_[camera] = std::make_unique( + observations_[camera].data(), extrinsics_[camera].data(), n); + reprojection_[camera]->SetNumActiveFactors(n); + problem_.AddFactorBatch(reprojection_[camera].get(), camera == 1 ? &cauchy_ : nullptr, + states); + } + + // Relative poses along each ring (the measured deltas of the true ring). + std::vector deltas; + std::vector between_states; + for (int k = 0; k < num_poses; ++k) { + const int next = (k / kNumPoses) * kNumPoses + (k + 1) % kNumPoses; + // residual = Log(Δ T_k⁻¹ T_next): Δ = T_next⁻¹ T_k. + deltas.push_back(Compose(Inverse(ring[next]), ring[k])); + between_states.push_back(pose_batch_->StateDevicePtr(k)); + between_states.push_back(pose_batch_->StateDevicePtr(next)); + } + deltas_ = dvector(deltas); + between_ = std::make_unique(deltas_.data(), deltas.size()); + between_->SetNumActiveFactors(deltas.size()); + problem_.AddFactorBatch(between_.get(), between_states); + + if (rings > 1) { + pose_problem_ = dvector(pose_problem); + point_problem_ = dvector(point_problem); + std::vector ids = {pose_problem_.data(), point_problem_.data()}; + if (!retained.empty()) { + retained_problem_ = dvector(retained_problem); + ids.push_back(retained_problem_.data()); + } + problem_.SetProblemPartition(rings, ids); + } + } + + Problem &GetProblem() { return problem_; } + const StateBatch *Points() const { return point_batch_.get(); } + VectorStateBatch<3> *PointBatch() { return point_batch_.get(); } + /** The free tangent columns (poses, points, retained points). */ + size_t NumColumns() const { + size_t n = (pose_batch_->NumActiveStates() - pose_batch_->NumConstStates()) * 6 + + (point_batch_->NumActiveStates() - point_batch_->NumConstStates()) * 3; + return retained_batch_ ? n + retained_batch_->NumActiveStates() * 3 : n; + } + + private: + dvector poses_; + dvector> points_, retained_; + dvector pose_const_, point_const_; + dvector pose_problem_, point_problem_, retained_problem_; + std::unique_ptr pose_batch_; + std::unique_ptr> point_batch_, retained_batch_; + dvector> observations_[3]; + dvector extrinsics_[3]; + std::unique_ptr reprojection_[3]; + CauchyLossFunctionBatch cauchy_{1.f, 0.01f}; + dvector deltas_; + std::unique_ptr between_; + Problem problem_; +}; + +std::vector ToHost(const dvector &v) { + std::vector out(v.size()); + v.CopyToHost(out.data(), out.size()); + return out; +} + +/** ‖a - b‖ / ‖b‖. */ +double RelativeError(const std::vector &a, const std::vector &b) { + double num = 0.0, den = 0.0; + for (size_t i = 0; i < b.size(); ++i) { + num += (static_cast(a[i]) - b[i]) * (static_cast(a[i]) - b[i]); + den += static_cast(b[i]) * b[i]; + } + return std::sqrt(num / std::max(den, 1e-300)); +} + +/** One structure setup and system build of the minimizer for `type`. */ +struct BuiltSystem { + BuiltSystem(cudaStream_t stream, Problem &problem, LinearSystemType type, ColumnScaling scaling) { + options.linear_system = type; + // The exact reference solve for Assembled (SchurComplement has its own solver). + options.sparse_linear_solver_type = SparseLinearSolverType::DenseCholesky; + options.column_scaling = scaling; + system.linear_system = internal::CreateLinearSystem(options); + internal::SetUpStructure(stream, problem, options, state_ops, system); + state.Recreate(stream, problem); + internal::BuildSystem(stream, problem, state, options, system, scratch); + system.step.resize(system.rhs.size()); + } + + MinimizerOptions options; + StateBatchOps state_ops; + internal::MinimizerSystem system; + internal::MinimizerScratch scratch; + MinimizerState state; +}; + +/** A bundle variant, by name. */ +struct Variant { + const char *name; + BundleOptions options; +}; + +const Variant kVariants[] = { + {"Plain", {}}, + {"TwoSubproblems", {2, 0, false}}, + {"RetainedPoints", {1, 30, false}}, +}; + +class SchurSystemTest : public ::testing::TestWithParam> {}; + +/** + * Same rhs, Hessian diagonal and damped (and masked) step as the assembled + * system, and the same step curvature. + */ +TEST_P(SchurSystemTest, MatchesAssembledSystem) { + const auto [variant, scaling] = GetParam(); + SmallBundle bundle(variant.options); + Problem &problem = bundle.GetProblem(); + ASSERT_TRUE(problem.CheckConsistency()); + problem.SetEliminatedStateBatches({bundle.Points()}); + CudaStream cuda_stream; + cudaStream_t stream = cuda_stream.GetStream(); + BuiltSystem assembled(stream, problem, LinearSystemType::Assembled, scaling); + BuiltSystem schur(stream, problem, LinearSystemType::SchurComplement, scaling); + const size_t n = assembled.system.rhs.size(); + ASSERT_EQ(n, schur.system.rhs.size()); + ASSERT_EQ(n, bundle.NumColumns()); + EXPECT_LT(RelativeError(ToHost(schur.system.rhs), ToHost(assembled.system.rhs)), 1e-5); + dvector diagonal_a, diagonal_s; + assembled.system.linear_system->ExtractScaledHessianDiagonal(stream, diagonal_a); + schur.system.linear_system->ExtractScaledHessianDiagonal(stream, diagonal_s); + const std::vector diagonal = ToHost(diagonal_a); + EXPECT_LT(RelativeError(ToHost(diagonal_s), diagonal), 1e-5); + + // LM-like damping, and a mask holding a few pose and point coordinates. + std::vector damping(n), mask(n, 1.f); + for (size_t i = 0; i < n; ++i) damping[i] = 1e-3f * diagonal[i] + 1e-6f; + for (size_t i : {size_t(3), size_t(20), size_t(66 + 4), size_t(66 + 100)}) mask[i] = 0.f; + dvector d_damping(damping), d_mask(mask); + for (bool masked : {false, true}) { + std::vector> steps; + for (BuiltSystem *built : {&assembled, &schur}) { + LinearSystem &system = *built->system.linear_system; + system.SetDamping(&d_damping); + system.SetMask(masked ? &d_mask : nullptr); + dvector rhs(n); + std::vector host_rhs = ToHost(built->system.rhs); + if (masked) { + for (size_t i = 0; i < n; ++i) host_rhs[i] *= mask[i]; + } + rhs.CopyFromHost(host_rhs.data(), n); + dvector step(n); + ASSERT_TRUE(system.PrepareSolver(stream, problem, rhs, step)); + ASSERT_TRUE(system.Solve(stream, rhs, step)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + steps.push_back(ToHost(step)); + } + EXPECT_LT(RelativeError(steps[1], steps[0]), 2e-3) << "masked " << masked; + if (masked) { + for (size_t i = 0; i < n; ++i) { + if (mask[i] == 0.f) EXPECT_EQ(steps[1][i], 0.f) << "held column " << i; + } + } + } + + // δᵀHδ per subproblem. + ProblemPartition partition; + MinimizerState state; + state.Recreate(stream, problem); + partition.Build(stream, problem, state, n); + std::vector host_step(n); + std::mt19937 rng(3); + std::normal_distribution g; + for (float &v : host_step) v = g(rng); + const size_t np = problem.NumProblems(); + dvector step(host_step), curvature(2 * np); + assembled.system.linear_system->StepCurvature(stream, step, partition, curvature.data()); + schur.system.linear_system->StepCurvature(stream, step, partition, curvature.data() + np); + const std::vector c = ToHost(curvature); + for (size_t p = 0; p < np; ++p) EXPECT_NEAR(c[np + p], c[p], 1e-4f * c[p]) << "subproblem " << p; +} + +INSTANTIATE_TEST_SUITE_P(Variants, SchurSystemTest, + ::testing::Combine(::testing::ValuesIn(kVariants), + ::testing::Values(ColumnScaling::None, + ColumnScaling::HessianDiagonal)), + [](const ::testing::TestParamInfo &info) { + return std::string(std::get<0>(info.param).name) + + (std::get<1>(info.param) == ColumnScaling::None ? "Unscaled" + : "Scaled"); + }); + +class SchurSystemMinimizeTest : public ::testing::TestWithParam {}; + +/** The same minimization as the assembled system. */ +TEST_P(SchurSystemMinimizeTest, ConvergesLikeAssembled) { + std::vector final_costs; + for (LinearSystemType type : {LinearSystemType::Assembled, LinearSystemType::SchurComplement}) { + SmallBundle bundle(GetParam().options); + Problem &problem = bundle.GetProblem(); + problem.SetEliminatedStateBatches({bundle.Points()}); + LevenbergMarquardtMinimizerOptions options; + options.base_options.max_num_iterations = 30; + options.base_options.linear_system = type; + options.base_options.sparse_linear_solver_type = SparseLinearSolverType::DenseCholesky; + options.base_options.column_scaling = ColumnScaling::HessianDiagonal; + LevenbergMarquardtMinimizer minimizer(options); + CudaStream stream; + const MinimizerSummary summary = minimizer.Minimize(stream.GetStream(), problem); + EXPECT_LT(summary.final_cost, 1e-2f * summary.initial_cost); + final_costs.push_back(summary.final_cost); + } + EXPECT_NEAR(final_costs[1], final_costs[0], 1e-3f * final_costs[0]); +} + +INSTANTIATE_TEST_SUITE_P(Variants, SchurSystemMinimizeTest, ::testing::ValuesIn(kVariants), + [](const ::testing::TestParamInfo &info) { + return std::string(info.param.name); + }); + +/** + * Runs `minimizer` on a bundle with a singular landmark. Returns whether all + * points stay finite; `moved` receives how far the singular landmark moved. + */ +bool MinimizeSingular(Minimizer &minimizer, MinimizerSummary *summary, float *moved) { + BundleOptions options; + options.singular_landmark = true; + SmallBundle bundle(options); + Problem &problem = bundle.GetProblem(); + problem.SetEliminatedStateBatches({bundle.Points()}); + const VectorStateBatch<3> *points = bundle.PointBatch(); + std::vector before(points->NumActiveStates() * 3), after(before.size()); + THROW_ON_CUDA_ERROR(cudaMemcpy(before.data(), points->StateDevicePtr(0), + before.size() * sizeof(float), cudaMemcpyDeviceToHost)); + CudaStream stream; + *summary = minimizer.Minimize(stream.GetStream(), problem); + THROW_ON_CUDA_ERROR(cudaMemcpy(after.data(), points->StateDevicePtr(0), + after.size() * sizeof(float), cudaMemcpyDeviceToHost)); + *moved = 0.f; + for (size_t i = after.size() - 3; i < after.size(); ++i) { + *moved = std::max(*moved, std::fabs(after[i] - before[i])); + } + return std::all_of(after.begin(), after.end(), [](float v) { return std::isfinite(v); }); +} + +MinimizerOptions SingularOptions(LinearSystemType type, bool safety_checks) { + MinimizerOptions options; + options.max_num_iterations = 10; + options.linear_system = type; + options.sparse_linear_solver_type = SparseLinearSolverType::DenseCholesky; + options.disable_safety_checks = !safety_checks; + return options; +} + +/** + * A landmark whose Hessian block is singular (an exactly zero Jacobian + * column, which no damping of the form λ diag(H) fixes) makes the full + * system singular: with safety checks Gauss-Newton throws; Levenberg-Marquardt + * rejects the failed step instead of throwing (with or without checks), and + * the states stay finite. + */ +TEST(SchurSystemFailureTest, AssembledFailedSolveIsRejected) { + MinimizerSummary summary; + float moved = 0.f; + { + GaussNewtonMinimizer gn(SingularOptions(LinearSystemType::Assembled, true)); + EXPECT_THROW(MinimizeSingular(gn, &summary, &moved), std::runtime_error); + } + for (bool checks : {true, false}) { + LevenbergMarquardtMinimizerOptions options; + options.base_options = SingularOptions(LinearSystemType::Assembled, checks); + LevenbergMarquardtMinimizer lm(options); + bool finite = false; + ASSERT_NO_THROW(finite = MinimizeSingular(lm, &summary, &moved)) << "checks " << checks; + EXPECT_TRUE(finite) << "checks " << checks; + EXPECT_TRUE(std::isfinite(summary.final_cost)); + EXPECT_LE(summary.final_cost, summary.initial_cost); + } +} + +/** + * The Schur complement factors a singular landmark block with a small + * diagonal shift (as Ceres uses a pseudo-inverse): the unobservable landmark + * stays where it is, and both methods converge on the rest of the problem. + */ +TEST(SchurSystemFailureTest, SchurRegularizesSingularLandmarks) { + for (bool checks : {true, false}) { + MinimizerSummary summary; + float moved = 1.f; + GaussNewtonMinimizer gn(SingularOptions(LinearSystemType::SchurComplement, checks)); + ASSERT_NO_THROW(EXPECT_TRUE(MinimizeSingular(gn, &summary, &moved))) << "checks " << checks; + EXPECT_LT(summary.final_cost, 1e-2f * summary.initial_cost); + EXPECT_LT(moved, 1e-3f); + LevenbergMarquardtMinimizerOptions options; + options.base_options = SingularOptions(LinearSystemType::SchurComplement, checks); + LevenbergMarquardtMinimizer lm(options); + ASSERT_NO_THROW(EXPECT_TRUE(MinimizeSingular(lm, &summary, &moved))) << "checks " << checks; + EXPECT_LT(summary.final_cost, 1e-2f * summary.initial_cost); + EXPECT_LT(moved, 1e-3f); + } +} + +/** Setup errors name what to fix. */ +TEST(SchurSystemSetupTest, RejectsInvalidPartitions) { + CudaStream stream; + LevenbergMarquardtMinimizerOptions options; + options.base_options.linear_system = LinearSystemType::SchurComplement; + { + SmallBundle bundle; // no eliminated batches declared + LevenbergMarquardtMinimizer minimizer(options); + EXPECT_THROW(minimizer.Minimize(stream.GetStream(), bundle.GetProblem()), + std::invalid_argument); + } + { + // A factor between two points: they cannot both be eliminated. + SmallBundle bundle; + Problem &problem = bundle.GetProblem(); + dvector> delta(std::vector>{{0.f, 0.f, 0.f}}); + VectorBetweenFactorBatch<3> between(delta.data(), 1); + between.SetNumActiveFactors(1); + VectorStateBatch<3> *points = bundle.PointBatch(); + problem.AddFactorBatch(&between, {points->StateDevicePtr(1), points->StateDevicePtr(2)}); + problem.SetEliminatedStateBatches({bundle.Points()}); + LevenbergMarquardtMinimizer minimizer(options); + EXPECT_THROW(minimizer.Minimize(stream.GetStream(), problem), std::invalid_argument); + } +} + +/** The Schur complement has its own reduced solver: sparse_linear_solver_type is not read. */ +TEST(SchurSystemSetupTest, IgnoresSparseLinearSolverType) { + MinimizerOptions options; + options.linear_system = LinearSystemType::SchurComplement; + options.sparse_linear_solver_type = SparseLinearSolverType::BlockTridiagonal; + EXPECT_NO_THROW(internal::CreateLinearSystem(options)); +} + +} // namespace +} // namespace cunls diff --git a/tests/stereo_loop_ba_test.cpp b/tests/stereo_loop_ba_test.cpp new file mode 100644 index 0000000..8059709 --- /dev/null +++ b/tests/stereo_loop_ba_test.cpp @@ -0,0 +1,774 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * @file stereo_loop_ba_test.cpp + * @brief Synthetic replica of a cuSFM Mapper bundle adjustment on a closed + * stereo loop. + * + * Mirrors the first Mapper bundle adjustment of a cuSFM stereo dataset + * (Galileo "figure8", October 2026; the cuSFM adapter's factor graph): + * 724 vehicle rig poses on two laps of a figure-eight in a cluttered room, a + * forward stereo pair (0.15 m baseline) with fixed camera_from_rig transforms, + * about 90k landmarks mostly 0.5 to 3 m away, 1.2M reprojection factors in + * normalized coordinates weighted by focal / 4 px with a Cauchy loss, a triple + * smooth factor on every three consecutive poses, and the first pose held + * constant. + * + * The initial state is what the front end hands the Mapper: poses integrated + * from smoothly drifting visual odometry (about 0.35 m mean, 0.65 m max error + * after alignment, as on the cuSFM data), every landmark triangulated in those + * poses from the run of frames that created it, and observations from matched + * revisits kept only within cuSFM's 25 px reprojection gate. The map is then + * locally consistent, and the few loop observations that survive the gate are + * almost switched off by the Cauchy loss. Closing the loop is a smooth + * deformation of the whole trajectory along a nearly flat valley of the cost: + * the slowest mode of the normal equations, which an inexact inner linear + * solve resolves only a little per Levenberg-Marquardt iteration. + * + * On the cuSFM data the default BlockSparsePCG (200 iterations, relative + * tolerance 1e-3) hits the 200-iteration cap with about 0.6 m of the + * correction left, while an exact linear solve closes the loop within one + * minimization. Here: cuDSS ends near 6 mm, BlockSparsePCG at 150 to 175 mm + * (max position error against the ground truth after a rigid alignment; no + * scale, the stereo baseline fixes it). + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "cunls/common/cuda_stream.h" +#include "cunls/common/device_vector.h" +#include "cunls/common/helper.h" +#include "cunls/common/types.h" +#include "cunls/factor/information/information_factor_batch.h" +#include "cunls/factor/reprojection_factor_batch.h" +#include "cunls/minimizer/levenberg_marquardt_minimizer.h" +#include "cunls/minimizer/problem.h" +#include "cunls/robustifier/cauchy_loss_function_batch.h" +#include "cunls/state/se3_state_batch.h" +#include "cunls/state/vector_state_batch.h" +#include "tests/triple_smooth_factor_batch.h" + +namespace cunls { +namespace { + +// ----------------------------------------------------------------------------- +// Host rigid-body math in double precision. +// ----------------------------------------------------------------------------- + +using Vec3 = std::array; +using Mat3 = std::array; // row-major + +/** p -> rotation * p + translation. */ +struct Rigid { + Mat3 rotation{1, 0, 0, 0, 1, 0, 0, 0, 1}; + Vec3 translation{0, 0, 0}; +}; + +Vec3 Apply(const Mat3 &m, const Vec3 &v) { + return {m[0] * v[0] + m[1] * v[1] + m[2] * v[2], m[3] * v[0] + m[4] * v[1] + m[5] * v[2], + m[6] * v[0] + m[7] * v[1] + m[8] * v[2]}; +} + +Mat3 Multiply(const Mat3 &a, const Mat3 &b) { + Mat3 out{}; + for (int r = 0; r < 3; ++r) { + for (int c = 0; c < 3; ++c) { + for (int k = 0; k < 3; ++k) out[r * 3 + c] += a[r * 3 + k] * b[k * 3 + c]; + } + } + return out; +} + +Mat3 Transpose(const Mat3 &m) { return {m[0], m[3], m[6], m[1], m[4], m[7], m[2], m[5], m[8]}; } + +Vec3 Apply(const Rigid &t, const Vec3 &p) { + Vec3 out = Apply(t.rotation, p); + for (int i = 0; i < 3; ++i) out[i] += t.translation[i]; + return out; +} + +Rigid Compose(const Rigid &a, const Rigid &b) { + return {Multiply(a.rotation, b.rotation), Apply(a, b.translation)}; +} + +Rigid Inverse(const Rigid &t) { + Rigid out; + out.rotation = Transpose(t.rotation); + const Vec3 rt = Apply(out.rotation, t.translation); + out.translation = {-rt[0], -rt[1], -rt[2]}; + return out; +} + +/** Rotation matrix of a rotation vector (Rodrigues). */ +Mat3 ExpSO3(const Vec3 &w) { + const double theta = std::sqrt(w[0] * w[0] + w[1] * w[1] + w[2] * w[2]); + const double a = theta > 1e-9 ? std::sin(theta) / theta : 1.0; + const double b = theta > 1e-9 ? (1.0 - std::cos(theta)) / (theta * theta) : 0.5; + const Mat3 k{0, -w[2], w[1], w[2], 0, -w[0], -w[1], w[0], 0}; + const Mat3 k2 = Multiply(k, k); + Mat3 out{}; + for (int i = 0; i < 9; ++i) out[i] = (i % 4 == 0 ? 1.0 : 0.0) + a * k[i] + b * k2[i]; + return out; +} + +/** Rotation from yaw (about z), then pitch (about y), then roll (about x). */ +Mat3 YawPitchRoll(double yaw, double pitch, double roll) { + return Multiply(ExpSO3({0, 0, yaw}), Multiply(ExpSO3({0, pitch, 0}), ExpSO3({roll, 0, 0}))); +} + +SE3Transform ToSE3(const Rigid &t) { + SE3Transform out{}; + for (int r = 0; r < 3; ++r) { + for (int c = 0; c < 3; ++c) out[r * 4 + c] = static_cast(t.rotation[r * 3 + c]); + out[r * 4 + 3] = static_cast(t.translation[r]); + } + out[15] = 1.f; + return out; +} + +Vec3 Position(const SE3Transform &t) { return {t[3], t[7], t[11]}; } + +/** + * Translation errors [m] of `estimate` against `truth` after the rigid + * alignment (Kabsch, no scale) that minimizes their squared sum. + */ +std::vector AlignedPositionErrors(const std::vector &estimate, + const std::vector &truth) { + const size_t n = truth.size(); + Vec3 ce{0, 0, 0}, ct{0, 0, 0}; + for (size_t i = 0; i < n; ++i) { + for (int k = 0; k < 3; ++k) { + ce[k] += estimate[i][k] / n; + ct[k] += truth[i][k] / n; + } + } + // Horn's method: the rotation is the unit quaternion of the largest + // eigenvalue of a symmetric 4x4 matrix built from the cross covariance + // C = sum (e - ce)(t - ct)^T; cyclic Jacobi rotations diagonalize it. + Mat3 c{}; + for (size_t i = 0; i < n; ++i) { + for (int r = 0; r < 3; ++r) { + for (int s = 0; s < 3; ++s) { + c[r * 3 + s] += (estimate[i][r] - ce[r]) * (truth[i][s] - ct[s]); + } + } + } + const double sxx = c[0], sxy = c[1], sxz = c[2], syx = c[3], syy = c[4], syz = c[5], szx = c[6], + szy = c[7], szz = c[8]; + double a[4][4] = {{sxx + syy + szz, syz - szy, szx - sxz, sxy - syx}, + {syz - szy, sxx - syy - szz, sxy + syx, szx + sxz}, + {szx - sxz, sxy + syx, -sxx + syy - szz, syz + szy}, + {sxy - syx, szx + sxz, syz + szy, -sxx - syy + szz}}; + double v[4][4] = {{1, 0, 0, 0}, {0, 1, 0, 0}, {0, 0, 1, 0}, {0, 0, 0, 1}}; + for (int sweep = 0; sweep < 50; ++sweep) { + for (int p = 0; p < 3; ++p) { + for (int q = p + 1; q < 4; ++q) { + if (std::fabs(a[p][q]) < 1e-300) continue; + const double theta = 0.5 * std::atan2(2.0 * a[p][q], a[q][q] - a[p][p]); + const double cs = std::cos(theta), sn = std::sin(theta); + for (int k = 0; k < 4; ++k) { // A <- A J + const double akp = a[k][p], akq = a[k][q]; + a[k][p] = cs * akp - sn * akq; + a[k][q] = sn * akp + cs * akq; + } + for (int k = 0; k < 4; ++k) { // A <- J^T A + const double apk = a[p][k], aqk = a[q][k]; + a[p][k] = cs * apk - sn * aqk; + a[q][k] = sn * apk + cs * aqk; + } + for (int k = 0; k < 4; ++k) { // V <- V J + const double vkp = v[k][p], vkq = v[k][q]; + v[k][p] = cs * vkp - sn * vkq; + v[k][q] = sn * vkp + cs * vkq; + } + } + } + } + int best = 0; + for (int k = 1; k < 4; ++k) { + if (a[k][k] > a[best][best]) best = k; + } + const double q[4] = {v[0][best], v[1][best], v[2][best], v[3][best]}; + const double w = q[0], x = q[1], y = q[2], z = q[3]; + const Mat3 rot{1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y), + 2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x), + 2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)}; + std::vector errors(n); + for (size_t i = 0; i < n; ++i) { + const Vec3 d{estimate[i][0] - ce[0], estimate[i][1] - ce[1], estimate[i][2] - ce[2]}; + const Vec3 aligned = Apply(rot, d); + double sq = 0.0; + for (int k = 0; k < 3; ++k) sq += std::pow(aligned[k] + ct[k] - truth[i][k], 2); + errors[i] = std::sqrt(sq); + } + return errors; +} + +// ----------------------------------------------------------------------------- +// Scene. +// ----------------------------------------------------------------------------- + +constexpr int kNumPoses = 724; +constexpr double kLaps = 2.0; +constexpr double kHalfWidth = 2.75; // x extent of the figure eight [m] +constexpr double kLength = 4.3; // y extent [m] +constexpr double kRigHeight = 1.0; +constexpr double kBaseline = 0.15; // stereo baseline [m] +// Camera model: pixels per normalized unit, image half extents in normalized +// coordinates (1280 x 800 at f = 830). +constexpr double kFocal = 830.0; +constexpr double kHalfFovX = 640.0 / kFocal; +constexpr double kHalfFovY = 400.0 / kFocal; +constexpr double kMinDepth = 0.3, kMaxDepth = 10.0; +constexpr double kPixelNoise = 0.6; // observation noise [px] +constexpr double kSigmaPixels = 4.0; // reprojection_error_stddev_pixels of cuSFM +constexpr double kOutlierRatio = 0.005; // observations moved by up to 30 px +// Room around the trajectory and boxes standing in it: landmarks on their +// surfaces, mostly 0.5 to 3 m from the cameras. +constexpr double kRoomX = 4.0, kRoomY = 3.4, kRoomZ = 2.6; +constexpr int kNumBoxes = 60; +constexpr double kBoxClearance = 0.35; // from the trajectory [m] +constexpr int kNumLandmarks = 45000; +// Feature tracks: observed for a geometric number of consecutive frames +// within each visibility run (one in five features tracks longer); a revisit +// is matched to the landmark with this probability, otherwise the front end +// creates a new landmark for it. +constexpr double kMeanTrackLength = 3.0; +constexpr double kLongTrackProbability = 0.2; // tracks of a well-textured feature +constexpr double kMeanLongTrackLength = 25.0; +constexpr double kRevisitMatchProbability = 0.5; +constexpr int kMinObservedFrames = 2; +// cuSFM's map initialization: observations further than this from the +// triangulated point are dropped, and a point needs this many observations. +constexpr double kInitialMaxReprojectionErrorPixels = 25.0; +constexpr size_t kMinObservations = 3; +// A landmark observed more than this many frames apart links a loop. +constexpr int kLoopFrameSpan = 100; +// Visual odometry drifts smoothly: its rotation bias wanders (random walk), +// its scale is slightly off, and the white noise per frame is small, so the +// initial map is locally consistent (about 1 px) while the loop is not. +constexpr double kDriftYawBias = 1.1e-3; // initial bias [rad / frame] +constexpr double kDriftBiasWalk = 3e-5; // bias random walk [rad / frame] +constexpr double kDriftScale = 4e-3; // translation scale error +constexpr double kDriftRotationNoise = 3e-5; // [rad] +constexpr double kDriftTranslationNoise = 2e-4; // [m] +// cuSFM relative_pose_*_stddev for the triple smooth factors. +constexpr double kSmoothRotationSigma = 1.0 * M_PI / 180.0; +constexpr double kSmoothTranslationSigma = 0.03; +constexpr float kZThreshold = 1e-3f; + +/** Ground truth, initial values and measurements of the synthetic problem. */ +struct StereoLoopScene { + std::vector truth_poses; ///< world_from_rig. + std::vector initial_poses; ///< Drifted world_from_rig. + std::vector> initial_points; + std::vector> observations; ///< Normalized coordinates. + std::vector camera_from_rig; ///< Per observation. + std::vector pose_ids, point_ids; ///< Per observation. + std::vector camera_ids; ///< Per observation: 0 left, 1 right. + size_t num_loop_points = 0; ///< Landmarks observed over more than kLoopFrameSpan frames. +}; + +/** Camera i (0 left, 1 right) of the forward stereo pair, camera_from_rig. */ +Rigid CameraFromRig(int camera) { + // Camera z forward (rig x), x right (rig -y), y down (rig -z). + Rigid rig_from_camera; + rig_from_camera.rotation = {0, 0, 1, -1, 0, 0, 0, -1, 0}; + rig_from_camera.translation = {0.1, camera == 0 ? 0.5 * kBaseline : -0.5 * kBaseline, 0.336}; + return Inverse(rig_from_camera); +} + +/** + * Least-squares intersection of the observation rays (midpoint method): + * minimizes the sum of squared distances of `point` to the rays through the + * camera centers. `track` entries need frame, camera, u and v. + */ +template +bool Triangulate(const Track &track, const std::vector (&world_from_camera)[2], + Vec3 *point) { + double a[9] = {}, b[3] = {}; + for (const auto &o : track) { + const Rigid &t = world_from_camera[o.camera][o.frame]; + Vec3 d = Apply(t.rotation, {o.u, o.v, 1.0}); + const double norm = std::sqrt(d[0] * d[0] + d[1] * d[1] + d[2] * d[2]); + for (double &x : d) x /= norm; + for (int r = 0; r < 3; ++r) { + for (int c = 0; c < 3; ++c) { + const double m = (r == c ? 1.0 : 0.0) - d[r] * d[c]; + a[r * 3 + c] += m; + b[r] += m * t.translation[c]; + } + } + } + const double det = a[0] * (a[4] * a[8] - a[5] * a[7]) - a[1] * (a[3] * a[8] - a[5] * a[6]) + + a[2] * (a[3] * a[7] - a[4] * a[6]); + if (!(std::fabs(det) > 1e-12)) return false; + const double inv[9] = { + a[4] * a[8] - a[5] * a[7], a[2] * a[7] - a[1] * a[8], a[1] * a[5] - a[2] * a[4], + a[5] * a[6] - a[3] * a[8], a[0] * a[8] - a[2] * a[6], a[2] * a[3] - a[0] * a[5], + a[3] * a[7] - a[4] * a[6], a[1] * a[6] - a[0] * a[7], a[0] * a[4] - a[1] * a[3]}; + for (int r = 0; r < 3; ++r) { + (*point)[r] = (inv[r * 3] * b[0] + inv[r * 3 + 1] * b[1] + inv[r * 3 + 2] * b[2]) / det; + } + return true; +} + +StereoLoopScene MakeScene() { + std::mt19937 rng(20261015); + std::normal_distribution gauss(0.0, 1.0); + std::uniform_real_distribution uniform(0.0, 1.0); + StereoLoopScene scene; + + // Two laps of a Gerono lemniscate, heading along the path, gentle pitch and + // roll from the floor. + for (int k = 0; k < kNumPoses; ++k) { + const double s = 2.0 * M_PI * kLaps * k / kNumPoses; + const double x = kHalfWidth * std::sin(s); + const double y = kLength * std::sin(s) * std::cos(s); + const double dx = kHalfWidth * std::cos(s), dy = kLength * std::cos(2.0 * s); + Rigid pose; + pose.rotation = + YawPitchRoll(std::atan2(dy, dx), 0.02 * std::sin(3.0 * s), 0.015 * std::cos(5.0 * s)); + pose.translation = {x, y, kRigHeight + 0.03 * std::sin(2.0 * s)}; + scene.truth_poses.push_back(pose); + } + + // Odometry: relative motions with a wandering rotation bias, a scale error + // and small white noise, integrated from the (exact) first pose. + scene.initial_poses.push_back(scene.truth_poses[0]); + Vec3 bias{0.2 * kDriftYawBias, -0.1 * kDriftYawBias, kDriftYawBias}; + for (int k = 1; k < kNumPoses; ++k) { + Rigid motion = Compose(Inverse(scene.truth_poses[k - 1]), scene.truth_poses[k]); + for (double &b : bias) b += kDriftBiasWalk * gauss(rng); + Vec3 w; + for (int i = 0; i < 3; ++i) w[i] = bias[i] + kDriftRotationNoise * gauss(rng); + motion.rotation = Multiply(motion.rotation, ExpSO3(w)); + for (double &t : motion.translation) { + t = (1.0 + kDriftScale) * t + kDriftTranslationNoise * gauss(rng); + } + scene.initial_poses.push_back(Compose(scene.initial_poses.back(), motion)); + } + + const Rigid camera_from_rig[2] = {CameraFromRig(0), CameraFromRig(1)}; + const SE3Transform camera_from_rig_se3[2] = {ToSE3(camera_from_rig[0]), + ToSE3(camera_from_rig[1])}; + std::vector camera_from_world[2]; + for (int c = 0; c < 2; ++c) { + for (const Rigid &pose : scene.truth_poses) { + camera_from_world[c].push_back(Compose(camera_from_rig[c], Inverse(pose))); + } + } + std::vector drifted_camera_from_world[2], drifted_world_from_camera[2]; + for (int c = 0; c < 2; ++c) { + for (const Rigid &pose : scene.initial_poses) { + drifted_camera_from_world[c].push_back(Compose(camera_from_rig[c], Inverse(pose))); + drifted_world_from_camera[c].push_back(Inverse(drifted_camera_from_world[c].back())); + } + } + auto project = [&](int camera, int frame, const Vec3 &p, Vec3 &out) { + const Vec3 pc = Apply(camera_from_world[camera][frame], p); + if (!(pc[2] > kMinDepth && pc[2] < kMaxDepth)) return false; + const double u = pc[0] / pc[2], v = pc[1] / pc[2]; + if (std::fabs(u) > kHalfFovX || std::fabs(v) > kHalfFovY) return false; + out = {u, v, pc[2]}; + return true; + }; + + // Landmark surfaces: rectangles origin + a u + b v, a, b in [0, 1]. + struct Surface { + Vec3 origin, u, v; + }; + std::vector surfaces = { + {{-kRoomX, -kRoomY, 0}, {2 * kRoomX, 0, 0}, {0, 2 * kRoomY, 0}}, // floor + {{-kRoomX, -kRoomY, kRoomZ}, {2 * kRoomX, 0, 0}, {0, 2 * kRoomY, 0}}, // ceiling + {{-kRoomX, -kRoomY, 0}, {2 * kRoomX, 0, 0}, {0, 0, kRoomZ}}, + {{-kRoomX, kRoomY, 0}, {2 * kRoomX, 0, 0}, {0, 0, kRoomZ}}, + {{-kRoomX, -kRoomY, 0}, {0, 2 * kRoomY, 0}, {0, 0, kRoomZ}}, + {{kRoomX, -kRoomY, 0}, {0, 2 * kRoomY, 0}, {0, 0, kRoomZ}}}; + // Boxes on the floor, clear of the trajectory: their top and four sides. + for (int placed = 0; placed < kNumBoxes;) { + const double sx = 0.3 + 0.7 * uniform(rng), sy = 0.3 + 0.7 * uniform(rng), + sz = 0.3 + 1.3 * uniform(rng); + const double cx = (2.0 * uniform(rng) - 1.0) * (kRoomX - sx), + cy = (2.0 * uniform(rng) - 1.0) * (kRoomY - sy); + double clearance = 1e9; + for (const Rigid &pose : scene.truth_poses) { + clearance = + std::min(clearance, std::hypot(pose.translation[0] - cx, pose.translation[1] - cy)); + } + if (clearance < kBoxClearance + 0.5 * std::hypot(sx, sy)) continue; + const Vec3 o{cx - 0.5 * sx, cy - 0.5 * sy, 0.0}; + surfaces.push_back({{o[0], o[1], sz}, {sx, 0, 0}, {0, sy, 0}}); + surfaces.push_back({o, {sx, 0, 0}, {0, 0, sz}}); + surfaces.push_back({{o[0], o[1] + sy, 0}, {sx, 0, 0}, {0, 0, sz}}); + surfaces.push_back({o, {0, sy, 0}, {0, 0, sz}}); + surfaces.push_back({{o[0] + sx, o[1], 0}, {0, sy, 0}, {0, 0, sz}}); + ++placed; + } + std::vector area_cdf; + for (const Surface &f : surfaces) { + const Vec3 &u = f.u, &v = f.v; + const Vec3 n{u[1] * v[2] - u[2] * v[1], u[2] * v[0] - u[0] * v[2], u[0] * v[1] - u[1] * v[0]}; + area_cdf.push_back((area_cdf.empty() ? 0.0 : area_cdf.back()) + + std::sqrt(n[0] * n[0] + n[1] * n[1] + n[2] * n[2])); + } + std::geometric_distribution track_extra(1.0 / kMeanTrackLength); + std::geometric_distribution long_track_extra(1.0 / kMeanLongTrackLength); + std::vector visible; + for (int i = 0; i < kNumLandmarks; ++i) { + const double pick = uniform(rng) * area_cdf.back(); + const Surface &f = + surfaces[std::lower_bound(area_cdf.begin(), area_cdf.end(), pick) - area_cdf.begin()]; + const double a = uniform(rng), b = uniform(rng); + const Vec3 p{f.origin[0] + a * f.u[0] + b * f.v[0], f.origin[1] + a * f.u[1] + b * f.v[1], + f.origin[2] + a * f.u[2] + b * f.v[2]}; + + // Frames whose left camera sees it, split into runs of consecutive frames. + visible.clear(); + Vec3 uvz; + for (int k = 0; k < kNumPoses; ++k) { + if (project(0, k, p, uvz)) visible.push_back(k); + } + std::vector> runs; // [first, last] + for (size_t j = 0; j < visible.size(); ++j) { + if (j == 0 || visible[j] != visible[j - 1] + 1) runs.push_back({visible[j], visible[j]}); + runs.back().second = visible[j]; + } + // One track per run: a random stretch of geometric length. Later runs + // (revisits) are matched to the landmark with kRevisitMatchProbability; + // the front end makes an unmatched one a new landmark of its own. + // Per landmark: its frames, the first `local` of them from its first run. + struct Landmark { + std::vector frames; + size_t local = 0; + }; + std::vector landmarks(1); + const bool long_track = uniform(rng) < kLongTrackProbability; + for (size_t r = 0; r < runs.size(); ++r) { + Landmark &landmark = r == 0 || uniform(rng) < kRevisitMatchProbability + ? landmarks.front() + : landmarks.emplace_back(); + const int run_length = runs[r].second - runs[r].first + 1; + const int length = + std::min(run_length, 2 + (long_track ? long_track_extra(rng) : track_extra(rng))); + const int start = runs[r].first + static_cast(uniform(rng) * (run_length - length + 1)); + for (int k = start; k < start + length; ++k) landmark.frames.push_back(k); + if (landmark.local == 0) landmark.local = landmark.frames.size(); + } + for (const Landmark &landmark : landmarks) { + if (static_cast(landmark.frames.size()) < kMinObservedFrames) continue; + struct Observation { + int frame, camera; + double u, v; + }; + std::vector track, local; + for (size_t j = 0; j < landmark.frames.size(); ++j) { + const int k = landmark.frames[j]; + for (int c = 0; c < 2; ++c) { + if (!project(c, k, p, uvz)) continue; + double u = uvz[0] + kPixelNoise / kFocal * gauss(rng); + double v = uvz[1] + kPixelNoise / kFocal * gauss(rng); + if (uniform(rng) < kOutlierRatio) { + u += 30.0 / kFocal * (2.0 * uniform(rng) - 1.0); + v += 30.0 / kFocal * (2.0 * uniform(rng) - 1.0); + } + track.push_back({k, c, u, v}); + if (j < landmark.local) local.push_back(track.back()); + } + } + // The front end's map: the point triangulated in the drifted poses from + // the run that created it (locally consistent), the observations of + // matched revisits added when within the reprojection gate. + Vec3 init{0, 0, 0}; + if (local.size() < 2 || !Triangulate(local, drifted_world_from_camera, &init)) continue; + std::vector kept; + for (const Observation &o : track) { + const Vec3 pc = Apply(drifted_camera_from_world[o.camera][o.frame], init); + if (pc[2] > kMinDepth && kFocal * std::hypot(pc[0] / pc[2] - o.u, pc[1] / pc[2] - o.v) <= + kInitialMaxReprojectionErrorPixels) { + kept.push_back(o); + } + } + track.swap(kept); + if (track.size() < kMinObservations) continue; + + const int point_id = static_cast(scene.initial_points.size()); + int first = kNumPoses, last = -1; + for (const Observation &o : track) { + scene.observations.push_back({static_cast(o.u), static_cast(o.v)}); + scene.camera_from_rig.push_back(camera_from_rig_se3[o.camera]); + scene.pose_ids.push_back(o.frame); + scene.camera_ids.push_back(o.camera); + scene.point_ids.push_back(point_id); + first = std::min(first, o.frame); + last = std::max(last, o.frame); + } + scene.initial_points.push_back( + {static_cast(init[0]), static_cast(init[1]), static_cast(init[2])}); + scene.num_loop_points += last - first > kLoopFrameSpan ? 1 : 0; + } + } + return scene; +} + +/** Device data, batches and problem of one StereoLoopScene. */ +class StereoLoopProblem { + public: + explicit StereoLoopProblem(const StereoLoopScene &scene) { + const size_t num_obs = scene.observations.size(); + std::vector poses; + for (const Rigid &pose : scene.initial_poses) poses.push_back(ToSE3(pose)); + poses_ = dvector(poses); + points_ = dvector>(scene.initial_points); + const_pose_ids_ = dvector(std::vector{0}); + pose_batch_ = std::make_unique(reinterpret_cast(poses_.data()), + poses.size(), const_pose_ids_.data(), 1); + pose_batch_->SetNumActiveStates(pose_batch_->Capacity(), pose_batch_->ConstCapacity()); + point_batch_ = std::make_unique>( + reinterpret_cast(points_.data()), scene.initial_points.size()); + point_batch_->SetNumActiveStates(point_batch_->Capacity()); + problem_.AddStateBatch(pose_batch_.get()); + problem_.AddStateBatch(point_batch_.get()); + + // Reprojection: residuals in normalized coordinates scaled to pixels / + // sigma, Cauchy loss of scale 1 (cuSFM's mapper defaults). + const float info = static_cast(kFocal / kSigmaPixels); + observations_ = dvector>(scene.observations); + camera_from_rig_ = dvector(scene.camera_from_rig); + reprojection_info_ = dvector>(std::vector>(num_obs, {info, 0, 0, info})); + reprojection_ = std::make_unique>( + reprojection_info_.data(), num_obs, observations_.data(), camera_from_rig_.data(), num_obs, + kZThreshold); + reprojection_->SetNumActiveFactors(num_obs); + std::vector reprojection_states; + reprojection_states.reserve(2 * num_obs); + for (size_t i = 0; i < num_obs; ++i) { + reprojection_states.push_back(pose_batch_->StateDevicePtr(scene.pose_ids[i])); + reprojection_states.push_back(point_batch_->StateDevicePtr(scene.point_ids[i])); + } + cauchy_ = std::make_unique(1.f, 1.f); + problem_.AddFactorBatch(reprojection_.get(), cauchy_.get(), reprojection_states); + + // Triple smooth factors over consecutive poses (uniform timestamps: + // ratio 1/2), information from cuSFM's relative pose sigmas. + const size_t num_smooth = kNumPoses - 2; + Matrix<6> smooth_info{}; + for (int i = 0; i < 3; ++i) { + smooth_info[i * 7] = static_cast(1.0 / kSmoothRotationSigma); + smooth_info[(i + 3) * 7] = static_cast(1.0 / kSmoothTranslationSigma); + } + smooth_ratios_ = dvector(std::vector(num_smooth, 0.5f)); + smooth_info_ = dvector>(std::vector>(num_smooth, smooth_info)); + smooth_ = std::make_unique>( + smooth_info_.data(), num_smooth, smooth_ratios_.data(), num_smooth); + smooth_->SetNumActiveFactors(num_smooth); + std::vector smooth_states; + for (size_t k = 0; k < num_smooth; ++k) { + for (size_t j = 0; j < 3; ++j) smooth_states.push_back(pose_batch_->StateDevicePtr(k + j)); + } + problem_.AddFactorBatch(smooth_.get(), smooth_states); + } + + Problem &GetProblem() { return problem_; } + + /** The landmark state batch (eliminated by LinearSystemType::SchurComplement). */ + const StateBatch *Points() const { return point_batch_.get(); } + + /** Current pose positions (synchronous copy). */ + std::vector Positions() const { + std::vector poses(poses_.size()); + poses_.CopyToHost(poses.data(), poses.size()); + std::vector out; + for (const SE3Transform &pose : poses) out.push_back(Position(pose)); + return out; + } + + private: + dvector poses_; + dvector> points_; + dvector const_pose_ids_; + dvector> observations_; + dvector camera_from_rig_; + dvector> reprojection_info_; + dvector smooth_ratios_; + dvector> smooth_info_; + std::unique_ptr pose_batch_; + std::unique_ptr> point_batch_; + std::unique_ptr> reprojection_; + std::unique_ptr cauchy_; + std::unique_ptr> smooth_; + Problem problem_; +}; + +/** Mean and max of a vector. */ +std::pair MeanMax(const std::vector &values) { + double sum = 0.0, max = 0.0; + for (double v : values) { + sum += v; + max = std::max(max, v); + } + return {sum / values.size(), max}; +} + +class StereoLoopBaTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { scene_ = std::make_unique(MakeScene()); } + static void TearDownTestSuite() { scene_.reset(); } + + static std::vector Positions(const std::vector &poses) { + std::vector out; + for (const Rigid &pose : poses) out.push_back(pose.translation); + return out; + } + std::vector TruthPositions() const { return Positions(scene_->truth_poses); } + std::vector InitialPositions() const { return Positions(scene_->initial_poses); } + + /** + * One cuSFM-like Levenberg-Marquardt solve (the adapter's options) with the + * given linear solver, from the drifted initial state. Returns the summary; + * `before` and `after` receive the aligned position errors [m]. + */ + MinimizerSummary Solve(SparseLinearSolverType solver, std::vector *before, + std::vector *after, + LinearSystemType linear_system = LinearSystemType::Assembled) { + StereoLoopProblem setup(*scene_); + Problem &problem = setup.GetProblem(); + EXPECT_TRUE(problem.CheckConsistency()); + problem.SetEliminatedStateBatches({setup.Points()}); + *before = AlignedPositionErrors(setup.Positions(), TruthPositions()); + + MinimizerOptions options; + options.max_num_iterations = 200; + options.cost_tolerance = 1e-9f; + options.state_tolerance = 1e-9f; + options.column_scaling = ColumnScaling::HessianDiagonal; + options.sparse_linear_solver_type = solver; + options.linear_system = linear_system; + LevenbergMarquardtMinimizerOptions lm_options; + lm_options.base_options = options; + lm_options.relative_reduction_tolerance = 1e-9f; + LevenbergMarquardtMinimizer minimizer(lm_options); + + CudaStream stream; + const auto start = std::chrono::steady_clock::now(); + const MinimizerSummary summary = minimizer.Minimize(stream.GetStream(), problem); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + const double seconds = + std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + *after = AlignedPositionErrors(setup.Positions(), TruthPositions()); + const auto [mean0, max0] = MeanMax(*before); + const auto [mean1, max1] = MeanMax(*after); + std::printf( + "[stereo loop BA] %zu poses, %zu points (%zu spanning > 100 frames), %zu observations\n" + " iterations %zu, cost %.6g -> %.6g, %.2f s\n" + " position error vs truth [mm]: mean %.1f max %.1f -> mean %.1f max %.1f\n", + scene_->truth_poses.size(), scene_->initial_points.size(), scene_->num_loop_points, + scene_->observations.size(), summary.num_iterations, summary.initial_cost, + summary.final_cost, seconds, 1e3 * mean0, 1e3 * max0, 1e3 * mean1, 1e3 * max1); + RecordProperty("iterations", static_cast(summary.num_iterations)); + RecordProperty("max_error_mm", static_cast(std::lround(1e3 * max1))); + return summary; + } + + static std::unique_ptr scene_; +}; + +std::unique_ptr StereoLoopBaTest::scene_; + +/** The scene has the scale and structure of the cuSFM problem. */ +TEST_F(StereoLoopBaTest, SceneMatchesCuSfmScale) { + EXPECT_EQ(scene_->truth_poses.size(), 724u); + EXPECT_GT(scene_->initial_points.size(), 50000u); + EXPECT_GT(scene_->observations.size(), 1000000u); + EXPECT_LT(scene_->observations.size(), 1500000u); + EXPECT_GT(scene_->num_loop_points, 100u) << "a few observations link the loop"; + std::vector per_frame(kNumPoses, 0); + for (int k : scene_->pose_ids) ++per_frame[k]; + EXPECT_GE(*std::min_element(per_frame.begin(), per_frame.end()), 50) + << "every frame needs a usable number of observations (cuSFM's fewest: 54)"; + // The front end's map is locally consistent: small initial reprojection + // errors (cuSFM's: median 1.05 px, p90 4.5 px). + std::vector errors; + for (size_t i = 0; i < scene_->observations.size(); ++i) { + const Vector<3> &x = scene_->initial_points[scene_->point_ids[i]]; + const Rigid camera_from_world = Compose(CameraFromRig(scene_->camera_ids[i]), + Inverse(scene_->initial_poses[scene_->pose_ids[i]])); + const Vec3 pc = Apply(camera_from_world, {x[0], x[1], x[2]}); + errors.push_back(kFocal * std::hypot(pc[0] / pc[2] - scene_->observations[i][0], + pc[1] / pc[2] - scene_->observations[i][1])); + } + std::sort(errors.begin(), errors.end()); + const double p50 = errors[errors.size() / 2], p90 = errors[errors.size() * 9 / 10]; + std::printf("[stereo loop BA] initial reprojection error [px]: p50 %.2f p90 %.2f\n", p50, p90); + EXPECT_LT(p50, 2.5); + EXPECT_LT(p90, 10.0); + const auto [mean, max] = MeanMax(AlignedPositionErrors(InitialPositions(), TruthPositions())); + EXPECT_GT(mean, 0.25) << "the initial drift must be loop-scale"; + EXPECT_GT(max, 0.5); + EXPECT_LT(max, 1.0); +} + +#ifdef CUNLS_ENABLE_CUDSS +/** An exact linear solve closes the loop within one solve. */ +TEST_F(StereoLoopBaTest, CuDssClosesTheLoop) { + std::vector before, after; + const MinimizerSummary summary = Solve(SparseLinearSolverType::cuDSS, &before, &after); + EXPECT_LT(summary.final_cost, summary.initial_cost); + EXPECT_LT(summary.num_iterations, 200u) << "LM should converge before the iteration cap"; + EXPECT_LT(MeanMax(after).second, 0.02) << "max position error [m]"; +} +#endif + +/** + * The landmarks eliminated, the reduced camera system (4338 columns) solved + * by the supernodal Cholesky: exact like cuDSS on the full system. + */ +TEST_F(StereoLoopBaTest, SchurClosesTheLoop) { + std::vector before, after; + const MinimizerSummary summary = Solve(SparseLinearSolverType::BlockSparsePCG, &before, &after, + LinearSystemType::SchurComplement); + EXPECT_LT(summary.final_cost, summary.initial_cost); + EXPECT_LT(summary.num_iterations, 200u) << "LM should converge before the iteration cap"; + EXPECT_LT(MeanMax(after).second, 0.02) << "max position error [m]"; +} + +/** + * The default inner solver (block-Jacobi PCG on the full system, 200 + * iterations, relative tolerance 1e-3) under-solves the loop closure mode and + * creeps: it reaches the iteration cap with most of the drift left (use the + * Schur complement or cuDSS for such problems). Disabled; run it with + * --gtest_also_run_disabled_tests to see the gap. + */ +TEST_F(StereoLoopBaTest, DISABLED_BlockSparsePcgClosesTheLoop) { + std::vector before, after; + const MinimizerSummary summary = Solve(SparseLinearSolverType::BlockSparsePCG, &before, &after); + EXPECT_LT(summary.final_cost, summary.initial_cost); + EXPECT_LT(MeanMax(after).second, 0.02) << "max position error [m]"; +} + +} // namespace +} // namespace cunls diff --git a/tests/synthetic_sba_test.cpp b/tests/synthetic_sba_test.cpp index 7e408b8..017586b 100644 --- a/tests/synthetic_sba_test.cpp +++ b/tests/synthetic_sba_test.cpp @@ -26,7 +26,9 @@ #include #include +#include #include +#include #include #include #include @@ -327,6 +329,9 @@ TEST_P(SyntheticSbaTest, Optimize) { test_utils::PCGMaxIterFromEnv(200); options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); + if (test_utils::UseSchurComplementFromEnv(options)) { + problem.SetEliminatedStateBatches({&point_batch}); + } LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -335,8 +340,14 @@ TEST_P(SyntheticSbaTest, Optimize) { MinimizerSummary summary; { auto range = profiler_domain_.CreateDomainRange("Minimize"); + const auto start = std::chrono::steady_clock::now(); summary = minimizer.Minimize(stream.GetStream(), problem); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + const double ms = + std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + std::printf( + "[synthetic SBA] %s: %zu observations, %zu iterations, cost %.6g -> %.6g, %.1f ms\n", + params.label, n_obs, summary.num_iterations, summary.initial_cost, summary.final_cost, ms); } EXPECT_TRUE(std::isfinite(summary.initial_cost)); diff --git a/tests/triple_smooth_factor_batch.cu b/tests/triple_smooth_factor_batch.cu new file mode 100644 index 0000000..9dbf85e --- /dev/null +++ b/tests/triple_smooth_factor_batch.cu @@ -0,0 +1,279 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include + +#include "tests/triple_smooth_factor_batch.h" + +namespace cunls { +namespace test_utils { +namespace { +constexpr int kThreadsPerBlock = 128; +// Step of the central differences, in tangent units (radians, meters). +constexpr double kDifferenceStep = 1e-6; + +// Rigid transform, double precision: p -> rotation * p + translation. +struct Pose { + double rotation[9]; + double translation[3]; +}; + +__device__ Pose LoadPose(const float* matrix) { + Pose pose; + for (int row = 0; row < 3; ++row) { + for (int col = 0; col < 3; ++col) { + pose.rotation[row * 3 + col] = matrix[row * 4 + col]; + } + pose.translation[row] = matrix[row * 4 + 3]; + } + return pose; +} + +__device__ Pose Multiply(const Pose& lhs, const Pose& rhs) { + Pose out; + for (int row = 0; row < 3; ++row) { + for (int col = 0; col < 3; ++col) { + double sum = 0.0; + for (int k = 0; k < 3; ++k) { + sum += lhs.rotation[row * 3 + k] * rhs.rotation[k * 3 + col]; + } + out.rotation[row * 3 + col] = sum; + } + double sum = lhs.translation[row]; + for (int k = 0; k < 3; ++k) { + sum += lhs.rotation[row * 3 + k] * rhs.translation[k]; + } + out.translation[row] = sum; + } + return out; +} + +__device__ Pose Inverse(const Pose& pose) { + Pose out; + for (int row = 0; row < 3; ++row) { + for (int col = 0; col < 3; ++col) { + out.rotation[row * 3 + col] = pose.rotation[col * 3 + row]; + } + } + for (int row = 0; row < 3; ++row) { + double sum = 0.0; + for (int k = 0; k < 3; ++k) { + sum -= out.rotation[row * 3 + k] * pose.translation[k]; + } + out.translation[row] = sum; + } + return out; +} + +// out = I + a * K + b * K^2 for the skew matrix K of `axis`. +__device__ void SkewSeries(const double* axis, double a, double b, double* out) { + const double k[9] = {0.0, -axis[2], axis[1], axis[2], 0.0, -axis[0], -axis[1], axis[0], 0.0}; + for (int row = 0; row < 3; ++row) { + for (int col = 0; col < 3; ++col) { + double square = 0.0; + for (int i = 0; i < 3; ++i) { + square += k[row * 3 + i] * k[i * 3 + col]; + } + out[row * 3 + col] = (row == col ? 1.0 : 0.0) + a * k[row * 3 + col] + b * square; + } + } +} + +// Rotation matrix of the rotation vector `w` (Rodrigues). +__device__ void ExpSO3(const double* w, double* rotation) { + const double theta2 = w[0] * w[0] + w[1] * w[1] + w[2] * w[2]; + const double theta = sqrt(theta2); + double a = 1.0 - theta2 / 6.0; + double b = 0.5 - theta2 / 24.0; + if (theta > 1e-4) { + a = sin(theta) / theta; + b = (1.0 - cos(theta)) / theta2; + } + SkewSeries(w, a, b, rotation); +} + +// Rotation vector of a rotation matrix, through the quaternion (the branch of +// its largest component), so that it is exact for every angle. +__device__ void LogSO3(const double* r, double* w) { + const double trace = r[0] + r[4] + r[8]; + double q[4]; // w, x, y, z + if (trace > 0.0) { + const double s = 2.0 * sqrt(trace + 1.0); + q[0] = 0.25 * s; + q[1] = (r[7] - r[5]) / s; + q[2] = (r[2] - r[6]) / s; + q[3] = (r[3] - r[1]) / s; + } else if (r[0] >= r[4] && r[0] >= r[8]) { + const double s = 2.0 * sqrt(1.0 + r[0] - r[4] - r[8]); + q[0] = (r[7] - r[5]) / s; + q[1] = 0.25 * s; + q[2] = (r[1] + r[3]) / s; + q[3] = (r[2] + r[6]) / s; + } else if (r[4] >= r[8]) { + const double s = 2.0 * sqrt(1.0 + r[4] - r[0] - r[8]); + q[0] = (r[2] - r[6]) / s; + q[1] = (r[1] + r[3]) / s; + q[2] = 0.25 * s; + q[3] = (r[5] + r[7]) / s; + } else { + const double s = 2.0 * sqrt(1.0 + r[8] - r[0] - r[4]); + q[0] = (r[3] - r[1]) / s; + q[1] = (r[2] + r[6]) / s; + q[2] = (r[5] + r[7]) / s; + q[3] = 0.25 * s; + } + if (q[0] < 0.0) { + for (int i = 0; i < 4; ++i) { + q[i] = -q[i]; + } + } + const double norm = sqrt(q[1] * q[1] + q[2] * q[2] + q[3] * q[3]); + const double scale = norm < 1e-12 ? 2.0 / q[0] : 2.0 * atan2(norm, q[0]) / norm; + for (int i = 0; i < 3; ++i) { + w[i] = scale * q[i + 1]; + } +} + +// SE(3) exponential of the twist [w; rho]: translation = J_l(w) rho. +__device__ Pose ExpSE3(const double* twist) { + Pose pose; + ExpSO3(twist, pose.rotation); + const double theta2 = twist[0] * twist[0] + twist[1] * twist[1] + twist[2] * twist[2]; + const double theta = sqrt(theta2); + double a = 0.5 - theta2 / 24.0; + double b = 1.0 / 6.0 - theta2 / 120.0; + if (theta > 1e-4) { + a = (1.0 - cos(theta)) / theta2; + b = (theta - sin(theta)) / (theta2 * theta); + } + double left_jacobian[9]; + SkewSeries(twist, a, b, left_jacobian); + for (int row = 0; row < 3; ++row) { + double sum = 0.0; + for (int k = 0; k < 3; ++k) { + sum += left_jacobian[row * 3 + k] * twist[3 + k]; + } + pose.translation[row] = sum; + } + return pose; +} + +// SE(3) logarithm: [w; rho] with w = Log(R), rho = J_l(w)^-1 t. +__device__ void LogSE3(const Pose& pose, double* twist) { + LogSO3(pose.rotation, twist); + const double theta2 = twist[0] * twist[0] + twist[1] * twist[1] + twist[2] * twist[2]; + const double theta = sqrt(theta2); + double b = 1.0 / 12.0 + theta2 / 720.0; + if (theta > 1e-4) { + b = (1.0 - 0.5 * theta * sin(theta) / (1.0 - cos(theta))) / theta2; + } + double inverse_jacobian[9]; + SkewSeries(twist, -0.5, b, inverse_jacobian); + for (int row = 0; row < 3; ++row) { + double sum = 0.0; + for (int k = 0; k < 3; ++k) { + sum += inverse_jacobian[row * 3 + k] * pose.translation[k]; + } + twist[3 + row] = sum; + } +} + +// The residual of a triple smooth factor, see TripleSmoothFactorBatch. +__device__ void Residual(const Pose& first, const Pose& middle, const Pose& last, double ratio, + double* residual) { + double twist[6]; + LogSE3(Multiply(Inverse(first), last), twist); + for (int i = 0; i < 6; ++i) { + twist[i] *= ratio; + } + const Pose error = Multiply(Inverse(middle), Multiply(first, ExpSE3(twist))); + LogSO3(error.rotation, residual); + for (int i = 0; i < 3; ++i) { + residual[3 + i] = error.translation[i]; + } +} + +__device__ int MeasurementIndex(int item, const int* factor_ids, int num_factors) { + if (factor_ids != nullptr) { + return factor_ids[item]; + } + return item < num_factors ? item : item % num_factors; +} + +__global__ void TripleSmoothKernel(float const* const* state_pointers, const float* ratios, + size_t num_items, const int* factor_ids, int num_factors, + float* residuals, float* jacobians) { + const int item = threadIdx.x + blockIdx.x * blockDim.x; + if (item >= static_cast(num_items)) { + return; + } + const double ratio = ratios[MeasurementIndex(item, factor_ids, num_factors)]; + Pose poses[3]; + for (int i = 0; i < 3; ++i) { + poses[i] = LoadPose(state_pointers[3 * item + i]); + } + + double residual[6]; + Residual(poses[0], poses[1], poses[2], ratio, residual); + for (int i = 0; i < 6; ++i) { + residuals[6 * item + i] = static_cast(residual[i]); + } + if (jacobians == nullptr) { + return; + } + + // Column 6 * state + k: the residual's response to T_state * Exp(h e_k). + float* jacobian = jacobians + static_cast(item) * 6 * 18; + for (int state = 0; state < 3; ++state) { + for (int k = 0; k < 6; ++k) { + double residuals_by_sign[2][6]; + for (int sign = 0; sign < 2; ++sign) { + double twist[6] = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0}; + twist[k] = sign == 0 ? kDifferenceStep : -kDifferenceStep; + Pose moved[3] = {poses[0], poses[1], poses[2]}; + moved[state] = Multiply(poses[state], ExpSE3(twist)); + Residual(moved[0], moved[1], moved[2], ratio, residuals_by_sign[sign]); + } + for (int row = 0; row < 6; ++row) { + jacobian[row * 18 + 6 * state + k] = static_cast( + (residuals_by_sign[0][row] - residuals_by_sign[1][row]) / (2.0 * kDifferenceStep)); + } + } + } +} + +} // namespace + +TripleSmoothFactorBatch::TripleSmoothFactorBatch(const float* ratios, size_t capacity) + : SizedFactorBatch(capacity), ratios_(ratios) {} + +bool TripleSmoothFactorBatch::Evaluate(float* residuals, float* jacobians, + float const* const* state_pointers, cudaStream_t stream, + const int* factor_ids, size_t num_factor_ids) const { + const size_t num_items = num_factor_ids == 0 ? NumActiveFactors() : num_factor_ids; + if (num_items == 0 || NumActiveFactors() == 0) return true; + const size_t num_blocks = (num_items + kThreadsPerBlock - 1) / kThreadsPerBlock; + TripleSmoothKernel<<>>( + state_pointers, ratios_, num_items, factor_ids, static_cast(NumActiveFactors()), + residuals, jacobians); + return cudaGetLastError() == cudaSuccess; +} + +} // namespace test_utils +} // namespace cunls diff --git a/tests/triple_smooth_factor_batch.h b/tests/triple_smooth_factor_batch.h new file mode 100644 index 0000000..b1a37fe --- /dev/null +++ b/tests/triple_smooth_factor_batch.h @@ -0,0 +1,69 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include + +#include "cunls/factor/sized_factor_batch.h" + +namespace cunls { +namespace test_utils { + +/** + * @brief Test copy of cuSFM's triple smooth factor (cuSFM + * `src/core/bundle_adjustment/cunls_triple_smooth_factor.cu`), so that the + * synthetic cuSFM-like bundle adjustment tests build the same factor graph as + * the cuSFM Mapper. + * + * The middle one of three timestamp-adjacent world_from_rig poses has to stay + * on the SE(3) interpolation of the outer two. With the interpolation ratio + * r = (t_1 - t_0) / (t_2 - t_0): + * + * xi = Log(T_0^-1 T_2), E = T_1^-1 T_0 Exp(r xi), + * residual = [Log_SO3(rotation(E)); translation(E)]. + * + * The Jacobians with respect to [T_0, T_1, T_2] (right perturbations, tangent + * [rotation; translation]) are central differences of the residual in double + * precision, as in cuSFM. Wrap it in InformationFactorBatch for a 6x6 + * square-root information matrix. + */ +class TripleSmoothFactorBatch : public SizedFactorBatch<6, 6, 6, 6> { + public: + /** + * @brief Constructs the batch. + * + * @param ratios Device pointer to the per-factor interpolation ratios, in + * (0, 1); at least `capacity` floats, not owned. + * @param capacity Number of factors the ratio buffer holds. The active count + * starts at 0: call SetNumActiveFactors(). + */ + TripleSmoothFactorBatch(const float *ratios, size_t capacity); + + /** @brief Residuals (6 per item) and optional 6x18 Jacobians; see FactorBatch::Evaluate. */ + bool Evaluate(float *residuals, float *jacobians, float const *const *state_pointers, + cudaStream_t stream, const int *factor_ids = nullptr, + size_t num_factor_ids = 0) const override; + + private: + const float *ratios_; ///< Device per-factor interpolation ratios (not owned). +}; + +} // namespace test_utils +} // namespace cunls diff --git a/tests/utils.h b/tests/utils.h index 6978647..73ecd7a 100644 --- a/tests/utils.h +++ b/tests/utils.h @@ -378,18 +378,34 @@ std::vector> CopyStateToHost(const VectorStateBatch &state_batc /** * @brief Reads CUNLS_SOLVER from the environment and returns the matching - * solver type. Defaults to BlockSparsePCG — the sweep on SBA and - * loop-closure PGO shows it dominates cuDSS at every nontrivial - * problem size (see profile/PCG_RESULTS.md). Set - * CUNLS_SOLVER=cuDSS to opt back into the sparse-direct path. - * Recognised values: "BlockSparsePCG", "cuDSS". + * solver type. Defaults to BlockSparseCholesky, the library default + * (MinimizerOptions::sparse_linear_solver_type). + * Recognised values: "BlockSparseCholesky", "BlockSparsePCG", "cuDSS". */ inline SparseLinearSolverType SolverTypeFromEnv() { const char *s = std::getenv("CUNLS_SOLVER"); if (s != nullptr && std::strcmp(s, "cuDSS") == 0) { return SparseLinearSolverType::cuDSS; } - return SparseLinearSolverType::BlockSparsePCG; + if (s != nullptr && std::strcmp(s, "BlockSparsePCG") == 0) { + return SparseLinearSolverType::BlockSparsePCG; + } + return SparseLinearSolverType::BlockSparseCholesky; +} + +/** + * @brief Applies CUNLS_LINEAR_SYSTEM: with "SchurComplement", sets the Schur + * complement linear system (reduced system solved by + * BlockSparseCholesky); the + * caller then declares the landmarks + * (Problem::SetEliminatedStateBatches). Otherwise leaves `options`. + * @return Whether the Schur complement is used. + */ +inline bool UseSchurComplementFromEnv(MinimizerOptions &options) { + const char *s = std::getenv("CUNLS_LINEAR_SYSTEM"); + if (s == nullptr || std::strcmp(s, "SchurComplement") != 0) return false; + options.linear_system = LinearSystemType::SchurComplement; + return true; } /** @brief True when cunls was built with CUNLS_ENABLE_CUDSS=ON. */ @@ -416,12 +432,14 @@ inline std::vector AvailableSolverTypes( * @brief Solver ids of the typed minimizer tests. * * Ids 0 and 1 are the cuDSSLinearSolverMode values (cuDSS backend); - * kBlockSparsePCGSolverId selects BlockSparsePCG, which needs no cuDSS. + * kBlockSparsePCGSolverId and kBlockSparseCholeskySolverId select + * BlockSparsePCG and BlockSparseCholesky, which need no cuDSS. */ inline constexpr int kBlockSparsePCGSolverId = 2; +inline constexpr int kBlockSparseCholeskySolverId = 3; -/** @brief Solver id for tests that need a working sparse solver of any kind. */ -inline constexpr int kDefaultTestSolverId = kCuDSSEnabled ? 0 : kBlockSparsePCGSolverId; +/** @brief Solver id for tests that need a working sparse solver of any kind (the default). */ +inline constexpr int kDefaultTestSolverId = kBlockSparseCholeskySolverId; /** @brief Points @p options at the solver selected by @p solver_id. */ inline void ConfigureTestSolver(int solver_id, MinimizerOptions &options) { @@ -429,6 +447,10 @@ inline void ConfigureTestSolver(int solver_id, MinimizerOptions &options) { options.sparse_linear_solver_type = SparseLinearSolverType::BlockSparsePCG; return; } + if (solver_id == kBlockSparseCholeskySolverId) { + options.sparse_linear_solver_type = SparseLinearSolverType::BlockSparseCholesky; + return; + } options.sparse_linear_solver_type = SparseLinearSolverType::cuDSS; options.sparse_linear_solver_config.cudss_solver_options = { .mode = static_cast(solver_id), diff --git a/tests/weighted_factor_batch_test.cpp b/tests/weighted_factor_batch_test.cpp index 1e70d1b..7c2d8ba 100644 --- a/tests/weighted_factor_batch_test.cpp +++ b/tests/weighted_factor_batch_test.cpp @@ -94,13 +94,15 @@ struct TestParam { }; constexpr int kPCG = test_utils::kBlockSparsePCGSolverId; +constexpr int kChol = test_utils::kBlockSparseCholeskySolverId; typedef ::testing::Types< #ifdef CUNLS_ENABLE_CUDSS TestParam<1, 0>, TestParam<2, 0>, TestParam<3, 0>, TestParam<4, 0>, TestParam<1, 1>, TestParam<2, 1>, TestParam<3, 1>, TestParam<4, 1>, #endif - TestParam<1, kPCG>, TestParam<2, kPCG>, TestParam<3, kPCG>, TestParam<4, kPCG>> + TestParam<1, kPCG>, TestParam<2, kPCG>, TestParam<3, kPCG>, TestParam<4, kPCG>, + TestParam<1, kChol>, TestParam<2, kChol>, TestParam<3, kChol>, TestParam<4, kChol>> WeightedTestParams; TYPED_TEST_CASE(WeightedFactorBatchTest, WeightedTestParams);