dynamic bounce buffer size

rapidsai · rapids-bot · Sep 2, 2024 · Sep 2, 2024 · Sep 2, 2024 · Sep 2, 2024
commit 7fdff4c17647748b6f2990fcf4b91bf22abed208
@@ -85,15 +85,29 @@ Set the environment variable `KVIKIO_COMPAT_MODE` to enable/disable compatibilit
   - when running in Windows Subsystem for Linux (WSL).
   - when `/run/udev` isn't readable, which typically happens when running inside a docker image not launched with `--volume /run/udev:/run/udev:ro`.
 
+This setting can also be controlled by `defaults::compat_mode()` and `defaults::compat_mode_reset()`.
+
+
 #### Thread Pool (KVIKIO_NTHREADS)
 KvikIO can use multiple threads for IO automatically. Set the environment variable `KVIKIO_NTHREADS` to the number of threads in the thread pool. If not set, the default value is 1.
 
+This setting can also be controlled by `defaults::thread_pool_nthreads()` and `defaults::thread_pool_nthreads_reset()`.
+
 #### Task Size (KVIKIO_TASK_SIZE)
 KvikIO splits parallel IO operations into multiple tasks. Set the environment variable `KVIKIO_TASK_SIZE` to the maximum task size (in bytes). If not set, the default value is 4194304 (4 MiB).
 
+This setting can also be controlled by `defaults::task_size()` and `defaults::task_size_reset()`.
+
 #### GDS Threshold (KVIKIO_GDS_THRESHOLD)
 In order to improve performance of small IO, `.pread()` and `.pwrite()` implement a shortcut that circumvent the threadpool and use the POSIX backend directly. Set the environment variable `KVIKIO_GDS_THRESHOLD` to the minimum size (in bytes) to use GDS. If not set, the default value is 1048576 (1 MiB).
 
+This setting can also be controlled by `defaults::gds_threshold()` and `defaults::gds_threshold_reset()`.
+
+#### Size of the Bounce Buffer (KVIKIO_GDS_THRESHOLD)
+KvikIO might have to use an intermediate host buffer when copying between file and device memory. Set the environment variable ``KVIKIO_BOUNCE_BUFFER_SIZE`` to size (in bytes) of this "bounce" buffer. If not set, the default value is 16777216 (16 MiB).
+
+This setting can also be controlled by `defaults::bounce_buffer_size()` and `defaults::bounce_buffer_size_reset()`.
+
 
 ## Example
 

@@ -44,14 +44,14 @@ class AllocRetain {
    private:
     AllocRetain* _manager;
     void* _alloc;
-    const std::size_t _size;
+    std::size_t const _size;
 
    public:
     Alloc(AllocRetain* manager, void* alloc, std::size_t size)
       : _manager(manager), _alloc{alloc}, _size{size}
     {
     }
-    Alloc(const Alloc&)            = delete;
+    Alloc(Alloc const&)            = delete;
     Alloc& operator=(Alloc const&) = delete;
     Alloc(Alloc&& o)               = delete;
     Alloc& operator=(Alloc&& o)    = delete;
@@ -65,11 +65,16 @@ class AllocRetain {
   {
     try {
       clear();
-    } catch (const CUfileException& e) {
+    } catch (CUfileException const& e) {
       std::cerr << "~AllocRetain(): " << e.what() << std::endl;
     }
   }
 
+  /**
+   * @brief Free all retained allocations
+   *
+   * NB: The `_mutex` must be taken prior to calling this function, if not called from the dtor.
+   */
   void clear()
   {
     while (!_free_allocs.empty()) {
@@ -78,12 +83,24 @@ class AllocRetain {
     }
   }
 
-  [[nodiscard]] Alloc get()
+  /**
+   * @brief Ensure the size of the retained allocations match `defaults::bounce_buffer_size()`
+   *
+   * NB: `_mutex` must be taken prior to calling this function.
+   */
+  void ensure_alloc_size()
   {
-    const std::lock_guard lock(_mutex);
-    if (_size != defaults::bounce_buffer_size()) {
+    auto const bounce_buffer_size = defaults::bounce_buffer_size();
+    if (_size != bounce_buffer_size) {
+      _size = bounce_buffer_size;
       clear();  // the desired allocation size has changed.
     }
+  }
+
+  [[nodiscard]] Alloc get()
+  {
+    std::lock_guard const lock(_mutex);
+    ensure_alloc_size();
 
     // Check if we have an allocation available
     if (!_free_allocs.empty()) {
@@ -101,10 +118,8 @@ class AllocRetain {
 
   void put(void* alloc, std::size_t size)
   {
-    const std::lock_guard lock(_mutex);
-    if (_size != defaults::bounce_buffer_size()) {
-      clear();  // the desired allocation size has changed.
-    }
+    std::lock_guard const lock(_mutex);
+    ensure_alloc_size();
 
     // If the size of `alloc` matches the sizes of the retained allocations,
     // it is added to the set of free allocation otherwise it is freed.
@@ -121,7 +136,7 @@ class AllocRetain {
     return _instance;
   }
 
-  AllocRetain(const AllocRetain&)            = delete;
+  AllocRetain(AllocRetain const&)            = delete;
   AllocRetain& operator=(AllocRetain const&) = delete;
   AllocRetain(AllocRetain&& o)               = delete;
   AllocRetain& operator=(AllocRetain&& o)    = delete;

diff --git a/docs/source/runtime_settings.rst b/docs/source/runtime_settings.rst
@@ -10,17 +10,29 @@ Set the environment variable ``KVIKIO_COMPAT_MODE`` to enable/disable compatibil
   * when running in Windows Subsystem for Linux (WSL).
   * when ``/run/udev`` isn't readable, which typically happens when running inside a docker image not launched with ``--volume /run/udev:/run/udev:ro``.
 
+This setting can also be controlled by :py:func:`kvikio.defaults.compat_mode`, :py:func:`kvikio.defaults.compat_mode_reset`, and :py:func:`kvikio.defaults.set_compat_mode`.
+
 
 Thread Pool ``KVIKIO_NTHREADS``
 -------------------------------
 KvikIO can use multiple threads for IO automatically. Set the environment variable ``KVIKIO_NTHREADS`` to the number of threads in the thread pool. If not set, the default value is 1.
 
+This setting can also be controlled by :py:func:`kvikio.defaults.get_num_threads`, :py:func:`kvikio.defaults.num_threads_reset`, and :py:func:`kvikio.defaults.set_num_threads`.
 
 Task Size ``KVIKIO_TASK_SIZE``
 ------------------------------
 KvikIO splits parallel IO operations into multiple tasks. Set the environment variable ``KVIKIO_TASK_SIZE`` to the maximum task size (in bytes). If not set, the default value is 4194304 (4 MiB).
 
+This setting can also be controlled by :py:func:`kvikio.defaults.task_size`, :py:func:`kvikio.defaults.task_size_reset`, and :py:func:`kvikio.defaults.set_task_size`.
 
 GDS Threshold ``KVIKIO_GDS_THRESHOLD``
 --------------------------------------
 In order to improve performance of small IO, ``.pread()`` and ``.pwrite()`` implement a shortcut that circumvent the threadpool and use the POSIX backend directly. Set the environment variable ``KVIKIO_GDS_THRESHOLD`` to the minimum size (in bytes) to use GDS. If not set, the default value is 1048576 (1 MiB).
+
+This setting can also be controlled by :py:func:`kvikio.defaults.gds_threshold`, :py:func:`kvikio.defaults.gds_threshold_reset`, and :py:func:`kvikio.defaults.set_gds_threshold`.
+
+Size of the Bounce Buffer ``KVIKIO_BOUNCE_BUFFER_SIZE``
+-------------------------------------------------------
+KvikIO might have to use an intermediate host buffer when copying between file and device memory. Set the environment variable ``KVIKIO_BOUNCE_BUFFER_SIZE`` to size (in bytes) of this "bounce" buffer. If not set, the default value is 16777216 (16 MiB).
+
+This setting can also be controlled by :py:func:`kvikio.defaults.bounce_buffer_size`, :py:func:`kvikio.defaults.bounce_buffer_size_reset`, and :py:func:`kvikio.defaults.set_bounce_buffer_size`.
@@ -23,8 +23,8 @@ def compat_mode() -> bool:
     - when `/run/udev` isn't readable, which typically happens when running inside
     a docker image not launched with `--volume /run/udev:/run/udev:ro`
 
-    Return
-    ------
+    Returns
+    -------
     bool
         Whether KvikIO is running in compatibility mode or not.
     """
@@ -68,8 +68,8 @@ def get_num_threads() -> int:
     Set the default value using `num_threads_reset()` or by setting the
     `KVIKIO_NTHREADS` environment variable. If not set, the default value is 1.
 
-    Return
-    ------
+    Returns
+    -------
     nthreads: int
         The number of threads in the current thread pool.
     """
@@ -119,8 +119,8 @@ def task_size() -> int:
     the `KVIKIO_TASK_SIZE` environment variable. If not set,
     the default value is 4 MiB.
 
-    Return
-    ------
+    Returns
+    -------
     nbytes: int
         The default task size in bytes.
     """
@@ -166,8 +166,8 @@ def gds_threshold() -> int:
     `KVIKIO_GDS_THRESHOLD` environment variable. If not set, the default
     value is 1 MiB.
 
-    Return
-    ------
+    Returns
+    -------
     nbytes : int
         The default GDS threshold size in bytes.
     """