rmm_resource_adaptor.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
 
#include <cuda/memory_resource>
 
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/detail/rmm_resource_adaptor_impl.hpp>
#include <rapidsmpf/memory/back_ref_mixin.hpp>
#include <rapidsmpf/memory/scoped_memory_record.hpp>
 
namespace rapidsmpf {
class BufferResource;
 
class RmmResourceAdaptor
    : public cuda::mr::shared_resource<detail::RmmResourceAdaptorImpl<
          cuda::mr::any_resource<cuda::mr::device_accessible>>>,
      public BackRefMixin<BufferResource> {
    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() = 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]] ScopedMemoryRecord get_main_record() const;
 
    [[nodiscard]] std::int64_t current_allocated() const noexcept;
 
    void begin_scoped_memory_record();
 
    ScopedMemoryRecord end_scoped_memory_record();
 
  private:
    // Only `BufferResource` may create the primary adaptor; all other instances
    // are copies obtained via `BufferResource::device_mr_adaptor()`. This keeps
    // the back-reference lifetime contract enforceable (see class docs).
    friend class BufferResource;
 
    explicit RmmResourceAdaptor(
        cuda::mr::any_resource<cuda::mr::device_accessible> primary_mr
    );
};
 
static_assert(cuda::mr::resource_with<RmmResourceAdaptor, cuda::mr::device_accessible>);
 
}  // namespace rapidsmpf