numba-vulkan 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- numba_vulkan-0.1.0/LICENSE +24 -0
- numba_vulkan-0.1.0/PKG-INFO +553 -0
- numba_vulkan-0.1.0/README.md +519 -0
- numba_vulkan-0.1.0/pyproject.toml +135 -0
- numba_vulkan-0.1.0/pyproject.toml.orig +120 -0
- numba_vulkan-0.1.0/src/numba_vulkan/__init__.py +79 -0
- numba_vulkan-0.1.0/src/numba_vulkan/_emit.py +98 -0
- numba_vulkan-0.1.0/src/numba_vulkan/arrayfuncs.py +467 -0
- numba_vulkan-0.1.0/src/numba_vulkan/autotune.py +173 -0
- numba_vulkan-0.1.0/src/numba_vulkan/buffers.py +819 -0
- numba_vulkan-0.1.0/src/numba_vulkan/codegen.py +2224 -0
- numba_vulkan-0.1.0/src/numba_vulkan/compiler.py +635 -0
- numba_vulkan-0.1.0/src/numba_vulkan/data/LICENSE-libclc.txt +299 -0
- numba_vulkan-0.1.0/src/numba_vulkan/data/clspv--.bc +0 -0
- numba_vulkan-0.1.0/src/numba_vulkan/data/libclc-version.txt +1 -0
- numba_vulkan-0.1.0/src/numba_vulkan/dispatcher.py +1178 -0
- numba_vulkan-0.1.0/src/numba_vulkan/errors.py +56 -0
- numba_vulkan-0.1.0/src/numba_vulkan/kernelcache.py +315 -0
- numba_vulkan-0.1.0/src/numba_vulkan/legalize.py +754 -0
- numba_vulkan-0.1.0/src/numba_vulkan/libclc.py +281 -0
- numba_vulkan-0.1.0/src/numba_vulkan/mathfuncs.py +697 -0
- numba_vulkan-0.1.0/src/numba_vulkan/mathimpl.py +869 -0
- numba_vulkan-0.1.0/src/numba_vulkan/models.py +112 -0
- numba_vulkan-0.1.0/src/numba_vulkan/narrowing.py +429 -0
- numba_vulkan-0.1.0/src/numba_vulkan/probes.py +230 -0
- numba_vulkan-0.1.0/src/numba_vulkan/runtime.py +3644 -0
- numba_vulkan-0.1.0/src/numba_vulkan/structurize.py +912 -0
- numba_vulkan-0.1.0/src/numba_vulkan/stubs.py +451 -0
- numba_vulkan-0.1.0/src/numba_vulkan/target.py +501 -0
- numba_vulkan-0.1.0/src/numba_vulkan/ufuncs.py +143 -0
- numba_vulkan-0.1.0/src/numba_vulkan/vectorizers.py +889 -0
- numba_vulkan-0.1.0/src/numba_vulkan/vkdecl.py +693 -0
- numba_vulkan-0.1.0/src/numba_vulkan/vkimpl.py +1915 -0
- numba_vulkan-0.1.0/src/numba_vulkan/vktypes.py +251 -0
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
BSD 2-Clause License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026, Mirza Arnaut
|
|
4
|
+
|
|
5
|
+
Redistribution and use in source and binary forms, with or without
|
|
6
|
+
modification, are permitted provided that the following conditions are met:
|
|
7
|
+
|
|
8
|
+
1. Redistributions of source code must retain the above copyright notice, this
|
|
9
|
+
list of conditions and the following disclaimer.
|
|
10
|
+
|
|
11
|
+
2. Redistributions in binary form must reproduce the above copyright notice,
|
|
12
|
+
this list of conditions and the following disclaimer in the documentation
|
|
13
|
+
and/or other materials provided with the distribution.
|
|
14
|
+
|
|
15
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
16
|
+
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
17
|
+
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
18
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
19
|
+
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
20
|
+
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
21
|
+
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
22
|
+
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
23
|
+
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
24
|
+
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
@@ -0,0 +1,553 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: numba-vulkan
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A Vulkan compute target for Numba: Python kernels on any GPU with a Vulkan driver
|
|
5
|
+
Keywords: numba,vulkan,gpu,spir-v,jit,compiler
|
|
6
|
+
Author: Mirza Arnaut
|
|
7
|
+
Author-email: Mirza Arnaut <mirza.arnaut@tu-dortmund.de>
|
|
8
|
+
License-Expression: BSD-2-Clause AND Apache-2.0 WITH LLVM-exception
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
License-File: src/numba_vulkan/data/LICENSE-libclc.txt
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: Intended Audience :: Science/Research
|
|
14
|
+
Classifier: Operating System :: MacOS
|
|
15
|
+
Classifier: Operating System :: Microsoft :: Windows
|
|
16
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
17
|
+
Classifier: Programming Language :: Python :: 3 :: Only
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering
|
|
23
|
+
Classifier: Topic :: Software Development :: Compilers
|
|
24
|
+
Requires-Dist: numba>=0.68
|
|
25
|
+
Requires-Dist: llvmlite>=0.50
|
|
26
|
+
Requires-Dist: numpy>=1.25
|
|
27
|
+
Requires-Dist: vulkan>=1.3.275
|
|
28
|
+
Requires-Python: >=3.11
|
|
29
|
+
Project-URL: Repository, https://github.com/arunoruto/numba-vulkan
|
|
30
|
+
Project-URL: Documentation, https://arunoruto.github.io/numba-vulkan/
|
|
31
|
+
Project-URL: Issues, https://github.com/arunoruto/numba-vulkan/issues
|
|
32
|
+
Project-URL: Changelog, https://github.com/arunoruto/numba-vulkan/blob/main/CHANGELOG.md
|
|
33
|
+
Description-Content-Type: text/markdown
|
|
34
|
+
|
|
35
|
+
<a id="readme-top"></a>
|
|
36
|
+
|
|
37
|
+
<div align="center">
|
|
38
|
+
<img src="https://raw.githubusercontent.com/arunoruto/numba-vulkan/main/docs/source/_static/logo.svg" alt="numba-vulkan logo" width="160" />
|
|
39
|
+
<h1 align="center">numba-vulkan</h1>
|
|
40
|
+
|
|
41
|
+
<p align="center">
|
|
42
|
+
A Vulkan compute target for Numba: write a kernel in Python, run it on any GPU with a Vulkan driver.
|
|
43
|
+
<br />
|
|
44
|
+
<a href="https://github.com/numba/numba/issues/10116"><strong>Background: numba/numba#10116 »</strong></a>
|
|
45
|
+
<br />
|
|
46
|
+
<br />
|
|
47
|
+
<a href="https://arunoruto.github.io/numba-vulkan/">Documentation</a>
|
|
48
|
+
·
|
|
49
|
+
<a href="#usage">Usage</a>
|
|
50
|
+
·
|
|
51
|
+
<a href="#benchmarks">Benchmarks</a>
|
|
52
|
+
·
|
|
53
|
+
<a href="#roadmap">Roadmap</a>
|
|
54
|
+
</p>
|
|
55
|
+
</div>
|
|
56
|
+
|
|
57
|
+
<details>
|
|
58
|
+
<summary>Table of Contents</summary>
|
|
59
|
+
<ol>
|
|
60
|
+
<li>
|
|
61
|
+
<a href="#about-the-project">About The Project</a>
|
|
62
|
+
<ul>
|
|
63
|
+
<li><a href="#built-with">Built With</a></li>
|
|
64
|
+
<li><a href="#how-it-works">How It Works</a></li>
|
|
65
|
+
</ul>
|
|
66
|
+
</li>
|
|
67
|
+
<li>
|
|
68
|
+
<a href="#getting-started">Getting Started</a>
|
|
69
|
+
<ul>
|
|
70
|
+
<li><a href="#prerequisites">Prerequisites</a></li>
|
|
71
|
+
<li><a href="#installation">Installation</a></li>
|
|
72
|
+
</ul>
|
|
73
|
+
</li>
|
|
74
|
+
<li><a href="#usage">Usage</a></li>
|
|
75
|
+
<li><a href="#benchmarks">Benchmarks</a></li>
|
|
76
|
+
<li><a href="#what-works-and-what-does-not">What Works and What Does Not</a></li>
|
|
77
|
+
<li><a href="#known-issues">Known Issues</a></li>
|
|
78
|
+
<li><a href="#roadmap">Roadmap</a></li>
|
|
79
|
+
<li><a href="#authorship-and-ai-disclosure">Authorship and AI Disclosure</a></li>
|
|
80
|
+
<li><a href="#contributing">Contributing</a></li>
|
|
81
|
+
<li><a href="#license">License</a></li>
|
|
82
|
+
<li><a href="#contact">Contact</a></li>
|
|
83
|
+
<li><a href="#acknowledgments">Acknowledgments</a></li>
|
|
84
|
+
</ol>
|
|
85
|
+
</details>
|
|
86
|
+
|
|
87
|
+
> [!NOTE]
|
|
88
|
+
> **AI disclosure.** The code, tests and documentation in this repository
|
|
89
|
+
> were written by an AI (Anthropic's Claude) under the direction of a human,
|
|
90
|
+
> who came up with the idea, set the design direction and reviewed the
|
|
91
|
+
> result. See [Authorship](#authorship-and-ai-disclosure).
|
|
92
|
+
|
|
93
|
+
## About The Project
|
|
94
|
+
|
|
95
|
+
Numba's only maintained GPU target is CUDA, which ties GPU-accelerated Numba
|
|
96
|
+
code to NVIDIA hardware. Supporting every vendor separately is a lot of work;
|
|
97
|
+
Vulkan is implemented by nearly all of them.
|
|
98
|
+
|
|
99
|
+
numba-vulkan registers a `vulkan` target through Numba's target extension API
|
|
100
|
+
and compiles `numba.cuda`-style kernels to SPIR-V compute shaders. The same
|
|
101
|
+
kernel has been run on an NVIDIA GPU, an Intel integrated GPU and a CPU
|
|
102
|
+
rasteriser (llvmpipe), with results checked against NumPy.
|
|
103
|
+
|
|
104
|
+
This is alpha software. Its test suite passes on NVIDIA and Intel GPUs and
|
|
105
|
+
llvmpipe on Linux, on an Apple M3 Pro through MoltenVK, and on Windows with
|
|
106
|
+
lavapipe, and the benchmarks run on an Intel Arc GPU under Windows; but its
|
|
107
|
+
interfaces may still change and it has not yet been proven on real
|
|
108
|
+
workloads.
|
|
109
|
+
|
|
110
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
111
|
+
|
|
112
|
+
### Built With
|
|
113
|
+
|
|
114
|
+
- [Numba](https://numba.pydata.org/): bytecode frontend, type inference,
|
|
115
|
+
lowering and the target extension API
|
|
116
|
+
- [libclc](https://libclc.llvm.org/), LLVM's OpenCL math library, linked into
|
|
117
|
+
kernels for accurate and double-precision math
|
|
118
|
+
- [llvmlite](https://llvmlite.readthedocs.io/) 0.50+ (LLVM 22), whose SPIR-V
|
|
119
|
+
backend emits Vulkan-flavoured SPIR-V
|
|
120
|
+
- [vulkan](https://pypi.org/project/vulkan/): Python bindings for the Vulkan API
|
|
121
|
+
- [devenv](https://devenv.sh/) and [uv](https://docs.astral.sh/uv/) for the
|
|
122
|
+
development environment
|
|
123
|
+
|
|
124
|
+
### How It Works
|
|
125
|
+
|
|
126
|
+
1. `@nv.jit` creates a dispatcher for the `vulkan` target. Numba's own pipeline
|
|
127
|
+
(bytecode analysis, type inference, lowering) produces LLVM IR.
|
|
128
|
+
2. Arrays are a custom Numba type that carries the descriptor binding of its
|
|
129
|
+
storage buffer, because shaders have no general pointers. Element access is
|
|
130
|
+
emitted as a placeholder call taking a binding and an index.
|
|
131
|
+
3. Everything is inlined into one entry point and simplified with a few LLVM
|
|
132
|
+
passes. Only then are the placeholders expanded into buffer accesses.
|
|
133
|
+
4. LLVM's SPIR-V backend turns the result into a compute shader. It runs in a
|
|
134
|
+
child process, so a backend failure raises an exception instead of
|
|
135
|
+
aborting Python.
|
|
136
|
+
5. The runtime uploads the arguments (binding 0: error status, binding
|
|
137
|
+
`1 + k`: array argument `k`; scalars and array shapes as push constants),
|
|
138
|
+
dispatches, and copies written buffers back. Device arrays are used in
|
|
139
|
+
place.
|
|
140
|
+
|
|
141
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
142
|
+
|
|
143
|
+
## Getting Started
|
|
144
|
+
|
|
145
|
+
### Prerequisites
|
|
146
|
+
|
|
147
|
+
- libclc's `clspv--.bc` from LLVM 22 or older, for float64 and accurate
|
|
148
|
+
float32 math. Wheels bundle it and the devenv shell provides it; only a
|
|
149
|
+
bare source checkout needs it installed (see the
|
|
150
|
+
[math library docs](https://arunoruto.github.io/numba-vulkan/math_library.html)).
|
|
151
|
+
- A Vulkan 1.2 driver. [lavapipe](https://docs.mesa3d.org/drivers/llvmpipe.html)
|
|
152
|
+
(Mesa's CPU implementation) is enough to try it without a GPU.
|
|
153
|
+
- Python 3.11 or newer.
|
|
154
|
+
- Optional: `spirv-val` from SPIRV-Tools to validate generated shaders.
|
|
155
|
+
|
|
156
|
+
### Installation
|
|
157
|
+
|
|
158
|
+
From PyPI, with the Vulkan loader (`libvulkan1` on Debian and Ubuntu,
|
|
159
|
+
`vulkan-loader` on Fedora and Arch) and a Vulkan driver for your GPU
|
|
160
|
+
installed:
|
|
161
|
+
|
|
162
|
+
```sh
|
|
163
|
+
pip install numba-vulkan # or: uv add numba-vulkan
|
|
164
|
+
python -c "import numba_vulkan as nv; print(nv.list_devices())"
|
|
165
|
+
```
|
|
166
|
+
|
|
167
|
+
The wheel bundles libclc for accurate and `float64` math. On macOS it also
|
|
168
|
+
needs MoltenVK and an `llc` of LLVM 22 named by `NUMBA_VULKAN_LLC` (KI-32);
|
|
169
|
+
see below.
|
|
170
|
+
|
|
171
|
+
To work on numba-vulkan itself: with devenv (NixOS or any system with Nix),
|
|
172
|
+
everything including the Vulkan loader and SPIRV-Tools is provided by the
|
|
173
|
+
shell:
|
|
174
|
+
|
|
175
|
+
```sh
|
|
176
|
+
devenv shell
|
|
177
|
+
uv run pytest
|
|
178
|
+
```
|
|
179
|
+
|
|
180
|
+
Without Nix, install the Vulkan loader from your distribution, then:
|
|
181
|
+
|
|
182
|
+
```sh
|
|
183
|
+
uv sync
|
|
184
|
+
uv run pytest
|
|
185
|
+
```
|
|
186
|
+
|
|
187
|
+
On macOS, Vulkan runs on MoltenVK, and llvmlite's SPIR-V backend has to be
|
|
188
|
+
replaced by an `llc` of LLVM 22 (KI-32). The devenv shell sets up both
|
|
189
|
+
(`devenv shell -- uv run pytest`); the test suite passes on an Apple M3 Pro.
|
|
190
|
+
See [Getting started](https://arunoruto.github.io/numba-vulkan/getting_started.html#macos)
|
|
191
|
+
for the details and a Homebrew alternative.
|
|
192
|
+
|
|
193
|
+
The test suite runs every test on every Vulkan device it finds.
|
|
194
|
+
|
|
195
|
+
### Documentation
|
|
196
|
+
|
|
197
|
+
The Sphinx documentation is published at
|
|
198
|
+
<https://arunoruto.github.io/numba-vulkan/> and covers usage, the design, the
|
|
199
|
+
limitations, the known issues and the benchmarks. To build it locally:
|
|
200
|
+
|
|
201
|
+
```sh
|
|
202
|
+
uv sync --group docs
|
|
203
|
+
cd docs && uv run sphinx-build -M html ./source ./build
|
|
204
|
+
```
|
|
205
|
+
|
|
206
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
207
|
+
|
|
208
|
+
## Usage
|
|
209
|
+
|
|
210
|
+
A kernel looks like a `numba.cuda` kernel. `global_id(axis)` plays the role of
|
|
211
|
+
`cuda.grid`, and `forall` launches it over a 1-d to 3-d grid:
|
|
212
|
+
|
|
213
|
+
```python
|
|
214
|
+
import numpy as np
|
|
215
|
+
import numba_vulkan as nv
|
|
216
|
+
|
|
217
|
+
@nv.jit
|
|
218
|
+
def saxpy(a, x, y):
|
|
219
|
+
i = nv.global_id(0)
|
|
220
|
+
if i < x.shape[0]:
|
|
221
|
+
y[i] = a * x[i] + y[i]
|
|
222
|
+
|
|
223
|
+
x = np.arange(1000, dtype=np.float32)
|
|
224
|
+
y = np.ones(1000, dtype=np.float32)
|
|
225
|
+
saxpy.forall(x.size)(np.float32(2.0), x, y)
|
|
226
|
+
```
|
|
227
|
+
|
|
228
|
+
NumPy arrays are copied to the device and back on every call. Device arrays
|
|
229
|
+
stay there, as with `numba.cuda`:
|
|
230
|
+
|
|
231
|
+
```python
|
|
232
|
+
dx, dy = nv.to_device(x), nv.to_device(y)
|
|
233
|
+
saxpy.forall(x.size)(np.float32(2.0), dx, dy) # no copies
|
|
234
|
+
saxpy.forall(x.size)(np.float32(0.5), dx, dy)
|
|
235
|
+
y = dy.copy_to_host()
|
|
236
|
+
```
|
|
237
|
+
|
|
238
|
+
Functions compiled with `@nv.jit` can be called from kernels, may return
|
|
239
|
+
values or tuples, and can take arrays. Functions compiled with `numba.njit`
|
|
240
|
+
are recompiled for Vulkan when a kernel calls them:
|
|
241
|
+
|
|
242
|
+
```python
|
|
243
|
+
import math
|
|
244
|
+
from numba import njit
|
|
245
|
+
|
|
246
|
+
@njit
|
|
247
|
+
def norm(x, y):
|
|
248
|
+
return math.sqrt(x * x + y * y)
|
|
249
|
+
|
|
250
|
+
@nv.jit
|
|
251
|
+
def radius(x, y, out):
|
|
252
|
+
i = nv.global_id(0)
|
|
253
|
+
if i < out.size:
|
|
254
|
+
out[i] = norm(x[i], y[i])
|
|
255
|
+
```
|
|
256
|
+
|
|
257
|
+
Numba's ufunc decorators work too, without writing a kernel:
|
|
258
|
+
|
|
259
|
+
```python
|
|
260
|
+
from numba import vectorize, guvectorize
|
|
261
|
+
|
|
262
|
+
@vectorize(["float32(float32, float32)"], target="vulkan")
|
|
263
|
+
def hypot(a, b):
|
|
264
|
+
return math.sqrt(a * a + b * b)
|
|
265
|
+
|
|
266
|
+
@guvectorize(["void(float32[:], float32[:], float32[:])"], "(n),(n)->()",
|
|
267
|
+
target="vulkan")
|
|
268
|
+
def dot(a, b, out):
|
|
269
|
+
out[0] = np.dot(a, b)
|
|
270
|
+
|
|
271
|
+
hypot(x, y) # element-wise, with broadcasting
|
|
272
|
+
dot(matrix, vector) # one dot product per row
|
|
273
|
+
```
|
|
274
|
+
|
|
275
|
+
Cooperating kernels use the same building blocks as in `numba.cuda`:
|
|
276
|
+
workgroup-shared arrays, barriers, atomics and an explicit launch
|
|
277
|
+
configuration:
|
|
278
|
+
|
|
279
|
+
```python
|
|
280
|
+
@nv.jit
|
|
281
|
+
def block_sums(x, out):
|
|
282
|
+
partial = nv.shared.array(256, np.float32)
|
|
283
|
+
t, i = nv.local_id(0), nv.global_id(0)
|
|
284
|
+
partial[t] = x[i] if i < x.shape[0] else np.float32(0)
|
|
285
|
+
nv.barrier()
|
|
286
|
+
step = 128
|
|
287
|
+
while step > 0:
|
|
288
|
+
if t < step:
|
|
289
|
+
partial[t] += partial[t + step]
|
|
290
|
+
nv.barrier()
|
|
291
|
+
step //= 2
|
|
292
|
+
if t == 0:
|
|
293
|
+
nv.atomic.add(out, 0, partial[0])
|
|
294
|
+
|
|
295
|
+
block_sums[groups, 256](x, out) # like kernel[blocks, threads] in CUDA
|
|
296
|
+
```
|
|
297
|
+
|
|
298
|
+
Because it is a real Numba target, `@overload` works with it:
|
|
299
|
+
|
|
300
|
+
```python
|
|
301
|
+
from numba.extending import overload
|
|
302
|
+
|
|
303
|
+
@overload(my_function, target="vulkan")
|
|
304
|
+
def ol_my_function(x):
|
|
305
|
+
return lambda x: x * 2
|
|
306
|
+
```
|
|
307
|
+
|
|
308
|
+
Choosing a device:
|
|
309
|
+
|
|
310
|
+
```python
|
|
311
|
+
nv.list_devices() # [<0: NVIDIA ... (discrete)>, <1: ...>, ...]
|
|
312
|
+
nv.select_device("intel") # by name, index, or NUMBA_VULKAN_DEVICE
|
|
313
|
+
saxpy.forall(n, device="llvmpipe")(a, x, y)
|
|
314
|
+
```
|
|
315
|
+
|
|
316
|
+
Inspecting what was generated:
|
|
317
|
+
|
|
318
|
+
```python
|
|
319
|
+
from numba import float32
|
|
320
|
+
compiled = saxpy.compile((float32, float32[::1], float32[::1]))
|
|
321
|
+
compiled.llvm_ir, compiled.spirv, compiled.capabilities
|
|
322
|
+
```
|
|
323
|
+
|
|
324
|
+
Set `NUMBA_VULKAN_VALIDATE=1` to run every generated shader through
|
|
325
|
+
`spirv-val` before it reaches a driver.
|
|
326
|
+
|
|
327
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
328
|
+
|
|
329
|
+
## Benchmarks
|
|
330
|
+
|
|
331
|
+
Two suites compare numba-vulkan with Numba's CPU target and numba-cuda: the
|
|
332
|
+
same scalar function compiled for every backend (`benchmarks/bench.py`), and
|
|
333
|
+
the same CUDA-style kernels with shared memory, barriers and atomics
|
|
334
|
+
(`benchmarks/kernels.py`). Results are collected in the repository, one file
|
|
335
|
+
per run, and charted in the
|
|
336
|
+
[documentation](https://arunoruto.github.io/numba-vulkan/benchmarks.html),
|
|
337
|
+
including every backend, device and variant and how results change over
|
|
338
|
+
time. The fastest result per backend, GPUs with data on the device (log
|
|
339
|
+
scale; further left is faster):
|
|
340
|
+
|
|
341
|
+
<p align="center">
|
|
342
|
+
<img src="https://arunoruto.github.io/numba-vulkan/_images/machines.svg" alt="Best time per workload and backend" width="100%" />
|
|
343
|
+
</p>
|
|
344
|
+
|
|
345
|
+
In the first collected run (Intel i9-9900K, NVIDIA TITAN X, Intel UHD 630):
|
|
346
|
+
|
|
347
|
+
- With data on the device, numba-vulkan takes 25 to 55 % longer than
|
|
348
|
+
numba-cuda on the same card on most workloads, and is about level on the
|
|
349
|
+
tiled matrix product (1.85 against 1.75 ms). Repeated runs vary by 10 to
|
|
350
|
+
20 %, so small differences are not meaningful.
|
|
351
|
+
- With NumPy arrays as arguments, the copies dominate: the memory-bound
|
|
352
|
+
`saxpy` takes 7 to 10 ms on every GPU backend and 1.8 ms on one CPU thread.
|
|
353
|
+
- A kernel's first call, which compiles it, takes 0.04 to 0.25 s on Vulkan,
|
|
354
|
+
comparable to numba-cuda.
|
|
355
|
+
|
|
356
|
+
On an Apple M3 Pro (MoltenVK 1.4.2, kernels translated by `llc`, KI-32),
|
|
357
|
+
with data on the device:
|
|
358
|
+
|
|
359
|
+
- The compute-bound Mandelbrot set takes 7 to 9 ms, about level with the
|
|
360
|
+
TITAN X; within a full run the GPU shares the chip's power budget with
|
|
361
|
+
the CPU benchmarks before it, which costs up to 20 %.
|
|
362
|
+
- A launch that waits for its result costs about 0.2 ms, so the small
|
|
363
|
+
workloads mostly measure that: `saxpy` takes 0.61 ms, about the same as the CPU.
|
|
364
|
+
- The tiled matrix product reaches only about 185 GFLOP/s (11.7 ms), the
|
|
365
|
+
clearest candidate for tuning.
|
|
366
|
+
|
|
367
|
+
Please add a run from your machine; it takes a few minutes:
|
|
368
|
+
|
|
369
|
+
```sh
|
|
370
|
+
uv run python benchmarks/collect.py --user <your GitHub name>
|
|
371
|
+
```
|
|
372
|
+
|
|
373
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
374
|
+
|
|
375
|
+
## What Works and What Does Not
|
|
376
|
+
|
|
377
|
+
Works:
|
|
378
|
+
|
|
379
|
+
- `int32`/`int64`/`float32`/`float64`/`bool` scalars and N-d arrays of
|
|
380
|
+
`bool`, 8 to 64-bit integers, `float32` and `float64`
|
|
381
|
+
- indexing with integers and slices; slices are views, as in NumPy, and can
|
|
382
|
+
be iterated over, passed to functions, assigned to and reduced
|
|
383
|
+
(`a[i, 1:-1].sum()`, `for v in a[i]`, `out[i, :3] = 0`, `np.dot`, `.T`)
|
|
384
|
+
- arithmetic, comparisons, bit operations, casts, tuples, complex numbers,
|
|
385
|
+
`if`/`while`/`for ... in range(...)`, early `return`s, `min`/`max`/`abs`
|
|
386
|
+
- the `math` module in float32 **and float64** (including `hypot`, `log1p`,
|
|
387
|
+
`erf`, `gamma`, `isnan`...), accurate to the last digit and consistent
|
|
388
|
+
across devices, and the `**` operator
|
|
389
|
+
- NumPy functions on scalars, such as `np.sqrt(x[i])` or `np.maximum(a, b)`
|
|
390
|
+
- calling other `@nv.jit` and `@njit` functions; `@overload(target="vulkan")`
|
|
391
|
+
|
|
392
|
+
Does not work:
|
|
393
|
+
|
|
394
|
+
- **float64 math without libclc.** Vulkan's math library is 32-bit only, so
|
|
395
|
+
`sin`, `exp`, `pow` and friends in float64 come from libclc, LLVM's OpenCL
|
|
396
|
+
math library. Wheels bundle it; in a source checkout without it such calls
|
|
397
|
+
raise, unless you opt in to float32 precision with
|
|
398
|
+
`@nv.jit(narrow_math=True)`.
|
|
399
|
+
- **Arrays of run-time size** in kernels (`a[mask]`, `a.copy()`,
|
|
400
|
+
`np.zeros(n)`), `try`/`except`, recursion. Arrays of constant shape,
|
|
401
|
+
array expressions such as `a * 2 + b` and reductions along an axis such
|
|
402
|
+
as `a.sum(axis=0)` work; on the host, device arrays support masks and
|
|
403
|
+
index arrays.
|
|
404
|
+
- **float64 precision on devices without float64.** Many mobile GPUs and
|
|
405
|
+
Apple devices lack float64 or int64. Kernels are narrowed to 32-bit types
|
|
406
|
+
there, with a warning; float narrowing is tested on an Apple M3 Pro,
|
|
407
|
+
integer narrowing only on desktop GPUs with the feature switched off.
|
|
408
|
+
- **Selecting between array elements by reference** in ways the optimiser
|
|
409
|
+
turns into pointer selects is rejected.
|
|
410
|
+
|
|
411
|
+
Things found along the way:
|
|
412
|
+
|
|
413
|
+
- NVIDIA's driver returns wrong results for 64-bit signed remainders with
|
|
414
|
+
negative operands; remainders are computed from the quotient instead.
|
|
415
|
+
- NVIDIA's shader compiler segfaults on some invalid SPIR-V, which is why
|
|
416
|
+
modules are checked before being handed to a driver.
|
|
417
|
+
- float32 `sin` on Intel/Mesa is only accurate to about 1e-4.
|
|
418
|
+
- LLVM's SPIR-V backend miscompiles nested aggregate inserts and aborts on
|
|
419
|
+
several unsupported inputs.
|
|
420
|
+
- LLVM's SPIR-V structurizer emits invalid shaders for early returns and
|
|
421
|
+
short-circuit conditions; control flow is restructured beforehand.
|
|
422
|
+
- All three drivers reassociate float arithmetic unless told not to, so
|
|
423
|
+
every float operation is marked exact.
|
|
424
|
+
- LLVM's SPIR-V backend mistakes `x = a * b; (x < 0) ? -y : y` for GLSL's
|
|
425
|
+
`faceforward` and crashes; such comparisons are rewritten.
|
|
426
|
+
|
|
427
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
428
|
+
|
|
429
|
+
## Known Issues
|
|
430
|
+
|
|
431
|
+
The full list, with causes, workarounds and pointers to where each would be
|
|
432
|
+
fixed, is in [known issues](https://arunoruto.github.io/numba-vulkan/known_issues.html). The
|
|
433
|
+
ones most likely to bite:
|
|
434
|
+
|
|
435
|
+
| | Issue | Workaround |
|
|
436
|
+
| --- | --- | --- |
|
|
437
|
+
| KI-01 | float64 `math.sin(x)`, `x ** 2.5`... need libclc, which a source checkout lacks | install a wheel, or set `NUMBA_VULKAN_LIBCLC` |
|
|
438
|
+
| KI-04 | no arrays of run-time size in kernels: `a[mask]`, `a.copy()`, `np.zeros(n)` | constant-shape arrays, array expressions, `axis=` reductions |
|
|
439
|
+
| KI-10 | indices and divisions by zero are only checked on request | `@nv.jit(boundscheck=True, error_model="python")` while debugging |
|
|
440
|
+
| KI-17 | running without int64 (mobile) is tested only with the feature switched off on desktop GPUs | report what you find |
|
|
441
|
+
| KI-32 | llvmlite's SPIR-V backend segfaults on macOS arm64 | set `NUMBA_VULKAN_LLC` to an LLVM 22 `llc`; the devenv shell does |
|
|
442
|
+
|
|
443
|
+
`tests/test_known_issues.py` reproduces the coverage gaps as expected
|
|
444
|
+
failures, so the list stays honest:
|
|
445
|
+
|
|
446
|
+
```sh
|
|
447
|
+
uv run pytest tests/test_known_issues.py -rxX
|
|
448
|
+
```
|
|
449
|
+
|
|
450
|
+
To continue the work, start with the
|
|
451
|
+
[development guide](https://arunoruto.github.io/numba-vulkan/development.html).
|
|
452
|
+
|
|
453
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
454
|
+
|
|
455
|
+
## Roadmap
|
|
456
|
+
|
|
457
|
+
- [x] Vulkan target registered with Numba's target extension API
|
|
458
|
+
- [x] Kernels, device functions, overloads
|
|
459
|
+
- [x] Test suite running on every available Vulkan device
|
|
460
|
+
- [x] Sphinx documentation and benchmark suite
|
|
461
|
+
- [x] Math functions, NumPy ufuncs on scalars, complex numbers
|
|
462
|
+
- [x] Early returns and short-circuit conditions (control-flow restructuring)
|
|
463
|
+
- [x] Loops with `break`/`return` anywhere, `while` loops (fuzz-tested)
|
|
464
|
+
- [x] Device arrays and buffer reuse, to avoid copying on every call
|
|
465
|
+
- [x] Asynchronous launches on device arrays (`nv.synchronize()`)
|
|
466
|
+
- [x] Indexing (basic and advanced), views and NumPy ufuncs on device arrays
|
|
467
|
+
- [x] Structured (record) arrays
|
|
468
|
+
- [x] Autotuning of kernel variants (`nv.autotune`)
|
|
469
|
+
- [x] Device-side timing with events (`nv.event`)
|
|
470
|
+
- [x] Streams overlapping copies and kernels (`nv.stream`, `nv.pinned_array`)
|
|
471
|
+
- [x] Launch plans: repeated launches skip typing
|
|
472
|
+
- [x] Scalars as push constants
|
|
473
|
+
- [x] Slices, array views, iteration and reductions, also along an axis
|
|
474
|
+
- [x] Array expressions and local arrays
|
|
475
|
+
- [x] Kernels narrowed to 32-bit types on devices without float64/int64
|
|
476
|
+
- [x] float64 math, through libclc
|
|
477
|
+
- [x] Bundle libclc in wheels
|
|
478
|
+
- [ ] Publish to PyPI
|
|
479
|
+
- [x] `numba.vectorize` / `numba.guvectorize` with `target="vulkan"`
|
|
480
|
+
- [x] `reduce` of `vectorize` functions
|
|
481
|
+
- [x] `print()` in kernels, `float16` arrays, `error_model="python"`
|
|
482
|
+
- [ ] Other ufunc methods (`accumulate`, `outer`)
|
|
483
|
+
- [x] Shared memory, atomics, barriers and CUDA-style launch configuration
|
|
484
|
+
- [x] Subgroup operations: reductions, scans, votes, ballots, shuffles
|
|
485
|
+
- [x] On-disk caching of compiled kernels
|
|
486
|
+
- [x] `cache=True`: compiled kernels reused without recompiling
|
|
487
|
+
- [x] Testing on Apple (MoltenVK)
|
|
488
|
+
- [ ] Testing on AMD and mobile GPUs
|
|
489
|
+
|
|
490
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
491
|
+
|
|
492
|
+
## Authorship and AI Disclosure
|
|
493
|
+
|
|
494
|
+
| | Who |
|
|
495
|
+
| --- | --- |
|
|
496
|
+
| Idea | Mirza Arnaut, in [numba/numba#10116](https://github.com/numba/numba/issues/10116) |
|
|
497
|
+
| Direction and design decisions | Mirza Arnaut |
|
|
498
|
+
| Implementation: code, tests, benchmarks, documentation | Claude (Anthropic's Claude Fable 5.1, and Claude Opus 5.5 from 2 October 2026 on, running in Claude Code) |
|
|
499
|
+
| Review | Mirza Arnaut |
|
|
500
|
+
|
|
501
|
+
Every line here was generated by a language model in October 2026; the human
|
|
502
|
+
author steered the work and reviewed it but did not write the code. Many
|
|
503
|
+
implementation-level choices were made by the model while solving problems it
|
|
504
|
+
ran into, and should be judged on their merits. Benchmark numbers and
|
|
505
|
+
hardware observations come from real runs on the author's machine.
|
|
506
|
+
|
|
507
|
+
This is an independent project and is not endorsed by the Numba maintainers.
|
|
508
|
+
Numba has an
|
|
509
|
+
[AI tools policy](https://numba.readthedocs.io/en/stable/reference/ai_tools_policy.html)
|
|
510
|
+
that requires attributing LLM-generated contributions; it applies if any part
|
|
511
|
+
of this is ever proposed upstream.
|
|
512
|
+
|
|
513
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
514
|
+
|
|
515
|
+
## Contributing
|
|
516
|
+
|
|
517
|
+
Issues and pull requests are welcome, in particular reports of how the test
|
|
518
|
+
suite behaves on other GPUs and drivers:
|
|
519
|
+
|
|
520
|
+
```sh
|
|
521
|
+
uv run pytest
|
|
522
|
+
uv run python benchmarks/bench.py
|
|
523
|
+
```
|
|
524
|
+
|
|
525
|
+
Contributions written with AI tools are welcome on the same terms as the
|
|
526
|
+
project itself: say which tool produced them, and have a human review them.
|
|
527
|
+
|
|
528
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
529
|
+
|
|
530
|
+
## License
|
|
531
|
+
|
|
532
|
+
Distributed under the BSD 2-Clause License, the same license as Numba. See
|
|
533
|
+
[LICENSE](https://github.com/arunoruto/numba-vulkan/blob/main/LICENSE).
|
|
534
|
+
|
|
535
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
536
|
+
|
|
537
|
+
## Contact
|
|
538
|
+
|
|
539
|
+
Mirza Arnaut - [@arunoruto](https://github.com/arunoruto)
|
|
540
|
+
|
|
541
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|
|
542
|
+
|
|
543
|
+
## Acknowledgments
|
|
544
|
+
|
|
545
|
+
- [Numba](https://github.com/numba/numba) and
|
|
546
|
+
[numba-cuda](https://github.com/NVIDIA/numba-cuda), whose CUDA target is the
|
|
547
|
+
template for this one
|
|
548
|
+
- [Taichi](https://github.com/taichi-dev/taichi), for showing that Python
|
|
549
|
+
kernels on Vulkan work
|
|
550
|
+
- The LLVM SPIR-V backend developers
|
|
551
|
+
- [Best-README-Template](https://github.com/othneildrew/Best-README-Template)
|
|
552
|
+
|
|
553
|
+
<p align="right">(<a href="#readme-top">back to top</a>)</p>
|