Skip to content

Commit 39dd4b7

Browse files
committed
Merge branch 'imp_remove_cuda_11' of github.com:cjnolet/cuvs into imp_remove_cuda_11
2 parents 06e36e7 + 0ed0b70 commit 39dd4b7

File tree

6 files changed

+538
-5
lines changed

6 files changed

+538
-5
lines changed

.github/workflows/pr.yaml

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,11 @@
11
name: pr
2-
32
on:
43
push:
54
branches:
65
- "pull-request/[0-9]+"
7-
86
concurrency:
97
group: ${{ github.workflow }}-${{ github.ref }}
108
cancel-in-progress: true
11-
129
jobs:
1310
pr-builder:
1411
needs:
@@ -232,7 +229,6 @@ jobs:
232229
sccache -z;
233230
build-all --verbose;
234231
sccache -s;
235-
236232
telemetry-summarize:
237233
# This job must use a self-hosted runner to record telemetry traces.
238234
runs-on: linux-amd64-cpu4
@@ -242,3 +238,5 @@ jobs:
242238
steps:
243239
- name: Telemetry summarize
244240
uses: rapidsai/shared-actions/telemetry-dispatch-summarize@main
241+
env:
242+
GH_TOKEN: ${{ github.token }}

cpp/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -490,6 +490,7 @@ if(BUILD_SHARED_LIBS)
490490
src/neighbors/vamana_serialize_int8.cu
491491
src/preprocessing/quantize/scalar.cu
492492
src/preprocessing/quantize/binary.cu
493+
src/preprocessing/spectral/spectral_embedding.cu
493494
src/selection/select_k_float_int64_t.cu
494495
src/selection/select_k_float_int32_t.cu
495496
src/selection/select_k_float_uint32_t.cu
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Copyright (c) 2025, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#pragma once
18+
19+
#include <raft/core/device_coo_matrix.hpp>
20+
#include <raft/core/device_csr_matrix.hpp>
21+
#include <raft/core/device_mdspan.hpp>
22+
#include <raft/core/resources.hpp>
23+
24+
namespace cuvs::preprocessing::spectral_embedding {
25+
26+
/**
27+
* @brief Parameters for spectral embedding algorithm
28+
*/
29+
struct params {
30+
/** @brief The number of components to reduce the data to. */
31+
int n_components;
32+
33+
/** @brief The number of neighbors to use for the nearest neighbors graph. */
34+
int n_neighbors;
35+
36+
/** @brief Whether to normalize the Laplacian matrix. */
37+
bool norm_laplacian;
38+
39+
/** @brief Whether to drop the first eigenvector. */
40+
bool drop_first;
41+
42+
/** @brief Random seed for reproducibility */
43+
uint64_t seed;
44+
};
45+
46+
void transform(raft::resources const& handle,
47+
params config,
48+
raft::device_matrix_view<float, int, raft::row_major> dataset,
49+
raft::device_matrix_view<float, int, raft::col_major> embedding);
50+
51+
} // namespace cuvs::preprocessing::spectral_embedding
Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
1+
/*
2+
* Copyright (c) 2025, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#include <cuvs/neighbors/brute_force.hpp>
18+
#include <cuvs/preprocessing/spectral_embedding.hpp>
19+
20+
#include <raft/core/device_coo_matrix.hpp>
21+
#include <raft/core/device_mdspan.hpp>
22+
#include <raft/core/handle.hpp>
23+
#include <raft/core/resources.hpp>
24+
#include <raft/linalg/matrix_vector_op.cuh>
25+
#include <raft/matrix/gather.cuh>
26+
#include <raft/matrix/init.cuh>
27+
#include <raft/sparse/coo.hpp>
28+
#include <raft/sparse/linalg/laplacian.cuh>
29+
#include <raft/sparse/linalg/symmetrize.cuh>
30+
#include <raft/sparse/op/filter.cuh>
31+
#include <raft/sparse/solver/lanczos.cuh>
32+
#include <raft/sparse/solver/lanczos_types.hpp>
33+
#include <raft/util/cudart_utils.hpp>
34+
#include <raft/util/integer_utils.hpp>
35+
36+
#include <thrust/sequence.h>
37+
#include <thrust/tabulate.h>
38+
39+
namespace cuvs::preprocessing::spectral_embedding {
40+
41+
void create_connectivity_graph(raft::resources const& handle,
42+
params spectral_embedding_config,
43+
raft::device_matrix_view<float, int, raft::row_major> dataset,
44+
raft::device_matrix_view<float, int, raft::col_major> embedding,
45+
raft::device_coo_matrix<float, int, int, int>& connectivity_graph)
46+
{
47+
const int n_samples = dataset.extent(0);
48+
const int n_features = dataset.extent(1);
49+
const int k_search = spectral_embedding_config.n_neighbors;
50+
const size_t nnz = n_samples * k_search;
51+
52+
auto stream = raft::resource::get_cuda_stream(handle);
53+
54+
cuvs::neighbors::brute_force::search_params search_params;
55+
cuvs::neighbors::brute_force::index_params index_params;
56+
index_params.metric = cuvs::distance::DistanceType::L2SqrtExpanded;
57+
58+
auto d_indices = raft::make_device_matrix<int64_t>(handle, n_samples, k_search);
59+
auto d_distances = raft::make_device_matrix<float>(handle, n_samples, k_search);
60+
61+
auto index =
62+
cuvs::neighbors::brute_force::build(handle, index_params, raft::make_const_mdspan(dataset));
63+
64+
cuvs::neighbors::brute_force::search(
65+
handle, search_params, index, dataset, d_indices.view(), d_distances.view());
66+
67+
auto knn_rows = raft::make_device_vector<int>(handle, nnz);
68+
auto knn_cols = raft::make_device_vector<int>(handle, nnz);
69+
70+
raft::linalg::unary_op(
71+
handle, make_const_mdspan(d_indices.view()), knn_cols.view(), [] __device__(int64_t x) {
72+
return static_cast<int>(x);
73+
});
74+
75+
thrust::tabulate(raft::resource::get_thrust_policy(handle),
76+
knn_rows.data_handle(),
77+
knn_rows.data_handle() + nnz,
78+
[k_search] __device__(int idx) { return idx / k_search; });
79+
80+
// set all distances to 1.0f (connectivity KNN graph)
81+
raft::matrix::fill(handle, raft::make_device_vector_view(d_distances.data_handle(), nnz), 1.0f);
82+
83+
auto coo_matrix_view = raft::make_device_coo_matrix_view<const float, int, int, int>(
84+
d_distances.data_handle(),
85+
raft::make_device_coordinate_structure_view<int, int, int>(
86+
knn_rows.data_handle(), knn_cols.data_handle(), n_samples, n_samples, nnz));
87+
88+
auto sym_coo1_matrix =
89+
raft::make_device_coo_matrix<float, int, int, int>(handle, n_samples, n_samples);
90+
raft::sparse::linalg::coo_symmetrize<128, float, int, int>(
91+
handle, coo_matrix_view, sym_coo1_matrix, [] __device__(int row, int col, float a, float b) {
92+
return 0.5f * (a + b);
93+
});
94+
95+
raft::sparse::op::coo_sort<float>(n_samples,
96+
n_samples,
97+
sym_coo1_matrix.structure_view().get_nnz(),
98+
sym_coo1_matrix.structure_view().get_rows().data(),
99+
sym_coo1_matrix.structure_view().get_cols().data(),
100+
sym_coo1_matrix.get_elements().data(),
101+
stream);
102+
103+
raft::sparse::op::coo_remove_scalar<128, float, int, int>(
104+
handle,
105+
raft::make_device_coo_matrix_view<const float, int, int, int>(
106+
sym_coo1_matrix.get_elements().data(), sym_coo1_matrix.structure_view()),
107+
raft::make_host_scalar<float>(0.0f).view(),
108+
connectivity_graph);
109+
}
110+
111+
raft::device_csr_matrix_view<float, int, int, int> coo_to_csr_matrix(
112+
raft::resources const& handle,
113+
const int n_samples,
114+
raft::device_vector_view<int> sym_coo_row_ind,
115+
raft::device_coo_matrix<float, int, int, int>& sym_coo_matrix)
116+
{
117+
auto stream = raft::resource::get_cuda_stream(handle);
118+
119+
raft::sparse::op::coo_sort<float>(n_samples,
120+
n_samples,
121+
sym_coo_matrix.structure_view().get_nnz(),
122+
sym_coo_matrix.structure_view().get_rows().data(),
123+
sym_coo_matrix.structure_view().get_cols().data(),
124+
sym_coo_matrix.get_elements().data(),
125+
stream);
126+
127+
raft::sparse::convert::sorted_coo_to_csr(sym_coo_matrix.structure_view().get_rows().data(),
128+
sym_coo_matrix.structure_view().get_nnz(),
129+
sym_coo_row_ind.data_handle(),
130+
n_samples,
131+
stream);
132+
133+
auto sym_coo_nnz = sym_coo_matrix.structure_view().get_nnz();
134+
raft::copy(sym_coo_row_ind.data_handle() + sym_coo_row_ind.size() - 1, &sym_coo_nnz, 1, stream);
135+
136+
auto csr_matrix_view = raft::make_device_csr_matrix_view<float, int, int, int>(
137+
const_cast<float*>(sym_coo_matrix.get_elements().data()),
138+
raft::make_device_compressed_structure_view<int, int, int>(
139+
const_cast<int*>(sym_coo_row_ind.data_handle()),
140+
const_cast<int*>(sym_coo_matrix.structure_view().get_cols().data()),
141+
n_samples,
142+
n_samples,
143+
sym_coo_matrix.structure_view().get_nnz()));
144+
return csr_matrix_view;
145+
}
146+
147+
raft::device_csr_matrix<float, int, int, int> create_laplacian(
148+
raft::resources const& handle,
149+
params spectral_embedding_config,
150+
raft::device_csr_matrix_view<float, int, int, int> csr_matrix_view,
151+
raft::device_vector_view<float, int> diagonal)
152+
{
153+
auto laplacian = spectral_embedding_config.norm_laplacian
154+
? raft::sparse::linalg::laplacian_normalized(handle, csr_matrix_view, diagonal)
155+
: raft::sparse::linalg::compute_graph_laplacian(handle, csr_matrix_view);
156+
157+
auto laplacian_elements_view = raft::make_device_vector_view<float, int>(
158+
laplacian.get_elements().data(), laplacian.structure_view().get_nnz());
159+
160+
raft::linalg::unary_op(handle,
161+
raft::make_const_mdspan(laplacian_elements_view),
162+
laplacian_elements_view,
163+
[] __device__(float x) { return -x; });
164+
165+
return laplacian;
166+
}
167+
168+
void compute_eigenpairs(raft::resources const& handle,
169+
params spectral_embedding_config,
170+
const int n_samples,
171+
raft::device_csr_matrix<float, int, int, int> laplacian,
172+
raft::device_vector_view<float, int> diagonal,
173+
raft::device_matrix_view<float, int, raft::col_major> embedding)
174+
{
175+
auto config = raft::sparse::solver::lanczos_solver_config<float>();
176+
config.n_components = spectral_embedding_config.n_components;
177+
config.max_iterations = 1000;
178+
config.ncv = std::min(n_samples, std::max(2 * config.n_components + 1, 20));
179+
config.tolerance = 1e-5;
180+
config.which = raft::sparse::solver::LANCZOS_WHICH::LA;
181+
config.seed = spectral_embedding_config.seed;
182+
183+
auto eigenvalues =
184+
raft::make_device_vector<float, int, raft::col_major>(handle, config.n_components);
185+
auto eigenvectors =
186+
raft::make_device_matrix<float, int, raft::col_major>(handle, n_samples, config.n_components);
187+
188+
raft::sparse::solver::lanczos_compute_smallest_eigenvectors<int, float>(
189+
handle,
190+
config,
191+
raft::make_device_csr_matrix_view<float, int, int, int>(laplacian.get_elements().data(),
192+
laplacian.structure_view()),
193+
std::nullopt,
194+
eigenvalues.view(),
195+
eigenvectors.view());
196+
197+
if (spectral_embedding_config.norm_laplacian) {
198+
raft::linalg::matrix_vector_op<raft::Apply::ALONG_COLUMNS>(
199+
handle,
200+
raft::make_const_mdspan(eigenvectors.view()), // input matrix view
201+
raft::make_const_mdspan(diagonal), // input vector view
202+
eigenvectors.view(), // output matrix view (in-place)
203+
[] __device__(float elem, float diag) { return elem / diag; });
204+
}
205+
206+
// Create a sequence of reversed column indices
207+
config.n_components =
208+
spectral_embedding_config.drop_first ? config.n_components - 1 : config.n_components;
209+
auto col_indices = raft::make_device_vector<int>(handle, config.n_components);
210+
211+
// TODO: https://github.com/rapidsai/raft/issues/2661
212+
thrust::sequence(thrust::device,
213+
col_indices.data_handle(),
214+
col_indices.data_handle() + config.n_components,
215+
config.n_components - 1, // Start from the last column index
216+
-1 // Decrement (move backward)
217+
);
218+
219+
// Create row-major views of the column-major matrices
220+
// This is just a view re-interpretation, no data movement
221+
auto eigenvectors_row_view = raft::make_device_matrix_view<float, int, raft::row_major>(
222+
eigenvectors.data_handle(),
223+
eigenvectors.extent(1), // Swap dimensions for the view
224+
eigenvectors.extent(0));
225+
226+
auto embedding_row_view = raft::make_device_matrix_view<float, int, raft::row_major>(
227+
embedding.data_handle(),
228+
embedding.extent(1), // Swap dimensions for the view
229+
embedding.extent(0));
230+
231+
raft::matrix::gather<float, int, int>(
232+
handle,
233+
raft::make_const_mdspan(eigenvectors_row_view), // Source matrix (as row-major view)
234+
raft::make_const_mdspan(col_indices.view()), // Column indices to gather
235+
embedding_row_view // Destination matrix (as row-major view)
236+
);
237+
}
238+
239+
void transform(raft::resources const& handle,
240+
params spectral_embedding_config,
241+
raft::device_matrix_view<float, int, raft::row_major> dataset,
242+
raft::device_matrix_view<float, int, raft::col_major> embedding)
243+
{
244+
const int n_samples = dataset.extent(0);
245+
246+
auto sym_coo_matrix =
247+
raft::make_device_coo_matrix<float, int, int, int>(handle, n_samples, n_samples);
248+
auto sym_coo_row_ind = raft::make_device_vector<int>(handle, n_samples + 1);
249+
auto diagonal = raft::make_device_vector<float, int>(handle, n_samples);
250+
251+
create_connectivity_graph(handle, spectral_embedding_config, dataset, embedding, sym_coo_matrix);
252+
auto csr_matrix_view =
253+
coo_to_csr_matrix(handle, n_samples, sym_coo_row_ind.view(), sym_coo_matrix);
254+
auto laplacian =
255+
create_laplacian(handle, spectral_embedding_config, csr_matrix_view, diagonal.view());
256+
compute_eigenpairs(
257+
handle, spectral_embedding_config, n_samples, laplacian, diagonal.view(), embedding);
258+
}
259+
260+
} // namespace cuvs::preprocessing::spectral_embedding

cpp/tests/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,7 @@ if(BUILD_TESTS)
273273

274274
ConfigureTest(
275275
NAME PREPROCESSING_TEST PATH preprocessing/scalar_quantization.cu
276-
preprocessing/binary_quantization.cu GPUS 1 PERCENT 100
276+
preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu GPUS 1 PERCENT 100
277277
)
278278

279279
ConfigureTest(

0 commit comments

Comments
 (0)