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 <[email protected]>

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)):

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
+    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
+    {
         const uint li = gl_LocalInvocationID.x;
         const uint px = li & 7u, py = li >> 3;
         const ivec2 at = ivec2(X0 + int(px), Y0 + int(py));

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()