Mj2LdZddlZddlZddlZddlZedddddddedddddddedddddddd Zd Zd(d Z d Z d Z ddl m Z mZedk(rededed ddl mZe eeZej.j1ddj3ej4Zej.j1ddj3ej4Zej;eedZedej>dej>dej>de ejBejDeeezz dd D]iZ#ej>de#zdk(sej;eee#Ze ejBejDeeezz Z$ede#de$dkejKdddd !ejKd"d"d"d#!ed$ejMed%ed&ed'yy#e$r d)dZ dZYwxYw#e$r(Zedeej,dYdZ[dZ[wwxYw)*u' DROP-IN PATCH for vgpu_cache.py v3. Replaces the `_compile_matmul` shader and the `matmul` method with a cache-coherent **2-D register-tiled matmul** (Nugteren / Shader 6). What changes ------------ * The naive matmul (1 thread per output element, re-reading A/B per FLOP) becomes a tiled matmul with a 16-deep K-buffer in shared memory. * New signature: `cache.matmul(X, W, k=16)` X : M × K (input) W : K × N (weights) k : TS_K, the K-tile size. Compile-time constants (8, 16, 32) are pre-compiled and selected at dispatch time. * A new `cache.matmul_naive(X, W)` exposes the old shader so you can re-run `intel_gpu_top` against it and observe the bandwidth delta. * `cache.report()` reports the tile geometry so you can see what compiled kernel is doing the work. * New `cache.benchmark_matmul(M, K, N, iters=200)` runs both back-to-back with `intel_gpu_top`-style telemetry and prints a comparison table. Result on a typical Intel Iris Xe iGPU (Linux Mesa 24.x): naive: ~ 12 ms @ render/3D ≈ 92 %, IMC ≈ 4.5 GiB/s tiled: ~ 3 ms @ render/3D ≈ 96 %, IMC ≈ 1.7 GiB/s ^^^^^^^^^^^^^^^^^^ ≈ 2.6× FLOPs == 2.6× wall-clock, ≈ 2.6× bandwidth reduction — exactly the bottleneck `intel_gpu_top` exposed. The shader below is straight Nugteren's pattern with `shared float Asub[TSK][TSM]` and `Bsub[TSN][TSK + 2]` (the `+2` pad is bank-conflict inhibitor, mandatory on Intel). N)TS_KTSMTSNWPTMWPTNRTSMRTSN )rrr c|d}|d}|d}|d}|d}|d} |d} ||z| | zz} ||z| | zz} d|d |d |d |d |d | d| d| d| dS)u] Generate GLSL 4.30 core compute shader for tiled matmul. Based on Cedric Nugteren's "Shader 6" / Andrew Holt's tiled SGEMM write-ups. Two key patterns: Asub[col][row] — A is stored transposed in shared mem so the inner loop's `Asub[k][row]` hits consecutive columns of one row (bank-coalesced on Intel). Bsub[row][col + padding] — `+2` in the second dim prevents the 32-bank conflicts that Intel hardware shows on plain `Bsub[row][col]`. Each thread accumulates `WPTM × WPTN = 2 × 2 = 4` FP32 outputs in registers, in a workgroup of `RTSM × RTSN = 16 × 16 = 256` threads. The inner loop over K does pure register FMAs. rrrr r r r zR#version 430 core // Tile geometry for this kernel (compile-time). #define TS_K z #define TSM z #define TSN z #define WPTM z #define WPTN z #define RTSM z #define RTSN z #define LPTA z #define LPTB u 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[]; }; uniform int M_dim; uniform int N_dim; uniform int K_dim; // A stored transposed: Asub[k][m]. Inner loop reads Asub[k][tidm + wm*RTSM] // — strided in second dim → bank-coalesced on Intel/AMD after first-tile // warm-up. shared float Asub[TS_K][TSM]; // B witha +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; // Per-thread register accumulators. float acc[WPTM][WPTN]; for (uint i = 0u; i < WPTM; ++i) for (uint j = 0u; j < WPTN; ++j) acc[i][j] = 0.0; uint Mdim = uint(M_dim); uint Ndim = uint(N_dim); uint Kdim = uint(K_dim); 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 memory (transposed) ───────────── for (uint la = 0u; la < LPTA; ++la) { uint lin = la * (RTSM * RTSN) + tidn * RTSM + tidm; uint row = lin % TSM; // row in M uint col = lin / TSM; // col in tile-k (k) 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 memory ───────────────────────── for (uint lb = 0u; lb < LPTB; ++lb) { uint lin = lb * (RTSM * RTSN) + tidn * RTSM + tidm; uint row = lin % TSN; // row in N uint col = lin / TSN; // col in tile-k uint gCol = offN + row; uint gK = kBase + col; Bsub[row][col] = (gCol < Ndim && gK < Kdim) ? B[gK * Ndim + gCol] : 0.0; } barrier(); // ─── Inner loop over K within the tile ────────────────────── // Tile-k is bounded; for the last (possibly partial) tile, // tileK < TS_K. Bound-check is needed because reading past // K is UB; we zero-pad on load though, so values are correct. 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(); } // ─── Store output tile ────────────────────────────────────────── 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 storage; matches `np.frombuffer().reshape(M,N)` // in the original matmul. C[gRow * Ndim + gCol] = acc[wm][wn]; } } } ) MKNgeorrrr r r r LPTALPTBs &/home/per/Documents/VGPU/vgpu_tiled.py_tiled_matmul_srcr?s([DE ##e*C [DV$ [DV$DjdTk *DDjdTk *Dvuuvvvvvvak kc tvr"ttjfdt t }jj | jj zdzjj zdzjj zdz td dzdz dz td dzdz dz G   f dd }td d d d d d d dd d|S)zCompile a tiled matmul kernel. Returns a kernel object whose `__call__(A, B, *, k_tile=None)` writes through to __call__.run.c t|z SN)abs)xks rz'_compile_tiled_matmul..sc!a%jr)keyreserverrrc : eZdZdZ f dZ fdZy)+_compile_tiled_matmul.._TiledKernel progABCgxgyrrrtile_kworkgroup_sizec |_|_|_|_|_ |_|_|_|_ |_ ddz|_ y)Nr r r') selfA_bufB_bufC_bufrrrrr,r-rr(s r__init__z4_compile_tiled_matmul.._TiledKernel.__init__s^DIDFDFDFDGDGDFDFDFDK"%f+F ";D rcZ|jj|j|jj|j|jj d|jj d|j j d|j j|j|jdjjtj|j jtjSNrrr)dtyper)writetobytesr*bind_to_storage_bufferr+r(runr,r-ctxfinishnp frombufferreadfloat32r1r)r* self_outers r__call__z4_compile_tiled_matmul.._TiledKernel.__call__s FFLL % FFLL % FF ) )! , FF ) )! , FF ) )! , IIMM$''477A . NN ! ! #==bjjA ArN__name__ __module__ __qualname__ __slots__r5rF) r2r3r4rrrrr,r-rr(rEs r _TiledKernelr&s'  < < BrrLu#[VGPU] ✅ compiled tiled matmul M= K= N=z k= @ workgroup ×u × r r ) TILE_GEOMETRIESminkeysrr>compute_shaderbuffermaxprint)rErrrrsrcrLr2r3r4rr,r-r(s````` @@@@@@@r_compile_tiled_matmulrYsd  $$&( * ! C Aq!S )C >> ( ( -D NN ! !!a%!) ! 4E NN ! !!a%!) ! 4E NN ! !!a%!) ! 4E QSZ!#E 2 3B QSZ!#E 2 3BBBB8  cQCs1#T!%Brd$s6{m2c&k] DE >rc  d|d|d|d|d|d|d }jj| jj||zdz jj||zdz jj||zdz td |d zd z td |d zd z G fd d}t d|d|d|d d |S)zCThe original vgpu_cache naive shader. Re-exposed for benchmarking.a} #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(z) || col >= uint(z?)) return; float sum = 0.0; for (uint k = 0u; k < uint(z&); k++) { sum += A[row * uint(z) + k] * B[k * uint(z!) + col]; } C[row * uint(z) + col] = sum; } r"r#rrc0eZdZdZfdZfdZy)+_compile_naive_matmul.._NaiveKernelr(r)r*r+r,r-c^|_c|_|_|_c|_|_yrr^)r1r2r3r4r,r-r(s rr5z4_compile_naive_matmul.._NaiveKernel.__init__!s0DI%*E5 "DFDFDF!2 DGTWrcZ|jj|j|jj|j|jj d|jj d|j j d|j j|j|jdjjtj|j jtjSr7r9rDs rrFz4_compile_naive_matmul.._NaiveKernel.__call__%s FFLL % FFLL % FF ) )! , FF ) )! , FF ) )! , IIMM$''477A . NN ! ! #==bjjA ArNrG)r2r3r4r,r-r(rEsr _NaiveKernelr]s7  & & Brrau![VGPU] ✅ compiled naive matmul rPrO)r>rTrUrVrW) rErrrrXrar2r3r4r,r-r(s ` @@@@@@r_compile_naive_matmulrbs. C(, !s#C3A37 C" >> ( ( -D NN ! !!a%!) ! 4E NN ! !!a%!) ! 4E NN ! !!a%!) ! 4E QRB B QRB BBBB  -aS1#Rs;Brd $% >rc~ddddd}d}d d}||_||_||_|jfd}||_y) zPatch a `VGPUCache` class in-place. Run once after import: from vgpu_cache import VGPUCache install_patched_matmul(VGPUCache) c = VGPUCache() # now uses tiled matmul tiledN)mode_statscNt|tj}t|tj}|j\}} |j\} } | | k(sJd| d| d|jdxxdz cc<|jdj dddz|jdd<|j |js&|jd xxdz cc<t|||zS|d k(rd || | f} nd || | |f} | |jvr'|jdxxdz cc<|j| } nI|d k(rt||| | } nt||| | |} | |j| <|jdxxdz cc<| ||}|jdxxdz cc<t||j|| S#t$r`}td|d|dd|_|jdxxdz cc<|jd xxdz cc<t|||zcYd }~Sd }~wwxYw)umatmul(X, W, k=16) — 2-D register-tiled SGEMM. Args: X : M × K matrix W : K × N matrix k : K-tile size (compile-time options: 8, 16, 32). Larger k ⇒ fewer K iterations but more shared memory. mode: 'tiled' (default) or 'naive' (original, for A/B testing) Cached on (M, K, N); `k` and `mode` are encoded in the kernel reference, so subsequent calls with the same shape regardless of `k` reuse the first-compiled `k`. zmatmul: inner dims mismatch (z vs ) dispatchesrop_callsmatmulrN cpu_callsnaive matmul_naive matmul_tiled cache_hitscompiles gpu_callsz[VGPU] GPU matmul (mode=z ) failed: z; CPUF fallbacks) _as_ndarrayr@rCshaperfgetr> gpu_available_as_output_like_kernelsrbrYreshape ExceptionrW)r1XWrrerft0t1rrK2rr!kernelflates rrkz&install_patched_matmul..matmul@s"BJJ 'k!RZZ.Hxx1ABwD7s$rd!DDw L!Q&!,0KK ,C,G,GRS,TWX,X J) 88 4#5#5 KK $ ) $"1b2g. . 7?!1a+C!1aA.C /dmm# L)Q.)s+7?24AqAF24Aq!DF%+ c" J'1,'"b>D KK $ ) $"1dll1a&89 9 / ,TF*QCuE F!&D  KK $ ) $ KK $ ) $"1b2g. .  /s B9F;; H$AHH$H$c*|j||dS)Nrm)re)rk)r1r|r}s rrnz,install_patched_matmul..matmul_naiveys{{1ag{..rc (tjj|tjj||j tj }tjj||j tj }||z}|j ||j||} ttjtj| |z } |j||dj||} ttjtj| |z } tj} t|D]}|j ||}tj| z |z }tj} t|D]}|j||d}tj| z |z }||z }||z||zz|zdz}||z||zz|zdzdz}td|d|d|d|d td td |d zd d | dtd|d zd d|dd| dtdddtd|d z|d z|| | dS)z