-
Notifications
You must be signed in to change notification settings - Fork 3
/
host_unique_ptr.h
76 lines (63 loc) · 3 KB
/
host_unique_ptr.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
#ifndef HeterogeneousCore_CUDAUtilities_interface_host_unique_ptr_h
#define HeterogeneousCore_CUDAUtilities_interface_host_unique_ptr_h
#include <memory>
#include <functional>
#include "allocate_host.h"
namespace cudautils {
namespace host {
namespace impl {
// Additional layer of types to distinguish from host::unique_ptr
class HostDeleter {
public:
void operator()(void *ptr) { cudautils::free_host(ptr); }
};
}
template <typename T>
using unique_ptr = std::unique_ptr<T, impl::HostDeleter>;
namespace impl {
template <typename T>
struct make_host_unique_selector { using non_array = cudautils::host::unique_ptr<T>; };
template <typename T>
struct make_host_unique_selector<T[]> { using unbounded_array = cudautils::host::unique_ptr<T[]>; };
template <typename T, size_t N>
struct make_host_unique_selector<T[N]> { struct bounded_array {}; };
}
}
// Allocate pinned host memory
template <typename T>
typename host::impl::make_host_unique_selector<T>::non_array
make_host_unique(cudaStream_t stream) {
static_assert(std::is_trivially_constructible<T>::value, "Allocating with non-trivial constructor on the pinned host memory is not supported");
void *mem = allocate_host(sizeof(T), stream);
return typename host::impl::make_host_unique_selector<T>::non_array{reinterpret_cast<T *>(mem)};
}
template <typename T>
typename host::impl::make_host_unique_selector<T>::unbounded_array
make_host_unique(size_t n, cudaStream_t stream) {
using element_type = typename std::remove_extent<T>::type;
static_assert(std::is_trivially_constructible<element_type>::value, "Allocating with non-trivial constructor on the pinned host memory is not supported");
void *mem = allocate_host(n*sizeof(element_type), stream);
return typename host::impl::make_host_unique_selector<T>::unbounded_array{reinterpret_cast<element_type *>(mem)};
}
template <typename T, typename ...Args>
typename host::impl::make_host_unique_selector<T>::bounded_array
make_host_unique(Args&&...) = delete;
// No check for the trivial constructor, make it clear in the interface
template <typename T>
typename host::impl::make_host_unique_selector<T>::non_array
make_host_unique_uninitialized(cudaStream_t stream) {
void *mem = allocate_host(sizeof(T), stream);
return typename host::impl::make_host_unique_selector<T>::non_array{reinterpret_cast<T *>(mem)};
}
template <typename T>
typename host::impl::make_host_unique_selector<T>::unbounded_array
make_host_unique_uninitialized(size_t n, cudaStream_t stream) {
using element_type = typename std::remove_extent<T>::type;
void *mem = allocate_host(n*sizeof(element_type), stream);
return typename host::impl::make_host_unique_selector<T>::unbounded_array{reinterpret_cast<element_type *>(mem)};
}
template <typename T, typename ...Args>
typename host::impl::make_host_unique_selector<T>::bounded_array
make_host_unique_uninitialized(Args&&...) = delete;
}
#endif