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/detail/rmm_resource_adaptor_impl.hpp>
#include <rapidsmpf/error.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 cuda::mr::shared_resource<
          detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>> {
    using shared_base = cuda::mr::shared_resource<
        detail::RmmResourceAdaptorImpl<cuda::pinned_memory_pool>>;
 
  public:
    static constexpr std::nullopt_t Disabled = std::nullopt;
 
    static constexpr bool EnabledByDefault = false;
 
    static constexpr std::string_view DefaultInitiPoolSizeFactor = "0%";
 
    static constexpr std::string_view DefaultMaxPoolSizeFactor = "80%";
 
    static std::optional<PinnedMemoryResource> make_if_available(
        int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {}
    );
 
    static std::optional<PinnedMemoryResource> from_options(config::Options options);
 
    [[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_;
    }
 
    [[nodiscard]] std::function<std::int64_t()> get_memory_available_cb() const;
 
    friend void get_property(
        PinnedMemoryResource const&, cuda::mr::host_accessible
    ) noexcept {}
 
  private:
    PinnedMemoryResource(
        int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {}
    );
 
    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