numba-vulkan#
A Vulkan compute target for Numba: write a kernel in Python and run it on any GPU with a Vulkan driver.
import numpy as np
import numba_vulkan as nv
@nv.jit
def saxpy(a, x, y):
i = nv.global_id(0)
if i < x.shape[0]:
y[i] = a * x[i] + y[i]
x = np.arange(1000, dtype=np.float32)
y = np.ones(1000, dtype=np.float32)
saxpy.forall(x.size)(np.float32(2.0), x, y)
Numba’s only maintained GPU target is CUDA, which ties GPU-accelerated Numba code to NVIDIA hardware. Vulkan is implemented by nearly every vendor, so one target could cover all of them. The idea was raised in numba/numba#10116; this project explores how far it can be taken.
Note
The code and this documentation were written by an AI (Claude) under the direction of, and reviewed by, a human author. See Authorship and AI disclosure.
Warning
This is alpha software. Its test suite passes on NVIDIA and Intel GPUs and llvmpipe on Linux, on an Apple M3 Pro through MoltenVK, and on Windows with lavapipe, and the benchmarks run on an Intel Arc GPU under Windows; but its interfaces may still change and it has not yet been proven on real workloads. See Limitations and Known issues.
Contents