CUDA 集成#

Arrow 不限于 CPU 缓冲区（位于计算机的主内存中，也称为“主机内存”）。它还提供访问位于 CUDA 支持的 GPU 设备（在“设备内存”中）的缓冲区的规定。

注意

此功能是可选的，必须在构建时启用。如果您的包管理器没有这样做，您可能需要自己构建 Arrow。

CUDA 上下文#

CUDA 上下文表示对特定 CUDA 支持的设备的访问。例如，这是创建访问 CUDA 设备编号 0 的 CUDA 上下文

>>> from pyarrow import cuda
>>> ctx = cuda.Context(0)
>>>

CUDA 缓冲区#

CUDA 缓冲区可以通过使用 Context.buffer_from_data() 方法将数据从主机内存复制到 CUDA 设备的内存来创建。源数据可以是任何 Python 缓冲区类对象，包括 Arrow 缓冲区

>>> import numpy as np
>>> arr = np.arange(4, dtype=np.int32)
>>> arr.nbytes
16
>>> cuda_buf = ctx.buffer_from_data(arr)
>>> type(cuda_buf)
pyarrow._cuda.CudaBuffer
>>> cuda_buf.size     # The buffer's size in bytes
16
>>> cuda_buf.address  # The buffer's address in device memory
30088364544
>>> cuda_buf.context.device_number
0

相反，您可以将 CUDA 缓冲区复制回设备内存，获取一个常规的 CPU 缓冲区

>>> buf = cuda_buf.copy_to_host()
>>> type(buf)
pyarrow.lib.Buffer
>>> np.frombuffer(buf, dtype=np.int32)
array([0, 1, 2, 3], dtype=int32)

警告

许多 Arrow 函数期望一个 CPU 缓冲区，但不会检查缓冲区的实际类型。如果您将 CUDA 缓冲区传递给此类函数，您将遇到崩溃

>>> pa.py_buffer(b"x" * 16).equals(cuda_buf)
Segmentation fault

Numba 集成#

虽然你无法直接在 Python 中对 Arrow CUDA 缓冲区进行太多操作，但它们支持与 Numba 交互，Numba 是一种 JIT 编译器，可以将 Python 代码转换为优化的 CUDA 内核。

Arrow 到 Numba#

首先，让我们定义一个在 int32 数组上运行的 Numba CUDA 内核。在这里，我们将简单地递增每个数组元素（假设数组是可写的）。

import numba.cuda

@numba.cuda.jit
def increment_by_one(an_array):
    pos = numba.cuda.grid(1)
    if pos < an_array.size:
        an_array[pos] += 1

然后，我们需要将我们的 CUDA 缓冲区包装到一个具有正确数组元数据（形状、步长和数据类型）的 Numba “设备数组” 中。这是必要的，以便 Numba 可以识别数组的特征并使用适当的类型声明编译内核。

在这种情况下，元数据可以简单地从原始 NumPy 数组中获取。请注意，GPU 数据没有被复制，只是被指向。

>>> from numba.cuda.cudadrv.devicearray import DeviceNDArray
>>> device_arr = DeviceNDArray(arr.shape, arr.strides, arr.dtype, gpu_data=cuda_buf.to_numba())

（理想情况下，我们可以在 CPU 内存中定义一个 Arrow 数组，将其复制到 CUDA 内存中而不会丢失类型信息，然后在它上面调用 Numba 内核，而无需手动构建 DeviceNDArray；这目前还不可行）

最后，我们可以使用 Numba 设备数组运行 Numba CUDA 内核（这里使用 16x16 网格大小）。

>>> increment_by_one[16, 16](device_arr)

然后可以通过将 CUDA 缓冲区复制回 CPU 内存来检查结果。

>>> np.frombuffer(cuda_buf.copy_to_host(), dtype=np.int32)
array([1, 2, 3, 4], dtype=int32)

Numba 到 Arrow#

相反，可以使用 CudaBuffer.from_numba() 工厂方法将 Numba 创建的设备数组视为 Arrow CUDA 缓冲区。

为了举例说明，让我们首先创建一个 Numba 设备数组。

>>> arr = np.arange(10, 14, dtype=np.int32)
>>> arr
array([10, 11, 12, 13], dtype=int32)
>>> device_arr = numba.cuda.to_device(arr)

然后，我们可以创建一个指向设备数组内存的 CUDA 缓冲区。这次我们不需要显式传递 CUDA 上下文：适当的 CUDA 上下文将自动从 Numba 对象中检索和适配。

>>> cuda_buf = cuda.CudaBuffer.from_numba(device_arr.gpu_data)
>>> cuda_buf.size
16
>>> cuda_buf.address
30088364032
>>> cuda_buf.context.device_number
0

当然，我们可以将 CUDA 缓冲区复制回主机内存。

>>> np.frombuffer(cuda_buf.copy_to_host(), dtype=np.int32)
array([10, 11, 12, 13], dtype=int32)

另请参阅

Numba 的 CUDA 支持文档。