All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
thrust_allocator_adaptor.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2019-2024, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <rmm/cuda_device.hpp>
20 #include <rmm/detail/export.hpp>
21 #include <rmm/detail/thrust_namespace.h>
23 #include <rmm/resource_ref.hpp>
24 
25 #include <cuda/memory_resource>
26 #include <thrust/device_malloc_allocator.h>
27 #include <thrust/device_ptr.h>
28 #include <thrust/memory.h>
29 
30 namespace RMM_NAMESPACE {
31 namespace mr {
50 template <typename T>
51 class thrust_allocator : public thrust::device_malloc_allocator<T> {
52  public:
53  using Base = thrust::device_malloc_allocator<T>;
54  using pointer = typename Base::pointer;
55  using size_type = typename Base::size_type;
56 
63  template <typename U>
64  struct rebind {
66  };
67 
72  thrust_allocator() = default;
73 
80  explicit thrust_allocator(cuda_stream_view stream) : _stream{stream} {}
81 
90  : _stream{stream}, _mr(mr)
91  {
92  }
93 
99  template <typename U>
101  : _mr(other.resource()), _stream{other.stream()}, _device{other._device}
102  {
103  }
104 
112  {
113  cuda_set_device_raii dev{_device};
114  return thrust::device_pointer_cast(
115  static_cast<T*>(_mr.allocate_async(num * sizeof(T), _stream)));
116  }
117 
125  void deallocate(pointer ptr, size_type num)
126  {
127  cuda_set_device_raii dev{_device};
128  return _mr.deallocate_async(thrust::raw_pointer_cast(ptr), num * sizeof(T), _stream);
129  }
130 
135  {
136  return _mr;
137  }
138 
142  [[nodiscard]] cuda_stream_view stream() const noexcept { return _stream; }
143 
149  friend void get_property(thrust_allocator const&, cuda::mr::device_accessible) noexcept {}
150 
151  private:
152  cuda_stream_view _stream{};
154  cuda_device_id _device{get_current_cuda_device()};
155 }; // end of group
157 } // namespace mr
158 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:41
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:96
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:148
Provides the type of a thrust_allocator instantiated with another type.
Definition: thrust_allocator_adaptor.hpp:64