buffer_resource.hpp

 
#pragma once
 
#include <array>
#include <atomic>
#include <cstdint>
#include <memory>
#include <mutex>
#include <optional>
#include <ranges>
#include <unordered_map>
#include <utility>
 
#include <cuda/memory_resource>
 
#include <rmm/cuda_stream_pool.hpp>
 
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/host_memory_resource.hpp>
#include <rapidsmpf/memory/memory_reservation.hpp>
#include <rapidsmpf/memory/pinned_memory_resource.hpp>
#include <rapidsmpf/memory/resource_types.hpp>
#include <rapidsmpf/memory/spill_manager.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/statistics.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf {
 
enum class AllowOverbooking : bool {
    NO,  
    YES,  
};
 
class BufferResource : public std::enable_shared_from_this<BufferResource> {
  public:
    [[nodiscard]] static std::shared_ptr<BufferResource> create(
        cuda::mr::any_resource<cuda::mr::device_accessible> device_mr,
        std::optional<PinnedPoolProperties> pinned_pool_properties = PinnedMemoryDisabled,
        std::unordered_map<MemoryType, std::int64_t> memory_limits = {},
        std::optional<Duration> periodic_spill_check = std::chrono::milliseconds{1},
        std::shared_ptr<rmm::cuda_stream_pool> stream_pool = std::make_shared<
            rmm::cuda_stream_pool>(16, rmm::cuda_stream::flags::non_blocking),
        std::shared_ptr<Statistics> statistics = Statistics::disabled()
    );
 
    static std::shared_ptr<BufferResource> from_options(
        cuda::mr::any_resource<cuda::mr::device_accessible> mr,
        config::Options options,
        std::shared_ptr<Statistics> statistics = Statistics::disabled()
    );
 
    ~BufferResource() noexcept = default;
 
    BufferResource(BufferResource const&) = delete;
    BufferResource(BufferResource&&) = delete;
    BufferResource& operator=(BufferResource const&) = delete;
    BufferResource& operator=(BufferResource&&) = delete;
 
    [[nodiscard]] rmm::device_async_resource_ref device_mr() noexcept;
 
    [[nodiscard]] RmmResourceAdaptor& device_mr_adaptor() noexcept;
 
    [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept;
 
    [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr();
 
    [[nodiscard]] std::optional<PinnedMemoryResource> try_pinned_mr() const;
 
    [[nodiscard]] std::int64_t memory_available(MemoryType mem_type) const noexcept;
 
    void set_memory_limit(MemoryType mem_type, std::int64_t limit) noexcept;
 
    [[nodiscard]] std::size_t memory_reserved(MemoryType mem_type) const {
        return memory_reserved_[static_cast<std::size_t>(mem_type)];
    }
 
    std::pair<MemoryReservation, std::size_t> reserve(
        MemoryType mem_type, std::size_t size, AllowOverbooking allow_overbooking
    );
 
    MemoryReservation reserve_device_memory_and_spill(
        std::size_t size, AllowOverbooking allow_overbooking
    );
 
    template <std::ranges::input_range Range>
        requires std::convertible_to<std::ranges::range_value_t<Range>, MemoryType>
    [[nodiscard]] MemoryReservation reserve_or_fail(std::size_t size, Range mem_types) {
        // try to reserve memory from the given order
        for (auto const& mem_type : mem_types) {
            if (mem_type == MemoryType::PINNED_HOST && !pinned_mr_.has_value()) {
                // Pinned host memory is only available if the memory resource is
                // available.
                continue;
            }
            auto [res, _] = reserve(mem_type, size, AllowOverbooking::NO);
            if (res.size() == size) {
                return std::move(res);
            }
        }
        RAPIDSMPF_FAIL("failed to reserve memory", std::runtime_error);
    }
 
    [[nodiscard]] MemoryReservation reserve_or_fail(
        std::size_t size, MemoryType mem_type
    ) {
        return reserve_or_fail(size, std::ranges::single_view{mem_type});
    }
 
    std::size_t release(MemoryReservation& reservation, std::size_t size);
 
    std::unique_ptr<Buffer> make_buffer(
        std::size_t size, rmm::cuda_stream_view stream, MemoryReservation& reservation
    );
 
    std::unique_ptr<Buffer> make_buffer(
        rmm::cuda_stream_view stream, MemoryReservation&& reservation
    );
 
    std::unique_ptr<Buffer> move(
        std::unique_ptr<rmm::device_buffer> data, rmm::cuda_stream_view stream
    );
 
    std::unique_ptr<Buffer> move(
        std::unique_ptr<Buffer> buffer, MemoryReservation& reservation
    );
 
    std::unique_ptr<rmm::device_buffer> move_to_device_buffer(
        std::unique_ptr<Buffer> buffer, MemoryReservation& reservation
    );
 
    std::unique_ptr<HostBuffer> move_to_host_buffer(
        std::unique_ptr<Buffer> buffer, MemoryReservation& reservation
    );
 
    std::shared_ptr<rmm::cuda_stream_pool> const& stream_pool() const;
 
    SpillManager& spill_manager();
 
    std::shared_ptr<Statistics> statistics() const noexcept;
 
  private:
    BufferResource(
        cuda::mr::any_resource<cuda::mr::device_accessible> device_mr,
        std::optional<PinnedMemoryResource> pinned_mr,
        std::unordered_map<MemoryType, std::int64_t> memory_limits,
        std::optional<Duration> periodic_spill_check,
        std::shared_ptr<rmm::cuda_stream_pool> stream_pool,
        std::shared_ptr<Statistics> statistics
    );
 
    std::mutex mutex_;
    RmmResourceAdaptor owning_mr_;
    std::optional<PinnedMemoryResource> pinned_mr_;
    HostMemoryResource host_mr_;
    std::array<std::atomic<std::int64_t>, MEMORY_TYPES.size()> memory_limits_;
    // Zero initialized reserved counters.
    std::array<std::size_t, MEMORY_TYPES.size()> memory_reserved_ = {};
    std::shared_ptr<rmm::cuda_stream_pool> stream_pool_;
    SpillManager spill_manager_;
    std::shared_ptr<Statistics> statistics_;
};
 
static_assert(StatisticsProvider<BufferResource>);
 
std::int64_t device_limit_from_options(config::Options options);
 
std::optional<Duration> periodic_spill_check_from_options(config::Options options);
 
std::shared_ptr<rmm::cuda_stream_pool> stream_pool_from_options(config::Options options);
 
 
}  // namespace rapidsmpf