pinned_memory_resource.hpp

 
#pragma once
 
#include <cstddef>
#include <memory>
#include <optional>
 
#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/detail/rmm_resource_adaptor_impl.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/back_ref_mixin.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 {
 
class BufferResource;
 
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;
 
    int numa_id = get_current_numa_node();
};
 
inline constexpr std::optional<PinnedPoolProperties> PinnedMemoryDisabled{};
 
std::optional<PinnedPoolProperties> pinned_pool_properties_from_options(
    config::Options options
);
 
class PinnedMemoryResource final
    : public cuda::mr::shared_resource<
          detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>,
      public BackRefMixin<BufferResource> {
    using shared_base = cuda::mr::shared_resource<
        detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>;
 
  public:
    [[nodiscard]] void* allocate(
        cuda::stream_ref stream,
        std::size_t size,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) {
        return get().allocate(stream, size, alignment);
    }
 
    void deallocate(
        cuda::stream_ref stream,
        void* ptr,
        std::size_t size,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept {
        get().deallocate(stream, ptr, size, alignment);
    }
 
    [[nodiscard]] bool operator==(PinnedMemoryResource const& other) const noexcept {
        return get() == other.get();
    }
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept {
        return get().current_allocated();
    }
 
    [[nodiscard]] ScopedMemoryRecord get_main_memory_record() const {
        return get().get_main_record();
    }
 
    [[nodiscard]] constexpr PinnedPoolProperties const& properties() const noexcept {
        return pool_properties_;
    }
 
    friend void get_property(
        PinnedMemoryResource const&, cuda::mr::host_accessible
    ) noexcept {}
 
  private:
    explicit PinnedMemoryResource(PinnedPoolProperties pool_properties);
 
    friend class BufferResource;
 
    PinnedPoolProperties pool_properties_;  
};
 
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