Native row-major PCA (#3036)

Currently PCA only supports col-major which is not ideal for consumers passing in row-major datasets.
This PR exposes row-major APIs for PCA and calls the corresponding row-major raft primitives based on templated row or col layout.

This PR is backward compatible with cuml and cuvs.

Authors:
  - Anupam (https://github.com/aamijar)
  - Divye Gala (https://github.com/divyegala)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - Micka (https://github.com/lowener)

URL: #3036

Author: aamijar

cpp/include/raft/linalg/detail/pca.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -15,7 +15,9 @@
 #include <raft/linalg/eig.cuh>
 #include <raft/linalg/eltwise.cuh>
 #include <raft/linalg/gemm.cuh>
+#include <raft/linalg/map.cuh>
 #include <raft/linalg/pca_types.hpp>
+#include <raft/linalg/reduce.cuh>
 #include <raft/linalg/rsvd.cuh>
 #include <raft/linalg/transpose.cuh>
 #include <raft/matrix/copy.cuh>
@@ -104,13 +106,19 @@ void cal_comp_exp_vars_svd(raft::resources const& handle,
     handle, explained_vars.data_handle(), explained_var_ratio.data_handle(), n_components, stream);
 }
 
-template <typename math_t, typename idx_t>
+template <typename math_t, typename idx_t, typename LayoutPolicy>
 void cal_eig(raft::resources const& handle,
              const paramsTSVD& prms,
-             raft::device_matrix_view<math_t, idx_t, raft::col_major> in,
-             raft::device_matrix_view<math_t, idx_t, raft::col_major> components,
+             raft::device_matrix_view<math_t, idx_t, LayoutPolicy> in,
+             raft::device_matrix_view<math_t, idx_t, LayoutPolicy> components,
              raft::device_vector_view<math_t, idx_t> explained_var)
 {
+  static_assert(
+    std::is_same_v<LayoutPolicy, raft::row_major> || std::is_same_v<LayoutPolicy, raft::col_major>,
+    "cal_eig: layout must be raft::row_major or raft::col_major");
+
+  constexpr bool is_row_major = std::is_same_v<LayoutPolicy, raft::row_major>;
+
   auto stream          = resource::get_cuda_stream(handle);
   auto cusolver_handle = raft::resource::get_cusolver_dn_handle(handle);
 
@@ -141,61 +149,68 @@ void cal_eig(raft::resources const& handle,
   raft::matrix::col_reverse(handle_stream_zero,
                             raft::make_device_matrix_view<math_t, idx_t, raft::col_major>(
                               components.data_handle(), n_cols, n_cols));
-  raft::linalg::transpose(components.data_handle(), n_cols, stream);
+  if constexpr (!is_row_major) {
+    raft::linalg::transpose(components.data_handle(), n_cols, stream);
+  }
 
   raft::matrix::row_reverse(handle_stream_zero,
                             raft::make_device_matrix_view<math_t, idx_t, raft::row_major>(
                               explained_var.data_handle(), n_cols, idx_t(1)));
 }
 
 /**
- * @brief sign flip for PCA and tSVD. Stabilizes the sign of column major eigenvectors.
+ * @brief sign flip for PCA and tSVD. Stabilizes the sign of the eigenvectors.
+ * @tparam math_t element type
+ * @tparam idx_t index type
+ * @tparam LayoutPolicy layout of the input and components matrices
  * @param handle: raft::resources
- * @param input: input data [n_samples x n_features] (col-major)
- * @param components: components matrix [n_components x n_features] (col-major)
+ * @param input: input data [n_samples x n_features]
+ * @param components: components matrix [n_components x n_features]
  * @param center whether to mean-center input before computing signs
  * @param flip_signs_based_on_U whether to determine signs by U (true) or V.T (false)
  */
-template <typename math_t, typename idx_t>
+template <typename math_t, typename idx_t, typename LayoutPolicy>
 void sign_flip_components(raft::resources const& handle,
-                          raft::device_matrix_view<math_t, idx_t, raft::col_major> input,
-                          raft::device_matrix_view<math_t, idx_t, raft::col_major> components,
+                          raft::device_matrix_view<math_t, idx_t, LayoutPolicy> input,
+                          raft::device_matrix_view<math_t, idx_t, LayoutPolicy> components,
                           bool center,
                           bool flip_signs_based_on_U = false)
 {
+  static_assert(
+    std::is_same_v<LayoutPolicy, raft::row_major> || std::is_same_v<LayoutPolicy, raft::col_major>,
+    "sign_flip_components: layout must be raft::row_major or raft::col_major");
+  constexpr bool is_row_major = std::is_same_v<LayoutPolicy, raft::row_major>;
+
   auto stream       = resource::get_cuda_stream(handle);
   auto n_samples    = input.extent(0);
   auto n_features   = input.extent(1);
   auto n_components = components.extent(0);
 
   rmm::device_uvector<math_t> max_vals(static_cast<std::size_t>(n_components), stream);
-  auto components_view = raft::make_device_matrix_view<math_t, idx_t, raft::col_major>(
+  auto components_view = raft::make_device_matrix_view<math_t, idx_t, LayoutPolicy>(
     components.data_handle(), n_components, n_features);
   auto max_vals_view = raft::make_device_vector_view<math_t, idx_t>(max_vals.data(), n_components);
 
   if (flip_signs_based_on_U) {
     if (center) {
       rmm::device_uvector<math_t> col_means(static_cast<std::size_t>(n_features), stream);
-      raft::stats::mean<false>(
+      raft::stats::mean<is_row_major>(
         col_means.data(), input.data_handle(), n_features, n_samples, stream);
-      raft::stats::meanCenter<false, true>(
+      raft::stats::meanCenter<is_row_major, true>(
         input.data_handle(), input.data_handle(), col_means.data(), n_features, n_samples, stream);
     }
     rmm::device_uvector<math_t> US(static_cast<std::size_t>(n_samples * n_components), stream);
-    raft::linalg::gemm<math_t, math_t, math_t, math_t>(handle,
-                                                       input.data_handle(),
-                                                       n_samples,
-                                                       n_features,
-                                                       components.data_handle(),
-                                                       US.data(),
-                                                       n_samples,
-                                                       n_components,
-                                                       CUBLAS_OP_N,
-                                                       CUBLAS_OP_T,
-                                                       math_t(1),
-                                                       math_t(0),
-                                                       stream);
-    raft::linalg::reduce<false, false>(
+    raft::linalg::gemm(handle,
+                       input,
+                       raft::make_device_matrix_view<
+                         math_t,
+                         idx_t,
+                         std::conditional_t<is_row_major, raft::col_major, raft::row_major>>(
+                         components.data_handle(), n_features, n_components),
+                       raft::make_device_matrix_view<math_t, idx_t, LayoutPolicy>(
+                         US.data(), n_samples, n_components));
+
+    raft::linalg::reduce<is_row_major, false>(
       max_vals.data(),
       US.data(),
       n_components,
@@ -211,7 +226,7 @@ void sign_flip_components(raft::resources const& handle,
       },
       raft::identity_op());
   } else {
-    raft::linalg::reduce<false, true>(
+    raft::linalg::reduce<is_row_major, true>(
       max_vals.data(),
       components.data_handle(),
       n_features,
@@ -232,8 +247,15 @@ void sign_flip_components(raft::resources const& handle,
     handle,
     components_view,
     [components_view, max_vals_view, n_components, n_features] __device__(auto idx) {
-      auto row    = idx % n_components;
-      auto column = idx / n_components;
+      idx_t row;
+      idx_t column;
+      if constexpr (is_row_major) {
+        row    = idx / n_features;
+        column = idx % n_features;
+      } else {
+        row    = idx % n_components;
+        column = idx / n_components;
+      }
       return (max_vals_view(row) < math_t(0)) ? (-components_view(row, column))
                                               : components_view(row, column);
     });
@@ -376,18 +398,21 @@ void tsvd_fit(raft::resources const& handle,
  * @brief performs transform operation for the tsvd. Transforms the data to eigenspace.
  * @param[in] handle raft::resources
  * @param[in] prms: data structure that includes all the parameters from input size to algorithm.
- * @param[in] input: the data to transform. Size n_rows x n_cols (col-major).
- * @param[in] components: principal components. Size n_components x n_cols (col-major).
- * @param[out] trans_input: transformed output. Size n_rows x n_components (col-major).
+ * @param[in] input: the data to transform. Size n_rows x n_cols.
+ * @param[in] components: principal components. Size n_components x n_cols.
+ * @param[out] trans_input: transformed output. Size n_rows x n_components.
  */
-template <typename math_t, typename idx_t>
+template <typename math_t, typename idx_t, typename LayoutPolicy = raft::col_major>
 void tsvd_transform(raft::resources const& handle,
                     const paramsTSVD& prms,
-                    raft::device_matrix_view<math_t, idx_t, raft::col_major> input,
-                    raft::device_matrix_view<math_t, idx_t, raft::col_major> components,
-                    raft::device_matrix_view<math_t, idx_t, raft::col_major> trans_input)
+                    raft::device_matrix_view<math_t, idx_t, LayoutPolicy> input,
+                    raft::device_matrix_view<math_t, idx_t, LayoutPolicy> components,
+                    raft::device_matrix_view<math_t, idx_t, LayoutPolicy> trans_input)
 {
-  auto stream = resource::get_cuda_stream(handle);
+  static_assert(
+    std::is_same_v<LayoutPolicy, raft::row_major> || std::is_same_v<LayoutPolicy, raft::col_major>,
+    "tsvd_transform: layout must be raft::row_major or raft::col_major");
+  constexpr bool is_row_major = std::is_same_v<LayoutPolicy, raft::row_major>;
 
   auto n_rows       = input.extent(0);
   auto n_cols       = input.extent(1);
@@ -397,39 +422,34 @@ void tsvd_transform(raft::resources const& handle,
   ASSERT(n_rows > 0, "Parameter n_rows: number of rows cannot be less than one");
   ASSERT(n_components > 0, "Parameter n_components: number of components cannot be less than one");
 
-  math_t alpha = math_t(1);
-  math_t beta  = math_t(0);
   raft::linalg::gemm(handle,
-                     input.data_handle(),
-                     n_rows,
-                     n_cols,
-                     components.data_handle(),
-                     trans_input.data_handle(),
-                     n_rows,
-                     n_components,
-                     CUBLAS_OP_N,
-                     CUBLAS_OP_T,
-                     alpha,
-                     beta,
-                     stream);
+                     input,
+                     raft::make_device_matrix_view<
+                       math_t,
+                       idx_t,
+                       std::conditional_t<is_row_major, raft::col_major, raft::row_major>>(
+                       components.data_handle(), n_cols, n_components),
+                     trans_input);
 }
 
 /**
  * @brief performs inverse transform operation for the tsvd.
  * @param[in] handle raft::resources
  * @param[in] prms: data structure that includes all the parameters from input size to algorithm.
- * @param[in] trans_input: the transformed data. Size n_rows x n_components (col-major).
- * @param[in] components: principal components. Size n_components x n_cols (col-major).
- * @param[out] output: reconstructed output. Size n_rows x n_cols (col-major).
+ * @param[in] trans_input: the transformed data. Size n_rows x n_components.
+ * @param[in] components: principal components. Size n_components x n_cols.
+ * @param[out] output: reconstructed output. Size n_rows x n_cols.
  */
-template <typename math_t, typename idx_t>
+template <typename math_t, typename idx_t, typename LayoutPolicy = raft::col_major>
 void tsvd_inverse_transform(raft::resources const& handle,
                             const paramsTSVD& prms,
-                            raft::device_matrix_view<math_t, idx_t, raft::col_major> trans_input,
-                            raft::device_matrix_view<math_t, idx_t, raft::col_major> components,
-                            raft::device_matrix_view<math_t, idx_t, raft::col_major> output)
+                            raft::device_matrix_view<math_t, idx_t, LayoutPolicy> trans_input,
+                            raft::device_matrix_view<math_t, idx_t, LayoutPolicy> components,
+                            raft::device_matrix_view<math_t, idx_t, LayoutPolicy> output)
 {
-  auto stream = resource::get_cuda_stream(handle);
+  static_assert(
+    std::is_same_v<LayoutPolicy, raft::row_major> || std::is_same_v<LayoutPolicy, raft::col_major>,
+    "tsvd_inverse_transform: layout must be raft::row_major or raft::col_major");
 
   auto n_rows       = output.extent(0);
   auto n_cols       = output.extent(1);
@@ -439,22 +459,7 @@ void tsvd_inverse_transform(raft::resources const& handle,
   ASSERT(n_rows > 0, "Parameter n_rows: number of rows cannot be less than one");
   ASSERT(n_components > 0, "Parameter n_components: number of components cannot be less than one");
 
-  math_t alpha = math_t(1);
-  math_t beta  = math_t(0);
-
-  raft::linalg::gemm(handle,
-                     trans_input.data_handle(),
-                     n_rows,
-                     n_components,
-                     components.data_handle(),
-                     output.data_handle(),
-                     n_rows,
-                     n_cols,
-                     CUBLAS_OP_N,
-                     CUBLAS_OP_N,
-                     alpha,
-                     beta,
-                     stream);
+  raft::linalg::gemm(handle, trans_input, components, output);
 }
 
 /**