""" vgpu_cache.py — Universal VGPU persistent kernel cache. Now covers: • MatMul (existing) • Binary elementwise: add, sub, mul, div, pow • Unary elementwise: neg, abs, sqrt, square • Activations: relu, sigmoid, tanh, clip • Reductions: sum (CPU-only fallback on GPU - see notes) • Transparent numpy ↔ torch.Tensor interop • Dynamic dispatch: cache.add(a,b), cache.relu(x), cache.clip(x,0,1) etc. Works on all platforms: RPi5 (numpy), Linux/macOS/Windows (OpenGL). Author: Per Lindholm License: MIT Run: python vgpu_cache.py """ import numpy as np import os import sys import platform import time # ==================================================================== # PLATFORM DETECTION # ==================================================================== IS_RASPBERRY_PI = platform.machine().startswith('arm') and 'linux' in sys.platform IS_MACOS = platform.system() == 'Darwin' IS_WINDOWS = platform.system() == 'Windows' print(f"[VGPU] Platform: {platform.system()} {platform.machine()}") print(f"[VGPU] Raspberry Pi: {IS_RASPBERRY_PI}") # ==================================================================== # TORCH (optional) # ==================================================================== try: import torch TORCH_AVAILABLE = True except ImportError: torch = None TORCH_AVAILABLE = False print("[VGPU] torch not installed; only numpy tensors supported") # ==================================================================== # GPU BACKEND (optional) # ==================================================================== GPU_AVAILABLE = False GPU_BACKEND = None try: import moderngl if IS_RASPBERRY_PI: print("[VGPU] RPi5 detected: using CPU (numpy) backend") else: try: ctx_probe = moderngl.create_standalone_context() if hasattr(ctx_probe, 'compute_shader') and ctx_probe.compute_shader: GPU_AVAILABLE = True GPU_BACKEND = 'moderngl' print("[VGPU] ✅ GPU available (moderngl)") else: print("[VGPU] ⚠️ No compute shaders; using CPU") except Exception as e: print(f"[VGPU] ❌ GPU init failed: {e}; using CPU") except ImportError: print("[VGPU] moderngl not installed; using CPU (numpy)") if IS_MACOS and not GPU_AVAILABLE: print("[VGPU] macOS: Metal support not implemented yet (future work)") # ==================================================================== # HELPERS: tensor ↔ ndarray # ==================================================================== def _as_ndarray(x, dtype=None): """Convert torch.Tensor | np.ndarray | python iterable → np.ndarray.""" if TORCH_AVAILABLE and isinstance(x, torch.Tensor): arr = x.detach().cpu().numpy() else: arr = np.asarray(x) return arr.astype(dtype) if dtype else arr def _as_output_like(template, ndarray): """Wrap result back in the same container flavour as template.""" if TORCH_AVAILABLE and isinstance(template, torch.Tensor): return torch.from_numpy(ndarray).to(template.device) return ndarray # ==================================================================== # GLSL SHADER TEMPLATES # ==================================================================== # One shader per (op, arity). Shapes are pre-baked into the source. SHADER_TEMPLATES = { # ----------------- binary elementwise ----------------- 'add': """ #version 430 core layout(local_size_x = {LX}, local_size_y = 1) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = A[i] + B[i]; }} """, 'sub': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = A[i] - B[i]; }} """, 'mul': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = A[i] * B[i]; }} """, 'div': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = A[i] / B[i]; }} """, 'pow': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = pow(A[i], B[i]); }} """, # ----------------- unary elementwise ----------------- 'neg': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = -A[i]; }} """, 'abs': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = abs(A[i]); }} """, 'sqrt': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = sqrt(A[i]); }} """, 'square': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = A[i] * A[i]; }} """, # ----------------- activations ----------------- 'relu': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; float v = A[i]; C[i] = v > 0.0 ? v : 0.0; }} """, 'sigmoid': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = 1.0 / (1.0 + exp(-A[i])); }} """, 'tanh': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; C[i] = tanh(A[i]); }} """, 'clip': """ #version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int N; uniform float LO; uniform float HI; void main() {{ uint i = gl_GlobalInvocationID.x; if (i >= uint(N)) return; float v = A[i]; C[i] = clamp(v, LO, HI); }} """, } # CPU fallbacks (numpy name → callable) CPU_FALLBACKS = { 'add': lambda a, b: a + b, 'sub': lambda a, b: a - b, 'mul': lambda a, b: a * b, 'div': lambda a, b: a / b, 'pow': lambda a, b: a ** b, 'neg': lambda a: -a, 'abs': lambda a: np.abs(a), 'sqrt': lambda a: np.sqrt(a), 'square': lambda a: a * a, 'relu': lambda a: np.maximum(a, 0.0), 'sigmoid': lambda a: 1.0 / (1.0 + np.exp(-a)), 'tanh': lambda a: np.tanh(a), 'clip': lambda a, lo, hi: np.clip(a, lo, hi), } # Op arity: 'binary' | 'unary' OP_ARITY = { 'add': 'binary', 'sub': 'binary', 'mul': 'binary', 'div': 'binary', 'pow': 'binary', 'neg': 'unary', 'abs': 'unary', 'sqrt': 'unary', 'square': 'unary', 'relu': 'unary', 'sigmoid': 'unary', 'tanh': 'unary', 'clip': 'unary', } # ==================================================================== # GPU COMPILED-OP HOLDER (factory result, sits in the cache) # ==================================================================== class GPUOp: """Compiled shader + bound buffers for one (op, shape) combo.""" __slots__ = ('op', 'program', 'A_buf', 'B_buf', 'C_buf', 'gx', 'N', 'is_binary', 'extra_uniforms', 'ctx') def __init__(self, ctx, op, program, A_buf, B_buf, C_buf, gx, N, is_binary, extra_uniforms): self.ctx = ctx self.op = op self.program = program self.A_buf = A_buf self.B_buf = B_buf self.C_buf = C_buf self.gx = gx self.N = N self.is_binary = is_binary self.extra_uniforms = extra_uniforms # {'LO': 0.0, 'HI': 1.0} etc. def run(self, A, B=None): self.A_buf.write(A.astype(np.float32).tobytes()) self.A_buf.bind_to_storage_buffer(0) if self.is_binary: self.B_buf.write(B.astype(np.float32).tobytes()) self.B_buf.bind_to_storage_buffer(1) self.C_buf.bind_to_storage_buffer(2) # uniform 'N' gets set implicitly because it's baked into the source, # but extra uniforms need explicit set: for k, v in self.extra_uniforms.items(): self.program[k] = v self.program.run(self.gx, 1, 1) return np.frombuffer(self.C_buf.read(), dtype=np.float32).reshape(-1).copy() # ==================================================================== # CPU Matmul with Caching (existing — preserved for backwards compat) # ==================================================================== class VGPUCache: """Persistent, shape-keyed kernel cache for a wide range of ops.""" def __init__(self): self._kernels = {} # type: dict[tuple, GPUOp] self._stats = { 'compiles': 0, 'cache_hits': 0, 'dispatches': 0, 'cpu_calls': 0, 'gpu_calls': 0, 'op_calls': {}, # per-op call counts } self.gpu_available = GPU_AVAILABLE self.gpu_backend = GPU_BACKEND if self.gpu_available and self.gpu_backend == 'moderngl': try: self.ctx = moderngl.create_standalone_context() except Exception: self.gpu_available = False self.gpu_backend = None self.ctx = None print("[VGPU] GPU context creation failed; using CPU") else: self.ctx = None print(f"[VGPU] Using {'GPU' if self.gpu_available else 'CPU'} backend" + (f" (torch: {'yes' if TORCH_AVAILABLE else 'no'})" if TORCH_AVAILABLE else "")) # ------------------------------------------------------------------ # Public matmul (preserved signature – backwards compatible) # ------------------------------------------------------------------ def matmul(self, A, B) -> np.ndarray: """A @ B. Cached on (M, K, N) shape. Returns same flavour as A.""" t0 = _as_ndarray(A, dtype=np.float32) t1 = _as_ndarray(B, dtype=np.float32) M, K = t0.shape K2, N = t1.shape assert K == K2, f"Inner dimensions must match: {K} vs {K2}" key = ('matmul', M, K, N) # ---- matmul keeps its bespoke optimised shader from the # original file so existing users see no regression --------- if self.ctx is not None and self.gpu_available: try: self._stats['dispatches'] += 1 if key not in self._kernels or self._kernels[key] is None: self._kernels[key] = self._compile_matmul(M, K, N) self._stats['compiles'] += 1 else: self._stats['cache_hits'] += 1 result = self._kernels[key](t0, t1) self._stats['gpu_calls'] += 1 self._stats['op_calls']['matmul'] = \ self._stats['op_calls'].get('matmul', 0) + 1 return _as_output_like(A, result.reshape(M, N).astype(np.float32)) except Exception as e: print(f"[VGPU] GPU matmul failed: {e}; falling back to CPU") self.gpu_available = False self._stats['dispatches'] += 1 self._stats['cpu_calls'] += 1 self._stats['op_calls']['matmul'] = self._stats['op_calls'].get('matmul', 0) + 1 return _as_output_like(A, (t0 @ t1).astype(np.float32)) def _compile_matmul(self, M, K, N): """Optimised matmul shader (preserved from original).""" shader_src = f""" #version 430 core layout(local_size_x = 16, local_size_y = 16) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }}; layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }}; uniform int M = {M}; uniform int K = {K}; uniform int N = {N}; void main() {{ uint row = gl_GlobalInvocationID.y; uint col = gl_GlobalInvocationID.x; if (row >= uint(M) || col >= uint(N)) return; float sum = 0.0; for (uint k = 0u; k < uint(K); k++) {{ sum += A[row * uint(K) + k] * B[k * uint(N) + col]; }} C[row * uint(N) + col] = sum; }} """ program = self.ctx.compute_shader(shader_src) A_buf = self.ctx.buffer(reserve=M * K * 4) B_buf = self.ctx.buffer(reserve=K * N * 4) C_buf = self.ctx.buffer(reserve=M * N * 4) gx = max(1, (N + 15) // 16) gy = max(1, (M + 15) // 16) class GPUKernel: __slots__ = ('ctx','program','A_buf','B_buf','C_buf','gx','gy') def __init__(self, ctx, program, A_buf, B_buf, C_buf, gx, gy): self.ctx, self.program = ctx, program self.A_buf, self.B_buf, self.C_buf = A_buf, B_buf, C_buf self.gx, self.gy = gx, gy def __call__(self, A, B): self.A_buf.write(A.tobytes()) self.B_buf.write(B.tobytes()) self.A_buf.bind_to_storage_buffer(0) self.B_buf.bind_to_storage_buffer(1) self.C_buf.bind_to_storage_buffer(2) self.program.run(self.gx, self.gy, 1) return np.frombuffer(self.C_buf.read(), dtype=np.float32) print(f"[VGPU] ✅ Compiled matmul for {M}x{K}x{N}") return GPUKernel(self.ctx, program, A_buf, B_buf, C_buf, gx, gy) # ------------------------------------------------------------------ # Universal elementwise / activation dispatcher # ------------------------------------------------------------------ def _dispatch(self, op_name, A, *args): """Run a registered op on `A` (and optional `B`/scalar args). Returns same flavour (numpy/torch) as `A`.""" if op_name not in SHADER_TEMPLATES: raise ValueError(f"Unknown op '{op_name}'") A_nd = _as_ndarray(A, dtype=np.float32) N = int(A_nd.size) flat = A_nd.reshape(-1) is_binary = OP_ARITY[op_name] == 'binary' # -- extra uniform values (e.g. clip low/high) extra = {} if op_name == 'clip': lo, hi = float(args[0]), float(args[1]) extra = {'LO': lo, 'HI': hi} key = (op_name, N, str(A_nd.dtype)) self._stats['dispatches'] += 1 self._stats['op_calls'][op_name] = self._stats['op_calls'].get(op_name, 0) + 1 # ---- GPU path --------------------------------------------------- if self.ctx is not None and self.gpu_available: try: if key in self._kernels: self._stats['cache_hits'] += 1 else: self._kernels[key] = self._compile_op(op_name, N, is_binary) self._stats['compiles'] += 1 B_nd = _as_ndarray(args[0], dtype=np.float32).reshape(-1) \ if is_binary else None if is_binary and B_nd.size != N: raise ValueError(f"{op_name}: shapes must broadcast flat; " f"got {N} vs {B_nd.size}") flat_out = self._kernels[key].run(flat, B_nd) self._stats['gpu_calls'] += 1 out = flat_out.reshape(A_nd.shape).astype(A_nd.dtype) return _as_output_like(A, out) except Exception as e: # GPU failed (rare); fall through to CPU print(f"[VGPU] GPU {op_name} failed: {e}; using CPU") self.gpu_available = False # one-shot demotion # ---- CPU fallback ---------------------------------------------- self._stats['cpu_calls'] += 1 fn = CPU_FALLBACKS[op_name] if is_binary: B_nd = _as_ndarray(args[0], dtype=np.float32) out = fn(A_nd, B_nd) else: if op_name == 'clip': out = fn(A_nd, *args) else: out = fn(A_nd) return _as_output_like(A, out.astype(A_nd.dtype)) def _compile_op(self, op_name, N, is_binary): """Compile a binary/unary op shader and pre-allocate buffers.""" LX = 64 # workgroup size gx = max(1, (N + LX - 1) // LX) src = SHADER_TEMPLATES[op_name].format(LX=LX) program = self.ctx.compute_shader(src) A_buf = self.ctx.buffer(reserve=N * 4) B_buf = self.ctx.buffer(reserve=N * 4) if is_binary else None C_buf = self.ctx.buffer(reserve=N * 4) extra = {'LO': 0.0, 'HI': 1.0} if op_name == 'clip' else {} op = GPUOp(self.ctx, op_name, program, A_buf, B_buf, C_buf, gx, N, is_binary, extra) print(f"[VGPU] ✅ Compiled {op_name} (N={N})") return op # ------------------------------------------------------------------ # Dynamic dispatch: cache.(x, ...) -> result # ------------------------------------------------------------------ def __getattr__(self, name): # Avoid recursion on dunder / private lookups used by pickle etc. if name.startswith('_') or name not in SHADER_TEMPLATES: raise AttributeError(name) op = name def _call(a, *args): return self._dispatch(op, a, *args) _call.__name__ = op return _call # ------------------------------------------------------------------ # Diagnostics # ------------------------------------------------------------------ def report(self) -> str: backend = "GPU" if self.gpu_available else "CPU" ops = ", ".join(f"{k}={v}" for k, v in sorted(self._stats['op_calls'].items())) return (f"VGPUCache [{backend}]: " f"{self._stats['compiles']} compiles, " f"{self._stats['cache_hits']} cache hits, " f"{self._stats['dispatches']} dispatches, " f"{self._stats['cpu_calls']} CPU, " f"{self._stats['gpu_calls']} GPU" + (f" | ops: {{{ops}}}" if ops else "")) def clear(self): self._kernels.clear() print("[VGPU] kernel cache cleared") @property def stats(self): return dict(self._stats) # ==================================================================== # SELF-TEST # ==================================================================== if __name__ == "__main__": print("=" * 70) print("VGPU Universal Cache — extended self-test") print("=" * 70) cache = VGPUCache() # ------------- 1. Matmul (existing behaviour) ------------- print("\n[1] matmul 16x16") np.random.seed(42) A = np.random.randn(16, 16).astype(np.float32) B = np.random.randn(16, 16).astype(np.float32) assert np.max(np.abs(cache.matmul(A, B) - A @ B)) < 1e-3 print("[2] matmul 32x64 @ 64x16") A = np.random.randn(32, 64).astype(np.float32) B = np.random.randn(64, 16).astype(np.float32) assert np.max(np.abs(cache.matmul(A, B) - A @ B)) < 1e-3 print("[3] matmul cache hit (same shape)") _ = cache.matmul(A, B) print(" ", cache.report()) # ------------- 4. Binary elementwise ------------- print("\n[4] elementwise add/sub/mul/div/pow") A = np.random.randn(1024).astype(np.float32) B = np.random.randn(1024).astype(np.float32) + 1.0 # avoid div-by-zero for op, ref in [ ('add', A + B), ('sub', A - B), ('mul', A * B), ('div', A / B), ('pow', np.abs(A) ** B), ]: # bypass __getattr__ to test _dispatch directly out = cache._dispatch(op, A, B) diff = np.max(np.abs(out - ref)) print(f" {op:6s} max-diff = {diff:.2e}") assert diff < 1e-3, f"{op} failed" # ------------- 5. Unary elementwise ------------- print("\n[5] unary neg/abs/sqrt/square") A = np.abs(np.random.randn(1024).astype(np.float32)) + 1e-3 for op, ref in [ ('neg', -A), ('abs', np.abs(A)), ('sqrt', np.sqrt(A)), ('square', A * A), ]: out = cache._dispatch(op, A) diff = np.max(np.abs(out - ref)) print(f" {op:6s} max-diff = {diff:.2e}") assert diff < 1e-3, f"{op} failed" # ------------- 6. Activations ------------- print("\n[6] activations relu/sigmoid/tanh/clip") A = np.random.randn(2048).astype(np.float32) for op, ref in [ ('relu', np.maximum(A, 0)), ('sigmoid', 1.0 / (1.0 + np.exp(-A))), ('tanh', np.tanh(A)), ('clip', np.clip(A, -0.5, 0.5)), ]: if op == 'clip': out = cache._dispatch(op, A, -0.5, 0.5) else: out = cache._dispatch(op, A) diff = np.max(np.abs(out - ref)) print(f" {op:6s} max-diff = {diff:.2e}") assert diff < 1e-3, f"{op} failed" # ------------- 7. Dynamic dispatch via getattr ------------- print("\n[7] dynamic dispatch (cache.relu(...))") A = np.random.randn(512).astype(np.float32) ref = np.maximum(A, 0) out = cache.relu(A) diff = np.max(np.abs(out - ref)) print(f" cache.relu max-diff = {diff:.2e}") assert diff < 1e-3 # ------------- 8. torch.Tensor interop ------------- if TORCH_AVAILABLE: print("\n[8] torch.Tensor round-trip") t = torch.randn(8, 8) out_np = cache.matmul(t.cpu(), t.cpu()) assert isinstance(out_np, torch.Tensor), "should return torch.Tensor" assert torch.allclose(out_np, t.cpu() @ t.cpu(), atol=1e-3) out_t = cache.relu(t) assert isinstance(out_t, torch.Tensor) assert torch.allclose(out_t, torch.clamp(t, min=0.0), atol=1e-3) print(" torch ✓ (returns torch.Tensor when input is torch)") else: print("\n[8] torch.Tensor round-trip skipped (torch not installed)") # ------------- 9. Cache hits across repeated shapes ------------- print("\n[9] repeat same shape 200× — should hit cache") A = np.ones((1024,), dtype=np.float32) B = np.ones((1024,), dtype=np.float32) for _ in range(200): _ = cache.add(A, B) print(" ", cache.report()) # ------------- 10. Timing benchmark ------------- print("\n[10] timing 200× add(64×64)") A = np.random.randn(64, 64).astype(np.float32) B = np.random.randn(64, 64).astype(np.float32) t0 = time.time() for _ in range(200): _ = cache.add(A, B) dt = (time.time() - t0) / 200 * 1e3 print(f" {dt:.3f} ms / call ({cache.report()})") print("\n" + "=" * 70) print("✅ All tests passed!") print("=" * 70)