CUDA 数组接口（版本 3）

CUDA 内置目标弃用通知

Numba 中内置的 CUDA 目标已被弃用，后续开发已移至 NVIDIA numba-cuda 包。请参阅 内置 CUDA 目标弃用和维护状态。

CUDA 数组接口（或 CAI）旨在实现不同项目中 CUDA 类数组对象不同实现之间的互操作性。该思想借鉴了 NumPy 数组接口。

注意

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

Python 接口规范

注意

实验性功能。规范可能更改。

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

• shape: (整型, ...)

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

• typestr: str

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

• data: (整型, 布尔型)

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

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

• version: 整型

  一个表示所导出接口版本的整数。当前版本是 3。

以下是可选条目

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

  如果未给出 strides，或者它是 None，则数组采用 C 连续布局。否则，明确给出由 int（或 long）组成的元组，表示在每个维度中访问下一个元素需要跳过的字节数。

• descr

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

• mask: None 或暴露 __cuda_array_interface__ 的对象

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

  注意

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

• stream: 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 设计以及生产者和消费者的设计和实现不指定或保证数据竞争方面的正确性。

   • 相反，用户负责确保数据竞争方面的正确性。

接口要求

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 Device Array 会将其默认流设置为提供的流。

• 当 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 数组时使用 numba.cuda.as_cuda_array() 或 numba.cuda.from_cuda_array_interface() 传递 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 传递给消费者，则存在它可能在 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 数组接口的对象创建设备数组。选择哪种机制取决于创建的设备数组是否应维护其创建来源对象的生命周期

• 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 数组接口描述创建一个 DeviceNDArray。owner 是底层内存的所有者。生成的 DeviceNDArray 将从它获取引用。

如果 sync 为 True，则导入的流（如果存在）将进行同步。

指针属性

可以使用 cuPointerGetAttribute 或 cudaPointerGetAttributes 检索有关数据指针的附加信息。这些信息包括

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

• 该指针是否可由主机访问？

• 该指针是否为托管内存？

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

CUDA 数组接口的 0 版没有可选的 mask 属性来支持带掩码的数组。

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

CUDA 数组接口的 0 版和 1 版既没有澄清 C 连续数组的 strides 属性，也没有指定零大小数组的处理方式。

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

CUDA 数组接口的先前版本没有对同步做出任何声明。

互操作性

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

• Numba

• CuPy

• PyTorch

• PyArrow

• mpi4py

• ArrayViews

• JAX

• PyCUDA

• DALI: NVIDIA 数据加载库 :

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

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

• RAPIDS 技术栈

  • cuDF

  • cuML

  • cuSignal

  • RMM

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