buffer_resource.hpp

 
#pragma once
 
#include <array>
#include <memory>
#include <mutex>
#include <optional>
#include <ranges>
#include <unordered_map>
#include <utility>
 
#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/spill_manager.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/statistics.hpp>
#include <rapidsmpf/utils.hpp>
 
namespace rapidsmpf {
 
class BufferResource {
  public:
    using MemoryAvailable = std::function<std::int64_t()>;
 
    BufferResource(
        rmm::device_async_resource_ref device_mr,
        std::shared_ptr<PinnedMemoryResource> pinned_mr = PinnedMemoryResource::Disabled,
        std::unordered_map<MemoryType, MemoryAvailable> memory_available = {},
        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()
    );
 
    ~BufferResource() noexcept = default;
 
    [[nodiscard]] rmm::device_async_resource_ref device_mr() const noexcept {
        return device_mr_;
    }
 
    [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept {
        return host_mr_;
    }
 
    [[nodiscard]] rmm::host_async_resource_ref pinned_mr() {
        RAPIDSMPF_EXPECTS(
            pinned_mr_, "no pinned memory resource is available", std::invalid_argument
        );
        return *pinned_mr_;
    }
 
    [[nodiscard]] MemoryAvailable const& memory_available(MemoryType mem_type) const {
        return memory_available_.at(mem_type);
    }
 
    [[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, size_t size, bool allow_overbooking
    );
 
    MemoryReservation reserve_device_memory_and_spill(
        size_t size, bool allow_overbooking
    );
 
    template <std::ranges::input_range Range>
        requires std::convertible_to<std::ranges::range_value_t<Range>, MemoryType>
    [[nodiscard]] MemoryReservation reserve_or_fail(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_ == PinnedMemoryResource::Disabled)
            {
                // Pinned host memory is only available if the memory resource is
                // available.
                continue;
            }
            auto [res, _] = reserve(mem_type, size, false);
            if (res.size() == size) {
                return std::move(res);
            }
        }
        RAPIDSMPF_FAIL("failed to reserve memory", std::runtime_error);
    }
 
    [[nodiscard]] MemoryReservation reserve_or_fail(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> allocate(
        std::size_t size, rmm::cuda_stream_view stream, MemoryReservation& reservation
    );
 
    std::unique_ptr<Buffer> allocate(
        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
    );
 
    rmm::cuda_stream_pool const& stream_pool() const;
 
    SpillManager& spill_manager();
 
    std::shared_ptr<Statistics> statistics();
 
  private:
    std::mutex mutex_;
    rmm::device_async_resource_ref device_mr_;
    std::shared_ptr<PinnedMemoryResource> pinned_mr_;
    HostMemoryResource host_mr_;
    std::unordered_map<MemoryType, MemoryAvailable> memory_available_;
    // 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_;
};
 
class LimitAvailableMemory {
  public:
    constexpr LimitAvailableMemory(RmmResourceAdaptor const* mr, std::int64_t limit)
        : limit{limit}, mr_{mr} {}
 
    std::int64_t operator()() const {
        return limit - static_cast<std::int64_t>(mr_->current_allocated());
    }
 
  public:
    std::int64_t const limit;  
 
  private:
    RmmResourceAdaptor const* mr_;
};
 
}  // namespace rapidsmpf