thrust_allocator_adaptor.hpp
1 /*
2  * Copyright (c) 2019, 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 <thrust/detail/type_traits/pointer_traits.h>
20 #include <thrust/device_malloc_allocator.h>
21 
22 #include <rmm/mr/device/device_memory_resource.hpp>
24 
25 namespace rmm {
26 namespace mr {
37 template <typename T>
38 class thrust_allocator : public thrust::device_malloc_allocator<T> {
39  public:
40  using Base = thrust::device_malloc_allocator<T>;
41  using pointer = typename Base::pointer;
42  using size_type = typename Base::size_type;
43 
50  template <typename U>
51  struct rebind {
52  using other = thrust_allocator<U>;
53  };
54 
59  thrust_allocator() = default;
60 
68 
77  {
78  }
79 
85  template <typename U>
87  : _mr(other.resource()), _stream{other.stream()}
88  {
89  }
90 
97  pointer allocate(size_type n)
98  {
99  return thrust::device_pointer_cast(static_cast<T*>(_mr->allocate(n * sizeof(T), _stream)));
100  }
101 
109  void deallocate(pointer p, size_type n)
110  {
111  return _mr->deallocate(thrust::raw_pointer_cast(p), n * sizeof(T), _stream);
112  }
113 
117  device_memory_resource* resource() const noexcept { return _mr; }
118 
122  cuda_stream_view stream() const noexcept { return _stream; }
123 
124  private:
125  cuda_stream_view _stream{};
126  device_memory_resource* _mr{rmm::mr::get_current_device_resource()};
127 };
128 } // namespace mr
129 } // namespace rmm
rmm::mr::thrust_allocator::thrust_allocator
thrust_allocator(cuda_stream_view stream, device_memory_resource *mr)
Constructs a thrust_allocator using a device memory resource and stream.
Definition: thrust_allocator_adaptor.hpp:76
rmm::mr::thrust_allocator
An allocator compatible with Thrust containers and algorithms using a device_memory_resource for memo...
Definition: thrust_allocator_adaptor.hpp:38
rmm::mr::device_memory_resource::allocate
void * allocate(std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Allocates memory of size at least bytes.
Definition: device_memory_resource.hpp:102
per_device_resource.hpp
Management of per-device device_memory_resources.
rmm::mr::thrust_allocator::allocate
pointer allocate(size_type n)
Allocate objects of type T
Definition: thrust_allocator_adaptor.hpp:97
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::mr::device_memory_resource::deallocate
void deallocate(void *p, std::size_t bytes, cuda_stream_view stream=cuda_stream_view{})
Deallocate memory pointed to by p.
Definition: device_memory_resource.hpp:125
rmm::mr::thrust_allocator::thrust_allocator
thrust_allocator(thrust_allocator< U > const &other)
Copy constructor. Copies the resource pointer and stream.
Definition: thrust_allocator_adaptor.hpp:86
rmm::mr::thrust_allocator::rebind
Provides the type of a thrust_allocator instantiated with another type.
Definition: thrust_allocator_adaptor.hpp:51
rmm::mr::thrust_allocator::stream
cuda_stream_view stream() const noexcept
Returns the stream used by this allocator.
Definition: thrust_allocator_adaptor.hpp:122
rmm::mr::thrust_allocator::deallocate
void deallocate(pointer p, size_type n)
Deallocates objects of type T
Definition: thrust_allocator_adaptor.hpp:109
rmm::mr::thrust_allocator::thrust_allocator
thrust_allocator()=default
Default constructor creates an allocator using the default memory resource and default stream.
rmm::mr::thrust_allocator::thrust_allocator
thrust_allocator(cuda_stream_view stream)
Constructs a thrust_allocator using the default device memory resource and specified stream.
Definition: thrust_allocator_adaptor.hpp:67
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:83
rmm::mr::thrust_allocator::resource
device_memory_resource * resource() const noexcept
Returns the device memory resource used by this allocator.
Definition: thrust_allocator_adaptor.hpp:117