[REVIEW] Removing FusedL2NN and sparse prims from cuml

cpp/bench/prims/fused_l2_nn.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,7 +15,7 @@
  */
 
 #include <raft/cudart_utils.h>
-#include <distance/fused_l2_nn.cuh>
+#include <raft/distance/fused_l2_nn.cuh>
 #include <limits>
 #include <raft/linalg/norm.cuh>
 #include <raft/random/rng.cuh>
@@ -52,7 +52,7 @@ struct FusedL2NN : public Fixture {
     raft::linalg::rowNorm(yn, y, params.k, params.n, raft::linalg::L2Norm, true,
                           stream);
     auto blks = raft::ceildiv(params.m, 256);
-    MLCommon::Distance::initKernel<T, cub::KeyValuePair<int, T>, int>
+    raft::distance::initKernel<T, cub::KeyValuePair<int, T>, int>
       <<<blks, 256, 0, stream>>>(out, params.m, std::numeric_limits<T>::max(),
                                  op);
   }
@@ -69,9 +69,9 @@ struct FusedL2NN : public Fixture {
   void runBenchmark(::benchmark::State& state) override {
     loopOnState(state, [this]() {
       // it is enough to only benchmark the L2-squared metric
-      MLCommon::Distance::fusedL2NN<T, cub::KeyValuePair<int, T>, int>(
+      raft::distance::fusedL2NN<T, cub::KeyValuePair<int, T>, int>(
         out, x, y, xn, yn, params.m, params.n, params.k, (void*)workspace, op,
-        false, false, stream);
+        pairRedOp, false, false, stream);
     });
   }
 
@@ -80,7 +80,8 @@ struct FusedL2NN : public Fixture {
   T *x, *y, *xn, *yn;
   cub::KeyValuePair<int, T>* out;
   int* workspace;
-  MLCommon::Distance::MinAndDistanceReduceOp<int, T> op;
+  raft::distance::KVPMinReduce<int, T> pairRedOp;
+  raft::distance::MinAndDistanceReduceOp<int, T> op;
 };  // struct FusedL2NN
 
 static std::vector<FLNParams> getInputs() {

cpp/cmake/Dependencies.cmake

@@ -39,7 +39,7 @@ else(DEFINED ENV{RAFT_PATH})
 
   ExternalProject_Add(raft
     GIT_REPOSITORY    https://github.com/rapidsai/raft.git
-    GIT_TAG           4a79adcb0c0e87964dcdc9b9122f242b5235b702
+    GIT_TAG           ffa8a8f5289623b29bb0fa4ed7eeb37b6dfd365c
     PREFIX            ${RAFT_DIR}
     CONFIGURE_COMMAND ""
     BUILD_COMMAND     ""

cpp/src/dbscan/adjgraph/algo.cuh

@@ -23,7 +23,7 @@
 #include "../common.cuh"
 #include "pack.h"
 
-#include <sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.cuh>
 
 using namespace thrust;
 

cpp/src/dbscan/runner.cuh

@@ -21,7 +21,7 @@
 #include <cuml/common/device_buffer.hpp>
 #include <label/classlabels.cuh>
 #include <raft/cuda_utils.cuh>
-#include <sparse/csr.cuh>
+#include <raft/sparse/csr.cuh>
 #include "adjgraph/runner.cuh"
 #include "corepoints/compute.cuh"
 #include "corepoints/exchange.cuh"

cpp/src/kmeans/common.cuh

@@ -16,10 +16,11 @@
 #pragma once
 
 #include <distance/distance.cuh>
-#include <distance/fused_l2_nn.cuh>
 #include <linalg/reduce_cols_by_key.cuh>
 #include <linalg/reduce_rows_by_key.cuh>
 #include <matrix/gather.cuh>
+
+#include <raft/distance/fused_l2_nn.cuh>
 #include <raft/linalg/binary_op.cuh>
 #include <raft/linalg/matrix_vector_op.cuh>
 #include <raft/linalg/mean_squared_error.cuh>
@@ -75,14 +76,14 @@ struct FusedL2NNReduceOp {
   FusedL2NNReduceOp(LabelT _offset) : offset(_offset){};
 
   typedef typename cub::KeyValuePair<LabelT, DataT> KVP;
-  DI void operator()(KVP *out, const KVP &other) {
+  DI void operator()(LabelT rit, KVP *out, const KVP &other) {
     if (other.value < out->value) {
       out->key = offset + other.key;
       out->value = other.value;
     }
   }
 
-  DI void operator()(DataT *out, const KVP &other) {
+  DI void operator()(LabelT rit, DataT *out, const KVP &other) {
     if (other.value < *out) {
       *out = other.value;
     }
@@ -335,12 +336,13 @@ void minClusterAndDistance(
         workspace.resize((sizeof(int)) * ns, stream);
 
         FusedL2NNReduceOp<IndexT, DataT> redOp(cIdx);
+        raft::distance::KVPMinReduce<IndexT, DataT> pairRedOp;
 
-        MLCommon::Distance::fusedL2NN<DataT, cub::KeyValuePair<IndexT, DataT>,
-                                      IndexT>(
+        raft::distance::fusedL2NN<DataT, cub::KeyValuePair<IndexT, DataT>,
+                                  IndexT>(
           minClusterAndDistanceView.data(), datasetView.data(),
           centroidsView.data(), L2NormXView.data(), centroidsNormView.data(),
-          ns, nc, n_features, (void *)workspace.data(), redOp,
+          ns, nc, n_features, (void *)workspace.data(), redOp, pairRedOp,
           (metric == raft::distance::DistanceType::L2Expanded) ? false : true,
           false, stream);
       } else {
@@ -452,10 +454,11 @@ void minClusterDistance(const raft::handle_t &handle,
         workspace.resize((sizeof(int)) * ns, stream);
 
         FusedL2NNReduceOp<IndexT, DataT> redOp(cIdx);
-        MLCommon::Distance::fusedL2NN<DataT, DataT, IndexT>(
+        raft::distance::KVPMinReduce<IndexT, DataT> pairRedOp;
+        raft::distance::fusedL2NN<DataT, DataT, IndexT>(
           minClusterDistanceView.data(), datasetView.data(),
           centroidsView.data(), L2NormXView.data(), centroidsNormView.data(),
-          ns, nc, n_features, (void *)workspace.data(), redOp,
+          ns, nc, n_features, (void *)workspace.data(), redOp, pairRedOp,
           (metric == raft::distance::DistanceType::L2Expanded) ? false : true,
           false, stream);
       } else {

cpp/src/knn/knn_sparse.cu

@@ -17,7 +17,7 @@
 #include <cuml/common/logger.hpp>
 #include <cuml/neighbors/knn_sparse.hpp>
 
-#include <sparse/selection/knn.cuh>
+#include <raft/sparse/selection/knn.cuh>
 
 #include <cusparse_v2.h>
 
@@ -32,8 +32,9 @@ void brute_force_knn(raft::handle_t &handle, const int *idx_indptr,
                      int n_query_rows, int n_query_cols, int *output_indices,
                      float *output_dists, int k,
                      size_t batch_size_index,  // approx 1M
-                     size_t batch_size_query, ML::MetricType metric,
-                     float metricArg, bool expanded_form) {
+                     size_t batch_size_query,
+                     raft::distance::DistanceType metric, float metricArg,
+                     bool expanded_form) {
   auto d_alloc = handle.get_device_allocator();
   cusparseHandle_t cusparse_handle = handle.get_cusparse_handle();
   cudaStream_t stream = handle.get_stream();

cpp/src/spectral/spectral.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,9 +15,9 @@
  */
 
 #include <cuml/cuml.hpp>
-#include <sparse/coo.cuh>
+#include <raft/sparse/coo.cuh>
 
-#include <sparse/linalg/spectral.cuh>
+#include <raft/sparse/linalg/spectral.cuh>
 
 namespace ML {
 

cpp/src/tsne/distances.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,12 +17,13 @@
 #pragma once
 
 #include <raft/cudart_utils.h>
+#include <raft/linalg/distance_type.h>
 #include <cuml/neighbors/knn_sparse.hpp>
 #include <raft/linalg/eltwise.cuh>
+#include <raft/sparse/coo.cuh>
+#include <raft/sparse/linalg/symmetrize.cuh>
+#include <raft/sparse/selection/knn.cuh>
 #include <selection/knn.cuh>
-#include <sparse/coo.cuh>
-#include <sparse/linalg/symmetrize.cuh>
-#include <sparse/selection/knn.cuh>
 
 #include <cuml/manifold/common.hpp>
 
@@ -89,7 +90,7 @@ void get_distances(const raft::handle_t &handle,
     k_graph.knn_indices, k_graph.knn_dists, k_graph.n_neighbors,
     handle.get_cusparse_handle(), handle.get_device_allocator(), stream,
     ML::Sparse::DEFAULT_BATCH_SIZE, ML::Sparse::DEFAULT_BATCH_SIZE,
-    ML::MetricType::METRIC_L2);
+    raft::distance::DistanceType::L2Expanded);
 }
 
 // sparse, int64

cpp/src/umap/fuzzy_simpl_set/naive.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,10 +23,10 @@
 #include <raft/cudart_utils.h>
 #include <raft/cuda_utils.cuh>
 
-#include <sparse/op/sort.h>
+#include <raft/sparse/op/sort.h>
+#include <raft/sparse/coo.cuh>
+#include <raft/sparse/linalg/symmetrize.cuh>
 #include <raft/stats/mean.cuh>
-#include <sparse/coo.cuh>
-#include <sparse/linalg/symmetrize.cuh>
 
 #include <cuda_runtime.h>
 

cpp/src/umap/fuzzy_simpl_set/runner.cuh

@@ -18,7 +18,7 @@
 #include <cuml/manifold/umapparams.h>
 #include "naive.cuh"
 
-#include <sparse/coo.cuh>
+#include <raft/sparse/coo.cuh>
 
 namespace UMAPAlgo {
 

cpp/src/umap/init_embed/runner.cuh

@@ -18,7 +18,7 @@
 
 #include <cuml/manifold/umapparams.h>
 
-#include <sparse/coo.cuh>
+#include <raft/sparse/coo.cuh>
 
 #include "random_algo.cuh"
 #include "spectral_algo.cuh"

cpp/src/umap/init_embed/spectral_algo.cuh

@@ -19,7 +19,7 @@
 #include <cuml/manifold/umapparams.h>
 #include <cuml/common/device_buffer.hpp>
 
-#include <sparse/coo.cuh>
+#include <raft/sparse/coo.cuh>
 
 #include <raft/linalg/add.cuh>
 

cpp/src/umap/knn_graph/algo.cuh

@@ -21,8 +21,8 @@
 #include <cuml/neighbors/knn_sparse.hpp>
 #include <iostream>
 #include <raft/linalg/unary_op.cuh>
+#include <raft/sparse/selection/knn.cuh>
 #include <selection/knn.cuh>
-#include <sparse/selection/knn.cuh>
 
 #include <raft/cudart_utils.h>
 
@@ -91,7 +91,7 @@ void launcher(const raft::handle_t &handle,
     inputsB.n, inputsB.d, out.knn_indices, out.knn_dists, n_neighbors,
     handle.get_cusparse_handle(), d_alloc, stream,
     ML::Sparse::DEFAULT_BATCH_SIZE, ML::Sparse::DEFAULT_BATCH_SIZE,
-    ML::MetricType::METRIC_L2);
+    raft::distance::DistanceType::L2Expanded);
 }
 
 template <>

cpp/src/umap/runner.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -36,11 +36,11 @@
 #include <thrust/scan.h>
 #include <thrust/system/cuda/execution_policy.h>
 
-#include <sparse/op/sort.h>
-#include <sparse/convert/csr.cuh>
-#include <sparse/coo.cuh>
-#include <sparse/linalg/norm.cuh>
-#include <sparse/op/filter.cuh>
+#include <raft/sparse/op/sort.h>
+#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/coo.cuh>
+#include <raft/sparse/linalg/norm.cuh>
+#include <raft/sparse/op/filter.cuh>
 
 #include <raft/cuda_utils.cuh>
 

cpp/src/umap/simpl_set_embed/algo.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -27,11 +27,11 @@
 #include <cstdlib>
 #include <cuml/common/logger.hpp>
 #include <raft/random/rng_impl.cuh>
-#include <sparse/coo.cuh>
+#include <raft/sparse/coo.cuh>
 #include <string>
 #include "optimize_batch_kernel.cuh"
 
-#include <sparse/op/filter.cuh>
+#include <raft/sparse/op/filter.cuh>
 
 #pragma once
 

cpp/src/umap/simpl_set_embed/runner.cuh

@@ -19,7 +19,7 @@
 #include <cuml/manifold/umapparams.h>
 #include "algo.cuh"
 
-#include <sparse/coo.cuh>
+#include <raft/sparse/coo.cuh>
 
 namespace UMAPAlgo {
 

cpp/src/umap/supervised.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -34,12 +34,12 @@
 #include <thrust/scan.h>
 #include <thrust/system/cuda/execution_policy.h>
 
-#include <sparse/convert/csr.cuh>
-#include <sparse/coo.cuh>
-#include <sparse/linalg/add.cuh>
-#include <sparse/linalg/norm.cuh>
-#include <sparse/linalg/symmetrize.cuh>
-#include <sparse/op/filter.cuh>
+#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/coo.cuh>
+#include <raft/sparse/linalg/add.cuh>
+#include <raft/sparse/linalg/norm.cuh>
+#include <raft/sparse/linalg/symmetrize.cuh>
+#include <raft/sparse/op/filter.cuh>
 
 #include <raft/cuda_utils.cuh>
 

cpp/src_prims/distance/epsilon_neighborhood.cuh

@@ -17,13 +17,14 @@
 #pragma once
 
 #include <common/device_utils.cuh>
-#include <linalg/contractions.cuh>
+#include <raft/linalg/contractions.cuh>
 
 namespace MLCommon {
 namespace Distance {
 
 template <typename DataT, typename IdxT, typename Policy,
-          typename BaseClass = LinAlg::Contractions_NT<DataT, IdxT, Policy>>
+          typename BaseClass =
+            raft::linalg::Contractions_NT<DataT, IdxT, Policy>>
 struct EpsUnexpL2SqNeighborhood : public BaseClass {
  private:
   typedef Policy P;
@@ -174,7 +175,7 @@ template <typename DataT, typename IdxT, int VecLen>
 void epsUnexpL2SqNeighImpl(bool* adj, IdxT* vd, const DataT* x, const DataT* y,
                            IdxT m, IdxT n, IdxT k, DataT eps,
                            cudaStream_t stream) {
-  typedef typename LinAlg::Policy4x4<DataT, VecLen>::Policy Policy;
+  typedef typename raft::linalg::Policy4x4<DataT, VecLen>::Policy Policy;
   dim3 grid(raft::ceildiv<int>(m, Policy::Mblk),
             raft::ceildiv<int>(n, Policy::Nblk));
   dim3 blk(Policy::Nthreads);