rmm_resource_adaptor_impl.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <mutex>
#include <optional>
#include <stack>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <utility>
 
#include <cuda_runtime_api.h>
 
#include <cuda/memory_resource>
 
#include <rmm/aligned.hpp>
#include <rmm/cuda_stream.hpp>
#include <rmm/error.hpp>
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/scoped_memory_record.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
namespace rapidsmpf::detail {
 
template <
    cuda::mr::resource_with<cuda::mr::device_accessible> PrimaryMR,
    cuda::mr::resource_with<cuda::mr::device_accessible> FallbackMR =
        cuda::mr::any_resource<cuda::mr::device_accessible>>
class RmmResourceAdaptorImpl {
  public:
    // NOLINTBEGIN(clang-analyzer-core.StackAddressEscape): false positive — primary_mr
    // and fallback_mr are moved into a heap-allocated control block inside
    // make_shared_resource; the analyzer incorrectly traces the forwarding
    // reference chain back to the outer caller's stack frame.
    RmmResourceAdaptorImpl(  // NOLINT(clang-analyzer-core.StackAddressEscape)
        PrimaryMR primary_mr,
        std::optional<FallbackMR> fallback_mr = std::nullopt
    )
        : primary_mr_{std::move(primary_mr)}, fallback_mr_{std::move(fallback_mr)} {}
 
    // NOLINTEND(clang-analyzer-core.StackAddressEscape)
 
    template <typename... Args>
    explicit RmmResourceAdaptorImpl(std::in_place_t, Args&&... args)
        : primary_mr_{std::forward<Args>(args)...}, fallback_mr_{std::nullopt} {}
 
    ~RmmResourceAdaptorImpl() = default;
 
    RmmResourceAdaptorImpl(RmmResourceAdaptorImpl const&) = delete;
    RmmResourceAdaptorImpl(RmmResourceAdaptorImpl&&) = delete;
    RmmResourceAdaptorImpl& operator=(RmmResourceAdaptorImpl const&) = delete;
    RmmResourceAdaptorImpl& operator=(RmmResourceAdaptorImpl&&) = delete;
 
    [[nodiscard]] bool operator==(RmmResourceAdaptorImpl const& other) const noexcept {
        return this == std::addressof(other);
    }
 
    [[nodiscard]] PrimaryMR const& get_upstream_resource() const noexcept {
        return primary_mr_;
    }
 
    [[nodiscard]] std::optional<FallbackMR> const&
    get_fallback_resource() const noexcept {
        return fallback_mr_;
    }
 
    [[nodiscard]] ScopedMemoryRecord get_main_record() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return main_record_;
    }
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept {
        std::lock_guard<std::mutex> lock(mutex_);
        return main_record_.current();
    }
 
    void begin_scoped_memory_record() {
        std::lock_guard<std::mutex> lock(mutex_);
        record_stacks_[std::this_thread::get_id()].emplace();
    }
 
    ScopedMemoryRecord end_scoped_memory_record() {
        std::lock_guard lock(mutex_);
        auto& stack = record_stacks_.at(std::this_thread::get_id());
        RAPIDSMPF_EXPECTS(
            !stack.empty(),
            "calling end_scoped_memory_record() on an empty stack.",
            std::out_of_range
        );
        auto ret = stack.top();
        stack.pop();
        if (!stack.empty()) {
            stack.top().add_subscope(ret);
        }
        return ret;
    }
 
    void* allocate(
        cuda::stream_ref stream,
        std::size_t bytes,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) {
        constexpr auto PRIMARY = ScopedMemoryRecord::AllocType::PRIMARY;
        constexpr auto FALLBACK = ScopedMemoryRecord::AllocType::FALLBACK;
 
        void* ret{};
        auto alloc_type = PRIMARY;
        try {
            ret = primary_mr_.allocate(stream, bytes, alignment);
        } catch (rmm::out_of_memory const& e) {
            if (fallback_mr_.has_value()) {
                alloc_type = FALLBACK;
                ret = fallback_mr_->allocate(stream, bytes, alignment);
                std::lock_guard<std::mutex> lock(mutex_);
                fallback_allocations_.insert(ret);
            } else {
                throw;
            }
        }
        std::lock_guard<std::mutex> lock(mutex_);
        main_record_.record_allocation(alloc_type, safe_cast<std::int64_t>(bytes));
        if (!record_stacks_.empty()) {
            auto const thread_id = std::this_thread::get_id();
            auto& record = record_stacks_[thread_id];
            if (!record.empty()) {
                record.top().record_allocation(
                    alloc_type, safe_cast<std::int64_t>(bytes)
                );
                RAPIDSMPF_EXPECTS(
                    allocating_threads_.insert({ret, thread_id}).second,
                    "duplicate memory pointer"
                );
            }
        }
        return ret;
    }
 
    void deallocate(
        cuda::stream_ref stream,
        void* ptr,
        std::size_t bytes,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept {
        constexpr auto PRIMARY = ScopedMemoryRecord::AllocType::PRIMARY;
        constexpr auto FALLBACK = ScopedMemoryRecord::AllocType::FALLBACK;
 
        ScopedMemoryRecord::AllocType alloc_type;
        {
            std::lock_guard<std::mutex> lock(mutex_);
            alloc_type = (fallback_allocations_.erase(ptr) == 0) ? PRIMARY : FALLBACK;
            main_record_.record_deallocation(alloc_type, safe_cast<std::int64_t>(bytes));
            if (!allocating_threads_.empty()) {
                auto const node = allocating_threads_.extract(ptr);
                if (node) {
                    auto thread_id = node.mapped();
                    auto& record = record_stacks_[thread_id];
                    if (!record.empty()) {
                        record.top().record_deallocation(
                            alloc_type, safe_cast<std::int64_t>(bytes)
                        );
                    }
                }
            }
        }
        if (alloc_type == PRIMARY) {
            primary_mr_.deallocate(stream, ptr, bytes, alignment);
        } else {
            fallback_mr_->deallocate(stream, ptr, bytes, alignment);
        }
    }
 
    void* allocate_sync(
        std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) {
        auto* ptr = allocate(sync_stream_, bytes, alignment);
        sync_stream_.synchronize();
        return ptr;
    }
 
    void deallocate_sync(
        void* ptr,
        std::size_t bytes,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept {
        deallocate(sync_stream_, ptr, bytes, alignment);
    }
 
    friend void get_property(
        RmmResourceAdaptorImpl const&, cuda::mr::device_accessible
    ) noexcept {}
 
  private:
    mutable std::mutex mutex_;
    PrimaryMR primary_mr_;
    std::optional<FallbackMR> fallback_mr_;
    std::unordered_set<void*> fallback_allocations_;
 
    ScopedMemoryRecord main_record_;
    std::unordered_map<std::thread::id, std::stack<ScopedMemoryRecord>> record_stacks_;
    std::unordered_map<void*, std::thread::id> allocating_threads_;
 
    rmm::cuda_stream sync_stream_{
        rmm::cuda_stream::flags::non_blocking
    };  
};
 
}  // namespace rapidsmpf::detail