Instantiate only specific RAFT reduction kernels

build.sh

@@ -450,7 +450,7 @@ if (( NUMARGS == 0 )) || hasArg libcuvs || hasArg docs || hasArg tests || hasArg
       MSG="${MSG}<br/>parallel setting: $PARALLEL_LEVEL"
       MSG="${MSG}<br/>parallel build time: $compile_total seconds"
       if [[ -f "${LIBCUVS_BUILD_DIR}/libcuvs.so" ]]; then
-          LIBCUVS_FS=$(ls -lh ${LIBCUVS_BUILD_DIR}/libcuvs.so | awk '{print $5}')
+          LIBCUVS_FS=$(stat -c %s ${LIBCUVS_BUILD_DIR}/libcuvs.so | awk '{printf "%.2f MB", $1/1024/1024}')
           MSG="${MSG}<br/>libcuvs.so size: $LIBCUVS_FS"
       fi
       BMR_DIR=${RAPIDS_ARTIFACTS_DIR:-"${LIBCUVS_BUILD_DIR}"}

cpp/cmake/patches/faiss.diff

@@ -0,0 +1,19 @@
+diff --git a/faiss/gpu/impl/CuvsIVFFlat.cu b/faiss/gpu/impl/CuvsIVFFlat.cu
+index 2cccee860..c4cb13f48 100644
+--- a/faiss/gpu/impl/CuvsIVFFlat.cu
++++ b/faiss/gpu/impl/CuvsIVFFlat.cu
+@@ -427,13 +427,11 @@ void CuvsIVFFlat::copyInvertedListsFrom(const InvertedLists* ivf) {
+     // Precompute the centers vector norms for L2Expanded distance
+     if (this->metric_ == faiss::METRIC_L2) {
+         cuvs_index->allocate_center_norms(raft_handle);
+-        raft::linalg::rowNorm(
++        raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+                 cuvs_index->center_norms().value().data_handle(),
+                 cuvs_index->centers().data_handle(),
+                 cuvs_index->dim(),
+                 (uint32_t)nlist,
+-                raft::linalg::L2Norm,
+-                true,
+                 raft_handle.get_stream());
+     }
+ }

cpp/cmake/patches/faiss_override.json

@@ -3,7 +3,14 @@
       "faiss" : {
         "version": "1.10.0",
         "git_url": "https://github.com/facebookresearch/faiss.git",
-        "git_tag": "main"
+        "git_tag": "main",
+        "patches" : [
+          {
+            "file" : "${current_json_dir}/faiss.diff",
+            "issue" : "Apply RAFT breaking changes",
+            "fixed_in" : ""
+          }
+        ]
       }
     }
   }

cpp/cmake/thirdparty/get_raft.cmake

@@ -35,6 +35,8 @@ function(find_and_configure_raft)
         string(APPEND RAFT_COMPONENTS " distributed")
     endif()
 
+    # set(CPM_raft_SOURCE /raid/dgala/raft)
+
     #-----------------------------------------------------
     # Invoke CPM find_package()
     #-----------------------------------------------------
@@ -61,8 +63,8 @@ endfunction()
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION  ${RAFT_VERSION}.00
-        FORK                     ${RAFT_FORK}
-        PINNED_TAG               ${RAFT_PINNED_TAG}
+        FORK                     divyegala
+        PINNED_TAG               reduction-kernels
         ENABLE_MNMG_DEPENDENCIES OFF
         ENABLE_NVTX              OFF
         USE_RAFT_STATIC ${CUVS_USE_RAFT_STATIC}

cpp/src/cluster/detail/kmeans.cuh

@@ -143,13 +143,8 @@ void kmeansPlusPlus(raft::resources const& handle,
 
   if (metric == cuvs::distance::DistanceType::L2Expanded ||
       metric == cuvs::distance::DistanceType::L2SqrtExpanded) {
-    raft::linalg::rowNorm(L2NormX.data_handle(),
-                          X.data_handle(),
-                          X.extent(1),
-                          X.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+      L2NormX.data_handle(), X.data_handle(), X.extent(1), X.extent(0), stream);
   }
 
   raft::random::RngState rng(params.rng_state.seed, params.rng_state.type);
@@ -216,14 +211,12 @@ void kmeansPlusPlus(raft::resources const& handle,
 
     // Calculate costPerCandidate[n_trials] where costPerCandidate[i] is the cluster cost when using
     // centroid candidate-i
-    raft::linalg::reduce(costPerCandidate.data_handle(),
-                         minDistBuf.data_handle(),
-                         minDistBuf.extent(1),
-                         minDistBuf.extent(0),
-                         static_cast<DataT>(0),
-                         true,
-                         true,
-                         stream);
+    raft::linalg::reduce<true, true>(costPerCandidate.data_handle(),
+                                     minDistBuf.data_handle(),
+                                     minDistBuf.extent(1),
+                                     minDistBuf.extent(0),
+                                     static_cast<DataT>(0),
+                                     stream);
 
     // Greedy Choice - Choose the candidate that has minimum cluster cost
     // ArgMin operation below identifies the index of minimum cost in costPerCandidate
@@ -406,13 +399,8 @@ void kmeans_fit_main(raft::resources const& handle,
 
   if (metric == cuvs::distance::DistanceType::L2Expanded ||
       metric == cuvs::distance::DistanceType::L2SqrtExpanded) {
-    raft::linalg::rowNorm(L2NormX.data_handle(),
-                          X.data_handle(),
-                          X.extent(1),
-                          X.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+      L2NormX.data_handle(), X.data_handle(), X.extent(1), X.extent(0), stream);
   }
 
   RAFT_LOG_DEBUG(
@@ -634,13 +622,8 @@ void initScalableKMeansPlusPlus(raft::resources const& handle,
   auto L2NormX = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
   if (metric == cuvs::distance::DistanceType::L2Expanded ||
       metric == cuvs::distance::DistanceType::L2SqrtExpanded) {
-    raft::linalg::rowNorm(L2NormX.data_handle(),
-                          X.data_handle(),
-                          X.extent(1),
-                          X.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+      L2NormX.data_handle(), X.data_handle(), X.extent(1), X.extent(0), stream);
   }
 
   auto minClusterDistanceVec = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
@@ -1049,13 +1032,8 @@ void kmeans_predict(raft::resources const& handle,
   auto L2NormX = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
   if (metric == cuvs::distance::DistanceType::L2Expanded ||
       metric == cuvs::distance::DistanceType::L2SqrtExpanded) {
-    raft::linalg::rowNorm(L2NormX.data_handle(),
-                          X.data_handle(),
-                          X.extent(1),
-                          X.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+      L2NormX.data_handle(), X.data_handle(), X.extent(1), X.extent(0), stream);
   }
 
   // computes minClusterAndDistance[0:n_samples) where  minClusterAndDistance[i]

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -113,8 +113,8 @@ inline std::enable_if_t<std::is_floating_point_v<MathT>> predict_core(
 
       auto centroidsNorm =
         raft::make_device_mdarray<MathT, IdxT>(handle, mr, raft::make_extents<IdxT>(n_clusters));
-      raft::linalg::rowNorm<MathT, IdxT>(
-        centroidsNorm.data_handle(), centers, dim, n_clusters, raft::linalg::L2Norm, true, stream);
+      raft::linalg::rowNorm<raft::linalg::L2Norm, true, MathT, IdxT>(
+        centroidsNorm.data_handle(), centers, dim, n_clusters, stream);
 
       cuvs::distance::fusedDistanceNNMinReduce<MathT, raft::KeyValuePair<IdxT, MathT>, IdxT>(
         minClusterAndDistance.data_handle(),
@@ -156,14 +156,8 @@ inline std::enable_if_t<std::is_floating_point_v<MathT>> predict_core(
 
       auto centroidsNorm =
         raft::make_device_mdarray<MathT, IdxT>(handle, mr, raft::make_extents<IdxT>(n_clusters));
-      raft::linalg::rowNorm<MathT, IdxT>(centroidsNorm.data_handle(),
-                                         centers,
-                                         dim,
-                                         n_clusters,
-                                         raft::linalg::L2Norm,
-                                         true,
-                                         stream,
-                                         raft::sqrt_op{});
+      raft::linalg::rowNorm<raft::linalg::L2Norm, true, MathT, IdxT>(
+        centroidsNorm.data_handle(), centers, dim, n_clusters, stream, raft::sqrt_op{});
 
       cuvs::distance::fusedDistanceNNMinReduce<MathT, raft::KeyValuePair<IdxT, MathT>, IdxT>(
         minClusterAndDistance.data_handle(),
@@ -395,8 +389,8 @@ void compute_norm(const raft::resources& handle,
     dataset_ptr = static_cast<const MathT*>(mapped_dataset.data());
   }
 
-  raft::linalg::rowNorm<MathT, IdxT>(
-    dataset_norm, dataset_ptr, dim, n_rows, raft::linalg::L2Norm, true, stream, norm_fin_op);
+  raft::linalg::rowNorm<raft::linalg::L2Norm, true, MathT, IdxT>(
+    dataset_norm, dataset_ptr, dim, n_rows, stream, norm_fin_op);
 }
 
 /**
@@ -732,8 +726,8 @@ void balancing_em_iters(const raft::resources& handle,
           cluster_centers, n_clusters, dim);
         auto clusters_out_view = raft::make_device_matrix_view<MathT, IdxT, raft::row_major>(
           cluster_centers, n_clusters, dim);
-        raft::linalg::row_normalize(
-          handle, clusters_in_view, clusters_out_view, raft::linalg::L2Norm);
+        raft::linalg::row_normalize<raft::linalg::L2Norm>(
+          handle, clusters_in_view, clusters_out_view);
         break;
       }
       default: break;

cpp/src/cluster/detail/kmeans_common.cuh

@@ -382,13 +382,11 @@ void minClusterAndDistanceCompute(
 
   if (is_fused) {
     L2NormBuf_OR_DistBuf.resize(n_clusters, stream);
-    raft::linalg::rowNorm(L2NormBuf_OR_DistBuf.data(),
-                          centroids.data_handle(),
-                          centroids.extent(1),
-                          centroids.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(L2NormBuf_OR_DistBuf.data(),
+                                                      centroids.data_handle(),
+                                                      centroids.extent(1),
+                                                      centroids.extent(0),
+                                                      stream);
   } else {
     // TODO: Unless pool allocator is used, passing in a workspace for this
     // isn't really increasing performance because this needs to do a re-allocation
@@ -518,13 +516,11 @@ void minClusterDistanceCompute(raft::resources const& handle,
 
   if (is_fused) {
     L2NormBuf_OR_DistBuf.resize(n_clusters, stream);
-    raft::linalg::rowNorm(L2NormBuf_OR_DistBuf.data(),
-                          centroids.data_handle(),
-                          centroids.extent(1),
-                          centroids.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(L2NormBuf_OR_DistBuf.data(),
+                                                      centroids.data_handle(),
+                                                      centroids.extent(1),
+                                                      centroids.extent(0),
+                                                      stream);
   } else {
     L2NormBuf_OR_DistBuf.resize(dataBatchSize * centroidsBatchSize, stream);
   }

cpp/src/cluster/detail/kmeans_mg.cuh

@@ -228,13 +228,8 @@ void initKMeansPlusPlus(const raft::resources& handle,
   auto L2NormX = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
   if (metric == cuvs::distance::DistanceType::L2Expanded ||
       metric == cuvs::distance::DistanceType::L2SqrtExpanded) {
-    raft::linalg::rowNorm(L2NormX.data_handle(),
-                          X.data_handle(),
-                          X.extent(1),
-                          X.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+      L2NormX.data_handle(), X.data_handle(), X.extent(1), X.extent(0), stream);
   }
 
   auto minClusterDistance = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
@@ -578,13 +573,8 @@ void fit(const raft::resources& handle,
   auto L2NormX = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
   if (metric == cuvs::distance::DistanceType::L2Expanded ||
       metric == cuvs::distance::DistanceType::L2SqrtExpanded) {
-    raft::linalg::rowNorm(L2NormX.data_handle(),
-                          X.data_handle(),
-                          X.extent(1),
-                          X.extent(0),
-                          raft::linalg::L2Norm,
-                          true,
-                          stream);
+    raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+      L2NormX.data_handle(), X.data_handle(), X.extent(1), X.extent(0), stream);
   }
 
   DataT priorClusteringCost = 0;

cpp/src/cluster/kmeans.cuh

@@ -481,13 +481,8 @@ void cluster_cost(raft::resources const& handle,
 
   auto x_norms = raft::make_device_vector<DataT>(handle, n_samples);
 
-  raft::linalg::rowNorm(x_norms.data_handle(),
-                        X.data_handle(),
-                        n_features,
-                        n_samples,
-                        raft::linalg::L2Norm,
-                        true,
-                        stream);
+  raft::linalg::rowNorm<raft::linalg::L2Norm, true>(
+    x_norms.data_handle(), X.data_handle(), n_features, n_samples, stream);
 
   auto min_cluster_distance = raft::make_device_vector<DataT>(handle, n_samples);
   rmm::device_uvector<DataT> l2_norm_or_distance_buffer(0, stream);