project_6/cccl_upstream/thrust/thrust/mr/allocator.h

// SPDX-FileCopyrightText: Copyright (c) 2018, NVIDIA Corporation. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

/*! \file
 *  \brief Allocator types usable with \ref memory_resources "Memory Resources".
 */

#pragma once

#include <thrust/detail/config.h>

#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
#  pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
#  pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
#  pragma system_header
#endif // no system header

#include <thrust/detail/config.h>

#include <thrust/detail/config/memory_resource.h>
#include <thrust/mr/polymorphic_adaptor.h>
#include <thrust/mr/validator.h>

#include <cuda/std/__iterator/iterator_traits.h>
#include <cuda/std/__memory/pointer_traits.h>
#include <cuda/std/limits>

THRUST_NAMESPACE_BEGIN
namespace mr
{
/*! \addtogroup allocators Allocators
 *  \ingroup memory_management
 *  \{
 */

/*! An \p mr::allocator is a template that fulfills the C++ requirements for Allocators,
 *  allowing to use the NPA-based memory resources where an Allocator is required. Unlike
 *  memory resources, but like other allocators, \p mr::allocator is typed and bound to
 *  allocate object of a specific type, however it can be freely rebound to other types.
 *
 *  \tparam T the type that will be allocated by this allocator.
 *  \tparam MR the upstream memory resource to use for memory allocation. Must derive from
 *      \p thrust::mr::memory_resource and must be \p final (in C++11 and beyond).
 */
template <typename T, class MR>
class allocator : private validator<MR>
{
public:
  /*! The pointer to void type of this allocator. */
  using void_pointer = typename MR::pointer;

  /*! The value type allocated by this allocator. Equivalent to \p T. */
  using value_type = T;
  /*! The pointer type allocated by this allocator. Equivaled to the pointer type of \p MR rebound to \p T. */
  using pointer = typename ::cuda::std::pointer_traits<void_pointer>::template rebind<T>;
  /*! The pointer to const type. Equivalent to a pointer type of \p MR rebound to <tt>const T</tt>. */
  using const_pointer = typename ::cuda::std::pointer_traits<void_pointer>::template rebind<const T>;
  /*! The reference to the type allocated by this allocator. Supports smart references. */
  using reference = typename ::cuda::std::iterator_traits<pointer>::reference;
  /*! The const reference to the type allocated by this allocator. Supports smart references. */
  using const_reference = typename ::cuda::std::iterator_traits<const_pointer>::reference;
  /*! The size type of this allocator. Always \p std::size_t. */
  using size_type = std::size_t;
  /*! The difference type between pointers allocated by this allocator. */
  using difference_type = typename ::cuda::std::pointer_traits<pointer>::difference_type;

  /*! Specifies that the allocator shall be propagated on container copy assignment. */
  using propagate_on_container_copy_assignment = ::cuda::std::true_type;
  /*! Specifies that the allocator shall be propagated on container move assignment. */
  using propagate_on_container_move_assignment = ::cuda::std::true_type;
  /*! Specifies that the allocator shall be propagated on container swap. */
  using propagate_on_container_swap = ::cuda::std::true_type;

  /*! The \p rebind metafunction provides the type of an \p allocator instantiated with another type.
   *
   *  \tparam U the other type to use for instantiation.
   */
  template <typename U>
  struct rebind
  {
    /*! The alias \p other gives the type of the rebound \p allocator.
     */
    using other = allocator<U, MR>;
  };

  /*! Calculates the maximum number of elements allocated by this allocator.
   *
   *  \return the maximum value of \p std::size_t, divided by the size of \p T.
   */
  _CCCL_EXEC_CHECK_DISABLE
  template <class SizeType = size_type>
  _CCCL_HOST_DEVICE size_type max_size() const
  {
    // Workaround for cudafe++ < 13.1 + gcc < 13 replacing `numeric_limits<size_t>` with
    // `numeric_limits<conditional<is_void_v<void>, __common_type2_imp<uint64_t, uint64_t>::type, void>::type>`
    return (::cuda::std::numeric_limits<SizeType>::max)() / sizeof(T);
  }

  /*! Constructor.
   *
   *  \param resource the resource to be used to allocate raw memory.
   */
  _CCCL_HOST_DEVICE allocator(MR* resource)
      : mem_res(resource)
  {}

  /*! Copy constructor. Copies the resource pointer. */
  template <typename U>
  _CCCL_HOST_DEVICE allocator(const allocator<U, MR>& other)
      : mem_res(other.resource())
  {}

  /*! Allocates objects of type \p T.
   *
   *  \param n number of elements to allocate
   *  \return a pointer to the newly allocated storage.
   */
  [[nodiscard]] _CCCL_HOST pointer allocate(size_type n)
  {
    return static_cast<pointer>(mem_res->do_allocate(n * sizeof(T), alignof(T)));
  }

  /*! Deallocates objects of type \p T.
   *
   *  \param p pointer returned by a previous call to \p allocate
   *  \param n number of elements, passed as an argument to the \p allocate call that produced \p p
   */
  _CCCL_HOST void deallocate(pointer p, size_type n) noexcept // NOLINT(bugprone-exception-escape)
  {
    return mem_res->do_deallocate(p, n * sizeof(T), alignof(T));
  }

  /*! Extracts the memory resource used by this allocator.
   *
   *  \return the memory resource used by this allocator.
   */
  _CCCL_HOST_DEVICE MR* resource() const
  {
    return mem_res;
  }

private:
  MR* mem_res;
};

/*! Compares the allocators for equality by comparing the underlying memory resources. */
template <typename T, typename MR>
_CCCL_HOST_DEVICE bool operator==(const allocator<T, MR>& lhs, const allocator<T, MR>& rhs) noexcept
{
  return *lhs.resource() == *rhs.resource();
}

/*! Compares the allocators for inequality by comparing the underlying memory resources. */
template <typename T, typename MR>
_CCCL_HOST_DEVICE bool operator!=(const allocator<T, MR>& lhs, const allocator<T, MR>& rhs) noexcept
{
  return !(lhs == rhs);
}

template <typename T, typename Pointer>
using polymorphic_allocator = allocator<T, polymorphic_adaptor_resource<Pointer>>;

/*! A helper allocator class that uses global instances of a given upstream memory resource. Requires the memory
 * resource to be default constructible.
 *
 *  \tparam T the type that will be allocated by this allocator.
 *  \tparam Upstream the upstream memory resource to use for memory allocation. Must derive from
 *      \p thrust::mr::memory_resource and must be \p final (in C++11 and beyond).
 */
template <typename T, typename Upstream>
class stateless_resource_allocator : public thrust::mr::allocator<T, Upstream>
{
  using base = thrust::mr::allocator<T, Upstream>;

public:
  /*! The \p rebind metafunction provides the type of an \p stateless_resource_allocator instantiated with another type.
   *
   *  \tparam U the other type to use for instantiation.
   */
  template <typename U>
  struct rebind
  {
    /*! The alias \p other gives the type of the rebound \p stateless_resource_allocator.
     */
    using other = stateless_resource_allocator<U, Upstream>;
  };

  /*! Default constructor. Uses \p get_global_resource to get the global instance of \p Upstream and initializes the
   *      \p allocator base subobject with that resource.
   */
  _CCCL_EXEC_CHECK_DISABLE
  _CCCL_HOST_DEVICE stateless_resource_allocator()
      : base(get_global_resource<Upstream>())
  {}

  /*! Copy constructor. Copies the memory resource pointer. */
  _CCCL_HOST_DEVICE stateless_resource_allocator(const stateless_resource_allocator& other)
      : base(other)
  {}

  /*! Conversion constructor from an allocator of a different type. Copies the memory resource pointer. */
  template <typename U>
  _CCCL_HOST_DEVICE stateless_resource_allocator(const stateless_resource_allocator<U, Upstream>& other)
      : base(other)
  {}

  stateless_resource_allocator& operator=(const stateless_resource_allocator&) = default;

  /*! Destructor. */
  ~stateless_resource_allocator() = default;
};

/*! \} // allocators
 */
} // namespace mr
THRUST_NAMESPACE_END