pinned_memory_resource.hpp

 
#pragma once
 
#include <cstddef>
#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/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;
 
    friend void get_property(
        PinnedMemoryResource const&, cuda::mr::device_accessible
    ) noexcept {}
 
  private:
    // We cannot assign cuda::pinned_memory_pool directly to device_async_resource_ref /
    // host_async_resource_ref: the ref only stores a pointer, but its constructor
    // requires the referenced type to be copyable and movable (CCCL __basic_any_ref
    // constraint). pinned_memory_pool is not copyable, so we wrap it in
    // PinnedMemoryResource, which holds the pool in a shared_resource and is copyable and
    // movable. Copies share the same pool (is_equal compares pool_ pointers).
    cuda::mr::shared_resource<cuda::pinned_memory_pool> pool_;
};
 
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