Use the new tuning API internally for detail::find::dispatch (#9240)

* add new dispatch tuning API to cub::find

* reviews

Author: gonidelis

cub/benchmarks/bench/find_if/base.cu

@@ -9,6 +9,31 @@
 
 #include <nvbench_helper.cuh>
 
+// %RANGE% TUNE_LOAD ld 0:2:1
+// %RANGE% TUNE_ITEMS_PER_THREAD ipt 7:24:1
+// %RANGE% TUNE_THREADS_PER_BLOCK_POW2 tpb 6:10:1
+
+#if !TUNE_BASE
+#  if TUNE_LOAD == 0
+#    define TUNE_LOAD_MODIFIER cub::LOAD_DEFAULT
+#  elif TUNE_LOAD == 1
+#    define TUNE_LOAD_MODIFIER cub::LOAD_LDG
+#  else // TUNE_LOAD == 2
+#    define TUNE_LOAD_MODIFIER cub::LOAD_CA
+#  endif // TUNE_LOAD
+
+template <typename T>
+struct bench_policy_selector
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(::cuda::compute_capability) const
+    -> cub::detail::find::find_policy
+  {
+    return cub::detail::find::find_policy{
+      (1 << TUNE_THREADS_PER_BLOCK_POW2), cub::Nominal4BItemsToItems<T>(TUNE_ITEMS_PER_THREAD), 4, TUNE_LOAD_MODIFIER};
+  }
+};
+#endif // !TUNE_BASE
+
 template <typename T, typename OffsetT>
 void find_if(nvbench::state& state, nvbench::type_list<T, OffsetT>)
 {
@@ -23,33 +48,27 @@ void find_if(nvbench::state& state, nvbench::type_list<T, OffsetT>)
   thrust::fill(dinput.begin() + mismatch_point, dinput.end(), val);
   thrust::device_vector<OffsetT> d_result(1, thrust::no_init);
 
-  void* d_temp_storage = nullptr;
-  size_t temp_storage_bytes{};
-
   state.add_global_memory_reads<T>(mismatch_point);
   state.add_global_memory_writes<OffsetT>(1);
 
-  cub::DeviceFind::FindIf(
-    d_temp_storage,
-    temp_storage_bytes,
-    thrust::raw_pointer_cast(dinput.data()),
-    thrust::raw_pointer_cast(d_result.data()),
-    cuda::equal_to_value<T>(val),
-    static_cast<OffsetT>(dinput.size()),
-    nullptr);
-
-  thrust::device_vector<uint8_t> temp_storage(temp_storage_bytes, thrust::no_init);
-  d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
-
-  state.exec(nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
-    cub::DeviceFind::FindIf(
-      d_temp_storage,
-      temp_storage_bytes,
+  caching_allocator_t alloc;
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env = cub_bench_env(
+      alloc,
+      launch
+#if !TUNE_BASE
+      ,
+      cuda::execution::tune(bench_policy_selector<T>{})
+#endif // !TUNE_BASE
+    );
+    _CCCL_TRY_CUDA_API(
+      cub::DeviceFind::FindIf,
+      "FindIf failed",
       thrust::raw_pointer_cast(dinput.data()),
       thrust::raw_pointer_cast(d_result.data()),
       cuda::equal_to_value<T>(val),
       static_cast<OffsetT>(dinput.size()),
-      launch.get_stream());
+      env);
   });
 }
 

cub/cub/device/device_find.cuh

@@ -428,9 +428,13 @@ struct DeviceFind
 
     using OffsetT = detail::choose_offset_t<NumItemsT>;
 
-    return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) {
-      return detail::find::dispatch(storage, bytes, d_in, d_out, static_cast<OffsetT>(num_items), scan_op, stream);
-    });
+    using default_policy_selector = detail::find::policy_selector_from_types<detail::it_value_t<InputIteratorT>>;
+
+    return detail::dispatch_with_env_and_tuning<default_policy_selector>(
+      env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) {
+        return detail::find::dispatch(
+          storage, bytes, d_in, d_out, static_cast<OffsetT>(num_items), scan_op, stream, policy_selector);
+      });
   }
 
   //! @rst

cub/test/catch2_test_device_find_env.cu

@@ -11,6 +11,7 @@ struct stream_registry_factory_t;
 
 #include <thrust/device_vector.h>
 
+#include <cuda/functional>
 #include <cuda/iterator>
 
 #include "catch2_test_env_launch_helper.h"
@@ -34,6 +35,21 @@ struct is_greater_than_t
   }
 };
 
+// A policy selector that forces a specific block size, so a test can verify the tuning was applied.
+template <int ThreadsPerBlock>
+struct find_tuning
+{
+  _CCCL_HOST_DEVICE_API constexpr auto operator()(cuda::compute_capability) const -> cub::detail::find::find_policy
+  {
+    return {ThreadsPerBlock, 4, 4, cub::LOAD_LDG};
+  }
+};
+
+using block_size_extracting_predicate_t = block_size_extracting_op<::cuda::always_false>;
+
+using block_sizes =
+  c2h::type_list<cuda::std::integral_constant<unsigned int, 64>, cuda::std::integral_constant<unsigned int, 128>>;
+
 #if TEST_LAUNCH == 0
 
 TEST_CASE("Device FindIf works with default environment", "[find][device]")
@@ -186,3 +202,24 @@ C2H_TEST("Device UpperBound uses environment", "[find][device]")
   c2h::device_vector<int> expected = {1, 2, 3, 4};
   REQUIRE(d_output == expected);
 }
+
+#if TEST_LAUNCH != 1
+C2H_TEST("Device FindIf can be tuned", "[find][device]", block_sizes)
+{
+  constexpr unsigned int target_block_size = c2h::get<0, TestType>::value;
+
+  constexpr int num_items = 1024;
+  auto d_in               = c2h::device_vector<int>(num_items, 0);
+  auto d_out              = c2h::device_vector<int>(1, thrust::no_init);
+  auto d_block_size       = c2h::device_vector<unsigned int>(1, 0);
+
+  block_size_extracting_predicate_t predicate{thrust::raw_pointer_cast(d_block_size.data())};
+
+  auto env = cuda::execution::tune(find_tuning<static_cast<int>(target_block_size)>{});
+
+  device_find_if(d_in.begin(), d_out.begin(), predicate, num_items, env);
+
+  REQUIRE(d_out[0] == num_items); // predicate never matches
+  REQUIRE(d_block_size[0] == target_block_size);
+}
+#endif // TEST_LAUNCH != 1