polymorphic_allocator.hpp

转到此文件的文档。

/*
 * 版权所有 (c) 2020-2025, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可，版本 2.0 (以下简称"许可") 获得许可；
 * 您不得使用此文件，除非符合许可规定。
 * 您可以在以下位置获得许可的副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则根据许可分发的软件按"原样"提供，
 * 没有任何明示或默示的担保或条件。
 * 有关许可下特定语言的管理权限和限制，请参阅许可。
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>

#include <cstddef>
#include <memory>

namespace RMM_NAMESPACE {
namespace mr {
template <typename T>
class polymorphic_allocator {
 public
 using value_type = T;
 polymorphic_allocator() = default;

 polymorphic_allocator(device_async_resource_ref mr) : mr_{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_view stream)
 {
 return static_cast<value_type*>(
 get_upstream_resource().allocate_async(num * sizeof(T), stream));
 }

 void deallocate(value_type* ptr, std::size_t num, cuda_stream_view stream)
 {
 get_upstream_resource().deallocate_async(ptr, num * sizeof(T), stream);
 }

 [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept
 {
 return mr_;
 }

 private
 rmm::device_async_resource_ref 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_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);
}
 // 组结束
} // 命名空间 mr
} // 命名空间 RMM_NAMESPACE