Fix target counting in strings char-parallel replace

cpp/src/strings/replace/replace.cu

@@ -238,6 +238,31 @@ struct replace_parallel_chars_fn {
   cudf::size_type maxrepl;
 };
 
+template <int64_t block_size, size_type bytes_per_thread>
+CUDF_KERNEL void count_targets_kernel(replace_parallel_chars_fn fn,
+                                      int64_t chars_bytes,
+                                      int64_t* d_output)
+{
+  auto const idx      = cudf::detail::grid_1d::global_thread_id();
+  auto const byte_idx = static_cast<int64_t>(idx) * bytes_per_thread;
+  auto const lane_idx = static_cast<cudf::size_type>(threadIdx.x);
+
+  using block_reduce = cub::BlockReduce<int64_t, block_size>;
+  __shared__ typename block_reduce::TempStorage temp_storage;
+
+  int64_t count = 0;
+  // each thread processes multiple bytes
+  for (auto i = byte_idx; (i < (byte_idx + bytes_per_thread)) && (i < chars_bytes); ++i) {
+    count += fn.has_target(i);
+  }
+  auto const total = block_reduce(temp_storage).Reduce(count, cub::Sum());
+
+  if ((lane_idx == 0) && (total > 0)) {
+    cuda::atomic_ref<int64_t, cuda::thread_scope_block> ref{*d_output};
+    ref.fetch_add(total, cuda::std::memory_order_relaxed);
+  }
+}
+
 std::unique_ptr<column> replace_character_parallel(strings_column_view const& input,
                                                    string_view const& d_target,
                                                    string_view const& d_replacement,
@@ -260,10 +285,14 @@ std::unique_ptr<column> replace_character_parallel(strings_column_view const& in
 
   // Count the number of targets in the entire column.
   // Note this may over-count in the case where a target spans adjacent strings.
-  auto target_count = thrust::count_if(rmm::exec_policy_nosync(stream),
-                                       thrust::make_counting_iterator<int64_t>(0),
-                                       thrust::make_counting_iterator<int64_t>(chars_bytes),
-                                       [fn] __device__(int64_t idx) { return fn.has_target(idx); });
+  rmm::device_scalar<int64_t> d_target_count(0, stream);
+  constexpr int64_t block_size         = 512;
+  constexpr size_type bytes_per_thread = 4;
+  auto const num_blocks                = util::div_rounding_up_safe(
+    util::div_rounding_up_safe(chars_bytes, static_cast<int64_t>(bytes_per_thread)), block_size);
+  count_targets_kernel<block_size, bytes_per_thread>
+    <<<num_blocks, block_size, 0, stream.value()>>>(fn, chars_bytes, d_target_count.data());
+  auto target_count = d_target_count.value(stream);
 
   // Create a vector of every target position in the chars column.
   // These may also include overlapping targets which will be resolved later.