Make query async

cpp/cmake/rapids_config.cmake

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2018-2024, NVIDIA CORPORATION.
+# Copyright (c) 2018-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
@@ -11,6 +11,9 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 # =============================================================================
+set(rapids-cmake-repo "kingcrimsontianyu/rapids-cmake")
+set(rapids-cmake-branch "bump-bs-threadpool-version-5.0.0")
+
 file(READ "${CMAKE_CURRENT_LIST_DIR}/../../VERSION" _rapids_version)
 if(_rapids_version MATCHES [[^([0-9][0-9])\.([0-9][0-9])\.([0-9][0-9])]])
   set(RAPIDS_VERSION_MAJOR "${CMAKE_MATCH_1}")

cpp/include/kvikio/defaults.hpp

@@ -81,13 +81,15 @@ bool getenv_or(std::string_view env_var_name, bool default_val);
 template <>
 CompatMode getenv_or(std::string_view env_var_name, CompatMode default_val);
 
+using BS_thread_pool = BS::thread_pool<BS::tp::none>;
+
 /**
  * @brief Singleton class of default values used throughout KvikIO.
  *
  */
 class defaults {
  private:
-  BS::thread_pool _thread_pool{get_num_threads_from_env()};
+  BS_thread_pool _thread_pool{get_num_threads_from_env()};
   CompatMode _compat_mode;
   std::size_t _task_size;
   std::size_t _gds_threshold;
@@ -183,7 +185,7 @@ class defaults {
    *
    * @return The the default thread pool instance.
    */
-  [[nodiscard]] static BS::thread_pool& thread_pool();
+  [[nodiscard]] static BS_thread_pool& thread_pool();
 
   /**
    * @brief Get the number of threads in the default thread pool.

cpp/include/kvikio/parallel_operation.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2024, NVIDIA CORPORATION.
+ * 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.
@@ -17,6 +17,7 @@
 
 #include <cassert>
 #include <future>
+#include <memory>
 #include <numeric>
 #include <system_error>
 #include <utility>
@@ -38,6 +39,13 @@ std::future<std::size_t> submit_task(
     [=] { return op(buf, size, file_offset, devPtr_offset); });
 }
 
+template <typename F>
+auto make_copyable_lambda(F&& f)
+{
+  auto sp = std::make_shared<F>(std::forward<F>(f));
+  return [sp]() -> decltype(auto) { return (*sp)(); };
+}
+
 }  // namespace detail
 
 /**
@@ -83,14 +91,15 @@ std::future<std::size_t> parallel_io(F op,
   if (size > 0) { tasks.push_back(detail::submit_task(op, buf, size, file_offset, devPtr_offset)); }
 
   // Finally, we sum the result of all tasks.
-  auto gather_tasks = [](std::vector<std::future<std::size_t>>&& tasks) -> std::size_t {
-    std::size_t ret = 0;
-    for (auto& task : tasks) {
-      ret += task.get();
-    }
-    return ret;
-  };
-  return std::async(std::launch::deferred, gather_tasks, std::move(tasks));
+  auto gather_tasks =
+    detail::make_copyable_lambda([tasks = std::move(tasks)]() mutable -> std::size_t {
+      std::size_t ret = 0;
+      for (auto& task : tasks) {
+        ret += task.get();
+      }
+      return ret;
+    });
+  return defaults::thread_pool().submit_task(std::move(gather_tasks));
 }
 
 }  // namespace kvikio

cpp/include/kvikio/utils.hpp

@@ -15,6 +15,7 @@
  */
 #pragma once
 
+#include <cassert>
 #include <chrono>
 #include <cstring>
 #include <future>
@@ -152,6 +153,7 @@ std::tuple<void*, std::size_t, std::size_t> get_alloc_info(const void* devPtr,
 template <typename T>
 bool is_future_done(const T& future)
 {
+  assert(future.valid());
   return future.wait_for(std::chrono::seconds(0)) != std::future_status::timeout;
 }
 

cpp/src/defaults.cpp

@@ -163,7 +163,7 @@ bool defaults::is_compat_mode_preferred(CompatMode compat_mode) noexcept
 
 bool defaults::is_compat_mode_preferred() { return is_compat_mode_preferred(compat_mode()); }
 
-BS::thread_pool& defaults::thread_pool() { return instance()->_thread_pool; }
+BS_thread_pool& defaults::thread_pool() { return instance()->_thread_pool; }
 
 unsigned int defaults::thread_pool_nthreads() { return thread_pool().get_thread_count(); }
 

cpp/src/file_handle.cpp

@@ -293,11 +293,12 @@ std::future<std::size_t> FileHandle::pread(void* buf,
 
   // Shortcut that circumvent the threadpool and use the POSIX backend directly.
   if (size < gds_threshold) {
-    auto task = [this, ctx, buf, size, file_offset]() -> std::size_t {
-      PushAndPopContext c(ctx);
-      return detail::posix_device_read(_fd_direct_off, buf, size, file_offset, 0);
-    };
-    return std::async(std::launch::deferred, task);
+    PushAndPopContext c(ctx);
+    auto bytes_read = detail::posix_device_read(_fd_direct_off, buf, size, file_offset, 0);
+    std::promise<std::size_t> read_promise;
+    auto read_future = read_promise.get_future();
+    read_promise.set_value(bytes_read);
+    return read_future;
   }
 
   // Let's synchronize once instead of in each task.
@@ -343,11 +344,12 @@ std::future<std::size_t> FileHandle::pwrite(const void* buf,
 
   // Shortcut that circumvent the threadpool and use the POSIX backend directly.
   if (size < gds_threshold) {
-    auto task = [this, ctx, buf, size, file_offset]() -> std::size_t {
-      PushAndPopContext c(ctx);
-      return detail::posix_device_write(_fd_direct_off, buf, size, file_offset, 0);
-    };
-    return std::async(std::launch::deferred, task);
+    PushAndPopContext c(ctx);
+    auto bytes_write = detail::posix_device_write(_fd_direct_off, buf, size, file_offset, 0);
+    std::promise<std::size_t> write_promise;
+    auto write_future = write_promise.get_future();
+    write_promise.set_value(bytes_write);
+    return write_future;
   }
 
   // Let's synchronize once instead of in each task.