pinned_memory_resource.hpp

 
#pragma once
 
#include <cstddef>
#include <functional>
#include <memory>
 
#include <cuda.h>
#include <cuda_runtime_api.h>
 
#include <cuda/memory_resource>
 
#include <rmm/aligned.hpp>
#include <rmm/cuda_device.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>
 
#include <rapidsmpf/config.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/host_memory_resource.hpp>
#include <rapidsmpf/rmm_resource_adaptor.hpp>
#include <rapidsmpf/system_info.hpp>
#include <rapidsmpf/utils/misc.hpp>
 
// NOLINTBEGIN(modernize-macro-to-enum)
#define RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION 12060
#define RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION_STR "v12.6"
 
// NOLINTEND(modernize-macro-to-enum)
 
namespace rapidsmpf {
 
inline bool is_pinned_memory_resources_supported() {
    static const bool supported = [] {
        // check if the device supports async memory pools
        int cuda_pool_supported{};
        auto attr_result = cudaDeviceGetAttribute(
            &cuda_pool_supported,
            cudaDevAttrMemoryPoolsSupported,
            rmm::get_current_cuda_device().value()
        );
        if (attr_result != cudaSuccess || cuda_pool_supported != 1) {
            return false;
        }
 
        int cuda_driver_version{};
        auto driver_result = cudaDriverGetVersion(&cuda_driver_version);
        int cuda_runtime_version{};
        auto runtime_result = cudaRuntimeGetVersion(&cuda_runtime_version);
        return driver_result == cudaSuccess && runtime_result == cudaSuccess
               && cuda_driver_version >= RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION
               && cuda_runtime_version >= RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION;
    }();
    return supported;
}
 
struct PinnedPoolProperties {
    std::size_t initial_pool_size = 0;
 
    std::optional<std::size_t> max_pool_size = std::nullopt;
};
 
class PinnedMemoryResource final : public HostMemoryResource {
  public:
    static constexpr auto Disabled = nullptr;
 
    PinnedMemoryResource(
        int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {}
    );
 
    static std::shared_ptr<PinnedMemoryResource> make_if_available(
        int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {}
    );
 
    static std::shared_ptr<PinnedMemoryResource> from_options(config::Options options);
 
    ~PinnedMemoryResource() override;
 
    void* allocate(
        rmm::cuda_stream_view stream,
        std::size_t size,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) override;
 
    void deallocate(
        rmm::cuda_stream_view stream,
        void* ptr,
        std::size_t size,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept override;
 
    [[nodiscard]] bool is_equal(HostMemoryResource const& other) const noexcept override;
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept {
        return pool_tracker_.current_allocated();
    }
 
    [[nodiscard]] ScopedMemoryRecord get_main_memory_record() const {
        return pool_tracker_.get_main_record();
    }
 
    [[nodiscard]] constexpr PinnedPoolProperties const& properties() const noexcept {
        return pool_properties_;
    }
 
    [[nodiscard]] std::function<std::int64_t()> get_memory_available_cb() const;
 
    friend void get_property(
        PinnedMemoryResource const&, cuda::mr::device_accessible
    ) noexcept {}
 
  private:
    PinnedPoolProperties pool_properties_;  
 
    // The cuda::pinned_memory_pool is moved into the RmmResourceAdaptor's primary_mr
    // (an any_resource<device_accessible>), which keeps it alive. RmmResourceAdaptor is
    // itself a shared_resource<RmmResourceAdaptorImpl>, so copies of PinnedMemoryResource
    // share the same underlying pool and adaptor state (memory statistics). Copies are
    // equal iff they share the same RmmResourceAdaptorImpl (is_equal compares
    // pool_tracker_).
    RmmResourceAdaptor pool_tracker_;  
};
 
static_assert(cuda::mr::resource<PinnedMemoryResource>);
static_assert(cuda::mr::resource_with<PinnedMemoryResource, cuda::mr::host_accessible>);
static_assert(cuda::mr::resource_with<PinnedMemoryResource, cuda::mr::device_accessible>);
 
}  // namespace rapidsmpf