polymorphic_allocator.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/aligned.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda/memory_resource>
#include <cuda/stream>
 
#include <cstddef>
#include <memory>
 
RMM_NAMESPACE_BEGIN
namespace mr {
template <typename T>
class polymorphic_allocator {
 public:
  using value_type = T;  
  polymorphic_allocator() = default;

  polymorphic_allocator(cuda::mr::any_resource<cuda::mr::device_accessible> mr) : mr_(std::move(mr))
  {
  }

  template <typename U>
  polymorphic_allocator(polymorphic_allocator<U> const& other) noexcept
    : mr_(other.get_upstream_resource())
  {
  }

  value_type* allocate(std::size_t num, cuda::stream_ref stream)
  {
    return static_cast<value_type*>(
      mr_.allocate(stream, num * sizeof(T), rmm::CUDA_ALLOCATION_ALIGNMENT));
  }

  void deallocate(value_type* ptr, std::size_t num, cuda::stream_ref stream) noexcept
  {
    mr_.deallocate(stream, ptr, num * sizeof(T), rmm::CUDA_ALLOCATION_ALIGNMENT);
  }

  [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
  {
    return rmm::device_async_resource_ref{mr_};
  }
 
 private:
  mutable cuda::mr::any_resource<cuda::mr::device_accessible> mr_{
    get_current_device_resource_ref()};  
};

template <typename T, typename U>
bool operator==(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
{
  return lhs.get_upstream_resource() == rhs.get_upstream_resource();
}

template <typename T, typename U>
bool operator!=(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
{
  return not(lhs == rhs);
}

template <typename Allocator>
class stream_allocator_adaptor {
 public:
  using value_type =
    typename std::allocator_traits<Allocator>::value_type;  
 
  stream_allocator_adaptor() = delete;

  stream_allocator_adaptor(Allocator const& allocator, cuda::stream_ref stream)
    : alloc_{allocator}, stream_{stream}
  {
  }

  template <typename OtherAllocator>
  stream_allocator_adaptor(stream_allocator_adaptor<OtherAllocator> const& other)
    : stream_allocator_adaptor{other.underlying_allocator(), other.stream()}
  {
  }

  template <typename T>
  struct rebind {
    using other = stream_allocator_adaptor<typename std::allocator_traits<
      Allocator>::template rebind_alloc<T>>;  
  };

  value_type* allocate(std::size_t num) { return alloc_.allocate(num, stream()); }

  void deallocate(value_type* ptr, std::size_t num) noexcept
  {
    alloc_.deallocate(ptr, num, stream());
  }

  [[nodiscard]] cuda::stream_ref stream() const noexcept { return stream_; }

  [[nodiscard]] Allocator underlying_allocator() const noexcept { return alloc_; }
 
 private:
  Allocator alloc_;  
  cuda::stream_ref stream_{
    cuda::invalid_stream_t{}};  
};

template <typename A, typename O>
bool operator==(stream_allocator_adaptor<A> const& lhs, stream_allocator_adaptor<O> const& rhs)
{
  return lhs.underlying_allocator() == rhs.underlying_allocator();
}

template <typename A, typename O>
bool operator!=(stream_allocator_adaptor<A> const& lhs, stream_allocator_adaptor<O> const& rhs)
{
  return not(lhs == rhs);
}
  // end of group
}  // namespace mr
RMM_NAMESPACE_END