-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathhamr_cuda_malloc_uva_allocator.h
150 lines (117 loc) · 5.02 KB
/
hamr_cuda_malloc_uva_allocator.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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
#ifndef hamr_cuda_malloc_uva_allocator_h
#define hamr_cuda_malloc_uva_allocator_h
#include "hamr_config.h"
#include <type_traits>
#include <memory>
namespace hamr
{
/// a deleter for arrays allocated with cuda_malloc_uva
template <typename T, typename E = void>
class cuda_malloc_uva_deleter {};
/// a deleter for arrays allocated with cuda_malloc_uva, specialized for objects
template <typename T>
class cuda_malloc_uva_deleter<T, typename std::enable_if<!std::is_arithmetic<T>::value>::type>
{
public:
/** constructs the deleter
* @param[in] ptr the pointer to the array to delete
* @param[in] n the number of elements in the array
*/
cuda_malloc_uva_deleter(cudaStream_t str, T *ptr, size_t n);
/** deletes the array
* @param[in] ptr the pointer to the array to delete. must be the same as
* that passed during construction.
*/
void operator()(T *ptr);
private:
T *m_ptr;
size_t m_elem;
cudaStream_t m_str;
};
/// a deleter for arrays allocated with cuda_malloc_uva, specialized for numbers
template <typename T>
class cuda_malloc_uva_deleter<T, typename std::enable_if<std::is_arithmetic<T>::value>::type>
{
public:
/** constructs the deleter
* @param[in] ptr the pointer to the array to delete
* @param[in] n the number of elements in the array
*/
cuda_malloc_uva_deleter(cudaStream_t str, T *ptr, size_t n);
/** deletes the array
* @param[in] ptr the pointer to the array to delete. must be the same as
* that passed during construction.
*/
void operator()(T *ptr);
private:
T *m_ptr;
size_t m_elem;
};
/// a class for allocating arrays with cuda_malloc_uva
template <typename T, typename E = void>
struct cuda_malloc_uva_allocator {};
/// a class for allocating arrays with cuda_malloc_uva, specialized for objects
template <typename T>
struct HAMR_EXPORT cuda_malloc_uva_allocator<T, typename std::enable_if<!std::is_arithmetic<T>::value>::type>
{
/** allocate an array of n elements.
* @param[in] str a stream used to order operations, or null for the
* default stream
* @param[in] n the number of elements to allocate
* @returns a shared point to the array that holds a deleter for the memory
*/
static std::shared_ptr<T> allocate(cudaStream_t str, size_t n);
/** allocate an array of n elements.
* @param[in] str a stream used to order operations, or null for the
* default stream
* @param[in] n the number of elements to allocate
* @param[in] val a value to initialize the elements to
* @returns a shared point to the array that holds a deleter for the memory
*/
static std::shared_ptr<T> allocate(cudaStream_t str, size_t n, const T &val);
/** allocate an array of n elements.
* @param[in] str a stream used to order operations, or null for the
* default stream
* @param[in] n the number of elements to allocate
* @param[in] vals an array of values to initialize the elements with
* @param[in] cudaVals a flag that is set to true if vals is accessible from codes running in CUDA
* @returns a shared point to the array that holds a deleter for the memory
*/
template <typename U>
static std::shared_ptr<T> allocate(cudaStream_t str, size_t n, const U *vals, bool cudaVals = false);
};
/// a class for allocating arrays with cuda_malloc_uva, specialized for numbers
template <typename T>
struct HAMR_EXPORT cuda_malloc_uva_allocator<T, typename std::enable_if<std::is_arithmetic<T>::value>::type>
{
/** allocate an array of n elements.
* @param[in] str a stream used to order operations, or null for the
* default stream
* @param[in] n the number of elements to allocate
* @returns a shared point to the array that holds a deleter for the memory
*/
static std::shared_ptr<T> allocate(cudaStream_t str, size_t n);
/** allocate an array of n elements.
* @param[in] str a stream used to order operations, or null for the
* default stream
* @param[in] n the number of elements to allocate
* @param[in] val a value to initialize the elements to
* @returns a shared point to the array that holds a deleter for the memory
*/
static std::shared_ptr<T> allocate(cudaStream_t str, size_t n, const T &val);
/** allocate an array of n elements.
* @param[in] str a stream used to order operations, or null for the
* default stream
* @param[in] n the number of elements to allocate
* @param[in] vals an array of values to initialize the elements with
* @param[in] cudaVals a flag set to true if vals is accessible from codes running in CUDA
* @returns a shared point to the array that holds a deleter for the memory
*/
template <typename U>
static std::shared_ptr<T> allocate(cudaStream_t str, size_t n, const U *vals, bool cudaVals = false);
};
}
#if !defined(HAMR_SEPARATE_IMPL)
#include "hamr_cuda_malloc_uva_allocator_impl.h"
#endif
#endif