// 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 {
/// \brief EXPERIMENTAL: Device type enum which matches up with C Data Device types
enum class DeviceAllocationType : char {
kCPU = 1,
kCUDA = 2,
kCUDA_HOST = 3,
kOPENCL = 4,
kVULKAN = 7,
kMETAL = 8,
kVPI = 9,
kROCM = 10,
kROCM_HOST = 11,
kEXT_DEV = 12,
kCUDA_MANAGED = 13,
kONEAPI = 14,
kWEBGPU = 15,
kHEXAGON = 16,
};
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 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* device_stream,
Stream::release_fn_t 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 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