概述

Numba 通过直接将 Python 代码的受限子集编译成遵循 CUDA 执行模型的 CUDA 内核和设备函数，支持 CUDA GPU 编程。在 Numba 中编写的内核似乎可以直接访问 NumPy 数组。NumPy 数组在 CPU 和 GPU 之间自动传输。

术语

CUDA 编程主题中的几个重要术语列于此处：

  • host: 中央处理器

  • device: GPU

  • host memory: 系统主内存

  • 设备内存：GPU 卡上的板载内存

  • kernels：由主机启动并在设备上执行的GPU函数

  • device function: 在设备上执行的GPU函数，只能从设备上调用（即从内核或其他设备函数调用）

编程模型

Numba 暴露的大多数 CUDA 编程功能直接映射到 NVidia 提供的 CUDA C 语言。因此，建议您阅读官方的 CUDA C 编程指南。

要求

支持的GPU

Numba 支持计算能力为 3.5 或更高的 CUDA 启用 GPU。对计算能力低于 5.0 的设备的支持已弃用，并将在未来的 Numba 版本中移除。

计算能力为5.0或更高的设备包括（但不限于）：

  • 嵌入式平台：NVIDIA Jetson Nano、Jetson Orin Nano、TX1、TX2、Xavier NX、AGX Xavier、AGX Orin。

  • 桌面 / 服务器 GPU: 所有采用 Maxwell 微架构或更高版本的 GPU。例如：GTX 9 / 10 / 16 系列，RTX 20 / 30 / 40 系列，Quadro / Tesla M / P / V / RTX 系列，RTX A 系列，RTX Ada / SFF，A / L 系列，H100。

  • 笔记本电脑GPU：所有采用Maxwell微架构或更高版本的GPU。例如：MX系列、Quadro M / P / T系列（移动版）、RTX 20 / 30系列（移动版）、RTX A系列（移动版）。

软件

Numba 旨在支持过去3年内发布的CUDA Toolkit版本。目前，最低要求的工具包版本是11.2。还需要一个足以支持该工具包版本的NVIDIA驱动程序（另见 CUDA 次要版本兼容性）。

Conda 用户可以将 CUDA Toolkit 安装到 conda 环境中。

对于 CUDA 12，需要 cuda-nvcc 和 cuda-nvrtc:

$ conda install -c conda-forge cuda-nvcc cuda-nvrtc "cuda-version>=12.0"

对于 CUDA 11，需要 cudatoolkit:

$ conda install -c conda-forge cudatoolkit "cuda-version>=11.2,<12.0"

如果你没有使用 Conda，或者你想使用不同版本的 CUDA 工具包，以下描述了 Numba 如何搜索 CUDA 工具包的安装。

CUDA 绑定

Numba 支持通过 NVIDIA CUDA Python 绑定 及其基于 ctypes 的绑定与 CUDA 驱动 API 进行交互。这两种绑定的功能是等效的。基于 ctypes 的绑定目前是默认的，但在未来的 Numba 版本中，如果环境中有 NVIDIA 绑定，将默认使用 NVIDIA 绑定。

你可以通过以下方式安装 NVIDIA 绑定:

$ conda install -c conda-forge cuda-python

如果你正在使用 Conda，或者:

$ pip install cuda-python

如果你正在使用 pip。

通过将环境变量 NUMBA_CUDA_USE_NVIDIA_BINDING 设置为 "1" 来启用 NVIDIA 绑定的使用。

设置 CUDA 安装路径

Numba 按以下顺序搜索 CUDA 工具包安装：

  1. Conda 安装的 CUDA Toolkit 包

  2. 环境变量 CUDA_HOME 指向已安装的 CUDA 工具包的目录（例如 /home/user/cuda-12）

  3. 在Linux平台上，系统范围的安装路径为 /usr/local/cuda 。版本化的安装路径（例如 /usr/local/cuda-12.0 ）被有意忽略。用户可以使用 CUDA_HOME 来选择特定版本。

除了可以通过conda安装到环境中的CUDA工具包库，或者通过 CUDA SDK安装程序 系统范围安装的CUDA工具包库之外，Numba中的CUDA目标还需要一个最新的NVIDIA图形驱动程序。更新的图形驱动程序也会由CUDA SDK安装程序安装，因此不需要两者都做。如果 libcuda 库位于非标准位置，用户可以将环境变量 NUMBA_CUDA_DRIVER 设置为共享库文件的路径（不是目录路径）。

缺失的 CUDA 功能

Numba 尚未实现 CUDA 的所有功能。以下列出了一些缺失的功能：

  • 动态并行性

  • 纹理内存