Numba CUDA Target

0.30.0 · active · verified Tue Apr 14

Numba-cuda provides a CUDA target for the Numba Python JIT compiler, enabling Python functions to be compiled and executed on NVIDIA GPUs. It allows users to write custom GPU kernels and device functions directly in a subset of Python. The library, currently at version 0.30.0, is actively developed by NVIDIA, with its release cycle now decoupled from the main Numba project to facilitate more frequent updates and new feature development.

Warnings

deprecated The built-in CUDA target in the main `numba` package is deprecated. New features and most bug fixes are now exclusively implemented in `numba-cuda`. While the old target remains for compatibility, it's strongly recommended to install `numba-cuda` for active development and to ensure access to the latest capabilities.
Fix: Always include `pip install numba-cuda` (or `conda install numba-cuda`) in your environment setup alongside `numba`.
breaking In `numba-cuda` v0.28.0, there was an attempt to shift error classes from `numba.core.errors.TypingError` to `numba.cuda.errors` namespaces. This caused compatibility issues with existing code that relied on catching the old error types and was subsequently reverted. Users should be aware that such internal error type changes can be breaking.
Fix: Ensure your `try-except` blocks are robust to potential changes in error types. For maximum compatibility, catch broader exception types or consult release notes if you encounter unexpected `TypingError` propagation.
breaking The internal `DeviceArray` implementation underwent refactoring, and certain internal `enums` and `ctypes` code were removed in `numba-cuda` v0.23.0 and v0.28.0 respectively. Code that directly interacted with these internal components or undocumented APIs may break.
Fix: Avoid relying on Numba's internal implementation details. Stick to the public API documented in `numba.cuda` for memory management (`cuda.to_device`, `cuda.device_array`), kernel launching, and device interactions.
gotcha Numba CUDA kernel functions cannot return values. Any results computed within a kernel must be written to arrays passed as arguments to the kernel. This is a common pattern in CUDA C/C++ and applies to Numba CUDA kernels as well.
Fix: Pass empty or pre-allocated device arrays as arguments to your kernel, and have the kernel write its output into these arrays. Copy the results back to the host after kernel execution if needed.
gotcha The first call to a Numba CUDA kernel includes the Just-In-Time (JIT) compilation overhead, which can be significant. For accurate performance benchmarking, always time subsequent calls to the kernel after the initial compilation has completed (e.g., by performing a 'warm-up' run).
Fix: Run your kernel once with dummy data to trigger compilation, then measure the execution time of subsequent calls. Use `cuda.synchronize()` to ensure all GPU operations have completed before measuring elapsed time.
breaking Support for NVIDIA GPUs with compute capability less than 5.0 is deprecated and will be removed in future releases. Additionally, Numba-CUDA requires a minimum CUDA Toolkit version of 11.2.
Fix: Ensure your target GPU has compute capability 5.0 or higher. Upgrade your NVIDIA drivers and CUDA Toolkit to version 11.2 or newer.

Install

pip install numba-cuda PyPI
conda install -c conda-forge numba-cuda Conda (conda-forge)

Imports

cuda
```
from numba import cuda
```
All CUDA-specific functionality is exposed through the `numba.cuda` module.

Quickstart

This quickstart demonstrates a basic vector addition using a Numba CUDA kernel. It covers defining a kernel with `@cuda.jit`, allocating and transferring data between host (CPU) and device (GPU) memory, configuring and launching the kernel, and copying results back to the host. Ensure you have a CUDA-enabled GPU and appropriate drivers installed.

import numpy as np
from numba import cuda
import os

# Check for CUDA availability (runtime dependency)
if not cuda.is_available():
    print("CUDA is not available. Please ensure you have an NVIDIA GPU and CUDA drivers installed.")
    exit()

# Define a CUDA kernel
@cuda.jit
def add_vectors(x, y, out):
    idx = cuda.grid(1)
    if idx < len(out):
        out[idx] = x[idx] + y[idx]

# Host-side code
N = 1000000
x_host = np.arange(N, dtype=np.float32)
y_host = np.arange(N, dtype=np.float32)
out_host = np.empty_like(x_host)

# Allocate memory on the device and copy data
x_device = cuda.to_device(x_host)
y_device = cuda.to_device(y_host)
out_device = cuda.device_array_like(out_host)

# Configure the kernel launch
threadsperblock = 256
blockspergrid = (N + (threadsperblock - 1)) // threadsperblock

# Launch the kernel
add_vectors[blockspergrid, threadsperblock](x_device, y_device, out_device)

# Copy the result back to the host
out_device.copy_to_host(out_host)

# Verify the result
expected_out = x_host + y_host
assert np.allclose(out_host, expected_out)
print("Vector addition on GPU successful!")

view raw JSON →