Skip to content

Cleanup UMAP API #6462

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: branch-25.06
Choose a base branch
from
Draft

Cleanup UMAP API #6462

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions cpp/bench/sg/umap.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2024, NVIDIA CORPORATION.
* Copyright (c) 2020-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.
Expand Down Expand Up @@ -167,7 +167,6 @@ class UmapTransform : public UmapBase {
this->params.nrows,
this->params.ncols,
this->data.X.data(),
this->params.nrows,
embeddings,
this->params.nrows,
&uParams,
Expand Down
42 changes: 19 additions & 23 deletions cpp/include/cuml/manifold/umap.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2024, NVIDIA CORPORATION.
* Copyright (c) 2019-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.
Expand Down Expand Up @@ -168,21 +168,19 @@ void fit_sparse(const raft::handle_t& handle,
* @param[in] X: pointer to input array to be inferred
* @param[in] n: n_samples of input array to be inferred
* @param[in] d: n_features of input array to be inferred
* @param[in] orig_X: pointer to original training array
* @param[in] orig_n: number of rows in original training array
* @param[in] embedding: pointer to embedding created during training
* @param[in] embedding_n: number of rows in embedding created during training
* @param[in] trained_X: pointer to original training array
* @param[in] trained_embeddings: pointer to embedding created during training
* @param[in] n_trained: number of rows in original training input
* @param[in] params: pointer to ML::UMAPParams object
* @param[out] transformed: pointer to embedding produced through projection
*/
void transform(const raft::handle_t& handle,
float* X,
int n,
int d,
float* orig_X,
int orig_n,
float* embedding,
int embedding_n,
float* trained_X,
float* trained_embeddings,
int n_trained,
UMAPParams* params,
float* transformed);

Expand All @@ -196,13 +194,12 @@ void transform(const raft::handle_t& handle,
* @param[in] nnz: number of stored values of input array to be inferred
* @param[in] n: n_samples of input array
* @param[in] d: n_features of input array
* @param[in] orig_x_indptr: pointer to index pointer array of original training array
* @param[in] orig_x_indices: pointer to index array of original training array
* @param[in] orig_x_data: pointer to data array of original training array
* @param[in] orig_nnz: number of stored values of original training array
* @param[in] orig_n: number of rows in original training array
* @param[in] embedding: pointer to embedding created during training
* @param[in] embedding_n: number of rows in embedding created during training
* @param[in] trained_indptr: pointer to index pointer array of original training array
* @param[in] trained_indices: pointer to index array of original training array
* @param[in] trained_data: pointer to data array of original training array
* @param[in] trained_nnz: number of stored values of original training array
* @param[in] trained_embedding: pointer to embedding created during training
* @param[in] n_trained: number of rows in original training input
* @param[in] params: pointer to ML::UMAPParams object
* @param[out] transformed: pointer to embedding produced through projection
*/
Expand All @@ -213,13 +210,12 @@ void transform_sparse(const raft::handle_t& handle,
size_t nnz,
int n,
int d,
int* orig_x_indptr,
int* orig_x_indices,
float* orig_x_data,
size_t orig_nnz,
int orig_n,
float* embedding,
int embedding_n,
int* trained_indptr,
int* trained_indices,
float* trained_data,
size_t trained_nnz,
float* trained_embedding,
int n_trained,
UMAPParams* params,
float* transformed);

Expand Down
15 changes: 6 additions & 9 deletions cpp/src/umap/runner.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,6 @@ CUML_KERNEL void init_transform(int* indices,
T* weights,
int n,
const T* embeddings,
int embeddings_n,
int n_components,
T* result,
int n_neighbors)
Expand Down Expand Up @@ -302,9 +301,8 @@ void _fit_supervised(const raft::handle_t& handle,
template <typename value_idx, typename value_t, typename umap_inputs, typename nnz_t, int TPB_X>
void _transform(const raft::handle_t& handle,
const umap_inputs& inputs,
umap_inputs& orig_x_inputs,
value_t* embedding,
int embedding_n,
umap_inputs& fitted_inputs,
value_t* fitted_embeddings,
UMAPParams* params,
value_t* transformed)
{
Expand Down Expand Up @@ -343,7 +341,7 @@ void _transform(const raft::handle_t& handle,
}

kNNGraph::run<value_idx, value_t, umap_inputs>(
handle, orig_x_inputs, inputs, knn_graph, k, params, stream);
handle, fitted_inputs, inputs, knn_graph, k, params, stream);

raft::common::nvtx::pop_range();

Expand Down Expand Up @@ -424,8 +422,7 @@ void _transform(const raft::handle_t& handle,
init_transform<TPB_X, value_t><<<grid_n, blk, 0, stream>>>(graph_coo.cols(),
vals_normed.data(),
graph_coo.n_rows,
embedding,
embedding_n,
fitted_embeddings,
params->n_components,
transformed,
params->n_neighbors);
Expand Down Expand Up @@ -490,8 +487,8 @@ void _transform(const raft::handle_t& handle,

SimplSetEmbedImpl::optimize_layout<value_t, nnz_t, TPB_X>(transformed,
inputs.n,
embedding,
embedding_n,
fitted_embeddings,
fitted_inputs.n,
comp_coo.rows(),
comp_coo.cols(),
comp_coo.nnz,
Expand Down
50 changes: 23 additions & 27 deletions cpp/src/umap/umap.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,18 +134,17 @@ void transform(const raft::handle_t& handle,
float* X,
int n,
int d,
float* orig_X,
int orig_n,
float* embedding,
int embedding_n,
float* trained_X,
float* trained_embeddings,
int n_trained,
UMAPParams* params,
float* transformed)
{
if (dispatch_to_uint64_t(n, params->n_neighbors, params->n_components))
_transform<uint64_t>(
handle, X, n, d, orig_X, orig_n, embedding, embedding_n, params, transformed);
handle, X, n, d, trained_X, trained_embeddings, n_trained, params, transformed);
else
_transform<int>(handle, X, n, d, orig_X, orig_n, embedding, embedding_n, params, transformed);
_transform<int>(handle, X, n, d, trained_X, trained_embeddings, n_trained, params, transformed);
}

void transform_sparse(const raft::handle_t& handle,
Expand All @@ -155,13 +154,12 @@ void transform_sparse(const raft::handle_t& handle,
size_t nnz,
int n,
int d,
int* orig_x_indptr,
int* orig_x_indices,
float* orig_x_data,
size_t orig_nnz,
int orig_n,
float* embedding,
int embedding_n,
int* trained_indptr,
int* trained_indices,
float* trained_data,
size_t trained_nnz,
float* trained_embedding,
int n_trained,
UMAPParams* params,
float* transformed)
{
Expand All @@ -173,13 +171,12 @@ void transform_sparse(const raft::handle_t& handle,
nnz,
n,
d,
orig_x_indptr,
orig_x_indices,
orig_x_data,
orig_nnz,
orig_n,
embedding,
embedding_n,
trained_indptr,
trained_indices,
trained_data,
trained_nnz,
trained_embedding,
n_trained,
params,
transformed);
else
Expand All @@ -190,13 +187,12 @@ void transform_sparse(const raft::handle_t& handle,
nnz,
n,
d,
orig_x_indptr,
orig_x_indices,
orig_x_data,
orig_nnz,
orig_n,
embedding,
embedding_n,
trained_indptr,
trained_indices,
trained_data,
trained_nnz,
trained_embedding,
n_trained,
params,
transformed);
}
Expand Down
30 changes: 14 additions & 16 deletions cpp/src/umap/umap.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,19 +219,18 @@ inline void _transform(const raft::handle_t& handle,
float* X,
int n,
int d,
float* orig_X,
int orig_n,
float* embedding,
int embedding_n,
float* trained_X,
float* trained_embeddings,
int n_trained,
UMAPParams* params,
float* transformed)
{
RAFT_EXPECTS(params->build_algo == ML::UMAPParams::graph_build_algo::BRUTE_FORCE_KNN,
"build algo nn_descent not supported for transform()");
manifold_dense_inputs_t<float> inputs(X, nullptr, n, d);
manifold_dense_inputs_t<float> orig_inputs(orig_X, nullptr, orig_n, d);
manifold_dense_inputs_t<float> trained_inputs(trained_X, nullptr, n_trained, d);
UMAPAlgo::_transform<knn_indices_dense_t, float, manifold_dense_inputs_t<float>, nnz_t, TPB_X>(
handle, inputs, orig_inputs, embedding, embedding_n, params, transformed);
handle, inputs, trained_inputs, trained_embeddings, params, transformed);
}

template <typename nnz_t>
Expand All @@ -242,26 +241,25 @@ inline void _transform_sparse(const raft::handle_t& handle,
size_t nnz,
int n,
int d,
int* orig_x_indptr,
int* orig_x_indices,
float* orig_x_data,
size_t orig_nnz,
int orig_n,
float* embedding,
int embedding_n,
int* trained_indptr,
int* trained_indices,
float* trained_data,
size_t trained_nnz,
float* trained_embedding,
int n_trained,
UMAPParams* params,
float* transformed)
{
RAFT_EXPECTS(params->build_algo == ML::UMAPParams::graph_build_algo::BRUTE_FORCE_KNN,
"build algo nn_descent not supported for transform()");
manifold_sparse_inputs_t<knn_indices_sparse_t, float> inputs(
indptr, indices, data, nullptr, nnz, n, d);
manifold_sparse_inputs_t<knn_indices_sparse_t, float> orig_x_inputs(
orig_x_indptr, orig_x_indices, orig_x_data, nullptr, orig_nnz, orig_n, d);
manifold_sparse_inputs_t<knn_indices_sparse_t, float> trained_inputs(
trained_indptr, trained_indices, trained_data, nullptr, trained_nnz, n_trained, d);

UMAPAlgo::
_transform<knn_indices_sparse_t, float, manifold_sparse_inputs_t<int, float>, nnz_t, TPB_X>(
handle, inputs, orig_x_inputs, embedding, embedding_n, params, transformed);
handle, inputs, trained_inputs, trained_embedding, params, transformed);
}

} // namespace UMAP
Expand Down
3 changes: 1 addition & 2 deletions cpp/tests/sg/umap_parametrizable_test.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020-2024, NVIDIA CORPORATION.
* Copyright (c) 2020-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.
Expand Down Expand Up @@ -209,7 +209,6 @@ class UMAPParametrizableTest : public ::testing::Test {
n_samples,
umap_params.n_components,
X,
n_samples,
model_embedding,
n_samples,
&umap_params,
Expand Down
22 changes: 9 additions & 13 deletions python/cuml/cuml/manifold/umap.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,10 +108,9 @@ IF GPUBUILD == 1:
float * X,
int n,
int d,
float * orig_X,
int orig_n,
float * trained_X,
float * embedding,
int embedding_n,
int n_trained,
UMAPParams * params,
float * out) except +

Expand All @@ -122,13 +121,12 @@ IF GPUBUILD == 1:
size_t nnz,
int n,
int d,
int *orig_x_indptr,
int *orig_x_indices,
float *orig_x_data,
size_t orig_nnz,
int orig_n,
float *embedding,
int embedding_n,
int *trained_x_indptr,
int *trained_x_indices,
float *trained_x_data,
size_t trained_nnz,
float *trained_embedding,
int n_trained,
UMAPParams *params,
float *transformed) except +

Expand Down Expand Up @@ -832,7 +830,6 @@ class UMAP(UniversalBase,
<int*><uintptr_t> self._raw_data.indices.ptr,
<float*><uintptr_t> self._raw_data.data.ptr,
<size_t> self._raw_data.nnz,
<int> self._raw_data.shape[0],
<float*> _embed_ptr,
<int> self._raw_data.shape[0],
<UMAPParams*> umap_params,
Expand All @@ -843,9 +840,8 @@ class UMAP(UniversalBase,
<int> n_rows,
<int> n_cols,
<float*><uintptr_t>self._raw_data.ptr,
<int> self._raw_data.shape[0],
<float*> _embed_ptr,
<int> n_rows,
<int> self._raw_data.shape[0],
<UMAPParams*> umap_params,
<float*> _xformed_ptr)
self.handle.sync()
Expand Down