Fix Laplacian calculation in spectral partitioning (#2568)

These changes enable the use of the new Lanczos solver for spectral partitioning. This solver gives the correct eigenvalues in some cases where the old solver did not, but a previous attempt to introduce this led to a downstream breakage because the Laplacian was not being calculated correctly in the spectral partitioning logic.

In this PR, we introduce a new `compute_graph_laplacian` function to correctly generate a new CSR matrix representing the graph Laplacian of the input CSR-formatted adjacency matrix. This resolves #2419. This function is used in the spectral partition function to obtain the correct partitioning using the new Lanczos solver. This eliminates all internal uses of `computeSmallestEigenvectors`, the old Lanczos implementation, which should allow it to be removed after a deprecation cycle for downstream consumers. This will in turn allow us to resolve #313.

Note that the originally-reported breakage from the previous attempt to switch to the new Lanczos solver was used to create the test confirming correct functionality of the current PR.

This PR is marked as a bugfix because it unblocks usage of the corrected Lanczos implementation. It is based on and requires PR #2541, and it will be marked as draft until that PR is merged.

Authors:
  - William Hicks (https://github.com/wphicks)

Approvers:
  - Victor Lafargue (https://github.com/viclafargue)
  - Divye Gala (https://github.com/divyegala)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #2568

cpp/include/raft/core/sparse_types.hpp

@@ -168,7 +168,7 @@ class sparse_matrix {
                 row_type n_rows,
                 col_type n_cols,
                 nnz_type nnz = 0) noexcept(std::is_nothrow_default_constructible_v<container_type>)
-    : structure_{handle, n_rows, n_cols, nnz}, cp_{}, c_elements_{cp_.create(handle, 0)} {};
+    : structure_{handle, n_rows, n_cols, nnz}, cp_{}, c_elements_{cp_.create(handle, nnz)} {};
 
   // Constructor that owns the data but not the structure
   // This constructor is only callable with a `structure_type == *_structure_view`

cpp/include/raft/sparse/linalg/detail/laplacian.cuh

@@ -0,0 +1,118 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+#include <raft/core/detail/macros.hpp>
+#include <raft/core/device_csr_matrix.hpp>
+#include <raft/core/resources.hpp>
+
+#include <type_traits>
+
+namespace raft {
+namespace sparse {
+namespace linalg {
+namespace detail {
+
+/* Compute the graph Laplacian of an adjacency matrix
+ *
+ * This kernel implements the necessary logic for computing a graph
+ * Laplacian for an adjacency matrix in CSR format. A custom kernel is
+ * required because cusparse does not conveniently implement matrix subtraction with 64-bit
+ * indices. The custom kernel also allows the computation to be completed
+ * with no extra allocations or compute.
+ */
+template <typename ElementType, typename IndptrType, typename IndicesType>
+RAFT_KERNEL compute_graph_laplacian_kernel(ElementType* output_values,
+                                           IndicesType* output_indices,
+                                           IndptrType* output_indptr,
+                                           IndptrType dim,
+                                           ElementType const* adj_values,
+                                           IndicesType const* adj_indices,
+                                           IndptrType const* adj_indptr)
+{
+  /* The graph Laplacian L of an adjacency matrix A is given by:
+   * L = D - A
+   * where D is the degree matrix of A. The degree matrix is itself defined
+   * as the sum of each row of A and represents the degree of the node
+   * indicated by the index of the row. */
+
+  for (auto row = threadIdx.x + blockIdx.x * blockDim.x; row < dim; row += blockDim.x * gridDim.x) {
+    auto row_begin = adj_indptr[row];
+    auto row_end   = adj_indptr[row + 1];
+    // All output indexes will need to be offset by the row, since every row will
+    // gain exactly one new non-zero element. degree_output_index is the index
+    // where we will store the degree of each row
+    auto degree_output_index = row_begin + row;
+    auto degree_value        = ElementType{};
+    // value_index indicates the index of the current value in the original
+    // adjacency matrix
+    for (auto value_index = row_begin; value_index < row_end; ++value_index) {
+      auto col_index         = adj_indices[value_index];
+      auto is_lower_diagonal = col_index < row;
+      auto output_index      = value_index + row + !is_lower_diagonal;
+      auto input_value       = adj_values[value_index];
+      degree_value += input_value;
+      output_values[output_index]  = ElementType{-1} * input_value;
+      output_indices[output_index] = col_index;
+      // Increment the index where we will store the degree for every non-zero
+      // element before we reach the diagonal
+      degree_output_index += is_lower_diagonal;
+    }
+    output_values[degree_output_index]  = degree_value;
+    output_indices[degree_output_index] = row;
+    output_indptr[row]                  = row_begin + row;
+    output_indptr[row + 1]              = row_end + row + 1;
+  }
+}
+
+template <typename ElementType, typename IndptrType, typename IndicesType, typename NZType>
+auto compute_graph_laplacian(
+  raft::resources const& res,
+  device_csr_matrix_view<ElementType, IndptrType, IndicesType, NZType> input)
+{
+  auto input_structure = input.structure_view();
+  auto dim             = input_structure.get_n_rows();
+  RAFT_EXPECTS(dim == input_structure.get_n_cols(),
+               "The graph Laplacian can only be computed on a square adjacency matrix");
+  auto result = make_device_csr_matrix<std::remove_const_t<ElementType>,
+                                       std::remove_const_t<IndptrType>,
+                                       std::remove_const_t<IndicesType>,
+                                       std::remove_const_t<NZType>>(
+    res,
+    dim,
+    dim,
+    /* The nnz for the result will be the dimension of the (square) input matrix plus the number of
+     * non-zero elements in the original matrix, since we introduce non-zero elements along the
+     * diagonal to represent the degree of each node. */
+    input_structure.get_nnz() + dim);
+  auto result_structure                         = result.structure_view();
+  auto static constexpr const threads_per_block = 256;
+  auto blocks = std::min(int((dim + threads_per_block - 1) / threads_per_block), 65535);
+  auto stream = resource::get_cuda_stream(res);
+  detail::compute_graph_laplacian_kernel<<<threads_per_block, blocks, 0, stream>>>(
+    result.get_elements().data(),
+    result_structure.get_indices().data(),
+    result_structure.get_indptr().data(),
+    dim,
+    input.get_elements().data(),
+    input_structure.get_indices().data(),
+    input_structure.get_indptr().data());
+  return result;
+}
+
+}  // namespace detail
+}  // namespace linalg
+}  // namespace sparse
+}  // namespace raft

cpp/include/raft/sparse/linalg/laplacian.cuh

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+#include <raft/core/device_csr_matrix.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/sparse/linalg/detail/laplacian.cuh>
+
+namespace raft {
+namespace sparse {
+namespace linalg {
+
+/** Given a CSR adjacency matrix, return the graph Laplacian
+ *
+ * Note that for non-symmetric matrices, the out-degree Laplacian is returned.
+ */
+template <typename ElementType, typename IndptrType, typename IndicesType, typename NZType>
+auto compute_graph_laplacian(
+  raft::resources const& res,
+  device_csr_matrix_view<ElementType, IndptrType, IndicesType, NZType> input)
+{
+  return detail::compute_graph_laplacian(res, input);
+}
+
+}  // namespace linalg
+}  // namespace sparse
+}  // namespace raft

cpp/include/raft/spectral/detail/matrix_wrappers.hpp

@@ -15,6 +15,8 @@
  */
 #pragma once
 
+#include <raft/core/device_csr_matrix.hpp>
+#include <raft/core/device_span.hpp>
 #include <raft/core/resource/cublas_handle.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resource/cusparse_handle.hpp>
@@ -33,6 +35,7 @@
 #include <thrust/system/cuda/execution_policy.h>
 
 #include <algorithm>
+#include <cstddef>
 
 // =========================================================
 // Useful macros
@@ -181,6 +184,19 @@ struct sparse_matrix_t {
   {
   }
 
+  auto to_csr_matrix_view() const
+  {
+    // The usage of sparse_matrix_t prior to introduction of this method
+    // assumed that all data was strictly on device. We will make the same
+    // assumption for construction of the csr_matrix_view
+    return device_csr_matrix_view<value_type const, index_type const, index_type const, index_type>{
+      device_span<value_type const>{values_, std::uint64_t(nnz_)},
+      device_compressed_structure_view<index_type const, index_type const, index_type>{
+        device_span<index_type const>{row_offsets_, std::uint64_t(nrows_ + 1)},
+        device_span<index_type const>{col_indices_, std::uint64_t(nnz_)},
+        ncols_}};
+  }
+
   virtual ~sparse_matrix_t(void) =
     default;  // virtual because used as base for following matrix types
 

cpp/include/raft/spectral/detail/partition.hpp

@@ -18,6 +18,7 @@
 #include <raft/core/resource/cublas_handle.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/detail/cublas_wrappers.hpp>
+#include <raft/sparse/linalg/laplacian.cuh>
 #include <raft/spectral/cluster_solvers.cuh>
 #include <raft/spectral/detail/spectral_util.cuh>
 #include <raft/spectral/eigen_solvers.cuh>
@@ -97,14 +98,15 @@ std::tuple<vertex_t, weight_t, vertex_t> partition(
   // Compute eigenvectors of Laplacian
 
   // Initialize Laplacian
-  /// sparse_matrix_t<vertex_t, weight_t> A{handle, graph};
-  spectral::matrix::laplacian_matrix_t<vertex_t, weight_t, nnz_t> L{handle, csr_m};
+  auto laplacian =
+    raft::sparse::linalg::compute_graph_laplacian(handle, csr_m.to_csr_matrix_view());
 
   auto eigen_config = eigen_solver.get_config();
   auto nEigVecs     = eigen_config.n_eigVecs;
 
   // Compute smallest eigenvalues and eigenvectors
-  std::get<0>(stats) = eigen_solver.solve_smallest_eigenvectors(handle, L, eigVals, eigVecs);
+  std::get<0>(stats) =
+    eigen_solver.solve_smallest_eigenvectors(handle, laplacian.view(), eigVals, eigVecs);
 
   // Whiten eigenvector matrix
   transform_eigen_matrix(handle, n, nEigVecs, eigVecs);

cpp/include/raft/spectral/eigen_solvers.cuh

@@ -18,6 +18,7 @@
 
 #pragma once
 
+#include <raft/core/mdspan.hpp>
 #include <raft/sparse/solver/lanczos.cuh>
 #include <raft/spectral/matrix_wrappers.hpp>
 
@@ -49,28 +50,52 @@ struct lanczos_solver_t {
   {
   }
 
-  index_type_t solve_smallest_eigenvectors(
+  auto solve_smallest_eigenvectors(
     raft::resources const& handle,
     matrix::sparse_matrix_t<index_type_t, value_type_t, size_type_t> const& A,
     value_type_t* __restrict__ eigVals,
     value_type_t* __restrict__ eigVecs) const
+  {
+    auto csr_structure =
+      raft::make_device_compressed_structure_view<index_type_t, index_type_t, index_type_t>(
+        const_cast<index_type_t*>(A.row_offsets_),
+        const_cast<index_type_t*>(A.col_indices_),
+        A.nrows_,
+        A.ncols_,
+        A.nnz_);
+
+    auto csr_matrix =
+      raft::make_device_csr_matrix_view<value_type_t, index_type_t, index_type_t, index_type_t>(
+        const_cast<value_type_t*>(A.values_), csr_structure);
+    return solve_smallest_eigenvectors(handle, csr_matrix, eigVals, eigVecs);
+  }
+
+  auto solve_smallest_eigenvectors(
+    raft::resources const& handle,
+    device_csr_matrix_view<value_type_t, index_type_t, index_type_t, index_type_t> input,
+    value_type_t* __restrict__ eigVals,
+    value_type_t* __restrict__ eigVecs) const
   {
     RAFT_EXPECTS(eigVals != nullptr, "Null eigVals buffer.");
     RAFT_EXPECTS(eigVecs != nullptr, "Null eigVecs buffer.");
-    index_type_t iters{};
-    sparse::solver::computeSmallestEigenvectors(handle,
-                                                A,
-                                                config_.n_eigVecs,
-                                                config_.maxIter,
-                                                config_.restartIter,
-                                                config_.tol,
-                                                config_.reorthogonalize,
-                                                iters,
-                                                eigVals,
-                                                eigVecs,
-                                                config_.seed);
 
-    return iters;
+    auto lanczos_config = raft::sparse::solver::lanczos_solver_config<value_type_t>{
+      config_.n_eigVecs, config_.maxIter, config_.restartIter, config_.tol, config_.seed};
+    auto v0_opt = std::optional<raft::device_vector_view<value_type_t, uint32_t, raft::row_major>>{
+      std::nullopt};
+    auto input_structure = input.structure_view();
+
+    auto solver_iterations = sparse::solver::lanczos_compute_smallest_eigenvectors(
+      handle,
+      lanczos_config,
+      input,
+      v0_opt,
+      raft::make_device_vector_view<value_type_t, uint32_t, raft::col_major>(eigVals,
+                                                                             config_.n_eigVecs),
+      raft::make_device_matrix_view<value_type_t, uint32_t, raft::col_major>(
+        eigVecs, input_structure.get_n_rows(), config_.n_eigVecs));
+
+    return index_type_t{solver_iterations};
   }
 
   index_type_t solve_largest_eigenvectors(

cpp/tests/CMakeLists.txt

@@ -259,6 +259,7 @@ if(BUILD_TESTS)
     sparse/csr_transpose.cu
     sparse/degree.cu
     sparse/filter.cu
+    sparse/laplacian.cu
     sparse/masked_matmul.cu
     sparse/norm.cu
     sparse/normalize.cu