rmm_resource_adaptor_impl.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <mutex>
#include <stack>
#include <thread>
#include <unordered_map>
#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>
class RmmResourceAdaptorImpl {
  public:
    // NOLINTBEGIN(clang-analyzer-core.StackAddressEscape): false positive — primary_mr
    // is 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.
    explicit RmmResourceAdaptorImpl(  // NOLINT(clang-analyzer-core.StackAddressEscape)
        PrimaryMR primary_mr
    )
        : primary_mr_{std::move(primary_mr)} {}
 
    // NOLINTEND(clang-analyzer-core.StackAddressEscape)
 
    template <typename... Args>
    explicit RmmResourceAdaptorImpl(std::in_place_t, Args&&... args)
        : primary_mr_{std::forward<Args>(args)...} {}
 
    ~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]] 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
    ) {
        void* ret = primary_mr_.allocate(stream, bytes, alignment);
        std::lock_guard<std::mutex> lock(mutex_);
        main_record_.record_allocation(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(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 {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            main_record_.record_deallocation(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(safe_cast<std::int64_t>(bytes));
                    }
                }
            }
        }
        primary_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_;
 
    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