From f70adc8202e32449dccfee30edf6118250d833ce Mon Sep 17 00:00:00 2001 From: itchyfeetleech <169226071+itchyfeetleech@users.noreply.github.com> Date: Mon, 19 Oct 2026 19:50:41 +1300 Subject: [PATCH] engine: preset L's expansion stage at two windows per wave layer_m.comp's L expansion (KIND 4) wrote a window's block output to shared memory and ran the per-pixel head after a workgroup barrier, both inside the windows' loop, so it could not walk two windows per wave (#error for KIND >= 3). Every window now writes its output first and the head runs once after the barrier; L's input stage (KIND 3) gathers its features once, before its windows' loop, the same way. At one window per wave RADV emits the same machine code for both stages as before. WPW > 2 stays an error: the head and the gather need the block's 64 pixels on 64 invocations. compile_m.py builds dec0_l (and dec0_l_planes) with two windows per wave, which load each weight tile once for both: 96 -> 216 vector registers, 14 -> 7 waves per SIMD. Not under M_EXACT=1, like the network layers (it would gain 7 us there, 516 -> 509); M_DEFS_DEC0L=-DWPW=1 builds the previous layout. The input stage at two windows per wave is identical too but slower (452 -> 465 us; 144 -> 256 registers, 10 -> 5 waves), so it stays at one. Every output frame byte-identical (L, HDR and LDR, 1280x720 -> 1440p, 1114x626 -> 1080p, 1920x1080 -> 1440p, the stress inputs) and so is everything D4R_M_DUMP writes - with the new default, with both stages at one window and at two, and in the exact build (one window, and the expansion at two). The windows' outputs swapped in the head's staging changes every frame. Preset M's shaders are unchanged. RX 9070 XT, d4r-m L 720p -> 1440p, interleaved: expansion 383 -> 366 us (1114x626 -> 1080p: 342 -> 326), frame 2205 -> 2191 us. Signed-off-by: itchyfeetleech <169226071+itchyfeetleech@users.noreply.github.com> --- engine/compile_m.py | 4 +++- engine/layer_m.comp | 18 +++++++++++------- engine/test_m.py | 2 +- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/engine/compile_m.py b/engine/compile_m.py index b35e38c..66b4ac7 100644 --- a/engine/compile_m.py +++ b/engine/compile_m.py @@ -32,6 +32,8 @@ from test_m import pack_layer, defines EXACT = os.environ.get('M_EXACT') == '1' FAST = [] if EXACT or '-DM_RELAXED' in os.environ.get('M_DEFS', '').split() else ['-DM_RELAXED'] FAST_WPW = [] if EXACT or os.environ.get('M_WPW') or 'WPW=' in os.environ.get('M_DEFS', '') else ['-DWPW=2'] +# preset L's expansion stage (dec0l) two windows per wave as well; its input stage (enc0l) is slower that way +L_WPW = [] if EXACT or 'WPW=' in os.environ.get('M_DEFS', '') + os.environ.get('M_DEFS_DEC0L', '') else ['-DWPW=2'] SHADERS = {'enc1': 'm_enc1', 'enc2': 'm_enc2', 'enc3_tube': 'm_tube', 'dec2': 'm_dec2', 'dec1': 'm_dec1'} @@ -159,7 +161,7 @@ def compile_l(a, one, R): for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW) for planes, suffix in ((0, ''), (1, '_planes')): glsl('layer_m.comp', f'enc0_l{suffix}.spv', defines('enc0l') + ldr + [f'-DPLANES={planes}'] + FAST) - glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST) + glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST + L_WPW) # the reconstruction one render texel wider per tier than M's (4x4 taps): 11 (reconstruction/render >= 3), # 16 (the whole post_3_2 range), 12 (the compact tier: render/reconstruction <= 0.45, jitter within half a texel) for out, tile in (('post_l', 11), ('post_l_q', 16), ('post_l_c', 12)): diff --git a/engine/layer_m.comp b/engine/layer_m.comp index 36bd42c..84e673c 100644 --- a/engine/layer_m.comp +++ b/engine/layer_m.comp @@ -64,8 +64,8 @@ #ifndef WPW #define WPW 1 #endif -#if WPW > 1 && KIND >= 3 -#error layer_m.comp: L's input and expansion stages stage their pixels per thread (WPW = 1) +#if WPW > 2 && KIND >= 3 +#error layer_m.comp: L's input and expansion stages handle their 64 pixels one per invocation (WPW <= 2) #endif layout(local_size_x = 128 / WPW) in; // o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8); @@ -496,8 +496,7 @@ void main() #endif #if KIND == 3 // ------------------------------------------------------------ L's input: features, embedding -> x0 - FORW { - WINV + { const float e1 = ep.expo == 0.0 ? 1.0 : ep.expo, e2 = ep.expoPrev == 0.0 ? 1.0 : ep.expoPrev; const uint li = gl_LocalInvocationID.x; if (li < 64u) @@ -510,7 +509,10 @@ void main() const uint at = (((py >> 2) * 2u + (px >> 2)) * 16u + (py & 3u) * 4u + (px & 3u)) * FS; for (uint k = 0u; k < 32u; ++k) FEW[at + k] = fe[k]; } - SYNC + } + SYNC + FORW { + WINV // the embedding has the shape of an MLP chunk: 32 features -> 32 hidden (ReLU, FP8) -> C B8 f0, f1; coopMatLoad(f0, FEW, mt * 16u * FS, FS, CM); @@ -532,8 +534,8 @@ void main() for (int i = 0; i < 8; ++i) X0S[xb + 16u * (N) + uint(i)] = float16_t(x_[i]); } EACH_T(EMB_N) #undef EMB_N - SYNC } + SYNC #endif // ------------------------------------------------------------ norm 1: sum of squares, residual and h1 FORW { @@ -768,7 +770,9 @@ void main() const uint hb = (mt * 16u + row) * HS + ch8; for (uint nt = 0u; nt < NT; ++nt) for (int i = 0; i < 8; ++i) HO[hb + 16u * nt + uint(i)] = float16_t(res[nt][i]); - SYNC + } + SYNC + { const uint li = gl_LocalInvocationID.x; const uint px = li & 7u, py = li >> 3; const ivec2 at = ivec2(X0 + int(px), Y0 + int(py)); diff --git a/engine/test_m.py b/engine/test_m.py index 257e4c0..ffa273e 100644 --- a/engine/test_m.py +++ b/engine/test_m.py @@ -90,7 +90,7 @@ def defines(name): # compile_m.py's arithmetic (ZLUDA's lowering); M_TEXL=0 gives the native accuracy kernels' for the network layers. # M_DEFS: extra defines for every layer_m.comp stage; M_DEFS_ENC1, M_DEFS_ENC3_TUBE, ...: for one layer texl = ['-DTEXL=1'] if kind >= 3 or os.environ.get('M_TEXL', '1') != '0' else [] - # M_WPW=2: two 4x4-token windows per wave in the network layers (layer_m.comp WPW; not L's input / expansion) + # M_WPW=2: two 4x4-token windows per wave in the network layers (layer_m.comp WPW; L's input / expansion: M_DEFS_*) wpw = [f"-DWPW={os.environ['M_WPW']}"] if kind < 3 and os.environ.get('M_WPW') else [] return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={kind}', f'-DXC={NPM or CIN or 16}'] + texl + wpw + os.environ.get('M_DEFS', '').split() + os.environ.get('M_DEFS_' + name.upper(), '').split()