Batched NNLS

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)

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 <cuml/common/export.hpp>
+
+#include <cstdint>
+
+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

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 <cuml/solvers/nnls.hpp>
+
+#include <raft/core/handle.hpp>
+
+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<float>(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<double>(handle, A, m, n, B, n_problems, masks, X, fitted, params);
+}
+
+}  // namespace Solver
+}  // namespace ML

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 <cuml/common/utils.hpp>
+#include <cuml/solvers/nnls.hpp>
+
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/handle.hpp>
+#include <raft/core/nvtx.hpp>
+#include <raft/linalg/gemm.cuh>
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+#include <optional>
+
+namespace ML {
+namespace Solver {
+namespace detail {
+
+template <typename T>
+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<T, int, raft::col_major>(handle, n, n);
+  auto C = raft::make_device_matrix<T, int, raft::col_major>(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<T*>(A);
+  auto At_view =
+    raft::make_device_matrix_view<T, int, raft::row_major>(A_mut, n, m);              // A^T (n x m)
+  auto A_view = raft::make_device_matrix_view<T, int, raft::col_major>(A_mut, m, n);  // A   (m x n)
+  auto B_view = raft::make_device_matrix_view<T, int, raft::col_major>(const_cast<T*>(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<T>(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<T, int, raft::col_major>(X, n, P);
+  std::optional<raft::device_matrix_view<const std::uint8_t, int, raft::col_major>> M_view;
+  if (masks != nullptr)
+    M_view = raft::make_device_matrix_view<const std::uint8_t, int, raft::col_major>(masks, n, P);
+  nnls_lawson_batched_dispatch<T>(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