Arrow C 设备数据接口#

警告

Arrow C 设备数据接口应被视为实验性

基本原理#

当前的C 数据接口,以及大多数对其的实现,都假设所有提供的数据缓冲区都是 CPU 缓冲区。由于 Apache Arrow 被设计为一种通用的内存中格式,用于表示表格(“列式”)数据,因此人们会希望在非 CPU 硬件(如 GPU)上利用这些数据。一个这样的例子是RAPIDS cuDF 库,它使用 Arrow 内存格式以及 CUDA 用于 NVIDIA GPU。由于在主机和设备之间复制数据代价很高,因此理想情况下,即使在运行时和库之间传递数据时,也应尽可能长时间地将数据保留在设备上。

Arrow C 设备数据接口在现有的 C 数据接口的基础上构建,向其中添加了一组非常小且稳定的 C 定义。这些定义等同于 C 数据接口中的ArrowArrayArrowArrayStream结构,它们添加了成员以允许指定设备类型并传递必要的信息以与生产者同步。对于非 C/C++ 语言和运行时,将 C 定义转换为相应的 C FFI 声明应该与当前 C 数据接口一样简单。

然后,应用程序和库可以使用 Arrow 模式和 Arrow 格式化的内存在非 CPU 设备上交换数据,就像现在使用 CPU 数据一样容易。这将使数据能够在这些设备上保留更长时间,并避免在主机和设备之间进行代价高昂的复制,仅仅是为了利用新的库和运行时。

目标#

  • 公开一个基于现有 C 数据接口构建的 ABI 稳定接口。

  • 使第三方项目能够轻松地实现支持,而无需大量的初始投入。

  • 允许在同一进程中运行的独立运行时和组件之间进行 Arrow 格式化设备内存的零拷贝共享。

  • 避免需要一对一的适配层,例如 Python 进程传递 CUDA 数据的CUDA 数组接口

  • 无需显式依赖(无论是编译时还是运行时)Arrow 软件项目本身即可实现集成。

Arrow C 设备数据接口的目的是扩展当前 C 数据接口的范围,使其也成为 GPU 或 FPGA 等设备上列式处理的标准底层构建块。

结构定义#

由于它是基于 C 数据接口构建的,因此 C 设备数据接口使用C 数据接口规范中定义的ArrowSchemaArrowArray结构。然后它添加以下独立定义。与 Arrow 项目的其余部分一样,它们在 Apache License 2.0 下可用。

#ifndef ARROW_C_DEVICE_DATA_INTERFACE
#define ARROW_C_DEVICE_DATA_INTERFACE

// Device type for the allocated memory
typedef int32_t ArrowDeviceType;

// CPU device, same as using ArrowArray directly
#define ARROW_DEVICE_CPU 1
// CUDA GPU Device
#define ARROW_DEVICE_CUDA 2
// Pinned CUDA CPU memory by cudaMallocHost
#define ARROW_DEVICE_CUDA_HOST 3
// OpenCL Device
#define ARROW_DEVICE_OPENCL 4
// Vulkan buffer for next-gen graphics
#define ARROW_DEVICE_VULKAN 7
// Metal for Apple GPU
#define ARROW_DEVICE_METAL 8
// Verilog simulator buffer
#define ARROW_DEVICE_VPI 9
// ROCm GPUs for AMD GPUs
#define ARROW_DEVICE_ROCM 10
// Pinned ROCm CPU memory allocated by hipMallocHost
#define ARROW_DEVICE_ROCM_HOST 11
// Reserved for extension
//
// used to quickly test extension devices, semantics
// can differ based on implementation
#define ARROW_DEVICE_EXT_DEV 12
// CUDA managed/unified memory allocated by cudaMallocManaged
#define ARROW_DEVICE_CUDA_MANAGED 13
// Unified shared memory allocated on a oneAPI
// non-partitioned device.
//
// A call to the oneAPI runtime is required to determine the
// device type, the USM allocation type and the sycl context
// that it is bound to.
#define ARROW_DEVICE_ONEAPI 14
// GPU support for next-gen WebGPU standard
#define ARROW_DEVICE_WEBGPU 15
// Qualcomm Hexagon DSP
#define ARROW_DEVICE_HEXAGON 16

struct ArrowDeviceArray {
  struct ArrowArray array;
  int64_t device_id;
  ArrowDeviceType device_type;
  void* sync_event;

  // reserved bytes for future expansion
  int64_t reserved[3];
};

#endif  // ARROW_C_DEVICE_DATA_INTERFACE

注意

规范保护ARROW_C_DEVICE_DATA_INTERFACE旨在避免如果两个项目在其自己的头文件中复制定义,而第三方项目包含这两个项目,则避免重复定义。因此,当复制这些定义时,务必保持此保护完全不变。

ArrowDeviceType#

ArrowDeviceType typedef 用于指示提供的内存缓冲区是在哪种类型的设备上分配的。这与device_id结合使用,应该足以引用正确的数据缓冲区。

然后我们使用宏来定义不同设备类型的值。提供的宏值与广泛使用的dlpackDLDeviceType定义值兼容,使用相同的值,因为每个值都等同于dlpack.h中的等效kDL<type>枚举。随着时间的推移,该列表将与这些等效的枚举值保持同步,以确保兼容性,而不是可能出现偏差。为了避免 Arrow 项目需要审查新的硬件设备,新的添加应该首先添加到 dlpack,然后再在此处添加相应的宏。

为了确保 ABI 的可预测性,我们使用宏而不是enum,因此存储类型不依赖于编译器。

ARROW_DEVICE_CPU#

CPU 设备,等效于直接使用ArrowArray而不是使用ArrowDeviceArray

ARROW_DEVICE_CUDA#

一个CUDA GPU 设备。这可能表示使用运行时库 (cudaMalloc) 或设备驱动程序 (cuMemAlloc) 分配的数据。

ARROW_DEVICE_CUDA_HOST#

通过使用cudaMallocHostcuMemAllocHost由 CUDA 固定和锁定页面的 CPU 内存。

ARROW_DEVICE_OPENCL#

通过使用OpenCL(开放计算语言)框架在设备上分配的数据。

ARROW_DEVICE_VULKAN#

Vulkan框架和库分配的数据。

ARROW_DEVICE_METAL#

使用Metal框架和库在 Apple GPU 设备上的数据。

ARROW_DEVICE_VPI#

表示使用 Verilog 仿真器缓冲区。

ARROW_DEVICE_ROCM#

使用ROCm堆栈的 AMD 设备。

ARROW_DEVICE_ROCM_HOST#

通过使用hipMallocHost由 ROCm 固定和锁定页面的 CPU 内存。

ARROW_DEVICE_EXT_DEV#

此值是用于扩展当前未以其他方式表示的设备的备用方案。如果使用此设备类型,则生产者需要提供与设备相关的其他信息/上下文。这用于快速测试扩展设备,并且语义可能因实现而异。

ARROW_DEVICE_CUDA_MANAGED#

cudaMallocManaged分配的 CUDA 托管/统一内存。

ARROW_DEVICE_ONEAPI#

在 Intel oneAPI非分区设备上分配的统一共享内存。需要调用oneAPI运行时才能确定特定的设备类型、USM 分配类型以及它绑定的 sycl 上下文。

ARROW_DEVICE_WEBGPU#

下一代 WebGPU 标准的 GPU 支持

ARROW_DEVICE_HEXAGON#

在 Qualcomm Hexagon DSP 设备上分配的数据。

ArrowDeviceArray 结构#

ArrowDeviceArray结构嵌入 C 数据ArrowArray结构并添加消费者使用数据所需的附加信息。它具有以下字段

struct ArrowArray ArrowDeviceArray.array#

必需。分配的数组数据。void**缓冲区中的值(以及任何子元素的缓冲区)是在设备上分配的内容。缓冲区值应为设备指针。其余结构应可供 CPU 访问。

此结构的private_datarelease回调应包含与根据其分配的设备释放数组相关的任何必要信息和结构,而不是在此处具有单独的释放回调和private_data指针。

int64_t ArrowDeviceArray.device_id#

必需。如果系统上存在此类型的多个设备,则用于标识特定设备的设备 ID。ID 的语义将取决于硬件,但我们使用int64_t来使 ID 随着时间的推移适应设备的变化。

对于没有设备标识符的内在概念的设备类型(例如ARROW_DEVICE_CPU),建议使用device_id-1 作为约定。

ArrowDeviceType ArrowDeviceArray.device_type#

必需。可以访问数组中缓冲区的设备的类型。

void *ArrowDeviceArray.sync_event#

可选。如果需要,用于同步的类似事件的对象。

许多设备(如 GPU)相对于 CPU 处理主要是异步的。因此,为了安全地访问设备内存,通常需要一个对象来同步处理。由于不同的设备将使用不同的类型来指定此对象,因此我们使用void*,它可以强制转换为设备适当类型的指针。

如果不需要同步,则可以为 null。如果此值不为 null,则在尝试访问缓冲区中的内存之前,必须使用它来调用设备的适当同步方法(例如cudaStreamWaitEventhipStreamWaitEvent)。

如果提供了事件,则生产者必须确保在触发事件之前,导出的数据在设备上可用。消费者在尝试访问导出的数据之前,应等待事件。

另请参阅

下面的同步事件类型部分。

int64_t ArrowDeviceArray.reserved[3]#

随着非 CPU 开发的扩展,可能需要扩展此结构。为了在不破坏 ABI 兼容性的情况下进行扩展,我们在对象末尾保留了 24 字节。为了确保将来 ABI 的安全演进,这些字节在初始化后**必须**被清零。

同步事件类型#

下表列出了每种设备类型预期的事件类型。如果设备不支持任何事件类型(“N/A”),则 sync_event 成员应始终为 null。

请记住,如果不需要同步即可访问数据,则事件**可以**为 null。

设备类型

实际事件类型

备注

ARROW_DEVICE_CPU

N/A

ARROW_DEVICE_CUDA

cudaEvent_t*

ARROW_DEVICE_CUDA_HOST

cudaEvent_t*

ARROW_DEVICE_OPENCL

cl_event*

ARROW_DEVICE_VULKAN

VkEvent*

ARROW_DEVICE_METAL

MTLEvent*

ARROW_DEVICE_VPI

N/A

ARROW_DEVICE_ROCM

hipEvent_t*

ARROW_DEVICE_ROCM_HOST

hipEvent_t*

ARROW_DEVICE_EXT_DEV

ARROW_DEVICE_CUDA_MANAGED

cudaEvent_t*

ARROW_DEVICE_ONEAPI

sycl::event*

ARROW_DEVICE_WEBGPU

N/A

ARROW_DEVICE_HEXAGON

N/A

备注

  • (1) 目前尚不清楚框架是否具有支持的事件类型。

  • (2) 扩展设备具有生产者定义的语义,因此如果需要扩展设备的同步,生产者应记录类型。

语义#

内存管理#

首先也是最重要的是:在此接口的所有内容中,**只有**数据缓冲区本身驻留在设备内存中(即 ArrowArray 结构的 buffers 成员)。其他所有内容都应位于 CPU 内存中。

ArrowDeviceArray 结构包含一个 ArrowArray 对象,该对象本身具有用于释放内存的特定语义。“基础结构”一词在下文中指的是在生产者和消费者之间直接传递的 ArrowDeviceArray 对象,而不是其任何子结构。

该基础结构旨在由**消费者**进行栈分配或堆分配。在这种情况下,生产者 API 应获取指向消费者分配的结构的指针。

但是,结构指向的任何数据**必须**由生产者分配和维护。这包括 sync_event 成员(如果它不为 null),以及 ArrowArray 对象中的任何指针(与往常一样)。数据生命周期通过 ArrowArray 成员的 release 回调进行管理。

对于 ArrowDeviceArray,已释放结构的语义和回调语义与ArrowArray 本身相同。除了任何已分配的事件之外,释放设备数据缓冲区所需的任何生产者特定上下文信息都应存储在 ArrowArrayprivate_data 成员中,并由 release 回调进行管理。

移动数组#

消费者可以通过按位复制或浅成员复制来**移动** ArrowDeviceArray 结构。然后,它**必须**通过将嵌入的 ArrowArray 结构的 release 成员设置为 NULL 来标记源结构已释放,但**不**调用该释放回调。这确保在任何给定时间只有一个活动的结构副本,并且生命周期正确地传达给生产者。

像往常一样,当不再需要目标结构时,将对其调用释放回调。

记录批次#

与 C 数据接口本身一样,记录批次可以被简单地视为等效的结构数组。在这种情况下,顶级 ArrowSchema 的元数据可用于记录批次的架构级元数据。

可变性#

生产者和消费者都**应该**将导出数据(即通过嵌入的 ArrowArraybuffers 成员在设备上可访问的数据)视为不可变的,因为否则任何一方都可能在另一方对其进行修改时看到不一致的数据。

同步#

如果 sync_event 成员不为 NULL,则消费者不应尝试访问或读取数据,直到他们在该事件上同步。如果 sync_event 成员为 NULL,则**必须**安全地访问数据,而无需消费者进行任何同步。

C 生产者示例#

导出简单的 int32 设备数组#

导出具有空元数据的不可为空的 int32 类型。这可以在C 数据接口文档中直接看到示例。

要导出数据本身,我们通过释放回调将所有权转移给消费者。此示例将使用 CUDA,但对于任何设备都可以使用等效的调用

static void release_int32_device_array(struct ArrowArray* array) {
    assert(array->n_buffers == 2);
    // destroy the event
    cudaEvent_t* ev_ptr = (cudaEvent_t*)(array->private_data);
    cudaError_t status = cudaEventDestroy(*ev_ptr);
    assert(status == cudaSuccess);
    free(ev_ptr);

    // free the buffers and the buffers array
    status = cudaFree(array->buffers[1]);
    assert(status == cudaSuccess);
    free(array->buffers);

    // mark released
    array->release = NULL;
}

void export_int32_device_array(void* cudaAllocedPtr,
                               cudaStream_t stream,
                               int64_t length,
                               struct ArrowDeviceArray* array) {
    // get device id
    int device;
    cudaError_t status;
    status = cudaGetDevice(&device);
    assert(status == cudaSuccess);

    cudaEvent_t* ev_ptr = (cudaEvent_t*)malloc(sizeof(cudaEvent_t));
    assert(ev_ptr != NULL);
    status = cudaEventCreate(ev_ptr);
    assert(status == cudaSuccess);

    // record event on the stream, assuming that the passed in
    // stream is where the work to produce the data will be processing.
    status = cudaEventRecord(*ev_ptr, stream);
    assert(status == cudaSuccess);

    memset(array, 0, sizeof(struct ArrowDeviceArray));
    // initialize fields
    *array = (struct ArrowDeviceArray) {
        .array = (struct ArrowArray) {
            .length = length,
            .null_count = 0,
            .offset = 0,
            .n_buffers = 2,
            .n_children = 0,
            .children = NULL,
            .dictionary = NULL,
            // bookkeeping
            .release = &release_int32_device_array,
            // store the event pointer as private data in the array
            // so that we can access it in the release callback.
            .private_data = (void*)(ev_ptr),
        },
        .device_id = (int64_t)(device),
        .device_type = ARROW_DEVICE_CUDA,
        // pass the event pointer to the consumer
        .sync_event = (void*)(ev_ptr),
    };

    // allocate list of buffers
    array->array.buffers = (const void**)malloc(sizeof(void*) * array->array.n_buffers);
    assert(array->array.buffers != NULL);
    array->array.buffers[0] = NULL;
    array->array.buffers[1] = cudaAllocedPtr;
}

// calling the release callback should be done using the array member
// of the device array.
static void release_device_array_helper(struct ArrowDeviceArray* arr) {
    arr->array.release(&arr->array);
}

设备流接口#

C 流接口类似,C 设备数据接口还指定了一个更高级别的结构,用于简化在单个进程内流式传输数据的通信。

语义#

Arrow C 设备流公开了数据块的流式源,每个数据块都具有相同的架构。通过调用阻塞式拉取式迭代函数获取块。预计所有块都应在相同的设备类型上提供数据(但不一定是相同的设备 ID)。如果需要在多个设备类型上提供数据流,则生产者应为每个设备类型提供单独的流对象。

结构定义#

C 设备流接口由单个 struct 定义定义。

#ifndef ARROW_C_DEVICE_STREAM_INTERFACE
#define ARROW_C_DEVICE_STREAM_INTERFACE

struct ArrowDeviceArrayStream {
    // device type that all arrays will be accessible from
    ArrowDeviceType device_type;
    // callbacks
    int (*get_schema)(struct ArrowDeviceArrayStream*, struct ArrowSchema*);
    int (*get_next)(struct ArrowDeviceArrayStream*, struct ArrowDeviceArray*);
    const char* (*get_last_error)(struct ArrowDeviceArrayStream*);

    // release callback
    void (*release)(struct ArrowDeviceArrayStream*);

    // opaque producer-specific data
    void* private_data;
};

#endif  // ARROW_C_DEVICE_STREAM_INTERFACE

注意

规范保护 ARROW_C_DEVICE_STREAM_INTERFACE 用于避免如果两个项目将 C 设备流接口定义复制到它们自己的头文件中,并且第三方项目从这两个项目中包含,则避免重复定义。因此,当复制这些定义时,务必保持此保护的原样。

ArrowDeviceArrayStream 结构#

ArrowDeviceArrayStream 提供了一个可以访问结果数据的设备类型以及与 Arrow 数组流式源交互所需的回调。它具有以下字段:

ArrowDeviceType device_type#

必填。此流在其中生成数据的设备类型。此流生成的全部 ArrowDeviceArray 都应具有与此处设置的相同的设备类型。这为消费者提供了一个便利,无需检查检索到的每个数组,而是允许用于流的高级编码结构。

int (*ArrowDeviceArrayStream.get_schema)(struct ArrowDeviceArrayStream*, struct ArrowSchema *out)#

必填。此回调允许消费者查询流中数据块的架构。所有数据块的架构都相同。

此回调不得在已释放的 ArrowDeviceArrayStream 上调用。

返回值:成功时返回 0,否则返回非零错误代码

int (*ArrowDeviceArrayStream.get_next)(struct ArrowDeviceArrayStream*, struct ArrowDeviceArray *out)#

必填。此回调允许消费者获取流中的下一个数据块。

此回调不得在已释放的 ArrowDeviceArrayStream 上调用。

下一个数据块**必须**可从与 ArrowDeviceArrayStream.device_type 匹配的设备类型访问。

返回值:成功时返回 0,否则返回非零错误代码

成功后,消费者必须检查 ArrowDeviceArray 的嵌入式 ArrowArray 是否被标记为已释放。如果嵌入式 ArrowDeviceArray.array 已释放,则表示已到达流的末尾。否则,ArrowDeviceArray 包含一个有效的数据块。

const char *(*ArrowDeviceArrayStream.get_last_error)(struct ArrowDeviceArrayStream*)#

必填。此回调允许消费者获取上次错误的文本描述。

此回调**仅**应在 ArrowDeviceArrayStream 上的上次操作返回错误时调用。不得在已释放的 ArrowDeviceArrayStream 上调用。

返回值:指向以 NULL 结尾的字符字符串(UTF8 编码)的指针。如果不可用详细描述,也可以返回 NULL。

返回的指针仅保证在流的下一个回调调用之前有效。如果字符字符串打算生存更长时间,则应将其复制到消费者管理的存储区中。

void (*ArrowDeviceArrayStream.release)(struct ArrowDeviceArrayStream*)#

必填。指向生产者提供的释放回调的指针。

void *ArrowDeviceArrayStream.private_data#

可选。指向生产者提供的私有数据的 opaque 指针。

消费者**不得**处理此成员。此成员的生命周期由生产者处理,特别是由释放回调处理。

结果生命周期#

get_schemaget_next 回调返回的数据必须独立释放。它们的生命周期与 ArrowDeviceArrayStream 的生命周期无关。

流生命周期#

C 流的生命周期使用释放回调进行管理,其用法与C 数据接口中的类似。

线程安全#

流源不假定是线程安全的。想要从多个线程调用 get_next 的消费者应确保这些调用被序列化。

与其他交换格式的互操作性#

其他交换 API(例如CUDA 数组接口)包含用于传递正在导出的数据缓冲区的形状和数据类型的成员。此信息对于解释正在共享的设备数据缓冲区中的原始字节是必需的。与其在 ArrowDeviceArray 旁边存储数据形状/类型,用户应利用现有的 ArrowSchema 结构来传递任何数据类型和形状信息。

更新此规范#

注意

由于此规范仍被视为实验性,因此存在(仍然很低)可能发生细微变化的可能性。将其标记为“实验性”的原因是我们不知道我们不知道什么。已经进行了工作和研究以确保与许多不同框架兼容的通用 ABI,但始终有可能遗漏了一些东西。一旦在正式的 Arrow 版本中支持此功能,并且观察到用法以确认不需要任何修改,则将删除“实验性”标签并冻结 ABI。

一旦此规范在正式的 Arrow 版本中得到支持,C ABI 将被冻结。这意味着ArrowDeviceArray结构定义不得以任何方式更改 - 包括添加新成员。

允许向后兼容的更改,例如ArrowDeviceType的新宏值,或将保留的 24 字节转换为不同的类型/成员,而无需更改结构的大小。

任何不兼容的更改都应作为新规范的一部分,例如ArrowDeviceArrayV2