From 2b0428d9ea4d702baec586ea063b35e6e735cfbe Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Sat, 1 Aug 2026 20:28:38 -0700 Subject: [PATCH 1/7] Hessian block assembly + BSR --- CMakeLists.txt | 4 +- cunls/common/llms.txt | 2 +- cunls/common/types.h | 85 +- .../linear_solver/block_sparse_pcg_solver.cu | 520 +++++-- cunls/linear_solver/block_sparse_pcg_solver.h | 51 +- .../linear_solver/csr_sparse_linear_solver.h | 30 +- .../cudss_sparse_linear_solver.h | 34 +- cunls/linear_solver/dense_cholesky_solver.h | 16 +- cunls/linear_solver/dense_linear_solver.h | 26 +- cunls/linear_solver/dense_qr_solver.h | 16 +- cunls/minimizer/CMakeLists.txt | 8 +- cunls/minimizer/block_hessian_assembler.cu | 371 +++++ cunls/minimizer/block_hessian_assembler.h | 147 ++ cunls/minimizer/bsr_matrix.cu | 532 +++++++ cunls/minimizer/bsr_matrix.h | 141 ++ .../minimizer/cusparse_matrix_multiplier.cpp | 209 --- cunls/minimizer/cusparse_matrix_multiplier.h | 94 -- cunls/minimizer/fast_matrix_multiplier.cu | 650 --------- cunls/minimizer/fast_matrix_multiplier.h | 56 - cunls/minimizer/gauss_newton_minimizer.cu | 98 +- cunls/minimizer/gauss_newton_minimizer.h | 95 +- cunls/minimizer/hessian_structure.cu | 682 +++++++++ cunls/minimizer/hessian_structure.h | 149 ++ cunls/minimizer/jacobian_ops.cu | 396 ----- .../levenberg_marquardt_minimizer.cpp | 92 +- .../minimizer/levenberg_marquardt_minimizer.h | 29 +- cunls/minimizer/llms.txt | 29 +- cunls/minimizer/minimizer_state.cu | 6 + cunls/minimizer/minimizer_state.h | 31 +- cunls/minimizer/normal_equations.cu | 120 ++ cunls/minimizer/normal_equations.h | 146 ++ cunls/minimizer/residual_batch.cu | 17 +- cunls/minimizer/sparse_matrix.cu | 540 ++----- cunls/minimizer/sparse_matrix.h | 133 +- cunls/minimizer/sparse_matrix_multiplier.cpp | 39 - cunls/minimizer/sparse_matrix_multiplier.h | 98 -- cunls/state/state_batch_ops.cu | 70 +- cunls/state/state_batch_ops.h | 32 +- docs/sphinx/api/common.rst | 17 +- docs/sphinx/api/minimizer.rst | 122 +- python/pycunls/__init__.py | 2 - python/pycunls/_pycunls_core.pyi | 6 - python/src/bind_types.cpp | 9 +- python/tests/test_minimizer.py | 5 +- tests/block_hessian_assembler_test.cpp | 1288 +++++++++++++++++ tests/cusparse_matrix_multiplier_test.cpp | 245 ---- tests/fast_matrix_multiplier_test.cpp | 521 ------- tests/gauss_newton_test.cpp | 26 - tests/jacobian_ops_test.cpp | 146 -- tests/sparse_matrix_test.cpp | 365 +---- tests/utils.h | 42 +- 51 files changed, 4658 insertions(+), 3930 deletions(-) create mode 100644 cunls/minimizer/block_hessian_assembler.cu create mode 100644 cunls/minimizer/block_hessian_assembler.h create mode 100644 cunls/minimizer/bsr_matrix.cu create mode 100644 cunls/minimizer/bsr_matrix.h delete mode 100644 cunls/minimizer/cusparse_matrix_multiplier.cpp delete mode 100644 cunls/minimizer/cusparse_matrix_multiplier.h delete mode 100644 cunls/minimizer/fast_matrix_multiplier.cu delete mode 100644 cunls/minimizer/fast_matrix_multiplier.h create mode 100644 cunls/minimizer/hessian_structure.cu create mode 100644 cunls/minimizer/hessian_structure.h delete mode 100644 cunls/minimizer/jacobian_ops.cu create mode 100644 cunls/minimizer/normal_equations.cu create mode 100644 cunls/minimizer/normal_equations.h delete mode 100644 cunls/minimizer/sparse_matrix_multiplier.cpp delete mode 100644 cunls/minimizer/sparse_matrix_multiplier.h create mode 100644 tests/block_hessian_assembler_test.cpp delete mode 100644 tests/cusparse_matrix_multiplier_test.cpp delete mode 100644 tests/fast_matrix_multiplier_test.cpp delete mode 100644 tests/jacobian_ops_test.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 30f8d82..36c0330 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -142,11 +142,9 @@ if(BUILD_TESTING) tests/prior_vector_prior_factor_test.cpp tests/residual_batch_test.cpp tests/loss_function_test.cpp - tests/jacobian_ops_test.cpp tests/state_batch_ops_test.cpp tests/sparse_matrix_test.cpp - tests/cusparse_matrix_multiplier_test.cpp - tests/fast_matrix_multiplier_test.cpp + tests/block_hessian_assembler_test.cpp tests/sparse_linear_solver_test.cpp tests/dense_linear_solver_test.cpp tests/dense_cholesky_solver_test.cpp diff --git a/cunls/common/llms.txt b/cunls/common/llms.txt index ad74e0f..2f86e43 100644 --- a/cunls/common/llms.txt +++ b/cunls/common/llms.txt @@ -12,7 +12,7 @@ Shared infrastructure used across cuNLS modules: ## Key files - `device_vector.h`: `DeviceVector` (`dvector` alias in `types.h`) -- `types.h`: `Vector`, `Matrix`, `SE3Transform`, `CSRSparseMatrix`, `SparseJacobian` +- `types.h`: `Vector`, `Matrix`, `SE3Transform`, `CSRSparseMatrix`, `BSRSparseMatrix`, `PerFactorJacobians` - `cuda_stream.h`: RAII CUDA stream wrapper - `cublas_helper.h`, `cusparse_helper.h`, `cusolver_helper.h`, `cudss_helper.h` - `profiler.h`: NVTX range/domain wrappers diff --git a/cunls/common/types.h b/cunls/common/types.h index 187126d..8510f6b 100644 --- a/cunls/common/types.h +++ b/cunls/common/types.h @@ -33,7 +33,8 @@ namespace cunls { * * @tparam Dim Number of elements in the vector. */ -template using Vector = cuda::std::array; +template +using Vector = cuda::std::array; /** * @brief Fixed-size square matrix of floats stored in row-major order. @@ -43,7 +44,8 @@ template using Vector = cuda::std::array; * * @tparam Dim Number of rows (and columns) in the square matrix. */ -template using Matrix = cuda::std::array; +template +using Matrix = cuda::std::array; /** * @brief SE(3) transformation matrix representation. @@ -68,19 +70,22 @@ using SL4Transform = Matrix<4>; * @brief Alias for a device (GPU) vector. * @tparam T Element type (must be trivially copyable). */ -template using dvector = DeviceVector; +template +using dvector = DeviceVector; /** * @brief Alias for a host (CPU) vector (std::vector). * @tparam T Element type. */ -template using hvector = std::vector; +template +using hvector = std::vector; /** * @brief Alias for a pinned (CPU) vector (PinnedVector). * @tparam T Element type. */ -template using pvector = PinnedVector; +template +using pvector = PinnedVector; /** * @brief Compressed Sparse Row (CSR) matrix stored in GPU memory. @@ -130,16 +135,14 @@ struct CSRMatrixDimensions { */ struct CSRSparseMatrix { dvector row_offsets; ///< Row offset array (num_rows + 1 entries). - dvector col_ids; ///< Column index array (num_nonzeros entries). - dvector values; ///< Non-zero value array (num_nonzeros entries). + dvector col_ids; ///< Column index array (num_nonzeros entries). + dvector values; ///< Non-zero value array (num_nonzeros entries). /** * @brief Returns the number of rows in the matrix. * @return Number of rows (row_offsets.size() - 1). */ - size_t NumRows() const { - return row_offsets.empty() ? 0 : row_offsets.size() - 1; - } + size_t NumRows() const { return row_offsets.empty() ? 0 : row_offsets.size() - 1; } /** * @brief Returns the number of non-zero entries in the matrix. @@ -149,25 +152,61 @@ struct CSRSparseMatrix { }; /** - * @brief COO-like sparse structure storing row and column indices. + * @brief Block Sparse Row (BSR) matrix stored in GPU memory. + * + * A square matrix partitioned into uniform `block_size x block_size` tiles. + * Only tiles containing at least one structural non-zero are stored: + * - row_offsets: index into col_ids/values-tiles per block row + * (size = num_block_rows + 1). + * - col_ids: block-column index of each stored tile (size = NumBlocks()). + * - values: tiles laid out row-major and contiguously, i.e. entry (k, l) of + * tile `t` lives at `values[t * block_size * block_size + k * block_size + l]` + * (CUSPARSE_DIRECTION_ROW). + * + * The Hessian of a factor graph is naturally block structured: every state + * block contributes a dense tile per neighbour. Storing it this way keeps one + * column index per tile instead of one per scalar entry, which is where the + * SpMV bandwidth saving comes from — the values are identical either way. * - * Used as a triplet/coordinate format for building sparse matrices - * before conversion to CSR. + * Requires every state block's tangent dimension to be a multiple of + * `block_size`; see ChooseHessianBlockSize(). */ -struct TripletSparseStructure { - dvector row_ids; ///< Row indices of non-zero entries. - dvector col_ids; ///< Column indices of non-zero entries. +struct BSRSparseMatrix { + dvector row_offsets; ///< Block-row offsets (num_block_rows + 1 entries). + dvector col_ids; ///< Block-column index per tile. + dvector values; ///< Tiles, row-major, block_size^2 floats each. + + int block_size = 1; ///< Tile edge length. + /** + * @brief Largest number of tiles in any block row. + * + * Selects the SpMV schedule. Factor-graph Hessians come in two shapes: a + * pose graph is near-uniform with a handful of tiles per row, while bundle + * adjustment is extremely skewed (a pose row holds one tile per observation + * of that camera, a landmark row a handful). One schedule cannot serve both. + */ + int max_tiles_per_row = 0; + int num_block_rows = 0; ///< Number of block rows (= block columns). + + /** @brief Number of stored tiles. */ + size_t NumBlocks() const { return col_ids.size(); } + + /** @brief Scalar row/column count. */ + int NumRows() const { return num_block_rows * block_size; } + + /** @brief Number of stored scalar entries (tiles are dense). */ + size_t NumNonZeros() const { return values.size(); } }; /** - * @brief Sparse Jacobian matrix in triplet format. + * @brief Per-factor dense Jacobian blocks, concatenated across residual + * batches. * - * Stores both the sparsity structure (row/column indices) and the - * non-zero values of the Jacobian. + * Each factor batch writes `NumFactors()` dense row-major blocks of + * `ResidualsSize() x sum(StateBlockSizes())` floats, and the batches are laid + * out back to back. There is no global sparse Jacobian: the Hessian is + * assembled from these blocks directly (see BlockHessianAssembler). */ -struct SparseJacobian { - TripletSparseStructure structure; ///< Row and column indices. - dvector values; ///< Non-zero Jacobian values. -}; +using PerFactorJacobians = dvector; } // namespace cunls diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index da519a8..b7857a2 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -133,13 +133,298 @@ constexpr int kMaxBlockSize = 16; * * @tparam B Compile-time tile side length. */ + +/** + * @brief Gathers one row of a diagonal tile out of scalar CSR. + * + * Column indices are sorted within a row, so the scan can stop as soon as it + * walks past the tile. For SBA pose rows, whose many landmark columns sort + * after the diagonal tile, that turns an O(nnz_per_row) scan into O(B). + */ +__device__ __forceinline__ void GatherTileRowCSR(const int *__restrict__ row_off, + const int *__restrict__ col_idx, + const float *__restrict__ values, int col_lo, + int B, int rr, float *tile_row) { + const int global_row = col_lo + rr; + const int col_hi = col_lo + B; + for (int k = row_off[global_row]; k < row_off[global_row + 1]; ++k) { + int c = col_idx[k]; + if (c >= col_hi) { + break; + } + if (c >= col_lo) { + tile_row[c - col_lo] = values[k]; + } + } +} + +/** + * @brief Same, out of uniform block storage. + * + * The scan walks the enclosing block row, which holds one index per tile + * instead of one per scalar column -- `block_size` times shorter than the CSR + * row it replaces. `B` need not be a multiple of `block_size`; the overlap + * test below handles partial tiles. + */ +__device__ __forceinline__ void GatherTileRowBSR(const int *__restrict__ row_off, + const int *__restrict__ col_idx, + const float *__restrict__ values, int block_size, + int col_lo, int B, int rr, float *tile_row) { + const int global_row = col_lo + rr; + const int block_row = global_row / block_size; + const int sub_row = global_row - block_row * block_size; + const int col_hi = col_lo + B; + const int tile_area = block_size * block_size; + for (int t = row_off[block_row]; t < row_off[block_row + 1]; ++t) { + const int c0 = col_idx[t] * block_size; + if (c0 >= col_hi) { + break; + } + if (c0 + block_size <= col_lo) { + continue; + } + const float *src = values + static_cast(t) * tile_area + sub_row * block_size; + for (int l = 0; l < block_size; ++l) { + const int c = c0 + l; + if (c >= col_lo && c < col_hi) { + tile_row[c - col_lo] = src[l]; + } + } + } +} + +/** + * @brief y = A * x for uniform block storage, one warp per block row. + * + * cuSPARSE's `cusparseSbsrmv` measured 3.6x slower than `csrmv_v3` on a + * bundle-adjustment Hessian with 3x3 tiles, which would negate the point of + * block storage, so the SpMV is written here. + * + * A warp per block row rather than a thread per row: bundle-adjustment Hessians + * are wildly non-uniform -- a pose block row holds one tile per observation of + * that camera (thousands) while a landmark row holds a handful -- so a + * thread-per-row schedule leaves the few pose threads serializing for as long + * as the whole kernel takes. Lanes stride over the row's tiles instead, and a + * butterfly reduction combines their partial `b`-vectors. + * + * @tparam kB Tile edge. + */ +template +__global__ void BsrMultiplyWarpKernel(int num_block_rows, const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, const float *__restrict__ x, + float *__restrict__ y) { + const int block_row = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + const int lane = threadIdx.x & 31; + if (block_row >= num_block_rows) { + return; + } + + const int end = row_offsets[block_row + 1]; + float acc[kB]; +#pragma unroll + for (int k = 0; k < kB; ++k) { + acc[k] = 0.f; + } + + for (int t = row_offsets[block_row] + lane; t < end; t += 32) { + const float *tile = values + static_cast(t) * kB * kB; + const float *xs = x + static_cast(col_ids[t]) * kB; + float xv[kB]; +#pragma unroll + for (int l = 0; l < kB; ++l) { + xv[l] = xs[l]; + } +#pragma unroll + for (int k = 0; k < kB; ++k) { +#pragma unroll + for (int l = 0; l < kB; ++l) { + acc[k] = fmaf(tile[k * kB + l], xv[l], acc[k]); + } + } + } + + // Butterfly rather than shfl_down so every lane ends with the totals and the + // first kB lanes can write the output run coalesced. +#pragma unroll + for (int k = 0; k < kB; ++k) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + acc[k] += __shfl_xor_sync(0xFFFFFFFFu, acc[k], offset); + } + } + if (lane < kB) { + y[block_row * kB + lane] = acc[lane]; + } +} + +/** + * @brief Runtime-tile-edge fallback for block sizes without a specialization. + * + * Same schedule as BsrMultiplyWarpKernel; the accumulator is sized to the + * largest edge ChooseHessianBlockSize can return. + */ +__global__ void BsrMultiplyWarpGenericKernel(int num_block_rows, int block_size, + const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, + const float *__restrict__ x, float *__restrict__ y) { + constexpr int kMaxBlockSize = 16; + const int block_row = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + const int lane = threadIdx.x & 31; + if (block_row >= num_block_rows) { + return; + } + + const int b = block_size; + const int end = row_offsets[block_row + 1]; + float acc[kMaxBlockSize]; + for (int k = 0; k < b; ++k) { + acc[k] = 0.f; + } + + for (int t = row_offsets[block_row] + lane; t < end; t += 32) { + const float *tile = values + static_cast(t) * b * b; + const float *xs = x + static_cast(col_ids[t]) * b; + for (int k = 0; k < b; ++k) { + float a = 0.f; + for (int l = 0; l < b; ++l) { + a = fmaf(tile[k * b + l], xs[l], a); + } + acc[k] += a; + } + } + + for (int k = 0; k < b; ++k) { + for (int offset = 16; offset > 0; offset >>= 1) { + acc[k] += __shfl_xor_sync(0xFFFFFFFFu, acc[k], offset); + } + } + if (lane < b) { + y[block_row * b + lane] = acc[lane]; + } +} + +/** + * @brief y = A * x, one thread per scalar row. + * + * The right schedule for near-uniform, short block rows -- a pose graph has a + * handful of tiles per row, so a whole warp per row would leave most lanes idle + * and pay for a reduction that spans mostly zeros. Thread `i` walks row + * `i % b` of every tile in block row `i / b`; the `b` consecutive threads + * sharing a block row read each tile as one contiguous run and broadcast their + * identical `col_ids` and `x` loads. + * + * @tparam kB Tile edge, or 0 to take it as a runtime argument. + */ +template +__global__ void BsrMultiplyRowKernel(int num_rows, int runtime_block_size, + const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, const float *__restrict__ x, + float *__restrict__ y) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= num_rows) { + return; + } + const int b = kB > 0 ? kB : runtime_block_size; + const int block_row = row / b; + const int sub_row = row - block_row * b; + + float acc = 0.f; + const int end = row_offsets[block_row + 1]; + for (int t = row_offsets[block_row]; t < end; ++t) { + const float *tile = values + static_cast(t) * b * b + sub_row * b; + const float *xs = x + static_cast(col_ids[t]) * b; + for (int l = 0; l < b; ++l) { + acc = fmaf(tile[l], xs[l], acc); + } + } + y[row] = acc; +} + +/** + * @brief Launches the BSR SpMV, choosing a schedule from the row lengths. + * + * A warp per block row tolerates skew but wastes lanes on short rows; a thread + * per scalar row is the opposite. Bundle adjustment needs the former (pose + * rows hold thousands of tiles), pose graphs the latter (every row holds a + * handful), so the peak row length decides. + */ +void LaunchBsrMultiply(cudaStream_t stream, int num_block_rows, int block_size, + int max_tiles_per_row, const int *row_offsets, const int *col_ids, + const float *values, const float *x, float *y) { + constexpr int kThreads = 128; + constexpr int kSkewThreshold = 32; + + if (max_tiles_per_row < kSkewThreshold) { + const int num_rows = num_block_rows * block_size; + const int grid = (num_rows + kThreads - 1) / kThreads; + switch (block_size) { +#define LAUNCH_ROW(BVAL) \ + case BVAL: \ + BsrMultiplyRowKernel \ + <<>>(num_rows, BVAL, row_offsets, col_ids, values, x, y); \ + break + LAUNCH_ROW(2); + LAUNCH_ROW(3); + LAUNCH_ROW(4); + LAUNCH_ROW(6); + LAUNCH_ROW(7); + LAUNCH_ROW(15); + LAUNCH_ROW(16); +#undef LAUNCH_ROW + default: + BsrMultiplyRowKernel<0><<>>(num_rows, block_size, row_offsets, + col_ids, values, x, y); + break; + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); + return; + } + + const int warps_per_block = kThreads / 32; + const int grid = (num_block_rows + warps_per_block - 1) / warps_per_block; + switch (block_size) { +#define LAUNCH_WARP(BVAL) \ + case BVAL: \ + BsrMultiplyWarpKernel \ + <<>>(num_block_rows, row_offsets, col_ids, values, x, y); \ + break + LAUNCH_WARP(2); + LAUNCH_WARP(3); + LAUNCH_WARP(4); + LAUNCH_WARP(5); + LAUNCH_WARP(6); + LAUNCH_WARP(7); + LAUNCH_WARP(8); +#undef LAUNCH_WARP + default: + BsrMultiplyWarpGenericKernel<<>>( + num_block_rows, block_size, row_offsets, col_ids, values, x, y); + break; + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +/** @brief Dispatches the gather on the storage layout. */ +__device__ __forceinline__ void GatherTileRow(const int *__restrict__ row_off, + const int *__restrict__ col_idx, + const float *__restrict__ values, bool block_storage, + int block_size, int col_lo, int B, int rr, + float *tile_row) { + if (block_storage) { + GatherTileRowBSR(row_off, col_idx, values, block_size, col_lo, B, rr, tile_row); + } else { + GatherTileRowCSR(row_off, col_idx, values, col_lo, B, rr, tile_row); + } +} + template -__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) { +__global__ void ExtractAndFactorBlockDiagonalsKernel( + const int *__restrict__ row_off, const int *__restrict__ col_idx, + const float *__restrict__ values, bool block_storage, int block_size, 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; @@ -157,20 +442,9 @@ __global__ void ExtractAndFactorBlockDiagonalsKernel(const int *__restrict__ row // 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]; - 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_hi) { - break; - } - if (c >= col_lo) { - tile[tid * B + (c - col_lo)] = values[k]; - } - } + GatherTileRow(row_off, col_idx, values, block_storage, block_size, col_lo, B, tid, + tile + tid * B); } __syncthreads(); @@ -221,9 +495,10 @@ __global__ void ExtractAndFactorBlockDiagonalsKernel(const int *__restrict__ row * 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) { + const float *__restrict__ values, bool block_storage, + int block_size, 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; @@ -238,20 +513,9 @@ __global__ void ExtractAndFactorGenericKernel(const int *__restrict__ row_off, __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_hi) { - break; - } - if (c >= col_lo) { - tile[tid * B + (c - col_lo)] = values[k]; - } - } + GatherTileRow(row_off, col_idx, values, block_storage, block_size, col_lo, B, tid, + tile + tid * B); } __syncthreads(); if (tid < B) { @@ -306,7 +570,8 @@ __global__ void ExtractAndFactorGenericKernel(const int *__restrict__ row_off, template __global__ void ExtractAndFactorPerThreadKernel(const int *__restrict__ row_off, const int *__restrict__ col_idx, - const float *__restrict__ values, int row_start, + const float *__restrict__ values, + bool block_storage, int block_size, int row_start, int num_blocks, int factor_offset, float pivot_floor, float *__restrict__ factors) { int block_row = blockIdx.x * blockDim.x + threadIdx.x; @@ -320,27 +585,10 @@ __global__ void ExtractAndFactorPerThreadKernel(const int *__restrict__ row_off, tile[i] = 0.f; } 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; - 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_hi) { - break; - } - if (c >= col_lo) { - tile[rr * B + (c - col_lo)] = values[k]; - } - } + GatherTileRow(row_off, col_idx, values, block_storage, block_size, col_lo, B, rr, + tile + rr * B); } // Symmetrize numerically. #pragma unroll @@ -383,24 +631,16 @@ __global__ void ExtractAndFactorPerThreadKernel(const int *__restrict__ row_off, /** 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) { + const float *__restrict__ values, bool block_storage, + int block_size, 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; + float d = 0.f; + GatherTileRow(row_off, col_idx, values, block_storage, block_size, row_start + idx, 1, 0, &d); + factors[factor_offset + idx] = fmaxf(fabsf(d), pivot_floor); } // ============================================================================= @@ -742,31 +982,31 @@ __global__ void DualDotKernel(const float *__restrict__ a, const float *__restri /** 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) { + int factor_offset, float pivot_floor, const int *row_off, + const int *col_idx, const float *vals, bool block_storage, + int block_size, 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(); if (B == 1) { int threads = 256; int blocks = (num_blocks + threads - 1) / threads; - ExtractScalarJacobi<<>>( - row_off, col_idx, vals, row_start, num_blocks, factor_offset, pivot_floor, factors); + ExtractScalarJacobi<<>>(row_off, col_idx, vals, block_storage, + block_size, row_start, num_blocks, + factor_offset, pivot_floor, factors); THROW_ON_CUDA_ERROR(cudaGetLastError()); return; } -#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; \ +#define LAUNCH_FACTOR_PER_THREAD(BVAL) \ + case BVAL: { \ + int threads = 256; \ + int blocks = (num_blocks + threads - 1) / threads; \ + ExtractAndFactorPerThreadKernel<<>>( \ + row_off, col_idx, vals, block_storage, block_size, row_start, num_blocks, factor_offset, \ + pivot_floor, factors); \ + break; \ } // Small-B path: one block per thread. @@ -787,10 +1027,11 @@ void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, int num // 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); \ +#define LAUNCH_FACTOR(BVAL) \ + case BVAL: \ + ExtractAndFactorBlockDiagonalsKernel<<>>( \ + row_off, col_idx, vals, block_storage, block_size, row_start, num_blocks, factor_offset, \ + pivot_floor, factors); \ break switch (B) { @@ -801,7 +1042,8 @@ void DispatchExtractAndFactor(cudaStream_t stream, int B, int row_start, int num 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); + row_off, col_idx, vals, block_storage, block_size, B, row_start, num_blocks, + factor_offset, pivot_floor, factors); break; } } @@ -909,7 +1151,7 @@ void DualDotAsync(cudaStream_t stream, const float *a, const float *b, const flo DualDotKernel<<>>(a, b, c, n, out_ab, out_ac); } -} // namespace +} // namespace // ============================================================================= // BlockSparsePCGSolver @@ -987,10 +1229,27 @@ bool BlockSparsePCGSolver::BuildSegmentTables(int matrix_dim) { return true; } +bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, const Problem &problem, + const BSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) { + bsr_view_ = &spd_matrix; + csr_view_ = nullptr; + return InitializeCommon(stream, problem, spd_matrix.NumRows(), rhs, result); +} + bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { - int n = static_cast(spd_matrix.NumRows()); + csr_view_ = &spd_matrix; + bsr_view_ = nullptr; + if (!InitializeCommon(stream, problem, static_cast(spd_matrix.NumRows()), rhs, result)) { + return false; + } + return InitializeCsrSpMV(stream, spd_matrix); +} + +bool BlockSparsePCGSolver::InitializeCommon(cudaStream_t stream, const Problem &problem, int n, + const dvector &rhs, dvector &result) { if (n != static_cast(rhs.size()) || n != static_cast(result.size())) { LogError("BlockSparsePCGSolver: dim mismatch (matrix={}, rhs={}, result={})", n, rhs.size(), result.size()); @@ -1048,6 +1307,12 @@ bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, const Problem &proble d_scratch_.resize(7); } + return true; +} + +bool BlockSparsePCGSolver::InitializeCsrSpMV(cudaStream_t stream, + const CSRSparseMatrix &spd_matrix) { + const int n = static_cast(spd_matrix.NumRows()); // cuSPARSE SpMV setup. The descriptor is reused across all PCG steps // and across all Solve calls until the matrix structure changes. mat_desc_ = @@ -1077,6 +1342,28 @@ bool BlockSparsePCGSolver::Initialize(cudaStream_t stream, const Problem &proble return true; } +bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) { + const int n = 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; + } + bsr_view_ = &spd_matrix; + csr_view_ = nullptr; + if (n != matrix_size_) { + Problem empty_problem; + if (!InitializeCommon(stream, empty_problem, n, rhs, result)) { + return false; + } + } + return SolveCommon(stream, n, rhs, result); +} + bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) { int n = static_cast(spd_matrix.NumRows()); @@ -1088,6 +1375,8 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd if (n == 0) { return true; } + csr_view_ = &spd_matrix; + bsr_view_ = nullptr; if (n != matrix_size_) { // Recovery path: matrix dim changed since the last Initialize. Use // a default-constructed Problem; the cached options_.block_layout @@ -1102,17 +1391,35 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd // Refresh the cuSPARSE descriptor's value pointer; the matrix's // structure is unchanged so no re-preprocess is needed. mat_desc_.UpdatePointers(spd_matrix); + return SolveCommon(stream, n, rhs, result); +} + +/** + * The CG recurrence is identical for both storage layouts; only the SpMV and + * the preconditioner gather read the matrix, and both dispatch on which view + * pointer is set. + */ +bool BlockSparsePCGSolver::SolveCommon(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) { + const bool block_storage = bsr_view_ != nullptr; + const int *row_off = + block_storage ? bsr_view_->row_offsets.data() : csr_view_->row_offsets.data(); + const int *col_idx = block_storage ? bsr_view_->col_ids.data() : csr_view_->col_ids.data(); + const float *vals = block_storage ? bsr_view_->values.data() : csr_view_->values.data(); + const int storage_block = block_storage ? bsr_view_->block_size : 1; // ----------------------------------------------------------------- // 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()); + options_.pivot_floor, row_off, col_idx, vals, block_storage, + storage_block, precond_factors_.data()); } auto handle = static_cast(cusparse_handle_.GetHandle(stream)); - auto matA = static_cast(mat_desc_.GetDescription()); + auto matA = + block_storage ? nullptr : static_cast(mat_desc_.GetDescription()); // ----------------------------------------------------------------- // 2. Initialize PCG with x_0 = 0 ⇒ r_0 = b. @@ -1158,8 +1465,10 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd // 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(&vecY, n, Ap_.data(), CUDA_R_32F)); + if (!block_storage) { + THROW_ON_CUSPARSE_ERROR(cusparseCreateDnVec(&vecX, n, p_.data(), CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR(cusparseCreateDnVec(&vecY, n, Ap_.data(), CUDA_R_32F)); + } const int threads = 256; const int blocks = (n + threads - 1) / threads; @@ -1167,12 +1476,19 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd int it = 0; for (; it < options_.max_iterations; ++it) { - // Ap = H * p. + // Ap = H * p. Block storage carries one column index per tile instead of + // one per scalar entry, which is where the bandwidth saving comes from. 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())); + if (block_storage) { + LaunchBsrMultiply(stream, bsr_view_->num_block_rows, bsr_view_->block_size, + bsr_view_->max_tiles_per_row, row_off, col_idx, vals, p_.data(), + Ap_.data()); + } else { + 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); @@ -1208,10 +1524,12 @@ bool BlockSparsePCGSolver::Solve(cudaStream_t stream, const CSRSparseMatrix &spd } } } - WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecX)); - WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecY)); + if (!block_storage) { + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecX)); + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(vecY)); + } last_iterations_ = it; return true; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h index c400608..5b64e65 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.h +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -142,7 +142,7 @@ struct BlockSparsePCGOptions { * dominates many iterative solver implementations. * * Implementation notes (see the `.cu` file for derivations): - * - SpMV is delegated to cuSPARSE (`cusparseSpMV` with the default + * - Scalar-CSR 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 @@ -206,9 +206,8 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { * @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; + 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. @@ -229,8 +228,19 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { * reset). * @return true on success, false on dimension mismatch. */ - bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + /** @copydoc CSRSparseLinearSolver::SupportsBlockStorage */ + bool SupportsBlockStorage() const override { return true; } + + /** @copydoc CSRSparseLinearSolver::Initialize */ + bool Initialize(cudaStream_t stream, const Problem &problem, const BSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) override; + + /** @copydoc CSRSparseLinearSolver::Solve */ + bool Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) override; + + 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 @@ -257,6 +267,16 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { */ bool BuildSegmentTables(int matrix_dim); + /** @brief Layout-independent part of Initialize (segments and scratch). */ + bool InitializeCommon(cudaStream_t stream, const Problem &problem, int n, + const dvector &rhs, dvector &result); + + /** @brief Builds the generic-API SpMV plan used by the CSR path only. */ + bool InitializeCsrSpMV(cudaStream_t stream, const CSRSparseMatrix &spd_matrix); + + /** @brief The CG recurrence, shared by both storage layouts. */ + bool SolveCommon(cudaStream_t stream, int n, const dvector &rhs, dvector &result); + BlockSparsePCGOptions options_; /** Cached host copy of `options_.block_layout`, normalized to a * single uniform segment when the user didn't supply one. */ @@ -271,10 +291,10 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { /** 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_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_; @@ -292,9 +312,9 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { // ------------------------------------------------------------------ // PCG scratch // ------------------------------------------------------------------ - dvector r_; ///< residual `r_k` - dvector z_; ///< preconditioned residual `z_k = M^{-1} r_k` - dvector p_; ///< search direction `p_k` + 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) /** Device-resident scalar slots: alpha, beta, , rz_old, rz_new, @@ -308,6 +328,11 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { cuSPARSEMatrixDescription mat_desc_; dvector spmv_buffer_; ///< work buffer for cuSPARSE SpMV + /** Non-owning view of the matrix passed to the current Solve; exactly one + * of the two is non-null and selects the storage layout. */ + const CSRSparseMatrix *csr_view_ = nullptr; + const BSRSparseMatrix *bsr_view_ = nullptr; + /** 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 e2963bb..5432683 100644 --- a/cunls/linear_solver/csr_sparse_linear_solver.h +++ b/cunls/linear_solver/csr_sparse_linear_solver.h @@ -23,7 +23,7 @@ namespace cunls { -class Problem; // forward declaration; defined in cunls/minimizer/problem.h. +class Problem; // forward declaration; defined in cunls/minimizer/problem.h. /** * @brief Base class for linear solvers operating on CSR matrices. @@ -65,8 +65,7 @@ class CSRSparseLinearSolver { * @return true on success, false if a dimension mismatch is detected. */ virtual bool Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) = 0; /** @@ -86,6 +85,29 @@ class CSRSparseLinearSolver { virtual bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) = 0; + /** + * @brief Whether this backend can consume a block-stored matrix directly. + * + * The Hessian of a factor graph is naturally block structured, and assembling + * it that way keeps one column index per tile instead of one per scalar + * entry. Backends that say yes get the block form; the rest are handed an + * expanded CSR copy, so no caller has to care. + */ + virtual bool SupportsBlockStorage() const { return false; } + + /** @brief BSR counterpart of Initialize; only called when supported. */ + virtual bool Initialize(cudaStream_t stream, const Problem &problem, + const BSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) { + return false; + } + + /** @brief BSR counterpart of Solve; only called when supported. */ + virtual bool Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) { + return false; + } + /** * @brief Disables post-factorization safety checks. * @@ -105,7 +127,7 @@ class CSRSparseLinearSolver { */ virtual ~CSRSparseLinearSolver() = default; -protected: + protected: bool safety_checks_enabled_ = true; }; } // namespace cunls diff --git a/cunls/linear_solver/cudss_sparse_linear_solver.h b/cunls/linear_solver/cudss_sparse_linear_solver.h index dacf420..a45d8b4 100644 --- a/cunls/linear_solver/cudss_sparse_linear_solver.h +++ b/cunls/linear_solver/cudss_sparse_linear_solver.h @@ -32,9 +32,9 @@ namespace cunls { * Controls the trade-off between initialization time and solve time. */ enum class cuDSSLinearSolverMode { - SlowInitFastSolve, ///< Slower initialization, faster subsequent solves. - FastInitSlowSolve, ///< Faster initialization, slower subsequent solves (uses - ///< refactorization). + SlowInitFastSolve, ///< Slower initialization, faster subsequent solves. + FastInitSlowSolve, ///< Faster initialization, slower subsequent solves (uses + ///< refactorization). }; /** @@ -46,8 +46,8 @@ enum class cuDSSLinearSolverMode { struct cuDSSLinearSolverOptions { cuDSSLinearSolverMode mode = cuDSSLinearSolverMode::SlowInitFastSolve; ///< Solver mode controlling - ///< the init/solve trade-off. - int nthreads = 1; ///< Number of threads for host-side operations. + ///< the init/solve trade-off. + int nthreads = 1; ///< Number of threads for host-side operations. std::string threading_lib_path = ""; ///< Path to the threading library (empty disables multi-threading). }; @@ -61,6 +61,13 @@ struct cuDSSLinearSolverOptions { */ class cuDSSLinearSolver : public CSRSparseLinearSolver { public: + // This backend consumes CSR only; SupportsBlockStorage() stays false, so the + // base class's block-storage overloads are never called on it. The + // using-declarations keep them visible rather than hidden by the CSR + // overrides below. + using CSRSparseLinearSolver::Initialize; + using CSRSparseLinearSolver::Solve; + /** * @brief Constructs a cuDSS linear solver. * @@ -69,8 +76,7 @@ class cuDSSLinearSolver : public CSRSparseLinearSolver { * @param options Solver configuration controlling the initialization/solve * trade-off and threading settings. */ - cuDSSLinearSolver( - cuDSSLinearSolverOptions options = cuDSSLinearSolverOptions()); + cuDSSLinearSolver(cuDSSLinearSolverOptions options = cuDSSLinearSolverOptions()); /** * @brief Performs setup for the sparse linear system. @@ -92,9 +98,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 Problem &problem, - 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 Solves a sparse SPD linear system Ax = b. @@ -116,16 +121,15 @@ class cuDSSLinearSolver : public CSRSparseLinearSolver { * (size must equal matrix rows). * @return true on success, false if any dimension mismatch is detected. */ - bool Solve(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; private: cuDSSLinearSolverOptions options_; ///< Solver configuration. - cuDSSHandle - cudss_handle_; ///< Owns the cuDSS handle used for all solver phases. + cuDSSHandle cudss_handle_; ///< Owns the cuDSS handle used for all solver phases. cuDSSDeviceMemPool device_mem_pool_; ///< Reusable pool for cuDSS allocations. - cuDSSData cudss_data_; ///< cuDSS data object storing internal solver state. + cuDSSData cudss_data_; ///< cuDSS data object storing internal solver state. cuDSSConfig cudss_config_; ///< cuDSS configuration for solver parameters. }; diff --git a/cunls/linear_solver/dense_cholesky_solver.h b/cunls/linear_solver/dense_cholesky_solver.h index 584595f..3821a67 100644 --- a/cunls/linear_solver/dense_cholesky_solver.h +++ b/cunls/linear_solver/dense_cholesky_solver.h @@ -42,6 +42,13 @@ namespace cunls { */ class DenseCholeskySolver : public CSRSparseLinearSolver { public: + // This backend consumes CSR only; SupportsBlockStorage() stays false, so the + // base class's block-storage overloads are never called on it. The + // using-declarations keep them visible rather than hidden by the CSR + // overrides below. + using CSRSparseLinearSolver::Initialize; + using CSRSparseLinearSolver::Solve; + /** * @brief Validates dimensions and pre-allocates internal buffers. * @@ -51,9 +58,8 @@ 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 Problem &problem, - 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. @@ -74,8 +80,8 @@ class DenseCholeskySolver : public CSRSparseLinearSolver { * (devInfo > 0 from potrf), or invalid parameter from potrs * (devInfo < 0). */ - bool Solve(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; private: void EnsureBuffersSize(cudaStream_t stream, size_t n); diff --git a/cunls/linear_solver/dense_linear_solver.h b/cunls/linear_solver/dense_linear_solver.h index ceb0441..e268baa 100644 --- a/cunls/linear_solver/dense_linear_solver.h +++ b/cunls/linear_solver/dense_linear_solver.h @@ -42,6 +42,13 @@ namespace cunls { */ class DenseLDLTSolver : public CSRSparseLinearSolver { public: + // This backend consumes CSR only; SupportsBlockStorage() stays false, so the + // base class's block-storage overloads are never called on it. The + // using-declarations keep them visible rather than hidden by the CSR + // overrides below. + using CSRSparseLinearSolver::Initialize; + using CSRSparseLinearSolver::Solve; + /** * @brief Allocates internal dense buffers for the given matrix size. * @@ -56,9 +63,8 @@ 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 Problem &problem, - 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. @@ -84,8 +90,8 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { * @return true on success, false on dimension mismatch, singular pivot, * or zero diagonal. */ - bool Solve(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; private: /** @@ -112,11 +118,11 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { void ConvertCSRToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, dvector &dense_matrix); - dvector dense_matrix_; ///< Dense row-major copy of A. - dvector ldlt_factor_; ///< In-place LDLT factor storage. - dvector permutation_; ///< Pivot permutation vector. - dvector permuted_rhs_; ///< P * b scratch vector. - dvector permuted_solution_; ///< Permuted solution scratch. + dvector dense_matrix_; ///< Dense row-major copy of A. + dvector ldlt_factor_; ///< In-place LDLT factor storage. + dvector permutation_; ///< Pivot permutation vector. + dvector permuted_rhs_; ///< P * b scratch vector. + dvector permuted_solution_; ///< Permuted solution scratch. dvector intermediate_solution_; ///< Intermediate solve scratch. /// Device-side kernel status flags (index 0 = factorize, index 1 = solve). diff --git a/cunls/linear_solver/dense_qr_solver.h b/cunls/linear_solver/dense_qr_solver.h index 154842e..fcfa1ef 100644 --- a/cunls/linear_solver/dense_qr_solver.h +++ b/cunls/linear_solver/dense_qr_solver.h @@ -44,6 +44,13 @@ namespace cunls { */ class DenseQRSolver : public CSRSparseLinearSolver { public: + // This backend consumes CSR only; SupportsBlockStorage() stays false, so the + // base class's block-storage overloads are never called on it. The + // using-declarations keep them visible rather than hidden by the CSR + // overrides below. + using CSRSparseLinearSolver::Initialize; + using CSRSparseLinearSolver::Solve; + /** * @brief Validates dimensions and pre-allocates internal buffers. * @@ -53,9 +60,8 @@ 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 Problem &problem, - 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. @@ -75,8 +81,8 @@ class DenseQRSolver : public CSRSparseLinearSolver { * @param result Output vector x (size must equal matrix rows). * @return true on success, false on dimension mismatch or singular matrix. */ - bool Solve(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; private: void EnsureBuffersSize(cudaStream_t stream, size_t n); diff --git a/cunls/minimizer/CMakeLists.txt b/cunls/minimizer/CMakeLists.txt index 63fe0f4..2937924 100644 --- a/cunls/minimizer/CMakeLists.txt +++ b/cunls/minimizer/CMakeLists.txt @@ -1,13 +1,13 @@ add_library(cunls_minimizer OBJECT + block_hessian_assembler.cu + bsr_matrix.cu + hessian_structure.cu device_reduction.cu gauss_newton_minimizer.cu - jacobian_ops.cu levenberg_marquardt_minimizer.cpp minimizer_state.cu + normal_equations.cu problem.cpp residual_batch.cu sparse_matrix.cu - cusparse_matrix_multiplier.cpp - sparse_matrix_multiplier.cpp - fast_matrix_multiplier.cu ) diff --git a/cunls/minimizer/block_hessian_assembler.cu b/cunls/minimizer/block_hessian_assembler.cu new file mode 100644 index 0000000..f4fb825 --- /dev/null +++ b/cunls/minimizer/block_hessian_assembler.cu @@ -0,0 +1,371 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include + +#include "cunls/common/helper.h" +#include "cunls/minimizer/block_hessian_assembler.h" +#include "cunls/minimizer/hessian_structure.h" +#include "cunls/minimizer/problem.h" + +namespace cunls { +namespace { + +constexpr int kWarpSize = 32; +constexpr int kMaxWarpsPerBlock = 8; +constexpr size_t kMaxSharedBytes = 48 * 1024; + +// ============================================================================ +// Assembly kernel +// ============================================================================ + +/** + * One warp per factor. + * + * Stages `J_f` (m x n, row-major) and `r_f` in shared memory, then has the 32 + * lanes cooperatively evaluate the n^2 entries of `H_f = J_f^T J_f` and the n + * entries of `b_f = -J_f^T r_f`, scattering each with a single atomicAdd. + * + * Per-factor global atomics drop from `m * n^2` (one per residual row per + * column pair) to `n^2 + n`, and consecutive lanes hit consecutive slots within + * a block-pair run. + * + * The two storage layouts differ only in how a `(row, column)` pair becomes an + * offset into the value array, so they share everything else: + * - scalar CSR: `row_offsets[row] + write_offset + column_in_block` + * - block BSR: the enclosing tile is + * `row_offsets[row / b] + write_offset + column_in_block / b`, and the entry + * sits at `(row % b) * b + column_in_block % b` inside it. + * `row / b`, `row % b` and their column counterparts are precomputed per local + * column, so the inner loop stays free of integer division. + * + * @tparam kBlockStorage Selects BSR addressing. + */ +template +__global__ void AssembleBlockHessianKernel( + int num_factors, int residual_dim, int tangent_dim, int num_blocks, int block_size, + const float *__restrict__ jacobians, const float *__restrict__ residuals, + const int *__restrict__ block_of_col, const int *__restrict__ offset_in_block, + const int *__restrict__ tile_of_col, const int *__restrict__ sub_of_col, + const int *__restrict__ triangular_pairs, const int *__restrict__ factor_cols, + const int *__restrict__ write_offsets, const int *__restrict__ row_offsets, + float *__restrict__ hessian_values, float *__restrict__ rhs) { + extern __shared__ int s_arena[]; + + const int lane = threadIdx.x & (kWarpSize - 1); + const int warp_in_block = threadIdx.x / kWarpSize; + const int warps_per_block = blockDim.x / kWarpSize; + const int num_pairs = tangent_dim * (tangent_dim + 1) / 2; + + // Block-wide: these tables depend only on the batch's block layout, not on + // the factor, so they are loaded once per thread block. + int *s_block_of_col = s_arena; + int *s_offset_in_block = s_block_of_col + tangent_dim; + int *s_tile_of_col = s_offset_in_block + tangent_dim; + int *s_sub_of_col = s_tile_of_col + tangent_dim; + int *s_triangular = s_sub_of_col + tangent_dim; + for (int i = threadIdx.x; i < tangent_dim; i += blockDim.x) { + s_block_of_col[i] = block_of_col[i]; + s_offset_in_block[i] = offset_in_block[i]; + if (kBlockStorage) { + s_tile_of_col[i] = tile_of_col[i]; + s_sub_of_col[i] = sub_of_col[i]; + } + } + for (int i = threadIdx.x; i < num_pairs; i += blockDim.x) { + s_triangular[i] = triangular_pairs[i]; + } + __syncthreads(); + + const int factor = blockIdx.x * warps_per_block + warp_in_block; + if (factor >= num_factors) { + return; + } + + // Per-warp arena: J (residual_dim*tangent_dim floats), r (residual_dim floats), row starts + // (tangent_dim ints), global columns (tangent_dim ints), block-pair write offsets + // (num_blocks*num_blocks ints). + const int per_warp_ints = + residual_dim * tangent_dim + residual_dim + 2 * tangent_dim + num_blocks * num_blocks; + int *s_warp = s_arena + 4 * tangent_dim + num_pairs + warp_in_block * per_warp_ints; + float *s_jacobian = reinterpret_cast(s_warp); + float *s_residual = s_jacobian + residual_dim * tangent_dim; + int *s_row_start = reinterpret_cast(s_residual + residual_dim); + int *s_global_col = s_row_start + tangent_dim; + int *s_write_offset = s_global_col + tangent_dim; + + const float *j_src = jacobians + static_cast(factor) * residual_dim * tangent_dim; + for (int i = lane; i < residual_dim * tangent_dim; i += kWarpSize) { + s_jacobian[i] = j_src[i]; + } + const float *r_src = residuals + static_cast(factor) * residual_dim; + for (int i = lane; i < residual_dim; i += kWarpSize) { + s_residual[i] = r_src[i]; + } + const int *wo_src = write_offsets + static_cast(factor) * num_blocks * num_blocks; + for (int i = lane; i < num_blocks * num_blocks; i += kWarpSize) { + s_write_offset[i] = wo_src[i]; + } + const int *cols_src = factor_cols + static_cast(factor) * num_blocks; + for (int p = lane; p < tangent_dim; p += kWarpSize) { + int col = cols_src[s_block_of_col[p]]; + if (col < 0) { + // Constant state block: this row and column of H_f are dropped, which + // is the block-level equivalent of the col_id == -1 triplet filter. + s_global_col[p] = -1; + s_row_start[p] = -1; + } else { + int global_col = col + s_offset_in_block[p]; + s_global_col[p] = global_col; + // `col` is a multiple of block_size, so the tile row is + // col / b + offset / b and the in-tile row is offset % b. + s_row_start[p] = kBlockStorage ? row_offsets[col / block_size + s_tile_of_col[p]] + : row_offsets[global_col]; + } + } + __syncwarp(); + + for (int p = lane; p < tangent_dim; p += kWarpSize) { + if (s_global_col[p] < 0) { + continue; + } + float acc = 0.f; + for (int k = 0; k < residual_dim; k++) { + acc = fmaf(s_jacobian[k * tangent_dim + p], s_residual[k], acc); + } + atomicAdd(&rhs[s_global_col[p]], -acc); + } + + for (int idx = lane; idx < num_pairs; idx += kWarpSize) { + const int packed = s_triangular[idx]; + const int p = packed >> 16; + const int q = packed & 0xFFFF; + if (s_row_start[p] < 0 || s_row_start[q] < 0) { + continue; + } + float acc = 0.f; + for (int k = 0; k < residual_dim; k++) { + acc = fmaf(s_jacobian[k * tangent_dim + p], s_jacobian[k * tangent_dim + q], acc); + } + const int block_p = s_block_of_col[p]; + const int block_q = s_block_of_col[q]; + if (kBlockStorage) { + const int tile_area = block_size * block_size; + atomicAdd(&hessian_values[static_cast(s_row_start[p] + + s_write_offset[block_p * num_blocks + block_q] + + s_tile_of_col[q]) * + tile_area + + s_sub_of_col[p] * block_size + s_sub_of_col[q]], + acc); + if (p != q) { + atomicAdd( + &hessian_values[static_cast(s_row_start[q] + + s_write_offset[block_q * num_blocks + block_p] + + s_tile_of_col[p]) * + tile_area + + s_sub_of_col[q] * block_size + s_sub_of_col[p]], + acc); + } + } else { + atomicAdd(&hessian_values[s_row_start[p] + s_write_offset[block_p * num_blocks + block_q] + + s_offset_in_block[q]], + acc); + if (p != q) { + atomicAdd(&hessian_values[s_row_start[q] + s_write_offset[block_q * num_blocks + block_p] + + s_offset_in_block[p]], + acc); + } + } + } +} + +/** + * Picks the largest warp count that fits the per-warp staging arena in shared + * memory, capped at kMaxWarpsPerBlock. + */ +int PickWarpsPerBlock(int residual_dim, int tangent_dim, int num_blocks, size_t &shared_bytes) { + const size_t block_ints = static_cast(4) * tangent_dim + + static_cast(tangent_dim) * (tangent_dim + 1) / 2; + const size_t warp_ints = static_cast(residual_dim) * tangent_dim + residual_dim + + 2 * static_cast(tangent_dim) + num_blocks * num_blocks; + const size_t warp_bytes = warp_ints * sizeof(int); + const size_t block_bytes = block_ints * sizeof(int); + + int warps = kMaxWarpsPerBlock; + if (warp_bytes > 0) { + size_t budget = + kMaxSharedBytes > block_bytes ? (kMaxSharedBytes - block_bytes) / warp_bytes : 0; + warps = static_cast(std::min(warps, budget)); + } + if (warps < 1) { + throw std::runtime_error("BlockHessianAssembler: factor too large for shared-memory assembly"); + } + shared_bytes = block_bytes + static_cast(warps) * warp_bytes; + return warps; +} + +} // namespace + +void BlockHessianAssembler::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + CSRSparseMatrix &hessian) { + auto range = profiler_domain_.CreateDomainRange("Initialize"); + num_cols_ = num_cols; + block_size_ = 1; + structure_builder_.Build(stream, problem, num_cols, hessian, /*want_scatter_maps=*/true); + BuildPlans(problem); +} + +void BlockHessianAssembler::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + int block_size, BSRSparseMatrix &hessian) { + auto range = profiler_domain_.CreateDomainRange("Initialize"); + num_cols_ = num_cols; + block_size_ = block_size; + structure_builder_.Build(stream, problem, num_cols, block_size, hessian, + /*want_scatter_maps=*/true); + BuildPlans(problem); +} + +void BlockHessianAssembler::BuildPlans(const Problem &problem) { + // The structure builder resolved each factor's state pointers to global + // columns and segmented the block pairs; the row-relative write offsets the + // assembler needs fell out of that segmentation, so nothing is recomputed + // here. + const auto &layout = structure_builder_.Layout(); + plans_.clear(); + plans_.resize(layout.size()); + + for (size_t i = 0; i < layout.size(); i++) { + BatchPlan &plan = plans_[i]; + plan.layout = layout[i]; + + auto block_sizes = problem.GetResidualBatches()[i].GetFactorBatch()->StateBlockSizes(); + + // Local column -> (block, offset within block). + std::vector block_of_col(plan.layout.tangent_dim); + std::vector offset_in_block(plan.layout.tangent_dim); + int cursor = 0; + for (int b = 0; b < plan.layout.num_blocks; b++) { + for (size_t k = 0; k < block_sizes[b]; k++) { + block_of_col[cursor] = b; + offset_in_block[cursor] = static_cast(k); + cursor++; + } + } + plan.block_of_col.resize(block_of_col.size()); + plan.block_of_col.CopyFromHost(block_of_col.data(), block_of_col.size()); + plan.offset_in_block.resize(offset_in_block.size()); + plan.offset_in_block.CopyFromHost(offset_in_block.data(), offset_in_block.size()); + + // Upper-triangle enumeration of H_f, packed as (p << 16) | q. + if (plan.layout.tangent_dim > 0xFFFF) { + throw std::runtime_error("BlockHessianAssembler: factor tangent dimension exceeds 65535"); + } + std::vector triangular; + triangular.reserve(static_cast(plan.layout.tangent_dim) * + (plan.layout.tangent_dim + 1) / 2); + for (int p = 0; p < plan.layout.tangent_dim; p++) { + for (int q = p; q < plan.layout.tangent_dim; q++) { + triangular.push_back((p << 16) | q); + } + } + plan.triangular_pairs.resize(triangular.size()); + plan.triangular_pairs.CopyFromHost(triangular.data(), triangular.size()); + + // Tile / in-tile decomposition of each local column, so the assembly + // kernel never divides by the block size. Meaningless (and unread) when + // the Hessian is stored as scalar CSR. + std::vector tile_of_col(plan.layout.tangent_dim, 0); + std::vector sub_of_col(plan.layout.tangent_dim, 0); + if (block_size_ > 1) { + for (int p = 0; p < plan.layout.tangent_dim; p++) { + tile_of_col[p] = offset_in_block[p] / block_size_; + sub_of_col[p] = offset_in_block[p] % block_size_; + } + } + plan.tile_of_col.resize(tile_of_col.size()); + plan.tile_of_col.CopyFromHost(tile_of_col.data(), tile_of_col.size()); + plan.sub_of_col.resize(sub_of_col.size()); + plan.sub_of_col.CopyFromHost(sub_of_col.data(), sub_of_col.size()); + } +} + +void BlockHessianAssembler::Assemble(cudaStream_t stream, const Problem &problem, + const float *jacobian_values, const float *residuals, + CSRSparseMatrix &hessian, dvector &rhs) { + auto range = profiler_domain_.CreateDomainRange("Assemble"); + PrepareOutputs(stream, hessian.values, rhs); + LaunchAssembly(stream, jacobian_values, residuals, hessian.row_offsets.data(), + hessian.values.data(), rhs.data(), /*block_storage=*/false); +} + +void BlockHessianAssembler::Assemble(cudaStream_t stream, const Problem &problem, + const float *jacobian_values, const float *residuals, + BSRSparseMatrix &hessian, dvector &rhs) { + auto range = profiler_domain_.CreateDomainRange("Assemble"); + PrepareOutputs(stream, hessian.values, rhs); + LaunchAssembly(stream, jacobian_values, residuals, hessian.row_offsets.data(), + hessian.values.data(), rhs.data(), /*block_storage=*/true); +} + +void BlockHessianAssembler::PrepareOutputs(cudaStream_t stream, dvector &values, + dvector &rhs) { + rhs.resize(static_cast(num_cols_)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(rhs.data(), 0, rhs.size() * sizeof(float), stream)); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(values.data(), 0, values.size() * sizeof(float), stream)); +} + +void BlockHessianAssembler::LaunchAssembly(cudaStream_t stream, const float *jacobian_values, + const float *residuals, const int *row_offsets, + float *hessian_values, float *rhs, bool block_storage) { + const int *factor_cols = structure_builder_.FactorCols().data(); + const int *write_offsets = structure_builder_.WriteOffsets().data(); + + for (const BatchPlan &plan : plans_) { + const HessianBatchLayout &layout = plan.layout; + if (layout.num_factors == 0 || layout.tangent_dim == 0) { + continue; + } + + size_t shared_bytes = 0; + int warps = + PickWarpsPerBlock(layout.residual_dim, layout.tangent_dim, layout.num_blocks, shared_bytes); + int threads = warps * kWarpSize; + int grid = (layout.num_factors + warps - 1) / warps; + + if (block_storage) { + AssembleBlockHessianKernel<<>>( + layout.num_factors, layout.residual_dim, layout.tangent_dim, layout.num_blocks, + block_size_, jacobian_values + layout.jacobian_offset, residuals + layout.residual_offset, + plan.block_of_col.data(), plan.offset_in_block.data(), plan.tile_of_col.data(), + plan.sub_of_col.data(), plan.triangular_pairs.data(), factor_cols + layout.col_offset, + write_offsets + layout.pair_offset, row_offsets, hessian_values, rhs); + } else { + AssembleBlockHessianKernel<<>>( + layout.num_factors, layout.residual_dim, layout.tangent_dim, layout.num_blocks, 1, + jacobian_values + layout.jacobian_offset, residuals + layout.residual_offset, + plan.block_of_col.data(), plan.offset_in_block.data(), plan.tile_of_col.data(), + plan.sub_of_col.data(), plan.triangular_pairs.data(), factor_cols + layout.col_offset, + write_offsets + layout.pair_offset, row_offsets, hessian_values, rhs); + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } +} + +} // namespace cunls diff --git a/cunls/minimizer/block_hessian_assembler.h b/cunls/minimizer/block_hessian_assembler.h new file mode 100644 index 0000000..aa00055 --- /dev/null +++ b/cunls/minimizer/block_hessian_assembler.h @@ -0,0 +1,147 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include + +#include "cunls/common/profiler.h" +#include "cunls/common/types.h" +#include "cunls/minimizer/hessian_structure.h" + +namespace cunls { + +class Problem; + +/** + * @brief Assembles the normal equations directly from per-factor Jacobian + * blocks. + * + * The classic path materializes the whole sparse Jacobian `J` in CSR, forms + * `H = J^T J` with one warp per *residual row* (`m * n^2` scattered global + * atomics per factor), and then runs a separate cuSPARSE SpMV for + * `b = -J^T r`. + * + * This class instead contracts each factor locally. A factor batch already + * writes its Jacobian as `NumFactors()` dense row-major `m x n` blocks, so a + * single kernel can read `J_f`, form `H_f = J_f^T J_f` and `b_f = -J_f^T r_f` + * in shared memory, and scatter both into the global CSR. That drops the + * atomic count by exactly a factor of `m` and removes the triplet-to-CSR + * conversion, the `J^T J` kernel and the RHS SpMV from every iteration. + * + * Storage stays plain CSR: the pattern produced by ComputeHessianStructure + * lays out each block pair contiguously within a row and at the same + * row-relative offset for every row of the block, so the scatter address is + * `row_offsets[col_a + i] + write_offset(a,b) + j`. No intermediate + * block-format Hessian is needed, and every downstream consumer (LM damping, + * column scaling, PCG, cuDSS) sees the matrix it already expects. + */ +class BlockHessianAssembler { +public: + /** + * @brief Builds the Hessian sparsity pattern and the per-factor scatter maps. + * + * Must be called whenever the problem structure changes. The maps cost + * `num_factors * (nb + nb^2)` integers per factor batch, where `nb` is the + * number of state blocks a factor touches. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem. + * @param num_cols Number of free tangent dimensions in the reduced system. + * @param[out] hessian CSR Hessian; structure filled, values left unset. + */ + void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + CSRSparseMatrix &hessian); + + /** + * @brief Same, but targeting uniform block storage. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem. + * @param num_cols Number of free tangent dimensions in the reduced system. + * @param block_size Tile edge; see ChooseHessianBlockSize(). + * @param[out] hessian BSR Hessian; structure filled, values left unset. + */ + void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, int block_size, + BSRSparseMatrix &hessian); + + /** + * @brief Scatter-accumulates `H = J^T J` and `rhs = -J^T r`. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem (must match Initialize). + * @param jacobian_values Per-factor dense Jacobian blocks, batches + * concatenated in problem order; see JacobianValuesSize(). + * @param residuals Residual vector, batches concatenated in problem order. + * @param[in,out] hessian CSR Hessian initialized by Initialize(); values are + * zeroed and then accumulated. + * @param[out] rhs Right-hand side `-J^T r`; resized to `num_cols`. + */ + void Assemble(cudaStream_t stream, const Problem &problem, const float *jacobian_values, + const float *residuals, CSRSparseMatrix &hessian, dvector &rhs); + + /** @brief Same, into the BSR Hessian prepared by the block-size overload. */ + void Assemble(cudaStream_t stream, const Problem &problem, const float *jacobian_values, + const float *residuals, BSRSparseMatrix &hessian, dvector &rhs); + + /** @brief Total floats needed for the per-factor Jacobian value buffer. */ + size_t JacobianValuesSize() const { return structure_builder_.JacobianValuesSize(); } + +private: + /** @brief Per-residual-batch constants uploaded once for the kernel. */ + struct BatchPlan { + HessianBatchLayout layout; ///< Geometry and flat-buffer offsets. + /// n entries: block index owning each local column. + dvector block_of_col; + /// n entries: offset of each local column inside its block. + dvector offset_in_block; + /// n entries: index of the enclosing tile column, for block storage. + dvector tile_of_col; + /// n entries: offset within that tile, for block storage. + dvector sub_of_col; + /// n(n+1)/2 entries: the upper-triangle (p, q) pairs of H_f packed as + /// `(p << 16) | q`. H_f is symmetric, so enumerating the triangle halves + /// the shared-memory traffic and the FLOPs; a precomputed table keeps + /// every lane busy, which a `q < p` skip inside the square loop would not. + dvector triangular_pairs; + }; + + /** @brief Fills plans_ from the structure builder's layout. */ + void BuildPlans(const Problem &problem); + + /** @brief Zeroes the value array and the right-hand side. */ + void PrepareOutputs(cudaStream_t stream, dvector &values, dvector &rhs); + + /** @brief Launches one assembly kernel per residual batch. */ + void LaunchAssembly(cudaStream_t stream, const float *jacobian_values, const float *residuals, + const int *row_offsets, float *hessian_values, float *rhs, + bool block_storage); + + /// Owns the Hessian sparsity pattern and the per-factor scatter maps. + HessianStructureBuilder structure_builder_; + std::vector plans_; + + int num_cols_ = 0; + /// Tile edge of the target storage; 1 means scalar CSR. + int block_size_ = 1; + + profiler::Domain profiler_domain_{"BlockHessianAssembler"}; +}; + +} // namespace cunls diff --git a/cunls/minimizer/bsr_matrix.cu b/cunls/minimizer/bsr_matrix.cu new file mode 100644 index 0000000..e491001 --- /dev/null +++ b/cunls/minimizer/bsr_matrix.cu @@ -0,0 +1,532 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "cunls/common/helper.h" +#include "cunls/minimizer/bsr_matrix.h" +#include "cunls/minimizer/device_reduction.h" +#include "cunls/minimizer/problem.h" +#include "cunls/state/state_batch.h" + +namespace cunls { +namespace { + +constexpr int kBlockSize = 256; + +int GridFor(size_t count) { return static_cast((count + kBlockSize - 1) / kBlockSize); } + +/** + * @brief Fills the CSR row offsets of the expanded matrix. + * + * Every tile in a block row contributes `block_size` columns to each of the + * block row's `block_size` scalar rows, so a row's length is known from the + * block row's tile count alone. + */ +__global__ void FillExpandedRowOffsetsKernel(int num_rows, int block_size, + const int *__restrict__ block_row_offsets, + int *__restrict__ row_offsets) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row > num_rows) { + return; + } + // Rows before `row` belong to whole block rows plus a partial one. + const int block_row = row / block_size; + const int sub_row = row - block_row * block_size; + const int whole = block_row_offsets[block_row] * block_size; + const int partial = sub_row * (block_row_offsets[block_row + 1] - block_row_offsets[block_row]); + row_offsets[row] = (whole + partial) * block_size; +} + +/** @brief Scatters each tile entry to its scalar CSR position. */ +__global__ void ExpandTilesToCSRKernel(size_t num_values, int block_size, + const int *__restrict__ row_of_tile, + const int *__restrict__ block_row_offsets, + const int *__restrict__ block_col_ids, + const float *__restrict__ values, + const int *__restrict__ row_offsets, + int *__restrict__ col_ids, float *__restrict__ out_values) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_values) { + return; + } + const int tile_area = block_size * block_size; + const size_t tile = idx / tile_area; + const int within = static_cast(idx - tile * tile_area); + const int row_in_tile = within / block_size; + const int col_in_tile = within - row_in_tile * block_size; + + const int block_row = row_of_tile[tile]; + const int tile_in_row = static_cast(tile) - block_row_offsets[block_row]; + const int row = block_row * block_size + row_in_tile; + const int slot = row_offsets[row] + tile_in_row * block_size + col_in_tile; + col_ids[slot] = block_col_ids[tile] * block_size + col_in_tile; + out_values[slot] = values[idx]; +} + +/** One thread per block row; locates the diagonal tile and reads its diagonal. */ +__global__ void ExtractBlockDiagonalKernel(int num_block_rows, int block_size, + const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, + float *__restrict__ diagonal) { + int block_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_row >= num_block_rows) { + return; + } + const int tile_area = block_size * block_size; + for (int t = row_offsets[block_row]; t < row_offsets[block_row + 1]; t++) { + if (col_ids[t] != block_row) { + continue; + } + const float *tile = values + static_cast(t) * tile_area; + for (int k = 0; k < block_size; k++) { + diagonal[block_row * block_size + k] = tile[k * block_size + k]; + } + return; + } + // Structurally absent diagonal tile: the diagonal is zero there. + for (int k = 0; k < block_size; k++) { + diagonal[block_row * block_size + k] = 0.f; + } +} + +/** One thread per block row; adds `scale * diagonal` to the diagonal tile. */ +__global__ void AddScaledBlockDiagonalKernel(int num_block_rows, int block_size, + const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, float scale, + const float *__restrict__ diagonal, + float *__restrict__ values) { + int block_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_row >= num_block_rows) { + return; + } + const int tile_area = block_size * block_size; + for (int t = row_offsets[block_row]; t < row_offsets[block_row + 1]; t++) { + if (col_ids[t] != block_row) { + continue; + } + float *tile = values + static_cast(t) * tile_area; + for (int k = 0; k < block_size; k++) { + tile[k * block_size + k] += scale * diagonal[block_row * block_size + k]; + } + return; + } +} + +/** + * One thread per stored scalar entry. `row_of_tile` maps a tile index to its + * block row so the entry's global (i, j) can be recovered. + */ +__global__ void ScaleSymmetricBSRKernel(size_t num_values, int block_size, + const int *__restrict__ row_of_tile, + const int *__restrict__ col_ids, + const float *__restrict__ scale, + float *__restrict__ values) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_values) { + return; + } + const int tile_area = block_size * block_size; + const size_t tile = idx / tile_area; + const int within = static_cast(idx - tile * tile_area); + const int row_in_tile = within / block_size; + const int col_in_tile = within - row_in_tile * block_size; + const int row = row_of_tile[tile] * block_size + row_in_tile; + const int col = col_ids[tile] * block_size + col_in_tile; + values[idx] *= scale[row] * scale[col]; +} + +/** One thread per tile; records the tile's block row. */ +__global__ void FillRowOfTileKernel(int num_block_rows, const int *__restrict__ row_offsets, + int *__restrict__ row_of_tile) { + int block_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_row >= num_block_rows) { + return; + } + for (int t = row_offsets[block_row]; t < row_offsets[block_row + 1]; t++) { + row_of_tile[t] = block_row; + } +} + +/** + * @brief y = A * x for uniform block storage, one warp per block row. + * + * cuSPARSE's `cusparseSbsrmv` measured 3.6x slower than `csrmv_v3` on a + * bundle-adjustment Hessian with 3x3 tiles, which would negate the point of + * block storage, so the SpMV is written here. + * + * A warp per block row rather than a thread per row: bundle-adjustment Hessians + * are wildly non-uniform -- a pose block row holds one tile per observation of + * that camera (thousands) while a landmark row holds a handful -- so a + * thread-per-row schedule leaves the few pose threads serializing for as long + * as the whole kernel takes. Lanes stride over the row's tiles instead, and a + * butterfly reduction combines their partial `b`-vectors. + * + * @tparam kB Tile edge. + */ +template +__global__ void BsrMultiplyWarpKernel(int num_block_rows, const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, const float *__restrict__ x, + float *__restrict__ y) { + const int block_row = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + const int lane = threadIdx.x & 31; + if (block_row >= num_block_rows) { + return; + } + + const int end = row_offsets[block_row + 1]; + float acc[kB]; +#pragma unroll + for (int k = 0; k < kB; ++k) { + acc[k] = 0.f; + } + + for (int t = row_offsets[block_row] + lane; t < end; t += 32) { + const float *tile = values + static_cast(t) * kB * kB; + const float *xs = x + static_cast(col_ids[t]) * kB; + float xv[kB]; +#pragma unroll + for (int l = 0; l < kB; ++l) { + xv[l] = xs[l]; + } +#pragma unroll + for (int k = 0; k < kB; ++k) { +#pragma unroll + for (int l = 0; l < kB; ++l) { + acc[k] = fmaf(tile[k * kB + l], xv[l], acc[k]); + } + } + } + + // Butterfly rather than shfl_down so every lane ends with the totals and the + // first kB lanes can write the output run coalesced. +#pragma unroll + for (int k = 0; k < kB; ++k) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + acc[k] += __shfl_xor_sync(0xFFFFFFFFu, acc[k], offset); + } + } + if (lane < kB) { + y[block_row * kB + lane] = acc[lane]; + } +} + +/** + * @brief Runtime-tile-edge fallback for block sizes without a specialization. + * + * Same schedule as BsrMultiplyWarpKernel; the accumulator is sized to the + * largest edge ChooseHessianBlockSize can return. + */ +__global__ void BsrMultiplyWarpGenericKernel(int num_block_rows, int block_size, + const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, + const float *__restrict__ x, float *__restrict__ y) { + constexpr int kMaxBlockSize = 16; + const int block_row = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + const int lane = threadIdx.x & 31; + if (block_row >= num_block_rows) { + return; + } + + const int end = row_offsets[block_row + 1]; + float acc[kMaxBlockSize]; + for (int k = 0; k < block_size; ++k) { + acc[k] = 0.f; + } + + for (int t = row_offsets[block_row] + lane; t < end; t += 32) { + const float *tile = values + static_cast(t) * block_size * block_size; + const float *xs = x + static_cast(col_ids[t]) * block_size; + for (int row_in_tile = 0; row_in_tile < block_size; ++row_in_tile) { + float sum = 0.f; + for (int col = 0; col < block_size; ++col) { + sum = fmaf(tile[row_in_tile * block_size + col], xs[col], sum); + } + acc[row_in_tile] += sum; + } + } + + for (int k = 0; k < block_size; ++k) { + for (int offset = 16; offset > 0; offset >>= 1) { + acc[k] += __shfl_xor_sync(0xFFFFFFFFu, acc[k], offset); + } + } + if (lane < block_size) { + y[block_row * block_size + lane] = acc[lane]; + } +} + +/** + * @brief y = A * x, one thread per scalar row. + * + * The right schedule for near-uniform, short block rows -- a pose graph has a + * handful of tiles per row, so a whole warp per row would leave most lanes idle + * and pay for a reduction that spans mostly zeros. Thread `i` walks row + * `i % b` of every tile in block row `i / b`; the `b` consecutive threads + * sharing a block row read each tile as one contiguous run and broadcast their + * identical `col_ids` and `x` loads. + * + * @tparam kB Tile edge, or 0 to take it as a runtime argument. + */ +template +__global__ void BsrMultiplyRowKernel(int num_rows, int runtime_block_size, + const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, const float *__restrict__ x, + float *__restrict__ y) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= num_rows) { + return; + } + const int edge = kB > 0 ? kB : runtime_block_size; + const int block_row = row / edge; + const int sub_row = row - block_row * edge; + + float acc = 0.f; + const int end = row_offsets[block_row + 1]; + for (int t = row_offsets[block_row]; t < end; ++t) { + const float *tile = values + static_cast(t) * edge * edge + sub_row * edge; + const float *xs = x + static_cast(col_ids[t]) * edge; + for (int l = 0; l < edge; ++l) { + acc = fmaf(tile[l], xs[l], acc); + } + } + y[row] = acc; +} + +/** + * @brief Launches the BSR SpMV, choosing a schedule from the row lengths. + * + * A warp per block row tolerates skew but wastes lanes on short rows; a thread + * per scalar row is the opposite. Bundle adjustment needs the former (pose + * rows hold thousands of tiles), pose graphs the latter (every row holds a + * handful), so the peak row length decides. + */ +void LaunchBsrMultiply(cudaStream_t stream, int num_block_rows, int block_size, + int max_tiles_per_row, const int *row_offsets, const int *col_ids, + const float *values, const float *x, float *y) { + constexpr int kThreads = 128; + constexpr int kSkewThreshold = 32; + + if (max_tiles_per_row < kSkewThreshold) { + const int num_rows = num_block_rows * block_size; + const int grid = (num_rows + kThreads - 1) / kThreads; + switch (block_size) { +#define LAUNCH_ROW(BVAL) \ + case BVAL: \ + BsrMultiplyRowKernel \ + <<>>(num_rows, BVAL, row_offsets, col_ids, values, x, y); \ + break + LAUNCH_ROW(2); + LAUNCH_ROW(3); + LAUNCH_ROW(4); + LAUNCH_ROW(6); + LAUNCH_ROW(7); + LAUNCH_ROW(15); + LAUNCH_ROW(16); +#undef LAUNCH_ROW + default: + BsrMultiplyRowKernel<0><<>>(num_rows, block_size, row_offsets, + col_ids, values, x, y); + break; + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); + return; + } + + const int warps_per_block = kThreads / 32; + const int grid = (num_block_rows + warps_per_block - 1) / warps_per_block; + switch (block_size) { +#define LAUNCH_WARP(BVAL) \ + case BVAL: \ + BsrMultiplyWarpKernel \ + <<>>(num_block_rows, row_offsets, col_ids, values, x, y); \ + break + LAUNCH_WARP(2); + LAUNCH_WARP(3); + LAUNCH_WARP(4); + LAUNCH_WARP(5); + LAUNCH_WARP(6); + LAUNCH_WARP(7); + LAUNCH_WARP(8); +#undef LAUNCH_WARP + default: + BsrMultiplyWarpGenericKernel<<>>( + num_block_rows, block_size, row_offsets, col_ids, values, x, y); + break; + } + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +} // namespace + +void MultiplyBSRByDenseVector(cudaStream_t stream, const BSRSparseMatrix &matrix, + const dvector &x, dvector &y) { + y.resize(static_cast(matrix.NumRows())); + if (matrix.NumRows() == 0) { + return; + } + LaunchBsrMultiply(stream, matrix.num_block_rows, matrix.block_size, matrix.max_tiles_per_row, + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), + x.data(), y.data()); +} + +int ChooseHessianBlockSize(const Problem &problem, int max_block_size) { + int block_size = 0; + for (const auto *state_batch : problem.GetStateBatches()) { + const size_t num_free = state_batch->NumStateBlocks() - state_batch->NumConstStateBlocks(); + if (num_free == 0) { + continue; // Contributes no columns, so its tangent size is irrelevant. + } + block_size = std::gcd(block_size, static_cast(state_batch->TangentSize())); + if (block_size == 1) { + return 1; + } + } + if (block_size <= 1) { + return 1; + } + while (block_size > max_block_size) { + // Fall back to the largest divisor within budget; halving keeps the + // divisibility invariant for the even sizes cuNLS actually uses and + // otherwise bails out to scalar storage. + if (block_size % 2 != 0) { + return 1; + } + block_size /= 2; + } + return block_size; +} + +void ExtractDiagonal(cudaStream_t stream, const BSRSparseMatrix &matrix, dvector &diagonal) { + diagonal.resize(static_cast(matrix.NumRows())); + if (matrix.num_block_rows == 0) { + return; + } + ExtractBlockDiagonalKernel<<>>( + matrix.num_block_rows, matrix.block_size, matrix.row_offsets.data(), matrix.col_ids.data(), + matrix.values.data(), diagonal.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +void AddScaledDiagonal(cudaStream_t stream, float scale, const dvector &diagonal, + const BSRSparseMatrix &matrix, BSRSparseMatrix &result) { + CopyBSRSparseMatrix(stream, matrix, result); + if (result.num_block_rows == 0) { + return; + } + AddScaledBlockDiagonalKernel<<>>( + result.num_block_rows, result.block_size, result.row_offsets.data(), result.col_ids.data(), + scale, diagonal.data(), result.values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +void ScaleSymmetric(cudaStream_t stream, BSRSparseMatrix &matrix, const dvector &scale, + dvector &row_of_tile) { + if (matrix.values.empty()) { + return; + } + // Scaling needs each entry's global row, which BSR does not store per tile. + // The map is caller-owned so this stays allocation- and sync-free on the + // solver's hot path. + row_of_tile.resize(matrix.NumBlocks()); + FillRowOfTileKernel<<>>( + matrix.num_block_rows, matrix.row_offsets.data(), row_of_tile.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + ScaleSymmetricBSRKernel<<>>( + matrix.values.size(), matrix.block_size, row_of_tile.data(), matrix.col_ids.data(), + scale.data(), matrix.values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +void CopyBSRSparseMatrix(cudaStream_t stream, const BSRSparseMatrix &input, + BSRSparseMatrix &output) { + if (&input == &output) { + return; + } + // Always copy the structure, never just the values. Matching sizes do not + // imply matching connectivity: a solver reused across a stream of problems + // hits pairs with the same block-row and tile counts but different sparsity, + // and reusing the previous pattern silently pairs new values with old + // columns. + output.block_size = input.block_size; + output.num_block_rows = input.num_block_rows; + output.max_tiles_per_row = input.max_tiles_per_row; + + output.row_offsets.resize(input.row_offsets.size()); + output.col_ids.resize(input.col_ids.size()); + if (!input.row_offsets.empty()) { + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(output.row_offsets.data(), input.row_offsets.data(), + input.row_offsets.size() * sizeof(int), + cudaMemcpyDeviceToDevice, stream)); + } + if (!input.col_ids.empty()) { + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(output.col_ids.data(), input.col_ids.data(), + input.col_ids.size() * sizeof(int), + cudaMemcpyDeviceToDevice, stream)); + } + + output.values.resize(input.values.size()); + if (!input.values.empty()) { + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(output.values.data(), input.values.data(), + input.values.size() * sizeof(float), + cudaMemcpyDeviceToDevice, stream)); + } +} + +void ConvertBSRToCSR(cudaStream_t stream, const BSRSparseMatrix &input, CSRSparseMatrix &output, + dvector &row_of_tile) { + const int num_rows = input.NumRows(); + output.row_offsets.resize(static_cast(num_rows) + 1); + output.col_ids.resize(input.NumNonZeros()); + output.values.resize(input.NumNonZeros()); + if (num_rows == 0) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(output.row_offsets.data(), 0, sizeof(int), stream)); + return; + } + + row_of_tile.resize(input.NumBlocks()); + FillRowOfTileKernel<<>>( + input.num_block_rows, input.row_offsets.data(), row_of_tile.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + FillExpandedRowOffsetsKernel<<(num_rows) + 1), kBlockSize, 0, + stream>>>(num_rows, input.block_size, input.row_offsets.data(), + output.row_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + ExpandTilesToCSRKernel<<>>( + input.values.size(), input.block_size, row_of_tile.data(), input.row_offsets.data(), + input.col_ids.data(), input.values.data(), output.row_offsets.data(), output.col_ids.data(), + output.values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const BSRSparseMatrix &matrix, + const dvector &step, dvector &scratch, + float *d_out, float *d_partials) { + MultiplyBSRByDenseVector(stream, matrix, step, scratch); + DotProductToDevice(stream, step.data(), scratch.data(), step.size(), d_out, d_partials); +} + +} // namespace cunls diff --git a/cunls/minimizer/bsr_matrix.h b/cunls/minimizer/bsr_matrix.h new file mode 100644 index 0000000..a5724e7 --- /dev/null +++ b/cunls/minimizer/bsr_matrix.h @@ -0,0 +1,141 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include + +#include "cunls/common/types.h" + +namespace cunls { + +class Problem; + +/** + * @brief Largest uniform tile size that divides every state block's tangent + * dimension. + * + * Uniform BSR needs one tile edge for the whole matrix, but cuNLS problems mix + * tangent dimensions (SE3 poses at 6, landmarks at 3, arbitrary vectors). The + * gcd is the largest edge that tiles every block pair exactly, so no padding is + * ever stored: PGO gives 6, bundle adjustment gives 3. A gcd of 1 means block + * storage would degenerate to scalar CSR, and the caller should stay on CSR. + * + * Block column offsets are cumulative sums of tangent dimensions, so they are + * automatically multiples of the gcd. + * + * @param problem The optimization problem. + * @param max_block_size Upper bound on the returned edge length. + * @return The tile edge, or 1 when block storage is not worthwhile. + */ +int ChooseHessianBlockSize(const Problem &problem, int max_block_size = 16); + +/** @brief Largest tile edge the block SpMV supports. */ +constexpr int kMaxHessianBlockSize = 16; + +/** + * @brief Extracts the main diagonal of a BSR matrix. + * + * @param stream CUDA stream for GPU operations. + * @param matrix BSR matrix. + * @param[out] diagonal Output of length matrix.NumRows(). + */ +void ExtractDiagonal(cudaStream_t stream, const BSRSparseMatrix &matrix, dvector &diagonal); + +/** + * @brief Computes result = matrix + scale * diag(diagonal) for BSR storage. + * + * @param stream CUDA stream for GPU operations. + * @param scale Scaling factor applied to the diagonal values. + * @param diagonal Diagonal values to add; length matrix.NumRows(). + * @param matrix Input BSR matrix. + * @param[out] result Output BSR matrix; may alias @p matrix for in-place. + */ +void AddScaledDiagonal(cudaStream_t stream, float scale, const dvector &diagonal, + const BSRSparseMatrix &matrix, BSRSparseMatrix &result); + +/** + * @brief Symmetric diagonal scaling: A_ij *= scale[i] * scale[j]. + * + * @param stream CUDA stream for GPU operations. + * @param[in,out] matrix BSR matrix updated in place. + * @param scale Length must equal matrix.NumRows(). + * @param[out] row_of_tile Caller-owned scratch mapping tile index to block row; + * resized as needed and rebuilt on each call. + */ +void ScaleSymmetric(cudaStream_t stream, BSRSparseMatrix &matrix, const dvector &scale, + dvector &row_of_tile); + +/** + * @brief Deep-copies a BSR matrix (structure and values). + * + * The structure is always copied, not just the values: equal sizes do not imply + * equal connectivity when one solver instance is reused across problems. + * + * @param stream CUDA stream for GPU operations. + * @param input Source matrix. + * @param[out] output Destination matrix. + */ +void CopyBSRSparseMatrix(cudaStream_t stream, const BSRSparseMatrix &input, + BSRSparseMatrix &output); + +/** + * @brief Expands a BSR matrix into scalar CSR. + * + * Used for solver backends that cannot consume block storage (cuDSS and the + * dense factorizations). Column indices come out sorted within each row. + * + * @param stream CUDA stream for GPU operations. + * @param input BSR matrix. + * @param[out] output CSR matrix; resized as needed. + * @param[out] row_of_tile Caller-owned scratch mapping tile index to block row. + */ +void ConvertBSRToCSR(cudaStream_t stream, const BSRSparseMatrix &input, CSRSparseMatrix &output, + dvector &row_of_tile); + +/** + * @brief Sparse matrix-vector product y = A * x for BSR storage. + * + * cuSPARSE's `cusparseSbsrmv` is not used: it measured 3.6x slower than + * `csrmv_v3` on a bundle-adjustment Hessian, which would negate the format's + * whole advantage. See bsr_matrix.cu. + * + * @param stream CUDA stream for GPU operations. + * @param matrix BSR matrix A. + * @param x Input vector of length matrix.NumRows(). + * @param[out] y Output vector; resized to matrix.NumRows(). + */ +void MultiplyBSRByDenseVector(cudaStream_t stream, const BSRSparseMatrix &matrix, + const dvector &x, dvector &y); + +/** + * @brief Async BSR-weighted squared step: d_out[0] = step^T A step. + * + * @param stream CUDA stream for GPU operations. + * @param matrix BSR matrix A. + * @param step Step vector. + * @param[out] scratch Scratch for the SpMV result. + * @param d_out Device destination for the scalar. + * @param d_partials Reduction scratch; see device_reduction.h. + */ +void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const BSRSparseMatrix &matrix, + const dvector &step, dvector &scratch, + float *d_out, float *d_partials); + +} // namespace cunls diff --git a/cunls/minimizer/cusparse_matrix_multiplier.cpp b/cunls/minimizer/cusparse_matrix_multiplier.cpp deleted file mode 100644 index 6cda14b..0000000 --- a/cunls/minimizer/cusparse_matrix_multiplier.cpp +++ /dev/null @@ -1,209 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "cunls/minimizer/cusparse_matrix_multiplier.h" - -#include - -#include - -#include "cunls/common/cusparse_helper.h" -#include "cunls/common/helper.h" -#include "cunls/common/types.h" -#include "cunls/minimizer/sparse_matrix.h" - -namespace cunls { - -constexpr cusparseOperation_t operation = CUSPARSE_OPERATION_NON_TRANSPOSE; - -cuSPARSESparseMatrixMultiplier::cuSPARSESparseMatrixMultiplier() { - cusparseSpGEMMDescr_t descr = nullptr; - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMM_createDescr(&descr)); - gemm_description_ = static_cast(descr); -} - -cuSPARSESparseMatrixMultiplier::~cuSPARSESparseMatrixMultiplier() { - WARN_ON_CUSPARSE_ERROR(cusparseSpGEMM_destroyDescr( - static_cast(gemm_description_))); -} - -void cuSPARSESparseMatrixMultiplier::Transpose(cudaStream_t stream, - const CSRSparseMatrix &matrix, - CSRSparseMatrix &transposed) { - int num_rows, num_cols, num_nonzeros; - ExtractMatrixMetadata(stream, matrix, num_rows, num_cols, num_nonzeros); - - if (transposed.values.size() != num_nonzeros) { - transposed.values.resize(num_nonzeros); - } - - if (transposed.col_ids.size() != num_nonzeros) { - transposed.col_ids.resize(num_nonzeros); - } - - if (transposed.row_offsets.size() != num_cols + 1) { - transposed.row_offsets.resize(num_cols + 1); - } - - auto cusparse_handle = - static_cast(handle_.GetHandle(stream)); - - size_t bufferSize = 0; - - THROW_ON_CUSPARSE_ERROR(cusparseCsr2cscEx2_bufferSize( - cusparse_handle, num_rows, num_cols, num_nonzeros, matrix.values.data(), - matrix.row_offsets.data(), matrix.col_ids.data(), - transposed.values.data(), transposed.row_offsets.data(), - transposed.col_ids.data(), CUDA_R_32F, CUSPARSE_ACTION_NUMERIC, - CUSPARSE_INDEX_BASE_ZERO, CUSPARSE_CSR2CSC_ALG_DEFAULT, &bufferSize)); - - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - - if (buffer1.size() < bufferSize) { - buffer1.resize(bufferSize); - } - - THROW_ON_CUSPARSE_ERROR(cusparseCsr2cscEx2( - cusparse_handle, num_rows, num_cols, num_nonzeros, matrix.values.data(), - matrix.row_offsets.data(), matrix.col_ids.data(), - transposed.values.data(), transposed.row_offsets.data(), - transposed.col_ids.data(), CUDA_R_32F, CUSPARSE_ACTION_NUMERIC, - CUSPARSE_INDEX_BASE_ZERO, CUSPARSE_CSR2CSC_ALG_DEFAULT, buffer1.data())); -} - -void cuSPARSESparseMatrixMultiplier::Initialize(cudaStream_t stream, - const Problem & /*problem*/, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) { - Transpose(stream, input, temp_matrix_); - - auto handle = handle_.GetHandle(stream); - - int num_rows, num_cols, num_nonzeros; - ExtractMatrixMetadata(stream, input, num_rows, num_cols, num_nonzeros); - - descrA_ = std::move(cuSPARSEMatrixDescription(num_cols, num_rows, - num_nonzeros, temp_matrix_)); - descrB_ = std::move( - cuSPARSEMatrixDescription(num_rows, num_cols, num_nonzeros, input)); - descrC_ = std::move(cuSPARSEMatrixDescription(num_cols, num_cols)); - - EstimateWork(handle); - ReuseNonzeros(handle); - - buffer1.clear(); - buffer2.clear(); - - int64_t C_num_rows, C_num_cols, C_nnz; - THROW_ON_CUSPARSE_ERROR(cusparseSpMatGetSize( - static_cast(descrC_.GetDescription()), &C_num_rows, - &C_num_cols, &C_nnz)); - - output.row_offsets.resize(C_num_rows + 1); - output.col_ids.resize(C_nnz); - output.values.resize(C_nnz, 0); - - descrC_.UpdatePointers(output); - ReuseCopy(handle); - - buffer3.clear(); -} - -void cuSPARSESparseMatrixMultiplier::ComputeSquaredMatrix( - cudaStream_t stream, const Problem & /*problem*/, - const CSRSparseMatrix &input, CSRSparseMatrix &output) { - Transpose(stream, input, temp_matrix_); - - auto handle = static_cast(handle_.GetHandle(stream)); - - constexpr float alpha = 1; - constexpr float beta = 0; - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_compute( - handle, operation, operation, &alpha, - static_cast(descrA_.GetDescription()), - static_cast(descrB_.GetDescription()), &beta, - static_cast(descrC_.GetDescription()), CUDA_R_32F, - CUSPARSE_SPGEMM_DEFAULT, - static_cast(gemm_description_))); -} - -void cuSPARSESparseMatrixMultiplier::EstimateWork(void *handle) { - auto h = static_cast(handle); - size_t bufferSize1 = 0; - - auto dA = static_cast(descrA_.GetDescription()); - auto dB = static_cast(descrB_.GetDescription()); - auto dC = static_cast(descrC_.GetDescription()); - auto gemm = static_cast(gemm_description_); - - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_workEstimation( - h, operation, operation, dA, dB, dC, CUSPARSE_SPGEMM_DEFAULT, gemm, - &bufferSize1, NULL)); - - buffer1.resize(bufferSize1); - - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_workEstimation( - h, operation, operation, dA, dB, dC, CUSPARSE_SPGEMM_DEFAULT, gemm, - &bufferSize1, buffer1.data())); -} - -void cuSPARSESparseMatrixMultiplier::ReuseNonzeros(void *handle) { - auto h = static_cast(handle); - size_t bufferSize2 = 0; - size_t bufferSize3 = 0; - size_t bufferSize4 = 0; - - auto dA = static_cast(descrA_.GetDescription()); - auto dB = static_cast(descrB_.GetDescription()); - auto dC = static_cast(descrC_.GetDescription()); - auto gemm = static_cast(gemm_description_); - - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_nnz( - h, operation, operation, dA, dB, dC, CUSPARSE_SPGEMM_DEFAULT, gemm, - &bufferSize2, NULL, &bufferSize3, NULL, &bufferSize4, NULL)); - - buffer2.resize(bufferSize2); - buffer3.resize(bufferSize3); - buffer4.resize(bufferSize4); - - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_nnz( - h, operation, operation, dA, dB, dC, CUSPARSE_SPGEMM_DEFAULT, gemm, - &bufferSize2, buffer2.data(), &bufferSize3, buffer3.data(), &bufferSize4, - buffer4.data())); -} - -void cuSPARSESparseMatrixMultiplier::ReuseCopy(void *handle) { - auto h = static_cast(handle); - size_t bufferSize5 = 0; - - auto dA = static_cast(descrA_.GetDescription()); - auto dB = static_cast(descrB_.GetDescription()); - auto dC = static_cast(descrC_.GetDescription()); - auto gemm = static_cast(gemm_description_); - - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_copy( - h, operation, operation, dA, dB, dC, CUSPARSE_SPGEMM_DEFAULT, gemm, - &bufferSize5, NULL)); - - buffer5.resize(bufferSize5); - - THROW_ON_CUSPARSE_ERROR(cusparseSpGEMMreuse_copy( - h, operation, operation, dA, dB, dC, CUSPARSE_SPGEMM_DEFAULT, gemm, - &bufferSize5, buffer5.data())); -} - -} // namespace cunls diff --git a/cunls/minimizer/cusparse_matrix_multiplier.h b/cunls/minimizer/cusparse_matrix_multiplier.h deleted file mode 100644 index c7107d0..0000000 --- a/cunls/minimizer/cusparse_matrix_multiplier.h +++ /dev/null @@ -1,94 +0,0 @@ -/* - * 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 "cunls/common/cusparse_helper.h" -#include "cunls/common/types.h" -#include "cunls/minimizer/sparse_matrix_multiplier.h" - -namespace cunls { - -/** - * @brief Computes A^T * A using cuSPARSE's SpGEMM reuse API. - * - * Caches the sparsity-pattern analysis across calls so that only the numeric - * phase is repeated when the pattern is unchanged. Internally transposes the - * input via CSR-to-CSC conversion and then multiplies A^T * A. - */ -class cuSPARSESparseMatrixMultiplier : public SparseMatrixMultiplier { -public: - /** @brief Constructs and initializes the cuSPARSE GEMM descriptor. */ - cuSPARSESparseMatrixMultiplier(); - - /** @brief Destroys the cuSPARSE GEMM descriptor and frees resources. */ - ~cuSPARSESparseMatrixMultiplier() override; - - /** - * @brief Initializes the GEMM structural analysis for A^T * A. - * - * Transposes the input matrix to obtain its structure, then runs the - * full three-phase cuSPARSE GEMM reuse setup (work estimation, nonzero - * analysis, copy preparation) and allocates the output matrix. - * Must be called once whenever the sparsity pattern changes. - * - * @param stream CUDA stream for GPU operations. - * @param problem Optimization problem (unused by this implementation). - * @param input Input sparse matrix A (typically the Jacobian in CSR). - * @param[out] output Output sparse matrix A^T * A (structure allocated). - */ - void Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) override; - - /** - * @brief Computes A^T * A for a sparse matrix A. - * - * Transposes the input matrix to update values, then performs the - * numeric phase of the cuSPARSE GEMM reuse API. - * - * @param stream CUDA stream for GPU operations. - * @param problem Optimization problem (unused by this implementation). - * @param input Input sparse matrix A (typically the Jacobian). - * @param[out] output Output sparse matrix A^T * A. - */ - void ComputeSquaredMatrix(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) override; - -private: - void Transpose(cudaStream_t stream, const CSRSparseMatrix &matrix, - CSRSparseMatrix &transposed); - void EstimateWork(void *handle); - void ReuseNonzeros(void *handle); - void ReuseCopy(void *handle); - - cuSPARSEHandle handle_; - CSRSparseMatrix temp_matrix_; - - void *gemm_description_ = nullptr; - - cuSPARSEMatrixDescription descrA_, descrB_, descrC_; - - dvector buffer1; - dvector buffer2; - dvector buffer3; - dvector buffer4; - dvector buffer5; -}; - -} // namespace cunls diff --git a/cunls/minimizer/fast_matrix_multiplier.cu b/cunls/minimizer/fast_matrix_multiplier.cu deleted file mode 100644 index 84c645b..0000000 --- a/cunls/minimizer/fast_matrix_multiplier.cu +++ /dev/null @@ -1,650 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include - -#include -#include -#include -#include -#include - -#include "cunls/common/helper.h" -#include "cunls/minimizer/fast_matrix_multiplier.h" -#include "cunls/minimizer/sparse_matrix.h" - -#define WARP_SIZE 32 - -namespace cunls { - -// ============================================================================ -// Device helpers -// ============================================================================ - -__device__ __forceinline__ int BinarySearchDevice(const int *__restrict__ arr, - int lo, int hi, int target) { - while (lo < hi) { - int mid = lo + ((hi - lo) >> 1); - if (arr[mid] < target) - lo = mid + 1; - else - hi = mid; - } - return lo; -} - -// ============================================================================ -// Position precomputation kernel (runs once during Initialize) -// ============================================================================ - -/** - * One warp per input row. Precomputes output positions for every (a, b) - * column pair using binary search, storing them in a flat map so that the - * compute kernel can avoid repeated searches. - * - * Layout: position_map[(start + a) * max_nnz_per_row + b] = output flat index. - */ -__global__ void -PrecomputePositionsKernel(const int *__restrict__ input_row_offsets, - const int *__restrict__ input_col_ids, - int num_input_rows, - const int *__restrict__ output_row_offsets, - const int *__restrict__ output_col_ids, - int *__restrict__ position_map, int max_nnz_per_row) { - extern __shared__ int s_cols[]; - - int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE; - int lane = threadIdx.x & (WARP_SIZE - 1); - int warp_in_block = threadIdx.x / WARP_SIZE; - - if (warp_id >= num_input_rows) - return; - - int *my_cols = s_cols + warp_in_block * max_nnz_per_row; - - int start = input_row_offsets[warp_id]; - int end = input_row_offsets[warp_id + 1]; - int L = end - start; - - for (int i = lane; i < L; i += WARP_SIZE) { - my_cols[i] = input_col_ids[start + i]; - } - __syncwarp(); - - for (int a = 0; a < L; a++) { - int col_a = my_cols[a]; - int out_start = output_row_offsets[col_a]; - int out_end = output_row_offsets[col_a + 1]; - int map_base = (start + a) * max_nnz_per_row; - - for (int b = lane; b < L; b += WARP_SIZE) { - int col_b = my_cols[b]; - int pos = BinarySearchDevice(output_col_ids, out_start, out_end, col_b); - position_map[map_base + b] = pos; - } - } -} - -// ============================================================================ -// Compute kernel -// ============================================================================ - -/** - * One warp per input row. Loads values into shared memory, then for every - * (a, b) pair looks up the precomputed output position from position_map - * and accumulates val_a * val_b with atomicAdd. - * - * Block-sparsity optimization: for block-sparse Jacobians (state blocks up - * to 16x16), the per-row nnz `L` is small (typically <= max_nnz_per_row). - * The 32 warp lanes are decomposed into `kAPerWarp = 32 / kThreadsPerA` - * sub-groups, each processing a different `a` value in parallel rather - * than serializing them in the outer loop. This halves (or more) the - * outer-loop iteration count vs. the naive one-`a`-at-a-time scheme. - * - * Shared memory layout (per block): - * [warp0_vals ... warpN_vals] - * - * @tparam kThreadsPerA Power-of-two number of lanes that cooperate on a - * single `a` value. Must be >= max_nnz_per_row - * (rounded up to a power of two), capped at 32. - */ -template -__global__ void ComputeJtJKernel(const int *__restrict__ input_row_offsets, - const float *__restrict__ input_values, - int num_input_rows, - const int *__restrict__ position_map, - float *__restrict__ output_values, - int max_nnz_per_row) { - static_assert(kThreadsPerA > 0 && (kThreadsPerA & (kThreadsPerA - 1)) == 0, - "kThreadsPerA must be a power of two"); - static_assert(kThreadsPerA <= WARP_SIZE, "kThreadsPerA must be <= 32"); - constexpr int kAPerWarp = WARP_SIZE / kThreadsPerA; - - extern __shared__ float s_vals[]; - - int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE; - int lane = threadIdx.x & (WARP_SIZE - 1); - int warp_in_block = threadIdx.x / WARP_SIZE; - - if (warp_id >= num_input_rows) { - return; - } - - float *my_vals = s_vals + warp_in_block * max_nnz_per_row; - - int start = input_row_offsets[warp_id]; - int end = input_row_offsets[warp_id + 1]; - int L = end - start; - - for (int i = lane; i < L; i += WARP_SIZE) { - my_vals[i] = input_values[start + i]; - } - __syncwarp(); - - // Lane decomposition: a_sub picks which `a` value within the kAPerWarp - // group this lane handles; b_lane is the lane's column within the - // kThreadsPerA group cooperating on that `a`. - const int a_sub = lane / kThreadsPerA; - const int b_lane = lane & (kThreadsPerA - 1); - - for (int a_base = 0; a_base < L; a_base += kAPerWarp) { - int a = a_base + a_sub; - if (a < L) { - float val_a = my_vals[a]; - int map_base = (start + a) * max_nnz_per_row; - for (int b = b_lane; b < L; b += kThreadsPerA) { - int pos = position_map[map_base + b]; - atomicAdd(&output_values[pos], val_a * my_vals[b]); - } - } - } -} - -// Picks the smallest power-of-two >= max_nnz_per_row, capped at WARP_SIZE. -// Lower values let multiple `a` columns be processed in parallel within -// one warp; for max_nnz_per_row <= 16 this gives a 2x or better reduction -// in serialized outer-loop iterations. -static int PickThreadsPerA(int max_nnz_per_row) { - int t = 1; - while (t < max_nnz_per_row && t < WARP_SIZE) { - t *= 2; - } - return t; -} - -// ============================================================================ -// Host helpers -// ============================================================================ - -static int GetMaxNnzPerRow(const Problem &problem) { - int max_nnz = 0; - for (const auto &batch : problem.GetResidualBatches()) { - auto sizes = batch.GetFactorBatch()->StateBlockSizes(); - int nnz = 0; - for (auto s : sizes) - nnz += static_cast(s); - max_nnz = std::max(max_nnz, nnz); - } - return max_nnz; -} - -static void ComputeBlockConfig(int max_nnz_per_row, int smem_per_element, - int &warps_per_block, int &block_size, - size_t &smem_bytes) { - constexpr int kMaxSharedMem = 48 * 1024; - int smem_per_warp = max_nnz_per_row * smem_per_element; - warps_per_block = - std::min(8, smem_per_warp > 0 ? kMaxSharedMem / smem_per_warp : 8); - warps_per_block = std::max(warps_per_block, 1); - block_size = warps_per_block * WARP_SIZE; - smem_bytes = static_cast(warps_per_block) * smem_per_warp; -} - -// ============================================================================ -// Structure-aware sparsity pattern computation -// ============================================================================ - -struct ExpandPair { - int row_offset; - int col_offset; - int row_tangent; - int col_tangent; - int write_offset; -}; - -/** - * One thread per block pair. Atomically increments per-row non-zero counts - * for all rows spanned by the block pair. - */ -__global__ void -ComputeBlockRowCountsKernel(const ExpandPair *__restrict__ pairs, int num_pairs, - int *__restrict__ row_counts) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_pairs) - return; - auto p = pairs[idx]; - for (int i = 0; i < p.row_tangent; i++) { - atomicAdd(&row_counts[p.row_offset + i], p.col_tangent); - } -} - -/** - * One thread block per block pair. Writes dense column indices for the - * sub-block into the CSR col_ids array using precomputed row offsets and - * per-pair write offsets. - */ -__global__ void ExpandBlockPairsKernel(const ExpandPair *__restrict__ pairs, - int num_pairs, - const int *__restrict__ row_offsets, - int *__restrict__ col_ids) { - int pair_idx = blockIdx.x; - if (pair_idx >= num_pairs) - return; - auto p = pairs[pair_idx]; - int total = p.row_tangent * p.col_tangent; - for (int k = threadIdx.x; k < total; k += blockDim.x) { - int i = k / p.col_tangent; - int j = k % p.col_tangent; - col_ids[row_offsets[p.row_offset + i] + p.write_offset + j] = - p.col_offset + j; - } -} - -/** - * Derives the J^T J sparsity pattern from the factor graph connectivity. - * Each factor connecting state blocks A,B produces dense sub-blocks - * (A,A), (A,B), (B,A), (B,B) in the Hessian. Block pairs are collected - * on the host, deduplicated via a hash set, and expanded into a CSR on - * the GPU. - * - * Block-sparsity optimization: rather than push every (row, col, row_tang, - * col_tang) tuple per factor (millions of entries) and then sort+unique, - * we deduplicate on-the-fly with an unordered_set keyed on - * (row_offset << 32 | col_offset). Tangent sizes are constant per - * column-offset and looked up after dedup. All const-state-id D2H copies - * are issued first and a single stream sync gates host processing. - * - * Uses buffer_ as scratch space to avoid runtime GPU memory allocation. - */ -void FastSparseMatrixMultiplier::ComputeOutputStructure(cudaStream_t stream, - const Problem &problem, - CSRSparseMatrix &output, - int num_cols) { - struct BlockInfo { - int col_offset; - int tangent_size; - }; - - // Per-batch descriptor for O(1) pointer → column lookup via arithmetic. - struct BatchDesc { - const float *base; - int ambient_size; - int tangent_size; - int num_blocks; - std::vector col_offsets; // [block_idx] → column, -1 if constant - }; - - const auto &state_batches = problem.GetStateBatches(); - std::vector batch_descs; - batch_descs.reserve(state_batches.size()); - - // Phase 1: issue all const-state-id D2H copies into one pinned buffer - // and synchronize the stream once instead of once per batch. - std::vector const_offsets(state_batches.size() + 1, 0); - for (size_t i = 0; i < state_batches.size(); i++) { - const_offsets[i + 1] = - const_offsets[i] + state_batches[i]->NumConstStateBlocks(); - } - size_t total_const = const_offsets.back(); - - if (total_const > 0) { - if (pinned_buf_.size() < total_const) { - pinned_buf_.resize(total_const); - } - for (size_t i = 0; i < state_batches.size(); i++) { - size_t nc = state_batches[i]->NumConstStateBlocks(); - if (nc > 0) { - THROW_ON_CUDA_ERROR(cudaMemcpyAsync( - pinned_buf_.data() + const_offsets[i], - state_batches[i]->ConstStateIds(), nc * sizeof(int), - cudaMemcpyDeviceToHost, stream)); - } - } - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - } - - // Phase 2: build per-batch column-offset tables. - int last_col = 0; - for (size_t i = 0; i < state_batches.size(); i++) { - auto *batch = state_batches[i]; - BatchDesc bd; - bd.base = batch->StateBlockDevicePtr(0); - bd.ambient_size = static_cast(batch->AmbientSize()); - bd.tangent_size = static_cast(batch->TangentSize()); - bd.num_blocks = static_cast(batch->NumStateBlocks()); - bd.col_offsets.assign(bd.num_blocks, 0); - - size_t nc = state_batches[i]->NumConstStateBlocks(); - std::vector is_const(bd.num_blocks, false); - for (size_t k = 0; k < nc; k++) { - int idx = pinned_buf_[const_offsets[i] + k]; - if (idx >= 0 && idx < bd.num_blocks) { - is_const[idx] = true; - } - } - - for (int j = 0; j < bd.num_blocks; j++) { - if (is_const[j]) { - bd.col_offsets[j] = -1; - } else { - bd.col_offsets[j] = last_col; - last_col += bd.tangent_size; - } - } - batch_descs.push_back(std::move(bd)); - } - - // Pre-build a col_offset -> tangent_size map. Tangent size is constant - // per state batch, so every column offset has a fixed tangent. - std::unordered_map col_to_tangent; - col_to_tangent.reserve(static_cast(num_cols)); - for (const auto &bd : batch_descs) { - for (int col : bd.col_offsets) { - if (col >= 0) { - col_to_tangent.emplace(col, bd.tangent_size); - } - } - } - - // Lambda: resolve device pointer → BlockInfo via pointer arithmetic. - auto resolve_ptr = [&](const float *ptr) -> BlockInfo { - for (const auto &bd : batch_descs) { - auto diff = ptr - bd.base; - if (diff >= 0 && - diff < static_cast(bd.num_blocks) * bd.ambient_size) { - int idx = static_cast(diff / bd.ambient_size); - int col = bd.col_offsets[idx]; - if (col >= 0) { - return {col, bd.tangent_size}; - } - return {-1, 0}; - } - } - return {-1, 0}; - }; - - const auto &res_batches = problem.GetResidualBatches(); - const auto &state_ptrs = problem.GetStatePointers(); - - // Estimate the number of unique (row, col) block pairs. We don't know - // the real count up front; reserving a healthy fraction of the worst - // case avoids most rehashes for typical workloads. - size_t total_pair_estimate = 0; - for (size_t rb = 0; rb < res_batches.size(); rb++) { - auto *factor = res_batches[rb].GetFactorBatch(); - size_t nb = factor->StateBlockSizes().size(); - total_pair_estimate += factor->NumFactors() * nb * nb; - } - - // Hash-set dedup: O(1) insert per (row, col) pair vs. the previous - // sort+unique on millions of 16-byte structs. Keys pack - // (row_offset, col_offset) into a single 64-bit integer. - std::unordered_set seen_pairs; - seen_pairs.reserve(std::min(total_pair_estimate, - static_cast(1) << 22)); - - // Per-factor scratch. 16x16 block sparse means at most ~16 blocks per - // factor; stack array avoids per-factor heap allocations. - constexpr int kMaxBlocksPerFactor = 64; - int local_cols[kMaxBlocksPerFactor]; - std::vector local_cols_dyn; - - for (size_t rb = 0; rb < res_batches.size(); rb++) { - auto *factor = res_batches[rb].GetFactorBatch(); - size_t nf = factor->NumFactors(); - size_t nb = factor->StateBlockSizes().size(); - const auto &h_ptrs = state_ptrs[rb]; - - int *cols_ptr; - if (nb <= kMaxBlocksPerFactor) { - cols_ptr = local_cols; - } else { - local_cols_dyn.resize(nb); - cols_ptr = local_cols_dyn.data(); - } - - for (size_t f = 0; f < nf; f++) { - int nblocks = 0; - const float *const *fp = h_ptrs.data() + f * nb; - for (size_t b = 0; b < nb; b++) { - auto info = resolve_ptr(fp[b]); - if (info.col_offset >= 0) { - cols_ptr[nblocks++] = info.col_offset; - } - } - for (int a = 0; a < nblocks; a++) { - uint64_t row_part = static_cast( - static_cast(cols_ptr[a])) - << 32; - for (int b = 0; b < nblocks; b++) { - uint64_t key = - row_part | static_cast(cols_ptr[b]); - seen_pairs.insert(key); - } - } - } - } - - // Materialize into a sorted BlockPair vector for offset computation. - // The unique count is O(num_state_blocks^2) in the worst case but - // typically much smaller, so this sort dominates only at <<1% scale - // compared to the previous sort over all (factor x nb x nb) entries. - struct BlockPair { - int row_offset; - int col_offset; - int row_tangent; - int col_tangent; - }; - - std::vector pairs; - pairs.reserve(seen_pairs.size()); - for (uint64_t key : seen_pairs) { - int row = static_cast(key >> 32); - int col = static_cast(key & 0xFFFFFFFFu); - auto row_it = col_to_tangent.find(row); - auto col_it = col_to_tangent.find(col); - int row_t = row_it != col_to_tangent.end() ? row_it->second : 0; - int col_t = col_it != col_to_tangent.end() ? col_it->second : 0; - pairs.push_back({row, col, row_t, col_t}); - } - std::sort(pairs.begin(), pairs.end(), - [](const BlockPair &a, const BlockPair &b) { - if (a.row_offset != b.row_offset) { - return a.row_offset < b.row_offset; - } - return a.col_offset < b.col_offset; - }); - - // Compute per-pair write offsets: within a block row, each successive - // block column starts after the previous one's tangent width. - std::vector gpu_pairs(pairs.size()); - int prev_row = -1; - int cum = 0; - for (size_t k = 0; k < pairs.size(); k++) { - if (pairs[k].row_offset != prev_row) { - cum = 0; - prev_row = pairs[k].row_offset; - } - gpu_pairs[k] = {pairs[k].row_offset, pairs[k].col_offset, - pairs[k].row_tangent, pairs[k].col_tangent, cum}; - cum += pairs[k].col_tangent; - } - - static_assert(sizeof(ExpandPair) == 5 * sizeof(int)); - int num_pairs = static_cast(gpu_pairs.size()); - size_t pairs_ints = static_cast(num_pairs) * 5; - - buffer_.resize(pairs_ints + num_cols); - auto *d_pairs = reinterpret_cast(buffer_.data()); - int *row_counts = buffer_.data() + pairs_ints; - - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(d_pairs, gpu_pairs.data(), - num_pairs * sizeof(ExpandPair), - cudaMemcpyHostToDevice, stream)); - THROW_ON_CUDA_ERROR( - cudaMemsetAsync(row_counts, 0, num_cols * sizeof(int), stream)); - - { - constexpr int kBlockSize = 256; - int grid = (num_pairs + kBlockSize - 1) / kBlockSize; - ComputeBlockRowCountsKernel<<>>( - d_pairs, num_pairs, row_counts); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - } - - output.row_offsets.resize(num_cols + 1); - THROW_ON_CUDA_ERROR( - cudaMemsetAsync(output.row_offsets.data(), 0, sizeof(int), stream)); - - { - auto stream_policy = thrust::cuda::par_nosync.on(stream); - thrust::device_ptr counts_ptr(row_counts); - thrust::device_ptr offsets_ptr(output.row_offsets.data()); - thrust::inclusive_scan(stream_policy, counts_ptr, counts_ptr + num_cols, - offsets_ptr + 1); - } - - if (pinned_buf_.size() < 1) { - pinned_buf_.resize(1); - } - THROW_ON_CUDA_ERROR( - cudaMemcpyAsync(pinned_buf_.data(), output.row_offsets.data() + num_cols, - sizeof(int), cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - int total_nnz = pinned_buf_[0]; - - output.col_ids.resize(total_nnz); - output.values.resize(total_nnz); - - { - constexpr int kBlockSize = 256; - ExpandBlockPairsKernel<<>>( - d_pairs, num_pairs, output.row_offsets.data(), output.col_ids.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - } -} - -// ============================================================================ -// FastSparseMatrixMultiplier -// ============================================================================ - -void FastSparseMatrixMultiplier::Initialize(cudaStream_t stream, - const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) { - if (input.row_offsets.size() <= 1 || input.values.empty()) { - max_nnz_per_row_ = 0; - return; - } - - int num_rows, num_cols, num_nonzeros; - ExtractMatrixMetadata(stream, input, num_rows, num_cols, num_nonzeros); - - max_nnz_per_row_ = GetMaxNnzPerRow(problem); - if (max_nnz_per_row_ == 0) { - return; - } - - ComputeOutputStructure(stream, problem, output, num_cols); - - // Precompute output positions for every input (a, b) pair so - // that ComputeSquaredMatrix avoids binary searches entirely. - { - size_t map_size = static_cast(num_nonzeros) * max_nnz_per_row_; - position_map_.resize(map_size); - - int warps_per_block, block_size; - size_t smem; - ComputeBlockConfig(max_nnz_per_row_, static_cast(sizeof(int)), - warps_per_block, block_size, smem); - int grid = (num_rows + warps_per_block - 1) / warps_per_block; - - PrecomputePositionsKernel<<>>( - input.row_offsets.data(), input.col_ids.data(), num_rows, - output.row_offsets.data(), output.col_ids.data(), position_map_.data(), - max_nnz_per_row_); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - } -} - -void FastSparseMatrixMultiplier::ComputeSquaredMatrix( - cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &input, - CSRSparseMatrix &output) { - int num_input_rows = static_cast(input.row_offsets.size()) - 1; - if (num_input_rows <= 0 || max_nnz_per_row_ == 0 || output.values.empty()) { - return; - } - - THROW_ON_CUDA_ERROR(cudaMemsetAsync( - output.values.data(), 0, output.values.size() * sizeof(float), stream)); - - int warps_per_block, block_size; - size_t smem; - ComputeBlockConfig(max_nnz_per_row_, static_cast(sizeof(float)), - warps_per_block, block_size, smem); - int grid = (num_input_rows + warps_per_block - 1) / warps_per_block; - - // Dispatch to a kernel specialization sized to the per-row nnz so that - // multiple `a` columns are processed in parallel within one warp. - int threads_per_a = PickThreadsPerA(max_nnz_per_row_); - switch (threads_per_a) { - case 1: - ComputeJtJKernel<1><<>>( - input.row_offsets.data(), input.values.data(), num_input_rows, - position_map_.data(), output.values.data(), max_nnz_per_row_); - break; - case 2: - ComputeJtJKernel<2><<>>( - input.row_offsets.data(), input.values.data(), num_input_rows, - position_map_.data(), output.values.data(), max_nnz_per_row_); - break; - case 4: - ComputeJtJKernel<4><<>>( - input.row_offsets.data(), input.values.data(), num_input_rows, - position_map_.data(), output.values.data(), max_nnz_per_row_); - break; - case 8: - ComputeJtJKernel<8><<>>( - input.row_offsets.data(), input.values.data(), num_input_rows, - position_map_.data(), output.values.data(), max_nnz_per_row_); - break; - case 16: - ComputeJtJKernel<16><<>>( - input.row_offsets.data(), input.values.data(), num_input_rows, - position_map_.data(), output.values.data(), max_nnz_per_row_); - break; - default: - ComputeJtJKernel<32><<>>( - input.row_offsets.data(), input.values.data(), num_input_rows, - position_map_.data(), output.values.data(), max_nnz_per_row_); - break; - } - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - -} // namespace cunls diff --git a/cunls/minimizer/fast_matrix_multiplier.h b/cunls/minimizer/fast_matrix_multiplier.h deleted file mode 100644 index 69b8eea..0000000 --- a/cunls/minimizer/fast_matrix_multiplier.h +++ /dev/null @@ -1,56 +0,0 @@ -/* - * 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 "cunls/common/profiler.h" -#include "cunls/common/types.h" -#include "cunls/minimizer/sparse_matrix_multiplier.h" - -namespace cunls { - -/** - * @brief Computes J^T * J using custom warp-efficient CUDA kernels. - * - * Uses the Problem's factor graph connectivity to derive the output sparsity - * pattern during initialization, then warp-cooperative scatter-multiply - * during compute. Requires a Problem with at least one residual batch. - */ -class FastSparseMatrixMultiplier : public SparseMatrixMultiplier { -public: - void Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) override; - - void ComputeSquaredMatrix(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) override; - -private: - void ComputeOutputStructure(cudaStream_t stream, const Problem &problem, - CSRSparseMatrix &output, int num_cols); - - int max_nnz_per_row_ = 0; - dvector position_map_; ///< Precomputed output positions, indexed as - ///< [input_nnz_idx * max_nnz_per_row + b]. - dvector buffer_; ///< Reusable scratch buffer for initialization. - pvector pinned_buf_; ///< Reusable pinned buffer for D2H readbacks. - profiler::Domain profiler_domain_{ - "FastSparseMatrixMultiplier"}; ///< Profiling domain. -}; - -} // namespace cunls diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index 981fe7f..22c6aa8 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -41,8 +41,7 @@ namespace cunls { GaussNewtonMinimizer::GaussNewtonMinimizer(const MinimizerOptions &options) : options_(options), solver_(CreateCSRSparseLinearSolver(options_.sparse_linear_solver_type, - options_.sparse_linear_solver_config)), - gemm_(CreateSparseMatrixMultiplier(options_.sparse_square_multiplier_type)) { + options_.sparse_linear_solver_config)) { if (options_.disable_safety_checks) { solver_->DisableSafetyChecks(); } @@ -68,15 +67,10 @@ void InitializeResiduals(const Problem &problem, dvector &residuals) { } } -void GaussNewtonMinimizer::InitializeJacobian(cudaStream_t stream, const Problem &problem) { - current_state_.BuildTripletSparseStructure(stream, problem, sparse_jacobian_.structure); - - // No sync needed: col_ids was resized synchronously on the host inside - // BuildTripletSparseStructure, so its size is already known, and the - // independent values buffer is only consumed later on the same stream. - size_t jacobian_size = sparse_jacobian_.structure.col_ids.size(); - if (sparse_jacobian_.values.size() != jacobian_size) { - sparse_jacobian_.values.resize(jacobian_size); +void GaussNewtonMinimizer::ResizeFactorJacobians() { + size_t num_floats = normal_equations_.JacobianValuesSize(); + if (factor_jacobians_.size() != num_floats) { + factor_jacobians_.resize(num_floats); } } @@ -154,19 +148,18 @@ float GaussNewtonMinimizer::ComputeCost(cudaStream_t stream, const Problem &prob /** * @brief Computes residuals and Jacobian for the current states. * - * Evaluates all factor batches to compute residual values and their - * Jacobian matrices. The residuals are stored in a dense vector, and the - * Jacobian is stored in COO (triplet) sparse format. + * Evaluates all factor batches to compute residual values and their Jacobian + * matrices. Both are dense per-factor blocks, concatenated across batches. * * @param stream CUDA stream for GPU operations. * @param problem The optimization problem. * @param minimizer_state Current minimizer state. * @param[out] residuals Output residual vector. - * @param[out] coo_jacobian Output Jacobian in COO format. + * @param[out] jacobians Output per-factor dense Jacobian blocks. */ void ComputeResidualAndJacobian(cudaStream_t stream, const Problem &problem, const MinimizerState &minimizer_state, dvector &residuals, - SparseJacobian &coo_jacobian, dvector &buffer) { + PerFactorJacobians &jacobians, dvector &buffer) { const auto &state_pointers = minimizer_state.GetStatePointers(); const auto &residual_batches = problem.GetResidualBatches(); size_t max_n = 0; @@ -180,7 +173,7 @@ void ComputeResidualAndJacobian(cudaStream_t stream, const Problem &problem, float *workspace_ptr = reinterpret_cast(buffer.data()); float *residuals_ptr = residuals.data(); - float *jacobian_ptr = coo_jacobian.values.data(); + float *jacobian_ptr = jacobians.data(); for (size_t i = 0; i < residual_batches.size(); i++) { const auto &rb = residual_batches[i]; @@ -215,21 +208,18 @@ void ComputeResidualAndJacobian(cudaStream_t stream, const Problem &problem, * @param[out] rhs Output right-hand side vector (-J^T r). */ void GaussNewtonMinimizer::ApplyColumnScalingToNormalEquations(cudaStream_t stream, - CSRSparseMatrix &lhs, dvector &rhs) { if (options_.column_scaling == ColumnScaling::None) { return; } - if (options_.column_scaling == ColumnScaling::HessianDiagonal) { - ExtractDiagonal(stream, hessian_, column_scale_); - InvertSqrtWithFloorInPlace(stream, column_scale_); - } else { - ComputeJacobianColumnScaling(stream, csr_jacobian_, jacobian_dims_.num_cols, - jacobian_dims_.num_nonzeros, column_scale_); - } + // H_jj = ||J_{:,j}||^2, so scaling by the Hessian diagonal is the same as + // scaling by the Jacobian column norms -- and no global Jacobian exists to + // read the latter from. + normal_equations_.ExtractHessianDiagonal(stream, column_scale_); + InvertSqrtWithFloorInPlace(stream, column_scale_); - ScaleSymmetricCSR(stream, lhs, column_scale_); + normal_equations_.ScaleLhsSymmetric(stream, column_scale_); ElementwiseMultiplyInPlace(stream, rhs.data(), column_scale_.data(), rhs.size()); } @@ -242,23 +232,13 @@ void GaussNewtonMinimizer::MapScaledLinearSolutionToTangentStep(cudaStream_t str } void GaussNewtonMinimizer::BuildSystem(cudaStream_t stream, const Problem &problem, - const MinimizerState &minimizer_state, CSRSparseMatrix &lhs, - dvector &rhs) { + const MinimizerState &minimizer_state) { auto range = profiler_domain_.CreateDomainRange("BuildSystem"); - ComputeResidualAndJacobian(stream, problem, minimizer_state, residuals_, sparse_jacobian_, + ComputeResidualAndJacobian(stream, problem, minimizer_state, residuals_, factor_jacobians_, buffer_); - - // Copy values from triplet Jacobian to precomputed CSR structure using - // mapping - ConvertTripletToCSRValues(stream, sparse_jacobian_, csr_mapping_, csr_jacobian_); - - gemm_->ComputeSquaredMatrix(stream, problem, csr_jacobian_, hessian_); - CopyCSRSparseMatrix(stream, hessian_, lhs); - auto handle = cusparse_handle_.GetHandle(stream); - ComputeRHS(stream, handle, csr_jacobian_, jacobian_dims_.num_rows, jacobian_dims_.num_cols, - jacobian_dims_.num_nonzeros, residuals_, rhs, buffer_); - - ApplyColumnScalingToNormalEquations(stream, lhs, rhs); + normal_equations_.Assemble(stream, problem, factor_jacobians_.data(), residuals_.data(), + rhs_work_); + ApplyColumnScalingToNormalEquations(stream, rhs_work_); } /** @@ -301,32 +281,14 @@ void GaussNewtonMinimizer::UpdateStates(cudaStream_t stream, const MinimizerStat */ void GaussNewtonMinimizer::Initialize(cudaStream_t stream, Problem &problem) { auto range = profiler_domain_.CreateDomainRange("Initialize"); - jacobian_dims_.Invalidate(); - hessian_dims_.Invalidate(); - InitializeResiduals(problem, residuals_); - InitializeJacobian(stream, problem); state_ops_.Preprocess(stream, problem.GetStateBatches()); - // Convert Jacobian triplet structure to CSR once. The structure doesn't - // change across iterations; only values are updated via the mapping. - auto handle = cusparse_handle_.GetHandle(stream); - - { - auto r1 = profiler_domain_.CreateDomainRange("ConvertTripletStructureToCSR"); - ConvertTripletStructureToCSR(stream, handle, sparse_jacobian_.structure, csr_jacobian_, - csr_mapping_, buffer_); - } - - gemm_->Initialize(stream, problem, csr_jacobian_, hessian_); - - { - int nr, nc, nnz; - ExtractMatrixMetadata(stream, csr_jacobian_, nr, nc, nnz); - jacobian_dims_.Set(nr, nc, nnz); - ExtractMatrixMetadata(stream, hessian_, nr, nc, nnz); - hessian_dims_.Set(nr, nc, nnz); - } + // The Hessian pattern comes straight from the factor graph; no global + // Jacobian is involved. + normal_equations_.Initialize(stream, problem, static_cast(state_ops_.NumReducedStates()), + solver_->SupportsBlockStorage()); + ResizeFactorJacobians(); } /** @@ -490,14 +452,14 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, Problem &pr } // Build initial linear system - BuildSystem(stream, problem, current_state_, lhs_work_, rhs_work_); + BuildSystem(stream, problem, current_state_); step_.resize(rhs_work_.size()); { // Perform symbolic analysis on the requested CUDA stream. auto sa_range = profiler_domain_.CreateDomainRange("PerformSymbolicAnalysis"); - bool success = solver_->Initialize(stream, problem, lhs_work_, rhs_work_, step_); + bool success = normal_equations_.InitializeSolver(stream, *solver_, problem, rhs_work_, step_); if (!success) { std::string str = "Failed to initialize linear solver"; LogError(str); @@ -517,7 +479,7 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, Problem &pr { auto solve_range = profiler_domain_.CreateDomainRange("LinearSolve"); - bool success = solver_->Solve(stream, lhs_work_, rhs_work_, step_); + bool success = normal_equations_.Solve(stream, *solver_, rhs_work_, step_); if (!success) { std::string str = "Failed to solve linear system"; LogError(str); @@ -563,7 +525,7 @@ MinimizerSummary GaussNewtonMinimizer::Minimize(cudaStream_t stream, Problem &pr current_state_.Copy(stream, updated_state_.GetStates()); } - BuildSystem(stream, problem, current_state_, lhs_work_, rhs_work_); + BuildSystem(stream, problem, current_state_); }; LogMessage("Optimization finished"); diff --git a/cunls/minimizer/gauss_newton_minimizer.h b/cunls/minimizer/gauss_newton_minimizer.h index d3b466e..a0f5535 100644 --- a/cunls/minimizer/gauss_newton_minimizer.h +++ b/cunls/minimizer/gauss_newton_minimizer.h @@ -27,9 +27,9 @@ #include "cunls/common/types.h" #include "cunls/linear_solver/sparse_linear_solver.h" #include "cunls/minimizer/minimizer_state.h" +#include "cunls/minimizer/normal_equations.h" #include "cunls/minimizer/problem.h" #include "cunls/minimizer/sparse_matrix.h" -#include "cunls/minimizer/sparse_matrix_multiplier.h" #include "cunls/state/state_batch_ops.h" namespace cunls { @@ -44,10 +44,12 @@ namespace cunls { enum class ColumnScaling { /** No scaling (identity S). */ None = 0, - /** \f$S_{ii} = 1 / \sqrt{H_{ii}}\f$ with a floor on the diagonal. */ + /** + * \f$S_{ii} = 1 / \sqrt{H_{ii}}\f$ with a floor on the diagonal. + * + * Equivalently \f$1 / \|J_{:,j}\|_2\f$, since \f$H_{jj} = \|J_{:,j}\|_2^2\f$. + */ HessianDiagonal = 1, - /** \f$S_{jj} = 1 / \|J_{:,j}\|_2\f$ from the CSR Jacobian. */ - JacobianColumnNorm = 2, }; /** @@ -124,8 +126,7 @@ struct MinimizerOptions { * * Default: BlockSparsePCG. */ - SparseLinearSolverType sparse_linear_solver_type = - SparseLinearSolverType::BlockSparsePCG; + SparseLinearSolverType sparse_linear_solver_type = SparseLinearSolverType::BlockSparsePCG; /** * @brief Configuration for the sparse linear solver. @@ -145,20 +146,6 @@ struct MinimizerOptions { */ SparseLinearSolverConfig sparse_linear_solver_config = {}; - /** - * @brief Strategy for computing the approximate Hessian J^T * J. - * - * - ``cuSPARSE``: uses cuSPARSE SpGEMM reuse API (transpose + multiply). - * Robust and well-tested; may allocate large internal work buffers. - * - ``Fast``: fast warp-efficient CUDA kernels with bitmap-based - * sparsity pattern discovery. Exploits the Problem's factor layout - * for kernel tuning. - * - * Default: Fast. - */ - SparseMatrixMultiplierType sparse_square_multiplier_type = - SparseMatrixMultiplierType::Fast; - /** * @brief Optional diagonal scaling of the GN/LM normal equations. * @@ -247,7 +234,7 @@ class GaussNewtonMinimizer { */ MinimizerSummary Minimize(cudaStream_t stream, Problem &problem); -protected: + protected: /** * @brief Checks if convergence criteria are satisfied. * @@ -263,9 +250,8 @@ class GaussNewtonMinimizer { * (updated_cost / current_cost). * @return True if converged, false otherwise. */ - virtual bool CheckConvergence(cudaStream_t stream, float updated_cost, - float current_cost, const dvector &step, - float &step_quality); + virtual bool CheckConvergence(cudaStream_t stream, float updated_cost, float current_cost, + const dvector &step, float &step_quality); /** * @brief Fused cost evaluation + convergence check with a single D2H + sync. @@ -283,11 +269,10 @@ class GaussNewtonMinimizer { * @param[out] step_quality Step quality metric. * @return True if converged. */ - virtual bool - EvaluateAndCheckConvergence(cudaStream_t stream, const Problem &problem, - const MinimizerState &updated_state, - float current_cost, const dvector &step, - float &updated_cost, float &step_quality); + virtual bool EvaluateAndCheckConvergence(cudaStream_t stream, const Problem &problem, + const MinimizerState &updated_state, float current_cost, + const dvector &step, float &updated_cost, + float &step_quality); /** * @brief Determines if a step should be accepted. @@ -331,12 +316,9 @@ class GaussNewtonMinimizer { * @param stream CUDA stream for GPU operations. * @param problem The optimization problem. * @param minimizer_state Current minimizer state. - * @param[out] lhs Output left-hand side matrix (H = J^T J). - * @param[out] rhs Output right-hand side vector (-J^T r). */ virtual void BuildSystem(cudaStream_t stream, const Problem &problem, - const MinimizerState &minimizer_state, - CSRSparseMatrix &lhs, dvector &rhs); + const MinimizerState &minimizer_state); /** * @brief Updates states with the computed step. @@ -368,26 +350,22 @@ class GaussNewtonMinimizer { * Caller must copy and sync to read the scalar. */ void ComputeCostAsync(cudaStream_t stream, const Problem &problem, - const MinimizerState &minimizer_state, - float *d_cost_out); + const MinimizerState &minimizer_state, float *d_cost_out); private: /** * @brief Applies diagonal column scaling to the normal-equation system. * * After the unscaled Hessian copy is in lhs and rhs holds b = -J^T r, - * this may replace the solve target with S H S z = S b: fills column_scale_, - * scales lhs symmetrically, and sets rhs_i *= S_i. When column_scaling is - * None, returns immediately (hessian_ is unchanged and already separate). + * After Assemble() the working left-hand side holds H and rhs holds + * b = -J^T r, this may replace the solve target with S H S z = S b: it fills + * column_scale_, scales the left-hand side symmetrically, and sets + * rhs_i *= S_i. No-op when column_scaling is None. * * @param stream CUDA stream for GPU work. - * @param[in,out] lhs Approximate Hessian H = J^T J (scaled in-place when - * enabled). * @param[in,out] rhs Right-hand side b (elementwise-scaled when enabled). */ - void ApplyColumnScalingToNormalEquations(cudaStream_t stream, - CSRSparseMatrix &lhs, - dvector &rhs); + void ApplyColumnScalingToNormalEquations(cudaStream_t stream, dvector &rhs); /** * @brief Maps the scaled linear unknown z to the manifold tangent step dx = S @@ -401,32 +379,26 @@ class GaussNewtonMinimizer { * @param[in,out] step Solution vector from the linear solver; overwritten by * dx. */ - void MapScaledLinearSolutionToTangentStep(cudaStream_t stream, - dvector &step); + void MapScaledLinearSolutionToTangentStep(cudaStream_t stream, dvector &step); - /** - * @brief Builds Jacobian COO structure and resizes value buffer. - */ - void InitializeJacobian(cudaStream_t stream, const Problem &problem); + /** @brief Sizes the per-factor Jacobian buffer for the problem. */ + void ResizeFactorJacobians(); -protected: + protected: const MinimizerOptions options_; ///< Optimizer configuration options. - SparseLinearSolverPtr solver_; ///< Linear solver for Gauss-Newton system. - SparseMatrixMultiplierPtr gemm_; ///< Matrix multiplication for H = J^T J. + SparseLinearSolverPtr solver_; ///< Linear solver for the normal equations. cuSPARSEHandle cusparse_handle_; ///< cuSPARSE handle for sparse operations. StateBatchOps state_ops_; ///< Operations on state batches. dvector residuals_; ///< Residual vector storage. - SparseJacobian sparse_jacobian_; ///< Jacobian in COO (triplet) format. - CSRSparseMatrix csr_jacobian_; ///< Jacobian in CSR format. - CSRMatrixDimensions jacobian_dims_; ///< Cached Jacobian dimensions. - dvector csr_mapping_; ///< Mapping from triplet to CSR indices. + /// Per-factor dense Jacobian blocks; the only Jacobian ever materialized. + PerFactorJacobians factor_jacobians_; - CSRSparseMatrix hessian_; ///< Approximate Hessian H = J^T J. - CSRMatrixDimensions hessian_dims_; ///< Cached Hessian dimensions. + /// The assembled system: Hessian, working left-hand side, and their storage. + NormalEquations normal_equations_; /// Diagonal S when column_scaling is enabled; size = number of tangent DOFs. dvector column_scale_; @@ -442,15 +414,12 @@ class GaussNewtonMinimizer { /// Scratch buffer for partial sums used by reduction kernels. dvector d_reduce_partials_; - /// Working normal-equation system; retained across Minimize calls to preserve - /// capacity. - CSRSparseMatrix lhs_work_; + /// Right-hand side; retained across Minimize calls to preserve capacity. dvector rhs_work_; MinimizerState current_state_; MinimizerState updated_state_; - profiler::Domain profiler_domain_{ - "GaussNewtonMinimizer"}; ///< Profiling domain. + profiler::Domain profiler_domain_{"GaussNewtonMinimizer"}; ///< Profiling domain. }; } // namespace cunls diff --git a/cunls/minimizer/hessian_structure.cu b/cunls/minimizer/hessian_structure.cu new file mode 100644 index 0000000..43fe3b7 --- /dev/null +++ b/cunls/minimizer/hessian_structure.cu @@ -0,0 +1,682 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "cunls/common/helper.h" +#include "cunls/minimizer/hessian_structure.h" +#include "cunls/minimizer/problem.h" +#include "cunls/state/state_batch_ops.h" + +namespace cunls { +namespace { + +constexpr int kBlockSize = 256; + +/// Sorts last, so candidate pairs touching a constant state land in a trailing +/// segment that the expansion simply ignores. +constexpr uint64_t kInvalidPairKey = ~uint64_t(0); + +int GridFor(size_t count) { return static_cast((count + kBlockSize - 1) / kBlockSize); } + +/** + * One thread per (factor, block). Resolves the factor's state pointer to a + * global column offset by testing it against one state batch's storage range, + * exactly like `col_ids_kernel` does for the triplet path. Threads whose + * pointer belongs to a different batch leave the entry untouched, so the caller + * runs this once per state batch over a `-1`-initialized output. + */ +__global__ void ResolveFactorColumnsKernel(int num_entries, int num_blocks_per_factor, + float const *const *__restrict__ state_pointers, + const int *__restrict__ block_sizes, + const float *__restrict__ batch_base, int ambient_dim, + int tangent_dim, int num_state_blocks, + const int *__restrict__ block_col_map, + int *__restrict__ out_cols) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_entries) { + return; + } + + // A factor's block slot only maps onto this batch if the tangent dims agree; + // mirrors the guard in FillColIdsInJacobianBlock. + int block = idx % num_blocks_per_factor; + if (block_sizes[block] != tangent_dim) { + return; + } + + const float *ptr = state_pointers[idx]; + ptrdiff_t diff = ptr - batch_base; + if (diff < 0 || diff >= static_cast(num_state_blocks) * ambient_dim) { + return; + } + + int block_index = static_cast(diff / ambient_dim); + int col = block_col_map[block_index]; + if (col >= 0) { + out_cols[idx] = col; + } +} + +/** + * One thread per (state block, tangent component). Records the owning block's + * tangent size at every column it spans, so a block pair can look up its tile + * dimensions from the two column indices alone. + */ +__global__ void FillTangentAtColKernel(int num_entries, int tangent_dim, + const int *__restrict__ block_col_map, + int *__restrict__ tangent_at_col) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_entries) { + return; + } + int block = idx / tangent_dim; + int component = idx - block * tangent_dim; + int col = block_col_map[block]; + if (col >= 0) { + tangent_at_col[col + component] = tangent_dim; + } +} + +/** One thread per candidate (factor, block_a, block_b). */ +__global__ void BuildPairKeysKernel(int num_entries, int num_blocks, + const int *__restrict__ factor_cols, + uint64_t *__restrict__ keys) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_entries) { + return; + } + + int pairs_per_factor = num_blocks * num_blocks; + int factor = idx / pairs_per_factor; + int pair = idx - factor * pairs_per_factor; + int block_a = pair / num_blocks; + int block_b = pair - block_a * num_blocks; + + int col_a = factor_cols[factor * num_blocks + block_a]; + int col_b = factor_cols[factor * num_blocks + block_b]; + keys[idx] = (col_a < 0 || col_b < 0) + ? kInvalidPairKey + : ((static_cast(col_a) << 32) | static_cast(col_b)); +} + +/** Marks the first candidate of each run of equal keys. */ +__global__ void MarkGroupStartsKernel(int num_valid, const uint64_t *__restrict__ sorted_keys, + int *__restrict__ flags) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_valid) { + return; + } + flags[idx] = (idx == 0 || sorted_keys[idx] != sorted_keys[idx - 1]) ? 1 : 0; +} + +/** Turns an inclusive scan of the group-start flags into 0-based group ids. */ +__global__ void DecrementKernel(int count, int *__restrict__ values) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < count) { + values[idx] -= 1; + } +} + +/** + * One thread per sorted candidate; the group's representative writes out the + * pair's row/column and tile dimensions. + */ +__global__ void GatherUniquePairsKernel(int num_valid, const uint64_t *__restrict__ sorted_keys, + const int *__restrict__ flags, + const int *__restrict__ group_id, + const int *__restrict__ tangent_at_col, + int *__restrict__ pair_row, int *__restrict__ pair_col, + int *__restrict__ pair_row_tangent, + int *__restrict__ pair_col_tangent) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_valid || flags[idx] == 0) { + return; + } + uint64_t key = sorted_keys[idx]; + int row = static_cast(key >> 32); + int col = static_cast(key & 0xFFFFFFFFu); + int g = group_id[idx]; + pair_row[g] = row; + pair_col[g] = col; + pair_row_tangent[g] = tangent_at_col[row]; + pair_col_tangent[g] = tangent_at_col[col]; +} + +/** + * One thread per block pair. Atomically accumulates per-row non-zero counts + * for every row the pair spans. + */ +__global__ void ComputeBlockRowCountsKernel(int num_pairs, const int *__restrict__ pair_row, + const int *__restrict__ row_tangent, + const int *__restrict__ col_tangent, + int *__restrict__ row_counts) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_pairs) { + return; + } + int row = pair_row[idx]; + int tiles = row_tangent[idx]; + int width = col_tangent[idx]; + for (int i = 0; i < tiles; i++) { + atomicAdd(&row_counts[row + i], width); + } +} + +/** + * One thread block per block pair. Writes the dense column indices of the tile + * into the CSR col_ids array. Every row of the block row shares the same + * `write_offset`, which is what makes the assembler's scatter map so small. + */ +__global__ void ExpandBlockPairsKernel(int num_pairs, const int *__restrict__ pair_row, + const int *__restrict__ pair_col, + const int *__restrict__ row_tangent, + const int *__restrict__ col_tangent, + const int *__restrict__ write_offset, + const int *__restrict__ row_offsets, + int *__restrict__ col_ids) { + int pair = blockIdx.x; + if (pair >= num_pairs) { + return; + } + int row = pair_row[pair]; + int col = pair_col[pair]; + int height = row_tangent[pair]; + int width = col_tangent[pair]; + int offset = write_offset[pair]; + int total = height * width; + for (int k = threadIdx.x; k < total; k += blockDim.x) { + int row_in_block = k / width; + int col_in_block = k - row_in_block * width; + col_ids[row_offsets[row + row_in_block] + offset + col_in_block] = col + col_in_block; + } +} + +/** out[i] = in[i] / divisor. */ +__global__ void DivideKernel(int count, int divisor, const int *__restrict__ in, + int *__restrict__ out) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < count) { + out[idx] = in[idx] / divisor; + } +} + +/** + * One thread per block pair. Accumulates per-block-row tile counts for every + * block row the pair spans. + */ +__global__ void ComputeBlockRowTileCountsKernel(int num_pairs, int block_size, + const int *__restrict__ pair_row, + const int *__restrict__ row_tangent, + const int *__restrict__ tiles_per_pair, + int *__restrict__ block_row_counts) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_pairs) { + return; + } + int block_row = pair_row[idx] / block_size; + int rows = row_tangent[idx] / block_size; + int width = tiles_per_pair[idx]; + for (int i = 0; i < rows; i++) { + atomicAdd(&block_row_counts[block_row + i], width); + } +} + +/** + * One thread block per block pair. Writes the tile column indices of the pair + * into the BSR col_ids array. As in the scalar case, every block row of the + * pair shares the same tile-valued `write_offset`. + */ +__global__ void ExpandBlockPairsBSRKernel( + int num_pairs, int block_size, const int *__restrict__ pair_row, + const int *__restrict__ pair_col, const int *__restrict__ row_tangent, + const int *__restrict__ col_tangent, const int *__restrict__ write_offset, + const int *__restrict__ row_offsets, int *__restrict__ col_ids) { + int pair = blockIdx.x; + if (pair >= num_pairs) { + return; + } + const int block_row = pair_row[pair] / block_size; + const int block_col = pair_col[pair] / block_size; + const int rows = row_tangent[pair] / block_size; + const int cols = col_tangent[pair] / block_size; + const int offset = write_offset[pair]; + const int total = rows * cols; + for (int k = threadIdx.x; k < total; k += blockDim.x) { + int tile_row = k / cols; + int tile_col = k - tile_row * cols; + col_ids[row_offsets[block_row + tile_row] + offset + tile_col] = block_col + tile_col; + } +} + +/** + * Pushes each group's row-relative write offset back to the candidate slot it + * came from, giving the assembler `write_offsets[f * nb * nb + a * nb + b]` + * with no search. + */ +__global__ void ScatterWriteOffsetsKernel(int num_candidates, int num_valid, + const int *__restrict__ pair_order, + const int *__restrict__ group_id, + const int *__restrict__ group_write_offset, + int *__restrict__ write_offsets) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_candidates) { + return; + } + // Candidates at or past num_valid touch a constant state; the sentinel key + // sorted them into the trailing segment. + int value = idx < num_valid ? group_write_offset[group_id[idx]] : -1; + write_offsets[pair_order[idx]] = value; +} + +} // namespace + +void HessianStructureBuilder::BuildLayout(const Problem &problem) { + const auto &residual_batches = problem.GetResidualBatches(); + layout_.assign(residual_batches.size(), HessianBatchLayout()); + + size_t jacobian_cursor = 0; + size_t residual_cursor = 0; + size_t col_cursor = 0; + size_t pair_cursor = 0; + + for (size_t i = 0; i < residual_batches.size(); i++) { + const auto *factor_batch = residual_batches[i].GetFactorBatch(); + auto block_sizes = factor_batch->StateBlockSizes(); + + HessianBatchLayout &layout = layout_[i]; + layout.num_factors = static_cast(factor_batch->NumFactors()); + layout.residual_dim = static_cast(factor_batch->ResidualsSize()); + layout.num_blocks = static_cast(block_sizes.size()); + layout.tangent_dim = + static_cast(std::accumulate(block_sizes.begin(), block_sizes.end(), size_t(0))); + layout.jacobian_offset = jacobian_cursor; + layout.residual_offset = residual_cursor; + layout.col_offset = col_cursor; + layout.pair_offset = pair_cursor; + + const size_t nf = layout.num_factors; + jacobian_cursor += nf * layout.residual_dim * layout.tangent_dim; + residual_cursor += nf * layout.residual_dim; + col_cursor += nf * layout.num_blocks; + pair_cursor += nf * layout.num_blocks * layout.num_blocks; + } + + jacobian_values_size_ = jacobian_cursor; + total_pairs_ = pair_cursor; + factor_cols_.resize(col_cursor); +} + +void HessianStructureBuilder::ResolveFactorColumns(cudaStream_t stream, const Problem &problem, + int num_cols, dvector &tangent_at_col) { + const auto &residual_batches = problem.GetResidualBatches(); + const auto &host_state_pointers = problem.GetStatePointers(); + + if (factor_cols_.empty()) { + return; + } + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(factor_cols_.data(), 0xFF, factor_cols_.size() * sizeof(int), stream)); + + tangent_at_col.resize(static_cast(num_cols)); + if (num_cols > 0) { + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(tangent_at_col.data(), 0, tangent_at_col.size() * sizeof(int), stream)); + } + + // Staging buffers live only for this call. State pointers are uploaded once + // and reused across the state-batch loop below. + dvector state_pointers(factor_cols_.size()); + std::vector> block_sizes_device(residual_batches.size()); + for (size_t i = 0; i < residual_batches.size(); i++) { + auto block_sizes = residual_batches[i].GetFactorBatch()->StateBlockSizes(); + std::vector sizes_host(block_sizes.begin(), block_sizes.end()); + block_sizes_device[i].resize(sizes_host.size()); + block_sizes_device[i].CopyFromHost(sizes_host.data(), sizes_host.size()); + + const size_t count = static_cast(layout_[i].num_factors) * layout_[i].num_blocks; + if (count > 0) { + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(state_pointers.data() + layout_[i].col_offset, + host_state_pointers[i].data(), count * sizeof(float *), + cudaMemcpyHostToDevice, stream)); + } + } + + // A (factor, block) slot is owned by exactly one state batch, so each batch + // overwrites only its own entries of the -1-initialized output. + dvector block_col_map; + int last_col = 0; + for (auto *state_batch : problem.GetStateBatches()) { + ComputeStateBlockColumnOffsets(stream, last_col, state_batch, block_col_map); + + const size_t tangent_entries = state_batch->NumStateBlocks() * state_batch->TangentSize(); + if (tangent_entries > 0) { + FillTangentAtColKernel<<>>( + static_cast(tangent_entries), static_cast(state_batch->TangentSize()), + block_col_map.data(), tangent_at_col.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + for (size_t i = 0; i < residual_batches.size(); i++) { + const size_t count = static_cast(layout_[i].num_factors) * layout_[i].num_blocks; + if (count == 0) { + continue; + } + ResolveFactorColumnsKernel<<>>( + static_cast(count), layout_[i].num_blocks, + state_pointers.data() + layout_[i].col_offset, block_sizes_device[i].data(), + state_batch->StateBlockDevicePtr(0), static_cast(state_batch->AmbientSize()), + static_cast(state_batch->TangentSize()), + static_cast(state_batch->NumStateBlocks()), block_col_map.data(), + factor_cols_.data() + layout_[i].col_offset); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + last_col += + static_cast((state_batch->NumStateBlocks() - state_batch->NumConstStateBlocks()) * + state_batch->TangentSize()); + } + + // state_pointers goes out of scope here; the kernels above are ordered behind + // its upload on `stream`, so wait before the allocation is released. + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); +} + +int HessianStructureBuilder::DiscoverBlockPairs(cudaStream_t stream, const Problem &problem, + int num_cols, dvector &pair_fields, + dvector &pair_order, dvector &pair_group, + size_t &num_valid) { + auto stream_policy = thrust::cuda::par_nosync.on(stream); + + BuildLayout(problem); + // Every buffer below is scratch for this call only. Holding them as members + // would keep ~20 bytes per candidate pair alive for the whole solve, which on + // a large SBA problem is hundreds of MiB against no reuse benefit. + dvector tangent_at_col; + ResolveFactorColumns(stream, problem, num_cols, tangent_at_col); + + num_valid = 0; + if (total_pairs_ == 0 || num_cols == 0) { + return 0; + } + + // ---- Candidate keys ----------------------------------------------------- + dvector pair_keys(total_pairs_); + pair_order.resize(total_pairs_); + for (const HessianBatchLayout &layout : layout_) { + const size_t count = + static_cast(layout.num_factors) * layout.num_blocks * layout.num_blocks; + if (count == 0) { + continue; + } + BuildPairKeysKernel<<>>( + static_cast(count), layout.num_blocks, factor_cols_.data() + layout.col_offset, + pair_keys.data() + layout.pair_offset); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + // ---- Sort and segment --------------------------------------------------- + thrust::device_ptr keys(pair_keys.data()); + thrust::device_ptr order(pair_order.data()); + thrust::sequence(stream_policy, order, order + total_pairs_); + thrust::sort_by_key(stream_policy, keys, keys + total_pairs_, order); + + // Constant-state candidates carry the sentinel key and sorted to the tail. + size_t num_invalid = + static_cast(thrust::count(stream_policy, keys, keys + total_pairs_, kInvalidPairKey)); + num_valid = total_pairs_ - num_invalid; + + int num_pairs = 0; + dvector group_flags(num_valid); + pair_group.resize(num_valid); + if (num_valid > 0) { + MarkGroupStartsKernel<<>>( + static_cast(num_valid), pair_keys.data(), group_flags.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + // Inclusive, not exclusive. An exclusive scan yields the correct id only + // for elements that *start* a group and hands every duplicate the id of the + // next group -- and duplicates are the whole point here, since every block + // pair shared by several factors must resolve to one group. Scan + // inclusively (so the last element holds the group count) and shift down. + thrust::device_ptr flags(group_flags.data()); + thrust::device_ptr groups(pair_group.data()); + thrust::inclusive_scan(stream_policy, flags, flags + num_valid, groups); + + if (pinned_buf_.size() < 2) { + pinned_buf_.resize(2); + } + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(pinned_buf_.data(), pair_group.data() + num_valid - 1, + sizeof(int), cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + num_pairs = pinned_buf_[0]; + + DecrementKernel<<>>(static_cast(num_valid), + pair_group.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + // ---- Per-pair descriptors ---------------------------------------------- + // Layout: [row | col | row_tangent | col_tangent | write_offset] + const size_t stride = static_cast(num_pairs); + pair_fields.resize(5 * stride); + if (num_pairs > 0) { + int *pair_row = pair_fields.data(); + int *pair_col = pair_row + stride; + int *row_tangent = pair_col + stride; + int *col_tangent = row_tangent + stride; + + GatherUniquePairsKernel<<>>( + static_cast(num_valid), pair_keys.data(), group_flags.data(), pair_group.data(), + tangent_at_col.data(), pair_row, pair_col, row_tangent, col_tangent); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + return num_pairs; +} + +void HessianStructureBuilder::Build(cudaStream_t stream, const Problem &problem, int num_cols, + CSRSparseMatrix &output, bool want_scatter_maps) { + auto stream_policy = thrust::cuda::par_nosync.on(stream); + + dvector pair_fields; + dvector pair_order; + dvector pair_group; + size_t num_valid = 0; + const int num_pairs = + DiscoverBlockPairs(stream, problem, num_cols, pair_fields, pair_order, pair_group, num_valid); + + output.row_offsets.resize(static_cast(num_cols) + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(output.row_offsets.data(), 0, + output.row_offsets.size() * sizeof(int), stream)); + write_offsets_.resize(want_scatter_maps ? total_pairs_ : 0); + + if (total_pairs_ == 0 || num_cols == 0) { + output.col_ids.resize(0); + output.values.resize(0); + return; + } + + const size_t stride = static_cast(num_pairs); + int *pair_row = pair_fields.data(); + int *pair_col = pair_row + stride; + int *row_tangent = pair_col + stride; + int *col_tangent = row_tangent + stride; + int *write_offset = col_tangent + stride; + + dvector row_counts(static_cast(num_cols)); + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(row_counts.data(), 0, static_cast(num_cols) * sizeof(int), stream)); + + if (num_pairs > 0) { + // Pairs are sorted by (row, col), so a segmented scan over the block row + // gives each tile's offset within the row. + thrust::exclusive_scan_by_key(stream_policy, thrust::device_pointer_cast(pair_row), + thrust::device_pointer_cast(pair_row) + num_pairs, + thrust::device_pointer_cast(col_tangent), + thrust::device_pointer_cast(write_offset)); + + ComputeBlockRowCountsKernel<<>>( + num_pairs, pair_row, row_tangent, col_tangent, row_counts.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + { + thrust::device_ptr counts(row_counts.data()); + thrust::device_ptr offsets(output.row_offsets.data()); + thrust::inclusive_scan(stream_policy, counts, counts + num_cols, offsets + 1); + } + + if (pinned_buf_.size() < 2) { + pinned_buf_.resize(2); + } + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(pinned_buf_.data(), output.row_offsets.data() + num_cols, + sizeof(int), cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + const int total_nnz = pinned_buf_[0]; + + output.col_ids.resize(static_cast(total_nnz)); + output.values.resize(static_cast(total_nnz)); + + if (num_pairs > 0) { + ExpandBlockPairsKernel<<>>( + num_pairs, pair_row, pair_col, row_tangent, col_tangent, write_offset, + output.row_offsets.data(), output.col_ids.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + if (want_scatter_maps) { + ScatterWriteOffsetsKernel<<>>( + static_cast(total_pairs_), static_cast(num_valid), pair_order.data(), + pair_group.data(), write_offset, write_offsets_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + // The scratch buffers go out of scope; the kernels above read them. + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); +} + +void HessianStructureBuilder::Build(cudaStream_t stream, const Problem &problem, int num_cols, + int block_size, BSRSparseMatrix &output, + bool want_scatter_maps) { + auto stream_policy = thrust::cuda::par_nosync.on(stream); + + if (block_size < 1 || num_cols % block_size != 0) { + throw std::runtime_error( + "HessianStructureBuilder: block size must divide the tangent dimension"); + } + + dvector pair_fields; + dvector pair_order; + dvector pair_group; + size_t num_valid = 0; + const int num_pairs = + DiscoverBlockPairs(stream, problem, num_cols, pair_fields, pair_order, pair_group, num_valid); + + const int num_block_rows = num_cols / block_size; + output.block_size = block_size; + output.num_block_rows = num_block_rows; + output.max_tiles_per_row = 0; + output.row_offsets.resize(static_cast(num_block_rows) + 1); + THROW_ON_CUDA_ERROR(cudaMemsetAsync(output.row_offsets.data(), 0, + output.row_offsets.size() * sizeof(int), stream)); + write_offsets_.resize(want_scatter_maps ? total_pairs_ : 0); + + if (total_pairs_ == 0 || num_cols == 0) { + output.col_ids.resize(0); + output.values.resize(0); + return; + } + + const size_t stride = static_cast(num_pairs); + int *pair_row = pair_fields.data(); + int *pair_col = pair_row + stride; + int *row_tangent = pair_col + stride; + int *col_tangent = row_tangent + stride; + int *write_offset = col_tangent + stride; + + // Tile counts, not scalar column counts: a (row_tangent x col_tangent) pair + // occupies (row_tangent / b) x (col_tangent / b) tiles. + dvector tiles_per_pair(stride); + dvector block_row_counts(static_cast(num_block_rows)); + THROW_ON_CUDA_ERROR( + cudaMemsetAsync(block_row_counts.data(), 0, block_row_counts.size() * sizeof(int), stream)); + + if (num_pairs > 0) { + DivideKernel<<>>(num_pairs, block_size, col_tangent, + tiles_per_pair.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + thrust::exclusive_scan_by_key(stream_policy, thrust::device_pointer_cast(pair_row), + thrust::device_pointer_cast(pair_row) + num_pairs, + thrust::device_pointer_cast(tiles_per_pair.data()), + thrust::device_pointer_cast(write_offset)); + + ComputeBlockRowTileCountsKernel<<>>( + num_pairs, block_size, pair_row, row_tangent, tiles_per_pair.data(), + block_row_counts.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + if (num_block_rows > 0) { + thrust::device_ptr counts(block_row_counts.data()); + // Peak row length before the scan destroys it; the SpMV picks its schedule + // from this. + output.max_tiles_per_row = *thrust::max_element(stream_policy, counts, counts + num_block_rows); + thrust::device_ptr offsets(output.row_offsets.data()); + thrust::inclusive_scan(stream_policy, counts, counts + num_block_rows, offsets + 1); + } + + if (pinned_buf_.size() < 2) { + pinned_buf_.resize(2); + } + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(pinned_buf_.data(), + output.row_offsets.data() + num_block_rows, sizeof(int), + cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + const int total_tiles = pinned_buf_[0]; + + output.col_ids.resize(static_cast(total_tiles)); + output.values.resize(static_cast(total_tiles) * block_size * block_size); + + if (num_pairs > 0) { + ExpandBlockPairsBSRKernel<<>>( + num_pairs, block_size, pair_row, pair_col, row_tangent, col_tangent, write_offset, + output.row_offsets.data(), output.col_ids.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + if (want_scatter_maps) { + ScatterWriteOffsetsKernel<<>>( + static_cast(total_pairs_), static_cast(num_valid), pair_order.data(), + pair_group.data(), write_offset, write_offsets_.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + } + + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); +} + +} // namespace cunls diff --git a/cunls/minimizer/hessian_structure.h b/cunls/minimizer/hessian_structure.h new file mode 100644 index 0000000..4e36fee --- /dev/null +++ b/cunls/minimizer/hessian_structure.h @@ -0,0 +1,149 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include + +#include "cunls/common/types.h" + +namespace cunls { + +class Problem; + +/** @brief Per-residual-batch geometry and offsets into the flat scatter maps. */ +struct HessianBatchLayout { + int num_factors = 0; ///< Factors in the batch. + int residual_dim = 0; ///< m: residual dimension of one factor. + int tangent_dim = 0; ///< n: sum of the factor's state block sizes. + int num_blocks = 0; ///< nb: state blocks a factor touches. + size_t jacobian_offset = 0; ///< Offset into the flat Jacobian value buffer. + size_t residual_offset = 0; ///< Offset into the flat residual vector. + size_t col_offset = 0; ///< Offset into FactorCols(), stride nb. + size_t pair_offset = 0; ///< Offset into WriteOffsets(), stride nb*nb. +}; + +/** + * @brief Derives the J^T J sparsity pattern from factor-graph connectivity. + * + * Every factor connecting state blocks A,B contributes dense sub-blocks + * (A,A), (A,B), (B,A), (B,B) to the Hessian. The set of distinct block pairs + * is the block-level sparsity pattern; expanding it gives the scalar CSR. + * + * The whole derivation runs on the GPU: resolve each factor's state pointers to + * global columns, pack every candidate pair into a 64-bit key, then sort and + * segment. The previous host implementation spent most of a large SBA solve + * inside an `unordered_set` of 8 M keys plus a `std::sort` of 5 M pairs. + * + * The resulting CSR has a property `BlockHessianAssembler` depends on: within + * one block row, the columns of a given block pair are contiguous, and the + * offset of that run relative to the row start is identical for every row of + * the block. That offset is exactly what `WriteOffsets()` returns, so the + * assembler gets its scatter map as a by-product of the segmentation instead of + * binary-searching the expanded column-index array. + * + * `output.values` is resized but left uninitialized. + */ +class HessianStructureBuilder { +public: + /** + * @brief Builds the scalar CSR pattern, and optionally the scatter maps. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem. + * @param num_cols Free tangent dimensions in the reduced system. + * @param[out] output CSR matrix; row_offsets/col_ids filled, values sized. + * @param want_scatter_maps When false, WriteOffsets() is left empty and the + * final scatter pass is skipped (the J^T J path does not need it). + */ + void Build(cudaStream_t stream, const Problem &problem, int num_cols, CSRSparseMatrix &output, + bool want_scatter_maps); + + /** + * @brief Builds the same pattern in uniform block storage. + * + * Block-pair discovery is already block-level, so this is the *cheaper* of + * the two expansions: each state-block pair emits + * `(row_tangent / b) x (col_tangent / b)` tiles instead of + * `row_tangent * col_tangent` scalar column indices. + * + * WriteOffsets() is then measured in tiles rather than scalar columns. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem. + * @param num_cols Free tangent dimensions; must be a multiple of block_size. + * @param block_size Tile edge; see ChooseHessianBlockSize(). + * @param[out] output BSR matrix; structure filled, values sized. + * @param want_scatter_maps As for the CSR overload. + */ + void Build(cudaStream_t stream, const Problem &problem, int num_cols, int block_size, + BSRSparseMatrix &output, bool want_scatter_maps); + + /** @brief Per-residual-batch geometry, indexed as the problem's batches. */ + const std::vector &Layout() const { return layout_; } + + /** @brief Global column of each (factor, block); -1 when constant. */ + const dvector &FactorCols() const { return factor_cols_; } + + /** + * @brief Row-relative offset of each (factor, block_a, block_b) run. + * + * -1 when either block is a constant state. Empty unless Build() was called + * with `want_scatter_maps`. + */ + const dvector &WriteOffsets() const { return write_offsets_; } + + /** @brief Total floats needed for the per-factor Jacobian value buffer. */ + size_t JacobianValuesSize() const { return jacobian_values_size_; } + +private: + /** + * @brief Shared front half: layout, column resolution, key sort, segmentation. + * + * Leaves the deduplicated block pairs in @p pair_fields and returns their + * count; both expansions differ only in what they do with them. + */ + int DiscoverBlockPairs(cudaStream_t stream, const Problem &problem, int num_cols, + dvector &pair_fields, dvector &pair_order, + dvector &pair_group, size_t &num_valid); + + /** @brief Fills layout_ and the flat-buffer offsets from the problem. */ + void BuildLayout(const Problem &problem); + + /** + * @brief Resolves every (factor, block) slot to a global column. + * + * @param[out] tangent_at_col Tangent size of the block owning each column. + */ + void ResolveFactorColumns(cudaStream_t stream, const Problem &problem, int num_cols, + dvector &tangent_at_col); + + std::vector layout_; + dvector factor_cols_; + dvector write_offsets_; + + /// Reusable pinned staging buffer for D2H readbacks. + pvector pinned_buf_; + + size_t total_pairs_ = 0; + size_t jacobian_values_size_ = 0; +}; + +} // namespace cunls diff --git a/cunls/minimizer/jacobian_ops.cu b/cunls/minimizer/jacobian_ops.cu deleted file mode 100644 index ab98a41..0000000 --- a/cunls/minimizer/jacobian_ops.cu +++ /dev/null @@ -1,396 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "cunls/common/helper.h" -#include "cunls/minimizer/minimizer_state.h" - -namespace cunls { - -namespace detail { - -/** - * @brief Calculates the number of values in a Jacobian block for a cost - * function. - * - * Each factor batch produces a dense Jacobian block of size - * (num_factors * residual_dim) x (sum of state block sizes). - * - * @param factor_batch Pointer to the factor batch. - * @return Total number of values in the Jacobian block. - */ -size_t CalculateJacobianBlockSize(const FactorBatch *factor_batch) { - size_t num_factors = factor_batch->NumFactors(); - size_t num_rows = num_factors * factor_batch->ResidualsSize(); - - auto block_sizes = factor_batch->StateBlockSizes(); - size_t num_cols = std::accumulate(block_sizes.begin(), block_sizes.end(), 0); - - return num_rows * num_cols; -} - -/** - * @brief Calculates the total number of values in the full Jacobian. - * - * Sums the Jacobian block sizes across all factor batches. - * - * @param factor_batches Vector of factor batch pointers. - * @return Total number of Jacobian values (including potential zeros for - * constant params). - */ -size_t -CalculateJacobianNonZeros(const std::vector &factor_batches) { - size_t size = 0; - for (const auto batch : factor_batches) { - size += CalculateJacobianBlockSize(batch); - } - return size; -} - -/** - * @brief Holds metadata about a state batch for Jacobian column indexing. - * - * Stores the base device pointer, block dimensions, and a column-ID mapping - * array used by the col_ids_kernel to compute Jacobian column indices. - */ -struct StateBatchDescription { - /** - * @brief Constructs from a state batch and its column-ID mapping. - * - * @param pbatch Pointer to the state batch. - * @param p_col_mapping Device pointer to column-ID mapping array. - */ - StateBatchDescription(const StateBatch *pbatch, const int *p_col_mapping) - : ptr(pbatch->StateBlockDevicePtr(0)), tangent_dim(pbatch->TangentSize()), - ambient_dim(pbatch->AmbientSize()), - num_blocks(pbatch->NumStateBlocks()), col_ids(p_col_mapping) {} - const float *ptr; ///< Base device pointer to the first state block. - size_t tangent_dim; ///< Tangent (local) dimension of each state block. - size_t ambient_dim; ///< Ambient dimension of each state block. - size_t num_blocks; ///< Number of state blocks in the batch. - const int *col_ids; ///< Column-ID mapping: block index -> starting column in - ///< reduced system. -}; - -/** - * @brief Holds metadata about a Jacobian block for a single factor batch. - * - * Describes the dimensions of the dense Jacobian sub-block produced by one - * factor batch: num_rows x num_cols, where rows correspond to residuals - * and columns correspond to state components. - */ -struct JacobianBlockDescription { - /** - * @brief Constructs from a factor batch. - * - * @param factor_batch Pointer to the factor batch. - */ - JacobianBlockDescription(const FactorBatch *factor_batch) - : residual_dim(factor_batch->ResidualsSize()) { - size_t num_factors = factor_batch->NumFactors(); - num_rows = num_factors * factor_batch->ResidualsSize(); - - auto block_sizes = factor_batch->StateBlockSizes(); - num_cols = std::accumulate(block_sizes.begin(), block_sizes.end(), 0); - num_state_blocks_in_res = block_sizes.size(); - } - size_t num_rows; ///< Total number of rows (num_factors * residual_dim). - size_t num_cols; ///< Total number of columns (sum of state block sizes). - size_t residual_dim; ///< Dimension of a single residual vector. - size_t num_state_blocks_in_res; ///< Number of state blocks per residual. -}; - -/** - * @brief CUDA kernel to assign row indices for a Jacobian block. - * - * Each thread writes the row index for one element of the dense Jacobian - * block. The row index is offset by the number of rows from previously - * processed factor batches. - * - * @param[out] start Output row index array (dense Jacobian block, row-major). - * @param num_rows Number of rows in this Jacobian block. - * @param num_cols Number of columns in this Jacobian block. - * @param offset Row offset from previously processed factor batches. - * - * Grid: ((num_cols+7)/8, (num_rows+63)/64), Block: (8, 64). - */ -__global__ void row_ids_kernel(int *start, size_t num_rows, size_t num_cols, - int offset) { - int x = threadIdx.x + blockIdx.x * blockDim.x; - int y = threadIdx.y + blockIdx.y * blockDim.y; - - if (x >= num_cols || y >= num_rows) { - return; - } - - start[y * num_cols + x] = offset + y; -} - -/** - * @brief Fills the row index array for the full triplet Jacobian structure. - * - * Launches the row_ids_kernel for each factor batch, advancing - * the output pointer and row offset across batches. - * - * @param stream CUDA stream for GPU operations. - * @param factor_batches Vector of factor batch pointers. - * @param[out] row_ids Output device vector of row indices. - */ -void FillRowIds(cudaStream_t stream, - const std::vector &factor_batches, - dvector &row_ids) { - dim3 block(8, 64); - - int *ptr = row_ids.data(); - int offset = 0; - for (auto batch : factor_batches) { - JacobianBlockDescription j_desc(batch); - dim3 grid((j_desc.num_cols + block.x - 1) / block.x, - (j_desc.num_rows + block.y - 1) / block.y); - - row_ids_kernel<<>>(ptr, j_desc.num_rows, - j_desc.num_cols, offset); - - THROW_ON_CUDA_ERROR(cudaGetLastError()); - - ptr += j_desc.num_rows * j_desc.num_cols; - offset += j_desc.num_rows; - } -} - -/** - * @brief CUDA kernel to assign column indices for a Jacobian block. - * - * For each element in the dense Jacobian block, determines which state - * block the column belongs to by matching the state pointer against the - * state batch's base pointer. Constant state pointers receive col_id = -1. - * - * @param[out] start Output column index array (dense Jacobian block, - * row-major). - * @param state_pointers Device pointer array mapping factor instances to - * state_pointers. - * @param col_block_index Index of the current state block within the residual. - * @param j_desc Jacobian block dimensions and metadata. - * @param pb_desc State batch metadata with column-ID mapping. - * - * Grid: ((tangent_dim+7)/8, (num_rows+63)/64), Block: (8, 64). Max 512 threads. - */ -__launch_bounds__(512) __global__ - void col_ids_kernel(int *__restrict__ start, - float const *const *__restrict__ state_pointers, - int col_block_index, JacobianBlockDescription j_desc, - StateBatchDescription pb_desc) { - int x = threadIdx.x + blockIdx.x * blockDim.x; - int y = threadIdx.y + blockIdx.y * blockDim.y; - - assert(col_block_index < j_desc.num_state_blocks_in_res); - - if (x >= pb_desc.tangent_dim || y >= j_desc.num_rows) { - return; - } - - int factor_id = y / j_desc.residual_dim; - - float const *ptr = state_pointers[factor_id * j_desc.num_state_blocks_in_res + - col_block_index]; - - int pidx = (ptr - pb_desc.ptr) / pb_desc.ambient_dim; - - bool valid = pidx >= 0 && pidx < pb_desc.num_blocks; - int param_col_id = valid ? pb_desc.col_ids[pidx] : -1; - - if (param_col_id != -1) { - start[y * j_desc.num_cols + x] = param_col_id + x; - } -} - -/** - * @brief Creates a mapping from state block indices to column IDs in the - * reduced system. - * - * For each state block in the batch, assigns a starting column index in - * the reduced (non-constant) system. Constant state_pointers are marked with - * -1. Uses exclusive scan to compute cumulative column offsets. - * - * @param stream CUDA stream for GPU operations. - * @param offset Starting column index for this state batch. - * @param state_batch The state batch. - * @param[out] col_ids Device vector of column IDs (one per state block). - */ -void CreateStateColIdMapping(cudaStream_t stream, int offset, - const StateBatch *state_batch, - thrust::device_vector &col_ids) { - const int *cparam_ids = state_batch->ConstStateIds(); - size_t num_const_params = state_batch->NumConstStateBlocks(); - size_t tangent_dim = state_batch->TangentSize(); - auto stream_policy = thrust::cuda::par_nosync.on(stream); - - col_ids.resize(state_batch->NumStateBlocks()); - thrust::fill(stream_policy, col_ids.begin(), col_ids.end(), - static_cast(tangent_dim)); - - if (cparam_ids != nullptr && num_const_params > 0) { - auto zero_it = thrust::make_constant_iterator(0); - thrust::device_ptr cparam_ids_ptr(cparam_ids); - thrust::scatter(stream_policy, zero_it, zero_it + num_const_params, - cparam_ids_ptr, col_ids.begin()); - thrust::exclusive_scan(stream_policy, col_ids.begin(), col_ids.end(), - col_ids.begin(), offset); - auto minus_one_it = thrust::make_constant_iterator(-1); - thrust::scatter(stream_policy, minus_one_it, - minus_one_it + num_const_params, cparam_ids_ptr, - col_ids.begin()); - } else { - thrust::exclusive_scan(stream_policy, col_ids.begin(), col_ids.end(), - col_ids.begin(), offset); - } -} - -/** - * @brief Fills column indices for a Jacobian block corresponding to one factor - * batch. - * - * Launches the col_ids_kernel for each state block referenced by the cost - * function. Skips state blocks whose tangent dimension doesn't match the - * expected block size. - * - * @param stream CUDA stream for GPU operations. - * @param[out] start Pointer into the column index array for this Jacobian - * block. - * @param factor_batch The factor batch. - * @param param_ptr Device state pointers for this factor batch. - * @param pb_desc State batch description with column-ID mapping. - */ -void FillColIdsInJacobianBlock(cudaStream_t stream, int *start, - const FactorBatch *factor_batch, - const DeviceVector ¶m_ptr, - const StateBatchDescription &pb_desc) { - auto block_sizes = factor_batch->StateBlockSizes(); - - float const *const *state_pointers = param_ptr.data(); - - JacobianBlockDescription j_desc(factor_batch); - - const dim3 block(8, 64); - const dim3 grid((pb_desc.tangent_dim + block.x - 1) / block.x, - (j_desc.num_rows + block.y - 1) / block.y); - - int *col_ptr = start; - for (size_t i = 0; i < block_sizes.size(); i++) { - if (pb_desc.tangent_dim != block_sizes[i]) { - // Mismatch in tangent dims, skip this col block - col_ptr += block_sizes[i]; - continue; - } - - col_ids_kernel<<>>(col_ptr, state_pointers, i, - j_desc, pb_desc); - - THROW_ON_CUDA_ERROR(cudaGetLastError()); - - col_ptr += block_sizes[i]; - } -} - -/** - * @brief Fills the column index array for the full triplet Jacobian structure. - * - * Iterates over state batches and factor batches to assign column - * indices in the reduced linear system. Initializes all column IDs to -1 - * (invalid), then populates valid entries via FillColIdsInJacobianBlock. - * - * @param stream CUDA stream for GPU operations. - * @param factor_batches Vector of factor batch pointers. - * @param state_pointers Device state pointers per factor batch. - * @param state_batches Vector of state batch pointers. - * @param[out] col_ids Output device vector of column indices. - */ -void FillColIds(cudaStream_t stream, - const std::vector &factor_batches, - const std::vector> &state_pointers, - const std::vector &state_batches, - dvector &col_ids) { - auto stream_policy = thrust::cuda::par_nosync.on(stream); - - // Set all the column indices to be invalid, i.e. -1 - thrust::device_ptr col_ids_ptr(col_ids.data()); - thrust::fill(stream_policy, col_ids_ptr, col_ids_ptr + col_ids.size(), -1); - - // Buffer to store the mapping between param pointers and - // column ids - thrust::device_vector pbatch_cols; - int last_col_id = 0; - for (auto pbatch : state_batches) { - // Create a mapping. Account for the number of already processed - // state_pointers. - CreateStateColIdMapping(stream, last_col_id, pbatch, pbatch_cols); - - StateBatchDescription pb_desc(pbatch, - thrust::raw_pointer_cast(pbatch_cols.data())); - - // A pointer to the start of the jacobian block for each factor batch. - int *start = col_ids.data(); - for (size_t j = 0; j < factor_batches.size(); j++) { - auto batch = factor_batches[j]; - - // Computes column indices for the jacobian block for this factor batch. - FillColIdsInJacobianBlock(stream, start, batch, state_pointers[j], - pb_desc); - - // Update the pointer s.t in points to the next jacobian block - start += CalculateJacobianBlockSize(batch); - } - - // Update the last column index - last_col_id += (pbatch->NumStateBlocks() - pbatch->NumConstStateBlocks()) * - pbatch->TangentSize(); - } -} - -} // namespace detail - -void MinimizerState::BuildTripletSparseStructure( - cudaStream_t stream, const Problem &problem, - TripletSparseStructure &structure) { - CopyProblemStatePointersFromHost(problem); - - std::vector factor_batches; - factor_batches.reserve(problem.GetResidualBatches().size()); - - for (const auto &rb : problem.GetResidualBatches()) { - factor_batches.push_back(rb.GetFactorBatch()); - } - - size_t jacobian_size = detail::CalculateJacobianNonZeros(factor_batches); - structure.row_ids.resize(jacobian_size); - structure.col_ids.resize(jacobian_size); - - detail::FillRowIds(stream, factor_batches, structure.row_ids); - detail::FillColIds(stream, factor_batches, problem_state_ptrs_device_, - problem.GetStateBatches(), structure.col_ids); -} -} // namespace cunls diff --git a/cunls/minimizer/levenberg_marquardt_minimizer.cpp b/cunls/minimizer/levenberg_marquardt_minimizer.cpp index 645a5e7..9135c59 100644 --- a/cunls/minimizer/levenberg_marquardt_minimizer.cpp +++ b/cunls/minimizer/levenberg_marquardt_minimizer.cpp @@ -16,6 +16,7 @@ */ #include "cunls/minimizer/levenberg_marquardt_minimizer.h" + #include "cunls/common/helper.h" #include "cunls/common/log.h" #include "cunls/common/types.h" @@ -38,16 +39,12 @@ namespace cunls { * @param stream CUDA stream for GPU operations. * @param problem The optimization problem. * @param minimizer_state Current minimizer state. - * @param[out] lhs Output left-hand side matrix (H + lambda * diag(H)). - * @param[out] rhs Output right-hand side vector (-J^T r). */ -void LevenbergMarquardtMinimizer::BuildSystem( - cudaStream_t stream, const Problem &problem, - const MinimizerState &minimizer_state, CSRSparseMatrix &lhs, - dvector &rhs) { - GaussNewtonMinimizer::BuildSystem(stream, problem, minimizer_state, lhs, rhs); - ExtractDiagonal(stream, lhs, diagonal_); - AddScaledDiagonal(stream, lambda_, diagonal_, lhs, lhs); +void LevenbergMarquardtMinimizer::BuildSystem(cudaStream_t stream, const Problem &problem, + const MinimizerState &minimizer_state) { + GaussNewtonMinimizer::BuildSystem(stream, problem, minimizer_state); + normal_equations_.ExtractLhsDiagonal(stream, diagonal_); + normal_equations_.AddScaledDiagonalToLhs(stream, lambda_, diagonal_); } /** @@ -70,36 +67,27 @@ void LevenbergMarquardtMinimizer::BuildSystem( * @param[out] step_quality Output rho metric (actual/predicted reduction). * @return True if converged, false otherwise. */ -bool LevenbergMarquardtMinimizer::CheckConvergence(cudaStream_t stream, - float updated_cost, - float current_cost, - const dvector &step, +bool LevenbergMarquardtMinimizer::CheckConvergence(cudaStream_t stream, float updated_cost, + float current_cost, const dvector &step, float &step_quality) { constexpr size_t kSlots = 3; - if (d_scalars_.size() < kSlots) - d_scalars_.resize(kSlots); - if (h_scalars_.size() < kSlots) - h_scalars_.resize(kSlots); + if (d_scalars_.size() < kSlots) d_scalars_.resize(kSlots); + if (h_scalars_.size() < kSlots) h_scalars_.resize(kSlots); size_t partials_needed = ReducePartialCount(step_.size()); if (d_reduce_partials_.size() < partials_needed) { d_reduce_partials_.resize(partials_needed); } // Enqueue all three reductions async - ComputeSquaredStepAsync(stream, step_, d_scalars_.data(), - d_reduce_partials_.data()); - ComputeWeightedSquaredStepAsync(stream, diagonal_, step_, - d_scalars_.data() + 1, + ComputeSquaredStepAsync(stream, step_, d_scalars_.data(), d_reduce_partials_.data()); + ComputeWeightedSquaredStepAsync(stream, diagonal_, step_, d_scalars_.data() + 1, d_reduce_partials_.data()); - auto handle = cusparse_handle_.GetHandle(stream); - ComputeWeightedSquaredStepAsync( - stream, handle, hessian_, hessian_dims_.num_rows, hessian_dims_.num_cols, - hessian_dims_.num_nonzeros, step_, buffer_, d_scalars_.data() + 2, - d_reduce_partials_.data()); + normal_equations_.WeightedSquaredStepAsync(stream, cusparse_handle_.GetHandle(stream), step_, + d_scalars_.data() + 2, d_reduce_partials_.data(), + buffer_); // Single D2H + single sync - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(h_scalars_.data(), d_scalars_.data(), - kSlots * sizeof(float), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(h_scalars_.data(), d_scalars_.data(), kSlots * sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); @@ -107,14 +95,12 @@ bool LevenbergMarquardtMinimizer::CheckConvergence(cudaStream_t stream, float diag_weight = h_scalars_[1]; float matrix_weight = h_scalars_[2]; - float predicted_relative_reduction = - (matrix_weight + 2.f * lambda_ * diag_weight) / current_cost; + float predicted_relative_reduction = (matrix_weight + 2.f * lambda_ * diag_weight) / current_cost; LogMessage("Predicted relative reduction = {}", predicted_relative_reduction); LogMessage("Step squared norm = {}", step_sq_norm); - float rho = - (1.f - updated_cost / current_cost) / predicted_relative_reduction; + float rho = (1.f - updated_cost / current_cost) / predicted_relative_reduction; step_quality = rho; @@ -128,15 +114,12 @@ bool LevenbergMarquardtMinimizer::CheckConvergence(cudaStream_t stream, } bool LevenbergMarquardtMinimizer::EvaluateAndCheckConvergence( - cudaStream_t stream, const Problem &problem, - const MinimizerState &updated_state, float current_cost, - const dvector &step, float &updated_cost, float &step_quality) { + cudaStream_t stream, const Problem &problem, const MinimizerState &updated_state, + float current_cost, const dvector &step, float &updated_cost, float &step_quality) { // 4 slots: [0]=cost, [1]=step_sq_norm, [2]=diag_weight, [3]=matrix_weight constexpr size_t kSlots = 4; - if (d_scalars_.size() < kSlots) - d_scalars_.resize(kSlots); - if (h_scalars_.size() < kSlots) - h_scalars_.resize(kSlots); + if (d_scalars_.size() < kSlots) d_scalars_.resize(kSlots); + if (h_scalars_.size() < kSlots) h_scalars_.resize(kSlots); // Enqueue cost reduction ComputeCostAsync(stream, problem, updated_state, d_scalars_.data()); @@ -147,24 +130,19 @@ bool LevenbergMarquardtMinimizer::EvaluateAndCheckConvergence( } // Enqueue squared step norm - ComputeSquaredStepAsync(stream, step_, d_scalars_.data() + 1, - d_reduce_partials_.data()); + ComputeSquaredStepAsync(stream, step_, d_scalars_.data() + 1, d_reduce_partials_.data()); // Enqueue diag-weighted step norm - ComputeWeightedSquaredStepAsync(stream, diagonal_, step_, - d_scalars_.data() + 2, + ComputeWeightedSquaredStepAsync(stream, diagonal_, step_, d_scalars_.data() + 2, d_reduce_partials_.data()); // Enqueue sparse-weighted step norm (SpMV + dot, all on stream) - auto handle = cusparse_handle_.GetHandle(stream); - ComputeWeightedSquaredStepAsync( - stream, handle, hessian_, hessian_dims_.num_rows, hessian_dims_.num_cols, - hessian_dims_.num_nonzeros, step_, buffer_, d_scalars_.data() + 3, - d_reduce_partials_.data()); + normal_equations_.WeightedSquaredStepAsync(stream, cusparse_handle_.GetHandle(stream), step_, + d_scalars_.data() + 3, d_reduce_partials_.data(), + buffer_); // Single D2H + single sync for all 4 scalars - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(h_scalars_.data(), d_scalars_.data(), - kSlots * sizeof(float), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(h_scalars_.data(), d_scalars_.data(), kSlots * sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); @@ -173,20 +151,17 @@ bool LevenbergMarquardtMinimizer::EvaluateAndCheckConvergence( float diag_weight = h_scalars_[2]; float matrix_weight = h_scalars_[3]; - float predicted_relative_reduction = - (matrix_weight + 2.f * lambda_ * diag_weight) / current_cost; + float predicted_relative_reduction = (matrix_weight + 2.f * lambda_ * diag_weight) / current_cost; LogMessage("Predicted relative reduction = {}", predicted_relative_reduction); LogMessage("Step squared norm = {}", step_sq_norm); - float rho = - (1.f - updated_cost / current_cost) / predicted_relative_reduction; + float rho = (1.f - updated_cost / current_cost) / predicted_relative_reduction; step_quality = rho; return (step_sq_norm < options_.base_options.state_tolerance || - predicted_relative_reduction < - options_.relative_reduction_tolerance || + predicted_relative_reduction < options_.relative_reduction_tolerance || updated_cost < options_.base_options.cost_tolerance); } @@ -246,9 +221,8 @@ bool LevenbergMarquardtMinimizer::AcceptStep(float step_quality) { * @param stream CUDA stream for GPU operations. * @param problem The optimization problem to initialize for. */ -void LevenbergMarquardtMinimizer::Initialize(cudaStream_t stream, - Problem &problem) { +void LevenbergMarquardtMinimizer::Initialize(cudaStream_t stream, Problem &problem) { GaussNewtonMinimizer::Initialize(stream, problem); lambda_ = options_.initial_lambda; } -} // namespace cunls \ No newline at end of file +} // namespace cunls \ No newline at end of file diff --git a/cunls/minimizer/levenberg_marquardt_minimizer.h b/cunls/minimizer/levenberg_marquardt_minimizer.h index cf7932d..6f5a720 100644 --- a/cunls/minimizer/levenberg_marquardt_minimizer.h +++ b/cunls/minimizer/levenberg_marquardt_minimizer.h @@ -122,18 +122,17 @@ struct LevenbergMarquardtMinimizerOptions { * while often converging faster than gradient descent. */ class LevenbergMarquardtMinimizer : public GaussNewtonMinimizer { -public: + public: /** * @brief Constructs a Levenberg-Marquardt optimizer. * * @param options Configuration options. Defaults to standard LM options. */ LevenbergMarquardtMinimizer( - const LevenbergMarquardtMinimizerOptions &options = - LevenbergMarquardtMinimizerOptions()) + const LevenbergMarquardtMinimizerOptions &options = LevenbergMarquardtMinimizerOptions()) : GaussNewtonMinimizer(options.base_options), options_(options) {} -private: + private: /** * @brief Initializes LM-specific data structures. * @@ -156,8 +155,7 @@ class LevenbergMarquardtMinimizer : public GaussNewtonMinimizer { * @param[out] rhs Output right-hand side vector (-J^T r). */ void BuildSystem(cudaStream_t stream, const Problem &problem, - const MinimizerState &minimizer_state, CSRSparseMatrix &lhs, - dvector &rhs) override; + const MinimizerState &minimizer_state) override; /** * @brief Checks convergence using LM-specific criteria. @@ -173,9 +171,8 @@ class LevenbergMarquardtMinimizer : public GaussNewtonMinimizer { * @param[out] step_quality Output rho metric (actual/predicted reduction). * @return True if converged, false otherwise. */ - bool CheckConvergence(cudaStream_t stream, float updated_cost, - float current_cost, const dvector &step, - float &step_quality) override; + bool CheckConvergence(cudaStream_t stream, float updated_cost, float current_cost, + const dvector &step, float &step_quality) override; /** * @brief Fused cost + LM convergence with a single D2H + sync. @@ -184,10 +181,8 @@ class LevenbergMarquardtMinimizer : public GaussNewtonMinimizer { * and sparse-weighted step norm into one memcpy and one sync. */ bool EvaluateAndCheckConvergence(cudaStream_t stream, const Problem &problem, - const MinimizerState &updated_state, - float current_cost, - const dvector &step, - float &updated_cost, + const MinimizerState &updated_state, float current_cost, + const dvector &step, float &updated_cost, float &step_quality) override; /** @@ -212,11 +207,11 @@ class LevenbergMarquardtMinimizer : public GaussNewtonMinimizer { */ bool RejectStep(float step_quality) override; - const LevenbergMarquardtMinimizerOptions options_; ///< LM-specific options. + const LevenbergMarquardtMinimizerOptions options_; ///< LM-specific options. - dvector diagonal_; ///< Diagonal of the Hessian matrix (J^T J). + dvector diagonal_; ///< Diagonal of the Hessian matrix (J^T J). - float lambda_; ///< Current damping factor. + float lambda_; ///< Current damping factor. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/llms.txt b/cunls/minimizer/llms.txt index 760d5e4..bbda3c9 100644 --- a/cunls/minimizer/llms.txt +++ b/cunls/minimizer/llms.txt @@ -14,23 +14,28 @@ optimization. - `levenberg_marquardt_minimizer.h`: `LevenbergMarquardtMinimizerOptions`, `LevenbergMarquardtMinimizer` - `minimizer_state.h`: `MinimizerState` snapshot/copy utilities -- `sparse_matrix.h`: sparse conversion, RHS, norms (free functions) -- `sparse_matrix_multiplier.h`: `SparseMatrixMultiplier` base class, enum, - factory, and `SparseMatrixMultiplierPtr` -- `cusparse_matrix_multiplier.h`: `cuSPARSESparseMatrixMultiplier` (cuSPARSE backend) -- `fast_matrix_multiplier.h`: `FastSparseMatrixMultiplier` - (custom CUDA kernels) +- `sparse_matrix.h`: CSR diagonal ops, scaling, norms (free functions) +- `normal_equations.h`: `NormalEquations` (owns H, the working LHS, and their + storage layout; the only place either layout is named) +- `bsr_matrix.h`: `BSRSparseMatrix` ops, `ChooseHessianBlockSize`, block SpMV +- `hessian_structure.h`: `HessianStructureBuilder` (GPU sparsity derivation) +- `block_hessian_assembler.h`: `BlockHessianAssembler` (one kernel for H and rhs) ## Solver flow -1. Evaluate residuals/Jacobians for all residual batches. -2. Build sparse Jacobian in triplet, convert/update CSR. -3. Form normal equations (`J^T J`, `-J^T r`). -4. Solve sparse linear system. -5. Apply manifold-aware state updates (`StateBatchOps`). -6. Accept/reject step and iterate until convergence. +1. Evaluate residuals and per-factor dense Jacobian blocks for all residual + batches. +2. Contract each factor locally and scatter-add into the normal equations + (`H = J^T J`, `-J^T r`); the global Jacobian is never materialized. +3. Solve the sparse linear system. +4. Apply manifold-aware state updates (`StateBatchOps`). +5. Accept/reject step and iterate until convergence. ## Important constraints - Problem pointers are non-owning; caller keeps lifetime. - Inputs and intermediate buffers are GPU-resident. +- The Hessian is stored as BSR when the state tangent dimensions share a common + factor and the solver supports it, otherwise as scalar CSR. This is chosen + automatically inside `NormalEquations`; no user-facing option, and no storage + branch outside that class. diff --git a/cunls/minimizer/minimizer_state.cu b/cunls/minimizer/minimizer_state.cu index 54e2f25..128ec2c 100644 --- a/cunls/minimizer/minimizer_state.cu +++ b/cunls/minimizer/minimizer_state.cu @@ -167,6 +167,12 @@ void MinimizerState::Create(cudaStream_t stream, const Problem &problem) { assert(param_ptrs.size() == new_ptrs.size()); + // A factor batch may legitimately hold zero factors; a zero-size grid is + // an invalid launch configuration. + if (param_ptrs.empty()) { + continue; + } + float **new_pointers = new_ptrs.data(); float *const *old_pointers = param_ptrs.data(); diff --git a/cunls/minimizer/minimizer_state.h b/cunls/minimizer/minimizer_state.h index d5f3e2f..ecf14f7 100644 --- a/cunls/minimizer/minimizer_state.h +++ b/cunls/minimizer/minimizer_state.h @@ -38,8 +38,8 @@ namespace cunls { * - state_pointers_: One device vector per residual batch containing pointers * to state blocks, remapped to point into the copied state * storage. - * - problem_state_ptrs_device_: Device copy of host problem pointer lists for - * Jacobian structure and remap kernels. + * - problem_state_ptrs_device_: Device copy of host problem pointer lists, + * used by the pointer-remap kernel. */ class MinimizerState { public: @@ -54,17 +54,13 @@ class MinimizerState { * @param stream CUDA stream for GPU operations. * @param problem The problem to create a state snapshot from. */ - MinimizerState(cudaStream_t stream, const Problem &problem) { - Create(stream, problem); - } + MinimizerState(cudaStream_t stream, const Problem &problem) { Create(stream, problem); } /** * @brief Refreshes storage from the problem (realloc only when capacity is * insufficient). */ - void Recreate(cudaStream_t stream, const Problem &problem) { - Create(stream, problem); - } + void Recreate(cudaStream_t stream, const Problem &problem) { Create(stream, problem); } /** * @brief Copies state values from another state. @@ -107,21 +103,7 @@ class MinimizerState { * * @return Const reference to vector of state pointer vectors. */ - const std::vector> &GetStatePointers() const { - return state_pointers_; - } - - /** - * @brief Builds the triplet (COO) Jacobian sparsity structure on the GPU. - * - * Definition (implementation) in jacobian_ops.cu. - * - * @param stream CUDA stream for GPU operations. - * @param problem The optimization problem. - * @param[out] structure Output row and column index arrays. - */ - void BuildTripletSparseStructure(cudaStream_t stream, const Problem &problem, - TripletSparseStructure &structure); + const std::vector> &GetStatePointers() const { return state_pointers_; } private: /** @@ -168,8 +150,7 @@ class MinimizerState { */ std::vector> state_pointers_; - /// Device copy of problem.GetStatePointers() for Jacobian FillColIds and - /// remap. + /// Device copy of problem.GetStatePointers(), used by the remap kernel. std::vector> problem_state_ptrs_device_; }; diff --git a/cunls/minimizer/normal_equations.cu b/cunls/minimizer/normal_equations.cu new file mode 100644 index 0000000..0d7e074 --- /dev/null +++ b/cunls/minimizer/normal_equations.cu @@ -0,0 +1,120 @@ +/* + * 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/common/log.h" +#include "cunls/minimizer/normal_equations.h" +#include "cunls/minimizer/problem.h" +#include "cunls/minimizer/sparse_matrix.h" + +namespace cunls { + +void NormalEquations::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + bool solver_supports_block_storage) { + csr_dims_.Invalidate(); + + // Block storage only pays off when the tangent dimensions share a factor and + // the solver can read tiles; a CSR-only backend would just have to expand + // them again. + block_size_ = solver_supports_block_storage ? ChooseHessianBlockSize(problem) : 1; + + if (UsesBlockStorage()) { + assembler_.Initialize(stream, problem, num_cols, block_size_, bsr_hessian_); + LogMessage("Hessian storage: BSR with block size {}", block_size_); + return; + } + + assembler_.Initialize(stream, problem, num_cols, csr_hessian_); + int num_rows = 0, num_matrix_cols = 0, num_nonzeros = 0; + ExtractMatrixMetadata(stream, csr_hessian_, num_rows, num_matrix_cols, num_nonzeros); + csr_dims_.Set(num_rows, num_matrix_cols, num_nonzeros); +} + +void NormalEquations::Assemble(cudaStream_t stream, const Problem &problem, const float *jacobians, + const float *residuals, dvector &rhs) { + if (UsesBlockStorage()) { + assembler_.Assemble(stream, problem, jacobians, residuals, bsr_hessian_, rhs); + CopyBSRSparseMatrix(stream, bsr_hessian_, bsr_lhs_); + return; + } + assembler_.Assemble(stream, problem, jacobians, residuals, csr_hessian_, rhs); + CopyCSRSparseMatrix(stream, csr_hessian_, csr_lhs_); +} + +void NormalEquations::ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const { + if (UsesBlockStorage()) { + ExtractDiagonal(stream, bsr_hessian_, diagonal); + return; + } + ExtractDiagonal(stream, csr_hessian_, diagonal); +} + +void NormalEquations::ExtractLhsDiagonal(cudaStream_t stream, dvector &diagonal) const { + if (UsesBlockStorage()) { + ExtractDiagonal(stream, bsr_lhs_, diagonal); + return; + } + ExtractDiagonal(stream, csr_lhs_, diagonal); +} + +void NormalEquations::AddScaledDiagonalToLhs(cudaStream_t stream, float scale, + const dvector &diagonal) { + if (UsesBlockStorage()) { + AddScaledDiagonal(stream, scale, diagonal, bsr_lhs_, bsr_lhs_); + return; + } + AddScaledDiagonal(stream, scale, diagonal, csr_lhs_, csr_lhs_); +} + +void NormalEquations::ScaleLhsSymmetric(cudaStream_t stream, const dvector &scale) { + if (UsesBlockStorage()) { + ScaleSymmetric(stream, bsr_lhs_, scale, tile_row_scratch_); + return; + } + ScaleSymmetricCSR(stream, csr_lhs_, scale); +} + +void NormalEquations::WeightedSquaredStepAsync(cudaStream_t stream, void *cusparse_handle, + const dvector &step, float *d_out, + float *d_partials, dvector &buffer) { + if (UsesBlockStorage()) { + ComputeWeightedSquaredStepAsync(stream, bsr_hessian_, step, block_spmv_scratch_, d_out, + d_partials); + return; + } + ComputeWeightedSquaredStepAsync(stream, cusparse_handle, csr_hessian_, csr_dims_.num_rows, + csr_dims_.num_cols, csr_dims_.num_nonzeros, step, buffer, d_out, + d_partials); +} + +bool NormalEquations::InitializeSolver(cudaStream_t stream, CSRSparseLinearSolver &solver, + const Problem &problem, const dvector &rhs, + dvector &step) { + if (UsesBlockStorage()) { + return solver.Initialize(stream, problem, bsr_lhs_, rhs, step); + } + return solver.Initialize(stream, problem, csr_lhs_, rhs, step); +} + +bool NormalEquations::Solve(cudaStream_t stream, CSRSparseLinearSolver &solver, + const dvector &rhs, dvector &step) { + if (UsesBlockStorage()) { + return solver.Solve(stream, bsr_lhs_, rhs, step); + } + return solver.Solve(stream, csr_lhs_, rhs, step); +} + +} // namespace cunls diff --git a/cunls/minimizer/normal_equations.h b/cunls/minimizer/normal_equations.h new file mode 100644 index 0000000..cf48e25 --- /dev/null +++ b/cunls/minimizer/normal_equations.h @@ -0,0 +1,146 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "cunls/common/types.h" +#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/minimizer/block_hessian_assembler.h" +#include "cunls/minimizer/bsr_matrix.h" + +namespace cunls { + +class Problem; + +/** + * @brief The assembled normal equations, and the storage they live in. + * + * Holds two matrices: the Hessian `H` as assembled, and a working copy the + * minimizer is free to damp and scale before handing it to the solver. Both + * are stored in one of two layouts, chosen once at Initialize(): + * + * - **Block (BSR)** when the state tangent dimensions share a common factor and + * the solver can read tiles. One column index per dense tile instead of one + * per scalar entry, which is bandwidth the solver's SpMV no longer moves. + * - **Scalar (CSR)** otherwise. + * + * Which layout is live is an implementation detail. Callers work in terms of + * "the Hessian" and "the working left-hand side"; every operation dispatches + * internally, so no branch on storage leaks into the minimizers. + * + * A BSR matrix with `block_size == 1` is byte-for-byte the same layout as CSR, + * so the scalar case could in principle be expressed as block storage too. It + * is kept separate deliberately: at `block_size == 1` every block operation + * degenerates to its scalar counterpart but runs through less-optimized code — + * notably cuSPARSE's `csrmv`, which reaches ~90% of peak bandwidth and which a + * one-tile-per-entry BSR kernel cannot beat. Block storage is a win only when + * the tiles are real. + */ +class NormalEquations { +public: + /** + * @brief Derives the sparsity pattern and picks the storage layout. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem. + * @param num_cols Number of free tangent dimensions in the reduced system. + * @param solver_supports_block_storage Whether the active solver can consume + * block storage; see CSRSparseLinearSolver::SupportsBlockStorage. + */ + void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, + bool solver_supports_block_storage); + + /** @brief Floats needed for the per-factor Jacobian buffer Assemble() reads. */ + size_t JacobianValuesSize() const { return assembler_.JacobianValuesSize(); } + + /** + * @brief Assembles `H` and `rhs`, then refreshes the working left-hand side. + * + * @param stream CUDA stream for GPU operations. + * @param problem The optimization problem (must match Initialize). + * @param jacobians Per-factor dense Jacobian blocks. + * @param residuals Residual vector. + * @param[out] rhs Right-hand side `-J^T r`. + */ + void Assemble(cudaStream_t stream, const Problem &problem, const float *jacobians, + const float *residuals, dvector &rhs); + + /** @brief Diagonal of the assembled (undamped) Hessian. */ + void ExtractHessianDiagonal(cudaStream_t stream, dvector &diagonal) const; + + /** @brief Diagonal of the working left-hand side. */ + void ExtractLhsDiagonal(cudaStream_t stream, dvector &diagonal) const; + + /** @brief Adds `scale * diag(diagonal)` to the working left-hand side. */ + void AddScaledDiagonalToLhs(cudaStream_t stream, float scale, const dvector &diagonal); + + /** @brief Applies `A_ij *= scale[i] * scale[j]` to the working left-hand side. */ + void ScaleLhsSymmetric(cudaStream_t stream, const dvector &scale); + + /** + * @brief Async `step^T H step` against the assembled (undamped) Hessian. + * + * @param stream CUDA stream for GPU operations. + * @param cusparse_handle Opaque cuSPARSE handle, used by the scalar path. + * @param step Step vector. + * @param[out] d_out Device destination for the scalar. + * @param[out] d_partials Reduction scratch; see device_reduction.h. + * @param[out] buffer Scratch for the scalar path's SpMV. + */ + void WeightedSquaredStepAsync(cudaStream_t stream, void *cusparse_handle, + const dvector &step, float *d_out, float *d_partials, + dvector &buffer); + + /** @brief Hands the working left-hand side to the solver for symbolic setup. */ + bool InitializeSolver(cudaStream_t stream, CSRSparseLinearSolver &solver, const Problem &problem, + const dvector &rhs, dvector &step); + + /** @brief Solves with the working left-hand side. */ + bool Solve(cudaStream_t stream, CSRSparseLinearSolver &solver, const dvector &rhs, + dvector &step); + + /** @brief True when the block layout is live for the current problem. */ + bool UsesBlockStorage() const { return block_size_ > 1; } + + /** @brief Tile edge of the block layout; 1 when the scalar layout is live. */ + int BlockSize() const { return block_size_; } + + /** @brief Working left-hand side in scalar storage; empty under block storage. */ + const CSRSparseMatrix &LhsCSR() const { return csr_lhs_; } + + /** @brief Working left-hand side in block storage; empty under scalar storage. */ + const BSRSparseMatrix &LhsBSR() const { return bsr_lhs_; } + +private: + BlockHessianAssembler assembler_; + + // Exactly one pair is populated, decided by block_size_. + CSRSparseMatrix csr_hessian_; + CSRSparseMatrix csr_lhs_; + BSRSparseMatrix bsr_hessian_; + BSRSparseMatrix bsr_lhs_; + + CSRMatrixDimensions csr_dims_; ///< Cached dims for the scalar SpMV. + dvector tile_row_scratch_; ///< Tile-to-block-row map for scaling. + dvector block_spmv_scratch_; ///< SpMV result for the block path. + + int block_size_ = 1; +}; + +} // namespace cunls diff --git a/cunls/minimizer/residual_batch.cu b/cunls/minimizer/residual_batch.cu index 43cc1d2..cc0d2c4 100644 --- a/cunls/minimizer/residual_batch.cu +++ b/cunls/minimizer/residual_batch.cu @@ -212,18 +212,23 @@ bool ResidualBatch::Evaluate(cudaStream_t stream, float *workspace, float *residuals, float const *const *state_pointers, float *cost, float *jacobians) const { - assert(residuals != nullptr); - assert(state_pointers != nullptr); - - factor_batch_->Evaluate(residuals, jacobians, state_pointers, stream); - int num_residuals = static_cast(factor_batch_->NumFactors()); int residual_dim = static_cast(factor_batch_->ResidualsSize()); - if (num_residuals == 0) + // Return before the preconditions below: an empty batch has nothing to + // evaluate, its buffers are legitimately null (a zero-size DeviceVector has + // no allocation), and the factor kernels would be launched with a zero-size + // grid. + if (num_residuals == 0) { return true; + } + + assert(residuals != nullptr); + assert(state_pointers != nullptr); assert(workspace != nullptr); + factor_batch_->Evaluate(residuals, jacobians, state_pointers, stream); + float *sq_err_ptr = nullptr; float3 *rho_ptr = nullptr; MapRobustWorkspace(workspace, num_residuals, &sq_err_ptr, &rho_ptr); diff --git a/cunls/minimizer/sparse_matrix.cu b/cunls/minimizer/sparse_matrix.cu index 34afa83..ad7fa31 100644 --- a/cunls/minimizer/sparse_matrix.cu +++ b/cunls/minimizer/sparse_matrix.cu @@ -50,18 +50,21 @@ namespace cunls { * The number of columns is determined by finding the maximum column index + 1. * Requires matrix to have at least one row and one non-zero element. */ -void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, - int &num_rows, int &num_cols, int &num_nonzeros) { - num_rows = matrix.row_offsets.size() - 1; - assert(num_rows > 0); - - num_nonzeros = matrix.values.size(); - assert(num_nonzeros > 0); +void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, int &num_rows, + int &num_cols, int &num_nonzeros) { + num_rows = matrix.row_offsets.empty() ? 0 : static_cast(matrix.row_offsets.size() - 1); + num_nonzeros = static_cast(matrix.values.size()); + + // A fully-constrained problem (every state block constant) yields an empty + // system; max_element over an empty range would dereference end(). + if (num_nonzeros == 0) { + num_cols = 0; + return; + } auto stream_policy = thrust::cuda::par_nosync.on(stream); thrust::device_ptr col_ids_ptr(matrix.col_ids.data()); - auto max_col_idx_it = thrust::max_element(stream_policy, col_ids_ptr, - col_ids_ptr + num_nonzeros); + auto max_col_idx_it = thrust::max_element(stream_policy, col_ids_ptr, col_ids_ptr + num_nonzeros); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); num_cols = *max_col_idx_it + 1; assert(num_cols > 0); @@ -73,7 +76,8 @@ void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, * indicates missing or invalid elements (e.g., due to constant states in * Jacobians). */ -template struct NotEqualOperator { +template +struct NotEqualOperator { __host__ __device__ bool operator()(const int &x) { return x != Value; } }; @@ -92,15 +96,14 @@ template struct NotEqualOperator { */ __global__ void extract_diagonal_kernel(const int *__restrict__ row_ptr, const int *__restrict__ col_ind, - const float *__restrict__ values, - float *__restrict__ diag, + const float *__restrict__ values, float *__restrict__ diag, int num_rows) { - int row = blockIdx.x * blockDim.y + threadIdx.y; // warp-level row assignment + int row = blockIdx.x * blockDim.y + threadIdx.y; // warp-level row assignment if (row >= num_rows) { return; } - int lane = threadIdx.x; // thread within warp + int lane = threadIdx.x; // thread within warp int start = row_ptr[row]; int end = row_ptr[row + 1]; @@ -138,10 +141,8 @@ __global__ void extract_diagonal_kernel(const int *__restrict__ row_ptr, */ __global__ void add_scaled_diagonal_kernel(const int *__restrict__ row_offsets, const int *__restrict__ col_indices, - float *__restrict__ values, - float scale, - const float *__restrict__ diagonal, - int num_rows) { + float *__restrict__ values, float scale, + const float *__restrict__ diagonal, int num_rows) { // One warp per row: warp_id identifies the row, lane_id is the thread within // warp const int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / WARP_SIZE; @@ -173,81 +174,11 @@ __global__ void add_scaled_diagonal_kernel(const int *__restrict__ row_offsets, if (is_diag) { values[idx] += scale * diagonal[row]; } - return; // All lanes exit together (uniform control flow) + return; // All lanes exit together (uniform control flow) } } } -/** - * CUDA kernel to build a mapping from triplet indices to CSR indices. - * For each entry in the original triplet structure, finds its position - * in the CSR format using linear search within the row. - * - * @param triplet_row_ids Original triplet row indices - * @param triplet_col_ids Original triplet column indices - * @param num_triplets Total number of triplet entries - * @param csr_row_offsets CSR row offsets array - * @param csr_col_ids CSR column indices array - * @param mapping Output mapping: mapping[triplet_idx] = csr_idx, or -1 if - * invalid - */ -__global__ void build_triplet_to_csr_mapping_kernel( - const int *__restrict__ triplet_row_ids, - const int *__restrict__ triplet_col_ids, int num_triplets, - const int *__restrict__ csr_row_offsets, - const int *__restrict__ csr_col_ids, int *__restrict__ mapping) { - int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_triplets) { - return; - } - - int col = triplet_col_ids[tid]; - if (col < 0) { - mapping[tid] = -1; - return; - } - - int row = triplet_row_ids[tid]; - int start = csr_row_offsets[row]; - int end = csr_row_offsets[row + 1]; - - // Linear search for col in csr_col_ids[start:end] - for (int i = start; i < end; i++) { - if (csr_col_ids[i] == col) { - mapping[tid] = i; - return; - } - } - - // Should not be reached for valid entries - mapping[tid] = -1; -} - -/** - * CUDA kernel to scatter triplet values into CSR format using a precomputed - * mapping. Each thread processes one triplet entry and writes its value to - * the corresponding CSR position. - * - * @param triplet_values Source triplet values - * @param mapping Precomputed triplet-to-CSR index mapping - * @param num_triplets Total number of triplet entries - * @param csr_values Destination CSR values array - */ -__global__ void -scatter_triplet_values_kernel(const float *__restrict__ triplet_values, - const int *__restrict__ mapping, int num_triplets, - float *__restrict__ csr_values) { - int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_triplets) { - return; - } - - int csr_idx = mapping[tid]; - if (csr_idx >= 0) { - csr_values[csr_idx] = triplet_values[tid]; - } -} - /** * Creates a deep copy of a CSR sparse matrix. * @@ -275,18 +206,16 @@ void CopyCSRSparseMatrix(cudaStream_t stream, const CSRSparseMatrix &input, thrust::device_ptr in_values_ptr(input.values.data()); thrust::device_ptr out_values_ptr(output.values.data()); - thrust::copy(stream_policy, in_values_ptr, - in_values_ptr + input.values.size(), out_values_ptr); + thrust::copy(stream_policy, in_values_ptr, in_values_ptr + input.values.size(), out_values_ptr); thrust::device_ptr in_col_ids_ptr(input.col_ids.data()); thrust::device_ptr out_col_ids_ptr(output.col_ids.data()); - thrust::copy(stream_policy, in_col_ids_ptr, - in_col_ids_ptr + input.col_ids.size(), out_col_ids_ptr); + thrust::copy(stream_policy, in_col_ids_ptr, in_col_ids_ptr + input.col_ids.size(), + out_col_ids_ptr); thrust::device_ptr in_row_offsets_ptr(input.row_offsets.data()); thrust::device_ptr out_row_offsets_ptr(output.row_offsets.data()); - thrust::copy(stream_policy, in_row_offsets_ptr, - in_row_offsets_ptr + input.row_offsets.size(), + thrust::copy(stream_policy, in_row_offsets_ptr, in_row_offsets_ptr + input.row_offsets.size(), out_row_offsets_ptr); } @@ -294,9 +223,10 @@ void CopyCSRSparseMatrix(cudaStream_t stream, const CSRSparseMatrix &input, * Symmetric diagonal scaling of CSR values: A_ij *= scale[i]*scale[j]. * One CUDA warp per row; lanes stride over that row's nnz for coalesced access. */ -__global__ void scale_symmetric_csr_rows_kernel( - const int *__restrict__ row_offsets, const int *__restrict__ col_ids, - float *__restrict__ values, const float *__restrict__ scale, int num_rows) { +__global__ void scale_symmetric_csr_rows_kernel(const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + float *__restrict__ values, + const float *__restrict__ scale, int num_rows) { int row = blockIdx.x * blockDim.y + threadIdx.y; if (row >= num_rows) { return; @@ -311,20 +241,18 @@ __global__ void scale_symmetric_csr_rows_kernel( } } -void ScaleSymmetricCSR(cudaStream_t stream, CSRSparseMatrix &matrix, - const dvector &scale) { +void ScaleSymmetricCSR(cudaStream_t stream, CSRSparseMatrix &matrix, const dvector &scale) { int num_rows = static_cast(matrix.row_offsets.size() - 1); assert(static_cast(scale.size()) == num_rows); dim3 block(WARP_SIZE, 8); dim3 grid((num_rows + block.y - 1) / block.y); scale_symmetric_csr_rows_kernel<<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - scale.data(), num_rows); + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), scale.data(), + num_rows); THROW_ON_CUDA_ERROR(cudaGetLastError()); } -__global__ void inv_sqrt_floor_kernel(float *__restrict__ v, int n, - float floor_value) { +__global__ void inv_sqrt_floor_kernel(float *__restrict__ v, int n, float floor_value) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= n) { return; @@ -336,69 +264,14 @@ __global__ void inv_sqrt_floor_kernel(float *__restrict__ v, int n, v[i] = rsqrtf(x); } -void InvertSqrtWithFloorInPlace(cudaStream_t stream, dvector &v, - float floor_value) { +void InvertSqrtWithFloorInPlace(cudaStream_t stream, dvector &v, float floor_value) { int n = static_cast(v.size()); if (n == 0) { return; } constexpr int block_size = 256; int grid = (n + block_size - 1) / block_size; - inv_sqrt_floor_kernel<<>>(v.data(), n, - floor_value); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - -/** - * Per-column sum of squares ||J_{:,j}||_2^2 from CSR Jacobian entries. - * One CUDA thread per nonzero; uses atomicAdd (O(nnz), no sort / no Thrust). - */ -__global__ void -jacobian_accum_col_sq_atomic_kernel(const int *__restrict__ col_ids, - const float *__restrict__ values, int nnz, - float *__restrict__ col_sums) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= nnz) { - return; - } - float t = values[i]; - atomicAdd(col_sums + col_ids[i], t * t); -} - -__global__ void col_sq_to_inv_norm_kernel(float *__restrict__ col_sums, - int num_cols, float eps) { - int j = blockIdx.x * blockDim.x + threadIdx.x; - if (j >= num_cols) { - return; - } - float s = col_sums[j]; - if (s < eps) { - s = eps; - } - col_sums[j] = rsqrtf(s); -} - -void ComputeJacobianColumnScaling(cudaStream_t stream, - const CSRSparseMatrix &jacobian, int num_cols, - int num_nonzeros, - dvector &column_scale) { - column_scale.resize(static_cast(num_cols)); - - THROW_ON_CUDA_ERROR( - cudaMemsetAsync(column_scale.data(), 0, - static_cast(num_cols) * sizeof(float), stream)); - - constexpr int block_size = 256; - int grid_nnz = (num_nonzeros + block_size - 1) / block_size; - jacobian_accum_col_sq_atomic_kernel<<>>( - jacobian.col_ids.data(), jacobian.values.data(), num_nonzeros, - column_scale.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - - int grid_cols = (num_cols + block_size - 1) / block_size; - const float floor_value = 1e-12f; - col_sq_to_inv_norm_kernel<<>>( - column_scale.data(), num_cols, floor_value); + inv_sqrt_floor_kernel<<>>(v.data(), n, floor_value); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -413,16 +286,15 @@ void ComputeJacobianColumnScaling(cudaStream_t stream, * Each warp processes one row, searching for the diagonal element (where row == * col). */ -void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, - dvector &diagonal) { +void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, dvector &diagonal) { size_t num_rows = matrix.row_offsets.size() - 1; diagonal.resize(num_rows); dim3 block(32, 4); dim3 grid((num_rows + block.y - 1) / block.y); - extract_diagonal_kernel<<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - diagonal.data(), num_rows); + extract_diagonal_kernel<<>>(matrix.row_offsets.data(), + matrix.col_ids.data(), matrix.values.data(), + diagonal.data(), num_rows); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -441,64 +313,43 @@ void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, * Uses cuSPARSE SpMV (Sparse Matrix-Vector multiplication) with preprocessing * for optimal performance. Buffer is automatically resized as needed. */ -static void SpMVImpl(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, int num_rows, int num_cols, - int num_nonzeros, bool transpose_matrix, - const dvector &x, dvector &result, - dvector &buffer) { +static void SpMVImpl(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, int num_rows, + int num_cols, int num_nonzeros, bool transpose_matrix, const dvector &x, + dvector &result, dvector &buffer) { auto cusparse_handle = static_cast(handle); result.resize(transpose_matrix ? num_cols : num_rows); - assert(x.size() == - static_cast(transpose_matrix ? num_rows : num_cols)); + assert(x.size() == static_cast(transpose_matrix ? num_rows : num_cols)); - cuSPARSEMatrixDescription matrix_description(num_rows, num_cols, num_nonzeros, - matrix); + cuSPARSEMatrixDescription matrix_description(num_rows, num_cols, num_nonzeros, matrix); cuSPARSEVectorDescription vec_x_description(x); cuSPARSEVectorDescription vec_result_description(result); - auto matA = - static_cast(matrix_description.GetDescription()); - auto vecX = - static_cast(vec_x_description.GetDescription()); - auto vecY = static_cast( - vec_result_description.GetDescription()); + auto matA = static_cast(matrix_description.GetDescription()); + auto vecX = static_cast(vec_x_description.GetDescription()); + auto vecY = static_cast(vec_result_description.GetDescription()); constexpr float alpha = 1; constexpr float beta = 0; - cusparseOperation_t operation = transpose_matrix - ? CUSPARSE_OPERATION_TRANSPOSE - : CUSPARSE_OPERATION_NON_TRANSPOSE; + cusparseOperation_t operation = + transpose_matrix ? CUSPARSE_OPERATION_TRANSPOSE : CUSPARSE_OPERATION_NON_TRANSPOSE; size_t bufferSize = 0; - THROW_ON_CUSPARSE_ERROR(cusparseSpMV_bufferSize( - cusparse_handle, operation, &alpha, matA, vecX, &beta, vecY, CUDA_R_32F, - CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize)); + THROW_ON_CUSPARSE_ERROR(cusparseSpMV_bufferSize(cusparse_handle, operation, &alpha, matA, vecX, + &beta, vecY, CUDA_R_32F, + CUSPARSE_SPMV_ALG_DEFAULT, &bufferSize)); buffer.resize(bufferSize); auto buffer_ptr = buffer.data(); - THROW_ON_CUSPARSE_ERROR(cusparseSpMV_preprocess( - cusparse_handle, operation, &alpha, matA, vecX, &beta, vecY, CUDA_R_32F, - CUSPARSE_SPMV_ALG_DEFAULT, buffer_ptr)); - - THROW_ON_CUSPARSE_ERROR(cusparseSpMV(cusparse_handle, operation, &alpha, matA, - vecX, &beta, vecY, CUDA_R_32F, - CUSPARSE_SPMV_ALG_DEFAULT, buffer_ptr)); -} + THROW_ON_CUSPARSE_ERROR(cusparseSpMV_preprocess(cusparse_handle, operation, &alpha, matA, vecX, + &beta, vecY, CUDA_R_32F, + CUSPARSE_SPMV_ALG_DEFAULT, buffer_ptr)); -void MultiplySparseMatrixByDenseVector(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, - bool transpose_matrix, - const dvector &x, - dvector &result, - dvector &buffer) { - int num_rows, num_cols, num_nonzeros; - ExtractMatrixMetadata(stream, matrix, num_rows, num_cols, num_nonzeros); - SpMVImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, - transpose_matrix, x, result, buffer); + THROW_ON_CUSPARSE_ERROR(cusparseSpMV(cusparse_handle, operation, &alpha, matA, vecX, &beta, vecY, + CUDA_R_32F, CUSPARSE_SPMV_ALG_DEFAULT, buffer_ptr)); } /** @@ -515,8 +366,7 @@ void MultiplySparseMatrixByDenseVector(cudaStream_t stream, void *handle, * First copies the input matrix, then uses a CUDA kernel to add the scaled * diagonal elements to the existing diagonal entries of the matrix. */ -void AddScaledDiagonal(cudaStream_t stream, float scale, - const dvector &diagonal, +void AddScaledDiagonal(cudaStream_t stream, float scale, const dvector &diagonal, const CSRSparseMatrix &matrix, CSRSparseMatrix &result) { int num_rows = diagonal.size(); assert(num_rows + 1 == matrix.row_offsets.size()); @@ -524,141 +374,12 @@ void AddScaledDiagonal(cudaStream_t stream, float scale, CopyCSRSparseMatrix(stream, matrix, result); // Launch one warp (32 threads) per row for warp-cooperative diagonal search - constexpr int block_size = 256; // Must be multiple of WARP_SIZE + constexpr int block_size = 256; // Must be multiple of WARP_SIZE const int total_threads = num_rows * WARP_SIZE; const int blocks = (total_threads + block_size - 1) / block_size; add_scaled_diagonal_kernel<<>>( - result.row_offsets.data(), result.col_ids.data(), result.values.data(), - scale, diagonal.data(), num_rows); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - -/** - * Converts a TripletSparseStructure to CSR format and builds a mapping from - * triplet indices to CSR indices. The mapping enables efficient value updates - * without re-computing the structure on each iteration. - * - * @param stream CUDA stream for asynchronous operations - * @param handle cuSPARSE library handle - * @param structure Triplet sparse structure (may contain -1 for invalid - * entries) - * @param csr Output CSR sparse matrix (structure filled, values zeroed) - * @param mapping Output mapping: mapping[triplet_idx] = csr_idx, or -1 - * @param buffer Temporary buffer for intermediate computations - */ -void ConvertTripletStructureToCSR(cudaStream_t stream, void *handle, - const TripletSparseStructure &structure, - CSRSparseMatrix &csr, dvector &mapping, - dvector &buffer) { - auto cusparse_handle = static_cast(handle); - auto stream_policy = thrust::cuda::par_nosync.on(stream); - - const auto &col_ids = structure.col_ids; - const auto &row_ids = structure.row_ids; - size_t num_triplets = col_ids.size(); - - if (num_triplets == 0) { - csr.row_offsets.resize(1); - csr.col_ids.resize(0); - csr.values.resize(0); - mapping.resize(0); - THROW_ON_CUDA_ERROR( - cudaMemsetAsync(csr.row_offsets.data(), 0, sizeof(int), stream)); - return; - } - - // Count valid entries (col_id != -1) and find dimensions - thrust::device_ptr col_ids_ptr(col_ids.data()); - int number_of_nonzeros = - thrust::count_if(stream_policy, col_ids_ptr, col_ids_ptr + num_triplets, - NotEqualOperator<-1>()); - - thrust::device_ptr row_ids_ptr(row_ids.data()); - auto max_row_idx_it = thrust::max_element(stream_policy, row_ids_ptr, - row_ids_ptr + num_triplets); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - - int num_rows = *max_row_idx_it + 1; - assert(num_rows > 0); - - // Allocate CSR arrays - csr.values.resize(number_of_nonzeros); - csr.col_ids.resize(number_of_nonzeros); - csr.row_offsets.resize(num_rows + 1); - - // Phase 1: Stream compaction to filter out invalid entries (col_id == -1), - // producing COO format directly without intermediate triplet representation. - size_t buffer_size_in_bytes = number_of_nonzeros * sizeof(int); - - if (buffer.size() < buffer_size_in_bytes) { - buffer.resize(buffer_size_in_bytes); - } - - auto coo_row_ids_ptr = reinterpret_cast(buffer.data()); - - { - auto input_begin = thrust::make_zip_iterator( - thrust::make_tuple(thrust::device_pointer_cast(row_ids.data()), - thrust::device_pointer_cast(col_ids.data()))); - auto output_begin = thrust::make_zip_iterator( - thrust::make_tuple(thrust::device_pointer_cast(coo_row_ids_ptr), - thrust::device_pointer_cast(csr.col_ids.data()))); - thrust::device_ptr stencil(col_ids.data()); - - thrust::copy_if(stream_policy, input_begin, input_begin + num_triplets, - stencil, output_begin, NotEqualOperator<-1>()); - - THROW_ON_CUDA_ERROR(cudaMemsetAsync( - csr.values.data(), 0, number_of_nonzeros * sizeof(float), stream)); - } - - // Convert COO row indices to CSR row offsets - THROW_ON_CUSPARSE_ERROR(cusparseXcoo2csr( - cusparse_handle, coo_row_ids_ptr, number_of_nonzeros, num_rows, - csr.row_offsets.data(), cusparseIndexBase_t::CUSPARSE_INDEX_BASE_ZERO)); - - // Phase 2: Build mapping from triplet indices to CSR indices. - // For each original triplet entry, linear search in the CSR - // structure to find the corresponding CSR index. - mapping.resize(num_triplets); - { - constexpr size_t block_size = 256; - size_t num_blocks = (num_triplets + block_size - 1) / block_size; - build_triplet_to_csr_mapping_kernel<<>>( - structure.row_ids.data(), structure.col_ids.data(), num_triplets, - csr.row_offsets.data(), csr.col_ids.data(), mapping.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - } -} - -/** - * Copies values from a SparseJacobian in triplet form to a CSRSparseMatrix - * using a precomputed mapping from ConvertTripletStructureToCSR. - * This is much faster than recomputing the full structure+values conversion - * since it only scatters values to their precomputed positions. - * - * @param stream CUDA stream for asynchronous operations - * @param jacobian Sparse Jacobian in triplet format with updated values - * @param mapping Precomputed mapping from triplet indices to CSR indices - * @param csr CSR sparse matrix with precomputed structure (values updated) - */ -void ConvertTripletToCSRValues(cudaStream_t stream, - const SparseJacobian &jacobian, - const dvector &mapping, - CSRSparseMatrix &csr) { - assert(jacobian.values.size() == jacobian.structure.col_ids.size()); - assert(jacobian.values.size() == jacobian.structure.row_ids.size()); - assert(jacobian.values.size() == mapping.size()); - - size_t num_triplets = jacobian.values.size(); - if (num_triplets == 0) { - return; - } - - constexpr size_t block_size = 256; - size_t num_blocks = (num_triplets + block_size - 1) / block_size; - scatter_triplet_values_kernel<<>>( - jacobian.values.data(), mapping.data(), num_triplets, csr.values.data()); + result.row_offsets.data(), result.col_ids.data(), result.values.data(), scale, + diagonal.data(), num_rows); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -678,58 +399,31 @@ void ConvertTripletToCSRValues(cudaStream_t stream, */ __global__ void negate_kernel(float *__restrict__ data, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) - data[i] = -data[i]; + if (i < n) data[i] = -data[i]; } void NegateVector(cudaStream_t stream, float *data, size_t n) { - if (n == 0) - return; + if (n == 0) return; constexpr int kBlock = 256; int grid = static_cast((n + kBlock - 1) / kBlock); negate_kernel<<>>(data, static_cast(n)); THROW_ON_CUDA_ERROR(cudaGetLastError()); } -__global__ void elementwise_multiply_kernel(float *__restrict__ a, - const float *__restrict__ b, +__global__ void elementwise_multiply_kernel(float *__restrict__ a, const float *__restrict__ b, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) - a[i] *= b[i]; + if (i < n) a[i] *= b[i]; } -void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, - size_t n) { - if (n == 0) - return; +void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, size_t n) { + if (n == 0) return; constexpr int kBlock = 256; int grid = static_cast((n + kBlock - 1) / kBlock); - elementwise_multiply_kernel<<>>(a, b, - static_cast(n)); + elementwise_multiply_kernel<<>>(a, b, static_cast(n)); THROW_ON_CUDA_ERROR(cudaGetLastError()); } -void ComputeRHS(cudaStream_t stream, void *handle, - const CSRSparseMatrix &jacobian, - const dvector &residuals, dvector &rhs, - dvector &buffer) { - constexpr bool transpose_matrix = true; - MultiplySparseMatrixByDenseVector(stream, handle, jacobian, transpose_matrix, - residuals, rhs, buffer); - NegateVector(stream, rhs.data(), rhs.size()); -} - -void ComputeRHS(cudaStream_t stream, void *handle, - const CSRSparseMatrix &jacobian, int num_rows, int num_cols, - int num_nonzeros, const dvector &residuals, - dvector &rhs, dvector &buffer) { - constexpr bool transpose_matrix = true; - SpMVImpl(stream, handle, jacobian, num_rows, num_cols, num_nonzeros, - transpose_matrix, residuals, rhs, buffer); - NegateVector(stream, rhs.data(), rhs.size()); -} - /** * Computes the weighted squared norm of a step vector: step^T * W * step, * where W is a diagonal weight matrix. @@ -744,31 +438,27 @@ void ComputeRHS(cudaStream_t stream, void *handle, * product with the weights: sum(step[i] * weights[i] * step[i]). * Used in trust region methods and optimization algorithms. */ -void ComputeWeightedSquaredStepAsync(cudaStream_t stream, - const dvector &weights, - const dvector &step, float *d_out, - float *d_partials) { +void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const dvector &weights, + const dvector &step, float *d_out, float *d_partials) { assert(step.size() == weights.size()); - WeightedDotProductToDevice(stream, step.data(), weights.data(), step.data(), - step.size(), d_out, d_partials); + WeightedDotProductToDevice(stream, step.data(), weights.data(), step.data(), step.size(), d_out, + d_partials); } -float ComputeWeightedSquaredStep(cudaStream_t stream, - const dvector &weights, - const dvector &step, - dvector &buffer) { +float ComputeWeightedSquaredStep(cudaStream_t stream, const dvector &weights, + const dvector &step, dvector &buffer) { assert(step.size() == weights.size()); size_t partials_count = ReducePartialCount(step.size()); buffer.resize((partials_count + 1) * sizeof(float)); float *d_out = reinterpret_cast(buffer.data()); float *d_partials = d_out + 1; - WeightedDotProductToDevice(stream, step.data(), weights.data(), step.data(), - step.size(), d_out, d_partials); + WeightedDotProductToDevice(stream, step.data(), weights.data(), step.data(), step.size(), d_out, + d_partials); float result; - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&result, d_out, sizeof(float), - cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); assert(result >= 0); return result; @@ -791,30 +481,28 @@ float ComputeWeightedSquaredStep(cudaStream_t stream, */ static thread_local dvector g_spmv_result; -static void WeightedSquaredStepSparseAsyncImpl( - cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - int num_rows, int num_cols, int num_nonzeros, const dvector &step, - dvector &buffer, float *d_out, float *d_partials) { +static void WeightedSquaredStepSparseAsyncImpl(cudaStream_t stream, void *handle, + const CSRSparseMatrix &matrix, int num_rows, + int num_cols, int num_nonzeros, + const dvector &step, dvector &buffer, + float *d_out, float *d_partials) { constexpr bool transpose_matrix = false; - SpMVImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, - transpose_matrix, step, g_spmv_result, buffer); - DotProductToDevice(stream, g_spmv_result.data(), step.data(), - g_spmv_result.size(), d_out, d_partials); + SpMVImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, transpose_matrix, step, + g_spmv_result, buffer); + DotProductToDevice(stream, g_spmv_result.data(), step.data(), g_spmv_result.size(), d_out, + d_partials); } -void ComputeWeightedSquaredStepAsync( - cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - int num_rows, int num_cols, int num_nonzeros, const dvector &step, - dvector &buffer, float *d_out, float *d_partials) { - WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, - num_nonzeros, step, buffer, d_out, - d_partials); +void ComputeWeightedSquaredStepAsync(cudaStream_t stream, void *handle, + const CSRSparseMatrix &matrix, int num_rows, int num_cols, + int num_nonzeros, const dvector &step, + dvector &buffer, float *d_out, float *d_partials) { + WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, step, + buffer, d_out, d_partials); } -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, - const dvector &step, - dvector &buffer) { +float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, + const dvector &step, dvector &buffer) { int num_rows, num_cols, num_nonzeros; ExtractMatrixMetadata(stream, matrix, num_rows, num_cols, num_nonzeros); @@ -823,35 +511,31 @@ float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, float *d_out = d_scratch.data(); float *d_partials = d_out + 1; - WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, - num_nonzeros, step, buffer, d_out, - d_partials); + WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, step, + buffer, d_out, d_partials); float result; - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&result, d_out, sizeof(float), - cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); assert(result >= 0); return result; } -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, int num_rows, - int num_cols, int num_nonzeros, - const dvector &step, - dvector &buffer) { +float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, + int num_rows, int num_cols, int num_nonzeros, + const dvector &step, dvector &buffer) { size_t partials_count = ReducePartialCount(step.size()); dvector d_scratch(partials_count + 1); float *d_out = d_scratch.data(); float *d_partials = d_out + 1; - WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, - num_nonzeros, step, buffer, d_out, - d_partials); + WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, step, + buffer, d_out, d_partials); float result; - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&result, d_out, sizeof(float), - cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); assert(result >= 0); return result; @@ -864,10 +548,9 @@ float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, * @param step Step vector. * @return The squared L2 norm (scalar value). */ -void ComputeSquaredStepAsync(cudaStream_t stream, const dvector &step, - float *d_out, float *d_partials) { - DotProductToDevice(stream, step.data(), step.data(), step.size(), d_out, - d_partials); +void ComputeSquaredStepAsync(cudaStream_t stream, const dvector &step, float *d_out, + float *d_partials) { + DotProductToDevice(stream, step.data(), step.data(), step.size(), d_out, d_partials); } float ComputeSquaredStep(cudaStream_t stream, const dvector &step) { @@ -876,12 +559,11 @@ float ComputeSquaredStep(cudaStream_t stream, const dvector &step) { float *d_out = d_scratch.data(); float *d_partials = d_out + 1; - DotProductToDevice(stream, step.data(), step.data(), step.size(), d_out, - d_partials); + DotProductToDevice(stream, step.data(), step.data(), step.size(), d_out, d_partials); float result; - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(&result, d_out, sizeof(float), - cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); assert(result >= 0); return result; diff --git a/cunls/minimizer/sparse_matrix.h b/cunls/minimizer/sparse_matrix.h index edde97f..4eb201f 100644 --- a/cunls/minimizer/sparse_matrix.h +++ b/cunls/minimizer/sparse_matrix.h @@ -36,8 +36,8 @@ namespace cunls { * @param[out] num_cols Number of columns in the matrix. * @param[out] num_nonzeros Number of nonzero elements. */ -void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, - int &num_rows, int &num_cols, int &num_nonzeros); +void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, int &num_rows, + int &num_cols, int &num_nonzeros); /** * @brief Extracts the diagonal elements from a CSR sparse matrix. @@ -49,8 +49,7 @@ void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, * @param matrix CSR sparse matrix to extract diagonal from. * @param[out] diagonal Output vector of diagonal elements. */ -void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, - dvector &diagonal); +void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, dvector &diagonal); /** * @brief Adds a scaled diagonal to a sparse matrix. @@ -64,8 +63,7 @@ void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, * @param matrix Input CSR sparse matrix. * @param[out] result Output CSR sparse matrix (may alias matrix for in-place). */ -void AddScaledDiagonal(cudaStream_t stream, float scale, - const dvector &diagonal, +void AddScaledDiagonal(cudaStream_t stream, float scale, const dvector &diagonal, const CSRSparseMatrix &matrix, CSRSparseMatrix &result); /** @@ -89,92 +87,14 @@ void CopyCSRSparseMatrix(cudaStream_t stream, const CSRSparseMatrix &input, * @param[in,out] matrix CSR matrix updated in-place. * @param scale Length must match matrix row/column dimension (square H). */ -void ScaleSymmetricCSR(cudaStream_t stream, CSRSparseMatrix &matrix, - const dvector &scale); +void ScaleSymmetricCSR(cudaStream_t stream, CSRSparseMatrix &matrix, const dvector &scale); /** * @brief Sets v[i] = 1 / sqrt(max(v[i], floor_value)) for all i (in-place). * * Used after ExtractDiagonal when building S from Hessian diagonal. */ -void InvertSqrtWithFloorInPlace(cudaStream_t stream, dvector &v, - float floor_value = 1e-12f); - -/** - * @brief Sets column_scale[j] = 1 / ||J_{:,j}||_2 from CSR Jacobian J. - * - * Accumulates J_ij^2 per column with one thread per nonzero (atomicAdd), then - * applies the epsilon floor for empty columns. No global sort and no Thrust. - * - * @param stream CUDA stream. - * @param jacobian CSR Jacobian (rows = residuals, cols = parameters). - * @param[out] column_scale Per-column scaling; length = number of columns of J. - */ -void ComputeJacobianColumnScaling(cudaStream_t stream, - const CSRSparseMatrix &jacobian, int num_cols, - int num_nonzeros, - dvector &column_scale); - -/** - * @brief Converts a triplet sparse structure to CSR format with index mapping. - * - * Filters out invalid entries (col_id == -1) from the triplet structure, - * converts valid entries to CSR format, and builds a mapping from triplet - * indices to CSR indices. The mapping enables efficient value-only updates - * on subsequent iterations without re-converting the structure. - * - * @param stream CUDA stream for GPU operations. - * @param handle Opaque cuSPARSE library handle (void*). - * @param structure Input triplet sparse structure (may contain -1 in col_ids). - * @param[out] csr Output CSR sparse matrix (structure filled, values zeroed). - * @param[out] mapping Output index mapping: mapping[triplet_idx] = csr_idx, or - * -1. - * @param[out] buffer Temporary buffer for intermediate computations. - */ -void ConvertTripletStructureToCSR(cudaStream_t stream, void *handle, - const TripletSparseStructure &structure, - CSRSparseMatrix &csr, dvector &mapping, - dvector &buffer); - -/** - * @brief Scatters Jacobian values from triplet format into CSR format. - * - * Uses the precomputed mapping from ConvertTripletStructureToCSR to copy - * updated Jacobian values from the triplet representation directly into - * their corresponding CSR positions. Much faster than full re-conversion. - * - * @param stream CUDA stream for GPU operations. - * @param jacobian Sparse Jacobian in triplet format with updated values. - * @param mapping Precomputed triplet-to-CSR index mapping. - * @param[out] csr CSR sparse matrix whose values are updated. - */ -void ConvertTripletToCSRValues(cudaStream_t stream, - const SparseJacobian &jacobian, - const dvector &mapping, - CSRSparseMatrix &csr); - -/** - * @brief Computes the right-hand side of the normal equations: rhs = -J^T * r. - * - * Performs sparse matrix-transpose-vector multiplication followed by negation - * to produce the negative gradient used in Gauss-Newton / LM solvers. - * - * @param stream CUDA stream for GPU operations. - * @param handle Opaque cuSPARSE library handle (void*). - * @param jacobian CSR sparse Jacobian matrix (J). - * @param residuals Dense residual vector (r). - * @param[out] rhs Output right-hand side vector (-J^T * r). - * @param[out] buffer Temporary buffer for cuSPARSE operations. - */ -void ComputeRHS(cudaStream_t stream, void *handle, - const CSRSparseMatrix &jacobian, - const dvector &residuals, dvector &rhs, - dvector &buffer); - -void ComputeRHS(cudaStream_t stream, void *handle, - const CSRSparseMatrix &jacobian, int num_rows, int num_cols, - int num_nonzeros, const dvector &residuals, - dvector &rhs, dvector &buffer); +void InvertSqrtWithFloorInPlace(cudaStream_t stream, dvector &v, float floor_value = 1e-12f); /** * @brief Computes the squared L2 norm of a step vector. @@ -199,10 +119,8 @@ float ComputeSquaredStep(cudaStream_t stream, const dvector &step); * @param[out] buffer Temporary buffer for intermediate computations. * @return The weighted squared norm (scalar value). */ -float ComputeWeightedSquaredStep(cudaStream_t stream, - const dvector &weights, - const dvector &step, - dvector &buffer); +float ComputeWeightedSquaredStep(cudaStream_t stream, const dvector &weights, + const dvector &step, dvector &buffer); /** * @brief Computes a sparse-matrix-weighted squared step norm. @@ -218,16 +136,12 @@ float ComputeWeightedSquaredStep(cudaStream_t stream, * @param[out] buffer Temporary buffer for cuSPARSE operations. * @return The weighted squared norm (scalar value). */ -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, - const dvector &step, - dvector &buffer); +float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, + const dvector &step, dvector &buffer); -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, int num_rows, - int num_cols, int num_nonzeros, - const dvector &step, - dvector &buffer); +float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, + int num_rows, int num_cols, int num_nonzeros, + const dvector &step, dvector &buffer); // ---- Async variants: write scalar result to device memory, no D2H or sync -- @@ -237,26 +151,24 @@ float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, * @param d_partials Scratch buffer with at least * ReducePartialCount(step.size()) floats (from device_reduction.h). */ -void ComputeSquaredStepAsync(cudaStream_t stream, const dvector &step, - float *d_out, float *d_partials); +void ComputeSquaredStepAsync(cudaStream_t stream, const dvector &step, float *d_out, + float *d_partials); /** * @brief Async diag-weighted squared step: d_out[0] = step^T diag(w) step. */ -void ComputeWeightedSquaredStepAsync(cudaStream_t stream, - const dvector &weights, - const dvector &step, float *d_out, - float *d_partials); +void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const dvector &weights, + const dvector &step, float *d_out, float *d_partials); /** * @brief Async sparse-weighted squared step: d_out[0] = step^T A step. * * Performs SpMV (A*step) then dot(step, A*step) into d_out, all on the stream. */ -void ComputeWeightedSquaredStepAsync( - cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - int num_rows, int num_cols, int num_nonzeros, const dvector &step, - dvector &buffer, float *d_out, float *d_partials); +void ComputeWeightedSquaredStepAsync(cudaStream_t stream, void *handle, + const CSRSparseMatrix &matrix, int num_rows, int num_cols, + int num_nonzeros, const dvector &step, + dvector &buffer, float *d_out, float *d_partials); /** * @brief Elementwise vector negation: out[i] = -in[i]. @@ -266,7 +178,6 @@ void NegateVector(cudaStream_t stream, float *data, size_t n); /** * @brief Elementwise multiply: out[i] = a[i] * b[i]. */ -void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, - size_t n); +void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, size_t n); } // namespace cunls diff --git a/cunls/minimizer/sparse_matrix_multiplier.cpp b/cunls/minimizer/sparse_matrix_multiplier.cpp deleted file mode 100644 index 437b29d..0000000 --- a/cunls/minimizer/sparse_matrix_multiplier.cpp +++ /dev/null @@ -1,39 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include "cunls/minimizer/sparse_matrix_multiplier.h" - -#include - -#include "cunls/minimizer/cusparse_matrix_multiplier.h" -#include "cunls/minimizer/fast_matrix_multiplier.h" - -namespace cunls { - -SparseMatrixMultiplierPtr -CreateSparseMatrixMultiplier(SparseMatrixMultiplierType type) { - switch (type) { - case SparseMatrixMultiplierType::cuSPARSE: - return std::make_unique(); - case SparseMatrixMultiplierType::Fast: - return std::make_unique(); - default: - throw std::invalid_argument("Invalid sparse square multiplier type"); - } -} - -} // namespace cunls diff --git a/cunls/minimizer/sparse_matrix_multiplier.h b/cunls/minimizer/sparse_matrix_multiplier.h deleted file mode 100644 index b4f0910..0000000 --- a/cunls/minimizer/sparse_matrix_multiplier.h +++ /dev/null @@ -1,98 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#pragma once - -#include - -#include - -#include "cunls/common/types.h" -#include "cunls/minimizer/problem.h" - -namespace cunls { - -/** - * @brief Abstract base class for computing A^T * A of a sparse CSR matrix. - * - * Provides a two-phase interface: Initialize() precomputes the output sparsity - * pattern once per problem structure, and ComputeSquaredMatrix() fills in - * numerical values and may be called repeatedly as the input values change. - * Concrete implementations can use different GPU strategies (e.g. cuSPARSE GEMM - * reuse or custom CUDA kernels). - */ -class SparseMatrixMultiplier { -public: - /** @brief Virtual destructor for proper cleanup of derived instances. */ - virtual ~SparseMatrixMultiplier() = default; - - /** - * @brief Analyzes the sparsity pattern of A^T * A and allocates the output. - * - * Must be called once whenever the sparsity pattern of @p input changes. - * Implementations may use @p problem to extract structural hints (e.g. max - * nonzeros per row) that accelerate the analysis. - * - * @param stream CUDA stream for GPU operations. - * @param problem Optimization problem providing structural information. - * @param input Input sparse matrix A (typically the Jacobian in CSR). - * @param[out] output Output sparse matrix A^T * A (structure allocated). - */ - virtual void Initialize(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) = 0; - - /** - * @brief Computes A^T * A for a sparse matrix A. - * - * The output sparsity structure must already be set by a prior call to - * Initialize(). Only the numerical values are recomputed. - * - * @param stream CUDA stream for GPU operations. - * @param problem Optimization problem (may be used for kernel tuning). - * @param input Input sparse matrix A (typically the Jacobian). - * @param[out] output Output sparse matrix A^T * A. - */ - virtual void ComputeSquaredMatrix(cudaStream_t stream, const Problem &problem, - const CSRSparseMatrix &input, - CSRSparseMatrix &output) = 0; -}; - -/** - * @brief Strategy for computing the approximate Hessian J^T * J. - */ -enum class SparseMatrixMultiplierType { - cuSPARSE, ///< cuSPARSE SpGEMM reuse API (transpose + multiply). - Fast, ///< Fast warp-efficient CUDA kernels with bitmap pattern - ///< discovery. -}; - -/** - * @brief Smart pointer type for sparse square multipliers. - */ -using SparseMatrixMultiplierPtr = std::unique_ptr; - -/** - * @brief Factory function to create a sparse square multiplier. - * - * @param type The strategy to use for computing A^T * A. - * @return A unique pointer to the created multiplier instance. - */ -SparseMatrixMultiplierPtr -CreateSparseMatrixMultiplier(SparseMatrixMultiplierType type); - -} // namespace cunls diff --git a/cunls/state/state_batch_ops.cu b/cunls/state/state_batch_ops.cu index 3f1e6de..1e8233a 100644 --- a/cunls/state/state_batch_ops.cu +++ b/cunls/state/state_batch_ops.cu @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include #include #include @@ -45,10 +47,8 @@ namespace cunls { * * Grid/block: launched with ceil(num_const_ids / 32) blocks of 32 threads. */ -__global__ void binary_pattern_kernel(bool *__restrict__ binary_states, - size_t tangent_dim, - const int *__restrict__ const_ids, - size_t num_const_ids) { +__global__ void binary_pattern_kernel(bool *__restrict__ binary_states, size_t tangent_dim, + const int *__restrict__ const_ids, size_t num_const_ids) { int tid = threadIdx.x + blockIdx.x * blockDim.x; if (tid >= num_const_ids) { return; @@ -76,21 +76,19 @@ __global__ void binary_pattern_kernel(bool *__restrict__ binary_states, * @param binary_pattern Output device vector of booleans (resized by * caller). */ -void fill_binary_pattern(cudaStream_t stream, size_t tangent_dim, - const int *const_state_block_ids, +void fill_binary_pattern(cudaStream_t stream, size_t tangent_dim, const int *const_state_block_ids, size_t num_const_state_blocks, thrust::device_vector &binary_pattern) { auto stream_policy = thrust::cuda::par_nosync.on(stream); // Fill pattern with trues - thrust::fill(stream_policy, binary_pattern.begin(), binary_pattern.end(), - true); + thrust::fill(stream_policy, binary_pattern.begin(), binary_pattern.end(), true); if (const_state_block_ids == nullptr || num_const_state_blocks == 0) { return; } bool *input_ptr = thrust::raw_pointer_cast(binary_pattern.data()); - size_t block_size = 32; // one WARP + size_t block_size = 32; // one WARP size_t num_blocks = (num_const_state_blocks + block_size - 1) / block_size; // Set the pattern to false for the constant state blocks @@ -122,14 +120,40 @@ struct CustomBinaryFunctor { /** @copydoc StateBatchOps::StateBatchOps(cudaStream_t, const * std::vector&) */ -StateBatchOps::StateBatchOps(cudaStream_t stream, - const std::vector &state_batches) { +StateBatchOps::StateBatchOps(cudaStream_t stream, const std::vector &state_batches) { Preprocess(stream, state_batches); } /** @copydoc StateBatchOps::InitUpdatesVector */ -void StateBatchOps::InitUpdatesVector( - const std::vector &state_batches) { +void ComputeStateBlockColumnOffsets(cudaStream_t stream, int first_column, + const StateBatch *state_batch, + DeviceVector &column_offsets) { + const int *const_state_ids = state_batch->ConstStateIds(); + const size_t num_const_state_blocks = state_batch->NumConstStateBlocks(); + const size_t tangent_size = state_batch->TangentSize(); + auto stream_policy = thrust::cuda::par_nosync.on(stream); + + column_offsets.resize(state_batch->NumStateBlocks()); + thrust::device_ptr begin(column_offsets.data()); + thrust::device_ptr end = begin + column_offsets.size(); + thrust::fill(stream_policy, begin, end, static_cast(tangent_size)); + + if (const_state_ids != nullptr && num_const_state_blocks > 0) { + // Give constant blocks width zero so the scan skips over them, then stamp + // them with -1 once the offsets are in place. + auto zero_it = thrust::make_constant_iterator(0); + thrust::device_ptr const_ids_ptr(const_state_ids); + thrust::scatter(stream_policy, zero_it, zero_it + num_const_state_blocks, const_ids_ptr, begin); + thrust::exclusive_scan(stream_policy, begin, end, begin, first_column); + auto minus_one_it = thrust::make_constant_iterator(-1); + thrust::scatter(stream_policy, minus_one_it, minus_one_it + num_const_state_blocks, + const_ids_ptr, begin); + } else { + thrust::exclusive_scan(stream_policy, begin, end, begin, first_column); + } +} + +void StateBatchOps::InitUpdatesVector(const std::vector &state_batches) { delta_ptrs_.clear(); size_t updates_size = 0; @@ -154,8 +178,8 @@ void StateBatchOps::InitUpdatesVector( } /** @copydoc StateBatchOps::InitMapping */ -void StateBatchOps::InitMapping( - cudaStream_t stream, const std::vector &state_batches) { +void StateBatchOps::InitMapping(cudaStream_t stream, + const std::vector &state_batches) { thrust::device_vector temp_buffer; map_.resize(state_updates_.size()); thrust::device_ptr map_ptr(map_.data()); @@ -170,14 +194,13 @@ void StateBatchOps::InitMapping( size_t num_elements = batch->NumStateBlocks() * tangent_dim; temp_buffer.resize(num_elements); - fill_binary_pattern(stream, tangent_dim, const_state_ids, num_const_blocks, - temp_buffer); + fill_binary_pattern(stream, tangent_dim, const_state_ids, num_const_blocks, temp_buffer); thrust::counting_iterator iter(static_cast(map_it - map_ptr)); // Make map contain valid ids for non-const states, and INT_MAX for // constant states (so they sort to the end) - thrust::transform(stream_policy, iter, iter + num_elements, - temp_buffer.begin(), map_it, CustomBinaryFunctor()); + thrust::transform(stream_policy, iter, iter + num_elements, temp_buffer.begin(), map_it, + CustomBinaryFunctor()); num_reduced_states_ += num_elements - num_const_blocks * tangent_dim; map_it += num_elements; @@ -197,8 +220,7 @@ void StateBatchOps::Preprocess(cudaStream_t stream, } /** @copydoc StateBatchOps::Plus */ -void StateBatchOps::Plus(cudaStream_t stream, - const std::vector &x_ptrs, +void StateBatchOps::Plus(cudaStream_t stream, const std::vector &x_ptrs, const DeviceVector &delta, std::vector &x_plus_delta_ptrs) { assert(delta.size() == num_reduced_states_); @@ -208,14 +230,12 @@ void StateBatchOps::Plus(cudaStream_t stream, // Zero out the updates thrust::device_ptr updates_ptr(state_updates_.data()); - thrust::fill(stream_policy, updates_ptr, updates_ptr + state_updates_.size(), - 0.0f); + thrust::fill(stream_policy, updates_ptr, updates_ptr + state_updates_.size(), 0.0f); // Scatter the delta values across the update vector thrust::device_ptr delta_ptr(delta.data()); thrust::device_ptr map_ptr(map_.data()); - thrust::scatter(stream_policy, delta_ptr, delta_ptr + delta.size(), map_ptr, - updates_ptr); + thrust::scatter(stream_policy, delta_ptr, delta_ptr + delta.size(), map_ptr, updates_ptr); for (size_t i = 0; i < x_ptrs.size(); i++) { auto state_batch = user_state_batches_[i]; diff --git a/cunls/state/state_batch_ops.h b/cunls/state/state_batch_ops.h index bee262c..29a2f11 100644 --- a/cunls/state/state_batch_ops.h +++ b/cunls/state/state_batch_ops.h @@ -25,6 +25,24 @@ namespace cunls { +/** + * @brief Computes each state block's starting column in the reduced system. + * + * The reduced system omits constant state blocks, so a block's column offset is + * the running sum of the tangent sizes of the non-constant blocks before it. + * Constant blocks are marked with -1 so callers can distinguish "column 0" from + * "no column at all". + * + * @param stream CUDA stream for GPU operations. + * @param first_column Column index at which this batch's blocks start, i.e. the + * total tangent size of all preceding batches' non-constant blocks. + * @param state_batch The state batch. + * @param[out] column_offsets One entry per state block; resized as needed. + */ +void ComputeStateBlockColumnOffsets(cudaStream_t stream, int first_column, + const StateBatch *state_batch, + DeviceVector &column_offsets); + /** * @brief Orchestrates manifold Plus operations across multiple state batches. * @@ -44,8 +62,7 @@ class StateBatchOps { * preprocessing. * @param state_batches Vector of pointers to state batches to manage. */ - StateBatchOps(cudaStream_t stream, - const std::vector &state_batches); + StateBatchOps(cudaStream_t stream, const std::vector &state_batches); /** @brief Default constructor. Call Preprocess() before use. */ StateBatchOps() = default; @@ -60,8 +77,7 @@ class StateBatchOps { * @param stream CUDA stream for asynchronous GPU operations. * @param state_batches Vector of pointers to state batches to manage. */ - void Preprocess(cudaStream_t stream, - const std::vector &state_batches); + void Preprocess(cudaStream_t stream, const std::vector &state_batches); /** * @brief Applies manifold Plus operations across all state batches. @@ -79,8 +95,7 @@ class StateBatchOps { * one per state batch. */ void Plus(cudaStream_t stream, const std::vector &x_ptrs, - const DeviceVector &delta, - std::vector &x_plus_delta_ptrs); + const DeviceVector &delta, std::vector &x_plus_delta_ptrs); /** * @brief Returns the number of reduced (non-constant) states. @@ -90,7 +105,7 @@ class StateBatchOps { size_t NumReducedStates() const { return num_reduced_states_; } // Protected for testing -protected: + protected: /** @brief Device vector storing the mapping from reduced state indices to * full (including constant) state indices. */ DeviceVector map_; @@ -112,8 +127,7 @@ class StateBatchOps { * @param stream CUDA stream for asynchronous GPU operations. * @param state_batches Vector of state batches. */ - void InitMapping(cudaStream_t stream, - const std::vector &state_batches); + void InitMapping(cudaStream_t stream, const std::vector &state_batches); /** @brief Cached pointers to the user-supplied state batches. */ std::vector user_state_batches_; diff --git a/docs/sphinx/api/common.rst b/docs/sphinx/api/common.rst index d93fd49..a0e1b7d 100644 --- a/docs/sphinx/api/common.rst +++ b/docs/sphinx/api/common.rst @@ -27,12 +27,17 @@ Sparse matrix structs - `col_ids` - [in/out] CSR column indices. - `values` - [in/out] CSR non-zero values. - methods: `NumRows()`, `NumNonZeros()`. -- `TripletSparseStructure` - - `row_ids` - [in/out] triplet row indices. - - `col_ids` - [in/out] triplet column indices. -- `SparseJacobian` - - `structure` - [in/out] sparse structure. - - `values` - [in/out] sparse values. +- `BSRSparseMatrix` + - `row_offsets` - [in/out] block-row offsets. + - `col_ids` - [in/out] block-column index of each stored tile. + - `values` - [in/out] dense `block_size` x `block_size` tiles, row-major. + - `block_size` - [in/out] tile edge length. + - `num_block_rows` - [in/out] number of block rows. + - `max_tiles_per_row` - [in/out] longest block row; selects the SpMV schedule. + - methods: `NumBlocks()`, `NumRows()`, `NumNonZeros()`. +- `PerFactorJacobians` + - alias for `dvector`: per-factor dense Jacobian blocks, concatenated + across residual batches. There is no global sparse Jacobian. DeviceVector --------------- diff --git a/docs/sphinx/api/minimizer.rst b/docs/sphinx/api/minimizer.rst index 4ea9d4d..c0fae59 100644 --- a/docs/sphinx/api/minimizer.rst +++ b/docs/sphinx/api/minimizer.rst @@ -146,13 +146,9 @@ Used when constructing a :code:`GaussNewtonMinimizer`. ``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). - Default: ``Fast``. - **column_scaling** [in]: Diagonal scaling of the normal equations; see the - column-scaling note above. Values: ``None``, ``HessianDiagonal``, - ``JacobianColumnNorm``. Default: ``None``. + column-scaling note above. Values: ``None``, ``HessianDiagonal``. + Default: ``None``. - **disable_safety_checks** [in]: When ``false``, the minimizer enables all optional runtime validation. Currently this covers post-factorization checks in the linear solver: Cholesky checks cuSOLVER ``devInfo`` after @@ -360,86 +356,49 @@ problem’s state storage. Useful for rollback or warm starts. :param ``problem``: [out] Problem whose state storage is overwritten with the copied values. :returns: [out] No return value. -.. _minimizer-state-build-triplet-label: +.. _hessian-assembly-label: -------------------------------------------------------------------------------- -:code:`MinimizerState::BuildTripletSparseStructure` +Hessian assembly -------------------------------------------------------------------------------- -**Purpose:** Fills the row and column index arrays of a :code:`TripletSparseStructure` -for the problem’s Jacobian (COO / triplet layout). Implementation lives in -``jacobian_ops.cu``; the minimizer uploads the problem’s host-held state-pointer -lists to device internally before building column indices. +The normal equations are assembled directly from the per-factor Jacobian blocks +that each factor batch writes; no global sparse Jacobian is ever materialized. -.. cpp:function:: void BuildTripletSparseStructure(cudaStream_t stream, const Problem& problem, TripletSparseStructure& structure) +:code:`HessianStructureBuilder` (``cunls/minimizer/hessian_structure.h``) derives +the sparsity pattern from factor-graph connectivity on the GPU: it resolves each +factor's state pointers to global columns, packs every candidate block pair into +a 64-bit key, then sorts and segments. It also returns, for each +``(factor, block_a, block_b)`` slot, the row-relative offset at which that tile +starts — the map the assembler scatters through. - :param ``stream``: [in] CUDA stream for GPU work. - :param ``problem``: [in] Factor graph and state-pointer mappings. - :param ``structure``: [out] Row and column index device buffers sized to the Jacobian nonzeros. - :returns: [out] No return value. - -.. _sparse-square-multiplier-label: +:code:`BlockHessianAssembler` (``cunls/minimizer/block_hessian_assembler.h``) +runs one kernel per residual batch, one warp per factor. Each warp stages +:math:`J_f` and :math:`r_f` in shared memory, forms +:math:`H_f = J_f^T J_f` and :math:`b_f = -J_f^T r_f`, and scatter-adds both into +the global system. -------------------------------------------------------------------------------- -:code:`SparseMatrixMultiplier` (base class) +Hessian storage -------------------------------------------------------------------------------- -Abstract interface for computing :math:`A^T A` of a sparse CSR matrix. -Two concrete implementations are provided: - -- **cuSPARSESparseMatrixMultiplier** — uses the cuSPARSE SpGEMM reuse API - (transpose + multiply). Select with - :code:`SparseMatrixMultiplierType::cuSPARSE`. -- **FastSparseMatrixMultiplier** — uses custom warp-efficient CUDA - kernels with bitmap-based sparsity pattern discovery. Select with - :code:`SparseMatrixMultiplierType::Fast`. - -.. cpp:function:: void SparseMatrixMultiplier::Initialize(cudaStream_t stream, const Problem& problem, const CSRSparseMatrix& input, CSRSparseMatrix& output) - - Analyzes the sparsity pattern of :math:`A^T A` and allocates the output - matrix. Must be called once whenever the sparsity pattern changes. - - :param ``stream``: [in] CUDA stream for GPU operations. - :param ``problem``: [in] Optimization problem providing structural hints. - :param ``input``: [in] Input sparse matrix :math:`A`. - :param ``output``: [out] Output sparse matrix :math:`A^T A` (structure allocated). - :returns: [out] No return value. - -.. cpp:function:: void SparseMatrixMultiplier::ComputeSquaredMatrix(cudaStream_t stream, const Problem& problem, const CSRSparseMatrix& input, CSRSparseMatrix& output) - - Computes the numerical values of :math:`A^T A`. The output structure must - already be set by a prior call to :cpp:func:`Initialize`. - - :param ``stream``: [in] CUDA stream for GPU operations. - :param ``problem``: [in] Optimization problem (may be used for kernel tuning). - :param ``input``: [in] Input sparse matrix :math:`A`. - :param ``output``: [out] Output sparse matrix :math:`A^T A`. - :returns: [out] No return value. +The Hessian of a factor graph is block structured, so it is stored as BSR — one +column index per dense tile instead of one per scalar entry. That is bandwidth +the iterative solver's SpMV no longer has to move. -.. _sparse-square-multiplier-type-label: - --------------------------------------------------------------------------------- -:code:`SparseMatrixMultiplierType` --------------------------------------------------------------------------------- +Both the Hessian and the working left-hand side are owned by +:code:`NormalEquations` (``cunls/minimizer/normal_equations.h``), which is the +only place either layout is named; the minimizers work in terms of "the Hessian" +and "the left-hand side" and never branch on storage. -Enum in ``cunls/minimizer/sparse_matrix_multiplier.h``: - -- ``cuSPARSE`` — cuSPARSE SpGEMM reuse API (transpose + multiply). -- ``Fast`` — fast warp-efficient CUDA kernels with bitmap pattern - discovery. - -.. _sparse-square-multiplier-factory-label: - --------------------------------------------------------------------------------- -:code:`CreateSparseMatrixMultiplier` --------------------------------------------------------------------------------- - -Factory function in ``cunls/minimizer/sparse_matrix_multiplier.h``: - -.. cpp:function:: SparseMatrixMultiplierPtr CreateSparseMatrixMultiplier(SparseMatrixMultiplierType type) - - :param ``type``: [in] Strategy to use for computing :math:`A^T A`. - :returns: [out] Heap-allocated multiplier instance. +The layout is chosen automatically, with no user-facing switch. Block storage +requires a tile edge dividing every state block's tangent dimension (the gcd of +the tangent sizes; see :code:`ChooseHessianBlockSize` in +``cunls/minimizer/bsr_matrix.h``) **and** a solver that reports +:code:`CSRSparseLinearSolver::SupportsBlockStorage`. When either does not hold — +a gcd of one, or a backend such as cuDSS or the dense factorizations that needs +CSR anyway — the minimizer falls back to scalar CSR with no behavioural change. +No conversion is ever performed on the solver path. ================================================================================ Python API (``pycunls``) @@ -495,10 +454,6 @@ values and then override individual fields. ``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; - ``Fast`` uses warp-efficient CUDA kernels with bitmap pattern discovery. - **column_scaling** (``ColumnScaling``, default ``ColumnScaling.none``) — optional diagonal scaling :math:`S` for the normal equations (:math:`S H S\, z = S b`, then :math:`\Delta x = S z`). See @@ -743,17 +698,6 @@ Integer enum selecting the linear-system backend. - ``SparseLinearSolverType.DenseQR`` — converts CSR to dense and solves via cuSOLVER QR factorization (works for any non-singular matrix). --------------------------------------------------------------------------------- -``pycunls.SparseMatrixMultiplierType`` --------------------------------------------------------------------------------- - -Integer enum selecting the strategy for computing the approximate Hessian -:math:`J^T J`. - -- ``SparseMatrixMultiplierType.cuSPARSE`` — cuSPARSE SpGEMM reuse API. -- ``SparseMatrixMultiplierType.Fast`` — warp-efficient CUDA kernels with - bitmap pattern discovery. - .. _py-minimizer-example-label: -------------------------------------------------------------------------------- diff --git a/python/pycunls/__init__.py b/python/pycunls/__init__.py index 5e1451e..e963de5 100644 --- a/python/pycunls/__init__.py +++ b/python/pycunls/__init__.py @@ -56,7 +56,6 @@ CublasHandle, # --- Enumerations --- SparseLinearSolverType, - SparseMatrixMultiplierType, ColumnScaling, # --- Minimizer options & summary --- MinimizerOptions, @@ -129,7 +128,6 @@ "CudaStream", "CublasHandle", "SparseLinearSolverType", - "SparseMatrixMultiplierType", "ColumnScaling", "MinimizerOptions", "MinimizerSummary", diff --git a/python/pycunls/_pycunls_core.pyi b/python/pycunls/_pycunls_core.pyi index 35f598a..0e85191 100644 --- a/python/pycunls/_pycunls_core.pyi +++ b/python/pycunls/_pycunls_core.pyi @@ -55,16 +55,11 @@ class SparseLinearSolverType(enum.IntEnum): DenseQR = ... BlockSparsePCG = ... -class SparseMatrixMultiplierType(enum.IntEnum): - cuSPARSE = ... - Fast = ... - class ColumnScaling(enum.IntEnum): """Diagonal scaling mode for the GN/LM normal equations.""" none = ... hessian_diagonal = ... - jacobian_column_norm = ... # =================================================================== # Options and summary @@ -78,7 +73,6 @@ class MinimizerOptions: cost_tolerance: float max_consecutive_rejected_steps: int sparse_linear_solver_type: SparseLinearSolverType - sparse_square_multiplier_type: SparseMatrixMultiplierType column_scaling: ColumnScaling disable_safety_checks: bool diff --git a/python/src/bind_types.cpp b/python/src/bind_types.cpp index 6edf47e..d0b93d9 100644 --- a/python/src/bind_types.cpp +++ b/python/src/bind_types.cpp @@ -73,14 +73,9 @@ void bind_types(nb::module_ &m) { .value("DenseQR", cunls::SparseLinearSolverType::DenseQR) .value("BlockSparsePCG", cunls::SparseLinearSolverType::BlockSparsePCG); - nb::enum_(m, "SparseMatrixMultiplierType") - .value("cuSPARSE", cunls::SparseMatrixMultiplierType::cuSPARSE) - .value("Fast", cunls::SparseMatrixMultiplierType::Fast); - nb::enum_(m, "ColumnScaling") .value("none", cunls::ColumnScaling::None) - .value("hessian_diagonal", cunls::ColumnScaling::HessianDiagonal) - .value("jacobian_column_norm", cunls::ColumnScaling::JacobianColumnNorm); + .value("hessian_diagonal", cunls::ColumnScaling::HessianDiagonal); // --- Minimizer configuration structs --- // All fields are read/write so users can tune convergence behaviour @@ -98,8 +93,6 @@ void bind_types(nb::module_ &m) { &cunls::MinimizerOptions::max_consecutive_rejected_steps) .def_rw("sparse_linear_solver_type", &cunls::MinimizerOptions::sparse_linear_solver_type) - .def_rw("sparse_square_multiplier_type", - &cunls::MinimizerOptions::sparse_square_multiplier_type) .def_rw("column_scaling", &cunls::MinimizerOptions::column_scaling) .def_rw("disable_safety_checks", &cunls::MinimizerOptions::disable_safety_checks, diff --git a/python/tests/test_minimizer.py b/python/tests/test_minimizer.py index 96c1477..edcf529 100644 --- a/python/tests/test_minimizer.py +++ b/python/tests/test_minimizer.py @@ -106,10 +106,7 @@ def test_converges(self, stream): @pytest.mark.parametrize( "scaling", - [ - pycunls.ColumnScaling.hessian_diagonal, - pycunls.ColumnScaling.jacobian_column_norm, - ], + [pycunls.ColumnScaling.hessian_diagonal], ) def test_converges_with_column_scaling(self, stream, scaling): problem, states_gpu, target = _make_prior_problem() diff --git a/tests/block_hessian_assembler_test.cpp b/tests/block_hessian_assembler_test.cpp new file mode 100644 index 0000000..41e9067 --- /dev/null +++ b/tests/block_hessian_assembler_test.cpp @@ -0,0 +1,1288 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * @file block_hessian_assembler_test.cpp + * @brief Correctness and performance tests for Hessian assembly and storage. + * + * Three suites: + * - HessianStructureTest checks the GPU-derived sparsity pattern against a + * naive CPU oracle, + * - BlockHessianAssemblerTest checks the assembled H and rhs against a CPU + * oracle that contracts the same per-factor Jacobians, + * - HessianStorageTest checks scalar CSR and block BSR against each other. + * + * The oracles share no code with the GPU paths they validate, which matters: + * both storage layouts run through one assembler, so comparing them against + * each other cannot catch an error common to both. + */ + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "cunls/common/cublas_helper.h" +#include "cunls/common/cuda_stream.h" +#include "cunls/common/device_vector.h" +#include "cunls/common/helper.h" +#include "cunls/common/types.h" +#include "cunls/factor/prior_vector_factor_batch.h" +#include "cunls/factor/reprojection_factor_batch.h" +#include "cunls/factor/se3_between_factor_batch.h" +#include "cunls/factor/vector_between_factor_batch.h" +#include "cunls/math/so_se_lie_math.h" +#include "cunls/minimizer/bsr_matrix.h" +#include "cunls/minimizer/device_reduction.h" +#include "cunls/minimizer/gauss_newton_minimizer.h" +#include "cunls/minimizer/normal_equations.h" +#include "cunls/minimizer/problem.h" +#include "cunls/minimizer/sparse_matrix.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 { +namespace { + +// ============================================================================ +// Harness +// ============================================================================ + +/** + * @brief Exposes one BuildSystem call and its outputs. + * + * BuildSystem is protected on GaussNewtonMinimizer, and hessian_/lhs_work_ + * live alongside it, so a thin subclass is the least invasive way to compare + * the two assembly paths on identical input. + */ +class SystemBuilder : public GaussNewtonMinimizer { +public: + /** + * @brief Builds with block storage (the default for a block-capable solver). + */ + SystemBuilder() : GaussNewtonMinimizer(MakeOptions(SparseLinearSolverType::BlockSparsePCG)) {} + + explicit SystemBuilder(const MinimizerOptions &options) : GaussNewtonMinimizer(options) {} + + /** + * @brief Builds with scalar CSR storage. + * + * There is no switch for this: storage is chosen from the problem's tangent + * dimensions and the solver's capability. Selecting a CSR-only backend is + * how a caller actually ends up on the scalar path, so that is what the + * comparisons here exercise. Only assembly is compared, never the solve, so + * the backend choice does not otherwise affect the result. + */ + static SystemBuilder Scalar() { + return SystemBuilder(MakeOptions(SparseLinearSolverType::cuDSS)); + } + + /** @brief Runs Initialize + one BuildSystem and syncs. */ + void Build(cudaStream_t stream, Problem &problem) { + Initialize(stream, problem); + current_state_.Recreate(stream, problem); + BuildSystem(stream, problem, current_state_); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + } + + const CSRSparseMatrix &HessianAsCSR(cudaStream_t stream) { + if (!normal_equations_.UsesBlockStorage()) { + return normal_equations_.LhsCSR(); + } + ConvertBSRToCSR(stream, normal_equations_.LhsBSR(), csr_mirror_, expand_scratch_); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + return csr_mirror_; + } + + const dvector &Rhs() const { return rhs_work_; } + const dvector &Residuals() const { return residuals_; } + const PerFactorJacobians &FactorJacobians() const { return factor_jacobians_; } + /** @brief step^T H step against the undamped Hessian, as LM computes it. */ + float WeightedSquaredStep(cudaStream_t stream, const dvector &step) { + d_scalars_.resize(1); + d_reduce_partials_.resize(ReducePartialCount(step.size())); + normal_equations_.WeightedSquaredStepAsync(stream, cusparse_handle_.GetHandle(stream), step, + d_scalars_.data(), d_reduce_partials_.data(), + buffer_); + float out = 0.f; + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&out, d_scalars_.data(), sizeof(float), cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + return out; + } + NormalEquations &Equations() { return normal_equations_; } + bool UsesBlockStorage() const { return normal_equations_.UsesBlockStorage(); } + +private: + static MinimizerOptions MakeOptions(SparseLinearSolverType solver) { + MinimizerOptions options; + options.sparse_linear_solver_type = solver; + return options; + } + + CSRSparseMatrix csr_mirror_; + dvector expand_scratch_; +}; + +/** @brief Host-side snapshot of an assembled system. */ +struct SystemSnapshot { + std::vector row_offsets; + std::vector col_ids; + std::vector values; + std::vector rhs; +}; + +SystemSnapshot Snapshot(SystemBuilder &builder, cudaStream_t stream) { + const CSRSparseMatrix &h = builder.HessianAsCSR(stream); + SystemSnapshot out; + out.row_offsets.resize(h.row_offsets.size()); + out.col_ids.resize(h.col_ids.size()); + out.values.resize(h.values.size()); + out.rhs.resize(builder.Rhs().size()); + h.row_offsets.CopyToHost(out.row_offsets.data(), out.row_offsets.size()); + h.col_ids.CopyToHost(out.col_ids.data(), out.col_ids.size()); + h.values.CopyToHost(out.values.data(), out.values.size()); + builder.Rhs().CopyToHost(out.rhs.data(), out.rhs.size()); + return out; +} + +/** @brief Largest magnitude in a vector, used to set a relative tolerance. */ +float MaxAbs(const std::vector &v) { + float m = 0.f; + for (float x : v) { + m = std::max(m, std::fabs(x)); + } + return m; +} + +// ============================================================================ +// Problem builders +// ============================================================================ + +/** @brief Owns the device data behind a Problem for the lifetime of a test. */ +struct VectorChainProblem { + static constexpr int kDim = 4; + + std::vector host_states; + std::vector> host_priors; + dvector states; + dvector> deltas; + dvector> priors; + dvector const_ids; + std::unique_ptr> state_batch; + std::unique_ptr> between_batch; + std::unique_ptr> prior_batch; + std::vector between_pointers; + std::vector prior_pointers; + Problem problem; +}; + +/** + * @brief Chain of `num_blocks` vector states with between + prior factors. + * + * @param num_blocks Number of state blocks. + * @param constant_indices State blocks held constant (may be empty). + * @param repeat_block When true, every between factor references the same + * state block on both sides, exercising the accumulate-twice path. + * @param stride Gap between the two blocks a between factor joins; larger + * values keep the block and factor counts but change the connectivity. + */ +std::unique_ptr MakeVectorChain(int num_blocks, + const std::vector &constant_indices, + bool repeat_block = false, int stride = 1) { + constexpr int kDim = VectorChainProblem::kDim; + auto data = std::make_unique(); + std::mt19937 rng(7); + std::uniform_real_distribution dist(-1.f, 1.f); + + std::vector &host_states = data->host_states; + host_states.resize(num_blocks * kDim); + for (float &x : host_states) { + x = dist(rng); + } + data->states.resize(host_states.size()); + data->states.CopyFromHost(host_states.data(), host_states.size()); + + const int num_between = num_blocks - 1; + std::vector> host_deltas(num_between); + for (auto &d : host_deltas) { + for (int k = 0; k < kDim; k++) { + d[k] = dist(rng); + } + } + data->deltas.resize(host_deltas.size()); + data->deltas.CopyFromHost(host_deltas.data(), host_deltas.size()); + + std::vector> &host_priors = data->host_priors; + host_priors.resize(num_blocks); + for (auto &p : host_priors) { + for (int k = 0; k < kDim; k++) { + p[k] = dist(rng); + } + } + data->priors.resize(host_priors.size()); + data->priors.CopyFromHost(host_priors.data(), host_priors.size()); + + if (constant_indices.empty()) { + data->state_batch = std::make_unique>(data->states.data(), num_blocks); + } else { + data->const_ids.resize(constant_indices.size()); + data->const_ids.CopyFromHost(constant_indices.data(), constant_indices.size()); + data->state_batch = std::make_unique>( + data->states.data(), num_blocks, data->const_ids.data(), constant_indices.size()); + } + + data->between_batch = + std::make_unique>(data->deltas.data(), num_between); + data->prior_batch = + std::make_unique>(data->priors.data(), num_blocks); + + for (int i = 0; i < num_between; i++) { + float *left = data->states.data() + static_cast(i) * kDim; + const int right_index = repeat_block ? i : (i + stride) % num_blocks; + float *right = data->states.data() + static_cast(right_index) * kDim; + data->between_pointers.push_back(left); + data->between_pointers.push_back(right); + } + for (int i = 0; i < num_blocks; i++) { + data->prior_pointers.push_back(data->states.data() + static_cast(i) * kDim); + } + + data->problem.AddStateBatch(data->state_batch.get()); + data->problem.AddFactorBatch(data->between_batch.get(), data->between_pointers); + data->problem.AddFactorBatch(data->prior_batch.get(), data->prior_pointers); + return data; +} + +/** @brief Owns the device data behind a mixed-tangent-dimension Problem. */ +struct MixedDimProblem { + dvector states_a; + dvector states_b; + dvector> priors_a; + dvector> priors_b; + std::unique_ptr> batch_a; + std::unique_ptr> batch_b; + std::unique_ptr> prior_a; + std::unique_ptr> prior_b; + std::vector pointers_a; + std::vector pointers_b; + Problem problem; +}; + +/** @brief Two state batches with different tangent dims in one problem. */ +std::unique_ptr MakeMixedDim(int count) { + auto data = std::make_unique(); + std::mt19937 rng(11); + std::uniform_real_distribution dist(-1.f, 1.f); + + std::vector host_a(count * 3); + std::vector host_b(count * 6); + for (float &x : host_a) { + x = dist(rng); + } + for (float &x : host_b) { + x = dist(rng); + } + data->states_a.resize(host_a.size()); + data->states_a.CopyFromHost(host_a.data(), host_a.size()); + data->states_b.resize(host_b.size()); + data->states_b.CopyFromHost(host_b.data(), host_b.size()); + + std::vector> pa(count); + std::vector> pb(count); + for (int i = 0; i < count; i++) { + for (int k = 0; k < 3; k++) { + pa[i][k] = dist(rng); + } + for (int k = 0; k < 6; k++) { + pb[i][k] = dist(rng); + } + } + data->priors_a.resize(pa.size()); + data->priors_a.CopyFromHost(pa.data(), pa.size()); + data->priors_b.resize(pb.size()); + data->priors_b.CopyFromHost(pb.data(), pb.size()); + + data->batch_a = std::make_unique>(data->states_a.data(), count); + data->batch_b = std::make_unique>(data->states_b.data(), count); + data->prior_a = std::make_unique>(data->priors_a.data(), count); + data->prior_b = std::make_unique>(data->priors_b.data(), count); + + for (int i = 0; i < count; i++) { + data->pointers_a.push_back(data->states_a.data() + static_cast(i) * 3); + data->pointers_b.push_back(data->states_b.data() + static_cast(i) * 6); + } + + data->problem.AddStateBatch(data->batch_a.get()); + data->problem.AddStateBatch(data->batch_b.get()); + data->problem.AddFactorBatch(data->prior_a.get(), data->pointers_a); + data->problem.AddFactorBatch(data->prior_b.get(), data->pointers_b); + return data; +} + +/** @brief Owns the device data behind an SE3 pose-graph Problem. */ +struct PoseGraphProblem { + cuBLASHandle cublas_handle; + dvector poses; + dvector deltas; + dvector const_ids; + std::unique_ptr pose_batch; + std::unique_ptr between_batch; + std::vector pointers; + Problem problem; +}; + +/** @brief Random SE3 poses joined by consecutive between factors. */ +std::unique_ptr MakePoseGraph(int num_poses, bool fix_first_pose) { + auto data = std::make_unique(); + CudaStream stream; + std::mt19937 rng(23); + std::uniform_real_distribution rot(-0.4f, 0.4f); + std::uniform_real_distribution trans(-2.f, 2.f); + + auto random_transforms = [&](int count, dvector &out) { + std::vector> twists(count); + for (auto &t : twists) { + for (int k = 0; k < 3; k++) { + t[k] = rot(rng); + } + for (int k = 3; k < 6; k++) { + t[k] = trans(rng); + } + } + dvector> twists_device(count); + twists_device.CopyFromHost(twists.data(), twists.size()); + out.resize(count); + ComputeExpSE3(stream.GetStream(), reinterpret_cast(twists_device.data()), 6, 4, + 16, count, reinterpret_cast(out.data())); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + }; + + random_transforms(num_poses, data->poses); + random_transforms(num_poses - 1, data->deltas); + + if (fix_first_pose) { + std::vector ids{0}; + data->const_ids.resize(1); + data->const_ids.CopyFromHost(ids.data(), 1); + data->pose_batch = std::make_unique( + data->cublas_handle, reinterpret_cast(data->poses.data()), num_poses, + data->const_ids.data(), 1); + } else { + data->pose_batch = std::make_unique( + data->cublas_handle, reinterpret_cast(data->poses.data()), num_poses); + } + + data->between_batch = std::make_unique(data->deltas.data(), num_poses - 1); + + auto *base = reinterpret_cast(data->poses.data()); + for (int i = 0; i + 1 < num_poses; i++) { + data->pointers.push_back(base + static_cast(i) * 16); + data->pointers.push_back(base + static_cast(i + 1) * 16); + } + + data->problem.AddStateBatch(data->pose_batch.get()); + data->problem.AddFactorBatch(data->between_batch.get(), data->pointers); + return data; +} + +/** @brief Owns the device data behind a small bundle-adjustment Problem. */ +struct BundleProblem { + cuBLASHandle cublas_handle; + dvector poses; + dvector points; + dvector> observations; + std::unique_ptr pose_batch; + std::unique_ptr> point_batch; + std::unique_ptr reprojection_batch; + std::unique_ptr huber; + std::vector pointers; + Problem problem; +}; + +/** + * @brief Reprojection factors over 6-dof poses and 3-dof points. + * + * Mixes tangent dims *inside* a single factor (6 + 3), which the vector + * problems above cannot exercise. + */ +std::unique_ptr MakeBundle(int num_poses, int num_points, bool robust_loss) { + auto data = std::make_unique(); + CudaStream stream; + std::mt19937 rng(31); + std::uniform_real_distribution small(-0.15f, 0.15f); + std::uniform_real_distribution obs_noise(-0.02f, 0.02f); + + std::vector> twists(num_poses); + for (auto &t : twists) { + for (int k = 0; k < 6; k++) { + t[k] = small(rng); + } + } + dvector> twists_device(num_poses); + twists_device.CopyFromHost(twists.data(), twists.size()); + data->poses.resize(num_poses); + ComputeExpSE3(stream.GetStream(), reinterpret_cast(twists_device.data()), 6, 4, 16, + num_poses, reinterpret_cast(data->poses.data())); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + + // Points sit well in front of the cameras so every projection is valid. + std::vector host_points(num_points * 3); + std::uniform_real_distribution lateral(-1.f, 1.f); + std::uniform_real_distribution depth(4.f, 8.f); + for (int i = 0; i < num_points; i++) { + host_points[i * 3 + 0] = lateral(rng); + host_points[i * 3 + 1] = lateral(rng); + host_points[i * 3 + 2] = depth(rng); + } + data->points.resize(host_points.size()); + data->points.CopyFromHost(host_points.data(), host_points.size()); + + std::vector> host_obs; + auto *pose_base = reinterpret_cast(data->poses.data()); + for (int p = 0; p < num_points; p++) { + for (int c = 0; c < num_poses; c++) { + float x = host_points[p * 3 + 0]; + float y = host_points[p * 3 + 1]; + float z = host_points[p * 3 + 2]; + host_obs.push_back({x / z + obs_noise(rng), y / z + obs_noise(rng)}); + data->pointers.push_back(pose_base + static_cast(c) * 16); + data->pointers.push_back(data->points.data() + static_cast(p) * 3); + } + } + data->observations.resize(host_obs.size()); + data->observations.CopyFromHost(host_obs.data(), host_obs.size()); + + data->pose_batch = std::make_unique( + data->cublas_handle, reinterpret_cast(data->poses.data()), num_poses); + data->point_batch = std::make_unique>(data->points.data(), num_points); + data->reprojection_batch = + std::make_unique(data->observations.data(), host_obs.size()); + + data->problem.AddStateBatch(data->pose_batch.get()); + data->problem.AddStateBatch(data->point_batch.get()); + if (robust_loss) { + data->huber = std::make_unique(0.02f); + data->problem.AddFactorBatch(data->reprojection_batch.get(), data->huber.get(), data->pointers); + } else { + data->problem.AddFactorBatch(data->reprojection_batch.get(), data->pointers); + } + return data; +} + +// ============================================================================ +// Structure oracle +// ============================================================================ + +/** @brief Host-computed CSR sparsity pattern. */ +struct ReferenceStructure { + std::vector row_offsets; + std::vector col_ids; + /// Per residual batch: the global column of each (factor, block) slot, -1 + /// when the block is a constant state. + std::vector> factor_columns; + int num_cols = 0; +}; + +/** + * @brief Independent CPU reference for the Hessian sparsity pattern. + * + * Both assembly paths now share HessianStructureBuilder, so comparing them + * against each other can no longer catch a structure bug. This is a + * deliberately naive transcription of the definition — resolve every factor's + * state pointers to global columns, collect the distinct block pairs in an + * ordered set, expand each tile — so it shares no code with the GPU sort and + * segmentation it validates. + */ +ReferenceStructure BuildReferenceStructure(const Problem &problem) { + struct BatchDesc { + const float *base = nullptr; + int ambient = 0; + int tangent = 0; + int num_blocks = 0; + std::vector col_of_block; + }; + + std::vector descs; + int last_col = 0; + for (const auto *state_batch : problem.GetStateBatches()) { + BatchDesc d; + d.base = state_batch->StateBlockDevicePtr(0); + d.ambient = static_cast(state_batch->AmbientSize()); + d.tangent = static_cast(state_batch->TangentSize()); + d.num_blocks = static_cast(state_batch->NumStateBlocks()); + + std::vector const_ids(state_batch->NumConstStateBlocks()); + if (!const_ids.empty()) { + THROW_ON_CUDA_ERROR(cudaMemcpy(const_ids.data(), state_batch->ConstStateIds(), + const_ids.size() * sizeof(int), cudaMemcpyDeviceToHost)); + } + std::vector is_const(d.num_blocks, false); + for (int id : const_ids) { + if (id >= 0 && id < d.num_blocks) { + is_const[id] = true; + } + } + + d.col_of_block.assign(d.num_blocks, -1); + for (int j = 0; j < d.num_blocks; j++) { + if (!is_const[j]) { + d.col_of_block[j] = last_col; + last_col += d.tangent; + } + } + descs.push_back(std::move(d)); + } + + ReferenceStructure out; + out.num_cols = last_col; + + std::vector tangent_at_col(last_col, 0); + for (const BatchDesc &d : descs) { + for (int j = 0; j < d.num_blocks; j++) { + if (d.col_of_block[j] >= 0) { + for (int k = 0; k < d.tangent; k++) { + tangent_at_col[d.col_of_block[j] + k] = d.tangent; + } + } + } + } + + // Mirrors the tangent-dim guard the resolve kernel applies. + auto resolve = [&](const float *ptr, int declared_tangent) -> int { + for (const BatchDesc &d : descs) { + ptrdiff_t diff = ptr - d.base; + if (diff >= 0 && diff < static_cast(d.num_blocks) * d.ambient) { + if (d.tangent != declared_tangent) { + return -1; + } + return d.col_of_block[diff / d.ambient]; + } + } + return -1; + }; + + std::set> pairs; + const auto &residual_batches = problem.GetResidualBatches(); + const auto &state_pointers = problem.GetStatePointers(); + out.factor_columns.resize(residual_batches.size()); + for (size_t i = 0; i < residual_batches.size(); i++) { + const auto *factor_batch = residual_batches[i].GetFactorBatch(); + auto block_sizes = factor_batch->StateBlockSizes(); + const size_t nb = block_sizes.size(); + for (size_t f = 0; f < factor_batch->NumFactors(); f++) { + std::vector cols; + for (size_t b = 0; b < nb; b++) { + int col = resolve(state_pointers[i][f * nb + b], static_cast(block_sizes[b])); + out.factor_columns[i].push_back(col); + if (col >= 0) { + cols.push_back(col); + } + } + for (int a : cols) { + for (int b : cols) { + pairs.emplace(a, b); + } + } + } + } + + // Expand: within a block row the tiles sit side by side in column order. + std::vector row_counts(last_col, 0); + std::vector write_offsets; + write_offsets.reserve(pairs.size()); + int prev_row = -1; + int cursor = 0; + for (const auto &p : pairs) { + if (p.first != prev_row) { + cursor = 0; + prev_row = p.first; + } + write_offsets.push_back(cursor); + cursor += tangent_at_col[p.second]; + for (int i = 0; i < tangent_at_col[p.first]; i++) { + row_counts[p.first + i] += tangent_at_col[p.second]; + } + } + + out.row_offsets.assign(last_col + 1, 0); + for (int r = 0; r < last_col; r++) { + out.row_offsets[r + 1] = out.row_offsets[r] + row_counts[r]; + } + out.col_ids.assign(out.row_offsets.back(), -1); + + size_t k = 0; + for (const auto &p : pairs) { + const int height = tangent_at_col[p.first]; + const int width = tangent_at_col[p.second]; + for (int i = 0; i < height; i++) { + for (int j = 0; j < width; j++) { + out.col_ids[out.row_offsets[p.first + i] + write_offsets[k] + j] = p.second + j; + } + } + k++; + } + return out; +} + +/** + * @brief Host reference for the assembled normal equations. + * + * Contracts the very same per-factor Jacobian blocks the GPU kernel reads, but + * with plain nested loops into a dense-per-row map, so it shares no addressing + * or accumulation logic with the code it checks. Constant state blocks are + * dropped exactly as the kernel drops them. + * + * @param problem The optimization problem. + * @param jacobians Per-factor dense Jacobian blocks, read back from the device. + * @param residuals Residual vector, read back from the device. + * @param structure Sparsity pattern the values must be laid out against. + * @param[out] values CSR values of H. + * @param[out] rhs Right-hand side -J^T r. + */ +void ComputeReferenceSystem(const Problem &problem, const std::vector &jacobians, + const std::vector &residuals, + const ReferenceStructure &structure, std::vector &values, + std::vector &rhs) { + values.assign(structure.col_ids.size(), 0.f); + rhs.assign(structure.num_cols, 0.f); + + // (row, col) -> index into values, built once from the reference pattern. + std::map, size_t> slot; + for (int row = 0; row + 1 < static_cast(structure.row_offsets.size()); row++) { + for (int k = structure.row_offsets[row]; k < structure.row_offsets[row + 1]; k++) { + slot[{row, structure.col_ids[k]}] = static_cast(k); + } + } + + size_t jacobian_cursor = 0; + size_t residual_cursor = 0; + const auto &residual_batches = problem.GetResidualBatches(); + for (size_t i = 0; i < residual_batches.size(); i++) { + const auto *factor_batch = residual_batches[i].GetFactorBatch(); + const auto block_sizes = factor_batch->StateBlockSizes(); + const size_t num_blocks = block_sizes.size(); + const size_t residual_dim = factor_batch->ResidualsSize(); + const size_t tangent_dim = std::accumulate(block_sizes.begin(), block_sizes.end(), size_t(0)); + + const std::vector &columns = structure.factor_columns[i]; + for (size_t f = 0; f < factor_batch->NumFactors(); f++) { + const float *J = jacobians.data() + jacobian_cursor + f * residual_dim * tangent_dim; + const float *r = residuals.data() + residual_cursor + f * residual_dim; + + // Local tangent index -> global column, or -1 for a constant block. + std::vector global(tangent_dim, -1); + size_t cursor = 0; + for (size_t b = 0; b < num_blocks; b++) { + const int base = columns[f * num_blocks + b]; + for (size_t k = 0; k < block_sizes[b]; k++, cursor++) { + global[cursor] = base < 0 ? -1 : base + static_cast(k); + } + } + + for (size_t p = 0; p < tangent_dim; p++) { + if (global[p] < 0) { + continue; + } + double b_acc = 0.0; + for (size_t k = 0; k < residual_dim; k++) { + b_acc += static_cast(J[k * tangent_dim + p]) * r[k]; + } + rhs[global[p]] -= static_cast(b_acc); + + for (size_t q = 0; q < tangent_dim; q++) { + if (global[q] < 0) { + continue; + } + double h_acc = 0.0; + for (size_t k = 0; k < residual_dim; k++) { + h_acc += static_cast(J[k * tangent_dim + p]) * J[k * tangent_dim + q]; + } + values[slot.at({global[p], global[q]})] += static_cast(h_acc); + } + } + } + jacobian_cursor += factor_batch->NumFactors() * residual_dim * tangent_dim; + residual_cursor += factor_batch->NumFactors() * residual_dim; + } +} + +/** + * @brief Asserts the assembled H and rhs match the CPU oracle. + * + * Structure is compared exactly; values to `rel_tol` relative to the largest + * entry, which is the right scale here because both sides sum the same products + * in a different order, so the error is bounded by the largest partial sum + * rather than by each entry's own magnitude. + */ +void ExpectMatchesReference(Problem &problem, float rel_tol = 1e-5f) { + CudaStream stream; + SystemBuilder builder; + builder.Build(stream.GetStream(), problem); + SystemSnapshot actual = Snapshot(builder, stream.GetStream()); + + std::vector jacobians(builder.FactorJacobians().size()); + std::vector residuals(builder.Residuals().size()); + builder.FactorJacobians().CopyToHost(jacobians.data(), jacobians.size()); + builder.Residuals().CopyToHost(residuals.data(), residuals.size()); + + ReferenceStructure structure = BuildReferenceStructure(problem); + std::vector expected_values; + std::vector expected_rhs; + ComputeReferenceSystem(problem, jacobians, residuals, structure, expected_values, expected_rhs); + + ASSERT_EQ(structure.row_offsets, actual.row_offsets); + ASSERT_EQ(structure.col_ids, actual.col_ids); + ASSERT_EQ(expected_values.size(), actual.values.size()); + ASSERT_EQ(expected_rhs.size(), actual.rhs.size()); + + const float h_tol = rel_tol * std::max(MaxAbs(expected_values), 1e-6f); + for (size_t i = 0; i < expected_values.size(); i++) { + ASSERT_NEAR(expected_values[i], actual.values[i], h_tol) << "Hessian mismatch at nnz " << i; + } + const float b_tol = rel_tol * std::max(MaxAbs(expected_rhs), 1e-6f); + for (size_t i = 0; i < expected_rhs.size(); i++) { + ASSERT_NEAR(expected_rhs[i], actual.rhs[i], b_tol) << "RHS mismatch at " << i; + } +} + +/** @brief Asserts the GPU-built pattern matches the CPU oracle exactly. */ +void ExpectStructureMatchesReference(Problem &problem) { + ReferenceStructure expected = BuildReferenceStructure(problem); + + CudaStream stream; + SystemBuilder builder = SystemBuilder::Scalar(); + builder.Build(stream.GetStream(), problem); + SystemSnapshot actual = Snapshot(builder, stream.GetStream()); + + ASSERT_EQ(expected.row_offsets.size(), actual.row_offsets.size()); + EXPECT_EQ(expected.row_offsets, actual.row_offsets); + ASSERT_EQ(expected.col_ids.size(), actual.col_ids.size()); + EXPECT_EQ(expected.col_ids, actual.col_ids); +} + +TEST(HessianStructureTest, MatchesReferenceOnVectorChain) { + // Between and prior factors both emit the (i, i) block pair, so this covers + // the duplicate-key segmentation that an exclusive scan gets wrong. + auto data = MakeVectorChain(64, {}); + ExpectStructureMatchesReference(data->problem); +} + +TEST(HessianStructureTest, MatchesReferenceWithConstantStates) { + auto data = MakeVectorChain(64, {0, 7, 8, 63}); + ExpectStructureMatchesReference(data->problem); +} + +TEST(HessianStructureTest, MatchesReferenceWithMixedTangentDims) { + auto data = MakeMixedDim(48); + ExpectStructureMatchesReference(data->problem); +} + +TEST(HessianStructureTest, MatchesReferenceOnPoseGraph) { + auto data = MakePoseGraph(512, /*fix_first_pose=*/true); + ExpectStructureMatchesReference(data->problem); +} + +TEST(HessianStructureTest, MatchesReferenceOnBundleAdjustment) { + auto data = MakeBundle(8, 200, /*robust_loss=*/false); + ExpectStructureMatchesReference(data->problem); +} + +TEST(HessianStructureTest, MatchesReferenceWithRepeatedStateBlock) { + auto data = MakeVectorChain(32, {}, /*repeat_block=*/true); + ExpectStructureMatchesReference(data->problem); +} + +// ============================================================================ +// Storage-layout tests +// ============================================================================ + +/** + * @brief Asserts scalar and block storage assemble the same system. + * + * Both go through the same assembler, so this isolates the addressing change: + * a block pair is exactly tiled by `b x b` tiles, so expanding the block form + * must reproduce the scalar pattern and values entry for entry. Accumulation + * order differs between the two, hence a tolerance on values. + */ +void ExpectStorageLayoutsAgree(Problem &problem, float rel_tol = 1e-5f) { + CudaStream stream; + + SystemBuilder scalar = SystemBuilder::Scalar(); + scalar.Build(stream.GetStream(), problem); + ASSERT_FALSE(scalar.UsesBlockStorage()); + SystemSnapshot expected = Snapshot(scalar, stream.GetStream()); + + SystemBuilder block; + block.Build(stream.GetStream(), problem); + ASSERT_TRUE(block.UsesBlockStorage()) << "problem should qualify for block storage"; + SystemSnapshot actual = Snapshot(block, stream.GetStream()); + + EXPECT_EQ(expected.row_offsets, actual.row_offsets); + EXPECT_EQ(expected.col_ids, actual.col_ids); + ASSERT_EQ(expected.values.size(), actual.values.size()); + + const float h_tol = rel_tol * std::max(MaxAbs(expected.values), 1e-6f); + for (size_t i = 0; i < expected.values.size(); i++) { + ASSERT_NEAR(expected.values[i], actual.values[i], h_tol) << "Hessian mismatch at nnz " << i; + } + const float b_tol = rel_tol * std::max(MaxAbs(expected.rhs), 1e-6f); + ASSERT_EQ(expected.rhs.size(), actual.rhs.size()); + for (size_t i = 0; i < expected.rhs.size(); i++) { + ASSERT_NEAR(expected.rhs[i], actual.rhs[i], b_tol) << "RHS mismatch at " << i; + } +} + +TEST(HessianStorageTest, BlockMatchesScalarOnPoseGraph) { + auto data = MakePoseGraph(512, /*fix_first_pose=*/true); + ExpectStorageLayoutsAgree(data->problem); +} + +TEST(HessianStorageTest, BlockMatchesScalarOnBundleAdjustment) { + // Mixed 6-dof poses and 3-dof points: gcd 3, so a pose tile spans a 2x2 grid + // of blocks while a point tile is a single block. + auto data = MakeBundle(8, 200, /*robust_loss=*/false); + ExpectStorageLayoutsAgree(data->problem); +} + +TEST(HessianStorageTest, BlockMatchesScalarWithConstantStates) { + auto data = MakeVectorChain(64, {0, 7, 8, 63}); + ExpectStorageLayoutsAgree(data->problem); +} + +TEST(HessianStorageTest, BlockMatchesScalarWithRepeatedStateBlock) { + auto data = MakeVectorChain(32, {}, /*repeat_block=*/true); + ExpectStorageLayoutsAgree(data->problem); +} + +/** + * @brief Checks the BSR diagonal operations against their CSR counterparts. + * + * These two are what Levenberg-Marquardt applies to the LHS on every + * iteration, and the assembly-equivalence tests above never damp, so nothing + * else covers them. + */ +TEST(HessianStorageTest, BlockDiagonalOpsMatchScalar) { + auto data = MakeBundle(8, 200, /*robust_loss=*/false); + CudaStream stream; + cudaStream_t s = stream.GetStream(); + + SystemBuilder scalar = SystemBuilder::Scalar(); + scalar.Build(s, data->problem); + SystemBuilder block; + block.Build(s, data->problem); + ASSERT_TRUE(block.UsesBlockStorage()); + + // 1. Diagonal extraction. + dvector csr_diag; + dvector bsr_diag; + scalar.Equations().ExtractLhsDiagonal(s, csr_diag); + block.Equations().ExtractLhsDiagonal(s, bsr_diag); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + std::vector expected_diag(csr_diag.size()); + std::vector actual_diag(bsr_diag.size()); + csr_diag.CopyToHost(expected_diag.data(), expected_diag.size()); + bsr_diag.CopyToHost(actual_diag.data(), actual_diag.size()); + ASSERT_EQ(expected_diag.size(), actual_diag.size()); + const float diag_tol = 1e-5f * std::max(MaxAbs(expected_diag), 1e-6f); + for (size_t i = 0; i < expected_diag.size(); i++) { + ASSERT_NEAR(expected_diag[i], actual_diag[i], diag_tol) << "diagonal mismatch at " << i; + } + + // 2. Damped update, in place, exactly as LM performs it. + constexpr float kLambda = 0.017f; + scalar.Equations().AddScaledDiagonalToLhs(s, kLambda, csr_diag); + block.Equations().AddScaledDiagonalToLhs(s, kLambda, bsr_diag); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + SystemSnapshot expected = Snapshot(scalar, s); + SystemSnapshot actual = Snapshot(block, s); + ASSERT_EQ(expected.values.size(), actual.values.size()); + const float tol = 1e-5f * std::max(MaxAbs(expected.values), 1e-6f); + for (size_t i = 0; i < expected.values.size(); i++) { + ASSERT_NEAR(expected.values[i], actual.values[i], tol) << "damped LHS mismatch at nnz " << i; + } +} + +/** + * @brief The PCG solver must behave identically on the two storage layouts. + * + * Same matrix, same right-hand side, same preconditioner -- so the solution and + * the iteration count should match. A divergence here points at the block + * SpMV or the block-Jacobi tile gather rather than at assembly. + */ +TEST(HessianStorageTest, PcgAgreesBetweenStorages) { + auto data = MakePoseGraph(2048, /*fix_first_pose=*/true); + CudaStream stream; + cudaStream_t s = stream.GetStream(); + + SystemBuilder block; + block.Build(s, data->problem); + ASSERT_TRUE(block.UsesBlockStorage()); + + CSRSparseMatrix csr; + dvector expand_scratch; + ConvertBSRToCSR(s, block.Equations().LhsBSR(), csr, expand_scratch); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + const size_t n = block.Rhs().size(); + dvector rhs(n); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(rhs.data(), block.Rhs().data(), n * sizeof(float), + cudaMemcpyDeviceToDevice, s)); + + BlockSparsePCGOptions pcg_options; + pcg_options.block_size = 6; + pcg_options.max_iterations = 500; + pcg_options.relative_tolerance = 1e-6f; + + dvector x_csr(n); + dvector x_bsr(n); + + BlockSparsePCGSolver csr_solver(pcg_options); + ASSERT_TRUE(csr_solver.Initialize(s, data->problem, csr, rhs, x_csr)); + ASSERT_TRUE(csr_solver.Solve(s, csr, rhs, x_csr)); + + BlockSparsePCGSolver bsr_solver(pcg_options); + ASSERT_TRUE(bsr_solver.Initialize(s, data->problem, block.Equations().LhsBSR(), rhs, x_bsr)); + ASSERT_TRUE(bsr_solver.Solve(s, block.Equations().LhsBSR(), rhs, x_bsr)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + std::vector a(n); + std::vector b(n); + x_csr.CopyToHost(a.data(), n); + x_bsr.CopyToHost(b.data(), n); + + // Same iteration count means the two are tracking the same recurrence. The + // solutions themselves only agree to the level PCG was asked for: it stops on + // a relative residual, so the reordered block SpMV moves x within that ball. + EXPECT_EQ(csr_solver.LastIterations(), bsr_solver.LastIterations()); + const float tol = 1e-3f * std::max(MaxAbs(a), 1e-6f); + for (size_t i = 0; i < n; i++) { + ASSERT_NEAR(a[i], b[i], tol) << "solution mismatch at " << i; + } +} + +/** + * @brief One PCG iteration must be bit-comparable across storage layouts. + * + * After a single iteration the iterate depends only on the preconditioner and + * one SpMV, so this separates a bad block-Jacobi tile gather (which would show + * up here) from ordinary float divergence accumulating over many iterations. + */ +TEST(HessianStorageTest, FirstPcgIterationAgreesBetweenStorages) { + auto data = MakePoseGraph(1024, /*fix_first_pose=*/true); + CudaStream stream; + cudaStream_t s = stream.GetStream(); + + SystemBuilder block; + block.Build(s, data->problem); + ASSERT_TRUE(block.UsesBlockStorage()); + + // Feed both solvers the *same* matrix, so only the reader differs. + CSRSparseMatrix csr; + dvector expand_scratch; + ConvertBSRToCSR(s, block.Equations().LhsBSR(), csr, expand_scratch); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + const size_t n = block.Rhs().size(); + dvector rhs(n); + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(rhs.data(), block.Rhs().data(), n * sizeof(float), + cudaMemcpyDeviceToDevice, s)); + + BlockSparsePCGOptions pcg_options; + pcg_options.block_size = 6; + pcg_options.max_iterations = 1; + pcg_options.check_period = 1; + + dvector x_csr(n); + dvector x_bsr(n); + BlockSparsePCGSolver csr_solver(pcg_options); + ASSERT_TRUE(csr_solver.Initialize(s, data->problem, csr, rhs, x_csr)); + ASSERT_TRUE(csr_solver.Solve(s, csr, rhs, x_csr)); + + BlockSparsePCGSolver bsr_solver(pcg_options); + ASSERT_TRUE(bsr_solver.Initialize(s, data->problem, block.Equations().LhsBSR(), rhs, x_bsr)); + ASSERT_TRUE(bsr_solver.Solve(s, block.Equations().LhsBSR(), rhs, x_bsr)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + std::vector a(n); + std::vector b(n); + x_csr.CopyToHost(a.data(), n); + x_bsr.CopyToHost(b.data(), n); + + const float tol = 1e-6f * std::max(MaxAbs(a), 1e-6f); + for (size_t i = 0; i < n; i++) { + ASSERT_NEAR(a[i], b[i], tol) << "first iterate mismatch at " << i; + } +} + +/** + * @brief `step^T H step` must agree across storage layouts. + * + * Levenberg-Marquardt divides by this to form rho, so an error here does not + * corrupt the solution directly -- it corrupts the accept/reject decision, and + * the minimizer walks off to a worse answer while every matrix-level test still + * passes. + */ +TEST(HessianStorageTest, WeightedSquaredStepAgreesBetweenStorages) { + auto data = MakeBundle(8, 300, /*robust_loss=*/false); + CudaStream stream; + cudaStream_t s = stream.GetStream(); + + SystemBuilder scalar = SystemBuilder::Scalar(); + scalar.Build(s, data->problem); + SystemBuilder block; + block.Build(s, data->problem); + ASSERT_TRUE(block.UsesBlockStorage()); + + const size_t n = scalar.Rhs().size(); + std::vector host_step(n); + std::mt19937 rng(97); + std::uniform_real_distribution dist(-1.f, 1.f); + for (float &v : host_step) { + v = dist(rng); + } + dvector step(n); + step.CopyFromHost(host_step.data(), n); + + const float expected = scalar.WeightedSquaredStep(s, step); + const float actual = block.WeightedSquaredStep(s, step); + EXPECT_NEAR(expected, actual, 1e-4f * std::max(std::fabs(expected), 1.f)); +} + +/** + * @brief One builder reused across problems must not carry structure over. + * + * Regression test. The BSR copy originally skipped the structure whenever the + * destination's sizes already matched, so a second problem with the same block + * and tile counts but different connectivity inherited the first one's column + * indices. Every matrix-level equivalence test still passed -- the damage only + * showed up as Levenberg-Marquardt converging to worse costs on a dataset of + * many similarly-sized problems. + */ +TEST(HessianStorageTest, ReusedBuilderRebuildsStructurePerProblem) { + // Same block count and same factor count, different connectivity: a chain + // versus a pairing that skips two. + auto chain = MakeVectorChain(32, {}); + auto skipped = MakeVectorChain(32, {}, /*repeat_block=*/false, /*stride=*/3); + ASSERT_TRUE(chain->problem.CheckConsistency()); + ASSERT_TRUE(skipped->problem.CheckConsistency()); + + CudaStream stream; + cudaStream_t s = stream.GetStream(); + + SystemBuilder fresh; + fresh.Build(s, skipped->problem); + SystemSnapshot expected = Snapshot(fresh, s); + + SystemBuilder reused; + reused.Build(s, chain->problem); + reused.Build(s, skipped->problem); + SystemSnapshot actual = Snapshot(reused, s); + + EXPECT_EQ(expected.row_offsets, actual.row_offsets); + EXPECT_EQ(expected.col_ids, actual.col_ids); + ASSERT_EQ(expected.values.size(), actual.values.size()); + const float tol = 1e-5f * std::max(MaxAbs(expected.values), 1e-6f); + for (size_t i = 0; i < expected.values.size(); i++) { + ASSERT_NEAR(expected.values[i], actual.values[i], tol) << "stale structure at nnz " << i; + } +} + +TEST(HessianStorageTest, BlockSizeIsTheTangentDimensionGcd) { + auto pgo = MakePoseGraph(16, /*fix_first_pose=*/false); + EXPECT_EQ(6, ChooseHessianBlockSize(pgo->problem)); + + auto bundle = MakeBundle(4, 20, /*robust_loss=*/false); + EXPECT_EQ(3, ChooseHessianBlockSize(bundle->problem)); + + // 4-dimensional vector states alone tile at 4. + auto chain = MakeVectorChain(8, {}); + EXPECT_EQ(4, ChooseHessianBlockSize(chain->problem)); + + // gcd(3, 6) = 3 stays, but a coprime mix must fall back to scalar storage. + auto mixed = MakeMixedDim(8); + EXPECT_EQ(3, ChooseHessianBlockSize(mixed->problem)); +} + +TEST(HessianStorageTest, FallsBackToScalarWhenSolverNeedsCSR) { + // cuDSS consumes CSR, so block storage would only have to be expanded again. + auto data = MakePoseGraph(64, /*fix_first_pose=*/true); + MinimizerOptions options; + options.sparse_linear_solver_type = SparseLinearSolverType::cuDSS; + + CudaStream stream; + SystemBuilder builder(options); + builder.Build(stream.GetStream(), data->problem); + EXPECT_FALSE(builder.UsesBlockStorage()); +} + +// ============================================================================ +// Equivalence tests +// ============================================================================ + +TEST(BlockHessianAssemblerTest, MatchesReferenceOnVectorChain) { + auto data = MakeVectorChain(64, {}); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceWithConstantStates) { + auto data = MakeVectorChain(64, {0, 7, 8, 63}); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, AccumulatesRepeatedStateBlock) { + // Both slots of every between factor point at the same state block, so all + // four H_f sub-blocks land on the same CSR entries and must accumulate. + // + // Checked against the analytic answer directly. The pre-rewrite path built a + // triplet Jacobian and resolved duplicate (row, col) entries to a single CSR + // slot, storing rather than accumulating into it, so it silently dropped one + // of the two Jacobian blocks and got this case wrong. + // + // r_f = x_i - x_i - delta = -delta with J_f = [I | -I], so the four + // sub-blocks sum to I - I - I + I = 0 and b_f = -(I - I)^T r_f = 0. Only + // the prior factor (residual x_i - p_i, Jacobian I) survives. + constexpr int kBlocks = 32; + constexpr int kDim = VectorChainProblem::kDim; + auto data = MakeVectorChain(kBlocks, {}, /*repeat_block=*/true); + ASSERT_TRUE(data->problem.CheckConsistency()); + + CudaStream stream; + SystemBuilder block; + block.Build(stream.GetStream(), data->problem); + SystemSnapshot snapshot = Snapshot(block, stream.GetStream()); + + ASSERT_EQ(snapshot.row_offsets.size(), size_t(kBlocks * kDim + 1)); + for (int row = 0; row < kBlocks * kDim; row++) { + for (int k = snapshot.row_offsets[row]; k < snapshot.row_offsets[row + 1]; k++) { + float expected = snapshot.col_ids[k] == row ? 1.f : 0.f; + EXPECT_NEAR(expected, snapshot.values[k], 1e-6f) + << "row " << row << ", col " << snapshot.col_ids[k]; + } + } + + ASSERT_EQ(snapshot.rhs.size(), size_t(kBlocks * kDim)); + for (int i = 0; i < kBlocks; i++) { + for (int k = 0; k < kDim; k++) { + float expected = data->host_priors[i][k] - data->host_states[i * kDim + k]; + EXPECT_NEAR(expected, snapshot.rhs[i * kDim + k], 1e-5f) + << "block " << i << ", component " << k; + } + } +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceWithMixedTangentDims) { + auto data = MakeMixedDim(48); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceOnPoseGraph) { + auto data = MakePoseGraph(512, /*fix_first_pose=*/false); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceOnPoseGraphWithFixedPose) { + auto data = MakePoseGraph(512, /*fix_first_pose=*/true); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceOnBundleAdjustment) { + auto data = MakeBundle(8, 200, /*robust_loss=*/false); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceWithRobustLoss) { + auto data = MakeBundle(8, 200, /*robust_loss=*/true); + ASSERT_TRUE(data->problem.CheckConsistency()); + ExpectMatchesReference(data->problem); +} + +TEST(BlockHessianAssemblerTest, MatchesReferenceWithColumnScaling) { + auto data = MakePoseGraph(256, /*fix_first_pose=*/true); + ASSERT_TRUE(data->problem.CheckConsistency()); + + CudaStream stream; + MinimizerOptions options; + options.column_scaling = ColumnScaling::HessianDiagonal; + + SystemBuilder reference(options); + reference.Build(stream.GetStream(), data->problem); + std::vector expected = Snapshot(reference, stream.GetStream()).values; + + SystemBuilder block(options); + block.Build(stream.GetStream(), data->problem); + std::vector actual = Snapshot(block, stream.GetStream()).values; + + ASSERT_EQ(expected.size(), actual.size()); + const float tol = 1e-5f * std::max(MaxAbs(expected), 1e-6f); + for (size_t i = 0; i < expected.size(); i++) { + ASSERT_NEAR(expected[i], actual[i], tol) << "mismatch at " << i; + } +} + +TEST(BlockHessianAssemblerTest, EmptyFactorBatchIsSkipped) { + // A zero-factor batch alongside a populated one must not crash or perturb + // the assembled system. + auto data = MakeVectorChain(16, {}); + dvector> empty_priors; + PriorVectorFactorBatch empty_batch(empty_priors.data(), 0); + std::vector no_pointers; + data->problem.AddFactorBatch(&empty_batch, no_pointers); + + CudaStream stream; + SystemBuilder block; + ASSERT_NO_THROW(block.Build(stream.GetStream(), data->problem)); + EXPECT_GT(block.HessianAsCSR(stream.GetStream()).NumNonZeros(), 0u); +} + +TEST(BlockHessianAssemblerTest, AllConstantStatesProduceZeroSystem) { + std::vector all_const(16); + for (int i = 0; i < 16; i++) { + all_const[i] = i; + } + auto data = MakeVectorChain(16, all_const); + + CudaStream stream; + SystemBuilder block; + ASSERT_NO_THROW(block.Build(stream.GetStream(), data->problem)); + + SystemSnapshot snapshot = Snapshot(block, stream.GetStream()); + for (float v : snapshot.values) { + EXPECT_EQ(v, 0.f); + } + for (float v : snapshot.rhs) { + EXPECT_EQ(v, 0.f); + } +} + +} // namespace +} // namespace cunls diff --git a/tests/cusparse_matrix_multiplier_test.cpp b/tests/cusparse_matrix_multiplier_test.cpp deleted file mode 100644 index ac0ae94..0000000 --- a/tests/cusparse_matrix_multiplier_test.cpp +++ /dev/null @@ -1,245 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/** - * @file sparse_matrix_multiplication_test.cpp - * @brief Unit tests for GPU-based sparse matrix multiplication (A^T * A). - * - * Generates a random sparse CSR matrix, computes its Hessian (A^T * A) on - * both CPU and GPU, and verifies that the GPU result matches the CPU reference. - */ - -#include "cunls/minimizer/cusparse_matrix_multiplier.h" - -#include - -#include -#include -#include - -#include "cunls/common/cuda_stream.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 { - -namespace { - -/** - * @brief Generates a random square sparse matrix in CSR format. - * - * Each row has ~10% non-zero entries with a guaranteed diagonal element. - * - * @param rng Random number generator. - * @param rows Number of rows (and columns) in the square matrix. - * @param csr_values Output non-zero values. - * @param csr_col_idx Output column indices. - * @param csr_row_offsets Output row offset array. - */ -void GenerateRandomCSRMatrix(std::mt19937 &rng, int rows, - std::vector &csr_values, - std::vector &csr_col_idx, - std::vector &csr_row_offsets) { - const int cols = rows; - // Value distribution: random floats between 0.1 and 1.0 - std::uniform_real_distribution val_dist(0.1, 1.0); - // Column distribution: uniformly random column indices - std::uniform_int_distribution col_dist(0, cols - 1); - - // Sparsity level: approximately 10% of entries will be non-zero - constexpr float sparsity = 0.1; - - // Clear output vectors to start fresh - csr_row_offsets.clear(); - csr_col_idx.clear(); - csr_values.clear(); - - // CSR format: row_offsets[i] indicates where row i starts in values/col_idx - // arrays - csr_row_offsets.push_back(0); - - for (int i = 0; i < rows; ++i) { - // Calculate number of non-zero entries for this row - int nnz_in_row = 1 + static_cast(cols * sparsity); - // Start with diagonal element to ensure matrix has full rank - std::unordered_set used_cols = {i}; - - // Randomly select additional column indices (avoiding duplicates) - while (used_cols.size() < nnz_in_row) { - int col = col_dist(rng); - used_cols.insert(col); - } - - // Sort columns to maintain CSR format - std::vector sorted_cols(used_cols.begin(), used_cols.end()); - std::sort(sorted_cols.begin(), sorted_cols.end()); - - // Add all non-zero entries for this row - for (int col : sorted_cols) { - csr_col_idx.push_back(col); - csr_values.push_back(val_dist(rng)); - } - - // Record where the next row starts - csr_row_offsets.push_back(static_cast(csr_col_idx.size())); - } -} - -/** - * @brief Computes A^T * A on the CPU as a reference implementation. - * - * For each row of A, accumulates the outer product of that row with itself. - * - * @param row_ptr CSR row offsets of A. - * @param col_idx CSR column indices of A. - * @param values CSR values of A. - * @param AtA_row_ptr Output CSR row offsets of A^T * A. - * @param AtA_col_idx Output CSR column indices of A^T * A. - * @param AtA_values Output CSR values of A^T * A. - */ -void ComputeHessianCPU(const std::vector &row_ptr, - const std::vector &col_idx, - const std::vector &values, - std::vector &AtA_row_ptr, - std::vector &AtA_col_idx, - std::vector &AtA_values) { - int rows = row_ptr.size() - 1; - int cols = rows; // Square matrix assumption - - // Accumulate entries row-wise for A^T * A - // row_acc[i][j] will store the (i,j) entry of A^T * A - std::vector> row_acc(cols); // A^T * A has 'cols' rows - - // For each row of A - for (int i = 0; i < rows; ++i) { - int start = row_ptr[i]; - int end = row_ptr[i + 1]; - - // Compute outer product: row[i] * row[i]^T - // This contributes to multiple entries of A^T * A - for (int j = start; j < end; ++j) { - int col_j = col_idx[j]; // Column index in A - float val_j = values[j]; // Value A[i, col_j] - - for (int k = start; k < end; ++k) { - int col_k = col_idx[k]; // Another column index in A - float val_k = values[k]; // Value A[i, col_k] - - // Add A[i, col_j] * A[i, col_k] to (A^T * A)[col_j, col_k] - row_acc[col_j][col_k] += val_j * val_k; - } - } - } - - // Convert accumulated sparse rows to CSR format - AtA_row_ptr.clear(); - AtA_col_idx.clear(); - AtA_values.clear(); - - AtA_row_ptr.push_back(0); - for (int i = 0; i < cols; ++i) { - // Add all non-zero entries in row i of A^T * A - for (const auto &[col, val] : row_acc[i]) { - if (val != 0.0) { - AtA_col_idx.push_back(col); - AtA_values.push_back(val); - } - } - // Record where the next row starts - AtA_row_ptr.push_back(static_cast(AtA_col_idx.size())); - } -} - -} // namespace - -/** @brief Verifies GPU A^T*A computation matches the CPU reference - * implementation. */ -TEST(cuSPARSESparseMatrixMultiplierTest, ComputeHessian) { - profiler::ScopedRange range("ComputeHessianTest"); - // Use fixed seed for reproducibility - unsigned int fixed_seed = 0; - std::mt19937 gen(fixed_seed); - - // Test with a moderately large sparse matrix - constexpr int matrix_size = 1000; - - // Generate random sparse matrix in CSR format - std::vector csr_values; - std::vector csr_col_idx; - std::vector csr_row_offsets; - - GenerateRandomCSRMatrix(gen, matrix_size, csr_values, csr_col_idx, - csr_row_offsets); - - // Compute reference result on CPU - std::vector AtA_row_ptr; - std::vector AtA_col_idx; - std::vector AtA_values; - - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, AtA_row_ptr, - AtA_col_idx, AtA_values); - - // Convert input matrix to device-compatible format - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - // Compute A^T * A on GPU - cuSPARSESparseMatrixMultiplier smm; - CudaStream stream; - Problem problem; - - smm.Initialize(stream.GetStream(), problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - { - // Warm up - smm.ComputeSquaredMatrix(stream.GetStream(), problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - } - - { - profiler::ScopedRange range("ComputeSquaredMatrix"); - smm.ComputeSquaredMatrix(stream.GetStream(), problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - } - - // Verify that the sparsity structure matches exactly - std::vector hessian_row_offsets(hessian.row_offsets.size()); - hessian.row_offsets.CopyToHost(hessian_row_offsets.data(), - hessian_row_offsets.size()); - ASSERT_EQ(hessian_row_offsets, AtA_row_ptr); - - std::vector hessian_col_ids(hessian.col_ids.size()); - hessian.col_ids.CopyToHost(hessian_col_ids.data(), hessian_col_ids.size()); - ASSERT_EQ(hessian_col_ids, AtA_col_idx); - - // Verify that all non-zero values match within tolerance - std::vector hessian_values(hessian.values.size()); - hessian.values.CopyToHost(hessian_values.data(), hessian_values.size()); - ASSERT_EQ(hessian_values.size(), AtA_values.size()); - for (size_t i = 0; i < AtA_values.size(); i++) { - ASSERT_NEAR(hessian_values[i], AtA_values[i], 1e-3); - } -} - -} // namespace cunls diff --git a/tests/fast_matrix_multiplier_test.cpp b/tests/fast_matrix_multiplier_test.cpp deleted file mode 100644 index 027878b..0000000 --- a/tests/fast_matrix_multiplier_test.cpp +++ /dev/null @@ -1,521 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/** - * @file fast_matrix_multiplier_test.cpp - * @brief Unit tests for the custom GPU-based J^T * J computation. - * - * Generates random sparse CSR matrices, builds a matching Problem with - * factor graph connectivity, computes the Hessian (A^T * A) on both CPU - * and GPU using FastSparseMatrixMultiplier, and verifies the results match. - */ - -#include "cunls/minimizer/fast_matrix_multiplier.h" - -#include - -#include -#include -#include -#include -#include - -#include "cunls/common/cuda_stream.h" -#include "cunls/common/helper.h" -#include "cunls/common/profiler.h" -#include "cunls/common/types.h" -#include "cunls/minimizer/problem.h" -#include "cunls/state/vector_state_batch.h" -#include "tests/utils.h" - -namespace cunls { - -namespace { - -// ============================================================================ -// Mock factor batch for testing -// ============================================================================ - -class MockFactorBatch : public FactorBatch { -public: - MockFactorBatch(size_t num_factors, size_t num_blocks_per_factor) - : num_factors_(num_factors), sizes_(num_blocks_per_factor, 1) {} - - bool Evaluate(float *, float *, float const *const *, - cudaStream_t) const override { - return true; - } - size_t ResidualsSize() const override { return 1; } - std::vector StateBlockSizes() const override { return sizes_; } - size_t NumFactors() const override { return num_factors_; } - -private: - size_t num_factors_; - std::vector sizes_; -}; - -// ============================================================================ -// Test Problem builder -// ============================================================================ - -/** - * Holds all resources needed to keep a test Problem alive. - * The Problem references the state/factor batches via raw pointers, - * so this struct ensures their lifetimes are properly managed. - */ -struct TestProblemData { - dvector state_data; - dvector const_ids; - std::unique_ptr> state_batch; - std::unique_ptr factor_batch; - Problem problem; -}; - -/** - * Builds a Problem whose factor graph structure matches the given CSR - * Jacobian. Each CSR row becomes one factor, each column becomes one - * scalar state block (tangent_size = ambient_size = 1). Rows with fewer - * nonzeros than the maximum are padded with pointers to a dummy constant - * state block that the structure-aware path skips. - */ -TestProblemData BuildProblemFromCSR(const std::vector &csr_row_offsets, - const std::vector &csr_col_idx, - int num_cols) { - TestProblemData data; - int num_rows = static_cast(csr_row_offsets.size()) - 1; - - int max_nnz = 0; - for (int r = 0; r < num_rows; r++) { - max_nnz = std::max(max_nnz, csr_row_offsets[r + 1] - csr_row_offsets[r]); - } - - int total_blocks = num_cols + 1; - data.state_data.resize(total_blocks); - - std::vector h_const_ids = {num_cols}; - data.const_ids = dvector(h_const_ids); - - data.state_batch = std::make_unique>( - data.state_data.data(), total_blocks, data.const_ids.data(), - data.const_ids.size()); - - data.factor_batch = std::make_unique(num_rows, max_nnz); - - float *dummy = data.state_batch->StateBlockDevicePtr(num_cols); - std::vector state_ptrs; - state_ptrs.reserve(static_cast(num_rows) * max_nnz); - for (int r = 0; r < num_rows; r++) { - int start = csr_row_offsets[r]; - int end = csr_row_offsets[r + 1]; - for (int j = start; j < end; j++) { - state_ptrs.push_back( - data.state_batch->StateBlockDevicePtr(csr_col_idx[j])); - } - for (int j = end - start; j < max_nnz; j++) { - state_ptrs.push_back(dummy); - } - } - - data.problem.AddStateBatch(data.state_batch.get()); - data.problem.AddFactorBatch(data.factor_batch.get(), state_ptrs); - - return data; -} - -// ============================================================================ -// CSR generation and CPU reference -// ============================================================================ - -void GenerateRandomCSRMatrix(std::mt19937 &rng, int rows, int cols, - float sparsity, std::vector &csr_values, - std::vector &csr_col_idx, - std::vector &csr_row_offsets) { - std::uniform_real_distribution val_dist(0.1f, 1.0f); - std::uniform_int_distribution col_dist(0, cols - 1); - - csr_row_offsets.clear(); - csr_col_idx.clear(); - csr_values.clear(); - csr_row_offsets.push_back(0); - - for (int i = 0; i < rows; ++i) { - int nnz_in_row = std::max(1, static_cast(cols * sparsity)); - std::unordered_set used_cols; - used_cols.insert(col_dist(rng)); - - while (static_cast(used_cols.size()) < nnz_in_row) { - used_cols.insert(col_dist(rng)); - } - - std::vector sorted_cols(used_cols.begin(), used_cols.end()); - std::sort(sorted_cols.begin(), sorted_cols.end()); - - for (int col : sorted_cols) { - csr_col_idx.push_back(col); - csr_values.push_back(val_dist(rng)); - } - csr_row_offsets.push_back(static_cast(csr_col_idx.size())); - } -} - -void ComputeHessianCPU(const std::vector &row_ptr, - const std::vector &col_idx, - const std::vector &values, int num_cols, - std::vector &AtA_row_ptr, - std::vector &AtA_col_idx, - std::vector &AtA_values) { - int rows = row_ptr.size() - 1; - - std::vector> row_acc(num_cols); - - for (int i = 0; i < rows; ++i) { - int start = row_ptr[i]; - int end = row_ptr[i + 1]; - - for (int j = start; j < end; ++j) { - int col_j = col_idx[j]; - float val_j = values[j]; - - for (int k = start; k < end; ++k) { - int col_k = col_idx[k]; - float val_k = values[k]; - row_acc[col_j][col_k] += val_j * val_k; - } - } - } - - AtA_row_ptr.clear(); - AtA_col_idx.clear(); - AtA_values.clear(); - AtA_row_ptr.push_back(0); - - for (int i = 0; i < num_cols; ++i) { - for (const auto &[col, val] : row_acc[i]) { - if (val != 0.0f) { - AtA_col_idx.push_back(col); - AtA_values.push_back(val); - } - } - AtA_row_ptr.push_back(static_cast(AtA_col_idx.size())); - } -} - -void VerifyHessian(const CSRSparseMatrix &hessian, - const std::vector &ref_row_ptr, - const std::vector &ref_col_idx, - const std::vector &ref_values) { - std::vector hessian_row_offsets(hessian.row_offsets.size()); - hessian.row_offsets.CopyToHost(hessian_row_offsets.data(), - hessian_row_offsets.size()); - ASSERT_EQ(hessian_row_offsets, ref_row_ptr); - - std::vector hessian_col_ids(hessian.col_ids.size()); - hessian.col_ids.CopyToHost(hessian_col_ids.data(), hessian_col_ids.size()); - ASSERT_EQ(hessian_col_ids, ref_col_idx); - - std::vector hessian_values(hessian.values.size()); - hessian.values.CopyToHost(hessian_values.data(), hessian_values.size()); - ASSERT_EQ(hessian_values.size(), ref_values.size()); - for (size_t i = 0; i < ref_values.size(); i++) { - ASSERT_NEAR(hessian_values[i], ref_values[i], 1e-3) - << "Mismatch at index " << i; - } -} - -} // namespace - -TEST(FastSparseMatrixMultiplierTest, SquareMatrix) { - profiler::ScopedRange range("SymmetricSquare_SquareMatrix"); - - std::mt19937 gen(42); - constexpr int size = 500; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - GenerateRandomCSRMatrix(gen, size, size, 0.05f, csr_values, csr_col_idx, - csr_row_offsets); - - std::vector ref_row_ptr, ref_col_idx; - std::vector ref_values; - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, size, ref_row_ptr, - ref_col_idx, ref_values); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, size); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - VerifyHessian(hessian, ref_row_ptr, ref_col_idx, ref_values); -} - -TEST(FastSparseMatrixMultiplierTest, TallMatrix) { - profiler::ScopedRange range("SymmetricSquare_TallMatrix"); - - std::mt19937 gen(123); - constexpr int rows = 2000; - constexpr int cols = 200; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - GenerateRandomCSRMatrix(gen, rows, cols, 0.05f, csr_values, csr_col_idx, - csr_row_offsets); - - std::vector ref_row_ptr, ref_col_idx; - std::vector ref_values; - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, cols, ref_row_ptr, - ref_col_idx, ref_values); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, cols); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - VerifyHessian(hessian, ref_row_ptr, ref_col_idx, ref_values); -} - -TEST(FastSparseMatrixMultiplierTest, SparseRowJacobian) { - profiler::ScopedRange range("SymmetricSquare_SparseRow"); - - std::mt19937 gen(7); - constexpr int rows = 5000; - constexpr int cols = 500; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - GenerateRandomCSRMatrix(gen, rows, cols, 0.02f, csr_values, csr_col_idx, - csr_row_offsets); - - std::vector ref_row_ptr, ref_col_idx; - std::vector ref_values; - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, cols, ref_row_ptr, - ref_col_idx, ref_values); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, cols); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - VerifyHessian(hessian, ref_row_ptr, ref_col_idx, ref_values); -} - -TEST(FastSparseMatrixMultiplierTest, ValueReuseAfterInitialize) { - profiler::ScopedRange range("SymmetricSquare_ValueReuse"); - - std::mt19937 gen(99); - constexpr int rows = 1000; - constexpr int cols = 100; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - GenerateRandomCSRMatrix(gen, rows, cols, 0.1f, csr_values, csr_col_idx, - csr_row_offsets); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, cols); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - // Generate new values (same structure) and recompute without re-initializing. - std::uniform_real_distribution val_dist(0.1f, 1.0f); - for (auto &v : csr_values) - v = val_dist(gen); - input_matrix.values = dvector(csr_values); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - std::vector ref_row_ptr, ref_col_idx; - std::vector ref_values; - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, cols, ref_row_ptr, - ref_col_idx, ref_values); - - VerifyHessian(hessian, ref_row_ptr, ref_col_idx, ref_values); -} - -TEST(FastSparseMatrixMultiplierTest, LargeSquareMatrix) { - profiler::ScopedRange range("SymmetricSquare_LargeSquare"); - - std::mt19937 gen(0); - constexpr int size = 1000; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - GenerateRandomCSRMatrix(gen, size, size, 0.02f, csr_values, csr_col_idx, - csr_row_offsets); - - std::vector ref_row_ptr, ref_col_idx; - std::vector ref_values; - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, size, ref_row_ptr, - ref_col_idx, ref_values); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, size); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - VerifyHessian(hessian, ref_row_ptr, ref_col_idx, ref_values); -} - -TEST(FastSparseMatrixMultiplierTest, SingleColumnPerRow) { - profiler::ScopedRange range("SymmetricSquare_SingleColumn"); - - constexpr int rows = 100; - constexpr int cols = 50; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - csr_row_offsets.push_back(0); - - std::mt19937 gen(55); - std::uniform_real_distribution val_dist(0.1f, 1.0f); - std::uniform_int_distribution col_dist(0, cols - 1); - - for (int i = 0; i < rows; ++i) { - csr_col_idx.push_back(col_dist(gen)); - csr_values.push_back(val_dist(gen)); - csr_row_offsets.push_back(static_cast(csr_col_idx.size())); - } - - std::vector ref_row_ptr, ref_col_idx_ref; - std::vector ref_values; - ComputeHessianCPU(csr_row_offsets, csr_col_idx, csr_values, cols, ref_row_ptr, - ref_col_idx_ref, ref_values); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, cols); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - VerifyHessian(hessian, ref_row_ptr, ref_col_idx_ref, ref_values); -} - -TEST(FastSparseMatrixMultiplierTest, PerformanceBenchmark) { - profiler::ScopedRange range("SymmetricSquare_Performance"); - - std::mt19937 gen(42); - constexpr int rows = 20000; - constexpr int cols = 2000; - constexpr float sparsity = 0.005f; - - std::vector csr_values; - std::vector csr_col_idx, csr_row_offsets; - GenerateRandomCSRMatrix(gen, rows, cols, sparsity, csr_values, csr_col_idx, - csr_row_offsets); - - CSRSparseMatrix input_matrix, hessian; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - auto test_data = BuildProblemFromCSR(csr_row_offsets, csr_col_idx, cols); - FastSparseMatrixMultiplier smm; - CudaStream stream; - - smm.Initialize(stream.GetStream(), test_data.problem, input_matrix, hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - // Warmup - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, input_matrix, - hessian); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - cudaEvent_t start, stop; - THROW_ON_CUDA_ERROR(cudaEventCreate(&start)); - THROW_ON_CUDA_ERROR(cudaEventCreate(&stop)); - - constexpr int kNumIterations = 50; - THROW_ON_CUDA_ERROR(cudaEventRecord(start, stream.GetStream())); - for (int i = 0; i < kNumIterations; i++) { - smm.ComputeSquaredMatrix(stream.GetStream(), test_data.problem, - input_matrix, hessian); - } - THROW_ON_CUDA_ERROR(cudaEventRecord(stop, stream.GetStream())); - THROW_ON_CUDA_ERROR(cudaEventSynchronize(stop)); - - float ms; - THROW_ON_CUDA_ERROR(cudaEventElapsedTime(&ms, start, stop)); - - std::cout << "[Performance] ComputeSquaredMatrix: " << ms / kNumIterations - << " ms/iteration (" << rows << "x" << cols - << ", input_nnz=" << csr_values.size() - << ", output_nnz=" << hessian.NumNonZeros() << ")" << std::endl; - - THROW_ON_CUDA_ERROR(cudaEventDestroy(start)); - THROW_ON_CUDA_ERROR(cudaEventDestroy(stop)); -} - -} // namespace cunls diff --git a/tests/gauss_newton_test.cpp b/tests/gauss_newton_test.cpp index 84cf968..d6c6f99 100644 --- a/tests/gauss_newton_test.cpp +++ b/tests/gauss_newton_test.cpp @@ -293,32 +293,6 @@ TYPED_TEST(GaussNewtonMinimizerTest, LMColumnScalingHessianDiagonal) { this->CheckConvergence(vector_states); } -/** - * @brief Levenberg-Marquardt with Jacobian column-norm scaling. - */ -TYPED_TEST(GaussNewtonMinimizerTest, LMColumnScalingJacobianColumnNorm) { - auto test_range = this->profiler_domain_.CreateDomainRange("LMColumnScalingJacobianColumnNorm"); - typename TestFixture::StateData state_data(this->state_values_); - auto &vector_states = state_data.get(); - auto device_pointers = test_utils::CollectStatePointers(vector_states); - typename TestFixture::FactorData factor_data(this->observations_); - auto &factor_batch = factor_data.get(); - - Problem problem; - problem.AddFactorBatch(&factor_batch, device_pointers); - problem.AddStateBatch(&vector_states); - - CudaStream stream; - LevenbergMarquardtMinimizerOptions lm_options; - lm_options.base_options = this->minimizer_options_; - lm_options.base_options.column_scaling = ColumnScaling::JacobianColumnNorm; - LevenbergMarquardtMinimizer minimizer(lm_options); - minimizer.Minimize(stream.GetStream(), problem); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - - this->CheckConvergence(vector_states); -} - /** * @brief Gauss-Newton with column scaling (shared MinimizerOptions path). */ diff --git a/tests/jacobian_ops_test.cpp b/tests/jacobian_ops_test.cpp deleted file mode 100644 index 9f3dca1..0000000 --- a/tests/jacobian_ops_test.cpp +++ /dev/null @@ -1,146 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. - * All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/** @file jacobian_ops_test.cpp - * @brief Tests for Jacobian sparse structure construction on GPU. - */ - -#include "cunls/minimizer/minimizer_state.h" - -#include - -#include -#include -#include - -#include "cunls/common/cuda_stream.h" -#include "cunls/common/device_vector.h" -#include "cunls/common/helper.h" -#include "cunls/common/profiler.h" -#include "cunls/common/types.h" -#include "cunls/factor/prior_vector_factor_batch.h" -#include "cunls/state/vector_state_batch.h" -#include "tests/utils.h" - -namespace cunls { - -/** - * @brief Test fixture for Jacobian operations with typed vector sizes. - * - * Sets up a randomly-sized optimization problem with constant and non-constant - * states, then verifies GPU-built triplet sparse structures are consistent - * across two MinimizerState builds. - * - * @tparam VectorSize Compile-time vector dimension. - */ -template class JacobianOpsTest : public ::testing::Test { -public: - static constexpr int kDim = VectorSize::size; - - /** @brief Generates a randomly-sized optimization problem with ~20% constant - * states. */ - void SetUp() override { - unsigned int fixed_seed = 0; - std::mt19937 gen(fixed_seed); - - std::uniform_int_distribution states_distrib(1000, 100000); - std::uniform_int_distribution cost_batches_distrib(1, 10); - num_vectors = states_distrib(gen); - num_factor_batches = cost_batches_distrib(gen); - - std::uniform_real_distribution unit_distr(0, 1); - for (size_t i = 0; i < num_vectors; i++) { - // 20% of the states are constant - if (unit_distr(gen) > 0.8) { - constant_state_ids.push_back(i); - } - } - } - - size_t num_vectors; - size_t num_factor_batches; - std::vector constant_state_ids; - - profiler::Domain profiler_domain_{"JacobianOpsTest"}; -}; - -typedef ::testing::Types, test_utils::Size<2>, - test_utils::Size<3>, test_utils::Size<4>> - VectorSizes; -TYPED_TEST_CASE(JacobianOpsTest, VectorSizes); - -/** @brief Verifies that two GPU triplet builds match (idempotent structure). */ -TYPED_TEST(JacobianOpsTest, BuildTripletSparseStructure) { - auto test_range = this->profiler_domain_.CreateDomainRange( - "BuildTripletSparseStructureTest"); - // Prepare inputs - auto seq_vecs = - test_utils::MakeSequentialVectors(this->num_vectors); - test_utils::VectorStateData state_data( - seq_vecs, this->constant_state_ids); - auto &vector_states = state_data.get(); - auto device_pointers = test_utils::CollectStatePointers(vector_states); - auto obs_vecs = test_utils::MakeConstantVectors( - this->num_vectors, 1.f); - test_utils::PriorFactorData factor_data(obs_vecs); - auto &factor_batch = factor_data.get(); - - // Build the optimization problem - Problem problem; - problem.AddStateBatch(&vector_states); - - for (size_t i = 0; i < this->num_factor_batches; i++) { - problem.AddFactorBatch(&factor_batch, device_pointers); - } - - ASSERT_TRUE(problem.CheckConsistency()); - - CudaStream stream; - - TripletSparseStructure structure_a; - TripletSparseStructure structure_b; - { - auto range = - this->profiler_domain_.CreateDomainRange("BuildTripletSparseStructure"); - MinimizerState ms_a; - ms_a.BuildTripletSparseStructure(stream.GetStream(), problem, structure_a); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - MinimizerState ms_b; - ms_b.BuildTripletSparseStructure(stream.GetStream(), problem, structure_b); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - } - - ASSERT_EQ(structure_a.row_ids.size(), structure_b.row_ids.size()); - ASSERT_EQ(structure_a.col_ids.size(), structure_b.col_ids.size()); - - std::vector row_ids(structure_a.row_ids.size()); - structure_a.row_ids.CopyToHost(row_ids.data(), row_ids.size()); - std::vector row_ids_b(structure_b.row_ids.size()); - structure_b.row_ids.CopyToHost(row_ids_b.data(), row_ids_b.size()); - for (size_t i = 0; i < structure_a.row_ids.size(); i++) { - ASSERT_EQ(row_ids[i], row_ids_b[i]); - } - - std::vector col_ids(structure_a.col_ids.size()); - structure_a.col_ids.CopyToHost(col_ids.data(), col_ids.size()); - std::vector col_ids_b(structure_b.col_ids.size()); - structure_b.col_ids.CopyToHost(col_ids_b.data(), col_ids_b.size()); - for (size_t i = 0; i < structure_a.col_ids.size(); i++) { - ASSERT_EQ(col_ids[i], col_ids_b[i]); - } -} -} // namespace cunls diff --git a/tests/sparse_matrix_test.cpp b/tests/sparse_matrix_test.cpp index 1c32687..67d6f5f 100644 --- a/tests/sparse_matrix_test.cpp +++ b/tests/sparse_matrix_test.cpp @@ -20,7 +20,7 @@ * @brief Unit tests for sparse matrix operations (COO-to-CSR, diagonal, copy, * SpMV, RHS). * - * Tests GPU implementations of triplet-to-CSR conversion, diagonal extraction, + * Tests GPU implementations of diagonal extraction, * scaled diagonal addition, matrix copy, weighted squared step norms, and * right-hand-side computation against CPU reference implementations. */ @@ -30,7 +30,7 @@ #include #include -#include +#include #include #include "cunls/common/cuda_stream.h" @@ -45,83 +45,6 @@ namespace cunls { namespace { -/** @brief Represents a sparse matrix entry in COO (Coordinate) format. */ -struct Triplet { - float value; - int row_id; - int col_id; -}; - -/** - * @brief Converts a COO sparse matrix to CSR format on the CPU. - * - * Handles -1 sentinel column values as invalid/missing entries. - * - * @param coo_values COO non-zero values. - * @param coo_col_indices COO column indices (-1 marks invalid entries). - * @param coo_row_indices COO row indices. - * @param csr_values Output CSR values. - * @param csr_col_indices Output CSR column indices. - * @param csr_row_pointers Output CSR row pointers. - */ -void ConvertTripletToCSRCPU(const std::vector &coo_values, - const std::vector &coo_col_indices, - const std::vector &coo_row_indices, - std::vector &csr_values, - std::vector &csr_col_indices, - std::vector &csr_row_pointers) { - // Create triplet representation for easier sorting - std::vector triplets; - triplets.reserve(coo_values.size()); - for (int i = 0; i < coo_values.size(); i++) { - triplets.push_back({coo_values[i], coo_row_indices[i], coo_col_indices[i]}); - } - - // Sort triplets, moving invalid entries (col_id == -1) to the end - std::stable_sort(triplets.begin(), triplets.end(), - [](const Triplet &left, const Triplet &right) { - return right.col_id < 0; - }); - - // Count the number of valid non-zero entries (excluding -1 column indices) - size_t num_nonzeros = - std::accumulate(coo_col_indices.begin(), coo_col_indices.end(), 0, - [](int acc, int val) { return acc + int(val != -1); }); - - // Determine the number of rows in the matrix - int num_rows = - *std::max_element(coo_row_indices.begin(), coo_row_indices.end()) + 1; - - // Step 1: Count non-zeros per row - std::vector row_counts(num_rows, 0); - for (size_t i = 0; i < num_nonzeros; i++) { - const Triplet &triplet = triplets[i]; - row_counts[triplet.row_id]++; - } - - // Step 2: Compute CSR row pointers (cumulative sum of row counts) - // row_pointers[i] indicates where row i starts in the values array - csr_row_pointers.resize(num_rows + 1); - csr_row_pointers[0] = 0; - for (int i = 0; i < num_rows; ++i) { - csr_row_pointers[i + 1] = csr_row_pointers[i] + row_counts[i]; - } - - // Step 3: Populate CSR values and column indices arrays - csr_values.resize(num_nonzeros); - csr_col_indices.resize(num_nonzeros); - - // Track the current position for each row as we fill the arrays - std::vector current_row_positions = csr_row_pointers; - for (size_t i = 0; i < num_nonzeros; i++) { - const Triplet &triplet = triplets[i]; - int target_index = current_row_positions[triplet.row_id]; - csr_values[target_index] = triplet.value; - csr_col_indices[target_index] = triplet.col_id; - current_row_positions[triplet.row_id]++; - } -} - /** * @brief Computes y = A * x on the CPU where A is in CSR format. * @@ -129,50 +52,19 @@ void ConvertTripletToCSRCPU(const std::vector &coo_values, * @param col_idx CSR column indices. * @param values CSR values. * @param x Input vector. - * @param y Output result vector. + * @param[out] y Output result vector. */ -void MultiplyCSRMatrixByVector(const std::vector &row_ptr, - const std::vector &col_idx, - const std::vector &values, - const std::vector &x, +void MultiplyCSRMatrixByVector(const std::vector &row_ptr, const std::vector &col_idx, + const std::vector &values, const std::vector &x, std::vector &y) { - int rows = row_ptr.size() - 1; - y.assign(rows, 0.0); - - // For each row, iterate through its non-zero elements - for (int i = 0; i < rows; ++i) { - for (int j = row_ptr[i]; j < row_ptr[i + 1]; ++j) { - y[i] += values[j] * x[col_idx[j]]; - } - } -} - -/** - * @brief Computes y = -A^T * x on the CPU (negative transpose SpMV). - * - * @param row_ptr CSR row offsets. - * @param col_idx CSR column indices. - * @param values CSR values. - * @param x Input vector. - * @param y Output result vector. - */ -void ComputeRHSonCPU(const std::vector &row_ptr, - const std::vector &col_idx, - const std::vector &values, - const std::vector &x, std::vector &y) { - int rows = row_ptr.size() - 1; - int cols = rows; - - y.assign(cols, 0.0); // A^T * x has size equal to number of columns in A - - // For transpose multiplication, iterate through rows but accumulate into - // columns - for (int i = 0; i < rows; ++i) { - float xi = x[i]; - for (int j = row_ptr[i]; j < row_ptr[i + 1]; ++j) { - int col = col_idx[j]; - y[col] -= values[j] * xi; // Note: negative sign + const size_t num_rows = row_ptr.size() - 1; + y.assign(num_rows, 0.f); + for (size_t row = 0; row < num_rows; ++row) { + float sum = 0.f; + for (int k = row_ptr[row]; k < row_ptr[row + 1]; ++k) { + sum += values[k] * x[col_idx[k]]; } + y[row] = sum; } } @@ -184,10 +76,8 @@ void ComputeRHSonCPU(const std::vector &row_ptr, * @param values CSR values. * @param diagonal Output diagonal vector. */ -void ExtractDiagonalCPU(const std::vector &row_ptr, - const std::vector &col_idx, - const std::vector &values, - std::vector &diagonal) { +void ExtractDiagonalCPU(const std::vector &row_ptr, const std::vector &col_idx, + const std::vector &values, std::vector &diagonal) { size_t rows = row_ptr.size() - 1; diagonal.clear(); @@ -198,7 +88,7 @@ void ExtractDiagonalCPU(const std::vector &row_ptr, for (int j = row_ptr[i]; j < row_ptr[i + 1]; ++j) { if (col_idx[j] == i) { diagonal[i] = values[j]; - break; // Found diagonal element for this row + break; // Found diagonal element for this row } } } @@ -213,8 +103,7 @@ void ExtractDiagonalCPU(const std::vector &row_ptr, * @param scale Scalar multiplier for the diagonal. * @param diagonal Diagonal vector to add. */ -void AddScaledDiagonalCPU(const std::vector &row_ptr, - const std::vector &col_idx, +void AddScaledDiagonalCPU(const std::vector &row_ptr, const std::vector &col_idx, std::vector &values, float scale, const std::vector &diagonal) { // For each row, find and update the diagonal element @@ -222,7 +111,7 @@ void AddScaledDiagonalCPU(const std::vector &row_ptr, for (int j = row_ptr[i]; j < row_ptr[i + 1]; ++j) { if (col_idx[j] == i) { values[j] += scale * diagonal[i]; - break; // Found and updated diagonal element for this row + break; // Found and updated diagonal element for this row } } } @@ -233,71 +122,50 @@ void AddScaledDiagonalCPU(const std::vector &row_ptr, /** * @brief Test fixture for sparse matrix operations. * - * Sets up a random sparse matrix in both COO (triplet) and CSR formats + * Sets up a random sparse matrix in CSR format * for use across multiple test cases. */ class SparseMatrixTest : public ::testing::Test { public: /** - * @brief Generates a random sparse matrix in COO (triplet) format. + * @brief Generates a random sparse matrix directly in CSR format. * - * ~10% of entries are non-zero, diagonal is always included, and - * 20% of column entries are set to -1 (invalid/sentinel). + * Roughly 10% of each row is non-zero and the diagonal is always present, so + * the diagonal-extraction and damping tests have something to find. Column + * indices are sorted within a row, which every consumer relies on. * * @param rng Random number generator. * @param rows Number of rows (and columns) in the square matrix. */ - void GenerateRandomTripletMatrix(std::mt19937 &rng, int rows) { - int cols = rows; - std::uniform_real_distribution val_dist(0.1, 1.0); + void GenerateRandomCSRMatrix(std::mt19937 &rng, int rows) { + const int cols = rows; + std::uniform_real_distribution value_dist(0.1f, 1.0f); std::uniform_int_distribution col_dist(0, cols - 1); - - std::uniform_real_distribution unit_distr(0, 1); - - constexpr float sparsity = 0.1; - - for (int i = 0; i < rows; ++i) { - // Each row has at least the diagonal element plus ~10% of columns - int nnz_in_row = 1 + static_cast(cols * sparsity); - std::unordered_set used_cols{i}; // Start with diagonal - - // Generate random column indices for this row - while (used_cols.size() < nnz_in_row) { - int col = col_dist(rng); - used_cols.insert(col); + constexpr float kSparsity = 0.1f; + + csr_row_offsets.assign(1, 0); + for (int row = 0; row < rows; ++row) { + const size_t nnz_in_row = 1 + static_cast(cols * kSparsity); + std::set row_cols{row}; // the diagonal is always present + while (row_cols.size() < nnz_in_row) { + row_cols.insert(col_dist(rng)); } - - // Add all non-zero entries for this row - for (int col : used_cols) { - triplet_row_idx.push_back(i); - triplet_values.push_back(val_dist(rng)); - - // 20% chance to mark column as invalid (-1) - col = unit_distr(rng) > 0.8 ? -1 : col; - triplet_col_idx.push_back(col); + for (int col : row_cols) { + csr_col_idx.push_back(col); + csr_values.push_back(value_dist(rng)); } + csr_row_offsets.push_back(static_cast(csr_col_idx.size())); } } - /** @brief Generates a random matrix and converts to CSR format for tests. */ + /** @brief Generates the random CSR matrix shared by the tests. */ void SetUp() override { - unsigned int fixed_seed = 0; - std::mt19937 gen(fixed_seed); - - GenerateRandomTripletMatrix(gen, matrix_size); - - // Convert from COO to CSR format for testing - ConvertTripletToCSRCPU(triplet_values, triplet_col_idx, triplet_row_idx, - csr_values, csr_col_idx, csr_row_offsets); + std::mt19937 rng(0); + GenerateRandomCSRMatrix(rng, matrix_size); } const size_t matrix_size = 1000; - // COO (Coordinate/Triplet) format data - std::vector triplet_row_idx; - std::vector triplet_col_idx; - std::vector triplet_values; - // CSR (Compressed Sparse Row) format data std::vector csr_row_offsets; std::vector csr_col_idx; @@ -309,72 +177,13 @@ class SparseMatrixTest : public ::testing::Test { profiler::Domain profiler_domain_{"SparseMatrixTest"}; }; -/** @brief Verifies GPU two-step COO-to-CSR conversion matches the CPU - * reference. */ -TEST_F(SparseMatrixTest, ConvertTripletToCSR) { - auto test_range = - this->profiler_domain_.CreateDomainRange("ConvertTripletToCSRTest"); - // Prepare GPU input in COO format - size_t num_nonzeros = triplet_values.size(); - - SparseJacobian sp_jacobian; - sp_jacobian.structure.row_ids.resize(num_nonzeros); - sp_jacobian.structure.col_ids.resize(num_nonzeros); - sp_jacobian.values.resize(num_nonzeros); - - // Copy COO data to device memory - sp_jacobian.structure.row_ids.CopyFromHost(triplet_row_idx.data(), - triplet_row_idx.size()); - sp_jacobian.structure.col_ids.CopyFromHost(triplet_col_idx.data(), - triplet_col_idx.size()); - sp_jacobian.values.CopyFromHost(triplet_values.data(), triplet_values.size()); - - // Step 1: Convert structure to CSR and build mapping - CudaStream stream; - CSRSparseMatrix csr_matrix; - dvector mapping; - cuSPARSEHandle cusparse_handle; - auto handle = cusparse_handle.GetHandle(stream.GetStream()); - { - auto range = this->profiler_domain_.CreateDomainRange( - "ConvertTripletStructureToCSR"); - ConvertTripletStructureToCSR(stream.GetStream(), handle, - sp_jacobian.structure, csr_matrix, mapping, - this->buffer); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - } - - // Step 2: Copy values using the mapping - { - auto range = - this->profiler_domain_.CreateDomainRange("ConvertTripletToCSRValues"); - ConvertTripletToCSRValues(stream.GetStream(), sp_jacobian, mapping, - csr_matrix); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - } - - // Verify that GPU result matches CPU reference implementation - hvector csr_values_host(csr_matrix.values.size()); - csr_matrix.values.CopyToHost(csr_values_host.data(), csr_values_host.size()); - hvector csr_row_offsets_host(csr_matrix.row_offsets.size()); - csr_matrix.row_offsets.CopyToHost(csr_row_offsets_host.data(), - csr_row_offsets_host.size()); - hvector csr_col_ids_host(csr_matrix.col_ids.size()); - csr_matrix.col_ids.CopyToHost(csr_col_ids_host.data(), - csr_col_ids_host.size()); - ASSERT_EQ(csr_values_host, csr_values); - ASSERT_EQ(csr_row_offsets_host, csr_row_offsets); - ASSERT_EQ(csr_col_ids_host, csr_col_idx); -} - /** @brief Verifies GPU diagonal extraction from a CSR matrix matches the CPU * reference. */ TEST_F(SparseMatrixTest, ExtractDiagonal) { auto test_range = this->profiler_domain_.CreateDomainRange("ExtractDiagonal"); // Prepare GPU matrix CSRSparseMatrix csr_matrix; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - csr_matrix); + test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, csr_matrix); // Compute expected result using CPU reference std::vector diagonal; @@ -392,20 +201,17 @@ TEST_F(SparseMatrixTest, ExtractDiagonal) { // Verify GPU result matches CPU reference hvector device_diagonal_host(device_diagonal.size()); - device_diagonal.CopyToHost(device_diagonal_host.data(), - device_diagonal_host.size()); + device_diagonal.CopyToHost(device_diagonal_host.data(), device_diagonal_host.size()); ASSERT_EQ(device_diagonal_host, diagonal); } /** @brief Verifies GPU AddScaledDiagonal (A + scale * diag(d)) matches the CPU * reference. */ TEST_F(SparseMatrixTest, AddScaledDiagonal) { - auto test_range = - this->profiler_domain_.CreateDomainRange("AddScaledDiagonalTest"); + auto test_range = this->profiler_domain_.CreateDomainRange("AddScaledDiagonalTest"); // Prepare GPU matrix CSRSparseMatrix input_matrix; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); + test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, input_matrix); // Generate random diagonal and scale factor float scale = 10; @@ -414,16 +220,14 @@ TEST_F(SparseMatrixTest, AddScaledDiagonal) { dvector device_diagonal(diagonal); // Compute expected result using CPU reference - AddScaledDiagonalCPU(csr_row_offsets, csr_col_idx, csr_values, scale, - diagonal); + AddScaledDiagonalCPU(csr_row_offsets, csr_col_idx, csr_values, scale, diagonal); // Perform operation on GPU CSRSparseMatrix result_matrix; CudaStream stream; { auto range = this->profiler_domain_.CreateDomainRange("AddScaledDiagonal"); - AddScaledDiagonal(stream.GetStream(), scale, device_diagonal, input_matrix, - result_matrix); + AddScaledDiagonal(stream.GetStream(), scale, device_diagonal, input_matrix, result_matrix); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); } @@ -434,8 +238,7 @@ TEST_F(SparseMatrixTest, AddScaledDiagonal) { result_matrix.row_offsets.CopyToHost(result_row_offsets_host.data(), result_row_offsets_host.size()); hvector result_col_ids_host(result_matrix.col_ids.size()); - result_matrix.col_ids.CopyToHost(result_col_ids_host.data(), - result_col_ids_host.size()); + result_matrix.col_ids.CopyToHost(result_col_ids_host.data(), result_col_ids_host.size()); ASSERT_EQ(result_row_offsets_host, csr_row_offsets); ASSERT_EQ(result_col_ids_host, csr_col_idx); @@ -453,8 +256,7 @@ TEST_F(SparseMatrixTest, Copy) { auto test_range = this->profiler_domain_.CreateDomainRange("CopyTest"); // Prepare source matrix on GPU CSRSparseMatrix input_matrix; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); + test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, input_matrix); // Copy matrix on GPU CSRSparseMatrix result_matrix; @@ -470,8 +272,7 @@ TEST_F(SparseMatrixTest, Copy) { result_matrix.row_offsets.CopyToHost(result_row_offsets_host.data(), result_row_offsets_host.size()); hvector result_col_ids_host(result_matrix.col_ids.size()); - result_matrix.col_ids.CopyToHost(result_col_ids_host.data(), - result_col_ids_host.size()); + result_matrix.col_ids.CopyToHost(result_col_ids_host.data(), result_col_ids_host.size()); ASSERT_EQ(result_row_offsets_host, csr_row_offsets); ASSERT_EQ(result_col_ids_host, csr_col_idx); @@ -487,8 +288,7 @@ TEST_F(SparseMatrixTest, Copy) { /** @brief Verifies GPU vector-form weighted squared step norm: sum(steps[i]^2 * * weights[i]). */ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepFirst) { - auto test_range = this->profiler_domain_.CreateDomainRange( - "ComputeWeightedSquaredStepFirstTest"); + auto test_range = this->profiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepFirstTest"); // Generate random weights and steps std::vector weights; test_utils::GenerateRandomVector(matrix_size, weights); @@ -509,10 +309,8 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepFirst) { CudaStream stream; float result; { - auto range = this->profiler_domain_.CreateDomainRange( - "ComputeWeightedSquaredStepFirst"); - result = ComputeWeightedSquaredStep(stream.GetStream(), dweights, dsteps, - buffer); + auto range = this->profiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepFirst"); + result = ComputeWeightedSquaredStep(stream.GetStream(), dweights, dsteps, buffer); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); } @@ -523,8 +321,8 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepFirst) { /** @brief Verifies GPU matrix-form weighted squared step norm: steps^T * A * * steps. */ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { - auto test_range = this->profiler_domain_.CreateDomainRange( - "ComputeWeightedSquaredStepSecondTest"); + auto test_range = + this->profiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepSecondTest"); // Generate random step vector std::vector steps; test_utils::GenerateRandomVector(matrix_size, steps); @@ -532,8 +330,7 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { // Compute expected result on CPU: steps^T * A * steps // First compute temp = A * steps std::vector temp; - MultiplyCSRMatrixByVector(csr_row_offsets, csr_col_idx, csr_values, steps, - temp); + MultiplyCSRMatrixByVector(csr_row_offsets, csr_col_idx, csr_values, steps, temp); // Then compute steps^T * temp float gt_value = 0; for (size_t i = 0; i < matrix_size; i++) { @@ -543,8 +340,7 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { // Prepare GPU matrix and vector CSRSparseMatrix input_matrix; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); + test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, input_matrix); dvector dsteps(steps); @@ -554,10 +350,8 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { auto handle = cusparse_handle.GetHandle(stream.GetStream()); float result; { - auto range = this->profiler_domain_.CreateDomainRange( - "ComputeWeightedSquaredStepSecond"); - result = ComputeWeightedSquaredStep(stream.GetStream(), handle, - input_matrix, dsteps, buffer); + auto range = this->profiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepSecond"); + result = ComputeWeightedSquaredStep(stream.GetStream(), handle, input_matrix, dsteps, buffer); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); } @@ -566,46 +360,6 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { ASSERT_NEAR(result, gt_value, 1e-3); } -/** @brief Verifies GPU RHS computation (y = -A^T * x) matches the CPU - * reference. */ -TEST_F(SparseMatrixTest, ComputeRHS) { - auto test_range = this->profiler_domain_.CreateDomainRange("ComputeRHSTest"); - // Generate random input vector - std::vector steps; - test_utils::GenerateRandomVector(matrix_size, steps); - - // Compute expected result using CPU reference - std::vector gt; - ComputeRHSonCPU(csr_row_offsets, csr_col_idx, csr_values, steps, gt); - - // Prepare GPU matrix and vector - CSRSparseMatrix input_matrix; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - input_matrix); - - dvector dsteps(steps); - dvector result; - - // Compute RHS on GPU - CudaStream stream; - cuSPARSEHandle cusparse_handle; - auto handle = cusparse_handle.GetHandle(stream.GetStream()); - { - auto range = this->profiler_domain_.CreateDomainRange("ComputeRHS"); - ComputeRHS(stream.GetStream(), handle, input_matrix, dsteps, result, - buffer); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); - } - - // Verify GPU result matches CPU reference (within tolerance) - hvector hresult(result.size()); - result.CopyToHost(hresult.data(), hresult.size()); - ASSERT_EQ(hresult.size(), gt.size()); - for (size_t i = 0; i < gt.size(); i++) { - ASSERT_NEAR(hresult[i], gt[i], 1e-3); - } -} - /** @brief SHS scaling matches hand-derived 2x2 normal equations. */ TEST(SparseMatrixColumnScaling, SymmetricScaling2x2) { std::vector csr_row_offsets = {0, 2, 4}; @@ -613,8 +367,7 @@ TEST(SparseMatrixColumnScaling, SymmetricScaling2x2) { std::vector csr_values = {4.f, 1.f, 1.f, 9.f}; CSRSparseMatrix h; - test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, - h); + test_utils::CreateCSRSparseMatrix(csr_row_offsets, csr_col_idx, csr_values, h); std::vector host_scale = {0.5f, 1.f / 3.f}; dvector scale(host_scale); diff --git a/tests/utils.h b/tests/utils.h index 79e3a29..c588d78 100644 --- a/tests/utils.h +++ b/tests/utils.h @@ -21,13 +21,12 @@ #pragma once +#include +#include #include #include #include -#include -#include - #include "cunls/common/device_vector.h" #include "cunls/common/helper.h" #include "cunls/common/types.h" @@ -48,7 +47,8 @@ namespace test_utils { * Used with ::testing::Types to parameterize tests across vector dimensions. * @tparam Value Compile-time integer value. */ -template struct Size { +template +struct Size { static constexpr int size = Value; }; @@ -58,7 +58,8 @@ template struct Size { * Used for matrix dimensions and other size_t-valued compile-time parameters. * @tparam Value Compile-time size_t value. */ -template struct SizeT { +template +struct SizeT { static constexpr size_t value = Value; }; @@ -77,10 +78,8 @@ template struct SizeT { * @param values Value array for non-zero elements. * @param matrix Output CSRSparseMatrix to populate with the data. */ -void CreateCSRSparseMatrix(const std::vector &row_ptr, - const std::vector &col_idx, - const std::vector &values, - CSRSparseMatrix &matrix); +void CreateCSRSparseMatrix(const std::vector &row_ptr, const std::vector &col_idx, + const std::vector &values, CSRSparseMatrix &matrix); /** * @brief Generates a random vector with values in [0.1, 1.0] using a fixed @@ -170,7 +169,8 @@ std::vector> MakeSequentialVectors(size_t count) { * @param count Number of vectors to generate. * @return Vector of zero-filled vectors. */ -template std::vector> MakeZeroVectors(size_t count) { +template +std::vector> MakeZeroVectors(size_t count) { std::vector> v(count); for (size_t i = 0; i < count; i++) { v[i].fill(0); @@ -224,7 +224,8 @@ inline std::vector MakeSequentialIds(size_t count) { * * @tparam Dim Dimension of each vector state. */ -template struct VectorStateData { +template +struct VectorStateData { DeviceVector> vectors; DeviceVector const_ids; std::unique_ptr> batch; @@ -245,8 +246,8 @@ template struct VectorStateData { const_ids = DeviceVector(const_state_ids); const float *data_ptr = reinterpret_cast(vectors.data()); const int *const_ids_ptr = const_ids.empty() ? nullptr : const_ids.data(); - batch = std::make_unique>( - data_ptr, num_vectors, const_ids_ptr, const_ids.size()); + batch = std::make_unique>(data_ptr, num_vectors, const_ids_ptr, + const_ids.size()); } /** @brief Returns a reference to the managed VectorStateBatch. */ @@ -268,7 +269,8 @@ template struct VectorStateData { * * @tparam Dim Dimension of each observation vector. */ -template struct PriorFactorData { +template +struct PriorFactorData { DeviceVector> observations_device; std::unique_ptr> factor_batch; @@ -279,8 +281,8 @@ template struct PriorFactorData { */ PriorFactorData(const std::vector> &observations) { observations_device = DeviceVector>(observations); - factor_batch = std::make_unique>( - observations_device.data(), observations.size()); + factor_batch = std::make_unique>(observations_device.data(), + observations.size()); } /** @brief Returns a reference to the managed factor batch. */ @@ -333,15 +335,13 @@ DeviceVector CollectStatePointersDevice(StateBatchType &state_batch) { * @return Host vector of state values. */ template -std::vector> -CopyStateToHost(const VectorStateBatch &state_batch) { +std::vector> CopyStateToHost(const VectorStateBatch &state_batch) { auto ptr = state_batch.StateBlockDevicePtr(0); size_t num_blocks = state_batch.NumStateBlocks(); auto vec_ptr = reinterpret_cast *>(ptr); std::vector> out(num_blocks); - THROW_ON_CUDA_ERROR(cudaMemcpy(out.data(), vec_ptr, - num_blocks * sizeof(Vector), - cudaMemcpyDeviceToHost)); + THROW_ON_CUDA_ERROR( + cudaMemcpy(out.data(), vec_ptr, num_blocks * sizeof(Vector), cudaMemcpyDeviceToHost)); return out; } From 0610db6636071c4a0ed6bdc0dfd05d7053f6f7f6 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Sat, 1 Aug 2026 20:49:32 -0700 Subject: [PATCH 2/7] Remove redundant code --- CMakeLists.txt | 1 + LICENSE | 1 - NOTICE | 2 +- cunls/common/cublas_helper.cpp | 5 +- cunls/common/cublas_helper.h | 16 +- cunls/common/cuda_stream.cpp | 5 +- cunls/common/cuda_stream.h | 10 +- cunls/common/cudss_helper.h | 48 +++--- cunls/common/cusolver_helper.cpp | 14 +- cunls/common/cusolver_helper.h | 22 ++- cunls/common/cusparse_helper.cpp | 50 +++--- cunls/common/cusparse_helper.h | 31 ++-- cunls/common/device_vector.h | 77 ++++----- cunls/common/helper.h | 32 ++-- cunls/common/log.cpp | 46 +++--- cunls/common/log.h | 16 +- cunls/common/pinned_vector.h | 15 +- cunls/common/profiler.cpp | 19 +-- cunls/common/profiler.h | 28 ++-- cunls/common/type_traits.h | 5 +- cunls/common/types.h | 12 +- cunls/common/utils.cpp | 28 ++-- cunls/common/utils.h | 8 +- cunls/factor/factor_batch.h | 7 +- cunls/factor/information_factor_batch.cpp | 33 ++-- cunls/factor/information_factor_batch.h | 41 ++--- cunls/factor/pnp_factor_batch.cu | 45 +++--- cunls/factor/pnp_factor_batch.h | 15 +- cunls/factor/point_to_plane_factor_batch.cu | 13 +- .../linear_solver/block_sparse_pcg_solver.cu | 4 +- cunls/linear_solver/block_sparse_pcg_solver.h | 16 +- .../linear_solver/csr_sparse_linear_solver.h | 4 +- .../cudss_sparse_linear_solver.h | 16 +- cunls/linear_solver/dense_cholesky_solver.h | 6 +- cunls/linear_solver/dense_linear_solver.h | 8 +- cunls/linear_solver/dense_qr_solver.h | 6 +- cunls/minimizer/block_hessian_assembler.cu | 4 +- cunls/minimizer/block_hessian_assembler.h | 8 +- cunls/minimizer/bsr_matrix.cu | 87 +--------- cunls/minimizer/bsr_matrix.h | 34 +--- cunls/minimizer/gauss_newton_minimizer.cu | 50 ------ cunls/minimizer/gauss_newton_minimizer.h | 38 ++--- cunls/minimizer/hessian_structure.cu | 4 +- cunls/minimizer/hessian_structure.h | 10 +- .../levenberg_marquardt_minimizer.cpp | 68 +------- .../minimizer/levenberg_marquardt_minimizer.h | 23 --- cunls/minimizer/minimizer_state.h | 6 +- cunls/minimizer/normal_equations.cu | 2 +- cunls/minimizer/normal_equations.h | 11 +- cunls/minimizer/residual_batch.cu | 92 +++++------ cunls/minimizer/sparse_matrix.cu | 151 ++---------------- cunls/minimizer/sparse_matrix.h | 57 +------ cunls/state/state_batch_ops.cu | 2 +- cunls/state/state_batch_ops.h | 6 +- python/src/bind_types.cpp | 54 +++---- tests/block_hessian_assembler_test.cpp | 15 +- tests/bsr_expansion.cu | 130 +++++++++++++++ tests/bsr_expansion.h | 43 +++++ tests/sparse_matrix_test.cpp | 42 ++++- tests/utils.h | 4 +- 60 files changed, 638 insertions(+), 1008 deletions(-) create mode 100644 tests/bsr_expansion.cu create mode 100644 tests/bsr_expansion.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 36c0330..9e2bb3a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -137,6 +137,7 @@ if(BUILD_TESTING) add_executable( nls_tests tests/utils.cpp + tests/bsr_expansion.cu tests/factor_batch_test.cpp tests/problem_test.cpp tests/prior_vector_prior_factor_test.cpp diff --git a/LICENSE b/LICENSE index a944d3f..decdb6a 100644 --- a/LICENSE +++ b/LICENSE @@ -200,4 +200,3 @@ 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. - \ No newline at end of file diff --git a/NOTICE b/NOTICE index 8dfd3c4..4332b2d 100644 --- a/NOTICE +++ b/NOTICE @@ -40,4 +40,4 @@ This product includes third-party software components: https://github.com/NVIDIA/warp Full license texts for third-party components are provided in the -third_party/LICENSES directory. \ No newline at end of file +third_party/LICENSES directory. diff --git a/cunls/common/cublas_helper.cpp b/cunls/common/cublas_helper.cpp index d61524c..e2f6496 100644 --- a/cunls/common/cublas_helper.cpp +++ b/cunls/common/cublas_helper.cpp @@ -15,9 +15,10 @@ * limitations under the License. */ +#include "cunls/common/cublas_helper.h" + #include -#include "cunls/common/cublas_helper.h" #include "cunls/common/log.h" namespace cunls { @@ -77,4 +78,4 @@ void *cuBLASHandle::GetHandle(cudaStream_t stream) { return handle_; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cublas_helper.h b/cunls/common/cublas_helper.h index f45bbff..26a1d18 100644 --- a/cunls/common/cublas_helper.h +++ b/cunls/common/cublas_helper.h @@ -38,8 +38,7 @@ const char *cublasGetErrorString(int status); * If the cuBLAS status indicates an error, throws an exception with * a descriptive error message. */ -#define THROW_ON_CUBLAS_ERROR(status) \ - CHECK_CUDA_ERROR(status, cublasGetErrorString, true) +#define THROW_ON_CUBLAS_ERROR(status) CHECK_CUDA_ERROR(status, cublasGetErrorString, true) /** * @brief Macro to log a warning on cuBLAS errors. @@ -47,8 +46,7 @@ const char *cublasGetErrorString(int status); * If the cuBLAS status indicates an error, logs a warning message * but does not throw an exception. */ -#define WARN_ON_CUBLAS_ERROR(status) \ - CHECK_CUDA_ERROR(status, cublasGetErrorString, false) +#define WARN_ON_CUBLAS_ERROR(status) CHECK_CUDA_ERROR(status, cublasGetErrorString, false) /** * @brief RAII wrapper for cuBLAS handle management. @@ -61,7 +59,7 @@ const char *cublasGetErrorString(int status); * Non-copyable: Prevents accidental handle duplication. */ class cuBLASHandle { -public: + public: cuBLASHandle() = default; cuBLASHandle(const cuBLASHandle &) = delete; @@ -89,9 +87,9 @@ class cuBLASHandle { */ void *GetHandle(cudaStream_t stream); -private: - cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream. - void *handle_ = nullptr; ///< The cuBLAS handle. + private: + cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream. + void *handle_ = nullptr; ///< The cuBLAS handle. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cuda_stream.cpp b/cunls/common/cuda_stream.cpp index 645dc87..0def8f3 100644 --- a/cunls/common/cuda_stream.cpp +++ b/cunls/common/cuda_stream.cpp @@ -30,8 +30,7 @@ namespace cunls { * @param sync_on_destroy If true, the stream will be synchronized before * destruction. */ -CudaStream::CudaStream(bool sync_on_destroy) - : sync_on_destroy_(sync_on_destroy) { +CudaStream::CudaStream(bool sync_on_destroy) : sync_on_destroy_(sync_on_destroy) { THROW_ON_CUDA_ERROR(cudaStreamCreate(&stream)); } @@ -51,4 +50,4 @@ CudaStream::~CudaStream() { /** @brief Returns a reference to the underlying cudaStream_t handle. */ cudaStream_t &CudaStream::GetStream() { return stream; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cuda_stream.h b/cunls/common/cuda_stream.h index 3bd53f6..a07847a 100644 --- a/cunls/common/cuda_stream.h +++ b/cunls/common/cuda_stream.h @@ -31,7 +31,7 @@ namespace cunls { * Non-copyable to prevent accidental sharing of stream ownership. */ class CudaStream { -public: + public: /** * @brief Constructs a new CUDA stream. * @@ -55,8 +55,8 @@ class CudaStream { */ cudaStream_t &GetStream(); -private: - cudaStream_t stream; ///< The underlying CUDA stream handle. - bool sync_on_destroy_; ///< Whether to synchronize on destruction. + private: + cudaStream_t stream; ///< The underlying CUDA stream handle. + bool sync_on_destroy_; ///< Whether to synchronize on destruction. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cudss_helper.h b/cunls/common/cudss_helper.h index 6474b25..a76f1d3 100644 --- a/cunls/common/cudss_helper.h +++ b/cunls/common/cudss_helper.h @@ -17,9 +17,9 @@ #pragma once -#include #include +#include #include #include @@ -42,8 +42,7 @@ const char *cudssGetErrorString(int status); * If the status indicates an error, this macro will throw an exception with * a descriptive error message. */ -#define THROW_ON_CUDSS_ERROR(status) \ - CHECK_CUDA_ERROR(status, cudssGetErrorString, true) +#define THROW_ON_CUDSS_ERROR(status) CHECK_CUDA_ERROR(status, cudssGetErrorString, true) /** * @brief Macro to check cuDSS status and log a warning on error. @@ -51,8 +50,7 @@ const char *cudssGetErrorString(int status); * If the status indicates an error, this macro will log a warning but will * not throw an exception. */ -#define WARN_ON_CUDSS_ERROR(status) \ - CHECK_CUDA_ERROR(status, cudssGetErrorString, false) +#define WARN_ON_CUDSS_ERROR(status) CHECK_CUDA_ERROR(status, cudssGetErrorString, false) /** * @brief Reusable device memory pool used by cuDSS callbacks. @@ -62,7 +60,7 @@ const char *cudssGetErrorString(int status); * the retained capacity of an available block. */ class cuDSSDeviceMemPool { -public: + public: cuDSSDeviceMemPool() = default; cuDSSDeviceMemPool(const cuDSSDeviceMemPool &) = delete; @@ -91,7 +89,7 @@ class cuDSSDeviceMemPool { */ int Dealloc(void *ptr, size_t size, cudaStream_t stream); -private: + private: struct Block { void *ptr = nullptr; size_t capacity = 0; @@ -105,14 +103,12 @@ class cuDSSDeviceMemPool { /** * @brief C callback wrapper for cuDSS device allocation. */ -int cuDSSDeviceMemPoolAlloc(void *ctx, void **ptr, size_t size, - cudaStream_t stream); +int cuDSSDeviceMemPoolAlloc(void *ctx, void **ptr, size_t size, cudaStream_t stream); /** * @brief C callback wrapper for cuDSS device deallocation. */ -int cuDSSDeviceMemPoolDealloc(void *ctx, void *ptr, size_t size, - cudaStream_t stream); +int cuDSSDeviceMemPoolDealloc(void *ctx, void *ptr, size_t size, cudaStream_t stream); /** * @brief Installs a cuDSS memory handler backed by a custom pool. @@ -139,7 +135,7 @@ void DetachcuDSSDeviceMemHandler(void *handle); * and automatically destroyed in the destructor. */ class cuDSSHandle { -public: + public: cuDSSHandle() = default; cuDSSHandle(const cuDSSHandle &) = delete; @@ -161,9 +157,9 @@ class cuDSSHandle { */ void *GetHandle(cudaStream_t stream); -private: - cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream. - void *handle_ = nullptr; ///< The cuDSS handle. + private: + cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream. + void *handle_ = nullptr; ///< The cuDSS handle. }; /** @@ -173,7 +169,7 @@ class cuDSSHandle { * or a dense vector and manages its lifecycle. */ class cuDSSDescription { -public: + public: /** * @brief Constructs a cuDSS matrix descriptor from a CSR sparse matrix. * @@ -200,8 +196,8 @@ class cuDSSDescription { */ void *GetDescription() { return matrix_; } -private: - void *matrix_; ///< The cuDSS matrix descriptor. + private: + void *matrix_; ///< The cuDSS matrix descriptor. }; /** @@ -211,7 +207,7 @@ class cuDSSDescription { * parameters and options. */ class cuDSSConfig { -public: + public: /** @brief Constructor that creates a cuDSS configuration object. */ cuDSSConfig(int reordering_algorithm = 0, int nthreads = 1); @@ -224,8 +220,8 @@ class cuDSSConfig { */ void *GetData() const { return config_; } -private: - void *config_ = nullptr; ///< The cuDSS configuration handle. + private: + void *config_ = nullptr; ///< The cuDSS configuration handle. }; /** @@ -235,7 +231,7 @@ class cuDSSConfig { * and working memory during the factorization and solve phases. */ class cuDSSData { -public: + public: cuDSSData() = default; /** @brief Destructor that releases the cuDSS data object. */ @@ -253,9 +249,9 @@ class cuDSSData { */ void *GetData(void *handle); -private: - void *handle_ = nullptr; ///< Associated cuDSS handle. - void *data_ = nullptr; ///< The cuDSS data handle. + private: + void *handle_ = nullptr; ///< Associated cuDSS handle. + void *data_ = nullptr; ///< The cuDSS data handle. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cusolver_helper.cpp b/cunls/common/cusolver_helper.cpp index fdec272..6bc8854 100644 --- a/cunls/common/cusolver_helper.cpp +++ b/cunls/common/cusolver_helper.cpp @@ -15,9 +15,10 @@ * limitations under the License. */ +#include "cunls/common/cusolver_helper.h" + #include -#include "cunls/common/cusolver_helper.h" #include "cunls/common/log.h" namespace cunls { @@ -85,8 +86,7 @@ cuSolverHandle::cuSolverHandle() { } cuSolverHandle::~cuSolverHandle() { - WARN_ON_CUSOLVER_ERROR( - cusolverDnDestroy(static_cast(handle_))); + WARN_ON_CUSOLVER_ERROR(cusolverDnDestroy(static_cast(handle_))); } void *cuSolverHandle::GetHandle(cudaStream_t stream) { @@ -101,8 +101,7 @@ void *cuSolverHandle::GetHandle(cudaStream_t stream) { } if (handle_ != nullptr) { - THROW_ON_CUSOLVER_ERROR( - cusolverDnDestroy(static_cast(handle_))); + THROW_ON_CUSOLVER_ERROR(cusolverDnDestroy(static_cast(handle_))); } stream_ = stream; @@ -120,8 +119,7 @@ cuSolverInfo::cuSolverInfo() { } cuSolverInfo::~cuSolverInfo() { - WARN_ON_CUSOLVER_ERROR( - cusolverDnDestroySyevjInfo(static_cast(info_))); + WARN_ON_CUSOLVER_ERROR(cusolverDnDestroySyevjInfo(static_cast(info_))); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cusolver_helper.h b/cunls/common/cusolver_helper.h index 4f9f9f0..f9fab1f 100644 --- a/cunls/common/cusolver_helper.h +++ b/cunls/common/cusolver_helper.h @@ -38,8 +38,7 @@ const char *cusolverGetErrorString(int status); * Checks the cuSolver status and throws std::runtime_error with a descriptive * message if the status indicates an error. */ -#define THROW_ON_CUSOLVER_ERROR(status) \ - CHECK_CUDA_ERROR(status, cusolverGetErrorString, true) +#define THROW_ON_CUSOLVER_ERROR(status) CHECK_CUDA_ERROR(status, cusolverGetErrorString, true) /** * @brief Macro that logs a warning if cuSolver operation fails. @@ -47,8 +46,7 @@ const char *cusolverGetErrorString(int status); * Checks the cuSolver status and logs a warning message if the status indicates * an error, but does not throw an exception. */ -#define WARN_ON_CUSOLVER_ERROR(status) \ - CHECK_CUDA_ERROR(status, cusolverGetErrorString, false) +#define WARN_ON_CUSOLVER_ERROR(status) CHECK_CUDA_ERROR(status, cusolverGetErrorString, false) /** * @brief RAII wrapper for cuSolver handle management. @@ -58,7 +56,7 @@ const char *cusolverGetErrorString(int status); * a specific CUDA stream when GetHandle is called. */ class cuSolverHandle { -public: + public: /// Constructs a cuSolver handle (handle is created lazily on first GetHandle /// call) cuSolverHandle(); @@ -88,9 +86,9 @@ class cuSolverHandle { */ void *GetHandle(cudaStream_t stream); -private: - cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream - void *handle_ = nullptr; ///< cuSolver handle + private: + cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream + void *handle_ = nullptr; ///< cuSolver handle }; /** @@ -100,7 +98,7 @@ class cuSolverHandle { * symmetric eigenvalue decomposition. */ class cuSolverInfo { -public: + public: /// Constructs a cuSolver info object cuSolverInfo(); /// Destroys the cuSolver info object @@ -121,8 +119,8 @@ class cuSolverInfo { */ void *GetInfo() const { return info_; } -private: - void *info_ = nullptr; ///< cuSolver eigenvalue solver info handle + private: + void *info_ = nullptr; ///< cuSolver eigenvalue solver info handle }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cusparse_helper.cpp b/cunls/common/cusparse_helper.cpp index 5b23632..4966dbe 100644 --- a/cunls/common/cusparse_helper.cpp +++ b/cunls/common/cusparse_helper.cpp @@ -15,10 +15,11 @@ * limitations under the License. */ +#include "cunls/common/cusparse_helper.h" + #include #include -#include "cunls/common/cusparse_helper.h" #include "cunls/common/log.h" namespace cunls { @@ -30,8 +31,7 @@ const char *cusparseGetErrorString(int status) { /** @copydoc cuSPARSEHandle::~cuSPARSEHandle */ cuSPARSEHandle::~cuSPARSEHandle() { if (handle_ != nullptr) { - WARN_ON_CUSPARSE_ERROR( - cusparseDestroy(static_cast(handle_))); + WARN_ON_CUSPARSE_ERROR(cusparseDestroy(static_cast(handle_))); } } @@ -48,8 +48,7 @@ void *cuSPARSEHandle::GetHandle(cudaStream_t stream) { } if (handle_ != nullptr) { - THROW_ON_CUSPARSE_ERROR( - cusparseDestroy(static_cast(handle_))); + THROW_ON_CUSPARSE_ERROR(cusparseDestroy(static_cast(handle_))); } stream_ = stream; @@ -69,8 +68,7 @@ cuSPARSEMatrixDescription &cuSPARSEMatrixDescription::operator=( } if (description_) { - WARN_ON_CUSPARSE_ERROR( - cusparseDestroySpMat(static_cast(description_))); + WARN_ON_CUSPARSE_ERROR(cusparseDestroySpMat(static_cast(description_))); } description_ = std::exchange(other.description_, nullptr); @@ -80,36 +78,32 @@ cuSPARSEMatrixDescription &cuSPARSEMatrixDescription::operator=( /** @copydoc * cuSPARSEMatrixDescription::cuSPARSEMatrixDescription(int,int,int,const * CSRSparseMatrix&) */ -cuSPARSEMatrixDescription::cuSPARSEMatrixDescription( - int num_rows, int num_cols, int num_nonzeros, - const CSRSparseMatrix &matrix) { +cuSPARSEMatrixDescription::cuSPARSEMatrixDescription(int num_rows, int num_cols, int num_nonzeros, + const CSRSparseMatrix &matrix) { auto rows_ptr = const_cast(matrix.row_offsets.data()); auto cols_ptr = const_cast(matrix.col_ids.data()); auto values_ptr = const_cast(matrix.values.data()); cusparseSpMatDescr_t descr = nullptr; THROW_ON_CUSPARSE_ERROR(cusparseCreateCsr( - &descr, num_rows, num_cols, num_nonzeros, rows_ptr, cols_ptr, values_ptr, - CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, - CUDA_R_32F)); + &descr, num_rows, num_cols, num_nonzeros, rows_ptr, cols_ptr, values_ptr, CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); description_ = static_cast(descr); }; /** @copydoc cuSPARSEMatrixDescription::cuSPARSEMatrixDescription(int,int) */ -cuSPARSEMatrixDescription::cuSPARSEMatrixDescription(int num_rows, - int num_cols) { +cuSPARSEMatrixDescription::cuSPARSEMatrixDescription(int num_rows, int num_cols) { cusparseSpMatDescr_t descr = nullptr; - THROW_ON_CUSPARSE_ERROR(cusparseCreateCsr( - &descr, num_rows, num_cols, 0, NULL, NULL, NULL, CUSPARSE_INDEX_32I, - CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR(cusparseCreateCsr(&descr, num_rows, num_cols, 0, NULL, NULL, NULL, + CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, + CUSPARSE_INDEX_BASE_ZERO, CUDA_R_32F)); description_ = static_cast(descr); }; /** @copydoc cuSPARSEMatrixDescription::~cuSPARSEMatrixDescription */ cuSPARSEMatrixDescription::~cuSPARSEMatrixDescription() { if (description_) { - WARN_ON_CUSPARSE_ERROR( - cusparseDestroySpMat(static_cast(description_))); + WARN_ON_CUSPARSE_ERROR(cusparseDestroySpMat(static_cast(description_))); } } @@ -119,34 +113,30 @@ void cuSPARSEMatrixDescription::UpdatePointers(const CSRSparseMatrix &matrix) { auto cols_ptr = const_cast(matrix.col_ids.data()); auto values_ptr = const_cast(matrix.values.data()); - THROW_ON_CUSPARSE_ERROR( - cusparseCsrSetPointers(static_cast(description_), - rows_ptr, cols_ptr, values_ptr)); + THROW_ON_CUSPARSE_ERROR(cusparseCsrSetPointers(static_cast(description_), + rows_ptr, cols_ptr, values_ptr)); } /** @copydoc cuSPARSEMatrixDescription::GetDescription */ void *cuSPARSEMatrixDescription::GetDescription() { return description_; } /** @copydoc cuSPARSEVectorDescription::cuSPARSEVectorDescription */ -cuSPARSEVectorDescription::cuSPARSEVectorDescription( - const dvector &vec) { +cuSPARSEVectorDescription::cuSPARSEVectorDescription(const dvector &vec) { auto ptr = const_cast(vec.data()); cusparseDnVecDescr_t descr = nullptr; - THROW_ON_CUSPARSE_ERROR( - cusparseCreateDnVec(&descr, vec.size(), ptr, CUDA_R_32F)); + THROW_ON_CUSPARSE_ERROR(cusparseCreateDnVec(&descr, vec.size(), ptr, CUDA_R_32F)); description_ = static_cast(descr); }; /** @copydoc cuSPARSEVectorDescription::~cuSPARSEVectorDescription */ cuSPARSEVectorDescription::~cuSPARSEVectorDescription() { if (description_) { - WARN_ON_CUSPARSE_ERROR( - cusparseDestroyDnVec(static_cast(description_))); + WARN_ON_CUSPARSE_ERROR(cusparseDestroyDnVec(static_cast(description_))); } } /** @copydoc cuSPARSEVectorDescription::GetDescription */ void *cuSPARSEVectorDescription::GetDescription() { return description_; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/cusparse_helper.h b/cunls/common/cusparse_helper.h index 3210310..a1d5fde 100644 --- a/cunls/common/cusparse_helper.h +++ b/cunls/common/cusparse_helper.h @@ -39,8 +39,7 @@ const char *cusparseGetErrorString(int status); * fails. It uses the cusparseGetErrorString function to provide detailed error * messages. */ -#define THROW_ON_CUSPARSE_ERROR(status) \ - CHECK_CUDA_ERROR(status, cusparseGetErrorString, true) +#define THROW_ON_CUSPARSE_ERROR(status) CHECK_CUDA_ERROR(status, cusparseGetErrorString, true) /** * @brief Macro to check cuSPARSE status and issue a warning on error @@ -49,8 +48,7 @@ const char *cusparseGetErrorString(int status); * throwing. Useful for cleanup operations where exceptions should not be * thrown. */ -#define WARN_ON_CUSPARSE_ERROR(status) \ - CHECK_CUDA_ERROR(status, cusparseGetErrorString, false) +#define WARN_ON_CUSPARSE_ERROR(status) CHECK_CUDA_ERROR(status, cusparseGetErrorString, false) /** * @class cuSPARSEHandle @@ -63,7 +61,7 @@ const char *cusparseGetErrorString(int status); * requested. */ class cuSPARSEHandle { -public: + public: /** * @brief Default constructor * @@ -94,9 +92,9 @@ class cuSPARSEHandle { */ void *GetHandle(cudaStream_t stream); -private: - cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream - void *handle_ = nullptr; ///< The cuSPARSE handle object + private: + cudaStream_t stream_ = nullptr; ///< Currently associated CUDA stream + void *handle_ = nullptr; ///< The cuSPARSE handle object }; /** @@ -109,7 +107,7 @@ class cuSPARSEHandle { * descriptor. */ class cuSPARSEMatrixDescription { -public: + public: /** * @brief Default constructor * @@ -155,8 +153,7 @@ class cuSPARSEMatrixDescription { * @param other The source object to move from * @return Reference to this object */ - cuSPARSEMatrixDescription & - operator=(cuSPARSEMatrixDescription &&other) noexcept; + cuSPARSEMatrixDescription &operator=(cuSPARSEMatrixDescription &&other) noexcept; /** * @brief Destructor @@ -188,8 +185,8 @@ class cuSPARSEMatrixDescription { */ void *GetDescription(); -private: - void *description_ = nullptr; ///< The cuSPARSE matrix descriptor + private: + void *description_ = nullptr; ///< The cuSPARSE matrix descriptor }; /** @@ -201,7 +198,7 @@ class cuSPARSEMatrixDescription { * of float values and is commonly used in sparse matrix-vector operations. */ class cuSPARSEVectorDescription { -public: + public: /** * @brief Constructor from device vector * @@ -231,8 +228,8 @@ class cuSPARSEVectorDescription { */ void *GetDescription(); -private: - void *description_ = nullptr; ///< The cuSPARSE vector descriptor + private: + void *description_ = nullptr; ///< The cuSPARSE vector descriptor }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/device_vector.h b/cunls/common/device_vector.h index eb519b9..d941f71 100644 --- a/cunls/common/device_vector.h +++ b/cunls/common/device_vector.h @@ -36,8 +36,9 @@ namespace cunls { * * @tparam T The element type stored in the vector. Must be trivially copyable. */ -template class DeviceVector { -public: +template +class DeviceVector { + public: /** * @brief Default constructor. Creates an empty vector with no allocated * memory. @@ -69,9 +70,8 @@ template class DeviceVector { if (num_elements > 0) { THROW_ON_CUDA_ERROR(cudaMalloc(&data_, num_elements * sizeof(T))); std::vector host_data(num_elements, fill_value); - THROW_ON_CUDA_ERROR(cudaMemcpy(data_, host_data.data(), - num_elements * sizeof(T), - cudaMemcpyHostToDevice)); + THROW_ON_CUDA_ERROR( + cudaMemcpy(data_, host_data.data(), num_elements * sizeof(T), cudaMemcpyHostToDevice)); } } @@ -82,11 +82,10 @@ template class DeviceVector { T *pinned = nullptr; THROW_ON_CUDA_ERROR(cudaMallocHost(&pinned, num_elements * sizeof(T))); std::fill(pinned, pinned + num_elements, fill_value); - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(data_, pinned, - num_elements * sizeof(T), - cudaMemcpyHostToDevice, stream)); THROW_ON_CUDA_ERROR( - cudaLaunchHostFunc(stream, [](void *p) { cudaFreeHost(p); }, pinned)); + cudaMemcpyAsync(data_, pinned, num_elements * sizeof(T), cudaMemcpyHostToDevice, stream)); + THROW_ON_CUDA_ERROR(cudaLaunchHostFunc( + stream, [](void *p) { cudaFreeHost(p); }, pinned)); } } @@ -100,28 +99,25 @@ template class DeviceVector { * @param host_vector The host vector to copy from. */ explicit DeviceVector(const std::vector &host_vector) - : data_(nullptr), size_(host_vector.size()), - capacity_(host_vector.size()) { + : data_(nullptr), size_(host_vector.size()), capacity_(host_vector.size()) { if (size_ > 0) { THROW_ON_CUDA_ERROR(cudaMalloc(&data_, size_ * sizeof(T))); - THROW_ON_CUDA_ERROR(cudaMemcpy(data_, host_vector.data(), - size_ * sizeof(T), - cudaMemcpyHostToDevice)); + THROW_ON_CUDA_ERROR( + cudaMemcpy(data_, host_vector.data(), size_ * sizeof(T), cudaMemcpyHostToDevice)); } } DeviceVector(const std::vector &host_vector, cudaStream_t stream) - : data_(nullptr), size_(host_vector.size()), - capacity_(host_vector.size()) { + : data_(nullptr), size_(host_vector.size()), capacity_(host_vector.size()) { if (size_ > 0) { THROW_ON_CUDA_ERROR(cudaMalloc(&data_, size_ * sizeof(T))); T *pinned = nullptr; THROW_ON_CUDA_ERROR(cudaMallocHost(&pinned, size_ * sizeof(T))); std::copy(host_vector.begin(), host_vector.end(), pinned); - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(data_, pinned, size_ * sizeof(T), - cudaMemcpyHostToDevice, stream)); THROW_ON_CUDA_ERROR( - cudaLaunchHostFunc(stream, [](void *p) { cudaFreeHost(p); }, pinned)); + cudaMemcpyAsync(data_, pinned, size_ * sizeof(T), cudaMemcpyHostToDevice, stream)); + THROW_ON_CUDA_ERROR(cudaLaunchHostFunc( + stream, [](void *p) { cudaFreeHost(p); }, pinned)); } } @@ -222,12 +218,10 @@ template class DeviceVector { */ void CopyFromHost(const T *src, size_t num_elements) { if (num_elements > capacity_) { - throw std::runtime_error( - "CopyFromHost: num_elements exceeds allocated capacity"); + throw std::runtime_error("CopyFromHost: num_elements exceeds allocated capacity"); } if (num_elements > 0) { - THROW_ON_CUDA_ERROR(cudaMemcpy(data_, src, num_elements * sizeof(T), - cudaMemcpyHostToDevice)); + THROW_ON_CUDA_ERROR(cudaMemcpy(data_, src, num_elements * sizeof(T), cudaMemcpyHostToDevice)); } size_ = num_elements; } @@ -244,8 +238,7 @@ template class DeviceVector { throw std::runtime_error("CopyToHost: num_elements exceeds stored size"); } if (num_elements > 0) { - THROW_ON_CUDA_ERROR(cudaMemcpy(dst, data_, num_elements * sizeof(T), - cudaMemcpyDeviceToHost)); + THROW_ON_CUDA_ERROR(cudaMemcpy(dst, data_, num_elements * sizeof(T), cudaMemcpyDeviceToHost)); } } @@ -257,16 +250,13 @@ template class DeviceVector { * @param stream CUDA stream for the async operation. * @throws std::runtime_error if num_elements exceeds capacity. */ - void CopyFromHostAsync(const T *src, size_t num_elements, - CudaStream &stream) { + void CopyFromHostAsync(const T *src, size_t num_elements, CudaStream &stream) { if (num_elements > capacity_) { - throw std::runtime_error( - "CopyFromHostAsync: num_elements exceeds allocated capacity"); + throw std::runtime_error("CopyFromHostAsync: num_elements exceeds allocated capacity"); } if (num_elements > 0) { THROW_ON_CUDA_ERROR(cudaMemcpyAsync(data_, src, num_elements * sizeof(T), - cudaMemcpyHostToDevice, - stream.GetStream())); + cudaMemcpyHostToDevice, stream.GetStream())); } size_ = num_elements; } @@ -281,13 +271,11 @@ template class DeviceVector { */ void CopyToHostAsync(T *dst, size_t num_elements, CudaStream &stream) const { if (num_elements > size_) { - throw std::runtime_error( - "CopyToHostAsync: num_elements exceeds stored size"); + throw std::runtime_error("CopyToHostAsync: num_elements exceeds stored size"); } if (num_elements > 0) { THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dst, data_, num_elements * sizeof(T), - cudaMemcpyDeviceToHost, - stream.GetStream())); + cudaMemcpyDeviceToHost, stream.GetStream())); } } @@ -314,8 +302,7 @@ template class DeviceVector { // Copy existing data if requested if (preserve_data && data_ != nullptr && size_ > 0) { - THROW_ON_CUDA_ERROR(cudaMemcpy(new_data, data_, size_ * sizeof(T), - cudaMemcpyDeviceToDevice)); + THROW_ON_CUDA_ERROR(cudaMemcpy(new_data, data_, size_ * sizeof(T), cudaMemcpyDeviceToDevice)); } // Free old memory @@ -346,8 +333,7 @@ template class DeviceVector { // Copy existing data if (data_ != nullptr && size_ > 0) { - THROW_ON_CUDA_ERROR(cudaMemcpy(new_data, data_, size_ * sizeof(T), - cudaMemcpyDeviceToDevice)); + THROW_ON_CUDA_ERROR(cudaMemcpy(new_data, data_, size_ * sizeof(T), cudaMemcpyDeviceToDevice)); } // Free old memory @@ -387,8 +373,7 @@ template class DeviceVector { THROW_ON_CUDA_ERROR(cudaMalloc(&new_data, size_ * sizeof(T))); // Copy data - THROW_ON_CUDA_ERROR(cudaMemcpy(new_data, data_, size_ * sizeof(T), - cudaMemcpyDeviceToDevice)); + THROW_ON_CUDA_ERROR(cudaMemcpy(new_data, data_, size_ * sizeof(T), cudaMemcpyDeviceToDevice)); // Free old memory WARN_ON_CUDA_ERROR(cudaFree(data_)); @@ -397,10 +382,10 @@ template class DeviceVector { capacity_ = size_; } -private: - T *data_; ///< Pointer to device memory - size_t size_; ///< Number of elements stored - size_t capacity_; ///< Number of elements allocated + private: + T *data_; ///< Pointer to device memory + size_t size_; ///< Number of elements stored + size_t capacity_; ///< Number of elements allocated }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/helper.h b/cunls/common/helper.h index 2316fea..d22adf1 100644 --- a/cunls/common/helper.h +++ b/cunls/common/helper.h @@ -38,31 +38,29 @@ namespace cunls { * @param is_throw Compile-time boolean: if true, throws on error; if false, * only logs. */ -#define CHECK_CUDA_ERROR(status, to_string_fn, is_throw) \ - do { \ - auto ret = (status); \ - if (ret != 0) { \ - std::stringstream msg; \ - msg << "[CUDA] error " << (to_string_fn)(ret) << "(" << ret << ")"; \ - msg << " in " << LOCATION << std::endl; \ - LogError(msg.str()); \ - if constexpr (is_throw) { \ - throw std::runtime_error(msg.str()); \ - } \ - } \ +#define CHECK_CUDA_ERROR(status, to_string_fn, is_throw) \ + do { \ + auto ret = (status); \ + if (ret != 0) { \ + std::stringstream msg; \ + msg << "[CUDA] error " << (to_string_fn)(ret) << "(" << ret << ")"; \ + msg << " in " << LOCATION << std::endl; \ + LogError(msg.str()); \ + if constexpr (is_throw) { \ + throw std::runtime_error(msg.str()); \ + } \ + } \ } while (0) /** * @brief Throws std::runtime_error on CUDA runtime API errors. * @param status cudaError_t value returned by a CUDA runtime call. */ -#define THROW_ON_CUDA_ERROR(status) \ - CHECK_CUDA_ERROR(status, cudaGetErrorString, true) +#define THROW_ON_CUDA_ERROR(status) CHECK_CUDA_ERROR(status, cudaGetErrorString, true) /** * @brief Logs a warning on CUDA runtime API errors without throwing. * @param status cudaError_t value returned by a CUDA runtime call. */ -#define WARN_ON_CUDA_ERROR(status) \ - CHECK_CUDA_ERROR(status, cudaGetErrorString, false) -} // namespace cunls +#define WARN_ON_CUDA_ERROR(status) CHECK_CUDA_ERROR(status, cudaGetErrorString, false) +} // namespace cunls diff --git a/cunls/common/log.cpp b/cunls/common/log.cpp index 94cc09b..32dc5b7 100644 --- a/cunls/common/log.cpp +++ b/cunls/common/log.cpp @@ -39,28 +39,28 @@ namespace { spdlog::level::level_enum GetLogLevel(Verbosity verbosity) { spdlog::level::level_enum log_level; switch (verbosity) { - case Verbosity::Error: - log_level = spdlog::level::err; - break; - case Verbosity::Warning: - log_level = spdlog::level::warn; - break; - case Verbosity::Message: - log_level = spdlog::level::info; - break; - case Verbosity::Debug: - log_level = spdlog::level::debug; - break; - case Verbosity::Silent: - log_level = spdlog::level::off; - break; - default: - log_level = spdlog::level::off; - break; + case Verbosity::Error: + log_level = spdlog::level::err; + break; + case Verbosity::Warning: + log_level = spdlog::level::warn; + break; + case Verbosity::Message: + log_level = spdlog::level::info; + break; + case Verbosity::Debug: + log_level = spdlog::level::debug; + break; + case Verbosity::Silent: + log_level = spdlog::level::off; + break; + default: + log_level = spdlog::level::off; + break; } return log_level; } -} // namespace +} // namespace /** @copydoc SetLoggerOptions */ void SetLoggerOptions(Verbosity verbosity, Sink sink, const std::string &path) { @@ -72,8 +72,7 @@ void SetLoggerOptions(Verbosity verbosity, Sink sink, const std::string &path) { } if (path.empty()) { - spdlog::warn( - "Empty path provided to the logger. Using the default logger."); + spdlog::warn("Empty path provided to the logger. Using the default logger."); spdlog::set_level(log_level); return; } @@ -90,8 +89,7 @@ void SetLoggerOptions(Verbosity verbosity, Sink sink, const std::string &path) { sinks.push_back(file_sink); } - auto logger = - std::make_shared("logger", begin(sinks), end(sinks)); + auto logger = std::make_shared("logger", begin(sinks), end(sinks)); logger->set_level(log_level); spdlog::set_default_logger(logger); @@ -124,4 +122,4 @@ void LogMessage(std::string_view msg) { Log(Verbosity::Message, msg); } /** @copydoc LogDebug(std::string_view) */ void LogDebug(std::string_view msg) { Log(Verbosity::Debug, msg); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/log.h b/cunls/common/log.h index eb8bdc5..b570b0b 100644 --- a/cunls/common/log.h +++ b/cunls/common/log.h @@ -115,8 +115,7 @@ namespace { * @param target_index Runtime index of the argument to output. */ template -void format_arg_at_index(std::ostringstream &oss, const Tuple &args, - size_t target_index) { +void format_arg_at_index(std::ostringstream &oss, const Tuple &args, size_t target_index) { if (Index == target_index && Index < std::tuple_size_v) { oss << std::get(args); } @@ -134,8 +133,7 @@ void format_arg_at_index(std::ostringstream &oss, const Tuple &args, * @param target_index Runtime index of the argument to output. */ template -void format_arg_recursive(std::ostringstream &oss, const Tuple &args, - size_t target_index, +void format_arg_recursive(std::ostringstream &oss, const Tuple &args, size_t target_index, std::index_sequence) { (format_arg_at_index(oss, args, target_index), ...); } @@ -178,14 +176,12 @@ std::string vformat_to(std::string_view fmt_str, Args &&...args) { } // Extract placeholder content - std::string placeholder = - fmt.substr(open_brace + 1, close_brace - open_brace - 1); + std::string placeholder = fmt.substr(open_brace + 1, close_brace - open_brace - 1); if (placeholder.empty()) { // Empty placeholder {} - use next argument in order if (arg_index < num_args) { - format_arg_recursive(result, args_tuple, arg_index, - std::make_index_sequence{}); + format_arg_recursive(result, args_tuple, arg_index, std::make_index_sequence{}); arg_index++; } else { result << "{}"; @@ -215,7 +211,7 @@ std::string vformat_to(std::string_view fmt_str, Args &&...args) { return result.str(); } -} // namespace +} // namespace /** * @brief Logs a formatted message at the given verbosity level (C++17 path). @@ -281,4 +277,4 @@ template inline void LogDebug(std::string_view fmt, Args &&...args) { Log(Verbosity::Debug, fmt, std::forward(args)...); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/pinned_vector.h b/cunls/common/pinned_vector.h index 3b7fd12..be0d19b 100644 --- a/cunls/common/pinned_vector.h +++ b/cunls/common/pinned_vector.h @@ -45,8 +45,9 @@ namespace cunls { * * @tparam T Element type. Must be trivially copyable. */ -template class PinnedVector { -public: +template +class PinnedVector { + public: /** @brief Default constructor. Creates an empty vector with no allocation. */ PinnedVector() : data_(nullptr), size_(0), capacity_(0) {} @@ -163,10 +164,10 @@ template class PinnedVector { size_ = new_size; } -private: - T *data_; ///< Pointer to page-locked host memory. - size_t size_; ///< Number of elements stored. - size_t capacity_; ///< Number of elements allocated. + private: + T *data_; ///< Pointer to page-locked host memory. + size_t size_; ///< Number of elements stored. + size_t capacity_; ///< Number of elements allocated. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/profiler.cpp b/cunls/common/profiler.cpp index 48c8ced..e0efd03 100644 --- a/cunls/common/profiler.cpp +++ b/cunls/common/profiler.cpp @@ -18,9 +18,8 @@ #include "cunls/common/profiler.h" #ifdef ENABLE_PROFILING -#include - #include +#include #endif namespace cunls::profiler { @@ -58,11 +57,9 @@ DomainRange::~DomainRange() { nvtxDomainRangePop((nvtxDomainHandle_t)handle_); } } -} // namespace internal +} // namespace internal -ScopedRange::ScopedRange(const std::string &name) : name_(name) { - nvtxRangePushA(name_.c_str()); -}; +ScopedRange::ScopedRange(const std::string &name) : name_(name) { nvtxRangePushA(name_.c_str()); }; /** * @brief Destructor that pops the range from the NVTX stack @@ -73,8 +70,7 @@ ScopedRange::~ScopedRange() { nvtxRangePop(); } * @brief Constructs a new profiling domain with a random color * @param name The name of the domain (visible in profiling tools) */ -Domain::Domain(const std::string &name) - : name_(name), handle_(nvtxDomainCreateA(name_.c_str())) { +Domain::Domain(const std::string &name) : name_(name), handle_(nvtxDomainCreateA(name_.c_str())) { std::random_device rd; std::mt19937 gen(rd()); @@ -104,14 +100,13 @@ namespace internal { * @param name The name to display in the profiler for this range * @param color The color for this range (ARGB format, default is 0) */ -DomainRange::DomainRange(void *handle, const std::string &name, - uint32_t color) {} +DomainRange::DomainRange(void *handle, const std::string &name, uint32_t color) {} /** * @brief Destructor that pops the range from the domain stack */ DomainRange::~DomainRange() {} -} // namespace internal +} // namespace internal ScopedRange::ScopedRange(const std::string &name) {} @@ -141,4 +136,4 @@ internal::DomainRange Domain::CreateDomainRange(const std::string &name) const { } #endif -} // namespace cunls::profiler +} // namespace cunls::profiler diff --git a/cunls/common/profiler.h b/cunls/common/profiler.h index 20504d0..5822696 100644 --- a/cunls/common/profiler.h +++ b/cunls/common/profiler.h @@ -29,7 +29,7 @@ namespace cunls::profiler { * and non-movable to ensure one-to-one scope-to-range mapping. */ class ScopedRange { -public: + public: ScopedRange(const ScopedRange &) = delete; ScopedRange &operator=(const ScopedRange &) = delete; ScopedRange(ScopedRange &&) = delete; @@ -47,8 +47,8 @@ class ScopedRange { */ ~ScopedRange(); -private: - std::string name_; ///< Label for the profiling range. + private: + std::string name_; ///< Label for the profiling range. }; namespace internal { @@ -60,7 +60,7 @@ namespace internal { * and ends it on destruction. Used internally by Domain::CreateDomainRange. */ class DomainRange { -public: + public: DomainRange(const DomainRange &) = delete; DomainRange &operator=(const DomainRange &) = delete; DomainRange(DomainRange &&) = delete; @@ -80,12 +80,12 @@ class DomainRange { */ ~DomainRange(); -private: - void *handle_ = nullptr; ///< Handle to the NVTX domain. - std::string name_; ///< Copy of the name string to ensure lifetime. + private: + void *handle_ = nullptr; ///< Handle to the NVTX domain. + std::string name_; ///< Copy of the name string to ensure lifetime. }; -} // namespace internal +} // namespace internal /** * @brief NVTX profiling domain with automatic color cycling. @@ -97,7 +97,7 @@ class DomainRange { * Non-copyable and non-movable to ensure unique domain ownership. */ class Domain { -public: + public: Domain(const Domain &) = delete; Domain &operator=(const Domain &) = delete; Domain(Domain &&) = delete; @@ -126,9 +126,9 @@ class Domain { */ internal::DomainRange CreateDomainRange(const std::string &name) const; -private: - std::string name_; ///< Copy of the name string to ensure lifetime. - void *handle_ = nullptr; ///< Handle to the NVTX domain. - uint32_t color_ = 0; ///< Current color counter for range cycling. + private: + std::string name_; ///< Copy of the name string to ensure lifetime. + void *handle_ = nullptr; ///< Handle to the NVTX domain. + uint32_t color_ = 0; ///< Current color counter for range cycling. }; -} // namespace cunls::profiler +} // namespace cunls::profiler diff --git a/cunls/common/type_traits.h b/cunls/common/type_traits.h index 1f23aae..c0adcb6 100644 --- a/cunls/common/type_traits.h +++ b/cunls/common/type_traits.h @@ -45,7 +45,6 @@ struct DerivedFromAnySizedFactorBatchHelper { */ template struct IsDerivedFromAnySizedFactorBatch - : decltype(DerivedFromAnySizedFactorBatchHelper::test( - std::declval())){}; + : decltype(DerivedFromAnySizedFactorBatchHelper::test(std::declval())) {}; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/types.h b/cunls/common/types.h index 8510f6b..d32223d 100644 --- a/cunls/common/types.h +++ b/cunls/common/types.h @@ -109,8 +109,6 @@ struct CSRMatrixDimensions { int num_cols = -1; int num_nonzeros = -1; - bool IsValid() const { return num_rows >= 0; } - void Set(int rows, int cols, int nnz) { num_rows = rows; num_cols = cols; @@ -134,7 +132,7 @@ struct CSRMatrixDimensions { * - values: non-zero values (size = num_nonzeros). */ struct CSRSparseMatrix { - dvector row_offsets; ///< Row offset array (num_rows + 1 entries). + dvector row_offsets; ///< Row offset array (num_rows + 1 entries). dvector col_ids; ///< Column index array (num_nonzeros entries). dvector values; ///< Non-zero value array (num_nonzeros entries). @@ -172,11 +170,11 @@ struct CSRSparseMatrix { * `block_size`; see ChooseHessianBlockSize(). */ struct BSRSparseMatrix { - dvector row_offsets; ///< Block-row offsets (num_block_rows + 1 entries). + dvector row_offsets; ///< Block-row offsets (num_block_rows + 1 entries). dvector col_ids; ///< Block-column index per tile. dvector values; ///< Tiles, row-major, block_size^2 floats each. - int block_size = 1; ///< Tile edge length. + int block_size = 1; ///< Tile edge length. /** * @brief Largest number of tiles in any block row. * @@ -186,7 +184,7 @@ struct BSRSparseMatrix { * of that camera, a landmark row a handful). One schedule cannot serve both. */ int max_tiles_per_row = 0; - int num_block_rows = 0; ///< Number of block rows (= block columns). + int num_block_rows = 0; ///< Number of block rows (= block columns). /** @brief Number of stored tiles. */ size_t NumBlocks() const { return col_ids.size(); } @@ -209,4 +207,4 @@ struct BSRSparseMatrix { */ using PerFactorJacobians = dvector; -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/utils.cpp b/cunls/common/utils.cpp index 7646cd4..2b2a11a 100644 --- a/cunls/common/utils.cpp +++ b/cunls/common/utils.cpp @@ -15,25 +15,23 @@ * limitations under the License. */ +#include "cunls/common/utils.h" + #include #include #include #include -#include "cunls/common/utils.h" - namespace cunls { /** @copydoc DumpCSRSparseMatrixToFile */ -void DumpCSRSparseMatrixToFile(const std::string &filename, - const CSRSparseMatrix &matrix) { +void DumpCSRSparseMatrixToFile(const std::string &filename, const CSRSparseMatrix &matrix) { // Copy device vectors to host vectors hvector host_row_offsets(matrix.row_offsets.size()); hvector host_col_ids(matrix.col_ids.size()); hvector host_values(matrix.values.size()); - matrix.row_offsets.CopyToHost(host_row_offsets.data(), - host_row_offsets.size()); + matrix.row_offsets.CopyToHost(host_row_offsets.data(), host_row_offsets.size()); matrix.col_ids.CopyToHost(host_col_ids.data(), host_col_ids.size()); matrix.values.CopyToHost(host_values.data(), host_values.size()); @@ -42,9 +40,7 @@ void DumpCSRSparseMatrixToFile(const std::string &filename, size_t num_nonzeros = matrix.NumNonZeros(); size_t num_cols = 0; if (num_nonzeros > 0) { - num_cols = *std::max_element(host_col_ids.begin(), - host_col_ids.begin() + num_nonzeros) + - 1; + num_cols = *std::max_element(host_col_ids.begin(), host_col_ids.begin() + num_nonzeros) + 1; } // Open file in append binary mode @@ -60,8 +56,7 @@ void DumpCSRSparseMatrixToFile(const std::string &filename, file.write(reinterpret_cast(&num_rows_u32), sizeof(uint32_t)); file.write(reinterpret_cast(&num_cols_u32), sizeof(uint32_t)); - file.write(reinterpret_cast(&num_nonzeros_u32), - sizeof(uint32_t)); + file.write(reinterpret_cast(&num_nonzeros_u32), sizeof(uint32_t)); // Convert and write row_offsets (int -> uint32_t) for (size_t i = 0; i < host_row_offsets.size(); ++i) { @@ -76,15 +71,13 @@ void DumpCSRSparseMatrixToFile(const std::string &filename, } // Write values (float) - file.write(reinterpret_cast(host_values.data()), - num_nonzeros * sizeof(float)); + file.write(reinterpret_cast(host_values.data()), num_nonzeros * sizeof(float)); file.close(); } /** @copydoc DumpVectorToFile */ -void DumpVectorToFile(const std::string &filename, - const dvector &vector) { +void DumpVectorToFile(const std::string &filename, const dvector &vector) { // Copy device vector to host vector hvector host_vector(vector.size()); vector.CopyToHost(host_vector.data(), host_vector.size()); @@ -98,10 +91,9 @@ void DumpVectorToFile(const std::string &filename, // Write binary data according to the format uint32_t size = static_cast(vector.size()); file.write(reinterpret_cast(&size), sizeof(uint32_t)); - file.write(reinterpret_cast(host_vector.data()), - size * sizeof(float)); + file.write(reinterpret_cast(host_vector.data()), size * sizeof(float)); file.close(); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/common/utils.h b/cunls/common/utils.h index 12ec7bb..7221d3b 100644 --- a/cunls/common/utils.h +++ b/cunls/common/utils.h @@ -40,8 +40,7 @@ namespace cunls { * @param filename Path to the output binary file (created or appended to). * @param matrix The CSR sparse matrix to dump. */ -void DumpCSRSparseMatrixToFile(const std::string &filename, - const CSRSparseMatrix &matrix); +void DumpCSRSparseMatrixToFile(const std::string &filename, const CSRSparseMatrix &matrix); /** * @brief Dumps a device vector to a binary file for debugging/analysis. @@ -56,7 +55,6 @@ void DumpCSRSparseMatrixToFile(const std::string &filename, * @param filename Path to the output binary file (created or appended to). * @param vector The device vector to dump. */ -void DumpVectorToFile(const std::string &filename, - const dvector &vector); +void DumpVectorToFile(const std::string &filename, const dvector &vector); -} // namespace cunls +} // namespace cunls diff --git a/cunls/factor/factor_batch.h b/cunls/factor/factor_batch.h index f31639b..6f9ca73 100644 --- a/cunls/factor/factor_batch.h +++ b/cunls/factor/factor_batch.h @@ -36,7 +36,7 @@ namespace cunls { * state block sizes at compile time. */ class FactorBatch { -public: + public: /** * @brief Evaluates residuals and optionally Jacobians for all factors * in the batch. @@ -51,8 +51,7 @@ class FactorBatch { * @param stream CUDA stream for asynchronous execution. * @return true if evaluation succeeded, false otherwise. */ - virtual bool Evaluate(float *residuals, float *jacobians, - float const *const *state_pointers, + virtual bool Evaluate(float *residuals, float *jacobians, float const *const *state_pointers, cudaStream_t stream) const = 0; /** @brief Virtual destructor for safe polymorphic deletion. */ @@ -77,4 +76,4 @@ class FactorBatch { virtual size_t NumFactors() const = 0; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/factor/information_factor_batch.cpp b/cunls/factor/information_factor_batch.cpp index 203ea8c..69dbf29 100644 --- a/cunls/factor/information_factor_batch.cpp +++ b/cunls/factor/information_factor_batch.cpp @@ -15,43 +15,40 @@ * limitations under the License. */ +#include "cunls/factor/information_factor_batch.h" + #include #include "cunls/common/cublas_helper.h" -#include "cunls/factor/information_factor_batch.h" namespace cunls { -void ApplyInformationToResiduals(void *cublas_handle, - const float *sqrt_information, - float *residuals, size_t residual_size, - size_t num_factors) { +void ApplyInformationToResiduals(void *cublas_handle, const float *sqrt_information, + float *residuals, size_t residual_size, size_t num_factors) { constexpr float alpha = 1.0f; constexpr float beta = 0.0f; const size_t stride = residual_size * residual_size; constexpr size_t inc = 1; THROW_ON_CUBLAS_ERROR(cublasSgemvStridedBatched( - static_cast(cublas_handle), CUBLAS_OP_N, residual_size, - residual_size, &alpha, sqrt_information, residual_size, stride, residuals, - inc, residual_size, &beta, residuals, inc, residual_size, num_factors)); + static_cast(cublas_handle), CUBLAS_OP_N, residual_size, residual_size, &alpha, + sqrt_information, residual_size, stride, residuals, inc, residual_size, &beta, residuals, inc, + residual_size, num_factors)); } -void ApplyInformationToJacobians(void *cublas_handle, - const float *sqrt_information, - float *jacobians, size_t residual_size, - size_t jacobian_pitch, size_t num_factors) { +void ApplyInformationToJacobians(void *cublas_handle, const float *sqrt_information, + float *jacobians, size_t residual_size, size_t jacobian_pitch, + size_t num_factors) { constexpr float alpha = 1.0f; constexpr float beta = 0.0f; const size_t info_stride = residual_size * residual_size; const size_t jacobian_stride = jacobian_pitch * residual_size; THROW_ON_CUBLAS_ERROR(cublasSgemmStridedBatched( - static_cast(cublas_handle), CUBLAS_OP_N, CUBLAS_OP_N, - jacobian_pitch, residual_size, residual_size, &alpha, jacobians, - jacobian_pitch, jacobian_stride, sqrt_information, residual_size, - info_stride, &beta, jacobians, jacobian_pitch, jacobian_stride, - num_factors)); + static_cast(cublas_handle), CUBLAS_OP_N, CUBLAS_OP_N, jacobian_pitch, + residual_size, residual_size, &alpha, jacobians, jacobian_pitch, jacobian_stride, + sqrt_information, residual_size, info_stride, &beta, jacobians, jacobian_pitch, + jacobian_stride, num_factors)); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/factor/information_factor_batch.h b/cunls/factor/information_factor_batch.h index 1a5a9fe..6760b4a 100644 --- a/cunls/factor/information_factor_batch.h +++ b/cunls/factor/information_factor_batch.h @@ -41,10 +41,8 @@ namespace cunls { * @param residual_size Dimension of each residual / information matrix. * @param num_factors Number of factors in the batch. */ -void ApplyInformationToResiduals(void *cublas_handle, - const float *sqrt_information, - float *residuals, size_t residual_size, - size_t num_factors); +void ApplyInformationToResiduals(void *cublas_handle, const float *sqrt_information, + float *residuals, size_t residual_size, size_t num_factors); /** * @brief Applies sqrt-information matrices to a batch of Jacobian matrices. @@ -60,10 +58,9 @@ void ApplyInformationToResiduals(void *cublas_handle, * Jacobian. * @param num_factors Number of factors in the batch. */ -void ApplyInformationToJacobians(void *cublas_handle, - const float *sqrt_information, - float *jacobians, size_t residual_size, - size_t jacobian_pitch, size_t num_factors); +void ApplyInformationToJacobians(void *cublas_handle, const float *sqrt_information, + float *jacobians, size_t residual_size, size_t jacobian_pitch, + size_t num_factors); /** * @brief Wrapper factor that applies square-root information matrices. @@ -82,10 +79,9 @@ void ApplyInformationToJacobians(void *cublas_handle, * remain valid for the lifetime of this object. The memory layout is: [mat0: * residual_size^2 floats][mat1: residual_size^2 floats]... */ -template ::value, int> = 0> +template ::value, int> = 0> class InformationFactorBatch : public T::sized_layout { -public: + public: using InformationMatrix = Matrix; /** @@ -115,8 +111,7 @@ class InformationFactorBatch : public T::sized_layout { if (num_matrices_ != factor_batch_.NumFactors()) { std::stringstream ss; ss << "Number of sqrt information matrices (" << num_matrices_ - << ") must match wrapped factor batch size (" - << factor_batch_.NumFactors() << ")"; + << ") must match wrapped factor batch size (" << factor_batch_.NumFactors() << ")"; LogError(ss.str()); throw std::invalid_argument(ss.str()); } @@ -145,19 +140,16 @@ class InformationFactorBatch : public T::sized_layout { * @param stream CUDA stream for asynchronous execution * @return true if evaluation succeeded, false otherwise */ - bool Evaluate(float *residuals, float *jacobians, - float const *const *state_pointers, + bool Evaluate(float *residuals, float *jacobians, float const *const *state_pointers, cudaStream_t stream) const final { factor_batch_.Evaluate(residuals, jacobians, state_pointers, stream); auto handle = cublas_handle_.GetHandle(stream); - auto info_ptr = - reinterpret_cast(sqrt_information_matrices_ptr_); + auto info_ptr = reinterpret_cast(sqrt_information_matrices_ptr_); const size_t rsize = T::residual_size_; const size_t num_factors = factor_batch_.NumFactors(); - ApplyInformationToResiduals(handle, info_ptr, residuals, rsize, - num_factors); + ApplyInformationToResiduals(handle, info_ptr, residuals, rsize, num_factors); if (jacobians == nullptr) { return true; @@ -167,14 +159,13 @@ class InformationFactorBatch : public T::sized_layout { const size_t jacobian_pitch = std::accumulate(state_block_sizes.begin(), state_block_sizes.end(), 0); - ApplyInformationToJacobians(handle, info_ptr, jacobians, rsize, - jacobian_pitch, num_factors); + ApplyInformationToJacobians(handle, info_ptr, jacobians, rsize, jacobian_pitch, num_factors); return true; } -private: - T factor_batch_; ///< Wrapped factor batch + private: + T factor_batch_; ///< Wrapped factor batch /// Pointer to user-managed device memory containing square-root information /// matrices. @@ -183,7 +174,7 @@ class InformationFactorBatch : public T::sized_layout { /// Number of per-factor square-root information matrices (equals batch size). size_t num_matrices_; - cuBLASHandle &cublas_handle_; ///< cuBLAS handle for matrix operations + cuBLASHandle &cublas_handle_; ///< cuBLAS handle for matrix operations }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/factor/pnp_factor_batch.cu b/cunls/factor/pnp_factor_batch.cu index fd29c5f..f8ed1a5 100644 --- a/cunls/factor/pnp_factor_batch.cu +++ b/cunls/factor/pnp_factor_batch.cu @@ -36,15 +36,12 @@ constexpr size_t kPnPBlockSize = 256; * Replaces: pnp_collect_poses_kernel + cuBLAS SGEMM (or memcpy) + * pnp_cost_kernel. */ -__global__ void pnp_fused_kernel(const Vector<2> *observations, - const Vector<3> *points_world, +__global__ void pnp_fused_kernel(const Vector<2> *observations, const Vector<3> *points_world, float const *const *state_pointers, - const SE3Transform *poses_camera_from_rig, - float *residuals, float *jacobians, - float z_threshold, int num_observations) { + const SE3Transform *poses_camera_from_rig, float *residuals, + float *jacobians, float z_threshold, int num_observations) { int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_observations) - return; + if (tid >= num_observations) return; constexpr int kResidualDim = 2; constexpr int kJacobianCols = 6; @@ -79,8 +76,7 @@ __global__ void pnp_fused_kernel(const Vector<2> *observations, pose[11] = e20 * r03 + e21 * r13 + e22 * r23 + e23; } else { #pragma unroll - for (int i = 0; i < 12; i++) - pose[i] = rig[i]; + for (int i = 0; i < 12; i++) pose[i] = rig[i]; } float point_cam[3]; @@ -110,8 +106,7 @@ __global__ void pnp_fused_kernel(const Vector<2> *observations, if (point_cam[2] < z_threshold) { #pragma unroll - for (int i = 0; i < kJacobianBlockSize; i++) - jac_ptr[i] = 0.0f; + for (int i = 0; i < kJacobianBlockSize; i++) jac_ptr[i] = 0.0f; return; } @@ -152,23 +147,25 @@ __global__ void pnp_fused_kernel(const Vector<2> *observations, } } -PnPFactorBatch::PnPFactorBatch(const Vector<2> *observations, - const Vector<3> *points_world, +PnPFactorBatch::PnPFactorBatch(const Vector<2> *observations, const Vector<3> *points_world, size_t num_observations, float z_threshold) - : observations_(observations), points_world_(points_world), - num_observations_(num_observations), z_threshold_(z_threshold) {} + : observations_(observations), + points_world_(points_world), + num_observations_(num_observations), + z_threshold_(z_threshold) {} PnPFactorBatch::PnPFactorBatch(const Vector<2> *observations, const SE3Transform *poses_camera_from_rig, - const Vector<3> *points_world, - size_t num_observations, float z_threshold) - : observations_(observations), points_world_(points_world), + const Vector<3> *points_world, size_t num_observations, + float z_threshold) + : observations_(observations), + points_world_(points_world), poses_camera_from_rig_(poses_camera_from_rig), - num_observations_(num_observations), z_threshold_(z_threshold) {} + num_observations_(num_observations), + z_threshold_(z_threshold) {} bool PnPFactorBatch::Evaluate(float *residuals, float *jacobians, - float const *const *state_pointers, - cudaStream_t stream) const { + float const *const *state_pointers, cudaStream_t stream) const { if (num_observations_ == 0) { return true; } @@ -176,11 +173,11 @@ bool PnPFactorBatch::Evaluate(float *residuals, float *jacobians, size_t num_blocks = (num_observations_ + kPnPBlockSize - 1) / kPnPBlockSize; pnp_fused_kernel<<>>( - observations_, points_world_, state_pointers, poses_camera_from_rig_, - residuals, jacobians, z_threshold_, static_cast(num_observations_)); + observations_, points_world_, state_pointers, poses_camera_from_rig_, residuals, jacobians, + z_threshold_, static_cast(num_observations_)); THROW_ON_CUDA_ERROR(cudaGetLastError()); return true; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/factor/pnp_factor_batch.h b/cunls/factor/pnp_factor_batch.h index c618b14..2063279 100644 --- a/cunls/factor/pnp_factor_batch.h +++ b/cunls/factor/pnp_factor_batch.h @@ -37,7 +37,7 @@ namespace cunls { * - 6: SE(3) pose tangent */ class PnPFactorBatch : public SizedFactorBatch<2, 6> { -public: + public: /** * @brief Constructs with identity camera-from-rig extrinsics. */ @@ -47,18 +47,15 @@ class PnPFactorBatch : public SizedFactorBatch<2, 6> { /** * @brief Constructs with per-correspondence camera-from-rig transforms. */ - PnPFactorBatch(const Vector<2> *observations, - const SE3Transform *poses_camera_from_rig, - const Vector<3> *points_world, size_t num_observations, - float z_threshold = 1e-3f); + PnPFactorBatch(const Vector<2> *observations, const SE3Transform *poses_camera_from_rig, + const Vector<3> *points_world, size_t num_observations, float z_threshold = 1e-3f); - bool Evaluate(float *residuals, float *jacobians, - float const *const *state_pointers, + bool Evaluate(float *residuals, float *jacobians, float const *const *state_pointers, cudaStream_t stream) const final; size_t NumFactors() const final { return num_observations_; } -private: + private: PnPFactorBatch() = delete; const Vector<2> *observations_; @@ -68,4 +65,4 @@ class PnPFactorBatch : public SizedFactorBatch<2, 6> { float z_threshold_ = 1e-3f; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/factor/point_to_plane_factor_batch.cu b/cunls/factor/point_to_plane_factor_batch.cu index b0d6ae2..55d6e62 100644 --- a/cunls/factor/point_to_plane_factor_batch.cu +++ b/cunls/factor/point_to_plane_factor_batch.cu @@ -135,12 +135,10 @@ constexpr size_t kBlockSize = 256; * @note Launch configuration: <<>> */ -__global__ void point_to_plane_cost_kernel(const float *p_observations, - const float *q_observations, +__global__ void point_to_plane_cost_kernel(const float *p_observations, const float *q_observations, const float *nq_observations, - float const *const *state_pointers, - float *residuals, float *jacobians, - int num_correspondences) { + float const *const *state_pointers, float *residuals, + float *jacobians, int num_correspondences) { int tid = threadIdx.x + blockIdx.x * blockDim.x; if (tid >= num_correspondences) { return; @@ -231,11 +229,10 @@ bool PointToPlaneFactorBatch::Evaluate(float *residuals, float *jacobians, size_t num_blocks = (num_factors + kBlockSize - 1) / kBlockSize; point_to_plane_cost_kernel<<>>( - p_data_ptr, q_data_ptr, nq_data_ptr, state_pointers, residuals, jacobians, - num_factors); + p_data_ptr, q_data_ptr, nq_data_ptr, state_pointers, residuals, jacobians, num_factors); THROW_ON_CUDA_ERROR(cudaGetLastError()); return true; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index b7857a2..68b7fd6 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -1151,7 +1151,7 @@ void DualDotAsync(cudaStream_t stream, const float *a, const float *b, const flo DualDotKernel<<>>(a, b, c, n, out_ab, out_ac); } -} // namespace +} // namespace // ============================================================================= // BlockSparsePCGSolver @@ -1532,4 +1532,4 @@ bool BlockSparsePCGSolver::SolveCommon(cudaStream_t stream, int n, const dvector return true; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h index 5b64e65..4924c68 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.h +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -156,7 +156,7 @@ struct BlockSparsePCGOptions { * and does a small LDLT entirely in shared memory). */ class BlockSparsePCGSolver : public CSRSparseLinearSolver { -public: + public: /** * @brief Constructs the solver with the given options. * @@ -186,7 +186,7 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { * `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 + * @c options_.block_layout takes precedence; passing an empty problem * (default-constructed) reverts to the uniform * @ref BlockSparsePCGOptions::block_size. * @@ -195,7 +195,7 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { * derive the block-Jacobi preconditioner layout * when @ref BlockSparsePCGOptions::block_layout is * empty and the caller hasn't explicitly invoked - * @ref SetBlockLayout. + * @c options_.block_layout. * @param spd_matrix Coefficient matrix `H` in CSR format. Only its * sparsity pattern is examined here; values are read * on every @ref Solve. @@ -248,7 +248,7 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { */ int LastIterations() const { return last_iterations_; } -private: + private: // ------------------------------------------------------------------ // Layout helpers // ------------------------------------------------------------------ @@ -295,7 +295,7 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { 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 + int block_row_start; ///< first block index in the global tile order }; std::vector segments_; @@ -315,7 +315,7 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { 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) + dvector Ap_; ///< `H p_k` (the SpMV output) /** Device-resident scalar slots: alpha, beta, , rz_old, rz_new, * ||r||^2, ||b||^2. Layout is fixed in the .cu file. */ @@ -326,7 +326,7 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { // ------------------------------------------------------------------ cuSPARSEHandle cusparse_handle_; cuSPARSEMatrixDescription mat_desc_; - dvector spmv_buffer_; ///< work buffer for cuSPARSE SpMV + dvector spmv_buffer_; ///< work buffer for cuSPARSE SpMV /** Non-owning view of the matrix passed to the current Solve; exactly one * of the two is non-null and selects the storage layout. */ @@ -337,4 +337,4 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { int last_iterations_ = 0; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/csr_sparse_linear_solver.h b/cunls/linear_solver/csr_sparse_linear_solver.h index 5432683..a47fb7d 100644 --- a/cunls/linear_solver/csr_sparse_linear_solver.h +++ b/cunls/linear_solver/csr_sparse_linear_solver.h @@ -40,7 +40,7 @@ class Problem; // forward declaration; defined in cunls/minimizer/problem.h. * site). Solvers that don't care can simply ignore the argument. */ class CSRSparseLinearSolver { -public: + public: /** * @brief Performs setup work for the linear system. * @@ -130,4 +130,4 @@ class CSRSparseLinearSolver { protected: bool safety_checks_enabled_ = true; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/cudss_sparse_linear_solver.h b/cunls/linear_solver/cudss_sparse_linear_solver.h index a45d8b4..91939d7 100644 --- a/cunls/linear_solver/cudss_sparse_linear_solver.h +++ b/cunls/linear_solver/cudss_sparse_linear_solver.h @@ -45,11 +45,11 @@ enum class cuDSSLinearSolverMode { */ struct cuDSSLinearSolverOptions { cuDSSLinearSolverMode mode = - cuDSSLinearSolverMode::SlowInitFastSolve; ///< Solver mode controlling + cuDSSLinearSolverMode::SlowInitFastSolve; ///< Solver mode controlling ///< the init/solve trade-off. int nthreads = 1; ///< Number of threads for host-side operations. std::string threading_lib_path = - ""; ///< Path to the threading library (empty disables multi-threading). + ""; ///< Path to the threading library (empty disables multi-threading). }; /** @@ -60,7 +60,7 @@ struct cuDSSLinearSolverOptions { * library for GPU-accelerated direct factorization. */ class cuDSSLinearSolver : public CSRSparseLinearSolver { -public: + public: // This backend consumes CSR only; SupportsBlockStorage() stays false, so the // base class's block-storage overloads are never called on it. The // using-declarations keep them visible rather than hidden by the CSR @@ -124,13 +124,13 @@ class cuDSSLinearSolver : public CSRSparseLinearSolver { bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; -private: - cuDSSLinearSolverOptions options_; ///< Solver configuration. + private: + cuDSSLinearSolverOptions options_; ///< Solver configuration. cuDSSHandle cudss_handle_; ///< Owns the cuDSS handle used for all solver phases. - cuDSSDeviceMemPool device_mem_pool_; ///< Reusable pool for cuDSS allocations. + cuDSSDeviceMemPool device_mem_pool_; ///< Reusable pool for cuDSS allocations. cuDSSData cudss_data_; ///< cuDSS data object storing internal solver state. - cuDSSConfig cudss_config_; ///< cuDSS configuration for solver parameters. + cuDSSConfig cudss_config_; ///< cuDSS configuration for solver parameters. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/dense_cholesky_solver.h b/cunls/linear_solver/dense_cholesky_solver.h index 3821a67..9f58f23 100644 --- a/cunls/linear_solver/dense_cholesky_solver.h +++ b/cunls/linear_solver/dense_cholesky_solver.h @@ -41,7 +41,7 @@ namespace cunls { * reports a non-zero devInfo from potrf). */ class DenseCholeskySolver : public CSRSparseLinearSolver { -public: + public: // This backend consumes CSR only; SupportsBlockStorage() stays false, so the // base class's block-storage overloads are never called on it. The // using-declarations keep them visible rather than hidden by the CSR @@ -83,7 +83,7 @@ class DenseCholeskySolver : public CSRSparseLinearSolver { bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; -private: + private: void EnsureBuffersSize(cudaStream_t stream, size_t n); void ConvertCSRToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, @@ -97,4 +97,4 @@ class DenseCholeskySolver : public CSRSparseLinearSolver { size_t last_n_ = 0; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/dense_linear_solver.h b/cunls/linear_solver/dense_linear_solver.h index e268baa..07cf63d 100644 --- a/cunls/linear_solver/dense_linear_solver.h +++ b/cunls/linear_solver/dense_linear_solver.h @@ -41,7 +41,7 @@ namespace cunls { * accurate bool return from Solve(). */ class DenseLDLTSolver : public CSRSparseLinearSolver { -public: + public: // This backend consumes CSR only; SupportsBlockStorage() stays false, so the // base class's block-storage overloads are never called on it. The // using-declarations keep them visible rather than hidden by the CSR @@ -93,7 +93,7 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; -private: + private: /** * @brief Ensures all internal buffers are (re-)allocated for an n x n system. * @@ -123,7 +123,7 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { dvector permutation_; ///< Pivot permutation vector. dvector permuted_rhs_; ///< P * b scratch vector. dvector permuted_solution_; ///< Permuted solution scratch. - dvector intermediate_solution_; ///< Intermediate solve scratch. + dvector intermediate_solution_; ///< Intermediate solve scratch. /// Device-side kernel status flags (index 0 = factorize, index 1 = solve). /// Each kernel writes 1 on success or 0 on failure. @@ -135,4 +135,4 @@ class DenseLDLTSolver : public CSRSparseLinearSolver { pvector status_pinned_; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/dense_qr_solver.h b/cunls/linear_solver/dense_qr_solver.h index fcfa1ef..661efb7 100644 --- a/cunls/linear_solver/dense_qr_solver.h +++ b/cunls/linear_solver/dense_qr_solver.h @@ -43,7 +43,7 @@ namespace cunls { * via devInfo. */ class DenseQRSolver : public CSRSparseLinearSolver { -public: + public: // This backend consumes CSR only; SupportsBlockStorage() stays false, so the // base class's block-storage overloads are never called on it. The // using-declarations keep them visible rather than hidden by the CSR @@ -84,7 +84,7 @@ class DenseQRSolver : public CSRSparseLinearSolver { bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; -private: + private: void EnsureBuffersSize(cudaStream_t stream, size_t n); void ConvertCSRToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, @@ -101,4 +101,4 @@ class DenseQRSolver : public CSRSparseLinearSolver { size_t last_n_ = 0; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/block_hessian_assembler.cu b/cunls/minimizer/block_hessian_assembler.cu index f4fb825..9397bda 100644 --- a/cunls/minimizer/block_hessian_assembler.cu +++ b/cunls/minimizer/block_hessian_assembler.cu @@ -221,7 +221,7 @@ int PickWarpsPerBlock(int residual_dim, int tangent_dim, int num_blocks, size_t return warps; } -} // namespace +} // namespace void BlockHessianAssembler::Initialize(cudaStream_t stream, const Problem &problem, int num_cols, CSRSparseMatrix &hessian) { @@ -368,4 +368,4 @@ void BlockHessianAssembler::LaunchAssembly(cudaStream_t stream, const float *jac } } -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/block_hessian_assembler.h b/cunls/minimizer/block_hessian_assembler.h index aa00055..c9201ee 100644 --- a/cunls/minimizer/block_hessian_assembler.h +++ b/cunls/minimizer/block_hessian_assembler.h @@ -53,7 +53,7 @@ class Problem; * column scaling, PCG, cuDSS) sees the matrix it already expects. */ class BlockHessianAssembler { -public: + public: /** * @brief Builds the Hessian sparsity pattern and the per-factor scatter maps. * @@ -103,10 +103,10 @@ class BlockHessianAssembler { /** @brief Total floats needed for the per-factor Jacobian value buffer. */ size_t JacobianValuesSize() const { return structure_builder_.JacobianValuesSize(); } -private: + private: /** @brief Per-residual-batch constants uploaded once for the kernel. */ struct BatchPlan { - HessianBatchLayout layout; ///< Geometry and flat-buffer offsets. + HessianBatchLayout layout; ///< Geometry and flat-buffer offsets. /// n entries: block index owning each local column. dvector block_of_col; /// n entries: offset of each local column inside its block. @@ -144,4 +144,4 @@ class BlockHessianAssembler { profiler::Domain profiler_domain_{"BlockHessianAssembler"}; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/bsr_matrix.cu b/cunls/minimizer/bsr_matrix.cu index e491001..e484bf2 100644 --- a/cunls/minimizer/bsr_matrix.cu +++ b/cunls/minimizer/bsr_matrix.cu @@ -31,54 +31,6 @@ constexpr int kBlockSize = 256; int GridFor(size_t count) { return static_cast((count + kBlockSize - 1) / kBlockSize); } -/** - * @brief Fills the CSR row offsets of the expanded matrix. - * - * Every tile in a block row contributes `block_size` columns to each of the - * block row's `block_size` scalar rows, so a row's length is known from the - * block row's tile count alone. - */ -__global__ void FillExpandedRowOffsetsKernel(int num_rows, int block_size, - const int *__restrict__ block_row_offsets, - int *__restrict__ row_offsets) { - int row = blockIdx.x * blockDim.x + threadIdx.x; - if (row > num_rows) { - return; - } - // Rows before `row` belong to whole block rows plus a partial one. - const int block_row = row / block_size; - const int sub_row = row - block_row * block_size; - const int whole = block_row_offsets[block_row] * block_size; - const int partial = sub_row * (block_row_offsets[block_row + 1] - block_row_offsets[block_row]); - row_offsets[row] = (whole + partial) * block_size; -} - -/** @brief Scatters each tile entry to its scalar CSR position. */ -__global__ void ExpandTilesToCSRKernel(size_t num_values, int block_size, - const int *__restrict__ row_of_tile, - const int *__restrict__ block_row_offsets, - const int *__restrict__ block_col_ids, - const float *__restrict__ values, - const int *__restrict__ row_offsets, - int *__restrict__ col_ids, float *__restrict__ out_values) { - size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx >= num_values) { - return; - } - const int tile_area = block_size * block_size; - const size_t tile = idx / tile_area; - const int within = static_cast(idx - tile * tile_area); - const int row_in_tile = within / block_size; - const int col_in_tile = within - row_in_tile * block_size; - - const int block_row = row_of_tile[tile]; - const int tile_in_row = static_cast(tile) - block_row_offsets[block_row]; - const int row = block_row * block_size + row_in_tile; - const int slot = row_offsets[row] + tile_in_row * block_size + col_in_tile; - col_ids[slot] = block_col_ids[tile] * block_size + col_in_tile; - out_values[slot] = values[idx]; -} - /** One thread per block row; locates the diagonal tile and reads its diagonal. */ __global__ void ExtractBlockDiagonalKernel(int num_block_rows, int block_size, const int *__restrict__ row_offsets, @@ -377,8 +329,11 @@ void LaunchBsrMultiply(cudaStream_t stream, int num_block_rows, int block_size, THROW_ON_CUDA_ERROR(cudaGetLastError()); } -} // namespace - +/** + * @brief y = A * x for block storage, into a caller-owned buffer. + * + * Internal: the only caller is ComputeWeightedSquaredStepAsync below. + */ void MultiplyBSRByDenseVector(cudaStream_t stream, const BSRSparseMatrix &matrix, const dvector &x, dvector &y) { y.resize(static_cast(matrix.NumRows())); @@ -390,6 +345,8 @@ void MultiplyBSRByDenseVector(cudaStream_t stream, const BSRSparseMatrix &matrix x.data(), y.data()); } +} // namespace + int ChooseHessianBlockSize(const Problem &problem, int max_block_size) { int block_size = 0; for (const auto *state_batch : problem.GetStateBatches()) { @@ -494,34 +451,6 @@ void CopyBSRSparseMatrix(cudaStream_t stream, const BSRSparseMatrix &input, } } -void ConvertBSRToCSR(cudaStream_t stream, const BSRSparseMatrix &input, CSRSparseMatrix &output, - dvector &row_of_tile) { - const int num_rows = input.NumRows(); - output.row_offsets.resize(static_cast(num_rows) + 1); - output.col_ids.resize(input.NumNonZeros()); - output.values.resize(input.NumNonZeros()); - if (num_rows == 0) { - THROW_ON_CUDA_ERROR(cudaMemsetAsync(output.row_offsets.data(), 0, sizeof(int), stream)); - return; - } - - row_of_tile.resize(input.NumBlocks()); - FillRowOfTileKernel<<>>( - input.num_block_rows, input.row_offsets.data(), row_of_tile.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - - FillExpandedRowOffsetsKernel<<(num_rows) + 1), kBlockSize, 0, - stream>>>(num_rows, input.block_size, input.row_offsets.data(), - output.row_offsets.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); - - ExpandTilesToCSRKernel<<>>( - input.values.size(), input.block_size, row_of_tile.data(), input.row_offsets.data(), - input.col_ids.data(), input.values.data(), output.row_offsets.data(), output.col_ids.data(), - output.values.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const BSRSparseMatrix &matrix, const dvector &step, dvector &scratch, float *d_out, float *d_partials) { @@ -529,4 +458,4 @@ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const BSRSparseMatrix DotProductToDevice(stream, step.data(), scratch.data(), step.size(), d_out, d_partials); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/bsr_matrix.h b/cunls/minimizer/bsr_matrix.h index a5724e7..76e455d 100644 --- a/cunls/minimizer/bsr_matrix.h +++ b/cunls/minimizer/bsr_matrix.h @@ -46,9 +46,6 @@ class Problem; */ int ChooseHessianBlockSize(const Problem &problem, int max_block_size = 16); -/** @brief Largest tile edge the block SpMV supports. */ -constexpr int kMaxHessianBlockSize = 16; - /** * @brief Extracts the main diagonal of a BSR matrix. * @@ -95,35 +92,6 @@ void ScaleSymmetric(cudaStream_t stream, BSRSparseMatrix &matrix, const dvector< void CopyBSRSparseMatrix(cudaStream_t stream, const BSRSparseMatrix &input, BSRSparseMatrix &output); -/** - * @brief Expands a BSR matrix into scalar CSR. - * - * Used for solver backends that cannot consume block storage (cuDSS and the - * dense factorizations). Column indices come out sorted within each row. - * - * @param stream CUDA stream for GPU operations. - * @param input BSR matrix. - * @param[out] output CSR matrix; resized as needed. - * @param[out] row_of_tile Caller-owned scratch mapping tile index to block row. - */ -void ConvertBSRToCSR(cudaStream_t stream, const BSRSparseMatrix &input, CSRSparseMatrix &output, - dvector &row_of_tile); - -/** - * @brief Sparse matrix-vector product y = A * x for BSR storage. - * - * cuSPARSE's `cusparseSbsrmv` is not used: it measured 3.6x slower than - * `csrmv_v3` on a bundle-adjustment Hessian, which would negate the format's - * whole advantage. See bsr_matrix.cu. - * - * @param stream CUDA stream for GPU operations. - * @param matrix BSR matrix A. - * @param x Input vector of length matrix.NumRows(). - * @param[out] y Output vector; resized to matrix.NumRows(). - */ -void MultiplyBSRByDenseVector(cudaStream_t stream, const BSRSparseMatrix &matrix, - const dvector &x, dvector &y); - /** * @brief Async BSR-weighted squared step: d_out[0] = step^T A step. * @@ -138,4 +106,4 @@ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const BSRSparseMatrix const dvector &step, dvector &scratch, float *d_out, float *d_partials); -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index 22c6aa8..5128aa2 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -291,56 +291,6 @@ void GaussNewtonMinimizer::Initialize(cudaStream_t stream, Problem &problem) { ResizeFactorJacobians(); } -/** - * @brief Checks if convergence criteria are satisfied. - * - * Convergence is determined by: - * 1. Step size: squared step norm < state_tolerance - * 2. Cost reduction: updated_cost < cost_tolerance - * 3. Step quality: step_quality >= 1.0 (cost increased) - * - * Also computes step_quality = updated_cost / current_cost as a metric for - * step acceptance/rejection. - * - * @param stream CUDA stream for GPU operations. - * @param updated_cost Cost after applying the step. - * @param current_cost Cost before applying the step. - * @param step State update step vector. - * @param[out] step_quality Output step quality metric (updated_cost / - * current_cost). - * @return True if converged, false otherwise. - */ -bool GaussNewtonMinimizer::CheckConvergence(cudaStream_t stream, float updated_cost, - float current_cost, const dvector &step, - float &step_quality) { - auto range = profiler_domain_.CreateDomainRange("CheckConvergence"); - - if (d_scalars_.size() < 1) d_scalars_.resize(1); - if (h_scalars_.size() < 1) h_scalars_.resize(1); - size_t partials_needed = ReducePartialCount(step.size()); - if (d_reduce_partials_.size() < partials_needed) { - d_reduce_partials_.resize(partials_needed); - } - - ComputeSquaredStepAsync(stream, step, d_scalars_.data(), d_reduce_partials_.data()); - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(h_scalars_.data(), d_scalars_.data(), sizeof(float), - cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - - float squared_step = h_scalars_[0]; - LogMessage("Squared step = {}", squared_step); - step_quality = updated_cost / current_cost; - - LogMessage("Step quality = {}", step_quality); - - if (squared_step < options_.state_tolerance || updated_cost < options_.cost_tolerance || - step_quality >= 1) { - return true; - } - - return false; -} - bool GaussNewtonMinimizer::EvaluateAndCheckConvergence(cudaStream_t stream, const Problem &problem, const MinimizerState &updated_state, float current_cost, diff --git a/cunls/minimizer/gauss_newton_minimizer.h b/cunls/minimizer/gauss_newton_minimizer.h index a0f5535..520ce39 100644 --- a/cunls/minimizer/gauss_newton_minimizer.h +++ b/cunls/minimizer/gauss_newton_minimizer.h @@ -187,7 +187,7 @@ struct MinimizerOptions { * @brief Gauss-Newton nonlinear least-squares optimizer. */ class GaussNewtonMinimizer { -public: + public: /** * @brief Constructs a Gauss-Newton optimizer. * @@ -235,24 +235,6 @@ class GaussNewtonMinimizer { MinimizerSummary Minimize(cudaStream_t stream, Problem &problem); protected: - /** - * @brief Checks if convergence criteria are satisfied. - * - * Determines whether the optimization has converged based on step size, - * cost reduction, and step quality. Also computes the step quality metric - * (ratio of updated cost to current cost). - * - * @param stream CUDA stream for GPU operations. - * @param updated_cost Cost after applying the step. - * @param current_cost Cost before applying the step. - * @param step State update step vector. - * @param[out] step_quality Output argument for step quality metric - * (updated_cost / current_cost). - * @return True if converged, false otherwise. - */ - virtual bool CheckConvergence(cudaStream_t stream, float updated_cost, float current_cost, - const dvector &step, float &step_quality); - /** * @brief Fused cost evaluation + convergence check with a single D2H + sync. * @@ -352,7 +334,7 @@ class GaussNewtonMinimizer { void ComputeCostAsync(cudaStream_t stream, const Problem &problem, const MinimizerState &minimizer_state, float *d_cost_out); -private: + private: /** * @brief Applies diagonal column scaling to the normal-equation system. * @@ -385,14 +367,14 @@ class GaussNewtonMinimizer { void ResizeFactorJacobians(); protected: - const MinimizerOptions options_; ///< Optimizer configuration options. + const MinimizerOptions options_; ///< Optimizer configuration options. SparseLinearSolverPtr solver_; ///< Linear solver for the normal equations. - cuSPARSEHandle cusparse_handle_; ///< cuSPARSE handle for sparse operations. + cuSPARSEHandle cusparse_handle_; ///< cuSPARSE handle for sparse operations. - StateBatchOps state_ops_; ///< Operations on state batches. + StateBatchOps state_ops_; ///< Operations on state batches. - dvector residuals_; ///< Residual vector storage. + dvector residuals_; ///< Residual vector storage. /// Per-factor dense Jacobian blocks; the only Jacobian ever materialized. PerFactorJacobians factor_jacobians_; @@ -403,9 +385,9 @@ class GaussNewtonMinimizer { /// Diagonal S when column_scaling is enabled; size = number of tangent DOFs. dvector column_scale_; - dvector step_; ///< State update step vector. + dvector step_; ///< State update step vector. - dvector buffer_; ///< Temporary buffer for sparse operations. + dvector buffer_; ///< Temporary buffer for sparse operations. /// Device staging buffer for async scalar reductions (cost, step norm, etc.). dvector d_scalars_; @@ -419,7 +401,7 @@ class GaussNewtonMinimizer { MinimizerState current_state_; MinimizerState updated_state_; - profiler::Domain profiler_domain_{"GaussNewtonMinimizer"}; ///< Profiling domain. + profiler::Domain profiler_domain_{"GaussNewtonMinimizer"}; ///< Profiling domain. }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/hessian_structure.cu b/cunls/minimizer/hessian_structure.cu index 43fe3b7..b1a66d3 100644 --- a/cunls/minimizer/hessian_structure.cu +++ b/cunls/minimizer/hessian_structure.cu @@ -293,7 +293,7 @@ __global__ void ScatterWriteOffsetsKernel(int num_candidates, int num_valid, write_offsets[pair_order[idx]] = value; } -} // namespace +} // namespace void HessianStructureBuilder::BuildLayout(const Problem &problem) { const auto &residual_batches = problem.GetResidualBatches(); @@ -679,4 +679,4 @@ void HessianStructureBuilder::Build(cudaStream_t stream, const Problem &problem, THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/hessian_structure.h b/cunls/minimizer/hessian_structure.h index 4e36fee..f3645e3 100644 --- a/cunls/minimizer/hessian_structure.h +++ b/cunls/minimizer/hessian_structure.h @@ -34,8 +34,8 @@ struct HessianBatchLayout { int residual_dim = 0; ///< m: residual dimension of one factor. int tangent_dim = 0; ///< n: sum of the factor's state block sizes. int num_blocks = 0; ///< nb: state blocks a factor touches. - size_t jacobian_offset = 0; ///< Offset into the flat Jacobian value buffer. - size_t residual_offset = 0; ///< Offset into the flat residual vector. + size_t jacobian_offset = 0; ///< Offset into the flat Jacobian value buffer. + size_t residual_offset = 0; ///< Offset into the flat residual vector. size_t col_offset = 0; ///< Offset into FactorCols(), stride nb. size_t pair_offset = 0; ///< Offset into WriteOffsets(), stride nb*nb. }; @@ -62,7 +62,7 @@ struct HessianBatchLayout { * `output.values` is resized but left uninitialized. */ class HessianStructureBuilder { -public: + public: /** * @brief Builds the scalar CSR pattern, and optionally the scatter maps. * @@ -113,7 +113,7 @@ class HessianStructureBuilder { /** @brief Total floats needed for the per-factor Jacobian value buffer. */ size_t JacobianValuesSize() const { return jacobian_values_size_; } -private: + private: /** * @brief Shared front half: layout, column resolution, key sort, segmentation. * @@ -146,4 +146,4 @@ class HessianStructureBuilder { size_t jacobian_values_size_ = 0; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/levenberg_marquardt_minimizer.cpp b/cunls/minimizer/levenberg_marquardt_minimizer.cpp index 9135c59..e6d1acc 100644 --- a/cunls/minimizer/levenberg_marquardt_minimizer.cpp +++ b/cunls/minimizer/levenberg_marquardt_minimizer.cpp @@ -47,72 +47,6 @@ void LevenbergMarquardtMinimizer::BuildSystem(cudaStream_t stream, const Problem normal_equations_.AddScaledDiagonalToLhs(stream, lambda_, diagonal_); } -/** - * @brief Checks convergence using LM-specific criteria. - * - * Computes the rho metric which measures the ratio of actual cost reduction to - * predicted cost reduction: - * rho = (actual_reduction) / (predicted_reduction) - * - * Convergence is determined by: - * 1. Step size: squared step norm < state_tolerance - * 2. Cost: updated_cost < cost_tolerance - * 3. Predicted reduction: predicted_relative_reduction < - * relative_reduction_tolerance - * - * @param stream CUDA stream for GPU operations. - * @param updated_cost Cost after applying the step. - * @param current_cost Cost before applying the step. - * @param step State update step vector. - * @param[out] step_quality Output rho metric (actual/predicted reduction). - * @return True if converged, false otherwise. - */ -bool LevenbergMarquardtMinimizer::CheckConvergence(cudaStream_t stream, float updated_cost, - float current_cost, const dvector &step, - float &step_quality) { - constexpr size_t kSlots = 3; - if (d_scalars_.size() < kSlots) d_scalars_.resize(kSlots); - if (h_scalars_.size() < kSlots) h_scalars_.resize(kSlots); - size_t partials_needed = ReducePartialCount(step_.size()); - if (d_reduce_partials_.size() < partials_needed) { - d_reduce_partials_.resize(partials_needed); - } - - // Enqueue all three reductions async - ComputeSquaredStepAsync(stream, step_, d_scalars_.data(), d_reduce_partials_.data()); - ComputeWeightedSquaredStepAsync(stream, diagonal_, step_, d_scalars_.data() + 1, - d_reduce_partials_.data()); - normal_equations_.WeightedSquaredStepAsync(stream, cusparse_handle_.GetHandle(stream), step_, - d_scalars_.data() + 2, d_reduce_partials_.data(), - buffer_); - - // Single D2H + single sync - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(h_scalars_.data(), d_scalars_.data(), kSlots * sizeof(float), - cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - - float step_sq_norm = h_scalars_[0]; - float diag_weight = h_scalars_[1]; - float matrix_weight = h_scalars_[2]; - - float predicted_relative_reduction = (matrix_weight + 2.f * lambda_ * diag_weight) / current_cost; - - LogMessage("Predicted relative reduction = {}", predicted_relative_reduction); - LogMessage("Step squared norm = {}", step_sq_norm); - - float rho = (1.f - updated_cost / current_cost) / predicted_relative_reduction; - - step_quality = rho; - - if (step_sq_norm < options_.base_options.state_tolerance || - predicted_relative_reduction < options_.relative_reduction_tolerance || - updated_cost < options_.base_options.cost_tolerance) { - return true; - } - - return false; -} - bool LevenbergMarquardtMinimizer::EvaluateAndCheckConvergence( cudaStream_t stream, const Problem &problem, const MinimizerState &updated_state, float current_cost, const dvector &step, float &updated_cost, float &step_quality) { @@ -225,4 +159,4 @@ void LevenbergMarquardtMinimizer::Initialize(cudaStream_t stream, Problem &probl GaussNewtonMinimizer::Initialize(stream, problem); lambda_ = options_.initial_lambda; } -} // namespace cunls \ No newline at end of file +} // namespace cunls diff --git a/cunls/minimizer/levenberg_marquardt_minimizer.h b/cunls/minimizer/levenberg_marquardt_minimizer.h index 6f5a720..83eacc5 100644 --- a/cunls/minimizer/levenberg_marquardt_minimizer.h +++ b/cunls/minimizer/levenberg_marquardt_minimizer.h @@ -157,29 +157,6 @@ class LevenbergMarquardtMinimizer : public GaussNewtonMinimizer { void BuildSystem(cudaStream_t stream, const Problem &problem, const MinimizerState &minimizer_state) override; - /** - * @brief Checks convergence using LM-specific criteria. - * - * Computes the rho metric (ratio of actual to predicted cost reduction) and - * checks convergence based on step size, cost, and predicted relative - * reduction. - * - * @param stream CUDA stream for GPU operations. - * @param updated_cost Cost after applying the step. - * @param current_cost Cost before applying the step. - * @param step State update step vector. - * @param[out] step_quality Output rho metric (actual/predicted reduction). - * @return True if converged, false otherwise. - */ - bool CheckConvergence(cudaStream_t stream, float updated_cost, float current_cost, - const dvector &step, float &step_quality) override; - - /** - * @brief Fused cost + LM convergence with a single D2H + sync. - * - * Batches cost reduction, squared step norm, diag-weighted step norm, - * and sparse-weighted step norm into one memcpy and one sync. - */ bool EvaluateAndCheckConvergence(cudaStream_t stream, const Problem &problem, const MinimizerState &updated_state, float current_cost, const dvector &step, float &updated_cost, diff --git a/cunls/minimizer/minimizer_state.h b/cunls/minimizer/minimizer_state.h index ecf14f7..a2abe66 100644 --- a/cunls/minimizer/minimizer_state.h +++ b/cunls/minimizer/minimizer_state.h @@ -42,7 +42,7 @@ namespace cunls { * used by the pointer-remap kernel. */ class MinimizerState { -public: + public: MinimizerState() = default; /** @@ -105,7 +105,7 @@ class MinimizerState { */ const std::vector> &GetStatePointers() const { return state_pointers_; } -private: + private: /** * @brief Creates minimizer state from a problem. * @@ -166,4 +166,4 @@ class MinimizerState { * @param[out] problem Destination problem to update. */ void Copy(cudaStream_t stream, const MinimizerState &state, Problem &problem); -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/normal_equations.cu b/cunls/minimizer/normal_equations.cu index 0d7e074..9170084 100644 --- a/cunls/minimizer/normal_equations.cu +++ b/cunls/minimizer/normal_equations.cu @@ -117,4 +117,4 @@ bool NormalEquations::Solve(cudaStream_t stream, CSRSparseLinearSolver &solver, return solver.Solve(stream, csr_lhs_, rhs, step); } -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/normal_equations.h b/cunls/minimizer/normal_equations.h index cf48e25..cab57f9 100644 --- a/cunls/minimizer/normal_equations.h +++ b/cunls/minimizer/normal_equations.h @@ -53,7 +53,7 @@ class Problem; * the tiles are real. */ class NormalEquations { -public: + public: /** * @brief Derives the sparsity pattern and picks the storage layout. * @@ -118,16 +118,13 @@ class NormalEquations { /** @brief True when the block layout is live for the current problem. */ bool UsesBlockStorage() const { return block_size_ > 1; } - /** @brief Tile edge of the block layout; 1 when the scalar layout is live. */ - int BlockSize() const { return block_size_; } - /** @brief Working left-hand side in scalar storage; empty under block storage. */ const CSRSparseMatrix &LhsCSR() const { return csr_lhs_; } /** @brief Working left-hand side in block storage; empty under scalar storage. */ const BSRSparseMatrix &LhsBSR() const { return bsr_lhs_; } -private: + private: BlockHessianAssembler assembler_; // Exactly one pair is populated, decided by block_size_. @@ -138,9 +135,9 @@ class NormalEquations { CSRMatrixDimensions csr_dims_; ///< Cached dims for the scalar SpMV. dvector tile_row_scratch_; ///< Tile-to-block-row map for scaling. - dvector block_spmv_scratch_; ///< SpMV result for the block path. + dvector block_spmv_scratch_; ///< SpMV result for the block path. int block_size_ = 1; }; -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/residual_batch.cu b/cunls/minimizer/residual_batch.cu index cc0d2c4..9666095 100644 --- a/cunls/minimizer/residual_batch.cu +++ b/cunls/minimizer/residual_batch.cu @@ -29,35 +29,30 @@ constexpr int kWarpSize = 32; namespace { -void MapRobustWorkspace(float *workspace, size_t num_residuals, - float **sq_err_out, float3 **rho_out) { +void MapRobustWorkspace(float *workspace, size_t num_residuals, float **sq_err_out, + float3 **rho_out) { *sq_err_out = workspace; const size_t sq_bytes = num_residuals * sizeof(float); const size_t align = alignof(float3); const size_t rho_byte_offset = (sq_bytes + align - 1u) / align * align; - *rho_out = reinterpret_cast(reinterpret_cast(workspace) + - rho_byte_offset); + *rho_out = reinterpret_cast(reinterpret_cast(workspace) + rho_byte_offset); } -} // namespace +} // namespace // ============================================================================ // Device helpers // ============================================================================ -__device__ __forceinline__ float jacobian_scaling_alpha(float sq_norm, - const float3 &rho) { - if ((sq_norm == 0.0f) || (rho.z <= 0.0f)) - return 0.0f; +__device__ __forceinline__ float jacobian_scaling_alpha(float sq_norm, const float3 &rho) { + if ((sq_norm == 0.0f) || (rho.z <= 0.0f)) return 0.0f; const float D = 1.0f + 2.0f * sq_norm * rho.z / rho.y; return (1.0f - sqrtf(D)) / sq_norm; } -__device__ __forceinline__ float residual_scaling(float sq_norm, - const float3 &rho) { +__device__ __forceinline__ float residual_scaling(float sq_norm, const float3 &rho) { float sqrt_rho1 = sqrtf(rho.y); - if ((sq_norm == 0.0f) || (rho.z <= 0.0f)) - return sqrt_rho1; + if ((sq_norm == 0.0f) || (rho.z <= 0.0f)) return sqrt_rho1; const float D = 1.0f + 2.0f * sq_norm * rho.z / rho.y; return sqrt_rho1 / (1.0f - (1.0f - sqrtf(D))); } @@ -69,13 +64,13 @@ __device__ __forceinline__ float residual_scaling(float sq_norm, // One thread per factor. Each thread computes ||r||^2 inline (loop over // residual_dim), writes rho = {s, 1, 0}, and optionally cost = 0.5*s. -__global__ void fused_trivial_sq_error_cost_kernel( - const float *__restrict__ residuals, float *__restrict__ sq_err, - float3 *__restrict__ rho, float *__restrict__ cost, int num_residuals, - int residual_dim) { +__global__ void fused_trivial_sq_error_cost_kernel(const float *__restrict__ residuals, + float *__restrict__ sq_err, + float3 *__restrict__ rho, + float *__restrict__ cost, int num_residuals, + int residual_dim) { int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_residuals) - return; + if (tid >= num_residuals) return; const float *r = residuals + tid * residual_dim; float sum = 0.0f; @@ -85,8 +80,7 @@ __global__ void fused_trivial_sq_error_cost_kernel( } sq_err[tid] = sum; rho[tid] = {sum, 1.0f, 0.0f}; - if (cost != nullptr) - cost[tid] = 0.5f * sum; + if (cost != nullptr) cost[tid] = 0.5f * sum; } // ============================================================================ @@ -97,12 +91,10 @@ __global__ void fused_trivial_sq_error_cost_kernel( // wasted warp lanes. __global__ void square_error_thread_kernel(const float *__restrict__ residuals, - float *__restrict__ squared_error, - int num_residuals, + float *__restrict__ squared_error, int num_residuals, int residual_dim) { int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_residuals) - return; + if (tid >= num_residuals) return; const float *r = residuals + tid * residual_dim; float sum = 0.0f; @@ -120,12 +112,12 @@ __global__ void square_error_thread_kernel(const float *__restrict__ residuals, // then applies it to all residual_dim elements. Eliminates warp-shuffle // overhead and wasted lanes for small residual dims. -__global__ void scale_residuals_thread_kernel( - float *__restrict__ residuals, const float *__restrict__ squared_error, - const float3 *__restrict__ rho, int num_residuals, int residual_dim) { +__global__ void scale_residuals_thread_kernel(float *__restrict__ residuals, + const float *__restrict__ squared_error, + const float3 *__restrict__ rho, int num_residuals, + int residual_dim) { int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_residuals) - return; + if (tid >= num_residuals) return; float scaling = residual_scaling(squared_error[tid], rho[tid]); float *r = residuals + tid * residual_dim; @@ -147,16 +139,14 @@ __global__ void scale_residuals_thread_kernel( // <= 32 (covers all practical NLS problems), the dot product r'*J_col is // computed via warp shuffle reduction, eliminating shared memory entirely. -__global__ void -scale_jacobians_warp_kernel(const float *__restrict__ residuals, - float *__restrict__ jacobians, - const float *__restrict__ squared_error, - const float3 *__restrict__ rho_coeffs, - int num_residuals, int residual_dim, int num_cols) { +__global__ void scale_jacobians_warp_kernel(const float *__restrict__ residuals, + float *__restrict__ jacobians, + const float *__restrict__ squared_error, + const float3 *__restrict__ rho_coeffs, + int num_residuals, int residual_dim, int num_cols) { const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / kWarpSize; const int lane = threadIdx.x % kWarpSize; - if (warp_id >= num_residuals) - return; + if (warp_id >= num_residuals) return; const int row_offset = warp_id * residual_dim; const float *res_base = residuals + row_offset; @@ -191,12 +181,10 @@ scale_jacobians_warp_kernel(const float *__restrict__ residuals, // ============================================================================ // Kernel 5: Extract cost // ============================================================================ -__global__ void extract_cost_kernel(float *__restrict__ cost, - const float3 *__restrict__ rho, +__global__ void extract_cost_kernel(float *__restrict__ cost, const float3 *__restrict__ rho, int num_residuals) { int tid = threadIdx.x + blockIdx.x * blockDim.x; - if (tid >= num_residuals) - return; + if (tid >= num_residuals) return; cost[tid] = 0.5f * rho[tid].x; } @@ -204,12 +192,10 @@ __global__ void extract_cost_kernel(float *__restrict__ cost, // ResidualBatch implementation // ============================================================================ -ResidualBatch::ResidualBatch(FactorBatch *factor_batch, - LossFunctionBatch *loss_function) +ResidualBatch::ResidualBatch(FactorBatch *factor_batch, LossFunctionBatch *loss_function) : factor_batch_(factor_batch), loss_function_(loss_function) {} -bool ResidualBatch::Evaluate(cudaStream_t stream, float *workspace, - float *residuals, +bool ResidualBatch::Evaluate(cudaStream_t stream, float *workspace, float *residuals, float const *const *state_pointers, float *cost, float *jacobians) const { int num_residuals = static_cast(factor_batch_->NumFactors()); @@ -238,8 +224,7 @@ bool ResidualBatch::Evaluate(cudaStream_t stream, float *workspace, if (loss_function_ == nullptr) { // Fast path: trivial loss. Fuse sq_error + trivial_loss + cost into 1 // kernel. - fused_trivial_sq_error_cost_kernel<<>>( + fused_trivial_sq_error_cost_kernel<<>>( residuals, sq_err_ptr, rho_ptr, cost, num_residuals, residual_dim); THROW_ON_CUDA_ERROR(cudaGetLastError()); return true; @@ -254,15 +239,13 @@ bool ResidualBatch::Evaluate(cudaStream_t stream, float *workspace, if (jacobians != nullptr) { int num_cols = 0; - for (const auto &d : factor_batch_->StateBlockSizes()) - num_cols += d; + for (const auto &d : factor_batch_->StateBlockSizes()) num_cols += d; int warps_per_block = kBlockSize / kWarpSize; int jac_blocks = (num_residuals + warps_per_block - 1) / warps_per_block; scale_jacobians_warp_kernel<<>>( - residuals, jacobians, sq_err_ptr, rho_ptr, num_residuals, residual_dim, - num_cols); + residuals, jacobians, sq_err_ptr, rho_ptr, num_residuals, residual_dim, num_cols); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -272,11 +255,10 @@ bool ResidualBatch::Evaluate(cudaStream_t stream, float *workspace, THROW_ON_CUDA_ERROR(cudaGetLastError()); if (cost != nullptr) { - extract_cost_kernel<<>>( - cost, rho_ptr, num_residuals); + extract_cost_kernel<<>>(cost, rho_ptr, num_residuals); THROW_ON_CUDA_ERROR(cudaGetLastError()); } return true; } -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/sparse_matrix.cu b/cunls/minimizer/sparse_matrix.cu index ad7fa31..9124d56 100644 --- a/cunls/minimizer/sparse_matrix.cu +++ b/cunls/minimizer/sparse_matrix.cu @@ -383,33 +383,6 @@ void AddScaledDiagonal(cudaStream_t stream, float scale, const dvector &d THROW_ON_CUDA_ERROR(cudaGetLastError()); } -/** - * Computes the right-hand side vector for the normal equations: rhs = -J^T * r. - * This is a key step in Gauss-Newton and Levenberg-Marquardt optimization - * algorithms. - * - * @param stream CUDA stream for asynchronous operations - * @param jacobian CSR sparse Jacobian matrix (J) - * @param residuals Dense residual vector (r) - * @param rhs Output right-hand side vector (-J^T * r) - * @param buffer Temporary buffer for sparse matrix operations - * - * First computes J^T * r using sparse matrix-vector multiplication, - * then negates the result to get -J^T * r. - */ -__global__ void negate_kernel(float *__restrict__ data, int n) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) data[i] = -data[i]; -} - -void NegateVector(cudaStream_t stream, float *data, size_t n) { - if (n == 0) return; - constexpr int kBlock = 256; - int grid = static_cast((n + kBlock - 1) / kBlock); - negate_kernel<<>>(data, static_cast(n)); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - __global__ void elementwise_multiply_kernel(float *__restrict__ a, const float *__restrict__ b, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; @@ -425,18 +398,7 @@ void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, s } /** - * Computes the weighted squared norm of a step vector: step^T * W * step, - * where W is a diagonal weight matrix. - * - * @param stream CUDA stream for asynchronous operations - * @param weights Diagonal weight values (W) - * @param step Step vector - * @param buffer Temporary buffer for intermediate computations - * @return The weighted squared norm (scalar value) - * - * Computes the inner product of the step vector with its element-wise - * product with the weights: sum(step[i] * weights[i] * step[i]). - * Used in trust region methods and optimization algorithms. + * @brief Async diagonally-weighted squared step: d_out[0] = step^T diag(w) step. */ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const dvector &weights, const dvector &step, float *d_out, float *d_partials) { @@ -445,100 +407,20 @@ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const dvector & d_partials); } -float ComputeWeightedSquaredStep(cudaStream_t stream, const dvector &weights, - const dvector &step, dvector &buffer) { - assert(step.size() == weights.size()); - size_t partials_count = ReducePartialCount(step.size()); - buffer.resize((partials_count + 1) * sizeof(float)); - float *d_out = reinterpret_cast(buffer.data()); - float *d_partials = d_out + 1; - - WeightedDotProductToDevice(stream, step.data(), weights.data(), step.data(), step.size(), d_out, - d_partials); - - float result; - THROW_ON_CUDA_ERROR( - cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - assert(result >= 0); - return result; -} - /** - * Computes the weighted squared norm using a sparse matrix: step^T * A * step. - * This overload uses a sparse matrix A instead of diagonal weights. - * - * @param stream CUDA stream for asynchronous operations - * @param matrix Sparse weight matrix (A) - * @param step Step vector - * @param buffer Temporary buffer for sparse matrix operations - * @return The weighted squared norm (scalar value) + * @brief Async sparse-weighted squared step: d_out[0] = step^T A step. * - * First computes A * step using sparse matrix-vector multiplication, - * then computes the inner product with the original step vector. - * Used when the weighting is represented as a full sparse matrix rather than - * just diagonal weights. + * Runs the SpMV into a slice of `buffer` and reduces against `step`, so the + * whole thing stays on the stream with no host synchronization. */ -static thread_local dvector g_spmv_result; - -static void WeightedSquaredStepSparseAsyncImpl(cudaStream_t stream, void *handle, - const CSRSparseMatrix &matrix, int num_rows, - int num_cols, int num_nonzeros, - const dvector &step, dvector &buffer, - float *d_out, float *d_partials) { - constexpr bool transpose_matrix = false; - SpMVImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, transpose_matrix, step, - g_spmv_result, buffer); - DotProductToDevice(stream, g_spmv_result.data(), step.data(), g_spmv_result.size(), d_out, - d_partials); -} - void ComputeWeightedSquaredStepAsync(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, int num_rows, int num_cols, int num_nonzeros, const dvector &step, dvector &buffer, float *d_out, float *d_partials) { - WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, step, - buffer, d_out, d_partials); -} - -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - const dvector &step, dvector &buffer) { - int num_rows, num_cols, num_nonzeros; - ExtractMatrixMetadata(stream, matrix, num_rows, num_cols, num_nonzeros); - - size_t partials_count = ReducePartialCount(step.size()); - dvector d_scratch(partials_count + 1); - float *d_out = d_scratch.data(); - float *d_partials = d_out + 1; - - WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, step, - buffer, d_out, d_partials); - - float result; - THROW_ON_CUDA_ERROR( - cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - assert(result >= 0); - return result; -} - -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - int num_rows, int num_cols, int num_nonzeros, - const dvector &step, dvector &buffer) { - size_t partials_count = ReducePartialCount(step.size()); - dvector d_scratch(partials_count + 1); - float *d_out = d_scratch.data(); - float *d_partials = d_out + 1; - - WeightedSquaredStepSparseAsyncImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, step, - buffer, d_out, d_partials); - - float result; - THROW_ON_CUDA_ERROR( - cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - assert(result >= 0); - return result; + static thread_local dvector spmv_result; + SpMVImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, /*transpose_matrix=*/false, + step, spmv_result, buffer); + DotProductToDevice(stream, step.data(), spmv_result.data(), step.size(), d_out, d_partials); } /** @@ -553,19 +435,4 @@ void ComputeSquaredStepAsync(cudaStream_t stream, const dvector &step, fl DotProductToDevice(stream, step.data(), step.data(), step.size(), d_out, d_partials); } -float ComputeSquaredStep(cudaStream_t stream, const dvector &step) { - size_t partials_count = ReducePartialCount(step.size()); - dvector d_scratch(partials_count + 1); - float *d_out = d_scratch.data(); - float *d_partials = d_out + 1; - - DotProductToDevice(stream, step.data(), step.data(), step.size(), d_out, d_partials); - - float result; - THROW_ON_CUDA_ERROR( - cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); - assert(result >= 0); - return result; -} -} // namespace cunls +} // namespace cunls diff --git a/cunls/minimizer/sparse_matrix.h b/cunls/minimizer/sparse_matrix.h index 4eb201f..6f3a795 100644 --- a/cunls/minimizer/sparse_matrix.h +++ b/cunls/minimizer/sparse_matrix.h @@ -96,53 +96,6 @@ void ScaleSymmetricCSR(cudaStream_t stream, CSRSparseMatrix &matrix, const dvect */ void InvertSqrtWithFloorInPlace(cudaStream_t stream, dvector &v, float floor_value = 1e-12f); -/** - * @brief Computes the squared L2 norm of a step vector. - * - * Computes step^T * step using a GPU inner product. - * - * @param stream CUDA stream for GPU operations. - * @param step Step vector. - * @return The squared L2 norm (scalar value). - */ -float ComputeSquaredStep(cudaStream_t stream, const dvector &step); - -/** - * @brief Computes a diagonally-weighted squared step norm. - * - * Computes step^T * diag(weights) * step, i.e., sum(weights[i] * step[i]^2). - * Used in Levenberg-Marquardt to evaluate predicted cost reduction. - * - * @param stream CUDA stream for GPU operations. - * @param weights Diagonal weight values. - * @param step Step vector. - * @param[out] buffer Temporary buffer for intermediate computations. - * @return The weighted squared norm (scalar value). - */ -float ComputeWeightedSquaredStep(cudaStream_t stream, const dvector &weights, - const dvector &step, dvector &buffer); - -/** - * @brief Computes a sparse-matrix-weighted squared step norm. - * - * Computes step^T * A * step using sparse matrix-vector multiplication - * followed by an inner product. Used when the weighting is a full sparse - * matrix rather than just diagonal weights. - * - * @param stream CUDA stream for GPU operations. - * @param handle Opaque cuSPARSE library handle (void*). - * @param matrix Sparse weight matrix (A). - * @param step Step vector. - * @param[out] buffer Temporary buffer for cuSPARSE operations. - * @return The weighted squared norm (scalar value). - */ -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - const dvector &step, dvector &buffer); - -float ComputeWeightedSquaredStep(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, - int num_rows, int num_cols, int num_nonzeros, - const dvector &step, dvector &buffer); - // ---- Async variants: write scalar result to device memory, no D2H or sync -- /** @@ -170,14 +123,6 @@ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, void *handle, int num_nonzeros, const dvector &step, dvector &buffer, float *d_out, float *d_partials); -/** - * @brief Elementwise vector negation: out[i] = -in[i]. - */ -void NegateVector(cudaStream_t stream, float *data, size_t n); - -/** - * @brief Elementwise multiply: out[i] = a[i] * b[i]. - */ void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, size_t n); -} // namespace cunls +} // namespace cunls diff --git a/cunls/state/state_batch_ops.cu b/cunls/state/state_batch_ops.cu index 1e8233a..54751ee 100644 --- a/cunls/state/state_batch_ops.cu +++ b/cunls/state/state_batch_ops.cu @@ -243,4 +243,4 @@ void StateBatchOps::Plus(cudaStream_t stream, const std::vector & } } -} // namespace cunls +} // namespace cunls diff --git a/cunls/state/state_batch_ops.h b/cunls/state/state_batch_ops.h index 29a2f11..d916201 100644 --- a/cunls/state/state_batch_ops.h +++ b/cunls/state/state_batch_ops.h @@ -53,7 +53,7 @@ void ComputeStateBlockColumnOffsets(cudaStream_t stream, int first_column, * batch. */ class StateBatchOps { -public: + public: /** * @brief Constructs and preprocesses the operator for the given state * batches. @@ -110,7 +110,7 @@ class StateBatchOps { * full (including constant) state indices. */ DeviceVector map_; -private: + private: /** * @brief Allocates the full-size state updates buffer and computes per-batch * delta pointers. @@ -143,4 +143,4 @@ class StateBatchOps { * constant blocks. */ size_t num_reduced_states_ = 0; }; -} // namespace cunls +} // namespace cunls diff --git a/python/src/bind_types.cpp b/python/src/bind_types.cpp index d0b93d9..9a2880c 100644 --- a/python/src/bind_types.cpp +++ b/python/src/bind_types.cpp @@ -22,11 +22,10 @@ // other bind_*.cpp file relies on it for constructor arguments that accept // CuPy arrays or plain integer addresses. -#include "bindings.h" - #include #include +#include "bindings.h" #include "cunls/common/cublas_helper.h" #include "cunls/common/cuda_stream.h" #include "cunls/linear_solver/sparse_linear_solver.h" @@ -42,16 +41,14 @@ // cupy.ndarray, which exposes the GPU pointer through its memory // descriptor. uintptr_t extract_device_ptr(nb::handle obj) { - if (nb::isinstance(obj)) - return nb::cast(obj); + if (nb::isinstance(obj)) return nb::cast(obj); return nb::cast(obj.attr("data").attr("ptr")); } void bind_types(nb::module_ &m) { // --- CUDA stream / cuBLAS handle wrappers --- - nb::class_(m, "CudaStream", - "RAII wrapper for a CUDA stream.") + nb::class_(m, "CudaStream", "RAII wrapper for a CUDA stream.") .def(nb::init(), nb::arg("sync_on_destroy") = false) .def( "get_stream", @@ -60,8 +57,7 @@ void bind_types(nb::module_ &m) { }, "Returns the underlying cudaStream_t as an integer handle."); - nb::class_(m, "CublasHandle", - "RAII wrapper for a cuBLAS handle.") + nb::class_(m, "CublasHandle", "RAII wrapper for a cuBLAS handle.") .def(nb::init<>()); // --- Enumerations for solver/multiplier strategy selection --- @@ -82,20 +78,16 @@ void bind_types(nb::module_ &m) { // from Python before passing the options to a minimizer constructor. nb::class_( - m, "MinimizerOptions", - "Options for Gauss-Newton and Levenberg-Marquardt minimizers.") + m, "MinimizerOptions", "Options for Gauss-Newton and Levenberg-Marquardt minimizers.") .def(nb::init<>()) - .def_rw("max_num_iterations", - &cunls::MinimizerOptions::max_num_iterations) + .def_rw("max_num_iterations", &cunls::MinimizerOptions::max_num_iterations) .def_rw("state_tolerance", &cunls::MinimizerOptions::state_tolerance) .def_rw("cost_tolerance", &cunls::MinimizerOptions::cost_tolerance) .def_rw("max_consecutive_rejected_steps", &cunls::MinimizerOptions::max_consecutive_rejected_steps) - .def_rw("sparse_linear_solver_type", - &cunls::MinimizerOptions::sparse_linear_solver_type) + .def_rw("sparse_linear_solver_type", &cunls::MinimizerOptions::sparse_linear_solver_type) .def_rw("column_scaling", &cunls::MinimizerOptions::column_scaling) - .def_rw("disable_safety_checks", - &cunls::MinimizerOptions::disable_safety_checks, + .def_rw("disable_safety_checks", &cunls::MinimizerOptions::disable_safety_checks, "When False, the minimizer enables all optional runtime " "validation. Currently this covers post-factorization " "checks in the linear solver: pivot/diagonal checks " @@ -110,38 +102,28 @@ void bind_types(nb::module_ &m) { "singular or ill-conditioned matrices may produce silently " "incorrect results."); - nb::class_(m, "MinimizerSummary", - "Summary of a minimization run.") + nb::class_(m, "MinimizerSummary", "Summary of a minimization run.") .def_ro("num_iterations", &cunls::MinimizerSummary::num_iterations) .def_ro("initial_cost", &cunls::MinimizerSummary::initial_cost) .def_ro("final_cost", &cunls::MinimizerSummary::final_cost) .def_ro("iteration_costs", &cunls::MinimizerSummary::iteration_costs) .def("__repr__", [](const cunls::MinimizerSummary &s) { - return "MinimizerSummary(iterations=" + - std::to_string(s.num_iterations) + + return "MinimizerSummary(iterations=" + std::to_string(s.num_iterations) + ", initial_cost=" + std::to_string(s.initial_cost) + ", final_cost=" + std::to_string(s.final_cost) + ")"; }); nb::class_( - m, "LevenbergMarquardtMinimizerOptions", - "Options for the Levenberg-Marquardt minimizer.") + m, "LevenbergMarquardtMinimizerOptions", "Options for the Levenberg-Marquardt minimizer.") .def(nb::init<>()) - .def_rw("base_options", - &cunls::LevenbergMarquardtMinimizerOptions::base_options) - .def_rw("initial_lambda", - &cunls::LevenbergMarquardtMinimizerOptions::initial_lambda) - .def_rw("lambda_upscale", - &cunls::LevenbergMarquardtMinimizerOptions::lambda_upscale) - .def_rw("lambda_downscale", - &cunls::LevenbergMarquardtMinimizerOptions::lambda_downscale) - .def_rw("lambda_max", - &cunls::LevenbergMarquardtMinimizerOptions::lambda_max) - .def_rw("lambda_min", - &cunls::LevenbergMarquardtMinimizerOptions::lambda_min) + .def_rw("base_options", &cunls::LevenbergMarquardtMinimizerOptions::base_options) + .def_rw("initial_lambda", &cunls::LevenbergMarquardtMinimizerOptions::initial_lambda) + .def_rw("lambda_upscale", &cunls::LevenbergMarquardtMinimizerOptions::lambda_upscale) + .def_rw("lambda_downscale", &cunls::LevenbergMarquardtMinimizerOptions::lambda_downscale) + .def_rw("lambda_max", &cunls::LevenbergMarquardtMinimizerOptions::lambda_max) + .def_rw("lambda_min", &cunls::LevenbergMarquardtMinimizerOptions::lambda_min) .def_rw("step_accept_threshold", &cunls::LevenbergMarquardtMinimizerOptions::step_accept_threshold) .def_rw("lambda_downscale_threshold", - &cunls::LevenbergMarquardtMinimizerOptions:: - lambda_downscale_threshold); + &cunls::LevenbergMarquardtMinimizerOptions::lambda_downscale_threshold); } diff --git a/tests/block_hessian_assembler_test.cpp b/tests/block_hessian_assembler_test.cpp index 41e9067..1a5f8a6 100644 --- a/tests/block_hessian_assembler_test.cpp +++ b/tests/block_hessian_assembler_test.cpp @@ -60,6 +60,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/bsr_expansion.h" #include "tests/utils.h" namespace cunls { @@ -77,7 +78,7 @@ namespace { * the two assembly paths on identical input. */ class SystemBuilder : public GaussNewtonMinimizer { -public: + public: /** * @brief Builds with block storage (the default for a block-capable solver). */ @@ -110,7 +111,7 @@ class SystemBuilder : public GaussNewtonMinimizer { if (!normal_equations_.UsesBlockStorage()) { return normal_equations_.LhsCSR(); } - ConvertBSRToCSR(stream, normal_equations_.LhsBSR(), csr_mirror_, expand_scratch_); + test_utils::ExpandBSRToCSR(stream, normal_equations_.LhsBSR(), csr_mirror_, expand_scratch_); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); return csr_mirror_; } @@ -134,7 +135,7 @@ class SystemBuilder : public GaussNewtonMinimizer { NormalEquations &Equations() { return normal_equations_; } bool UsesBlockStorage() const { return normal_equations_.UsesBlockStorage(); } -private: + private: static MinimizerOptions MakeOptions(SparseLinearSolverType solver) { MinimizerOptions options; options.sparse_linear_solver_type = solver; @@ -944,7 +945,7 @@ TEST(HessianStorageTest, PcgAgreesBetweenStorages) { CSRSparseMatrix csr; dvector expand_scratch; - ConvertBSRToCSR(s, block.Equations().LhsBSR(), csr, expand_scratch); + test_utils::ExpandBSRToCSR(s, block.Equations().LhsBSR(), csr, expand_scratch); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); const size_t n = block.Rhs().size(); @@ -1003,7 +1004,7 @@ TEST(HessianStorageTest, FirstPcgIterationAgreesBetweenStorages) { // Feed both solvers the *same* matrix, so only the reader differs. CSRSparseMatrix csr; dvector expand_scratch; - ConvertBSRToCSR(s, block.Equations().LhsBSR(), csr, expand_scratch); + test_utils::ExpandBSRToCSR(s, block.Equations().LhsBSR(), csr, expand_scratch); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); const size_t n = block.Rhs().size(); @@ -1284,5 +1285,5 @@ TEST(BlockHessianAssemblerTest, AllConstantStatesProduceZeroSystem) { } } -} // namespace -} // namespace cunls +} // namespace +} // namespace cunls diff --git a/tests/bsr_expansion.cu b/tests/bsr_expansion.cu new file mode 100644 index 0000000..7d23ef1 --- /dev/null +++ b/tests/bsr_expansion.cu @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/** + * @file bsr_expansion.cu + * @brief Expands a block Hessian into scalar CSR, for tests only. + * + * cuNLS never needs this: a solver that cannot read tiles reports + * SupportsBlockStorage() == false and is handed scalar CSR directly, so no + * conversion happens on any path. It lives here rather than in the library + * because its only purpose is to make the two storage layouts directly + * comparable in the equivalence tests. + */ + +#include "cunls/common/helper.h" +#include "tests/bsr_expansion.h" + +namespace cunls { +namespace test_utils { +namespace { + +constexpr int kBlockSize = 256; + +int GridFor(size_t count) { return static_cast((count + kBlockSize - 1) / kBlockSize); } + +/** @brief Records, for each stored tile, the block row it belongs to. */ +__global__ void FillRowOfTileKernel(int num_block_rows, const int *__restrict__ row_offsets, + int *__restrict__ row_of_tile) { + int block_row = blockIdx.x * blockDim.x + threadIdx.x; + if (block_row >= num_block_rows) { + return; + } + for (int tile = row_offsets[block_row]; tile < row_offsets[block_row + 1]; tile++) { + row_of_tile[tile] = block_row; + } +} + +/** + * @brief Fills the CSR row offsets of the expanded matrix. + * + * Every tile in a block row contributes `block_size` columns to each of that + * block row's `block_size` scalar rows, so a row's length follows from the + * block row's tile count alone. + */ +__global__ void FillExpandedRowOffsetsKernel(int num_rows, int block_size, + const int *__restrict__ block_row_offsets, + int *__restrict__ row_offsets) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row > num_rows) { + return; + } + const int block_row = row / block_size; + const int sub_row = row - block_row * block_size; + const int whole = block_row_offsets[block_row] * block_size; + const int partial = sub_row * (block_row_offsets[block_row + 1] - block_row_offsets[block_row]); + row_offsets[row] = (whole + partial) * block_size; +} + +/** @brief Scatters each tile entry to its scalar CSR position. */ +__global__ void ExpandTilesToCSRKernel(size_t num_values, int block_size, + const int *__restrict__ row_of_tile, + const int *__restrict__ block_row_offsets, + const int *__restrict__ block_col_ids, + const float *__restrict__ values, + const int *__restrict__ row_offsets, + int *__restrict__ col_ids, float *__restrict__ out_values) { + size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_values) { + return; + } + const int tile_area = block_size * block_size; + const size_t tile = idx / tile_area; + const int within = static_cast(idx - tile * tile_area); + const int row_in_tile = within / block_size; + const int col_in_tile = within - row_in_tile * block_size; + + const int block_row = row_of_tile[tile]; + const int tile_in_row = static_cast(tile) - block_row_offsets[block_row]; + const int row = block_row * block_size + row_in_tile; + const int slot = row_offsets[row] + tile_in_row * block_size + col_in_tile; + col_ids[slot] = block_col_ids[tile] * block_size + col_in_tile; + out_values[slot] = values[idx]; +} + +} // namespace + +void ExpandBSRToCSR(cudaStream_t stream, const BSRSparseMatrix &input, CSRSparseMatrix &output, + dvector &row_of_tile) { + const int num_rows = input.NumRows(); + output.row_offsets.resize(static_cast(num_rows) + 1); + output.col_ids.resize(input.NumNonZeros()); + output.values.resize(input.NumNonZeros()); + if (num_rows == 0) { + THROW_ON_CUDA_ERROR(cudaMemsetAsync(output.row_offsets.data(), 0, sizeof(int), stream)); + return; + } + + row_of_tile.resize(input.NumBlocks()); + FillRowOfTileKernel<<>>( + input.num_block_rows, input.row_offsets.data(), row_of_tile.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + FillExpandedRowOffsetsKernel<<(num_rows) + 1), kBlockSize, 0, + stream>>>(num_rows, input.block_size, input.row_offsets.data(), + output.row_offsets.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); + + ExpandTilesToCSRKernel<<>>( + input.values.size(), input.block_size, row_of_tile.data(), input.row_offsets.data(), + input.col_ids.data(), input.values.data(), output.row_offsets.data(), output.col_ids.data(), + output.values.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +} // namespace test_utils +} // namespace cunls diff --git a/tests/bsr_expansion.h b/tests/bsr_expansion.h new file mode 100644 index 0000000..b9f3839 --- /dev/null +++ b/tests/bsr_expansion.h @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "cunls/common/types.h" + +namespace cunls { +namespace test_utils { + +/** + * @brief Expands a block Hessian into scalar CSR so the two storage layouts can + * be compared entry for entry. + * + * Test-only; cuNLS itself never converts between layouts. Column indices come + * out sorted within each row. + * + * @param stream CUDA stream for GPU operations. + * @param input BSR matrix. + * @param[out] output CSR matrix; resized as needed. + * @param[out] row_of_tile Caller-owned scratch mapping tile index to block row. + */ +void ExpandBSRToCSR(cudaStream_t stream, const BSRSparseMatrix &input, CSRSparseMatrix &output, + dvector &row_of_tile); + +} // namespace test_utils +} // namespace cunls diff --git a/tests/sparse_matrix_test.cpp b/tests/sparse_matrix_test.cpp index 67d6f5f..b9612b3 100644 --- a/tests/sparse_matrix_test.cpp +++ b/tests/sparse_matrix_test.cpp @@ -39,12 +39,32 @@ #include "cunls/common/helper.h" #include "cunls/common/profiler.h" #include "cunls/common/types.h" +#include "cunls/minimizer/device_reduction.h" #include "tests/utils.h" namespace cunls { namespace { +/** + * @brief Runs an async reduction to completion and returns the scalar. + * + * The minimizers only ever use the async forms, so the tests do too; this wraps + * the device-side result the way Levenberg-Marquardt reads it. + */ +template +float RunAsyncReduction(cudaStream_t stream, size_t length, AsyncFn &&enqueue) { + dvector scratch(ReducePartialCount(length) + 1); + float *d_out = scratch.data(); + float *d_partials = d_out + 1; + enqueue(d_out, d_partials); + float result = 0.f; + THROW_ON_CUDA_ERROR( + cudaMemcpyAsync(&result, d_out, sizeof(float), cudaMemcpyDeviceToHost, stream)); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); + return result; +} + /** * @brief Computes y = A * x on the CPU where A is in CSR format. * @@ -117,7 +137,7 @@ void AddScaledDiagonalCPU(const std::vector &row_ptr, const std::vector &row_ptr, const std::vectorprofiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepFirst"); - result = ComputeWeightedSquaredStep(stream.GetStream(), dweights, dsteps, buffer); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + result = + RunAsyncReduction(stream.GetStream(), dsteps.size(), [&](float *d_out, float *d_partials) { + ComputeWeightedSquaredStepAsync(stream.GetStream(), dweights, dsteps, d_out, d_partials); + }); } // Verify GPU result matches CPU computation @@ -351,8 +373,14 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { float result; { auto range = this->profiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepSecond"); - result = ComputeWeightedSquaredStep(stream.GetStream(), handle, input_matrix, dsteps, buffer); - THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + int num_rows = 0, num_cols = 0, num_nonzeros = 0; + ExtractMatrixMetadata(stream.GetStream(), input_matrix, num_rows, num_cols, num_nonzeros); + result = + RunAsyncReduction(stream.GetStream(), dsteps.size(), [&](float *d_out, float *d_partials) { + ComputeWeightedSquaredStepAsync(stream.GetStream(), handle, input_matrix, num_rows, + num_cols, num_nonzeros, dsteps, buffer, d_out, + d_partials); + }); } result /= matrix_size; @@ -384,4 +412,4 @@ TEST(SparseMatrixColumnScaling, SymmetricScaling2x2) { ASSERT_NEAR(out[3], 1.f, 1e-4f); } -} // namespace cunls +} // namespace cunls diff --git a/tests/utils.h b/tests/utils.h index c588d78..cd88934 100644 --- a/tests/utils.h +++ b/tests/utils.h @@ -395,5 +395,5 @@ inline float PCGTolFromEnv(float fallback) { return (v > 0.f) ? v : fallback; } -} // namespace test_utils -} // namespace cunls +} // namespace test_utils +} // namespace cunls From ced495b32993e5d7051c005e6a331dee2e46010e Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Sat, 1 Aug 2026 21:20:54 -0700 Subject: [PATCH 3/7] Fix discussions v1 --- .../linear_solver/block_sparse_pcg_solver.cu | 20 +++++-- cunls/minimizer/bsr_matrix.cu | 18 +++++-- cunls/minimizer/normal_equations.cu | 7 ++- cunls/minimizer/normal_equations.h | 6 +-- cunls/minimizer/sparse_matrix.cu | 36 +++++++++---- cunls/minimizer/sparse_matrix.h | 5 +- tests/block_hessian_assembler_test.cpp | 41 ++++++++++---- tests/sparse_matrix_test.cpp | 54 ++++++++++++++++++- 8 files changed, 151 insertions(+), 36 deletions(-) diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index 68b7fd6..510beb3 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -89,6 +89,7 @@ #include #include #include +#include #include "cunls/common/cusparse_helper.h" #include "cunls/common/helper.h" @@ -258,18 +259,23 @@ __global__ void BsrMultiplyWarpKernel(int num_block_rows, const int *__restrict_ } } +/// Largest tile edge the generic BSR SpMV can accumulate per lane. +/// Distinct from the preconditioner's block-size limit; this one is set by +/// the fixed-size accumulator in BsrMultiplyWarpGenericKernel. +constexpr int kMaxSpMVBlockSize = 16; + /** * @brief Runtime-tile-edge fallback for block sizes without a specialization. * - * Same schedule as BsrMultiplyWarpKernel; the accumulator is sized to the - * largest edge ChooseHessianBlockSize can return. + * Same schedule as BsrMultiplyWarpKernel, but the per-lane accumulator is a + * fixed-size array, so the caller must reject edges above + * kMaxSpMVBlockSize before dispatching here. */ __global__ void BsrMultiplyWarpGenericKernel(int num_block_rows, int block_size, const int *__restrict__ row_offsets, const int *__restrict__ col_ids, const float *__restrict__ values, const float *__restrict__ x, float *__restrict__ y) { - constexpr int kMaxBlockSize = 16; const int block_row = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; const int lane = threadIdx.x & 31; if (block_row >= num_block_rows) { @@ -278,7 +284,7 @@ __global__ void BsrMultiplyWarpGenericKernel(int num_block_rows, int block_size, const int b = block_size; const int end = row_offsets[block_row + 1]; - float acc[kMaxBlockSize]; + float acc[kMaxSpMVBlockSize]; for (int k = 0; k < b; ++k) { acc[k] = 0.f; } @@ -400,6 +406,12 @@ void LaunchBsrMultiply(cudaStream_t stream, int num_block_rows, int block_size, LAUNCH_WARP(8); #undef LAUNCH_WARP default: + // The generic kernel accumulates into a fixed-size per-lane array; a + // larger edge would write past it, silently, so refuse instead. + if (block_size > kMaxSpMVBlockSize) { + throw std::runtime_error("BSR SpMV: block size exceeds " + + std::to_string(kMaxSpMVBlockSize)); + } BsrMultiplyWarpGenericKernel<<>>( num_block_rows, block_size, row_offsets, col_ids, values, x, y); break; diff --git a/cunls/minimizer/bsr_matrix.cu b/cunls/minimizer/bsr_matrix.cu index e484bf2..c548e98 100644 --- a/cunls/minimizer/bsr_matrix.cu +++ b/cunls/minimizer/bsr_matrix.cu @@ -17,6 +17,7 @@ #include #include +#include #include "cunls/common/helper.h" #include "cunls/minimizer/bsr_matrix.h" @@ -181,18 +182,21 @@ __global__ void BsrMultiplyWarpKernel(int num_block_rows, const int *__restrict_ } } +/// Largest tile edge the generic BSR SpMV can accumulate per lane. +constexpr int kMaxSpMVBlockSize = 16; + /** * @brief Runtime-tile-edge fallback for block sizes without a specialization. * - * Same schedule as BsrMultiplyWarpKernel; the accumulator is sized to the - * largest edge ChooseHessianBlockSize can return. + * Same schedule as BsrMultiplyWarpKernel, but the per-lane accumulator is a + * fixed-size array, so the caller must reject edges above + * kMaxSpMVBlockSize before dispatching here. */ __global__ void BsrMultiplyWarpGenericKernel(int num_block_rows, int block_size, const int *__restrict__ row_offsets, const int *__restrict__ col_ids, const float *__restrict__ values, const float *__restrict__ x, float *__restrict__ y) { - constexpr int kMaxBlockSize = 16; const int block_row = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; const int lane = threadIdx.x & 31; if (block_row >= num_block_rows) { @@ -200,7 +204,7 @@ __global__ void BsrMultiplyWarpGenericKernel(int num_block_rows, int block_size, } const int end = row_offsets[block_row + 1]; - float acc[kMaxBlockSize]; + float acc[kMaxSpMVBlockSize]; for (int k = 0; k < block_size; ++k) { acc[k] = 0.f; } @@ -322,6 +326,12 @@ void LaunchBsrMultiply(cudaStream_t stream, int num_block_rows, int block_size, LAUNCH_WARP(8); #undef LAUNCH_WARP default: + // The generic kernel accumulates into a fixed-size per-lane array; a + // larger edge would write past it, silently, so refuse instead. + if (block_size > kMaxSpMVBlockSize) { + throw std::runtime_error("BSR SpMV: block size exceeds " + + std::to_string(kMaxSpMVBlockSize)); + } BsrMultiplyWarpGenericKernel<<>>( num_block_rows, block_size, row_offsets, col_ids, values, x, y); break; diff --git a/cunls/minimizer/normal_equations.cu b/cunls/minimizer/normal_equations.cu index 9170084..22e1d45 100644 --- a/cunls/minimizer/normal_equations.cu +++ b/cunls/minimizer/normal_equations.cu @@ -91,13 +91,12 @@ void NormalEquations::WeightedSquaredStepAsync(cudaStream_t stream, void *cuspar const dvector &step, float *d_out, float *d_partials, dvector &buffer) { if (UsesBlockStorage()) { - ComputeWeightedSquaredStepAsync(stream, bsr_hessian_, step, block_spmv_scratch_, d_out, - d_partials); + ComputeWeightedSquaredStepAsync(stream, bsr_hessian_, step, spmv_scratch_, d_out, d_partials); return; } ComputeWeightedSquaredStepAsync(stream, cusparse_handle, csr_hessian_, csr_dims_.num_rows, - csr_dims_.num_cols, csr_dims_.num_nonzeros, step, buffer, d_out, - d_partials); + csr_dims_.num_cols, csr_dims_.num_nonzeros, step, spmv_scratch_, + buffer, d_out, d_partials); } bool NormalEquations::InitializeSolver(cudaStream_t stream, CSRSparseLinearSolver &solver, diff --git a/cunls/minimizer/normal_equations.h b/cunls/minimizer/normal_equations.h index cab57f9..87adb6c 100644 --- a/cunls/minimizer/normal_equations.h +++ b/cunls/minimizer/normal_equations.h @@ -133,9 +133,9 @@ class NormalEquations { BSRSparseMatrix bsr_hessian_; BSRSparseMatrix bsr_lhs_; - CSRMatrixDimensions csr_dims_; ///< Cached dims for the scalar SpMV. - dvector tile_row_scratch_; ///< Tile-to-block-row map for scaling. - dvector block_spmv_scratch_; ///< SpMV result for the block path. + CSRMatrixDimensions csr_dims_; ///< Cached dims for the scalar SpMV. + dvector tile_row_scratch_; ///< Tile-to-block-row map for scaling. + dvector spmv_scratch_; ///< SpMV result, either layout. int block_size_ = 1; }; diff --git a/cunls/minimizer/sparse_matrix.cu b/cunls/minimizer/sparse_matrix.cu index 9124d56..0f18d63 100644 --- a/cunls/minimizer/sparse_matrix.cu +++ b/cunls/minimizer/sparse_matrix.cu @@ -48,7 +48,12 @@ namespace cunls { * @param num_nonzeros Output argument for the number of non-zero elements * * The number of columns is determined by finding the maximum column index + 1. - * Requires matrix to have at least one row and one non-zero element. + * + * An empty system is well-formed, not an error: a fully-constrained problem + * leaves nothing to solve, which is encoded as `row_offsets == {0}` (one more + * offset than rows, with no rows) and reported here as 0 x 0 with no nonzeros. + * Because the column count is a reduction over the column indices, it cannot be + * recovered when there are no nonzeros, and is reported as 0. */ void ExtractMatrixMetadata(cudaStream_t stream, const CSRSparseMatrix &matrix, int &num_rows, int &num_cols, int &num_nonzeros) { @@ -242,8 +247,11 @@ __global__ void scale_symmetric_csr_rows_kernel(const int *__restrict__ row_offs } void ScaleSymmetricCSR(cudaStream_t stream, CSRSparseMatrix &matrix, const dvector &scale) { - int num_rows = static_cast(matrix.row_offsets.size() - 1); + int num_rows = static_cast(matrix.NumRows()); assert(static_cast(scale.size()) == num_rows); + if (num_rows == 0) { + return; + } dim3 block(WARP_SIZE, 8); dim3 grid((num_rows + block.y - 1) / block.y); scale_symmetric_csr_rows_kernel<<>>( @@ -287,9 +295,12 @@ void InvertSqrtWithFloorInPlace(cudaStream_t stream, dvector &v, float fl * col). */ void ExtractDiagonal(cudaStream_t stream, const CSRSparseMatrix &matrix, dvector &diagonal) { - size_t num_rows = matrix.row_offsets.size() - 1; + size_t num_rows = matrix.NumRows(); diagonal.resize(num_rows); + if (num_rows == 0) { + return; + } dim3 block(32, 4); dim3 grid((num_rows + block.y - 1) / block.y); extract_diagonal_kernel<<>>(matrix.row_offsets.data(), @@ -371,7 +382,12 @@ void AddScaledDiagonal(cudaStream_t stream, float scale, const dvector &d int num_rows = diagonal.size(); assert(num_rows + 1 == matrix.row_offsets.size()); + // The copy is the caller-visible half of the contract, so it happens even + // when there is no diagonal left to damp. CopyCSRSparseMatrix(stream, matrix, result); + if (num_rows == 0) { + return; + } // Launch one warp (32 threads) per row for warp-cooperative diagonal search constexpr int block_size = 256; // Must be multiple of WARP_SIZE @@ -410,17 +426,19 @@ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const dvector & /** * @brief Async sparse-weighted squared step: d_out[0] = step^T A step. * - * Runs the SpMV into a slice of `buffer` and reduces against `step`, so the - * whole thing stays on the stream with no host synchronization. + * Runs the SpMV into `scratch` and reduces against `step`, so the whole thing + * stays on the stream with no host synchronization. `scratch` is caller-owned + * and resized here; it ties the buffer's lifetime to the object driving the + * stream rather than to the thread. */ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, int num_rows, int num_cols, int num_nonzeros, const dvector &step, - dvector &buffer, float *d_out, float *d_partials) { - static thread_local dvector spmv_result; + dvector &scratch, dvector &buffer, + float *d_out, float *d_partials) { SpMVImpl(stream, handle, matrix, num_rows, num_cols, num_nonzeros, /*transpose_matrix=*/false, - step, spmv_result, buffer); - DotProductToDevice(stream, step.data(), spmv_result.data(), step.size(), d_out, d_partials); + step, scratch, buffer); + DotProductToDevice(stream, step.data(), scratch.data(), step.size(), d_out, d_partials); } /** diff --git a/cunls/minimizer/sparse_matrix.h b/cunls/minimizer/sparse_matrix.h index 6f3a795..6527c82 100644 --- a/cunls/minimizer/sparse_matrix.h +++ b/cunls/minimizer/sparse_matrix.h @@ -117,11 +117,14 @@ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, const dvector & * @brief Async sparse-weighted squared step: d_out[0] = step^T A step. * * Performs SpMV (A*step) then dot(step, A*step) into d_out, all on the stream. + * `scratch` holds the SpMV result and is resized as needed; the caller owns it + * so that one buffer per driver object cannot be shared across streams. */ void ComputeWeightedSquaredStepAsync(cudaStream_t stream, void *handle, const CSRSparseMatrix &matrix, int num_rows, int num_cols, int num_nonzeros, const dvector &step, - dvector &buffer, float *d_out, float *d_partials); + dvector &scratch, dvector &buffer, + float *d_out, float *d_partials); void ElementwiseMultiplyInPlace(cudaStream_t stream, float *a, const float *b, size_t n); diff --git a/tests/block_hessian_assembler_test.cpp b/tests/block_hessian_assembler_test.cpp index 1a5f8a6..374b5a0 100644 --- a/tests/block_hessian_assembler_test.cpp +++ b/tests/block_hessian_assembler_test.cpp @@ -933,6 +933,10 @@ TEST(HessianStorageTest, BlockDiagonalOpsMatchScalar) { * Same matrix, same right-hand side, same preconditioner -- so the solution and * the iteration count should match. A divergence here points at the block * SpMV or the block-Jacobi tile gather rather than at assembly. + * + * The tolerance is one this preconditioner actually reaches on a pose graph. + * Asking for more only drives both solvers into the iteration cap, where the + * counts match trivially and the comparison stops testing anything. */ TEST(HessianStorageTest, PcgAgreesBetweenStorages) { auto data = MakePoseGraph(2048, /*fix_first_pose=*/true); @@ -956,7 +960,7 @@ TEST(HessianStorageTest, PcgAgreesBetweenStorages) { BlockSparsePCGOptions pcg_options; pcg_options.block_size = 6; pcg_options.max_iterations = 500; - pcg_options.relative_tolerance = 1e-6f; + pcg_options.relative_tolerance = 1e-3f; dvector x_csr(n); dvector x_bsr(n); @@ -975,10 +979,20 @@ TEST(HessianStorageTest, PcgAgreesBetweenStorages) { x_csr.CopyToHost(a.data(), n); x_bsr.CopyToHost(b.data(), n); - // Same iteration count means the two are tracking the same recurrence. The - // solutions themselves only agree to the level PCG was asked for: it stops on - // a relative residual, so the reordered block SpMV moves x within that ball. - EXPECT_EQ(csr_solver.LastIterations(), bsr_solver.LastIterations()); + // Both must stop on the residual test, not on the cap -- otherwise the counts + // agree only because they were clamped to the same number. + ASSERT_LT(csr_solver.LastIterations(), pcg_options.max_iterations); + ASSERT_LT(bsr_solver.LastIterations(), pcg_options.max_iterations); + + // Near-equal iteration counts mean the two are tracking the same recurrence. + // Not exactly equal: the block SpMV sums each row in a different order, so the + // residual differs in its last bits and can cross the threshold an iteration + // early or late. A real fault in the SpMV or the tile gather changes the + // count by far more than this, or stops it converging at all. + EXPECT_NEAR(csr_solver.LastIterations(), bsr_solver.LastIterations(), 2); + + // The solutions themselves only agree to the level PCG was asked for: it stops + // on a relative residual, so the reordered block SpMV moves x within that ball. const float tol = 1e-3f * std::max(MaxAbs(a), 1e-6f); for (size_t i = 0; i < n; i++) { ASSERT_NEAR(a[i], b[i], tol) << "solution mismatch at " << i; @@ -1232,15 +1246,24 @@ TEST(BlockHessianAssemblerTest, MatchesReferenceWithColumnScaling) { ASSERT_TRUE(data->problem.CheckConsistency()); CudaStream stream; - MinimizerOptions options; - options.column_scaling = ColumnScaling::HessianDiagonal; + MinimizerOptions block_options; + block_options.column_scaling = ColumnScaling::HessianDiagonal; + + // Column scaling reads the Hessian diagonal and rescales the LHS in place, and + // both of those are implemented once per layout. The reference therefore has + // to be the scalar path -- selecting a CSR-only backend is how a caller gets + // there -- or the comparison is block storage against itself. + MinimizerOptions scalar_options = block_options; + scalar_options.sparse_linear_solver_type = SparseLinearSolverType::cuDSS; - SystemBuilder reference(options); + SystemBuilder reference(scalar_options); reference.Build(stream.GetStream(), data->problem); + ASSERT_FALSE(reference.UsesBlockStorage()); std::vector expected = Snapshot(reference, stream.GetStream()).values; - SystemBuilder block(options); + SystemBuilder block(block_options); block.Build(stream.GetStream(), data->problem); + ASSERT_TRUE(block.UsesBlockStorage()); std::vector actual = Snapshot(block, stream.GetStream()).values; ASSERT_EQ(expected.size(), actual.size()); diff --git a/tests/sparse_matrix_test.cpp b/tests/sparse_matrix_test.cpp index b9612b3..ad7787f 100644 --- a/tests/sparse_matrix_test.cpp +++ b/tests/sparse_matrix_test.cpp @@ -375,11 +375,12 @@ TEST_F(SparseMatrixTest, ComputeWeightedSquaredStepSecond) { auto range = this->profiler_domain_.CreateDomainRange("ComputeWeightedSquaredStepSecond"); int num_rows = 0, num_cols = 0, num_nonzeros = 0; ExtractMatrixMetadata(stream.GetStream(), input_matrix, num_rows, num_cols, num_nonzeros); + dvector spmv_scratch; result = RunAsyncReduction(stream.GetStream(), dsteps.size(), [&](float *d_out, float *d_partials) { ComputeWeightedSquaredStepAsync(stream.GetStream(), handle, input_matrix, num_rows, - num_cols, num_nonzeros, dsteps, buffer, d_out, - d_partials); + num_cols, num_nonzeros, dsteps, spmv_scratch, buffer, + d_out, d_partials); }); } @@ -412,4 +413,53 @@ TEST(SparseMatrixColumnScaling, SymmetricScaling2x2) { ASSERT_NEAR(out[3], 1.f, 1e-4f); } +/** + * @brief A 0x0 system is a valid CSR matrix, and every op over it is a no-op. + * + * `row_offsets == {0}` is the well-formed encoding of an empty matrix -- one + * more offset than rows, with no rows. Every one of these ops derives its grid + * from the row count, so a zero row count must be recognized before launch + * rather than turned into an empty grid. + */ +TEST(SparseMatrixEmptySystem, OperationsOnZeroRowMatrixAreNoOps) { + CSRSparseMatrix empty; + test_utils::CreateCSRSparseMatrix({0}, {}, {}, empty); + ASSERT_EQ(empty.NumRows(), 0); + + CudaStream stream; + dvector scale; + dvector diagonal; + + ScaleSymmetricCSR(stream.GetStream(), empty, scale); + ExtractDiagonal(stream.GetStream(), empty, diagonal); + EXPECT_EQ(diagonal.size(), 0u); + + // Damping still has to produce the (empty) output matrix, since callers read + // `damped` afterwards regardless of size. + CSRSparseMatrix damped; + AddScaledDiagonal(stream.GetStream(), 1e-3f, diagonal, empty, damped); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream.GetStream())); + + EXPECT_EQ(damped.NumRows(), 0); + EXPECT_EQ(damped.NumNonZeros(), 0); + ASSERT_EQ(damped.row_offsets.size(), 1u); + hvector offsets(1); + damped.row_offsets.CopyToHost(offsets.data(), offsets.size()); + EXPECT_EQ(offsets[0], 0); +} + +/** @brief Metadata of an empty-but-well-formed CSR is 0x0 with no nonzeros. */ +TEST(SparseMatrixEmptySystem, MetadataOfZeroRowMatrixIsAllZero) { + CSRSparseMatrix empty; + test_utils::CreateCSRSparseMatrix({0}, {}, {}, empty); + + CudaStream stream; + int num_rows = -1, num_cols = -1, num_nonzeros = -1; + ExtractMatrixMetadata(stream.GetStream(), empty, num_rows, num_cols, num_nonzeros); + + EXPECT_EQ(num_rows, 0); + EXPECT_EQ(num_cols, 0); + EXPECT_EQ(num_nonzeros, 0); +} + } // namespace cunls From 21f19be04be492897051ad21698616f3c9895bb5 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Sat, 1 Aug 2026 21:33:23 -0700 Subject: [PATCH 4/7] Fix discussions v2 --- .../linear_solver/block_sparse_pcg_solver.cu | 8 ++++++- cunls/linear_solver/block_sparse_pcg_solver.h | 22 +++++++++---------- cunls/state/state_batch_ops.cu | 2 +- docs/sphinx/api/minimizer.rst | 12 +++++----- 4 files changed, 26 insertions(+), 18 deletions(-) diff --git a/cunls/linear_solver/block_sparse_pcg_solver.cu b/cunls/linear_solver/block_sparse_pcg_solver.cu index 510beb3..a6f4273 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.cu +++ b/cunls/linear_solver/block_sparse_pcg_solver.cu @@ -652,7 +652,13 @@ __global__ void ExtractScalarJacobi(const int *__restrict__ row_off, } float d = 0.f; GatherTileRow(row_off, col_idx, values, block_storage, block_size, row_start + idx, 1, 0, &d); - factors[factor_offset + idx] = fmaxf(fabsf(d), pivot_floor); + // Floor the magnitude but keep the sign, as the B > 1 pivots do: this is the + // same preconditioner at B == 1, so a diagonal entry must not be scaled one + // way here and the other way one block size up. + if (fabsf(d) < pivot_floor) { + d = (d >= 0.f) ? pivot_floor : -pivot_floor; + } + factors[factor_offset + idx] = d; } // ============================================================================= diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h index 4924c68..ebfc1f3 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.h +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -209,6 +209,17 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; + /** @copydoc CSRSparseLinearSolver::SupportsBlockStorage */ + bool SupportsBlockStorage() const override { return true; } + + /** @copydoc CSRSparseLinearSolver::Initialize */ + bool Initialize(cudaStream_t stream, const Problem &problem, const BSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) override; + + /** @copydoc CSRSparseLinearSolver::Solve */ + bool Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) override; + /** * @brief Runs the PCG loop on `H x = b`, writing into @p result. * @@ -228,17 +239,6 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { * reset). * @return true on success, false on dimension mismatch. */ - /** @copydoc CSRSparseLinearSolver::SupportsBlockStorage */ - bool SupportsBlockStorage() const override { return true; } - - /** @copydoc CSRSparseLinearSolver::Initialize */ - bool Initialize(cudaStream_t stream, const Problem &problem, const BSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) override; - - /** @copydoc CSRSparseLinearSolver::Solve */ - bool Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, const dvector &rhs, - dvector &result) override; - bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; diff --git a/cunls/state/state_batch_ops.cu b/cunls/state/state_batch_ops.cu index 54751ee..c320e8d 100644 --- a/cunls/state/state_batch_ops.cu +++ b/cunls/state/state_batch_ops.cu @@ -124,7 +124,7 @@ StateBatchOps::StateBatchOps(cudaStream_t stream, const std::vector &column_offsets) { diff --git a/docs/sphinx/api/minimizer.rst b/docs/sphinx/api/minimizer.rst index c0fae59..ae950a5 100644 --- a/docs/sphinx/api/minimizer.rst +++ b/docs/sphinx/api/minimizer.rst @@ -88,9 +88,10 @@ Wikipedia links above for convergence and damping strategies. :code:`MinimizerOptions::column_scaling` can re-scale the normal equations with a diagonal :math:`S`: the linear solve uses :math:`S H S \, z = S b` with :math:`H = J^T J` and :math:`b = -J^T r`, then applies the physical tangent step -:math:`\Delta x = S z`. Modes are: no scaling (default); :math:`S_{ii} = 1/\sqrt{H_{ii}}` -(with a small floor for stability); or :math:`S_{jj} = 1/\|J_{:,j}\|_2` from the -CSR Jacobian. For Levenberg-Marquardt, damping uses the diagonal of the **scaled** +:math:`\Delta x = S z`. Modes are: no scaling (default); or +:math:`S_{ii} = 1/\sqrt{H_{ii}}` with a small floor for stability, which is +equivalently :math:`1/\|J_{:,j}\|_2` since :math:`H_{jj} = \|J_{:,j}\|_2^2`. +For Levenberg-Marquardt, damping uses the diagonal of the **scaled** Hessian: :math:`S H S + \lambda \operatorname{diag}(S H S)`. ================================================================================ @@ -491,8 +492,9 @@ values and then override individual fields. Enum used by ``MinimizerOptions.column_scaling`` (and ``LevenbergMarquardtMinimizerOptions.base_options.column_scaling``): - **none** — identity scaling (standard :math:`H \Delta x = -J^T r`). -- **hessian_diagonal** — :math:`S_{ii} = 1 / \sqrt{H_{ii}}` with a numerical floor. -- **jacobian_column_norm** — :math:`S_{jj} = 1 / \|J_{:,j}\|_2` from the CSR Jacobian. +- **hessian_diagonal** — :math:`S_{ii} = 1 / \sqrt{H_{ii}}` with a numerical + floor. Equivalently :math:`1 / \|J_{:,j}\|_2`, since + :math:`H_{jj} = \|J_{:,j}\|_2^2`. For LM, damping uses the diagonal of the **scaled** Hessian. See the C++ column-scaling theory note earlier in this page. From b7a8f50667428242dbb513035c114ad0e1abdd39 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Sun, 2 Aug 2026 09:33:56 -0700 Subject: [PATCH 5/7] Fix CI --- python/src/bind_types.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/python/src/bind_types.cpp b/python/src/bind_types.cpp index 9a2880c..56f9fef 100644 --- a/python/src/bind_types.cpp +++ b/python/src/bind_types.cpp @@ -31,7 +31,6 @@ #include "cunls/linear_solver/sparse_linear_solver.h" #include "cunls/minimizer/gauss_newton_minimizer.h" #include "cunls/minimizer/levenberg_marquardt_minimizer.h" -#include "cunls/minimizer/sparse_matrix_multiplier.h" // Convert a Python object to a raw device pointer (uintptr_t). // From a2e3c0cbf74f252ec44d6cf88a99deea3153189d Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 12 Aug 2026 17:22:55 -0700 Subject: [PATCH 6/7] Improve consistency --- cunls/linear_solver/CMakeLists.txt | 1 + cunls/linear_solver/block_sparse_pcg_solver.h | 10 +- .../cudss_sparse_linear_solver.h | 8 +- cunls/linear_solver/dense_cholesky_solver.cu | 141 ++----------- cunls/linear_solver/dense_cholesky_solver.h | 68 ++----- cunls/linear_solver/dense_linear_solver.cu | 161 +++------------ cunls/linear_solver/dense_linear_solver.h | 85 ++------ .../linear_solver/dense_linear_solver_base.cu | 188 ++++++++++++++++++ .../linear_solver/dense_linear_solver_base.h | 132 ++++++++++++ cunls/linear_solver/dense_qr_solver.cu | 157 +++------------ cunls/linear_solver/dense_qr_solver.h | 58 ++---- cunls/linear_solver/llms.txt | 25 ++- cunls/linear_solver/sparse_linear_solver.cpp | 35 ++-- cunls/linear_solver/sparse_linear_solver.h | 40 ++-- ...r_solver.h => sparse_linear_solver_base.h} | 61 +++--- cunls/minimizer/block_hessian_assembler.h | 18 +- cunls/minimizer/gauss_newton_minimizer.cu | 4 +- cunls/minimizer/normal_equations.cu | 26 ++- cunls/minimizer/normal_equations.h | 8 +- docs/sphinx/api/linear_solver.rst | 29 ++- docs/sphinx/api/minimizer.rst | 14 +- tests/block_hessian_assembler_test.cpp | 139 ++++++++++++- tests/bsr_expansion.cu | 8 +- 23 files changed, 776 insertions(+), 640 deletions(-) create mode 100644 cunls/linear_solver/dense_linear_solver_base.cu create mode 100644 cunls/linear_solver/dense_linear_solver_base.h rename cunls/linear_solver/{csr_sparse_linear_solver.h => sparse_linear_solver_base.h} (65%) diff --git a/cunls/linear_solver/CMakeLists.txt b/cunls/linear_solver/CMakeLists.txt index ff0d5af..18ac580 100644 --- a/cunls/linear_solver/CMakeLists.txt +++ b/cunls/linear_solver/CMakeLists.txt @@ -3,6 +3,7 @@ add_library(cunls_linear_solver OBJECT cudss_sparse_linear_solver.cpp dense_cholesky_solver.cu dense_linear_solver.cu + dense_linear_solver_base.cu dense_qr_solver.cu sparse_linear_solver.cpp ) diff --git a/cunls/linear_solver/block_sparse_pcg_solver.h b/cunls/linear_solver/block_sparse_pcg_solver.h index ebfc1f3..b33e6da 100644 --- a/cunls/linear_solver/block_sparse_pcg_solver.h +++ b/cunls/linear_solver/block_sparse_pcg_solver.h @@ -23,7 +23,7 @@ #include #include "cunls/common/cusparse_helper.h" -#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" namespace cunls { @@ -155,7 +155,7 @@ struct BlockSparsePCGOptions { * 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 { +class BlockSparsePCGSolver : public SparseLinearSolver { public: /** * @brief Constructs the solver with the given options. @@ -209,14 +209,14 @@ class BlockSparsePCGSolver : public CSRSparseLinearSolver { bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) final; - /** @copydoc CSRSparseLinearSolver::SupportsBlockStorage */ + /** @copydoc SparseLinearSolver::SupportsBlockStorage */ bool SupportsBlockStorage() const override { return true; } - /** @copydoc CSRSparseLinearSolver::Initialize */ + /** @copydoc SparseLinearSolver::Initialize */ bool Initialize(cudaStream_t stream, const Problem &problem, const BSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) override; - /** @copydoc CSRSparseLinearSolver::Solve */ + /** @copydoc SparseLinearSolver::Solve */ bool Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, const dvector &rhs, dvector &result) override; diff --git a/cunls/linear_solver/cudss_sparse_linear_solver.h b/cunls/linear_solver/cudss_sparse_linear_solver.h index 91939d7..cd2a433 100644 --- a/cunls/linear_solver/cudss_sparse_linear_solver.h +++ b/cunls/linear_solver/cudss_sparse_linear_solver.h @@ -22,7 +22,7 @@ #include #include "cunls/common/cudss_helper.h" -#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" namespace cunls { @@ -59,14 +59,14 @@ struct cuDSSLinearSolverOptions { * Ax = b where A is symmetric positive definite. It uses NVIDIA's cuDSS * library for GPU-accelerated direct factorization. */ -class cuDSSLinearSolver : public CSRSparseLinearSolver { +class cuDSSLinearSolver : public SparseLinearSolver { public: // This backend consumes CSR only; SupportsBlockStorage() stays false, so the // base class's block-storage overloads are never called on it. The // using-declarations keep them visible rather than hidden by the CSR // overrides below. - using CSRSparseLinearSolver::Initialize; - using CSRSparseLinearSolver::Solve; + using SparseLinearSolver::Initialize; + using SparseLinearSolver::Solve; /** * @brief Constructs a cuDSS linear solver. diff --git a/cunls/linear_solver/dense_cholesky_solver.cu b/cunls/linear_solver/dense_cholesky_solver.cu index 87d2fbb..b45bb6b 100644 --- a/cunls/linear_solver/dense_cholesky_solver.cu +++ b/cunls/linear_solver/dense_cholesky_solver.cu @@ -23,122 +23,48 @@ #include "cunls/linear_solver/dense_cholesky_solver.h" namespace cunls { -namespace { -constexpr int kWarpSize = 32; +bool DenseCholeskySolver::FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) { + auto handle = static_cast(cusolver_handle_.GetHandle(stream)); -// Writes the CSR matrix into a dense buffer in row-major order. cuSOLVER -// (potrf, potrs) expects column-major layout, but this solver is only used -// for symmetric matrices (the Gauss-Newton Hessian J^T J), for which -// row-major == column-major (A == A^T). Do NOT use this kernel for -// non-symmetric inputs without first transposing the layout. -__global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, - const int *__restrict__ col_ids, - const float *__restrict__ values, - int num_rows, - float *__restrict__ dense_matrix) { - const int row = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpSize; - if (row >= num_rows) { - return; - } - const int lane = threadIdx.x % kWarpSize; - const int row_start = row_offsets[row]; - const int row_end = row_offsets[row + 1]; - float *const dense_row = dense_matrix + row * num_rows; - for (int idx = row_start + lane; idx < row_end; idx += kWarpSize) { - dense_row[col_ids[idx]] = values[idx]; - } -} - -} // namespace - -bool DenseCholeskySolver::Initialize(cudaStream_t stream, - const Problem & /*problem*/, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, - dvector &result) { - const size_t matrix_size = spd_matrix.NumRows(); - if (matrix_size != rhs.size()) { - LogError("LHS size: {} does not match RHS size: {}", matrix_size, - rhs.size()); - return false; - } - if (matrix_size != result.size()) { - LogError("LHS size: {} does not match result size: {}", matrix_size, - result.size()); - return false; - } - EnsureBuffersSize(stream, matrix_size); - return true; -} - -bool DenseCholeskySolver::Solve(cudaStream_t stream, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, - dvector &result) { - const size_t matrix_size = spd_matrix.NumRows(); - if (matrix_size != rhs.size()) { - LogError("LHS size: {} does not match RHS size: {}", matrix_size, - rhs.size()); - return false; - } - if (matrix_size != result.size()) { - LogError("LHS size: {} does not match result size: {}", matrix_size, - result.size()); - return false; - } - if (matrix_size == 0) { - return true; - } - - EnsureBuffersSize(stream, matrix_size); - ConvertCSRToDense(stream, spd_matrix, dense_matrix_); - - const int n = static_cast(matrix_size); - auto handle = - static_cast(cusolver_handle_.GetHandle(stream)); - - THROW_ON_CUSOLVER_ERROR( - cusolverDnSpotrf(handle, CUBLAS_FILL_MODE_LOWER, n, dense_matrix_.data(), - n, workspace_.data(), - static_cast(workspace_.size()), dev_info_.data())); + THROW_ON_CUSOLVER_ERROR(cusolverDnSpotrf(handle, CUBLAS_FILL_MODE_LOWER, n, dense_matrix_.data(), + n, workspace_.data(), + static_cast(workspace_.size()), dev_info_.data())); if (safety_checks_enabled_) { // Check devInfo from potrf before proceeding to potrs, since potrs would // overwrite it. devInfo > 0 means the leading minor of order devInfo is // not positive-definite; devInfo < 0 means the devInfo-th parameter was // invalid. - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dev_info_pinned_.data(), - dev_info_.data(), sizeof(int), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dev_info_pinned_.data(), dev_info_.data(), sizeof(int), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); if (dev_info_pinned_[0] != 0) { - LogError("Cholesky factorization failed (devInfo = {}). Matrix is likely " - "not positive-definite.", - dev_info_pinned_[0]); + LogError( + "Cholesky factorization failed (devInfo = {}). Matrix is likely " + "not positive-definite.", + dev_info_pinned_[0]); return false; } } // potrs solves in-place on the RHS buffer, so copy rhs -> result first. - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(result.data(), rhs.data(), - n * sizeof(float), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(result.data(), rhs.data(), n * sizeof(float), cudaMemcpyDeviceToDevice, stream)); THROW_ON_CUSOLVER_ERROR(cusolverDnSpotrs(handle, CUBLAS_FILL_MODE_LOWER, n, 1, - dense_matrix_.data(), n, - result.data(), n, dev_info_.data())); + dense_matrix_.data(), n, result.data(), n, + dev_info_.data())); if (safety_checks_enabled_) { - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dev_info_pinned_.data(), - dev_info_.data(), sizeof(int), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dev_info_pinned_.data(), dev_info_.data(), sizeof(int), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); if (dev_info_pinned_[0] < 0) { - LogError("cusolverDnSpotrs reported invalid parameter at index {}.", - -dev_info_pinned_[0]); + LogError("cusolverDnSpotrs reported invalid parameter at index {}.", -dev_info_pinned_[0]); return false; } if (dev_info_pinned_[0] > 0) { @@ -151,45 +77,20 @@ bool DenseCholeskySolver::Solve(cudaStream_t stream, } void DenseCholeskySolver::EnsureBuffersSize(cudaStream_t stream, size_t n) { - const size_t matrix_elements = n * n; - if (dense_matrix_.size() != matrix_elements) { - dense_matrix_.resize(matrix_elements); - } if (dev_info_.size() != 1) { dev_info_.resize(1); dev_info_pinned_.resize(1); } if (n != last_n_ && n > 0) { - auto handle = - static_cast(cusolver_handle_.GetHandle(stream)); + auto handle = static_cast(cusolver_handle_.GetHandle(stream)); int lwork = 0; - THROW_ON_CUSOLVER_ERROR(cusolverDnSpotrf_bufferSize( - handle, CUBLAS_FILL_MODE_LOWER, static_cast(n), - dense_matrix_.data(), static_cast(n), &lwork)); + THROW_ON_CUSOLVER_ERROR(cusolverDnSpotrf_bufferSize(handle, CUBLAS_FILL_MODE_LOWER, + static_cast(n), dense_matrix_.data(), + static_cast(n), &lwork)); workspace_.resize(static_cast(lwork)); last_n_ = n; } } -void DenseCholeskySolver::ConvertCSRToDense(cudaStream_t stream, - const CSRSparseMatrix &matrix, - dvector &dense_matrix) { - const int num_rows = static_cast(matrix.NumRows()); - if (num_rows == 0) { - return; - } - THROW_ON_CUDA_ERROR(cudaMemsetAsync( - dense_matrix.data(), 0, - static_cast(num_rows) * num_rows * sizeof(float), stream)); - - constexpr int kThreads = 256; - constexpr int kWarpsPerBlock = kThreads / kWarpSize; - const int blocks = (num_rows + kWarpsPerBlock - 1) / kWarpsPerBlock; - csr_to_dense_kernel<<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_rows, dense_matrix.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/dense_cholesky_solver.h b/cunls/linear_solver/dense_cholesky_solver.h index 9f58f23..72e8ebb 100644 --- a/cunls/linear_solver/dense_cholesky_solver.h +++ b/cunls/linear_solver/dense_cholesky_solver.h @@ -21,76 +21,52 @@ #include "cunls/common/cusolver_helper.h" #include "cunls/common/types.h" -#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/linear_solver/dense_linear_solver_base.h" namespace cunls { /** * @brief Dense GPU linear solver based on Cholesky factorization via cuSOLVER. * - * Converts the input CSR symmetric positive-definite matrix to a dense matrix - * and solves A x = b via: + * Densifies the symmetric positive-definite coefficient matrix (from either + * sparse layout; see DenseLinearSolverBase) and solves A x = b via: * 1) Cholesky factorization: A = L L^T (cusolverDnSpotrf) * 2) Triangular solve using the factor (cusolverDnSpotrs) * * Since the input matrix is symmetric, the row-major dense representation - * produced by CSR conversion is identical to column-major, so no transpose - * is required for the column-major cuSOLVER API. + * produced by the scatter is identical to column-major, so no transpose is + * required for the column-major cuSOLVER API. * * Returns false from Solve() if the matrix is not positive-definite (cuSOLVER * reports a non-zero devInfo from potrf). */ -class DenseCholeskySolver : public CSRSparseLinearSolver { - public: - // This backend consumes CSR only; SupportsBlockStorage() stays false, so the - // base class's block-storage overloads are never called on it. The - // using-declarations keep them visible rather than hidden by the CSR - // overrides below. - using CSRSparseLinearSolver::Initialize; - using CSRSparseLinearSolver::Solve; +class DenseCholeskySolver : public DenseLinearSolverBase { + protected: + /** @copydoc DenseLinearSolverBase::EnsureBuffersSize */ + void EnsureBuffersSize(cudaStream_t stream, size_t n) final; /** - * @brief Validates dimensions and pre-allocates internal buffers. - * - * @param stream CUDA stream used to query the cuSOLVER workspace size. - * @param spd_matrix The SPD 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). - * @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; - - /** - * @brief Converts CSR to dense and solves via Cholesky factorization. + * @brief Factorizes the dense matrix via Cholesky and solves. * * The pipeline is: - * 1. CSR -> dense conversion. - * 2. cusolverDnSpotrf (in-place Cholesky factorization). - * 3. devInfo check after potrf (if safety checks enabled). - * 4. Copy rhs into result (potrs works in-place on B). - * 5. cusolverDnSpotrs (triangular solve). - * 6. devInfo check after potrs (if safety checks enabled). + * 1. cusolverDnSpotrf (in-place Cholesky factorization). + * 2. devInfo check after potrf (if safety checks enabled). + * 3. Copy rhs into result (potrs works in-place on B). + * 4. cusolverDnSpotrs (triangular solve). + * 5. devInfo check after potrs (if safety checks enabled). * * @param stream CUDA stream for asynchronous GPU operations. - * @param spd_matrix The SPD 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). - * @return true on success, false on dimension mismatch, non-SPD matrix - * (devInfo > 0 from potrf), or invalid parameter from potrs - * (devInfo < 0). + * @param n Matrix dimension. + * @param rhs The right-hand side vector b. + * @param result Output vector x. + * @return true on success, false on a non-SPD matrix (devInfo > 0 from + * potrf) or an invalid parameter from potrs (devInfo < 0). */ - bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, - dvector &result) final; + bool FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) final; private: - void EnsureBuffersSize(cudaStream_t stream, size_t n); - - void ConvertCSRToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, - dvector &dense_matrix); - cuSolverHandle cusolver_handle_; - dvector dense_matrix_; dvector workspace_; dvector dev_info_; pvector dev_info_pinned_; diff --git a/cunls/linear_solver/dense_linear_solver.cu b/cunls/linear_solver/dense_linear_solver.cu index 3cf2be5..77e4a9f 100644 --- a/cunls/linear_solver/dense_linear_solver.cu +++ b/cunls/linear_solver/dense_linear_solver.cu @@ -41,12 +41,9 @@ constexpr int kMaxBlockSize = kMaxWarpsPerBlock * kWarpSize; /// is >= n. Keeping the block small when n is small reduces wasted /// threads that would just idle in the synchronization barriers. inline int SelectBlockSize(int n) { - if (n <= 32) - return 32; - if (n <= 64) - return 64; - if (n <= 128) - return 128; + if (n <= 32) return 32; + if (n <= 64) return 64; + if (n <= 128) return 128; return kMaxBlockSize; } @@ -75,34 +72,6 @@ __device__ __forceinline__ float WarpReduceSum(float val) { return val; } -// --------------------------------------------------------------------------- -// CSR -> Dense conversion kernel -// --------------------------------------------------------------------------- - -/// @brief Scatters CSR values into a dense row-major matrix. -/// -/// Each warp processes one row: threads in the warp iterate over the row's -/// non-zero entries in parallel and write them to the corresponding column -/// position in the dense output. The caller must zero-initialize -/// @p dense_matrix before launch. -__global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, - const int *__restrict__ col_ids, - const float *__restrict__ values, - int num_rows, - float *__restrict__ dense_matrix) { - const int row = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpSize; - if (row >= num_rows) { - return; - } - const int lane = threadIdx.x % kWarpSize; - const int row_start = row_offsets[row]; - const int row_end = row_offsets[row + 1]; - float *const dense_row = dense_matrix + row * num_rows; - for (int idx = row_start + lane; idx < row_end; idx += kWarpSize) { - dense_row[col_ids[idx]] = values[idx]; - } -} - // --------------------------------------------------------------------------- // Pivoted LDLT factorization kernel // --------------------------------------------------------------------------- @@ -129,9 +98,11 @@ __global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, /// @p check_status != 0). /// @param check_status When non-zero, enable pivot-value checks and status /// reporting; when zero, skip them for lower latency. -__global__ void factorize_symmetric_pivoted_ldlt_kernel( - const float *__restrict__ A, int n, float *__restrict__ ldlt, - int *__restrict__ permutation, int *__restrict__ status, int check_status) { +__global__ void factorize_symmetric_pivoted_ldlt_kernel(const float *__restrict__ A, int n, + float *__restrict__ ldlt, + int *__restrict__ permutation, + int *__restrict__ status, + int check_status) { if (blockIdx.x != 0) { return; } @@ -228,8 +199,7 @@ __global__ void factorize_symmetric_pivoted_ldlt_kernel( } __syncthreads(); if (!factorization_ok) { - if (tid == 0) - *status = 0; + if (tid == 0) *status = 0; return; } } @@ -387,7 +357,7 @@ __global__ void solve_from_pivoted_ldlt_kernel( } } -} // namespace +} // namespace // --------------------------------------------------------------------------- // Host-side launcher functions @@ -400,18 +370,16 @@ __global__ void solve_from_pivoted_ldlt_kernel( /// @p check_status is true. /// @param check_status When true, the kernel checks pivots for near-zero /// values and reports status; when false, skips checks. -void FactorizeSymmetricPivotedLDLT(cudaStream_t stream, - const float *dense_symmetric_matrix, int n, - float *ldlt_factor, int *permutation, - int *status, bool check_status) { +void FactorizeSymmetricPivotedLDLT(cudaStream_t stream, const float *dense_symmetric_matrix, int n, + float *ldlt_factor, int *permutation, int *status, + bool check_status) { if (n == 0) { return; } const int threads = SelectBlockSize(n); factorize_symmetric_pivoted_ldlt_kernel<<<1, threads, 0, stream>>>( - dense_symmetric_matrix, n, ldlt_factor, permutation, status, - check_status ? 1 : 0); + dense_symmetric_matrix, n, ldlt_factor, permutation, status, check_status ? 1 : 0); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -422,18 +390,17 @@ void FactorizeSymmetricPivotedLDLT(cudaStream_t stream, /// @p check_status is true. /// @param check_status When true, the kernel checks diagonal elements and /// reports status; when false, skips checks. -void SolveFromPivotedLDLT(cudaStream_t stream, const float *ldlt_factor, - const int *permutation, const float *rhs, int n, - float *permuted_rhs, float *intermediate_solution, - float *permuted_solution, float *solution, +void SolveFromPivotedLDLT(cudaStream_t stream, const float *ldlt_factor, const int *permutation, + const float *rhs, int n, float *permuted_rhs, + float *intermediate_solution, float *permuted_solution, float *solution, int *status, bool check_status) { if (n == 0) { return; } const int threads = SelectBlockSize(n); solve_from_pivoted_ldlt_kernel<<<1, threads, 0, stream>>>( - ldlt_factor, permutation, rhs, n, permuted_rhs, intermediate_solution, - permuted_solution, solution, status, check_status ? 1 : 0); + ldlt_factor, permutation, rhs, n, permuted_rhs, intermediate_solution, permuted_solution, + solution, status, check_status ? 1 : 0); THROW_ON_CUDA_ERROR(cudaGetLastError()); } @@ -441,65 +408,20 @@ void SolveFromPivotedLDLT(cudaStream_t stream, const float *ldlt_factor, // DenseLDLTSolver public API // --------------------------------------------------------------------------- -bool DenseLDLTSolver::Initialize(cudaStream_t stream, - const Problem & /*problem*/, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, - dvector &result) { - (void)stream; - const size_t matrix_size = spd_matrix.NumRows(); - if (matrix_size != rhs.size()) { - LogError("LHS size: {} does not match RHS size: {}", matrix_size, - rhs.size()); - return false; - } - if (matrix_size != result.size()) { - LogError("LHS size: {} does not match result size: {}", matrix_size, - result.size()); - return false; - } - EnsureBuffersSize(matrix_size); - return true; -} +bool DenseLDLTSolver::FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) { + FactorizeSymmetricPivotedLDLT(stream, dense_matrix_.data(), n, ldlt_factor_.data(), + permutation_.data(), status_.data(), safety_checks_enabled_); -bool DenseLDLTSolver::Solve(cudaStream_t stream, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) { - const size_t matrix_size = spd_matrix.NumRows(); - if (matrix_size != rhs.size()) { - LogError("LHS size: {} does not match RHS size: {}", matrix_size, - rhs.size()); - return false; - } - if (matrix_size != result.size()) { - LogError("LHS size: {} does not match result size: {}", matrix_size, - result.size()); - return false; - } - if (matrix_size == 0) { - return true; - } - - EnsureBuffersSize(matrix_size); - - ConvertCSRToDense(stream, spd_matrix, dense_matrix_); - - const int n = static_cast(matrix_size); - - FactorizeSymmetricPivotedLDLT(stream, dense_matrix_.data(), n, - ldlt_factor_.data(), permutation_.data(), - status_.data(), safety_checks_enabled_); - - SolveFromPivotedLDLT(stream, ldlt_factor_.data(), permutation_.data(), - rhs.data(), n, permuted_rhs_.data(), - intermediate_solution_.data(), permuted_solution_.data(), - result.data(), status_.data() + 1, + SolveFromPivotedLDLT(stream, ldlt_factor_.data(), permutation_.data(), rhs.data(), n, + permuted_rhs_.data(), intermediate_solution_.data(), + permuted_solution_.data(), result.data(), status_.data() + 1, safety_checks_enabled_); if (safety_checks_enabled_) { THROW_ON_CUDA_ERROR(cudaMemcpyAsync(status_pinned_.data(), status_.data(), - kNumStatuses * sizeof(int), - cudaMemcpyDeviceToHost, stream)); + kNumStatuses * sizeof(int), cudaMemcpyDeviceToHost, + stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); if (status_pinned_[0] == 0) { @@ -518,11 +440,8 @@ bool DenseLDLTSolver::Solve(cudaStream_t stream, // DenseLDLTSolver private helpers // --------------------------------------------------------------------------- -void DenseLDLTSolver::EnsureBuffersSize(size_t n) { +void DenseLDLTSolver::EnsureBuffersSize(cudaStream_t /*stream*/, size_t n) { const size_t matrix_elements = n * n; - if (dense_matrix_.size() != matrix_elements) { - dense_matrix_.resize(matrix_elements); - } if (ldlt_factor_.size() != matrix_elements) { ldlt_factor_.resize(matrix_elements); } @@ -544,24 +463,4 @@ void DenseLDLTSolver::EnsureBuffersSize(size_t n) { } } -void DenseLDLTSolver::ConvertCSRToDense(cudaStream_t stream, - const CSRSparseMatrix &matrix, - dvector &dense_matrix) { - const int num_rows = static_cast(matrix.NumRows()); - if (num_rows == 0) { - return; - } - THROW_ON_CUDA_ERROR(cudaMemsetAsync( - dense_matrix.data(), 0, - static_cast(num_rows) * num_rows * sizeof(float), stream)); - - constexpr int kThreads = 256; - constexpr int kWarpsPerBlock = kThreads / kWarpSize; - const int blocks = (num_rows + kWarpsPerBlock - 1) / kWarpsPerBlock; - csr_to_dense_kernel<<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_rows, dense_matrix.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/dense_linear_solver.h b/cunls/linear_solver/dense_linear_solver.h index 07cf63d..1637ecf 100644 --- a/cunls/linear_solver/dense_linear_solver.h +++ b/cunls/linear_solver/dense_linear_solver.h @@ -20,15 +20,15 @@ #include #include "cunls/common/types.h" -#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/linear_solver/dense_linear_solver_base.h" namespace cunls { /** * @brief Dense GPU linear solver based on pivoted LDLT factorization. * - * Converts the input CSR symmetric matrix to a dense row-major matrix and - * solves A x = b via: + * Densifies the symmetric coefficient matrix (from either sparse layout; see + * DenseLinearSolverBase) and solves A x = b via: * 1) Symmetric pivoted LDLT factorization: P^T A P = L D L^T * 2) Triangular/diagonal solves in the permuted system * 3) Permutation back to the original variable ordering @@ -40,85 +40,42 @@ namespace cunls { * synchronization. This avoids per-kernel sync and gives the caller an * accurate bool return from Solve(). */ -class DenseLDLTSolver : public CSRSparseLinearSolver { - public: - // This backend consumes CSR only; SupportsBlockStorage() stays false, so the - // base class's block-storage overloads are never called on it. The - // using-declarations keep them visible rather than hidden by the CSR - // overrides below. - using CSRSparseLinearSolver::Initialize; - using CSRSparseLinearSolver::Solve; - +class DenseLDLTSolver : public DenseLinearSolverBase { + protected: /** - * @brief Allocates internal dense buffers for the given matrix size. - * - * Validates that dimensions of @p spd_matrix, @p rhs, and @p result are - * consistent, then pre-allocates all working buffers (dense matrix, LDLT - * factors, permutation, scratch vectors, and status buffers) so that - * subsequent Solve() calls do not allocate. + * @copydoc DenseLinearSolverBase::EnsureBuffersSize * - * @param stream CUDA stream (unused; buffers are allocated synchronously). - * @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). - * @return true on success, false if a dimension mismatch is detected. + * Pre-allocates the LDLT factors, permutation, scratch vectors and status + * buffers so that subsequent Solve() calls do not allocate. */ - bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) final; + void EnsureBuffersSize(cudaStream_t stream, size_t n) final; /** - * @brief Converts CSR to dense and solves via pivoted LDLT factorization. + * @brief Factorizes the dense matrix via pivoted LDLT and solves. * * Handles symmetric matrices including indefinite ones (not limited to SPD). * * The pipeline is: - * 1. CSR -> dense conversion (one kernel). - * 2. Pivoted LDLT factorization kernel -> writes status_[0]. - * 3. Triangular/diagonal solve kernel -> writes status_[1]. - * 4. Single async copy of status_[0..1] to status_pinned_[0..1]. - * 5. Stream synchronization. - * 6. Host-side check of both status flags. + * 1. Pivoted LDLT factorization kernel -> writes status_[0]. + * 2. Triangular/diagonal solve kernel -> writes status_[1]. + * 3. Single async copy of status_[0..1] to status_pinned_[0..1]. + * 4. Stream synchronization. + * 5. Host-side check of both status flags. * * If the factorization encounters a (near-)singular pivot, status_[0] is * set to 0 and the function returns false. If the solve encounters a zero * diagonal element, status_[1] is set to 0 and the function returns false. * * @param stream CUDA stream for asynchronous GPU operations. - * @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). - * @return true on success, false on dimension mismatch, singular pivot, - * or zero diagonal. + * @param n Matrix dimension. + * @param rhs The right-hand side vector b. + * @param result Output vector x. + * @return true on success, false on a singular pivot or zero diagonal. */ - bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, - dvector &result) final; + bool FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) final; private: - /** - * @brief Ensures all internal buffers are (re-)allocated for an n x n system. - * - * Called from both Initialize() and Solve(). Each buffer is resized only - * when its current size differs from the required size, so repeated calls - * with the same n are essentially free. - * - * @param n Number of rows (and columns) of the dense system. - */ - void EnsureBuffersSize(size_t n); - - /** - * @brief Converts a CSR matrix to dense row-major format on the GPU. - * - * Zeroes the dense output, then scatters CSR values into the correct - * positions using a warp-per-row kernel. - * - * @param stream CUDA stream for the memset and kernel launch. - * @param matrix Input CSR matrix. - * @param dense_matrix Output dense buffer (must be pre-allocated to n*n). - */ - void ConvertCSRToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, - dvector &dense_matrix); - - dvector dense_matrix_; ///< Dense row-major copy of A. dvector ldlt_factor_; ///< In-place LDLT factor storage. dvector permutation_; ///< Pivot permutation vector. dvector permuted_rhs_; ///< P * b scratch vector. diff --git a/cunls/linear_solver/dense_linear_solver_base.cu b/cunls/linear_solver/dense_linear_solver_base.cu new file mode 100644 index 0000000..5345a5b --- /dev/null +++ b/cunls/linear_solver/dense_linear_solver_base.cu @@ -0,0 +1,188 @@ +/* + * 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/common/helper.h" +#include "cunls/common/log.h" +#include "cunls/linear_solver/dense_linear_solver_base.h" + +namespace cunls { +namespace { + +constexpr int kWarpSize = 32; +constexpr int kScatterThreads = 256; + +/// @brief Scatters CSR values into a dense row-major matrix. +/// +/// One warp per row; the warp's lanes stride over the row's non-zeros. The +/// caller zeroes @p dense_matrix first, so entries not stored stay zero. +__global__ void ScatterCSRToDenseKernel(const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, int num_rows, + float *__restrict__ dense_matrix) { + const int row = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpSize; + if (row >= num_rows) { + return; + } + const int lane = threadIdx.x % kWarpSize; + const int row_start = row_offsets[row]; + const int row_end = row_offsets[row + 1]; + float *const dense_row = dense_matrix + static_cast(row) * num_rows; + for (int idx = row_start + lane; idx < row_end; idx += kWarpSize) { + dense_row[col_ids[idx]] = values[idx]; + } +} + +/// @brief Scatters BSR tiles into a dense row-major matrix. +/// +/// One CUDA block per block row. The row's tiles and their entries are walked +/// as one flat range so a short row cannot leave most of the block idle, which +/// matters for bundle-adjustment Hessians where tile counts per row differ by +/// orders of magnitude. Within a tile the flat index runs fastest over the +/// tile's columns, so consecutive threads write consecutive dense columns. +__global__ void ScatterBSRToDenseKernel(const int *__restrict__ row_offsets, + const int *__restrict__ col_ids, + const float *__restrict__ values, int block_size, + int num_rows, float *__restrict__ dense_matrix) { + const int block_row = blockIdx.x; + const int row_begin = row_offsets[block_row]; + const int row_end = row_offsets[block_row + 1]; + const int tile_area = block_size * block_size; + const int total = (row_end - row_begin) * tile_area; + + for (int idx = threadIdx.x; idx < total; idx += blockDim.x) { + const int tile = row_begin + idx / tile_area; + const int entry = idx % tile_area; + const int k = entry / block_size; + const int l = entry - k * block_size; + const size_t dense_row = static_cast(block_row) * block_size + k; + const int dense_col = col_ids[tile] * block_size + l; + dense_matrix[dense_row * num_rows + dense_col] = + values[static_cast(tile) * tile_area + entry]; + } +} + +/** + * @brief Scatters a CSR matrix into a zeroed dense row-major buffer. + * + * @param stream CUDA stream for the memset and kernel launch. + * @param matrix Input CSR matrix; must be square and stored in full (both + * triangles), which is what HessianStructureBuilder produces. + * @param[out] dense_matrix Destination, pre-allocated to `NumRows()^2`. + */ +void ScatterToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, + dvector &dense_matrix) { + const int num_rows = static_cast(matrix.NumRows()); + if (num_rows == 0) { + return; + } + THROW_ON_CUDA_ERROR(cudaMemsetAsync( + dense_matrix.data(), 0, static_cast(num_rows) * num_rows * sizeof(float), stream)); + + constexpr int kWarpsPerBlock = kScatterThreads / kWarpSize; + const int blocks = (num_rows + kWarpsPerBlock - 1) / kWarpsPerBlock; + ScatterCSRToDenseKernel<<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), num_rows, + dense_matrix.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +/** + * @brief Scatters a BSR matrix into a zeroed dense row-major buffer. + * + * One CUDA block per block row walks that row's tiles and writes their + * `block_size^2` entries out. Consecutive threads cover consecutive columns + * within a tile row, so the dense writes coalesce. + * + * @param stream CUDA stream for the memset and kernel launch. + * @param matrix Input BSR matrix; must be square and stored in full. + * @param[out] dense_matrix Destination, pre-allocated to `NumRows()^2`. + */ +void ScatterToDense(cudaStream_t stream, const BSRSparseMatrix &matrix, + dvector &dense_matrix) { + const int num_rows = matrix.NumRows(); + if (num_rows == 0 || matrix.num_block_rows == 0) { + return; + } + THROW_ON_CUDA_ERROR(cudaMemsetAsync( + dense_matrix.data(), 0, static_cast(num_rows) * num_rows * sizeof(float), stream)); + + ScatterBSRToDenseKernel<<>>( + matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), matrix.block_size, + num_rows, dense_matrix.data()); + THROW_ON_CUDA_ERROR(cudaGetLastError()); +} + +} // namespace + +bool DenseLinearSolverBase::InitializeCommon(cudaStream_t stream, size_t num_rows, + const dvector &rhs, + const dvector &result) { + if (num_rows != rhs.size()) { + LogError("LHS size: {} does not match RHS size: {}", num_rows, rhs.size()); + return false; + } + if (num_rows != result.size()) { + LogError("LHS size: {} does not match result size: {}", num_rows, result.size()); + return false; + } + const size_t matrix_elements = num_rows * num_rows; + if (dense_matrix_.size() != matrix_elements) { + dense_matrix_.resize(matrix_elements); + } + EnsureBuffersSize(stream, num_rows); + return true; +} + +bool DenseLinearSolverBase::Initialize(cudaStream_t stream, const Problem & /*problem*/, + const CSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) { + return InitializeCommon(stream, spd_matrix.NumRows(), rhs, result); +} + +bool DenseLinearSolverBase::Initialize(cudaStream_t stream, const Problem & /*problem*/, + const BSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) { + return InitializeCommon(stream, static_cast(spd_matrix.NumRows()), rhs, result); +} + +bool DenseLinearSolverBase::Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) { + const size_t num_rows = spd_matrix.NumRows(); + if (!InitializeCommon(stream, num_rows, rhs, result)) { + return false; + } + if (num_rows == 0) { + return true; // vacuously solved; nothing to launch + } + ScatterToDense(stream, spd_matrix, dense_matrix_); + return FactorizeAndSolve(stream, static_cast(num_rows), rhs, result); +} + +bool DenseLinearSolverBase::Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) { + const size_t num_rows = static_cast(spd_matrix.NumRows()); + if (!InitializeCommon(stream, num_rows, rhs, result)) { + return false; + } + if (num_rows == 0) { + return true; // vacuously solved; nothing to launch + } + ScatterToDense(stream, spd_matrix, dense_matrix_); + return FactorizeAndSolve(stream, static_cast(num_rows), rhs, result); +} + +} // namespace cunls diff --git a/cunls/linear_solver/dense_linear_solver_base.h b/cunls/linear_solver/dense_linear_solver_base.h new file mode 100644 index 0000000..7fa886e --- /dev/null +++ b/cunls/linear_solver/dense_linear_solver_base.h @@ -0,0 +1,132 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. + * All rights reserved. SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include "cunls/common/types.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" + +namespace cunls { + +/** + * @brief Shared machinery for the backends that densify before factorizing. + * + * DenseLDLTSolver, DenseCholeskySolver and DenseQRSolver all begin by + * scattering the sparse coefficient matrix into an `n x n` dense buffer and + * then never look at the sparse form again. The sparse layout is therefore + * invisible to them past the first kernel, which is why they accept *both* + * scalar CSR and block BSR: scattering tiles is if anything simpler than + * scattering CSR rows, since a tile carries its own `block_size x block_size` + * geometry and needs no per-row offset indirection. + * + * Accepting both matters beyond tidiness. The Hessian layout is chosen for + * the problem, and a backend that declines block storage forces the whole + * assembly onto scalar CSR (see NormalEquations::Initialize) — giving up the + * block layout's smaller index array and the `BuildSystem` time that goes with + * it. A solver that discards the layout one kernel later has no reason to + * impose that cost. After this class, cuDSS is the only backend that does. + * + * The dense buffer holds the matrix row-major. cuSOLVER and cuBLAS want + * column-major, but every matrix reaching these solvers is a Gauss-Newton + * Hessian and therefore symmetric, so the two layouts coincide. Do not reuse + * this class for a non-symmetric operator without transposing. + * + * Subclasses supply three things: workspace sizing, the factorize-and-solve + * body, and nothing else. Dimension validation, the dense scatter and the + * layout dispatch all live here. + */ +class DenseLinearSolverBase : public SparseLinearSolver { + public: + /** + * @copydoc SparseLinearSolver::SupportsBlockStorage + * + * Always true: the sparse layout does not survive the scatter into the dense + * buffer, so there is nothing for either layout to be better at. + */ + bool SupportsBlockStorage() const final { return true; } + + /** @copydoc SparseLinearSolver::Initialize */ + bool Initialize(cudaStream_t stream, const Problem &problem, const CSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) final; + + /** @copydoc SparseLinearSolver::Initialize */ + bool Initialize(cudaStream_t stream, const Problem &problem, const BSRSparseMatrix &spd_matrix, + const dvector &rhs, dvector &result) final; + + /** @copydoc SparseLinearSolver::Solve */ + bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) final; + + /** @copydoc SparseLinearSolver::Solve */ + bool Solve(cudaStream_t stream, const BSRSparseMatrix &spd_matrix, const dvector &rhs, + dvector &result) final; + + /** @copydoc SparseLinearSolver::DisableSafetyChecks */ + void DisableSafetyChecks() final { safety_checks_enabled_ = false; } + + protected: + /** + * @brief Allocates the backend's working buffers for an `n x n` system. + * + * Called from both Initialize and Solve. Implementations should resize only + * when the requested size differs, so repeated calls at a fixed `n` are free. + * @ref dense_matrix_ is sized by this class before the call, so subclasses + * only need their own buffers. + * + * @param stream CUDA stream, for backends that query a library workspace size. + * @param n Number of rows (and columns) of the dense system. + */ + virtual void EnsureBuffersSize(cudaStream_t stream, size_t n) = 0; + + /** + * @brief Factorizes @ref dense_matrix_ and solves against it. + * + * Called with @ref dense_matrix_ already populated with the full symmetric + * matrix in row-major order and all buffers sized for `n`. + * + * @param stream CUDA stream for asynchronous GPU operations. + * @param n Matrix dimension; guaranteed positive. + * @param rhs Right-hand side `b`. + * @param result Output vector `x`, caller-allocated. + * @return true on success, false on a detected numerical failure. + */ + virtual bool FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) = 0; + + /** @brief Dense row-major copy of the coefficient matrix; `n * n` floats. */ + dvector dense_matrix_; + + /** @brief Whether to run post-factorization numerical checks. */ + bool safety_checks_enabled_ = true; + + private: + /** + * @brief Common Initialize body: validate, size buffers. + * + * @param stream CUDA stream for asynchronous GPU operations. + * @param num_rows Scalar row count of the coefficient matrix. + * @param rhs Right-hand side, checked against @p num_rows. + * @param result Output vector, checked against @p num_rows. + * @return true on success, false on dimension mismatch. + */ + bool InitializeCommon(cudaStream_t stream, size_t num_rows, const dvector &rhs, + const dvector &result); +}; + +} // namespace cunls diff --git a/cunls/linear_solver/dense_qr_solver.cu b/cunls/linear_solver/dense_qr_solver.cu index ae5ff2d..066a4fe 100644 --- a/cunls/linear_solver/dense_qr_solver.cu +++ b/cunls/linear_solver/dense_qr_solver.cu @@ -25,15 +25,13 @@ namespace cunls { namespace { -constexpr int kWarpSize = 32; constexpr float kDiagonalEpsilonAbs = 1e-7f; /// Checks whether any diagonal element of an n x n column-major matrix has /// absolute value below the threshold. Sets *status = 0 if any diagonal is /// near-zero, 1 otherwise. -__global__ void check_diagonal_kernel(const float *__restrict__ matrix, int n, - int lda, float threshold, - int *__restrict__ status) { +__global__ void check_diagonal_kernel(const float *__restrict__ matrix, int n, int lda, + float threshold, int *__restrict__ status) { *status = 1; for (int i = threadIdx.x; i < n; i += blockDim.x) { if (fabsf(matrix[i * lda + i]) <= threshold) { @@ -42,132 +40,62 @@ __global__ void check_diagonal_kernel(const float *__restrict__ matrix, int n, } } -// Writes the CSR matrix into a dense buffer in row-major order. cuSOLVER -// (geqrf, ormqr) and cuBLAS (trsm) expect column-major layout, but this -// solver is only used for symmetric matrices (the Gauss-Newton Hessian -// J^T J), for which row-major == column-major (A == A^T). Do NOT use this -// kernel for non-symmetric inputs without first transposing the layout. -__global__ void csr_to_dense_kernel(const int *__restrict__ row_offsets, - const int *__restrict__ col_ids, - const float *__restrict__ values, - int num_rows, - float *__restrict__ dense_matrix) { - const int row = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpSize; - if (row >= num_rows) { - return; - } - const int lane = threadIdx.x % kWarpSize; - const int row_start = row_offsets[row]; - const int row_end = row_offsets[row + 1]; - float *const dense_row = dense_matrix + row * num_rows; - for (int idx = row_start + lane; idx < row_end; idx += kWarpSize) { - dense_row[col_ids[idx]] = values[idx]; - } -} +} // namespace -} // namespace - -bool DenseQRSolver::Initialize(cudaStream_t stream, const Problem & /*problem*/, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, - dvector &result) { - const size_t matrix_size = spd_matrix.NumRows(); - if (matrix_size != rhs.size()) { - LogError("LHS size: {} does not match RHS size: {}", matrix_size, - rhs.size()); - return false; - } - if (matrix_size != result.size()) { - LogError("LHS size: {} does not match result size: {}", matrix_size, - result.size()); - return false; - } - EnsureBuffersSize(stream, matrix_size); - return true; -} - -bool DenseQRSolver::Solve(cudaStream_t stream, - const CSRSparseMatrix &spd_matrix, - const dvector &rhs, dvector &result) { - const size_t matrix_size = spd_matrix.NumRows(); - if (matrix_size != rhs.size()) { - LogError("LHS size: {} does not match RHS size: {}", matrix_size, - rhs.size()); - return false; - } - if (matrix_size != result.size()) { - LogError("LHS size: {} does not match result size: {}", matrix_size, - result.size()); - return false; - } - if (matrix_size == 0) { - return true; - } - - EnsureBuffersSize(stream, matrix_size); - ConvertCSRToDense(stream, spd_matrix, dense_matrix_); - - const int n = static_cast(matrix_size); - auto cusolver = - static_cast(cusolver_handle_.GetHandle(stream)); +bool DenseQRSolver::FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) { + auto cusolver = static_cast(cusolver_handle_.GetHandle(stream)); // 1. QR factorization: A = Q R (in-place, R in upper triangle, Householder // reflectors stored below the diagonal, scalars in tau). - THROW_ON_CUSOLVER_ERROR(cusolverDnSgeqrf( - cusolver, n, n, dense_matrix_.data(), n, tau_.data(), workspace_.data(), - static_cast(workspace_.size()), dev_info_.data())); + THROW_ON_CUSOLVER_ERROR(cusolverDnSgeqrf(cusolver, n, n, dense_matrix_.data(), n, tau_.data(), + workspace_.data(), static_cast(workspace_.size()), + dev_info_.data())); if (safety_checks_enabled_) { // 2. Check R's diagonal for rank deficiency. geqrf always sets devInfo = 0 // on success, so we must explicitly inspect the diagonal of R. - check_diagonal_kernel<<<1, 256, 0, stream>>>( - dense_matrix_.data(), n, n, kDiagonalEpsilonAbs, dev_info_.data()); + check_diagonal_kernel<<<1, 256, 0, stream>>>(dense_matrix_.data(), n, n, kDiagonalEpsilonAbs, + dev_info_.data()); THROW_ON_CUDA_ERROR(cudaGetLastError()); - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dev_info_pinned_.data(), - dev_info_.data(), sizeof(int), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(dev_info_pinned_.data(), dev_info_.data(), sizeof(int), cudaMemcpyDeviceToHost, stream)); THROW_ON_CUDA_ERROR(cudaStreamSynchronize(stream)); if (dev_info_pinned_[0] == 0) { - LogError("QR factorization detected a (near-)zero diagonal in R. " - "Matrix is rank-deficient."); + LogError( + "QR factorization detected a (near-)zero diagonal in R. " + "Matrix is rank-deficient."); return false; } } // 3. Copy rhs into work vector (ormqr and trsm operate in-place on it). - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(rhs_copy_.data(), rhs.data(), - n * sizeof(float), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(rhs_copy_.data(), rhs.data(), n * sizeof(float), cudaMemcpyDeviceToDevice, stream)); // 4. Apply Q^T to the rhs: y = Q^T * b. - THROW_ON_CUSOLVER_ERROR(cusolverDnSormqr( - cusolver, CUBLAS_SIDE_LEFT, CUBLAS_OP_T, n, 1, n, dense_matrix_.data(), n, - tau_.data(), rhs_copy_.data(), n, workspace_.data(), - static_cast(workspace_.size()), dev_info_.data())); + THROW_ON_CUSOLVER_ERROR(cusolverDnSormqr(cusolver, CUBLAS_SIDE_LEFT, CUBLAS_OP_T, n, 1, n, + dense_matrix_.data(), n, tau_.data(), rhs_copy_.data(), + n, workspace_.data(), + static_cast(workspace_.size()), dev_info_.data())); // 5. Solve the upper-triangular system R x = y via cuBLAS trsm. auto cublas = static_cast(cublas_handle_.GetHandle(stream)); const float alpha = 1.0f; - THROW_ON_CUBLAS_ERROR( - cublasStrsm(cublas, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, - CUBLAS_DIAG_NON_UNIT, n, 1, &alpha, dense_matrix_.data(), n, - rhs_copy_.data(), n)); + THROW_ON_CUBLAS_ERROR(cublasStrsm(cublas, CUBLAS_SIDE_LEFT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, + CUBLAS_DIAG_NON_UNIT, n, 1, &alpha, dense_matrix_.data(), n, + rhs_copy_.data(), n)); // 6. Copy solution out. - THROW_ON_CUDA_ERROR(cudaMemcpyAsync(result.data(), rhs_copy_.data(), - n * sizeof(float), + THROW_ON_CUDA_ERROR(cudaMemcpyAsync(result.data(), rhs_copy_.data(), n * sizeof(float), cudaMemcpyDeviceToDevice, stream)); return true; } void DenseQRSolver::EnsureBuffersSize(cudaStream_t stream, size_t n) { - const size_t matrix_elements = n * n; - if (dense_matrix_.size() != matrix_elements) { - dense_matrix_.resize(matrix_elements); - } if (tau_.size() != n) { tau_.resize(n); } @@ -180,18 +108,17 @@ void DenseQRSolver::EnsureBuffersSize(cudaStream_t stream, size_t n) { } if (n != last_n_ && n > 0) { - auto handle = - static_cast(cusolver_handle_.GetHandle(stream)); + auto handle = static_cast(cusolver_handle_.GetHandle(stream)); int geqrf_lwork = 0; - THROW_ON_CUSOLVER_ERROR(cusolverDnSgeqrf_bufferSize( - handle, static_cast(n), static_cast(n), dense_matrix_.data(), - static_cast(n), &geqrf_lwork)); + THROW_ON_CUSOLVER_ERROR(cusolverDnSgeqrf_bufferSize(handle, static_cast(n), + static_cast(n), dense_matrix_.data(), + static_cast(n), &geqrf_lwork)); int ormqr_lwork = 0; THROW_ON_CUSOLVER_ERROR(cusolverDnSormqr_bufferSize( - handle, CUBLAS_SIDE_LEFT, CUBLAS_OP_T, static_cast(n), 1, - static_cast(n), dense_matrix_.data(), static_cast(n), - tau_.data(), rhs_copy_.data(), static_cast(n), &ormqr_lwork)); + handle, CUBLAS_SIDE_LEFT, CUBLAS_OP_T, static_cast(n), 1, static_cast(n), + dense_matrix_.data(), static_cast(n), tau_.data(), rhs_copy_.data(), + static_cast(n), &ormqr_lwork)); int lwork = (geqrf_lwork > ormqr_lwork) ? geqrf_lwork : ormqr_lwork; workspace_.resize(static_cast(lwork)); @@ -199,24 +126,4 @@ void DenseQRSolver::EnsureBuffersSize(cudaStream_t stream, size_t n) { } } -void DenseQRSolver::ConvertCSRToDense(cudaStream_t stream, - const CSRSparseMatrix &matrix, - dvector &dense_matrix) { - const int num_rows = static_cast(matrix.NumRows()); - if (num_rows == 0) { - return; - } - THROW_ON_CUDA_ERROR(cudaMemsetAsync( - dense_matrix.data(), 0, - static_cast(num_rows) * num_rows * sizeof(float), stream)); - - constexpr int kThreads = 256; - constexpr int kWarpsPerBlock = kThreads / kWarpSize; - const int blocks = (num_rows + kWarpsPerBlock - 1) / kWarpsPerBlock; - csr_to_dense_kernel<<>>( - matrix.row_offsets.data(), matrix.col_ids.data(), matrix.values.data(), - num_rows, dense_matrix.data()); - THROW_ON_CUDA_ERROR(cudaGetLastError()); -} - -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/dense_qr_solver.h b/cunls/linear_solver/dense_qr_solver.h index 661efb7..c8bdd93 100644 --- a/cunls/linear_solver/dense_qr_solver.h +++ b/cunls/linear_solver/dense_qr_solver.h @@ -22,77 +22,55 @@ #include "cunls/common/cublas_helper.h" #include "cunls/common/cusolver_helper.h" #include "cunls/common/types.h" -#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/linear_solver/dense_linear_solver_base.h" namespace cunls { /** * @brief Dense GPU linear solver based on QR factorization via cuSOLVER. * - * Converts the input CSR matrix to a dense matrix and solves A x = b via: + * Densifies the coefficient matrix (from either sparse layout; see + * DenseLinearSolverBase) and solves A x = b via: * 1) QR factorization: A = Q R (cusolverDnSgeqrf) * 2) Apply Q^T to rhs: y = Q^T b (cusolverDnSormqr) * 3) Triangular solve: R x = y (cublasStrsm) * * Since the input matrix is symmetric, the row-major dense representation - * produced by CSR conversion is identical to column-major, so no transpose - * is required for the column-major cuSOLVER/cuBLAS APIs. + * produced by the scatter is identical to column-major, so no transpose is + * required for the column-major cuSOLVER/cuBLAS APIs. * * QR factorization works for any non-singular square matrix (not limited * to SPD). Returns false from Solve() if the factorization reports an error * via devInfo. */ -class DenseQRSolver : public CSRSparseLinearSolver { - public: - // This backend consumes CSR only; SupportsBlockStorage() stays false, so the - // base class's block-storage overloads are never called on it. The - // using-declarations keep them visible rather than hidden by the CSR - // overrides below. - using CSRSparseLinearSolver::Initialize; - using CSRSparseLinearSolver::Solve; +class DenseQRSolver : public DenseLinearSolverBase { + protected: + /** @copydoc DenseLinearSolverBase::EnsureBuffersSize */ + void EnsureBuffersSize(cudaStream_t stream, size_t n) final; /** - * @brief Validates dimensions and pre-allocates internal buffers. - * - * @param stream CUDA stream used to query the cuSOLVER workspace size. - * @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). - * @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; - - /** - * @brief Converts CSR to dense and solves via QR factorization. + * @brief Factorizes the dense matrix via QR and solves. * * The pipeline is: - * 1. CSR -> dense conversion. - * 2. cusolverDnSgeqrf (in-place QR factorization). + * 1. cusolverDnSgeqrf (in-place QR factorization). + * 2. Rank-deficiency check on R's diagonal (if safety checks enabled). * 3. Copy rhs into a work vector. * 4. cusolverDnSormqr (Q^T * b in-place). * 5. cublasStrsm (R x = Q^T b upper-triangular solve). * 6. Copy result from work vector. - * 7. Async copy of devInfo to pinned host, stream sync, host check. * * @param stream CUDA stream for asynchronous GPU operations. - * @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). - * @return true on success, false on dimension mismatch or singular matrix. + * @param n Matrix dimension. + * @param rhs The right-hand side vector b. + * @param result Output vector x. + * @return true on success, false on a singular or rank-deficient matrix. */ - bool Solve(cudaStream_t stream, const CSRSparseMatrix &spd_matrix, const dvector &rhs, - dvector &result) final; + bool FactorizeAndSolve(cudaStream_t stream, int n, const dvector &rhs, + dvector &result) final; private: - void EnsureBuffersSize(cudaStream_t stream, size_t n); - - void ConvertCSRToDense(cudaStream_t stream, const CSRSparseMatrix &matrix, - dvector &dense_matrix); - cuSolverHandle cusolver_handle_; cuBLASHandle cublas_handle_; - dvector dense_matrix_; dvector tau_; dvector workspace_; dvector rhs_copy_; diff --git a/cunls/linear_solver/llms.txt b/cunls/linear_solver/llms.txt index ee1abfe..649019d 100644 --- a/cunls/linear_solver/llms.txt +++ b/cunls/linear_solver/llms.txt @@ -6,14 +6,27 @@ Sparse linear system abstraction and implementation used by minimizers. ## Key types -- `csr_sparse_linear_solver.h`: abstract `CSRSparseLinearSolver` +- `sparse_linear_solver_base.h`: abstract `SparseLinearSolver` — CSR + `Initialize`/`Solve` are pure virtual; BSR counterparts are defaulted and + guarded by `SupportsBlockStorage()`. - `sparse_linear_solver.h`: - `SparseLinearSolverType` - `SparseLinearSolverConfig` - - `CreateCSRSparseLinearSolver(...)` + - `CreateSparseLinearSolver(...)` +- `dense_linear_solver_base.h`: + - `DenseLinearSolverBase` — shared machinery for the backends that densify + before factorizing. Owns `dense_matrix_`, dimension validation, the + CSR/BSR scatter and the layout dispatch; subclasses supply only + `EnsureBuffersSize` and `FactorizeAndSolve`. `SupportsBlockStorage()` is + `true`: the sparse layout does not survive the scatter, so neither layout + can be better for them. +- `dense_linear_solver.h`: `DenseLDLTSolver` — custom pivoted-LDLT kernels; + handles indefinite symmetric matrices. +- `dense_cholesky_solver.h`: `DenseCholeskySolver` — cuSOLVER potrf/potrs. +- `dense_qr_solver.h`: `DenseQRSolver` — cuSOLVER geqrf/ormqr + cuBLAS trsm. - `cudss_sparse_linear_solver.h`: - `cuDSSLinearSolverOptions` - - `cuDSSLinearSolver` + - `cuDSSLinearSolver` — the only backend that declines block storage. - `block_sparse_pcg_solver.h`: - `BlockSparsePCGOptions` — `block_size` / `block_layout`, `max_iterations`, `relative_tolerance`, `absolute_tolerance`, @@ -28,8 +41,10 @@ Sparse linear system abstraction and implementation used by minimizers. ## Expected system form -- Solves SPD systems in CSR format (`CSRSparseMatrix`), typically normal - equations assembled by `GaussNewtonMinimizer` / `LevenbergMarquardtMinimizer`. +- Solves SPD systems given as `CSRSparseMatrix` or, for backends that report + `SupportsBlockStorage()`, `BSRSparseMatrix` — typically normal equations + assembled by `GaussNewtonMinimizer` / `LevenbergMarquardtMinimizer`. The + layout is chosen in `NormalEquations`; no conversion runs on the solve path. ## Performance modes diff --git a/cunls/linear_solver/sparse_linear_solver.cpp b/cunls/linear_solver/sparse_linear_solver.cpp index b4277ec..8c3d01c 100644 --- a/cunls/linear_solver/sparse_linear_solver.cpp +++ b/cunls/linear_solver/sparse_linear_solver.cpp @@ -16,31 +16,30 @@ */ #include "cunls/linear_solver/sparse_linear_solver.h" + #include "cunls/common/cudss_helper.h" namespace cunls { -/** @copydoc CreateCSRSparseLinearSolver +/** @copydoc CreateSparseLinearSolver * @throws std::invalid_argument If an unsupported solver type is specified. */ -SparseLinearSolverPtr -CreateCSRSparseLinearSolver(SparseLinearSolverType type, - const SparseLinearSolverConfig &config) { +SparseLinearSolverPtr CreateSparseLinearSolver(SparseLinearSolverType type, + const SparseLinearSolverConfig &config) { switch (type) { - case SparseLinearSolverType::cuDSS: - return std::make_unique(config.cudss_solver_options); - case SparseLinearSolverType::DenseLDLT: - return std::make_unique(); - case SparseLinearSolverType::DenseCholesky: - 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"); + case SparseLinearSolverType::cuDSS: + return std::make_unique(config.cudss_solver_options); + case SparseLinearSolverType::DenseLDLT: + return std::make_unique(); + case SparseLinearSolverType::DenseCholesky: + 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"); } } -} // namespace cunls \ No newline at end of file +} // namespace cunls \ No newline at end of file diff --git a/cunls/linear_solver/sparse_linear_solver.h b/cunls/linear_solver/sparse_linear_solver.h index 11bd931..007b854 100644 --- a/cunls/linear_solver/sparse_linear_solver.h +++ b/cunls/linear_solver/sparse_linear_solver.h @@ -22,11 +22,11 @@ #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" #include "cunls/linear_solver/dense_linear_solver.h" #include "cunls/linear_solver/dense_qr_solver.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" namespace cunls { @@ -34,20 +34,21 @@ 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. - 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. + cuDSS, ///< Sparse direct solver using NVIDIA's cuDSS library. The + ///< only backend that cannot consume block storage. + DenseLDLT, ///< Densifies the coefficient matrix and solves with a custom + ///< CUDA pivoted LDLT kernel. + DenseCholesky, ///< Densifies the coefficient matrix and solves with cuSOLVER + ///< Cholesky factorization (cusolverDnSpotrf / + ///< cusolverDnSpotrs). Requires SPD matrix. + DenseQR, ///< Densifies the coefficient matrix 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. }; /** @@ -66,7 +67,7 @@ struct SparseLinearSolverConfig { /** * @brief Smart pointer type for sparse linear solvers. */ -using SparseLinearSolverPtr = std::unique_ptr; +using SparseLinearSolverPtr = std::unique_ptr; /** * @brief Factory function to create a sparse linear solver. @@ -78,8 +79,7 @@ using SparseLinearSolverPtr = std::unique_ptr; * @param config Solver-specific configuration options. * @return A unique pointer to the created solver instance. */ -SparseLinearSolverPtr -CreateCSRSparseLinearSolver(SparseLinearSolverType type, - const SparseLinearSolverConfig &config); +SparseLinearSolverPtr CreateSparseLinearSolver(SparseLinearSolverType type, + const SparseLinearSolverConfig &config); -} // namespace cunls +} // namespace cunls diff --git a/cunls/linear_solver/csr_sparse_linear_solver.h b/cunls/linear_solver/sparse_linear_solver_base.h similarity index 65% rename from cunls/linear_solver/csr_sparse_linear_solver.h rename to cunls/linear_solver/sparse_linear_solver_base.h index a47fb7d..842ede1 100644 --- a/cunls/linear_solver/csr_sparse_linear_solver.h +++ b/cunls/linear_solver/sparse_linear_solver_base.h @@ -26,12 +26,18 @@ namespace cunls { class Problem; // forward declaration; defined in cunls/minimizer/problem.h. /** - * @brief Base class for linear solvers operating on CSR matrices. + * @brief Base class for solvers of the sparse symmetric system `A x = b`. * - * 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, block-Jacobi PCG). + * Derived classes implement specific strategies (cuDSS direct factorization, + * dense pivoted LDLT / Cholesky / QR, block-Jacobi PCG). + * + * `A` arrives in one of two layouts, and every backend must accept at least + * scalar CSR: Initialize and Solve are pure virtual for @ref CSRSparseMatrix + * and defaulted for @ref BSRSparseMatrix. A backend opts into the block form + * by overriding @ref SupportsBlockStorage together with the two BSR overloads; + * one that does not is never handed a BSR matrix, so its defaults are dead. + * The layout is decided once per problem in @ref NormalEquations, and no + * conversion between the two ever runs on the solve path. * * Initialize receives the originating @ref Problem so solvers can adapt * to its block / factor-graph structure (e.g. @@ -39,7 +45,7 @@ class Problem; // forward declaration; defined in cunls/minimizer/problem.h. * build its block-Jacobi preconditioner without a downcast at the call * site). Solvers that don't care can simply ignore the argument. */ -class CSRSparseLinearSolver { +class SparseLinearSolver { public: /** * @brief Performs setup work for the linear system. @@ -90,8 +96,20 @@ class CSRSparseLinearSolver { * * The Hessian of a factor graph is naturally block structured, and assembling * it that way keeps one column index per tile instead of one per scalar - * entry. Backends that say yes get the block form; the rest are handed an - * expanded CSR copy, so no caller has to care. + * entry. Backends that say yes get the block form. + * + * Backends that say no are not handed a converted copy — no BSR-to-CSR + * expansion runs anywhere on the solve path. Instead the Hessian is + * assembled *natively* in scalar CSR for them (see NormalEquations), which is + * their optimum: the block layout's saving is the index array, and + * materializing scalar indices for a CSR-only backend would give that saving + * straight back plus a per-iteration value permutation. + * + * The layout is therefore a property of the problem, and this method is a + * veto, not a request. Vetoing costs the block-storage delta only — the + * smaller index array and the assembly time that comes with it — never the + * much larger block-wise *assembly* win, which is layout-independent and + * which every backend gets unconditionally. */ virtual bool SupportsBlockStorage() const { return false; } @@ -109,25 +127,22 @@ class CSRSparseLinearSolver { } /** - * @brief Disables post-factorization safety checks. + * @brief Advisory request to skip post-factorization safety checks. * - * By default, dense solvers copy a device-side status flag back to the - * host after factorization and synchronize the stream to detect singular - * or non-positive-definite matrices. Calling this method skips the extra - * device-to-host memcpy, stream synchronization, and (for the LDLT solver) - * in-kernel pivot/diagonal checks, which can be a significant fraction of - * the total solve time for small systems. + * The dense backends copy a device-side status flag back to the host after + * factorization and synchronize the stream to detect singular or + * non-positive-definite matrices. Skipping that removes a device-to-host + * memcpy, a stream synchronization, and (for LDLT) in-kernel pivot checks, + * which can be a significant fraction of the solve time for small systems. + * + * Backends with no such phase — cuDSS, which reports through its own status, + * and the PCG solver, which has no factorization — ignore this. It is a + * hint, not a contract; see DenseLinearSolverBase for the implementation. */ - void DisableSafetyChecks() { safety_checks_enabled_ = false; } - - /** @brief Returns whether post-factorization safety checks are enabled. */ - bool SafetyChecksEnabled() const { return safety_checks_enabled_; } + virtual void DisableSafetyChecks() {} /** @brief Virtual destructor for proper cleanup of derived solver instances. */ - virtual ~CSRSparseLinearSolver() = default; - - protected: - bool safety_checks_enabled_ = true; + virtual ~SparseLinearSolver() = default; }; } // namespace cunls diff --git a/cunls/minimizer/block_hessian_assembler.h b/cunls/minimizer/block_hessian_assembler.h index c9201ee..1417d24 100644 --- a/cunls/minimizer/block_hessian_assembler.h +++ b/cunls/minimizer/block_hessian_assembler.h @@ -45,12 +45,18 @@ class Problem; * atomic count by exactly a factor of `m` and removes the triplet-to-CSR * conversion, the `J^T J` kernel and the RHS SpMV from every iteration. * - * Storage stays plain CSR: the pattern produced by ComputeHessianStructure - * lays out each block pair contiguously within a row and at the same - * row-relative offset for every row of the block, so the scatter address is - * `row_offsets[col_a + i] + write_offset(a,b) + j`. No intermediate - * block-format Hessian is needed, and every downstream consumer (LM damping, - * column scaling, PCG, cuDSS) sees the matrix it already expects. + * The saving above is layout-independent: it comes from contracting per factor + * rather than per residual row, and it is the same whether the target is scalar + * CSR or block BSR. Both are supported, selected by NormalEquations and passed + * in via the `block_size` overload of Initialize(). There is no intermediate + * Hessian and no conversion between the two — the assembler scatters straight + * into whichever layout it was initialized for. + * + * Scattering is direct in either case because the pattern produced by + * HessianStructureBuilder lays out each block pair contiguously within a row + * and at the same row-relative offset for every row of the block. For scalar + * CSR the address is `row_offsets[col_a + i] + write_offset(a,b) + j`; for BSR + * the same offsets index tiles instead of entries. */ class BlockHessianAssembler { public: diff --git a/cunls/minimizer/gauss_newton_minimizer.cu b/cunls/minimizer/gauss_newton_minimizer.cu index 5128aa2..5b17e07 100644 --- a/cunls/minimizer/gauss_newton_minimizer.cu +++ b/cunls/minimizer/gauss_newton_minimizer.cu @@ -40,8 +40,8 @@ namespace cunls { */ GaussNewtonMinimizer::GaussNewtonMinimizer(const MinimizerOptions &options) : options_(options), - solver_(CreateCSRSparseLinearSolver(options_.sparse_linear_solver_type, - options_.sparse_linear_solver_config)) { + solver_(CreateSparseLinearSolver(options_.sparse_linear_solver_type, + options_.sparse_linear_solver_config)) { if (options_.disable_safety_checks) { solver_->DisableSafetyChecks(); } diff --git a/cunls/minimizer/normal_equations.cu b/cunls/minimizer/normal_equations.cu index 22e1d45..c2c7100 100644 --- a/cunls/minimizer/normal_equations.cu +++ b/cunls/minimizer/normal_equations.cu @@ -26,10 +26,12 @@ void NormalEquations::Initialize(cudaStream_t stream, const Problem &problem, in bool solver_supports_block_storage) { csr_dims_.Invalidate(); - // Block storage only pays off when the tangent dimensions share a factor and - // the solver can read tiles; a CSR-only backend would just have to expand - // them again. - block_size_ = solver_supports_block_storage ? ChooseHessianBlockSize(problem) : 1; + // Two independent conditions have to hold. The tangent dimensions must share + // a factor, or there are no tiles to form; and the solver must be able to read + // tiles, since a CSR-only backend is better served by a natively assembled CSR + // than by one expanded from blocks. + const int problem_block_size = ChooseHessianBlockSize(problem); + block_size_ = solver_supports_block_storage ? problem_block_size : 1; if (UsesBlockStorage()) { assembler_.Initialize(stream, problem, num_cols, block_size_, bsr_hessian_); @@ -37,6 +39,18 @@ void NormalEquations::Initialize(cudaStream_t stream, const Problem &problem, in return; } + // Say which of the two conditions failed. Falling back to scalar costs the + // block-storage delta, and a silent fallback is the kind of thing that only + // shows up as an unexplained regression in a profile. + if (problem_block_size <= 1) { + LogMessage("Hessian storage: CSR (tangent dimensions share no common factor)"); + } else { + LogMessage( + "Hessian storage: CSR (solver does not accept block storage; " + "problem would have supported block size {})", + problem_block_size); + } + assembler_.Initialize(stream, problem, num_cols, csr_hessian_); int num_rows = 0, num_matrix_cols = 0, num_nonzeros = 0; ExtractMatrixMetadata(stream, csr_hessian_, num_rows, num_matrix_cols, num_nonzeros); @@ -99,7 +113,7 @@ void NormalEquations::WeightedSquaredStepAsync(cudaStream_t stream, void *cuspar buffer, d_out, d_partials); } -bool NormalEquations::InitializeSolver(cudaStream_t stream, CSRSparseLinearSolver &solver, +bool NormalEquations::InitializeSolver(cudaStream_t stream, SparseLinearSolver &solver, const Problem &problem, const dvector &rhs, dvector &step) { if (UsesBlockStorage()) { @@ -108,7 +122,7 @@ bool NormalEquations::InitializeSolver(cudaStream_t stream, CSRSparseLinearSolve return solver.Initialize(stream, problem, csr_lhs_, rhs, step); } -bool NormalEquations::Solve(cudaStream_t stream, CSRSparseLinearSolver &solver, +bool NormalEquations::Solve(cudaStream_t stream, SparseLinearSolver &solver, const dvector &rhs, dvector &step) { if (UsesBlockStorage()) { return solver.Solve(stream, bsr_lhs_, rhs, step); diff --git a/cunls/minimizer/normal_equations.h b/cunls/minimizer/normal_equations.h index 87adb6c..1286df8 100644 --- a/cunls/minimizer/normal_equations.h +++ b/cunls/minimizer/normal_equations.h @@ -20,7 +20,7 @@ #include #include "cunls/common/types.h" -#include "cunls/linear_solver/csr_sparse_linear_solver.h" +#include "cunls/linear_solver/sparse_linear_solver_base.h" #include "cunls/minimizer/block_hessian_assembler.h" #include "cunls/minimizer/bsr_matrix.h" @@ -61,7 +61,7 @@ class NormalEquations { * @param problem The optimization problem. * @param num_cols Number of free tangent dimensions in the reduced system. * @param solver_supports_block_storage Whether the active solver can consume - * block storage; see CSRSparseLinearSolver::SupportsBlockStorage. + * block storage; see SparseLinearSolver::SupportsBlockStorage. */ void Initialize(cudaStream_t stream, const Problem &problem, int num_cols, bool solver_supports_block_storage); @@ -108,11 +108,11 @@ class NormalEquations { dvector &buffer); /** @brief Hands the working left-hand side to the solver for symbolic setup. */ - bool InitializeSolver(cudaStream_t stream, CSRSparseLinearSolver &solver, const Problem &problem, + bool InitializeSolver(cudaStream_t stream, SparseLinearSolver &solver, const Problem &problem, const dvector &rhs, dvector &step); /** @brief Solves with the working left-hand side. */ - bool Solve(cudaStream_t stream, CSRSparseLinearSolver &solver, const dvector &rhs, + bool Solve(cudaStream_t stream, SparseLinearSolver &solver, const dvector &rhs, dvector &step); /** @brief True when the block layout is live for the current problem. */ diff --git a/docs/sphinx/api/linear_solver.rst b/docs/sphinx/api/linear_solver.rst index e507b9c..4634e6d 100644 --- a/docs/sphinx/api/linear_solver.rst +++ b/docs/sphinx/api/linear_solver.rst @@ -4,7 +4,9 @@ Linear Solver API `cunls/linear_solver` hosts linear-system abstractions (block-Jacobi PCG, cuDSS integration, dense pivoted LDLT, dense Cholesky, and dense QR -solvers) behind a common CSR-based interface. +solvers) behind a common interface. Every backend accepts scalar CSR; all +but cuDSS also accept block BSR directly, which is what +:cpp:func:`SupportsBlockStorage` reports. SparseLinearSolverType ---------------------- @@ -18,11 +20,11 @@ Enum in `cunls/linear_solver/sparse_linear_solver.h`: - `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. +- `DenseLDLT` — densifies the coefficient matrix and solves with a custom + CUDA pivoted LDLT kernel. +- `DenseCholesky` — densifies the coefficient matrix and solves with + cuSOLVER Cholesky (requires SPD). +- `DenseQR` — densifies the coefficient matrix and solves with cuSOLVER QR. SparseLinearSolverConfig ------------------------ @@ -36,10 +38,10 @@ chosen ``SparseLinearSolverType`` is used: different backend is selected. - Dense backends take no extra configuration. -CSRSparseLinearSolver +SparseLinearSolver --------------------- -Abstract base (`cunls/linear_solver/csr_sparse_linear_solver.h`). +Abstract base (`cunls/linear_solver/sparse_linear_solver_base.h`). .. cpp:function:: bool Initialize(cudaStream_t stream, const Problem& problem, const CSRSparseMatrix& spd_matrix, const dvector& rhs, dvector& result) @@ -84,7 +86,9 @@ Abstract base (`cunls/linear_solver/csr_sparse_linear_solver.h`). .. cpp:function:: void DisableSafetyChecks() - Disables runtime safety checks in the solver. By default (safety checks + Advisory request to skip post-factorization safety checks; backends with no + such phase (cuDSS, which reports through its own status, and PCG, which has + no factorization) ignore it. By default (safety checks enabled), dense solvers validate every factorization and solve step: Cholesky checks cuSOLVER ``devInfo`` after ``potrf`` and ``potrs``; QR inspects the diagonal of ``R`` for rank deficiency; LDLT performs in-kernel @@ -96,11 +100,6 @@ Abstract base (`cunls/linear_solver/csr_sparse_linear_solver.h`). or ill-conditioned matrices. Normally called by the minimizer when ``MinimizerOptions::disable_safety_checks`` is ``true``. -.. cpp:function:: bool SafetyChecksEnabled() const - - :returns: ``true`` when post-factorization safety checks are enabled - (the default). - BlockSparsePCGOptions --------------------- @@ -360,7 +359,7 @@ Factory Function in `sparse_linear_solver.h`: -.. cpp:function:: SparseLinearSolverPtr CreateCSRSparseLinearSolver(SparseLinearSolverType type, const SparseLinearSolverConfig& config) +.. cpp:function:: SparseLinearSolverPtr CreateSparseLinearSolver(SparseLinearSolverType type, const SparseLinearSolverConfig& config) :param ``type``: [in] Backend type to instantiate. :param ``config``: [in] Backend-specific configuration blob. diff --git a/docs/sphinx/api/minimizer.rst b/docs/sphinx/api/minimizer.rst index ae950a5..ab895c5 100644 --- a/docs/sphinx/api/minimizer.rst +++ b/docs/sphinx/api/minimizer.rst @@ -396,10 +396,16 @@ The layout is chosen automatically, with no user-facing switch. Block storage requires a tile edge dividing every state block's tangent dimension (the gcd of the tangent sizes; see :code:`ChooseHessianBlockSize` in ``cunls/minimizer/bsr_matrix.h``) **and** a solver that reports -:code:`CSRSparseLinearSolver::SupportsBlockStorage`. When either does not hold — -a gcd of one, or a backend such as cuDSS or the dense factorizations that needs -CSR anyway — the minimizer falls back to scalar CSR with no behavioural change. -No conversion is ever performed on the solver path. +:code:`SparseLinearSolver::SupportsBlockStorage`. When either does not hold — +a gcd of one, or a backend such as cuDSS that needs CSR anyway — the minimizer +falls back to scalar CSR with no behavioural change. No conversion is ever +performed on the solver path; the fallback is assembled natively in CSR. + +The dense factorizations (:code:`DenseLDLT`, :code:`DenseCholesky`, +:code:`DenseQR`) accept either layout. They scatter the coefficient matrix into +an ``n x n`` dense buffer and never consult the sparse form again, so the layout +is invisible to them past the first kernel. cuDSS is the only backend that +declines block storage. ================================================================================ Python API (``pycunls``) diff --git a/tests/block_hessian_assembler_test.cpp b/tests/block_hessian_assembler_test.cpp index 374b5a0..8cf3337 100644 --- a/tests/block_hessian_assembler_test.cpp +++ b/tests/block_hessian_assembler_test.cpp @@ -50,6 +50,10 @@ #include "cunls/factor/reprojection_factor_batch.h" #include "cunls/factor/se3_between_factor_batch.h" #include "cunls/factor/vector_between_factor_batch.h" +#include "cunls/linear_solver/dense_cholesky_solver.h" +#include "cunls/linear_solver/dense_linear_solver.h" +#include "cunls/linear_solver/dense_qr_solver.h" +#include "cunls/linear_solver/sparse_linear_solver.h" #include "cunls/math/so_se_lie_math.h" #include "cunls/minimizer/bsr_matrix.h" #include "cunls/minimizer/device_reduction.h" @@ -1143,7 +1147,10 @@ TEST(HessianStorageTest, BlockSizeIsTheTangentDimensionGcd) { } TEST(HessianStorageTest, FallsBackToScalarWhenSolverNeedsCSR) { - // cuDSS consumes CSR, so block storage would only have to be expanded again. + // cuDSS consumes CSR, so the Hessian is assembled natively in CSR for it. + // Assembling in BSR and expanding would be strictly worse: the block layout's + // saving is the index array, which an expansion puts straight back, plus a + // per-iteration value permutation on top. auto data = MakePoseGraph(64, /*fix_first_pose=*/true); MinimizerOptions options; options.sparse_linear_solver_type = SparseLinearSolverType::cuDSS; @@ -1154,6 +1161,136 @@ TEST(HessianStorageTest, FallsBackToScalarWhenSolverNeedsCSR) { EXPECT_FALSE(builder.UsesBlockStorage()); } +/** + * @brief The storage layout must follow the problem, not the solver choice. + * + * A backend that densifies (or otherwise discards the sparse layout) has no + * reason to force the whole assembly onto scalar CSR, because doing so costs + * the block layout's smaller index array and the BuildSystem time with it. + * cuDSS is the one backend that genuinely cannot read tiles; every other one + * must take the block path on a problem that qualifies for it. + * + * This is the regression guard for that invariant: it is exactly the property + * that silently degrades if a new backend forgets to declare block support. + */ +TEST(HessianStorageTest, OnlyCuDSSDeclinesBlockStorage) { + auto data = MakePoseGraph(64, /*fix_first_pose=*/true); + ASSERT_EQ(6, ChooseHessianBlockSize(data->problem)) << "fixture must qualify for block storage"; + + const std::vector> kBackends = { + {SparseLinearSolverType::cuDSS, "cuDSS"}, + {SparseLinearSolverType::DenseLDLT, "DenseLDLT"}, + {SparseLinearSolverType::DenseCholesky, "DenseCholesky"}, + {SparseLinearSolverType::DenseQR, "DenseQR"}, + {SparseLinearSolverType::BlockSparsePCG, "BlockSparsePCG"}, + }; + + CudaStream stream; + for (const auto &[type, name] : kBackends) { + MinimizerOptions options; + options.sparse_linear_solver_type = type; + SystemBuilder builder(options); + builder.Build(stream.GetStream(), data->problem); + + const bool expect_block = type != SparseLinearSolverType::cuDSS; + EXPECT_EQ(expect_block, builder.UsesBlockStorage()) + << name << " selected the wrong Hessian storage layout"; + } +} + +/** + * @brief Every dense backend must read BSR and CSR to the same answer. + * + * The dense solvers scatter into an `n x n` buffer and never look at the sparse + * form again, so both layouts have to land identically in that buffer. Solving + * the same system twice — once from tiles, once from the expanded scalar copy — + * compares the two scatter kernels through the factorization that consumes + * them, which is the only place a mis-scattered entry would actually matter. + */ +void ExpectDenseBackendAgreesAcrossLayouts(SparseLinearSolver &solver, Problem &problem) { + CudaStream stream; + cudaStream_t s = stream.GetStream(); + + SystemBuilder builder; + builder.Build(s, problem); + ASSERT_TRUE(builder.UsesBlockStorage()) << "fixture must qualify for block storage"; + + const size_t n = static_cast(builder.Equations().LhsBSR().NumRows()); + + // Damp before solving, exactly as Levenberg-Marquardt does. An undamped + // Gauss-Newton Hessian is only positive *semi*-definite whenever the problem + // has gauge freedom — a bundle with no fixed pose has seven such directions — + // and Cholesky rightly refuses it. Damping is applied to the block LHS + // before the scalar copy is expanded from it, so both layouts hold the same + // matrix and the comparison stays about the scatter. + dvector ones(1.f, n); + builder.Equations().AddScaledDiagonalToLhs(s, 1e-2f, ones); + + const BSRSparseMatrix &bsr = builder.Equations().LhsBSR(); + CSRSparseMatrix csr; + dvector scratch; + test_utils::ExpandBSRToCSR(s, bsr, csr, scratch); + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + ASSERT_EQ(n, csr.NumRows()); + + const dvector &rhs = builder.Rhs(); + ASSERT_EQ(n, rhs.size()); + dvector x_block(n); + dvector x_scalar(n); + + // Scalar first: it is the reference, and if the fixture's Hessian is not + // solvable by this backend at all, the failure should point at the fixture + // rather than at the block scatter under test. + Problem empty; + ASSERT_TRUE(solver.Initialize(s, empty, csr, rhs, x_scalar)); + ASSERT_TRUE(solver.Solve(s, csr, rhs, x_scalar)) << "reference (scalar CSR) solve failed"; + ASSERT_TRUE(solver.Initialize(s, empty, bsr, rhs, x_block)); + ASSERT_TRUE(solver.Solve(s, bsr, rhs, x_block)) << "block (BSR) solve failed"; + THROW_ON_CUDA_ERROR(cudaStreamSynchronize(s)); + + std::vector block_host(n); + std::vector scalar_host(n); + x_block.CopyToHost(block_host.data(), n); + x_scalar.CopyToHost(scalar_host.data(), n); + + // The dense buffer is bit-identical either way; the factorization is the same + // code on the same input, so this is a tight bound rather than a solve-quality + // tolerance. + const float tol = 1e-5f * std::max(MaxAbs(scalar_host), 1e-6f); + for (size_t i = 0; i < block_host.size(); i++) { + ASSERT_NEAR(scalar_host[i], block_host[i], tol) << "solution mismatch at " << i; + } +} + +TEST(HessianStorageTest, DenseLDLTAgreesAcrossLayouts) { + auto data = MakePoseGraph(48, /*fix_first_pose=*/true); + DenseLDLTSolver solver; + ExpectDenseBackendAgreesAcrossLayouts(solver, data->problem); +} + +TEST(HessianStorageTest, DenseCholeskyAgreesAcrossLayouts) { + auto data = MakePoseGraph(48, /*fix_first_pose=*/true); + DenseCholeskySolver solver; + ExpectDenseBackendAgreesAcrossLayouts(solver, data->problem); +} + +TEST(HessianStorageTest, DenseQRAgreesAcrossLayouts) { + auto data = MakePoseGraph(48, /*fix_first_pose=*/true); + DenseQRSolver solver; + ExpectDenseBackendAgreesAcrossLayouts(solver, data->problem); +} + +/** + * @brief Bundle-adjustment tiles are 3x3 and the block rows are skewed, which + * the pose-graph fixtures above do not exercise. + */ +TEST(HessianStorageTest, DenseCholeskyAgreesAcrossLayoutsOnBundle) { + auto data = MakeBundle(6, 120, /*robust_loss=*/false); + DenseCholeskySolver solver; + ExpectDenseBackendAgreesAcrossLayouts(solver, data->problem); +} + // ============================================================================ // Equivalence tests // ============================================================================ diff --git a/tests/bsr_expansion.cu b/tests/bsr_expansion.cu index 7d23ef1..daa18ba 100644 --- a/tests/bsr_expansion.cu +++ b/tests/bsr_expansion.cu @@ -66,7 +66,13 @@ __global__ void FillExpandedRowOffsetsKernel(int num_rows, int block_size, const int block_row = row / block_size; const int sub_row = row - block_row * block_size; const int whole = block_row_offsets[block_row] * block_size; - const int partial = sub_row * (block_row_offsets[block_row + 1] - block_row_offsets[block_row]); + // The sentinel row `num_rows` lands on block_row == num_block_rows, where + // block_row_offsets[block_row + 1] is out of bounds. Its sub_row is 0, so + // the partial term is zero there anyway; skip the read rather than multiply + // it away. + const int partial = + (sub_row == 0) ? 0 + : sub_row * (block_row_offsets[block_row + 1] - block_row_offsets[block_row]); row_offsets[row] = (whole + partial) * block_size; } From 9de040e3da4c8a17d348a0d7682ceafcdefc63b1 Mon Sep 17 00:00:00 2001 From: Alex Korovko Date: Wed, 12 Aug 2026 17:31:52 -0700 Subject: [PATCH 7/7] Fix linter --- cunls/linear_solver/sparse_linear_solver.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cunls/linear_solver/sparse_linear_solver.cpp b/cunls/linear_solver/sparse_linear_solver.cpp index 8c3d01c..e1fce3b 100644 --- a/cunls/linear_solver/sparse_linear_solver.cpp +++ b/cunls/linear_solver/sparse_linear_solver.cpp @@ -42,4 +42,4 @@ SparseLinearSolverPtr CreateSparseLinearSolver(SparseLinearSolverType type, } } -} // namespace cunls \ No newline at end of file +} // namespace cunls