Apache Arrow C++ CUDA 支持：GPU 内存管理、缓冲区与 IPC 编程完整指南

Apache Arrow C CUDA 支持GPU 内存管理、缓冲区与 IPC 编程完整指南【免费下载链接】arrowApache Arrow is a multi-language toolbox for accelerated data interchange and in-memory processing项目地址: https://gitcode.com/gh_mirrors/arrow13/arrow导读Apache Arrow C 在arrow::cuda命名空间中提供了完整的 CUDA 扩展模块让开发者可以在 Arrow 的统一内存抽象之上直接操作 GPU 显存从设备/上下文管理、显存缓冲区分配与读写到跨进程的 CUDA IPC 共享再到把 RecordBatch 直接序列化到显存。本文以官方 API 参考文档 docs/source/cpp/api/cuda.rst 为骨架结合仓库中 cpp/src/arrow/gpu 目录下的真实源码实现与测试用例系统讲解每个核心类的职责、关键 API 与底层原理并给出可落地的构建与编码示例帮助读者在 C 项目中正确启用并驾驭 Arrow 的 GPU 能力。模块概览Arrow C 的 CUDA 扩展层docs/source/cpp/api/cuda.rst是 Arrow C API Reference见 docs/source/cpp/api.rst 的api/cuda条目中专门面向 CUDA 支持的章节其将整个 GPU 扩展划分为五组核心抽象分组核心类 / 函数职责Contextsarrow::cuda::CudaDeviceManager、arrow::cuda::CudaContext设备与 CUDA driver context 的获取、生命周期管理Devicesarrow::cuda::CudaDevice、arrow::cuda::CudaMemoryManager接入 Arrow 的Device/MemoryManager设备抽象体系Buffersarrow::cuda::CudaBuffer、arrow::cuda::CudaHostBufferGPU 显存缓冲区与设备可访问的锁页主机内存Memory I/Oarrow::cuda::CudaBufferReader、arrow::cuda::CudaBufferWriter面向缓冲区Buffer的文件式读写接口IPCarrow::cuda::CudaIpcMemHandle、cuda-ipc-functions函数组跨进程共享显存、RecordBatch 的显存序列化/反序列化这些类全部位于 cpp/src/arrow/gpu/cuda_context.h 与 cpp/src/arrow/gpu/cuda_memory.h 两个头文件中IPC 函数组则在 cpp/src/arrow/gpu/cuda_arrow_ipc.h编译为一个独立的动态/静态库arrow_cuda因此 GPU 支持完全可选、按需加载。构建与启用ARROW_CUDACUDA 模块默认不参与构建。需要在 CMake 配置时显式打开ARROW_CUDA选项cmake -DARROW_CUDAON -DARROW_IPCON ...该选项在 cpp/cmake_modules/DefineOptions.cmake#L314-L318 中定义说明文字为 Build the Arrow CUDA extensions (requires CUDA toolkit)且强制依赖ARROW_IPC——因为 CUDA 模块不仅要提供裸显存缓冲区还要支持 IPC 消息的序列化见下文 IPC 章节。从 cpp/src/arrow/gpu/CMakeLists.txt 可以还原出完整的构建细节支持通过环境变量CUDA_HOME指定 CUDA 工具包根目录set(CUDA_TOOLKIT_ROOT_DIR $ENV{CUDA_HOME})CMake 版本 ≥ 3.17 时使用find_package(CUDAToolkit REQUIRED)并链接CUDA::cuda_driver老版本退化为find_package(CUDA)因为旧 CMake 无法识别 stubs 目录下的 driver 库源码由 4 个文件组成cuda_arrow_ipc.cc、cuda_context.cc、cuda_internal.cc、cuda_memory.cc产出libarrow_cudaArrowCUDA包 /arrow-cuda.pcpkg-config 名同时生成cuda_version.h并安装到 include 目录头文件统一安装到arrow/gpu/附带测试目标cuda_test见 cpp/src/arrow/gpu/cuda_test.cc与基准arrow-gpu-cuda-benchmark见 cpp/src/arrow/gpu/cuda_benchmark.cc。编译宏ARROW_CUDA会写入 cpp/src/arrow/util/config.h.cmake上层代码如 Arrow Flight 的 perf_server.cc据此做条件编译。ContextsCudaDeviceManager 与 CudaContextCudaDeviceManager一切 GPU 操作的入口CudaDeviceManager是 CUDA 模块的全局入口其完整 API 定义在 cpp/src/arrow/gpu/cuda_context.h#L44-L85static ResultCudaDeviceManager* Instance()获取单例实例Resultstd::shared_ptrCudaDevice GetDevice(int device_number)按逻辑设备号获取设备实例Resultstd::shared_ptrCudaContext GetContext(int device_number)获取该设备缓存的 CUDA driver contextResultstd::shared_ptrCudaContext GetSharedContext(int device_number, void* handle)基于其他库创建的既有 CUDA context handle 构造共享 context句柄不转移所有权Resultstd::shared_ptrCudaHostBuffer AllocateHost(int device_number, int64_t nbytes)在 CPU 侧分配对指定 GPU 快速可访问的锁页内存Status FreeHost(void* data, int64_t nbytes)释放上述内存int num_devices() const返回可用 GPU 数量。CudaContextCUDA driver API 的面向对象封装CudaContext的头文件注释cuda_context.h#L317-L318直接点明其定位Object-oriented interface to the low-level CUDA driver API。它基于 CUDAdriver APIcu*函数族而非 runtime API 构建核心能力包括显存分配与释放Allocate(nbytes)返回CudaBufferFree(device_ptr, nbytes)手动释放显存视图View(data, nbytes)对外部拥有的设备地址创建CudaBuffer视图不拥有、不释放底层内存调用方需保证内存属于该 context同步Synchronize()阻塞至所有设备任务完成bytes_allocated()返回本 context 已分配的字节数互操作handle()暴露原始 CUDA context 句柄供其他库使用GetDeviceAddress()把一个内存地址解析为内核可达的设备地址该地址可能是设备显存、cudaMallocHost/cudaMallocManaged分配的主机内存或cudaHostRegister锁页的内存关系查询device()返回关联的CudaDevicememory_manager()返回关联的CudaMemoryManagerdevice_number()返回逻辑设备号生命周期Close()显式关闭 context。底层实现context 默认通过cuDevicePrimaryCtxRetain获取设备的主 contextprimary contextClose()时调用cuDevicePrimaryCtxRelease释放而通过GetSharedContext创建的共享 context 不拥有底层句柄析构时不会释放——这两条路径分别对应 cpp/src/arrow/gpu/cuda_context.cc#L80-L106 中的Impl::Init与Impl::InitShared。文件头部的struct DevicePropertiescuda_context.cc#L45-L63展示了设备元数据设备号、CUdevice句柄、总显存、GPU 名称的获取方式依次调用cuDeviceGet、cuDeviceTotalMem、cuDeviceGetName。此外cpp/src/arrow/gpu/cuda_internal.h 提供了模块内部的错误处理工具StatusFromCuda把CUresult转换为arrow::Status宏CU_RETURN_NOT_OK快速包装调用ContextSaver则在构造时cuCtxPushCurrent、析构时cuCtxPopCurrent保证在 CUDA 的当前 context 栈上安全切换。DevicesCudaDevice 与 CudaMemoryManager接入 Arrow 设备抽象体系Arrow C 用arrow::Device表示 CPU/GPU 等计算设备用arrow::MemoryManager描述在某设备上的分配策略这一设计在 docs/source/cpp/memory.rst 的 Devices 一节有系统说明非 CPU 设备如 GPU 显存正是通过该抽象统一建模的。CudaDevice继承arrow::DeviceCudaMemoryManager继承arrow::MemoryManager从而让 GPU 内存参与 Arrow 统一的缓冲区生命周期与拷贝体系。CudaDevice绑定单一逻辑 GPUCudaDevice的每个实例绑定一个逻辑设备号APIcuda_context.h#L91-L240包括static Make(int device_number)工厂方法device_number()、device_name()、total_memory()返回设备元信息device_type()返回DeviceAllocationType::kCUDAdevice_id()即设备号handle()返回原始CUdevice句柄可用于向其他库暴露设备GetContext()/GetSharedContext(handle)获取该设备的主 context 或基于外部句柄的共享 contextAllocateHostBuffer(size)用该设备主 context 分配 GPU 可访问的主机缓冲区流Stream支持MakeStream()/MakeStream(flags)创建Device::Stream包装器WrapStream(device_stream, release_fn)包装既有CUstream包装器不拥有流对象需调用方自行cuStreamCreate/cuStreamDestroyStream::Synchronize()与WaitEvent提供同步能力事件支持SyncEvent类包装CUeventWait()阻塞等待事件完成、Record(stream)把事件记录到流上。CudaMemoryManager显存分配策略CudaMemoryManagercuda_context.h#L253-L304实现arrow::MemoryManager接口AllocateBuffer(size)在显存上分配缓冲区GetBufferReader(buf)/GetBufferWriter(buf)把缓冲区包装为可读/可写的文件接口对应下文 Memory I/Ocuda_device()便捷返回具体类型的CudaDevice避免device()的强制转换MakeDeviceSyncEvent()/WrapDeviceSyncEvent(...)创建或包装 CUDA 事件CopyBufferFrom/CopyBufferTo/CopyNonOwnedFrom/CopyNonOwnedTo与ViewBufferFrom/ViewBufferTo实现跨设备如 CPU↔GPU、GPU↔GPU的拷贝与零拷贝视图是上层Buffer::Copy/Buffer::View的设备级落点。arrow::cuda还导出了类型判断/转换辅助函数IsCudaDevice、AsCudaDevice、IsCudaMemoryManager、AsCudaMemoryManagercuda_context.h#L242-L250用于在通用Device/MemoryManager接口上安全地恢复具体 CUDA 类型。BuffersCudaBuffer 与 CudaHostBufferCudaBuffer位于 GPU 显存上的 Arrow 缓冲区CudaBuffer继承自arrow::Buffer是位于 GPU 设备上的 Arrow 缓冲区cuda_memory.h#L34-L108。头文件特别警告Be careful using this in any Arrow code which may not be GPU-aware——多数 Arrow 代码假设缓冲区指向 CPU 内存因此在把CudaBuffer传入通用计算/IO 管道前必须确认其 GPU 感知能力。核心 API构造支持从设备地址uintptr_t或字节指针构造own_data控制析构时是否释放底层显存is_ipc标记 IPC 映射的内存也支持从父缓冲区切出子区间offset/size 形式static FromBuffer(std::shared_ptrBuffer)把通用缓冲区转回CudaBuffer若底层不是 GPU 内存则返回错误CopyToHost(position, nbytes, out)显存 → 主机内存CopyFromHost(position, data, nbytes)主机内存 → 显存CopyFromDevice(position, data, nbytes)同 context 内的显存 → 显存CopyFromAnotherDevice(src_ctx, position, data, nbytes)跨设备拷贝对应CudaContext::CopyDeviceToAnotherDevice的调用链ExportForIpc()导出为CudaIpcMemHandle供其他进程共享注意调用后该显存在CudaBuffer析构时不再自动释放context()返回所属的CudaContext。CudaHostBuffer设备可访问的 CPU 锁页内存CudaHostBuffercuda_memory.h#L110-L120是使用cudaHostAlloc创建的设备可访问 CPU 内存继承自arrow::MutableBuffer即可写。它让 GPU 能快速访问 CPU 侧数据从而显著提升主机↔设备拷贝效率。其GetDeviceAddress(ctx)返回 GPU 读取该内存所用的设备地址。与之对应的模块级函数AllocateCudaHostBuffer(device_number, size)cuda_memory.h#L242-L252是推荐的分配入口CudaDeviceManager::AllocateHost内部即走此路径。同文件还提供两个低层地址转换函数GetDeviceAddress(cpu_data, ctx)获取 CPU 数据对应的设备地址GetHostAddress(device_ptr)反向获取设备地址对应的 CPU 地址cuda_memory.h#L254-L261。Memory Input / OutputCudaBufferReader 与 CudaBufferWriter把显存缓冲区包装为 Arrow 的io文件接口后即可无缝复用 IPC 读取、序列化等基于文件抽象的既有代码。CudaBufferReader显存零拷贝读取CudaBufferReadercuda_memory.h#L153-L199基于RandomAccessFileConcurrencyWrapper实现arrow::io::RandomAccessFile支持Read、ReadAt、Seek、Tell、GetSize等标准文件操作。文档注释中有两条关键行为约定读入Buffer时返回指向设备内存的缓冲区零拷贝此时结果通常不能直接交给期望 CPU 缓冲区的 Arrow 代码读入裸指针时则执行显存 → 主机内存的真实拷贝。supports_zero_copy()用于通告该文件接口是否支持零拷贝读取。构造函数接受任意Buffer内部会将其转换为CudaBuffer并持有 context 与设备地址。CudaBufferWriter带可选缓冲的显存写入CudaBufferWritercuda_memory.h#L201-L240实现arrow::io::WritableFile向CudaBuffer写入数据默认非缓冲每次Write直接触发设备拷贝。关键扩展方法SetBufferSize(buffer_size)设置 CPU 侧缓冲大小以限制cudaMemcpy的调用次数——小片写入先聚合在主机缓冲中Flush()或Close()时一次性刷入显存buffer_size()返回主机缓冲大小0 表示未缓冲num_bytes_buffered()返回当前主机侧缓存的字节数。这组 API 允许在低延迟直写与批量拷贝减少系统调用之间按写入模式取舍。IPC跨进程共享显存CUDA IPC 让多个进程或同一进程内多个 context共享同一块显存是 GPU 集群/多进程场景的基础能力。文档将该部分分为句柄类CudaIpcMemHandle与函数组cuda-ipc-functions。CudaIpcMemHandleIPC 句柄的封装与传输CudaIpcMemHandlecuda_memory.h#L122-L151封装 CUDA driver API 的CUipcMemHandlestatic FromBuffer(const void* opaque_handle)从例如其他进程传来的不透明字节流构造句柄对象Serialize(MemoryPool* pool)把句柄序列化为arrow::Buffer便于通过网络或本地通道传递给其他进程。句柄本身由CudaBuffer::ExportForIpc()产生。导出后接收方调用CudaContext::OpenIpcBuffer(handle)把 IPC 显存段映射为本进程的CudaBuffer用完以CloseIpcBuffer(buffer)关闭映射对应 cuda_context.h#L346-L354。cuda-ipc-functionsRecordBatch 的显存直通该函数组cpp/src/arrow/gpu/cuda_arrow_ipc.h#L44-L69只有两个函数却打通了Arrow 列式数据 ↔ GPU 显存的完整链路// 把 RecordBatch 序列化写入 GPU 设备内存 Resultstd::shared_ptrCudaBuffer SerializeRecordBatch( const RecordBatch batch, CudaContext* ctx); // 从显存中的完整 IPC 消息读取 RecordBatch Resultstd::shared_ptrRecordBatch ReadRecordBatch( const std::shared_ptrSchema schema, const ipc::DictionaryMemo* dictionary_memo, // 可传 nullptr无字典字段时 const std::shared_ptrCudaBuffer buffer, MemoryPool* pool default_memory_pool()); // 仅用于元数据分配实现细节cpp/src/arrow/gpu/cuda_arrow_ipc.ccSerializeRecordBatch内部调用通用的ipc::SerializeRecordBatch(batch, ctx-memory_manager())——即直接以CudaMemoryManager作为分配器完成消息写入再经CudaBuffer::FromBuffer转回显存缓冲区ReadRecordBatch先用CudaBufferReader包装显存缓冲区调用ipc::ReadMessage解析消息pool只用于元数据分配随后ipc::ReadRecordBatch在设备内存上执行零拷贝读取数据体不经过主机中转。因此生产者进程SerializeRecordBatch→ExportForIpc→ 传输句柄消费者进程FromBuffer→OpenIpcBuffer→ReadRecordBatch即可实现跨进程的显存到显存数据交换。实战完整的使用流程源自仓库测试用例cpp/src/arrow/gpu/cuda_test.cc#L74-L120 的TestCudaBase展示了标准的初始化与内存操作范式可作为实际编码模板#include arrow/gpu/cuda_api.h // 统一入口头文件 using arrow::cuda::CudaDeviceManager; using arrow::cuda::AsCudaMemoryManager; // 1. 获取全局管理器与设备 ASSERT_OK_AND_ASSIGN(auto manager, CudaDeviceManager::Instance()); ASSERT_OK_AND_ASSIGN(auto device, manager-GetDevice(0)); // 逻辑设备号 0 // 2. 获取该设备的主 CUDA context ASSERT_OK_AND_ASSIGN(auto context, device-GetContext()); // 3. 通过默认 MemoryManager 在显存上分配缓冲区 ASSERT_OK_AND_ASSIGN(auto mm, AsCudaMemoryManager(device-default_memory_manager())); ASSERT_OK_AND_ASSIGN(auto device_buffer, mm-AllocateBuffer(size)); // 4. 主机 → 显存写入显存 → 主机回读 // (AllocateBuffer 分配 CPU 缓冲CopyFromHost/CopyToHost 完成双向拷贝)测试中还覆盖了以下进阶场景均可作为理解 API 语义的参考多 GPUkGpuNumber 0、kOtherGpuNumber 1通过CudaDevice::Make(1)探测第二块 GPU 是否存在OtherGpuDevice()cuda_test.cc#L104-L111再配合CopyFromAnotherDevice做跨设备拷贝非主 context 互操作NonPrimaryContext()通过cuCtxCreate创建非主 context 后用device-GetSharedContext(cuctx)包装cuda_test.cc#L92-L102缓冲区内容断言AssertCudaBufferEquals先用CopyToHost把显存拷回 CPU 缓冲再比较cuda_test.cc#L60-L72验证了CopyToHost的正确用法。在 CPU 侧使用这些对象时建议遵循 docs/source/cpp/memory.rst 提出的设备无关编程原则用Buffer::is_cpu()判断缓冲区是否 CPU 可读用Buffer::View/ViewOrCopy进行跨设备视图或拷贝用Buffer::GetReader/GetWriter统一获取文件接口——CudaBuffer正是通过CudaMemoryManager全面接入这套机制的。生态联动与性能基准CUDA 模块并非孤立存在Arrow Flight 集成cpp/src/arrow/flight/perf_server.cc与cpp/src/arrow/flight/flight_benchmark.cc在#ifdef ARROW_CUDA下提供 GPU 相关的性能测试与基准路径说明显存缓冲区可直接承载 Flight RPC 的大数据搬运基准测试cpp/src/arrow/gpu/cuda_benchmark.cc构建目标arrow-gpu-cuda-benchmark针对显存分配、主机↔设备拷贝等操作提供可复现的基准运行该库的基准与测试需要在配置时打开ARROW_BUILD_BENCHMARKS/ARROW_BUILD_TESTS并保证机器上存在可用 GPU测试默认使用设备 0多 GPU 用例需要第二块 GPU 才会执行语言绑定仓库的 GLib/Ruby 层c_glib/arrow-cuda-glib与ruby/red-arrow-cuda基于同一套 C CUDA API 提供了更上层的封装可作为阅读官方 C 用法的对照实现。总结arrow::cuda把 CUDA driver API 的复杂度收敛进 Arrow 的Device/MemoryManager/Buffer/io抽象之中CudaDeviceManager与CudaContext负责设备与上下文生命周期CudaDevice/CudaMemoryManager让显存参与统一的缓冲区与拷贝体系CudaBuffer/CudaHostBuffer提供显存与锁页主机内存两种载体CudaBufferReader/CudaBufferWriter赋予显存标准文件式读写能力而CudaIpcMemHandle与cuda-ipc-functions则把跨进程显存共享和 RecordBatch 显存直通串成完整链路。需要 GPU 支持时以-DARROW_CUDAON依赖ARROW_IPC构建即可随后从 cpp/src/arrow/gpu/cuda_test.cc 中的模式出发即可快速构建出高效、可移植的 GPU 数据管道。【免费下载链接】arrowApache Arrow is a multi-language toolbox for accelerated data interchange and in-memory processing项目地址: https://gitcode.com/gh_mirrors/arrow13/arrow创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考