pygpgpu 0.2.0__tar.gz → 0.2.2__tar.gz

{pygpgpu-0.2.0/pygpgpu.egg-info → pygpgpu-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pygpgpu
-Version: 0.2.0
+Version: 0.2.2
 Summary: Run Python functions as GPU kernels with the @kernel decorator (OpenCL / CUDA / HIP); handwritten .cl/.cu/.hip kernels import like Python modules
 Author-email: "王炳辉 (Bing-Hui WANG)" <binghui.wang@foxmail.com>
 License: MIT License
@@ -32,6 +32,7 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.7
 Classifier: Programming Language :: Python :: 3.8
@@ -50,9 +51,9 @@ License-File: LICENSE
 Requires-Dist: numpy>=1.21
 Requires-Dist: imageio>=2.30
 Requires-Dist: tree-sitter>=0.20.0
-Requires-Dist: tree-sitter-opencl>=0.0.2
-Requires-Dist: tree-sitter-cuda37>=0.21.2
-Requires-Dist: pygf>=0.1.4
+Requires-Dist: tree-sitter-opencl>=0.0.3
+Requires-Dist: tree-sitter-cuda37>=0.21.3
+Requires-Dist: pygf>=0.2.0
 Requires-Dist: typing_extensions>=4.4.0
 Provides-Extra: dev
 Requires-Dist: pytest>=8; extra == "dev"
@@ -64,7 +65,7 @@ Dynamic: license-file
 
 **English** | [Chinese](README.zh-CN.md)
 
-Decorate a Python function and it runs on the GPU — one codebase, three backends (**OpenCL / CUDA / HIP**).
+Decorate a Python function and it runs on the GPU — one codebase, four backends (**OpenCL / CUDA / HIP / Metal**).
 
 ```python
 import numpy as np
@@ -78,7 +79,7 @@ def saxpy(x: np.ndarray, y: np.ndarray, out: np.ndarray, alpha: float) -> None:
 x = np.arange(1_000_000, dtype=np.float32)
 out = np.zeros_like(x)
 
-saxpy(x, x, out, 2.0)              # no backend given: probe cuda -> hip -> opencl, take the first available
+saxpy(x, x, out, 2.0)              # no backend given: probe cuda -> hip -> metal -> opencl, take the first available
 saxpy["cuda"](x, x, out, 2.0)      # bracket switches the backend
 saxpy["hip:gfx1036"](x, x, out, 2.0)   # backend:device
 saxpy["python"](x, x, out, 2.0)    # the same function, serially in the interpreter (no GPU)
@@ -90,25 +91,26 @@ No C code to write, no build system to configure; the first call translates, com
 
 - **`@kernel` decorator** — decorate a Python function directly (class methods and static methods included); at call time it is specialized per argument type, translated to backend source and compiled (monomorphization).
 - **Plain-Python helper functions** — a helper called from a kernel needs no decorator, and its annotations are optional: it is translated into a device function and monomorphized from the call site, recursively for helpers that call helpers. The same holds for helpers reached through a name: `module.func(...)`, `SomeClass.static_method(...)`, `self.static_method(...)`, and `self.other_method(...)` between methods of one class.
-- **All three backends in one process** — the bracket chooses the backend per call, so a single process (and even a single kernel object) can mix OpenCL, CUDA and HIP; each backend/device pair is specialized and cached separately.
+- **Every backend in one process** — the bracket chooses the backend per call, so a single process (and even a single kernel object) can mix OpenCL, CUDA, HIP and Metal; each backend/device pair is specialized and cached separately.
 - **Pure-Python backend** — `k["python"]` runs the same `@kernel` body in the interpreter, serially on the CPU: no GPU, no compiler, no runtime. Useful for testing, for small inputs and as a reference for what the GPU semantics mean.
-- **Handwritten kernel file import** — `.cl` / `.cu` / `.hip` files import like Python modules: kernels and struct types become module attributes, and a `.pyi` type stub is generated automatically.
+- **Handwritten kernel file import** — `.cl` / `.cu` / `.hip` / `.metal` files import like Python modules: kernels and struct types become module attributes, and a `.pyi` type stub is generated automatically.
 - **Low-level OOP API** — Device/Context/Buffer/Program/Kernel, for when you need full control.
 - **One bracket syntax** for backend/device/work size (identical to handwritten kernels, order does not matter):
   ```python
   k(args)                          # fully automatic
-  k["cuda"](args)                  # backend (opencl/cuda/hip/python, aliases ocl/cu/rocm/amd/py)
+  k["cuda"](args)                  # backend (opencl/cuda/hip/metal/python, aliases ocl/cu/rocm/amd/mtl/py)
   k["opencl:gfx1036"](args)        # backend:device
-  k[(1024, 256)](args)             # work size: (total, per group), None = automatic
-  k["cuda"][(1024, 256)](args)     # combined
+  k[(4, 256)](args)                # work size: (blocks per grid, threads per block)
+  k["cuda"][(4, 256)](args)        # combined
   ```
-- **Total-work semantics** (OpenCL/SYCL style): in `(global, local)`, `global` is the **total number of work items**; CUDA/HIP grid dimensions are derived as `global / local`, so all three backends behave exactly the same.
+- **CUDA work-size brackets on every backend**: (blocks_per_grid, threads_per_block), the spelling a CUDA or HIP reader already writes. k[(4, 256)] launches 4 blocks of 256 threads (1024 work items); k[(n,)] is a bare item count with an automatic block size, and None means automatic. The same bracket works on every backend, in @kernel and in the low-level oop API, and a block count never has to divide anything: the total is locks * threads by construction.
 - **Static read/write analysis**: the kernel body is parsed to infer whether each pointer argument is read-only, write-only or read-write — write-only arguments skip the upload, read-only arguments skip the read-back, and data stays on the device across calls (zero intermediate copies for gnp arrays).
 - **Vector types**: 50 vector types (float4/int3/...) with `.x/.y/.z/.w` members, swizzles, construction, arithmetic (scalar broadcast), comparisons and component-wise math.
 - **Dual-sided function library** `pygpgpu.funcs` (aka `pygpgpu.vec_types.funcs`): `dot/cross/normalize/clamp/...` run as-is in host Python and are translated to device builtins inside `@kernel`.
-- **Device builtin namespace** `gpu.*`: `get_global_id/barrier/atomic_add/dot/...` covers GPU builtins that have no Python counterpart, mapped per backend.
+- **Device builtins, exported at the top level**: `get_global_id`, `barrier`, `atomic_add`, `dot`, ... are plain functions imported from `pygpgpu`. Inside a `@kernel` body the translator recognizes each one by the object it resolves to and maps it to the call of every backend, so the same kernel body runs on OpenCL, CUDA, HIP and in the interpreter.
 - **Multi-dimensional iteration**: `np.ndindex(shape)` / `np.ndenumerate(arr)` magic loops (2D/3D), `enumerate(arr)` for a 1D index plus value snapshot.
 - **Three cache layers**: parsed metadata / compiled artifacts / generated sources, all invalidated by mtime — unchanged sources mean zero recompilation.
+- **Host-side CUDA libraries bound with ctypes** — cuBLAS, cuBLASLt, cuSOLVER, cuFFT, cuSPARSE and the CUDA runtime, in the same shape as the driver (lazy DLL, signature table, the library's own error messages), with a numpy-friendly layer over them (`gemm`, `axpy`, `scal`, `dot`, `rfft`, `csrmv`). See "Host-side CUDA libraries" below.
 
 ## Installation
 
@@ -122,31 +124,32 @@ Python ≥ 3.7. Bring your own GPU runtime (any one of them is enough):
 |---|---|---|
 | OpenCL | GPU driver (usually already installed) | Nothing extra to install; works across NVIDIA/AMD/Intel |
 | CUDA | NVIDIA driver (includes NVRTC) | The Toolkit is not required — the driver-provided nvrtc is enough |
-| HIP | AMD HIP SDK | On Windows compilation runs through the SDK's clang; if MSVC/Windows SDK headers are not on the default paths, point `PYGPGPU_HIP_INCLUDES` at the include directories (os.pathsep-separated) |
+| HIP | AMD HIP SDK | Compilation runs through the SDK's clang: on Windows that is `<ROCm>/bin/clang.exe` with the MSVC/Windows SDK headers auto-discovered (point `PYGPGPU_HIP_INCLUDES` at the include directories, os.pathsep-separated, if they are not on the default paths); on Linux it is `<ROCm>/bin/clang` (or `hipcc` on `PATH`), which finds the system headers itself |
+| Metal | macOS on Apple hardware | The Metal backend compiles through the framework's own runtime compiler, so no Xcode command-line tools are needed. It requires a real Apple GPU: a Mac virtual machine has no Metal device, so this backend is unavailable there |
 | Python | nothing | Runs the `@kernel` body serially in the interpreter; selected explicitly with `k["python"]` and never probed for |
 
-When `@kernel` has no explicit `backend`, backends are probed in **cuda -> hip -> opencl** order and the first one that enumerates a device wins.
+When `@kernel` has no explicit `backend`, backends are probed in **cuda -> hip -> metal -> opencl** order and the first one that enumerates a device wins. The order puts Metal ahead of OpenCL on macOS, where OpenCL is deprecated and only exposes CPU devices.
 
 ## Quick start
 
 Every commonly used API is importable straight from the `pygpgpu` top level — no need to remember submodule names:
 
 ```python
-from pygpgpu import kernel, gpu, funcs, float4       # @kernel / device builtins / dual-sided funcs / vector types
+from pygpgpu import kernel, funcs, float4             # @kernel / dual-sided funcs / vector types
 from pygpgpu import list_devices, get_device
-from pygpgpu import climport, cuimport, hipimport    # import hooks for handwritten .cl/.cu/.hip files
+from pygpgpu import kernel_import                    # the import hook for handwritten .cl/.cu/.hip files
 ```
 
 The submodule paths (`pygpgpu.kernel` / `pygpgpu.vec_types` / `pygpgpu.opencl` / ...) still work. Note that top-level `pygpgpu.kernel` is the `@kernel` decorator; to reach things inside the submodule write `from pygpgpu.kernel import ...`.
 
-The three import hooks load lazily: the corresponding backend is only loaded when you actually use the hook, so `import pygpgpu` itself needs no GPU runtime at all.
+`kernel_import` needs no GPU runtime to be imported or installed, and a merged kernel loads only its own backend's runtime when it is first called.
 
 ### The @kernel decorator
 
 ```python
 import math
 import numpy as np
-from pygpgpu import kernel, gpu, funcs, float4
+from pygpgpu import kernel, funcs, float4
 
 def norm2(v: float4) -> float:
     return funcs.dot(v, v)
@@ -161,18 +164,69 @@ def matrix_scale(m: np.ndarray, out: np.ndarray, k: float) -> None:
     for y, x in np.ndindex(out.shape):    # 2D NDRange
         out[y, x] = m[y, x] * k
 
-@kernel(backend="cuda")
+@kernel(backend="cuda", domain="data")
 def histogram(data: np.ndarray, counts: np.ndarray) -> None:
-    gpu.work_size(data)                   # iteration domain for a gid-style kernel
-    i = gpu.get_global_id(0)
-    gpu.atomic_add(counts, data[i], 1)    # atomics, results read back automatically
+    i = get_global_id(0)
+    atomic_add(counts, data[i], 1)    # atomics, results read back automatically
 ```
 
 Supported: scalar/array/vector arguments, if/while/for-range, arithmetic/comparison/bit ops, `math.*` and `np.*` math functions, class methods (`self.xxx` attributes become arguments automatically), static methods, and calling helpers through a name — a module-level function, `module.func(...)`, `SomeClass.static_method(...)`, and `self.method(...)` for another method of the same class. Unsupported syntax raises a `SyntaxError` carrying the line number.
 
+### `+=` into an array element
+
+`arr[i] op= value` follows Python's semantics even when several work items meet at the same element: `bins[image[i, j, k]] += 1` accumulates every pixel, without writing `atomic_add` by hand. The translator emits a device atomic whenever it cannot prove the element belongs to one work item — the proof being that the subscript is built from exactly the magic loop's own index variables, so `out[i] += x[i]` stays a plain (faster) read-modify-write.
+
+```python
+@kernel
+def histogram(image: np.ndarray):
+    bins: np.ndarray = np.zeros((3, 256), dtype=np.int32)   # created on the host
+    for i, j, k in np.ndindex(image.shape):
+        bins[k, image[i, j, k]] += 1    # one atomic add per pixel: no lost counts
+    return bins
+```
+
+Only operators with a device atomic counterpart are automated (`+=`, `-=`, `&=`, `|=`, `^=`). A racy `*=`/`/=`/`**=`, or a racy accumulation into a float element, raises a `SyntaxError` instead of losing data: there is no portable float atomic on the OpenCL backend, so use an integer buffer or accumulate per work item and reduce afterwards (`examples/09_reduction.py`).
+
+### Accumulating into a single scalar
+
+A handwritten kernel takes a pointer to one element (`__global int* total`) and passes it to `atomicAdd`. There are three equivalent shapes here:
+
+```python
+from pygpgpu import atomic_add, kernel
+
+# 1) a 1-element array IS that pointer
+@kernel
+def count(x: np.ndarray, total: np.ndarray):
+    for i in range(x.shape[0]):
+        if x[i] > 0.0:
+            atomic_add(total, 0, 1)     # or simply: total[0] += 1
+
+# 2) an outer scalar declared nonlocal/global becomes a device pointer,
+#    and the result is written back to the enclosing scope after the call
+@kernel
+def count_global(x: np.ndarray):
+    global TOTAL
+    for i in range(x.shape[0]):
+        if x[i] > 0.0:
+            TOTAL += 1                  # atomic on the device, native in the interpreter
+
+# 3) a counter created in the body can be returned as the result
+@kernel
+def count_and_return(x: np.ndarray):
+    total: np.ndarray = np.zeros(1, dtype=np.int32)
+    for i in range(x.shape[0]):
+        if x[i] > 0.0:
+            total[0] += 1
+    return total                        # a 1-element array: the caller reads [0]
+```
+
+All three are exact on every backend — thousands of work items incrementing one location lose nothing. An outer scalar is always a *shared* location, so `TOTAL += 1` there is always atomic; `atomic_add(scalar, 0, 1)` is refused with a message pointing at that spelling, since an atomic needs an array to index.
+
+
+
 Type annotations are optional: the snippets in this README carry them for readability, but the device types come from the arguments at call time, so the same code works unannotated (`self` is not a kernel argument, so annotating it changes nothing).
 
-A kernel needs an iteration domain for its work size to be derived. Either use a magic loop (`for i in range(x.shape[0])`, whose bound is the domain) or, when dispatching on `gpu.get_global_id`, declare `gpu.work_size(x)` — that statement is host-side only, erased by the translator, and contributes all of `x`'s dimensions. A `k[(total, local)]` bracket also works. Without any of these the call raises instead of guessing.
+A kernel needs an iteration domain for its work size to be derived. Either use a magic loop (`for i in range(x.shape[0])`, whose bound is the domain) or, when dispatching on `get_global_id`, name the array in the decorator: `@kernel(domain="x")`. The domain contributes all of `x`'s dimensions, and a `k[(blocks, threads)]` bracket also works. Without any of these the call raises instead of guessing.
 
 ### Printing from a kernel
 
@@ -245,14 +299,30 @@ The two kernel styles keep their meanings, which is where the interpreter differ
 
 | Kernel style | What `k["python"](...)` does |
 |---|---|
-| magic loop (`for i in range(x.shape[0])`, `enumerate(x)`, `np.ndindex`/`np.ndenumerate`) | runs the function once: the loop already covers every element. A `(total, local)` bracket changes nothing here, because the loop, not the launch, decides the extent |
-| gid-style (`gpu.get_global_id`), domain from `gpu.work_size(x)` | runs the body once per work item of the declared domain (`x`'s full shape), with `gpu.get_global_id` reporting the current index — one serial step per element, in `np.ndindex` order |
-| gid-style, domain from a `k[(total, local)]` bracket | same, one step per work item of the bracket, so the body must clamp itself against the real element count exactly as it does on a GPU |
+| magic loop (`for i in range(x.shape[0])`, `enumerate(x)`, `np.ndindex`/`np.ndenumerate`) | runs the function once: the loop already covers every element. A `(blocks, threads)` bracket changes nothing here, because the loop, not the launch, decides the extent |
+| gid-style (`get_global_id`), domain from `@kernel(domain="x")` | runs the body once per work item of the declared domain (`x`'s full shape), with `get_global_id` reporting the current index — one serial step per element, in `np.ndindex` order |
+| gid-style, domain from a `k[(blocks, threads)]` bracket | same, one step per work item of the bracket, so the body must clamp itself against the real element count exactly as it does on a GPU |
 
-Inside such a body `gpu.*` resolves to host implementations: the indexing builtins, the atomics (plain read-modify-writes — serial execution makes them exact), `gpu.barrier()` (a no-op) and the vector math family, which uses `pygpgpu.funcs`' host implementation. Two things cannot be emulated and raise instead: `gpu.local` (group-local memory is shared by parallel work items) and a `gpu.get_global_id` call inside a magic-loop body, where there is no current work item — use the loop variable there.
+Inside such a body the same functions resolve to their host implementations: the indexing builtins, the atomics (plain read-modify-writes — serial execution makes them exact), `barrier()` (a no-op) and the vector math family, which uses `pygpgpu.funcs`' host implementation. Two things cannot be emulated and raise instead: `local` (group-local memory is shared by parallel work items) and a `get_global_id` call inside a magic-loop body, where there is no current work item — use the loop variable there.
 
 Vector-typed arrays work unchanged: reading `v = points[i]` yields a vector instance, and `out[i] = v.xyzy * 2.0` writes its components.
 
+### One body, two behaviours: `KERNEL_TRANSLATING`
+
+The interpreter runs the body as written, so a kernel that needs *different* code on each side has nowhere to put the difference — and a plain `if` on a value the device cannot see would be refused. `KERNEL_TRANSLATING` is the escape hatch, borrowed from `typing.TYPE_CHECKING`: it is `False` in Python and `True` while the body is being translated, so the condition is decided at **translation time** and exactly one branch is emitted:
+
+```python
+@kernel(domain="out")
+def measure(data: np.ndarray, out: np.ndarray):
+    i = get_global_id(0)
+    if KERNEL_TRANSLATING:
+        out[i] = data[i] * 2.0          # device: the real computation
+    else:
+        out[i] = len(str(i))            # python: a host-only expression the device could not run
+```
+
+The branch is inlined as a macro — no `if` survives into the generated source — and the dropped branch is never translated, so Python-only work costs the device nothing. Only the name is recognized, so a kernel may declare its own `KERNEL_TRANSLATING` instead of importing it.
+
 ### Handwritten kernel files
 
 Already have C-style kernels? `.cl` / `.cu` / `.hip` files (native syntax of the matching backend) import like Python modules. Write a `vecadd.cl`:
@@ -268,14 +338,16 @@ __kernel void vecadd(__global const float* a, __global const float* b,
 ```
 
 ```python
-from pygpgpu import climport             # cuimport / hipimport work the same way
+from pygpgpu import kernel_import
 import numpy as np
-import vecadd                          # just like importing a Python module
+
+kernel_import.install()
+import vecadd                          # one .cl file = a module with one backend
 
 a = np.random.rand(1_000_000).astype(np.float32)
 out = np.zeros_like(a)
-vecadd.vecadd(a, a, out, np.int32(len(a)))    # work size inferred
-vecadd.vecadd[(1_000_192, 256)](a, a, out, np.int32(len(a)))   # explicit (total, per group)
+vecadd.vecadd(a, a, out, np.int32(len(a)))        # work size inferred
+vecadd.vecadd[(3907, 256)](a, a, out, np.int32(len(a)))   # (blocks, threads)
 ```
 
 Note the `n` argument: a handwritten kernel must clamp against the real element count, because the
@@ -284,6 +356,114 @@ size, not the buffer length, so it is not a usable bound.
 
 Struct definitions map to ctypes types; a `.pyi` stub is generated next to the kernel so IDE completion works.
 
+### One kernel, several backends
+
+Same-named `.cl` / `.cu` / `.hip` files are ONE module. Put `vecops.cl`, `vecops.cu` and
+`vecops.hip` in a package directory (with an `__init__.py`), then:
+
+```python
+from pygpgpu import kernel_import
+
+kernel_import.install()
+from kernel_import import vecops        # the package that holds the three sources
+
+vecops.add_offset(a, b, out, offset, n)      # default device, like a @kernel call
+vecops.add_offset["cuda"](a, b, out, ...)    # a bracket names the backend
+vecops.add_offset["opencl:1"][(blocks, threads)](...)   # device + work size, either order
+```
+
+What makes one module out of three files:
+
+* a kernel declared with the **same parameter list** in several files is one function with one
+  implementation per backend -- the bracket picks the implementation, the call keeps the kernel's
+  own signature;
+* declared with **different parameter lists** it is a set of overloads -- the stub declares one
+  `@overload` per list, and the argument types pick the implementation at call time (a bare call
+  walks the same `cuda -> hip -> metal -> opencl` probe a `@kernel` uses);
+* declared in only one file it still exists, and a bracket naming a backend without an
+  implementation raises `KeyError` at the bracket.
+
+A `.pyi` stub is generated for the merged module: one `class Kernel_<name>` per kernel, one
+`__call__` per distinct parameter list, and the structs from the first file that declares each.
+
+### The scientific-computing vocabulary
+
+The device-side math functions are plain names, like every other builtin:
+
+```python
+from pygpgpu import erf, tgamma, fast_exp, normal_cdf
+
+@kernel
+def integrate(x: np.ndarray, out: np.ndarray):
+    for i in range(out.shape[0]):
+        out[i] = erf(x[i]) * tgamma(x[i] + 1.0)
+```
+
+Three groups, and the difference matters:
+
+* the functions **all three backends have** -- `erf`, `erfc`, `tgamma`, `lgamma`, `rint`,
+  `remainder`, `fmin`, `fmax`, `fdim`, `copysign`, `nextafter`, `ldexp`, `fract`, `sign`,
+  and the `native_*` family. These run on OpenCL, CUDA and HIP alike.
+* the functions **CUDA and HIP have** and OpenCL does not -- the fast approximations
+  (`fast_exp`, `fast_log`, `fast_sin`, ...), the correctly-rounded forms (`exact_recip`,
+  `exact_divide`, `exact_sqrt`, ...), the special functions (`erf_inverse`, `normal_cdf`,
+  `reciprocal_hypot`, `sin_pi`, `cos_pi`, the Bessel functions) and the conversions
+  (`float_to_int_nearest`, `int_to_float`, ...). A kernel that uses one of these is a
+  **CUDA or HIP kernel**: on OpenCL the translator refuses it *by name*, so the message
+  says which function has no OpenCL counterpart rather than leaving you with an
+  "undeclared identifier" from the compiler.
+* the `mad` / `fma` family and the vector operations, which pygf already publishes and
+  which are pygf's own objects here (see `pygpgpu.funcs`).
+
+CUDA and HIP answer the same thing by the same name, but they are not the same device and
+the difference is worth knowing:
+
+* HIP has no `__exp2f`, so `native_exp2` is spelled `exp2f` there -- the ordinary name,
+  which HIP does have.
+* HIP has **none** of the round-toward-zero intrinsics (`__frcp_rz`, `__fdiv_rz`,
+  `__fsqrt_rz`), so `exact_recip_rz`, `exact_divide_rz` and `exact_sqrt_rz` are computed
+  on HIP by a device helper the package emits: it takes the correctly rounded value, uses
+  the fused multiply-add to learn which side of it the exact result is on, and steps one
+  unit toward zero when it is the other way. The result agrees with CUDA's own intrinsic
+  bit for bit (a division is the one exception, because the two hardware divide units
+  round differently in the last bit -- the rounding *direction* is the same either way).
+
+Two names are worth a warning, because the device's documentation is the authority and
+the spelling does not tell you:
+
+* CUDA's `rhypotf` is the **reciprocal** of the hypotenuse (`1/sqrt(x²+y²)`), so the name
+  here is `reciprocal_hypot`; the plain hypotenuse is `hypot`.
+* CUDA's `sinpif` is `sin(πx)`, **not** the sinc function -- so the name here is `sin_pi`.
+
+### Naming a device function that already exists
+
+A device function that is not one of the builtins -- a `__device__` function in a header
+you ship, or a compiler intrinsic with a spelling of your own -- is registered:
+
+```python
+from pygpgpu import device_function
+
+@device_function(cuda="my_fast_rcp", opencl="native_recip",
+                 headers={"cuda": "myclib.h"})
+def my_fast_rcp(x: float) -> float:
+    """The host's implementation; the device gets the symbol the registry names."""
+    return 1.0 / x
+
+@kernel
+def use_it(x: np.ndarray, out: np.ndarray):
+    for i in range(out.shape[0]):
+        out[i] = my_fast_rcp(x[i])
+```
+
+The device calls the symbol the table names; the header is copied into the cache and
+included in the generated source, so you do not have to know where the cache is (a header
+given as a full `#include` line is carried into the source verbatim instead, which is what
+a system header such as `#include <cuda_fp16.h>` needs). A backend with no symbol for the
+function is refused by name, the same way the CUDA-only builtins are.
+
+The body is the **host's** implementation -- the `python` backend runs it -- so the two
+sides have to agree, exactly as they do for the builtins.
+
 ### Class methods
 
 ```python
@@ -459,7 +639,7 @@ wp.launch(saxpy, dim=n, inputs=[x, y, out, 2.0])
 | Input data | numpy arrays directly | wrap in `ti.field`, explicit `from_numpy`/`to_numpy` | `cuda.to_device` + `copy_to_host` by hand | wrap in `wp.array` |
 | Helper functions | **plain Python — no decorator, annotations optional** | `@ti.func` + annotated args | `@cuda.jit(device=True)` | `@wp.func` + annotated args |
 | Backends in one process | **all three, mixed freely** | one, fixed by `ti.init(arch=...)` | CUDA only | CUDA only |
-| Launch size | inferred from the loop bound (or `gpu.work_size(x)` / a bracket) | inferred from `range` | blocks × threads by hand, plus a guard | pass `dim=` |
+| Launch size | inferred from the loop bound (or `domain="x"` / a bracket) | inferred from `range` | blocks × threads by hand, plus a guard | pass `dim=` |
 | Bounds guard | emitted for you | emitted for you | you write it | you write it |
 | Backends | OpenCL, CUDA, HIP | CPU, CUDA, Vulkan and others | CPU, CUDA | CPU, CUDA |
 | AMD GPU | HIP (native AMD codegen) and OpenCL | via Vulkan | no | no |
@@ -484,7 +664,7 @@ apply_all["cuda"](x, out)      # NVIDIA
 apply_all["hip"](x, out)       # AMD
 apply_all["opencl"](x, out)    # whichever device OpenCL enumerates
 apply_all["python"](x, out)    # the interpreter, on the CPU (no device needed)
-apply_all(x, out)              # unqualified: probe cuda -> hip -> opencl
+apply_all(x, out)              # unqualified: probe cuda -> hip -> metal -> opencl
 ```
 
 Each backend/device pair is specialized and cached separately, so alternating between them is just
@@ -496,11 +676,86 @@ Where the others are ahead — worth knowing before you pick:
 - **Automatic differentiation**: Taichi and Warp can differentiate kernels; this library cannot.
 - **Sparse and irregular data structures**: Taichi has sparse grids and SNode trees; here you pass dense numpy arrays (or handle sparsity yourself).
 - **CPU fallback**: all three can run on the CPU; this library runs on the CPU only through the explicit `k["python"]` interpreter backend, which is serial and meant for testing rather than throughput.
-- **Minimum interpreter version**: this library runs on Python 3.7 and newer, on both 64- and 32-bit Windows builds; the others require a newer interpreter.
+- **Minimum interpreter version**: this library runs on Python 3.7 and newer, on 64- and 32-bit Windows and on Linux x86-64; the others require a newer interpreter. Linux is exercised end to end under WSL (CUDA through the `/usr/lib/wsl` driver passthrough, OpenCL through a runtime the distribution ships), so the driver-name and library-discovery paths there are the real ones rather than an untested branch.
 
 This table was written from each project's documentation and the snippets above; their APIs move
 fast, so confirm against current docs before relying on a specific cell.
 
+### Return values and outer variables
+
+A kernel runs once per work item, so there is no single C return value to hand back. The
+translator gives every returned position a device parameter of its own — a one-element
+buffer for a scalar or vector, the array object itself for an array — and the first work
+item to reach a `return` claims the result with an atomic compare-and-swap. The host reads
+the slots back after the launch and returns exactly what the winning work item returned:
+
+```python
+@kernel
+def first_peak(x: np.ndarray):
+    for i in range(x.shape[0]):
+        if x[i] > 10.0:
+            return x[i], i          # a tuple
+    return None                     # a launch where nothing returns is None
+
+value, index = first_peak[backend](x)
+```
+
+The supported return types are exactly the supported parameter types, every `return`
+statement may have its own shape (`return x[i]` from one branch, `return x[i], i` from
+another), and a bare `return` gives `None`. `submit()` and `async_call()` deliver the same
+value once their future resolves, after read-back. A value-returning `return` is refused
+below the magic loop: there it would return one work item's local, not the function's
+result.
+
+Variables the body reads from outside become parameters too, so a changed outer value is
+visible on the next call — inlining it would leave the kernel stale:
+
+```python
+SCALE = 2.0
+
+@kernel
+def scaled(x: np.ndarray, out: np.ndarray):
+    for i in range(out.shape[0]):
+        out[i] = x[i] * SCALE
+
+scaled[backend](x, out)     # SCALE is passed as an argument
+SCALE = 100.0               # the next call uses 100.0
+scaled[backend](x, out)
+```
+
+A `nonlocal`/`global` declaration makes the outer scalar writable: it travels as a pointer,
+every use inside is a dereference, and the value is copied back into its scope after the
+call — a closure cell for `nonlocal`, the module globals for `global`. Outer arrays are
+pointer parameters as well, and are read back when the kernel writes them. Rebinding an
+outer array (as opposed to writing its elements) is refused, because the object is shared
+with the caller. A helper function gets the same treatment as the kernel that calls it.
+
+Statements above the magic loop may also *create* data with numpy — `np.zeros`, `np.ones`,
+`np.empty`, `np.full`, `np.arange`, `np.linspace`, `*_like`, `np.float32(...)` and the rest
+of the allocating constructors. The interpreter runs those statements, so the array has
+exactly the shape and dtype numpy gives it, and it reaches the kernel as a parameter:
+
+```python
+@kernel
+def normalized(x: np.ndarray, out: np.ndarray):
+    total = np.zeros(1, dtype=np.float32)      # built on the host
+    for i in range(x.shape[0]):
+        total[0] = total[0] + x[i]
+    for i in range(out.shape[0]):
+        out[i] = x[i] / total[0]
+
+@kernel
+def doubled(x: np.ndarray):
+    result = np.zeros_like(x)                  # created, then returned as itself
+    for i in range(x.shape[0]):
+        result[i] = x[i] * 2.0
+    return result
+```
+
+A created object the body returns reuses the parameter it already travels through, so the
+caller receives the host's own object rather than a copy. Creation below the magic loop is
+refused: it would build one object per work item instead of one.
+
 ### Images
 
 `Image` is one image type for every backend. A kernel indexes it like an array, and the
@@ -561,6 +816,97 @@ The HIP backend refuses images: it reuses CUDA code generation, which has no tex
 object type, and the driver rejects `hipTexObjectCreate` on hardware without texture
 support.
 
+## Host-side CUDA libraries
+
+The CUDA *driver* (`pygpgpu/cuda/driver`) is the low-level API: contexts, streams,
+modules, memory. The CUDA *math libraries* ship as separate DLLs next to it and do
+the actual mathematics. They are bound with ctypes under `pygpgpu/cuda/libraries`, in
+the same shape as the driver — a lazily loaded DLL, a signature table per library, and
+an error code turned into a `RuntimeError` carrying the library's own message:
+
+```python
+from pygpgpu.cuda.libraries import driver_for
+
+CUBLAS = driver_for("cublas")          # the cuBLAS entry points
+CUFFT  = driver_for("cufft")           # the cuFFT entry points
+```
+
+| Library | Bound operations |
+|---|---|
+| **cuBLAS** | `gemm` (float32/float64), `axpy`, `scal`, `dot`, `GemmEx`, handle and stream management |
+| **cuBLASLt** | the generic `matmul` with a heuristically chosen algorithm, matrix layouts, epilogues |
+| **cuFFT** | `plan1d`, `exec` (R2C / C2R / C2C), plan destruction, stream binding |
+| **cuSPARSE** | `csrmv` (createCsr / createDnVec / SpMV), dense and sparse descriptors |
+| **cuSOLVER** | `potrf` (Cholesky), `gesvd` (SVD), `getrf` (LU), with their workspace queries |
+| **CUDA runtime** | `cudaMalloc`, `cudaFree`, `cudaMemcpy`, `cudaMemset`, the device count and properties |
+
+A thin, numpy-friendly layer sits on top of the raw entry points, so a call takes
+arrays rather than device pointers:
+
+```python
+import numpy as np
+from pygpgpu.cuda_libraries import gemm, axpy, dot, rfft, csrmv
+
+a = np.random.rand(3, 4).astype(np.float32)
+b = np.random.rand(4, 2).astype(np.float32)
+
+c = gemm(a, b)                  # a @ b, on the GPU
+d = axpy(2.0, a[0], b[0])       # elementwise
+e = dot(a[0], b[0])             # a scalar
+f = rfft(a[0])                  # n/2+1 complex numbers, like np.fft.rfft
+g = csrmv(row_offsets, column_indices, values, x)   # a sparse product
+```
+
+The same names are reachable as attributes of the module, so both spellings work:
+
+```python
+from pygpgpu import cuda_libraries
+cuda_libraries.gemm(a, b)
+
+from pygpgpu.cuda_libraries import gemm
+gemm(a, b)
+```
+
+The wrappers take numpy's **row-major** layout and return the same layout; cuBLAS
+reads column-major, and the difference is handled inside, so `gemm(a, b)` is just
+`a @ b`. The 1-D cases follow numpy's own rules for `@`, so a matrix and a vector
+multiply without the caller reshaping anything.
+
+**A CUDA context has to be current.** A library handle belongs to the context that
+was current when it was created, so the wrappers make one current if none is — but
+if you use the raw entry points, make one current yourself first:
+
+```python
+from pygpgpu import get_device
+
+device = get_device("cuda:0")
+device.default_context.make_current()      # before the first library call
+```
+
+Without it the call fails with `CUDA_ERROR_INVALID_CONTEXT`, and the message says
+so.
+
+**These are host-side operations, not device builtins.** A `gemm` moves whole arrays
+to the device and runs one library call; a `@kernel` body is per-thread arithmetic.
+That is why a `@kernel` cannot call `cuda_libraries.gemm` — the two are different
+layers, and mixing them would mean a kernel launching a library call per work item.
+The device-side counterparts (a per-element multiply, a per-element dot) are the
+device builtins the top level already exports.
+
+The libraries are **loaded lazily**: `import pygpgpu` loads no CUDA library at all,
+and each library's DLL is loaded on the first call that needs it. A machine without
+the CUDA toolkit still imports the package, and the failure surfaces at the call with
+the loader's own error.
+
+**cuSOLVER needs kernels the consumer cards do not have.** NVIDIA does not ship
+cuSOLVER's dense factorization kernels for every architecture, and a card without
+them fails at execution rather than at setup: the call returns a failure code and the
+context reports an illegal memory access. Measured on the machine this was developed
+on (an RTX 4070 Laptop, compute capability 8.9, CUDA 12.6), `potrf`, `getrf`,
+`gesvd`, `gesvdj` and `syevj` all fail that way, with `compute-sanitizer` putting the
+fault inside cuSOLVER's own DLL. The bindings are correct and the test skips with
+that reason instead of failing.
+
 ## Device selection
 
 Three forms of device identifier inside the brackets (same on every backend):
@@ -583,7 +929,7 @@ dev = get_device("cuda:0")   # the Device object of that backend, directly
 
 ### Choosing the default device
 
-A call without a bracket runs on the first usable backend in `cuda -> hip -> opencl`.
+A call without a bracket runs on the first usable backend in `cuda -> hip -> metal -> opencl`.
 `set_default_device()` replaces that order:
 
 ```python
@@ -609,6 +955,7 @@ is. A kernel decorated with an explicit `@kernel(backend=...)` keeps that backen
 | OpenCL | `__clcache__/` | parsed meta + per-device binaries |
 | CUDA | `__cucache__/` | parsed meta + per-device cubin/ptx |
 | HIP | `__hipcache__/` | parsed meta + per-device code objects |
+| Metal | `__metalcache__/` | generated MSL source, compiled at first launch |
 | Python | none | nothing is generated or compiled; the body runs as-is |
 
 Caches are invalidated by the mtime of the source file (and every `#include`); sources generated by `@kernel` are keyed by a hash of the Python source. These directories should be in `.gitignore`.
@@ -637,15 +984,15 @@ set_print_info(True)       # verbose again (the default)
 
 ## Known limitations
 
-- `@kernel` does not support: strings/lists/dicts, closures, recursion, kernel return values, dynamic multi-dimensional shape indexing;
+- `@kernel` does not support: strings/lists/dicts, recursion, dynamic multi-dimensional shape indexing. Closures and module globals *are* supported: an outer variable becomes a kernel parameter, and one declared `nonlocal`/`global` is written back after the call (see "Return values and outer variables" above);
 - `print()` inside a kernel prints one `True`/`False` pair, integer, real or vector at a time. `sep` and `end` are translated, but must be string literals (the format string is fixed in the device code); `file` and `flush` are refused; it also cannot print an array (print an element instead). A line containing many floats expands into a `switch` over their digit counts, so an extremely long line can hit the code-size guard and should be split into several calls. The HIP backend rejects `print()` outright: its driver drops the device on such kernels on the tested runtime;
 - vector comparison, writing through a swizzle, and nesting beyond three operands are restricted in favour of cross-backend consistency (the error message suggests a workaround);
-- `Image` is unsupported on the HIP backend: HIP reuses CUDA code generation, which has no texture object type, and `hipTexObjectCreate` is rejected with "operation not supported" on hardware without texture support (an integrated Radeon under ROCm 7.2 Windows, for instance). The translator raises a `PyKernelSyntaxError` naming the backend; use cuda, opencl or python for image kernels;
+- `Image` is unsupported on the HIP backend: HIP reuses CUDA code generation, which has no texture object type, and `hipTexObjectCreate` is rejected with "operation not supported" on hardware without texture support (an integrated Radeon under ROCm 7.2 Windows, for instance). The translator raises a `SyntaxError` naming the backend; use cuda, opencl or python for image kernels;
 - type annotations on a `@kernel` function are accepted and ignored: the device type comes from the actual argument at call time, since one kernel is specialized per argument-type signature. Annotations are there for editors and type checkers, not for translation;
 - a `@kernel` signature may carry default values, and a call may pass arguments by keyword: the call is bound to positional order before launch, with omitted defaults filled in. `*args`, `**kwargs` and keyword-only parameters remain unsupported;
 - fp16 (`np.float16`) arguments are unavailable on the CUDA/HIP backends (use float32);
 - the HIP backend compiles with the clang shipped by the AMD HIP SDK; hipRTC on this machine (ROCm 7.2 Windows) fails on any input, so it is not used;
-- the pure-Python backend (`k["python"]`) does not emulate `gpu.local` (group-shared memory): it runs serially and has no parallel work items to share it.
+- the pure-Python backend (`k["python"]`) does not emulate `local` (group-shared memory): it runs serially and has no parallel work items to share it.
 
 On CUDA, the array APIs (`Array` upload/read-back through `cuMemcpy2D`) work on every supported interpreter. `cuSurfObjectCreate` and `cuTexObjectCreate` occasionally return `CUDA_ERROR_INVALID_VALUE` on a descriptor that is byte-for-byte valid; the failures cluster in the first few calls of a process and then stop. Measured on the machine this was developed on (driver 13.4, RTX 4070): about one in a hundred creations fails, retrying an already-failed descriptor does not recover it, and a fresh descriptor for a newly created array usually does. Because an `Image` binds one surface and one texture per launch, a scenario that launches several image kernels in one process can therefore lose a few launches. The `python` and OpenCL backends are unaffected: OpenCL images are created per launch on a code path that has not shown the failure, and the `python` backend touches no device at all. `examples/11_texture2d.py` reports a driver refusal instead of raising.
 ## License