Merge remote-tracking branch 'upstream/branch-25.04' into tom/retry-http

cpp/include/kvikio/defaults.hpp

@@ -22,11 +22,10 @@
 #include <stdexcept>
 #include <string>
 
-#include <BS_thread_pool.hpp>
-
 #include <kvikio/compat_mode.hpp>
 #include <kvikio/http_status_codes.hpp>
 #include <kvikio/shim/cufile.hpp>
+#include <kvikio/threadpool_wrapper.hpp>
 
 /**
  * @brief KvikIO namespace.
@@ -64,7 +63,7 @@ std::vector<int> getenv_or(std::string_view env_var_name, std::vector<int> defau
  */
 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;
@@ -162,7 +161,7 @@ class defaults {
    *
    * @return 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/file_handle.hpp

@@ -56,7 +56,9 @@ class FileHandle {
   friend class CompatModeManager;
 
  public:
-  static constexpr mode_t m644 = S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH;
+  // 644 is a common setting of Unix file permissions: read and write for owner, read-only for group
+  // and others.
+  static constexpr mode_t m644 = S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH;
   FileHandle() noexcept        = default;
 
   /**

cpp/include/kvikio/nvtx.hpp

@@ -62,7 +62,7 @@ using nvtx_registered_string_type = nvtx3::registered_string_in<libkvikio_domain
     }                                                                                             \
   }
 #define KVIKIO_NVTX_SCOPED_RANGE_IMPL_2(message, payload) \
-  KVIKIO_NVTX_SCOPED_RANGE_IMPL_3(message, payload, kvikio::nvtx_manager::default_color())
+  KVIKIO_NVTX_SCOPED_RANGE_IMPL_3(message, payload, kvikio::NvtxManager::default_color())
 #define KVIKIO_NVTX_SCOPED_RANGE_SELECTOR(_1, _2, _3, NAME, ...) NAME
 #define KVIKIO_NVTX_SCOPED_RANGE_IMPL(...)                                         \
   KVIKIO_NVTX_SCOPED_RANGE_SELECTOR(                                               \
@@ -85,9 +85,9 @@ using nvtx_color_type = int;
 /**
  * @brief Utility singleton class for NVTX annotation.
  */
-class nvtx_manager {
+class NvtxManager {
  public:
-  static nvtx_manager& instance() noexcept;
+  static NvtxManager& instance() noexcept;
 
   /**
    * @brief Return the default color.
@@ -114,13 +114,13 @@ class nvtx_manager {
    */
   static void rename_current_thread(std::string_view new_name) noexcept;
 
-  nvtx_manager(nvtx_manager const&)            = delete;
-  nvtx_manager& operator=(nvtx_manager const&) = delete;
-  nvtx_manager(nvtx_manager&&)                 = delete;
-  nvtx_manager& operator=(nvtx_manager&&)      = delete;
+  NvtxManager(NvtxManager const&)            = delete;
+  NvtxManager& operator=(NvtxManager const&) = delete;
+  NvtxManager(NvtxManager&&)                 = delete;
+  NvtxManager& operator=(NvtxManager&&)      = delete;
 
  private:
-  nvtx_manager() = default;
+  NvtxManager() = default;
 };
 
 /**

cpp/include/kvikio/parallel_operation.hpp

@@ -46,7 +46,7 @@ inline const std::pair<const nvtx_color_type&, std::uint64_t> get_next_color_and
 {
   static std::atomic_uint64_t call_counter{1ull};
   auto call_idx    = call_counter.fetch_add(1ull, std::memory_order_relaxed);
-  auto& nvtx_color = nvtx_manager::get_color_by_index(call_idx);
+  auto& nvtx_color = NvtxManager::get_color_by_index(call_idx);
   return {nvtx_color, call_idx};
 }
 
@@ -57,17 +57,10 @@ std::future<std::size_t> submit_task(F op,
                                      std::size_t file_offset,
                                      std::size_t devPtr_offset,
                                      std::uint64_t nvtx_payload = 0ull,
-                                     nvtx_color_type nvtx_color = nvtx_manager::default_color())
+                                     nvtx_color_type nvtx_color = NvtxManager::default_color())
 {
   return defaults::thread_pool().submit_task([=] {
     KVIKIO_NVTX_SCOPED_RANGE("task", nvtx_payload, nvtx_color);
-
-    // Rename the worker thread in the thread pool to improve clarity from nsys-ui.
-    // Note: This NVTX feature is currently not supported by nsys-ui.
-    thread_local std::once_flag call_once_per_thread;
-    std::call_once(call_once_per_thread,
-                   [] { nvtx_manager::rename_current_thread("thread pool"); });
-
     return op(buf, size, file_offset, devPtr_offset);
   });
 }
@@ -94,7 +87,7 @@ std::future<std::size_t> parallel_io(F op,
                                      std::size_t task_size,
                                      std::size_t devPtr_offset,
                                      std::uint64_t call_idx     = 0,
-                                     nvtx_color_type nvtx_color = nvtx_manager::default_color())
+                                     nvtx_color_type nvtx_color = NvtxManager::default_color())
 {
   if (task_size == 0) { throw std::invalid_argument("`task_size` cannot be zero"); }
 

cpp/include/kvikio/threadpool_wrapper.hpp

@@ -0,0 +1,57 @@
+/*
+ * Copyright (c) 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 <functional>
+
+#include <BS_thread_pool.hpp>
+
+#include <kvikio/nvtx.hpp>
+
+namespace kvikio {
+
+template <typename pool_type>
+class thread_pool_wrapper : public pool_type {
+ public:
+  /**
+   * @brief Construct a new thread pool wrapper, and invoke a pre-defined initialization function in
+   * each worker thread.
+   *
+   * @param nthreads The number of threads to use.
+   */
+  thread_pool_wrapper(unsigned int nthreads) : pool_type{nthreads, worker_thread_init_func} {}
+
+  /**
+   * @brief Reset the number of threads in the thread pool, and invoke a pre-defined initialization
+   * function in each worker thread.
+   *
+   * @param nthreads The number of threads to use.
+   */
+  void reset(unsigned int nthreads) { pool_type::reset(nthreads, worker_thread_init_func); }
+
+ private:
+  inline static std::function<void()> worker_thread_init_func{[] {
+    KVIKIO_NVTX_SCOPED_RANGE("worker thread init", 0, NvtxManager::default_color());
+    // Rename the worker thread in the thread pool to improve clarity from nsys-ui.
+    // Note: This NVTX feature is currently not supported by nsys-ui.
+    NvtxManager::rename_current_thread("thread pool");
+  }};
+};
+
+using BS_thread_pool = thread_pool_wrapper<BS::thread_pool>;
+
+}  // namespace kvikio

cpp/src/bounce_buffer.cpp

@@ -74,7 +74,8 @@ AllocRetain::Alloc AllocRetain::get()
   // If no available allocation, allocate and register a new one
   void* alloc{};
   // Allocate page-locked host memory
-  CUDA_DRIVER_TRY(cudaAPI::instance().MemHostAlloc(&alloc, _size, CU_MEMHOSTREGISTER_PORTABLE));
+  // Under unified addressing, host memory allocated this way is automatically portable and mapped.
+  CUDA_DRIVER_TRY(cudaAPI::instance().MemHostAlloc(&alloc, _size, CU_MEMHOSTALLOC_PORTABLE));
   return Alloc(this, alloc, _size);
 }
 

cpp/src/defaults.cpp

@@ -165,7 +165,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/nvtx.cpp

@@ -26,13 +26,13 @@
 
 namespace kvikio {
 
-nvtx_manager& nvtx_manager::instance() noexcept
+NvtxManager& NvtxManager::instance() noexcept
 {
-  static nvtx_manager _instance;
+  static NvtxManager _instance;
   return _instance;
 }
 
-const nvtx_color_type& nvtx_manager::default_color() noexcept
+const nvtx_color_type& NvtxManager::default_color() noexcept
 {
 #ifdef KVIKIO_CUDA_FOUND
   static nvtx_color_type default_color{nvtx3::argb{0, 255, 255, 255}};
@@ -43,7 +43,7 @@ const nvtx_color_type& nvtx_manager::default_color() noexcept
 #endif
 }
 
-const nvtx_color_type& nvtx_manager::get_color_by_index(std::uint64_t idx) noexcept
+const nvtx_color_type& NvtxManager::get_color_by_index(std::uint64_t idx) noexcept
 {
 #ifdef KVIKIO_CUDA_FOUND
   constexpr std::size_t num_color{16};
@@ -72,7 +72,7 @@ const nvtx_color_type& nvtx_manager::get_color_by_index(std::uint64_t idx) noexc
 #endif
 }
 
-void nvtx_manager::rename_current_thread(std::string_view new_name) noexcept
+void NvtxManager::rename_current_thread(std::string_view new_name) noexcept
 {
 #ifdef KVIKIO_CUDA_FOUND
   auto tid = syscall(SYS_gettid);

cpp/tests/test_basic_io.cpp

@@ -14,8 +14,8 @@
  * limitations under the License.
  */
 
+#include <kvikio/defaults.hpp>
 #include <kvikio/file_handle.hpp>
-#include "kvikio/defaults.hpp"
 #include "utils.hpp"
 
 using namespace kvikio::test;

docs/source/install.rst

@@ -36,6 +36,18 @@ Install the **nightly release** from the ``rapidsai-nightly`` channel like:
 
   If the nightly install doesn't work, set ``channel_priority: flexible`` in your ``.condarc``.
 
+
+PyPI
+----
+
+KvikIO is also available on PyPI. Install the latest release like:
+
+.. code-block::
+
+  pip install kvikio-cu11 # for CUDA 11
+  pip install kvikio-cu12 # for CUDA 12
+
+
 Build from source
 -----------------