parallel_operation.hpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <atomic>
#include <cassert>
#include <future>
#include <memory>
#include <numeric>
#include <system_error>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <kvikio/defaults.hpp>
#include <kvikio/detail/nvtx.hpp>
#include <kvikio/error.hpp>
#include <kvikio/threadpool_wrapper.hpp>
#include <kvikio/utils.hpp>
 
namespace kvikio {
 
namespace detail {
 
template <typename F>
auto make_copyable_lambda(F op)
{
  // Create the callable on the heap by moving from op. Use a shared pointer to manage its lifetime.
  auto sp = std::make_shared<F>(std::move(op));
 
  // Use the copyable closure as the proxy of the move-only callable.
  return
    [sp](auto&&... args) -> decltype(auto) { return (*sp)(std::forward<decltype(args)>(args)...); };
}
 
inline const std::pair<const nvtx_color_type&, std::uint64_t> get_next_color_and_call_idx() noexcept
{
  static std::atomic_uint64_t call_counter{1ull};
  auto call_idx    = call_counter.fetch_add(1ull, std::memory_order_relaxed);
  auto& nvtx_color = NvtxManager::get_color_by_index(call_idx);
  return {nvtx_color, call_idx};
}
 
template <typename F, typename T>
std::future<std::size_t> submit_task(F op,
                                     T buf,
                                     std::size_t size,
                                     std::size_t file_offset,
                                     std::size_t devPtr_offset,
                                     ThreadPool* thread_pool    = &defaults::thread_pool(),
                                     std::uint64_t nvtx_payload = 0ull,
                                     nvtx_color_type nvtx_color = NvtxManager::default_color())
{
  static_assert(std::is_invocable_r_v<std::size_t,
                                      decltype(op),
                                      decltype(buf),
                                      decltype(size),
                                      decltype(file_offset),
                                      decltype(devPtr_offset)>);
 
  return thread_pool->submit_task([=] {
    KVIKIO_NVTX_SCOPED_RANGE("task", nvtx_payload, nvtx_color);
    return op(buf, size, file_offset, devPtr_offset);
  });
}
 
template <typename F>
std::future<std::size_t> submit_move_only_task(
  F op_move_only,
  ThreadPool* thread_pool    = &defaults::thread_pool(),
  std::uint64_t nvtx_payload = 0ull,
  nvtx_color_type nvtx_color = NvtxManager::default_color())
{
  static_assert(std::is_invocable_r_v<std::size_t, F>);
  auto op_copyable = make_copyable_lambda(std::move(op_move_only));
  return thread_pool->submit_task([=] {
    KVIKIO_NVTX_SCOPED_RANGE("task", nvtx_payload, nvtx_color);
    return op_copyable();
  });
}
 
}  // namespace detail
 
template <typename F, typename T>
std::future<std::size_t> parallel_io(F op,
                                     T buf,
                                     std::size_t size,
                                     std::size_t file_offset,
                                     std::size_t task_size,
                                     std::size_t devPtr_offset,
                                     ThreadPool* thread_pool    = &defaults::thread_pool(),
                                     std::uint64_t call_idx     = 0,
                                     nvtx_color_type nvtx_color = NvtxManager::default_color())
{
  KVIKIO_EXPECT(task_size > 0, "`task_size` must be positive", std::invalid_argument);
  KVIKIO_EXPECT(thread_pool != nullptr, "The thread pool must not be nullptr");
  static_assert(std::is_invocable_r_v<std::size_t,
                                      decltype(op),
                                      decltype(buf),
                                      decltype(size),
                                      decltype(file_offset),
                                      decltype(devPtr_offset)>);
 
  // Single-task guard
  if (task_size >= size || get_page_size() >= size) {
    return detail::submit_task(
      op, buf, size, file_offset, devPtr_offset, thread_pool, call_idx, nvtx_color);
  }
 
  std::vector<std::future<std::size_t>> tasks;
  tasks.reserve(size / task_size);
 
  // 1) Submit all tasks but the last one. These are all `task_size` sized tasks.
  while (size > task_size) {
    tasks.push_back(detail::submit_task(
      op, buf, task_size, file_offset, devPtr_offset, thread_pool, call_idx, nvtx_color));
    file_offset += task_size;
    devPtr_offset += task_size;
    size -= task_size;
  }
 
  // 2) Submit the last task, which consists of performing the last I/O and waiting the previous
  // tasks.
  auto last_task = [=, tasks = std::move(tasks)]() mutable -> std::size_t {
    auto ret = op(buf, size, file_offset, devPtr_offset);
    for (auto& task : tasks) {
      ret += task.get();
    }
    return ret;
  };
  return detail::submit_move_only_task(std::move(last_task), thread_pool, call_idx, nvtx_color);
}
 
}  // namespace kvikio