diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ae6e90a49f..128aa2e34b 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -462,7 +462,9 @@ if(BUILD_CUML_CPP_LIBRARY) # todo: separate solvers better if(all_algo OR solvers_algo) - target_sources(cuml_objs PRIVATE src/solver/lars.cu src/solver/solver.cu) + target_sources( + cuml_objs PRIVATE src/solver/lars.cu src/solver/solver.cu src/solver/nnls_batched.cu + ) endif() if(all_algo OR spectralclustering_algo) diff --git a/cpp/include/cuml/solvers/nnls.hpp b/cpp/include/cuml/solvers/nnls.hpp new file mode 100644 index 0000000000..2b1982931b --- /dev/null +++ b/cpp/include/cuml/solvers/nnls.hpp @@ -0,0 +1,91 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include + +namespace raft { +class handle_t; +} + +namespace CUML_EXPORT ML { +namespace Solver { + +/** + * @brief Solver backend selector for the batched NNLS entry point. + * + * Only the Lawson-Hanson active-set method is currently exposed. The selector + * is kept so that additional backends can be added later without changing the + * call signature. + */ +enum class NnlsBatchedSolver { + LAWSON = 0 ///< Lawson-Hanson active-set (exact, best for small n_cols). +}; + +/** + * Parameters for the batched NNLS solver. + */ +struct NnlsBatchedParams { + NnlsBatchedSolver solver = NnlsBatchedSolver::LAWSON; + int max_iter = 0; ///< 0 => per-solver default (3 * n_cols + 1 for Lawson). + double tol = 1e-6; ///< Dual-feasibility (KKT) tolerance on the projected gradient. +}; + +/** + * Solve a batch of Non-Negative Least Squares problems that share the same + * coefficient matrix but differ by right-hand side and active-column support: + * + * for p in [0, n_problems): + * X[:, p] = argmin_{x >= 0, x[j]=0 for masks[j,p]==0} + * 1/2 || A[:, support_p] x[support_p] - B[:, p] ||_2^2 + * + * The shared matrix A stays resident; its Gram matrix G = A^T A and the RHS + * projections C = A^T B are formed once (via cuBLAS) and reused by every + * problem. masks selects the active support per problem; masked-out + * coordinates of X are pinned to zero. + * + * @param handle raft handle (all work on its main stream). + * @param A column-major coefficient matrix, shape (m, n). + * @param m number of rows of A (length of each B column). + * @param n number of columns of A (length of each X column). + * @param B column-major RHS matrix, shape (m, n_problems). + * @param n_problems number of problems / columns of B and X. + * @param masks column-major uint8 matrix, shape (n, n_problems), + * F-contiguous; element (j, p) lives at masks[p*n + j] and + * is nonzero iff column j is active for problem p. May be + * null, meaning every column is active for every problem. + * @param X output solutions, column-major (n, n_problems). Masked-out + * rows are written as 0. + * @param fitted optional output A @ X, column-major (m, n_problems). May + * be null to skip the final gemm. + * @param params solver selection and per-backend knobs. + */ +void nnlsBatched(raft::handle_t& handle, + const float* A, + int m, + int n, + const float* B, + int n_problems, + const std::uint8_t* masks, + float* X, + float* fitted, + const NnlsBatchedParams& params); + +void nnlsBatched(raft::handle_t& handle, + const double* A, + int m, + int n, + const double* B, + int n_problems, + const std::uint8_t* masks, + double* X, + double* fitted, + const NnlsBatchedParams& params); + +} // namespace Solver +} // end namespace CUML_EXPORT ML diff --git a/cpp/src/solver/nnls_batched.cu b/cpp/src/solver/nnls_batched.cu new file mode 100644 index 0000000000..c0bda31006 --- /dev/null +++ b/cpp/src/solver/nnls_batched.cu @@ -0,0 +1,44 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "nnls_batched.cuh" + +#include + +#include + +namespace ML { +namespace Solver { + +void nnlsBatched(raft::handle_t& handle, + const float* A, + int m, + int n, + const float* B, + int n_problems, + const std::uint8_t* masks, + float* X, + float* fitted, + const NnlsBatchedParams& params) +{ + detail::nnls_batched_impl(handle, A, m, n, B, n_problems, masks, X, fitted, params); +} + +void nnlsBatched(raft::handle_t& handle, + const double* A, + int m, + int n, + const double* B, + int n_problems, + const std::uint8_t* masks, + double* X, + double* fitted, + const NnlsBatchedParams& params) +{ + detail::nnls_batched_impl(handle, A, m, n, B, n_problems, masks, X, fitted, params); +} + +} // namespace Solver +} // namespace ML diff --git a/cpp/src/solver/nnls_batched.cuh b/cpp/src/solver/nnls_batched.cuh new file mode 100644 index 0000000000..c997594f20 --- /dev/null +++ b/cpp/src/solver/nnls_batched.cuh @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "nnls_lawson.cuh" // detail::nnls_lawson_batched_dispatch + +#include +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace ML { +namespace Solver { +namespace detail { + +template +void nnls_batched_impl(raft::handle_t& handle, + const T* A, + int m, + int n, + const T* B, + int P, + const std::uint8_t* masks, + T* X, + T* fitted, + const NnlsBatchedParams& params) +{ + raft::common::nvtx::range fun_scope("ML::Solver::nnlsBatched(%d, %d, %d)", m, n, P); + ASSERT(m >= 1, "ML::Solver::nnlsBatched: m must be >= 1."); + ASSERT(n >= 1, "ML::Solver::nnlsBatched: n must be >= 1."); + ASSERT(P >= 1, "ML::Solver::nnlsBatched: n_problems must be >= 1."); + + cudaStream_t stream = handle.get_stream(); + + // Precompute the resident Gram matrix and RHS projections once, then reuse + // them across every problem in the batch: G = A^T A (n x n), C = A^T B (n x P). + // A col-major (m, n) buffer viewed as row-major (n, m) is exactly A^T, so the + // transpose is expressed through the operand layout rather than a flag. + auto G = raft::make_device_matrix(handle, n, n); + auto C = raft::make_device_matrix(handle, n, P); + + // gemm's mdspan overload shares one ElementType across all operands, so the + // read-only inputs are wrapped in non-const views (gemm never writes them). + auto* A_mut = const_cast(A); + auto At_view = + raft::make_device_matrix_view(A_mut, n, m); // A^T (n x m) + auto A_view = raft::make_device_matrix_view(A_mut, m, n); // A (m x n) + auto B_view = raft::make_device_matrix_view(const_cast(B), m, P); + raft::linalg::gemm(handle, At_view, A_view, G.view()); + raft::linalg::gemm(handle, At_view, B_view, C.view()); + + // Solve every problem with the batched Lawson-Hanson kernel. A max_iter of 0 + // selects the tight active-set cap of 3 * n + 1 outer steps. + int max_iter = params.max_iter; + if (max_iter <= 0) max_iter = 3 * n + 1; + const T tol = static_cast(params.tol); + + auto G_view = raft::make_const_mdspan(G.view()); + auto C_view = raft::make_const_mdspan(C.view()); + auto X_view = raft::make_device_matrix_view(X, n, P); + std::optional> M_view; + if (masks != nullptr) + M_view = raft::make_device_matrix_view(masks, n, P); + nnls_lawson_batched_dispatch(handle, G_view, C_view, M_view, X_view, max_iter, tol); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + + // Optional fitted = A @ X (m x P). + if (fitted != nullptr) { + const T one = T(1); + const T zero = T(0); + raft::linalg::gemm(handle, + /*trans_a=*/false, + /*trans_b=*/false, + m, + P, + n, + &one, + A, + m, + X, + n, + &zero, + fitted, + m, + stream); + } +} + +} // namespace detail +} // namespace Solver +} // namespace ML diff --git a/cpp/src/solver/nnls_lawson.cuh b/cpp/src/solver/nnls_lawson.cuh new file mode 100644 index 0000000000..1e3882af66 --- /dev/null +++ b/cpp/src/solver/nnls_lawson.cuh @@ -0,0 +1,1073 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include // device_matrix_view / col_major +#include // ASSERT +#include // raft::resource::get_cuda_stream +#include // raft::resource::get_custom_resource +#include // raft::resource::get_device_properties +#include // raft::resources +#include // raft::cache::lru +#include // RAFT_CUDA_TRY +#include // raft::WarpSize +#include // raft::div_rounding_up_unsafe +#include // raft::blockReduce / blockRankedReduce / warpReduce + +#include // rmm::device_uvector (L scratch) + +#include + +#include +#include +#include +#include +#include + +namespace ML { +namespace Solver { +namespace detail { + +/** + * Cholesky ridge added to a pivot to keep the factorisation positive-definite in + * the presence of round-off. It scales with the pivot magnitude and with the + * working precision (looser for float, tighter for double). + */ +template +__device__ inline T lawson_ridge_eps(T diag) +{ + const T rel = (sizeof(T) == 4 ? T(1e-7) : T(1e-14)); + return rel * (diag > T(0) ? diag : T(1)); +} + +/** + * Threshold below which an active coordinate driven down by the line search is + * treated as exactly zero and dropped from the active set. + */ +template +__device__ inline T lawson_zero_eps() +{ + return sizeof(T) == 4 ? T(1e-12) : T(1e-15); +} + +/** + * Finalise a bordering Cholesky append on lane 0: form the new pivot from the + * Gram diagonal `a22` and the accumulated dots, and, when it stays positive, + * write the diagonal of `L` and extend the forward-solve state `y`. Publishes + * +1 (accepted) or -1 (non-positive pivot) through `red_val[0]`. + */ +template +__device__ inline void lawson_finish_pivot( + T* L, int ld, int np, T a22, const T* c, int j_star, T* y, T* red_val, T dot_ll, T dot_ly) +{ + const T d2 = a22 + lawson_ridge_eps(a22) - dot_ll; + if (d2 > T(0)) { + const T d = std::sqrt(d2); + L[(np - 1) + (np - 1) * ld] = d; + y[np - 1] = (c[j_star] - dot_ly) / d; + red_val[0] = T(1); + } else { + red_val[0] = T(-1); + } +} + +/** + * Compute the dynamic shared-memory footprint of the Lawson-Hanson kernel for + * a single problem with `n` columns and a block of `BlockSize` threads. Layout + * (T and int arrays first, then int8): + * T c[n] A^T b + * T x[n] current solution + * T w[n] gradient (also scratch for the removed-index list) + * T s[n] trial solution (also RHS / downdate scratch) + * T y[n] forward-solve state L^-1 c_P (maintained + * incrementally; back-solved into s each iteration) + * T red_val[WarpSize] reduction scratch (also used to broadcast scalars) + * int red_idx[WarpSize] reduction scratch (also used to broadcast scalars) + * int idx[n] compact list of active-set column indices + * int8 act[n] 1 if column is in active set, 0 otherwise + * + * red_val/red_idx back the RAFT block reductions (raft::blockRankedReduce needs + * a value slot followed by an index slot per warp lane, i.e. WarpSize of each, + * laid out contiguously with red_idx immediately after red_val); slot 0 of each + * doubles as the scalar-broadcast channel. + * + * Neither the Gram matrix G nor the Cholesky factor L is staged into shared + * memory: G is read directly from global memory (L2-cached across the grid) and + * L lives in a per-block global scratch slab (also L2-cached). Shared memory is + * therefore only O(n), so occupancy is not bound by the n*n factor. + */ +template +inline std::size_t lawson_smem_bytes(int n) +{ + std::size_t bytes = 0; + bytes += sizeof(T) * static_cast(n) * 5; // c, x, w, s, y + bytes += sizeof(T) * raft::WarpSize; // red_val + bytes += sizeof(int) * raft::WarpSize; // red_idx + bytes += sizeof(int) * static_cast(n); // idx + bytes += sizeof(std::int8_t) * static_cast(n); // act + return bytes; +} + +template +struct LawsonSmem { + T* c; + T* x; + T* w; + T* s; + T* y; + T* red_val; + int* red_idx; + int* idx; + std::int8_t* act; +}; + +template +__device__ LawsonSmem lawson_smem_layout(unsigned char* smem, int n) +{ + LawsonSmem L; + // Wider arrays first (double is 8-byte aligned off the 16-byte-aligned base), + // then the 1-byte act[]; see lawson_smem_bytes. + L.c = reinterpret_cast(smem); + L.x = L.c + n; + L.w = L.x + n; + L.s = L.w + n; + L.y = L.s + n; + L.red_val = L.y + n; + // red_idx sits immediately after red_val so a single pointer (red_val) backs + // raft::blockRankedReduce, which expects the index slots at &shbuf[WarpSize]. + L.red_idx = reinterpret_cast(L.red_val + raft::WarpSize); + L.idx = L.red_idx + raft::WarpSize; + L.act = reinterpret_cast(L.idx + n); + return L; +} + +/** + * Block-wide argmax of `w[i]` restricted to indices where `act[i] == 0` and, + * when `mask` is non-empty, column `i` is enabled by the support (`mask(i) != 0`). + * Passing an empty `mask` view means every inactive column is eligible. Returns + * (max_value, argmax_index) on tid 0; the result is communicated to all threads + * via the supplied `red_val[0]` / `red_idx[0]` slots. + */ +template +__device__ inline void block_argmax_inactive(const T* w, + const std::int8_t* act, + raft::device_vector_view mask, + int n, + T* red_val, + int* red_idx) +{ + const int tid = threadIdx.x; + const bool has_mask = mask.data_handle() != nullptr; + + T thread_max = -std::numeric_limits::infinity(); + int thread_idx = -1; + for (int i = tid; i < n; i += BlockSize) { + if (act[i] == 0 && (!has_mask || mask(i) != 0)) { + T v = w[i]; + if (v > thread_max) { + thread_max = v; + thread_idx = i; + } + } + } + + // red_val backs the reduction scratch (red_idx is contiguous right after it); + // ties resolve to whichever lane RAFT keeps, which is fine for the algorithm. + auto res = raft::blockRankedReduce(thread_max, red_val, thread_idx, raft::max_op{}); + if (tid == 0) { + red_val[0] = res.first; + red_idx[0] = res.second; + } + __syncthreads(); +} + +/** + * Block-wide minimum of x[idx[jj]] / (x[idx[jj]] - s[jj]) over jj in [0, np) + * where s[jj] <= 0. Used to compute the alpha step. Also returns the count + * of binding indices (those with s[jj] <= 0) in `red_idx[0]`. If no binding + * indices are found, alpha is +inf. + */ +template +__device__ inline void block_min_alpha( + const T* x, const T* s, const int* idx, int np, T* red_val, int* red_idx) +{ + const int tid = threadIdx.x; + + T thread_min = std::numeric_limits::infinity(); + int thread_cnt = 0; + for (int jj = tid; jj < np; jj += BlockSize) { + T s_jj = s[jj]; + if (s_jj <= T(0)) { + T x_jj = x[idx[jj]]; + T denom = x_jj - s_jj; // strictly positive when x_jj >= 0 and s_jj <= 0 + // Numerical safety: if denom is tiny (x_jj == 0 and s_jj == 0) treat as + // non-binding (no contribution). + if (denom > T(0)) { + T alpha = x_jj / denom; + if (alpha < thread_min) thread_min = alpha; + ++thread_cnt; + } + } + } + // Two block reductions over the same warp scratch, run back to back: the min + // (blockRankedReduce pads absent lanes with +inf) and the binding count (a + // plain sum, for which blockReduce's zero padding is the correct identity). + auto min_res = raft::blockRankedReduce(thread_min, red_val, tid, raft::min_op{}); + int n_bind = raft::blockReduce(thread_cnt, reinterpret_cast(red_idx), raft::add_op{}); + if (tid == 0) { + red_val[0] = min_res.first; + red_idx[0] = n_bind; + } + __syncthreads(); +} + +/** + * Block-wide minimum of `s[0..np)`. Result on `red_val[0]`. + */ +template +__device__ inline void block_min(const T* s, int np, T* red_val) +{ + const int tid = threadIdx.x; + T thread_min = std::numeric_limits::infinity(); + for (int i = tid; i < np; i += BlockSize) { + T v = s[i]; + if (v < thread_min) thread_min = v; + } + // red_val's WarpSize index scratch (red_idx) is reserved contiguously after + // it, so blockRankedReduce is safe even though the index is unused here. + auto res = raft::blockRankedReduce(thread_min, red_val, tid, raft::min_op{}); + if (tid == 0) red_val[0] = res.first; + __syncthreads(); +} + +/** + * Projected gradient w = c - G x, reading the Gram matrix G directly from global + * memory (L2-cached across the grid). Because x is zero outside the active set, + * only the np active columns contribute: + * w[j] = c[j] - sum_{kk +__device__ inline void block_matvec_gradient( + T* w, + const T* c, + raft::device_matrix_view G, + const int* idx, + const T* x, + int np, + int n) +{ + const int tid = threadIdx.x; + for (int j = tid; j < n; j += BlockSize) { + T acc = c[j]; + for (int kk = 0; kk < np; ++kk) { + const int k = idx[kk]; + acc -= G(j, k) * x[k]; + } + w[j] = acc; + } + __syncthreads(); +} + +/** + * Incremental "bordering" Cholesky update. On entry L (a per-block global + * scratch slab, column-major, leading dimension `ld`, L2-cached) holds the + * (np-1)x(np-1) lower factor of the previously active submatrix; the newly + * activated column is idx[np-1]. The new Gram column + * is read directly from global memory G and the factor is extended in place: + * solve L_11 l = a_12 for the new row l = L(np-1, 0:np-1), + * new diagonal L(np-1, np-1) = sqrt(a_22 - l . l). + * Returns false (leaving the leading (np-1) block untouched) when the new pivot + * would be non-positive, i.e. activating idx[np-1] breaks positive-definiteness. + * + * Device analogue of raft::linalg::choleskyRank1Update, which is a host/cuBLAS + * routine and so cannot be called from within a block. `scratch` is O(np) + * working space (the new column / row l), and positive-definiteness is kept by a + * per-pivot ridge on a_22 (lawson_ridge_eps) rather than a global regulariser. + * + * Called by the whole block: the a_12 gather from global memory uses every + * thread for full memory throughput. Small factors (m = np-1 <= WarpSize) take a + * single-warp fast path -- the sequential forward solve, l.l/l.y dots and pivot + * test run on warp 0 alone (cheap `__syncwarp`/shuffle instead of block + * barriers). Larger factors use a blocked (panel) forward solve: panels of + * WarpSize columns are solved top-down, the O(m^2) update of the rows below each + * panel is applied by the WHOLE block (coalesced column reads), only the small + * diagonal panel solve stays on warp 0, and the closing (only O(m)) l.l/l.y dots + * stay on warp 0 with the same warpReduce as the fast path (so the pivot is + * bit-identical) -- block barriers are O(m / WarpSize) rather than one warp-0 + * serial section that idles the other warps. A single closing `__syncthreads` + * publishes the extended factor and the pivot-ok flag (`red_val[0]`). + * + * The forward-solve state y = L^-1 c_P is extended in the same pass: the leading + * block of L and the prefix of c_P are unchanged, so y[0:m] is untouched and the + * only new component is y[m] = (c[j_star] - l . y[0:m]) / L_22 (for m == 0 this + * degenerates to c[j_star] / L_22). It is written only when the pivot is + * accepted, so a rejected activation leaves y intact. + */ +template +__device__ inline bool block_chol_append(T* L, + int ld, + int np, + raft::device_matrix_view G, + const int* idx, + const T* c, + T* y, + T* red_val, + T* scratch) +{ + const int tid = threadIdx.x; + const int m = np - 1; // size of the existing factor L_11 + const int j_star = idx[np - 1]; + + // a_12[i] = G(idx[i], j_star): block-wide gather from global memory. + for (int i = tid; i < m; i += BlockSize) + scratch[i] = G(idx[i], j_star); + __syncthreads(); + + const int lane = tid % raft::WarpSize; + + // Fast path: one warp performs the whole forward solve L_11 l = a_12 (and the + // l.l / l.y dots) when the existing factor fits in a single panel. + if (m <= raft::WarpSize) { + if (tid < raft::WarpSize) { + // Forward solve in place in scratch (pivot folded into row i). + for (int i = 0; i < m; ++i) { + T y_i = scratch[i] / L[i + i * ld]; + __syncwarp(); + for (int j = i + lane; j < m; j += raft::WarpSize) + scratch[j] = (j > i) ? scratch[j] - L[j + i * ld] * y_i : y_i; + __syncwarp(); + } + + // Store l as the new row (np-1) of L; accumulate dot_ll = l . l and + // dot_ly = l . y[0:m] (the latter extends the forward-solve state). + T thread_dot_ll = T(0); + T thread_dot_ly = T(0); + for (int j = lane; j < m; j += raft::WarpSize) { + T lj = scratch[j]; + L[(np - 1) + j * ld] = lj; + thread_dot_ll += lj * lj; + thread_dot_ly += lj * y[j]; + } + thread_dot_ll = raft::warpReduce(thread_dot_ll, raft::add_op{}); + thread_dot_ly = raft::warpReduce(thread_dot_ly, raft::add_op{}); + if (lane == 0) + lawson_finish_pivot( + L, ld, np, G(j_star, j_star), c, j_star, y, red_val, thread_dot_ll, thread_dot_ly); + } + __syncthreads(); + return red_val[0] > T(0); + } + + // Blocked forward solve for larger factors: panels of WarpSize rows are solved + // top-down, and the contribution of each solved panel to the rows below it (the + // O(m^2) bulk) is applied by the WHOLE block, while only the small diagonal + // panel solve stays on warp 0. Block barriers are O(m / WarpSize) instead of a + // single warp-0 serial section that idles the other warps. The scheme is + // right-looking, so the between-panel update reads column j of the panel + // (`L[i + j*ld]`, contiguous over the target rows i) -- a coalesced global read. + { + constexpr int b = raft::WarpSize; + const int warp = tid / raft::WarpSize; + const int n_panel = raft::div_rounding_up_unsafe(m, b); + + for (int k = 0; k < n_panel; ++k) { + const int lo = k * b; + const int hi = (lo + b < m) ? (lo + b) : m; + + // Within-panel solve (warp 0), right-looking over its own columns [lo,hi); + // scratch[lo:hi] becomes the solved l entries for this panel. + if (warp == 0) { + for (int i = lo; i < hi; ++i) { + T y_i = scratch[i] / L[i + i * ld]; + __syncwarp(); + for (int j = i + 1 + lane; j < hi; j += raft::WarpSize) + scratch[j] -= L[j + i * ld] * y_i; + __syncwarp(); + if (lane == 0) scratch[i] = y_i; + __syncwarp(); + } + } + __syncthreads(); + + // Between-panel update (whole block): rows [hi,m) subtract the contribution + // of the just-solved columns [lo,hi). Consecutive threads own consecutive + // target rows i, so L[i + j*ld] is coalesced; scratch[j] is the solved l_j. + for (int i = hi + tid; i < m; i += BlockSize) { + T acc = scratch[i]; + for (int j = lo; j < hi; ++j) + acc -= L[i + j * ld] * scratch[j]; + scratch[i] = acc; + } + __syncthreads(); + } + + // Store l as the new row (np-1) of L and form the closing dots on warp 0. + // These are only O(m) (the O(m^2) work was the forward solve above), so + // keeping them single-warp costs little; warp 0 reuses the same warpReduce + // as the fast path so both paths produce a bit-identical pivot. + if (warp == 0) { + T thread_dot_ll = T(0); + T thread_dot_ly = T(0); + for (int j = lane; j < m; j += raft::WarpSize) { + T lj = scratch[j]; + L[(np - 1) + j * ld] = lj; + thread_dot_ll += lj * lj; + thread_dot_ly += lj * y[j]; + } + thread_dot_ll = raft::warpReduce(thread_dot_ll, raft::add_op{}); + thread_dot_ly = raft::warpReduce(thread_dot_ly, raft::add_op{}); + if (lane == 0) + lawson_finish_pivot( + L, ld, np, G(j_star, j_star), c, j_star, y, red_val, thread_dot_ll, thread_dot_ly); + } + } + __syncthreads(); + return red_val[0] > T(0); +} + +/** + * Remove active-set position `p` (0-based, in the current np-ordering) from the + * np x np lower Cholesky factor L (per-block global scratch slab, column-major, + * leading dimension `ld`, L2-cached), producing the (np-1)x(np-1) factor of the + * submatrix with row/column p deleted. + * + * Deleting an interior row/column reduces to a positive rank-1 Cholesky update + * of the trailing diagonal block by the below-diagonal part of column p: + * M M^T = L33 L33^T + l3p l3p^T, + * applied with a sequence of Givens rotations (Golub & Van Loan, "deleting a + * column"). `v` is O(np) scratch holding l3p and the rotated residual. + * + * Called by the whole block. The compaction is split into two independent + * pieces: the columns left of p only shift rows up within a column, so the + * whole block drives them in parallel (one thread per column, no sync); the + * trailing block shift L(i,j) <- L(i+1,j+1) sends every element to its + * lower-left neighbour, so it decomposes into independent diagonals (constant + * i-j) -- one lane walks one diagonal top-down, giving a single sync-free + * warp-wide pass on warp 0. The Givens sweep (an inherently sequential angle + * recurrence) also runs on warp 0 with `__syncwarp`/shuffle. Only a single + * closing `__syncthreads` (no per-row block barrier) exposes the downdated + * factor to the block. `v` is O(np) scratch holding l3p and the rotated + * residual. + * + * The forward-solve state y = L^-1 c_P is downdated in lock-step: y[0:p] is + * unchanged, the deleted component y[p] is saved as the rotation partner, the + * trailing y[p+1:] is compacted into y[p:], and the same Givens (c,s) that + * retriangularise L are applied to (y[p+k], partner) so y stays L^-1 c_P for the + * shrunken active set. These are scalar recurrences carried in a lane-0 + * register, so they piggyback on the existing sweep at no extra sync cost. + */ +template +__device__ inline void block_chol_delete_one(T* L, int ld, int np, int p, T* y, T* v) +{ + const int tid = threadIdx.x; + + // Region 1 -- columns [0, p): drop row p by shifting rows [p, np-1) up one + // (L(i,j) <- L(i+1,j)). Each destination reads the row directly below it in + // the SAME column, so a thread owning a whole column and sweeping rows + // ascending is race-free without any sync; columns are independent, so the + // whole block runs this in parallel (one thread per column). + for (int j = tid; j < p; j += BlockSize) + for (int i = p; i < np - 1; ++i) + L[i + j * ld] = L[(i + 1) + j * ld]; + + // The trailing shift (diagonal-parallel), the y compaction and the Givens + // downdate (a sequential angle recurrence) stay on warp 0, using only + // warp-level sync; the closing block barrier below is the only __syncthreads. + // Region 1 (other warps) writes disjoint columns [0, p), so no barrier here. + if (tid < raft::WarpSize) { + const int lane = tid; + const int q = np - 1 - p; // trailing block size + + // Save l3p = L(p+1 : np-1, p) before the trailing shift overwrites column p. + for (int i = lane; i < q; i += raft::WarpSize) + v[i] = L[(p + 1 + i) + p * ld]; + + // Compact the forward-solve state y (drop y[p]); keep the old y[p] as the + // initial Givens partner for the trailing sweep below. + T partner = T(0); + if (lane == 0) { + partner = y[p]; + for (int i = p; i < np - 1; ++i) + y[i] = y[i + 1]; + } + __syncwarp(); // l3p captured before column p is overwritten below + + // Region 2 -- trailing block up-left shift (drop row p AND column p): + // L(i,j) <- L(i+1, j+1). The move sends every element to its lower-left + // neighbour, so it splits into independent diagonals (constant i-j): one + // lane walks one diagonal (d = i-j) from the top, reading each source before + // this same lane later overwrites it. Diagonals never alias across lanes, + // so the whole shift is a single sync-free warp-wide pass. + for (int d = lane; d < q; d += raft::WarpSize) + for (int k = 0; k <= q - 1 - d; ++k) + L[(p + d + k) + (p + k) * ld] = L[(p + d + k + 1) + (p + k + 1) * ld]; + __syncwarp(); // trailing block published before the Givens sweep reads it + + // Positive rank-1 update of the trailing block (rows/cols p..np-2) by v. + for (int k = 0; k < q; ++k) { + T c = T(0), s = T(0); + if (lane == 0) { + T Lkk = L[(p + k) + (p + k) * ld]; + T vk = v[k]; + T r = std::sqrt(Lkk * Lkk + vk * vk); + c = (r > T(0)) ? (Lkk / r) : T(1); + s = (r > T(0)) ? (vk / r) : T(0); + L[(p + k) + (p + k) * ld] = r; + // Rotate the forward-solve state with the same (c,s): the trailing + // component y[p+k] pairs with the running partner just like L's + // column p+k pairs with the v-column. + T yk = y[p + k]; + y[p + k] = c * yk + s * partner; + partner = c * partner - s * yk; + } + c = __shfl_sync(0xffffffffu, c, 0); + s = __shfl_sync(0xffffffffu, s, 0); + for (int t = k + 1 + lane; t < q; t += raft::WarpSize) { + const int row = p + t; + T lik = L[row + (p + k) * ld]; + T vi = v[t]; + L[row + (p + k) * ld] = c * lik + s * vi; + v[t] = c * vi - s * lik; + } + __syncwarp(); + } + } + __syncthreads(); +} + +/** + * Back substitution for the system L^T s = y, where y = L^-1 c_P is the + * incrementally-maintained forward-solve state. This completes the solve of + * L L^T s = c_P: because the forward half is kept up to date across appends and + * downdates, each inner iteration only needs this back pass (half the sequential + * depth of a full forward+back solve). L is the lower triangular factor stored + * column-major with leading dimension `ld`; only the np x np leading block + * (lower triangle) is read, so the incrementally-updated factor (fixed ld = n, + * current size np) can be solved without repacking. `y` is left intact (it must + * survive into the next iteration); the solution is written to `s`. + * + * Called by the whole block. After copying y -> s (coalesced), small active + * sets (np <= WarpSize) take a single-warp fast path -- the sequential + * substitution on warp 0 with cheap per-row `__syncwarp` and one closing + * `__syncthreads` -- which covers the common late-stage case with minimal + * barriers. Larger active sets use a blocked (panel) scheme: panels of + * WarpSize rows are solved bottom-up, and the contribution of the already-solved + * rows below each panel (the O(np^2) bulk) is applied by the WHOLE block, while + * only the small diagonal panel solve stays on warp 0. Block barriers are then + * O(np / WarpSize) instead of one giant warp-0 serial section that idles the + * other warps. The pass is left-looking, + * x_i = (y_i - sum_{j>i} L[j,i] * x_j) / L_ii, + * so each row reads column `i` below the diagonal (`L[j + i*ld]`) -- contiguous + * in the column-major factor (coalesced global read) -- and reduces the partial + * products across the warp. `y` is left intact for the next iteration. + */ +template +__device__ inline void block_chol_backsolve(const T* L, int ld, int np, const T* y, T* s) +{ + const int tid = threadIdx.x; + const int lane = tid % raft::WarpSize; + + for (int j = tid; j < np; j += BlockSize) + s[j] = y[j]; + __syncthreads(); + + // Fast path: one warp solves the whole system when it fits in a single panel. + if (np <= raft::WarpSize) { + if (tid < raft::WarpSize) { + for (int i = np - 1; i >= 0; --i) { + T partial = T(0); + for (int j = i + 1 + lane; j < np; j += raft::WarpSize) + partial += L[j + i * ld] * s[j]; + partial = raft::warpReduce(partial, raft::add_op{}); + T x_i = (s[i] - partial) / L[i + i * ld]; + __syncwarp(); + if (lane == 0) s[i] = x_i; + __syncwarp(); + } + } + __syncthreads(); + return; + } + + // Blocked back-substitution for larger active sets. + constexpr int b = raft::WarpSize; + const int nwarps = BlockSize / raft::WarpSize; + const int warp = tid / raft::WarpSize; + const int n_panel = raft::div_rounding_up_unsafe(np, b); + + for (int k = n_panel - 1; k >= 0; --k) { + const int lo = k * b; + const int hi = (lo + b < np) ? (lo + b) : np; + + // Between-panel update (whole block): apply the already-solved rows below the + // panel, s[i] -= sum_{j>=hi} L[j,i]*s[j] for i in [lo,hi). Each panel row is + // owned by one warp (self-contained warp reduction). Bottom panel has none. + if (hi < np) { + for (int i = lo + warp; i < hi; i += nwarps) { + T partial = T(0); + for (int j = hi + lane; j < np; j += raft::WarpSize) + partial += L[j + i * ld] * s[j]; + partial = raft::warpReduce(partial, raft::add_op{}); + if (lane == 0) s[i] -= partial; + } + __syncthreads(); + } + + // Within-panel solve (warp 0): back-substitution over [lo,hi), reducing only + // over the panel's own columns (higher columns were applied above). + if (warp == 0) { + for (int i = hi - 1; i >= lo; --i) { + T partial = T(0); + for (int j = i + 1 + lane; j < hi; j += raft::WarpSize) + partial += L[j + i * ld] * s[j]; + partial = raft::warpReduce(partial, raft::add_op{}); + T x_i = (s[i] - partial) / L[i + i * ld]; + __syncwarp(); + if (lane == 0) s[i] = x_i; + __syncwarp(); + } + } + __syncthreads(); + } +} + +/** + * Batched, masked Lawson-Hanson NNLS kernel -- the single solver kernel used + * for both batched and single-problem solves. The grid is persistent: it is + * launched with `min(P, resident)` blocks that stride over the problems + * (`for p = blockIdx.x; p < P; p += gridDim.x`), so the per-block global scratch + * for the Cholesky factor L is bounded by hardware occupancy rather than by the + * batch size. Each block reads the Gram matrix `G = A^T A` directly from global + * memory (shared by the whole grid, L2-cached) and its own RHS projection from + * column `p` of `C = A^T B` and active-column support from column `p` of `masks`. + * + * The active-set Cholesky factor L is maintained incrementally in a per-block + * global scratch slab (L2-cached; slab `blockIdx.x` of `L_scratch`): each outer + * iteration appends the entering column with a bordering update + * (block_chol_append), and the inner line search shrinks it with Givens + * downdates (block_chol_delete_one). Consequently G is read only once per outer + * iteration -- the active columns for the projected gradient plus the single new + * column for the append -- and never re-gathered inside the inner loop. + * + * The forward-solve state y = L^-1 c_P is maintained alongside L (extended by + * block_chol_append, downdated by block_chol_delete_one), so the inner line + * search only runs a back substitution (block_chol_backsolve) to get the trial + * solution s = L^-T y instead of a full forward+back solve each iteration. + * + * A "non-batched" solve is simply the P == 1, empty-`masks` case: the caller + * forms G and C = A^T b once with cuBLAS and launches this kernel with a single + * block (see nnls_batched_impl). Columns disabled by a problem's mask are + * excluded from the argmax so they can never enter the solution. + * + * @tparam BlockSize number of threads per block (a multiple of raft::WarpSize; + * chosen at launch by LawsonBlockDispatch). + * @param G (n, n) Gram matrix (column-major), shared by all problems. + * @param C (n, P) matrix of A^T B (column-major). + * @param masks (n, P) column-major uint8 support; column p is problem p's + * support. May be an empty view, meaning "all columns eligible". + * @param X (n, P) solutions (column-major); written on exit. + * @param L_scratch global scratch for the Cholesky factor L; `gridDim.x` slabs + * of `n*n` (column-major, ld = n). Block `blockIdx.x` owns + * slab `blockIdx.x`; contents are rebuilt per problem. + * @param max_iter outer-iteration cap. + * @param tol optimality tolerance on the projected gradient. + */ +template +__global__ __launch_bounds__(BlockSize) void nnls_lawson_batched_kernel( + raft::device_matrix_view G, + raft::device_matrix_view C, + raft::device_matrix_view masks, + raft::device_matrix_view X, + T* L_scratch, + int max_iter, + T tol) +{ + const int n = G.extent(0); + const int P = C.extent(1); + + extern __shared__ unsigned char smem_raw[]; + LawsonSmem S = lawson_smem_layout(smem_raw, n); + // The active-set Cholesky factor L lives in global memory (per-block scratch + // slab, L2-cached) with a fixed leading dimension n so incremental + // append/downdate never repack it. Each block owns the slab at blockIdx.x. + T* L = L_scratch + static_cast(blockIdx.x) * n * n; + + __shared__ int sm_n_active; + __shared__ int sm_j_star; + __shared__ int sm_new_n; + __shared__ int sm_n_removed; + + const int tid = threadIdx.x; + + const int inner_budget_total = 3 * n + 1; + + // Persistent grid: gridDim.x resident blocks stride over the P problems. + for (int p = blockIdx.x; p < P; p += gridDim.x) { + // Column p of the (possibly empty) support, as a 1-D view for the argmax. + auto mask_col = (masks.size() != 0) + ? raft::make_device_vector_view(&masks(0, p), n) + : raft::device_vector_view{}; + + // ---- Phase 1+2: load c = C[:, p]; init x and active set ------------------ + for (int j = tid; j < n; j += BlockSize) { + S.c[j] = C(j, p); + S.x[j] = T(0); + S.act[j] = 0; + } + if (tid == 0) sm_n_active = 0; + __syncthreads(); + + // ---- Phase 3: outer loop (active-set growth) ----------------------------- + // Each pass brings in the inactive column with the largest projected + // gradient; the KKT conditions hold once that gradient drops to `tol`. + for (int outer = 0; outer < max_iter; ++outer) { + // Projected gradient w = c - G x, reading active columns of G from global. + block_matvec_gradient(S.w, S.c, G, S.idx, S.x, sm_n_active, n); + + block_argmax_inactive(S.w, S.act, mask_col, n, S.red_val, S.red_idx); + T max_w = S.red_val[0]; + int j_star = S.red_idx[0]; + if (j_star < 0 || max_w <= tol) break; + + // Activate j_star (append it at the end of the compact active set). + if (tid == 0) { + S.act[j_star] = 1; + S.idx[sm_n_active] = j_star; + sm_n_active = sm_n_active + 1; + sm_j_star = j_star; + } + __syncthreads(); + + // Incremental bordering append: extend L with the new column read from + // global G. A non-positive pivot means idx[np-1] breaks positive- + // definiteness, so undo the activation and stop the outer loop. + bool ok = block_chol_append(L, n, sm_n_active, G, S.idx, S.c, S.y, S.red_val, S.s); + if (!ok) { + if (tid == 0) { + sm_n_active = sm_n_active - 1; + S.act[sm_j_star] = 0; + } + __syncthreads(); + break; + } + + // Inner loop: solve the unconstrained active-set problem and, while any + // coordinate is negative, take the largest feasible step and drop the + // coordinates that hit zero until the trial solution is non-negative. + for (int inner = 0; inner < inner_budget_total; ++inner) { + const int np = sm_n_active; + + // Complete the solve L L^T s = c_P using the incrementally-maintained + // forward-solve state y = L^-1 c_P: only the back substitution is needed. + block_chol_backsolve(L, n, np, S.y, S.s); + + block_min(S.s, np, S.red_val); + T min_s = S.red_val[0]; + if (min_s > T(0)) { + for (int j = tid; j < n; j += BlockSize) + S.x[j] = T(0); + __syncthreads(); + for (int j = tid; j < np; j += BlockSize) + S.x[S.idx[j]] = S.s[j]; + __syncthreads(); + break; + } + + block_min_alpha(S.x, S.s, S.idx, np, S.red_val, S.red_idx); + T alpha = S.red_val[0]; + int n_binding = S.red_idx[0]; + if (n_binding == 0) break; + + for (int jj = tid; jj < np; jj += BlockSize) { + int j_idx = S.idx[jj]; + T xi = S.x[j_idx]; + T si = S.s[jj]; + S.x[j_idx] = xi + alpha * (si - xi); + } + __syncthreads(); + + // Compact the active set, recording the removed local positions (ascending) + // in the free scratch backed by S.w (unused inside the inner loop). + int* rem = reinterpret_cast(S.w); + if (tid == 0) { + const T zero_eps = lawson_zero_eps(); + int new_n = 0; + int n_rem = 0; + for (int jj = 0; jj < np; ++jj) { + int j_idx = S.idx[jj]; + if (S.x[j_idx] > zero_eps) { + S.idx[new_n++] = j_idx; + } else { + S.act[j_idx] = 0; + S.x[j_idx] = T(0); + rem[n_rem++] = jj; + } + } + sm_new_n = new_n; + sm_n_removed = n_rem; + } + __syncthreads(); + + // Downdate L for each removed position, deleting in descending order so + // earlier (lower-index) deletions stay valid as np shrinks. + int cur_np = np; + for (int r = sm_n_removed - 1; r >= 0; --r) { + block_chol_delete_one(L, n, cur_np, rem[r], S.y, S.s); + --cur_np; + } + if (tid == 0) sm_n_active = sm_new_n; + __syncthreads(); + + if (sm_n_active == 0) break; + } + } + + for (int j = tid; j < n; j += BlockSize) + X(j, p) = S.x[j]; + // Barrier before the next problem reuses shared memory, so a fast thread's + // Phase-1 re-init of S.x cannot clobber a slow thread's writeback read. + __syncthreads(); + } +} + +/** + * Opt a kernel in to more dynamic shared memory than the device's default + * per-block budget, but only when the requirement actually exceeds it. Returns + * true if the kernel is allowed to use `smem_bytes` of dynamic shared memory + * (either because it already fits the default budget, or because the opt-in + * succeeded); returns false otherwise, after resetting the pending CUDA error + * so a later query is not misattributed. The caller decides how to react to a + * false result: the Lawson selector defers to the occupancy query + * (`blocks_per_sm <= 0`). + * + * Both thresholds come from the device (`sharedMemPerBlock` / + * `sharedMemPerBlockOptin`) rather than hardcoded 48 KB / 96 KB constants, so + * the limits track the actual architecture. The shmem/L1 carveout split is + * never touched: no `cudaFuncAttributePreferredSharedMemoryCarveout` is issued, + * and the opt-in is skipped entirely when it is not needed. + */ +template +inline bool nnls_set_smem_attr(raft::resources const& handle, Kernel kernel, std::size_t smem_bytes) +{ + const cudaDeviceProp& dev_props = raft::resource::get_device_properties(handle); + if (smem_bytes <= static_cast(dev_props.sharedMemPerBlock)) return true; + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_bytes)); + if (err == cudaSuccess) return true; + // Reset the sticky error so a subsequent RAFT_CUDA_TRY doesn't observe it. + (void)cudaGetLastError(); + return false; +} + +// Desired number of co-resident blocks per problem before we keep a larger +// (lower-occupancy) block size instead of shrinking it. +constexpr int kResidenceMultiple = 2; + +/** + * One block-size candidate for a batched Lawson solve: the kernel to launch, + * its block size, its dynamic-shared-memory requirement, and its resident-block + * count. A chain of these forms a `lawson_plan` (see below). The kernel + * signature does not depend on the block size, so a single function pointer can + * hold any instantiation. + */ +template +struct lawson_selected { + using GView = raft::device_matrix_view; + using XView = raft::device_matrix_view; + using MView = raft::device_matrix_view; + using kernel_t = void (*)(GView, GView, MView, XView, T*, int, T); + + kernel_t kernel = nullptr; + int block_size = 0; + std::size_t smem = 0; + // Number of co-resident blocks of this instantiation across the whole GPU + // (max active blocks per SM * number of SMs). It is the sole channel through + // which `n_problems` influences the choice of block size: a launch saturates + // the device once `n_problems` exceeds `resident * kResidenceMultiple`, i.e. + // it takes more than `kResidenceMultiple` waves to drain the batch. + long long resident = 0; +}; + +/** + * Per-`n` plan for the batched Lawson kernel: the fixed chain of block-size + * candidates (largest first), each tagged with its resident-block count. The + * chain depends only on `(T, n)` -- the shared-memory footprint and the per-SM + * occupancy do not depend on `n_problems` -- so it is built once (a handful of + * CUDA API calls) and cached per handle keyed by `n` alone. `n_problems` then + * selects a step by cheap arithmetic (`pick`), which means every batch size for + * a given `n` shares one cache entry -- the widest possible equivalence class. + */ +template +struct lawson_plan { + // Block sizes 1024 -> 512 -> ... -> 32 (raft::WarpSize): at most 6 steps. + static constexpr int kMaxSteps = 6; + lawson_selected steps[kMaxSteps]; + int count = 0; + + /** + * Pick the block size for `n_problems`. The chain is ordered from the + * largest block (fewest resident blocks) to the smallest, and the occupancy + * gate that admits each smaller step is `n_problems`-independent, so we simply + * walk to a smaller block while the current one cannot already saturate the + * device (`resident * kResidenceMultiple < n_problems`). No CUDA API calls + * are made here; the chain was built once at plan-construction time. + */ + const lawson_selected& pick(int n_problems) const + { + int i = 0; + while (i + 1 < count && + steps[i].resident * static_cast(kResidenceMultiple) < n_problems) { + ++i; + } + return steps[i]; + } +}; + +/** Per-handle custom resource holding the LRU of per-`n` Lawson plans. */ +template +struct lawson_kernel_cache { + static constexpr std::size_t kDefaultSize = 32; + raft::cache::lru, std::equal_to<>, lawson_plan> value{kDefaultSize}; +}; + +/** + * Occupancy-driven block-size planner for the batched Lawson kernel. + * + * For each block size (largest first) we ask the driver how many blocks of this + * kernel instantiation can be co-resident across the whole GPU + * (R = max active blocks per SM * number of SMs) and record it as a plan step. + * We then consider halving the block size -- which usually raises the block + * count -- and continue the chain. + * + * The halving is gated by a condition that does *not* depend on `n_problems`: + * we only extend the chain if a smaller block keeps at least the same per-SM + * occupancy, measured as resident threads (blocks_per_sm * BlockSize), or still + * has enough threads for the work (>= 2n). Near a hardware blocks-per-SM cap, + * or when the kernel is shared-memory bound, a smaller block can fail to raise + * the block count enough to compensate for the fewer threads each block carries, + * which would trade device saturation for lower utilisation -- so in that case + * we stop the chain. It bottoms out at a single warp (raft::WarpSize). + * + * The result is a `lawson_plan` (never launched), so the caller can cache it + * per `n` and pick a step for any `n_problems` without repeating CUDA API calls. + */ +template +struct LawsonBlockDispatch { + // Append this block size to the plan, then -- if the (n_problems-independent) + // occupancy gate allows -- the smaller ones. `blocks_per_sm` is this level's + // occupancy, already measured by the caller, so each level performs at most + // one new occupancy query (the one for the half-sized candidate). + static void build(raft::resources const& handle, int n, int blocks_per_sm, lawson_plan& plan) + { + const int n_sm = raft::resource::get_device_properties(handle).multiProcessorCount; + const std::size_t smem = lawson_smem_bytes(n); + plan.steps[plan.count++] = lawson_selected{&nnls_lawson_batched_kernel, + BlockSize, + smem, + static_cast(blocks_per_sm) * n_sm}; + + if constexpr (BlockSize > raft::WarpSize) { + constexpr int half_block = BlockSize / 2; + const std::size_t smem_half = lawson_smem_bytes(n); + // Raise the smem cap before querying; otherwise the driver reports 0 + // active blocks for a kernel whose smem exceeds the device default. + nnls_set_smem_attr(handle, nnls_lawson_batched_kernel, smem_half); + int blocks_per_sm_half = 0; + RAFT_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_sm_half, nnls_lawson_batched_kernel, half_block, smem_half)); + const int occ_full = blocks_per_sm * BlockSize; + const int occ_half = blocks_per_sm_half * half_block; + // Reuse this query as the next level's occupancy. + if ((occ_half >= occ_full) || (half_block >= n)) { + LawsonBlockDispatch::build(handle, n, blocks_per_sm_half, plan); + } + } + } + + /** Entry point: query the largest block's occupancy once, then build the chain. */ + static lawson_plan start(raft::resources const& handle, int n) + { + const std::size_t smem = lawson_smem_bytes(n); + // The occupancy query returns 0 for a kernel whose smem exceeds the device + // default unless its max-dynamic-smem attribute is raised first. + nnls_set_smem_attr(handle, nnls_lawson_batched_kernel, smem); + int blocks_per_sm = 0; + RAFT_CUDA_TRY(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks_per_sm, nnls_lawson_batched_kernel, BlockSize, smem)); + + lawson_plan plan; + build(handle, n, blocks_per_sm, plan); + + // The occupancy query is the definitive "can even one block be placed" test. + // Because our smem is independent of the block size, a zero top-level + // occupancy means no block size can run, so it is a hard failure. + RAFT_EXPECTS( + plan.count > 0 && plan.steps[0].resident > 0, + "ML::Solver::nnlsBatched: no block of the Lawson kernel can be placed on an SM for " + "n=%d (dynamic shared memory %zu B exceeds the device per-block opt-in limit %zu B); " + "reduce n_cols or pick a different solver.", + n, + smem, + static_cast( + raft::resource::get_device_properties(handle).sharedMemPerBlockOptin)); + return plan; + } +}; + +/** Dispatch a batched Lawson solve, choosing the block size from the kernel's + * occupancy and the batch size (see LawsonBlockDispatch). The per-`n` plan is + * cached per-handle keyed by `n` alone, so the CUDA API calls behind it run + * once per distinct `n`; the batch size `n_problems` then selects a plan step + * by cheap arithmetic on every dispatch. */ +template +inline void nnls_lawson_batched_dispatch( + raft::resources const& handle, + raft::device_matrix_view G, + raft::device_matrix_view C, + std::optional> masks, + raft::device_matrix_view X, + int max_iter, + T tol) +{ + using MView = raft::device_matrix_view; + + const int n = X.extent(0); + const int n_problems = X.extent(1); + + auto& cache = raft::resource::get_custom_resource>(handle)->value; + lawson_plan plan; + if (!cache.get(n, &plan)) { + // Cache miss: build the (CUDA-API-heavy) plan once for this `n` and memoise + // it. On a hit no cudaFuncSetAttribute / occupancy query runs -- the smem + // attribute set while first building this entry persists at CUDA-context + // scope. + plan = LawsonBlockDispatch::start(handle, n); + cache.set(n, plan); + } + const lawson_selected& sel = plan.pick(n_problems); + + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + + // Persistent grid: launch min(n_problems, resident) co-resident blocks that + // stride over the problems, so the per-block global scratch for the Cholesky + // factor L is bounded by hardware occupancy (resident * n*n) rather than by + // the batch size. Each block owns the slab at blockIdx.x. + int grid = static_cast(std::min(n_problems, sel.resident)); + grid = std::max(grid, 1); + rmm::device_uvector L_scratch(static_cast(grid) * n * n, stream); + + // Empty view => "all columns eligible" inside the kernel. + MView mv = masks.has_value() ? *masks : MView{}; + sel.kernel<<>>( + G, C, mv, X, L_scratch.data(), max_iter, tol); +} + +} // namespace detail +} // namespace Solver +} // namespace ML diff --git a/python/cuml/cuml/solvers/CMakeLists.txt b/python/cuml/cuml/solvers/CMakeLists.txt index 89025491c6..e9260cbccd 100644 --- a/python/cuml/cuml/solvers/CMakeLists.txt +++ b/python/cuml/cuml/solvers/CMakeLists.txt @@ -1,11 +1,11 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= -set(cython_sources cd.pyx qn.pyx sgd.pyx) +set(cython_sources cd.pyx qn.pyx sgd.pyx nnls.pyx) if(NOT SINGLEGPU) list(APPEND cython_sources cd_mg.pyx) diff --git a/python/cuml/cuml/solvers/__init__.py b/python/cuml/cuml/solvers/__init__.py index bc5b907ad3..b408358111 100644 --- a/python/cuml/cuml/solvers/__init__.py +++ b/python/cuml/cuml/solvers/__init__.py @@ -1,8 +1,11 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # from cuml.solvers.cd import CD +from cuml.solvers.nnls import nnls, nnls_batched from cuml.solvers.qn import QN from cuml.solvers.sgd import SGD + +__all__ = ["CD", "QN", "SGD", "nnls", "nnls_batched"] diff --git a/python/cuml/cuml/solvers/nnls.pyx b/python/cuml/cuml/solvers/nnls.pyx new file mode 100644 index 0000000000..5af51b94fb --- /dev/null +++ b/python/cuml/cuml/solvers/nnls.pyx @@ -0,0 +1,536 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +import cupy as cp +import numpy as np + +import cuml.internals.nvtx as nvtx +from cuml.internals.base import get_handle + +from libc.stdint cimport uint8_t, uintptr_t +from pylibraft.common.handle cimport handle_t + +__all__ = ( + "nnls", + "nnls_batched", + "fit_nnls_batched", +) + +_NVTX_DOMAIN = "cuml_python" +_NVTX_CATEGORY = "solvers.nnls" + +# The solver argument is retained on the public API so more backends can be +# added later, but Lawson-Hanson is the only option available for now. +_VALID_SOLVERS = {"lawson"} +_SUPPORTED_DTYPES = (np.dtype(np.float32), np.dtype(np.float64)) + + +cdef extern from "cuml/solvers/nnls.hpp" namespace "ML::Solver" nogil: + cdef enum class NnlsBatchedSolver(int): + LAWSON = 0 + + cdef cppclass NnlsBatchedParams: + NnlsBatchedSolver solver + int max_iter + double tol + NnlsBatchedParams() except + + + cdef void nnlsBatched( + handle_t& handle, + const float* A, + int m, + int n, + const float* B, + int n_problems, + const uint8_t* masks, + float* X, + float* fitted, + const NnlsBatchedParams& params, + ) except + + + cdef void nnlsBatched( + handle_t& handle, + const double* A, + int m, + int n, + const double* B, + int n_problems, + const uint8_t* masks, + double* X, + double* fitted, + const NnlsBatchedParams& params, + ) except + + + +_BATCHED_SOLVERS = { + "lawson": NnlsBatchedSolver.LAWSON, +} + + +def fit_nnls_batched( + A, + B, + masks=None, + *, + convert_dtype=False, + str solver="lawson", + int max_iter=0, + double tol=1e-6, + bint compute_fitted=True, +): + """Solve a batch of masked, shared-``A`` Non-Negative Least Squares + problems in a single kernel launch. + + For every column ``p`` of ``B`` this solves + ``argmin_{x >= 0, x[j]=0 where masks[j, p]==0} || A @ x - B[:, p] ||_2`` + where ``A`` is shared across the whole batch. The Gram matrix + ``G = A.T @ A`` and the projections ``C = A.T @ B`` are formed once and + reused by every problem. + + Parameters + ---------- + A : array-like, shape=(m, n) + Shared coefficient matrix, dtype ``float32`` or ``float64``. + B : array-like, shape=(m, n_problems) + Right-hand-side matrix (one column per problem). + masks : array-like, shape=(n, n_problems), optional + Column-major (``(n_signatures, n_problems)``) boolean/uint8 support; + ``masks[j, p]`` is nonzero iff column ``j`` is active for problem ``p``. + Pass an **F-contiguous** array for a zero-copy device path: its memory + layout (signature index contiguous) then matches the kernel's per-block + access exactly. ``None`` means every column is active for every problem. + convert_dtype : bool, default=False + If True, convert ``A``/``B`` to a supported dtype. + solver : {'lawson'}, default='lawson' + Backend to use. Lawson-Hanson is the exact active-set method and the + only backend currently available. + max_iter : int, default=0 + Iteration cap. ``0`` selects the active-set default (``3 * n + 1``). + tol : double, default=1e-6 + Dual-feasibility (KKT) tolerance on the projected gradient. + compute_fitted : bool, default=True + Whether to also return ``fitted = A @ X``. + + Returns + ------- + X : cupy.ndarray, shape=(n, n_problems) + Non-negative solutions (column-major), masked-out rows set to 0. + fitted : cupy.ndarray, shape=(m, n_problems) or None + ``A @ X`` when ``compute_fitted`` else ``None``. + """ + if solver not in _BATCHED_SOLVERS: + raise ValueError( + f"Unknown solver {solver!r}. " + f"Expected one of {sorted(_BATCHED_SOLVERS)}." + ) + + handle = get_handle() + + cdef int m, n + A = cp.asarray(A) + if A.dtype not in _SUPPORTED_DTYPES: + if convert_dtype: + A = A.astype(np.float32) + else: + raise ValueError( + f"Unsupported A dtype {A.dtype}; expected float32 or float64." + ) + A = cp.asfortranarray(A) + m = A.shape[0] + n = A.shape[1] if A.ndim > 1 else 1 + + if m < 1: + raise ValueError( + f"Found array with {m} sample(s) (shape={A.shape}) while a " + f"minimum of 1 is required." + ) + if n < 1: + raise ValueError( + f"Found array with {n} feature(s) (shape={A.shape}) while " + f"a minimum of 1 is required." + ) + + cdef int n_problems + B = cp.asarray(B) + if B.dtype != A.dtype: + if convert_dtype: + B = B.astype(A.dtype) + else: + raise ValueError( + f"B dtype {B.dtype} does not match A dtype {A.dtype}." + ) + B = cp.asfortranarray(B) + if B.shape[0] != m: + raise ValueError( + f"Expected B with {m} rows to match A, got {B.shape[0]}." + ) + n_problems = B.shape[1] if B.ndim > 1 else 1 + + cdef uintptr_t masks_ptr = 0 + masks_arr = None + if masks is not None: + # Column-major (n, n_problems): an F-contiguous uint8 input is used in + # place, and its raw layout (signature index fastest) matches the + # kernel's ``masks[p*n + j]`` per-block access. A non-conforming input + # is copied to this layout (still correct, just not zero-copy). + masks_arr = cp.asfortranarray( + cp.asarray(masks).astype(np.uint8, copy=False) + ) + masks_rows = masks_arr.shape[0] + masks_cols = masks_arr.shape[1] if masks_arr.ndim > 1 else 1 + if masks_rows != n or masks_cols != n_problems: + raise ValueError( + f"Expected masks of shape ({n}, {n_problems}), got " + f"({masks_rows}, {masks_cols})." + ) + masks_ptr = masks_arr.data.ptr + + X = cp.zeros((n, n_problems), dtype=A.dtype, order="F") + + fitted = None + cdef uintptr_t fitted_ptr = 0 + if compute_fitted: + fitted = cp.zeros((m, n_problems), dtype=A.dtype, order="F") + fitted_ptr = fitted.data.ptr + + cdef NnlsBatchedParams params + params.solver = _BATCHED_SOLVERS[solver] + params.max_iter = max_iter + params.tol = tol + + cdef uintptr_t A_ptr = A.data.ptr + cdef uintptr_t B_ptr = B.data.ptr + cdef uintptr_t X_ptr = X.data.ptr + cdef handle_t* handle_ = handle.getHandle() + cdef bint is_float32 = A.dtype == np.float32 + + with nogil: + if is_float32: + nnlsBatched( + handle_[0], + A_ptr, + m, + n, + B_ptr, + n_problems, + masks_ptr, + X_ptr, + fitted_ptr, + params, + ) + else: + nnlsBatched( + handle_[0], + A_ptr, + m, + n, + B_ptr, + n_problems, + masks_ptr, + X_ptr, + fitted_ptr, + params, + ) + handle.sync() + + return X, fitted + + +def nnls( + A, + b, + *, + maxiter=None, + solver="lawson", + compute_rnorm=True, + tol=None, + check_every=10, +): + """Solve ``argmin_x || A @ x - b ||_2`` for ``x >= 0``. + + This is a GPU-accelerated equivalent of :func:`scipy.optimize.nnls`. + + Parameters + ---------- + A : array-like, shape (m, n) + Coefficient matrix. Accepts NumPy arrays, CuPy arrays, or any + ``__cuda_array_interface__`` object. Must have dtype ``float32`` or + ``float64``; other dtypes are cast to ``float32``. + b : array-like, shape (m,) + Right-hand side vector. Will be cast to ``A.dtype`` if needed. + maxiter : int, optional + Maximum number of active-set iterations. Defaults to the Lawson + active-set cap ``3 * n + 1``. + solver : {'lawson'}, default='lawson' + Which solver backend to use. Lawson-Hanson is the exact active-set + method that solves the whole problem in a single CUDA kernel using + normal equations and shared-memory Cholesky; it is the only backend + currently available. Best for small problems (n_cols up to roughly 90 + in double precision) where launch latency dominates. + compute_rnorm : bool, default=True + Whether to compute the residual 2-norm. When ``False``, ``rnorm`` + is returned as ``None`` and an extra matmul + reduction + device sync + per call is avoided. Useful in tight loops where the caller does + not need the residual norm. + tol : float, optional + Dual-feasibility (KKT) tolerance. Defaults to ``1e-4``. + check_every : int, default=10 + Retained for API compatibility; unused by the Lawson backend. + + Returns + ------- + x : cupy.ndarray, shape (n,) + Solution vector with all entries >= 0. + rnorm : float or None + The 2-norm of the residual, ``|| A @ x - b ||_2``, or ``None`` when + ``compute_rnorm=False``. + + Examples + -------- + >>> import cupy as cp + >>> from cuml.solvers import nnls + >>> A = cp.array([[1, 0], [1, 0], [0, 1]], dtype=cp.float32) + >>> b = cp.array([2, 1, 1], dtype=cp.float32) + >>> x, rnorm = nnls(A, b) + """ + if solver not in _VALID_SOLVERS: + raise ValueError( + f"Unknown solver {solver!r}. Expected one of {sorted(_VALID_SOLVERS)}" + ) + + with nvtx.annotate( + message=f"nnls[{solver}]", + domain=_NVTX_DOMAIN, + category=_NVTX_CATEGORY, + ): + with nvtx.annotate( + message="nnls.prepare_inputs", + domain=_NVTX_DOMAIN, + category=_NVTX_CATEGORY, + ): + A_gpu = cp.asarray(A) + b_gpu = cp.asarray(b) + + if A_gpu.ndim != 2: + raise ValueError(f"Expected 2-D array for A, got {A_gpu.ndim}-D") + if b_gpu.ndim != 1: + raise ValueError(f"Expected 1-D array for b, got {b_gpu.ndim}-D") + if A_gpu.shape[0] != b_gpu.shape[0]: + raise ValueError( + f"Incompatible dimensions: A has {A_gpu.shape[0]} rows, " + f"b has {b_gpu.shape[0]} elements" + ) + + # Ensure dtype is supported by the underlying solvers without + # forcing an unconditional float64 -> float32 conversion (which + # costs a kernel plus a device sync per call inside + # ``input_to_cuml_array`` whenever ``convert_dtype=True``). + # By doing the cast explicitly here we can then pass + # ``convert_dtype=False`` and skip that overhead entirely. + if A_gpu.dtype not in _SUPPORTED_DTYPES: + A_gpu = A_gpu.astype(np.float32, copy=False) + if b_gpu.dtype != A_gpu.dtype: + b_gpu = b_gpu.astype(A_gpu.dtype, copy=False) + + with nvtx.annotate( + message=f"nnls.solve[{solver}]", + domain=_NVTX_DOMAIN, + category=_NVTX_CATEGORY, + ): + # A single-RHS problem is just a batch of one column with no mask, so + # reuse the batched Gram-form engine rather than maintaining a + # parallel single-problem kernel. max_iter of 0 selects the tight + # active-set cap (3 * n_cols + 1) and the tolerance default is the + # active-set gradient tolerance (1e-4). + X_batched, _ = fit_nnls_batched( + A_gpu, + b_gpu.reshape(-1, 1), + None, + convert_dtype=False, + solver=solver, + max_iter=(0 if maxiter is None else int(maxiter)), + tol=(1e-4 if tol is None else float(tol)), + compute_fitted=False, + ) + coef = cp.asarray(X_batched)[:, 0] + + x = cp.asarray(coef).ravel() + if compute_rnorm: + with nvtx.annotate( + message="nnls.compute_rnorm", + domain=_NVTX_DOMAIN, + category=_NVTX_CATEGORY, + ): + residual = A_gpu @ x - b_gpu + rnorm = float(cp.linalg.norm(residual)) + else: + rnorm = None + + return x, rnorm + + +def nnls_batched( + A, + B, + masks=None, + b_index=None, + *, + solver="lawson", + compute_fitted=True, + maxiter=None, + tol=1e-6, +): + """Solve a batch of masked, shared-``A`` NNLS problems on the GPU. + + For every problem ``j`` this solves:: + + argmin_x || A[:, masks[:, j]] @ x - B[:, b_index[j]] ||_2 + subject to x >= 0 + + The design matrix ``A`` is shared across the whole batch; only the active + column set (``masks``) and the target (``b_index``) vary per problem. + + This is a device-native cuML primitive: it returns device (cupy) arrays and + keeps everything on the GPU. Passing cupy inputs avoids any host transfer, + so a caller can keep ``A``/``B`` resident across repeated calls simply by + holding on to cupy arrays. + + Parameters + ---------- + A : array-like, shape (m, n) + Shared coefficient matrix (signatures). NumPy or cupy; cupy inputs are + used in place (no copy) when already ``float32``/``float64``. + B : array-like, shape (m, n_targets) + Distinct target vectors. Not duplicated per problem; ``b_index`` + selects which column of ``B`` each problem uses. + masks : array-like, shape (n, n_problems), bool/uint8, optional + Column ``j`` selects the active columns of ``A`` for problem ``j`` + (``masks[i, j]`` nonzero iff signature ``i`` is active for problem + ``j``). For a zero-copy device path this should be an **F-contiguous** + ``(n, n_problems)`` array. ``None`` means every column is active for + every problem (``n_problems == n_targets``). + b_index : array-like, shape (n_problems,), int, optional + Target column of ``B`` for each problem, gathered on the device. + Defaults to the identity mapping (requires ``n_targets == n_problems``). + solver : {'lawson'}, default='lawson' + Backend to use. Lawson-Hanson is the only backend currently available. + compute_fitted : bool, default=True + Whether to also return ``fitted = A @ X``. When ``False``, ``fitted`` + is ``None`` and the extra matmul is skipped. + maxiter : int, optional + Iteration cap. Defaults to the active-set default (selected by 0). + tol : float, default=1e-6 + Dual-feasibility (KKT) tolerance. + + Returns + ------- + X : cupy.ndarray, shape (n, n_problems) + Non-negative solutions (column-major), masked-out rows set to 0. + fitted : cupy.ndarray, shape (m, n_problems) or None + ``A @ X`` per problem when ``compute_fitted``, else ``None``. + """ + if solver not in _VALID_SOLVERS: + raise ValueError( + f"Unknown solver {solver!r}. Expected one of {sorted(_VALID_SOLVERS)}" + ) + + with nvtx.annotate( + message=f"nnls_batched[{solver}]", + domain=_NVTX_DOMAIN, + category=_NVTX_CATEGORY, + ): + with nvtx.annotate( + message="nnls_batched.prepare_inputs", + domain=_NVTX_DOMAIN, + category=_NVTX_CATEGORY, + ): + A_gpu = cp.asarray(A) + B_gpu = cp.asarray(B) + + if A_gpu.ndim != 2: + raise ValueError(f"Expected 2-D array for A, got {A_gpu.ndim}-D") + if B_gpu.ndim != 2: + raise ValueError(f"Expected 2-D array for B, got {B_gpu.ndim}-D") + if A_gpu.shape[0] != B_gpu.shape[0]: + raise ValueError( + f"Incompatible dimensions: A has {A_gpu.shape[0]} rows, " + f"B has {B_gpu.shape[0]} rows" + ) + + # Cast to a supported dtype once, up front, so the Cython layer can + # run with convert_dtype=False (no extra kernel + sync per call). + if A_gpu.dtype not in _SUPPORTED_DTYPES: + A_gpu = A_gpu.astype(np.float32, copy=False) + if B_gpu.dtype != A_gpu.dtype: + B_gpu = B_gpu.astype(A_gpu.dtype, copy=False) + + n_cols = A_gpu.shape[1] + + masks_gpu = None + n_problems = None + if masks is not None: + masks_gpu = cp.asarray(masks) + if masks_gpu.ndim != 2: + raise ValueError( + f"Expected 2-D masks, got {masks_gpu.ndim}-D" + ) + if masks_gpu.shape[0] != n_cols: + raise ValueError( + f"Expected masks with {n_cols} rows (signatures), " + f"got {masks_gpu.shape[0]}" + ) + n_problems = masks_gpu.shape[1] + # uint8, F-contiguous (n, n_problems): a no-op when the caller + # already supplies that layout, so the Cython/kernel path is + # zero-copy and the mask reads are coalesced. + masks_gpu = cp.asfortranarray( + masks_gpu.astype(np.uint8, copy=False) + ) + + # Per-problem RHS gather on the device (no host copy, no duplication + # of distinct targets on the host side). + if b_index is not None: + b_index_gpu = ( + cp.asarray(b_index).astype(cp.int64, copy=False).ravel() + ) + n_targets = B_gpu.shape[1] + if b_index_gpu.size and ( + int(b_index_gpu.max()) >= n_targets + or int(b_index_gpu.min()) < 0 + ): + raise ValueError( + f"b_index entries must lie in [0, {n_targets})" + ) + if n_problems is not None and b_index_gpu.shape[0] != n_problems: + raise ValueError( + f"Expected b_index of length {n_problems}, " + f"got {b_index_gpu.shape[0]}" + ) + B_gpu = B_gpu[:, b_index_gpu] + + if n_problems is None: + n_problems = B_gpu.shape[1] + elif B_gpu.shape[1] != n_problems: + raise ValueError( + f"masks has {n_problems} problems but B (after b_index " + f"gather) has {B_gpu.shape[1]} columns" + ) + + # The batched kernel consumes column-major (F) B. + B_gpu = cp.asfortranarray(B_gpu) + + X, fitted = fit_nnls_batched( + A_gpu, + B_gpu, + masks_gpu, + convert_dtype=False, + solver=str(solver), + max_iter=(0 if maxiter is None else int(maxiter)), + tol=(1e-6 if tol is None else float(tol)), + compute_fitted=compute_fitted, + ) + + X = cp.asarray(X) + fitted = cp.asarray(fitted) if fitted is not None else None + return X, fitted diff --git a/python/cuml/tests/test_nnls.py b/python/cuml/tests/test_nnls.py new file mode 100644 index 0000000000..355aac5ce3 --- /dev/null +++ b/python/cuml/tests/test_nnls.py @@ -0,0 +1,579 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +"""Tests for cuml.solvers.nnls / nnls_batched. + +Both the single-problem ``nnls`` wrapper and the masked, shared-``A`` batched +``nnls_batched`` primitive are exercised against ``scipy.optimize.nnls`` on +three accuracy facets: + +* the residual 2-norm is no worse than scipy's (within a small slack), +* the KKT residual ``max_j |min(x_j, g_j)|`` is small relative to the scale of + ``A^T b``, and +* every coefficient is non-negative and masked-out coordinates are zero. + +Only the Lawson-Hanson backend is currently available; the ``solver`` argument +is retained on the public API but ``"lawson"`` is the sole valid value. + +Dimensions cover the sizes seen in Mutation Signature Analysis (the primary +consumer) as well as larger ``n`` and larger batch counts. The heavier shapes +are guarded behind cuML's ``quality_param`` / ``stress_param`` tiers so the +default unit run stays fast. +""" + +import cupy as cp +import numpy as np +import pytest +from scipy.optimize import nnls as scipy_nnls + +from cuml.solvers import nnls as cuml_nnls +from cuml.solvers import nnls_batched as cuml_nnls_batched +from cuml.testing.utils import quality_param, stress_param, unit_param + + +def _lawson_kw(dtype): + """Solver kwargs for the Lawson backend (max_iter=0 -> active-set cap).""" + return dict(maxiter=0, tol=(1e-4 if dtype == np.float32 else 1e-8)) + + +# --------------------------------------------------------------------------- +# Single-problem NNLS (cuml.solvers.nnls) +# --------------------------------------------------------------------------- + + +def _kkt_residual(A, x, b): + """``max_j |min(x_j, g_j)|`` for the smooth NNLS objective.""" + g = A.T @ (A @ x - b) + return float(np.max(np.abs(np.minimum(x, g)))) + + +def _kkt_scale(A, b): + return max(1.0, float(np.max(np.abs(A.T @ b)))) + + +def _check_solution(A, x, b, *, dtype, residual_slack=1.5e-3, kkt_rel=1e-2): + assert x.shape == (A.shape[1],) + assert np.all(x >= -1e-6), "negative entries detected" + x = np.maximum(x, 0.0) + + _, rnorm_ref = scipy_nnls(A.astype(np.float64), b.astype(np.float64)) + rnorm = float(np.linalg.norm(A @ x - b)) + + # The residual should match scipy's to within float-precision slack, scaled + # by ||b|| so the bound stays meaningful for under-determined problems where + # rnorm_ref == 0 and the achievable residual is dominated by round-off. + b_scale = max(1.0, float(np.linalg.norm(b))) + abs_slack = (1e-3 if dtype == np.float32 else 1e-6) * b_scale + assert rnorm <= rnorm_ref * (1.0 + residual_slack) + abs_slack, ( + f"residual {rnorm:.6g} > scipy residual {rnorm_ref:.6g} " + f"(slack {residual_slack:.1%} + {abs_slack:.1e})" + ) + + kkt = _kkt_residual( + A.astype(np.float64), x.astype(np.float64), b.astype(np.float64) + ) + scale = _kkt_scale(A.astype(np.float64), b.astype(np.float64)) + assert kkt <= kkt_rel * scale, ( + f"KKT residual {kkt:.3e} > {kkt_rel:.1e} * scale {scale:.3e}" + ) + + +def _make_tall(n_rows, n_cols, *, seed, dtype, sparsity=0.5): + """Random tall problem with a known non-negative ground truth.""" + rng = np.random.default_rng(seed) + A = rng.standard_normal((n_rows, n_cols)).astype(dtype) + x_true = rng.uniform(0.0, 2.0, n_cols).astype(dtype) + mask = rng.random(n_cols) < sparsity + x_true[mask] = 0.0 + noise = (0.01 if dtype == np.float64 else 0.05) * rng.standard_normal( + n_rows + ).astype(dtype) + b = (A @ x_true + noise).astype(dtype) + return A, b + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize( + "n_rows,n_cols", + [ + unit_param(200, 50), # mildly tall + unit_param(96, 40), # MSA-ish + quality_param(2000, 200), # tall + quality_param(200, 200), # square + stress_param(4000, 400), # large n + ], +) +def test_nnls_random_dense(dtype, n_rows, n_cols): + A, b = _make_tall(n_rows, n_cols, seed=0, dtype=dtype) + + x, rnorm = cuml_nnls(A, b, **_lawson_kw(dtype)) + x = cp.asnumpy(x) + + _check_solution(A, x, b, dtype=dtype) + assert rnorm == pytest.approx( + float(np.linalg.norm(A @ x - b)), rel=1e-3, abs=1e-4 + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_nnls_rank_deficient(dtype): + """A near-duplicated column makes the Gram matrix ill-conditioned; the + solver should still converge to a valid (possibly non-unique) solution.""" + rng = np.random.default_rng(7) + n_rows, n_cols = 400, 30 + A = rng.standard_normal((n_rows, n_cols)).astype(dtype) + A[:, 5] = A[:, 4] * dtype(0.5) + dtype(1e-6) * rng.standard_normal( + n_rows + ).astype(dtype) + x_true = np.maximum(rng.standard_normal(n_cols), 0.0).astype(dtype) + b = (A @ x_true + dtype(0.01) * rng.standard_normal(n_rows)).astype(dtype) + + x, _ = cuml_nnls(A, b, **_lawson_kw(dtype)) + x = cp.asnumpy(x) + + # Larger residual slack: rank-deficient problems have a continuum of optima. + _check_solution(A, x, b, dtype=dtype, residual_slack=5e-3, kkt_rel=5e-2) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_nnls_zero_input(dtype): + """An all-zero ``A`` is degenerate; the solver should return x == 0.""" + A = np.zeros((10, 4), dtype=dtype) + b = np.ones(10, dtype=dtype) + + x, rnorm = cuml_nnls(A, b) + x = cp.asnumpy(x) + assert np.allclose(x, 0.0) + assert rnorm == pytest.approx(float(np.linalg.norm(b)), rel=1e-5) + + +def test_nnls_compute_rnorm_false(): + """compute_rnorm=False skips the residual and returns None for it.""" + A, b = _make_tall(128, 20, seed=1, dtype=np.float64) + x, rnorm = cuml_nnls(A, b, compute_rnorm=False) + assert rnorm is None + _check_solution(A, cp.asnumpy(x), b, dtype=np.float64) + + +def test_nnls_default_solver_is_lawson(): + """The solver argument defaults to (and only accepts) 'lawson'.""" + A, b = _make_tall(64, 12, seed=2, dtype=np.float64) + x_default, _ = cuml_nnls(A, b) + x_explicit, _ = cuml_nnls(A, b, solver="lawson") + np.testing.assert_allclose( + cp.asnumpy(x_default), cp.asnumpy(x_explicit), rtol=1e-9, atol=1e-9 + ) + + +def test_nnls_unknown_solver(): + A = np.eye(3, dtype=np.float32) + b = np.ones(3, dtype=np.float32) + with pytest.raises(ValueError): + cuml_nnls(A, b, solver="apg") + + +# --------------------------------------------------------------------------- +# Batched, masked, shared-A NNLS (cuml.solvers.nnls_batched) +# --------------------------------------------------------------------------- + + +def _msa_reference(A, B, masks, b_index): + """Per-problem scipy reference matching MSA run_NNLS.nnls_batched.""" + n_cols, n_problems = masks.shape + out_weights = np.zeros((n_cols, n_problems), dtype=np.float64) + out_fitted = np.zeros((A.shape[0], n_problems), dtype=np.float64) + Ad = A.astype(np.float64) + Bd = B.astype(np.float64) + for j in range(n_problems): + cols = np.flatnonzero(masks[:, j]) + if cols.size == 0: + continue + w, _ = scipy_nnls(Ad[:, cols], Bd[:, b_index[j]]) + out_weights[cols, j] = w + out_fitted[:, j] = Ad[:, cols] @ w + return out_weights, out_fitted + + +def _check_batched( + A, + B, + masks, + b_index, + out_weights, + out_fitted, + *, + dtype, + residual_slack=5e-3, + kkt_rel=5e-2, + ref_subset=None, +): + """Validate a batched solve. Non-negativity and masked-out zeros are + checked on every problem; the (expensive) scipy residual/KKT comparison is + limited to ``ref_subset`` problem indices when given, keeping large-batch + runs bounded.""" + n_cols, P = masks.shape + assert out_weights.shape == (n_cols, P) + assert np.all(out_weights >= -1e-5), "negative weights detected" + + Ad = A.astype(np.float64) + Bd = B.astype(np.float64) + b_scale_all = max( + 1.0, float(np.max(np.linalg.norm(Bd[:, b_index], axis=0))) + ) + abs_slack = (2e-3 if dtype == np.float32 else 1e-6) * b_scale_all + + if ref_subset is None: + ref_subset = range(P) + ref_subset = set(int(j) for j in ref_subset) + + for j in range(P): + cols = np.flatnonzero(masks[:, j]) + off = np.ones(n_cols, dtype=bool) + off[cols] = False + assert np.allclose(out_weights[off, j], 0.0, atol=1e-5), ( + f"problem {j}: masked-out weights are nonzero" + ) + if j not in ref_subset or cols.size == 0: + continue + + b_j = Bd[:, b_index[j]] + w_ref, _ = scipy_nnls(Ad[:, cols], b_j) + w = out_weights[cols, j].astype(np.float64) + r = float(np.linalg.norm(Ad[:, cols] @ w - b_j)) + r_ref = float(np.linalg.norm(Ad[:, cols] @ w_ref - b_j)) + assert r <= r_ref * (1.0 + residual_slack) + abs_slack, ( + f"problem {j}: residual {r:.6g} > scipy {r_ref:.6g}" + ) + g = Ad[:, cols].T @ (Ad[:, cols] @ w - b_j) + kkt = float(np.max(np.abs(np.minimum(w, g)))) + scale = max(1.0, float(np.max(np.abs(Ad[:, cols].T @ b_j)))) + assert kkt <= kkt_rel * scale, ( + f"problem {j}: KKT {kkt:.3e} > {kkt_rel:.1e} * {scale:.3e}" + ) + + assert out_fitted.shape == (A.shape[0], P) + + +def _run_batched(A, B, masks, **kw): + """Call the device-native nnls_batched and mirror (X, fitted) to host.""" + X, fitted = cuml_nnls_batched(A, B, masks, **kw) + X = cp.asnumpy(X) + fitted = None if fitted is None else cp.asnumpy(fitted) + return X, fitted + + +def _make_batched( + m, n, P, *, seed, dtype, signature_like=False, mask_prob=None +): + """Random batched problem with a known non-negative ground truth.""" + rng = np.random.default_rng(seed) + if signature_like: + A = np.abs(rng.standard_normal((m, n))).astype(dtype) + else: + A = rng.standard_normal((m, n)).astype(dtype) + x_true = rng.uniform(0.0, 3.0, (n, P)).astype(dtype) + x_true[rng.random((n, P)) < 0.5] = 0.0 + B = (A @ x_true + dtype(0.01) * rng.standard_normal((m, P))).astype(dtype) + if mask_prob is None: + masks = x_true > 0.0 + else: + masks = rng.random((n, P)) < mask_prob + # Guarantee every problem keeps at least a few active columns. + for p in range(P): + if masks[:, p].sum() < 3: + masks[rng.choice(n, min(3, n), replace=False), p] = True + b_index = np.arange(P) + return A, B, masks, b_index + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("P", [1, 4, 64]) +def test_batched_random_dense(dtype, P): + A, B, masks, b_index = _make_batched(128, 24, P, seed=P, dtype=dtype) + masks = np.ones((24, P), dtype=bool) # unmasked full support + + out_weights, out_fitted = _run_batched( + A, B, masks, b_index=b_index, **_lawson_kw(dtype) + ) + _check_batched(A, B, masks, b_index, out_weights, out_fitted, dtype=dtype) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_batched_masked(dtype): + A, B, masks, b_index = _make_batched( + 96, 20, 16, seed=123, dtype=dtype, mask_prob=0.6 + ) + + out_weights, out_fitted = _run_batched( + A, B, masks, b_index=b_index, **_lawson_kw(dtype) + ) + _check_batched(A, B, masks, b_index, out_weights, out_fitted, dtype=dtype) + + +# MSA mutation-signature dimensions (m rows = mutation contexts, n = signatures). +_MSA_SHAPES = [ + unit_param(96, 65), # SBS-96 default + unit_param(78, 11), # DBS + unit_param(83, 17), # ID + quality_param(192, 54), # SBS-192 + quality_param(288, 10), # SBS-288 + quality_param(1536, 10), # SBS-1536 (Gram setup dominates) + stress_param(4608, 10), # SBS-4608 +] + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("m,n", _MSA_SHAPES) +def test_batched_msa_shapes(dtype, m, n): + """Shapes mimicking a mutation-signature-attribution round.""" + P = 128 + A, B, masks, b_index = _make_batched( + m, n, P, seed=(m * 131 + n), dtype=dtype, signature_like=True + ) + + out_weights, out_fitted = _run_batched( + A, B, masks, b_index=b_index, **_lawson_kw(dtype) + ) + _check_batched(A, B, masks, b_index, out_weights, out_fitted, dtype=dtype) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize( + "n", [unit_param(64), quality_param(128), stress_param(256)] +) +def test_batched_larger_n(dtype, n): + """Larger n exercises the multi-panel Cholesky and back-solve paths.""" + m, P = max(2 * n, 128), 32 + A, B, masks, b_index = _make_batched( + m, n, P, seed=n, dtype=dtype, mask_prob=0.7 + ) + + out_weights, out_fitted = _run_batched( + A, B, masks, b_index=b_index, **_lawson_kw(dtype) + ) + _check_batched(A, B, masks, b_index, out_weights, out_fitted, dtype=dtype) + + +def test_batched_b_index_gather(): + """Leave-one-out style: many problems share a few distinct B columns.""" + dtype = np.float64 + rng = np.random.default_rng(99) + m, n = 96, 20 + n_targets, n_problems = 8, 64 + A = np.abs(rng.standard_normal((m, n))).astype(dtype) + B = np.abs(rng.standard_normal((m, n_targets))).astype(dtype) + masks = rng.random((n, n_problems)) < 0.7 + for p in range(n_problems): + if masks[:, p].sum() < 3: + masks[rng.choice(n, 3, replace=False), p] = True + b_index = rng.integers(0, n_targets, size=n_problems) + + out_weights, out_fitted = _run_batched( + A, B, masks, b_index=b_index, **_lawson_kw(dtype) + ) + _check_batched(A, B, masks, b_index, out_weights, out_fitted, dtype=dtype) + + +def test_batched_b_index_out_of_range(): + A = np.abs(np.random.default_rng(0).standard_normal((16, 4))) + B = np.abs(np.random.default_rng(1).standard_normal((16, 3))) + masks = np.ones((4, 2), dtype=bool) + with pytest.raises(ValueError): + cuml_nnls_batched(A, B, masks, b_index=np.array([0, 5])) + + +def test_batched_b_index_length_mismatch(): + A = np.abs(np.random.default_rng(0).standard_normal((16, 4))) + B = np.abs(np.random.default_rng(1).standard_normal((16, 3))) + masks = np.ones((4, 2), dtype=bool) # 2 problems + with pytest.raises(ValueError): + cuml_nnls_batched(A, B, masks, b_index=np.array([0, 1, 2])) + + +def test_batched_fitted_matches_AX(): + rng = np.random.default_rng(0) + m, n, P = 64, 12, 8 + A = np.abs(rng.standard_normal((m, n))).astype(np.float64) + B = np.abs(rng.standard_normal((m, P))).astype(np.float64) + masks = np.ones((n, P), dtype=bool) + + out_weights, out_fitted = _run_batched(A, B, masks) + np.testing.assert_allclose( + out_fitted, A @ out_weights, rtol=1e-6, atol=1e-6 + ) + + +def test_batched_returns_device_arrays(): + """nnls_batched is device-native: it returns cupy arrays, and fitted is + None when compute_fitted=False (skipping the extra matmul).""" + rng = np.random.default_rng(3) + A = np.abs(rng.standard_normal((32, 6))).astype(np.float64) + B = np.abs(rng.standard_normal((32, 5))).astype(np.float64) + masks = np.ones((6, 5), dtype=bool) + + X, fitted = cuml_nnls_batched(A, B, masks) + assert isinstance(X, cp.ndarray) + assert isinstance(fitted, cp.ndarray) + assert X.shape == (6, 5) + assert fitted.shape == (32, 5) + + X2, fitted2 = cuml_nnls_batched(A, B, masks, compute_fitted=False) + assert isinstance(X2, cp.ndarray) + assert fitted2 is None + np.testing.assert_allclose( + cp.asnumpy(X2), cp.asnumpy(X), rtol=1e-6, atol=1e-6 + ) + + +@pytest.mark.parametrize("order", ["C", "F"]) +def test_batched_mask_order_parity(order): + """The result must be identical whether the mask is C- or F-contiguous. + F-contiguous is the zero-copy device path; C-contiguous is copied to F + inside the Cython layer. Both must produce the same solution.""" + rng = np.random.default_rng(order == "F") + m, n, P = 96, 20, 32 + A = np.abs(rng.standard_normal((m, n))).astype(np.float64) + x_true = rng.uniform(0.0, 3.0, (n, P)).astype(np.float64) + x_true[rng.random((n, P)) < 0.5] = 0.0 + B = (A @ x_true).astype(np.float64) + + masks_bool = rng.random((n, P)) < 0.6 + for p in range(P): + if masks_bool[:, p].sum() < 3: + masks_bool[rng.choice(n, 3, replace=False), p] = True + + masks = np.asarray(masks_bool, order=order) + assert masks.flags["%s_CONTIGUOUS" % order] + b_index = np.arange(P) + + X, _ = cuml_nnls_batched( + A, B, masks, b_index=b_index, **_lawson_kw(np.float64) + ) + other = "C" if order == "F" else "F" + X_other, _ = cuml_nnls_batched( + A, + B, + np.asarray(masks_bool, order=other), + b_index=b_index, + **_lawson_kw(np.float64), + ) + np.testing.assert_allclose( + cp.asnumpy(X), cp.asnumpy(X_other), rtol=1e-6, atol=1e-6 + ) + + +def test_batched_all_false_mask_column(): + """A problem with no active columns yields an all-zero solution column.""" + rng = np.random.default_rng(11) + m, n, P = 48, 10, 6 + A = np.abs(rng.standard_normal((m, n))).astype(np.float64) + B = np.abs(rng.standard_normal((m, P))).astype(np.float64) + masks = np.ones((n, P), dtype=bool) + masks[:, 2] = False # empty support for problem 2 + + X, fitted = _run_batched(A, B, masks) + assert np.allclose(X[:, 2], 0.0) + assert np.allclose(fitted[:, 2], 0.0) + b_index = np.arange(P) + _check_batched(A, B, masks, b_index, X, fitted, dtype=np.float64) + + +def test_batched_repeated_calls_resident(): + """Repeated calls that reuse the same device-resident A/B (MSA cache + pattern, which also re-hits the per-n dispatch plan cache) are stable.""" + rng = np.random.default_rng(2024) + m, n, P = 96, 24, 40 + A = cp.asarray(np.abs(rng.standard_normal((m, n))).astype(np.float64)) + B = cp.asfortranarray( + cp.asarray(np.abs(rng.standard_normal((m, P))).astype(np.float64)) + ) + masks = cp.asarray(np.ones((n, P), dtype=np.uint8)) + + X0, _ = cuml_nnls_batched(A, B, masks, **_lawson_kw(np.float64)) + for _ in range(3): + Xi, _ = cuml_nnls_batched(A, B, masks, **_lawson_kw(np.float64)) + np.testing.assert_allclose( + cp.asnumpy(Xi), cp.asnumpy(X0), rtol=1e-9, atol=1e-9 + ) + + +def test_batched_on_device_l2_scoring_parity(): + """The GPU boundary (MSA batched_solve_and_score) reduces fitted to the + L2_normalised_by_first similarity on the device. Verify that the cupy + reduction matches the numpy reference for the same fitted vectors.""" + rng = np.random.default_rng(2024) + m, n = 96, 24 + n_targets, P = 10, 48 + A = np.abs(rng.standard_normal((m, n))).astype(np.float64) + B = np.abs(rng.standard_normal((m, n_targets))).astype(np.float64) + masks_bool = rng.random((n, P)) < 0.7 + for p in range(P): + if masks_bool[:, p].sum() < 3: + masks_bool[rng.choice(n, 3, replace=False), p] = True + masks = np.asarray(masks_bool, order="F") + b_index = rng.integers(0, n_targets, size=P) + norm_obs = np.linalg.norm(B, axis=0) + + A_dev = cp.asarray(A) + B_dev = cp.asfortranarray(cp.asarray(B)) + bi_dev = cp.asarray(b_index) + _, fitted = cuml_nnls_batched( + A_dev, + B_dev, + cp.asarray(masks), + b_index=bi_dev, + **_lawson_kw(np.float64), + ) + resid = cp.linalg.norm(B_dev[:, bi_dev] - fitted, axis=0) + sims_dev = cp.asnumpy(1.0 - resid / cp.asarray(norm_obs)[bi_dev]) + + fitted_np = cp.asnumpy(fitted) + sims_ref = ( + 1.0 + - np.linalg.norm(B[:, b_index] - fitted_np, axis=0) / norm_obs[b_index] + ) + np.testing.assert_allclose(sims_dev, sims_ref, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize("P", [quality_param(4096), stress_param(65536)]) +def test_batched_large_batch(P): + """MSA-scale batch counts. All GPU invariants are checked on every + problem; the scipy residual/KKT comparison is limited to a deterministic + subset so the run stays bounded.""" + dtype = np.float32 + m, n = 96, 65 + A, B, masks, b_index = _make_batched( + m, n, P, seed=7, dtype=dtype, signature_like=True, mask_prob=0.6 + ) + + out_weights, out_fitted = _run_batched( + A, B, masks, b_index=b_index, **_lawson_kw(dtype) + ) + + subset = np.linspace(0, P - 1, 64, dtype=int) + _check_batched( + A, + B, + masks, + b_index, + out_weights, + out_fitted, + dtype=dtype, + ref_subset=subset, + ) + + +def test_batched_unknown_solver(): + A = np.eye(4, dtype=np.float32) + B = np.ones((4, 2), dtype=np.float32) + masks = np.ones((4, 2), dtype=bool) + with pytest.raises(ValueError): + cuml_nnls_batched(A, B, masks, solver="apg") + + +def test_batched_bad_mask_shape(): + A = np.eye(4, dtype=np.float32) + B = np.ones((4, 2), dtype=np.float32) + with pytest.raises(ValueError): + cuml_nnls_batched(A, B, np.ones((3, 4), dtype=bool))