From e94c4018e7f415d26922235df4d75b9f8eab12b6 Mon Sep 17 00:00:00 2001 From: itchyfeetleech <169226071+itchyfeetleech@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:22:03 +1300 Subject: [PATCH] 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 <169226071+itchyfeetleech@users.noreply.github.com> --- engine/compile_m.py | 13 +++++++++---- engine/layer_m.comp | 22 +++++++++++++++++++++- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/engine/compile_m.py b/engine/compile_m.py index 66b4ac7..a5293bf 100644 --- a/engine/compile_m.py +++ b/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'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST + L_WPW) + 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 + 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)): diff --git a/engine/layer_m.comp b/engine/layer_m.comp index 84e673c..2d85faa 100644 --- a/engine/layer_m.comp +++ b/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