bounce_buffer.hpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <stack>
#include <utility>
 
#include <kvikio/defaults.hpp>
 
namespace kvikio {
 
class PageAlignedAllocator {
 public:
  void* allocate(std::size_t size);
 
  void deallocate(void* buffer, std::size_t size);
};
 
class CudaPinnedAllocator {
 public:
  void* allocate(std::size_t size);
 
  void deallocate(void* buffer, std::size_t size);
};
 
class CudaPageAlignedPinnedAllocator {
 public:
  void* allocate(std::size_t size);
 
  void deallocate(void* buffer, std::size_t size);
};
 
template <typename Allocator = CudaPinnedAllocator>
class BounceBufferPool {
 private:
  mutable std::mutex _mutex{};
  // Stack of free allocations (LIFO for cache locality)
  std::stack<void*> _free_buffers{};
  // The size of each allocation in `_free_buffers`
  std::size_t _buffer_size{defaults::bounce_buffer_size()};
  Allocator _allocator{};
 
 public:
  class Buffer {
   private:
    BounceBufferPool* _pool;
    void* _buffer;
    std::size_t _size;
 
   public:
    Buffer(BounceBufferPool<Allocator>* pool, void* buffer, std::size_t size);
    Buffer(Buffer const&)            = delete;
    Buffer& operator=(Buffer const&) = delete;
    Buffer(Buffer&& o) noexcept;
    Buffer& operator=(Buffer&& o) noexcept;
    ~Buffer() noexcept;
    void* get() const noexcept;
    void* get(std::ptrdiff_t offset) const noexcept;
    std::size_t size() const noexcept;
  };
 
  BounceBufferPool() = default;
 
  // Notice, we do not clear the allocations at destruction thus the allocations leaks
  // at exit. We do this because `BounceBufferPool::instance()` stores the allocations in a
  // static stack that are destructed below main, which is not allowed in CUDA:
  // <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#initialization>
  ~BounceBufferPool() noexcept = default;
 
 private:
  void _deallocate_buffers(std::stack<void*>& buffers, std::size_t buffer_size);
 
  [[nodiscard]] std::pair<std::stack<void*>, std::size_t> _detach_free_buffers();
 
  [[nodiscard]] std::pair<std::stack<void*>, std::size_t> _ensure_buffer_size();
 
 public:
  [[nodiscard]] Buffer get();
 
  void put(void* buffer, std::size_t size) noexcept;
 
  std::size_t clear();
 
  std::size_t num_free_buffers() const;
 
  std::size_t buffer_size() const;
 
  KVIKIO_EXPORT static BounceBufferPool& instance();
 
  BounceBufferPool(BounceBufferPool const&)            = delete;
  BounceBufferPool& operator=(BounceBufferPool const&) = delete;
  BounceBufferPool(BounceBufferPool&& o)               = delete;
  BounceBufferPool& operator=(BounceBufferPool&& o)    = delete;
};
 
using PageAlignedBounceBufferPool = BounceBufferPool<PageAlignedAllocator>;
 
using CudaPinnedBounceBufferPool = BounceBufferPool<CudaPinnedAllocator>;
 
using CudaPageAlignedPinnedBounceBufferPool = BounceBufferPool<CudaPageAlignedPinnedAllocator>;
}  // namespace kvikio