Skip to content

Commit c36d890

Browse files
authored
Native row-major PCA (#2125)
Currently PCA only supports col-major in the cpp layer. If a consumer's input dataset is in row-major it must be copied as col-major. With the introduction of RAFT native row-major support in PCA NVIDIA/raft#3036, users can seamlessly pass in row-major datasets. This PR exposes the cpp PCA row-major APIs and updates the C and Python APIs accordingly. Resolves #2115. Depends on NVIDIA/raft#3036. Authors: - Anupam (https://github.com/aamijar) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #2125
1 parent 4d901dd commit c36d890

9 files changed

Lines changed: 624 additions & 211 deletions

File tree

c/include/cuvs/preprocessing/pca.h

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -85,6 +85,9 @@ CUVS_EXPORT cuvsError_t cuvsPcaParamsDestroy(cuvsPcaParams_t params);
8585
* Computes the principal components, explained variances, singular values, and column means
8686
* from the input data.
8787
*
88+
* The layout of `input` (C-contiguous / row-major or F-contiguous / col-major) is detected
89+
* from its DLPack strides; `components` must use the same layout as `input`.
90+
*
8891
* @code {.c}
8992
* #include <cuvs/core/c_api.h>
9093
* #include <cuvs/preprocessing/pca.h>
@@ -98,9 +101,9 @@ CUVS_EXPORT cuvsError_t cuvsPcaParamsDestroy(cuvsPcaParams_t params);
98101
* cuvsPcaParamsCreate(&params);
99102
* params->n_components = 2;
100103
*
101-
* // Assume populated DLManagedTensor objects (col-major, float32, device memory)
102-
* DLManagedTensor input; // [n_rows x n_cols]
103-
* DLManagedTensor components; // [n_components x n_cols]
104+
* // Assume populated DLManagedTensor objects (float32, device memory)
105+
* DLManagedTensor input; // [n_rows x n_cols] (C- or F-contiguous)
106+
* DLManagedTensor components; // [n_components x n_cols] (same layout as input)
104107
* DLManagedTensor explained_var; // [n_components]
105108
* DLManagedTensor explained_var_ratio; // [n_components]
106109
* DLManagedTensor singular_vals; // [n_components]
@@ -117,8 +120,8 @@ CUVS_EXPORT cuvsError_t cuvsPcaParamsDestroy(cuvsPcaParams_t params);
117120
*
118121
* @param[in] res cuvsResources_t opaque C handle
119122
* @param[in] params PCA parameters
120-
* @param[inout] input input data [n_rows x n_cols] (col-major, float32, device)
121-
* @param[out] components principal components [n_components x n_cols] (col-major, float32, device)
123+
* @param[inout] input input data [n_rows x n_cols] (C- or F-contiguous, float32, device)
124+
* @param[out] components principal components [n_components x n_cols] (same layout as input)
122125
* @param[out] explained_var explained variances [n_components] (float32, device)
123126
* @param[out] explained_var_ratio explained variance ratios [n_components] (float32, device)
124127
* @param[out] singular_vals singular values [n_components] (float32, device)
@@ -142,12 +145,14 @@ CUVS_EXPORT cuvsError_t cuvsPcaFit(cuvsResources_t res,
142145
* @brief Perform PCA fit and transform in a single operation.
143146
*
144147
* Computes the principal components and transforms the input data into the eigenspace.
148+
* The layout of `input` (C- or F-contiguous) is detected from its DLPack strides; all
149+
* other matrix tensors must use the same layout.
145150
*
146151
* @param[in] res cuvsResources_t opaque C handle
147152
* @param[in] params PCA parameters
148-
* @param[inout] input input data [n_rows x n_cols] (col-major, float32, device)
149-
* @param[out] trans_input transformed data [n_rows x n_components] (col-major, float32, device)
150-
* @param[out] components principal components [n_components x n_cols] (col-major, float32, device)
153+
* @param[inout] input input data [n_rows x n_cols] (C- or F-contiguous, float32, device)
154+
* @param[out] trans_input transformed data [n_rows x n_components] (same layout as input)
155+
* @param[out] components principal components [n_components x n_cols] (same layout as input)
151156
* @param[out] explained_var explained variances [n_components] (float32, device)
152157
* @param[out] explained_var_ratio explained variance ratios [n_components] (float32, device)
153158
* @param[out] singular_vals singular values [n_components] (float32, device)
@@ -172,14 +177,16 @@ CUVS_EXPORT cuvsError_t cuvsPcaFitTransform(cuvsResources_t res,
172177
* @brief Perform PCA transform operation.
173178
*
174179
* Transforms the input data into the eigenspace using previously computed principal components.
180+
* The layout of `input` (C- or F-contiguous) is detected from its DLPack strides; all other
181+
* matrix tensors must use the same layout.
175182
*
176183
* @param[in] res cuvsResources_t opaque C handle
177184
* @param[in] params PCA parameters
178-
* @param[inout] input data to transform [n_rows x n_cols] (col-major, float32, device)
179-
* @param[in] components principal components [n_components x n_cols] (col-major, float32, device)
185+
* @param[inout] input data to transform [n_rows x n_cols] (C- or F-contiguous, float32, device)
186+
* @param[in] components principal components [n_components x n_cols] (same layout as input)
180187
* @param[in] singular_vals singular values [n_components] (float32, device)
181188
* @param[in] mu column means [n_cols] (float32, device)
182-
* @param[out] trans_input transformed data [n_rows x n_components] (col-major, float32, device)
189+
* @param[out] trans_input transformed data [n_rows x n_components] (same layout as input)
183190
* @return cuvsError_t
184191
*/
185192
CUVS_EXPORT cuvsError_t cuvsPcaTransform(cuvsResources_t res,
@@ -194,14 +201,17 @@ CUVS_EXPORT cuvsError_t cuvsPcaTransform(cuvsResources_t res,
194201
* @brief Perform PCA inverse transform operation.
195202
*
196203
* Transforms data from the eigenspace back to the original space.
204+
* The layout of `trans_input` (C- or F-contiguous) is detected from its DLPack strides;
205+
* all other matrix tensors must use the same layout.
197206
*
198207
* @param[in] res cuvsResources_t opaque C handle
199208
* @param[in] params PCA parameters
200-
* @param[in] trans_input transformed data [n_rows x n_components] (col-major, float32, device)
201-
* @param[in] components principal components [n_components x n_cols] (col-major, float32, device)
209+
* @param[in] trans_input transformed data [n_rows x n_components] (C- or F-contiguous,
210+
* float32, device)
211+
* @param[in] components principal components [n_components x n_cols] (same layout as trans_input)
202212
* @param[in] singular_vals singular values [n_components] (float32, device)
203213
* @param[in] mu column means [n_cols] (float32, device)
204-
* @param[out] output reconstructed data [n_rows x n_cols] (col-major, float32, device)
214+
* @param[out] output reconstructed data [n_rows x n_cols] (same layout as trans_input)
205215
* @return cuvsError_t
206216
*/
207217
CUVS_EXPORT cuvsError_t cuvsPcaInverseTransform(cuvsResources_t res,

c/src/preprocessing/pca.cpp

Lines changed: 84 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -40,6 +40,7 @@ cuvs::preprocessing::pca::params to_cpp_params(const cuvsPcaParams& c_params)
4040
return cpp_params;
4141
}
4242

43+
template <typename LayoutT>
4344
void _fit(cuvsResources_t res,
4445
const cuvsPcaParams& params,
4546
DLManagedTensor* input_tensor,
@@ -54,12 +55,12 @@ void _fit(cuvsResources_t res,
5455
auto res_ptr = reinterpret_cast<raft::resources*>(res);
5556
auto cpp_params = to_cpp_params(params);
5657

57-
using matrix_type = raft::device_matrix_view<float, int64_t, raft::col_major>;
58+
using matrix_type = raft::device_matrix_view<float, int64_t, LayoutT>;
5859
using vector_type = raft::device_vector_view<float, int64_t>;
5960
using scalar_type = raft::device_scalar_view<float, int64_t>;
6061

61-
auto input = cuvs::core::from_dlpack<matrix_type>(input_tensor);
62-
auto components = cuvs::core::from_dlpack<matrix_type>(components_tensor);
62+
auto input = cuvs::core::from_dlpack<matrix_type>(input_tensor);
63+
auto components = cuvs::core::from_dlpack<matrix_type>(components_tensor);
6364
auto explained_var = cuvs::core::from_dlpack<vector_type>(explained_var_tensor);
6465
auto explained_var_ratio = cuvs::core::from_dlpack<vector_type>(explained_var_ratio_tensor);
6566
auto singular_vals = cuvs::core::from_dlpack<vector_type>(singular_vals_tensor);
@@ -78,6 +79,7 @@ void _fit(cuvsResources_t res,
7879
flip_signs_based_on_U);
7980
}
8081

82+
template <typename LayoutT>
8183
void _fit_transform(cuvsResources_t res,
8284
const cuvsPcaParams& params,
8385
DLManagedTensor* input_tensor,
@@ -93,13 +95,13 @@ void _fit_transform(cuvsResources_t res,
9395
auto res_ptr = reinterpret_cast<raft::resources*>(res);
9496
auto cpp_params = to_cpp_params(params);
9597

96-
using matrix_type = raft::device_matrix_view<float, int64_t, raft::col_major>;
98+
using matrix_type = raft::device_matrix_view<float, int64_t, LayoutT>;
9799
using vector_type = raft::device_vector_view<float, int64_t>;
98100
using scalar_type = raft::device_scalar_view<float, int64_t>;
99101

100-
auto input = cuvs::core::from_dlpack<matrix_type>(input_tensor);
101-
auto trans_input = cuvs::core::from_dlpack<matrix_type>(trans_input_tensor);
102-
auto components = cuvs::core::from_dlpack<matrix_type>(components_tensor);
102+
auto input = cuvs::core::from_dlpack<matrix_type>(input_tensor);
103+
auto trans_input = cuvs::core::from_dlpack<matrix_type>(trans_input_tensor);
104+
auto components = cuvs::core::from_dlpack<matrix_type>(components_tensor);
103105
auto explained_var = cuvs::core::from_dlpack<vector_type>(explained_var_tensor);
104106
auto explained_var_ratio = cuvs::core::from_dlpack<vector_type>(explained_var_ratio_tensor);
105107
auto singular_vals = cuvs::core::from_dlpack<vector_type>(singular_vals_tensor);
@@ -119,6 +121,7 @@ void _fit_transform(cuvsResources_t res,
119121
flip_signs_based_on_U);
120122
}
121123

124+
template <typename LayoutT>
122125
void _transform(cuvsResources_t res,
123126
const cuvsPcaParams& params,
124127
DLManagedTensor* input_tensor,
@@ -130,7 +133,7 @@ void _transform(cuvsResources_t res,
130133
auto res_ptr = reinterpret_cast<raft::resources*>(res);
131134
auto cpp_params = to_cpp_params(params);
132135

133-
using matrix_type = raft::device_matrix_view<float, int64_t, raft::col_major>;
136+
using matrix_type = raft::device_matrix_view<float, int64_t, LayoutT>;
134137
using vector_type = raft::device_vector_view<float, int64_t>;
135138

136139
auto input = cuvs::core::from_dlpack<matrix_type>(input_tensor);
@@ -143,6 +146,7 @@ void _transform(cuvsResources_t res,
143146
*res_ptr, cpp_params, input, components, singular_vals, mu, trans_input);
144147
}
145148

149+
template <typename LayoutT>
146150
void _inverse_transform(cuvsResources_t res,
147151
const cuvsPcaParams& params,
148152
DLManagedTensor* trans_input_tensor,
@@ -154,7 +158,7 @@ void _inverse_transform(cuvsResources_t res,
154158
auto res_ptr = reinterpret_cast<raft::resources*>(res);
155159
auto cpp_params = to_cpp_params(params);
156160

157-
using matrix_type = raft::device_matrix_view<float, int64_t, raft::col_major>;
161+
using matrix_type = raft::device_matrix_view<float, int64_t, LayoutT>;
158162
using vector_type = raft::device_vector_view<float, int64_t>;
159163

160164
auto trans_input = cuvs::core::from_dlpack<matrix_type>(trans_input_tensor);
@@ -205,19 +209,32 @@ extern "C" cuvsError_t cuvsPcaFit(cuvsResources_t res,
205209
"PCA input must be float32 (kDLFloat, 32 bits)");
206210
RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(input->dl_tensor),
207211
"PCA input must be device-accessible memory");
208-
RAFT_EXPECTS(cuvs::core::is_f_contiguous(input),
209-
"PCA input must be col-major (Fortran-contiguous)");
210-
211-
_fit(res,
212-
*params,
213-
input,
214-
components,
215-
explained_var,
216-
explained_var_ratio,
217-
singular_vals,
218-
mu,
219-
noise_vars,
220-
flip_signs_based_on_U);
212+
213+
if (cuvs::core::is_f_contiguous(input)) {
214+
_fit<raft::col_major>(res,
215+
*params,
216+
input,
217+
components,
218+
explained_var,
219+
explained_var_ratio,
220+
singular_vals,
221+
mu,
222+
noise_vars,
223+
flip_signs_based_on_U);
224+
} else if (cuvs::core::is_c_contiguous(input)) {
225+
_fit<raft::row_major>(res,
226+
*params,
227+
input,
228+
components,
229+
explained_var,
230+
explained_var_ratio,
231+
singular_vals,
232+
mu,
233+
noise_vars,
234+
flip_signs_based_on_U);
235+
} else {
236+
RAFT_FAIL("PCA input must be contiguous (C- or F-order)");
237+
}
221238
});
222239
}
223240

@@ -239,20 +256,34 @@ extern "C" cuvsError_t cuvsPcaFitTransform(cuvsResources_t res,
239256
"PCA input must be float32 (kDLFloat, 32 bits)");
240257
RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(input->dl_tensor),
241258
"PCA input must be device-accessible memory");
242-
RAFT_EXPECTS(cuvs::core::is_f_contiguous(input),
243-
"PCA input must be col-major (Fortran-contiguous)");
244-
245-
_fit_transform(res,
246-
*params,
247-
input,
248-
trans_input,
249-
components,
250-
explained_var,
251-
explained_var_ratio,
252-
singular_vals,
253-
mu,
254-
noise_vars,
255-
flip_signs_based_on_U);
259+
260+
if (cuvs::core::is_f_contiguous(input)) {
261+
_fit_transform<raft::col_major>(res,
262+
*params,
263+
input,
264+
trans_input,
265+
components,
266+
explained_var,
267+
explained_var_ratio,
268+
singular_vals,
269+
mu,
270+
noise_vars,
271+
flip_signs_based_on_U);
272+
} else if (cuvs::core::is_c_contiguous(input)) {
273+
_fit_transform<raft::row_major>(res,
274+
*params,
275+
input,
276+
trans_input,
277+
components,
278+
explained_var,
279+
explained_var_ratio,
280+
singular_vals,
281+
mu,
282+
noise_vars,
283+
flip_signs_based_on_U);
284+
} else {
285+
RAFT_FAIL("PCA input must be contiguous (C- or F-order)");
286+
}
256287
});
257288
}
258289

@@ -270,10 +301,14 @@ extern "C" cuvsError_t cuvsPcaTransform(cuvsResources_t res,
270301
"PCA input must be float32 (kDLFloat, 32 bits)");
271302
RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(input->dl_tensor),
272303
"PCA input must be device-accessible memory");
273-
RAFT_EXPECTS(cuvs::core::is_f_contiguous(input),
274-
"PCA input must be col-major (Fortran-contiguous)");
275304

276-
_transform(res, *params, input, components, singular_vals, mu, trans_input);
305+
if (cuvs::core::is_f_contiguous(input)) {
306+
_transform<raft::col_major>(res, *params, input, components, singular_vals, mu, trans_input);
307+
} else if (cuvs::core::is_c_contiguous(input)) {
308+
_transform<raft::row_major>(res, *params, input, components, singular_vals, mu, trans_input);
309+
} else {
310+
RAFT_FAIL("PCA input must be contiguous (C- or F-order)");
311+
}
277312
});
278313
}
279314

@@ -291,9 +326,15 @@ extern "C" cuvsError_t cuvsPcaInverseTransform(cuvsResources_t res,
291326
"PCA trans_input must be float32 (kDLFloat, 32 bits)");
292327
RAFT_EXPECTS(cuvs::core::is_dlpack_device_compatible(trans_input->dl_tensor),
293328
"PCA trans_input must be device-accessible memory");
294-
RAFT_EXPECTS(cuvs::core::is_f_contiguous(trans_input),
295-
"PCA trans_input must be col-major (Fortran-contiguous)");
296329

297-
_inverse_transform(res, *params, trans_input, components, singular_vals, mu, output);
330+
if (cuvs::core::is_f_contiguous(trans_input)) {
331+
_inverse_transform<raft::col_major>(
332+
res, *params, trans_input, components, singular_vals, mu, output);
333+
} else if (cuvs::core::is_c_contiguous(trans_input)) {
334+
_inverse_transform<raft::row_major>(
335+
res, *params, trans_input, components, singular_vals, mu, output);
336+
} else {
337+
RAFT_FAIL("PCA trans_input must be contiguous (C- or F-order)");
338+
}
298339
});
299340
}

cpp/cmake/thirdparty/get_raft.cmake

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# =============================================================================
22
# cmake-format: off
3-
# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
3+
# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
44
# SPDX-License-Identifier: Apache-2.0
55
# cmake-format: on
66

0 commit comments

Comments
 (0)