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 
52  Status CopyToHost(const int64_t position, const int64_t nbytes, uint8_t* out) const;
53 
59  Status CopyFromHost(const int64_t position, const uint8_t* data, int64_t nbytes);
60 
67  virtual Status ExportForIpc(std::unique_ptr<CudaIpcMemHandle>* handle);
68 
69  std::shared_ptr<CudaContext> context() const { return context_; }
70 
71  protected:
72  std::shared_ptr<CudaContext> context_;
73  bool own_data_;
74  bool is_ipc_;
75 
76  virtual Status Close();
77 };
78 
81 class ARROW_EXPORT CudaHostBuffer : public MutableBuffer {
82  public:
84  ~CudaHostBuffer();
85 };
86 
89 class ARROW_EXPORT CudaIpcMemHandle {
90  public:
92 
97  static Status FromBuffer(const void* opaque_handle,
98  std::unique_ptr<CudaIpcMemHandle>* handle);
99 
104  Status Serialize(MemoryPool* pool, std::shared_ptr<Buffer>* out) const;
105 
106  private:
107  explicit CudaIpcMemHandle(const void* handle);
108 
109  struct CudaIpcMemHandleImpl;
110  std::unique_ptr<CudaIpcMemHandleImpl> impl_;
111 
112  const void* handle() const;
113 
114  friend CudaBuffer;
115  friend CudaContext;
116 };
117 
124 class ARROW_EXPORT CudaBufferReader : public io::BufferReader {
125  public:
126  explicit CudaBufferReader(const std::shared_ptr<CudaBuffer>& buffer);
127  ~CudaBufferReader();
128 
133  Status Read(int64_t nbytes, int64_t* bytes_read, uint8_t* buffer) override;
134 
139  Status Read(int64_t nbytes, std::shared_ptr<Buffer>* out) override;
140 
141  private:
142  std::shared_ptr<CudaBuffer> cuda_buffer_;
143  std::shared_ptr<CudaContext> context_;
144 };
145 
148 class ARROW_EXPORT CudaBufferWriter : public io::WriteableFile {
149  public:
150  explicit CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer);
151  ~CudaBufferWriter();
152 
154  Status Close() override;
155 
157  Status Flush() override;
158 
159  Status Seek(int64_t position) override;
160 
161  Status Write(const uint8_t* data, int64_t nbytes) override;
162 
163  Status WriteAt(int64_t position, const uint8_t* data, int64_t nbytes) override;
164 
165  Status Tell(int64_t* position) const override;
166 
172  Status SetBufferSize(const int64_t buffer_size);
173 
175  int64_t buffer_size() const;
176 
178  int64_t num_bytes_buffered() const;
179 
180  private:
181  class CudaBufferWriterImpl;
182  std::unique_ptr<CudaBufferWriterImpl> impl_;
183 };
184 
189 ARROW_EXPORT
190 Status AllocateCudaHostBuffer(const int64_t size, std::shared_ptr<CudaHostBuffer>* out);
191 
192 } // namespace gpu
193 } // namespace arrow
194 
195 #endif // ARROW_GPU_CUDA_MEMORY_H
MutableBuffer()
Definition: buffer.h:148
std::shared_ptr< CudaContext > context_
Definition: cuda_memory.h:72
Definition: interfaces.h:163
File interface for writing to CUDA buffers, with optional buffering.
Definition: cuda_memory.h:148
A Buffer whose contents can be mutated. May or may not own its data.
Definition: buffer.h:137
bool is_ipc_
Definition: cuda_memory.h:74
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:106
bool own_data_
Definition: cuda_memory.h:73
Definition: cuda_memory.h:89
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:124
Device-accessible CPU memory created using cudaHostAlloc.
Definition: cuda_memory.h:81
Top-level namespace for Apache Arrow C++ API.
Definition: allocator.h:29
std::shared_ptr< CudaContext > context() const
Definition: cuda_memory.h:69
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.
Immutable API for a chunk of bytes which may or may not be owned by the class instance.
Definition: buffer.h:48