CUDA 支持#

上下文 (Contexts)#

class CudaDeviceManager#

公共函数

Result<std::shared_ptr<CudaDevice>> GetDevice(int device_number)#

获取特定设备的 CudaDevice 实例。

参数:

device_number – [in] CUDA 设备编号

Result<std::shared_ptr<CudaContext>> GetContext(int device_number)#

获取特定设备的 CUDA 驱动程序上下文。

参数:

device_number – [in] CUDA 设备编号

返回：

缓存的上下文

Result<std::shared_ptr<CudaContext>> GetSharedContext(int device_number, void *handle)#

获取特定设备的共享 CUDA 驱动程序上下文。

参数:

  • device_number – [in] CUDA 设备编号

  • handle – [in] 由另一个库创建的 CUDA 上下文句柄

返回：

共享上下文

Result<std::shared_ptr<CudaHostBuffer>> AllocateHost(int device_number, int64_t nbytes)#

分配可对指定 GPU 设备进行快速访问的主机内存。

参数:

  • device_number – [in] CUDA 设备编号

  • nbytes – [in] 字节数

返回：

主机缓冲区或 Status

Status FreeHost(void *data, int64_t nbytes)#

释放主机内存。

给定的内存指针必须是之前通过 AllocateHost 分配的。

class CudaContext : public std::enable_shared_from_this<CudaContext>#

低级 CUDA 驱动程序 API 的面向对象接口。

公共函数

Result<std::unique_ptr<CudaBuffer>> Allocate(int64_t nbytes)#

为该上下文在 GPU 设备上分配 CUDA 内存。

参数:

nbytes – [in] 字节数

返回：

分配的缓冲区

Status Free(void *device_ptr, int64_t nbytes)#

为该上下文释放 GPU 设备上的 CUDA 内存。

参数:

  • device_ptr – [in] 缓冲区地址

  • nbytes – [in] 字节数

返回：

Status（状态）

Result<std::shared_ptr<CudaBuffer>> View(uint8_t *data, int64_t nbytes)#

为此上下文在 GPU 设备上创建 CUDA 内存的视图。

注意

调用者负责分配和释放内存，并确保内存属于此 CudaContext 实例所持有的 CUDA 上下文。

参数:

  • data – [in] 起始设备地址

  • nbytes – [in] 字节数

返回：

视图缓冲区

Result<std::shared_ptr<CudaBuffer>> OpenIpcBuffer(const CudaIpcMemHandle &ipc_handle)#

打开现有的 CUDA IPC 内存句柄。

参数:

ipc_handle – [in] 指向 CUipcMemHandle 的不透明指针 (驱动程序 API)

返回：

引用 IPC 片段的 CudaBuffer

Status CloseIpcBuffer(CudaBuffer *buffer)#

关闭通过 IPC 缓冲区内存映射。

参数:

buffer – [in] 引用该缓冲区的 CudaBuffer

返回：

Status（状态）

Status Synchronize(void)#

阻塞直至所有设备任务完成。

void *handle() const#

向其他库公开 CUDA 上下文句柄。

std::shared_ptr<CudaMemoryManager> memory_manager() const#

返回绑定到此上下文设备的默认内存管理器。

std::shared_ptr<CudaDevice> device() const#

返回与此上下文关联的设备实例。

int device_number() const#

返回逻辑设备编号。

Result<uintptr_t> GetDeviceAddress(uint8_t *addr)#

返回在该上下文中运行的内核可访问的设备地址。

设备地址定义为设备可访问的内存地址。虽然通常是设备内存地址，但它也可以是主机内存地址，例如，当内存被分配为托管内存 (使用 cudaMallocManaged)、主机内存 (使用 cudaMallocHost 或 cudaHostAlloc) 或者主机内存被分页锁定 (使用 cudaHostRegister) 时。

参数:

addr – [in] 设备或主机内存地址

返回：

设备地址

设备 (Devices)#

class CudaDevice : public arrow::Device#

CUDA 的 Device 实现。

每个 CudaDevice 实例都绑定到特定的 CUDA 设备（通过其逻辑设备编号标识）。

公共函数

virtual const char *type_name() const override#

此设备类型的简写。

返回的值对于每个设备类都是不同的，但对于给定类的所有实例都是相同的。它可以用作 RTTI 的替代品。

virtual std::string ToString() const override#

设备的易读描述。

返回的值应足够详细，以便在必要时区分不同的实例。

virtual bool Equals(const Device&) const override#

此实例是否指向与另一个相同的设备。

virtual std::shared_ptr<MemoryManager> default_memory_manager() override#

返回绑定到此设备的 MemoryManager 实例。

返回的实例使用此设备类型的 MemoryManager 实现的默认参数。有些设备也允许构建具有非默认参数的 MemoryManager 实例。

inline virtual DeviceAllocationType device_type() const override#

返回此设备的 DeviceAllocationType。

inline virtual int64_t device_id() const override#

如果存在多个此类型的设备，用于标识该设备的设备 ID。

如果没有等效的 “device_id”（例如非 NUMA 系统上的主 CPU 设备），则返回 -1。

int device_number() const#

返回设备逻辑编号。

std::string device_name() const#

返回 GPU 型号名称。

int64_t total_memory() const#

返回此设备的总内存。

int handle() const#

返回原始 CUDA 设备句柄。

返回的值可用于向其他库公开此设备。它应被解释为 CUdevice。

Result<std::shared_ptr<CudaContext>> GetContext()#

获取此设备的 CUDA 驱动程序上下文。

返回的上下文与设备的初始 CUDA 上下文相关联。这是获取设备上下文的推荐方式，因为它允许与任何使用初始 CUDA 上下文 API 的库透明地进行互操作。

Result<std::shared_ptr<CudaContext>> GetSharedContext(void *handle)#

使用现有句柄获取此设备的 CUDA 驱动程序上下文。

该句柄不归此实例所有：当 CudaContext 被销毁时，它不会被释放。此函数仅在需要与使用非初始上下文的库进行互操作时使用。

参数:

handle – [in] 由另一个库创建的 CUDA 上下文句柄

Result<std::shared_ptr<CudaHostBuffer>> AllocateHostBuffer(int64_t size)#

分配一个驻留在主机上、GPU 可访问的缓冲区。

该缓冲区是使用此设备的初始上下文分配的。

参数:

size – [in] 缓冲区大小（字节）

virtual Result<std::shared_ptr<Device::Stream>> MakeStream(unsigned int flags) override#

在当前上下文中创建 CUstream 包装器。

virtual Result<std::shared_ptr<Device::Stream>> WrapStream(void *device_stream, Stream::release_fn_t release_fn) override#

包装一个指向现有流的指针。

参数:

  • device_stream – 传入的流 (应该是 CUstream*)

  • release_fn – 用于释放流的析构函数。可以传递 nullptr 表示不需要销毁/释放。

公共静态函数

static Result<std::shared_ptr<CudaDevice>> Make(int device_number)#

返回特定设备的 CudaDevice 实例。

参数:

device_number – [in] CUDA 设备编号

class Stream : public arrow::Device::Stream#

实验性功能：CUstreams 的包装器。

不拥有 CUstream 对象，该对象必须使用 cuStreamCreate 和 cuStreamDestroy（或等效函数）单独构建和释放。默认构造将使用 cuda 默认流，且不允许通过字面量 0 或 nullptr 进行构造。

公共函数

virtual Status WaitEvent(const Device::SyncEvent&) override#

使流等待提供的事件。

告诉流在不同步事件完成之前进行等待，且不会阻塞 CPU。

virtual Status Synchronize() const override#

阻塞当前线程，直至流中剩余任务完成。

class SyncEvent : public arrow::Device::SyncEvent#

公共函数

virtual Status Wait() override#

阻塞直至同步事件被标记为完成。

virtual Status Record(const Device::Stream&) override#

在流上记录包装的事件。

一旦流完成了先前添加到其中的任务，它将触发该事件。

class CudaMemoryManager : public arrow::MemoryManager#

CUDA 的 MemoryManager 实现。

公共函数

virtual Result<std::shared_ptr<io::RandomAccessFile>> GetBufferReader(std::shared_ptr<Buffer> buf) override#

创建 RandomAccessFile 以读取特定缓冲区。

给定的缓冲区必须绑定到此 MemoryManager。

另请参阅 Buffer::GetReader 简写。

virtual Result<std::shared_ptr<io::OutputStream>> GetBufferWriter(std::shared_ptr<Buffer> buf) override#

创建 OutputStream 以写入特定缓冲区。

给定的缓冲区必须是可变的且绑定到此 MemoryManager。返回的流对象会写入缓冲区的基础内存（但不会调整其大小）。

另请参阅 Buffer::GetWriter 简写。

virtual Result<std::unique_ptr<Buffer>> AllocateBuffer(int64_t size) override#

分配 (可变的) Buffer。

缓冲区将在设备内存中分配。

std::shared_ptr<CudaDevice> cuda_device() const#

绑定到此 MemoryManager 的 CudaDevice 实例。

这是一个有用的简写，返回具体类型的指针，无需转换 device() 的结果。

virtual Result<std::shared_ptr<Device::SyncEvent>> MakeDeviceSyncEvent() override#

创建包装的 CUevent。

将在内部调用 cuEventCreate，并在事件被销毁时调用 cuEventDestroy。

virtual Result<std::shared_ptr<Device::SyncEvent>> WrapDeviceSyncEvent(void *sync_event, Device::SyncEvent::release_fn_t release_sync_event) override#

将现有事件包装为同步事件。

参数:

  • sync_event – 要包装的事件，必须是 CUevent*

  • release_sync_event – 销毁期间调用的函数，可以传递 nullptr 或无操作函数表示所有权由外部维护

缓冲区 (Buffers)#

class CudaBuffer : public arrow::Buffer#

位于 GPU 设备上的 Arrow 缓冲区。

在任何可能不支持 GPU 的 Arrow 代码中使用时需格外小心。

公共函数

Status CopyToHost(const int64_t position, const int64_t nbytes, void *out) const#

将内存从 GPU 设备复制到 CPU 主机。

参数:

  • position – [in] 缓冲区内开始复制的起始位置

  • nbytes – [in] 要复制的字节数

  • out – [out] 复制到的主机内存区域起始地址

返回：

Status（状态）

Status CopyFromHost(const int64_t position, const void *data, int64_t nbytes)#

将内存从主机复制到设备的指定位置。

参数:

  • position – [in] 复制到的起始位置

  • data – [in] 要复制的主机数据

  • nbytes – [in] 要复制的字节数

返回：

Status（状态）

Status CopyFromDevice(const int64_t position, const void *data, int64_t nbytes)#

将内存从设备复制到设备的指定位置。

注意

假定源和目标设备内存都在同一个上下文中分配。

参数:

  • position – [in] 缓冲区内要复制到的起始位置

  • data – [in] 要复制的设备内存区域起始地址

  • nbytes – [in] 要复制的字节数

返回：

Status（状态）

Status CopyFromAnotherDevice(const std::shared_ptr<CudaContext> &src_ctx, const int64_t position, const void *data, int64_t nbytes)#

将内存从另一个设备复制到当前设备的指定位置。

参数:

  • src_ctx – [in] 源设备内存的上下文

  • position – [in] 缓冲区内要复制到的起始位置

  • data – [in] 要复制的另一个设备内存区域起始地址

  • nbytes – [in] 要复制的字节数

返回：

Status（状态）

virtual Result<std::shared_ptr<CudaIpcMemHandle>> ExportForIpc()#

将此设备缓冲区公开为可在其他进程中使用的 IPC 内存。

注意

调用此函数后，当 CudaBuffer 被析构时，此设备内存将不会被释放。

返回：

句柄或 Status

公共静态函数

static Result<std::shared_ptr<CudaBuffer>> FromBuffer(std::shared_ptr<Buffer> buffer)#

将通用缓冲区转换回 CudaBuffer。

注意

如果缓冲区不是由 GPU 内存支持，则此函数返回错误。

参数:

buffer – [in] 要转换的缓冲区

返回：

CudaBuffer 或 Status

class CudaHostBuffer : public arrow::MutableBuffer#

使用 cudaHostAlloc 创建的可被设备访问的 CPU 内存。

公共函数

Result<uintptr_t> GetDeviceAddress(const std::shared_ptr<CudaContext> &ctx)#

返回 GPU 可用于读取此内存的设备地址。

内存输入/输出 (Memory Input / Output)#

class CudaBufferReader : public arrow::io::internal::RandomAccessFileConcurrencyWrapper<CudaBufferReader>#

用于从 CUDA 缓冲区进行零拷贝读取的文件接口。

注意：读取到 Buffer 会返回指向设备内存的 Buffer。它通常与期望缓冲区指向 CPU 内存的 Arrow 代码不兼容。然而，读取到原始指针会将设备内存复制到所指向的主机内存中。

公共函数

virtual bool closed() const override#

返回流是否已关闭。

virtual bool supports_zero_copy() const override#

如果 InputStream 支持零拷贝 Buffer 读取，则返回 true。

零拷贝读取意味着使用返回缓冲区的 Read() 重载。

class CudaBufferWriter : public arrow::io::WritableFile#

用于写入 CUDA 缓冲区的文件接口，支持可选的缓冲功能。

公共函数

virtual Status Close() override#

关闭写入器并将缓冲的字节刷新到 GPU。

virtual bool closed() const override#

返回流是否已关闭。

virtual Status Flush() override#

将缓冲的字节刷新到 GPU。

virtual Status Write(const void *data, int64_t nbytes) override#

将给定数据写入流中。

此方法总是完整地处理字节。根据流的语义，数据可能会立即写入、保存在缓冲区中或异步写入。如果流对数据进行缓冲，则会进行数据拷贝。为避免潜在的大规模拷贝，请使用接受已拥有的 Buffer 的 Write 变体。

virtual Result<int64_t> Tell() const override#

返回此流中的当前位置。

Status SetBufferSize(const int64_t buffer_size)#

设置 CPU 缓冲区大小以限制 cudaMemcpy 的调用次数。

默认情况下，写入是不缓冲的。

参数:

buffer_size – [in] 要分配的 CPU 缓冲区大小

返回：

Status（状态）

int64_t buffer_size() const#

返回主机 (CPU) 缓冲区大小，0 表示未缓冲。

int64_t num_bytes_buffered() const#

返回在主机上缓冲的字节数。

IPC#

class CudaIpcMemHandle#

公共函数

Result<std::shared_ptr<Buffer>> Serialize(MemoryPool *pool = default_memory_pool()) const#

将 CudaIpcMemHandle 写入 Buffer。

参数:

pool – [in] 用于分配内存的 MemoryPool

返回：

Buffer 或 Status

公共静态函数

static Result<std::shared_ptr<CudaIpcMemHandle>> FromBuffer(const void *opaque_handle)#

从不透明缓冲区（例如来自另一个进程）创建 CudaIpcMemHandle

（例如来自另一个进程）

参数:

opaque_handle – [in] 作为 const void* 的 CUipcMemHandle

返回：

句柄或 Status

Result<std::shared_ptr<CudaBuffer>> SerializeRecordBatch(const RecordBatch &batch, CudaContext *ctx)#

将记录批次（RecordBatch）消息写入 GPU 设备内存。

参数:

  • batch – [in] 要写入的记录批次

  • ctx – [in] 用于分配设备内存的 CudaContext

返回：

CudaBuffer 或 Status

Result<std::shared_ptr<RecordBatch>> ReadRecordBatch(const std::shared_ptr<Schema> &schema, const ipc::DictionaryMemo *dictionary_memo, const std::shared_ptr<CudaBuffer> &buffer, MemoryPool *pool = default_memory_pool())#

专门处理 CUDA 设备上元数据的 ReadRecordBatch。

参数:

  • schema – [in] 记录批次的 Schema

  • dictionary_memo – [in] 包含所有字典的 DictionaryMemo。如果您确定没有字典编码字段，则可以为 nullptr

  • buffer – [in] 包含完整 IPC 消息的 CudaBuffer

  • pool – [in] 用于为元数据分配空间的 MemoryPool

返回：

RecordBatch 或 Status