forked from apache/arrow
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcuda_memory.h
More file actions
195 lines (155 loc) · 6.33 KB
/
Copy pathcuda_memory.h
File metadata and controls
195 lines (155 loc) · 6.33 KB
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#ifndef ARROW_GPU_CUDA_MEMORY_H
#define ARROW_GPU_CUDA_MEMORY_H
#include <cstdint>
#include <memory>
#include "arrow/buffer.h"
#include "arrow/io/memory.h"
#include "arrow/memory_pool.h"
#include "arrow/status.h"
namespace arrow {
namespace gpu {
class CudaContext;
class CudaIpcMemHandle;
/// \class CudaBuffer
/// \brief An Arrow buffer located on a GPU device
///
/// Be careful using this in any Arrow code which may not be GPU-aware
class ARROW_EXPORT CudaBuffer : public Buffer {
public:
CudaBuffer(uint8_t* data, int64_t size, const std::shared_ptr<CudaContext>& context,
bool own_data = false, bool is_ipc = false);
CudaBuffer(const std::shared_ptr<CudaBuffer>& parent, const int64_t offset,
const int64_t size);
~CudaBuffer();
/// \brief Copy memory from GPU device to CPU host
/// \param[out] out a pre-allocated output buffer
/// \return Status
Status CopyToHost(const int64_t position, const int64_t nbytes, uint8_t* out) const;
/// \brief Copy memory to device at position
/// \param[in] position start position to copy bytes
/// \param[in] data the host data to copy
/// \param[in] nbytes number of bytes to copy
/// \return Status
Status CopyFromHost(const int64_t position, const uint8_t* data, int64_t nbytes);
/// \brief Expose this device buffer as IPC memory which can be used in other processes
/// \param[out] handle the exported IPC handle
/// \return Status
///
/// \note After calling this function, this device memory will not be freed
/// when the CudaBuffer is destructed
virtual Status ExportForIpc(std::unique_ptr<CudaIpcMemHandle>* handle);
std::shared_ptr<CudaContext> context() const { return context_; }
protected:
std::shared_ptr<CudaContext> context_;
bool own_data_;
bool is_ipc_;
virtual Status Close();
};
/// \class CudaHostBuffer
/// \brief Device-accessible CPU memory created using cudaHostAlloc
class ARROW_EXPORT CudaHostBuffer : public MutableBuffer {
public:
using MutableBuffer::MutableBuffer;
~CudaHostBuffer();
};
/// \class CudaIpcHandle
/// \brief A container for a CUDA IPC handle
class ARROW_EXPORT CudaIpcMemHandle {
public:
~CudaIpcMemHandle();
/// \brief Create CudaIpcMemHandle from opaque buffer (e.g. from another process)
/// \param[in] opaque_handle a CUipcMemHandle as a const void*
/// \param[out] handle the CudaIpcMemHandle instance
/// \return Status
static Status FromBuffer(const void* opaque_handle,
std::unique_ptr<CudaIpcMemHandle>* handle);
/// \brief Write CudaIpcMemHandle to a Buffer
/// \param[in] pool a MemoryPool to allocate memory from
/// \param[out] out the serialized buffer
/// \return Status
Status Serialize(MemoryPool* pool, std::shared_ptr<Buffer>* out) const;
private:
explicit CudaIpcMemHandle(const void* handle);
struct CudaIpcMemHandleImpl;
std::unique_ptr<CudaIpcMemHandleImpl> impl_;
const void* handle() const;
friend CudaBuffer;
friend CudaContext;
};
/// \class CudaBufferReader
/// \brief File interface for zero-copy read from CUDA buffers
///
/// Note: Reads return pointers to device memory. This means you must be
/// careful using this interface with any Arrow code which may expect to be
/// able to do anything other than pointer arithmetic on the returned buffers
class ARROW_EXPORT CudaBufferReader : public io::BufferReader {
public:
explicit CudaBufferReader(const std::shared_ptr<CudaBuffer>& buffer);
~CudaBufferReader();
/// \brief Read bytes into pre-allocated host memory
/// \param[in] nbytes number of bytes to read
/// \param[out] bytes_read actual number of bytes read
/// \param[out] buffer pre-allocated memory to write into
Status Read(int64_t nbytes, int64_t* bytes_read, uint8_t* buffer) override;
/// \brief Zero-copy read from device memory
/// \param[in] nbytes number of bytes to read
/// \param[out] out a Buffer referencing device memory
/// \return Status
Status Read(int64_t nbytes, std::shared_ptr<Buffer>* out) override;
private:
std::shared_ptr<CudaBuffer> cuda_buffer_;
std::shared_ptr<CudaContext> context_;
};
/// \class CudaBufferWriter
/// \brief File interface for writing to CUDA buffers, with optional buffering
class ARROW_EXPORT CudaBufferWriter : public io::WriteableFile {
public:
explicit CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer);
~CudaBufferWriter();
/// \brief Close writer and flush buffered bytes to GPU
Status Close() override;
/// \brief Flush buffered bytes to GPU
Status Flush() override;
Status Seek(int64_t position) override;
Status Write(const uint8_t* data, int64_t nbytes) override;
Status WriteAt(int64_t position, const uint8_t* data, int64_t nbytes) override;
Status Tell(int64_t* position) const override;
/// \brief Set CPU buffer size to limit calls to cudaMemcpy
/// \param[in] buffer_size the size of CPU buffer to allocate
/// \return Status
///
/// By default writes are unbuffered
Status SetBufferSize(const int64_t buffer_size);
/// \brief Returns size of host (CPU) buffer, 0 for unbuffered
int64_t buffer_size() const;
/// \brief Returns number of bytes buffered on host
int64_t num_bytes_buffered() const;
private:
class CudaBufferWriterImpl;
std::unique_ptr<CudaBufferWriterImpl> impl_;
};
/// \brief Allocate CUDA-accessible memory on CPU host
/// \param[in] size number of bytes
/// \param[out] out the allocated buffer
/// \return Status
ARROW_EXPORT
Status AllocateCudaHostBuffer(const int64_t size, std::shared_ptr<CudaHostBuffer>* out);
} // namespace gpu
} // namespace arrow
#endif // ARROW_GPU_CUDA_MEMORY_H