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 #pragma once 19 20 #include <cstdint> 21 #include <memory> 22 #include <string> 23 24 #include "arrow/io/type_fwd.h" 25 #include "arrow/type_fwd.h" 26 #include "arrow/util/compare.h" 27 #include "arrow/util/macros.h" 28 #include "arrow/util/visibility.h" 29 30 namespace arrow { 31 32 class MemoryManager; 33 34 /// \brief EXPERIMENTAL: Abstract interface for hardware devices 35 /// 36 /// This object represents a device with access to some memory spaces. 37 /// When handling a Buffer or raw memory address, it allows deciding in which 38 /// context the raw memory address should be interpreted 39 /// (e.g. CPU-accessible memory, or embedded memory on some particular GPU). 40 class ARROW_EXPORT Device : public std::enable_shared_from_this<Device>, 41 public util::EqualityComparable<Device> { 42 public: 43 virtual ~Device(); 44 45 /// \brief A shorthand for this device's type. 46 /// 47 /// The returned value is different for each device class, but is the 48 /// same for all instances of a given class. It can be used as a replacement 49 /// for RTTI. 50 virtual const char* type_name() const = 0; 51 52 /// \brief A human-readable description of the device. 53 /// 54 /// The returned value should be detailed enough to distinguish between 55 /// different instances, where necessary. 56 virtual std::string ToString() const = 0; 57 58 /// \brief Whether this instance points to the same device as another one. 59 virtual bool Equals(const Device&) const = 0; 60 61 /// \brief Whether this device is the main CPU device. 62 /// 63 /// This shorthand method is very useful when deciding whether a memory address 64 /// is CPU-accessible. is_cpu()65 bool is_cpu() const { return is_cpu_; } 66 67 /// \brief Return a MemoryManager instance tied to this device 68 /// 69 /// The returned instance uses default parameters for this device type's 70 /// MemoryManager implementation. Some devices also allow constructing 71 /// MemoryManager instances with non-default parameters. 72 virtual std::shared_ptr<MemoryManager> default_memory_manager() = 0; 73 74 protected: 75 ARROW_DISALLOW_COPY_AND_ASSIGN(Device); is_cpu_(is_cpu)76 explicit Device(bool is_cpu = false) : is_cpu_(is_cpu) {} 77 78 bool is_cpu_; 79 }; 80 81 /// \brief EXPERIMENTAL: An object that provides memory management primitives 82 /// 83 /// A MemoryManager is always tied to a particular Device instance. 84 /// It can also have additional parameters (such as a MemoryPool to 85 /// allocate CPU memory). 86 class ARROW_EXPORT MemoryManager : public std::enable_shared_from_this<MemoryManager> { 87 public: 88 virtual ~MemoryManager(); 89 90 /// \brief The device this MemoryManager is tied to device()91 const std::shared_ptr<Device>& device() const { return device_; } 92 93 /// \brief Whether this MemoryManager is tied to the main CPU device. 94 /// 95 /// This shorthand method is very useful when deciding whether a memory address 96 /// is CPU-accessible. is_cpu()97 bool is_cpu() const { return device_->is_cpu(); } 98 99 /// \brief Create a RandomAccessFile to read a particular buffer. 100 /// 101 /// The given buffer must be tied to this MemoryManager. 102 /// 103 /// See also the Buffer::GetReader shorthand. 104 virtual Result<std::shared_ptr<io::RandomAccessFile>> GetBufferReader( 105 std::shared_ptr<Buffer> buf) = 0; 106 107 /// \brief Create a OutputStream to write to a particular buffer. 108 /// 109 /// The given buffer must be mutable and tied to this MemoryManager. 110 /// The returned stream object writes into the buffer's underlying memory 111 /// (but it won't resize it). 112 /// 113 /// See also the Buffer::GetWriter shorthand. 114 virtual Result<std::shared_ptr<io::OutputStream>> GetBufferWriter( 115 std::shared_ptr<Buffer> buf) = 0; 116 117 /// \brief Allocate a (mutable) Buffer 118 /// 119 /// The buffer will be allocated in the device's memory. 120 virtual Result<std::shared_ptr<Buffer>> AllocateBuffer(int64_t size) = 0; 121 122 // XXX Should this take a `const Buffer&` instead 123 /// \brief Copy a Buffer to a destination MemoryManager 124 /// 125 /// See also the Buffer::Copy shorthand. 126 static Result<std::shared_ptr<Buffer>> CopyBuffer( 127 const std::shared_ptr<Buffer>& source, const std::shared_ptr<MemoryManager>& to); 128 129 /// \brief Make a no-copy Buffer view in a destination MemoryManager 130 /// 131 /// See also the Buffer::View shorthand. 132 static Result<std::shared_ptr<Buffer>> ViewBuffer( 133 const std::shared_ptr<Buffer>& source, const std::shared_ptr<MemoryManager>& to); 134 135 protected: 136 ARROW_DISALLOW_COPY_AND_ASSIGN(MemoryManager); 137 MemoryManager(const std::shared_ptr<Device> & device)138 explicit MemoryManager(const std::shared_ptr<Device>& device) : device_(device) {} 139 140 // Default implementations always return nullptr, should be overridden 141 // by subclasses that support data transfer. 142 // (returning nullptr means unsupported copy / view) 143 // In CopyBufferFrom and ViewBufferFrom, the `from` parameter is guaranteed to 144 // be equal to `buf->memory_manager()`. 145 virtual Result<std::shared_ptr<Buffer>> CopyBufferFrom( 146 const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& from); 147 virtual Result<std::shared_ptr<Buffer>> CopyBufferTo( 148 const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& to); 149 virtual Result<std::shared_ptr<Buffer>> ViewBufferFrom( 150 const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& from); 151 virtual Result<std::shared_ptr<Buffer>> ViewBufferTo( 152 const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& to); 153 154 std::shared_ptr<Device> device_; 155 }; 156 157 // ---------------------------------------------------------------------- 158 // CPU backend implementation 159 160 class ARROW_EXPORT CPUDevice : public Device { 161 public: 162 const char* type_name() const override; 163 std::string ToString() const override; 164 bool Equals(const Device&) const override; 165 166 std::shared_ptr<MemoryManager> default_memory_manager() override; 167 168 /// \brief Return the global CPUDevice instance 169 static std::shared_ptr<Device> Instance(); 170 171 /// \brief Create a MemoryManager 172 /// 173 /// The returned MemoryManager will use the given MemoryPool for allocations. 174 static std::shared_ptr<MemoryManager> memory_manager(MemoryPool* pool); 175 176 protected: CPUDevice()177 CPUDevice() : Device(true) {} 178 }; 179 180 class ARROW_EXPORT CPUMemoryManager : public MemoryManager { 181 public: 182 Result<std::shared_ptr<io::RandomAccessFile>> GetBufferReader( 183 std::shared_ptr<Buffer> buf) override; 184 Result<std::shared_ptr<io::OutputStream>> GetBufferWriter( 185 std::shared_ptr<Buffer> buf) override; 186 187 Result<std::shared_ptr<Buffer>> AllocateBuffer(int64_t size) override; 188 189 /// \brief Return the MemoryPool associated with this MemoryManager. pool()190 MemoryPool* pool() const { return pool_; } 191 192 protected: CPUMemoryManager(const std::shared_ptr<Device> & device,MemoryPool * pool)193 CPUMemoryManager(const std::shared_ptr<Device>& device, MemoryPool* pool) 194 : MemoryManager(device), pool_(pool) {} 195 196 static std::shared_ptr<MemoryManager> Make(const std::shared_ptr<Device>& device, 197 MemoryPool* pool = default_memory_pool()); 198 199 Result<std::shared_ptr<Buffer>> CopyBufferFrom( 200 const std::shared_ptr<Buffer>& buf, 201 const std::shared_ptr<MemoryManager>& from) override; 202 Result<std::shared_ptr<Buffer>> CopyBufferTo( 203 const std::shared_ptr<Buffer>& buf, 204 const std::shared_ptr<MemoryManager>& to) override; 205 Result<std::shared_ptr<Buffer>> ViewBufferFrom( 206 const std::shared_ptr<Buffer>& buf, 207 const std::shared_ptr<MemoryManager>& from) override; 208 Result<std::shared_ptr<Buffer>> ViewBufferTo( 209 const std::shared_ptr<Buffer>& buf, 210 const std::shared_ptr<MemoryManager>& to) override; 211 212 MemoryPool* pool_; 213 214 friend std::shared_ptr<MemoryManager> CPUDevice::memory_manager(MemoryPool* pool); 215 friend ARROW_EXPORT std::shared_ptr<MemoryManager> default_cpu_memory_manager(); 216 }; 217 218 /// \brief Return the default CPU MemoryManager instance 219 /// 220 /// The returned singleton instance uses the default MemoryPool. 221 /// This function is a faster spelling of 222 /// `CPUDevice::Instance()->default_memory_manager()`. 223 ARROW_EXPORT 224 std::shared_ptr<MemoryManager> default_cpu_memory_manager(); 225 226 } // namespace arrow 227