383 lines
15 KiB
C
383 lines
15 KiB
C
|
// 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.
|
||
|
|
|
||
|
|
#pragma once
|
||
|
|
|
||
|
|
#include <cstdint>
|
||
|
|
#include <functional>
|
||
|
|
#include <memory>
|
||
|
|
#include <string>
|
||
|
|
|
||
|
|
#include "arrow/io/type_fwd.h"
|
||
|
|
#include "arrow/result.h"
|
||
|
|
#include "arrow/status.h"
|
||
|
|
#include "arrow/type_fwd.h"
|
||
|
|
#include "arrow/util/compare.h"
|
||
|
|
#include "arrow/util/macros.h"
|
||
|
|
#include "arrow/util/visibility.h"
|
||
|
|
|
||
|
|
namespace arrow {
|
||
|
|
|
||
|
|
class MemoryManager;
|
||
|
|
|
||
|
|
/// \brief EXPERIMENTAL: Abstract interface for hardware devices
|
||
|
|
///
|
||
|
|
/// This object represents a device with access to some memory spaces.
|
||
|
|
/// When handling a Buffer or raw memory address, it allows deciding in which
|
||
|
|
/// context the raw memory address should be interpreted
|
||
|
|
/// (e.g. CPU-accessible memory, or embedded memory on some particular GPU).
|
||
|
|
class ARROW_EXPORT Device : public std::enable_shared_from_this<Device>,
|
||
|
|
public util::EqualityComparable<Device> {
|
||
|
|
public:
|
||
|
|
virtual ~Device();
|
||
|
|
|
||
|
|
/// \brief A shorthand for this device's type.
|
||
|
|
///
|
||
|
|
/// The returned value is different for each device class, but is the
|
||
|
|
/// same for all instances of a given class. It can be used as a replacement
|
||
|
|
/// for RTTI.
|
||
|
|
virtual const char* type_name() const = 0;
|
||
|
|
|
||
|
|
/// \brief A human-readable description of the device.
|
||
|
|
///
|
||
|
|
/// The returned value should be detailed enough to distinguish between
|
||
|
|
/// different instances, where necessary.
|
||
|
|
virtual std::string ToString() const = 0;
|
||
|
|
|
||
|
|
/// \brief Whether this instance points to the same device as another one.
|
||
|
|
virtual bool Equals(const Device&) const = 0;
|
||
|
|
|
||
|
|
/// \brief A device ID to identify this device if there are multiple of this type.
|
||
|
|
///
|
||
|
|
/// If there is no "device_id" equivalent (such as for the main CPU device on
|
||
|
|
/// non-numa systems) returns -1.
|
||
|
|
virtual int64_t device_id() const { return -1; }
|
||
|
|
|
||
|
|
/// \brief Whether this device is the main CPU device.
|
||
|
|
///
|
||
|
|
/// This shorthand method is very useful when deciding whether a memory address
|
||
|
|
/// is CPU-accessible.
|
||
|
|
bool is_cpu() const { return is_cpu_; }
|
||
|
|
|
||
|
|
/// \brief Return a MemoryManager instance tied to this device
|
||
|
|
///
|
||
|
|
/// The returned instance uses default parameters for this device type's
|
||
|
|
/// MemoryManager implementation. Some devices also allow constructing
|
||
|
|
/// MemoryManager instances with non-default parameters.
|
||
|
|
virtual std::shared_ptr<MemoryManager> default_memory_manager() = 0;
|
||
|
|
|
||
|
|
/// \brief Return the DeviceAllocationType of this device
|
||
|
|
virtual DeviceAllocationType device_type() const = 0;
|
||
|
|
|
||
|
|
class SyncEvent;
|
||
|
|
|
||
|
|
/// \brief EXPERIMENTAL: An opaque wrapper for Device-specific streams
|
||
|
|
///
|
||
|
|
/// In essence this is just a wrapper around a void* to represent the
|
||
|
|
/// standard concept of a stream/queue on a device. Derived classes
|
||
|
|
/// should be trivially constructible from it's device-specific counterparts.
|
||
|
|
class ARROW_EXPORT Stream {
|
||
|
|
public:
|
||
|
|
using release_fn_t = std::function<void(void*)>;
|
||
|
|
|
||
|
|
virtual ~Stream() = default;
|
||
|
|
|
||
|
|
virtual const void* get_raw() const { return stream_.get(); }
|
||
|
|
|
||
|
|
/// \brief Make the stream wait on the provided event.
|
||
|
|
///
|
||
|
|
/// Tells the stream that it should wait until the synchronization
|
||
|
|
/// event is completed without blocking the CPU.
|
||
|
|
virtual Status WaitEvent(const SyncEvent&) = 0;
|
||
|
|
|
||
|
|
/// \brief Blocks the current thread until a stream's remaining tasks are completed
|
||
|
|
virtual Status Synchronize() const = 0;
|
||
|
|
|
||
|
|
protected:
|
||
|
|
explicit Stream(void* stream, release_fn_t release_stream)
|
||
|
|
: stream_{stream, release_stream} {}
|
||
|
|
|
||
|
|
std::unique_ptr<void, release_fn_t> stream_;
|
||
|
|
};
|
||
|
|
|
||
|
|
virtual Result<std::shared_ptr<Stream>> MakeStream() { return NULLPTR; }
|
||
|
|
|
||
|
|
/// \brief Create a new device stream
|
||
|
|
///
|
||
|
|
/// This should create the appropriate stream type for the device,
|
||
|
|
/// derived from Device::Stream to allow for stream ordered events
|
||
|
|
/// and memory allocations.
|
||
|
|
virtual Result<std::shared_ptr<Stream>> MakeStream(
|
||
|
|
unsigned int ARROW_ARG_UNUSED(flags)) {
|
||
|
|
return NULLPTR;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// @brief Wrap an existing device stream alongside a release function
|
||
|
|
///
|
||
|
|
/// @param device_stream a pointer to the stream to wrap
|
||
|
|
/// @param release_fn a function to call during destruction, `nullptr` or
|
||
|
|
/// a no-op function can be passed to indicate ownership is maintained
|
||
|
|
/// externally
|
||
|
|
virtual Result<std::shared_ptr<Stream>> WrapStream(
|
||
|
|
void* ARROW_ARG_UNUSED(device_stream),
|
||
|
|
Stream::release_fn_t ARROW_ARG_UNUSED(release_fn)) {
|
||
|
|
return NULLPTR;
|
||
|
|
}
|
||
|
|
|
||
|
|
/// \brief EXPERIMENTAL: An object that provides event/stream sync primitives
|
||
|
|
class ARROW_EXPORT SyncEvent {
|
||
|
|
public:
|
||
|
|
using release_fn_t = std::function<void(void*)>;
|
||
|
|
|
||
|
|
virtual ~SyncEvent() = default;
|
||
|
|
|
||
|
|
void* get_raw() { return sync_event_.get(); }
|
||
|
|
|
||
|
|
/// @brief Block until sync event is completed.
|
||
|
|
virtual Status Wait() = 0;
|
||
|
|
|
||
|
|
/// @brief Record the wrapped event on the stream so it triggers
|
||
|
|
/// the event when the stream gets to that point in its queue.
|
||
|
|
virtual Status Record(const Stream&) = 0;
|
||
|
|
|
||
|
|
protected:
|
||
|
|
/// If creating this with a passed in event, the caller must ensure
|
||
|
|
/// that the event lives until clear_event is called on this as it
|
||
|
|
/// won't own it.
|
||
|
|
explicit SyncEvent(void* sync_event, release_fn_t release_sync_event)
|
||
|
|
: sync_event_{sync_event, release_sync_event} {}
|
||
|
|
|
||
|
|
std::unique_ptr<void, release_fn_t> sync_event_;
|
||
|
|
};
|
||
|
|
|
||
|
|
protected:
|
||
|
|
ARROW_DISALLOW_COPY_AND_ASSIGN(Device);
|
||
|
|
explicit Device(bool is_cpu = false) : is_cpu_(is_cpu) {}
|
||
|
|
|
||
|
|
bool is_cpu_;
|
||
|
|
};
|
||
|
|
|
||
|
|
/// \brief EXPERIMENTAL: An object that provides memory management primitives
|
||
|
|
///
|
||
|
|
/// A MemoryManager is always tied to a particular Device instance.
|
||
|
|
/// It can also have additional parameters (such as a MemoryPool to
|
||
|
|
/// allocate CPU memory).
|
||
|
|
class ARROW_EXPORT MemoryManager : public std::enable_shared_from_this<MemoryManager> {
|
||
|
|
public:
|
||
|
|
virtual ~MemoryManager();
|
||
|
|
|
||
|
|
/// \brief The device this MemoryManager is tied to
|
||
|
|
const std::shared_ptr<Device>& device() const { return device_; }
|
||
|
|
|
||
|
|
/// \brief Whether this MemoryManager is tied to the main CPU device.
|
||
|
|
///
|
||
|
|
/// This shorthand method is very useful when deciding whether a memory address
|
||
|
|
/// is CPU-accessible.
|
||
|
|
bool is_cpu() const { return device_->is_cpu(); }
|
||
|
|
|
||
|
|
/// \brief Create a RandomAccessFile to read a particular buffer.
|
||
|
|
///
|
||
|
|
/// The given buffer must be tied to this MemoryManager.
|
||
|
|
///
|
||
|
|
/// See also the Buffer::GetReader shorthand.
|
||
|
|
virtual Result<std::shared_ptr<io::RandomAccessFile>> GetBufferReader(
|
||
|
|
std::shared_ptr<Buffer> buf) = 0;
|
||
|
|
|
||
|
|
/// \brief Create a OutputStream to write to a particular buffer.
|
||
|
|
///
|
||
|
|
/// The given buffer must be mutable and tied to this MemoryManager.
|
||
|
|
/// The returned stream object writes into the buffer's underlying memory
|
||
|
|
/// (but it won't resize it).
|
||
|
|
///
|
||
|
|
/// See also the Buffer::GetWriter shorthand.
|
||
|
|
virtual Result<std::shared_ptr<io::OutputStream>> GetBufferWriter(
|
||
|
|
std::shared_ptr<Buffer> buf) = 0;
|
||
|
|
|
||
|
|
/// \brief Allocate a (mutable) Buffer
|
||
|
|
///
|
||
|
|
/// The buffer will be allocated in the device's memory.
|
||
|
|
virtual Result<std::unique_ptr<Buffer>> AllocateBuffer(int64_t size) = 0;
|
||
|
|
|
||
|
|
/// \brief Copy a Buffer to a destination MemoryManager
|
||
|
|
///
|
||
|
|
/// See also the Buffer::Copy shorthand.
|
||
|
|
static Result<std::shared_ptr<Buffer>> CopyBuffer(
|
||
|
|
const std::shared_ptr<Buffer>& source, const std::shared_ptr<MemoryManager>& to);
|
||
|
|
|
||
|
|
/// \brief Copy a non-owned Buffer to a destination MemoryManager
|
||
|
|
///
|
||
|
|
/// This is useful for cases where the source memory area is externally managed
|
||
|
|
/// (its lifetime not tied to the source Buffer), otherwise please use CopyBuffer().
|
||
|
|
static Result<std::unique_ptr<Buffer>> CopyNonOwned(
|
||
|
|
const Buffer& source, const std::shared_ptr<MemoryManager>& to);
|
||
|
|
|
||
|
|
/// \brief Make a no-copy Buffer view in a destination MemoryManager
|
||
|
|
///
|
||
|
|
/// See also the Buffer::View shorthand.
|
||
|
|
static Result<std::shared_ptr<Buffer>> ViewBuffer(
|
||
|
|
const std::shared_ptr<Buffer>& source, const std::shared_ptr<MemoryManager>& to);
|
||
|
|
|
||
|
|
/// \brief Copy a slice of a buffer into a CPU pointer
|
||
|
|
static Status CopyBufferSliceToCPU(const std::shared_ptr<Buffer>& buf, int64_t offset,
|
||
|
|
int64_t length, uint8_t* out_data);
|
||
|
|
|
||
|
|
/// \brief Create a new SyncEvent.
|
||
|
|
///
|
||
|
|
/// This version should construct the appropriate event for the device and
|
||
|
|
/// provide the unique_ptr with the correct deleter for the event type.
|
||
|
|
/// If the device does not require or work with any synchronization, it is
|
||
|
|
/// allowed for it to return a nullptr.
|
||
|
|
virtual Result<std::shared_ptr<Device::SyncEvent>> MakeDeviceSyncEvent();
|
||
|
|
|
||
|
|
/// \brief Wrap an event into a SyncEvent.
|
||
|
|
///
|
||
|
|
/// @param sync_event passed in sync_event (should be a pointer to the appropriate type)
|
||
|
|
/// @param release_sync_event destructor to free sync_event. `nullptr` may be
|
||
|
|
/// passed to indicate that no destruction/freeing is necessary
|
||
|
|
virtual Result<std::shared_ptr<Device::SyncEvent>> WrapDeviceSyncEvent(
|
||
|
|
void* sync_event, Device::SyncEvent::release_fn_t release_sync_event);
|
||
|
|
|
||
|
|
protected:
|
||
|
|
ARROW_DISALLOW_COPY_AND_ASSIGN(MemoryManager);
|
||
|
|
|
||
|
|
explicit MemoryManager(const std::shared_ptr<Device>& device) : device_(device) {}
|
||
|
|
|
||
|
|
// Default implementations always return nullptr, should be overridden
|
||
|
|
// by subclasses that support data transfer.
|
||
|
|
// (returning nullptr means unsupported copy / view)
|
||
|
|
// In CopyBufferFrom and ViewBufferFrom, the `from` parameter is guaranteed to
|
||
|
|
// be equal to `buf->memory_manager()`.
|
||
|
|
virtual Result<std::shared_ptr<Buffer>> CopyBufferFrom(
|
||
|
|
const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& from);
|
||
|
|
virtual Result<std::shared_ptr<Buffer>> CopyBufferTo(
|
||
|
|
const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& to);
|
||
|
|
virtual Result<std::unique_ptr<Buffer>> CopyNonOwnedFrom(
|
||
|
|
const Buffer& buf, const std::shared_ptr<MemoryManager>& from);
|
||
|
|
virtual Result<std::unique_ptr<Buffer>> CopyNonOwnedTo(
|
||
|
|
const Buffer& buf, const std::shared_ptr<MemoryManager>& to);
|
||
|
|
virtual Result<std::shared_ptr<Buffer>> ViewBufferFrom(
|
||
|
|
const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& from);
|
||
|
|
virtual Result<std::shared_ptr<Buffer>> ViewBufferTo(
|
||
|
|
const std::shared_ptr<Buffer>& buf, const std::shared_ptr<MemoryManager>& to);
|
||
|
|
|
||
|
|
std::shared_ptr<Device> device_;
|
||
|
|
};
|
||
|
|
|
||
|
|
// ----------------------------------------------------------------------
|
||
|
|
// CPU backend implementation
|
||
|
|
|
||
|
|
class ARROW_EXPORT CPUDevice : public Device {
|
||
|
|
public:
|
||
|
|
const char* type_name() const override;
|
||
|
|
std::string ToString() const override;
|
||
|
|
bool Equals(const Device&) const override;
|
||
|
|
DeviceAllocationType device_type() const override { return DeviceAllocationType::kCPU; }
|
||
|
|
|
||
|
|
std::shared_ptr<MemoryManager> default_memory_manager() override;
|
||
|
|
|
||
|
|
/// \brief Return the global CPUDevice instance
|
||
|
|
static std::shared_ptr<Device> Instance();
|
||
|
|
|
||
|
|
/// \brief Create a MemoryManager
|
||
|
|
///
|
||
|
|
/// The returned MemoryManager will use the given MemoryPool for allocations.
|
||
|
|
static std::shared_ptr<MemoryManager> memory_manager(MemoryPool* pool);
|
||
|
|
|
||
|
|
protected:
|
||
|
|
CPUDevice() : Device(true) {}
|
||
|
|
};
|
||
|
|
|
||
|
|
class ARROW_EXPORT CPUMemoryManager : public MemoryManager {
|
||
|
|
public:
|
||
|
|
Result<std::shared_ptr<io::RandomAccessFile>> GetBufferReader(
|
||
|
|
std::shared_ptr<Buffer> buf) override;
|
||
|
|
Result<std::shared_ptr<io::OutputStream>> GetBufferWriter(
|
||
|
|
std::shared_ptr<Buffer> buf) override;
|
||
|
|
|
||
|
|
Result<std::unique_ptr<Buffer>> AllocateBuffer(int64_t size) override;
|
||
|
|
|
||
|
|
/// \brief Return the MemoryPool associated with this MemoryManager.
|
||
|
|
MemoryPool* pool() const { return pool_; }
|
||
|
|
|
||
|
|
protected:
|
||
|
|
CPUMemoryManager(const std::shared_ptr<Device>& device, MemoryPool* pool)
|
||
|
|
: MemoryManager(device), pool_(pool) {}
|
||
|
|
|
||
|
|
static std::shared_ptr<MemoryManager> Make(const std::shared_ptr<Device>& device,
|
||
|
|
MemoryPool* pool = default_memory_pool());
|
||
|
|
|
||
|
|
Result<std::shared_ptr<Buffer>> CopyBufferFrom(
|
||
|
|
const std::shared_ptr<Buffer>& buf,
|
||
|
|
const std::shared_ptr<MemoryManager>& from) override;
|
||
|
|
Result<std::shared_ptr<Buffer>> CopyBufferTo(
|
||
|
|
const std::shared_ptr<Buffer>& buf,
|
||
|
|
const std::shared_ptr<MemoryManager>& to) override;
|
||
|
|
Result<std::unique_ptr<Buffer>> CopyNonOwnedFrom(
|
||
|
|
const Buffer& buf, const std::shared_ptr<MemoryManager>& from) override;
|
||
|
|
Result<std::unique_ptr<Buffer>> CopyNonOwnedTo(
|
||
|
|
const Buffer& buf, const std::shared_ptr<MemoryManager>& to) override;
|
||
|
|
Result<std::shared_ptr<Buffer>> ViewBufferFrom(
|
||
|
|
const std::shared_ptr<Buffer>& buf,
|
||
|
|
const std::shared_ptr<MemoryManager>& from) override;
|
||
|
|
Result<std::shared_ptr<Buffer>> ViewBufferTo(
|
||
|
|
const std::shared_ptr<Buffer>& buf,
|
||
|
|
const std::shared_ptr<MemoryManager>& to) override;
|
||
|
|
|
||
|
|
MemoryPool* pool_;
|
||
|
|
|
||
|
|
friend std::shared_ptr<MemoryManager> CPUDevice::memory_manager(MemoryPool* pool);
|
||
|
|
ARROW_FRIEND_EXPORT friend std::shared_ptr<MemoryManager> default_cpu_memory_manager();
|
||
|
|
};
|
||
|
|
|
||
|
|
/// \brief Return the default CPU MemoryManager instance
|
||
|
|
///
|
||
|
|
/// The returned singleton instance uses the default MemoryPool.
|
||
|
|
/// This function is a faster spelling of
|
||
|
|
/// `CPUDevice::Instance()->default_memory_manager()`.
|
||
|
|
ARROW_EXPORT
|
||
|
|
std::shared_ptr<MemoryManager> default_cpu_memory_manager();
|
||
|
|
|
||
|
|
using DeviceMapper =
|
||
|
|
std::function<Result<std::shared_ptr<MemoryManager>>(int64_t device_id)>;
|
||
|
|
|
||
|
|
/// \brief Register a function to retrieve a MemoryManager for a Device type
|
||
|
|
///
|
||
|
|
/// This registers the device type globally. A specific device type can only
|
||
|
|
/// be registered once. This method is thread-safe.
|
||
|
|
///
|
||
|
|
/// Currently, this registry is only used for importing data through the C Device
|
||
|
|
/// Data Interface (for the default Device to MemoryManager mapper in
|
||
|
|
/// arrow::ImportDeviceArray/ImportDeviceRecordBatch).
|
||
|
|
///
|
||
|
|
/// \param[in] device_type the device type for which to register a MemoryManager
|
||
|
|
/// \param[in] mapper function that takes a device id and returns the appropriate
|
||
|
|
/// MemoryManager for the registered device type and given device id
|
||
|
|
/// \return Status
|
||
|
|
ARROW_EXPORT
|
||
|
|
Status RegisterDeviceMapper(DeviceAllocationType device_type, DeviceMapper mapper);
|
||
|
|
|
||
|
|
/// \brief Get the registered function to retrieve a MemoryManager for the
|
||
|
|
/// given Device type
|
||
|
|
///
|
||
|
|
/// \param[in] device_type the device type
|
||
|
|
/// \return function that takes a device id and returns the appropriate
|
||
|
|
/// MemoryManager for the registered device type and given device id
|
||
|
|
ARROW_EXPORT
|
||
|
|
Result<DeviceMapper> GetDeviceMapper(DeviceAllocationType device_type);
|
||
|
|
|
||
|
|
} // namespace arrow
|