[FEA] support of prefiltered brute force

cpp/CMakeLists.txt

@@ -334,6 +334,14 @@ if(RAFT_COMPILE_LIBRARY)
     src/matrix/detail/select_k_float_int32.cu
     src/matrix/detail/select_k_half_int64_t.cu
     src/matrix/detail/select_k_half_uint32_t.cu
+    src/sparse/matrix/detail/select_k_half_uint32_t.cu
+    src/sparse/matrix/detail/select_k_double_int64_t.cu
+    src/sparse/matrix/detail/select_k_double_uint32_t.cu
+    src/sparse/matrix/detail/select_k_float_int64_t.cu
+    src/sparse/matrix/detail/select_k_float_uint32_t.cu
+    src/sparse/matrix/detail/select_k_float_int32.cu
+    src/sparse/matrix/detail/select_k_half_int64_t.cu
+    src/sparse/matrix/detail/select_k_half_uint32_t.cu
     src/neighbors/ball_cover.cu
     src/neighbors/brute_force_fused_l2_knn_float_int64_t.cu
     src/neighbors/brute_force_knn_int64_t_float_int64_t.cu

cpp/include/raft/core/bitmap.cuh

@@ -16,112 +16,31 @@
 
 #pragma once
 
+#include <raft/core/bitmap.hpp>
 #include <raft/core/bitset.cuh>
 #include <raft/core/detail/mdspan_util.cuh>
 #include <raft/core/device_container_policy.hpp>
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/resources.hpp>
 
-namespace raft::core {
-/**
- * @defgroup bitmap Bitmap
- * @{
- */
-/**
- * @brief View of a RAFT Bitmap.
- *
- * This lightweight structure which represents and manipulates a two-dimensional bitmap matrix view
- * with row major order. This class provides functionality for handling a matrix where each element
- * is represented as a bit in a bitmap.
- *
- * @tparam bitmap_t Underlying type of the bitmap array. Default is uint32_t.
- * @tparam index_t Indexing type used. Default is uint32_t.
- */
-template <typename bitmap_t = uint32_t, typename index_t = uint32_t>
-struct bitmap_view : public bitset_view<bitmap_t, index_t> {
-  static_assert((std::is_same<bitmap_t, uint32_t>::value ||
-                 std::is_same<bitmap_t, uint64_t>::value),
-                "The bitmap_t must be uint32_t or uint64_t.");
-  /**
-   * @brief Create a bitmap view from a device raw pointer.
-   *
-   * @param bitmap_ptr Device raw pointer
-   * @param rows Number of row in the matrix.
-   * @param cols Number of col in the matrix.
-   */
-  _RAFT_HOST_DEVICE bitmap_view(bitmap_t* bitmap_ptr, index_t rows, index_t cols)
-    : bitset_view<bitmap_t, index_t>(bitmap_ptr, rows * cols), rows_(rows), cols_(cols)
-  {
-  }
-
-  /**
-   * @brief Create a bitmap view from a device vector view of the bitset.
-   *
-   * @param bitmap_span Device vector view of the bitmap
-   * @param rows Number of row in the matrix.
-   * @param cols Number of col in the matrix.
-   */
-  _RAFT_HOST_DEVICE bitmap_view(raft::device_vector_view<bitmap_t, index_t> bitmap_span,
-                                index_t rows,
-                                index_t cols)
-    : bitset_view<bitmap_t, index_t>(bitmap_span, rows * cols), rows_(rows), cols_(cols)
-  {
-  }
+#include <type_traits>
 
- private:
-  // Hide the constructors of bitset_view.
-  _RAFT_HOST_DEVICE bitmap_view(bitmap_t* bitmap_ptr, index_t bitmap_len)
-    : bitset_view<bitmap_t, index_t>(bitmap_ptr, bitmap_len)
-  {
-  }
-
-  _RAFT_HOST_DEVICE bitmap_view(raft::device_vector_view<bitmap_t, index_t> bitmap_span,
-                                index_t bitmap_len)
-    : bitset_view<bitmap_t, index_t>(bitmap_span, bitmap_len)
-  {
-  }
-
- public:
-  /**
-   * @brief Device function to test if a given row and col are set in the bitmap.
-   *
-   * @param row Row index of the bit to test
-   * @param col Col index of the bit to test
-   * @return bool True if index has not been unset in the bitset
-   */
-  inline _RAFT_DEVICE auto test(const index_t row, const index_t col) const -> bool
-  {
-    return test(row * cols_ + col);
-  }
-
-  /**
-   * @brief Device function to set a given row and col to set_value in the bitset.
-   *
-   * @param row Row index of the bit to set
-   * @param col Col index of the bit to set
-   * @param new_value Value to set the bit to (true or false)
-   */
-  inline _RAFT_DEVICE void set(const index_t row, const index_t col, bool new_value) const
-  {
-    set(row * cols_ + col, &new_value);
-  }
-
-  /**
-   * @brief Get the total number of rows
-   * @return index_t The total number of rows
-   */
-  inline _RAFT_HOST_DEVICE index_t get_n_rows() const { return rows_; }
-
-  /**
-   * @brief Get the total number of columns
-   * @return index_t The total number of columns
-   */
-  inline _RAFT_HOST_DEVICE index_t get_n_cols() const { return cols_; }
+namespace raft::core {
 
- private:
-  index_t rows_;
-  index_t cols_;
-};
+template <typename bitmap_t, typename index_t>
+_RAFT_HOST_DEVICE inline bool bitmap_view<bitmap_t, index_t>::test(const index_t row,
+                                                                   const index_t col) const
+{
+  return test(row * cols_ + col);
+}
+
+template <typename bitmap_t, typename index_t>
+_RAFT_HOST_DEVICE void bitmap_view<bitmap_t, index_t>::set(const index_t row,
+                                                           const index_t col,
+                                                           bool new_value) const
+{
+  set(row * cols_ + col, &new_value);
+}
 
 /** @} */
 }  // end namespace raft::core

cpp/include/raft/core/bitmap.hpp

@@ -0,0 +1,123 @@
+/*
+ * Copyright (c) 2024, 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/bitset.hpp>
+#include <raft/core/detail/mdspan_util.cuh>
+#include <raft/core/device_container_policy.hpp>
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/resources.hpp>
+
+#include <type_traits>
+
+namespace raft::core {
+/**
+ * @defgroup bitmap Bitmap
+ * @{
+ */
+/**
+ * @brief View of a RAFT Bitmap.
+ *
+ * This lightweight structure which represents and manipulates a two-dimensional bitmap matrix view
+ * with row major order. This class provides functionality for handling a matrix where each element
+ * is represented as a bit in a bitmap.
+ *
+ * @tparam bitmap_t Underlying type of the bitmap array. Default is uint32_t.
+ * @tparam index_t Indexing type used. Default is uint32_t.
+ */
+template <typename bitmap_t = uint32_t, typename index_t = uint32_t>
+struct bitmap_view : public bitset_view<bitmap_t, index_t> {
+  static_assert((std::is_same<typename std::remove_const<bitmap_t>::type, uint32_t>::value ||
+                 std::is_same<typename std::remove_const<bitmap_t>::type, uint64_t>::value),
+                "The bitmap_t must be uint32_t or uint64_t.");
+  /**
+   * @brief Create a bitmap view from a device raw pointer.
+   *
+   * @param bitmap_ptr Device raw pointer
+   * @param rows Number of row in the matrix.
+   * @param cols Number of col in the matrix.
+   */
+  _RAFT_HOST_DEVICE bitmap_view(bitmap_t* bitmap_ptr, index_t rows, index_t cols)
+    : bitset_view<bitmap_t, index_t>(bitmap_ptr, rows * cols), rows_(rows), cols_(cols)
+  {
+  }
+
+  /**
+   * @brief Create a bitmap view from a device vector view of the bitset.
+   *
+   * @param bitmap_span Device vector view of the bitmap
+   * @param rows Number of row in the matrix.
+   * @param cols Number of col in the matrix.
+   */
+  _RAFT_HOST_DEVICE bitmap_view(raft::device_vector_view<bitmap_t, index_t> bitmap_span,
+                                index_t rows,
+                                index_t cols)
+    : bitset_view<bitmap_t, index_t>(bitmap_span, rows * cols), rows_(rows), cols_(cols)
+  {
+  }
+
+ private:
+  // Hide the constructors of bitset_view.
+  _RAFT_HOST_DEVICE bitmap_view(bitmap_t* bitmap_ptr, index_t bitmap_len)
+    : bitset_view<bitmap_t, index_t>(bitmap_ptr, bitmap_len)
+  {
+  }
+
+  _RAFT_HOST_DEVICE bitmap_view(raft::device_vector_view<bitmap_t, index_t> bitmap_span,
+                                index_t bitmap_len)
+    : bitset_view<bitmap_t, index_t>(bitmap_span, bitmap_len)
+  {
+  }
+
+ public:
+  /**
+   * @brief Device function to test if a given row and col are set in the bitmap.
+   *
+   * @param row Row index of the bit to test
+   * @param col Col index of the bit to test
+   * @return bool True if index has not been unset in the bitset
+   */
+  inline _RAFT_HOST_DEVICE bool test(const index_t row, const index_t col) const;
+
+  /**
+   * @brief Device function to set a given row and col to set_value in the bitset.
+   *
+   * @param row Row index of the bit to set
+   * @param col Col index of the bit to set
+   * @param new_value Value to set the bit to (true or false)
+   */
+  inline _RAFT_HOST_DEVICE void set(const index_t row, const index_t col, bool new_value) const;
+
+  /**
+   * @brief Get the total number of rows
+   * @return index_t The total number of rows
+   */
+  inline _RAFT_HOST_DEVICE index_t get_n_rows() const { return rows_; }
+
+  /**
+   * @brief Get the total number of columns
+   * @return index_t The total number of columns
+   */
+  inline _RAFT_HOST_DEVICE index_t get_n_cols() const { return cols_; }
+
+ private:
+  index_t rows_;
+  index_t cols_;
+};
+
+/** @} */
+}  // end namespace raft::core

cpp/include/raft/core/bitset.cuh

@@ -17,7 +17,7 @@
 #pragma once
 
 #include <raft/core/bitset.hpp>
-#include <raft/core/detail/mdspan_util.cuh>  // native_popc
+#include <raft/core/detail/popc.cuh>
 #include <raft/core/device_container_policy.hpp>
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
@@ -60,6 +60,12 @@ _RAFT_HOST_DEVICE void bitset_view<bitset_t, index_t>::set(const index_t sample_
   }
 }
 
+template <typename bitset_t, typename index_t>
+_RAFT_HOST_DEVICE inline index_t bitset_view<bitset_t, index_t>::n_elements() const
+{
+  return raft::ceildiv(bitset_len_, bitset_element_size);
+}
+
 template <typename bitset_t, typename index_t>
 bitset<bitset_t, index_t>::bitset(const raft::resources& res,
                                   raft::device_vector_view<const index_t, index_t> mask_index,
@@ -161,37 +167,9 @@ template <typename bitset_t, typename index_t>
 void bitset<bitset_t, index_t>::count(const raft::resources& res,
                                       raft::device_scalar_view<index_t> count_gpu_scalar)
 {
-  auto n_elements_ = n_elements();
-  auto count_gpu =
-    raft::make_device_vector_view<index_t, index_t>(count_gpu_scalar.data_handle(), 1);
-  auto bitset_matrix_view = raft::make_device_matrix_view<const bitset_t, index_t, raft::col_major>(
-    bitset_.data(), n_elements_, 1);
-
-  bitset_t n_last_element = (bitset_len_ % bitset_element_size);
-  bitset_t last_element_mask =
-    n_last_element ? (bitset_t)((bitset_t{1} << n_last_element) - bitset_t{1}) : ~bitset_t{0};
-  raft::linalg::coalesced_reduction(
-    res,
-    bitset_matrix_view,
-    count_gpu,
-    index_t{0},
-    false,
-    [last_element_mask, n_elements_] __device__(bitset_t element, index_t index) {
-      index_t result = 0;
-      if constexpr (bitset_element_size == 64) {
-        if (index == n_elements_ - 1)
-          result = index_t(raft::detail::popc(element & last_element_mask));
-        else
-          result = index_t(raft::detail::popc(element));
-      } else {  // Needed because popc is not overloaded for 16 and 8 bit elements
-        if (index == n_elements_ - 1)
-          result = index_t(raft::detail::popc(uint32_t{element} & last_element_mask));
-        else
-          result = index_t(raft::detail::popc(uint32_t{element}));
-      }
-
-      return result;
-    });
+  auto values =
+    raft::make_device_vector_view<const bitset_t, index_t>(bitset_.data(), n_elements());
+  raft::detail::popc(res, values, bitset_len_, count_gpu_scalar);
 }
 
 }  // end namespace raft::core