CUDA 数组接口 (版本 3)

The CUDA Array Interface (或 CAI) 是为了在不同项目的 CUDA 数组类对象之间实现互操作性而创建的。这个想法借鉴自 NumPy 数组接口。

备注

目前，我们仅定义了 Python 端的接口。未来，我们可能会添加一个 C 端的接口，以便在编译代码中高效地交换信息。

Python 接口规范

备注

实验性功能。规范可能会有所变化。

__cuda_array_interface__ 属性返回一个字典 (dict)，该字典必须包含以下条目：

• 形状: (整数, ...)

  一个由 int (或 long) 组成的元组，表示每个维度的大小。

• typestr: str

  类型字符串。这与 NumPy 数组接口 中的 typestr 定义相同。

• 数据: (整数, 布尔值)

  数据 是一个 2 元组。第一个元素是作为 Python int``（或 ``long）的数据指针。数据必须是设备可访问的。对于零大小的数组，在此处使用 0。第二个元素是作为 Python bool 的只读标志。

  由于接口的用户可能处于相同或不同的上下文中，最常见的情况是使用 CUDA 驱动 API 中的 cuPointerGetAttribute 和 ``CU_POINTER_ATTRIBUTE_DEVICE_POINTER``（或等效的 CUDA 运行时 API）来检索当前活动上下文中可用的设备指针。

• 版本：整数

  导出接口的版本号。当前版本是 3。

以下是可选条目：

• strides: None 或 (整数, ...)

  如果 strides 未给出，或为 None ，则数组采用 C 连续布局。否则，将显式给出一个 int （或 long） 的元组，用于表示在每个维度上访问下一个元素时需要跳过的字节数。

• 描述

  这是用于描述更复杂类型的。它遵循与 NumPy 数组接口 相同的规范。

• 掩码：None 或暴露 __cuda_array_interface__ 的对象

  如果 None ，则 data 中的所有值都是有效的。掩码数组的所有元素应仅解释为真或非真，以指示此数组中的哪些元素是有效的。这与 NumPy 数组接口 中的 mask 定义相同。

  备注

  Numba 目前不支持处理带有掩码的 CUDA 数组，如果将此类数组传递给 GPU 函数，将会引发 NotImplementedError 异常。

• 流: None 或 整数

  一个可选的流，在消费时必须进行同步，可以通过同步流或在给定流上对数据操作进行排队。此条目中的整数值如下：

  • 0: 这是不允许的，因为它会在 None 和默认流之间，以及在传统和每个线程的默认流之间产生歧义。任何可能使用 0 的情况都应该改为使用 None、1 或 2 以确保清晰。

  • 1: 传统默认流。

  • 2: 每个线程的默认流。

  • 任何其他整数：一个表示为Python整数的 cudaStream_t。

  当 None 时，不需要同步。有关更多详细信息，请参见下面的 同步 部分。

  在接口的未来修订版中，此条目可能会扩展（或添加另一个条目），以便可以指定一个事件来进行同步，而不是一个流。

同步

定义

在讨论同步时，使用以下定义：

• 生产者: 访问 __cuda_array_interface__ 的库 / 对象。

• 消费者: 访问生产者 __cuda_array_interface__ 的库 / 函数。

• 用户代码: 通过CAI使生产者和消费者共享数据的代码。

• 用户：编写或维护用户代码的人。用户可以在不了解CAI的情况下实现用户代码，因为CAI的访问可以对他们的视图隐藏。

在以下示例中：

import cupy
from numba import cuda

@cuda.jit
def add(x, y, out):
    start = cuda.grid(1)
    stride = cuda.gridsize(1)
    for i in range(start, x.shape[0], stride):
        out[i] = x[i] + y[i]

a = cupy.arange(10)
b = a * 2
out = cupy.zeros_like(a)

add[1, 32](a, b, out)

当 add 内核被启动时：

• a, b, out 是生产者。

• add 内核是消费者。

• 用户代码是 add[1, 32](a, b, out)。

• 代码的作者是用户。

设计动机

CAI 设计中与同步相关的元素旨在满足以下要求：

1. 通过CAI交换数据的生成者和消费者必须能够做到这一点而不会发生数据竞争。

2. 要求1应在不要求用户了解CAI的任何细节的情况下得到满足 - 换句话说，生产者和消费者之间异步操作数据的交换应默认正确。

   • 此要求的一个例外是针对明确记录用户需要采取额外步骤以确保同步正确性的生产者和消费者。在这种情况下，用户需要理解CUDA数组接口的细节，并且生产者/消费者库文档必须指定用户需要采取的步骤。

     应尽可能避免使用此异常，因为它是为了那些无法在没有用户参与的情况下实现同步语义的库而提供的——例如，那些与不了解CUDA数组接口的第三方库进行接口的库。

3. 当用户了解CAI的细节以及生产者和消费者的实现细节时，他们应该能够在其判断下，覆盖接口的一些同步语义以减少同步开销。覆盖同步语义意味着：

   • CAI 设计，以及 Producer 和 Consumer 的设计与实现，并未指定或保证关于数据竞争的正确性。

   • 相反，用户有责任确保数据竞争的正确性。

接口要求

stream 条目使生产者和消费者在交换数据时能够避免危险。消费者的预期行为如下：

• 当 stream 不存在或为 None 时：

  • 消费者端不需要进行同步。

  • 消费者可以在任何流上立即对底层数据进行操作排队。

• 当 stream 是一个整数时，其值表示生产者可能在数据上进行未完成操作的流，而消费者预期会执行以下操作之一：

  • 在访问数据之前进行同步，或

  • 在访问数据时进行入队操作。

  消费者可以选择使用哪种机制，并考虑以下因素：

  • 如果消费者在访问数据之前在提供的流上进行同步，那么它必须确保在其选择的流中的操作完成之前，提供的流中不会发生任何计算。这可以通过以下任一方式实现：

    • 在提供的流中放置一个等待事件，该事件在消费者对数据的所有操作完成后发生，或者

    • 在其自身流上的操作完成之前，避免将控制权返回给用户代码。

  • 如果消费者选择仅在提供的流中对数据进行操作排队，那么在将其工作排队后，它可能会立即将控制权返回给用户代码，因为所有工作都将在导出数组的流上序列化。即使用户代码随后诱导生产者开始在同一流上排队更多工作，这也足以确保正确性。

• 如果用户已将消费者设置为忽略CAI同步语义，消费者可能会假设它可以在没有任何进一步同步的情况下立即操作任何流中的数据，即使 stream 成员具有整数值。

通过CAI导出数组时，生产者必须确保：

• 如果在一个或多个流中排队了数据处理工作，那么对提供的 stream 进行同步就足以确保与所有待处理工作的同步。

  • 如果生产者没有排队的工作，或者只有排队在由 stream 标识的流上的工作，那么这个条件就满足了。

  • 如果生产者在多个流上对数据进行了排队处理，那么它必须在那些流上对排队的工作之后的活动进行排队，然后在提供的 stream 上等待这些活动。例如：

    1. 工作由生产者在流 7、9 和 15 上排队。

    2. 事件随后被排入每个流 7、9 和 15 中。

    3. 然后，生产者告诉流 3 等待步骤 2 的事件，并将 stream 条目设置为 3。

• 如果在数据上没有排队的工作，那么 stream 条目可能是 None，或者没有提供。

可选地，为了便于用户放松对同步语义的遵守：

• 生产者可以提供一个配置选项，以始终将 stream 设置为 None。

• 消费者可以提供一个配置选项来忽略 stream 的值，并表现得好像它是 None 或未提供。这省略了在生产者提供的流上的同步，并允许在生产者提供的流以外的流上排队工作。

这些选项不应在生产者或消费者中默认设置。CAI 规范没有规定设置这些选项的确切机制，也没有规定生产者或消费者可能提供的相关选项，以允许用户进一步控制同步行为。

Numba 中的同步

Numba 既不是严格的生产者也不是消费者 - 它可以被用户用来实现其中任何一个。为了便于正确实现同步语义，Numba 在接口同步方面表现出以下行为：

• 当 Numba 作为消费者时（例如当一个类似数组的对象传递给内核启动时）：如果 stream 是一个整数，那么 Numba 将立即在提供的 stream 上同步。从类似数组的对象创建的 Numba 设备数组 将其*默认流*设置为提供的流。

• 当 Numba 作为生产者时（当访问 Numba CUDA 数组的 __cuda_array_interface__ 属性时）：如果导出的 CUDA 数组有一个 默认流，那么它将作为 stream 条目给出。否则，stream 设置为 None。

备注

在 Numba 的术语中，数组的 默认流 是一个属性，指定如果没有提供其他流作为函数调用的参数，Numba 将在其中排队异步传输的流。这与正常 CUDA 术语中的 默认流 不同。

Numba 的同步行为会产生以下预期结果：

• 无论是作为生产者还是消费者交换数据，只要交互的另一方也遵循CAI同步语义，用户无需采取任何进一步行动即可正确进行。

• 期望用户执行以下任一操作：

  • 避免在不是其参数默认流的流上启动内核或其他操作，或者

  • 当在不是给定参数的默认流的流上启动操作时，他们应该在正在操作的流中插入一个事件，并在参数的默认流中等待该事件。有关此示例，请参见 下文。

用户可以通过将环境变量 NUMBA_CUDA_ARRAY_INTERFACE_SYNC 或配置变量 CUDA_ARRAY_INTERFACE_SYNC 设置为 0 来覆盖 Numba 的同步行为（参见 GPU 支持环境变量）。设置后，Numba 不会在导入数组的流上进行同步，用户需确保流同步的正确性。在从导出 CUDA 数组接口的对象创建 Numba CUDA 数组时，也可以通过在调用 numba.cuda.as_cuda_array() 或 numba.cuda.from_cuda_array_interface() 创建 Numba CUDA 数组时传递 sync=False 来省略同步。

未来，Numba 的同步实现有优化的空间，可以通过在同一流上启动内核或驱动 API 操作（例如内存复制或内存设置）时消除同步来实现。

在数组的非默认流上启动的示例

这个示例展示了如何确保当一个消费者在不同的流中启动内核时，能够安全地消费一个带有默认流的数组。

首先，我们需要导入 Numba 和一个消费库（本例中是一个名为 other_cai_library 的虚构库）：

from numba import cuda, int32, void
import other_cai_library

现在我们将定义一个核 - 这将初始化数组的元素，将每个条目设置为其索引：

@cuda.jit(void, int32[::1])
def initialize_array(x):
    i = cuda.grid(1)
    if i < len(x):
        x[i] = i

接下来我们将创建两个流：

array_stream = cuda.stream()
kernel_stream = cuda.stream()

然后创建一个数组，其中一个流作为其默认流：

N = 16384
x = cuda.device_array(N, stream=array_stream)

现在我们在另一个流中启动内核：

nthreads = 256
nblocks = N // nthreads

initialize_array[nthreads, nblocks, kernel_stream](x)

如果我们现在将 x 传递给一个 Consumer，存在一个风险，即它可能在 array_stream 中对其进行操作，而内核仍在 kernel_stream 中运行。为了防止 array_stream 中的操作在内核启动完成之前开始，我们创建一个事件并等待它：

# Create event
evt = cuda.event()
# Record the event after the kernel launch in kernel_stream
evt.record(kernel_stream)
# Wait for the event in array_stream
evt.wait(array_stream)

现在 other_cai_library 可以安全地使用 x 了：

other_cai_library.consume(x)

生命周期管理

数据

获取任何对象的 __cuda_array_interface__ 属性的值不会影响从中创建该对象的生命周期。特别要注意的是，该接口没有数据所有者的插槽。

用户代码必须确保拥有数据的对象在其可能被使用者使用期间保持生命周期。

流

与数据类似，CUDA 流也有有限的寿命。因此，要求在接口上导出数据的生产者确保导出的流的寿命等于或超过从中导出接口的对象的寿命。

Numba中的生命周期管理

生成数组

Numba 不会采取任何措施来维持从其导出接口的对象的生命周期 - 确保底层对象在使用导出接口期间保持活动状态是用户的责任。

在接口上导出的任何由Numba管理的流的生存期保证等于或超过底层对象的生存期，因为底层对象持有对该流的引用。

备注

Numba 管理的流是通过 cuda.default_stream()、cuda.legacy_default_stream() 或 cuda.per_thread_default_stream() 创建的。不由 Numba 管理的流是通过 cuda.external_stream() 从外部流创建的。

使用数组

Numba 提供了两种机制，用于从导出 CUDA Array Interface 的对象创建设备数组。使用哪种机制取决于创建的设备数组是否应保持其创建对象的生命周期：

• as_cuda_array: 这会创建一个持有对所属对象引用的设备数组。只要设备数组有引用，其底层数据也将保持活动状态，即使所有其他对原始所属对象的引用已被丢弃。

• from_cuda_array_interface: 默认情况下，这将创建一个不引用所属对象的设备数组。所属对象，或被视为所有者的其他对象可以通过 owner 参数传递。

这些函数的接口是：

cuda.as_cuda_array(sync=True)

从任何实现 CUDA 数组接口 的对象创建一个 DeviceNDArray。

创建了底层 GPU 缓冲区的视图。不会进行数据的复制。生成的 DeviceNDArray 将从 obj 获取引用。

如果 sync 是 True，那么导入的流（如果存在）将被同步。

cuda.from_cuda_array_interface(owner=None, sync=True)

从 cuda-array-interface 描述创建一个 DeviceNDArray。owner 是底层内存的所有者。生成的 DeviceNDArray 将从它那里获取一个引用。

如果 sync 是 True，那么导入的流（如果存在）将被同步。

指针属性

可以使用 cuPointerGetAttribute 或 cudaPointerGetAttributes 获取有关数据指针的更多信息。此类信息包括：

• 拥有该指针的 CUDA 上下文；

• 指针主机可访问吗？

• 指针是托管内存吗？

与CUDA数组接口（版本0）的区别

CUDA Array Interface 的 0 版本没有可选的 mask 属性来支持掩码数组。

与CUDA数组接口（版本1）的区别

CUDA Array Interface 的 0 和 1 版本既未明确 C 连续数组的 strides 属性，也未指定对零大小数组的处理方式。

与CUDA数组接口（版本2）的差异

CUDA Array Interface 的早期版本没有关于同步的声明。

互操作性

以下Python库已采用CUDA数组接口：

• Numba

• CuPy 是一个与NumPy API兼容的GPU加速库。

• PyTorch 是一个开源的深度学习框架。

• PyArrow 是一个用于处理CUDA上下文的Python库。

• mpi4py 是一个支持 CUDA 感知 MPI 的 Python 库。

• ArrayViews

• JAX 是一个强大的数值计算库。

• PyCUDA 是一个与使用 CUDA 数组接口的其他库进行互操作的教程。

• DALI: NVIDIA 数据加载库 :

  • TensorGPU 对象 暴露了 CUDA 数组接口。

  • 外部源操作符 消耗导出 CUDA 数组接口的对象。

• RAPIDS 堆栈：

  • cuDF 是一个用于处理表格数据的库。

  • cuML

  • cuSignal 是一个用于信号处理的库。

  • RMM

如果你的项目不在此列表中，请随时在 Numba 问题跟踪器 上报告。