Implement host/device/constexpr infimum/supremum functions

include/cuco/detail/utils.hpp

@@ -18,16 +18,16 @@
 #include <cuco/detail/error.hpp>
 #include <cuco/detail/utility/cuda.hpp>
 
-#include <iterator>
-#include <type_traits>
+#include <cuda/std/iterator>
+#include <cuda/std/type_traits>
 
 namespace cuco {
 namespace detail {
 
 template <typename Iterator>
 constexpr inline index_type distance(Iterator begin, Iterator end)
 {
-  using category = typename std::iterator_traits<Iterator>::iterator_category;
+  using category = typename cuda::std::iterator_traits<Iterator>::iterator_category;
   static_assert(std::is_base_of_v<std::random_access_iterator_tag, category>,
                 "Input iterator should be a random access iterator.");
   // `int64_t` instead of arch-dependant `long int`
@@ -48,28 +48,78 @@ constexpr inline index_type distance(Iterator begin, Iterator end)
  * element < value
  */
 template <class ForwardIt, class T>
-constexpr ForwardIt lower_bound(ForwardIt first, ForwardIt last, const T& value)
+__host__ __device__ constexpr ForwardIt lower_bound(ForwardIt first, ForwardIt last, T const& value)
 {
-  using diff_type = typename std::iterator_traits<ForwardIt>::difference_type;
+  auto const ff   = first;
+  using diff_type = typename cuda::std::iterator_traits<ForwardIt>::difference_type;
 
   ForwardIt it{};
-  diff_type count = std::distance(first, last);
+  diff_type count = cuda::std::distance(first, last);
   diff_type step{};
 
   while (count > 0) {
     it   = first;
     step = count / 2;
-    std::advance(it, step);
+    cuda::std::advance(it, step);
 
-    if (static_cast<T>(*it) < value) {
+    if (*it < value) {
       first = ++it;
       count -= step + 1;
-    } else
+    } else {
       count = step;
+    }
   }
 
   return first;
 }
 
+/**
+ * @brief Finds the largest element in the ordered range [first, last) smaller than or equal to
+ * `value`.
+ *
+ * @tparam ForwardIt Type of input iterator
+ * @tparam T Type of `value`
+ *
+ * @param first Iterator defining the start of the range to examine
+ * @param last Iterator defining the start of the range to examine
+ * @param value Value to compare the elements to
+ *
+ * @return Iterator pointing to the infimum value, else `last`
+ */
+template <class ForwardIt, class T>
+__host__ __device__ constexpr ForwardIt infimum(ForwardIt first, ForwardIt last, T const& value)
+{
+  auto it = lower_bound(first, last, value);
+
+  // If lower_bound returns the beginning, and it's not equal to value, then the value is smaller
+  // than all elements.
+  if (it == first && *it != value) { return last; }
+
+  // If lower_bound returned an iterator pointing to a value larger than the given value,
+  // we need to step back to get the next smallest.
+  if (it == last || *it != value) { --it; }
+
+  return it;
+}
+
+/**
+ * @brief Finds the smallest element in the ordered range [first, last) larger than or equal to
+ * `value`.
+ *
+ * @tparam ForwardIt Type of input iterator
+ * @tparam T Type of `value`
+ *
+ * @param first Iterator defining the start of the range to examine
+ * @param last Iterator defining the start of the range to examine
+ * @param value Value to compare the elements to
+ *
+ * @return Iterator pointing to the supremum value, else `last`
+ */
+template <class ForwardIt, class T>
+__host__ __device__ constexpr ForwardIt supremum(ForwardIt first, ForwardIt last, T const& value)
+{
+  return lower_bound(first, last, value);
+}
+
 }  // namespace detail
 }  // namespace cuco

tests/CMakeLists.txt

@@ -53,7 +53,8 @@ ConfigureTest(UTILITY_TEST
     utility/extent_test.cu
     utility/storage_test.cu
     utility/fast_int_test.cu
-    utility/hash_test.cu)
+    utility/hash_test.cu
+    utility/inf_sup_test.cu)
 
 ###################################################################################################
 # - static_set tests ------------------------------------------------------------------------------

tests/utility/inf_sup_test.cu

@@ -0,0 +1,94 @@
+/*
+ * 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.
+ */
+
+#include <utils.hpp>
+
+#include <cuco/detail/__config>
+#include <cuco/detail/utils.hpp>
+#include <cuco/extent.cuh>
+
+#include <catch2/catch_test_macros.hpp>
+
+#include <cstdint>
+#include <cuda/std/array>
+#include <limits>
+
+using T                        = uint64_t;
+__device__ auto constexpr data = cuda::std::array<T, 6>{1, 2, 3, 5, 6, 7};
+
+constexpr T none = std::numeric_limits<T>::max();  // denotes a missing value
+
+template <size_t N>
+constexpr T compute_host_infimum(cuco::experimental::extent<T, N> extent)
+{
+  auto const res = cuco::detail::infimum(data.begin(), data.end(), extent);
+
+  return (res != data.end()) ? *res : none;
+}
+
+template <size_t N>
+constexpr T compute_host_supremum(cuco::experimental::extent<T, N> extent)
+{
+  auto const res = cuco::detail::supremum(data.begin(), data.end(), extent);
+
+  return (res != data.end()) ? *res : none;
+}
+
+TEST_CASE("Infimum computation", "")
+{
+  SECTION("Check if host-generated infimum is correct.")
+  {
+    CHECK(compute_host_infimum(cuco::experimental::extent<T>{0}) == none);
+    CHECK(compute_host_infimum(cuco::experimental::extent<T>{1}) == 1);
+    CHECK(compute_host_infimum(cuco::experimental::extent<T>{4}) == 3);
+    CHECK(compute_host_infimum(cuco::experimental::extent<T>{7}) == 7);
+    CHECK(compute_host_infimum(cuco::experimental::extent<T>{8}) == 7);
+  }
+
+  SECTION("Check if constexpr infimum is correct.")
+  {
+    STATIC_REQUIRE(compute_host_infimum(cuco::experimental::extent<T, 0>{}) == none);
+    STATIC_REQUIRE(compute_host_infimum(cuco::experimental::extent<T, 1>{}) == 1);
+    STATIC_REQUIRE(compute_host_infimum(cuco::experimental::extent<T, 4>{}) == 3);
+    STATIC_REQUIRE(compute_host_infimum(cuco::experimental::extent<T, 7>{}) == 7);
+    STATIC_REQUIRE(compute_host_infimum(cuco::experimental::extent<T, 8>{}) == 7);
+  }
+
+  // TODO device test
+}
+
+TEST_CASE("Supremum computation", "")
+{
+  SECTION("Check if host-generated supremum is correct.")
+  {
+    CHECK(compute_host_supremum(cuco::experimental::extent<T>{0}) == 1);
+    CHECK(compute_host_supremum(cuco::experimental::extent<T>{1}) == 1);
+    CHECK(compute_host_supremum(cuco::experimental::extent<T>{4}) == 5);
+    CHECK(compute_host_supremum(cuco::experimental::extent<T>{7}) == 7);
+    CHECK(compute_host_supremum(cuco::experimental::extent<T>{8}) == none);
+  }
+
+  SECTION("Check if constexpr supremum is correct.")
+  {
+    STATIC_REQUIRE(compute_host_supremum(cuco::experimental::extent<T, 0>{}) == 1);
+    STATIC_REQUIRE(compute_host_supremum(cuco::experimental::extent<T, 1>{}) == 1);
+    STATIC_REQUIRE(compute_host_supremum(cuco::experimental::extent<T, 4>{}) == 5);
+    STATIC_REQUIRE(compute_host_supremum(cuco::experimental::extent<T, 7>{}) == 7);
+    STATIC_REQUIRE(compute_host_supremum(cuco::experimental::extent<T, 8>{}) == none);
+  }
+
+  // TODO device test
+}