From e9a2a108ad37b425535404dd730d0baac1fc1477 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 13 May 2026 09:41:55 -0700 Subject: [PATCH 1/4] PCG --- cunls/linear_solver/CMakeLists.txt | 1 + .../linear_solver/block_sparse_pcg_solver.cu | 676 ++++++++++++++++++ cunls/linear_solver/block_sparse_pcg_solver.h | 116 +++ cunls/linear_solver/llms.txt | 12 + cunls/linear_solver/sparse_linear_solver.cpp | 3 + cunls/linear_solver/sparse_linear_solver.h | 25 +- cunls/minimizer/gauss_newton_minimizer.cu | 8 +- tests/pgo_minimizer_test.cpp | 6 + tests/sba_minimizer_test.cpp | 10 + tests/sparse_linear_solver_test.cpp | 118 +++ tests/synthetic_pgo_test.cpp | 21 + tests/utils.h | 51 ++ 12 files changed, 1037 insertions(+), 10 deletions(-) create mode 100644 cunls/linear_solver/block_sparse_pcg_solver.cu create mode 100644 cunls/linear_solver/block_sparse_pcg_solver.h diff --git a/cunls/linear_solver/CMakeLists.txt b/cunls/linear_solver/CMakeLists.txt index eb9c008..003ee3f 100644 --- a/cunls/linear_solver/CMakeLists.txt +++ b/cunls/linear_solver/CMakeLists.txt @@ -1,4 +1,5 @@ add_library(cunls_linear_solver OBJECT + block_sparse_pcg_solver.cu cudss_sparse_linear_solver.cpp dense_cholesky_solver.cu dense_linear_solver.cu diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu new file mode 100644 index 0000000..f412249 --- /dev/null +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -0,0 +1,676 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cunls/linear_solver/block_sparse_pcg_solver.h" + +#include +#include + +#include +#include +#include + +#include "cunls/common/cusparse_helper.h" +#include "cunls/common/helper.h" +#include "cunls/common/log.h" + +namespace cunls { + +namespace { + +constexpr int kMaxBlockSize = 16; + +// ----------------------------------------------------------------------------- +// Device kernels +// ----------------------------------------------------------------------------- + +/** + * @brief Extracts the dense B x B diagonal tiles from a symmetric CSR matrix + * and stores their LDLT factors. + * + * One thread block handles one diagonal block. Threads cooperate to scan the + * rows in [block_row * B, block_row * B + B), pull the entries with column id + * in the same range, and assemble a dense B x B tile in shared memory. A + * single-thread LDLT then runs over the symmetric tile (sizes encountered here + * are 3-6, so the serial cost is negligible compared to the global I/O). + * + * The factored block is written back to ``factors`` in row-major layout: + * the strict lower-triangle holds L (1's on the diagonal implicitly), and the + * diagonal of the stored tile holds D. Upper triangle is undefined and not + * read. + */ +template +__global__ void ExtractAndFactorBlockDiagonals(const int *__restrict__ row_off, + const int *__restrict__ col_idx, + const float *__restrict__ values, + int num_blocks, float pivot_floor, + float *__restrict__ factors) { + int block_row = blockIdx.x; + if (block_row >= num_blocks) { + return; + } + __shared__ float tile[B * B]; + + int tid = threadIdx.x; + for (int i = tid; i < B * B; i += blockDim.x) { + tile[i] = 0.f; + } + __syncthreads(); + + // Each thread takes one of the B rows. + if (tid < B) { + int global_row = block_row * B + tid; + int start = row_off[global_row]; + int end = row_off[global_row + 1]; + int col_lo = block_row * B; + int col_hi = col_lo + B; + for (int k = start; k < end; ++k) { + int c = col_idx[k]; + if (c >= col_lo && c < col_hi) { + tile[tid * B + (c - col_lo)] = values[k]; + } + } + } + __syncthreads(); + + // Symmetrize so the LDLT below can read either triangle. CSR is + // symmetric for J^T J and Levenberg-Marquardt damping is on the diagonal, + // so off-diagonal entries should already match; symmetrizing keeps the + // factorization stable when only the lower triangle is stored. + if (tid < B) { + for (int j = tid + 1; j < B; ++j) { + float a = tile[tid * B + j]; + float b = tile[j * B + tid]; + float s = 0.5f * (a + b); + tile[tid * B + j] = s; + tile[j * B + tid] = s; + } + } + __syncthreads(); + + // Serial LDLT in shared memory. B is small (3..16); a single thread is the + // simplest correct implementation. Outer-product update with diagonal pivot + // and a small floor to keep the preconditioner stable on near-singular tiles. + if (tid == 0) { + for (int k = 0; k < B; ++k) { + float d = tile[k * B + k]; + if (fabsf(d) < pivot_floor) { + d = pivot_floor; + } + tile[k * B + k] = d; + float inv_d = 1.f / d; + for (int i = k + 1; i < B; ++i) { + float lik = tile[i * B + k] * inv_d; + tile[i * B + k] = lik; + for (int j = k + 1; j <= i; ++j) { + tile[i * B + j] -= lik * tile[j * B + k] * d; + } + } + } + } + __syncthreads(); + + // Write factors back (row-major). + float *out = factors + block_row * B * B; + for (int i = tid; i < B * B; i += blockDim.x) { + out[i] = tile[i]; + } +} + +/** + * @brief Applies the precomputed block-LDLT preconditioner: z = M^{-1} r. + * + * One block per diagonal tile, B threads per block. Reads r into registers, + * solves L y = r, D w = y, L^T z = w in three serial sweeps. Sizes are tiny + * (B in {3,6}) so the work fits in registers; the kernel is bandwidth-bound on + * the factor read. + */ +template +__global__ void ApplyBlockJacobiPreconditioner( + const float *__restrict__ factors, const float *__restrict__ r, + float *__restrict__ z, int num_blocks) { + int block_row = blockIdx.x; + if (block_row >= num_blocks) { + return; + } + __shared__ float L[B * B]; + __shared__ float v[B]; + + int tid = threadIdx.x; + if (tid < B) { + v[tid] = r[block_row * B + tid]; + } + for (int i = tid; i < B * B; i += blockDim.x) { + L[i] = factors[block_row * B * B + i]; + } + __syncthreads(); + + // Forward solve L y = r (unit diagonal); D w = y; L^T z = w. Serialized on + // thread 0 — the per-block work is 3*B^2 FMAs which dominates over any + // attempt to parallelize across B threads. + if (tid == 0) { + for (int i = 1; i < B; ++i) { + float s = v[i]; +#pragma unroll + for (int j = 0; j < B; ++j) { + if (j < i) { + s -= L[i * B + j] * v[j]; + } + } + v[i] = s; + } + for (int i = 0; i < B; ++i) { + v[i] /= L[i * B + i]; + } + for (int i = B - 2; i >= 0; --i) { + float s = v[i]; +#pragma unroll + for (int j = 0; j < B; ++j) { + if (j > i) { + s -= L[j * B + i] * v[j]; + } + } + v[i] = s; + } + } + __syncthreads(); + + if (tid < B) { + z[block_row * B + tid] = v[tid]; + } +} + +/** Scalar Jacobi fallback (B == 1) — z = r / diag. */ +__global__ void ApplyScalarJacobi(const float *__restrict__ factors, + const float *__restrict__ r, + float *__restrict__ z, int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) { + return; + } + float d = factors[i]; + z[i] = r[i] / d; +} + +/** Scalar Jacobi extractor (B == 1) — pulls diag(H) with a pivot floor. */ +__global__ void ExtractScalarJacobi(const int *__restrict__ row_off, + const int *__restrict__ col_idx, + const float *__restrict__ values, int n, + float pivot_floor, + float *__restrict__ factors) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) { + return; + } + int start = row_off[i]; + int end = row_off[i + 1]; + float d = pivot_floor; + for (int k = start; k < end; ++k) { + if (col_idx[k] == i) { + d = fmaxf(fabsf(values[k]), pivot_floor); + break; + } + } + factors[i] = d; +} + +/** y = a*x + b*z, all length n. Used for PCG vector updates. */ +__global__ void AxpyKernel(float a, const float *__restrict__ x, float b, + const float *__restrict__ z, float *__restrict__ y, + int n) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) { + return; + } + y[i] = a * x[i] + b * z[i]; +} + +/** x = a*p + x, r = r - a*Ap, plus partial-sum dot products of (new_r, z') — + * this kernel is the hot path of PCG and is intentionally minimal. + * ``alpha`` is read once from device memory so that the host can keep + * enqueueing kernels without waiting on the previous dot product. */ +__global__ void PcgUpdateKernel(const float *__restrict__ alpha_ptr, + const float *__restrict__ p, + const float *__restrict__ Ap, + float *__restrict__ x, float *__restrict__ r, + int n) { + __shared__ float a; + if (threadIdx.x == 0) { + a = alpha_ptr[0]; + } + __syncthreads(); + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) { + return; + } + x[i] += a * p[i]; + r[i] -= a * Ap[i]; +} + +/** p = z + beta * p, length n. ``beta`` is fetched from device memory. */ +__global__ void PcgDirectionKernel(const float *__restrict__ beta_ptr, + const float *__restrict__ z, + float *__restrict__ p, int n) { + __shared__ float b; + if (threadIdx.x == 0) { + b = beta_ptr[0]; + } + __syncthreads(); + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) { + return; + } + p[i] = z[i] + b * p[i]; +} + +/** Single-thread device kernel: alpha = rz_old / pAp (with guard). */ +__global__ void ComputeAlphaKernel(const float *__restrict__ rz_old, + const float *__restrict__ pAp, + float *__restrict__ alpha) { + float denom = pAp[0]; + alpha[0] = (denom > 0.f) ? rz_old[0] / denom : 0.f; +} + +/** Single-thread device kernel: beta = rz_new / rz_old; rz_old <- rz_new. */ +__global__ void ComputeBetaKernel(const float *__restrict__ rz_new, + float *__restrict__ rz_old, + float *__restrict__ beta) { + float num = rz_new[0]; + float denom = rz_old[0]; + beta[0] = (fabsf(denom) > 0.f) ? num / denom : 0.f; + rz_old[0] = num; +} + +/** + * @brief Single-pass dot product writing the result to ``out`` on device. + * + * Uses one block-stride loop per CTA, a warp reduction, and an atomicAdd into + * out[0]. Caller must zero out[0] before launch. Fine for the sizes we hit + * here (a few thousand to ~100k floats). + */ +__global__ void DotKernelZero(float *__restrict__ out) { out[0] = 0.f; } + +__global__ void DotKernel(const float *__restrict__ a, + const float *__restrict__ b, int n, + float *__restrict__ out) { + float s = 0.f; + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += blockDim.x * gridDim.x) { + s += a[i] * b[i]; + } + // Warp reduction. + for (int off = warpSize / 2; off > 0; off >>= 1) { + s += __shfl_down_sync(0xffffffffu, s, off); + } + __shared__ float warp_sums[32]; + int lane = threadIdx.x & 31; + int wid = threadIdx.x >> 5; + if (lane == 0) { + warp_sums[wid] = s; + } + __syncthreads(); + if (wid == 0) { + int nw = (blockDim.x + 31) >> 5; + s = (lane < nw) ? warp_sums[lane] : 0.f; + for (int off = warpSize / 2; off > 0; off >>= 1) { + s += __shfl_down_sync(0xffffffffu, s, off); + } + if (lane == 0) { + atomicAdd(out, s); + } + } +} + +// ----------------------------------------------------------------------------- +// Host helpers +// ----------------------------------------------------------------------------- + +void LaunchExtractAndFactor(cudaStream_t stream, int B, + const CSRSparseMatrix &matrix, int num_blocks, + float pivot_floor, dvector &factors) { + if (B == 1) { + int n = num_blocks; + int threads = 256; + int blocks = (n + threads - 1) / threads; + ExtractScalarJacobi<<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), + n, pivot_floor, factors.data()); + return; + } + + int threads = ((B + 31) / 32) * 32; // round up to a full warp + switch (B) { + case 2: + ExtractAndFactorBlockDiagonals<2><<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), + num_blocks, pivot_floor, factors.data()); + break; + case 3: + ExtractAndFactorBlockDiagonals<3><<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), + num_blocks, pivot_floor, factors.data()); + break; + case 6: + ExtractAndFactorBlockDiagonals<6><<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), + num_blocks, pivot_floor, factors.data()); + break; + case 7: + ExtractAndFactorBlockDiagonals<7><<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), + num_blocks, pivot_floor, factors.data()); + break; + default: + throw std::invalid_argument( + "BlockSparsePCGSolver: unsupported block_size (must be 1,2,3,6,7)"); + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +void LaunchApplyPrecond(cudaStream_t stream, int B, const float *factors, + const float *r, float *z, int num_blocks) { + if (B == 1) { + int n = num_blocks; + int threads = 256; + int blocks = (n + threads - 1) / threads; + ApplyScalarJacobi<<>>(factors, r, z, n); + return; + } + int threads = ((B + 31) / 32) * 32; + switch (B) { + case 2: + ApplyBlockJacobiPreconditioner<2> + <<>>(factors, r, z, num_blocks); + break; + case 3: + ApplyBlockJacobiPreconditioner<3> + <<>>(factors, r, z, num_blocks); + break; + case 6: + ApplyBlockJacobiPreconditioner<6> + <<>>(factors, r, z, num_blocks); + break; + case 7: + ApplyBlockJacobiPreconditioner<7> + <<>>(factors, r, z, num_blocks); + break; + default: + throw std::invalid_argument("BlockSparsePCGSolver: unsupported block_size"); + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +void DotAsync(cudaStream_t stream, const float *a, const float *b, int n, + float *d_out) { + DotKernelZero<<<1, 1, 0, stream>>>(d_out); + int threads = 256; + int blocks = std::min(1024, (n + threads - 1) / threads); + DotKernel<<>>(a, b, n, d_out); +} + +} // namespace + +// ----------------------------------------------------------------------------- +// BlockSparsePCGSolver +// ----------------------------------------------------------------------------- + +BlockSparsePCGSolver::BlockSparsePCGSolver(BlockSparsePCGOptions options) + : options_(options) { + if (options_.block_size < 1 || options_.block_size > kMaxBlockSize) { + throw std::invalid_argument( + "BlockSparsePCGSolver: block_size must be in [1, 16]"); + } +} + +BlockSparsePCGSolver::~BlockSparsePCGSolver() = default; + +bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, + const CSRSparseMatrix &spd_matrix, + const dvector &rhs, + dvector &result) { + int n = static_cast(spd_matrix.NumRows()); + if (n != static_cast(rhs.size()) || + n != static_cast(result.size())) { + LogError( + "BlockSparsePCGSolver: dim mismatch (matrix={}, rhs={}, result={})", n, + rhs.size(), result.size()); + return false; + } + int B = options_.block_size; + if (n % B != 0) { + LogError("BlockSparsePCGSolver: matrix size {} not divisible by block " + "size {}", + n, B); + return false; + } + matrix_size_ = n; + num_blocks_ = n / B; + + size_t factors_size = + static_cast(num_blocks_) * static_cast(B * B); + if (B == 1) { + factors_size = static_cast(n); + } + if (precond_factors_.size() < factors_size) { + precond_factors_.resize(factors_size); + } + if (r_.size() < static_cast(n)) { + r_.resize(n); + z_.resize(n); + p_.resize(n); + Ap_.resize(n); + } + if (d_scratch_.size() < 2) { + d_scratch_.resize(2); + } + + // Build the cuSPARSE descriptor for SpMV. Use raw cusparseDnVec with + // explicit size = n so capacity-vs-size drift in p_/Ap_ across successive + // Solves on differently-sized matrices doesn't break the SpMV preprocess. + mat_desc_ = + cuSPARSEMatrixDescription(n, n, static_cast(spd_matrix.NumNonZeros()), + spd_matrix); + auto handle = + static_cast(cusparse_handle_.GetHandle(stream)); + + auto matA = static_cast(mat_desc_.GetDescription()); + cusparseDnVecDescr_t vecX = nullptr; + cusparseDnVecDescr_t vecY = nullptr; + THROW_ON_CUSPARSE_ERROR( + cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR( + cusparseCreateDnVec(&vecY, n, Ap_.data(), CUDA_R_32F)); + + float alpha = 1.f; + float beta = 0.f; + size_t buffer_size = 0; + THROW_ON_CUSPARSE_ERROR(cusparseSpMV_bufferSize( + handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY, + CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, &buffer_size)); + if (buffer_size > spmv_buffer_.size()) { + spmv_buffer_.resize(buffer_size); + } + THROW_ON_CUSPARSE_ERROR(cusparseSpMV_preprocess( + handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &alpha, matA, vecX, &beta, vecY, + CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, spmv_buffer_.data())); + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecX)); + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecY)); + spmv_buffer_ready_ = true; + return true; +} + +bool BlockSparsePCGSolver::Solve(cudaStream_t stream, + const CSRSparseMatrix &spd_matrix, + const dvector &rhs, + dvector &result) { + int n = static_cast(spd_matrix.NumRows()); + if (n != static_cast(rhs.size()) || + n != static_cast(result.size())) { + LogError( + "BlockSparsePCGSolver: dim mismatch (matrix={}, rhs={}, result={})", n, + rhs.size(), result.size()); + return false; + } + if (n == 0) { + return true; + } + if (n != matrix_size_) { + if (!Initialize(stream, spd_matrix, rhs, result)) { + return false; + } + } + int B = options_.block_size; + + // Refresh the cuSPARSE descriptor's value pointer (structure is fixed across + // calls; the CSR matrix can move in memory between Solves). + mat_desc_.UpdatePointers(spd_matrix); + + // Refresh preconditioner from the current matrix values. + LaunchExtractAndFactor(stream, B, spd_matrix, num_blocks_, + options_.pivot_floor, precond_factors_); + + auto handle = + static_cast(cusparse_handle_.GetHandle(stream)); + auto matA = static_cast(mat_desc_.GetDescription()); + + // PCG with zero initial guess. The Gauss-Newton step is reset every outer + // iteration so warm-starting from the prior step would seed PCG far from the + // new solution; zeroing is both faster and more robust here. + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(result.data(), 0, n * sizeof(float), stream)); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(r_.data(), rhs.data(), n * sizeof(float), + cudaMemcpyDeviceToDevice, stream)); + + // z = M^{-1} r + LaunchApplyPrecond(stream, B, precond_factors_.data(), r_.data(), z_.data(), + num_blocks_); + // p = z + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(p_.data(), z_.data(), n * sizeof(float), + cudaMemcpyDeviceToDevice, stream)); + + // Scratch layout (all device-resident, never copied to host during the inner + // loop): [0] alpha, [1] beta, [2] pAp, [3] rz_old, [4] rz_new, [5] r_norm2, + // [6] b_norm2. + enum : int { + kAlpha = 0, + kBeta = 1, + kPAp = 2, + kRzOld = 3, + kRzNew = 4, + kRnorm2 = 5, + kBnorm2 = 6, + kScalarCount = 7 + }; + if (d_scratch_.size() < kScalarCount) { + d_scratch_.resize(kScalarCount); + } + + // rz_old = ; b_norm2 = + DotAsync(stream, r_.data(), z_.data(), n, d_scratch_.data() + kRzOld); + DotAsync(stream, rhs.data(), rhs.data(), n, d_scratch_.data() + kBnorm2); + + // Pull b_norm2 once for the host-side stopping threshold (the dot still + // runs async; the sync below is paid only once, not per iteration). + float b_norm2 = 0.f; + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&b_norm2, d_scratch_.data() + kBnorm2, + sizeof(float), cudaMemcpyDeviceToHost, + stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + float abs_tol2 = options_.absolute_tolerance * options_.absolute_tolerance; + float rel_tol2 = options_.relative_tolerance * options_.relative_tolerance; + float stop_thresh = fmaxf(abs_tol2, rel_tol2 * fmaxf(b_norm2, 1e-30f)); + + // Build dense-vector descriptors with explicit size = n. dvector::resize + // is grow-only, so p_/Ap_ may have capacity > n across successive Solves on + // problems with different dimensions; using p_.size() in the descriptor + // would mismatch matA's row count. + cusparseDnVecDescr_t vecX = nullptr; + cusparseDnVecDescr_t vecY = nullptr; + THROW_ON_CUSPARSE_ERROR( + cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR( + cusparseCreateDnVec(&vecY, n, Ap_.data(), CUDA_R_32F)); + + int threads = 256; + int blocks = (n + threads - 1) / threads; + + // Convergence is polled every kCheckPeriod iterations. Polling more often + // turns the inner loop back into a sequence of host syncs; polling less + // often risks doing extra work after the iteration has already converged. + // 2 is a safe default — small PGO systems converge in a handful of iterates + // and don't tolerate a long fixed period, while SBA still benefits from + // the avoided alpha/beta-on-host syncs that the rest of the inner loop now + // sidesteps. + constexpr int kCheckPeriod = 2; + int it = 0; + for (; it < options_.max_iterations; ++it) { + // Ap = A * p + float spmv_alpha = 1.f; + float spmv_beta = 0.f; + THROW_ON_CUSPARSE_ERROR(cusparseSpMV(handle, + CUSPARSE_OPERATION_NON_TRANSPOSE, + &spmv_alpha, matA, vecX, &spmv_beta, + vecY, CUDA_R_32F, + CUSPARSE_SPMV_ALG_DEFAULT, + spmv_buffer_.data())); + + // pAp = + DotAsync(stream, p_.data(), Ap_.data(), n, d_scratch_.data() + kPAp); + // alpha = rz_old / pAp on device. + ComputeAlphaKernel<<<1, 1, 0, stream>>>(d_scratch_.data() + kRzOld, + d_scratch_.data() + kPAp, + d_scratch_.data() + kAlpha); + PcgUpdateKernel<<>>( + d_scratch_.data() + kAlpha, p_.data(), Ap_.data(), result.data(), + r_.data(), n); + + // z = M^{-1} r ; rz_new = ; r_norm2 = + LaunchApplyPrecond(stream, B, precond_factors_.data(), r_.data(), z_.data(), + num_blocks_); + DotAsync(stream, r_.data(), r_.data(), n, d_scratch_.data() + kRnorm2); + DotAsync(stream, r_.data(), z_.data(), n, d_scratch_.data() + kRzNew); + + // beta = rz_new / rz_old, then rz_old <- rz_new (single device kernel). + ComputeBetaKernel<<<1, 1, 0, stream>>>(d_scratch_.data() + kRzNew, + d_scratch_.data() + kRzOld, + d_scratch_.data() + kBeta); + PcgDirectionKernel<<>>( + d_scratch_.data() + kBeta, z_.data(), p_.data(), n); + + // Convergence poll every kCheckPeriod iterations: one D2H of one float + // plus a stream sync. Cheap relative to a full inner-loop sync. + if (((it + 1) % kCheckPeriod) == 0 || + (it + 1) == options_.max_iterations) { + float r_norm2 = 0.f; + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&r_norm2, + d_scratch_.data() + kRnorm2, + sizeof(float), + cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + if (r_norm2 <= stop_thresh) { + ++it; + break; + } + } + } + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecX)); + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecY)); + last_iterations_ = it; + return true; +} + +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h new file mode 100644 index 0000000..b8aadce --- /dev/null +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -0,0 +1,116 @@ +/* + * 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/cusparse_helper.h" +#include "cunls/linear_solver/csr_sparse_linear_solver.h" + +namespace cunls { + +/** + * @brief Configuration options for the block-sparse PCG solver. + * + * The solver implements Preconditioned Conjugate Gradient with block-Jacobi + * preconditioning on a symmetric-positive-definite CSR matrix. Block size is + * inferred from the matrix structure: the user supplies a fixed block dimension + * B and the diagonal B x B tiles of the matrix are factored once per outer + * call to ``Solve`` and applied at every PCG iteration. + */ +struct BlockSparsePCGOptions { + /** Block dimension used by the block-Jacobi preconditioner. + * Must divide the matrix size. B = 1 falls back to scalar Jacobi. */ + int block_size = 6; + + /** Maximum number of PCG iterations. */ + int max_iterations = 200; + + /** Convergence threshold on the relative residual ||r_k|| / ||b||. */ + float relative_tolerance = 1e-3f; + + /** Absolute threshold on ||r_k|| for early exit. */ + float absolute_tolerance = 1e-30f; + + /** Floor added to LDLT diagonal pivots to keep the preconditioner + * numerically invertible when the diagonal block is near-singular. */ + float pivot_floor = 1e-12f; +}; + +/** + * @brief Block-Jacobi preconditioned conjugate gradient solver. + * + * Solves H x = b for symmetric positive (semi-)definite H stored in CSR + * format. The preconditioner consists of the dense ``B x B`` diagonal tiles of + * H, factored independently with one warp per block using a small in-register + * LDLT. SpMV is delegated to cuSPARSE (CSR Hermitian SpMV). + * + * The implementation is intended for normal equations from Gauss-Newton / + * Levenberg-Marquardt where: + * - H has natural block structure aligned with state blocks (e.g. 6 for SE3, + * 3 for Vector<3>, ...); + * - the matrix structure does not change between ``Solve`` calls inside one + * Minimize, so allocations and the cuSPARSE descriptor are reused; + * - the previous step is a good warm start for the next iteration. + */ +class BlockSparsePCGSolver : public CSRSparseLinearSolver { +public: + explicit BlockSparsePCGSolver(BlockSparsePCGOptions options = {}); + ~BlockSparsePCGSolver() override; + + bool Initialize(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) final; + + bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) final; + + /** @brief Number of PCG iterations consumed by the most recent ``Solve``. */ + int LastIterations() const { return last_iterations_; } + +private: + BlockSparsePCGOptions options_; + + /** Number of rows in the system (cached on Initialize / refreshed on Solve). + */ + int matrix_size_ = 0; + /** Number of block rows = matrix_size_ / block_size. */ + int num_blocks_ = 0; + + /** Stored LDLT factors of every B x B diagonal tile, packed contiguously + * (block i begins at index i * B * B; lower-triangle is L, diagonal is D). */ + dvector precond_factors_; + + /** Scratch device vectors used by PCG. */ + dvector r_; + dvector z_; + dvector p_; + dvector Ap_; + + /** Two-slot device scalar buffer used for fused reductions. */ + dvector d_scratch_; + + /** Persistent SpMV state. */ + cuSPARSEHandle cusparse_handle_; + cuSPARSEMatrixDescription mat_desc_; + dvector spmv_buffer_; + bool spmv_buffer_ready_ = false; + + int last_iterations_ = 0; +}; + +} // namespace cunls diff --git a/cunls/linear_solver/llms.txt b/cunls/linear_solver/llms.txt index 3b6bd47..2ba8c18 100644 --- a/cunls/linear_solver/llms.txt +++ b/cunls/linear_solver/llms.txt @@ -14,6 +14,18 @@ Sparse linear system abstraction and implementation used by minimizers. - `cudss_sparse_linear_solver.h`: - `cuDSSLinearSolverOptions` - `cuDSSLinearSolver` +- `block_sparse_pcg_solver.h`: + - `BlockSparsePCGOptions` — `block_size`, `max_iterations`, `relative_tolerance`, + `absolute_tolerance`, `pivot_floor`. + - `BlockSparsePCGSolver` — iterative SPD solver: cuSPARSE CSR SpMV + + block-Jacobi LDLT preconditioner (dense B x B diagonal tiles factored in + shared memory). Inner loop computes alpha/beta on device; convergence is + polled every 2 iterations to keep CPU launches ahead of the GPU. + Best fit: PGO (block_size = 6) and PGO-like normal equations. + Mixed-block problems (e.g. SBA with 6x6 pose + 3x3 landmark blocks) work + correctly with block_size = 3 but converge more slowly than cuDSS — + Schur-complement elimination of the landmarks would be the natural + follow-up. ## Expected system form diff --git a/cunls/linear_solver/sparse_linear_solver.cpp b/cunls/linear_solver/sparse_linear_solver.cpp index 01072f0..b4277ec 100644 --- a/cunls/linear_solver/sparse_linear_solver.cpp +++ b/cunls/linear_solver/sparse_linear_solver.cpp @@ -35,6 +35,9 @@ CreateCSRSparseLinearSolver(SparseLinearSolverType type, return std::make_unique(); case SparseLinearSolverType::DenseQR: return std::make_unique(); + case SparseLinearSolverType::BlockSparsePCG: + return std::make_unique( + config.block_sparse_pcg_options); default: throw std::invalid_argument("Invalid sparse linear solver type"); } diff --git a/cunls/linear_solver/sparse_linear_solver.h b/cunls/linear_solver/sparse_linear_solver.h index 7397a01..11bd931 100644 --- a/cunls/linear_solver/sparse_linear_solver.h +++ b/cunls/linear_solver/sparse_linear_solver.h @@ -21,6 +21,7 @@ #include +#include "cunls/linear_solver/block_sparse_pcg_solver.h" #include "cunls/linear_solver/csr_sparse_linear_solver.h" #include "cunls/linear_solver/cudss_sparse_linear_solver.h" #include "cunls/linear_solver/dense_cholesky_solver.h" @@ -33,15 +34,20 @@ namespace cunls { * @brief Selects the linear solver backend for the Gauss-Newton system. */ enum class SparseLinearSolverType { - cuDSS, ///< Sparse direct solver using NVIDIA's cuDSS library. - DenseLDLT, ///< Converts CSR to dense and solves with a custom CUDA - ///< pivoted LDLT kernel. - DenseCholesky, ///< Converts CSR to dense and solves with cuSOLVER Cholesky - ///< factorization (cusolverDnSpotrf / cusolverDnSpotrs). - ///< Requires SPD matrix. - DenseQR, ///< Converts CSR to dense and solves with cuSOLVER QR - ///< factorization (cusolverDnSgeqrf / cusolverDnSormqr / - ///< cublasStrsm). Works for any non-singular square matrix. + cuDSS, ///< Sparse direct solver using NVIDIA's cuDSS library. + DenseLDLT, ///< Converts CSR to dense and solves with a custom CUDA + ///< pivoted LDLT kernel. + DenseCholesky, ///< Converts CSR to dense and solves with cuSOLVER Cholesky + ///< factorization (cusolverDnSpotrf / cusolverDnSpotrs). + ///< Requires SPD matrix. + DenseQR, ///< Converts CSR to dense and solves with cuSOLVER QR + ///< factorization (cusolverDnSgeqrf / cusolverDnSormqr / + ///< cublasStrsm). Works for any non-singular square matrix. + BlockSparsePCG, ///< Block-Jacobi preconditioned CG. Iterative solver tuned + ///< for SPD normal equations with uniform diagonal block + ///< structure (e.g. 6x6 for SE3). Skips the sparse direct + ///< factorization cost; the preconditioner is refactored on + ///< every Solve from the current diagonal tiles. }; /** @@ -54,6 +60,7 @@ enum class SparseLinearSolverType { */ struct SparseLinearSolverConfig { cuDSSLinearSolverOptions cudss_solver_options; + BlockSparsePCGOptions block_sparse_pcg_options; }; /** diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index 55ad618..f3d9410 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -540,6 +540,7 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, LogError(str); throw std::runtime_error(str); } + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); } // Main optimization loop @@ -553,13 +554,18 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, summary.iteration_costs.push_back(summary.final_cost); { - auto solve_range = profiler_domain_.CreateDomainRange("Solve"); + auto solve_range = profiler_domain_.CreateDomainRange("LinearSolve"); bool success = solver_->Solve(stream, lhs_work_, rhs_work_, step_); if (!success) { std::string str = "Failed to solve linear system"; LogError(str); throw std::runtime_error(str); } + // Fair-measurement barrier: ensure the linear-solve kernels finish before + // the NVTX range ends. nsys reports CPU-side NVTX durations, so without a + // sync the recorded interval excludes asynchronous GPU work that the + // solver enqueued. + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); } MapScaledLinearSolutionToTangentStep(stream, step_); diff --git a/tests/pgo_minimizer_test.cpp b/tests/pgo_minimizer_test.cpp index 9ceca1c..c37d418 100644 --- a/tests/pgo_minimizer_test.cpp +++ b/tests/pgo_minimizer_test.cpp @@ -58,6 +58,7 @@ #include "cunls/factor/information_factor_batch.h" #include "cunls/factor/se3_between_factor_batch.h" #include "cunls/minimizer/levenberg_marquardt_minimizer.h" +#include "tests/utils.h" #include "cunls/minimizer/problem.h" #include "cunls/state/se3_state_batch.h" @@ -298,6 +299,11 @@ TEST_F(PgoMinimizerTestFixture, Optimize) { options.state_tolerance = 1e-10f; options.cost_tolerance = 1e-2f; options.disable_safety_checks = false; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options = { + test_utils::PCGBlockSizeFromEnv(6), + test_utils::PCGMaxIterFromEnv(200), + test_utils::PCGTolFromEnv(1e-3f)}; LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1000.0; diff --git a/tests/sba_minimizer_test.cpp b/tests/sba_minimizer_test.cpp index b3900bf..d04b7bf 100644 --- a/tests/sba_minimizer_test.cpp +++ b/tests/sba_minimizer_test.cpp @@ -54,6 +54,7 @@ #include "cunls/robustifier/huber_loss_function_batch.h" #include "cunls/state/se3_state_batch.h" #include "cunls/state/vector_state_batch.h" +#include "tests/utils.h" namespace cunls { @@ -371,6 +372,15 @@ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { options.state_tolerance = 1e-8f; options.cost_tolerance = 1e4; options.disable_safety_checks = false; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + // SBA has both 6x6 (poses) and 3x3 (landmarks) diagonal blocks. Using + // block_size=3 keeps the preconditioner cheap and well-conditioned for the + // landmark blocks while still capturing useful structure inside the 6x6 + // pose tiles (every 6x6 splits into a 2x2 grid of 3x3 sub-blocks). + options.sparse_linear_solver_config.block_sparse_pcg_options = { + test_utils::PCGBlockSizeFromEnv(3), + test_utils::PCGMaxIterFromEnv(400), + test_utils::PCGTolFromEnv(1e-3f)}; LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; diff --git a/tests/sparse_linear_solver_test.cpp b/tests/sparse_linear_solver_test.cpp index b101755..e963b8f 100644 --- a/tests/sparse_linear_solver_test.cpp +++ b/tests/sparse_linear_solver_test.cpp @@ -27,6 +27,8 @@ #include +#include +#include #include #include #include @@ -215,4 +217,120 @@ TEST(SparseLinearSolverTest, Solve) { ASSERT_NEAR(squared_error / matrix_size, 0, 1e-1); } +namespace { + +// Generates a block-SPD matrix: dense 6x6 diagonal blocks plus a sparse pattern +// of symmetric 6x6 off-diagonal couplings. Mirrors the structure of a +// Hessian from an SE3 pose graph and is the natural fit for the block-Jacobi +// preconditioner. +void GenerateBlockSPDMatrix(std::mt19937 &gen, int num_blocks, int block_size, + float off_diag_strength, + std::vector &csr_values, + std::vector &csr_col_idx, + std::vector &csr_row_offsets) { + int n = num_blocks * block_size; + std::uniform_real_distribution off_dist(-off_diag_strength, + off_diag_strength); + std::uniform_real_distribution prob_dist(0.f, 1.f); + // Dense per-row map for ordered CSR assembly. + std::vector> rows(n); + + // Strongly diagonal-dominant diagonal blocks for SPD guarantee. + for (int b = 0; b < num_blocks; ++b) { + for (int i = 0; i < block_size; ++i) { + for (int j = 0; j < block_size; ++j) { + int gi = b * block_size + i; + int gj = b * block_size + j; + if (i == j) { + rows[gi][gj] = 50.f; + } else { + float v = off_dist(gen) * 0.1f; + rows[gi][gj] = (rows[gi].count(gj) ? rows[gi][gj] : 0.f) + v; + } + } + } + } + // Add symmetric block off-diagonal couplings to a few neighbour blocks. + for (int b = 0; b < num_blocks; ++b) { + for (int nb = b + 1; nb < num_blocks; ++nb) { + if (prob_dist(gen) > 5.f / num_blocks) { + continue; + } + for (int i = 0; i < block_size; ++i) { + for (int j = 0; j < block_size; ++j) { + float v = off_dist(gen); + int gi = b * block_size + i; + int gj = nb * block_size + j; + rows[gi][gj] = (rows[gi].count(gj) ? rows[gi][gj] : 0.f) + v; + rows[gj][gi] = (rows[gj].count(gi) ? rows[gj][gi] : 0.f) + v; + } + } + } + } + + csr_row_offsets.clear(); + csr_col_idx.clear(); + csr_values.clear(); + csr_row_offsets.push_back(0); + for (int i = 0; i < n; ++i) { + for (const auto &[j, v] : rows[i]) { + csr_col_idx.push_back(j); + csr_values.push_back(v); + } + csr_row_offsets.push_back(static_cast(csr_col_idx.size())); + } +} + +} // namespace + +TEST(SparseLinearSolverTest, BlockSparsePCGSolve) { + std::mt19937 gen(7); + constexpr int num_blocks = 200; + constexpr int block_size = 6; + constexpr int matrix_size = num_blocks * block_size; + + std::vector csr_values; + std::vector csr_col_idx; + std::vector csr_row_offsets; + GenerateBlockSPDMatrix(gen, num_blocks, block_size, 1.0f, csr_values, + csr_col_idx, csr_row_offsets); + + CSRSparseMatrix mat; + test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, + mat); + + std::vector rhs_cpu; + test_utils::GenerateRandomVector(matrix_size, rhs_cpu); + dvector rhs(rhs_cpu); + dvector result(matrix_size); + + CudaStream stream; + BlockSparsePCGOptions opts; + opts.block_size = block_size; + opts.relative_tolerance = 1e-5f; + opts.max_iterations = 500; + BlockSparsePCGSolver solver(opts); + ASSERT_TRUE(solver.Initialize(stream.GetStream(), mat, rhs, result)); + ASSERT_TRUE(solver.Solve(stream.GetStream(), mat, rhs, result)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + + std::vector x(matrix_size); + result.CopyToHost(x.data(), matrix_size); + + std::vector Ax; + MultiplySymmetricCSRMatrixByVector(csr_row_offsets, csr_col_idx, csr_values, + x, Ax); + float sq_err = 0.f; + float b_sq = 0.f; + for (int i = 0; i < matrix_size; ++i) { + float d = Ax[i] - rhs_cpu[i]; + sq_err += d * d; + b_sq += rhs_cpu[i] * rhs_cpu[i]; + } + float rel_err = std::sqrt(sq_err / std::max(b_sq, 1e-30f)); + EXPECT_LT(rel_err, 1e-3f) << "PCG residual too large (relative): " << rel_err; + EXPECT_GT(solver.LastIterations(), 0); + EXPECT_LE(solver.LastIterations(), opts.max_iterations); +} + } // namespace cunls diff --git a/tests/synthetic_pgo_test.cpp b/tests/synthetic_pgo_test.cpp index 70e1b6a..d0c8b9f 100644 --- a/tests/synthetic_pgo_test.cpp +++ b/tests/synthetic_pgo_test.cpp @@ -44,6 +44,7 @@ #include "cunls/minimizer/levenberg_marquardt_minimizer.h" #include "cunls/minimizer/problem.h" #include "cunls/state/se3_state_batch.h" +#include "tests/utils.h" namespace cunls { @@ -261,6 +262,11 @@ TEST_F(SyntheticPGOTest, OptimizeConsecutiveBetweenConstraints) { options.state_tolerance = 1e-6f; options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options = { + test_utils::PCGBlockSizeFromEnv(6), + test_utils::PCGMaxIterFromEnv(400), + test_utils::PCGTolFromEnv(1e-4f)}; // GaussNewtonMinimizer minimizer(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; @@ -352,6 +358,11 @@ TEST_F(SyntheticPGOTest, InformationBetweenFactorBatch) { options.state_tolerance = 1e-6f; options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options = { + test_utils::PCGBlockSizeFromEnv(6), + test_utils::PCGMaxIterFromEnv(400), + test_utils::PCGTolFromEnv(1e-4f)}; // GaussNewtonMinimizer minimizer(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; @@ -428,6 +439,11 @@ TEST_F(SyntheticPGOTest, WeightedWrapsInformationBetweenFactorBatch) { options.state_tolerance = 1e-6f; options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options = { + test_utils::PCGBlockSizeFromEnv(6), + test_utils::PCGMaxIterFromEnv(400), + test_utils::PCGTolFromEnv(1e-4f)}; LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -502,6 +518,11 @@ TEST_F(SyntheticPGOTest, InformationWrapsWeightedBetweenFactorBatch) { options.state_tolerance = 1e-6f; options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options = { + test_utils::PCGBlockSizeFromEnv(6), + test_utils::PCGMaxIterFromEnv(400), + test_utils::PCGTolFromEnv(1e-4f)}; LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; diff --git a/tests/utils.h b/tests/utils.h index 07bfca5..c0e9313 100644 --- a/tests/utils.h +++ b/tests/utils.h @@ -25,10 +25,14 @@ #include #include +#include +#include + #include "cunls/common/device_vector.h" #include "cunls/common/helper.h" #include "cunls/common/types.h" #include "cunls/factor/prior_vector_factor_batch.h" +#include "cunls/linear_solver/sparse_linear_solver.h" #include "cunls/state/vector_state_batch.h" namespace cunls { @@ -341,5 +345,52 @@ CopyStateToHost(const VectorStateBatch &state_batch) { return out; } +// ============================================================================ +// Solver-selection helpers (driven by CUNLS_SOLVER env var) +// ============================================================================ + +/** + * @brief Reads CUNLS_SOLVER from the environment and returns the matching + * solver type. Defaults to cuDSS so existing tests keep their + * baseline behaviour. Recognised values: "cuDSS", "BlockSparsePCG". + */ +inline SparseLinearSolverType SolverTypeFromEnv() { + const char *s = std::getenv("CUNLS_SOLVER"); + if (s != nullptr && std::strcmp(s, "BlockSparsePCG") == 0) { + return SparseLinearSolverType::BlockSparsePCG; + } + return SparseLinearSolverType::cuDSS; +} + +/** Reads CUNLS_PCG_BLOCK_SIZE; defaults to ``fallback`` if unset/invalid. */ +inline int PCGBlockSizeFromEnv(int fallback) { + const char *s = std::getenv("CUNLS_PCG_BLOCK_SIZE"); + if (s == nullptr) { + return fallback; + } + int v = std::atoi(s); + return (v >= 1 && v <= 16) ? v : fallback; +} + +/** Reads CUNLS_PCG_MAX_ITER; defaults to ``fallback`` if unset/invalid. */ +inline int PCGMaxIterFromEnv(int fallback) { + const char *s = std::getenv("CUNLS_PCG_MAX_ITER"); + if (s == nullptr) { + return fallback; + } + int v = std::atoi(s); + return (v > 0) ? v : fallback; +} + +/** Reads CUNLS_PCG_TOL; defaults to ``fallback`` if unset/invalid. */ +inline float PCGTolFromEnv(float fallback) { + const char *s = std::getenv("CUNLS_PCG_TOL"); + if (s == nullptr) { + return fallback; + } + float v = static_cast(std::atof(s)); + return (v > 0.f) ? v : fallback; +} + } // namespace test_utils } // namespace cunls From c7de5d7f9fa2b877cb88c4166bb4b4281cf2ccb7 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 13 May 2026 12:49:44 -0700 Subject: [PATCH 2/4] Improve PCG --- CMakeLists.txt | 1 + .../linear_solver/block_sparse_pcg_solver.cu | 1078 +++++++++++++---- cunls/linear_solver/block_sparse_pcg_solver.h | 283 ++++- .../linear_solver/csr_sparse_linear_solver.h | 24 +- .../cudss_sparse_linear_solver.cpp | 1 + .../cudss_sparse_linear_solver.h | 3 +- cunls/linear_solver/dense_cholesky_solver.cu | 1 + cunls/linear_solver/dense_cholesky_solver.h | 6 +- cunls/linear_solver/dense_linear_solver.cu | 1 + cunls/linear_solver/dense_linear_solver.h | 6 +- cunls/linear_solver/dense_qr_solver.cu | 1 + cunls/linear_solver/dense_qr_solver.h | 6 +- cunls/linear_solver/llms.txt | 19 +- cunls/minimizer/gauss_newton_minimizer.cu | 8 +- cunls/minimizer/problem.cpp | 3 +- tests/dense_cholesky_solver_test.cpp | 12 +- tests/dense_linear_solver_test.cpp | 12 +- tests/dense_qr_solver_test.cpp | 12 +- tests/pgo_minimizer_test.cpp | 7 +- tests/sba_minimizer_test.cpp | 7 +- tests/sparse_linear_solver_test.cpp | 5 +- tests/synthetic_pgo_test.cpp | 266 +++- tests/synthetic_sba_test.cpp | 365 ++++++ tests/utils.h | 13 +- 24 files changed, 1774 insertions(+), 366 deletions(-) create mode 100644 tests/synthetic_sba_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index d42c9b1..22c7b24 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -172,6 +172,7 @@ if(BUILD_TESTING) tests/pnp_factor_batch_test.cpp tests/decreasing_scale_minimizer_test.cpp tests/synthetic_pgo_test.cpp + tests/synthetic_sba_test.cpp tests/dense_matrix_ops_test.cpp tests/utils_test.cpp tests/device_vector_test.cpp diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index f412249..b098a3d 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -15,11 +15,78 @@ * limitations under the License. */ +/** + * @file block_sparse_pcg_solver.cu + * + * Implementation of the block-Jacobi preconditioned conjugate gradient + * solver declared in @ref block_sparse_pcg_solver.h. + * + * ## Math summary + * + * Given a symmetric positive (semi-)definite matrix `H ∈ R^{N×N}` in CSR + * format and a right-hand side `b ∈ R^N`, the algorithm computes a + * sequence of approximations `x_k` minimizing the H-norm error + * `||x - x_*||_H` over the Krylov subspace + * `K_k(M^{-1} H, M^{-1} r_0)` where `M` is the block-Jacobi + * preconditioner. + * + * Block-Jacobi preconditioner. Define a partition of the index set + * `{0,...,N-1}` into contiguous diagonal blocks `B_1, B_2, ...`. The + * preconditioner is the block-diagonal matrix whose restriction to + * `B_i × B_i` equals the corresponding diagonal tile `H[B_i, B_i]`. + * Equivalently, `M^{-1}` is block-diagonal and applying it amounts to + * solving `|B_i|`-sized SPD systems independently per block. Each + * tile is factored once per @ref BlockSparsePCGSolver::Solve via LDLT + * (`H_d = L D L^T`); applying `M^{-1}` is then a triangular solve plus + * diagonal scaling. + * + * The PCG recurrence (Saad, *Iterative Methods for Sparse Linear + * Systems*, 2nd ed., Algorithm 9.1): + * r_0 = b - H x_0 + * z_0 = M^{-1} r_0 + * p_0 = z_0 + * rz_0 = + * for k = 0, 1, ...: + * q_k = H p_k // SpMV + * alpha_k = rz_k / // scalar + * x_{k+1} = x_k + alpha_k p_k // axpy + * r_{k+1} = r_k - alpha_k q_k // axpy + * z_{k+1} = M^{-1} r_{k+1} // block-Jacobi apply + * rz_{k+1}= // dot + * beta_k = rz_{k+1} / rz_k // scalar + * p_{k+1} = z_{k+1} + beta_k p_k // axpy + * + * We initialize `x_0 = 0`, which makes `r_0 = b` (saves one SpMV). + * + * ## Sync strategy + * + * The inner loop computes both alpha and beta on the device (single + * thread kernels reading the dot-product slots), so the host never has + * to read them. Convergence is checked every + * @ref BlockSparsePCGOptions::check_period iterations by copying + * `||r_k||^2` (one float) from device to host. This keeps the host + * roughly `check_period` iterations ahead of the GPU, which is the + * sweet spot — frequent syncs (period 1) re-stall the launch queue, + * infrequent syncs (period > 8) waste work after the residual already + * dipped below the tolerance. + * + * ## Variable block sizes + * + * Each segment in @ref BlockSparsePCGOptions::block_layout gets its + * own templated kernel launch (one per `Factor` and one per `Apply`). + * Per-segment dispatch is a small host-side branch on the templated + * block-size specialization (B ∈ {1, 2, 3, 4, 6, 7, 15}) and falls + * back to a runtime-B generic kernel for any other value. This keeps + * the hot kernels register-resident for the common cases (SE3 = 6, + * Vector<3> = 3) while still supporting arbitrary user state batches. + */ + #include "cunls/linear_solver/block_sparse_pcg_solver.h" #include #include +#include #include #include #include @@ -27,6 +94,8 @@ #include "cunls/common/cusparse_helper.h" #include "cunls/common/helper.h" #include "cunls/common/log.h" +#include "cunls/minimizer/problem.h" +#include "cunls/state/state_batch.h" namespace cunls { @@ -34,31 +103,41 @@ namespace { constexpr int kMaxBlockSize = 16; -// ----------------------------------------------------------------------------- -// Device kernels -// ----------------------------------------------------------------------------- +// ============================================================================= +// Device kernels — preconditioner construction (`Factor`) +// ============================================================================= /** - * @brief Extracts the dense B x B diagonal tiles from a symmetric CSR matrix - * and stores their LDLT factors. + * @brief Factors one segment of block-diagonal tiles, in place into the + * factor buffer. * - * One thread block handles one diagonal block. Threads cooperate to scan the - * rows in [block_row * B, block_row * B + B), pull the entries with column id - * in the same range, and assemble a dense B x B tile in shared memory. A - * single-thread LDLT then runs over the symmetric tile (sizes encountered here - * are 3-6, so the serial cost is negligible compared to the global I/O). + * One CTA per tile. Threads cooperatively: + * 1. Gather the dense `B × B` tile from `H`'s CSR rows + * `[segment_row_start + b*B, segment_row_start + b*B + B)`, + * filtering column indices to the tile's column range. Any + * entries outside are zero (the tile is dense in the + * preconditioner's view of H; CSR sparsity within the tile is + * irrelevant for Jacobi). + * 2. Symmetrize numerically: `H_d := (H_d + H_d^T) / 2`. Algebraically + * a no-op for `J^T J`, but cheap insurance against FP drift when + * LM damping or column scaling has been applied in place. + * 3. Run a serial LDLT on thread 0: + * `for k in 0..B: D_{kk} = H_{kk}; L_{ik} = H_{ik}/D_{kk}; + * H_{ij} -= L_{ik} D_{kk} L_{jk}` for i, j > k. + * A pivot floor (sign-preserving) guards against singular tiles. + * 4. Write the (lower-triangle L, diagonal D) result to the global + * factor buffer. * - * The factored block is written back to ``factors`` in row-major layout: - * the strict lower-triangle holds L (1's on the diagonal implicitly), and the - * diagonal of the stored tile holds D. Upper triangle is undefined and not - * read. + * The serial LDLT is fine for the sizes here (B ≤ 16, dominant cost is + * the global I/O for the tile, not the O(B³) arithmetic). + * + * @tparam B Compile-time tile side length. */ template -__global__ void ExtractAndFactorBlockDiagonals(const int *__restrict__ row_off, - const int *__restrict__ col_idx, - const float *__restrict__ values, - int num_blocks, float pivot_floor, - float *__restrict__ factors) { +__global__ void ExtractAndFactorBlockDiagonalsKernel( + const int *__restrict__ row_off, const int *__restrict__ col_idx, + const float *__restrict__ values, int row_start, int num_blocks, + int factor_offset, float pivot_floor, float *__restrict__ factors) { int block_row = blockIdx.x; if (block_row >= num_blocks) { return; @@ -71,12 +150,12 @@ __global__ void ExtractAndFactorBlockDiagonals(const int *__restrict__ row_off, } __syncthreads(); - // Each thread takes one of the B rows. + // Step 1: gather dense tile from CSR. if (tid < B) { - int global_row = block_row * B + tid; + int global_row = row_start + block_row * B + tid; int start = row_off[global_row]; int end = row_off[global_row + 1]; - int col_lo = block_row * B; + int col_lo = row_start + block_row * B; int col_hi = col_lo + B; for (int k = start; k < end; ++k) { int c = col_idx[k]; @@ -87,10 +166,7 @@ __global__ void ExtractAndFactorBlockDiagonals(const int *__restrict__ row_off, } __syncthreads(); - // Symmetrize so the LDLT below can read either triangle. CSR is - // symmetric for J^T J and Levenberg-Marquardt damping is on the diagonal, - // so off-diagonal entries should already match; symmetrizing keeps the - // factorization stable when only the lower triangle is stored. + // Step 2: numerical symmetrization. if (tid < B) { for (int j = tid + 1; j < B; ++j) { float a = tile[tid * B + j]; @@ -102,14 +178,14 @@ __global__ void ExtractAndFactorBlockDiagonals(const int *__restrict__ row_off, } __syncthreads(); - // Serial LDLT in shared memory. B is small (3..16); a single thread is the - // simplest correct implementation. Outer-product update with diagonal pivot - // and a small floor to keep the preconditioner stable on near-singular tiles. + // Step 3: serial LDLT (one thread). Diagonal D is stored on the + // tile's diagonal, strict lower triangle gets L (unit-diagonal + // implicit). if (tid == 0) { for (int k = 0; k < B; ++k) { float d = tile[k * B + k]; if (fabsf(d) < pivot_floor) { - d = pivot_floor; + d = (d >= 0.f) ? pivot_floor : -pivot_floor; } tile[k * B + k] = d; float inv_d = 1.f / d; @@ -124,25 +200,213 @@ __global__ void ExtractAndFactorBlockDiagonals(const int *__restrict__ row_off, } __syncthreads(); - // Write factors back (row-major). - float *out = factors + block_row * B * B; + // Step 4: write back. + float *out = factors + factor_offset + block_row * B * B; + for (int i = tid; i < B * B; i += blockDim.x) { + out[i] = tile[i]; + } +} + +/** Generic-B fallback for non-templated block sizes. Uses dynamic + * shared memory of size `B*B` floats. Slightly slower than the + * templated version (no constant-B unrolling) but works for any + * block size up to @c kMaxBlockSize. */ +__global__ void ExtractAndFactorGenericKernel( + const int *__restrict__ row_off, const int *__restrict__ col_idx, + const float *__restrict__ values, int B, int row_start, int num_blocks, + int factor_offset, float pivot_floor, float *__restrict__ factors) { + int block_row = blockIdx.x; + if (block_row >= num_blocks) { + return; + } + extern __shared__ float smem[]; + float *tile = smem; + + int tid = threadIdx.x; + for (int i = tid; i < B * B; i += blockDim.x) { + tile[i] = 0.f; + } + __syncthreads(); + + if (tid < B) { + int global_row = row_start + block_row * B + tid; + int start = row_off[global_row]; + int end = row_off[global_row + 1]; + int col_lo = row_start + block_row * B; + int col_hi = col_lo + B; + for (int k = start; k < end; ++k) { + int c = col_idx[k]; + if (c >= col_lo && c < col_hi) { + tile[tid * B + (c - col_lo)] = values[k]; + } + } + } + __syncthreads(); + if (tid < B) { + for (int j = tid + 1; j < B; ++j) { + float a = tile[tid * B + j]; + float b = tile[j * B + tid]; + float s = 0.5f * (a + b); + tile[tid * B + j] = s; + tile[j * B + tid] = s; + } + } + __syncthreads(); + if (tid == 0) { + for (int k = 0; k < B; ++k) { + float d = tile[k * B + k]; + if (fabsf(d) < pivot_floor) { + d = (d >= 0.f) ? pivot_floor : -pivot_floor; + } + tile[k * B + k] = d; + float inv_d = 1.f / d; + for (int i = k + 1; i < B; ++i) { + float lik = tile[i * B + k] * inv_d; + tile[i * B + k] = lik; + for (int j = k + 1; j <= i; ++j) { + tile[i * B + j] -= lik * tile[j * B + k] * d; + } + } + } + } + __syncthreads(); + float *out = factors + factor_offset + block_row * B * B; for (int i = tid; i < B * B; i += blockDim.x) { out[i] = tile[i]; } } /** - * @brief Applies the precomputed block-LDLT preconditioner: z = M^{-1} r. + * @brief One-thread-per-block factor for small B (≤ 6). + * + * Symmetric counterpart to @ref ApplyBlockJacobiPerThreadKernel. Each + * thread loads the dense `B × B` diagonal tile from CSR into thread-local + * registers (no shared memory), runs the LDLT in registers, and writes + * the factored tile back. Packs ~256 tiles per CTA — for SBA's 1 M + * landmark tiles this is a >200× reduction in CTA count vs the + * one-CTA-per-tile path. + * + * Note: the gather still has to scan the CSR row of length `nnz/B` + * per thread (one row per thread). For SBA's pose / landmark rows + * this is short (~5 entries for a landmark, ~25 for a pose) so the + * scan stays in register loops. + */ +template +__global__ void ExtractAndFactorPerThreadKernel( + const int *__restrict__ row_off, const int *__restrict__ col_idx, + const float *__restrict__ values, int row_start, int num_blocks, + int factor_offset, float pivot_floor, float *__restrict__ factors) { + int block_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_row >= num_blocks) { + return; + } + // Gather the dense B×B tile into registers. + float tile[B * B]; +#pragma unroll + for (int i = 0; i < B * B; ++i) { + tile[i] = 0.f; + } + int col_lo = row_start + block_row * B; + int col_hi = col_lo + B; +#pragma unroll + for (int rr = 0; rr < B; ++rr) { + int global_row = col_lo + rr; + int start = row_off[global_row]; + int end = row_off[global_row + 1]; + for (int k = start; k < end; ++k) { + int c = col_idx[k]; + if (c >= col_lo && c < col_hi) { + tile[rr * B + (c - col_lo)] = values[k]; + } + } + } + // Symmetrize numerically. +#pragma unroll + for (int i = 0; i < B; ++i) { +#pragma unroll + for (int j = i + 1; j < B; ++j) { + float s = 0.5f * (tile[i * B + j] + tile[j * B + i]); + tile[i * B + j] = s; + tile[j * B + i] = s; + } + } + // LDLT. +#pragma unroll + for (int k = 0; k < B; ++k) { + float d = tile[k * B + k]; + if (fabsf(d) < pivot_floor) { + d = (d >= 0.f) ? pivot_floor : -pivot_floor; + } + tile[k * B + k] = d; + float inv_d = 1.f / d; +#pragma unroll + for (int i = k + 1; i < B; ++i) { + float lik = tile[i * B + k] * inv_d; + tile[i * B + k] = lik; +#pragma unroll + for (int j = k + 1; j < B; ++j) { + if (j <= i) { + tile[i * B + j] -= lik * tile[j * B + k] * d; + } + } + } + } + float *out = factors + factor_offset + block_row * B * B; +#pragma unroll + for (int i = 0; i < B * B; ++i) { + out[i] = tile[i]; + } +} + +/** Scalar-Jacobi extractor for B == 1: `M^{-1}[i] = 1 / H[i,i]`. */ +__global__ void ExtractScalarJacobi(const int *__restrict__ row_off, + const int *__restrict__ col_idx, + const float *__restrict__ values, + int row_start, int n, int factor_offset, + float pivot_floor, + float *__restrict__ factors) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n) { + return; + } + int i = row_start + idx; + int start = row_off[i]; + int end = row_off[i + 1]; + float d = pivot_floor; + for (int k = start; k < end; ++k) { + if (col_idx[k] == i) { + d = fmaxf(fabsf(values[k]), pivot_floor); + break; + } + } + factors[factor_offset + idx] = d; +} + +// ============================================================================= +// Device kernels — preconditioner application (`Apply`) +// ============================================================================= + +/** + * @brief Computes `z[block] = M^{-1} r[block]` for one tile in a + * segment, where `M^{-1}` is the inverse of the LDLT factor. * - * One block per diagonal tile, B threads per block. Reads r into registers, - * solves L y = r, D w = y, L^T z = w in three serial sweeps. Sizes are tiny - * (B in {3,6}) so the work fits in registers; the kernel is bandwidth-bound on - * the factor read. + * Each CTA owns a single tile. We pull the tile's `B` residual entries + * into shared memory, then perform three serial sweeps on thread 0: + * forward solve `L y = r` (L is unit-lower triangular) + * diagonal scale `D w = y` (D is stored on the tile's diagonal) + * back solve `L^T z = w` + * giving the block `M^{-1} r`. Serial because B is at most 16; the + * per-block arithmetic is ~3*B² FMAs, fully dominated by the global + * factor read. + * + * @tparam B Compile-time tile side length. */ template -__global__ void ApplyBlockJacobiPreconditioner( - const float *__restrict__ factors, const float *__restrict__ r, - float *__restrict__ z, int num_blocks) { +__global__ void ApplyBlockJacobiKernel(const float *__restrict__ factors, + int factor_offset, + const float *__restrict__ r, + int row_start, int num_blocks, + float *__restrict__ z) { int block_row = blockIdx.x; if (block_row >= num_blocks) { return; @@ -152,17 +416,15 @@ __global__ void ApplyBlockJacobiPreconditioner( int tid = threadIdx.x; if (tid < B) { - v[tid] = r[block_row * B + tid]; + v[tid] = r[row_start + block_row * B + tid]; } for (int i = tid; i < B * B; i += blockDim.x) { - L[i] = factors[block_row * B * B + i]; + L[i] = factors[factor_offset + block_row * B * B + i]; } __syncthreads(); - // Forward solve L y = r (unit diagonal); D w = y; L^T z = w. Serialized on - // thread 0 — the per-block work is 3*B^2 FMAs which dominates over any - // attempt to parallelize across B threads. if (tid == 0) { + // Forward: L y = r (unit diagonal). for (int i = 1; i < B; ++i) { float s = v[i]; #pragma unroll @@ -173,9 +435,11 @@ __global__ void ApplyBlockJacobiPreconditioner( } v[i] = s; } + // Diagonal: D w = y. for (int i = 0; i < B; ++i) { v[i] /= L[i * B + i]; } + // Back: L^T z = w. for (int i = B - 2; i >= 0; --i) { float s = v[i]; #pragma unroll @@ -190,59 +454,154 @@ __global__ void ApplyBlockJacobiPreconditioner( __syncthreads(); if (tid < B) { - z[block_row * B + tid] = v[tid]; + z[row_start + block_row * B + tid] = v[tid]; } } -/** Scalar Jacobi fallback (B == 1) — z = r / diag. */ +/** Generic-B apply for non-templated tile sizes; dynamic shared mem. */ +__global__ void ApplyBlockJacobiGenericKernel( + const float *__restrict__ factors, int factor_offset, int B, + const float *__restrict__ r, int row_start, int num_blocks, + float *__restrict__ z) { + int block_row = blockIdx.x; + if (block_row >= num_blocks) { + return; + } + extern __shared__ float smem[]; + float *L = smem; + float *v = L + B * B; + + int tid = threadIdx.x; + if (tid < B) { + v[tid] = r[row_start + block_row * B + tid]; + } + for (int i = tid; i < B * B; i += blockDim.x) { + L[i] = factors[factor_offset + block_row * B * B + i]; + } + __syncthreads(); + if (tid == 0) { + for (int i = 1; i < B; ++i) { + float s = v[i]; + for (int j = 0; j < i; ++j) { + s -= L[i * B + j] * v[j]; + } + v[i] = s; + } + for (int i = 0; i < B; ++i) { + v[i] /= L[i * B + i]; + } + for (int i = B - 2; i >= 0; --i) { + float s = v[i]; + for (int j = i + 1; j < B; ++j) { + s -= L[j * B + i] * v[j]; + } + v[i] = s; + } + } + __syncthreads(); + if (tid < B) { + z[row_start + block_row * B + tid] = v[tid]; + } +} + +/** Scalar-Jacobi apply (B == 1). */ __global__ void ApplyScalarJacobi(const float *__restrict__ factors, - const float *__restrict__ r, - float *__restrict__ z, int n) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) { + int factor_offset, + const float *__restrict__ r, int row_start, + int n, float *__restrict__ z) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n) { return; } - float d = factors[i]; - z[i] = r[i] / d; + z[row_start + idx] = r[row_start + idx] / factors[factor_offset + idx]; } -/** Scalar Jacobi extractor (B == 1) — pulls diag(H) with a pivot floor. */ -__global__ void ExtractScalarJacobi(const int *__restrict__ row_off, - const int *__restrict__ col_idx, - const float *__restrict__ values, int n, - float pivot_floor, - float *__restrict__ factors) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) { +/** + * @brief One-thread-per-block apply, optimized for small B (≤ 8). + * + * The per-block work (3*B² FMAs ≈ 27 for B=3, 108 for B=6) fits in a + * handful of registers — well below the per-thread state budget that + * limits occupancy. Holding `L` and `v` in thread-local registers + * eliminates the per-CTA shared-memory store/load round-trip and lets a + * full CTA process 256 different blocks instead of 256 threads ganging + * up on one block. For SBA with 1 M landmark tiles (B=3), this drops + * the grid from 1 M CTAs to ~4 K — a >200× reduction in launch and + * scheduling overhead. + * + * @tparam B Compile-time tile side length, must be small enough for the + * per-thread register file (we instantiate for {2, 3, 4, 6}). + */ +template +__global__ void +ApplyBlockJacobiPerThreadKernel(const float *__restrict__ factors, + int factor_offset, + const float *__restrict__ r, int row_start, + int num_blocks, float *__restrict__ z) { + int block_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_row >= num_blocks) { return; } - int start = row_off[i]; - int end = row_off[i + 1]; - float d = pivot_floor; - for (int k = start; k < end; ++k) { - if (col_idx[k] == i) { - d = fmaxf(fabsf(values[k]), pivot_floor); - break; + const float *f = factors + factor_offset + block_row * B * B; + const float *r_in = r + row_start + block_row * B; + float *z_out = z + row_start + block_row * B; + + float L[B * B]; + float v[B]; +#pragma unroll + for (int i = 0; i < B * B; ++i) { + L[i] = f[i]; + } +#pragma unroll + for (int i = 0; i < B; ++i) { + v[i] = r_in[i]; + } + + // Forward solve L y = r (unit-lower triangular). +#pragma unroll + for (int i = 1; i < B; ++i) { + float s = v[i]; +#pragma unroll + for (int j = 0; j < B; ++j) { + if (j < i) { + s -= L[i * B + j] * v[j]; + } } + v[i] = s; + } + // Diagonal: D w = y. +#pragma unroll + for (int i = 0; i < B; ++i) { + v[i] /= L[i * B + i]; + } + // Back: L^T z = w. +#pragma unroll + for (int i = B - 2; i >= 0; --i) { + float s = v[i]; +#pragma unroll + for (int j = 0; j < B; ++j) { + if (j > i) { + s -= L[j * B + i] * v[j]; + } + } + v[i] = s; } - factors[i] = d; -} -/** y = a*x + b*z, all length n. Used for PCG vector updates. */ -__global__ void AxpyKernel(float a, const float *__restrict__ x, float b, - const float *__restrict__ z, float *__restrict__ y, - int n) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) { - return; +#pragma unroll + for (int i = 0; i < B; ++i) { + z_out[i] = v[i]; } - y[i] = a * x[i] + b * z[i]; } -/** x = a*p + x, r = r - a*Ap, plus partial-sum dot products of (new_r, z') — - * this kernel is the hot path of PCG and is intentionally minimal. - * ``alpha`` is read once from device memory so that the host can keep - * enqueueing kernels without waiting on the previous dot product. */ +// ============================================================================= +// Device kernels — PCG vector / scalar ops +// ============================================================================= + +/** + * @brief Fused `x ← x + α p ; r ← r − α q` for the PCG step. + * + * Reads `α` from device memory so the host doesn't have to wait on the + * preceding `` dot product. One thread per coordinate. + */ __global__ void PcgUpdateKernel(const float *__restrict__ alpha_ptr, const float *__restrict__ p, const float *__restrict__ Ap, @@ -261,7 +620,9 @@ __global__ void PcgUpdateKernel(const float *__restrict__ alpha_ptr, r[i] -= a * Ap[i]; } -/** p = z + beta * p, length n. ``beta`` is fetched from device memory. */ +/** + * @brief Updates the search direction `p ← z + β p` (in place into p). + */ __global__ void PcgDirectionKernel(const float *__restrict__ beta_ptr, const float *__restrict__ z, float *__restrict__ p, int n) { @@ -277,7 +638,7 @@ __global__ void PcgDirectionKernel(const float *__restrict__ beta_ptr, p[i] = z[i] + b * p[i]; } -/** Single-thread device kernel: alpha = rz_old / pAp (with guard). */ +/** `alpha = rz_old / ` (single-thread, device-side). */ __global__ void ComputeAlphaKernel(const float *__restrict__ rz_old, const float *__restrict__ pAp, float *__restrict__ alpha) { @@ -285,7 +646,7 @@ __global__ void ComputeAlphaKernel(const float *__restrict__ rz_old, alpha[0] = (denom > 0.f) ? rz_old[0] / denom : 0.f; } -/** Single-thread device kernel: beta = rz_new / rz_old; rz_old <- rz_new. */ +/** `beta = rz_new / rz_old; rz_old <- rz_new` (single-thread). */ __global__ void ComputeBetaKernel(const float *__restrict__ rz_new, float *__restrict__ rz_old, float *__restrict__ beta) { @@ -295,150 +656,339 @@ __global__ void ComputeBetaKernel(const float *__restrict__ rz_new, rz_old[0] = num; } -/** - * @brief Single-pass dot product writing the result to ``out`` on device. - * - * Uses one block-stride loop per CTA, a warp reduction, and an atomicAdd into - * out[0]. Caller must zero out[0] before launch. Fine for the sizes we hit - * here (a few thousand to ~100k floats). - */ -__global__ void DotKernelZero(float *__restrict__ out) { out[0] = 0.f; } +/** Zero one or two scalar slots. */ +__global__ void ZeroScalarKernel(float *out) { out[0] = 0.f; } +__global__ void Zero2ScalarKernel(float *out_a, float *out_b) { + out_a[0] = 0.f; + out_b[0] = 0.f; +} -__global__ void DotKernel(const float *__restrict__ a, - const float *__restrict__ b, int n, - float *__restrict__ out) { - float s = 0.f; - for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; - i += blockDim.x * gridDim.x) { - s += a[i] * b[i]; - } - // Warp reduction. +/** Performs `dst += local` on a single device scalar via a per-block warp + * reduction + atomicAdd. Helper for fused dot kernels below. */ +__device__ inline void ReduceAndAddToScalar(float local, float *dst) { for (int off = warpSize / 2; off > 0; off >>= 1) { - s += __shfl_down_sync(0xffffffffu, s, off); + local += __shfl_down_sync(0xffffffffu, local, off); } __shared__ float warp_sums[32]; int lane = threadIdx.x & 31; int wid = threadIdx.x >> 5; if (lane == 0) { - warp_sums[wid] = s; + warp_sums[wid] = local; } __syncthreads(); if (wid == 0) { int nw = (blockDim.x + 31) >> 5; - s = (lane < nw) ? warp_sums[lane] : 0.f; + float s = (lane < nw) ? warp_sums[lane] : 0.f; for (int off = warpSize / 2; off > 0; off >>= 1) { s += __shfl_down_sync(0xffffffffu, s, off); } if (lane == 0) { - atomicAdd(out, s); + atomicAdd(dst, s); } } } -// ----------------------------------------------------------------------------- -// Host helpers -// ----------------------------------------------------------------------------- +/** + * @brief Block-stride dot product `out[0] += ` (atomic add). + * + * Each CTA does a register-level reduction over its slice, then a + * single-warp reduction inside shared memory, then one atomicAdd into + * the scalar slot. Caller must zero `out[0]` first (the companion + * @ref ZeroScalarKernel kernel above does this). + */ +__global__ void DotKernel(const float *__restrict__ a, + const float *__restrict__ b, int n, + float *__restrict__ out) { + float s = 0.f; + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += blockDim.x * gridDim.x) { + s += a[i] * b[i]; + } + ReduceAndAddToScalar(s, out); +} + +/** + * @brief Two-output dot kernel: writes `` and `` in a + * single pass over `a`. + * + * Halves the number of vector reads and saves one kernel launch per + * PCG iteration (only one `ZeroScalarKernel` pair instead of two). + * Each thread accumulates two partial sums in registers and the + * per-block reduction is identical to @ref DotKernel. + */ +__global__ void DualDotKernel(const float *__restrict__ a, + const float *__restrict__ b, + const float *__restrict__ c, int n, + float *__restrict__ out_ab, + float *__restrict__ out_ac) { + float sab = 0.f; + float sac = 0.f; + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; + i += blockDim.x * gridDim.x) { + float ai = a[i]; + sab += ai * b[i]; + sac += ai * c[i]; + } + ReduceAndAddToScalar(sab, out_ab); + ReduceAndAddToScalar(sac, out_ac); +} + +// ============================================================================= +// Host dispatch helpers +// ============================================================================= + +/** Picks the right ExtractAndFactor specialization for B. */ +void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, + int num_blocks, int factor_offset, + float pivot_floor, + const CSRSparseMatrix &matrix, + float *factors) { + if (num_blocks == 0) { + return; + } + const int *row_off = matrix.row_offsets.data(); + const int *col_idx = matrix.col_ids.data(); + const float *vals = matrix.values.data(); -void LaunchExtractAndFactor(cudaStream_t stream, int B, - const CSRSparseMatrix &matrix, int num_blocks, - float pivot_floor, dvector &factors) { if (B == 1) { - int n = num_blocks; int threads = 256; - int blocks = (n + threads - 1) / threads; + int blocks = (num_blocks + threads - 1) / threads; ExtractScalarJacobi<<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - n, pivot_floor, factors.data()); + row_off, col_idx, vals, row_start, num_blocks, factor_offset, + pivot_floor, factors); + THROW_ON_CUDA_ERROR(cudaGetLastError()); return; } - int threads = ((B + 31) / 32) * 32; // round up to a full warp +#define LAUNCH_FACTOR_PER_THREAD(BVAL) \ + case BVAL: { \ + int threads = 256; \ + int blocks = (num_blocks + threads - 1) / threads; \ + ExtractAndFactorPerThreadKernel \ + <<>>(row_off, col_idx, vals, row_start, \ + num_blocks, factor_offset, \ + pivot_floor, factors); \ + break; \ + } + + // Small-B path: one block per thread. switch (B) { - case 2: - ExtractAndFactorBlockDiagonals<2><<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_blocks, pivot_floor, factors.data()); - break; - case 3: - ExtractAndFactorBlockDiagonals<3><<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_blocks, pivot_floor, factors.data()); - break; - case 6: - ExtractAndFactorBlockDiagonals<6><<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_blocks, pivot_floor, factors.data()); + LAUNCH_FACTOR_PER_THREAD(2); + LAUNCH_FACTOR_PER_THREAD(3); + LAUNCH_FACTOR_PER_THREAD(4); + LAUNCH_FACTOR_PER_THREAD(5); + LAUNCH_FACTOR_PER_THREAD(6); + default: break; - case 7: - ExtractAndFactorBlockDiagonals<7><<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_blocks, pivot_floor, factors.data()); + } +#undef LAUNCH_FACTOR_PER_THREAD + if (B >= 2 && B <= 6) { + THROW_ON_CUDA_ERROR(cudaGetLastError()); + return; + } + + // Large-B path: one CTA per block. + int threads = ((B + 31) / 32) * 32; +#define LAUNCH_FACTOR(BVAL) \ + case BVAL: \ + ExtractAndFactorBlockDiagonalsKernel \ + <<>>(row_off, col_idx, vals, \ + row_start, num_blocks, \ + factor_offset, pivot_floor, \ + factors); \ + break + + switch (B) { + LAUNCH_FACTOR(7); + LAUNCH_FACTOR(8); + LAUNCH_FACTOR(15); + LAUNCH_FACTOR(16); + default: { + size_t shared_bytes = static_cast(B) * B * sizeof(float); + ExtractAndFactorGenericKernel<<>>( + row_off, col_idx, vals, B, row_start, num_blocks, factor_offset, + pivot_floor, factors); break; - default: - throw std::invalid_argument( - "BlockSparsePCGSolver: unsupported block_size (must be 1,2,3,6,7)"); } + } +#undef LAUNCH_FACTOR THROW_ON_CUDA_ERROR(cudaGetLastError()); } -void LaunchApplyPrecond(cudaStream_t stream, int B, const float *factors, - const float *r, float *z, int num_blocks) { +/** Picks the right ApplyBlockJacobi specialization for B. + * + * For small B (≤ 6) we use a "one block per thread" path which moves + * the per-tile factor + residual into thread-local registers and + * packs ~256 tiles per CTA. This minimizes launch / scheduling + * overhead for problems with millions of small tiles (e.g. SBA + * landmarks). For larger B we fall back to the per-CTA path that + * stages the factor through shared memory. */ +void DispatchApplyPrecond(cudaStream_t stream, int B, int row_start, + int num_blocks, int factor_offset, + const float *factors, const float *r, float *z) { + if (num_blocks == 0) { + return; + } if (B == 1) { - int n = num_blocks; int threads = 256; - int blocks = (n + threads - 1) / threads; - ApplyScalarJacobi<<>>(factors, r, z, n); + int blocks = (num_blocks + threads - 1) / threads; + ApplyScalarJacobi<<>>(factors, factor_offset, r, + row_start, num_blocks, z); + THROW_ON_CUDA_ERROR(cudaGetLastError()); return; } + +#define LAUNCH_APPLY_PER_THREAD(BVAL) \ + case BVAL: { \ + int threads = 256; \ + int blocks = (num_blocks + threads - 1) / threads; \ + ApplyBlockJacobiPerThreadKernel \ + <<>>(factors, factor_offset, r, \ + row_start, num_blocks, z); \ + break; \ + } + + // Small-B path: one block per thread, lots of blocks per CTA. + switch (B) { + LAUNCH_APPLY_PER_THREAD(2); + LAUNCH_APPLY_PER_THREAD(3); + LAUNCH_APPLY_PER_THREAD(4); + LAUNCH_APPLY_PER_THREAD(5); + LAUNCH_APPLY_PER_THREAD(6); + default: + break; // fall through to per-CTA path below + } +#undef LAUNCH_APPLY_PER_THREAD + if (B >= 2 && B <= 6) { + THROW_ON_CUDA_ERROR(cudaGetLastError()); + return; + } + + // Large-B path: one CTA per block, factor staged through shared mem. int threads = ((B + 31) / 32) * 32; +#define LAUNCH_APPLY(BVAL) \ + case BVAL: \ + ApplyBlockJacobiKernel<<>>( \ + factors, factor_offset, r, row_start, num_blocks, z); \ + break + switch (B) { - case 2: - ApplyBlockJacobiPreconditioner<2> - <<>>(factors, r, z, num_blocks); + LAUNCH_APPLY(7); + LAUNCH_APPLY(8); + LAUNCH_APPLY(15); + LAUNCH_APPLY(16); + default: { + size_t shared_bytes = (static_cast(B) * B + B) * sizeof(float); + ApplyBlockJacobiGenericKernel<<>>(factors, factor_offset, B, r, + row_start, num_blocks, z); break; - case 3: - ApplyBlockJacobiPreconditioner<3> - <<>>(factors, r, z, num_blocks); - break; - case 6: - ApplyBlockJacobiPreconditioner<6> - <<>>(factors, r, z, num_blocks); - break; - case 7: - ApplyBlockJacobiPreconditioner<7> - <<>>(factors, r, z, num_blocks); - break; - default: - throw std::invalid_argument("BlockSparsePCGSolver: unsupported block_size"); } + } +#undef LAUNCH_APPLY THROW_ON_CUDA_ERROR(cudaGetLastError()); } +/** Enqueues `out[0] = ` (clear-then-reduce). */ void DotAsync(cudaStream_t stream, const float *a, const float *b, int n, - float *d_out) { - DotKernelZero<<<1, 1, 0, stream>>>(d_out); + float *out) { + ZeroScalarKernel<<<1, 1, 0, stream>>>(out); + int threads = 256; + int blocks = std::min(1024, (n + threads - 1) / threads); + DotKernel<<>>(a, b, n, out); +} + +/** Enqueues `out_ab = ` and `out_ac = ` in a single pass. + * Equivalent to two @ref DotAsync calls but with half the global reads. */ +void DualDotAsync(cudaStream_t stream, const float *a, const float *b, + const float *c, int n, float *out_ab, float *out_ac) { + Zero2ScalarKernel<<<1, 1, 0, stream>>>(out_ab, out_ac); int threads = 256; int blocks = std::min(1024, (n + threads - 1) / threads); - DotKernel<<>>(a, b, n, d_out); + DualDotKernel<<>>(a, b, c, n, out_ab, out_ac); } } // namespace -// ----------------------------------------------------------------------------- +// ============================================================================= // BlockSparsePCGSolver -// ----------------------------------------------------------------------------- +// ============================================================================= BlockSparsePCGSolver::BlockSparsePCGSolver(BlockSparsePCGOptions options) - : options_(options) { + : options_(std::move(options)) { if (options_.block_size < 1 || options_.block_size > kMaxBlockSize) { throw std::invalid_argument( "BlockSparsePCGSolver: block_size must be in [1, 16]"); } + for (const auto &seg : options_.block_layout) { + if (seg.second < 1 || seg.second > kMaxBlockSize) { + throw std::invalid_argument( + "BlockSparsePCGSolver: block_layout segment size must be in " + "[1, 16]"); + } + } + if (options_.check_period < 1) { + options_.check_period = 1; + } } BlockSparsePCGSolver::~BlockSparsePCGSolver() = default; +bool BlockSparsePCGSolver::BuildSegmentTables(int matrix_dim) { + segments_.clear(); + // Normalize the layout: use the user-supplied block_layout if it's + // non-empty; otherwise treat the whole matrix as a single segment of + // size `options_.block_size`. + layout_.clear(); + if (!options_.block_layout.empty()) { + layout_ = options_.block_layout; + } else { + if (matrix_dim % options_.block_size != 0) { + LogError("BlockSparsePCGSolver: matrix dim {} not divisible by uniform " + "block_size {}", + matrix_dim, options_.block_size); + return false; + } + layout_.push_back({matrix_dim / options_.block_size, options_.block_size}); + } + + int row_cursor = 0; + int block_cursor = 0; + size_t factor_cursor = 0; + for (const auto &[count, size] : layout_) { + if (count <= 0) { + continue; + } + if (size < 1 || size > kMaxBlockSize) { + LogError("BlockSparsePCGSolver: invalid segment size {}", size); + return false; + } + Segment s; + s.block_size = size; + s.num_blocks = count; + s.row_start = row_cursor; + s.factor_offset = static_cast(factor_cursor); + s.block_row_start = block_cursor; + segments_.push_back(s); + + row_cursor += count * size; + block_cursor += count; + factor_cursor += static_cast(count) * + (size == 1 ? 1ull + : static_cast(size) * size); + } + if (row_cursor != matrix_dim) { + LogError("BlockSparsePCGSolver: layout covers {} rows but matrix has {}", + row_cursor, matrix_dim); + return false; + } + total_blocks_ = block_cursor; + total_factor_floats_ = factor_cursor; + return true; +} + bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, + const Problem &problem, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { @@ -450,23 +1000,48 @@ bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, rhs.size(), result.size()); return false; } - int B = options_.block_size; - if (n % B != 0) { - LogError("BlockSparsePCGSolver: matrix size {} not divisible by block " - "size {}", - n, B); + // When @p problem carries state batches, the block-Jacobi layout is + // ALWAYS derived afresh from them — each non-empty batch contributes + // a segment with `size = TangentSize()` and + // `count = NumStateBlocks() - NumConstStateBlocks()`. Consecutive + // segments of equal size are merged so the dispatch loop has one + // entry per distinct-size group (typical: 1 entry for PGO, 2 for + // SBA). Re-deriving on every Initialize is what makes the same + // solver instance work across a stream of differently-dimensioned + // problems (e.g. the binary SBA test fixture iterates over many + // such problems with one minimizer). + // + // Falls back to the user-supplied @c options_.block_layout (or the + // uniform @c options_.block_size) only when @p problem has no + // registered state batches — the path tests exercise when they call + // the solver directly on a raw matrix. + std::vector> derived_layout; + for (const auto *sb : problem.GetStateBatches()) { + if (sb == nullptr) { + continue; + } + int t = static_cast(sb->TangentSize()); + int count = + static_cast(sb->NumStateBlocks() - sb->NumConstStateBlocks()); + if (count == 0) { + continue; + } + if (!derived_layout.empty() && derived_layout.back().second == t) { + derived_layout.back().first += count; + } else { + derived_layout.emplace_back(count, t); + } + } + if (!derived_layout.empty()) { + options_.block_layout = std::move(derived_layout); + } + if (!BuildSegmentTables(n)) { return false; } matrix_size_ = n; - num_blocks_ = n / B; - size_t factors_size = - static_cast(num_blocks_) * static_cast(B * B); - if (B == 1) { - factors_size = static_cast(n); - } - if (precond_factors_.size() < factors_size) { - precond_factors_.resize(factors_size); + if (precond_factors_.size() < total_factor_floats_) { + precond_factors_.resize(total_factor_floats_); } if (r_.size() < static_cast(n)) { r_.resize(n); @@ -474,16 +1049,14 @@ bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, p_.resize(n); Ap_.resize(n); } - if (d_scratch_.size() < 2) { - d_scratch_.resize(2); + if (d_scratch_.size() < 7) { + d_scratch_.resize(7); } - // Build the cuSPARSE descriptor for SpMV. Use raw cusparseDnVec with - // explicit size = n so capacity-vs-size drift in p_/Ap_ across successive - // Solves on differently-sized matrices doesn't break the SpMV preprocess. - mat_desc_ = - cuSPARSEMatrixDescription(n, n, static_cast(spd_matrix.NumNonZeros()), - spd_matrix); + // cuSPARSE SpMV setup. The descriptor is reused across all PCG steps + // and across all Solve calls until the matrix structure changes. + mat_desc_ = cuSPARSEMatrixDescription( + n, n, static_cast(spd_matrix.NumNonZeros()), spd_matrix); auto handle = static_cast(cusparse_handle_.GetHandle(stream)); @@ -509,7 +1082,6 @@ bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, spmv_buffer_.data())); WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecX)); WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecY)); - spmv_buffer_ready_ = true; return true; } @@ -529,42 +1101,55 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, return true; } if (n != matrix_size_) { - if (!Initialize(stream, spd_matrix, rhs, result)) { + // Recovery path: matrix dim changed since the last Initialize. Use + // a default-constructed Problem; the cached options_.block_layout + // (set by the prior Initialize) is reused if non-empty, otherwise + // the uniform options_.block_size path takes over. + Problem empty_problem; + if (!Initialize(stream, empty_problem, spd_matrix, rhs, result)) { return false; } } - int B = options_.block_size; - // Refresh the cuSPARSE descriptor's value pointer (structure is fixed across - // calls; the CSR matrix can move in memory between Solves). + // Refresh the cuSPARSE descriptor's value pointer; the matrix's + // structure is unchanged so no re-preprocess is needed. mat_desc_.UpdatePointers(spd_matrix); - // Refresh preconditioner from the current matrix values. - LaunchExtractAndFactor(stream, B, spd_matrix, num_blocks_, - options_.pivot_floor, precond_factors_); + // ----------------------------------------------------------------- + // 1. Rebuild the block-Jacobi preconditioner from current H values. + // ----------------------------------------------------------------- + for (const auto &s : segments_) { + DispatchExtractAndFactor(stream, s.block_size, s.row_start, s.num_blocks, + s.factor_offset, options_.pivot_floor, spd_matrix, + precond_factors_.data()); + } auto handle = static_cast(cusparse_handle_.GetHandle(stream)); auto matA = static_cast(mat_desc_.GetDescription()); - // PCG with zero initial guess. The Gauss-Newton step is reset every outer - // iteration so warm-starting from the prior step would seed PCG far from the - // new solution; zeroing is both faster and more robust here. + // ----------------------------------------------------------------- + // 2. Initialize PCG with x_0 = 0 ⇒ r_0 = b. + // ----------------------------------------------------------------- THROW_ON_CUDA_ERROR( cudaMemsetAsync(result.data(), 0, n * sizeof(float), stream)); THROW_ON_CUDA_ERROR(cudaMemcpyAsync(r_.data(), rhs.data(), n * sizeof(float), cudaMemcpyDeviceToDevice, stream)); - // z = M^{-1} r - LaunchApplyPrecond(stream, B, precond_factors_.data(), r_.data(), z_.data(), - num_blocks_); - // p = z + // z_0 = M^{-1} r_0, p_0 = z_0. + for (const auto &s : segments_) { + DispatchApplyPrecond(stream, s.block_size, s.row_start, s.num_blocks, + s.factor_offset, precond_factors_.data(), r_.data(), + z_.data()); + } THROW_ON_CUDA_ERROR(cudaMemcpyAsync(p_.data(), z_.data(), n * sizeof(float), cudaMemcpyDeviceToDevice, stream)); - // Scratch layout (all device-resident, never copied to host during the inner - // loop): [0] alpha, [1] beta, [2] pAp, [3] rz_old, [4] rz_new, [5] r_norm2, - // [6] b_norm2. + // ----------------------------------------------------------------- + // 3. Compute initial inner products on device. + // ----------------------------------------------------------------- + // Scratch slots: [0] alpha, [1] beta, [2] , [3] rz_old (= ), + // [4] rz_new, [5] ||r||^2, [6] ||b||^2. enum : int { kAlpha = 0, kBeta = 1, @@ -572,19 +1157,15 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, kRzOld = 3, kRzNew = 4, kRnorm2 = 5, - kBnorm2 = 6, - kScalarCount = 7 + kBnorm2 = 6 }; - if (d_scratch_.size() < kScalarCount) { - d_scratch_.resize(kScalarCount); - } - // rz_old = ; b_norm2 = DotAsync(stream, r_.data(), z_.data(), n, d_scratch_.data() + kRzOld); DotAsync(stream, rhs.data(), rhs.data(), n, d_scratch_.data() + kBnorm2); - // Pull b_norm2 once for the host-side stopping threshold (the dot still - // runs async; the sync below is paid only once, not per iteration). + // Pull ||b||^2 once for the host-side convergence threshold. This is + // the ONLY guaranteed host sync; the rest of the loop polls residual + // norm only every `check_period` iterations. float b_norm2 = 0.f; THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&b_norm2, d_scratch_.data() + kBnorm2, sizeof(float), cudaMemcpyDeviceToHost, @@ -594,10 +1175,12 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, float rel_tol2 = options_.relative_tolerance * options_.relative_tolerance; float stop_thresh = fmaxf(abs_tol2, rel_tol2 * fmaxf(b_norm2, 1e-30f)); - // Build dense-vector descriptors with explicit size = n. dvector::resize - // is grow-only, so p_/Ap_ may have capacity > n across successive Solves on - // problems with different dimensions; using p_.size() in the descriptor - // would mismatch matA's row count. + // ----------------------------------------------------------------- + // 4. Per-iteration loop. + // ----------------------------------------------------------------- + // p_/Ap_ may have capacity > n across successive Solves on differently + // sized matrices, so we construct fresh dense-vector descriptors of + // explicit size n each Solve. Cheap host calls. cusparseDnVecDescr_t vecX = nullptr; cusparseDnVecDescr_t vecY = nullptr; THROW_ON_CUSPARSE_ERROR( @@ -605,59 +1188,56 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, THROW_ON_CUSPARSE_ERROR( cusparseCreateDnVec(&vecY, n, Ap_.data(), CUDA_R_32F)); - int threads = 256; - int blocks = (n + threads - 1) / threads; - - // Convergence is polled every kCheckPeriod iterations. Polling more often - // turns the inner loop back into a sequence of host syncs; polling less - // often risks doing extra work after the iteration has already converged. - // 2 is a safe default — small PGO systems converge in a handful of iterates - // and don't tolerate a long fixed period, while SBA still benefits from - // the avoided alpha/beta-on-host syncs that the rest of the inner loop now - // sidesteps. - constexpr int kCheckPeriod = 2; + const int threads = 256; + const int blocks = (n + threads - 1) / threads; + const int check_period = options_.check_period; + int it = 0; for (; it < options_.max_iterations; ++it) { - // Ap = A * p + // Ap = H * p. float spmv_alpha = 1.f; float spmv_beta = 0.f; - THROW_ON_CUSPARSE_ERROR(cusparseSpMV(handle, - CUSPARSE_OPERATION_NON_TRANSPOSE, - &spmv_alpha, matA, vecX, &spmv_beta, - vecY, CUDA_R_32F, - CUSPARSE_SPMV_ALG_DEFAULT, - spmv_buffer_.data())); - - // pAp = + THROW_ON_CUSPARSE_ERROR(cusparseSpMV( + handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &spmv_alpha, matA, vecX, + &spmv_beta, vecY, CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, + spmv_buffer_.data())); + + // and alpha = rz_old / . DotAsync(stream, p_.data(), Ap_.data(), n, d_scratch_.data() + kPAp); - // alpha = rz_old / pAp on device. ComputeAlphaKernel<<<1, 1, 0, stream>>>(d_scratch_.data() + kRzOld, d_scratch_.data() + kPAp, d_scratch_.data() + kAlpha); + + // x ← x + α p; r ← r − α Ap. PcgUpdateKernel<<>>( d_scratch_.data() + kAlpha, p_.data(), Ap_.data(), result.data(), r_.data(), n); - // z = M^{-1} r ; rz_new = ; r_norm2 = - LaunchApplyPrecond(stream, B, precond_factors_.data(), r_.data(), z_.data(), - num_blocks_); - DotAsync(stream, r_.data(), r_.data(), n, d_scratch_.data() + kRnorm2); - DotAsync(stream, r_.data(), z_.data(), n, d_scratch_.data() + kRzNew); + // z = M^{-1} r. + for (const auto &s : segments_) { + DispatchApplyPrecond(stream, s.block_size, s.row_start, s.num_blocks, + s.factor_offset, precond_factors_.data(), r_.data(), + z_.data()); + } + // Fused: rz_new = , ||r||^2 = . Single pass over r. + DualDotAsync(stream, r_.data(), z_.data(), r_.data(), n, + d_scratch_.data() + kRzNew, d_scratch_.data() + kRnorm2); - // beta = rz_new / rz_old, then rz_old <- rz_new (single device kernel). + // beta = rz_new / rz_old; rz_old ← rz_new. ComputeBetaKernel<<<1, 1, 0, stream>>>(d_scratch_.data() + kRzNew, d_scratch_.data() + kRzOld, d_scratch_.data() + kBeta); + + // p ← z + β p. PcgDirectionKernel<<>>( d_scratch_.data() + kBeta, z_.data(), p_.data(), n); - // Convergence poll every kCheckPeriod iterations: one D2H of one float - // plus a stream sync. Cheap relative to a full inner-loop sync. - if (((it + 1) % kCheckPeriod) == 0 || + // Convergence poll every `check_period` iters: one D2H of one float + // + one stream sync. Far cheaper than a host sync per iteration. + if (((it + 1) % check_period) == 0 || (it + 1) == options_.max_iterations) { float r_norm2 = 0.f; - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&r_norm2, - d_scratch_.data() + kRnorm2, + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&r_norm2, d_scratch_.data() + kRnorm2, sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h index b8aadce..c27ed31 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.h +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -19,97 +19,296 @@ #include +#include +#include + #include "cunls/common/cusparse_helper.h" #include "cunls/linear_solver/csr_sparse_linear_solver.h" namespace cunls { /** - * @brief Configuration options for the block-sparse PCG solver. + * @brief Configuration options for @ref BlockSparsePCGSolver. + * + * Two block-Jacobi layouts are supported: + * - **Uniform** (default): every diagonal tile is @ref block_size × @ref + * block_size. Requires @c block_size to divide the matrix dimension. + * - **Variable** (preferred for SBA-like problems): @ref block_layout + * describes a sequence of contiguous segments + * `[(count_0, size_0), (count_1, size_1), ...]`, each contributing + * `count_i * size_i` rows to the matrix. When non-empty, + * @c block_layout overrides @c block_size. Use this when the matrix + * contains tiles of different sizes (e.g. SBA: 6×6 pose tiles followed + * by 3×3 landmark tiles). * - * The solver implements Preconditioned Conjugate Gradient with block-Jacobi - * preconditioning on a symmetric-positive-definite CSR matrix. Block size is - * inferred from the matrix structure: the user supplies a fixed block dimension - * B and the diagonal B x B tiles of the matrix are factored once per outer - * call to ``Solve`` and applied at every PCG iteration. + * All other fields are convergence / numerical knobs for the inner + * Preconditioned Conjugate Gradient loop. See @ref BlockSparsePCGSolver + * for the math. */ struct BlockSparsePCGOptions { - /** Block dimension used by the block-Jacobi preconditioner. - * Must divide the matrix size. B = 1 falls back to scalar Jacobi. */ + /** + * @brief Uniform diagonal tile size (rows = cols), used when + * @ref block_layout is empty. + * + * Must divide the matrix dimension. `block_size = 1` reduces the + * preconditioner to scalar Jacobi (M = diag(H)^{-1}). + */ int block_size = 6; - /** Maximum number of PCG iterations. */ + /** + * @brief Heterogeneous block-diagonal layout for the preconditioner. + * + * When non-empty, each element `(count_i, size_i)` describes a + * contiguous segment of `count_i` consecutive diagonal blocks of + * dimension `size_i × size_i`. Segments are concatenated in order: + * segment 0 covers rows `[0, count_0 * size_0)`, segment 1 covers + * the next `count_1 * size_1` rows, etc. The sum + * `Σ count_i * size_i` must equal the matrix dimension. + * + * For a Gauss-Newton Hessian, the natural layout follows the order of + * @ref Problem state batches with the tangent dimension of each batch + * as the segment block size. @ref GaussNewtonMinimizer derives this + * layout automatically from the problem when the active solver is a + * @ref BlockSparsePCGSolver. + */ + std::vector> block_layout; + + /** @brief Maximum number of PCG iterations per @ref Solve call. */ int max_iterations = 200; - /** Convergence threshold on the relative residual ||r_k|| / ||b||. */ + /** + * @brief Relative-residual convergence threshold. + * + * The iteration stops when `||r_k|| <= relative_tolerance * ||b||`. + * The squared residual is polled on the GPU every + * `check_period` iterations to avoid one host sync per step. + */ float relative_tolerance = 1e-3f; - /** Absolute threshold on ||r_k|| for early exit. */ + /** + * @brief Absolute-residual convergence threshold. + * + * The iteration also stops when `||r_k|| <= absolute_tolerance`. + * The effective stopping rule is + * `||r||^2 <= max(absolute_tolerance^2, relative_tolerance^2 * ||b||^2)`. + */ float absolute_tolerance = 1e-30f; - /** Floor added to LDLT diagonal pivots to keep the preconditioner - * numerically invertible when the diagonal block is near-singular. */ + /** + * @brief Floor on |D_{kk}| during the per-block LDLT pivoting. + * + * Any pivot smaller in magnitude is replaced by @c pivot_floor with + * the original sign. Keeps the preconditioner invertible on + * near-singular diagonal tiles (e.g. an unconstrained landmark); + * does not affect correctness of the outer iteration because the + * preconditioner is just a convergence accelerator. + */ float pivot_floor = 1e-12f; + + /** + * @brief Number of PCG iterations between host-side convergence + * polls. Higher values reduce host syncs but may do up to + * `check_period - 1` extra iterations after convergence. + */ + int check_period = 4; }; /** - * @brief Block-Jacobi preconditioned conjugate gradient solver. + * @brief Block-Jacobi preconditioned conjugate gradient solver for + * symmetric positive (semi-)definite CSR systems. + * + * Solves `H x = b` with the standard PCG recurrence + * (Saad, *Iterative Methods for Sparse Linear Systems*, 2nd ed., §9.2): + * @code{.unparsed} + * r_0 = b - H x_0, z_0 = M^{-1} r_0, p_0 = z_0, rz_0 = + * for k = 0, 1, ... + * q_k = H p_k + * alpha_k = rz_k / + * x_{k+1} = x_k + alpha_k p_k + * r_{k+1} = r_k - alpha_k q_k + * z_{k+1} = M^{-1} r_{k+1} + * rz_{k+1}= + * beta_k = rz_{k+1} / rz_k + * p_{k+1} = z_{k+1} + beta_k p_k + * @endcode + * where `M` is the block-Jacobi preconditioner formed from the dense + * diagonal tiles of `H`. Each diagonal tile is factored independently + * via in-shared-memory LDLT (`H_d = L D L^T`); applying `M^{-1}` is a + * batched triangular solve plus diagonal scaling. * - * Solves H x = b for symmetric positive (semi-)definite H stored in CSR - * format. The preconditioner consists of the dense ``B x B`` diagonal tiles of - * H, factored independently with one warp per block using a small in-register - * LDLT. SpMV is delegated to cuSPARSE (CSR Hermitian SpMV). + * The CG search direction satisfies ` = 0` for `i != j`, which + * makes alpha and beta computable in closed form from the recurrence's + * own inner products and removes the host's reorientation work that + * dominates many iterative solver implementations. * - * The implementation is intended for normal equations from Gauss-Newton / - * Levenberg-Marquardt where: - * - H has natural block structure aligned with state blocks (e.g. 6 for SE3, - * 3 for Vector<3>, ...); - * - the matrix structure does not change between ``Solve`` calls inside one - * Minimize, so allocations and the cuSPARSE descriptor are reused; - * - the previous step is a good warm start for the next iteration. + * Implementation notes (see the `.cu` file for derivations): + * - SpMV is delegated to cuSPARSE (`cusparseSpMV` with the default + * algorithm and an up-front `preprocess` pass) — the same matrix + * structure is reused across all PCG steps of one @ref Solve and + * typically across multiple @ref Solve calls inside a single + * Levenberg-Marquardt @ref Minimize. + * - All scalar quantities (`alpha`, `beta`, ``, ``, + * ``) live on the device for the whole inner loop. Only the + * residual norm is copied back to the host, and only once every + * `check_period` iterations (default 4). + * - The block-Jacobi factor `M` is rebuilt from the current `H` values + * on every @ref Solve (cheap: one CTA per tile reads a few floats + * and does a small LDLT entirely in shared memory). */ class BlockSparsePCGSolver : public CSRSparseLinearSolver { public: + /** + * @brief Constructs the solver with the given options. + * + * Validates the uniform `block_size` (must be in `[1, 16]`) and the + * `block_layout` per-segment sizes (same range). Allocation of + * device buffers is deferred to @ref Initialize. + * + * @param options Convergence and preconditioner-layout knobs. + */ explicit BlockSparsePCGSolver(BlockSparsePCGOptions options = {}); + + /** @brief Releases owned resources. */ ~BlockSparsePCGSolver() override; - bool Initialize(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, + /** + * @brief Prepares working buffers and the cuSPARSE SpMV plan for a + * system of the given size. + * + * Performs symbolic setup only — no PCG steps run here. Must be + * called once before the first @ref Solve. Subsequent @ref Solve + * calls may invoke this implicitly when the matrix dimension changes. + * + * When @p problem has registered state batches, the block-Jacobi + * layout is derived automatically from + * `problem.GetStateBatches()`: each batch contributes one layout + * segment with `size = TangentSize()` and + * `count = NumStateBlocks() - NumConstStateBlocks()`. Consecutive + * segments of equal size are merged so the dispatch loop only sees + * distinct-size groups. An explicit layout previously set via + * @ref SetBlockLayout takes precedence; passing an empty problem + * (default-constructed) reverts to the uniform + * @ref BlockSparsePCGOptions::block_size. + * + * @param stream CUDA stream used for all device work. + * @param problem The originating optimization problem. Used to + * derive the block-Jacobi preconditioner layout + * when @ref BlockSparsePCGOptions::block_layout is + * empty and the caller hasn't explicitly invoked + * @ref SetBlockLayout. + * @param spd_matrix Coefficient matrix `H` in CSR format. Only its + * sparsity pattern is examined here; values are read + * on every @ref Solve. + * @param rhs Right-hand side vector `b`. Used only to validate + * dimensions; never read here. + * @param result Solution vector `x`. Used only to validate + * dimensions; never written here. + * @return true on success, false on dimension mismatch or invalid + * layout. + */ + bool Initialize(cudaStream_t stream, const Problem &problem, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; + /** + * @brief Runs the PCG loop on `H x = b`, writing into @p result. + * + * Recomputes the preconditioner from the current values of `H`, runs + * up to @ref BlockSparsePCGOptions::max_iterations iterations of PCG, + * and stops when the residual norm satisfies the tolerance rule + * documented on @ref BlockSparsePCGOptions::relative_tolerance. + * + * @param stream CUDA stream used for all device work. + * @param spd_matrix Coefficient matrix `H` in CSR format (same + * sparsity pattern as in @ref Initialize, but + * values are read fresh on every call). + * @param rhs Right-hand side `b`. + * @param result Output vector `x`. Caller-allocated. The + * solver initializes `x_0 = 0` (a zero warm start is + * the right choice for Gauss-Newton's per-step + * reset). + * @return true on success, false on dimension mismatch. + */ bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; - /** @brief Number of PCG iterations consumed by the most recent ``Solve``. */ + /** + * @brief Number of PCG iterations consumed by the most recent + * @ref Solve call. Useful for convergence diagnostics. + */ int LastIterations() const { return last_iterations_; } private: - BlockSparsePCGOptions options_; + // ------------------------------------------------------------------ + // Layout helpers + // ------------------------------------------------------------------ - /** Number of rows in the system (cached on Initialize / refreshed on Solve). + /** + * @brief Populates @ref segment_* device arrays from the current + * layout (either @c options_.block_layout or a single + * uniform segment derived from @c options_.block_size). + * + * Computes per-segment offsets into the factor buffer and the + * required preconditioner-factor storage size. Called by + * @ref Initialize whenever the matrix dimension or layout changes. + * + * @param matrix_dim Total matrix size N. + * @return false if the layout doesn't tile N exactly, true otherwise. */ + bool BuildSegmentTables(int matrix_dim); + + BlockSparsePCGOptions options_; + /** Cached host copy of `options_.block_layout`, normalized to a + * single uniform segment when the user didn't supply one. */ + std::vector> layout_; + + /** Total matrix dimension N (cached from the last @ref Initialize). */ int matrix_size_ = 0; - /** Number of block rows = matrix_size_ / block_size. */ - int num_blocks_ = 0; + /** Total number of diagonal tiles across all segments. */ + int total_blocks_ = 0; + /** Total entries in @ref precond_factors_. */ + size_t total_factor_floats_ = 0; - /** Stored LDLT factors of every B x B diagonal tile, packed contiguously - * (block i begins at index i * B * B; lower-triangle is L, diagonal is D). */ + /** Per-segment data, kept on the host for the launch loop. */ + struct Segment { + int block_size; ///< side length of each tile (rows = cols) + int num_blocks; ///< number of tiles in this segment + int row_start; ///< first matrix row covered by this segment + int factor_offset; ///< first index in @ref precond_factors_ + int block_row_start; ///< first block index in the global tile order + }; + std::vector segments_; + + /** + * @brief LDLT factors of every diagonal tile, packed contiguously. + * + * Tile @c b of segment @c s lives at `precond_factors_[segment_factor_offset(s) + + * b * size_s * size_s ..]` in row-major order. Lower triangle holds + * `L` with unit diagonal (implicit); the stored diagonal holds `D`; + * upper triangle is unused. + */ dvector precond_factors_; - /** Scratch device vectors used by PCG. */ - dvector r_; - dvector z_; - dvector p_; - dvector Ap_; + // ------------------------------------------------------------------ + // PCG scratch + // ------------------------------------------------------------------ + dvector r_; ///< residual `r_k` + dvector z_; ///< preconditioned residual `z_k = M^{-1} r_k` + dvector p_; ///< search direction `p_k` + dvector Ap_; ///< `H p_k` (the SpMV output) - /** Two-slot device scalar buffer used for fused reductions. */ + /** Device-resident scalar slots: alpha, beta, , rz_old, rz_new, + * ||r||^2, ||b||^2. Layout is fixed in the .cu file. */ dvector d_scratch_; - /** Persistent SpMV state. */ + // ------------------------------------------------------------------ + // cuSPARSE SpMV state + // ------------------------------------------------------------------ cuSPARSEHandle cusparse_handle_; cuSPARSEMatrixDescription mat_desc_; - dvector spmv_buffer_; - bool spmv_buffer_ready_ = false; + dvector spmv_buffer_; ///< work buffer for cuSPARSE SpMV + /** Iteration count reported by the last @ref Solve. */ int last_iterations_ = 0; }; diff --git a/cunls/linear_solver/csr_sparse_linear_solver.h b/cunls/linear_solver/csr_sparse_linear_solver.h index af7cf72..e2963bb 100644 --- a/cunls/linear_solver/csr_sparse_linear_solver.h +++ b/cunls/linear_solver/csr_sparse_linear_solver.h @@ -23,13 +23,21 @@ namespace cunls { +class Problem; // forward declaration; defined in cunls/minimizer/problem.h. + /** * @brief Base class for linear solvers operating on CSR matrices. * * Provides a common interface for solving sparse symmetric linear systems * Ax = b where the matrix A is stored in CSR (Compressed Sparse Row) format. * Derived classes implement specific solver strategies (e.g. cuDSS direct - * factorization, dense pivoted LDLT). + * factorization, dense pivoted LDLT, block-Jacobi PCG). + * + * Initialize receives the originating @ref Problem so solvers can adapt + * to its block / factor-graph structure (e.g. + * @ref BlockSparsePCGSolver reads each state batch's @c TangentSize to + * build its block-Jacobi preconditioner without a downcast at the call + * site). Solvers that don't care can simply ignore the argument. */ class CSRSparseLinearSolver { public: @@ -44,13 +52,19 @@ class CSRSparseLinearSolver { * elements as the number of rows in @p spd_matrix; the solver does not * resize them. * - * @param stream CUDA stream for asynchronous GPU operations. + * @param stream CUDA stream for asynchronous GPU operations. + * @param problem The originating optimization problem. Solvers may use + * its state-batch structure to specialize their setup + * (e.g. derive a block-Jacobi preconditioner layout). + * Pass a default-constructed @ref Problem when calling + * on a raw matrix that wasn't produced by cuNLS factors. * @param spd_matrix The coefficient matrix A in CSR format. - * @param rhs The right-hand side vector b (size must equal matrix rows). - * @param result Output vector x (size must equal matrix rows). + * @param rhs The right-hand side vector b (size must equal matrix + * rows). + * @param result Output vector x (size must equal matrix rows). * @return true on success, false if a dimension mismatch is detected. */ - virtual bool Initialize(cudaStream_t stream, + virtual bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) = 0; diff --git a/cunls/linear_solver/cudss_sparse_linear_solver.cpp b/cunls/linear_solver/cudss_sparse_linear_solver.cpp index 1d2233e..2d87d13 100644 --- a/cunls/linear_solver/cudss_sparse_linear_solver.cpp +++ b/cunls/linear_solver/cudss_sparse_linear_solver.cpp @@ -68,6 +68,7 @@ cuDSSLinearSolver::cuDSSLinearSolver(cuDSSLinearSolverOptions options) /** @copydoc cuDSSLinearSolver::Initialize */ bool cuDSSLinearSolver::Initialize(cudaStream_t stream, + const Problem & /*problem*/, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { diff --git a/cunls/linear_solver/cudss_sparse_linear_solver.h b/cunls/linear_solver/cudss_sparse_linear_solver.h index c9389f3..26e2b47 100644 --- a/cunls/linear_solver/cudss_sparse_linear_solver.h +++ b/cunls/linear_solver/cudss_sparse_linear_solver.h @@ -92,7 +92,8 @@ class cuDSSLinearSolver : public CSRSparseLinearSolver { * @param result Output vector x (size must equal matrix rows). * @return true on success, false if a dimension mismatch is detected. */ - bool Initialize(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, + bool Initialize(cudaStream_t stream, const Problem &problem, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; /** diff --git a/cunls/linear_solver/dense_cholesky_solver.cu b/cunls/linear_solver/dense_cholesky_solver.cu index a140424..87d2fbb 100644 --- a/cunls/linear_solver/dense_cholesky_solver.cu +++ b/cunls/linear_solver/dense_cholesky_solver.cu @@ -53,6 +53,7 @@ __global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, } // namespace bool DenseCholeskySolver::Initialize(cudaStream_t stream, + const Problem & /*problem*/, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { diff --git a/cunls/linear_solver/dense_cholesky_solver.h b/cunls/linear_solver/dense_cholesky_solver.h index 90ce24b..6a34afd 100644 --- a/cunls/linear_solver/dense_cholesky_solver.h +++ b/cunls/linear_solver/dense_cholesky_solver.h @@ -51,8 +51,10 @@ class DenseCholeskySolver : public CSRSparseLinearSolver { * @param result Output vector x (size must equal matrix rows). * @return true on success, false if a dimension mismatch is detected. */ - bool Initialize(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + bool Initialize(cudaStream_t stream, const Problem &problem, + const CSRSparseMatrix &spd_matrix, + const dvector &rhs, + dvector &result) final; /** * @brief Converts CSR to dense and solves via Cholesky factorization. diff --git a/cunls/linear_solver/dense_linear_solver.cu b/cunls/linear_solver/dense_linear_solver.cu index da2901f..3cf2be5 100644 --- a/cunls/linear_solver/dense_linear_solver.cu +++ b/cunls/linear_solver/dense_linear_solver.cu @@ -442,6 +442,7 @@ void SolveFromPivotedLDLT(cudaStream_t stream, const float *ldlt_factor, // --------------------------------------------------------------------------- bool DenseLDLTSolver::Initialize(cudaStream_t stream, + const Problem & /*problem*/, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { diff --git a/cunls/linear_solver/dense_linear_solver.h b/cunls/linear_solver/dense_linear_solver.h index 854854e..796d5ec 100644 --- a/cunls/linear_solver/dense_linear_solver.h +++ b/cunls/linear_solver/dense_linear_solver.h @@ -56,8 +56,10 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { * @param result Output vector x (size must equal matrix rows). * @return true on success, false if a dimension mismatch is detected. */ - bool Initialize(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + bool Initialize(cudaStream_t stream, const Problem &problem, + const CSRSparseMatrix &spd_matrix, + const dvector &rhs, + dvector &result) final; /** * @brief Converts CSR to dense and solves via pivoted LDLT factorization. diff --git a/cunls/linear_solver/dense_qr_solver.cu b/cunls/linear_solver/dense_qr_solver.cu index 91774d5..557a572 100644 --- a/cunls/linear_solver/dense_qr_solver.cu +++ b/cunls/linear_solver/dense_qr_solver.cu @@ -68,6 +68,7 @@ __global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, } // namespace bool DenseQRSolver::Initialize(cudaStream_t stream, + const Problem & /*problem*/, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { diff --git a/cunls/linear_solver/dense_qr_solver.h b/cunls/linear_solver/dense_qr_solver.h index 70750dd..5616d43 100644 --- a/cunls/linear_solver/dense_qr_solver.h +++ b/cunls/linear_solver/dense_qr_solver.h @@ -53,8 +53,10 @@ class DenseQRSolver : public CSRSparseLinearSolver { * @param result Output vector x (size must equal matrix rows). * @return true on success, false if a dimension mismatch is detected. */ - bool Initialize(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + bool Initialize(cudaStream_t stream, const Problem &problem, + const CSRSparseMatrix &spd_matrix, + const dvector &rhs, + dvector &result) final; /** * @brief Converts CSR to dense and solves via QR factorization. diff --git a/cunls/linear_solver/llms.txt b/cunls/linear_solver/llms.txt index 2ba8c18..ee1abfe 100644 --- a/cunls/linear_solver/llms.txt +++ b/cunls/linear_solver/llms.txt @@ -15,17 +15,16 @@ Sparse linear system abstraction and implementation used by minimizers. - `cuDSSLinearSolverOptions` - `cuDSSLinearSolver` - `block_sparse_pcg_solver.h`: - - `BlockSparsePCGOptions` — `block_size`, `max_iterations`, `relative_tolerance`, - `absolute_tolerance`, `pivot_floor`. + - `BlockSparsePCGOptions` — `block_size` / `block_layout`, + `max_iterations`, `relative_tolerance`, `absolute_tolerance`, + `pivot_floor`. - `BlockSparsePCGSolver` — iterative SPD solver: cuSPARSE CSR SpMV + - block-Jacobi LDLT preconditioner (dense B x B diagonal tiles factored in - shared memory). Inner loop computes alpha/beta on device; convergence is - polled every 2 iterations to keep CPU launches ahead of the GPU. - Best fit: PGO (block_size = 6) and PGO-like normal equations. - Mixed-block problems (e.g. SBA with 6x6 pose + 3x3 landmark blocks) work - correctly with block_size = 3 but converge more slowly than cuDSS — - Schur-complement elimination of the landmarks would be the natural - follow-up. + block-Jacobi LDLT preconditioner (dense B x B diagonal tiles factored + in shared memory). Supports a uniform block size or a layout vector + of `(count, block_size)` segments matching the host state-batch + structure (e.g. `{{N_poses, 6}, {N_points, 3}}` for SBA). Inner loop + computes alpha/beta on device; convergence is polled every K + iterations to keep CPU launches ahead of the GPU. ## Expected system form diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index f3d9410..8b28074 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -534,13 +534,14 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, // Perform symbolic analysis on the requested CUDA stream. auto sa_range = profiler_domain_.CreateDomainRange("PerformSymbolicAnalysis"); - bool success = solver_->Initialize(stream, lhs_work_, rhs_work_, step_); + bool success = + solver_->Initialize(stream, problem, lhs_work_, rhs_work_, step_); if (!success) { std::string str = "Failed to initialize linear solver"; LogError(str); throw std::runtime_error(str); } - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + // THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); } // Main optimization loop @@ -555,7 +556,8 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, { auto solve_range = profiler_domain_.CreateDomainRange("LinearSolve"); - bool success = solver_->Solve(stream, lhs_work_, rhs_work_, step_); + bool success = + solver_->Solve(stream, lhs_work_, rhs_work_, step_); if (!success) { std::string str = "Failed to solve linear system"; LogError(str); diff --git a/cunls/minimizer/problem.cpp b/cunls/minimizer/problem.cpp index fddcb30..f15628f 100644 --- a/cunls/minimizer/problem.cpp +++ b/cunls/minimizer/problem.cpp @@ -16,6 +16,7 @@ */ #include "cunls/minimizer/problem.h" + #include "cunls/common/helper.h" #include "cunls/common/log.h" namespace cunls { @@ -165,7 +166,7 @@ bool Problem::CheckGraphConnectivity() const { /** * @brief Validates the complete problem structure. * - * Runs both input validation and graph connectivity checks. + * Runs input validation and graph-connectivity checks. * * @return True if the problem is well-formed, false otherwise. */ diff --git a/tests/dense_cholesky_solver_test.cpp b/tests/dense_cholesky_solver_test.cpp index 994ba4d..fcb07b6 100644 --- a/tests/dense_cholesky_solver_test.cpp +++ b/tests/dense_cholesky_solver_test.cpp @@ -166,7 +166,7 @@ TEST_F(DenseCholeskySolverTestFixture, SolveDenseSystemAcrossDifferentSizes) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -195,7 +195,7 @@ TEST_F(DenseCholeskySolverTestFixture, SolveReturnsFalseForZeroMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -219,7 +219,7 @@ TEST_F(DenseCholeskySolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -243,7 +243,7 @@ TEST_F(DenseCholeskySolverTestFixture, SolveReturnsFalseForIndefiniteMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -265,7 +265,7 @@ TEST_F(DenseCholeskySolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -287,7 +287,7 @@ TEST_F(DenseCholeskySolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); diff --git a/tests/dense_linear_solver_test.cpp b/tests/dense_linear_solver_test.cpp index b71b048..7484d8b 100644 --- a/tests/dense_linear_solver_test.cpp +++ b/tests/dense_linear_solver_test.cpp @@ -244,7 +244,7 @@ TEST_F(DenseLDLTSolverTestFixture, SolveDenseSystemAcrossDifferentSizes) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -286,7 +286,7 @@ TEST_F(DenseLDLTSolverTestFixture, SolveSymmetricIndefiniteSystem) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -399,7 +399,7 @@ TEST_F(DenseLDLTSolverTestFixture, SolveReturnsFalseForZeroMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -425,7 +425,7 @@ TEST_F(DenseLDLTSolverTestFixture, SolveReturnsFalseForRankDeficientMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -448,7 +448,7 @@ TEST_F(DenseLDLTSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -472,7 +472,7 @@ TEST_F(DenseLDLTSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); diff --git a/tests/dense_qr_solver_test.cpp b/tests/dense_qr_solver_test.cpp index 2689de6..e05672e 100644 --- a/tests/dense_qr_solver_test.cpp +++ b/tests/dense_qr_solver_test.cpp @@ -166,7 +166,7 @@ TEST_F(DenseQRSolverTestFixture, SolveDenseSystemAcrossDifferentSizes) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -199,7 +199,7 @@ TEST_F(DenseQRSolverTestFixture, SolveSymmetricIndefiniteSystem) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -227,7 +227,7 @@ TEST_F(DenseQRSolverTestFixture, SolveReturnsFalseForZeroMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -250,7 +250,7 @@ TEST_F(DenseQRSolverTestFixture, SolveReturnsFalseForRankDeficientMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -272,7 +272,7 @@ TEST_F(DenseQRSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -294,7 +294,7 @@ TEST_F(DenseQRSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), matrix, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); diff --git a/tests/pgo_minimizer_test.cpp b/tests/pgo_minimizer_test.cpp index c37d418..0c377f0 100644 --- a/tests/pgo_minimizer_test.cpp +++ b/tests/pgo_minimizer_test.cpp @@ -300,10 +300,9 @@ TEST_F(PgoMinimizerTestFixture, Optimize) { options.cost_tolerance = 1e-2f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options = { - test_utils::PCGBlockSizeFromEnv(6), - test_utils::PCGMaxIterFromEnv(200), - test_utils::PCGTolFromEnv(1e-3f)}; + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(200); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1000.0; diff --git a/tests/sba_minimizer_test.cpp b/tests/sba_minimizer_test.cpp index d04b7bf..2970e0f 100644 --- a/tests/sba_minimizer_test.cpp +++ b/tests/sba_minimizer_test.cpp @@ -377,10 +377,9 @@ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { // block_size=3 keeps the preconditioner cheap and well-conditioned for the // landmark blocks while still capturing useful structure inside the 6x6 // pose tiles (every 6x6 splits into a 2x2 grid of 3x3 sub-blocks). - options.sparse_linear_solver_config.block_sparse_pcg_options = { - test_utils::PCGBlockSizeFromEnv(3), - test_utils::PCGMaxIterFromEnv(400), - test_utils::PCGTolFromEnv(1e-3f)}; + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(3); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; diff --git a/tests/sparse_linear_solver_test.cpp b/tests/sparse_linear_solver_test.cpp index e963b8f..c111199 100644 --- a/tests/sparse_linear_solver_test.cpp +++ b/tests/sparse_linear_solver_test.cpp @@ -34,6 +34,7 @@ #include #include "cunls/common/cuda_stream.h" +#include "cunls/minimizer/problem.h" #include "cunls/common/helper.h" #include "cunls/common/profiler.h" #include "cunls/common/types.h" @@ -186,7 +187,7 @@ TEST(SparseLinearSolverTest, Solve) { cuDSSLinearSolver solver(cudss_solver_options); { profiler::ScopedRange range("Warm up"); - solver.Initialize(stream.GetStream(), input_matrix, rhs, result); + solver.Initialize(stream.GetStream(), Problem(), input_matrix, rhs, result); solver.Solve(stream.GetStream(), input_matrix, rhs, result); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); } @@ -310,7 +311,7 @@ TEST(SparseLinearSolverTest, BlockSparsePCGSolve) { opts.relative_tolerance = 1e-5f; opts.max_iterations = 500; BlockSparsePCGSolver solver(opts); - ASSERT_TRUE(solver.Initialize(stream.GetStream(), mat, rhs, result)); + ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), mat, rhs, result)); ASSERT_TRUE(solver.Solve(stream.GetStream(), mat, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); diff --git a/tests/synthetic_pgo_test.cpp b/tests/synthetic_pgo_test.cpp index d0c8b9f..29e2af1 100644 --- a/tests/synthetic_pgo_test.cpp +++ b/tests/synthetic_pgo_test.cpp @@ -263,10 +263,9 @@ TEST_F(SyntheticPGOTest, OptimizeConsecutiveBetweenConstraints) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options = { - test_utils::PCGBlockSizeFromEnv(6), - test_utils::PCGMaxIterFromEnv(400), - test_utils::PCGTolFromEnv(1e-4f)}; + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); // GaussNewtonMinimizer minimizer(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; @@ -359,10 +358,9 @@ TEST_F(SyntheticPGOTest, InformationBetweenFactorBatch) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options = { - test_utils::PCGBlockSizeFromEnv(6), - test_utils::PCGMaxIterFromEnv(400), - test_utils::PCGTolFromEnv(1e-4f)}; + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); // GaussNewtonMinimizer minimizer(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; @@ -440,10 +438,9 @@ TEST_F(SyntheticPGOTest, WeightedWrapsInformationBetweenFactorBatch) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options = { - test_utils::PCGBlockSizeFromEnv(6), - test_utils::PCGMaxIterFromEnv(400), - test_utils::PCGTolFromEnv(1e-4f)}; + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -519,10 +516,9 @@ TEST_F(SyntheticPGOTest, InformationWrapsWeightedBetweenFactorBatch) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options = { - test_utils::PCGBlockSizeFromEnv(6), - test_utils::PCGMaxIterFromEnv(400), - test_utils::PCGTolFromEnv(1e-4f)}; + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -542,4 +538,242 @@ TEST_F(SyntheticPGOTest, InformationWrapsWeightedBetweenFactorBatch) { stream.GetStream()); } +// ============================================================================= +// Loop-closure PGO benchmark +// ============================================================================= + +/** + * @brief Parameterized PGO with random loop closures. + * + * Builds a single set of @p n_poses SE(3) poses with sequential between + * factors (i, i+1) and adds @p n_lc random non-sequential loop-closure + * factors at randomly chosen pairs (i, j). All deltas are sampled so + * the ground-truth solution matches a small-jitter perturbation of the + * initial chain — i.e. the problem has a clear minimum. + * + * The first pose is held constant to fix the gauge. All test sizes + * pick LC counts in the user-specified 300-10000 range and grow Nposes + * accordingly. + */ +struct LcPgoParams { + int n_poses; + int n_lc; + const char *label; +}; + +inline std::ostream &operator<<(std::ostream &os, const LcPgoParams &p) { + return os << p.label; +} + +class LoopClosurePGOTest : public ::testing::TestWithParam { +protected: + /** + * @brief Generates a chain of poses by composing small random twists. + * The returned @p gt_poses is the ground truth; @p init_poses is + * a perturbed copy that the solver starts from. + */ + void GenerateChain(int n_poses, std::vector >_poses, + std::vector &init_poses) { + std::mt19937 rng(7); + std::uniform_real_distribution r(-0.1f, 0.1f); + std::uniform_real_distribution t(-0.4f, 0.4f); + std::normal_distribution jitter(0.f, 0.02f); + + // Build a chain of twists. + std::vector> twists(n_poses); + for (int i = 0; i < n_poses; ++i) { + twists[i] = {r(rng), r(rng), r(rng), t(rng), t(rng), t(rng)}; + } + dvector> twists_d(twists); + dvector seg_d(n_poses); + CudaStream stream; + ComputeExpSE3(stream.GetStream(), + reinterpret_cast(twists_d.data()), 6, 4, 16, + n_poses, reinterpret_cast(seg_d.data())); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + std::vector segs(n_poses); + seg_d.CopyToHost(segs.data(), n_poses); + + auto mul = [](const SE3Transform &a, + const SE3Transform &b) -> SE3Transform { + SE3Transform c{}; + for (int rr = 0; rr < 4; ++rr) { + for (int cc = 0; cc < 4; ++cc) { + float s = 0.f; + for (int k = 0; k < 4; ++k) { + s += a[rr * 4 + k] * b[k * 4 + cc]; + } + c[rr * 4 + cc] = s; + } + } + return c; + }; + + gt_poses.resize(n_poses); + gt_poses[0] = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; + for (int i = 1; i < n_poses; ++i) { + gt_poses[i] = mul(gt_poses[i - 1], segs[i]); + } + + // Initial guess: small per-pose noise. + init_poses = gt_poses; + std::vector> djitters(n_poses); + for (int i = 1; i < n_poses; ++i) { + for (int k = 0; k < 6; ++k) { + djitters[i][k] = jitter(rng); + } + } + dvector> djit_d(djitters); + dvector delta_d(n_poses); + ComputeExpSE3(stream.GetStream(), + reinterpret_cast(djit_d.data()), 6, 4, 16, + n_poses, reinterpret_cast(delta_d.data())); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + std::vector deltas(n_poses); + delta_d.CopyToHost(deltas.data(), n_poses); + for (int i = 1; i < n_poses; ++i) { + init_poses[i] = mul(gt_poses[i], deltas[i]); + } + } + + cuBLASHandle cublas_handle_; + profiler::Domain profiler_domain_ = + profiler::Domain("LoopClosurePGOTest"); +}; + +TEST_P(LoopClosurePGOTest, Optimize) { + auto p = GetParam(); + SCOPED_TRACE(std::string("LoopClosurePGOTest: ") + p.label); + + std::vector gt, init; + GenerateChain(p.n_poses, gt, init); + + // Sequential between factors (i, i+1). Delta = T_{i+1}^{-1} * T_i would + // make the residual zero at ground truth; cuNLS expects the user to + // supply that delta, but for a chain we know it equals the segment we + // sampled. Computing on device for exactness. + // Simpler: residual = Log(Delta * T_left^{-1} * T_right) per + // SE3BetweenFactorBatch. Setting Delta = T_left * T_right^{-1} = (T_i) + // * inv(T_{i+1}) makes the residual zero at GT. + auto inv_se3 = [](const SE3Transform &T) -> SE3Transform { + // [R t; 0 1]^-1 = [R^T -R^T t; 0 1] (row-major SE(3)). + SE3Transform inv{}; + inv[0] = T[0]; + inv[1] = T[4]; + inv[2] = T[8]; + inv[4] = T[1]; + inv[5] = T[5]; + inv[6] = T[9]; + inv[8] = T[2]; + inv[9] = T[6]; + inv[10] = T[10]; + inv[3] = -(inv[0] * T[3] + inv[1] * T[7] + inv[2] * T[11]); + inv[7] = -(inv[4] * T[3] + inv[5] * T[7] + inv[6] * T[11]); + inv[11] = -(inv[8] * T[3] + inv[9] * T[7] + inv[10] * T[11]); + inv[15] = 1.f; + return inv; + }; + auto mul = [](const SE3Transform &a, + const SE3Transform &b) -> SE3Transform { + SE3Transform c{}; + for (int rr = 0; rr < 4; ++rr) { + for (int cc = 0; cc < 4; ++cc) { + float s = 0.f; + for (int k = 0; k < 4; ++k) { + s += a[rr * 4 + k] * b[k * 4 + cc]; + } + c[rr * 4 + cc] = s; + } + } + return c; + }; + + // Sequential edges (i, i+1). + std::vector edges_left, edges_right; + std::vector deltas; + edges_left.reserve(p.n_poses - 1 + p.n_lc); + edges_right.reserve(p.n_poses - 1 + p.n_lc); + deltas.reserve(p.n_poses - 1 + p.n_lc); + for (int i = 0; i + 1 < p.n_poses; ++i) { + edges_left.push_back(i); + edges_right.push_back(i + 1); + deltas.push_back(mul(gt[i], inv_se3(gt[i + 1]))); + } + // Random loop closures (i, j) with |i - j| > 1. + std::mt19937 rng(11); + std::uniform_int_distribution pose_pick(0, p.n_poses - 1); + int added = 0; + int attempts = 0; + while (added < p.n_lc && attempts < p.n_lc * 20) { + ++attempts; + int i = pose_pick(rng); + int j = pose_pick(rng); + if (std::abs(i - j) <= 1) { + continue; + } + edges_left.push_back(i); + edges_right.push_back(j); + deltas.push_back(mul(gt[i], inv_se3(gt[j]))); + ++added; + } + ASSERT_EQ(added, p.n_lc) << "Could not allocate requested loop closures"; + + // Build problem. + dvector poses_d(init); + auto poses_ptr = reinterpret_cast(poses_d.data()); + std::vector const_ids = {0}; + dvector const_ids_d(const_ids); + SE3StateBatch pose_batch(cublas_handle_, poses_ptr, p.n_poses, + const_ids_d.data(), const_ids.size()); + + dvector deltas_d(deltas); + SE3BetweenFactorBatch between_batch(deltas_d.data(), deltas.size()); + + std::vector state_pointers; + state_pointers.reserve(deltas.size() * 2); + for (size_t e = 0; e < deltas.size(); ++e) { + state_pointers.push_back( + pose_batch.StateBlockDevicePtr(edges_left[e])); + state_pointers.push_back( + pose_batch.StateBlockDevicePtr(edges_right[e])); + } + + Problem problem; + problem.AddStateBatch(&pose_batch); + problem.AddFactorBatch(&between_batch, state_pointers); + ASSERT_TRUE(problem.CheckConsistency()); + + CudaStream stream; + MinimizerOptions options; + options.max_num_iterations = 15; + options.state_tolerance = 1e-10f; + options.cost_tolerance = 1e-6f; + options.disable_safety_checks = true; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); + LevenbergMarquardtMinimizerOptions lm_options; + lm_options.base_options = options; + lm_options.initial_lambda = 1e-3f; + LevenbergMarquardtMinimizer minimizer(lm_options); + + MinimizerSummary summary; + { + auto range = profiler_domain_.CreateDomainRange("Minimize"); + summary = minimizer.Minimize(stream.GetStream(), problem); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + } + EXPECT_TRUE(std::isfinite(summary.initial_cost)); + EXPECT_TRUE(std::isfinite(summary.final_cost)); + EXPECT_LE(summary.final_cost, summary.initial_cost + 1e-3f); +} + +INSTANTIATE_TEST_SUITE_P( + Sizes, LoopClosurePGOTest, + ::testing::Values(LcPgoParams{100, 300, "P100_LC300"}, + LcPgoParams{500, 1000, "P500_LC1k"}, + LcPgoParams{1000, 3000, "P1k_LC3k"}, + LcPgoParams{5000, 10000, "P5k_LC10k"})); + } // namespace cunls diff --git a/tests/synthetic_sba_test.cpp b/tests/synthetic_sba_test.cpp new file mode 100644 index 0000000..bd2fce0 --- /dev/null +++ b/tests/synthetic_sba_test.cpp @@ -0,0 +1,365 @@ +/* + * 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 + */ + +/** + * @file synthetic_sba_test.cpp + * @brief Parameterized synthetic SBA benchmark. + * + * Generates a synthetic bundle-adjustment problem (random poses, random + * landmarks, reprojection observations with Gaussian noise) at a grid of + * problem sizes (poses x points), runs LM, and lets nsys measure + * per-stage costs. Solver / Schur / PCG block size are picked from + * environment variables — see tests/utils.h. + * + * The first pose and the first landmark are held fixed to remove the + * SBA gauge degeneracy. + */ + +#include + +#include +#include +#include +#include +#include + +#include "cunls/common/cublas_helper.h" +#include "cunls/common/cuda_stream.h" +#include "cunls/common/helper.h" +#include "cunls/common/profiler.h" +#include "cunls/common/types.h" +#include "cunls/factor/information_factor_batch.h" +#include "cunls/factor/reprojection_factor_batch.h" +#include "cunls/math/so_se_lie_math.h" +#include "cunls/minimizer/levenberg_marquardt_minimizer.h" +#include "cunls/minimizer/problem.h" +#include "cunls/robustifier/huber_loss_function_batch.h" +#include "cunls/state/se3_state_batch.h" +#include "cunls/state/vector_state_batch.h" +#include "tests/utils.h" + +namespace cunls { + +struct SyntheticSbaParams { + int n_poses; + int n_points; + int obs_per_landmark; // visibility per landmark; total obs = n_points * obs_per_landmark. + const char *label; +}; + +inline std::ostream &operator<<(std::ostream &os, + const SyntheticSbaParams &p) { + return os << p.label; +} + +class SyntheticSbaTest : public ::testing::TestWithParam { +protected: + /** + * @brief Generates a random SE3 pose by sampling a small twist and + * composing with a forward translation, so all poses look at + * roughly the same scene region. + */ + SE3Transform RandomPose(std::mt19937 &rng) { + std::uniform_real_distribution r(-0.3f, 0.3f); + std::uniform_real_distribution t_xy(-2.f, 2.f); + std::uniform_real_distribution t_z(-1.f, 1.f); + Vector<6> twist{r(rng), r(rng), r(rng), t_xy(rng), t_xy(rng), t_z(rng)}; + dvector> d_twist({twist}); + dvector d_pose(1); + CudaStream stream; + constexpr size_t pitch = 4; + constexpr size_t stride = 16; + ComputeExpSE3(stream.GetStream(), + reinterpret_cast(d_twist.data()), 6, pitch, + stride, 1, reinterpret_cast(d_pose.data())); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + SE3Transform out; + d_pose.CopyToHost(&out, 1); + return out; + } + + /** Projects a world point through a world-from-camera pose to get a + * normalized image observation. Returns false if the point is behind + * the camera or too close to it. */ + static bool Project(const SE3Transform &pose_world_from_cam_or_cam_from_world, + const Vector<3> &p_world, Vector<2> &out, bool cam_from_world) { + const SE3Transform &T = pose_world_from_cam_or_cam_from_world; + Vector<3> p_cam{}; + if (cam_from_world) { + // T transforms world->cam. + p_cam[0] = T[0] * p_world[0] + T[1] * p_world[1] + T[2] * p_world[2] + T[3]; + p_cam[1] = T[4] * p_world[0] + T[5] * p_world[1] + T[6] * p_world[2] + T[7]; + p_cam[2] = T[8] * p_world[0] + T[9] * p_world[1] + T[10] * p_world[2] + T[11]; + } else { + // T transforms cam->world; compute inverse. + // Inverse of [R t; 0 1] is [R^T -R^T t; 0 1]. + float dx = p_world[0] - T[3]; + float dy = p_world[1] - T[7]; + float dz = p_world[2] - T[11]; + p_cam[0] = T[0] * dx + T[4] * dy + T[8] * dz; + p_cam[1] = T[1] * dx + T[5] * dy + T[9] * dz; + p_cam[2] = T[2] * dx + T[6] * dy + T[10] * dz; + } + if (!(p_cam[2] > 0.05f)) { + return false; + } + out[0] = p_cam[0] / p_cam[2]; + out[1] = p_cam[1] / p_cam[2]; + return true; + } + + /** Returns (host) ground-truth poses, perturbed poses, ground-truth + * landmarks, perturbed landmarks, observations, and per-obs ids. */ + void GenerateProblem(int n_poses, int n_points, int obs_per_landmark, + std::vector >_poses, + std::vector &init_poses, + std::vector> >_points, + std::vector> &init_points, + std::vector> &observations, + std::vector &pose_ids, + std::vector &point_ids) { + std::mt19937 rng(42); + + // Ground-truth poses: random small perturbations. Treat the pose + // state as world-from-cam (cuNLS's SE3 state). + gt_poses.resize(n_poses); + for (int i = 0; i < n_poses; ++i) { + gt_poses[i] = RandomPose(rng); + } + + // Ground-truth landmarks: uniform cloud in front of the cameras. + gt_points.resize(n_points); + std::uniform_real_distribution xy_dist(-10.f, 10.f); + std::uniform_real_distribution z_dist(3.f, 15.f); + for (int i = 0; i < n_points; ++i) { + gt_points[i] = {xy_dist(rng), xy_dist(rng), z_dist(rng)}; + } + + // Observations: each landmark gets up to `obs_per_landmark` valid + // observations from random poses. Landmarks that fail to obtain at + // least one valid observation (point behind every randomly sampled + // camera) are dropped from the problem to keep + // Problem::CheckConsistency happy — every registered state must be + // constrained by some factor. + std::uniform_int_distribution pose_pick(0, n_poses - 1); + std::normal_distribution noise(0.f, 1e-3f); + observations.clear(); + pose_ids.clear(); + point_ids.clear(); + observations.reserve(static_cast(n_points) * + static_cast(obs_per_landmark)); + pose_ids.reserve(observations.capacity()); + point_ids.reserve(observations.capacity()); + std::vector keep_landmark; + keep_landmark.reserve(n_points); + int next_remap = 0; + std::vector remap(n_points, -1); + for (int j = 0; j < n_points; ++j) { + int got = 0; + int attempts = 0; + while (got < obs_per_landmark && attempts < obs_per_landmark * 8) { + ++attempts; + int pose_id = pose_pick(rng); + Vector<2> obs{}; + if (!Project(gt_poses[pose_id], gt_points[j], obs, + /*cam_from_world=*/false)) { + continue; + } + if (remap[j] < 0) { + remap[j] = next_remap++; + keep_landmark.push_back(j); + } + obs[0] += noise(rng); + obs[1] += noise(rng); + observations.push_back(obs); + pose_ids.push_back(pose_id); + point_ids.push_back(remap[j]); // remap to the kept-landmark index + ++got; + } + } + // Rebuild gt_points to contain only kept landmarks, in remapped order. + std::vector> gt_pts_kept(keep_landmark.size()); + for (size_t i = 0; i < keep_landmark.size(); ++i) { + gt_pts_kept[i] = gt_points[keep_landmark[i]]; + } + gt_points.swap(gt_pts_kept); + + // Perturbed initial state: GT + small noise, except first pose / first + // landmark which stay at GT (and will be held constant to remove gauge). + // gt_points has been resized to only the kept landmarks; init_points + // matches. + std::normal_distribution pose_jitter_twist(0.f, 0.05f); + std::normal_distribution point_jitter(0.f, 0.05f); + init_poses = gt_poses; + init_points = gt_points; + { + std::vector> jitters(n_poses); + for (int i = 1; i < n_poses; ++i) { + Vector<6> &j = jitters[i]; + for (int k = 0; k < 6; ++k) { + j[k] = pose_jitter_twist(rng); + } + } + // Compose init_pose[i] = gt_pose[i] * exp(jitter) on device. + dvector> d_jit(jitters); + dvector d_delta(n_poses); + CudaStream stream; + ComputeExpSE3(stream.GetStream(), + reinterpret_cast(d_jit.data()), 6, 4, 16, + n_poses, reinterpret_cast(d_delta.data())); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + std::vector deltas(n_poses); + d_delta.CopyToHost(deltas.data(), n_poses); + // 4x4 matrix multiplication on host. + auto mul = [](const SE3Transform &a, + const SE3Transform &b) -> SE3Transform { + SE3Transform c{}; + for (int r = 0; r < 4; ++r) { + for (int cc = 0; cc < 4; ++cc) { + float s = 0.f; + for (int k = 0; k < 4; ++k) { + s += a[r * 4 + k] * b[k * 4 + cc]; + } + c[r * 4 + cc] = s; + } + } + return c; + }; + for (int i = 1; i < n_poses; ++i) { + init_poses[i] = mul(gt_poses[i], deltas[i]); + } + } + for (size_t j = 1; j < init_points.size(); ++j) { + Vector<3> &p = init_points[j]; + p[0] += point_jitter(rng); + p[1] += point_jitter(rng); + p[2] += point_jitter(rng); + } + } + + cuBLASHandle cublas_handle_; + profiler::Domain profiler_domain_ = profiler::Domain("SyntheticSbaTest"); +}; + +/** + * @brief Builds and runs LM on a synthetic SBA problem of the configured + * size. Holds pose 0 and landmark 0 constant. Time measurement + * is via NVTX ranges installed by the minimizer + the test + * fixture. + */ +TEST_P(SyntheticSbaTest, Optimize) { + auto params = GetParam(); + SCOPED_TRACE(std::string("SyntheticSbaTest: ") + params.label); + + std::vector gt_poses, init_poses; + std::vector> gt_points, init_points; + std::vector> observations; + std::vector pose_ids, point_ids; + GenerateProblem(params.n_poses, params.n_points, params.obs_per_landmark, + gt_poses, init_poses, gt_points, init_points, observations, + pose_ids, point_ids); + + const size_t n_obs = observations.size(); + ASSERT_GT(n_obs, 0u); + + // Identity 2x2 sqrt information per observation. + std::vector> sqrt_info(n_obs); + for (size_t i = 0; i < n_obs; ++i) { + sqrt_info[i] = {1.f, 0.f, 0.f, 1.f}; + } + + // Camera-from-rig: identity (single camera, no rig). + std::vector cam_from_rig(n_obs); + for (size_t i = 0; i < n_obs; ++i) { + cam_from_rig[i] = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; + } + + dvector poses_d(init_poses); + dvector> points_d(init_points); + dvector> obs_d(observations); + dvector> info_d(sqrt_info); + dvector cam_from_rig_d(cam_from_rig); + + // Hold pose 0 + landmark 0 constant to fix the gauge. + std::vector const_pose_ids = {0}; + std::vector const_point_ids = {0}; + dvector const_pose_ids_d(const_pose_ids); + dvector const_point_ids_d(const_point_ids); + + auto poses_ptr = reinterpret_cast(poses_d.data()); + auto points_ptr = reinterpret_cast(points_d.data()); + + SE3StateBatch pose_batch(cublas_handle_, poses_ptr, + static_cast(params.n_poses), + const_pose_ids_d.data(), const_pose_ids.size()); + VectorStateBatch<3> point_batch(points_ptr, init_points.size(), + const_point_ids_d.data(), + const_point_ids.size()); + + InformationFactorBatch info_factor( + cublas_handle_, info_d.data(), n_obs, obs_d.data(), + cam_from_rig_d.data(), n_obs, 1e-3f); + + std::vector state_pointers; + state_pointers.reserve(n_obs * 2); + for (size_t i = 0; i < n_obs; ++i) { + state_pointers.push_back( + pose_batch.StateBlockDevicePtr(static_cast(pose_ids[i]))); + state_pointers.push_back( + point_batch.StateBlockDevicePtr(static_cast(point_ids[i]))); + } + + HuberLossFunctionBatch huber(1.0f); + Problem problem; + problem.AddStateBatch(&pose_batch); + problem.AddStateBatch(&point_batch); + problem.AddFactorBatch(&info_factor, &huber, state_pointers); + ASSERT_TRUE(problem.CheckConsistency()); + + CudaStream stream; + MinimizerOptions options; + options.max_num_iterations = 10; + options.state_tolerance = 1e-10f; + options.cost_tolerance = 1e-6f; + options.disable_safety_checks = true; + options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(200); + options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); + LevenbergMarquardtMinimizerOptions lm_options; + lm_options.base_options = options; + lm_options.initial_lambda = 1e-3f; + LevenbergMarquardtMinimizer minimizer(lm_options); + + MinimizerSummary summary; + { + auto range = profiler_domain_.CreateDomainRange("Minimize"); + summary = minimizer.Minimize(stream.GetStream(), problem); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + } + + EXPECT_TRUE(std::isfinite(summary.initial_cost)); + EXPECT_TRUE(std::isfinite(summary.final_cost)); + EXPECT_LE(summary.final_cost, summary.initial_cost + 1e-3f); +} + +INSTANTIATE_TEST_SUITE_P( + Sizes, SyntheticSbaTest, + ::testing::Values( + SyntheticSbaParams{10, 1000, 5, "P10_L1k_obs5"}, + SyntheticSbaParams{50, 10000, 5, "P50_L10k_obs5"}, + SyntheticSbaParams{100, 50000, 5, "P100_L50k_obs5"}, + SyntheticSbaParams{250, 200000, 5, "P250_L200k_obs5"}, + // 1M landmarks: obs_per_landmark=2 keeps the Jacobian row count + // under cuNLS's per-batch grid-dim limit (~4M rows). Total obs ≈ 2M. + SyntheticSbaParams{500, 1000000, 2, "P500_L1M_obs2"})); + +} // namespace cunls diff --git a/tests/utils.h b/tests/utils.h index c0e9313..79e3a29 100644 --- a/tests/utils.h +++ b/tests/utils.h @@ -351,15 +351,18 @@ CopyStateToHost(const VectorStateBatch &state_batch) { /** * @brief Reads CUNLS_SOLVER from the environment and returns the matching - * solver type. Defaults to cuDSS so existing tests keep their - * baseline behaviour. Recognised values: "cuDSS", "BlockSparsePCG". + * 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". */ inline SparseLinearSolverType SolverTypeFromEnv() { const char *s = std::getenv("CUNLS_SOLVER"); - if (s != nullptr && std::strcmp(s, "BlockSparsePCG") == 0) { - return SparseLinearSolverType::BlockSparsePCG; + if (s != nullptr && std::strcmp(s, "cuDSS") == 0) { + return SparseLinearSolverType::cuDSS; } - return SparseLinearSolverType::cuDSS; + return SparseLinearSolverType::BlockSparsePCG; } /** Reads CUNLS_PCG_BLOCK_SIZE; defaults to ``fallback`` if unset/invalid. */ From 205bb4a74c1dc539c92da8a92356ee0b8233a8f2 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 13 May 2026 13:20:25 -0700 Subject: [PATCH 3/4] Improve perf --- .../linear_solver/block_sparse_pcg_solver.cu | 121 +++++++++++------- 1 file changed, 73 insertions(+), 48 deletions(-) diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index b098a3d..13cd7b9 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -87,6 +87,7 @@ #include #include +#include #include #include #include @@ -150,7 +151,10 @@ __global__ void ExtractAndFactorBlockDiagonalsKernel( } __syncthreads(); - // Step 1: gather dense tile from CSR. + // Step 1: gather dense tile from CSR. Early-break exploits the + // sorted-cols invariant — for SBA pose rows where the many + // landmark columns sit after the diagonal-tile range, this turns + // an O(nnz_per_row) scan into O(B). if (tid < B) { int global_row = row_start + block_row * B + tid; int start = row_off[global_row]; @@ -159,7 +163,10 @@ __global__ void ExtractAndFactorBlockDiagonalsKernel( int col_hi = col_lo + B; for (int k = start; k < end; ++k) { int c = col_idx[k]; - if (c >= col_lo && c < col_hi) { + if (c >= col_hi) { + break; + } + if (c >= col_lo) { tile[tid * B + (c - col_lo)] = values[k]; } } @@ -236,7 +243,10 @@ __global__ void ExtractAndFactorGenericKernel( int col_hi = col_lo + B; for (int k = start; k < end; ++k) { int c = col_idx[k]; - if (c >= col_lo && c < col_hi) { + if (c >= col_hi) { + break; + } + if (c >= col_lo) { tile[tid * B + (c - col_lo)] = values[k]; } } @@ -308,6 +318,12 @@ __global__ void ExtractAndFactorPerThreadKernel( } int col_lo = row_start + block_row * B; int col_hi = col_lo + B; + // CSR column indices are sorted within a row, so as soon as we walk + // past `col_hi` we are guaranteed never to see a column in the + // tile's range again — break out of the inner loop. Crucially, for + // SBA's pose rows where the non-diagonal cols are the (many) + // landmark cols sorted *after* the diagonal-tile cols, this turns a + // per-row scan of ~hundreds of entries into a scan of ~B entries. #pragma unroll for (int rr = 0; rr < B; ++rr) { int global_row = col_lo + rr; @@ -315,7 +331,10 @@ __global__ void ExtractAndFactorPerThreadKernel( int end = row_off[global_row + 1]; for (int k = start; k < end; ++k) { int c = col_idx[k]; - if (c >= col_lo && c < col_hi) { + if (c >= col_hi) { + break; + } + if (c >= col_lo) { tile[rr * B + (c - col_lo)] = values[k]; } } @@ -597,19 +616,25 @@ ApplyBlockJacobiPerThreadKernel(const float *__restrict__ factors, // ============================================================================= /** - * @brief Fused `x ← x + α p ; r ← r − α q` for the PCG step. + * @brief Fused: compute `α = rz_old / ` and apply + * `x ← x + α p ; r ← r − α q` in a single kernel. * - * Reads `α` from device memory so the host doesn't have to wait on the - * preceding `` dot product. One thread per coordinate. + * Thread 0 reads the two scalar slots, computes α, stores it in + * shared memory; all threads then perform the axpy on their slice. + * This eliminates the separate `ComputeAlphaKernel` launch + * (~1.2 µs / call × thousands of iters = several percent of total + * runtime on the small-Minimize SBA fixture). */ -__global__ void PcgUpdateKernel(const float *__restrict__ alpha_ptr, +__global__ void PcgUpdateKernel(const float *__restrict__ rz_old_ptr, + const float *__restrict__ pAp_ptr, const float *__restrict__ p, const float *__restrict__ Ap, float *__restrict__ x, float *__restrict__ r, int n) { __shared__ float a; if (threadIdx.x == 0) { - a = alpha_ptr[0]; + float denom = pAp_ptr[0]; + a = (denom > 0.f) ? rz_old_ptr[0] / denom : 0.f; } __syncthreads(); int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -621,14 +646,28 @@ __global__ void PcgUpdateKernel(const float *__restrict__ alpha_ptr, } /** - * @brief Updates the search direction `p ← z + β p` (in place into p). + * @brief Fused: compute `β = rz_new / rz_old`, store `rz_old ← rz_new`, + * and update the search direction `p ← z + β p` in a single + * kernel. Eliminates the separate `ComputeBetaKernel` launch + * and writes the new `rz_old` from the first CTA's thread 0 + * (subsequent CTAs don't read the slot in this kernel). */ -__global__ void PcgDirectionKernel(const float *__restrict__ beta_ptr, +__global__ void PcgDirectionKernel(const float *__restrict__ rz_new_ptr, + float *__restrict__ rz_old_ptr, const float *__restrict__ z, float *__restrict__ p, int n) { __shared__ float b; if (threadIdx.x == 0) { - b = beta_ptr[0]; + float num = rz_new_ptr[0]; + float denom = rz_old_ptr[0]; + b = (fabsf(denom) > 0.f) ? num / denom : 0.f; + // Only the first CTA's thread 0 writes rz_old; all other CTAs + // skip the write (the rz_old slot is read by the *next* + // iteration's PcgUpdate, which happens after this kernel + // completes, so any one writer is enough). + if (blockIdx.x == 0) { + rz_old_ptr[0] = num; + } } __syncthreads(); int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -638,24 +677,6 @@ __global__ void PcgDirectionKernel(const float *__restrict__ beta_ptr, p[i] = z[i] + b * p[i]; } -/** `alpha = rz_old / ` (single-thread, device-side). */ -__global__ void ComputeAlphaKernel(const float *__restrict__ rz_old, - const float *__restrict__ pAp, - float *__restrict__ alpha) { - float denom = pAp[0]; - alpha[0] = (denom > 0.f) ? rz_old[0] / denom : 0.f; -} - -/** `beta = rz_new / rz_old; rz_old <- rz_new` (single-thread). */ -__global__ void ComputeBetaKernel(const float *__restrict__ rz_new, - float *__restrict__ rz_old, - float *__restrict__ beta) { - float num = rz_new[0]; - float denom = rz_old[0]; - beta[0] = (fabsf(denom) > 0.f) ? num / denom : 0.f; - rz_old[0] = num; -} - /** Zero one or two scalar slots. */ __global__ void ZeroScalarKernel(float *out) { out[0] = 0.f; } __global__ void Zero2ScalarKernel(float *out_a, float *out_b) { @@ -889,20 +910,31 @@ void DispatchApplyPrecond(cudaStream_t stream, int B, int row_start, THROW_ON_CUDA_ERROR(cudaGetLastError()); } -/** Enqueues `out[0] = ` (clear-then-reduce). */ +/** Enqueues `out[0] = ` (clear-then-reduce). + * + * We use `cudaMemsetAsync` instead of a one-thread kernel because for + * a single-float clear the runtime can short-circuit it to a tiny + * memory operation, saving the per-kernel-launch dispatch overhead + * (~1 µs) that dominates on small-matrix workloads where PCG + * performs hundreds of dots per Solve. */ void DotAsync(cudaStream_t stream, const float *a, const float *b, int n, float *out) { - ZeroScalarKernel<<<1, 1, 0, stream>>>(out); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(out, 0, sizeof(float), stream)); int threads = 256; int blocks = std::min(1024, (n + threads - 1) / threads); DotKernel<<>>(a, b, n, out); } /** Enqueues `out_ab = ` and `out_ac = ` in a single pass. - * Equivalent to two @ref DotAsync calls but with half the global reads. */ + * Equivalent to two @ref DotAsync calls but with half the global reads + * AND a single 8-byte memset (vs two 4-byte memsets) — the two output + * slots are required to be adjacent in memory, which the caller + * guarantees by laying them out next to each other in @c d_scratch_. */ void DualDotAsync(cudaStream_t stream, const float *a, const float *b, const float *c, int n, float *out_ab, float *out_ac) { - Zero2ScalarKernel<<<1, 1, 0, stream>>>(out_ab, out_ac); + // Verified adjacent at the call site (kRzNew, kRnorm2 in d_scratch_). + assert(out_ac == out_ab + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(out_ab, 0, 2 * sizeof(float), stream)); int threads = 256; int blocks = std::min(1024, (n + threads - 1) / threads); DualDotKernel<<>>(a, b, c, n, out_ab, out_ac); @@ -1202,16 +1234,13 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, &spmv_beta, vecY, CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, spmv_buffer_.data())); - // and alpha = rz_old / . + // . DotAsync(stream, p_.data(), Ap_.data(), n, d_scratch_.data() + kPAp); - ComputeAlphaKernel<<<1, 1, 0, stream>>>(d_scratch_.data() + kRzOld, - d_scratch_.data() + kPAp, - d_scratch_.data() + kAlpha); - // x ← x + α p; r ← r − α Ap. + // α = rz_old / ; x ← x + α p; r ← r − α Ap (fused kernel). PcgUpdateKernel<<>>( - d_scratch_.data() + kAlpha, p_.data(), Ap_.data(), result.data(), - r_.data(), n); + d_scratch_.data() + kRzOld, d_scratch_.data() + kPAp, p_.data(), + Ap_.data(), result.data(), r_.data(), n); // z = M^{-1} r. for (const auto &s : segments_) { @@ -1223,14 +1252,10 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, DualDotAsync(stream, r_.data(), z_.data(), r_.data(), n, d_scratch_.data() + kRzNew, d_scratch_.data() + kRnorm2); - // beta = rz_new / rz_old; rz_old ← rz_new. - ComputeBetaKernel<<<1, 1, 0, stream>>>(d_scratch_.data() + kRzNew, - d_scratch_.data() + kRzOld, - d_scratch_.data() + kBeta); - - // p ← z + β p. + // β = rz_new / rz_old; rz_old ← rz_new; p ← z + β p (fused). PcgDirectionKernel<<>>( - d_scratch_.data() + kBeta, z_.data(), p_.data(), n); + d_scratch_.data() + kRzNew, d_scratch_.data() + kRzOld, z_.data(), + p_.data(), n); // Convergence poll every `check_period` iters: one D2H of one float // + one stream sync. Far cheaper than a host sync per iteration. From e8d2c3296ec54f7a746c1c3294aca81a368b65c1 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 13 May 2026 14:31:28 -0700 Subject: [PATCH 4/4] Fix docs and formatting --- README.md | 2 +- .../linear_solver/block_sparse_pcg_solver.cu | 73 +++++----- cunls/linear_solver/block_sparse_pcg_solver.h | 12 +- .../cudss_sparse_linear_solver.h | 4 +- cunls/linear_solver/dense_cholesky_solver.h | 3 +- cunls/linear_solver/dense_linear_solver.h | 3 +- cunls/linear_solver/dense_qr_solver.cu | 3 +- cunls/linear_solver/dense_qr_solver.h | 3 +- cunls/llms.txt | 3 +- cunls/minimizer/gauss_newton_minimizer.cu | 3 +- cunls/minimizer/gauss_newton_minimizer.h | 44 +++--- docs/sphinx/api/linear_solver.rst | 127 ++++++++++++++++-- docs/sphinx/api/minimizer.rst | 38 ++++-- llms.txt | 3 +- python/pycunls/_pycunls_core.pyi | 1 + python/src/bind_types.cpp | 3 +- python/tests/test_minimizer.py | 3 +- tests/dense_cholesky_solver_test.cpp | 18 ++- tests/dense_linear_solver_test.cpp | 18 ++- tests/dense_qr_solver_test.cpp | 18 ++- tests/pgo_minimizer_test.cpp | 11 +- tests/sba_minimizer_test.cpp | 9 +- tests/sparse_linear_solver_test.cpp | 8 +- tests/synthetic_pgo_test.cpp | 69 ++++++---- tests/synthetic_sba_test.cpp | 31 +++-- 25 files changed, 337 insertions(+), 173 deletions(-) diff --git a/README.md b/README.md index a069ffd..ff075d2 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ $\left\|v\right\|^2_{\Sigma} = v^T \Sigma^{-1} v$ is the Mahalanobis norm. | **Robust losses** | Huber, Cauchy, Arctan, SoftL1, Tolerant, Tukey, Scaled | | **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 | | **Custom factors** | User-defined CUDA kernels via `SizedFactorBatch` | -| **Linear solver** | NVIDIA cuDSS, dense LDLT, dense Cholesky (cuSOLVER), dense QR (cuSOLVER) | +| **Linear solver** | Block-sparse PCG (variable block-Jacobi preconditioner, default), NVIDIA cuDSS, dense LDLT, dense Cholesky (cuSOLVER), dense QR (cuSOLVER) | | **Safety checks** | Optional runtime validation (linear-solver diagnostics and more) — disable via `MinimizerOptions::disable_safety_checks` for low-latency solves | | **Execution model** | Fully asynchronous via CUDA streams | diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index 13cd7b9..3d54982 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -421,11 +421,10 @@ __global__ void ExtractScalarJacobi(const int *__restrict__ row_off, * @tparam B Compile-time tile side length. */ template -__global__ void ApplyBlockJacobiKernel(const float *__restrict__ factors, - int factor_offset, - const float *__restrict__ r, - int row_start, int num_blocks, - float *__restrict__ z) { +__global__ void +ApplyBlockJacobiKernel(const float *__restrict__ factors, int factor_offset, + const float *__restrict__ r, int row_start, + int num_blocks, float *__restrict__ z) { int block_row = blockIdx.x; if (block_row >= num_blocks) { return; @@ -478,10 +477,11 @@ __global__ void ApplyBlockJacobiKernel(const float *__restrict__ factors, } /** Generic-B apply for non-templated tile sizes; dynamic shared mem. */ -__global__ void ApplyBlockJacobiGenericKernel( - const float *__restrict__ factors, int factor_offset, int B, - const float *__restrict__ r, int row_start, int num_blocks, - float *__restrict__ z) { +__global__ void ApplyBlockJacobiGenericKernel(const float *__restrict__ factors, + int factor_offset, int B, + const float *__restrict__ r, + int row_start, int num_blocks, + float *__restrict__ z) { int block_row = blockIdx.x; if (block_row >= num_blocks) { return; @@ -553,9 +553,9 @@ __global__ void ApplyScalarJacobi(const float *__restrict__ factors, template __global__ void ApplyBlockJacobiPerThreadKernel(const float *__restrict__ factors, - int factor_offset, - const float *__restrict__ r, int row_start, - int num_blocks, float *__restrict__ z) { + int factor_offset, const float *__restrict__ r, + int row_start, int num_blocks, + float *__restrict__ z) { int block_row = blockIdx.x * blockDim.x + threadIdx.x; if (block_row >= num_blocks) { return; @@ -761,8 +761,7 @@ __global__ void DualDotKernel(const float *__restrict__ a, /** Picks the right ExtractAndFactor specialization for B. */ void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, int num_blocks, int factor_offset, - float pivot_floor, - const CSRSparseMatrix &matrix, + float pivot_floor, const CSRSparseMatrix &matrix, float *factors) { if (num_blocks == 0) { return; @@ -785,10 +784,9 @@ void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, case BVAL: { \ int threads = 256; \ int blocks = (num_blocks + threads - 1) / threads; \ - ExtractAndFactorPerThreadKernel \ - <<>>(row_off, col_idx, vals, row_start, \ - num_blocks, factor_offset, \ - pivot_floor, factors); \ + ExtractAndFactorPerThreadKernel<<>>( \ + row_off, col_idx, vals, row_start, num_blocks, factor_offset, \ + pivot_floor, factors); \ break; \ } @@ -813,10 +811,9 @@ void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, #define LAUNCH_FACTOR(BVAL) \ case BVAL: \ ExtractAndFactorBlockDiagonalsKernel \ - <<>>(row_off, col_idx, vals, \ - row_start, num_blocks, \ - factor_offset, pivot_floor, \ - factors); \ + <<>>( \ + row_off, col_idx, vals, row_start, num_blocks, factor_offset, \ + pivot_floor, factors); \ break switch (B) { @@ -826,7 +823,8 @@ void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, LAUNCH_FACTOR(16); default: { size_t shared_bytes = static_cast(B) * B * sizeof(float); - ExtractAndFactorGenericKernel<<>>( + ExtractAndFactorGenericKernel<<>>( row_off, col_idx, vals, B, row_start, num_blocks, factor_offset, pivot_floor, factors); break; @@ -863,9 +861,8 @@ void DispatchApplyPrecond(cudaStream_t stream, int B, int row_start, case BVAL: { \ int threads = 256; \ int blocks = (num_blocks + threads - 1) / threads; \ - ApplyBlockJacobiPerThreadKernel \ - <<>>(factors, factor_offset, r, \ - row_start, num_blocks, z); \ + ApplyBlockJacobiPerThreadKernel<<>>( \ + factors, factor_offset, r, row_start, num_blocks, z); \ break; \ } @@ -1006,8 +1003,7 @@ bool BlockSparsePCGSolver::BuildSegmentTables(int matrix_dim) { row_cursor += count * size; block_cursor += count; factor_cursor += static_cast(count) * - (size == 1 ? 1ull - : static_cast(size) * size); + (size == 1 ? 1ull : static_cast(size) * size); } if (row_cursor != matrix_dim) { LogError("BlockSparsePCGSolver: layout covers {} rows but matrix has {}", @@ -1095,8 +1091,7 @@ bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, auto matA = static_cast(mat_desc_.GetDescription()); cusparseDnVecDescr_t vecX = nullptr; cusparseDnVecDescr_t vecY = nullptr; - THROW_ON_CUSPARSE_ERROR( - cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR(cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); THROW_ON_CUSPARSE_ERROR( cusparseCreateDnVec(&vecY, n, Ap_.data(), CUDA_R_32F)); @@ -1215,8 +1210,7 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, // explicit size n each Solve. Cheap host calls. cusparseDnVecDescr_t vecX = nullptr; cusparseDnVecDescr_t vecY = nullptr; - THROW_ON_CUSPARSE_ERROR( - cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR(cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); THROW_ON_CUSPARSE_ERROR( cusparseCreateDnVec(&vecY, n, Ap_.data(), CUDA_R_32F)); @@ -1229,10 +1223,10 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, // Ap = H * p. float spmv_alpha = 1.f; float spmv_beta = 0.f; - THROW_ON_CUSPARSE_ERROR(cusparseSpMV( - handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &spmv_alpha, matA, vecX, - &spmv_beta, vecY, CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, - spmv_buffer_.data())); + THROW_ON_CUSPARSE_ERROR( + cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &spmv_alpha, + matA, vecX, &spmv_beta, vecY, CUDA_R_32F, + CUSPARSE_SPMV_ALG_DEFAULT, spmv_buffer_.data())); // . DotAsync(stream, p_.data(), Ap_.data(), n, d_scratch_.data() + kPAp); @@ -1259,12 +1253,11 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, // Convergence poll every `check_period` iters: one D2H of one float // + one stream sync. Far cheaper than a host sync per iteration. - if (((it + 1) % check_period) == 0 || - (it + 1) == options_.max_iterations) { + if (((it + 1) % check_period) == 0 || (it + 1) == options_.max_iterations) { float r_norm2 = 0.f; THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&r_norm2, d_scratch_.data() + kRnorm2, - sizeof(float), - cudaMemcpyDeviceToHost, stream)); + sizeof(float), cudaMemcpyDeviceToHost, + stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); if (r_norm2 <= stop_thresh) { ++it; diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h index c27ed31..c400608 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.h +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -207,8 +207,8 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { * layout. */ bool Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + const CSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) final; /** * @brief Runs the PCG loop on `H x = b`, writing into @p result. @@ -282,10 +282,10 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { /** * @brief LDLT factors of every diagonal tile, packed contiguously. * - * Tile @c b of segment @c s lives at `precond_factors_[segment_factor_offset(s) + - * b * size_s * size_s ..]` in row-major order. Lower triangle holds - * `L` with unit diagonal (implicit); the stored diagonal holds `D`; - * upper triangle is unused. + * Tile @c b of segment @c s lives at + * `precond_factors_[segment_factor_offset(s) + b * size_s * size_s ..]` in + * row-major order. Lower triangle holds `L` with unit diagonal (implicit); + * the stored diagonal holds `D`; upper triangle is unused. */ dvector precond_factors_; diff --git a/cunls/linear_solver/cudss_sparse_linear_solver.h b/cunls/linear_solver/cudss_sparse_linear_solver.h index 26e2b47..dacf420 100644 --- a/cunls/linear_solver/cudss_sparse_linear_solver.h +++ b/cunls/linear_solver/cudss_sparse_linear_solver.h @@ -93,8 +93,8 @@ class cuDSSLinearSolver : public CSRSparseLinearSolver { * @return true on success, false if a dimension mismatch is detected. */ bool Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + const CSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) final; /** * @brief Solves a sparse SPD linear system Ax = b. diff --git a/cunls/linear_solver/dense_cholesky_solver.h b/cunls/linear_solver/dense_cholesky_solver.h index 6a34afd..584595f 100644 --- a/cunls/linear_solver/dense_cholesky_solver.h +++ b/cunls/linear_solver/dense_cholesky_solver.h @@ -52,8 +52,7 @@ class DenseCholeskySolver : public CSRSparseLinearSolver { * @return true on success, false if a dimension mismatch is detected. */ bool Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; /** diff --git a/cunls/linear_solver/dense_linear_solver.h b/cunls/linear_solver/dense_linear_solver.h index 796d5ec..ceb0441 100644 --- a/cunls/linear_solver/dense_linear_solver.h +++ b/cunls/linear_solver/dense_linear_solver.h @@ -57,8 +57,7 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { * @return true on success, false if a dimension mismatch is detected. */ bool Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; /** diff --git a/cunls/linear_solver/dense_qr_solver.cu b/cunls/linear_solver/dense_qr_solver.cu index 557a572..ae5ff2d 100644 --- a/cunls/linear_solver/dense_qr_solver.cu +++ b/cunls/linear_solver/dense_qr_solver.cu @@ -67,8 +67,7 @@ __global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, } // namespace -bool DenseQRSolver::Initialize(cudaStream_t stream, - const Problem & /*problem*/, +bool DenseQRSolver::Initialize(cudaStream_t stream, const Problem & /*problem*/, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { diff --git a/cunls/linear_solver/dense_qr_solver.h b/cunls/linear_solver/dense_qr_solver.h index 5616d43..154842e 100644 --- a/cunls/linear_solver/dense_qr_solver.h +++ b/cunls/linear_solver/dense_qr_solver.h @@ -54,8 +54,7 @@ class DenseQRSolver : public CSRSparseLinearSolver { * @return true on success, false if a dimension mismatch is detected. */ bool Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; /** diff --git a/cunls/llms.txt b/cunls/llms.txt index 7b514b8..1b8ee7b 100644 --- a/cunls/llms.txt +++ b/cunls/llms.txt @@ -14,7 +14,8 @@ This directory contains all core library modules and the public umbrella header - `common/llms.txt`: CUDA handles, device vectors, sparse/dense type aliases. - `factor/llms.txt`: factor interfaces and built-in factor implementations. -- `linear_solver/llms.txt`: sparse solver interface and cuDSS implementation. +- `linear_solver/llms.txt`: sparse solver interface, BlockSparsePCG + iterative solver (default), cuDSS direct solver, and dense fallbacks. - `math/llms.txt`: Lie and matrix math kernels/utilities. - `minimizer/llms.txt`: `Problem`, residual wrappers, GN/LM solvers. - `robustifier/llms.txt`: robust loss batches. diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index 8b28074..16a4c7a 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -556,8 +556,7 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, { auto solve_range = profiler_domain_.CreateDomainRange("LinearSolve"); - bool success = - solver_->Solve(stream, lhs_work_, rhs_work_, step_); + bool success = solver_->Solve(stream, lhs_work_, rhs_work_, step_); if (!success) { std::string str = "Failed to solve linear system"; LogError(str); diff --git a/cunls/minimizer/gauss_newton_minimizer.h b/cunls/minimizer/gauss_newton_minimizer.h index 26e3ff1..d3b466e 100644 --- a/cunls/minimizer/gauss_newton_minimizer.h +++ b/cunls/minimizer/gauss_newton_minimizer.h @@ -111,31 +111,39 @@ struct MinimizerOptions { * @brief Type of sparse linear solver to use. * * Supported options: - * - cuDSS: GPU-accelerated sparse direct solver via NVIDIA's cuDSS library. - * - DenseLDLT: Converts the CSR matrix to dense and solves via a custom - * CUDA pivoted LDLT factorization. Suitable for small-to-medium systems - * where the matrix fits in dense form. - * - * Default: cuDSS + * - BlockSparsePCG (default): block-Jacobi preconditioned CG; derives + * the per-state-batch block layout automatically. Recommended for + * most SBA and PGO workloads — see `profile/PCG_RESULTS.md` and + * `profile/benchmark_a6000.png`. + * - cuDSS: NVIDIA cuDSS sparse direct solver. Preferred when many + * near-identical small Hessians are solved back-to-back and the + * per-iteration kernel-launch overhead of PCG dominates. + * - DenseLDLT: converts CSR to dense, solves with a custom pivoted + * LDLT kernel. Suitable for small dense Hessians. + * - DenseCholesky / DenseQR: cuSOLVER-backed dense solves. + * + * Default: BlockSparsePCG. */ SparseLinearSolverType sparse_linear_solver_type = - SparseLinearSolverType::cuDSS; + SparseLinearSolverType::BlockSparsePCG; /** * @brief Configuration for the sparse linear solver. * - * Contains solver-specific configuration options. Defaults to cuDSS solver - * with SlowInitFastSolve configuration and 1 thread. - * To enable mutiple threads in cuDSS, please provide the full path to the - * threading library in the cudss_solver_options. - * e.g. .cudss_solver_options = cuDSSLinearSolverOptions{ - * .mode = cuDSSLinearSolverMode::SlowInitFastSolve, - * .nthreads = 12, - * .threading_lib_path = "/path/to/libcudss_mtlayer_gomp.so" - * } + * Contains backend-specific options. Only the field corresponding to + * @ref sparse_linear_solver_type is read: + * - `block_sparse_pcg_options` for `BlockSparsePCG`, + * - `cudss_solver_options` for `cuDSS`, + * - Dense backends have no extra options. + * + * Defaults: `BlockSparsePCG` with the layout auto-derived from the + * problem's state batches at @ref GaussNewtonMinimizer::Initialize time. + * + * To enable multi-threaded cuDSS, set + * `cudss_solver_options.threading_lib_path` to the full path of + * `libcudss_mtlayer_gomp.so` (or equivalent). */ - SparseLinearSolverConfig sparse_linear_solver_config = { - .cudss_solver_options = cuDSSLinearSolverOptions()}; + SparseLinearSolverConfig sparse_linear_solver_config = {}; /** * @brief Strategy for computing the approximate Hessian J^T * J. diff --git a/docs/sphinx/api/linear_solver.rst b/docs/sphinx/api/linear_solver.rst index 4778169..e507b9c 100644 --- a/docs/sphinx/api/linear_solver.rst +++ b/docs/sphinx/api/linear_solver.rst @@ -2,19 +2,27 @@ Linear Solver API ################################################################################ -`cunls/linear_solver` hosts linear-system abstractions (cuDSS integration, -dense pivoted LDLT, dense Cholesky, and dense QR solvers) behind a common -CSR-based interface. +`cunls/linear_solver` hosts linear-system abstractions (block-Jacobi PCG, +cuDSS integration, dense pivoted LDLT, dense Cholesky, and dense QR +solvers) behind a common CSR-based interface. SparseLinearSolverType ---------------------- Enum in `cunls/linear_solver/sparse_linear_solver.h`: -- `cuDSS` -- `DenseLDLT` -- `DenseCholesky` -- `DenseQR` +- `BlockSparsePCG` — iterative block-Jacobi preconditioned conjugate + gradient solver. Default backend for Gauss-Newton and + Levenberg-Marquardt: outperforms ``cuDSS`` on most SBA / PGO + workloads (see ``profile/PCG_RESULTS.md``). +- `cuDSS` — NVIDIA cuDSS sparse direct solver. Pick when each Solve + sees a tiny system and PCG's per-iter kernel-launch overhead + dominates. +- `DenseLDLT` — converts CSR to dense and solves with a custom CUDA + pivoted LDLT kernel. +- `DenseCholesky` — converts CSR to dense and solves with cuSOLVER + Cholesky (requires SPD). +- `DenseQR` — converts CSR to dense and solves with cuSOLVER QR. SparseLinearSolverConfig ------------------------ @@ -22,25 +30,37 @@ SparseLinearSolverConfig Struct in `sparse_linear_solver.h`. Only the member corresponding to the chosen ``SparseLinearSolverType`` is used: -- `cudss_solver_options` - [in] cuDSS-specific backend options (ignored when - using ``DenseLDLT``, ``DenseCholesky``, or ``DenseQR``). +- `block_sparse_pcg_options` - [in] BlockSparsePCG-specific knobs; + 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. CSRSparseLinearSolver --------------------- Abstract base (`cunls/linear_solver/csr_sparse_linear_solver.h`). -.. cpp:function:: bool Initialize(cudaStream_t stream, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) +.. cpp:function:: bool Initialize(cudaStream_t stream, const Problem& problem, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) Performs setup work (at minimum symbolic analysis; some modes also perform an initial factorization) for the given sparsity pattern. Must be called before :cpp:func:`Solve` and re-called whenever the matrix structure changes. + Solvers may inspect ``problem`` to specialize their setup (e.g. + ``BlockSparsePCGSolver`` reads each state batch's ``TangentSize`` to + build the block-Jacobi layout). Solvers that do not need it (cuDSS, + dense backends) simply ignore it; for callers that operate on a raw + matrix without a cuNLS problem context, pass a default-constructed + ``Problem``. + Both ``rhs`` and ``result`` must be pre-allocated to the same number of elements as the number of rows in ``spd_matrix``; the solver does **not** resize them. :param ``stream``: [in] CUDA stream for asynchronous GPU operations. + :param ``problem``: [in] Originating problem; used by problem-aware + backends to derive per-solver structure. :param ``spd_matrix``: [in] Symmetric matrix in CSR format (backend-specific definiteness requirements may apply). :param ``rhs``: [in] Right-hand-side vector ``b`` (size must equal matrix rows). @@ -81,6 +101,93 @@ Abstract base (`cunls/linear_solver/csr_sparse_linear_solver.h`). :returns: ``true`` when post-factorization safety checks are enabled (the default). +BlockSparsePCGOptions +--------------------- + +Struct in `cunls/linear_solver/block_sparse_pcg_solver.h`. Convergence +and preconditioner-layout knobs for ``BlockSparsePCGSolver``. + +- ``block_size`` - [in] Uniform diagonal tile size; used only when + ``block_layout`` is empty. Must divide the matrix dimension. + ``block_size = 1`` reduces the preconditioner to scalar Jacobi. + Default: ``6``. +- ``block_layout`` - [in] Optional ``std::vector>`` + describing a heterogeneous block-diagonal layout: each + ``(count_i, size_i)`` element contributes ``count_i * size_i`` rows + with ``size_i``-square diagonal tiles. When empty (default), + ``Initialize`` derives the layout automatically from the + ``Problem``'s state batches (segment per non-empty batch with + ``size = TangentSize()``, ``count = NumStateBlocks() - + NumConstStateBlocks()``). +- ``max_iterations`` - [in] PCG iteration cap. Default: ``200``. +- ``relative_tolerance`` - [in] Stop when + ``||r_k|| <= relative_tolerance * ||b||``. Default: ``1e-3``. +- ``absolute_tolerance`` - [in] Stop also when + ``||r_k|| <= absolute_tolerance``. Default: ``1e-30``. +- ``pivot_floor`` - [in] Floor on ``|D_{kk}|`` during the per-tile + LDLT pivoting; keeps the preconditioner invertible on near-singular + diagonal tiles. Default: ``1e-12``. +- ``check_period`` - [in] Number of PCG iterations between host-side + convergence polls. Higher values reduce host syncs but may do up to + ``check_period - 1`` extra iterations after convergence. Default: + ``4``. + +BlockSparsePCGSolver +-------------------- + +Concrete implementation in `block_sparse_pcg_solver.h`. Solves +``H x = b`` for symmetric positive-(semi-)definite ``H`` with the +standard preconditioned conjugate gradient recurrence (Saad, +*Iterative Methods for Sparse Linear Systems*, §9.2) and a +block-Jacobi preconditioner formed from the dense diagonal tiles of +``H``. Each tile is factored independently with an in-shared-memory +LDLT. SpMV is delegated to cuSPARSE. + +All scalar quantities (``alpha``, ``beta``, ````, ````, +````) live on the device for the whole inner loop. Only the +residual norm is copied to the host, and only once every +``check_period`` iterations. + +.. cpp:function:: BlockSparsePCGSolver(BlockSparsePCGOptions options = BlockSparsePCGOptions()) + + :param ``options``: [in] Convergence / preconditioner-layout knobs. + See ``BlockSparsePCGOptions``. + :returns: Constructor has no return value. + +.. cpp:function:: bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, const Problem& problem, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) + + Builds the cuSPARSE SpMV plan, allocates the PCG scratch vectors, + and (when ``options.block_layout`` is empty) derives the block-Jacobi + layout from ``problem.GetStateBatches()``. + + :param ``stream``: [in] CUDA stream for asynchronous GPU operations. + :param ``problem``: [in] Source of the block-Jacobi layout when + ``options.block_layout`` is empty. Pass a default-constructed + ``Problem`` when calling on a raw matrix. + :param ``spd_matrix``: [in] Coefficient matrix ``H`` in CSR format. + :param ``rhs``: [in] Right-hand-side vector ``b``. + :param ``result``: [out] Solution vector ``x``. + :returns: ``true`` on success; ``false`` on dimension mismatch or + invalid layout. + +.. cpp:function:: bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) + + Rebuilds the block-Jacobi factor from the current values of ``H`` + and runs up to ``options.max_iterations`` PCG iterations with the + zero initial guess ``x_0 = 0``. + + :param ``stream``: [in] CUDA stream for asynchronous GPU operations. + :param ``spd_matrix``: [in] Coefficient matrix ``H`` (same structure + as in ``Initialize``). + :param ``rhs``: [in] Right-hand-side vector ``b``. + :param ``result``: [out] Solution vector ``x``. + :returns: ``true`` on success. + +.. cpp:function:: int BlockSparsePCGSolver::LastIterations() const + + :returns: Number of PCG iterations consumed by the most recent + ``Solve`` call. Useful for convergence diagnostics. + cuDSSLinearSolverMode --------------------- diff --git a/docs/sphinx/api/minimizer.rst b/docs/sphinx/api/minimizer.rst index 727b43a..4ea9d4d 100644 --- a/docs/sphinx/api/minimizer.rst +++ b/docs/sphinx/api/minimizer.rst @@ -130,17 +130,22 @@ Used when constructing a :code:`GaussNewtonMinimizer`. (cost increases or step quality below acceptance threshold) this many times in a row, the minimizer treats the current solution as converged. Set to 0 to disable. Default: 5. -- **sparse_linear_solver_type** [in]: Linear backend; options are ``cuDSS`` - (sparse direct solver via NVIDIA's cuDSS library), ``DenseLDLT`` (converts - CSR to dense and solves with a custom CUDA pivoted LDLT factorization), +- **sparse_linear_solver_type** [in]: Linear backend; options are + ``BlockSparsePCG`` (block-Jacobi preconditioned conjugate gradient; + layout auto-derived from the problem's state batches), ``cuDSS`` (sparse + direct solver via NVIDIA's cuDSS library), ``DenseLDLT`` (converts CSR + to dense and solves with a custom CUDA pivoted LDLT factorization), ``DenseCholesky`` (converts CSR to dense and solves via cuSOLVER Cholesky; requires SPD matrix), and ``DenseQR`` (converts CSR to dense and solves via cuSOLVER QR factorization; works for any non-singular matrix). - Default: ``cuDSS``. -- **sparse_linear_solver_config** [in]: Backend-specific options. For cuDSS, - contains :code:`cudss_solver_options` (mode, e.g. SlowInitFastSolve; - :code:`nthreads`; optional :code:`threading_lib_path` for multi-threaded - cuDSS). + Default: ``BlockSparsePCG``. +- **sparse_linear_solver_config** [in]: Backend-specific options. For + ``BlockSparsePCG`` contains :code:`block_sparse_pcg_options` + (``block_size`` / ``block_layout``, ``max_iterations``, + ``relative_tolerance``, ``absolute_tolerance``, ``pivot_floor``, + ``check_period``). For ``cuDSS`` contains :code:`cudss_solver_options` + (mode, ``nthreads``, optional ``threading_lib_path`` for multi-threaded + cuDSS). Dense backends take no extra configuration. - **sparse_square_multiplier_type** [in]: Strategy for computing the approximate Hessian :math:`J^T J`; options are ``cuSPARSE`` (cuSPARSE SpGEMM reuse API) and ``Fast`` (warp-efficient CUDA kernels with bitmap pattern discovery). @@ -482,11 +487,14 @@ values and then override individual fields. acceptance threshold) are allowed before the minimizer treats the current estimate as converged. Set to ``0`` to disable this criterion. - **sparse_linear_solver_type** (``SparseLinearSolverType``, default - ``cuDSS``) — selects the linear-system backend. ``cuDSS`` uses NVIDIA's - sparse direct solver; ``DenseLDLT`` converts to dense and factorizes with - a custom pivoted LDLT kernel; ``DenseCholesky`` converts to dense and uses - cuSOLVER Cholesky (requires SPD); ``DenseQR`` converts to dense and uses - cuSOLVER QR factorization (works for any non-singular matrix). + ``BlockSparsePCG``) — selects the linear-system backend. + ``BlockSparsePCG`` runs block-Jacobi preconditioned conjugate gradient + with the block layout derived automatically from the problem's state + batches. ``cuDSS`` uses NVIDIA's sparse direct solver; ``DenseLDLT`` + converts to dense and factorizes with a custom pivoted LDLT kernel; + ``DenseCholesky`` converts to dense and uses cuSOLVER Cholesky + (requires SPD); ``DenseQR`` converts to dense and uses cuSOLVER QR + factorization (works for any non-singular matrix). - **sparse_square_multiplier_type** (``SparseMatrixMultiplierType``, default ``Fast``) — strategy for computing the approximate Hessian :math:`J^T J`. ``cuSPARSE`` uses the cuSPARSE SpGEMM reuse API; @@ -723,6 +731,10 @@ Creates an empty problem with no states or factors. Integer enum selecting the linear-system backend. +- ``SparseLinearSolverType.BlockSparsePCG`` (default) — block-Jacobi + preconditioned conjugate gradient solver. The block layout is derived + automatically from the problem's state batches at minimizer-initialize + time. - ``SparseLinearSolverType.cuDSS`` — sparse direct solver via NVIDIA cuDSS. - ``SparseLinearSolverType.DenseLDLT`` — converts CSR to dense and solves with a custom CUDA pivoted LDLT kernel. diff --git a/llms.txt b/llms.txt index 9548419..4ba9f96 100644 --- a/llms.txt +++ b/llms.txt @@ -10,7 +10,8 @@ cuNLS organizes optimization into: - `FactorBatch` implementations for batched residual/Jacobian evaluation. - `Problem` for factor-graph composition. - `GaussNewtonMinimizer` and `LevenbergMarquardtMinimizer` for solving. -- Sparse backends (`cuDSS`) and robust losses. +- Sparse backends (`BlockSparsePCG` default, plus `cuDSS` direct and + dense LDLT/Cholesky/QR fallbacks) and robust losses. ## Entry points diff --git a/python/pycunls/_pycunls_core.pyi b/python/pycunls/_pycunls_core.pyi index 901b15a..35f598a 100644 --- a/python/pycunls/_pycunls_core.pyi +++ b/python/pycunls/_pycunls_core.pyi @@ -53,6 +53,7 @@ class SparseLinearSolverType(enum.IntEnum): DenseLDLT = ... DenseCholesky = ... DenseQR = ... + BlockSparsePCG = ... class SparseMatrixMultiplierType(enum.IntEnum): cuSPARSE = ... diff --git a/python/src/bind_types.cpp b/python/src/bind_types.cpp index 97e7cb3..6edf47e 100644 --- a/python/src/bind_types.cpp +++ b/python/src/bind_types.cpp @@ -70,7 +70,8 @@ void bind_types(nb::module_ &m) { .value("cuDSS", cunls::SparseLinearSolverType::cuDSS) .value("DenseLDLT", cunls::SparseLinearSolverType::DenseLDLT) .value("DenseCholesky", cunls::SparseLinearSolverType::DenseCholesky) - .value("DenseQR", cunls::SparseLinearSolverType::DenseQR); + .value("DenseQR", cunls::SparseLinearSolverType::DenseQR) + .value("BlockSparsePCG", cunls::SparseLinearSolverType::BlockSparsePCG); nb::enum_(m, "SparseMatrixMultiplierType") .value("cuSPARSE", cunls::SparseMatrixMultiplierType::cuSPARSE) diff --git a/python/tests/test_minimizer.py b/python/tests/test_minimizer.py index c154225..96c1477 100644 --- a/python/tests/test_minimizer.py +++ b/python/tests/test_minimizer.py @@ -168,7 +168,8 @@ def test_defaults(self): assert opts.max_num_iterations == 50 assert opts.state_tolerance == pytest.approx(1e-6) assert opts.cost_tolerance == pytest.approx(1e-6) - assert opts.sparse_linear_solver_type == pycunls.SparseLinearSolverType.cuDSS + assert (opts.sparse_linear_solver_type + == pycunls.SparseLinearSolverType.BlockSparsePCG) assert opts.column_scaling == pycunls.ColumnScaling.none assert opts.disable_safety_checks is True diff --git a/tests/dense_cholesky_solver_test.cpp b/tests/dense_cholesky_solver_test.cpp index fcb07b6..fcadd07 100644 --- a/tests/dense_cholesky_solver_test.cpp +++ b/tests/dense_cholesky_solver_test.cpp @@ -166,7 +166,8 @@ TEST_F(DenseCholeskySolverTestFixture, SolveDenseSystemAcrossDifferentSizes) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -195,7 +196,8 @@ TEST_F(DenseCholeskySolverTestFixture, SolveReturnsFalseForZeroMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -219,7 +221,8 @@ TEST_F(DenseCholeskySolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -243,7 +246,8 @@ TEST_F(DenseCholeskySolverTestFixture, SolveReturnsFalseForIndefiniteMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -265,7 +269,8 @@ TEST_F(DenseCholeskySolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -287,7 +292,8 @@ TEST_F(DenseCholeskySolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); diff --git a/tests/dense_linear_solver_test.cpp b/tests/dense_linear_solver_test.cpp index 7484d8b..0c5044a 100644 --- a/tests/dense_linear_solver_test.cpp +++ b/tests/dense_linear_solver_test.cpp @@ -244,7 +244,8 @@ TEST_F(DenseLDLTSolverTestFixture, SolveDenseSystemAcrossDifferentSizes) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -286,7 +287,8 @@ TEST_F(DenseLDLTSolverTestFixture, SolveSymmetricIndefiniteSystem) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -399,7 +401,8 @@ TEST_F(DenseLDLTSolverTestFixture, SolveReturnsFalseForZeroMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -425,7 +428,8 @@ TEST_F(DenseLDLTSolverTestFixture, SolveReturnsFalseForRankDeficientMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -448,7 +452,8 @@ TEST_F(DenseLDLTSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -472,7 +477,8 @@ TEST_F(DenseLDLTSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); diff --git a/tests/dense_qr_solver_test.cpp b/tests/dense_qr_solver_test.cpp index e05672e..4082600 100644 --- a/tests/dense_qr_solver_test.cpp +++ b/tests/dense_qr_solver_test.cpp @@ -166,7 +166,8 @@ TEST_F(DenseQRSolverTestFixture, SolveDenseSystemAcrossDifferentSizes) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -199,7 +200,8 @@ TEST_F(DenseQRSolverTestFixture, SolveSymmetricIndefiniteSystem) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); @@ -227,7 +229,8 @@ TEST_F(DenseQRSolverTestFixture, SolveReturnsFalseForZeroMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -250,7 +253,8 @@ TEST_F(DenseQRSolverTestFixture, SolveReturnsFalseForRankDeficientMatrix) { dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -272,7 +276,8 @@ TEST_F(DenseQRSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_FALSE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); } @@ -294,7 +299,8 @@ TEST_F(DenseQRSolverTestFixture, dvector rhs(rhs_host); dvector result(n); - ASSERT_TRUE(solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream_.GetStream(), Problem(), matrix, rhs, result)); ASSERT_TRUE(solver.Solve(stream_.GetStream(), matrix, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream_.GetStream())); diff --git a/tests/pgo_minimizer_test.cpp b/tests/pgo_minimizer_test.cpp index 0c377f0..6ef2876 100644 --- a/tests/pgo_minimizer_test.cpp +++ b/tests/pgo_minimizer_test.cpp @@ -58,9 +58,9 @@ #include "cunls/factor/information_factor_batch.h" #include "cunls/factor/se3_between_factor_batch.h" #include "cunls/minimizer/levenberg_marquardt_minimizer.h" -#include "tests/utils.h" #include "cunls/minimizer/problem.h" #include "cunls/state/se3_state_batch.h" +#include "tests/utils.h" namespace cunls { @@ -300,9 +300,12 @@ TEST_F(PgoMinimizerTestFixture, Optimize) { options.cost_tolerance = 1e-2f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(200); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(200); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1000.0; diff --git a/tests/sba_minimizer_test.cpp b/tests/sba_minimizer_test.cpp index 2970e0f..26a37d1 100644 --- a/tests/sba_minimizer_test.cpp +++ b/tests/sba_minimizer_test.cpp @@ -377,9 +377,12 @@ TEST_F(SbaMinimizerTestFixture, OptimizeAndCheckConvergence) { // block_size=3 keeps the preconditioner cheap and well-conditioned for the // landmark blocks while still capturing useful structure inside the 6x6 // pose tiles (every 6x6 splits into a 2x2 grid of 3x3 sub-blocks). - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(3); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(3); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; diff --git a/tests/sparse_linear_solver_test.cpp b/tests/sparse_linear_solver_test.cpp index c111199..e59ae52 100644 --- a/tests/sparse_linear_solver_test.cpp +++ b/tests/sparse_linear_solver_test.cpp @@ -34,10 +34,10 @@ #include #include "cunls/common/cuda_stream.h" -#include "cunls/minimizer/problem.h" #include "cunls/common/helper.h" #include "cunls/common/profiler.h" #include "cunls/common/types.h" +#include "cunls/minimizer/problem.h" #include "tests/utils.h" namespace cunls { @@ -187,7 +187,8 @@ TEST(SparseLinearSolverTest, Solve) { cuDSSLinearSolver solver(cudss_solver_options); { profiler::ScopedRange range("Warm up"); - solver.Initialize(stream.GetStream(), Problem(), input_matrix, rhs, result); + solver.Initialize(stream.GetStream(), Problem(), input_matrix, rhs, + result); solver.Solve(stream.GetStream(), input_matrix, rhs, result); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); } @@ -311,7 +312,8 @@ TEST(SparseLinearSolverTest, BlockSparsePCGSolve) { opts.relative_tolerance = 1e-5f; opts.max_iterations = 500; BlockSparsePCGSolver solver(opts); - ASSERT_TRUE(solver.Initialize(stream.GetStream(), Problem(), mat, rhs, result)); + ASSERT_TRUE( + solver.Initialize(stream.GetStream(), Problem(), mat, rhs, result)); ASSERT_TRUE(solver.Solve(stream.GetStream(), mat, rhs, result)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); diff --git a/tests/synthetic_pgo_test.cpp b/tests/synthetic_pgo_test.cpp index 29e2af1..4d1ac34 100644 --- a/tests/synthetic_pgo_test.cpp +++ b/tests/synthetic_pgo_test.cpp @@ -263,9 +263,12 @@ TEST_F(SyntheticPGOTest, OptimizeConsecutiveBetweenConstraints) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); // GaussNewtonMinimizer minimizer(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; @@ -358,9 +361,12 @@ TEST_F(SyntheticPGOTest, InformationBetweenFactorBatch) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); // GaussNewtonMinimizer minimizer(options); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; @@ -438,9 +444,12 @@ TEST_F(SyntheticPGOTest, WeightedWrapsInformationBetweenFactorBatch) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -516,9 +525,12 @@ TEST_F(SyntheticPGOTest, InformationWrapsWeightedBetweenFactorBatch) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = false; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -637,8 +649,7 @@ class LoopClosurePGOTest : public ::testing::TestWithParam { } cuBLASHandle cublas_handle_; - profiler::Domain profiler_domain_ = - profiler::Domain("LoopClosurePGOTest"); + profiler::Domain profiler_domain_ = profiler::Domain("LoopClosurePGOTest"); }; TEST_P(LoopClosurePGOTest, Optimize) { @@ -673,8 +684,7 @@ TEST_P(LoopClosurePGOTest, Optimize) { inv[15] = 1.f; return inv; }; - auto mul = [](const SE3Transform &a, - const SE3Transform &b) -> SE3Transform { + auto mul = [](const SE3Transform &a, const SE3Transform &b) -> SE3Transform { SE3Transform c{}; for (int rr = 0; rr < 4; ++rr) { for (int cc = 0; cc < 4; ++cc) { @@ -732,10 +742,8 @@ TEST_P(LoopClosurePGOTest, Optimize) { std::vector state_pointers; state_pointers.reserve(deltas.size() * 2); for (size_t e = 0; e < deltas.size(); ++e) { - state_pointers.push_back( - pose_batch.StateBlockDevicePtr(edges_left[e])); - state_pointers.push_back( - pose_batch.StateBlockDevicePtr(edges_right[e])); + state_pointers.push_back(pose_batch.StateBlockDevicePtr(edges_left[e])); + state_pointers.push_back(pose_batch.StateBlockDevicePtr(edges_right[e])); } Problem problem; @@ -750,9 +758,12 @@ TEST_P(LoopClosurePGOTest, Optimize) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = true; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(400); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(400); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-4f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f; @@ -769,11 +780,11 @@ TEST_P(LoopClosurePGOTest, Optimize) { EXPECT_LE(summary.final_cost, summary.initial_cost + 1e-3f); } -INSTANTIATE_TEST_SUITE_P( - Sizes, LoopClosurePGOTest, - ::testing::Values(LcPgoParams{100, 300, "P100_LC300"}, - LcPgoParams{500, 1000, "P500_LC1k"}, - LcPgoParams{1000, 3000, "P1k_LC3k"}, - LcPgoParams{5000, 10000, "P5k_LC10k"})); +INSTANTIATE_TEST_SUITE_P(Sizes, LoopClosurePGOTest, + ::testing::Values(LcPgoParams{100, 300, "P100_LC300"}, + LcPgoParams{500, 1000, "P500_LC1k"}, + LcPgoParams{1000, 3000, "P1k_LC3k"}, + LcPgoParams{5000, 10000, + "P5k_LC10k"})); } // namespace cunls diff --git a/tests/synthetic_sba_test.cpp b/tests/synthetic_sba_test.cpp index bd2fce0..c894cad 100644 --- a/tests/synthetic_sba_test.cpp +++ b/tests/synthetic_sba_test.cpp @@ -51,12 +51,12 @@ namespace cunls { struct SyntheticSbaParams { int n_poses; int n_points; - int obs_per_landmark; // visibility per landmark; total obs = n_points * obs_per_landmark. + int obs_per_landmark; // visibility per landmark; total obs = n_points * + // obs_per_landmark. const char *label; }; -inline std::ostream &operator<<(std::ostream &os, - const SyntheticSbaParams &p) { +inline std::ostream &operator<<(std::ostream &os, const SyntheticSbaParams &p) { return os << p.label; } @@ -90,14 +90,18 @@ class SyntheticSbaTest : public ::testing::TestWithParam { * normalized image observation. Returns false if the point is behind * the camera or too close to it. */ static bool Project(const SE3Transform &pose_world_from_cam_or_cam_from_world, - const Vector<3> &p_world, Vector<2> &out, bool cam_from_world) { + const Vector<3> &p_world, Vector<2> &out, + bool cam_from_world) { const SE3Transform &T = pose_world_from_cam_or_cam_from_world; Vector<3> p_cam{}; if (cam_from_world) { // T transforms world->cam. - p_cam[0] = T[0] * p_world[0] + T[1] * p_world[1] + T[2] * p_world[2] + T[3]; - p_cam[1] = T[4] * p_world[0] + T[5] * p_world[1] + T[6] * p_world[2] + T[7]; - p_cam[2] = T[8] * p_world[0] + T[9] * p_world[1] + T[10] * p_world[2] + T[11]; + p_cam[0] = + T[0] * p_world[0] + T[1] * p_world[1] + T[2] * p_world[2] + T[3]; + p_cam[1] = + T[4] * p_world[0] + T[5] * p_world[1] + T[6] * p_world[2] + T[7]; + p_cam[2] = + T[8] * p_world[0] + T[9] * p_world[1] + T[10] * p_world[2] + T[11]; } else { // T transforms cam->world; compute inverse. // Inverse of [R t; 0 1] is [R^T -R^T t; 0 1]. @@ -305,8 +309,8 @@ TEST_P(SyntheticSbaTest, Optimize) { const_point_ids.size()); InformationFactorBatch info_factor( - cublas_handle_, info_d.data(), n_obs, obs_d.data(), - cam_from_rig_d.data(), n_obs, 1e-3f); + cublas_handle_, info_d.data(), n_obs, obs_d.data(), cam_from_rig_d.data(), + n_obs, 1e-3f); std::vector state_pointers; state_pointers.reserve(n_obs * 2); @@ -331,9 +335,12 @@ TEST_P(SyntheticSbaTest, Optimize) { options.cost_tolerance = 1e-6f; options.disable_safety_checks = true; options.sparse_linear_solver_type = test_utils::SolverTypeFromEnv(); - options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = test_utils::PCGBlockSizeFromEnv(6); - options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = test_utils::PCGMaxIterFromEnv(200); - options.sparse_linear_solver_config.block_sparse_pcg_options.relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); + options.sparse_linear_solver_config.block_sparse_pcg_options.block_size = + test_utils::PCGBlockSizeFromEnv(6); + options.sparse_linear_solver_config.block_sparse_pcg_options.max_iterations = + test_utils::PCGMaxIterFromEnv(200); + options.sparse_linear_solver_config.block_sparse_pcg_options + .relative_tolerance = test_utils::PCGTolFromEnv(1e-3f); LevenbergMarquardtMinimizerOptions lm_options; lm_options.base_options = options; lm_options.initial_lambda = 1e-3f;