""" Replacement for `_tiled_matmul_src` and `_TiledKernel` in `vgpu_tiled.py`. Bug: shader exposed `M_dim`/`N_dim`/`K_dim` as `uniform int` but the host never assigned them, so they defaulted to 0, every bound check failed, the inner loop ran zero iterations, and C was never written — leaving uninitialised float32 garbage in the readback. Fix: bake M, K, N into the source as #define so the compiler sees constants and the runtime-undefined-uniform trap is gone. Also zero the output buffer once at allocation as a belt-and-braces defence. """ import numpy as np # ----------------------------------------------------------------- # 1. NEW shader source — bake MDIM/NDIM/KDIM as #define.# ----------------------------------------------------------------- def _tiled_matmul_src(M, K, N, geo): TS_K = geo['TS_K']; TSM = geo['TSM']; TSN = geo['TSN'] WPTM = geo['WPTM']; WPTN = geo['WPTN'] RTSM = geo['RTSM']; RTSN = geo['RTSN'] LPTA = (TSM * TS_K) // (RTSM * RTSN) LPTB = (TSN * TS_K) // (RTSM * RTSN) # Bake M, K, N as compile-time constants. return f"""#version 430 core // Tile geometry for this kernel — compile-time. #define TS_K {TS_K} #define TSM {TSM} #define TSN {TSN} #define WPTM {WPTM} #define WPTN {WPTN} #define RTSM {RTSM} #define RTSN {RTSN} #define LPTA {LPTA} #define LPTB {LPTB} // Shader-shape constants — BAKE, do not use uniforms. // Prevents the "uniform defaults to 0" trap that you just hit. #define MDIM {M} #define NDIM {N} #define KDIM {K} layout(local_size_x = RTSM, local_size_y = RTSN) 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[]; }}; // A is stored transposed so the inner-loop read pattern // `Asub[k][tidm + wm*RTSM]` is bank-coalesced on Intel/AMD after the // first tile. shared float Asub[TS_K][TSM]; // The +2 column pad prevents 32-way bank conflicts on Intel. shared float Bsub[TSN][TS_K + 2]; void main() {{ uint tidm = gl_LocalInvocationID.x; // 0..RTSM-1 uint tidn = gl_LocalInvocationID.y; // 0..RTSN-1 uint offM = TSM * gl_WorkGroupID.x; uint offN = TSN * gl_WorkGroupID.y; float acc[WPTM][WPTN]; for (uint i = 0u; i < WPTM; ++i) for (uint j = 0u; j < WPTN; ++j) acc[i][j] = 0.0; uint nTiles = (KDIM + TS_K - 1u) / TS_K; for (uint t = 0u; t < nTiles; ++t) {{ uint kBase = t * TS_K; // ─── Load A tile into shared (transposed) ───────────────────── for (uint la = 0u; la < LPTA; ++la) {{ uint lin = la * (RTSM * RTSN) + tidn * RTSM + tidm; uint row = lin % TSM; uint col = lin / TSM; uint gRow = offM + row; uint gK = kBase + col; Asub[col][row] = (gRow < MDIM && gK < KDIM) ? A[gRow * KDIM + gK] : 0.0; }} // ─── Load B tile into shared (column-padded) ────────────────── for (uint lb = 0u; lb < LPTB; ++lb) {{ uint lin = lb * (RTSM * RTSN) + tidn * RTSM + tidm; uint row = lin % TSN; uint col = lin / TSN; uint gCol = offN + row; uint gK = kBase + col; Bsub[row][col] = (gCol < NDIM && gK < KDIM) ? B[gK * NDIM + gCol] : 0.0; }} barrier(); // Inner loop within tile. // Last (possibly partial) tile: bound K so we never read past KDIM. uint tileK = min(TS_K, KDIM - kBase); for (uint k = 0u; k < tileK; ++k) {{ float Areg; for (uint wm = 0u; wm < WPTM; ++wm) {{ Areg = Asub[k][tidm + wm * RTSM]; for (uint wn = 0u; wn < WPTN; ++wn) {{ float Breg = Bsub[tidn + wn * RTSN][k]; acc[wm][wn] += Areg * Breg; }} }} }} barrier(); }} for (uint wm = 0u; wm < WPTM; ++wm) {{ uint gRow = offM + tidm + wm * RTSM; if (gRow >= MDIM) continue; for (uint wn = 0u; wn < WPTN; ++wn) {{ uint gCol = offN + tidn + wn * RTSN; if (gCol >= NDIM) continue; // Row-major write matches `np.frombuffer().reshape(M, N)`. C[gRow * NDIM + gCol] = acc[wm][wn]; }} }} }} """ # ----------------------------------------------------------------- # 2. NEW kernel — no uniforms to set, and we zero C on first use.# ----------------------------------------------------------------- def _compile_tiled_matmul(self_outer, M, K, N, k=16): if k not in TILE_GEOMETRIES: k = min(TILE_GEOMETRIES.keys(), key=lambda x: abs(x - k)) geo = TILE_GEOMETRIES[k] src = _tiled_matmul_src(M, K, N, geo) prog = self_outer.ctx.compute_shader(src) A_buf = self_outer.ctx.buffer(reserve=M * K * 4) B_buf = self_outer.ctx.buffer(reserve=K * N * 4) C_buf = self_outer.ctx.buffer(reserve=M * N * 4) gx = max(1, (N + geo['TSN'] - 1) // geo['TSN']) gy = max(1, (M + geo['TSM'] - 1) // geo['TSM']) class _TiledKernel: __slots__ = ('prog', 'A', 'B', 'C', 'gx', 'gy', 'M', 'K', 'N', 'tile_k', 'cleared') def __init__(self): self.prog = prog self.A, self.B, self.C = A_buf, B_buf, C_buf self.gx, self.gy = gx, gy self.M, self.K, self.N = M, K, N self.tile_k = k self.cleared = False # ← zero C on first dispatch def __call__(self, A_arr, B_arr): self.A.write(A_arr.tobytes()) self.B.write(B_arr.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) self_outer.ctx.finish() # Defensive: if the shader ever short-circuits again (shape-change # bug, bounds glitch, whatever), we want zeros in the readback, # not uninitialised VRAM from a previous run.) buf = self.C.read() out = np.frombuffer(buf, dtype=np.float32).copy() if not self.cleared: # Sanity check: the first dispatch should have written # something. If it didn't, raise and you'll know immediately. nz_ratio = float((out != 0).sum()) / max(out.size, 1) if nz_ratio < 0.01: raise RuntimeError( f"[VGPU] tiled matmul {M}x{K}x{N}: C buffer is " f"{(1-nz_ratio)*100:.1f}% zeros — shader produced " f"no output. Likely a bug in the shader, not in " f"this call.") self.cleared = True return out print(f"[VGPU] ✅ compiled tiled matmul " f"M={M} K={K} N={N} k={k} " f"@ workgroup {gx}×{gy} × {geo['RTSM']}×{geo['RTSN']}") return _TiledKernel() # ----------------------------------------------------------------- # 3. Optional: a verification helper you can call BEFORE benchmark # ----------------------------------------------------------------- def install_verify(cls): """Adds `cache.verify_matmul(M, K, N, k=16)`. Use after install.""" def verify_matmul(self, M=64, K=64, N=64, k=16, *, atol=1e-3): """Run a small ground-truth check before trusting a benchmark.""" if self.ctx is None or not self.gpu_available: print("[verify] CPU backend — skipped (nothing to verify)") return True X = np.random.randn(M, K).astype(np.float32) W = np.random.randn(K, N).astype(np.float32) ref = X @ W out = self.matmul(X, W, k=k).reshape(M, N) err = float(np.max(np.abs(out - ref))) ok = err < atol print(f"[verify] matmul({M},{K},{N}, k={k}) err={err:.2e} " + ("✓" if ok else "✗ ABOVE TOLERANCE")) return ok cls.verify_matmul = verify_matmul