20 #include <rmm/detail/export.hpp>
21 #include <rmm/detail/thrust_namespace.h>
25 #include <cuda/memory_resource>
26 #include <thrust/device_malloc_allocator.h>
27 #include <thrust/device_ptr.h>
28 #include <thrust/memory.h>
30 namespace RMM_NAMESPACE {
53 using Base = thrust::device_malloc_allocator<T>;
90 : _stream{stream}, _mr(mr)
101 : _mr(other.resource()), _stream{other.stream()}, _device{other._device}
114 return thrust::device_pointer_cast(
115 static_cast<T*
>(_mr.allocate_async(num *
sizeof(T), _stream)));
128 return _mr.deallocate_async(thrust::raw_pointer_cast(ptr), num *
sizeof(T), _stream);
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
An allocator compatible with Thrust containers and algorithms using a device_async_resource_ref for m...
Definition: thrust_allocator_adaptor.hpp:51
thrust_allocator(cuda_stream_view stream, rmm::device_async_resource_ref mr)
Constructs a thrust_allocator using a device memory resource and stream.
Definition: thrust_allocator_adaptor.hpp:89
typename Base::pointer pointer
The pointer type.
Definition: thrust_allocator_adaptor.hpp:54
thrust_allocator(thrust_allocator< U > const &other)
Copy constructor. Copies the resource pointer and stream.
Definition: thrust_allocator_adaptor.hpp:100
pointer allocate(size_type num)
Allocate objects of type T
Definition: thrust_allocator_adaptor.hpp:111
cuda_stream_view stream() const noexcept
The stream used by this allocator.
Definition: thrust_allocator_adaptor.hpp:142
rmm::device_async_resource_ref get_upstream_resource() const noexcept
rmm::device_async_resource_ref to the upstream resource
Definition: thrust_allocator_adaptor.hpp:134
thrust_allocator()=default
Default constructor creates an allocator using the default memory resource and default stream.
typename Base::size_type size_type
The size type.
Definition: thrust_allocator_adaptor.hpp:55
thrust::device_malloc_allocator< T > Base
The base type of this allocator.
Definition: thrust_allocator_adaptor.hpp:53
friend void get_property(thrust_allocator const &, cuda::mr::device_accessible) noexcept
Enables the cuda::mr::device_accessible property.
Definition: thrust_allocator_adaptor.hpp:149
void deallocate(pointer ptr, size_type num)
Deallocates objects of type T
Definition: thrust_allocator_adaptor.hpp:125
thrust_allocator(cuda_stream_view stream)
Constructs a thrust_allocator using the default device memory resource and specified stream.
Definition: thrust_allocator_adaptor.hpp:80
cuda_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:99
cuda::mr::async_resource_ref< cuda::mr::device_accessible > device_async_resource_ref
Alias for a cuda::mr::async_resource_ref with the property cuda::mr::device_accessible.
Definition: resource_ref.hpp:41
device_async_resource_ref get_current_device_resource_ref()
Get the device_async_resource_ref for the current device.
Definition: per_device_resource.hpp:411
Management of per-device device_memory_resources.
RAII class that sets the current CUDA device to the specified device on construction and restores the...
Definition: cuda_device.hpp:151
Provides the type of a thrust_allocator instantiated with another type.
Definition: thrust_allocator_adaptor.hpp:64