Function copy_if

include/hipSYCL/algorithms/algorithm.hpp:192–256

template <class ForwardIt1, class ForwardIt2, class UnaryPredicate>
sycl::event copy_if(sycl::queue &q, util::allocation_group &scratch_allocations,
                    ForwardIt1 first, ForwardIt1 last, ForwardIt2 d_first,
                    UnaryPredicate pred,
                    std::size_t *num_elements_copied = nullptr,
                    const std::vector<sycl::event> &deps = {}) {
  if(first == last) {
    if(num_elements_copied)
      *num_elements_copied = 0;
    return sycl::event{};
  }

  // TODO: We could optimize by switching between 32/64 bit types
  // depending on problem size
  using ScanT = std::size_t;

  auto generator = [=](auto idx, auto effective_group_id,
                       auto effective_global_id, auto problem_size) {
    if(effective_global_id >= problem_size)
      return ScanT{0};

    ForwardIt1 it = first;
    std::advance(it, effective_global_id);
    if(pred(*it))
      return ScanT{1};

    return ScanT{0};
  };

  auto result_processor = [=](auto idx, auto effective_group_id,
                       auto effective_global_id, auto problem_size,
                       auto value) {
    if (effective_global_id < problem_size) {
      ForwardIt2 output = d_first;
      ForwardIt1 input = first;
      std::advance(input, effective_global_id);
      std::advance(output, value);

      bool needs_copy = false;

      if(effective_global_id < problem_size) {
        auto input_value = *input;
        needs_copy = pred(input_value);
        if(needs_copy)
          *output = *input;
      }

      if (effective_global_id == problem_size - 1 && num_elements_copied) {
        ScanT inclusive_scan_result = value;
        // We did an exclusive scan, so if the last element also was copied,
        // we need to add that.
        if(needs_copy)
          ++inclusive_scan_result;
        
        *num_elements_copied = static_cast<std::size_t>(inclusive_scan_result);
      }
    }
  };