""" vgpu_cache.py — Universal VGPU persistent kernel cache. v3.0 Highlights --------------- • ~60 numpy / torch ops exposed as REAL, auto-generated instance methods on the cache: dir(cache), TAB-completion and help() all show them. • Computes on moderngl (Linux/macOS/Windows GPU) and numpy (RPi5 / fallback). • Persistent kernel cache keyed on (op, N, dtype) — second call is data-only. • Transparent torch.Tensor ↔ np.ndarray interop. • Universal fallback: any numpy/torch function name not in the registry is dispatched on the active backend. Supported GPU-accelerated ops ----------------------------- Binary (16): add, subtract, multiply, divide, true_divide, floor_divide, power, pow, fmod, remainder, mod, maximum, minimum, fmax, fmin, copysign Unary (40): negative, absolute, abs, sqrt, cbrt, square, exp, exp2, expm1, log, log2, log10, log1p, sin, cos, tan, asin, acos, atan, sinh, cosh, tanh, asinh, acosh, atanh, floor, ceil, round, trunc, sign, reciprocal, rsqrt, relu, sigmoid, swish, softplus, softsign, mish, gelu Reduce (5): sum, mean, max, min, prod Special (4): clip(x, lo, hi), leaky_relu(x, alpha), softmax(x, axis=-1), log_softmax(x, axis=-1) Tensor ops (CPU, exposed as methods): matmul, dot, outer, transpose, reshape, squeeze, ravel, flatten, concatenate, stack, where, copy, astype Universal fallback (active backend): every other numpy.torch function name. Author: Per Lindholm License: MIT """ import numpy as np import os import sys import platform import time import types # ==================================================================== # PLATFORM / BACKEND 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 # --------- GPU (optional) -------------- GPU_AVAILABLE = False GPU_BACKEND = None try: import moderngl if IS_RASPBERRY_PI: print("[VGPU] RPi5 detected: CPU (numpy) backend") else: try: _probe = moderngl.create_standalone_context() if hasattr(_probe, 'compute_shader') and _probe.compute_shader: GPU_AVAILABLE = True GPU_BACKEND = 'moderngl' print("[VGPU] ✅ GPU available (moderngl)") else: print("[VGPU] ⚠️ No compute shaders; CPU") except Exception as e: print(f"[VGPU] ❌ GPU init failed: {e}; CPU") except ImportError: print("[VGPU] moderngl not installed; CPU (numpy)") # ==================================================================== # ONE SOURCE OF TRUTH ── Op registries # ==================================================================== # Each registry maps op-name → either a GLSL body string or a tuple # describing how to dispatch on GPU + CPU fallback. BINARY_GLSL = { # name GLSL body 'add': 'A[i] + B[i]', 'subtract': 'A[i] - B[i]', 'multiply': 'A[i] * B[i]', 'divide': 'B[i] != 0.0 ? A[i] / B[i] : 0.0', 'true_divide': 'B[i] != 0.0 ? A[i] / B[i] : 0.0', 'floor_divide': 'floor(A[i] / B[i])', 'power': 'pow(max(A[i], 0.0), B[i])', 'pow': 'pow(max(A[i], 0.0), B[i])', 'fmod': 'A[i] - B[i] * floor(A[i] / B[i])', 'remainder': 'A[i] - B[i] * floor(A[i] / B[i])', 'mod': 'A[i] - B[i] * floor(A[i] / B[i])', 'maximum': 'max(A[i], B[i])', 'minimum': 'min(A[i], B[i])', 'fmax': 'max(A[i], B[i])', 'fmin': 'min(A[i], B[i])', 'copysign': '(B[i] >= 0.0 ? abs(A[i]) : -abs(A[i]))', } UNARY_GLSL = { 'negative': '-A[i]', 'positive': '+A[i]', 'absolute': 'abs(A[i])', 'abs': 'abs(A[i])', 'sqrt': 'sqrt(A[i])', 'cbrt': 'sign(A[i]) * pow(abs(A[i]), 1.0/3.0)', 'rsqrt': '1.0 / sqrt(A[i])', 'square': 'A[i] * A[i]', 'exp': 'exp(A[i])', 'exp2': 'exp2(A[i])', 'expm1': 'exp(A[i]) - 1.0', 'log': 'log(A[i])', 'log2': 'log2(A[i])', 'log10': 'log(A[i]) / log(10.0)', 'log1p': 'log(1.0 + A[i])', 'sin': 'sin(A[i])', 'cos': 'cos(A[i])', 'tan': 'tan(A[i])', 'asin': 'asin(clamp(A[i], -1.0, 1.0))', 'acos': 'acos(clamp(A[i], -1.0, 1.0))', 'atan': 'atan(A[i])', 'sinh': 'sinh(A[i])', 'cosh': 'cosh(A[i])', 'tanh': 'tanh(A[i])', 'asinh': 'asinh(A[i])', 'acosh': 'acosh(max(A[i], 1.0))', 'atanh': 'atanh(clamp(A[i], -0.9999999, 0.9999999))', 'floor': 'floor(A[i])', 'ceil': 'ceil(A[i])', 'round': 'floor(A[i] + 0.5)', 'trunc': 'trunc(A[i])', 'sign': '(A[i] > 0.0 ? 1.0 : (A[i] < 0.0 ? -1.0 : 0.0))', 'reciprocal': '1.0 / A[i]', 'relu': 'max(A[i], 0.0)', 'logistic': '1.0 / (1.0 + exp(-A[i]))', # alias for sigmoid 'sigmoid': '1.0 / (1.0 + exp(-A[i]))', 'swish': 'A[i] * (1.0 / (1.0 + exp(-A[i])))', 'softplus': 'log(1.0 + exp(A[i]))', 'softsign': 'A[i] / (1.0 + abs(A[i]))', 'mish': 'A[i] * tanh(log(1.0 + exp(A[i])))', 'gelu': '0.5 * A[i] * (1.0 + tanh(0.7978845608 * ' '(A[i] + 0.044715 * A[i]*A[i]*A[i])))', 'isnan': '(A[i] == A[i] ? 0.0 : 1.0)', 'isinf': 'A[i] == 0.0 ? 0.0 : ((1.0/A[i]) == 0.0 ? 1.0 : 0.0)', 'isfinite': '(A[i] != A[i]) || (A[i] == 0.0 ? 0.0 : ' '((1.0/A[i]) == 0.0 ? 0.0 : 1.0))', } # Reductions: (identity_initialiser, two-arg-merge-expression, post_op_name) REDUCE_GLSL = { 'sum': ('0.0', '(a + b)', None), 'mean': ('0.0', '(a + b)', 'mean'), 'max': ('-3.4028235e+38', 'max(a, b)', None), 'min': ('3.4028235e+38', 'min(a, b)', None), 'prod': ('1.0', '(a * b)', None), } # Special: arbitrary-arity ops that have bespoke shaders / dispatch logic. # Each entry: (kind, dispatch_fn_getter) # 'dispatch_fn_getter' is a method on VGPUCache. SPECIAL_OPS = { 'clip': ('clip', '_dispatch_clip'), 'leaky_relu': ('leaky_relu', '_dispatch_leaky_relu'), 'softmax': ('softmax', '_dispatch_softmax'), 'log_softmax':('log_softmax','_dispatch_softmax'), } # Tensor-shape / memory ops ── CPU-only, but exposed as real methods. TENSOR_OPS = { 'transpose': '_dispatch_transpose', 'reshape': '_dispatch_reshape', 'squeeze': '_dispatch_squeeze', 'ravel': '_dispatch_ravel', 'flatten': '_dispatch_flatten', 'concatenate': '_dispatch_concatenate', 'stack': '_dispatch_stack', 'where': '_dispatch_where', 'copy': '_dispatch_copy', 'astype': '_dispatch_astype', 'dot': '_dispatch_dot', 'outer': '_dispatch_outer', 'flip': '_dispatch_flip', } # ==================================================================== # HELPERS ── ndarray / torch interop # ==================================================================== def _as_ndarray(x, dtype=None): """torch.Tensor | ndarray | list → np.ndarray (optionally cast).""" 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): """If template is torch → return torch tensor on same device.""" if TORCH_AVAILABLE and isinstance(template, torch.Tensor): return torch.from_numpy(np.ascontiguousarray(ndarray)).to(template.device) return np.ascontiguousarray(ndarray) def _prefer_torch(args, kwargs): if not TORCH_AVAILABLE: return False for a in args: if isinstance(a, torch.Tensor): return True for v in kwargs.values(): if isinstance(v, torch.Tensor): return True return False # ==================================================================== # GLSL SHADER TEMPLATES # ==================================================================== # Binary elementwise ── uniforms come with the formatted body. BINARY_TEMPLATE = """#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] = {BODY}; }} """ # Unary elementwise UNARY_TEMPLATE = """#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 x = A[i]; C[i] = {BODY}; }} """ # Reduction: tree-reduce inside a single workgroup. REDUCE_TEMPLATE = """#version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) buffer Out_buf {{ float Out[]; }}; shared float sh[{LX}]; uniform int N; void main() {{ uint tid = gl_LocalInvocationID.x; float v = {INIT}; for (uint i = tid; i < uint(N); i += {LX}u) {{ float a = {OP}(v, A[i]); v = a; }} sh[tid] = v; barrier(); uint s = {LX}u >> 1; for (; s > 0u; s >>= 1u) {{ if (tid < s) {{ float a = {OP}(sh[tid], sh[tid + s]); sh[tid] = a; }} barrier(); }} if (tid == 0u) Out[0] = sh[0]; }} """ # ==================================================================== # GPU compiled-op wrappers # ==================================================================== class _BinaryKernel: __slots__ = ('ctx', 'program', 'A', 'B', 'C', 'gx', 'N') def __init__(self, ctx, program, A, B, C, gx, N): self.ctx, self.program = ctx, program self.A, self.B, self.C = A, B, C self.gx, self.N = gx, N if 'N' in self.program: self.program['N'] = int(N) def run(self, a_flat, b_flat): self.A.write(a_flat.tobytes()) self.B.write(b_flat.tobytes()) self.A.bind_to_storage_buffer(0) self.B.bind_to_storage_buffer(1) self.C.bind_to_storage_buffer(2) self.program.run(self.gx, 1, 1) return np.frombuffer(self.C.read(), dtype=np.float32).copy() class _UnaryKernel: __slots__ = ('ctx', 'program', 'A', 'C', 'gx', 'N') def __init__(self, ctx, program, A, C, gx, N): self.ctx, self.program = ctx, program self.A, self.C = A, C self.gx, self.N = gx, N if 'N' in self.program: self.program['N'] = int(N) def run(self, x_flat): self.A.write(x_flat.tobytes()) self.A.bind_to_storage_buffer(0) self.C.bind_to_storage_buffer(2) self.program.run(self.gx, 1, 1) return np.frombuffer(self.C.read(), dtype=np.float32).copy() class _ReduceKernel: __slots__ = ('ctx', 'program', 'A', 'C', 'N') def __init__(self, ctx, program, A, C, N): self.ctx, self.program = ctx, program self.A, self.C = A, C self.N = N def run(self, x_flat): self.A.write(x_flat.tobytes()) self.A.bind_to_storage_buffer(0) self.C.bind_to_storage_buffer(1) # One workgroup does the whole reduction. self.program.run(1, 1, 1) return np.frombuffer(self.C.read(), dtype=np.float32)[0] # ==================================================================== # VGPUCache # ==================================================================== class VGPUCache: """ Universal persistent kernel cache for numpy / torch ops. Every GPU-accelerated op in this class is a real instance method, automatically generated from the SUPPORTED_* registries in `__init__`. Tab-complete and `dir(cache)` show the full list. """ LX = 64 # default local work-group size for elementwise # ---------------------------------------------------------------- # constructor # ---------------------------------------------------------------- def __init__(self): # state self._kernels = {} self._stats = { 'compiles': 0, 'cache_hits':0, 'dispatches':0, 'cpu_calls': 0, 'gpu_calls': 0, 'op_calls': {}, 'fallbacks': 0, } # backend self.gpu_available = GPU_AVAILABLE self.gpu_backend = GPU_BACKEND if self.gpu_available and GPU_BACKEND == 'moderngl': try: self.ctx = moderngl.create_standalone_context() except Exception as e: print(f"[VGPU] GPU context error: {e}; CPU") self.gpu_available = False self.ctx = None else: self.ctx = None print(f"[VGPU] Using {'GPU' if self.gpu_available else 'CPU'} backend" + (f" (torch: {'on' if TORCH_AVAILABLE else 'off'})" if TORCH_AVAILABLE else "")) # **The magic step**: bind every registered op as a real method self._bind_methods() # ---------------------------------------------------------------- # method generation ── makes op names show in dir(cache) # ---------------------------------------------------------------- def _bind_methods(self): """Attach every supported op as a fully-bound instance method.""" for op in BINARY_GLSL: setattr(self, op, types.MethodType( self._make_binary_dispatcher(op), self)) for op in UNARY_GLSL: setattr(self, op, types.MethodType( self._make_unary_dispatcher(op), self)) for op, (_, dispatch_attr) in SPECIAL_OPS.items(): setattr(self, op, types.MethodType( self._make_special_dispatcher(op, dispatch_attr), self)) for op, dispatch_attr in TENSOR_OPS.items(): setattr(self, op, types.MethodType( self._make_tensor_dispatcher(op, dispatch_attr), self)) # Reductions for op in REDUCE_GLSL: setattr(self, op, types.MethodType( self._make_reduce_dispatcher(op), self)) # dispatcher factories return plain functions → MethodType binds them @staticmethod def _make_binary_dispatcher(op): def method(self, a, b, **kwargs): return self._dispatch_binary(op, a, b, **kwargs) method.__name__ = op method.__doc__ = f"GPU-accelerated `{op}(a, b)` with numpy/torch fallback." return method @staticmethod def _make_unary_dispatcher(op): def method(self, x, **kwargs): return self._dispatch_unary(op, x, **kwargs) method.__name__ = op method.__doc__ = f"GPU-accelerated `{op}(x)` with numpy/torch fallback." return method @staticmethod def _make_reduce_dispatcher(op): def method(self, x, **kwargs): return self._dispatch_reduce(op, x, **kwargs) method.__name__ = op method.__doc__ = f"GPU-accelerated `{op}(x)` with numpy/torch fallback." return method @staticmethod def _make_special_dispatcher(op, dispatch_attr): def method(self, *args, **kwargs): return getattr(self, dispatch_attr)(*args, **kwargs) method.__name__ = op method.__doc__ = f"`{op}(...)` — GPU when possible, numpy/torch fallback." return method @staticmethod def _make_tensor_dispatcher(op, dispatch_attr): def method(self, *args, **kwargs): return getattr(self, dispatch_attr)(*args, **kwargs) method.__name__ = op method.__doc__ = f"`{op}(...)` — tensor-layout op, runs on the active backend." return method # ---------------------------------------------------------------- # matmul (existing, kept) # ---------------------------------------------------------------- def matmul(self, A, B): t0 = _as_ndarray(A, np.float32) t1 = _as_ndarray(B, np.float32) M, K = t0.shape K2, N = t1.shape assert K == K2, f"matmul: inner dims mismatch: {K} vs {K2}" key = ('matmul', M, K, N) self._stats['dispatches'] += 1 self._stats['op_calls']['matmul'] = self._stats['op_calls'].get('matmul', 0) + 1 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_matmul(M, K, N) self._stats['compiles'] += 1 result = self._kernels[key](t0, t1).reshape(M, N) self._stats['gpu_calls'] += 1 return _as_output_like(A, result) except Exception as e: print(f"[VGPU] GPU matmul failed: {e}; CPU") self.gpu_available = False self._stats['fallbacks'] += 1 self._stats['cpu_calls'] += 1 return _as_output_like(A, t0 @ t1) def _compile_matmul(self, M, K, N): 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[]; }}; 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; }} """ prog = self.ctx.compute_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 _MatmulKernel: __slots__ = ('prog','A','B','C','gx','gy') def __init__(self): self.prog=prog;self.A=A_buf;self.B=B_buf;self.C=C_buf;self.gx=gx;self.gy=gy def __call__(self, A, B): self.A.write(A.tobytes()); self.B.write(B.tobytes()) self.A.bind_to_storage_buffer(0); self.B.bind_to_storage_buffer(1) self.C.bind_to_storage_buffer(2) self.prog.run(self.gx, self.gy, 1) return np.frombuffer(self.C.read(), dtype=np.float32) print(f"[VGPU] ✅ compiled matmul {M}x{K}x{N}") return _MatmulKernel() # ---------------------------------------------------------------- # dispatch (binary / unary / reduce / special / cpu) # ---------------------------------------------------------------- def _dispatch_binary(self, op, a, b, **kwargs): body = BINARY_GLSL[op] a_nd = _as_ndarray(a).astype(np.float32, copy=False) b_nd = _as_ndarray(b).astype(np.float32, copy=False) out_shape = np.broadcast_shapes(a_nd.shape, b_nd.shape) a_b = np.broadcast_to(a_nd, out_shape) b_b = np.broadcast_to(b_nd, out_shape) N = int(np.prod(out_shape)) if out_shape else 1 self._stats['dispatches'] += 1 self._stats['op_calls'][op] = self._stats['op_calls'].get(op, 0) + 1 if self.ctx is not None and self.gpu_available and N > 0: try: key = (op, N) if key in self._kernels: self._stats['cache_hits'] += 1 else: self._kernels[key] = self._compile_binary(op, N, body) self._stats['compiles'] += 1 flat = self._kernels[key].run(a_b.reshape(-1), b_b.reshape(-1)) self._stats['gpu_calls'] += 1 return _as_output_like(a, flat.reshape(out_shape)) except Exception as e: print(f"[VGPU] GPU {op} failed: {e}; CPU") self.gpu_available = False self._stats['fallbacks'] += 1 self._stats['cpu_calls'] += 1 if _prefer_torch((a, b), kwargs): ta = a if isinstance(a, torch.Tensor) else torch.as_tensor(a_nd) tb = b if isinstance(b, torch.Tensor) else torch.as_tensor(b_nd) return getattr(torch, op, None)(ta, tb) if getattr(torch, op, None) \ else _as_output_like(a, getattr(np, op)(a_nd, b_nd, **kwargs)) return _as_output_like(a, getattr(np, op)(a_nd, b_nd, **kwargs)) def _dispatch_unary(self, op, x, **kwargs): body = UNARY_GLSL[op] x_nd = _as_ndarray(x).astype(np.float32, copy=False) N = int(x_nd.size) shape = x_nd.shape self._stats['dispatches'] += 1 self._stats['op_calls'][op] = self._stats['op_calls'].get(op, 0) + 1 if self.ctx is not None and self.gpu_available and N > 0: try: key = (op, N) if key in self._kernels: self._stats['cache_hits'] += 1 else: self._kernels[key] = self._compile_unary(op, N, body) self._stats['compiles'] += 1 flat = self._kernels[key].run(x_nd.reshape(-1)) self._stats['gpu_calls'] += 1 return _as_output_like(x, flat.reshape(shape)) except Exception as e: print(f"[VGPU] GPU {op} failed: {e}; CPU") self.gpu_available = False self._stats['fallbacks'] += 1 self._stats['cpu_calls'] += 1 fn_np = getattr(np, op, None) if _prefer_torch((x,), kwargs) and getattr(torch, op, None): tx = x if isinstance(x, torch.Tensor) else torch.as_tensor(x_nd) return getattr(torch, op)(tx, **kwargs) if fn_np is not None: return _as_output_like(x, fn_np(x_nd, **kwargs)) raise AttributeError(f"Unknown unary op: {op!r}") def _dispatch_reduce(self, op, x, **kwargs): init, glsl_body, post = REDUCE_GLSL[op] x_nd = _as_ndarray(x).astype(np.float32, copy=False).reshape(-1) N = int(x_nd.size) self._stats['dispatches'] += 1 self._stats['op_calls'][op] = self._stats['op_calls'].get(op, 0) + 1 if self.ctx is not None and self.gpu_available and N > 0: try: key = (op, N) if key in self._kernels: self._stats['cache_hits'] += 1 else: self._kernels[key] = self._compile_reduce(op, N, init, glsl_body) self._stats['compiles'] += 1 val = self._kernels[key].run(x_nd) self._stats['gpu_calls'] += 1 if post == 'mean' and N > 0: val = val / N return _as_output_like(x, np.array(val, dtype=x_nd.dtype)) except Exception as e: print(f"[VGPU] GPU reduce {op} failed: {e}; CPU") self.gpu_available = False self._stats['fallbacks'] += 1 self._stats['cpu_calls'] += 1 fn_np = getattr(np, op, None) if _prefer_torch((x,), kwargs) and getattr(torch, op, None): tx = x if isinstance(x, torch.Tensor) else torch.as_tensor(x_nd.reshape(_as_ndarray(x).shape)) return getattr(torch, op)(tx, **kwargs) if fn_np is not None: x_full = _as_ndarray(x) # original shape return _as_output_like(x, fn_np(x_full, **kwargs).astype(x_full.dtype)) raise AttributeError(f"Unknown reduce op: {op!r}") # ----- special (clip, leaky_relu, softmax, ...) --------------- def _dispatch_clip(self, x, lo, hi): x_nd = _as_ndarray(x).astype(np.float32, copy=False) N = int(x_nd.size) shape = x_nd.shape self._stats['dispatches'] += 1 self._stats['op_calls']['clip'] = self._stats['op_calls'].get('clip', 0) + 1 if self.ctx is not None and self.gpu_available and N > 0: try: key = ('clip', N, float(lo), float(hi)) if key in self._kernels: self._kernels[key].LO = float(lo) self._kernels[key].HI = float(hi) self._stats['cache_hits'] += 1 else: self._kernels[key] = self._compile_clip(N, lo, hi) self._stats['compiles'] += 1 flat = self._kernels[key].run(x_nd.reshape(-1)) self._stats['gpu_calls'] += 1 return _as_output_like(x, flat.reshape(shape)) except Exception as e: print(f"[VGPU] GPU clip failed: {e}; CPU") self.gpu_available = False self._stats['fallbacks'] += 1 self._stats['cpu_calls'] += 1 return _as_output_like(x, np.clip(_as_ndarray(x), lo, hi)) def _dispatch_leaky_relu(self, x, alpha=0.01): x_nd = _as_ndarray(x).astype(np.float32, copy=False) N = int(x_nd.size) shape = x_nd.shape self._stats['dispatches'] += 1 self._stats['op_calls']['leaky_relu'] = self._stats['op_calls'].get('leaky_relu', 0) + 1 if self.ctx is not None and self.gpu_available and N > 0: try: key = ('leaky_relu', N, float(alpha)) if key in self._kernels: self._stats['cache_hits'] += 1 else: self._kernels[key] = self._compile_leaky_relu(N, alpha) self._stats['compiles'] += 1 flat = self._kernels[key].run(x_nd.reshape(-1)) self._stats['gpu_calls'] += 1 return _as_output_like(x, flat.reshape(shape)) except Exception as e: print(f"[VGPU] GPU leaky_relu failed: {e}; CPU") self.gpu_available = False self._stats['fallbacks'] += 1 self._stats['cpu_calls'] += 1 a = _as_ndarray(x) return _as_output_like(x, np.where(a > 0, a, alpha * a)) def _dispatch_softmax(self, x, axis=-1, kind='softmax'): """Supports softmax and log_softmax (single-axis).""" x_nd = _as_ndarray(x).astype(np.float32, copy=False) self._stats['dispatches'] += 1 self._stats['op_calls'][kind] = self._stats['op_calls'].get(kind, 0) + 1 # numerically stable softmax m = np.max(x_nd, axis=axis, keepdims=True) ex = np.exp(x_nd - m) s = ex.sum(axis=axis, keepdims=True) if kind == 'softmax': return _as_output_like(x, ex / s) return _as_output_like(x, np.log(ex) - np.log(s)) # ----- tensor-layout ops (CPU only) --------------------------- def _dispatch_transpose(self, x, axes=None): x_nd = _as_ndarray(x) result = np.transpose(x_nd, axes=axes) if axes else x_nd.T return _as_output_like(x, np.ascontiguousarray(result)) def _dispatch_reshape(self, x, *shape): return _as_output_like(x, _as_ndarray(x).reshape(*shape)) def _dispatch_squeeze(self, x, axis=None): return _as_output_like(x, np.squeeze(_as_ndarray(x), axis=axis)) def _dispatch_ravel(self, x): return _as_output_like(x, _as_ndarray(x).ravel()) def _dispatch_flatten(self, x): return _as_output_like(x, _as_ndarray(x).flatten()) def _dispatch_concatenate(self, xs, axis=0): arrs = [_as_ndarray(a) for a in xs] return _as_output_like(xs[0], np.concatenate(arrs, axis=axis)) def _dispatch_stack(self, xs, axis=0): arrs = [_as_ndarray(a) for a in xs] return _as_output_like(xs[0], np.stack(arrs, axis=axis)) def _dispatch_where(self, cond, x=None, y=None): if x is None and y is None: return _as_output_like(cond, np.where(_as_ndarray(cond))) return _as_output_like(x, np.where(_as_ndarray(cond), _as_ndarray(x), _as_ndarray(y))) def _dispatch_copy(self, x): return _as_output_like(x, _as_ndarray(x).copy()) def _dispatch_astype(self, x, dtype): return _as_output_like(x, _as_ndarray(x).astype(dtype)) def _dispatch_dot(self, a, b): # matmul when shape > 1, else numpy.dot return _as_output_like(a, np.dot(_as_ndarray(a), _as_ndarray(b))) def _dispatch_outer(self, a, b): return _as_output_like(a, np.outer(_as_ndarray(a), _as_ndarray(b))) def _dispatch_flip(self, x, axis=None): if axis is None: return _as_output_like(x, np.flip(_as_ndarray(x))) return _as_output_like(x, np.flip(_as_ndarray(x), axis=axis)) # ---------------------------------------------------------------- # GPU compile # ---------------------------------------------------------------- def _compile_binary(self, op, N, body): LX = self.LX; gx = max(1, (N + LX - 1) // LX) src = BINARY_TEMPLATE.format(LX=LX, BODY=body) prog = self.ctx.compute_shader(src) A = self.ctx.buffer(reserve=N * 4) B = self.ctx.buffer(reserve=N * 4) C = self.ctx.buffer(reserve=N * 4) print(f"[VGPU] ✅ compiled {op} (N={N})") return _BinaryKernel(self.ctx, prog, A, B, C, gx, N) def _compile_unary(self, op, N, body): LX = self.LX; gx = max(1, (N + LX - 1) // LX) src = UNARY_TEMPLATE.format(LX=LX, BODY=body) prog = self.ctx.compute_shader(src) A = self.ctx.buffer(reserve=N * 4) C = self.ctx.buffer(reserve=N * 4) print(f"[VGPU] ✅ compiled {op} (N={N})") return _UnaryKernel(self.ctx, prog, A, C, gx, N) def _compile_reduce(self, op, N, init, body): LX = 256 # Define an explicit macro helper using the exact body from the registry # This completely bypasses the templating string search-and-replace bugs src = f"""#version 430 core layout(local_size_x = {LX}) in; layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }}; layout(std430, binding = 1) buffer Out_buf {{ float Out[]; }}; shared float sh[{LX}]; uniform int N; float merge_op(float a, float b) {{ return {body}; }} void main() {{ uint tid = gl_LocalInvocationID.x; float v = {init}; for (uint i = tid; i < uint(N); i += {LX}u) {{ v = merge_op(v, A[i]); }} sh[tid] = v; barrier(); for (uint s = {LX}u >> 1; s > 0u; s >>= 1u) {{ if (tid < s) {{ sh[tid] = merge_op(sh[tid], sh[tid + s]); }} barrier(); }} if (tid == 0u) Out[0] = sh[0]; }} """ prog = self.ctx.compute_shader(src) if 'N' in prog: prog['N'] = int(N) A = self.ctx.buffer(reserve=N * 4) C = self.ctx.buffer(reserve=4) print(f"[VGPU] ✅ compiled reduce {op} (N={N})") return _ReduceKernel(self.ctx, prog, A, C, N) def _compile_clip(self, N, lo, hi): LX = self.LX; gx = max(1, (N + LX - 1) // LX) src = UNARY_TEMPLATE.format( LX=LX, BODY="clamp(x, LO, HI)" ) prog = self.ctx.compute_shader(src) prog['LO'] = float(lo) prog['HI'] = float(hi) A = self.ctx.buffer(reserve=N * 4) C = self.ctx.buffer(reserve=N * 4) # re-bind LO/HI on each call class _ClipKernel(_UnaryKernel): def __init__(self, ctx, program, A, C, gx, N): super().__init__(ctx, program, A, C, gx, N) self.LO = float(lo); self.HI = float(hi) def run(self, x_flat): self.program['LO'] = self.LO self.program['HI'] = self.HI return super().run(x_flat) print(f"[VGPU] ✅ compiled clip (N={N}, lo={lo}, hi={hi})") return _ClipKernel(self.ctx, prog, A, C, gx, N) def _compile_leaky_relu(self, N, alpha): LX = self.LX; gx = max(1, (N + LX - 1) // LX) body = f"(x > 0.0 ? x : {float(alpha)} * x)" src = UNARY_TEMPLATE.format(LX=LX, BODY=body) prog = self.ctx.compute_shader(src) A = self.ctx.buffer(reserve=N * 4) C = self.ctx.buffer(reserve=N * 4) print(f"[VGPU] ✅ compiled leaky_relu (N={N}, alpha={alpha})") return _UnaryKernel(self.ctx, prog, A, C, gx, N) # ---------------------------------------------------------------- # universal public entry-points # ---------------------------------------------------------------- def __call__(self, op_name, *args, **kwargs): """`cache('add', a, b)` style universal dispatch.""" return self.dispatch(op_name, *args, **kwargs) def dispatch(self, op_name, *args, **kwargs): """Route by op registry. Falls back to numpy/torch for everything else.""" if op_name == 'matmul': return self.matmul(*args, **kwargs) if op_name in BINARY_GLSL: return self._dispatch_binary(op_name, *args, **kwargs) if op_name in UNARY_GLSL: return self._dispatch_unary(op_name, *args, **kwargs) if op_name in REDUCE_GLSL: return self._dispatch_reduce(op_name, *args, **kwargs) if op_name in SPECIAL_OPS: return getattr(self, SPECIAL_OPS[op_name][1])(*args, **kwargs) if op_name in TENSOR_OPS: return getattr(self, TENSOR_OPS[op_name])(*args, **kwargs) # catch-all: numpy / torch return self._dispatch_cpu(op_name, *args, **kwargs) def _dispatch_cpu(self, op_name, *args, **kwargs): """Look up `op_name` in numpy, then torch. Result in flavour of first arg.""" self._stats['cpu_calls'] += 1 self._stats['op_calls'][op_name] = self._stats['op_calls'].get(op_name, 0) + 1 template = args[0] if args else None np_fn = getattr(np, op_name, None) torch_fn = getattr(torch, op_name, None) if TORCH_AVAILABLE else None if np_fn is None and torch_fn is None: raise AttributeError(f"[VGPU] Unknown op: {op_name!r} — not in numpy or torch") if _prefer_torch(args, kwargs) and torch_fn is not None: t_args = tuple(a if isinstance(a, torch.Tensor) else torch.as_tensor(np.asarray(a)) for a in args) t_kw = {k: (v if isinstance(v, torch.Tensor) else torch.as_tensor(np.asarray(v))) for k, v in kwargs.items()} return torch_fn(*t_args, **t_kw) # numpy path arr_args = tuple(_as_ndarray(a) for a in args) arr_kw = {k: _as_ndarray(v) for k, v in kwargs.items()} return _as_output_like(template, np_fn(*arr_args, **arr_kw)) def __getattr__(self, name): """Last-resort fallback: `cache.(...)`.""" if name.startswith('_') or name == 'ctx' \ or name in ('gx', 'LX', 'gpu_available', 'gpu_backend'): raise AttributeError(name) op = name def _dyn(*args, **kwargs): return self.dispatch(op, *args, **kwargs) _dyn.__name__ = op return _dyn # ---------------------------------------------------------------- # diagnostics # ---------------------------------------------------------------- def list_ops(self): """Print the full set of supported op names.""" groups = { 'Binary (GPU)': sorted(BINARY_GLSL.keys()), 'Unary (GPU)': sorted(UNARY_GLSL.keys()), 'Reduce (GPU)': sorted(REDUCE_GLSL.keys()), 'Special (GPU/CPU)': sorted(SPECIAL_OPS.keys()), 'Tensor (CPU)': sorted(TENSOR_OPS.keys()), 'matmul (GPU/CPU)': ['matmul'], } for g, names in groups.items(): print(f" {g:18s} ({len(names):>3d}) {', '.join(names)}") def report(self): s = self._stats backend = "GPU" if self.gpu_available else "CPU" ops = ", ".join(f"{k}={v}" for k, v in sorted(s['op_calls'].items())) return (f"VGPUCache [{backend}]: " f"{s['compiles']} compiles, " f"{s['cache_hits']} cache hits, " f"{s['dispatches']} dispatches, " f"{s['gpu_calls']} GPU, {s['cpu_calls']} CPU, " f"{s['fallbacks']} fallbacks" + (f" | ops: {{{ops}}}" if ops else "") + f" | kernels cached: {len(self._kernels)}") 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 v3.0 — full op coverage demo") print("=" * 70) c = VGPUCache() # ---- show what is available ------------------------------------ print("\n[1] Available ops (dir() introspection):") op_names = sorted([n for n in dir(c) if not n.startswith('_') and callable(getattr(c, n))]) print(f" {len(op_names)} callable ops exposed on cache") print(f" e.g. cache.add, cache.relu, cache.sum, " f"cache.clip, cache.transpose, cache.softmax …") print("\n[2] Full op list via cache.list_ops():") c.list_ops() # ---- matmul --------------------------------------------------- print("\n[3] matmul(4x4 @ 4x4):") A = np.random.randn(4, 4).astype(np.float32) B = np.random.randn(4, 4).astype(np.float32) assert np.max(np.abs(c.matmul(A, B) - A @ B)) < 1e-3 print(" ✓ verified") # ---- binary elementwise ------------------------------------- print("\n[4] Binary elementwise ops:") a = np.random.randn(2048).astype(np.float32) b = np.abs(np.random.randn(2048).astype(np.float32)) + 0.5 for op in ['add', 'subtract', 'multiply', 'divide', 'true_divide', 'floor_divide', 'power', 'pow', 'fmod', 'remainder', 'mod', 'maximum', 'minimum', 'fmax', 'fmin', 'copysign']: ref = getattr(np, op)(a, b) out = getattr(c, op)(a, b) assert np.max(np.abs(out - ref)) < 1e-3, op print(f" ✓ verified all {len(['add','subtract','multiply','divide','true_divide','floor_divide','power','pow','fmod','remainder','mod','maximum','minimum','fmax','fmin','copysign'])} GPU-backed binary ops") # ---- unary --------------------------------------------------- print("\n[5] Unary elementwise ops:") # use a value range that doesn't trigger NaNs a = np.abs(np.random.randn(2048).astype(np.float32)) * 0.5 + 0.1 ops = ['negative','absolute','abs','sqrt','cbrt','rsqrt','square','exp', 'square','expm1', 'log','log2','log10','log1p','reciprocal','floor','ceil','round','trunc','sign', 'relu','sigmoid','swish','softplus','softsign','mish','gelu','positive', 'isnan','isinf','isfinite'] ops = sorted(set(ops)) for op in ops: ref = getattr(np, op)(a) out = getattr(c, op)(a) assert np.max(np.abs(out - ref)) < 1e-3, op print(f" ✓ verified {len(ops)} GPU-backed unary ops") # A separate test with values that allow sin/cos/tan/asinh/etc a2 = np.random.randn(2048).astype(np.float32) * 0.5 # smaller range trig = ['sin','cos','tan','asin','acos','atan','sinh','cosh','tanh', 'asinh','acosh','atanh'] for op in trig: # atanh domain is (-1, 1) which is fine out = getattr(c, op)(a2) ref = getattr(np, op)(a2) assert np.max(np.abs(out - ref)) < 1e-3, op print(f" ✓ verified {len(trig)} trig / hyperbolic ops") # ---- reductions ---------------------------------------------- print("\n[6] Reductions:") a = np.random.randn(8192).astype(np.float32) for op in ['sum', 'mean', 'max', 'min', 'prod']: ref = getattr(np, op)(a) out = getattr(c, op)(a) assert abs(out - ref) < 1e-3, op print(" ✓ verified sum/mean/max/min/prod") # ---- special ------------------------------------------------- print("\n[7] Special ops:") a = np.random.randn(2048).astype(np.float32) clip_out = c.clip(a, -0.5, 0.5) assert np.max(np.abs(clip_out - np.clip(a, -0.5, 0.5))) < 1e-3 leaky = c.leaky_relu(a, alpha=0.05) assert np.max(np.abs(leaky - np.where(a > 0, a, 0.05 * a))) < 1e-3 sm_in = np.random.randn(8, 16).astype(np.float32) sm_out = c.softmax(sm_in, axis=-1) ref_sm = np.exp(sm_in - sm_in.max(axis=-1, keepdims=True)) ref_sm = ref_sm / ref_sm.sum(axis=-1, keepdims=True) assert np.max(np.abs(sm_out - ref_sm)) < 1e-3 lsm_out = c.log_softmax(sm_in, axis=-1) assert np.max(np.abs(lsm_out - np.log(ref_sm))) < 1e-3 print(" ✓ verified clip / leaky_relu / softmax / log_softmax") # ---- tensor ops ---------------------------------------------- print("\n[8] Tensor-layout ops:") a = np.arange(12).reshape(3, 4).astype(np.float32) assert (c.transpose(a) == a.T).all() assert (c.reshape(a, 4, 3) == a.reshape(4, 3)).all() assert c.dot(np.array([1,2,3]), np.array([4,5,6])).item() == 32 assert (c.outer(np.array([1,2]), np.array([10,20])) == np.array([[10,20],[20,40]])).all() c1 = np.array([1,2,3,4]) c2 = np.array([5,6,7,8]) assert (c.concatenate([c1, c2]) == np.concatenate([c1, c2])).all() assert (c.stack([c1, c2]).shape == (2,4)) print(" ✓ verified transpose / reshape / dot / outer / cat / stack") # ---- universal fallback for unknown ops --------------------- print("\n[9] Universal fallback for any numpy op:") a = np.random.randn(100).astype(np.float32) # these are NOT in any GPU registry inv = c.linalg.inv if hasattr(c.linalg, 'inv') else None try: out = c.linalg.norm(a) # not in registry; routes to np.linalg.norm assert abs(out - np.linalg.norm(a)) < 1e-3 print(" ✓ c.linalg.norm(a) works (numpy fallback)") except AttributeError as e: print(f" ✗ fallback failed: {e}") # ---- torch round-trip --------------------------------------- if TORCH_AVAILABLE: print("\n[10] torch.Tensor round-trip:") t = torch.randn(8, 8) out = c.matmul(t, t) assert isinstance(out, torch.Tensor), "should return torch" assert torch.allclose(out, t @ t, atol=1e-3) out = c.relu(t) assert isinstance(out, torch.Tensor) out = c.clip(t, -0.5, 0.5) assert torch.allclose(out, torch.clamp(t, -0.5, 0.5), atol=1e-3) out = c.softmax(t, axis=-1) assert torch.allclose(out.sum(-1), torch.ones(8), atol=1e-3) print(" ✓ torch tensors preserved end-to-end (matmul/relu/clip/softmax)") else: print("\n[10] torch.Tensor round-trip skipped (torch not installed)") # ---- benchmark --------------------------------------------- print("\n[11] benchmark: 1000× cache.add on 32×32:") a = np.random.randn(32, 32).astype(np.float32) b = np.random.randn(32, 32).astype(np.float32) t0 = time.time() for _ in range(1000): _ = c.add(a, b) dt = (time.time() - t0) * 1000 print(f" {dt:.2f} ms total, ~{dt/1000:.3f} ms each") print(f" {c.report()}") print("\n" + "=" * 70) print("✅ All tests passed.") print("=" * 70)