CUDA 整合#

Arrow 不僅限於 CPU 緩衝區（位於電腦主記憶體中，也稱為「主機記憶體」）。它也提供存取位於支援 CUDA 的 GPU 裝置上的緩衝區（在「裝置記憶體」中）的功能。

注意

此功能為選用性質，必須在建置時啟用。如果您的套件管理器未執行此操作，您可能需要自行建置 Arrow。

CUDA 環境#

CUDA 環境代表對特定支援 CUDA 裝置的存取。例如，以下程式碼建立一個存取 CUDA 裝置編號 0 的 CUDA 環境

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

CUDA 緩衝區#

CUDA 緩衝區可以透過將資料從主機記憶體複製到 CUDA 裝置的記憶體來建立，使用 Context.buffer_from_data() 方法。來源資料可以是任何類似 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 記憶體而不會遺失類型資訊，然後在不手動建構 DeviceNDArray 的情況下，對其調用 Numba 核心；這目前尚不可能）

最後，我們可以在 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 支援的文件。