host_memory_resource.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdlib>
#include <cstring>
 
#include <rmm/aligned.hpp>
#include <rmm/cuda_stream_view.hpp>
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/error.hpp>
 
namespace rapidsmpf {
 
class HostMemoryResource {
  public:
    HostMemoryResource() = default;
 
    virtual ~HostMemoryResource() = default;
 
    HostMemoryResource(HostMemoryResource const&) = default;  
    HostMemoryResource(HostMemoryResource&&) = default;  
 
    HostMemoryResource& operator=(HostMemoryResource const&) = default;
 
    HostMemoryResource& operator=(HostMemoryResource&&) = default;
 
    void* allocate_sync(std::size_t, std::size_t) {
        RAPIDSMPF_FAIL(
            "only async stream-ordered allocation must be used in RapidsMPF",
            std::invalid_argument
        );
    }
 
    void deallocate_sync(void*, std::size_t, std::size_t) {
        RAPIDSMPF_FAIL(
            "only async stream-ordered allocation must be used in RapidsMPF",
            std::invalid_argument
        );
    }
 
    virtual void* allocate(
        rmm::cuda_stream_view stream,
        std::size_t size,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    );
 
    virtual void deallocate(
        rmm::cuda_stream_view stream,
        void* ptr,
        std::size_t size,
        std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT
    ) noexcept;
 
    [[nodiscard]] virtual bool is_equal(
        [[maybe_unused]] HostMemoryResource const& other
    ) const noexcept {
        return true;
    }
 
    [[nodiscard]] bool operator==(HostMemoryResource const& other) const noexcept {
        return is_equal(other);
    }
 
    [[nodiscard]] bool operator!=(HostMemoryResource const& other) const noexcept {
        return !is_equal(other);
    }
 
    friend void get_property(
        HostMemoryResource const&, cuda::mr::host_accessible
    ) noexcept {}
};
 
static_assert(cuda::mr::resource<HostMemoryResource>);
static_assert(cuda::mr::resource_with<HostMemoryResource, cuda::mr::host_accessible>);
static_assert(!cuda::mr::resource_with<HostMemoryResource, cuda::mr::device_accessible>);
 
}  // namespace rapidsmpf