rmm_resource_adaptor.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <optional>
 
#include <cuda/memory_resource>
 
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/detail/rmm_resource_adaptor_impl.hpp>
#include <rapidsmpf/memory/scoped_memory_record.hpp>
 
namespace rapidsmpf {
 
class RmmResourceAdaptor
    : public cuda::mr::shared_resource<detail::RmmResourceAdaptorImpl<
          cuda::mr::any_resource<cuda::mr::device_accessible>>> {
    using any_device_resource = cuda::mr::any_resource<cuda::mr::device_accessible>;
    using shared_base =
        cuda::mr::shared_resource<detail::RmmResourceAdaptorImpl<any_device_resource>>;
 
  public:
    friend void get_property(
        RmmResourceAdaptor const&, cuda::mr::device_accessible
    ) noexcept {}
 
    RmmResourceAdaptor(
        cuda::mr::any_resource<cuda::mr::device_accessible> primary_mr,
        std::optional<cuda::mr::any_resource<cuda::mr::device_accessible>> fallback_mr =
            std::nullopt
    );
 
    ~RmmResourceAdaptor() = default;
 
    [[nodiscard]] bool operator==(RmmResourceAdaptor const& other) const noexcept {
        return get() == other.get();
    }
 
    [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept;
 
    [[nodiscard]] std::optional<rmm::device_async_resource_ref>
    get_fallback_resource() const noexcept;
 
    [[nodiscard]] ScopedMemoryRecord get_main_record() const;
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept;
 
    void begin_scoped_memory_record();
 
    ScopedMemoryRecord end_scoped_memory_record();
};
 
static_assert(cuda::mr::resource_with<RmmResourceAdaptor, cuda::mr::device_accessible>);
 
}  // namespace rapidsmpf