Metadata-Version: 2.4
Name: pygpgpu
Version: 0.1.1
Summary: Import GPU kernels like Python functions across OpenCL / CUDA / HIP backends
Author: Time-Coder
License: MIT License
        
        Copyright (c) 2026 Time-Coder
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://github.com/Time-Coder/PyGPGPU
Keywords: gpu,opencl,cuda,hip,rocm,kernel,jit
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: Microsoft :: Windows
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.7
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.7
Description-Content-Type: text/markdown
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: typing_extensions>=4.4.0
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: ruff>=0.9; extra == "dev"
Requires-Dist: ty>=0.0.60; extra == "dev"
Dynamic: license-file

# PyGPGPU

**English** | [Chinese](README.zh-CN.md)

Decorate a Python function and it runs on the GPU — one codebase, three backends (**OpenCL / CUDA / HIP**).

```python
import numpy as np
from pygpgpu import kernel

@kernel
def saxpy(x, y, out, alpha):
    for i in range(out.shape[0]):
        out[i] = x[i] * alpha + y[i]

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["cuda"](x, x, out, 2.0)      # bracket switches the backend
saxpy["hip:gfx1036"](x, x, out, 2.0)   # backend:device
```

No C code to write, no build system to configure; the first call translates, compiles and caches, every later call is free. Existing handwritten kernels are fine too — `.cl`/`.cu`/`.hip` files can be `import`ed like Python modules.

## Features

- **`@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 no type annotations: it is translated into a device function and monomorphized from the call site, recursively for helpers that call helpers.
- **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.
- **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.
- **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, aliases ocl/cu/rocm/amd)
  k["opencl:gfx1036"](args)        # backend:device
  k[(1024, 256)](args)             # work size: (total, per group), None = automatic
  k["cuda"][(1024, 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.
- **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.
- **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.

## Installation

```bash
pip install pygpgpu
```

Python ≥ 3.7. Bring your own GPU runtime (any one of them is enough):

| Backend | What you need | Notes |
|---|---|---|
| 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) |

When `@kernel` has no explicit `backend`, backends are probed in **cuda -> hip -> opencl** order and the first one that enumerates a device wins.

## 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 available_devices, get_device
from pygpgpu import climport, cuimport, hipimport    # import hooks 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.

### The @kernel decorator

```python
import math
import numpy as np
from pygpgpu import kernel, gpu, funcs, float4

def norm2(v):
    return funcs.dot(v, v)

@kernel
def normalize_all(vecs, out):
    for i, v in enumerate(vecs):          # 1D: index + float4 value snapshot
        out[i] = funcs.normalize(v)       # expanded component-wise on device

@kernel
def matrix_scale(m, out, k):
    for y, x in np.ndindex(out.shape):    # 2D NDRange
        out[y, x] = m[y, x] * k

@kernel(backend="cuda")
def histogram(data, counts):
    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
```

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. Unsupported syntax raises a `SyntaxError` carrying the line number.

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.

### 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`:

```c
__kernel void vecadd(__global const float* a, __global const float* b,
                     __global float* out, const int n)
{
    int i = get_global_id(0);
    if (i >= n) return;
    out[i] = a[i] + b[i];
}
```

```python
from pygpgpu import climport             # cuimport / hipimport work the same way
import numpy as np
import vecadd                          # just like importing a Python module

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)
```

Note the `n` argument: a handwritten kernel must clamp against the real element count, because the
launch size is rounded up to a multiple of the local size. `get_global_size(0)` reports that rounded
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.

### Class methods

```python
class Model:
    def __init__(self):
        self.weights = np.random.rand(1024).astype(np.float32)
        self.lr = 0.01

    @kernel
    def apply(self, x, out):
        for i in range(out.shape[0]):
            out[i] = x[i] * self.weights[i] + self.lr

m = Model()
m.apply(x, out)                     # self.weights/self.lr are passed as kernel arguments
m.apply["cuda"](x, out)             # brackets still available
```

### Copy optimization

Each pointer argument goes through function-body-level static analysis (`points[i] = ...` is a write, `x = a[i]` is a read, plus pointer-alias propagation and interprocedural analysis across functions):

```python
@kernel
def k(a, b, out):
    for i in range(out.shape[0]):
        out[i] = a[i] + b[i] * 2.0   # a read-only, b read-only, out write-only
# at runtime: a/b skip the read-back after upload, out skips the upload and only reads back
```

Combined with `pygpgpu.numpy` (gnp) arrays, data stays entirely on the device between consecutive kernel calls and is synchronized only when the host touches it.

## How it compares

The same computation, written in this library and in three other GPU-from-Python projects
([Taichi](https://www.taichi-lang.org/), [Numba](https://numba.pydata.org/),
[NVIDIA Warp](https://github.com/NVIDIA/warp)). Snippets are cut down to the parts that
differ — check each project's docs for exact syntax. The examples also show a shared helper
function, which is where the API gap is widest: here it stays ordinary Python, while the
others need a second decorator and typed signatures.

**PyGPGPU** — arrays in, arrays out; no launch geometry, no buffer management, and helper
functions need no decorator at all:

```python
from pygpgpu import kernel

def square(v):                          # plain Python: no decorator, no annotations
    return v * v

def clamp01(v):
    if v < 0.0:
        return 0.0
    if v > 1.0:
        return 1.0
    return v

def weight(v):                          # helpers may call other helpers
    return clamp01(square(v)) * 0.5 + 0.25

@kernel
def apply_all(x, out):
    for i in range(out.shape[0]):       # loop bound = work size
        out[i] = weight(x[i])

apply_all(x, out)                       # numpy arrays as-is
apply_all["hip"](x, out)                # the same function, another backend
```

**Taichi** — declare fields first, then copy data in and out by hand; a helper needs `@ti.func`
and every argument must be annotated:

```python
import taichi as ti
ti.init(arch=ti.gpu)

x = ti.field(ti.f32, n)
y = ti.field(ti.f32, n)
out = ti.field(ti.f32, n)
x.from_numpy(nx)
y.from_numpy(ny)

@ti.func                                  # helpers need their own decorator
def weight(v: ti.f32) -> ti.f32:
    return ti.max(0.0, ti.min(1.0, v * v)) * 0.5 + 0.25

@ti.kernel
def saxpy(a: ti.f32):
    for i in range(n):
        out[i] = x[i] * a + y[i] + weight(x[i])

saxpy(2.0)
res = out.to_numpy()
```

**Numba** — compute blocks/threads yourself, guard inside, manage device memory; callable helpers
are compiled separately with `@cuda.jit(device=True)`:

```python
from numba import cuda

@cuda.jit(device=True)                    # helpers need device=True
def weight(v):
    return min(1.0, max(0.0, v * v)) * 0.5 + 0.25

@cuda.jit
def saxpy(x, y, out, a):
    i = cuda.grid(1)
    if i < out.size:                      # bounds guard is your job
        out[i] = x[i] * a + y[i] + weight(x[i])

d_x = cuda.to_device(nx)
d_y = cuda.to_device(ny)
d_out = cuda.device_array_like(nx)
threads = 256
saxpy[(n + threads - 1) // threads, threads](d_x, d_y, d_out, 2.0)
res = d_out.copy_to_host()
```

**Warp** — typed array parameters, the launch size is passed explicitly, and helpers need `@wp.func`:

```python
import warp as wp
wp.init()

@wp.func                                  # helpers need their own decorator
def weight(v: float) -> float:
    return wp.min(1.0, wp.max(0.0, v * v)) * 0.5 + 0.25

@wp.kernel
def saxpy(x: wp.array(dtype=float), y: wp.array(dtype=float),
          out: wp.array(dtype=float), a: float):
    i = wp.tid()
    out[i] = x[i] * a + y[i] + weight(x[i])

wp.launch(saxpy, dim=n, inputs=[x, y, out, 2.0])
```

| | PyGPGPU | Taichi | Numba | Warp |
|---|---|---|---|---|
| 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, no annotations** | `@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=` |
| 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 |
| Switching backend | bracket `k["hip"]`, same source | `ti.init(arch=...)`, same source | not applicable | not applicable |
| Install footprint | Python + ctypes only: no LLVM, no CUDA Toolkit | ships its own runtime | needs LLVM | needs the CUDA Toolkit |
| Handwritten `.cl`/`.cu`/`.hip` | imported as Python modules, `.pyi` generated | no | via the `cuda` module's PTX/cubin path | no |
| Static read/write analysis | yes — uploads/read-backs skipped per argument | no | no | no |

### Several backends at the same time

Taichi fixes its backend once, at `ti.init(arch=...)`; Numba and Warp only target CUDA. Here the
bracket picks the backend per call, so one process — and even one kernel object — can use every
backend the machine has:

```python
@kernel
def apply_all(x, out):
    for i in range(out.shape[0]):
        out[i] = square(x[i]) + 1.0

apply_all["cuda"](x, out)      # NVIDIA
apply_all["hip"](x, out)       # AMD
apply_all["opencl"](x, out)    # whichever device OpenCL enumerates
apply_all(x, out)              # unqualified: probe cuda -> hip -> opencl
```

Each backend/device pair is specialized and cached separately, so alternating between them is just
an extra dictionary lookup — no recompilation after the first call. This is useful for comparing a
kernel across vendors on one machine (`examples/08_compare_backends.py`), or for falling back to a
second GPU when the first is busy.

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 needs a GPU.
- **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.

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.

## Device selection

Three forms of device identifier inside the brackets (same on every backend):

| Form | Meaning |
|---|---|
| `"0"`, `"1"` | index in the backend's device list (OpenCL uses flat numbering across platforms) |
| `"gfx1036"`, `"geforce"` | case-insensitive substring match on name/vendor/**architecture name** |
| `"geforce:1"` | when a name matches several devices, take the n-th one |

Helpers:

```python
from pygpgpu import available_devices, get_device

available_devices()          # every available "backend:device" string, e.g. ['cuda:0', 'hip:0', 'opencl:0']
dev = get_device("cuda:0")   # the Device object of that backend, directly
```

## Caching

| Backend | Directory | Contents |
|---|---|---|
| OpenCL | `__clcache__/` | parsed meta + per-device binaries |
| CUDA | `__cucache__/` | parsed meta + per-device cubin/ptx |
| HIP | `__hipcache__/` | parsed meta + per-device code objects |

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`.

## Known limitations

- `@kernel` does not support: strings/lists/dicts, closures, recursion, kernel return values, dynamic multi-dimensional shape indexing;
- 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);
- 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.

## License

MIT
