[FEA] Add support for bitmap_view & the API of bitmap_to_csr

- This PR is one part of the Feature of [FEA] Pre-filtered brute-force KNN #1969

Authors:
  - James Rong (https://github.com/rhdong)

Approvers:
  - Ben Frederickson (https://github.com/benfred)
  - Micka (https://github.com/lowener)
  - Corey J. Nolet (https://github.com/cjnolet)

cpp/bench/prims/CMakeLists.txt

@@ -139,7 +139,14 @@ if(BUILD_PRIMS_BENCH)
     bench/prims/random/rng.cu bench/prims/main.cpp
   )
 
-  ConfigureBench(NAME SPARSE_BENCH PATH bench/prims/sparse/convert_csr.cu bench/prims/main.cpp)
+  ConfigureBench(
+    NAME
+    SPARSE_BENCH
+    PATH
+    bench/prims/sparse/bitmap_to_csr.cu
+    bench/prims/sparse/convert_csr.cu
+    bench/prims/main.cpp
+  )
 
   ConfigureBench(
     NAME

cpp/bench/prims/sparse/bitmap_to_csr.cu

@@ -0,0 +1,155 @@
+/*
+ * 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.
+ */
+#include <common/benchmark.hpp>
+#include <raft/sparse/convert/csr.cuh>
+#include <rmm/device_uvector.hpp>
+
+#include <raft/core/device_resources.hpp>
+#include <raft/util/itertools.hpp>
+
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resources.hpp>
+
+#include <sstream>
+#include <vector>
+
+namespace raft::bench::sparse {
+
+template <typename index_t>
+struct bench_param {
+  index_t n_rows;
+  index_t n_cols;
+  float sparsity;
+};
+
+template <typename index_t>
+inline auto operator<<(std::ostream& os, const bench_param<index_t>& params) -> std::ostream&
+{
+  os << " rows*cols=" << params.n_rows << "*" << params.n_cols << "\tsparsity=" << params.sparsity;
+  return os;
+}
+
+template <typename bitmap_t, typename index_t, typename value_t = float>
+struct BitmapToCsrTest : public fixture {
+  BitmapToCsrTest(const bench_param<index_t>& p)
+    : fixture(true),
+      params(p),
+      handle(stream),
+      bitmap_d(0, stream),
+      nnz(0),
+      indptr_d(0, stream),
+      indices_d(0, stream),
+      values_d(0, stream)
+  {
+    index_t element = raft::ceildiv(params.n_rows * params.n_cols, index_t(sizeof(bitmap_t) * 8));
+    std::vector<bitmap_t> bitmap_h(element);
+    nnz = create_sparse_matrix(params.n_rows, params.n_cols, params.sparsity, bitmap_h);
+
+    bitmap_d.resize(bitmap_h.size(), stream);
+    indptr_d.resize(params.n_rows + 1, stream);
+    indices_d.resize(nnz, stream);
+    values_d.resize(nnz, stream);
+
+    update_device(bitmap_d.data(), bitmap_h.data(), bitmap_h.size(), stream);
+
+    resource::sync_stream(handle);
+  }
+
+  index_t create_sparse_matrix(index_t m, index_t n, float sparsity, std::vector<bitmap_t>& bitmap)
+  {
+    index_t total    = static_cast<index_t>(m * n);
+    index_t num_ones = static_cast<index_t>((total * 1.0f) * sparsity);
+    index_t res      = num_ones;
+
+    for (auto& item : bitmap) {
+      item = static_cast<bitmap_t>(0);
+    }
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<index_t> dis(0, total - 1);
+
+    while (num_ones > 0) {
+      index_t index = dis(gen);
+
+      bitmap_t& element    = bitmap[index / (8 * sizeof(bitmap_t))];
+      index_t bit_position = index % (8 * sizeof(bitmap_t));
+
+      if (((element >> bit_position) & 1) == 0) {
+        element |= (static_cast<index_t>(1) >> bit_position);
+        num_ones--;
+      }
+    }
+    return res;
+  }
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    std::ostringstream label_stream;
+    label_stream << params;
+    state.SetLabel(label_stream.str());
+
+    auto bitmap =
+      raft::core::bitmap_view<bitmap_t, index_t>(bitmap_d.data(), params.n_rows, params.n_cols);
+
+    auto csr_view = raft::make_device_compressed_structure_view<index_t, index_t, index_t>(
+      indptr_d.data(), indices_d.data(), params.n_rows, params.n_cols, nnz);
+    auto csr = raft::make_device_csr_matrix_view<value_t, index_t>(values_d.data(), csr_view);
+
+    raft::sparse::convert::bitmap_to_csr<bitmap_t, index_t>(handle, bitmap, csr);
+
+    resource::sync_stream(handle);
+    loop_on_state(state, [this, &bitmap, &csr]() {
+      raft::sparse::convert::bitmap_to_csr<bitmap_t, index_t>(handle, bitmap, csr);
+    });
+  }
+
+ protected:
+  const raft::device_resources handle;
+
+  bench_param<index_t> params;
+
+  rmm::device_uvector<bitmap_t> bitmap_d;
+  rmm::device_uvector<index_t> indptr_d;
+  rmm::device_uvector<index_t> indices_d;
+  rmm::device_uvector<value_t> values_d;
+
+  index_t nnz;
+};  // struct BitmapToCsrTest
+
+template <typename index_t>
+const std::vector<bench_param<index_t>> getInputs()
+{
+  std::vector<bench_param<index_t>> param_vec;
+  struct TestParams {
+    index_t m;
+    index_t n;
+    float sparsity;
+  };
+
+  const std::vector<TestParams> params_group = raft::util::itertools::product<TestParams>(
+    {index_t(10), index_t(1024)}, {index_t(1024 * 1024)}, {0.01f, 0.1f, 0.2f, 0.5f});
+
+  param_vec.reserve(params_group.size());
+  for (TestParams params : params_group) {
+    param_vec.push_back(bench_param<index_t>({params.m, params.n, params.sparsity}));
+  }
+  return param_vec;
+}
+
+RAFT_BENCH_REGISTER((BitmapToCsrTest<uint32_t, int, float>), "", getInputs<int>());
+
+}  // namespace raft::bench::sparse

cpp/include/raft/core/bitmap.hpp

@@ -0,0 +1,112 @@
+/*
+ * 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.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> {
+  /**
+   * @brief Create a bitmap view from a device raw pointer.
+   *
+   * @param bitset_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_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);
+  }
+
+ private:
+  index_t rows_;
+  index_t cols_;
+};
+
+/** @} */
+}  // end namespace raft::core

cpp/include/raft/sparse/convert/csr.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-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.
@@ -18,7 +18,10 @@
 
 #pragma once
 
+#include <raft/core/bitmap.hpp>
+#include <raft/core/device_csr_matrix.hpp>
 #include <raft/sparse/convert/detail/adj_to_csr.cuh>
+#include <raft/sparse/convert/detail/bitmap_to_csr.cuh>
 #include <raft/sparse/convert/detail/csr.cuh>
 #include <raft/sparse/csr.hpp>
 
@@ -102,6 +105,32 @@ void adj_to_csr(raft::resources const& handle,
   detail::adj_to_csr(handle, adj, row_ind, num_rows, num_cols, tmp, out_col_ind);
 }
 
+/**
+ * @brief Converts a bitmap matrix into unsorted CSR format matrix.
+ *
+ * @tparam       bitmap_t    Underlying type of the bitmap.
+ * @tparam       index_t     Indexing type used.
+ * @tparam       value_t     Data type of CSR
+ * @tparam       nnz_t       Type of CSR
+ *
+ * @param[in]    handle      RAFT handle
+ * @param[in]    bitmap      input raft::bitmap_view
+ * @param[inout] csr         output raft::device_csr_matrix_view
+ */
+template <typename bitmap_t, typename index_t, typename value_t, typename nnz_t>
+void bitmap_to_csr(raft::resources const& handle,
+                   raft::core::bitmap_view<bitmap_t, index_t> bitmap,
+                   raft::device_csr_matrix_view<value_t, index_t, index_t, nnz_t> csr)
+{
+  auto csr_view = csr.structure_view();
+  detail::bitmap_to_csr(handle,
+                        bitmap.data(),
+                        csr_view.get_n_rows(),
+                        csr_view.get_n_cols(),
+                        csr_view.get_indptr().data(),
+                        csr_view.get_indices().data());
+}
+
 };  // end NAMESPACE convert
 };  // end NAMESPACE sparse
 };  // end NAMESPACE raft