PyGPGPU

English | Chinese

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

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 imported 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):

  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

pip install -e .

Python ≥ 3.12. 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:

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

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:

__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];
}

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

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

@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, Numba, 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:

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:

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

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:

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:

@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 requires Python 3.12+ (it runs on 3.14); the others support a wider range of older versions.

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:

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