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/buffer/buffer.hpp>
#include <rapidsmpf/buffer/spill_manager.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/statistics.hpp>
#include <rapidsmpf/utils.hpp>
 
namespace rapidsmpf {
 
class MemoryReservation {
    friend class BufferResource;
 
  public:
    ~MemoryReservation() noexcept;
 
    void clear() noexcept;
 
    MemoryReservation(MemoryReservation&& o)
        : MemoryReservation{
              o.mem_type_, std::exchange(o.br_, nullptr), std::exchange(o.size_, 0)
          } {}
 
    MemoryReservation& operator=(MemoryReservation&& o) noexcept {
        clear();
        mem_type_ = o.mem_type_;
        br_ = std::exchange(o.br_, nullptr);
        size_ = std::exchange(o.size_, 0);
        return *this;
    }
 
    MemoryReservation(MemoryReservation const&) = delete;
    MemoryReservation& operator=(MemoryReservation const&) = delete;
 
    [[nodiscard]] constexpr std::size_t size() const noexcept {
        return size_;
    }
 
    [[nodiscard]] constexpr MemoryType mem_type() const noexcept {
        return mem_type_;
    }
 
    [[nodiscard]] constexpr BufferResource* br() const noexcept {
        return br_;
    }
 
  private:
    constexpr MemoryReservation(MemoryType mem_type, BufferResource* br, std::size_t size)
        : mem_type_{mem_type}, br_{br}, size_{size} {}
 
  private:
    MemoryType mem_type_;  
    BufferResource* br_;  
    std::size_t size_;  
};
 
class BufferResource {
  public:
    using MemoryAvailable = std::function<std::int64_t()>;
 
    BufferResource(
        rmm::device_async_resource_ref device_mr,
        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]] 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)];
    }
 
    [[nodiscard]] std::size_t& memory_reserved(MemoryType mem_type) {
        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_and_spill(
        MemoryType mem_type, 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) {
            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<std::vector<uint8_t>> move_to_host_vector(
        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::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