Replace ceildiv with div_round_up

rapidsai · Sep 6, 2024 · fefe4a8 · fefe4a8
1 parent 22f289d
commit fefe4a8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 19 deletions.
diff --git a/cpp/include/raft/core/bitset.cuh b/cpp/include/raft/core/bitset.cuh
@@ -47,7 +47,7 @@ _RAFT_HOST_DEVICE bool bitset_view<bitset_t, index_t>::operator[](const index_t
 
 template <typename bitset_t, typename index_t>
 _RAFT_DEVICE void bitset_view<bitset_t, index_t>::set(const index_t sample_index,
-                                                           bool set_value) const
+                                                      bool set_value) const
 {
   const index_t bit_element = sample_index / bitset_element_size;
   const index_t bit_index   = sample_index % bitset_element_size;
@@ -60,18 +60,12 @@ _RAFT_DEVICE void bitset_view<bitset_t, index_t>::set(const index_t sample_index
   }
 }
 
-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,
                                   index_t bitset_len,
                                   bool default_value)
-  : bitset_{std::size_t(raft::ceildiv(bitset_len, bitset_element_size)),
+  : bitset_{std::size_t(raft::div_rounding_up_safe(bitset_len, bitset_element_size)),
             raft::resource::get_cuda_stream(res)},
     bitset_len_{bitset_len}
 {
@@ -83,26 +77,20 @@ template <typename bitset_t, typename index_t>
 bitset<bitset_t, index_t>::bitset(const raft::resources& res,
                                   index_t bitset_len,
                                   bool default_value)
-  : bitset_{std::size_t(raft::ceildiv(bitset_len, bitset_element_size)),
+  : bitset_{std::size_t(raft::div_rounding_up_safe(bitset_len, bitset_element_size)),
             raft::resource::get_cuda_stream(res)},
     bitset_len_{bitset_len}
 {
   reset(res, default_value);
 }
 
-template <typename bitset_t, typename index_t>
-index_t bitset<bitset_t, index_t>::n_elements() const
-{
-  return raft::ceildiv(bitset_len_, bitset_element_size);
-}
-
 template <typename bitset_t, typename index_t>
 void bitset<bitset_t, index_t>::resize(const raft::resources& res,
                                        index_t new_bitset_len,
                                        bool default_value)
 {
-  auto old_size = raft::ceildiv(bitset_len_, bitset_element_size);
-  auto new_size = raft::ceildiv(new_bitset_len, bitset_element_size);
+  auto old_size = raft::div_rounding_up_safe(bitset_len_, bitset_element_size);
+  auto new_size = raft::div_rounding_up_safe(new_bitset_len, bitset_element_size);
   bitset_.resize(new_size);
   bitset_len_ = new_bitset_len;
   if (old_size < new_size) {

diff --git a/cpp/include/raft/core/bitset.hpp b/cpp/include/raft/core/bitset.hpp
@@ -20,6 +20,7 @@
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
+#include <raft/util/integer_utils.hpp>
 
 namespace raft::core {
 /**
@@ -89,7 +90,10 @@ struct bitset_view {
   /**
    * @brief Get the number of elements used by the bitset representation.
    */
-  inline _RAFT_HOST_DEVICE auto n_elements() const -> index_t;
+  inline _RAFT_HOST_DEVICE auto n_elements() const -> index_t
+  {
+    return raft::div_rounding_up_safe(bitset_len_, bitset_element_size);
+  }
 
   inline auto to_mdspan() -> raft::device_vector_view<bitset_t, index_t>
   {
@@ -173,7 +177,10 @@ struct bitset {
   /**
    * @brief Get the number of elements used by the bitset representation.
    */
-  inline auto n_elements() const -> index_t;
+  inline auto n_elements() const -> index_t
+  {
+    return raft::div_rounding_up_safe(bitset_len_, bitset_element_size);
+  }
 
   /** @brief Get an mdspan view of the current bitset */
   inline auto to_mdspan() -> raft::device_vector_view<bitset_t, index_t>