bounce_buffer.hpp

/*
 * 版权所有 (c) 2024-2025, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可证 2.0 版（以下简称“许可证”）获得许可；
 * 除非遵守许可证的规定，否则不得使用此文件。
 * 您可以在以下位置获取许可证副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则软件
 * 根据许可证分发是基于“按原样”的基础上进行的，
 * 不带任何形式的明示或暗示保证或条件。
 * 有关管理权限和限制的特定语言，请参见许可证。
 *
 */
#pragma once

#include <stack>

#include <kvikio/defaults.hpp>

namespace kvikio {

class AllocRetain {
 private
 std::mutex _mutex{};
 // Stack of free allocations
 std::stack<void*> _free_allocs{};
 // The size of each allocation in `_free_allocs`
 std::size_t _size{defaults::bounce_buffer_size()};

 public
 class Alloc {
 private
 AllocRetain* _manager;
 void* _alloc;
 std::size_t const _size;

 public
 Alloc(AllocRetain* manager, void* alloc, std::size_t size);
 Alloc(Alloc const&) = delete;
 Alloc& operator=(Alloc const&) = delete;
 Alloc(Alloc&& o) = delete;
 Alloc& operator=(Alloc&& o) = delete;
 ~Alloc() noexcept;
 void* get() noexcept;
 void* get(std::ptrdiff_t offset) noexcept;
 std::size_t size() noexcept;
 };

 AllocRetain() = default;

 // Notice, we do not clear the allocations at destruction thus the allocations leaks
 // at exit. We do this because `AllocRetain::instance()` stores the allocations in a
 // static stack that are destructed below main, which is not allowed in CUDA:
 // <https://docs.nvda.net.cn/cuda/cuda-c-programming-guide/index.html#initialization>
 ~AllocRetain() noexcept = default;

 private
 std::size_t _clear();

 void _ensure_alloc_size();

 public
 [[nodiscard]] Alloc get();

 void put(void* alloc, std::size_t size);

 std::size_t clear();

 KVIKIO_EXPORT static AllocRetain& instance();

 AllocRetain(AllocRetain const&) = delete;
 AllocRetain& operator=(AllocRetain const&) = delete;
 AllocRetain(AllocRetain&& o) = delete;
 AllocRetain& operator=(AllocRetain&& o) = delete;
};

} // namespace kvikio