transform.hpp

// SPDX-FileCopyrightText: Intel Corporation
//
// SPDX-License-Identifier: BSD-3-Clause
#pragma once
 
#include <dr/sp/detail.hpp>
#include <dr/sp/init.hpp>
#include <dr/sp/util.hpp>
 
namespace dr::sp {
 
template <class ExecutionPolicy>
auto transform(ExecutionPolicy &&policy, dr::distributed_range auto &&in,
               dr::distributed_iterator auto out, auto &&fn) {
 
  static_assert( // currently only one policy supported
      std::is_same_v<std::remove_cvref_t<ExecutionPolicy>, device_policy>);
 
  std::vector<sycl::event> events;
  using OutT = typename decltype(out)::value_type;
  std::vector<void *> buffers;
  const auto out_end = out + rng::size(in);
 
  for (auto &&[in_seg, out_seg] :
       views::zip(in, rng::subrange(out, out_end)).zipped_segments()) {
    auto in_device = policy.get_devices()[in_seg.rank()];
    auto &&q = __detail::queue(dr::ranges::rank(in_seg));
    const std::size_t seg_size = rng::size(in_seg);
    assert(seg_size == rng::size(out_seg));
    auto local_in_seg = __detail::local(in_seg);
 
    if (in_seg.rank() == out_seg.rank()) {
      auto local_out_seg = __detail::local(out_seg);
      events.emplace_back(q.parallel_for(seg_size, [=](auto idx) {
        local_out_seg[idx] = fn(local_in_seg[idx]);
      }));
    } else {
      OutT *buffer =
          sycl::malloc_device<OutT>(seg_size, in_device, dr::sp::context());
      buffers.push_back(buffer);
 
      sycl::event compute_event = q.parallel_for(
          seg_size, [=](auto idx) { buffer[idx] = fn(local_in_seg[idx]); });
      events.emplace_back(q.copy(buffer, __detail::local(out_seg.begin()),
                                 seg_size, compute_event));
    }
  }
  __detail::wait(events);
 
  for (auto *b : buffers)
    sycl::free(b, dr::sp::context());
 
  return rng::unary_transform_result<decltype(rng::end(in)), decltype(out_end)>{
      rng::end(in), out_end};
}
 
template <dr::distributed_range R, dr::distributed_iterator Iter, typename Fn>
auto transform(R &&in, Iter out, Fn &&fn) {
  return transform(dr::sp::par_unseq, std::forward<R>(in),
                   std::forward<Iter>(out), std::forward<Fn>(fn));
}
 
template <typename ExecutionPolicy, dr::distributed_iterator Iter1,
          dr::distributed_iterator Iter2, typename Fn>
auto transform(ExecutionPolicy &&policy, Iter1 in_begin, Iter1 in_end,
               Iter2 out_end, Fn &&fn) {
  return transform(
      std::forward<ExecutionPolicy>(policy),
      rng::subrange(std::forward<Iter1>(in_begin), std::forward<Iter1>(in_end)),
      std::forward<Iter2>(out_end), std::forward<Fn>(fn));
}
 
template <dr::distributed_iterator Iter1, dr::distributed_iterator Iter2,
          typename Fn>
auto transform(Iter1 in_begin, Iter1 in_end, Iter2 out_end, Fn &&fn) {
  return transform(dr::sp::par_unseq, std::forward<Iter1>(in_begin),
                   std::forward<Iter1>(in_end), std::forward<Iter2>(out_end),
                   std::forward<Fn>(fn));
}
 
} // namespace dr::sp