Apache Arrow (C++)
A columnar in-memory analytics layer designed to accelerate big data.
cuda_memory.h
Go to the documentation of this file.
1 // Licensed to the Apache Software Foundation (ASF) under one
2 // or more contributor license agreements. See the NOTICE file
3 // distributed with this work for additional information
4 // regarding copyright ownership. The ASF licenses this file
5 // to you under the Apache License, Version 2.0 (the
6 // "License"); you may not use this file except in compliance
7 // with the License. You may obtain a copy of the License at
8 //
9 // http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing,
12 // software distributed under the License is distributed on an
13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 // KIND, either express or implied. See the License for the
15 // specific language governing permissions and limitations
16 // under the License.
17 
18 #ifndef ARROW_GPU_CUDA_MEMORY_H
19 #define ARROW_GPU_CUDA_MEMORY_H
20 
21 #include <cstdint>
22 #include <memory>
23 
24 #include "arrow/buffer.h"
25 #include "arrow/io/memory.h"
26 #include "arrow/memory_pool.h"
27 #include "arrow/status.h"
28 
29 namespace arrow {
30 namespace gpu {
31 
32 class CudaContext;
33 class CudaIpcMemHandle;
34 
39 class ARROW_EXPORT CudaBuffer : public Buffer {
40  public:
41  CudaBuffer(uint8_t* data, int64_t size, const std::shared_ptr<CudaContext>& context,
42  bool own_data = false, bool is_ipc = false);
43 
44  CudaBuffer(const std::shared_ptr<CudaBuffer>& parent, const int64_t offset,
45  const int64_t size);
46 
47  ~CudaBuffer();
48 
55  static Status FromBuffer(std::shared_ptr<Buffer> buffer,
56  std::shared_ptr<CudaBuffer>* out);
57 
61  Status CopyToHost(const int64_t position, const int64_t nbytes, void* out) const;
62 
68  Status CopyFromHost(const int64_t position, const void* data, int64_t nbytes);
69 
76  virtual Status ExportForIpc(std::shared_ptr<CudaIpcMemHandle>* handle);
77 
78  std::shared_ptr<CudaContext> context() const { return context_; }
79 
80  protected:
81  std::shared_ptr<CudaContext> context_;
82  bool own_data_;
83  bool is_ipc_;
84 
85  virtual Status Close();
86 };
87 
90 class ARROW_EXPORT CudaHostBuffer : public MutableBuffer {
91  public:
93  ~CudaHostBuffer();
94 };
95 
98 class ARROW_EXPORT CudaIpcMemHandle {
99  public:
100  ~CudaIpcMemHandle();
101 
106  static Status FromBuffer(const void* opaque_handle,
107  std::shared_ptr<CudaIpcMemHandle>* handle);
108 
113  Status Serialize(MemoryPool* pool, std::shared_ptr<Buffer>* out) const;
114 
115  private:
116  explicit CudaIpcMemHandle(const void* handle);
117 
118  struct CudaIpcMemHandleImpl;
119  std::unique_ptr<CudaIpcMemHandleImpl> impl_;
120 
121  const void* handle() const;
122 
123  friend CudaBuffer;
124  friend CudaContext;
125 };
126 
133 class ARROW_EXPORT CudaBufferReader : public io::BufferReader {
134  public:
135  explicit CudaBufferReader(const std::shared_ptr<Buffer>& buffer);
136  ~CudaBufferReader() override;
137 
142  Status Read(int64_t nbytes, int64_t* bytes_read, void* buffer) override;
143 
148  Status Read(int64_t nbytes, std::shared_ptr<Buffer>* out) override;
149 
150  private:
151  std::shared_ptr<CudaBuffer> cuda_buffer_;
152  std::shared_ptr<CudaContext> context_;
153 };
154 
157 class ARROW_EXPORT CudaBufferWriter : public io::WriteableFile {
158  public:
159  explicit CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer);
160  ~CudaBufferWriter() override;
161 
163  Status Close() override;
164 
166  Status Flush() override;
167 
168  Status Seek(int64_t position) override;
169 
170  Status Write(const void* data, int64_t nbytes) override;
171 
172  Status WriteAt(int64_t position, const void* data, int64_t nbytes) override;
173 
174  Status Tell(int64_t* position) const override;
175 
181  Status SetBufferSize(const int64_t buffer_size);
182 
184  int64_t buffer_size() const;
185 
187  int64_t num_bytes_buffered() const;
188 
189  private:
190  class CudaBufferWriterImpl;
191  std::unique_ptr<CudaBufferWriterImpl> impl_;
192 };
193 
198 ARROW_EXPORT
199 Status AllocateCudaHostBuffer(const int64_t size, std::shared_ptr<CudaHostBuffer>* out);
200 
201 } // namespace gpu
202 } // namespace arrow
203 
204 #endif // ARROW_GPU_CUDA_MEMORY_H
MutableBuffer()
Definition: buffer.h:189
std::shared_ptr< CudaContext > context_
Definition: cuda_memory.h:81
Definition: interfaces.h:165
File interface for writing to CUDA buffers, with optional buffering.
Definition: cuda_memory.h:157
A Buffer whose contents can be mutated.
Definition: buffer.h:178
bool is_ipc_
Definition: cuda_memory.h:83
Friendlier interface to the CUDA driver API.
Definition: cuda_context.h:67
Random access zero-copy reads on an arrow::Buffer.
Definition: memory.h:107
Definition: status.h:94
bool own_data_
Definition: cuda_memory.h:82
Definition: cuda_memory.h:98
An Arrow buffer located on a GPU device.
Definition: cuda_memory.h:39
File interface for zero-copy read from CUDA buffers.
Definition: cuda_memory.h:133
Device-accessible CPU memory created using cudaHostAlloc.
Definition: cuda_memory.h:90
Top-level namespace for Apache Arrow C++ API.
Definition: adapter.h:32
std::shared_ptr< CudaContext > context() const
Definition: cuda_memory.h:78
Base class for memory allocation.
Definition: memory_pool.h:34
Status AllocateCudaHostBuffer(const int64_t size, std::shared_ptr< CudaHostBuffer > *out)
Allocate CUDA-accessible memory on CPU host.
Object containing a pointer to a piece of contiguous memory with a particular size.
Definition: buffer.h:48