[FEA] Masked NN for connect_components

cpp/bench/prims/distance/masked_nn.cu

@@ -128,7 +128,6 @@ struct masked_l2_nn : public fixture {
     dim3 block(32, 32);
     dim3 grid(10, 10);
     init_adj<<<grid, block, 0, stream>>>(p.pattern, p.n, adj.view(), group_idxs.view());
-
     RAFT_CUDA_TRY(cudaGetLastError());
   }
 

cpp/include/raft/cluster/detail/mst.cuh

@@ -20,7 +20,7 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
-#include <raft/sparse/neighbors/connect_components.cuh>
+#include <raft/sparse/neighbors/cross_component_nn.cuh>
 #include <raft/sparse/op/sort.cuh>
 #include <raft/sparse/solver/mst.cuh>
 #include <rmm/device_uvector.hpp>
@@ -81,15 +81,20 @@ void connect_knn_graph(
 
   raft::sparse::COO<value_t, value_idx> connected_edges(stream);
 
-  raft::sparse::neighbors::connect_components<value_idx, value_t>(handle,
+  // default row and column batch sizes are chosen for computing cross component nearest neighbors.
+  // Reference: PR #1445
+  static constexpr size_t default_row_batch_size = 4096;
+  static constexpr size_t default_col_batch_size = 16;
+
+  raft::sparse::neighbors::cross_component_nn<value_idx, value_t>(handle,
                                                                   connected_edges,
                                                                   X,
                                                                   color,
                                                                   m,
                                                                   n,
                                                                   reduction_op,
-                                                                  min(m, (size_t)4096),
-                                                                  min(n, (size_t)16));
+                                                                  min(m, default_row_batch_size),
+                                                                  min(n, default_col_batch_size));
 
   rmm::device_uvector<value_idx> indptr2(m + 1, stream);
   raft::sparse::convert::sorted_coo_to_csr(

cpp/include/raft/cluster/detail/single_linkage.cuh

@@ -81,7 +81,7 @@ void single_linkage(raft::resources const& handle,
    * 2. Construct MST, sorted by weights
    */
   rmm::device_uvector<value_idx> color(m, stream);
-  raft::sparse::neighbors::FixConnectivitiesRedOp<value_idx, value_t> op(color.data(), m);
+  raft::sparse::neighbors::FixConnectivitiesRedOp<value_idx, value_t> op(m);
   detail::build_sorted_mst<value_idx, value_t>(handle,
                                                X,
                                                indptr.data(),

cpp/include/raft/matrix/detail/gather.cuh

@@ -16,14 +16,9 @@
 
 #pragma once
 
-#include <raft/core/device_mdarray.hpp>
+#include <functional>
 #include <raft/core/operators.hpp>
-#include <raft/core/resource/thrust_policy.hpp>
-#include <raft/linalg/map.cuh>
-#include <raft/util/cuda_dev_essentials.cuh>
 #include <raft/util/cudart_utils.hpp>
-#include <raft/util/fast_int_div.cuh>
-#include <thrust/iterator/counting_iterator.h>
 
 namespace raft {
 namespace matrix {
@@ -141,16 +136,6 @@ void gatherImpl(const InputIteratorT in,
   // stencil value type
   typedef typename std::iterator_traits<StencilIteratorT>::value_type StencilValueT;
 
-  // return type of MapTransformOp, must be convertible to IndexT
-  typedef typename std::result_of<decltype(transform_op)(MapValueT)>::type MapTransformOpReturnT;
-  static_assert((std::is_convertible<MapTransformOpReturnT, IndexT>::value),
-                "MapTransformOp's result type must be convertible to signed integer");
-
-  // return type of UnaryPredicateOp, must be convertible to bool
-  typedef typename std::result_of<decltype(pred_op)(StencilValueT)>::type PredicateOpReturnT;
-  static_assert((std::is_convertible<PredicateOpReturnT, bool>::value),
-                "UnaryPredicateOp's result type must be convertible to bool type");
-
   IndexT len        = map_length * D;
   constexpr int TPB = 128;
   const int n_sm    = raft::getMultiProcessorCount();
@@ -350,92 +335,6 @@ void gather_if(const InputIteratorT in,
   gatherImpl(in, D, N, map, stencil, map_length, out, pred_op, transform_op, stream);
 }
 
-template <typename MatrixT, typename MapT, typename MapTransformOp, typename IndexT>
-void gatherInplaceImpl(raft::resources const& handle,
-                       raft::device_matrix_view<MatrixT, IndexT, raft::layout_c_contiguous> inout,
-                       raft::device_vector_view<const MapT, IndexT, raft::layout_c_contiguous> map,
-                       MapTransformOp transform_op,
-                       IndexT batch_size)
-{
-  // return type of MapTransformOp, must be convertible to IndexT
-  typedef typename std::result_of<decltype(transform_op)(MapT)>::type MapTransformOpReturnT;
-  RAFT_EXPECTS((std::is_convertible<MapTransformOpReturnT, IndexT>::value),
-               "MapTransformOp's result type must be convertible to signed integer");
-
-  IndexT m          = inout.extent(0);
-  IndexT n          = inout.extent(1);
-  IndexT map_length = map.extent(0);
-
-  // skip in case of 0 length input
-  if (map_length <= 0 || m <= 0 || n <= 0 || batch_size < 0) return;
-
-  RAFT_EXPECTS(map_length <= m, "Length of map should be <= number of rows for inplace gather");
-
-  // re-assign batch_size for default case
-  if (batch_size == 0) batch_size = n;
-
-  RAFT_EXPECTS(batch_size <= n, "batch size should be <= number of columns");
-
-  auto exec_policy = resource::get_thrust_policy(handle);
-  IndexT n_batches = raft::ceildiv(n, batch_size);
-  for (IndexT bid = 0; bid < n_batches; bid++) {
-    IndexT batch_offset   = bid * batch_size;
-    IndexT cols_per_batch = min(batch_size, n - batch_offset);
-    auto scratch_space =
-      raft::make_device_vector<MatrixT, IndexT>(handle, map_length * cols_per_batch);
-
-    auto gather_op = [inout = inout.data_handle(),
-                      map   = map.data_handle(),
-                      transform_op,
-                      batch_offset,
-                      map_length,
-                      cols_per_batch = raft::util::FastIntDiv(cols_per_batch),
-                      n] __device__(auto idx) {
-      IndexT row   = idx / cols_per_batch;
-      IndexT col   = idx % cols_per_batch;
-      MapT map_val = map[row];
-
-      IndexT i_src = transform_op(map_val);
-      return inout[i_src * n + batch_offset + col];
-    };
-    raft::linalg::map_offset(handle, scratch_space.view(), gather_op);
-
-    auto copy_op = [inout         = inout.data_handle(),
-                    map           = map.data_handle(),
-                    scratch_space = scratch_space.data_handle(),
-                    batch_offset,
-                    map_length,
-                    cols_per_batch = raft::util::FastIntDiv(cols_per_batch),
-                    n] __device__(auto idx) {
-      IndexT row                          = idx / cols_per_batch;
-      IndexT col                          = idx % cols_per_batch;
-      inout[row * n + batch_offset + col] = scratch_space[idx];
-      return;
-    };
-    auto counting = thrust::make_counting_iterator<IndexT>(0);
-    thrust::for_each(exec_policy, counting, counting + map_length * cols_per_batch, copy_op);
-  }
-}
-
-template <typename MatrixT, typename MapT, typename MapTransformOp, typename IndexT>
-void gather(raft::resources const& handle,
-            raft::device_matrix_view<MatrixT, IndexT, raft::layout_c_contiguous> inout,
-            raft::device_vector_view<const MapT, IndexT, raft::layout_c_contiguous> map,
-            MapTransformOp transform_op,
-            IndexT batch_size)
-{
-  gatherInplaceImpl(handle, inout, map, transform_op, batch_size);
-}
-
-template <typename MatrixT, typename MapT, typename IndexT>
-void gather(raft::resources const& handle,
-            raft::device_matrix_view<MatrixT, IndexT, raft::layout_c_contiguous> inout,
-            raft::device_vector_view<const MapT, IndexT, raft::layout_c_contiguous> map,
-            IndexT batch_size)
-{
-  gatherInplaceImpl(handle, inout, map, raft::identity_op(), batch_size);
-}
-
 }  // namespace detail
 }  // namespace matrix
 }  // namespace raft

cpp/include/raft/matrix/detail/gather_inplace.cuh

@@ -0,0 +1,116 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/resource/thrust_policy.hpp>
+#include <raft/linalg/map.cuh>
+#include <raft/util/fast_int_div.cuh>
+#include <thrust/iterator/counting_iterator.h>
+
+namespace raft {
+namespace matrix {
+namespace detail {
+
+template <typename MatrixT, typename MapT, typename MapTransformOp, typename IndexT>
+void gatherInplaceImpl(raft::resources const& handle,
+                       raft::device_matrix_view<MatrixT, IndexT, raft::layout_c_contiguous> inout,
+                       raft::device_vector_view<const MapT, IndexT, raft::layout_c_contiguous> map,
+                       MapTransformOp transform_op,
+                       IndexT batch_size)
+{
+  IndexT m          = inout.extent(0);
+  IndexT n          = inout.extent(1);
+  IndexT map_length = map.extent(0);
+
+  // skip in case of 0 length input
+  if (map_length <= 0 || m <= 0 || n <= 0 || batch_size < 0) return;
+
+  RAFT_EXPECTS(map_length <= m, "Length of map should be <= number of rows for inplace gather");
+
+  RAFT_EXPECTS(batch_size >= 0, "batch size should be >= 0");
+
+  // re-assign batch_size for default case
+  if (batch_size == 0 || batch_size > n) batch_size = n;
+
+  auto exec_policy = resource::get_thrust_policy(handle);
+
+  IndexT n_batches = raft::ceildiv(n, batch_size);
+
+  auto scratch_space = raft::make_device_vector<MatrixT, IndexT>(handle, map_length * batch_size);
+
+  for (IndexT bid = 0; bid < n_batches; bid++) {
+    IndexT batch_offset   = bid * batch_size;
+    IndexT cols_per_batch = min(batch_size, n - batch_offset);
+
+    auto gather_op = [inout = inout.data_handle(),
+                      map   = map.data_handle(),
+                      transform_op,
+                      batch_offset,
+                      map_length,
+                      cols_per_batch = raft::util::FastIntDiv(cols_per_batch),
+                      n] __device__(auto idx) {
+      IndexT row   = idx / cols_per_batch;
+      IndexT col   = idx % cols_per_batch;
+      MapT map_val = map[row];
+
+      IndexT i_src = transform_op(map_val);
+      return inout[i_src * n + batch_offset + col];
+    };
+    raft::linalg::map_offset(
+      handle,
+      raft::make_device_vector_view(scratch_space.data_handle(), map_length * cols_per_batch),
+      gather_op);
+
+    auto copy_op = [inout         = inout.data_handle(),
+                    map           = map.data_handle(),
+                    scratch_space = scratch_space.data_handle(),
+                    batch_offset,
+                    map_length,
+                    cols_per_batch = raft::util::FastIntDiv(cols_per_batch),
+                    n] __device__(auto idx) {
+      IndexT row                          = idx / cols_per_batch;
+      IndexT col                          = idx % cols_per_batch;
+      inout[row * n + batch_offset + col] = scratch_space[idx];
+      return;
+    };
+    auto counting = thrust::make_counting_iterator<IndexT>(0);
+    thrust::for_each(exec_policy, counting, counting + map_length * cols_per_batch, copy_op);
+  }
+}
+
+template <typename MatrixT, typename MapT, typename MapTransformOp, typename IndexT>
+void gather(raft::resources const& handle,
+            raft::device_matrix_view<MatrixT, IndexT, raft::layout_c_contiguous> inout,
+            raft::device_vector_view<const MapT, IndexT, raft::layout_c_contiguous> map,
+            MapTransformOp transform_op,
+            IndexT batch_size)
+{
+  gatherInplaceImpl(handle, inout, map, transform_op, batch_size);
+}
+
+template <typename MatrixT, typename MapT, typename IndexT>
+void gather(raft::resources const& handle,
+            raft::device_matrix_view<MatrixT, IndexT, raft::layout_c_contiguous> inout,
+            raft::device_vector_view<const MapT, IndexT, raft::layout_c_contiguous> map,
+            IndexT batch_size)
+{
+  gatherInplaceImpl(handle, inout, map, raft::identity_op(), batch_size);
+}
+
+}  // namespace detail
+}  // namespace matrix
+}  // namespace raft