From a67e88e8ab6f49d8f0783c16eecaa45bfc4c6981 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 14 Oct 2026 21:19:47 -0700 Subject: [PATCH 1/2] Add generalized Schur solver with automatic state elimination --- .gitignore | 7 + CMakeLists.txt | 1 + README.md | 2 +- cunls/linear_solver/CMakeLists.txt | 1 + .../linear_solver/generalized_schur_solver.cu | 1232 +++++++++++++++++ .../linear_solver/generalized_schur_solver.h | 276 ++++ cunls/linear_solver/llms.txt | 26 + cunls/linear_solver/sparse_linear_solver.cpp | 3 + cunls/linear_solver/sparse_linear_solver.h | 1 + docs/sphinx/api/linear_solver.rst | 23 +- docs/sphinx/generalized_schur.rst | 205 +++ docs/sphinx/index.rst | 1 + examples/sparse_bundle_adjustment/README.md | 6 + examples/sparse_bundle_adjustment/main.cpp | 2 + python/examples/imu_bundle_adjustment.py | 2 + python/examples/sparse_bundle_adjustment.py | 2 + python/pycunls/_pycunls_core.pyi | 1 + python/src/bind_types.cpp | 1 + python/tests/test_minimizer.py | 34 + scripts/benchmark_schur_in_docker.sh | 101 ++ scripts/plot_schur_benchmark.py | 69 + tests/generalized_schur_solver_test.cpp | 284 ++++ tests/sba_minimizer_test.cpp | 6 + tests/synthetic_sba_test.cpp | 134 +- tests/utils.h | 15 +- 25 files changed, 2423 insertions(+), 12 deletions(-) create mode 100644 cunls/linear_solver/generalized_schur_solver.cu create mode 100644 cunls/linear_solver/generalized_schur_solver.h create mode 100644 docs/sphinx/generalized_schur.rst create mode 100755 scripts/benchmark_schur_in_docker.sh create mode 100644 scripts/plot_schur_benchmark.py create mode 100644 tests/generalized_schur_solver_test.cpp diff --git a/.gitignore b/.gitignore index 32214a7..f0541d3 100644 --- a/.gitignore +++ b/.gitignore @@ -4,5 +4,12 @@ viz/ # Rendered video outputs (only the GIFs in assets/ are committed, via Git LFS) *.mp4 +*.pyc + +build/ +build_tests/ + +.cache/ + # Design notes are kept locally, not in the repository docs/design/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 15ec8ca..7eeebde 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -214,6 +214,7 @@ if(BUILD_TESTING) tests/motion_prior_perf_test.cpp tests/factor_manifold_facade_test.cpp tests/dense_linear_solver_test.cpp + tests/generalized_schur_solver_test.cpp tests/dense_cholesky_solver_test.cpp tests/dense_qr_solver_test.cpp tests/gauss_newton_test.cpp diff --git a/README.md b/README.md index 142e3fd..64ca851 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** | Block-sparse PCG (variable block-Jacobi preconditioner, default), [generalized Schur complement](docs/sphinx/generalized_schur.rst) (automatic state elimination), NVIDIA cuDSS (optional, loaded via `dlopen()` at runtime — see [C++ Installation](docs/sphinx/installation.rst)), dense LDLT, dense Cholesky (cuSOLVER), dense QR (cuSOLVER) | | **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/linear_solver/CMakeLists.txt b/cunls/linear_solver/CMakeLists.txt index f97ad20..b6960d9 100644 --- a/cunls/linear_solver/CMakeLists.txt +++ b/cunls/linear_solver/CMakeLists.txt @@ -1,4 +1,5 @@ add_library(cunls_linear_solver OBJECT + generalized_schur_solver.cu block_sparse_pcg_solver.cu block_tridiagonal_solver.cu dense_cholesky_solver.cu diff --git a/cunls/linear_solver/generalized_schur_solver.cu b/cunls/linear_solver/generalized_schur_solver.cu new file mode 100644 index 0000000..8d1374d --- /dev/null +++ b/cunls/linear_solver/generalized_schur_solver.cu @@ -0,0 +1,1232 @@ +/* + * 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. + */ + +// Select through SparseLinearSolverType::GeneralizedSchur. +#include + +#include +#include +#include +#include + +#include "cunls/common/cublas_helper.h" +#include "cunls/common/log.h" +#include "cunls/linear_solver/generalized_schur_solver.h" +#include "cunls/minimizer/problem.h" + +namespace cunls { + +namespace internal { +// Logical state rows have identical block sparsity. Initialize checks that +// contract before the fast relative-offset gathers used by the Schur kernels. +struct SchurMatrixView { + int num_rows; + int tile_size; + const int *row_offsets; + const int *column_ids; + const float *values; + const double *transformed_values = nullptr; + const size_t *transformed_offsets = nullptr; + const int *row_state_ids = nullptr; + const int *state_offsets = nullptr; + + __device__ int RowBegin(int row) const { return row_offsets[row / tile_size] * tile_size; } + + __device__ int RowEnd(int row) const { return row_offsets[row / tile_size + 1] * tile_size; } + + __device__ int ColumnIndex(int entry) const { + return column_ids[entry / tile_size] * tile_size + entry % tile_size; + } + + __device__ size_t ValueIndex(int row, int entry) const { + return size_t(entry / tile_size) * tile_size * tile_size + (row % tile_size) * tile_size + + entry % tile_size; + } + + __device__ int FindColumn(int row, int column) const { + int lower = row_offsets[row / tile_size], upper = row_offsets[row / tile_size + 1]; + while (lower < upper) { + int middle = lower + (upper - lower) / 2; + if (column_ids[middle] < column / tile_size) { + lower = middle + 1; + } else { + upper = middle; + } + } + return lower < row_offsets[row / tile_size + 1] && column_ids[lower] == column / tile_size + ? lower * tile_size + column % tile_size + : -1; + } + + __device__ double ReadEntry(int row, int entry) const { + if (!transformed_values) { + return values[ValueIndex(row, entry)]; + } + int state_id = row_state_ids[row]; + return transformed_values[transformed_offsets[state_id] + + size_t(row - state_offsets[state_id]) * + (RowEnd(row) - RowBegin(row)) + + entry - RowBegin(row)]; + } + + __device__ double ReadCoefficient(int row, int column) const { + int entry = FindColumn(row, column); + return entry < 0 ? 0. : ReadEntry(row, entry); + } +}; + +} // namespace internal + +namespace { +using MatrixView = internal::SchurMatrixView; +// Internal policy; selecting the enum is the complete public configuration. +constexpr float kRelativeTolerance = 1e-3f; +constexpr float kAbsoluteTolerance = 1e-30f; +constexpr int kMaxIterations = 200; +constexpr int kMaxRefinementPasses = 3; +constexpr int kMaxExplicitDimension = 4096; +constexpr size_t kMaxCliqueContributions = 512000000; + +__global__ void ValidateStructure(MatrixView matrix, int num_blocks, const int *state_offsets, + int *failure_flag) { + int block_id = blockIdx.x * blockDim.x + threadIdx.x; + if (block_id >= num_blocks) { + return; + } + int block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + for (int row = block_begin; row < block_begin + block_dimension; ++row) { + int begin = matrix.RowBegin(row), end = matrix.RowEnd(row); + if (end - begin != matrix.RowEnd(block_begin) - matrix.RowBegin(block_begin)) { + atomicExch(failure_flag, 1); + return; + } + for (int entry = begin; entry < end; ++entry) { + int column = matrix.ColumnIndex(entry); + if (column < 0 || column >= matrix.num_rows || + (entry > begin && matrix.ColumnIndex(entry - 1) >= column) || + column != matrix.ColumnIndex(matrix.RowBegin(block_begin) + entry - begin)) { + atomicExch(failure_flag, 1); + return; + } + } + } +} + +__global__ void ComputeStateRowWidths(MatrixView matrix, int num_blocks, const int *state_offsets, + int *state_row_widths) { + int block_id = blockIdx.x * blockDim.x + threadIdx.x; + if (block_id < num_blocks) { + state_row_widths[block_id] = + matrix.RowEnd(state_offsets[block_id]) - matrix.RowBegin(state_offsets[block_id]); + } +} + +__device__ unsigned HashStateId(unsigned state_id) { + state_id ^= state_id >> 16; + state_id *= 0x7feb352du; + state_id ^= state_id >> 15; + state_id *= 0x846ca68bu; + return state_id ^ (state_id >> 16); +} + +// Parallel local minima form an independent set. Stable tie-breaking avoids a +// sequential ordering on regular graphs; no graph edges leave device memory. +__global__ void SelectIndependentStates(MatrixView matrix, int num_blocks, const int *state_offsets, + const int *row_state_ids, const int *state_row_widths, + const int *status, int *chosen) { + int block_id = blockIdx.x * blockDim.x + threadIdx.x; + if (block_id >= num_blocks || status[block_id]) { + return; + } + long long score = (long long)state_row_widths[block_id] * + (state_offsets[block_id + 1] - state_offsets[block_id]); + for (int entry = matrix.RowBegin(state_offsets[block_id]); + entry < matrix.RowEnd(state_offsets[block_id]); ++entry) { + int index = row_state_ids[matrix.ColumnIndex(entry)]; + if (index == block_id || status[index]) { + continue; + } + long long neighbor_score = + (long long)state_row_widths[index] * (state_offsets[index + 1] - state_offsets[index]); + if (neighbor_score < score || + (neighbor_score == score && + (HashStateId(index) < HashStateId(block_id) || + (HashStateId(index) == HashStateId(block_id) && index < block_id)))) { + return; + } + } + chosen[block_id] = 1; +} + +__global__ void MarkSelectedNeighbors(MatrixView matrix, int num_blocks, const int *state_offsets, + const int *row_state_ids, const int *chosen, int *status) { + int block_id = blockIdx.x * blockDim.x + threadIdx.x; + if (block_id >= num_blocks || status[block_id]) { + return; + } + if (chosen[block_id]) { + status[block_id] = 1; + return; + } + for (int entry = matrix.RowBegin(state_offsets[block_id]); + entry < matrix.RowEnd(state_offsets[block_id]); ++entry) { + if (chosen[row_state_ids[matrix.ColumnIndex(entry)]]) { + status[block_id] = 2; + return; + } + } +} + +// Stable block prefixes are built with the dimension buckets on the host. +// One coalesced pass expands them into scalar row lists; no device selection +// library, scratch allocations, or four separate compaction synchronizations. +__global__ void ExpandPartitionRows(int num_rows, const int *row_state_ids, + const int *state_offsets, const int *status, const int *base, + int *erows, int *rrows, int *retained_indices) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= num_rows) { + return; + } + int row_index = row_state_ids[row], index = base[row_index] + row - state_offsets[row_index]; + bool eliminated = status[row_index] == 1; + if (eliminated) { + erows[index] = row; + } else { + rrows[index] = row; + } + retained_indices[row] = eliminated ? -1 : index; +} + +__global__ void ExtractDiagonalBlocks(MatrixView matrix, int num_blocks, const int *state_offsets, + const size_t *diagonal_offsets, double *diagonal, + int *unobservable_rows) { + int block_id = blockIdx.x * blockDim.x + threadIdx.x; + if (block_id >= num_blocks) { + return; + } + int block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + for (int row = 0; row < block_dimension; ++row) { + for (int column = 0; column < block_dimension; ++column) { + diagonal[diagonal_offsets[block_id] + row * block_dimension + column] = + matrix.ReadCoefficient(block_begin + row, block_begin + column); + } + } + // An exactly empty numerical row is unobservable (e.g. a landmark behind + // every camera). Choose zero in that direction without damping any observed + // equation. A nonzero RHS there is still rejected by the full residual check. + for (int row = 0; row < block_dimension; ++row) { + bool empty = diagonal[diagonal_offsets[block_id] + row * block_dimension + row] == 0.; + if (empty) { + for (int entry = matrix.RowBegin(block_begin + row); entry < matrix.RowEnd(block_begin + row); + ++entry) { + if (matrix.ReadEntry(block_begin + row, entry) != 0.) { + empty = false; + break; + } + } + } + unobservable_rows[block_begin + row] = empty; + if (empty) { + diagonal[diagonal_offsets[block_id] + row * block_dimension + row] = 1.; + } + } +} + +__global__ void RegularizeUnobservableRows(int num_blocks, const int *block_ids, + const int *state_offsets, const size_t *diagonal_offsets, + const int *unobservable_rows, double *diagonal) { + int work_index = blockIdx.x * blockDim.x + threadIdx.x; + if (work_index >= num_blocks) { + return; + } + int block_id = block_ids[work_index], block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + for (int row = 0; row < block_dimension; ++row) { + if (unobservable_rows[block_begin + row]) { + diagonal[diagonal_offsets[block_id] + row * block_dimension + row] = 1.; + } + } +} + +// Tiny diagonal blocks are inverted once per numeric solve, not per CG step. +// Cholesky and accumulation use double; no pivot flooring changes the system. +template +__global__ void InvertDiagonalBlocks(int num_blocks, const int *block_ids, const int *state_offsets, + const size_t *diagonal_offsets, double *diagonal, + double *inverse_blocks, int *failure_flag) { + int work_index = blockIdx.x * blockDim.x + threadIdx.x; + if (work_index >= num_blocks) { + return; + } + int block_id = block_ids[work_index], + block_dimension = + FixedDimension ? FixedDimension : state_offsets[block_id + 1] - state_offsets[block_id]; + double local_factor[FixedDimension ? FixedDimension * FixedDimension : 1]; + double *cholesky_factor = FixedDimension ? local_factor : diagonal + diagonal_offsets[block_id]; + if (FixedDimension) { + for (int index = 0; index < block_dimension * block_dimension; ++index) { + cholesky_factor[index] = diagonal[diagonal_offsets[block_id] + index]; + } + } + double *block_inverse = inverse_blocks + diagonal_offsets[block_id]; + for (int row = 0; row < block_dimension; ++row) { + for (int column = 0; column <= row; ++column) { + double pivot_value = cholesky_factor[row * block_dimension + column]; + for (int inner = 0; inner < column; ++inner) { + pivot_value -= cholesky_factor[row * block_dimension + inner] * + cholesky_factor[column * block_dimension + inner]; + } + if (column == row) { + if (!(pivot_value > 0.) || !isfinite(pivot_value)) { + atomicExch(failure_flag, 1); + return; + } + cholesky_factor[row * block_dimension + column] = sqrt(pivot_value); + } else { + cholesky_factor[row * block_dimension + column] = + pivot_value / cholesky_factor[column * block_dimension + column]; + } + } + } + for (int column = 0; column < block_dimension; ++column) { + for (int row = 0; row < block_dimension; ++row) { + double pivot_value = row == column ? 1. : 0.; + for (int inner = 0; inner < row; ++inner) { + pivot_value -= cholesky_factor[row * block_dimension + inner] * + block_inverse[inner * block_dimension + column]; + } + block_inverse[row * block_dimension + column] = + pivot_value / cholesky_factor[row * block_dimension + row]; + } + for (int row = block_dimension - 1; row >= 0; --row) { + double pivot_value = block_inverse[row * block_dimension + column]; + for (int inner = row + 1; inner < block_dimension; ++inner) { + pivot_value -= cholesky_factor[inner * block_dimension + row] * + block_inverse[inner * block_dimension + column]; + } + block_inverse[row * block_dimension + column] = + pivot_value / cholesky_factor[row * block_dimension + row]; + } + } + for (int row = 0; row < block_dimension; ++row) { + for (int column = 0; column < row; ++column) { + double symmetric_value = .5 * (block_inverse[row * block_dimension + column] + + block_inverse[column * block_dimension + row]); + block_inverse[row * block_dimension + column] = + block_inverse[column * block_dimension + row] = symmetric_value; + } + } +} + +__global__ void TransformCrossBlocks(MatrixView matrix, int num_eliminated_blocks, + const int *block_ids, const int *state_offsets, + const int *retained_indices, const size_t *diagonal_offsets, + const size_t *transformed_offsets, + const double *inverse_blocks, double *transformed_values) { + int work_index = blockIdx.x * 4 + threadIdx.x / 32; + if (work_index >= num_eliminated_blocks) { + return; + } + int block_id = block_ids[work_index], block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + for (int entry = matrix.RowBegin(block_begin) + threadIdx.x % 32; + entry < matrix.RowEnd(block_begin); entry += 32) { + int column = matrix.ColumnIndex(entry); + if (retained_indices[column] < 0) { + continue; + } + for (int row = 0; row < block_dimension; ++row) { + double transformed_value = 0.; + for (int inner = 0; inner < block_dimension; ++inner) { + transformed_value += + (double)inverse_blocks[diagonal_offsets[block_id] + row * block_dimension + inner] * + matrix.values[matrix.ValueIndex( + block_begin + inner, + matrix.RowBegin(block_begin + inner) + entry - matrix.RowBegin(block_begin))]; + } + transformed_values[transformed_offsets[block_id] + + size_t(row) * (matrix.RowEnd(block_begin) - matrix.RowBegin(block_begin)) + + entry - matrix.RowBegin(block_begin)] = transformed_value; + } + } +} + +__global__ void AssembleImplicitSchurDiagonal(MatrixView matrix, MatrixView transformed_matrix, + const int *block_ids, const int *state_offsets, + const int *retained_indices, + const size_t *diagonal_offsets, + double *diagonal_blocks) { + int block_id = block_ids[blockIdx.x], block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + // One CTA per retained block; each warp owns one diagonal-block entry. + int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + for (int entry_index = warp; entry_index < block_dimension * block_dimension; + entry_index += blockDim.x / 32) { + int row = block_begin + entry_index / block_dimension, + column = block_begin + entry_index % block_dimension; + double sum = 0.; + for (int entry = matrix.RowBegin(row) + lane; entry < matrix.RowEnd(row); entry += 32) { + int inner = matrix.ColumnIndex(entry); + if (retained_indices[inner] < 0) { + sum += (double)matrix.values[matrix.ValueIndex(row, entry)] * + transformed_matrix.ReadCoefficient(inner, column); + } + } + for (int delta = 16; delta; delta /= 2) { + sum += __shfl_down_sync(0xffffffff, sum, delta); + } + if (!lane) { + diagonal_blocks[diagonal_offsets[block_id] + entry_index] = + (double)matrix.ReadCoefficient(row, column) - sum; + } + } +} + +// For compact retained systems, pay for the Schur clique updates once and +// reuse a dense operator. Double atomic accumulation prevents high-degree +// camera rows losing small contributions; Krylov vectors remain float32. +__global__ void CopyRetainedMatrix(MatrixView matrix, int num_retained_rows, const int *row_ids, + const int *retained_indices, double *schur_matrix) { + int row = blockIdx.x, original_row = row_ids[row]; + for (int entry = matrix.RowBegin(original_row) + threadIdx.x; entry < matrix.RowEnd(original_row); + entry += blockDim.x) { + int column = retained_indices[matrix.ColumnIndex(entry)]; + if (column >= 0) { + schur_matrix[size_t(row) * num_retained_rows + column] = + matrix.values[matrix.ValueIndex(original_row, entry)]; + } + } +} + +// The common three-dimensional pivot path stages each cross block once in +// shared memory. Adjacent lanes read adjacent columns; a warp broadcasts the +// left operand while reading the right operand without repeated sparse gathers. +// Work is restricted to the upper triangle: symmetry halves both FP64 work and +// contended atomics. The generic path has no degree or state-size bound. +__global__ void AccumulateSchurUpdates(MatrixView matrix, MatrixView transformed_matrix, + int num_retained_rows, const int *block_ids, + const int *state_offsets, const int *retained_indices, + double *schur_matrix) { + int block_id = block_ids[blockIdx.x], block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + int begin = matrix.RowBegin(block_begin), row_width = matrix.RowEnd(block_begin) - begin; + __shared__ double cross[3 * 64], transformed[3 * 64]; + __shared__ int retained[64]; + if (block_dimension == 3 && row_width <= 64) { + for (int index = threadIdx.x; index < row_width; index += blockDim.x) { + retained[index] = retained_indices[matrix.ColumnIndex(begin + index)]; + } + for (int index = threadIdx.x; index < 3 * row_width; index += blockDim.x) { + int inner = index / row_width, column = index % row_width; + bool active = retained_indices[matrix.ColumnIndex(begin + column)] >= 0; + cross[inner * 64 + column] = + active + ? matrix.ReadEntry(block_begin + inner, matrix.RowBegin(block_begin + inner) + column) + : 0.; + transformed[inner * 64 + column] = + active ? transformed_matrix.ReadEntry(block_begin + inner, + matrix.RowBegin(block_begin + inner) + column) + : 0.; + } + __syncthreads(); + for (int pair = threadIdx.x; pair < row_width * row_width; pair += blockDim.x) { + int row = pair / row_width, column = pair % row_width; + int retained_row = retained[row], retained_column = retained[column]; + if (retained_row < 0 || retained_column < retained_row) { + continue; + } + double sum = 0.; +#pragma unroll + for (int inner = 0; inner < 3; ++inner) { + sum += cross[inner * 64 + row] * transformed[inner * 64 + column]; + } + atomicAdd(schur_matrix + size_t(retained_row) * num_retained_rows + retained_column, -sum); + } + } else { + for (size_t pair = threadIdx.x; pair < size_t(row_width) * row_width; pair += blockDim.x) { + int row = pair / row_width, column = pair % row_width; + int retained_row = retained_indices[matrix.ColumnIndex(begin + row)], + retained_column = retained_indices[matrix.ColumnIndex(begin + column)]; + if (retained_row < 0 || retained_column < retained_row) { + continue; + } + double sum = 0.; + for (int inner = 0; inner < block_dimension; ++inner) { + sum += matrix.ReadEntry(block_begin + inner, matrix.RowBegin(block_begin + inner) + row) * + transformed_matrix.ReadEntry(block_begin + inner, + matrix.RowBegin(block_begin + inner) + column); + } + atomicAdd(schur_matrix + size_t(retained_row) * num_retained_rows + retained_column, -sum); + } + } +} + +__global__ void MirrorUpperTriangle(int num_rows, double *matrix) { + size_t row_index = size_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (row_index < size_t(num_rows) * num_rows) { + int row = row_index / num_rows, column = row_index % num_rows; + if (row > column) { + matrix[row_index] = matrix[size_t(column) * num_rows + row]; + } + } +} + +// Direct terminal case for tiny retained systems; no nested solver object. +__global__ void FactorSmallSchurMatrix(int num_rows, const double *matrix, const int *row_ids, + const int *unobservable_rows, double *factor, + int *failure_flag) { + __shared__ double cholesky_factor[64 * 64]; + __shared__ int factorization_failed; + if (!threadIdx.x) { + factorization_failed = 0; + } + for (int block_id = threadIdx.x; block_id < num_rows * num_rows; block_id += blockDim.x) { + cholesky_factor[block_id] = matrix[block_id]; + if (block_id / num_rows == block_id % num_rows && + unobservable_rows[row_ids[block_id / num_rows]]) { + cholesky_factor[block_id] = 1.; + } + } + __syncthreads(); + for (int inner = 0; inner < num_rows; ++inner) { + if (!threadIdx.x) { + if (!(cholesky_factor[inner * num_rows + inner] > 0.) || + !isfinite(cholesky_factor[inner * num_rows + inner])) { + factorization_failed = 1; + atomicExch(failure_flag, 1); + } else { + cholesky_factor[inner * num_rows + inner] = sqrt(cholesky_factor[inner * num_rows + inner]); + } + } + __syncthreads(); + if (factorization_failed) { + return; + } + for (int block_id = inner + 1 + threadIdx.x; block_id < num_rows; block_id += blockDim.x) { + cholesky_factor[block_id * num_rows + inner] /= cholesky_factor[inner * num_rows + inner]; + } + __syncthreads(); + for (int entry_index = threadIdx.x; entry_index < num_rows * num_rows; + entry_index += blockDim.x) { + int block_id = entry_index / num_rows, index = entry_index % num_rows; + if (block_id >= index && index > inner) { + cholesky_factor[entry_index] -= cholesky_factor[block_id * num_rows + inner] * + cholesky_factor[index * num_rows + inner]; + } + } + __syncthreads(); + } + for (int block_id = threadIdx.x; block_id < num_rows * num_rows; block_id += blockDim.x) { + factor[block_id] = cholesky_factor[block_id]; + } +} + +__global__ void SolveSmallSchurSystem(int num_rows, const double *cholesky_factor, const float *rhs, + float *solution) { + double intermediate_solution[64]; + for (int row = 0; row < num_rows; ++row) { + double substitution_value = rhs[row]; + for (int column = 0; column < row; ++column) { + substitution_value -= + cholesky_factor[row * num_rows + column] * intermediate_solution[column]; + } + intermediate_solution[row] = substitution_value / cholesky_factor[row * num_rows + row]; + } + for (int row = num_rows - 1; row >= 0; --row) { + double substitution_value = intermediate_solution[row]; + for (int column = row + 1; column < num_rows; ++column) { + substitution_value -= + cholesky_factor[column * num_rows + row] * intermediate_solution[column]; + } + intermediate_solution[row] = substitution_value / cholesky_factor[row * num_rows + row]; + solution[row] = intermediate_solution[row]; + } +} + +__global__ void MultiplyDenseSchurMatrix(int num_rows, const double *matrix, const float *input, + float *output) { + int row = blockIdx.x; + double sum = 0; + for (int column = threadIdx.x; column < num_rows; column += 128) { + sum += matrix[size_t(row) * num_rows + column] * input[column]; + } + for (int shuffle_offset = 16; shuffle_offset; shuffle_offset /= 2) { + sum += __shfl_down_sync(0xffffffff, sum, shuffle_offset); + } + __shared__ double warp_sums[4]; + if (threadIdx.x % 32 == 0) { + warp_sums[threadIdx.x / 32] = sum; + } + __syncthreads(); + if (!threadIdx.x) { + output[row] = warp_sums[0] + warp_sums[1] + warp_sums[2] + warp_sums[3]; + } +} + +__global__ void ExtractExplicitSchurDiagonal(int num_blocks, int num_retained_rows, + const double *schur_matrix, const int *block_ids, + const int *state_offsets, const int *retained_indices, + const size_t *diagonal_offsets, + double *diagonal_blocks) { + int work_index = blockIdx.x * blockDim.x + threadIdx.x; + if (work_index >= num_blocks) { + return; + } + int block_id = block_ids[work_index], block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + for (int row = 0; row < block_dimension; ++row) { + for (int column = 0; column < block_dimension; ++column) { + diagonal_blocks[diagonal_offsets[block_id] + row * block_dimension + column] = + schur_matrix[size_t(retained_indices[block_begin + row]) * num_retained_rows + + retained_indices[block_begin + column]]; + } + } +} + +__global__ void ApplyBlockInverse(int num_rows, const int *row_ids, const int *row_state_ids, + const int *state_offsets, const size_t *diagonal_offsets, + const double *inverse_blocks, const int *retained_indices, + const float *input, float *output, bool input_uses_full_order) { + int work_index = blockIdx.x * blockDim.x + threadIdx.x; + if (work_index >= num_rows) { + return; + } + int row = row_ids[work_index], block_id = row_state_ids[row], + block_begin = state_offsets[block_id], + block_dimension = state_offsets[block_id + 1] - block_begin; + double sum = 0.; + for (int index = 0; index < block_dimension; ++index) { + sum += + (double)inverse_blocks[diagonal_offsets[block_id] + (row - block_begin) * block_dimension + + index] * + input[input_uses_full_order ? block_begin + index : retained_indices[block_begin + index]]; + } + output[input_uses_full_order ? row : work_index] = sum; +} + +// Retained vectors are compact. EliminatedRows work stays in original row order. +// CTA-per-row for long retained rows, warp-per-row for short pivot rows. +template +__global__ void MultiplyPartitionedMatrix(MatrixView matrix, int num_rows, const int *row_ids, + const int *retained_indices, const float *input, + const float *eliminated_values, float *output) { + int work_index = ThreadsPerRow == 32 ? blockIdx.x * 4 + threadIdx.x / 32 : blockIdx.x; + if (work_index >= num_rows) { + return; + } + int row = row_ids[work_index]; + double sum = 0.; + int lane = ThreadsPerRow == 32 ? threadIdx.x % 32 : threadIdx.x; + for (int entry = matrix.RowBegin(row) + lane; entry < matrix.RowEnd(row); + entry += ThreadsPerRow) { + int column = matrix.ColumnIndex(entry), index = retained_indices[column]; + if constexpr (EliminatedRows) { + if (index < 0) { + continue; + } + } + float input_value = + index >= 0 ? input[index] : (EliminatedRows ? 0.f : -eliminated_values[column]); + sum += (double)matrix.ReadEntry(row, entry) * input_value; + } + for (int shuffle_offset = 16; shuffle_offset; shuffle_offset /= 2) { + sum += __shfl_down_sync(0xffffffff, sum, shuffle_offset); + } + if constexpr (ThreadsPerRow == 32) { + if (threadIdx.x % 32 == 0) { + output[EliminatedRows ? row : work_index] = sum; + } + return; + } + __shared__ double warp_sums[ThreadsPerRow / 32]; + if (threadIdx.x % 32 == 0) { + warp_sums[threadIdx.x / 32] = sum; + } + __syncthreads(); + if (threadIdx.x == 0) { + sum = 0.; + for (int index = 0; index < ThreadsPerRow / 32; ++index) { + sum += warp_sums[index]; + } + output[EliminatedRows ? row : work_index] = sum; + } +} + +__global__ void AddRetainedRhs(int num_rows, const int *row_ids, const float *rhs, + float *reduced_rhs) { + int row_index = blockIdx.x * blockDim.x + threadIdx.x; + if (row_index < num_rows) { + reduced_rhs[row_index] += rhs[row_ids[row_index]]; + } +} + +__global__ void RecoverFullSolution(int num_rows, const int *retained_indices, + const float *retained_solution, + const float *eliminated_rhs_solution, + const float *eliminated_product, float *solution) { + int row_index = blockIdx.x * blockDim.x + threadIdx.x; + if (row_index < num_rows) { + solution[row_index] = retained_indices[row_index] >= 0 + ? retained_solution[retained_indices[row_index]] + : eliminated_rhs_solution[row_index] - eliminated_product[row_index]; + } +} + +__global__ void UpdateCgSolutionAndResidual(int num_rows, const float *residual_preconditioned_dot, + const float *search_direction_operator_dot, + const float *search_direction, + const float *operator_product, float *solution, + float *residual) { + int row_index = blockIdx.x * blockDim.x + threadIdx.x; + if (row_index >= num_rows) { + return; + } + float alpha = (*search_direction_operator_dot > 0.f && *residual_preconditioned_dot > 0.f) + ? *residual_preconditioned_dot / *search_direction_operator_dot + : 0.f; + solution[row_index] += alpha * search_direction[row_index]; + residual[row_index] -= alpha * operator_product[row_index]; +} + +__global__ void UpdateCgDirection(int num_rows, const float *current_dot, const float *previous_dot, + const float *preconditioned_residual, float *search_direction) { + int row_index = blockIdx.x * blockDim.x + threadIdx.x; + if (row_index < num_rows) { + search_direction[row_index] = + preconditioned_residual[row_index] + + (*previous_dot > 0.f ? *current_dot / *previous_dot : 0.f) * search_direction[row_index]; + } +} + +__global__ void AddCorrection(int num_rows, const float *correction, float *solution) { + int row_index = blockIdx.x * blockDim.x + threadIdx.x; + if (row_index < num_rows) { + solution[row_index] += correction[row_index]; + } +} + +template +__global__ void ComputeFullResidual(MatrixView matrix, const float *rhs, const float *solution, + float *residual, const int *long_rows) { + int row = ThreadsPerRow == 32 ? blockIdx.x * 4 + threadIdx.x / 32 : long_rows[blockIdx.x]; + int lane = threadIdx.x % 32; + if (row >= matrix.num_rows) { + return; + } + if constexpr (ThreadsPerRow == 32) { + if (matrix.RowEnd(row) - matrix.RowBegin(row) > 128) { + return; + } + } + double sum = 0.; + for (int entry = matrix.RowBegin(row) + (ThreadsPerRow == 32 ? lane : threadIdx.x); + entry < matrix.RowEnd(row); entry += ThreadsPerRow) { + sum += + (double)matrix.values[matrix.ValueIndex(row, entry)] * solution[matrix.ColumnIndex(entry)]; + } + for (int shuffle_offset = 16; shuffle_offset; shuffle_offset /= 2) { + sum += __shfl_down_sync(0xffffffff, sum, shuffle_offset); + } + if constexpr (ThreadsPerRow == 32) { + if (!lane) { + residual[row] = rhs[row] - sum; + } + } else { + __shared__ double warp_sums[ThreadsPerRow / 32]; + if (!lane) { + warp_sums[threadIdx.x / 32] = sum; + } + __syncthreads(); + if (!threadIdx.x) { + sum = 0.; + for (int index = 0; index < ThreadsPerRow / 32; ++index) { + sum += warp_sums[index]; + } + residual[row] = rhs[row] - sum; + } + } +} + +} // namespace + +bool GeneralizedSchurSolver::SupportsBlockStorage() const { return true; } + +bool GeneralizedSchurSolver::Initialize(cudaStream_t stream, const Problem &problem, + const CSRSparseMatrix &matrix, const dvector &rhs, + dvector &solution) { + return Initialize(stream, problem, + {int(matrix.NumRows()), 1, matrix.row_offsets.data(), matrix.col_ids.data(), + matrix.values.data()}, + rhs, solution); +} + +bool GeneralizedSchurSolver::Initialize(cudaStream_t stream, const Problem &problem, + const BSRSparseMatrix &matrix, const dvector &rhs, + dvector &solution) { + return Initialize(stream, problem, + {matrix.NumRows(), matrix.block_size, matrix.row_offsets.data(), + matrix.col_ids.data(), matrix.values.data()}, + rhs, solution); +} + +bool GeneralizedSchurSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &matrix, + const dvector &rhs, dvector &solution) { + return Solve(stream, + {int(matrix.NumRows()), 1, matrix.row_offsets.data(), matrix.col_ids.data(), + matrix.values.data()}, + rhs, solution); +} + +bool GeneralizedSchurSolver::Solve(cudaStream_t stream, const BSRSparseMatrix &matrix, + const dvector &rhs, dvector &solution) { + return Solve(stream, + {matrix.NumRows(), matrix.block_size, matrix.row_offsets.data(), + matrix.col_ids.data(), matrix.values.data()}, + rhs, solution); +} + +bool GeneralizedSchurSolver::Initialize(cudaStream_t stream, const Problem &problem, + MatrixView matrix, const dvector &rhs, + dvector &solution) { + if (matrix.num_rows != int(rhs.size()) || matrix.num_rows != int(solution.size()) || + matrix.num_rows <= 0) { + return false; + } + num_rows_ = matrix.num_rows; + std::vector state_offsets{0}, row_state_ids; + std::vector diagonal_offsets{0}; + for (auto *batch : problem.GetStateBatches()) { + int block_dimension = batch->TangentSize(); + size_t num_free_states = batch->NumActiveStates() - batch->NumConstStates(); + for (size_t block_id = 0; block_id < num_free_states; ++block_id) { + state_offsets.push_back(state_offsets.back() + block_dimension); + diagonal_offsets.push_back(diagonal_offsets.back() + + size_t(block_dimension) * block_dimension); + } + } + if (state_offsets.size() == 1) { + for (int row = 0; row < num_rows_; row += matrix.tile_size) { + state_offsets.push_back(row + matrix.tile_size); + diagonal_offsets.push_back(diagonal_offsets.back() + + size_t(matrix.tile_size) * matrix.tile_size); + } + } + if (state_offsets.back() != num_rows_) { + return false; + } + num_state_blocks_ = state_offsets.size() - 1; + row_state_ids.resize(num_rows_); + for (int block_id = 0; block_id < num_state_blocks_; ++block_id) { + std::fill(row_state_ids.begin() + state_offsets[block_id], + row_state_ids.begin() + state_offsets[block_id + 1], block_id); + } + state_row_offsets_ = dvector(state_offsets); + row_state_ids_ = dvector(row_state_ids); + diagonal_block_offsets_ = dvector(diagonal_offsets); + unobservable_rows_.resize(num_rows_); + device_failure_flag_.resize(1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(device_failure_flag_.data(), 0, sizeof(int), stream)); + ValidateStructure<<<(num_state_blocks_ + 127) / 128, 128, 0, stream>>>( + matrix, num_state_blocks_, state_row_offsets_.data(), device_failure_flag_.data()); + if (!CheckDeviceSuccess(stream)) { + LogError("GeneralizedSchur: inconsistent logical block structure"); + return false; + } + selection_status_.resize(num_state_blocks_); + selected_in_round_.resize(num_state_blocks_); + state_row_widths_.resize(num_state_blocks_); + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(selection_status_.data(), 0, num_state_blocks_ * sizeof(int), stream)); + ComputeStateRowWidths<<<(num_state_blocks_ + 127) / 128, 128, 0, stream>>>( + matrix, num_state_blocks_, state_row_offsets_.data(), state_row_widths_.data()); + for (int round = 0; round < 4; ++round) { + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(selected_in_round_.data(), 0, num_state_blocks_ * sizeof(int), stream)); + SelectIndependentStates<<<(num_state_blocks_ + 127) / 128, 128, 0, stream>>>( + matrix, num_state_blocks_, state_row_offsets_.data(), row_state_ids_.data(), + state_row_widths_.data(), selection_status_.data(), selected_in_round_.data()); + MarkSelectedNeighbors<<<(num_state_blocks_ + 127) / 128, 128, 0, stream>>>( + matrix, num_state_blocks_, state_row_offsets_.data(), row_state_ids_.data(), + selected_in_round_.data(), selection_status_.data()); + } + // Selection metadata is also needed for dimension buckets and the fill + // estimate. Transfer it once, and build all stable block prefixes in one pass. + std::vector statuses(num_state_blocks_), degrees(num_state_blocks_), + partition_offsets(num_state_blocks_), eliminated_blocks, retained_blocks, long_rows; + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(statuses.data(), selection_status_.data(), + num_state_blocks_ * sizeof(int), cudaMemcpyDeviceToHost, + stream)); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(degrees.data(), state_row_widths_.data(), + num_state_blocks_ * sizeof(int), cudaMemcpyDeviceToHost, + stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + std::map, std::vector> buckets; + std::vector transformed_offsets(num_state_blocks_); + size_t transformed_size = 0; + num_eliminated_rows_ = num_retained_rows_ = 0; + size_t estimated_clique_pairs = 0; + for (int block_id = 0; block_id < num_state_blocks_; ++block_id) { + int block_dimension = state_offsets[block_id + 1] - state_offsets[block_id]; + bool is_eliminated = statuses[block_id] == 1; + partition_offsets[block_id] = is_eliminated ? num_eliminated_rows_ : num_retained_rows_; + (is_eliminated ? num_eliminated_rows_ : num_retained_rows_) += block_dimension; + (is_eliminated ? eliminated_blocks : retained_blocks).push_back(block_id); + transformed_offsets[block_id] = transformed_size; + if (is_eliminated) { + estimated_clique_pairs += size_t(degrees[block_id]) * degrees[block_id]; + transformed_size += size_t(block_dimension) * degrees[block_id]; + } + if (degrees[block_id] > 128) { + for (int row = state_offsets[block_id]; row < state_offsets[block_id + 1]; ++row) { + long_rows.push_back(row); + } + } + bool specialized = block_dimension == 1 || block_dimension == 2 || block_dimension == 3 || + block_dimension == 4 || block_dimension == 6 || block_dimension == 7 || + block_dimension == 9 || block_dimension == 15 || block_dimension == 16; + buckets[{is_eliminated ? BlockRole::Eliminated : BlockRole::Retained, + specialized ? block_dimension : 0}] + .push_back(block_id); + } + num_eliminated_blocks_ = eliminated_blocks.size(); + num_retained_blocks_ = retained_blocks.size(); + eliminated_block_ids_ = dvector(eliminated_blocks); + retained_block_ids_ = dvector(retained_blocks); + long_residual_row_ids_ = dvector(long_rows); + partition_row_offsets_ = dvector(partition_offsets); + eliminated_row_ids_.resize(num_eliminated_rows_); + retained_row_ids_.resize(num_retained_rows_); + retained_row_indices_.resize(num_rows_); + ExpandPartitionRows<<<(num_rows_ + 255) / 256, 256, 0, stream>>>( + num_rows_, row_state_ids_.data(), state_row_offsets_.data(), selection_status_.data(), + partition_row_offsets_.data(), eliminated_row_ids_.data(), retained_row_ids_.data(), + retained_row_indices_.data()); + dimension_groups_.clear(); + for (auto &bucket : buckets) { + dimension_groups_.push_back( + {bucket.first.first, bucket.first.second, dvector(bucket.second)}); + } + diagonal_blocks_.resize(diagonal_offsets.back()); + inverse_blocks_.resize(diagonal_offsets.back()); + transformed_offsets_ = dvector(transformed_offsets); + transformed_cross_blocks_.resize(transformed_size); + device_failure_flag_.resize(1); + eliminated_rhs_solution_.resize(num_rows_); + eliminated_product_.resize(num_rows_); + full_residual_.resize(num_rows_); + correction_.resize(num_rows_); + retained_solution_.resize(num_retained_rows_); + reduced_rhs_.resize(num_retained_rows_); + reduced_residual_.resize(num_retained_rows_); + preconditioned_residual_.resize(num_retained_rows_); + search_direction_.resize(num_retained_rows_); + operator_product_.resize(num_retained_rows_); + cg_dot_products_.resize(4); + size_t free_bytes, total_bytes; + THROW_ON_CUDA_ERROR(cudaMemGetInfo(&free_bytes, &total_bytes)); + use_explicit_schur_ = num_retained_rows_ > 0 && num_retained_rows_ <= kMaxExplicitDimension && + estimated_clique_pairs <= kMaxCliqueContributions && + size_t(num_retained_rows_) * num_retained_rows_ * 12 < free_bytes / 4; + if (use_explicit_schur_) { + schur_matrix_.resize(size_t(num_retained_rows_) * num_retained_rows_); + } + use_direct_solve_ = use_explicit_schur_ && num_retained_rows_ <= 64; + if (use_direct_solve_) { + terminal_factor_.resize(size_t(num_retained_rows_) * num_retained_rows_); + } + LogMessage( + "GeneralizedSchur: {} eliminated / {} retained dimensions, {} / {} state blocks, " + "explicit={}", + num_eliminated_rows_, num_retained_rows_, num_eliminated_blocks_, num_retained_blocks_, + use_explicit_schur_); + return true; +} + +void GeneralizedSchurSolver::InvertDimensionGroups(cudaStream_t stream, BlockRole role) { + for (auto &group : dimension_groups_) { + if (group.role == role) { + int num_group_blocks = group.block_ids.size(); +#define LAUNCH_BLOCK_INVERSE(D) \ + InvertDiagonalBlocks<<<(num_group_blocks + 127) / 128, 128, 0, stream>>>( \ + num_group_blocks, group.block_ids.data(), state_row_offsets_.data(), \ + diagonal_block_offsets_.data(), diagonal_blocks_.data(), inverse_blocks_.data(), \ + device_failure_flag_.data()) + switch (group.dimension) { + case 1: + LAUNCH_BLOCK_INVERSE(1); + break; + case 2: + LAUNCH_BLOCK_INVERSE(2); + break; + case 3: + LAUNCH_BLOCK_INVERSE(3); + break; + case 4: + LAUNCH_BLOCK_INVERSE(4); + break; + case 6: + LAUNCH_BLOCK_INVERSE(6); + break; + case 7: + LAUNCH_BLOCK_INVERSE(7); + break; + case 9: + LAUNCH_BLOCK_INVERSE(9); + break; + case 15: + LAUNCH_BLOCK_INVERSE(15); + break; + case 16: + LAUNCH_BLOCK_INVERSE(16); + break; + default: + LAUNCH_BLOCK_INVERSE(0); + } +#undef LAUNCH_BLOCK_INVERSE + } + } +} + +bool GeneralizedSchurSolver::CheckDeviceSuccess(cudaStream_t stream) { + int failure_flag = 0; + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&failure_flag, device_failure_flag_.data(), sizeof(int), + cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + return failure_flag == 0; +} + +GeneralizedSchurSolver::MatrixView GeneralizedSchurSolver::MakeTransformedView( + MatrixView matrix) const { + matrix.transformed_values = transformed_cross_blocks_.data(); + matrix.transformed_offsets = transformed_offsets_.data(); + matrix.row_state_ids = row_state_ids_.data(); + matrix.state_offsets = state_row_offsets_.data(); + return matrix; +} + +void GeneralizedSchurSolver::ApplySchurOperator(cudaStream_t stream, MatrixView matrix, + const float *input, float *output) { + if (use_explicit_schur_) { + MultiplyDenseSchurMatrix<<>>( + num_retained_rows_, schur_matrix_.data(), input, output); + return; + } + MatrixView transformed_matrix = MakeTransformedView(matrix); + if (num_eliminated_rows_) { + MultiplyPartitionedMatrix<<<(num_eliminated_rows_ + 3) / 4, 128, 0, stream>>>( + transformed_matrix, num_eliminated_rows_, eliminated_row_ids_.data(), + retained_row_indices_.data(), input, nullptr, eliminated_product_.data()); + } + if (num_retained_rows_) { + MultiplyPartitionedMatrix<<>>( + matrix, num_retained_rows_, retained_row_ids_.data(), retained_row_indices_.data(), input, + eliminated_product_.data(), output); + } +} + +void GeneralizedSchurSolver::ApplyPreconditioner(cudaStream_t stream) { + ApplyBlockInverse<<<(num_retained_rows_ + 127) / 128, 128, 0, stream>>>( + num_retained_rows_, retained_row_ids_.data(), row_state_ids_.data(), + state_row_offsets_.data(), diagonal_block_offsets_.data(), inverse_blocks_.data(), + retained_row_indices_.data(), reduced_residual_.data(), preconditioned_residual_.data(), + false); +} + +bool GeneralizedSchurSolver::Solve(cudaStream_t stream, MatrixView matrix, + const dvector &rhs, dvector &solution) { + if (matrix.num_rows != num_rows_ || int(rhs.size()) != num_rows_ || + int(solution.size()) != num_rows_) { + return false; + } + auto handle = static_cast(cublas_handle_.GetHandle(stream)); + THROW_ON_CUBLAS_ERROR(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST)); + float rhs_norm = 0; + THROW_ON_CUBLAS_ERROR(cublasSnrm2(handle, num_rows_, rhs.data(), 1, &rhs_norm)); + if (!std::isfinite(rhs_norm)) { + return false; + } + if (rhs_norm == 0) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(solution.data(), 0, num_rows_ * sizeof(float), stream)); + return true; + } + THROW_ON_CUDA_ERROR(cudaMemsetAsync(device_failure_flag_.data(), 0, sizeof(int), stream)); + ExtractDiagonalBlocks<<<(num_state_blocks_ + 127) / 128, 128, 0, stream>>>( + matrix, num_state_blocks_, state_row_offsets_.data(), diagonal_block_offsets_.data(), + diagonal_blocks_.data(), unobservable_rows_.data()); + InvertDimensionGroups(stream, BlockRole::Eliminated); + if (!CheckDeviceSuccess(stream)) { + LogError("GeneralizedSchur: non-positive eliminated pivot"); + return false; + } + if (num_eliminated_blocks_) { + TransformCrossBlocks<<<(num_eliminated_blocks_ + 3) / 4, 128, 0, stream>>>( + matrix, num_eliminated_blocks_, eliminated_block_ids_.data(), state_row_offsets_.data(), + retained_row_indices_.data(), diagonal_block_offsets_.data(), transformed_offsets_.data(), + inverse_blocks_.data(), transformed_cross_blocks_.data()); + } + MatrixView transformed_matrix = MakeTransformedView(matrix); + if (use_explicit_schur_) { + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(schur_matrix_.data(), 0, + size_t(num_retained_rows_) * num_retained_rows_ * sizeof(double), stream)); + CopyRetainedMatrix<<>>( + matrix, num_retained_rows_, retained_row_ids_.data(), retained_row_indices_.data(), + schur_matrix_.data()); + if (num_eliminated_blocks_) { + AccumulateSchurUpdates<<>>( + matrix, transformed_matrix, num_retained_rows_, eliminated_block_ids_.data(), + state_row_offsets_.data(), retained_row_indices_.data(), schur_matrix_.data()); + } + MirrorUpperTriangle<<<(size_t(num_retained_rows_) * num_retained_rows_ + 255) / 256, 256, 0, + stream>>>(num_retained_rows_, schur_matrix_.data()); + ExtractExplicitSchurDiagonal<<<(num_retained_blocks_ + 127) / 128, 128, 0, stream>>>( + num_retained_blocks_, num_retained_rows_, schur_matrix_.data(), retained_block_ids_.data(), + state_row_offsets_.data(), retained_row_indices_.data(), diagonal_block_offsets_.data(), + diagonal_blocks_.data()); + } else if (num_retained_blocks_) { + AssembleImplicitSchurDiagonal<<>>( + matrix, transformed_matrix, retained_block_ids_.data(), state_row_offsets_.data(), + retained_row_indices_.data(), diagonal_block_offsets_.data(), diagonal_blocks_.data()); + } + if (num_retained_blocks_) { + RegularizeUnobservableRows<<<(num_retained_blocks_ + 127) / 128, 128, 0, stream>>>( + num_retained_blocks_, retained_block_ids_.data(), state_row_offsets_.data(), + diagonal_block_offsets_.data(), unobservable_rows_.data(), diagonal_blocks_.data()); + } + if (use_direct_solve_) { + FactorSmallSchurMatrix<<<1, 128, 0, stream>>>( + num_retained_rows_, schur_matrix_.data(), retained_row_ids_.data(), + unobservable_rows_.data(), terminal_factor_.data(), device_failure_flag_.data()); + } else { + InvertDimensionGroups(stream, BlockRole::Retained); + } + if (!CheckDeviceSuccess(stream)) { + LogError("GeneralizedSchur: non-positive reduced preconditioner"); + return false; + } + SolveWithCurrentFactors(stream, matrix, rhs.data(), solution.data(), rhs_norm); + for (int refinement = 0; refinement <= kMaxRefinementPasses; ++refinement) { + ComputeFullResidual<32><<<(num_rows_ + 3) / 4, 128, 0, stream>>>( + matrix, rhs.data(), solution.data(), full_residual_.data(), nullptr); + if (!long_residual_row_ids_.empty()) { + ComputeFullResidual<128><<>>( + matrix, rhs.data(), solution.data(), full_residual_.data(), + long_residual_row_ids_.data()); + } + float residual_norm; + THROW_ON_CUBLAS_ERROR(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST)); + THROW_ON_CUBLAS_ERROR(cublasSnrm2(handle, num_rows_, full_residual_.data(), 1, &residual_norm)); + if (residual_norm <= std::max(kAbsoluteTolerance, kRelativeTolerance * rhs_norm)) { + return true; + } + if (!std::isfinite(residual_norm)) { + return false; + } + if (refinement == kMaxRefinementPasses) { + LogError("GeneralizedSchur: residual {} exceeds tolerance (rhs {})", residual_norm, rhs_norm); + return false; + } + SolveWithCurrentFactors(stream, matrix, full_residual_.data(), correction_.data(), + residual_norm); + AddCorrection<<<(num_rows_ + 127) / 128, 128, 0, stream>>>(num_rows_, correction_.data(), + solution.data()); + } + return false; +} + +void GeneralizedSchurSolver::SolveWithCurrentFactors(cudaStream_t stream, MatrixView matrix, + const float *rhs, float *solution, + float rhs_norm) { + auto handle = static_cast(cublas_handle_.GetHandle(stream)); + MatrixView transformed_matrix = MakeTransformedView(matrix); + if (num_eliminated_rows_) { + ApplyBlockInverse<<<(num_eliminated_rows_ + 127) / 128, 128, 0, stream>>>( + num_eliminated_rows_, eliminated_row_ids_.data(), row_state_ids_.data(), + state_row_offsets_.data(), diagonal_block_offsets_.data(), inverse_blocks_.data(), + retained_row_indices_.data(), rhs, eliminated_rhs_solution_.data(), true); + } + if (num_retained_rows_) { + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(retained_solution_.data(), 0, num_retained_rows_ * sizeof(float), stream)); + } + if (num_retained_rows_) { + MultiplyPartitionedMatrix<<>>( + matrix, num_retained_rows_, retained_row_ids_.data(), retained_row_indices_.data(), + retained_solution_.data(), eliminated_rhs_solution_.data(), reduced_rhs_.data()); + AddRetainedRhs<<<(num_retained_rows_ + 127) / 128, 128, 0, stream>>>( + num_retained_rows_, retained_row_ids_.data(), rhs, reduced_rhs_.data()); + if (use_direct_solve_) { + SolveSmallSchurSystem<<<1, 1, 0, stream>>>(num_retained_rows_, terminal_factor_.data(), + reduced_rhs_.data(), retained_solution_.data()); + } else { + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(reduced_residual_.data(), reduced_rhs_.data(), + num_retained_rows_ * sizeof(float), + cudaMemcpyDeviceToDevice, stream)); + ApplyPreconditioner(stream); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(search_direction_.data(), preconditioned_residual_.data(), + num_retained_rows_ * sizeof(float), + cudaMemcpyDeviceToDevice, stream)); + THROW_ON_CUBLAS_ERROR(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_DEVICE)); + THROW_ON_CUBLAS_ERROR(cublasSdot(handle, num_retained_rows_, reduced_residual_.data(), 1, + preconditioned_residual_.data(), 1, + cg_dot_products_.data())); + for (int iteration = 0; iteration < kMaxIterations; ++iteration) { + ApplySchurOperator(stream, matrix, search_direction_.data(), operator_product_.data()); + THROW_ON_CUBLAS_ERROR(cublasSdot(handle, num_retained_rows_, search_direction_.data(), 1, + operator_product_.data(), 1, cg_dot_products_.data() + 1)); + UpdateCgSolutionAndResidual<<<(num_retained_rows_ + 127) / 128, 128, 0, stream>>>( + num_retained_rows_, cg_dot_products_.data(), cg_dot_products_.data() + 1, + search_direction_.data(), operator_product_.data(), retained_solution_.data(), + reduced_residual_.data()); + ApplyPreconditioner(stream); + THROW_ON_CUBLAS_ERROR(cublasSdot(handle, num_retained_rows_, reduced_residual_.data(), 1, + preconditioned_residual_.data(), 1, + cg_dot_products_.data() + 2)); + UpdateCgDirection<<<(num_retained_rows_ + 127) / 128, 128, 0, stream>>>( + num_retained_rows_, cg_dot_products_.data() + 2, cg_dot_products_.data(), + preconditioned_residual_.data(), search_direction_.data()); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(cg_dot_products_.data(), cg_dot_products_.data() + 2, + sizeof(float), cudaMemcpyDeviceToDevice, stream)); + if ((iteration + 1) % 4 == 0) { + float residual_norm; + THROW_ON_CUBLAS_ERROR(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST)); + THROW_ON_CUBLAS_ERROR( + cublasSnrm2(handle, num_retained_rows_, reduced_residual_.data(), 1, &residual_norm)); + THROW_ON_CUBLAS_ERROR(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_DEVICE)); + if (!std::isfinite(residual_norm)) { + break; + } + if (residual_norm <= std::max(kAbsoluteTolerance, .2f * kRelativeTolerance * rhs_norm)) { + break; + } + } + } + } + } + if (num_eliminated_rows_) { + MultiplyPartitionedMatrix<<<(num_eliminated_rows_ + 3) / 4, 128, 0, stream>>>( + transformed_matrix, num_eliminated_rows_, eliminated_row_ids_.data(), + retained_row_indices_.data(), retained_solution_.data(), nullptr, + eliminated_product_.data()); + } + RecoverFullSolution<<<(num_rows_ + 127) / 128, 128, 0, stream>>>( + num_rows_, retained_row_indices_.data(), retained_solution_.data(), + eliminated_rhs_solution_.data(), eliminated_product_.data(), solution); +} +} // namespace cunls diff --git a/cunls/linear_solver/generalized_schur_solver.h b/cunls/linear_solver/generalized_schur_solver.h new file mode 100644 index 0000000..3341563 --- /dev/null +++ b/cunls/linear_solver/generalized_schur_solver.h @@ -0,0 +1,276 @@ +/* + * 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/cublas_helper.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" + +namespace cunls { + +namespace internal { +struct SchurMatrixView; +} // namespace internal + +/** + * @brief Automatic state-block Schur elimination for SPD CSR/BSR systems. + * + * Selects independent free tangent states across batches, solves the retained + * system internally, and recovers the complete update. No factor-type hints or + * additional options are needed. Does not delegate to another solver backend. + * Initialize must be repeated when the state layout or matrix pattern changes. + * Solve refreshes all numeric factors and checks the original-system residual. + * See docs/sphinx/generalized_schur.rst for equations and numerical limits. + */ +class GeneralizedSchurSolver final : public SparseLinearSolver { + public: + /** @brief Accept BSR matrices without expanding them to scalar CSR storage. */ + bool SupportsBlockStorage() const override; + + /** + * @brief Build the symbolic elimination plan for a scalar CSR matrix. + * + * Free tangent-state dimensions come from problem. An empty problem instead + * treats each scalar row as one state. Rows in each logical dense state block + * must share a sorted column pattern. No numeric factorization is cached here. + * + * @param stream Stream used for validation, selection, and workspace setup. + * @param problem Source of the free-state layout; constant states are excluded. + * @param matrix Symmetric working matrix with a fixed block sparsity pattern. + * @param rhs Right-hand side; used to validate the required vector length. + * @param solution Preallocated output with the same length as rhs. + * @return False for incompatible dimensions or an invalid block pattern. + * + * Call again when the state layout or sparsity changes. Numeric matrix values + * may change between Solve calls without rebuilding this plan. + */ + bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &matrix, + const dvector &rhs, dvector &solution) override; + + /** + * @brief Build the same plan directly from BSR storage. + * + * Logical state blocks may span several BSR tiles. If problem has no state + * metadata, each BSR tile is treated as one state block. The vector-size, + * symmetry, and reinitialization requirements match the CSR overload. + */ + bool Initialize(cudaStream_t stream, const Problem &problem, const BSRSparseMatrix &matrix, + const dvector &rhs, dvector &solution) override; + + /** + * @brief Solve matrix * solution = rhs using the initialized CSR pattern. + * + * Refreshes pivot factors and the Schur preconditioner, solves the retained + * system internally, and recovers eliminated updates in the original ordering. + * The input solution is not used as a warm start. Local factors and Schur + * accumulation use double precision; vectors use float. The backend checks + * the original-system residual and may refine the update with the same factors. + * + * @return True when the full residual meets the internal tolerance; false for + * dimension mismatches, invalid pivots, nonfinite values, or failed accuracy. + * On failure, solution may contain an incomplete update. CUDA/cuBLAS errors + * follow the library's exception-based error handling. Host convergence checks + * synchronize stream; this method is not a fully asynchronous enqueue. + */ + bool Solve(cudaStream_t stream, const CSRSparseMatrix &matrix, const dvector &rhs, + dvector &solution) override; + + /** @brief BSR overload with the same numeric and failure contract as CSR Solve. */ + bool Solve(cudaStream_t stream, const BSRSparseMatrix &matrix, const dvector &rhs, + dvector &solution) override; + + private: + using MatrixView = internal::SchurMatrixView; + + /** @brief Select which diagonal factors a dimension group represents. */ + enum class BlockRole { Eliminated, Retained }; + + /** + * @brief State blocks sharing a factorization role and a kernel specialization. + * dimension == 0 selects the generic kernel; actual sizes remain in the state + * offsets. block_ids contains original logical-state indices on the GPU. + */ + struct DimensionGroup { + BlockRole role; + int dimension; + dvector block_ids; + }; + + /** @brief Common symbolic setup for the CSR and BSR adapters. */ + bool Initialize(cudaStream_t stream, const Problem &problem, MatrixView matrix, + const dvector &rhs, dvector &solution); + + /** @brief Factor and invert all diagonal blocks of the requested role. */ + void InvertDimensionGroups(cudaStream_t stream, BlockRole role); + + /** @brief Synchronize stream and read the device factorization/validation flag. */ + bool CheckDeviceSuccess(cudaStream_t stream); + + /** @brief View packed D^-1 B values using the original matrix's column graph. */ + MatrixView MakeTransformedView(MatrixView matrix) const; + + /** @brief Compute output = S * input in retained order, explicitly or implicitly. */ + void ApplySchurOperator(cudaStream_t stream, MatrixView matrix, const float *input, + float *output); + + /** @brief Apply inverse Schur diagonal blocks to the current reduced residual. */ + void ApplyPreconditioner(cudaStream_t stream); + + /** @brief Common numeric factorization, solve, and full-residual verification. */ + bool Solve(cudaStream_t stream, MatrixView matrix, const dvector &rhs, + dvector &solution); + + /** + * @brief Reduce rhs, solve S, and recover the full solution with existing factors. + * + * Used for the initial solve and refinement corrections. rhs_norm scales the + * reduced stopping criterion. The caller verifies the full-system residual; + * this helper does not independently establish convergence. + */ + void SolveWithCurrentFactors(cudaStream_t stream, MatrixView matrix, const float *rhs, + float *solution, float rhs_norm); + + /** @brief Factorization buckets built during symbolic initialization. */ + std::vector dimension_groups_; + + /** @brief Total number of free scalar tangent coordinates. */ + int num_rows_ = 0; + + /** @brief Number of logical free-state blocks. */ + int num_state_blocks_ = 0; + + /** @brief Scalar coordinates selected for elimination. */ + int num_eliminated_rows_ = 0; + + /** @brief Scalar coordinates in the reduced system. */ + int num_retained_rows_ = 0; + + /** @brief Number of independent eliminated state blocks. */ + int num_eliminated_blocks_ = 0; + + /** @brief Number of retained state blocks. */ + int num_retained_blocks_ = 0; + + /** @brief Prefix offsets of each logical state in the original row order. */ + dvector state_row_offsets_; + + /** @brief Original scalar row to logical state-block index. */ + dvector row_state_ids_; + + /** @brief Selection state: 0 undecided, 1 eliminated, 2 excluded/retained. */ + dvector selection_status_; + + /** @brief Temporary flags for a single independent-set selection round. */ + dvector selected_in_round_; + + /** @brief Stored scalar columns per state row, including its diagonal block. */ + dvector state_row_widths_; + + /** @brief Eliminated block indices in stable original order. */ + dvector eliminated_block_ids_; + + /** @brief Retained block indices in stable original order. */ + dvector retained_block_ids_; + + /** @brief Compact eliminated index to original scalar row. */ + dvector eliminated_row_ids_; + + /** @brief Compact retained index to original scalar row. */ + dvector retained_row_ids_; + + /** @brief Original row to compact retained index; -1 marks eliminated rows. */ + dvector retained_row_indices_; + + /** @brief One device integer set by validation or failed Cholesky pivots. */ + dvector device_failure_flag_; + + /** @brief Flags for exactly zero numerical rows in the working matrix. */ + dvector unobservable_rows_; + + /** @brief State prefixes within the selected or retained scalar partition. */ + dvector partition_row_offsets_; + + /** @brief Original rows assigned a full CTA for residual evaluation. */ + dvector long_residual_row_ids_; + + /** @brief Offsets of row-major dense blocks in diagonal/inverse storage. */ + dvector diagonal_block_offsets_; + + /** @brief Offsets of packed transformed rows for eliminated state blocks. */ + dvector transformed_offsets_; + + /** @brief Whether the plan materializes the dense reduced matrix. */ + bool use_explicit_schur_ = false; + + /** @brief Whether the retained system uses the small Cholesky terminal path. */ + bool use_direct_solve_ = false; + + /** @brief Pivot blocks and Schur diagonal blocks; generic inversion scratch. */ + dvector diagonal_blocks_; + + /** @brief Explicit row-major reduced matrix, when selected by the planner. */ + dvector schur_matrix_; + + /** @brief Dense inverses of eliminated pivots and retained preconditioner blocks. */ + dvector inverse_blocks_; + + /** @brief Packed D^-1 B values; retained rows have no allocation here. */ + dvector transformed_cross_blocks_; + + /** @brief Small reduced-system Cholesky factor for the direct path. */ + dvector terminal_factor_; + + /** @brief D^-1 b_e in original row order; only eliminated entries are defined. */ + dvector eliminated_rhs_solution_; + + /** @brief D^-1 B times a retained vector, stored in original row order. */ + dvector eliminated_product_; + + /** @brief Current reduced solution in compact retained order. */ + dvector retained_solution_; + + /** @brief b_r - B^T D^-1 b_e in compact retained order. */ + dvector reduced_rhs_; + + /** @brief CG residual g - S*x_r. */ + dvector reduced_residual_; + + /** @brief Block-Jacobi preconditioned reduced residual. */ + dvector preconditioned_residual_; + + /** @brief Current CG search direction. */ + dvector search_direction_; + + /** @brief Schur operator applied to the CG search direction. */ + dvector operator_product_; + + /** @brief Device dot products: current r^Tz, p^TSp, next r^Tz, spare. */ + dvector cg_dot_products_; + + /** @brief Original-system residual used for validation and refinement. */ + dvector full_residual_; + + /** @brief Full-order refinement update before adding it to the solution. */ + dvector correction_; + + /** @brief Lazily initialized handle for vector norms and dot products. */ + cuBLASHandle cublas_handle_; +}; + +} // namespace cunls diff --git a/cunls/linear_solver/llms.txt b/cunls/linear_solver/llms.txt index 9ade221..748e123 100644 --- a/cunls/linear_solver/llms.txt +++ b/cunls/linear_solver/llms.txt @@ -60,3 +60,29 @@ Sparse linear system abstraction and implementation used by minimizers. - `SlowInitFastSolve`: expensive setup, fast repeated solves. - `FastInitSlowSolve`: faster setup, slower repeated solves. + +## Generalized Schur complement + +- `generalized_schur_solver.h` / `.cu`: backend selected by + `SparseLinearSolverType::GeneralizedSchur`; no extra public options. +- GPU independent-set selection across all free logical state blocks, four + deterministic rounds. Single elimination level, automatic dense explicit + Schur matrix when affordable; otherwise implicit operator. No backend fallback. +- FP64 local Cholesky/inverses, cross-block transforms, Schur accumulation; + FP32 Krylov vectors. Internal direct terminal solve for <=64 retained rows, + otherwise internal reduced block-Jacobi CG. Full-system residual validation + at relative tolerance 1e-3, with at most three refinement corrections. +- Exactly zero-information rows choose zero motion; general singular or + indefinite systems are not supported. Symbolic plan must be rebuilt for + changed connectivity/constant IDs via the existing reuse_structure contract. +- Tests: generalized_schur_solver_test.cpp; SolverTypeFromEnv accepts + CUNLS_SOLVER=GeneralizedSchur in existing SBA/PGO tests. +- Reproducible SBA timing: scripts/benchmark_schur_in_docker.sh; user guide + and algorithm details in docs/sphinx/generalized_schur.rst. + +- No Thrust: stable block prefixes and dimension buckets in one host pass, + scalar row maps expanded by a coalesced CUDA kernel. Numeric FP64 cross + blocks packed only for eliminated rows; shared-memory three-dimensional + clique updates accumulate the upper triangle, then mirror it. Long residual + rows use a CTA, short rows a warp. SM 86 and SM 87 compilation verified; + runtime performance on Jetson remains unmeasured. diff --git a/cunls/linear_solver/sparse_linear_solver.cpp b/cunls/linear_solver/sparse_linear_solver.cpp index 6b1e552..ad3cdf2 100644 --- a/cunls/linear_solver/sparse_linear_solver.cpp +++ b/cunls/linear_solver/sparse_linear_solver.cpp @@ -20,6 +20,7 @@ #include #include "cunls/common/cudss_helper.h" +#include "cunls/linear_solver/generalized_schur_solver.h" namespace cunls { @@ -48,6 +49,8 @@ SparseLinearSolverPtr CreateSparseLinearSolver(SparseLinearSolverType type, return std::make_unique(); case SparseLinearSolverType::BlockSparsePCG: return std::make_unique(config.block_sparse_pcg_options); + case SparseLinearSolverType::GeneralizedSchur: + return std::make_unique(); case SparseLinearSolverType::BlockTridiagonal: return std::make_unique(); default: diff --git a/cunls/linear_solver/sparse_linear_solver.h b/cunls/linear_solver/sparse_linear_solver.h index 6cb5e3e..ce9c578 100644 --- a/cunls/linear_solver/sparse_linear_solver.h +++ b/cunls/linear_solver/sparse_linear_solver.h @@ -54,6 +54,7 @@ enum class SparseLinearSolverType { ///< in time (trajectories, MPC): needs the states' ///< stages (Problem::SetStateStages). One warp per ///< subproblem, no host synchronization. + GeneralizedSchur, ///< Automatic block elimination with an internal reduced-system solve. }; /** diff --git a/docs/sphinx/api/linear_solver.rst b/docs/sphinx/api/linear_solver.rst index 7060ad4..2441f6f 100644 --- a/docs/sphinx/api/linear_solver.rst +++ b/docs/sphinx/api/linear_solver.rst @@ -23,6 +23,10 @@ Enum in `cunls/linear_solver/sparse_linear_solver.h`: gradient solver. Default backend for Gauss-Newton and Levenberg-Marquardt: outperforms ``cuDSS`` on most SBA / PGO workloads (see ``profile/PCG_RESULTS.md``). +- `GeneralizedSchur` — automatically eliminates independent state blocks, + solves the reduced system internally, and recovers all updates. Accepts + heterogeneous state dimensions and CSR/BSR; no additional configuration + and no fallback to another solver backend. See :doc:`../generalized_schur`. - `cuDSS` — NVIDIA cuDSS sparse direct solver. Pick when each Solve sees a tiny system and PCG's per-iter kernel-launch overhead dominates. @@ -42,7 +46,7 @@ chosen ``SparseLinearSolverType`` is used: ignored when a different backend is selected. - `cudss_solver_options` - [in] cuDSS-specific options; ignored when a different backend is selected. -- Dense backends take no extra configuration. +- GeneralizedSchur and dense backends take no extra configuration. SparseLinearSolver --------------------- @@ -360,6 +364,23 @@ Works for any non-singular square matrix (not limited to SPD). :param ``result``: [out] Solution vector ``x`` (size must equal matrix rows). :returns: ``true`` on success; ``false`` if a dimension mismatch is detected. +GeneralizedSchurSolver +---------------------- + +Declared in ``cunls/linear_solver/generalized_schur_solver.h`` with its CUDA +implementation in ``generalized_schur_solver.cu``. The class implements both +CSR and BSR overloads of ``SparseLinearSolver::Initialize`` and ``Solve``. +Prefer selecting ``SparseLinearSolverType::GeneralizedSchur`` through the +existing factory or minimizer options; no additional configuration is required. + +``Initialize`` derives the free tangent-state partition and allocates the +symbolic plan. ``Solve`` refreshes numeric factors, solves the Schur system, +recovers eliminated updates, and checks the residual against the original +matrix. Changed state layout or sparsity requires reinitialization. Invalid +pivots or failure to reach the residual tolerance return ``false``. +The complete algorithm, GPU implementation, numerical limits, and benchmarks +are documented in :doc:`../generalized_schur`. + Factory ------- diff --git a/docs/sphinx/generalized_schur.rst b/docs/sphinx/generalized_schur.rst new file mode 100644 index 0000000..24bf59d --- /dev/null +++ b/docs/sphinx/generalized_schur.rst @@ -0,0 +1,205 @@ +Generalized Schur solver +======================== + +Select the backend using the existing minimizer options:: + + options.sparse_linear_solver_type = cunls::SparseLinearSolverType::GeneralizedSchur; + +Python exposes the same ``SparseLinearSolverType.GeneralizedSchur`` value. +There are no additional public options or required state annotations. + +For example, with Python Levenberg--Marquardt: + +.. code-block:: python + + options = pycunls.LevenbergMarquardtMinimizerOptions() + options.base_options.sparse_linear_solver_type = ( + pycunls.SparseLinearSolverType.GeneralizedSchur + ) + summary = pycunls.LevenbergMarquardtMinimizer(options).minimize(stream, problem) + +The C++ and Python sparse bundle adjustment examples, and the Python +visual-inertial bundle adjustment example, select this backend explicitly. + +How it works +------------ + +The solver builds a graph of free tangent-state blocks and automatically selects +an independent set using connectivity, block dimension, and deterministic +tie-breaking. Selection crosses state-batch boundaries: batches are storage +groups, not elimination levels. The implementation uses one elimination level; +it does not recursively eliminate each batch. It accepts general block-sparse +graphs and mixed state dimensions, without identifying cameras or landmarks. + +For the partitioned system ``[D B; B^T C]``, it factors the independent diagonal +blocks of ``D``, solves the reduced system ``S = C - B^T D^-1 B``, and recovers +the eliminated variables. Small reduced systems use an internal Cholesky solve. +Larger systems use an internal block-preconditioned conjugate-gradient iteration. +This backend owns these operations and never calls another cuNLS solver backend. + +The reduced matrix is explicitly assembled when its dimension, estimated +assembly work, and available memory permit. Otherwise, an implicit operator +avoids materializing fill. Local factors, transformed cross blocks, and explicit +Schur accumulation use double precision; the input matrix and Krylov vectors +remain float. Symbolic work can be reused through the existing +``reuse_structure`` option when the state layout and sparsity pattern are fixed. +Numeric factors are refreshed for each solve, including damping changes. + +The intended input is the symmetric positive-definite working system supplied +by the minimizer. General indefinite systems and unresolved gauge singularities +are not supported. Exactly zero-information rows with a consistent zero right +hand side receive zero updates. Invalid pivots cause failure rather than a +fallback to another backend. A full-system residual check requires relative +residual at most ``1e-3``; up to three refinement corrections are available. +The internal reduced iteration limit is 200. These are internal constants, +independent of PCG configuration options. + +Elimination and recovery equations +---------------------------------- + +Let ``H x = b`` be the working linear system after the minimizer has applied +damping and column scaling. A permutation places the selected states first: + +.. math:: + + \begin{bmatrix}D & B\\ B^T & C\end{bmatrix} + \begin{bmatrix}x_e\\x_r\end{bmatrix} + = \begin{bmatrix}b_e\\b_r\end{bmatrix}. + +No two selected states share a structural Hessian edge, so ``D`` is block +diagonal. Each selected state can therefore be factored independently on the +GPU. Substitution gives the reduced right hand side and recovery formula: + +.. math:: + + g = b_r - B^T D^{-1}b_e,\qquad + S x_r = g,\qquad S = C - B^T D^{-1}B, + +.. math:: + + x_e = D^{-1}(b_e-Bx_r). + +This is an algebraic elimination within a linear solve. It does not remove +states from ``Problem`` or create persistent marginalization priors. The full +update is returned in the caller's original ordering. When ``H`` is positive +definite, both the pivot blocks and the Schur complement are positive definite. + +Automatic selection and multiple batches +---------------------------------------- + +The planner uses complete tangent states as vertices, not individual BSR tiles +or entire batches. For example, a six-dimensional pose remains one vertex even +when stored using three-dimensional BSR tiles. Constant states have already +been removed from the free-variable system. Raw matrix solves without problem +metadata use scalar CSR entries or BSR tiles as vertices instead. + +Four GPU selection rounds favor low-connectivity, small-dimensional states. +The score is the number of stored scalar columns in the state's first row +multiplied by its tangent dimension, with a deterministic hash to break ties. +Each round selects local priority +minima and excludes their neighbors from subsequent selection. Any undecided +states are retained after the fourth round; maximality is not required for a +valid Schur decomposition. The heuristic does not promise minimum fill. + +In ordinary SBA, landmarks tend to be selected because they are small and have +few neighbors. On other graphs, poses or other state types may be selected. +States from several batches can be eliminated together whenever they are +independent. Connected states cannot both enter this diagonal elimination set. +Repeated elimination per batch would introduce dependencies and extra fill; +this implementation instead solves all retained states together in one reduced +system. Multiple elimination levels are not implemented. + +Storage, GPU work, and numerical checks +--------------------------------------- + +Symbolic initialization validates the block pattern, builds the partition and +row maps, and groups pivot work by dimension. Common sizes have specialized +kernels; a generic path handles other dimensions. CSR rows belonging to a +logical dense block must have consistent sorted column patterns, as produced +by the library's Hessian assembler. + +Each numeric solve refreshes local Cholesky factors, their small dense inverses, +and transformed cross blocks ``W = D^-1 B``. Explicit assembly adds each +eliminated state's contribution to its retained-neighbor clique. It is chosen +only for at most 4096 retained scalar dimensions, an estimated clique-work +budget of 512 million scalar pairs, and a conservative allocation estimate +below one quarter of currently free GPU memory. These internal heuristics +avoid allocating a large dense reduced matrix for general sparse graphs. + +The implicit path evaluates ``S v = C v - B^T (W v)`` from sparse storage. +Both iterative paths use block Jacobi formed from the actual diagonal blocks +of ``S``, including the eliminated-state contributions. Explicit systems of +at most 64 retained dimensions use the internal direct Cholesky path. An empty +retained system only needs independent local solves and recovery. + +Double-precision local work and Schur accumulation reduce cancellation when +forming ``C - B^T D^-1 B``. The reduced iterative stopping target is tighter +than the final full-system target. After recovery, the backend recomputes +``b - Hx`` using the original matrix. If necessary, it solves for a correction +with the same factors and checks again. Failure to meet the original-system +residual target is reported as failure, including when advisory safety checks +are disabled. This is a global residual criterion, not a separate accuracy +guarantee for each disconnected component. + +Reusing structure reuses the partition and allocations, not stale numerical +factors. Changing the state layout or sparsity requires symbolic +reinitialization under the existing solver lifecycle. Small or strongly damped +problems can favor BlockSparsePCG because Schur's numeric setup costs more than +a few full-system iterations; the measurements below expose that tradeoff. + +GPU implementation and Jetson Orin +---------------------------------- + +The backend uses CUDA kernels directly; it does not use Thrust. Selection +metadata is copied once for a single host pass that builds stable state prefixes, +size buckets, and the fill estimate. One coalesced GPU pass expands the prefixes +into scalar row maps. Graph edges stay on the GPU. This avoids separate +selection-library compactions and their scratch allocations and synchronizations. + +The numeric path is organized around the following costs: + +* Three-dimensional pivots with at most 64 stored scalar columns stage their + cross blocks in shared memory. Threads reuse those values for the clique + update instead of repeatedly loading and indexing sparse rows. Wider pivots + and other dimensions retain a general path. +* Explicit assembly computes only the upper triangle, then mirrors it. This + approximately halves its dot products and atomic updates while preserving + symmetry. Double-precision accumulation remains in place. +* Transformed cross blocks are packed by eliminated row. There is no unused + transformed storage for retained rows, and neighboring lanes read contiguous + transformed values. Dimension-specialized local inversions use compile-time + loop bounds. +* Sparse products use warp reductions. The original-system residual assigns a + warp to short rows and a full thread block to rows wider than 128 scalar + columns, avoiding a single warp serializing a highly connected state. + +The implementation uses 32-thread warps, ordinary shared memory, shuffle +reductions, and native double atomics. It does not depend on SM 90 features, +thread-block clusters, or desktop-only CUDA libraries. The optimized build was +compiled for both SM 86 and Orin's SM 87; its binary contains both targets. +The largest static shared-memory allocation is 32 KiB for the small terminal +factorization. Orin runtime correctness and performance still require the +Jetson runner; compiling an SM 87 target on x86 does not establish either. +FP64 work is deliberately retained for numerical stability and remains an +important cost on GPUs with limited FP64 throughput. + +Validation and reproduction +--------------------------- + +The dedicated generalized-Schur tests cover mixed state dimensions, CSR/BSR +storage, explicit and implicit operators, numeric reuse, zero right hand sides, +invalid pivots, wide cliques, and long residual rows. Python tests exercise the +backend enum, scaling, projected bounds, and numeric Jacobians. The C++ SBA, +Python SBA, and Python visual-inertial BA examples select the solver explicitly. + +Build and test using the standard repository scripts:: + + EXTRA_CMAKE_ARGS="-DBUILD_TESTING=ON" ./scripts/build_cunls_in_docker.sh Release ./build_docker + ./scripts/test_cunls_in_docker.sh ./build_docker + +The optional ``scripts/benchmark_schur_in_docker.sh`` script compares this +backend with BlockSparsePCG and saves its results under the build output +folder. It records the actual CUDA device and checks matching problem sizes, +linear residuals, and final nonlinear costs. Performance depends on the graph, +damping, GPU, and reuse of symbolic work; Schur setup can cost more than a few +PCG iterations on small or strongly damped systems. diff --git a/docs/sphinx/index.rst b/docs/sphinx/index.rst index 25789ae..cde9f17 100644 --- a/docs/sphinx/index.rst +++ b/docs/sphinx/index.rst @@ -50,6 +50,7 @@ Guides imu mpc custom_factors_and_states + generalized_schur numeric_jacobians =============================================================================== diff --git a/examples/sparse_bundle_adjustment/README.md b/examples/sparse_bundle_adjustment/README.md index 2a837f2..4ceb7be 100644 --- a/examples/sparse_bundle_adjustment/README.md +++ b/examples/sparse_bundle_adjustment/README.md @@ -95,3 +95,9 @@ cmake --build build/examples/all -j Output binary: - `./artifacts/examples/sparse_bundle_adjustment_example` + +The example selects `SparseLinearSolverType::GeneralizedSchur`. It automatically +eliminates independent state blocks (typically landmarks), solves the retained +system, and recovers the complete update. See the +[Schur solver guide](../../docs/sphinx/generalized_schur.rst) for the algorithm +and numerical and performance tradeoffs. diff --git a/examples/sparse_bundle_adjustment/main.cpp b/examples/sparse_bundle_adjustment/main.cpp index 614119c..9fcab62 100644 --- a/examples/sparse_bundle_adjustment/main.cpp +++ b/examples/sparse_bundle_adjustment/main.cpp @@ -105,6 +105,8 @@ int main() { // 6. Solve with Levenberg-Marquardt; the result is written into poses / points. cunls::LevenbergMarquardtMinimizerOptions options; + // Schur automatically eliminates independent states (typically landmarks in SBA). + options.base_options.sparse_linear_solver_type = cunls::SparseLinearSolverType::GeneralizedSchur; options.base_options.max_num_iterations = 80; options.base_options.state_tolerance = 1e-8f; options.base_options.cost_tolerance = 1e-8f; diff --git a/python/examples/imu_bundle_adjustment.py b/python/examples/imu_bundle_adjustment.py index cf73779..7e21296 100644 --- a/python/examples/imu_bundle_adjustment.py +++ b/python/examples/imu_bundle_adjustment.py @@ -183,6 +183,8 @@ def main(): problem.add_factor_batch(reproj, reproj_ptrs) options = pycunls.LevenbergMarquardtMinimizerOptions() + # Schur automatically eliminates independent states (typically landmarks in SBA). + options.base_options.sparse_linear_solver_type = pycunls.SparseLinearSolverType.GeneralizedSchur options.base_options.max_num_iterations = 30 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..3b6e37e 100644 --- a/python/examples/sparse_bundle_adjustment.py +++ b/python/examples/sparse_bundle_adjustment.py @@ -82,6 +82,8 @@ def main(): # 6. Solve with Levenberg-Marquardt. options = pycunls.LevenbergMarquardtMinimizerOptions() + # Schur automatically eliminates independent states (typically landmarks in SBA). + options.base_options.sparse_linear_solver_type = pycunls.SparseLinearSolverType.GeneralizedSchur options.base_options.max_num_iterations = 80 options.base_options.state_tolerance = 1e-8 options.base_options.cost_tolerance = 1e-8 diff --git a/python/pycunls/_pycunls_core.pyi b/python/pycunls/_pycunls_core.pyi index 45bd50b..a471f49 100644 --- a/python/pycunls/_pycunls_core.pyi +++ b/python/pycunls/_pycunls_core.pyi @@ -67,6 +67,7 @@ class SparseLinearSolverType(enum.Enum): DenseQR = ... BlockSparsePCG = ... BlockTridiagonal = ... + GeneralizedSchur = ... class ColumnScaling(enum.Enum): """Diagonal scaling mode for the GN/LM normal equations.""" diff --git a/python/src/bind_types.cpp b/python/src/bind_types.cpp index 0abdd43..5f4294d 100644 --- a/python/src/bind_types.cpp +++ b/python/src/bind_types.cpp @@ -94,6 +94,7 @@ void bind_types(nb::module_ &m) { .value("DenseCholesky", cunls::SparseLinearSolverType::DenseCholesky) .value("DenseQR", cunls::SparseLinearSolverType::DenseQR) .value("BlockSparsePCG", cunls::SparseLinearSolverType::BlockSparsePCG) + .value("GeneralizedSchur", cunls::SparseLinearSolverType::GeneralizedSchur) .value("BlockTridiagonal", cunls::SparseLinearSolverType::BlockTridiagonal); nb::enum_(m, "ColumnScaling") diff --git a/python/tests/test_minimizer.py b/python/tests/test_minimizer.py index 223dc03..1b0dd0d 100644 --- a/python/tests/test_minimizer.py +++ b/python/tests/test_minimizer.py @@ -55,6 +55,40 @@ def _make_prior_problem(): return problem, states_gpu, target +@pytest.mark.parametrize("scaled", [False, True]) +def test_generalized_schur_enum_and_solve(stream, scaled): + problem, states_gpu, target = _make_prior_problem() + opts = pycunls.MinimizerOptions() + opts.sparse_linear_solver_type = pycunls.SparseLinearSolverType.GeneralizedSchur + if scaled: + opts.column_scaling = pycunls.ColumnScaling.hessian_diagonal + summary = pycunls.GaussNewtonMinimizer(opts).minimize(stream, problem) + cp.cuda.runtime.streamSynchronize(stream.get_stream()) + assert summary.final_cost < 1e-6 + np.testing.assert_allclose(cp.asnumpy(states_gpu), target, atol=1e-4) + + +def test_generalized_schur_projected_bounds(stream): + x = cp.asarray([-3.0, 4.0, -1.0], dtype=cp.float32) + target = cp.asarray([1.0, 2.0, 3.0], dtype=cp.float32) + states = pycunls.VectorStateBatch3(x, 1) + states.set_num_active_states(1) + states.set_bounds(cp.asarray([0., 0., 0.], dtype=cp.float32), + cp.asarray([0.5, 3., 5.], dtype=cp.float32)) + factor = pycunls.PriorVectorFactorBatch3(target, 1) + factor.set_num_active_factors(1) + problem = pycunls.Problem() + problem.add_state_batch(states) + problem.add_factor_batch(factor, [states.state_device_ptr(0)]) + opts = pycunls.MinimizerOptions() + opts.sparse_linear_solver_type = pycunls.SparseLinearSolverType.GeneralizedSchur + opts.jacobian_mode = pycunls.JacobianMode.numeric + solver = pycunls.AugmentedLagrangianMinimizer(pycunls.GaussNewtonMinimizer(opts)) + solver.minimize(stream, problem) + cp.cuda.runtime.streamSynchronize(stream.get_stream()) + np.testing.assert_allclose(cp.asnumpy(x), [0.5, 2.0, 3.0], atol=1e-4) + + class TestGaussNewtonMinimizer: """Gauss-Newton convergence on the simple prior problem.""" def test_converges(self, stream): diff --git a/scripts/benchmark_schur_in_docker.sh b/scripts/benchmark_schur_in_docker.sh new file mode 100755 index 0000000..417eafc --- /dev/null +++ b/scripts/benchmark_schur_in_docker.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# Reproduce the Schur/PCG SBA measurements using an existing Release build. +set -euo pipefail +if [ "$#" -lt 1 ] || [ "$#" -gt 2 ]; then + echo "Usage: $0 [repetitions=20]" + exit 1 +fi +OUTPUT_DIR=$(realpath "$1") +REPEATS=${2:-20} +if ! [[ "$REPEATS" =~ ^[1-9][0-9]*$ ]]; then + echo "Repetitions must be a positive integer." + exit 1 +fi +if [ ! -x "$OUTPUT_DIR/build_shared/bin/nls_tests" ]; then + echo "Missing test binary. Run EXTRA_CMAKE_ARGS=-DBUILD_TESTING=ON ./scripts/build_cunls_in_docker.sh Release $OUTPUT_DIR" + exit 1 +fi +TTY_FLAG="" +[ -t 0 ] && TTY_FLAG="-it" +docker run --runtime=nvidia --gpus all --rm $TTY_FLAG \ + -v "$(pwd):/cunls:ro" -v "$OUTPUT_DIR:/cunls_install" \ + -e CUDA_DEVICE_ORDER=PCI_BUS_ID \ + -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" \ + -e CUNLS_BENCHMARK=1 -e CUNLS_BENCHMARK_REPEATS="$REPEATS" \ + -e CUNLS_PCG_TOL=0.001 -e CUNLS_PCG_MAX_ITER=200 \ + cunls:local bash -c ' + set -euo pipefail + nvidia-smi + mkdir -p /cunls_install/schur-benchmarks + for solver in BlockSparsePCG GeneralizedSchur; do + CUNLS_SOLVER=$solver /cunls_install/build_shared/bin/nls_tests \ + --gtest_filter="*SyntheticSba*" \ + --gtest_output="xml:/cunls_install/schur-benchmarks/$solver.xml" \ + | tee "/cunls_install/schur-benchmarks/$solver.log" + done + ' + +# Standard-library-only summary; preserve raw measurements beside the CSV. +docker run --rm -i --entrypoint python3 \ + -v "$OUTPUT_DIR:/cunls_install" cunls:local - /cunls_install/schur-benchmarks <<'PY' +import csv +import math +from pathlib import Path +import statistics +import sys + +root = Path(sys.argv[1]) +records = {} +for solver in ("BlockSparsePCG", "GeneralizedSchur"): + rows = [] + for line in (root / f"{solver}.log").read_text().splitlines(): + if line.startswith("SBA_"): + parts = line.split() + rows.append((parts[0], dict(part.split("=", 1) for part in parts[1:]))) + records[solver] = rows +scales = [record["scale"] for kind, record in records["BlockSparsePCG"] if kind == "SBA_SHAPE"] +assert scales, "No benchmark records found" +devices = [] +for solver, rows in records.items(): + found = [record for kind, record in rows if kind == "SBA_DEVICE"] + assert found and all(record == found[0] for record in found), (solver, "Missing or inconsistent GPU identity") + devices.append(found[0]) +assert devices[0] == devices[1], "Backends ran on different GPUs" +device = devices[0] +print(f"CUDA device: {device['name']} (SM {device['sm']}, PCI bus {device['pci_bus']})") +summary = [] +for scale in scales: + for stage in ("SBA_LINEAR", "SBA_WARM"): + stats = {} + for solver, rows in records.items(): + samples = sorted(float(record["ms"]) for kind, record in rows + if kind == stage and record["scale"] == scale) + assert samples, (solver, scale, stage) + stats[solver] = (statistics.median(samples), samples[math.ceil(.95 * len(samples))-1], len(samples)) + pcg, schur = stats["BlockSparsePCG"], stats["GeneralizedSchur"] + summary.append(dict(scale=scale, stage=stage, samples=pcg[2], + device=device["name"], compute_capability=device["sm"], + pcg_median_ms=pcg[0], schur_median_ms=schur[0], + pcg_p95_ms=pcg[1], schur_p95_ms=schur[1], speedup=pcg[0]/schur[0])) + assert pcg[2] == schur[2] + final_costs = {} + shapes = [] + for solver, rows in records.items(): + shape = next(record for kind, record in rows if kind == "SBA_SHAPE" and record["scale"] == scale) + shapes.append(shape) + residual = next(float(record["relative"]) for kind, record in rows if kind == "SBA_RESIDUAL" and record["scale"] == scale) + assert math.isfinite(residual) and residual <= 1.05e-3, (solver, scale, residual) + final_costs[solver] = statistics.median(float(record["final"]) for kind, record in rows + if kind == "SBA_WARM" and record["scale"] == scale) + assert shapes[0] == shapes[1], "Solvers must use identical input sizes" + pcg_cost, sc_cost = final_costs["BlockSparsePCG"], final_costs["GeneralizedSchur"] + assert abs(pcg_cost-sc_cost) <= max(1e-7, .001*abs(pcg_cost)), (scale, final_costs) +with (root / "summary.csv").open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=summary[0].keys()) + writer.writeheader() + writer.writerows(summary) +for row in summary: + print(f"{row['scale']:20} {row['stage']:11} PCG {row['pcg_median_ms']:9.3f} ms " + f"Schur {row['schur_median_ms']:9.3f} ms {row['speedup']:5.2f}x") +print(f"Raw logs, JUnit results, and summary: {root}") +PY diff --git a/scripts/plot_schur_benchmark.py b/scripts/plot_schur_benchmark.py new file mode 100644 index 0000000..865dd0e --- /dev/null +++ b/scripts/plot_schur_benchmark.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Plot benchmark_schur_in_docker.sh's CSV; requires matplotlib and numpy.""" + +import argparse +import csv +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("csv", type=Path) + parser.add_argument("output", type=Path, help="Output filename stem (PNG and SVG)") + parser.add_argument("--hardware", help="Override the GPU name recorded in the CSV") + args = parser.parse_args() + with args.csv.open() as source: + rows = list(csv.DictReader(source)) + + hardware = args.hardware or rows[0].get("device", "GPU not recorded").replace("_", " ") + plt.style.use("default") + fig, axes = plt.subplots(2, 1, figsize=(11, 9), layout="constrained") + for ax, stage, title in zip( + axes, + ("SBA_WARM", "SBA_LINEAR"), + ("Complete warm SBA optimization", "Frozen initial damped linear solve"), + ): + data = [row for row in rows if row["stage"] == stage] + if not data: + raise ValueError(f"Missing measurements for {stage}") + positions = np.arange(len(data)) + for offset, key, label, color in ( + (-0.2, "pcg", "BlockSparsePCG", "tab:blue"), + (0.2, "schur", "GeneralizedSchur", "tab:orange"), + ): + values = [float(row[f"{key}_median_ms"]) for row in data] + bars = ax.bar(positions + offset, values, 0.38, label=label, color=color) + ax.bar_label(bars, labels=[f"{time_ms:.2f}" for time_ms in values], padding=3, fontsize=9) + labels = [] + for row in data: + poses, landmarks, _ = row["scale"].split("_") + labels.append(f"{poses[1:]} / {landmarks[1:]}\n{float(row['speedup']):.2f}× speedup") + ax.set_xticks(positions, labels) + ax.set_xlabel("Poses / nominal landmarks (speedup = PCG time ÷ Schur time)") + ax.set_ylabel("Median wall time (ms, logarithmic scale)") + ax.set_yscale("log") + ax.set_ylim(top=max(float(row[f"{solver_key}_median_ms"]) for row in data for solver_key in ("pcg", "schur")) * 2) + ax.set_title(title) + ax.set_axisbelow(True) + ax.grid(axis="y", which="major", linestyle="--", alpha=0.45) + ax.legend(loc="upper left") + repeats = sorted({row["samples"] for row in rows}) + fig.suptitle( + f"SBA performance: GeneralizedSchur vs BlockSparsePCG\n" + f"{hardware} · {'/'.join(repeats)} repeats · structure reused", + fontsize=15, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.output.with_suffix(".png"), dpi=180) + fig.savefig(args.output.with_suffix(".svg")) + plt.close(fig) + + +if __name__ == "__main__": + main() diff --git a/tests/generalized_schur_solver_test.cpp b/tests/generalized_schur_solver_test.cpp new file mode 100644 index 0000000..3e9a0a1 --- /dev/null +++ b/tests/generalized_schur_solver_test.cpp @@ -0,0 +1,284 @@ +/* + * 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 "cunls/common/cuda_stream.h" +#include "cunls/linear_solver/sparse_linear_solver.h" +#include "cunls/minimizer/problem.h" +#include "cunls/state/vector_state_batch.h" + +namespace cunls { +namespace { +// Exercise existing fixed-size states without adding a public state type. +template +std::unique_ptr MakeVectorState(float *state_data, int dimension) { + if (dimension == Dimension) { + return std::make_unique>(state_data, 1); + } + if constexpr (Dimension < 22) { + return MakeVectorState(state_data, dimension); + } + throw std::invalid_argument("Unsupported test state dimension"); +} + +// Construct a block graph with dense blocks, then make it strictly diagonally +// dominant SPD. A known solution and double host residual provide an oracle +// independent of the Schur implementation and the competing GPU solvers. +void CheckGraphSolution(const std::vector &state_dimensions, int tile_size, int graph_shape, + bool use_state_metadata = true) { + int num_blocks = state_dimensions.size(), + num_rows = std::accumulate(state_dimensions.begin(), state_dimensions.end(), 0); + std::vector state_offsets{0}, row_state_ids; + for (int state_dimension : state_dimensions) { + state_offsets.push_back(state_offsets.back() + state_dimension); + } + for (int block_id = 0; block_id < num_blocks; ++block_id) { + for (int component = 0; component < state_dimensions[block_id]; ++component) { + row_state_ids.push_back(block_id); + } + } + std::vector dense_matrix(num_rows * num_rows, 0), expected_solution(num_rows), + host_rhs(num_rows, 0); + std::mt19937 random_generator(87 + graph_shape); + std::uniform_real_distribution random_value(-.2, .2); + for (int row = 0; row < num_rows; ++row) { + expected_solution[row] = random_value(random_generator); + } + for (int row = 0; row < num_rows; ++row) { + for (int column = row + 1; column < num_rows; ++column) { + int row_state = row_state_ids[row], column_state = row_state_ids[column]; + bool edge = row_state == column_state || + (graph_shape == 0 && abs(row_state - column_state) == 1) || + (graph_shape == 1 && (row_state < 2 || column_state < 2)) || graph_shape == 2; + if (edge) { + dense_matrix[row * num_rows + column] = dense_matrix[column * num_rows + row] = + random_value(random_generator); + } + } + } + for (int row = 0; row < num_rows; ++row) { + double sum = 1; + for (int column = 0; column < num_rows; ++column) { + sum += fabs(dense_matrix[row * num_rows + column]); + } + dense_matrix[row * num_rows + row] = sum; + } + if (graph_shape == 4) { + int unobserved = num_rows / 2; + for (int column = 0; column < num_rows; ++column) { + dense_matrix[unobserved * num_rows + column] = dense_matrix[column * num_rows + unobserved] = + 0.f; + } + expected_solution[unobserved] = 0.f; + } + for (int row = 0; row < num_rows; ++row) { + double sum = 0; + for (int column = 0; column < num_rows; ++column) { + sum += (double)dense_matrix[row * num_rows + column] * expected_solution[column]; + } + host_rhs[row] = sum; + } + std::vector row_offsets{0}, column_ids; + std::vector matrix_values; + for (int row = 0; row < num_rows; row += tile_size) { + for (int column = 0; column < num_rows; column += tile_size) { + bool nonzero = false; + for (int tile_row = 0; tile_row < tile_size; ++tile_row) { + for (int tile_column = 0; tile_column < tile_size; ++tile_column) { + nonzero |= dense_matrix[(row + tile_row) * num_rows + column + tile_column] != 0; + } + } + if (tile_size == 1 && row_state_ids[row] == row_state_ids[column]) { + nonzero = true; + } + if (nonzero) { + column_ids.push_back(column / tile_size); + for (int tile_row = 0; tile_row < tile_size; ++tile_row) { + for (int tile_column = 0; tile_column < tile_size; ++tile_column) { + matrix_values.push_back( + dense_matrix[(row + tile_row) * num_rows + column + tile_column]); + } + } + } + } + row_offsets.push_back(column_ids.size()); + } + CudaStream stream; + auto cuda_stream = stream.GetStream(); + Problem problem; + dvector state_data(num_rows), rhs(host_rhs), solution(num_rows); + std::vector> states; + if (use_state_metadata) { + for (int row = 0; row < num_blocks; ++row) { + states.push_back( + MakeVectorState(state_data.data() + state_offsets[row], state_dimensions[row])); + states.back()->SetNumActiveStates(1); + problem.AddStateBatch(states.back().get()); + } + } + auto solver = CreateSparseLinearSolver(SparseLinearSolverType::GeneralizedSchur, {}); + auto verify = [&]() { + std::vector actual_solution(num_rows); + solution.CopyToHost(actual_solution.data(), num_rows); + double residual_squared_norm = 0, rhs_squared_norm = 0; + for (int row = 0; row < num_rows; ++row) { + double residual = host_rhs[row]; + for (int column = 0; column < num_rows; ++column) { + residual -= (double)dense_matrix[row * num_rows + column] * actual_solution[column]; + } + residual_squared_norm += residual * residual; + rhs_squared_norm += (double)host_rhs[row] * host_rhs[row]; + EXPECT_NEAR(actual_solution[row], expected_solution[row], 5e-4); + } + EXPECT_LT(sqrt(residual_squared_norm / rhs_squared_norm), 1.05e-3); + }; + if (tile_size == 1) { + CSRSparseMatrix sparse_matrix; + sparse_matrix.row_offsets = dvector(row_offsets); + sparse_matrix.col_ids = dvector(column_ids); + sparse_matrix.values = dvector(matrix_values); + ASSERT_TRUE(solver->Initialize(cuda_stream, problem, sparse_matrix, rhs, solution)); + ASSERT_TRUE(solver->Solve(cuda_stream, sparse_matrix, rhs, solution)); + verify(); + // Reuse the symbolic plan for changed values and RHS, then zero RHS. + for (float &value : matrix_values) { + value *= 2; + } + for (float &value : host_rhs) { + value *= 2; + } + for (float &value : dense_matrix) { + value *= 2; + } + sparse_matrix.values.CopyFromHost(matrix_values.data(), matrix_values.size()); + rhs.CopyFromHost(host_rhs.data(), num_rows); + ASSERT_TRUE(solver->Solve(cuda_stream, sparse_matrix, rhs, solution)); + verify(); + std::fill(host_rhs.begin(), host_rhs.end(), 0); + rhs.CopyFromHost(host_rhs.data(), num_rows); + ASSERT_TRUE(solver->Solve(cuda_stream, sparse_matrix, rhs, solution)); + } else { + BSRSparseMatrix sparse_matrix; + sparse_matrix.block_size = tile_size; + sparse_matrix.num_block_rows = num_rows / tile_size; + sparse_matrix.row_offsets = dvector(row_offsets); + sparse_matrix.col_ids = dvector(column_ids); + sparse_matrix.values = dvector(matrix_values); + ASSERT_TRUE(solver->Initialize(cuda_stream, problem, sparse_matrix, rhs, solution)); + ASSERT_TRUE(solver->Solve(cuda_stream, sparse_matrix, rhs, solution)); + verify(); + } +} + +TEST(GeneralizedSchur, MixedDimensionsAndGeneralGraphs) { + for (int graph_shape = 0; graph_shape < 5; ++graph_shape) { + CheckGraphSolution({1, 3, 6, 7, 15, 16, 19, 2, 4}, 1, graph_shape); + } +} + +TEST(GeneralizedSchur, BlockStorage) { + for (int graph_shape = 0; graph_shape < 5; ++graph_shape) { + CheckGraphSolution({6, 6, 3, 3, 3, 3, 3, 3}, 3, graph_shape); + } +} + +TEST(GeneralizedSchur, RawMatrices) { + CheckGraphSolution({1, 1, 1, 1, 1, 1, 1}, 1, 0, false); + CheckGraphSolution({3, 3, 3, 3}, 3, 1, false); +} + +TEST(GeneralizedSchur, OnlyLargeBlocks) { CheckGraphSolution({19, 20, 21, 22}, 1, 0); } + +TEST(GeneralizedSchur, SingleState) { CheckGraphSolution({7}, 1, 3); } + +TEST(GeneralizedSchur, WideConnectedGraph) { + // A small pivot connected to more than 64 scalar columns exercises the + // generic clique path; rows wider than 128 use the cooperative residual. + std::vector state_dimensions(12, 16); + state_dimensions.push_back(3); + CheckGraphSolution(state_dimensions, 1, 2); +} + +TEST(GeneralizedSchur, LargeImplicitChainAndReinitialization) { + CudaStream stream; + auto solver = CreateSparseLinearSolver(SparseLinearSolverType::GeneralizedSchur, {}); + for (int num_rows : {10000, 31, 12000}) { + std::vector row_offsets{0}, column_ids; + std::vector values, expected_solution(num_rows), host_rhs(num_rows); + for (int row = 0; row < num_rows; ++row) { + expected_solution[row] = std::sin(.017 * row); + } + for (int row = 0; row < num_rows; ++row) { + if (row) { + column_ids.push_back(row - 1); + values.push_back(-1); + } + column_ids.push_back(row); + values.push_back(4); + if (row + 1 < num_rows) { + column_ids.push_back(row + 1); + values.push_back(-1); + } + row_offsets.push_back(column_ids.size()); + host_rhs[row] = 4 * expected_solution[row] - (row ? expected_solution[row - 1] : 0) - + (row + 1 < num_rows ? expected_solution[row + 1] : 0); + } + CSRSparseMatrix matrix; + matrix.row_offsets = dvector(row_offsets); + matrix.col_ids = dvector(column_ids); + matrix.values = dvector(values); + dvector rhs(host_rhs), solution(num_rows); + Problem problem; + ASSERT_TRUE(solver->Initialize(stream.GetStream(), problem, matrix, rhs, solution)); + ASSERT_TRUE(solver->Solve(stream.GetStream(), matrix, rhs, solution)); + std::vector actual_solution(num_rows); + solution.CopyToHost(actual_solution.data(), num_rows); + double residual_squared_norm = 0, rhs_squared_norm = 0; + for (int row = 0; row < num_rows; ++row) { + double residual = host_rhs[row] - 4 * actual_solution[row] + + (row ? actual_solution[row - 1] : 0) + + (row + 1 < num_rows ? actual_solution[row + 1] : 0); + residual_squared_norm += residual * residual; + rhs_squared_norm += double(host_rhs[row]) * host_rhs[row]; + } + EXPECT_LT(std::sqrt(residual_squared_norm / rhs_squared_norm), 1e-3); + } +} + +TEST(GeneralizedSchur, RejectsIndefinitePivot) { + CSRSparseMatrix matrix; + matrix.row_offsets = dvector(std::vector{0, 1}); + matrix.col_ids = dvector(std::vector{0}); + matrix.values = dvector(std::vector{-1}); + dvector host_rhs(std::vector{1}), solution(1); + CudaStream stream; + Problem problem; + auto solver = CreateSparseLinearSolver(SparseLinearSolverType::GeneralizedSchur, {}); + ASSERT_TRUE(solver->Initialize(stream.GetStream(), problem, matrix, host_rhs, solution)); + EXPECT_FALSE(solver->Solve(stream.GetStream(), matrix, host_rhs, solution)); +} + +} // namespace + +} // namespace cunls diff --git a/tests/sba_minimizer_test.cpp b/tests/sba_minimizer_test.cpp index d562a05..e9faec4 100644 --- a/tests/sba_minimizer_test.cpp +++ b/tests/sba_minimizer_test.cpp @@ -336,6 +336,9 @@ class SbaMinimizerTestFixture : public ::testing::Test { * about 0.08 (it was 0.12 to 0.15 when LM gave up after a few rejected steps). */ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { + if (std::getenv("CUNLS_BENCHMARK")) { + SetLoggerOptions(Verbosity::Message); + } std::string path = std::string(kTestDataDir) + "/filtered_ba_problems.bin"; std::ifstream in(path, std::ios::binary); if (!in.is_open()) { @@ -390,6 +393,9 @@ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { } SCOPED_TRACE("Problem index " + std::to_string(current_index)); + if (std::getenv("CUNLS_BENCHMARK")) { + std::printf("RECORDED_SBA index=%d\n", current_index); + } Problem problem; BuildProblemFromHost(host, &problem); diff --git a/tests/synthetic_sba_test.cpp b/tests/synthetic_sba_test.cpp index 7e408b8..d69cf7f 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 @@ -249,6 +251,9 @@ class SyntheticSbaTest : public ::testing::TestWithParam { * fixture. */ TEST_P(SyntheticSbaTest, Optimize) { + if (std::getenv("CUNLS_BENCHMARK")) { + SetLoggerOptions(Verbosity::Message); + } auto params = GetParam(); SCOPED_TRACE(std::string("SyntheticSbaTest: ") + params.label); @@ -261,6 +266,10 @@ TEST_P(SyntheticSbaTest, Optimize) { const size_t n_obs = observations.size(); ASSERT_GT(n_obs, 0u); + if (std::getenv("CUNLS_BENCHMARK")) { + std::printf("SBA_SHAPE scale=%s poses=%zu landmarks=%zu observations=%zu\n", params.label, + init_poses.size(), init_points.size(), n_obs); + } // Identity 2x2 sqrt information per observation. std::vector> sqrt_info(n_obs); @@ -320,6 +329,17 @@ TEST_P(SyntheticSbaTest, Optimize) { options.state_tolerance = 1e-10f; options.cost_tolerance = 1e-6f; options.disable_safety_checks = true; + if (std::getenv("CUNLS_BENCHMARK")) { + options.reuse_structure = true; + int device; + cudaDeviceProp properties{}; + THROW_ON_CUDA_ERROR(cudaGetDevice(&device)); + THROW_ON_CUDA_ERROR(cudaGetDeviceProperties(&properties, device)); + std::string name(properties.name); + std::replace(name.begin(), name.end(), ' ', '_'); + std::printf("SBA_DEVICE name=%s sm=%d%d pci_bus=%d pci_device=%d\n", name.c_str(), + properties.major, properties.minor, properties.pciBusID, properties.pciDeviceID); + } options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); @@ -330,8 +350,86 @@ TEST_P(SyntheticSbaTest, Optimize) { LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; + if (const char *repeat_env = std::getenv("CUNLS_BENCHMARK_REPEATS")) { + int repeats = std::max(1, std::atoi(repeat_env)); + internal::MinimizerSystem system; + internal::MinimizerScratch scratch; + StateBatchOps state_ops; + system.solver = CreateSparseLinearSolver(options.sparse_linear_solver_type, + options.sparse_linear_solver_config); + internal::SetUpStructure(stream.GetStream(), problem, options, state_ops, system); + MinimizerState state(stream.GetStream(), problem); + internal::BuildSystem(stream.GetStream(), problem, state, options, system, scratch); + dvector diagonal; + system.normal_equations.ExtractHessianDiagonal(stream.GetStream(), diagonal); + system.normal_equations.AddScaledDiagonalToLhs(stream.GetStream(), 1e-3f, diagonal); + system.step.resize(system.rhs.size()); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + size_t free_before, total_memory; + THROW_ON_CUDA_ERROR(cudaMemGetInfo(&free_before, &total_memory)); + auto init_start = std::chrono::steady_clock::now(); + ASSERT_TRUE(system.normal_equations.InitializeSolver(stream.GetStream(), *system.solver, + problem, system.rhs, system.step)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + double init_ms = + std::chrono::duration(std::chrono::steady_clock::now() - init_start) + .count(); + std::printf("SBA_SYMBOLIC scale=%s solver=%s ms=%.6f\n", params.label, + std::getenv("CUNLS_SOLVER"), init_ms); + SetLoggerOptions(Verbosity::Error); + for (int repeat = -1; repeat < repeats; ++repeat) { + auto start = std::chrono::steady_clock::now(); + ASSERT_TRUE(system.normal_equations.Solve(stream.GetStream(), *system.solver, system.rhs, + system.step)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + double ms = + std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); + if (repeat >= 0) { + std::printf("SBA_LINEAR scale=%s solver=%s repeat=%d ms=%.6f\n", params.label, + std::getenv("CUNLS_SOLVER"), repeat, ms); + } + } + size_t free_after; + THROW_ON_CUDA_ERROR(cudaMemGetInfo(&free_after, &total_memory)); + std::printf("SBA_MEMORY scale=%s solver=%s additional_mib=%.3f\n", params.label, + std::getenv("CUNLS_SOLVER"), + double(free_before > free_after ? free_before - free_after : 0) / (1024 * 1024)); + // Read the original working matrix, not Schur's recursive residual, and + // evaluate b-Ax independently in float64 outside the timed region. + const auto &working_matrix = system.normal_equations.LhsBSR(); + ASSERT_TRUE(system.normal_equations.UsesBlockStorage()); + int tile_size = working_matrix.block_size, num_rows = working_matrix.NumRows(); + std::vector row_offsets(working_matrix.row_offsets.size()), + column_ids(working_matrix.col_ids.size()); + std::vector values(working_matrix.values.size()), rhs(num_rows), step(num_rows); + working_matrix.row_offsets.CopyToHost(row_offsets.data(), row_offsets.size()); + working_matrix.col_ids.CopyToHost(column_ids.data(), column_ids.size()); + working_matrix.values.CopyToHost(values.data(), values.size()); + system.rhs.CopyToHost(rhs.data(), num_rows); + system.step.CopyToHost(step.data(), num_rows); + double residual_squared_norm = 0, rhs_squared_norm = 0; + for (int row = 0; row < num_rows; ++row) { + double residual = rhs[row]; + for (int entry = row_offsets[row / tile_size]; entry < row_offsets[row / tile_size + 1]; + ++entry) { + for (int tile_column = 0; tile_column < tile_size; ++tile_column) { + residual -= double(values[size_t(entry) * tile_size * tile_size + + (row % tile_size) * tile_size + tile_column]) * + step[column_ids[entry] * tile_size + tile_column]; + } + } + residual_squared_norm += residual * residual; + rhs_squared_norm += double(rhs[row]) * rhs[row]; + } + std::printf("SBA_RESIDUAL scale=%s solver=%s relative=%.9g\n", params.label, + std::getenv("CUNLS_SOLVER"), std::sqrt(residual_squared_norm / rhs_squared_norm)); + EXPECT_LE(std::sqrt(residual_squared_norm / rhs_squared_norm), 1.05e-3); + } LevenbergMarquardtMinimizer minimizer(lm_options); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + const auto solve_start = std::chrono::steady_clock::now(); MinimizerSummary summary; { auto range = profiler_domain_.CreateDomainRange("Minimize"); @@ -339,6 +437,36 @@ TEST_P(SyntheticSbaTest, Optimize) { THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); } + const double solve_ms = + std::chrono::duration(std::chrono::steady_clock::now() - solve_start) + .count(); + if (std::getenv("CUNLS_BENCHMARK")) { + std::printf("SBA_BENCH scale=%s solver=%s ms=%.6f iterations=%zu initial=%.9g final=%.9g\n", + params.label, + std::getenv("CUNLS_SOLVER") ? std::getenv("CUNLS_SOLVER") : "BlockSparsePCG", + solve_ms, summary.num_iterations, summary.initial_cost, summary.final_cost); + } + + if (const char *repeat_env = std::getenv("CUNLS_BENCHMARK_REPEATS")) { + int repeats = std::max(1, std::atoi(repeat_env)); + SetLoggerOptions(Verbosity::Error); + for (int repeat = 0; repeat < repeats; ++repeat) { + poses_d.CopyFromHost(init_poses.data(), init_poses.size()); + points_d.CopyFromHost(init_points.data(), init_points.size()); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + auto start = std::chrono::steady_clock::now(); + summary = minimizer.Minimize(stream.GetStream(), problem); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + double ms = + std::chrono::duration(std::chrono::steady_clock::now() - start) + .count(); + std::printf( + "SBA_WARM scale=%s solver=%s repeat=%d ms=%.6f iterations=%zu initial=%.9g final=%.9g\n", + params.label, std::getenv("CUNLS_SOLVER"), repeat, ms, summary.num_iterations, + summary.initial_cost, summary.final_cost); + } + } + EXPECT_TRUE(std::isfinite(summary.initial_cost)); EXPECT_TRUE(std::isfinite(summary.final_cost)); // Observation noise sigma = 1e-3 on both coordinates: at the optimum the cost @@ -348,9 +476,11 @@ TEST_P(SyntheticSbaTest, Optimize) { std::vector final_poses(params.n_poses); poses_d.CopyToHost(final_poses.data(), final_poses.size()); float max_err = 0.f; - for (int i = 0; i < params.n_poses; ++i) - for (int k : {3, 7, 11}) + for (int i = 0; i < params.n_poses; ++i) { + for (int k : {3, 7, 11}) { max_err = std::max(max_err, std::fabs(final_poses[i][k] - gt_poses[i][k])); + } + } EXPECT_LT(max_err, 0.05f); // ~1.5 cm after the 10 iterations } diff --git a/tests/utils.h b/tests/utils.h index 6978647..a604df4 100644 --- a/tests/utils.h +++ b/tests/utils.h @@ -377,16 +377,15 @@ 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". + * @brief Reads the test-only CUNLS_SOLVER selector. Defaults to BlockSparsePCG. + * Recognized alternatives: cuDSS and GeneralizedSchur. */ inline SparseLinearSolverType SolverTypeFromEnv() { - const char *s = std::getenv("CUNLS_SOLVER"); - if (s != nullptr && std::strcmp(s, "cuDSS") == 0) { + const char *solver_name = std::getenv("CUNLS_SOLVER"); + if (solver_name != nullptr && std::strcmp(solver_name, "GeneralizedSchur") == 0) { + return SparseLinearSolverType::GeneralizedSchur; + } + if (solver_name != nullptr && std::strcmp(solver_name, "cuDSS") == 0) { return SparseLinearSolverType::cuDSS; } return SparseLinearSolverType::BlockSparsePCG; From a979ae92c5c8e8979e032285761fbb6d4406e21d Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 14 Oct 2026 21:21:40 -0700 Subject: [PATCH 2/2] Remove useless files --- scripts/benchmark_schur_in_docker.sh | 101 --------------------------- scripts/plot_schur_benchmark.py | 69 ------------------ 2 files changed, 170 deletions(-) delete mode 100755 scripts/benchmark_schur_in_docker.sh delete mode 100644 scripts/plot_schur_benchmark.py diff --git a/scripts/benchmark_schur_in_docker.sh b/scripts/benchmark_schur_in_docker.sh deleted file mode 100755 index 417eafc..0000000 --- a/scripts/benchmark_schur_in_docker.sh +++ /dev/null @@ -1,101 +0,0 @@ -#!/bin/bash -# Reproduce the Schur/PCG SBA measurements using an existing Release build. -set -euo pipefail -if [ "$#" -lt 1 ] || [ "$#" -gt 2 ]; then - echo "Usage: $0 [repetitions=20]" - exit 1 -fi -OUTPUT_DIR=$(realpath "$1") -REPEATS=${2:-20} -if ! [[ "$REPEATS" =~ ^[1-9][0-9]*$ ]]; then - echo "Repetitions must be a positive integer." - exit 1 -fi -if [ ! -x "$OUTPUT_DIR/build_shared/bin/nls_tests" ]; then - echo "Missing test binary. Run EXTRA_CMAKE_ARGS=-DBUILD_TESTING=ON ./scripts/build_cunls_in_docker.sh Release $OUTPUT_DIR" - exit 1 -fi -TTY_FLAG="" -[ -t 0 ] && TTY_FLAG="-it" -docker run --runtime=nvidia --gpus all --rm $TTY_FLAG \ - -v "$(pwd):/cunls:ro" -v "$OUTPUT_DIR:/cunls_install" \ - -e CUDA_DEVICE_ORDER=PCI_BUS_ID \ - -e CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" \ - -e CUNLS_BENCHMARK=1 -e CUNLS_BENCHMARK_REPEATS="$REPEATS" \ - -e CUNLS_PCG_TOL=0.001 -e CUNLS_PCG_MAX_ITER=200 \ - cunls:local bash -c ' - set -euo pipefail - nvidia-smi - mkdir -p /cunls_install/schur-benchmarks - for solver in BlockSparsePCG GeneralizedSchur; do - CUNLS_SOLVER=$solver /cunls_install/build_shared/bin/nls_tests \ - --gtest_filter="*SyntheticSba*" \ - --gtest_output="xml:/cunls_install/schur-benchmarks/$solver.xml" \ - | tee "/cunls_install/schur-benchmarks/$solver.log" - done - ' - -# Standard-library-only summary; preserve raw measurements beside the CSV. -docker run --rm -i --entrypoint python3 \ - -v "$OUTPUT_DIR:/cunls_install" cunls:local - /cunls_install/schur-benchmarks <<'PY' -import csv -import math -from pathlib import Path -import statistics -import sys - -root = Path(sys.argv[1]) -records = {} -for solver in ("BlockSparsePCG", "GeneralizedSchur"): - rows = [] - for line in (root / f"{solver}.log").read_text().splitlines(): - if line.startswith("SBA_"): - parts = line.split() - rows.append((parts[0], dict(part.split("=", 1) for part in parts[1:]))) - records[solver] = rows -scales = [record["scale"] for kind, record in records["BlockSparsePCG"] if kind == "SBA_SHAPE"] -assert scales, "No benchmark records found" -devices = [] -for solver, rows in records.items(): - found = [record for kind, record in rows if kind == "SBA_DEVICE"] - assert found and all(record == found[0] for record in found), (solver, "Missing or inconsistent GPU identity") - devices.append(found[0]) -assert devices[0] == devices[1], "Backends ran on different GPUs" -device = devices[0] -print(f"CUDA device: {device['name']} (SM {device['sm']}, PCI bus {device['pci_bus']})") -summary = [] -for scale in scales: - for stage in ("SBA_LINEAR", "SBA_WARM"): - stats = {} - for solver, rows in records.items(): - samples = sorted(float(record["ms"]) for kind, record in rows - if kind == stage and record["scale"] == scale) - assert samples, (solver, scale, stage) - stats[solver] = (statistics.median(samples), samples[math.ceil(.95 * len(samples))-1], len(samples)) - pcg, schur = stats["BlockSparsePCG"], stats["GeneralizedSchur"] - summary.append(dict(scale=scale, stage=stage, samples=pcg[2], - device=device["name"], compute_capability=device["sm"], - pcg_median_ms=pcg[0], schur_median_ms=schur[0], - pcg_p95_ms=pcg[1], schur_p95_ms=schur[1], speedup=pcg[0]/schur[0])) - assert pcg[2] == schur[2] - final_costs = {} - shapes = [] - for solver, rows in records.items(): - shape = next(record for kind, record in rows if kind == "SBA_SHAPE" and record["scale"] == scale) - shapes.append(shape) - residual = next(float(record["relative"]) for kind, record in rows if kind == "SBA_RESIDUAL" and record["scale"] == scale) - assert math.isfinite(residual) and residual <= 1.05e-3, (solver, scale, residual) - final_costs[solver] = statistics.median(float(record["final"]) for kind, record in rows - if kind == "SBA_WARM" and record["scale"] == scale) - assert shapes[0] == shapes[1], "Solvers must use identical input sizes" - pcg_cost, sc_cost = final_costs["BlockSparsePCG"], final_costs["GeneralizedSchur"] - assert abs(pcg_cost-sc_cost) <= max(1e-7, .001*abs(pcg_cost)), (scale, final_costs) -with (root / "summary.csv").open("w", newline="") as f: - writer = csv.DictWriter(f, fieldnames=summary[0].keys()) - writer.writeheader() - writer.writerows(summary) -for row in summary: - print(f"{row['scale']:20} {row['stage']:11} PCG {row['pcg_median_ms']:9.3f} ms " - f"Schur {row['schur_median_ms']:9.3f} ms {row['speedup']:5.2f}x") -print(f"Raw logs, JUnit results, and summary: {root}") -PY diff --git a/scripts/plot_schur_benchmark.py b/scripts/plot_schur_benchmark.py deleted file mode 100644 index 865dd0e..0000000 --- a/scripts/plot_schur_benchmark.py +++ /dev/null @@ -1,69 +0,0 @@ -#!/usr/bin/env python3 -"""Plot benchmark_schur_in_docker.sh's CSV; requires matplotlib and numpy.""" - -import argparse -import csv -from pathlib import Path - -import matplotlib - -matplotlib.use("Agg") -import matplotlib.pyplot as plt -import numpy as np - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("csv", type=Path) - parser.add_argument("output", type=Path, help="Output filename stem (PNG and SVG)") - parser.add_argument("--hardware", help="Override the GPU name recorded in the CSV") - args = parser.parse_args() - with args.csv.open() as source: - rows = list(csv.DictReader(source)) - - hardware = args.hardware or rows[0].get("device", "GPU not recorded").replace("_", " ") - plt.style.use("default") - fig, axes = plt.subplots(2, 1, figsize=(11, 9), layout="constrained") - for ax, stage, title in zip( - axes, - ("SBA_WARM", "SBA_LINEAR"), - ("Complete warm SBA optimization", "Frozen initial damped linear solve"), - ): - data = [row for row in rows if row["stage"] == stage] - if not data: - raise ValueError(f"Missing measurements for {stage}") - positions = np.arange(len(data)) - for offset, key, label, color in ( - (-0.2, "pcg", "BlockSparsePCG", "tab:blue"), - (0.2, "schur", "GeneralizedSchur", "tab:orange"), - ): - values = [float(row[f"{key}_median_ms"]) for row in data] - bars = ax.bar(positions + offset, values, 0.38, label=label, color=color) - ax.bar_label(bars, labels=[f"{time_ms:.2f}" for time_ms in values], padding=3, fontsize=9) - labels = [] - for row in data: - poses, landmarks, _ = row["scale"].split("_") - labels.append(f"{poses[1:]} / {landmarks[1:]}\n{float(row['speedup']):.2f}× speedup") - ax.set_xticks(positions, labels) - ax.set_xlabel("Poses / nominal landmarks (speedup = PCG time ÷ Schur time)") - ax.set_ylabel("Median wall time (ms, logarithmic scale)") - ax.set_yscale("log") - ax.set_ylim(top=max(float(row[f"{solver_key}_median_ms"]) for row in data for solver_key in ("pcg", "schur")) * 2) - ax.set_title(title) - ax.set_axisbelow(True) - ax.grid(axis="y", which="major", linestyle="--", alpha=0.45) - ax.legend(loc="upper left") - repeats = sorted({row["samples"] for row in rows}) - fig.suptitle( - f"SBA performance: GeneralizedSchur vs BlockSparsePCG\n" - f"{hardware} · {'/'.join(repeats)} repeats · structure reused", - fontsize=15, - ) - args.output.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(args.output.with_suffix(".png"), dpi=180) - fig.savefig(args.output.with_suffix(".svg")) - plt.close(fig) - - -if __name__ == "__main__": - main()