polymorphic_allocator.hpp

/*
 * Copyright (c) 2020-2021, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
 
#include <cstddef>
#include <memory>
#include <type_traits>
 
namespace rmm::mr {
 
template <typename T>
class polymorphic_allocator {
 public:
  using value_type = T;
  polymorphic_allocator() = default;
 
  polymorphic_allocator(device_memory_resource* mr) : mr_{mr} {}
 
  template <typename U>
  polymorphic_allocator(polymorphic_allocator<U> const& other) noexcept : mr_{other.resource()}
  {
  }
 
  value_type* allocate(std::size_t num, cuda_stream_view stream)
  {
    return static_cast<value_type*>(resource()->allocate(num * sizeof(T), stream));
  }
 
  void deallocate(value_type* ptr, std::size_t num, cuda_stream_view stream)
  {
    resource()->deallocate(ptr, num * sizeof(T), stream);
  }
 
  [[nodiscard]] device_memory_resource* resource() const noexcept { return mr_; }
 
 private:
  device_memory_resource* mr_{
    get_current_device_resource()};  
};
 
template <typename T, typename U>
bool operator==(polymorphic_allocator<T> const& lhs, polymorphic_allocator<U> const& rhs)
{
  return lhs.resource()->is_equal(*rhs.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_view 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) { alloc_.deallocate(ptr, num, stream()); }
 
  [[nodiscard]] cuda_stream_view stream() const noexcept { return stream_; }
 
  [[nodiscard]] Allocator underlying_allocator() const noexcept { return alloc_; }
 
 private:
  Allocator alloc_;          
  cuda_stream_view stream_;  
};
 
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);
}
 
template <typename Allocator>
auto make_stream_allocator_adaptor(Allocator const& allocator, cuda_stream_view stream)
{
  return stream_allocator_adaptor<Allocator>{allocator, stream};
}
 
}  // namespace rmm::mr