engine: opt-in M_RESIDUAL_F32, the SR residual stream in f32 across the MLP

With M_RELAXED, layer_m.comp's MLP adds every 32-channel chunk's products to
the binary16 residual stream and rounds it again: one f32 -> f16 -> f32 round
trip per chunk and element, 4 to 16 chunks a layer. M_RESIDUAL_F32 seeds an
f32 accumulator from the residual before the MLP, adds every chunk to it
with plain multiply-adds and rounds to binary16 once after the last chunk.
It needs M_RELAXED (#error otherwise); compile_m.py builds it into every
layer_m.comp stage with M_RESIDUAL_F32=1 and refuses it with M_EXACT=1.

Off by default: the default build's shaders (M, L, HDR and LDR) and the
exact build's are byte-identical to the previous ones.

Against the default build, final RGB, peak 1, 1280x720 -> 1440p: six smooth
temporal frames M 43.3 - 45.9 dB (L 44.6 - 49.2), at most 1.5 off; four
frames of the noise stress input M 17 - 18 dB (L 19 - 22). The exact build is
as far from the default (M 43.2 - 44.6 dB, L 44.5 - 47.0; noise M 16.7 - 17.9,
L 19 - 20) and as far from this one (M 43.4 - 44.8, L 44.6 - 47.0): no
further from NVIDIA's arithmetic than the default.

RX 9070 XT, d4r-m 720p -> 1440p, interleaved: M 1521 -> 1410 us (network
953 -> 841), L 2195 -> 2077 us (network 944 -> 829).

Signed-off-by: itchyfeetleech <[email protected]>

engine/compile_m.py

@@ -34,6 +34,11 @@ FAST = [] if EXACT or '-DM_RELAXED' in os.environ.get('M_DEFS', '').split() else
 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']
+# M_RESIDUAL_F32=1 (opt-in, not with M_EXACT=1): every layer_m.comp stage keeps the residual stream f32 across the MLP's
+# 32-channel chunks, one binary16 rounding after them instead of one per chunk (engine/README.md)
+RESF32 = ['-DM_RESIDUAL_F32'] if os.environ.get('M_RESIDUAL_F32') == '1' else []
+if RESF32 and EXACT:
+    raise SystemExit('M_RESIDUAL_F32=1 is part of the relaxed arithmetic: not with M_EXACT=1')
 SHADERS = {'enc1': 'm_enc1', 'enc2': 'm_enc2', 'enc3_tube': 'm_tube', 'dec2': 'm_dec2', 'dec1': 'm_dec1'}
 
 
@@ -107,7 +112,7 @@ def main():
     # the network in the arithmetic of ZLUDA's lowering of NVIDIA's PTX (TEXL): byte-identical to the CUDA backend
     # with NVIDIA's translated kernels in the harness (2026-10-11), where the native accuracy kernels' K32_MIX order
     # drifted to 36-41 dB through the FP8 tube
-    for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW)
+    for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW + RESF32)
     for s in ('enc0_m', 'dec0_m', 'post_m', 'down_m'): glsl(s + '.comp', s + '.spv', ['-DLDR=1'] if a.ldr and s != 'dec0_m' else [])
     glsl('down_m.comp', 'down_m_rgb10.spv', ['-DOUTPUT_RGB10=1'] + (['-DLDR=1'] if a.ldr else []))
     glsl('post_m.comp', 'post_m_q.spv', ['-DTILE_W=15', '-DTILE_H=15'] + (['-DLDR=1'] if a.ldr else []))
@@ -158,10 +163,10 @@ def compile_l(a, one, R):
         subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, str(HERE / src), '-o', str(a.output / out)], check=True, stdout=subprocess.DEVNULL)
     ldr = ['-DLDR=1'] if a.ldr else []
     # the network in the arithmetic of ZLUDA's lowering of NVIDIA's PTX (TEXL), as for preset M (L_NOTES.md)
-    for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW)
+    for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW + RESF32)
     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'enc0_l{suffix}.spv', defines('enc0l') + ldr + [f'-DPLANES={planes}'] + FAST + RESF32)
-        glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST + L_WPW)
+        glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST + L_WPW + RESF32)
     # 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

@@ -725,6 +725,20 @@ void main()
     EACH_T(N2)
     }
     // ------------------------------------------------------------ the MLP, 32 hidden channels at a time
+    // M_RESIDUAL_F32 (opt-in, part of M_RELAXED): the residual stays f32 across the MLP's chunks and is rounded to
+    // binary16 once after them, instead of after every chunk
+#ifdef M_RESIDUAL_F32
+#ifndef M_RELAXED
+#error layer_m.comp: M_RESIDUAL_F32 needs M_RELAXED
+#endif
+    ACC RF_[WPW][NT];
+#define RF_IN(N) RF_[w][N] = ACC(res[N]);
+    FORW { EACH_T(RF_IN) }
+#undef RF_IN
+#define K32R(N, a0, b0, a1, b1) { RF_[w][N] = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), RF_[w][N])); }
+#else
+#define K32R(N, a0, b0, a1, b1) K32H(res[N], a0, b0, a1, b1)
+#endif
     for (uint c = 0u; c < UC / 8u; ++c)
     {
         const uint ow1 = p.o[20u + 3u * c], ob1 = p.o[21u + 3u * c], ow2 = p.o[22u + 3u * c];
@@ -759,9 +773,15 @@ void main()
         A8 a0, a1; \
         LDA8(a0, ow2, 0u, UC, 16u * (N)); \
         LDA8(a1, ow2, 1u, UC, 16u * (N)); \
-        FORW { K32H(res[N], a0, G0_[w], a1, G1_[w]) } }
+        FORW { K32R(N, a0, G0_[w], a1, G1_[w]) } }
         EACH_T(MLP2_N)
     }
+#undef K32R
+#ifdef M_RESIDUAL_F32
+#define RF_OUT(N) res[N] = HAC(RF_[w][N]);
+    FORW { EACH_T(RF_OUT) }
+#undef RF_OUT
+#endif
 #if KIND == 4
     // ------------------------------------------------------------ L's expansion head: per-pixel reconstruction inputs
     SYNC   // every wave is done with the V tiles in HO