parallel_operation.hpp

/*
 * Copyright (c) 2021-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#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/error.hpp>
#include <kvikio/nvtx.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,
                                     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 defaults::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,
  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 defaults::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,
                                     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);
  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, 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, 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), call_idx, nvtx_color);
}
 
}  // namespace kvikio