engine/rr: run dec0_1's translated tail on two waves

dec0_1's translated output tail ran on wave 0 while the block's other three
waves had finished. It is a prologue (y = r + MLP * s2 and the shuffles that
give each lane its two pixels) followed by the same per-pixel code twice.

gen_rr2_enc0.py --tail-waves N gives each of N waves its own copy of the
tail with only its N-th of the stores (in order: the first pixel's six, then
the second's), and the compiler drops what a copy's stores do not need.
Every stored value comes from the same instructions on the same values as on
one wave. With one window per workgroup, the window index no longer adds
gl_SubgroupID (it was 0 for the only valid wave).

compile_rr2.py uses two waves (D4R_RR2_TAIL_WAVES=1 for one): 120 -> 96
vector registers, 12 -> 16 waves per SIMD; RX 9070 XT, 720p -> 1440p:
dec0_1 565 -> 474 us. Four waves, splitting each pixel's stores, is slower
(529 us). Every image byte-identical on the replay inputs (also robustness2,
no robust access and the split gather build).

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

engine/rr/gen_rr2_enc0.py

@@ -295,7 +295,25 @@ def main():
             '    // the gather\'s index arithmetic the tail reads\n' + '\n'.join(recomputed) + '\n' +
             '\n'.join(load + keep) + '\n' + '\n'.join(tail) + '\n    }\n}\n')
         return
-    wx = (lambda t: t) if one_wave else (lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)'))
+    # --tail-waves N (dec0_1, one window per workgroup): the translated tail on N of the block's waves instead of one.
+    # Each wave runs its own copy of the tail with only the N-th part of its stores (in order: the first pixel's,
+    # then the second's); the compiler drops what a copy's stores do not need. Every stored value comes from the
+    # same instructions on the same values as on one wave, so the outputs are the same.
+    TW = int(sys.argv[sys.argv.index('--tail-waves') + 1]) if '--tail-waves' in sys.argv else 1
+    if TW > 1 and not (dec and G == 1):
+        raise SystemExit('--tail-waves is for dec0_1 with one window per workgroup')
+    wx = (lambda t: t) if one_wave or TW > 1 else (lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)'))
+    STORE = re.compile(r'^\s*(?!if\b|for\b|while\b|switch\b|return\b)[A-Za-z_]\w*\(')
+    stores = [i for i, l in enumerate(tail) if STORE.match(l)]
+    if TW > 1 and len(stores) % TW:
+        raise SystemExit(f'--tail-waves {TW}: {len(stores)} stores do not split evenly')
+    def tail_part(w):
+        mine = set(stores[w * len(stores) // TW:(w + 1) * len(stores) // TW])
+        return [l for i, l in enumerate(tail) if i not in stores or i in mine]
+    tails = ''.join(f'    if (rr2_valid && gl_SubgroupID == {w}u) {{\n' + '\n'.join(load + keep) + '\n' + '\n'.join(tail_part(w)) + '\n    }\n'
+                    for w in range(TW)) if TW > 1 else None
+    valid = (f'gl_SubgroupID < {TW}u && gl_WorkGroupID.x < rr2_gx' if TW > 1 else
+             f'gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx')
     text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' +
             ('#define RR2_XF_SHARED 1\n' if one_wave else '') + block + '\n' +
             ('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(rr2_msh(row, col))); }\n'
@@ -304,15 +322,16 @@ def main():
             'void main()\n{\n' + '\n'.join(body[:nd]) +
             '\n    // this wave\'s window, when the level has it\n'
             '    const uint rr2_gx = (param.w[0] + (param.w[20] & 0xffffu) + 7u) / 8u;\n'
-            f'    const bool rr2_valid = gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx;\n'
+            f'    const bool rr2_valid = {valid};\n'
             '    if (rr2_valid) {   // NVIDIA\'s warp: the translation before the first product\n' + wx('\n'.join(pre)) + '\n    }\n' +
             ('' if dec else '\n'.join(store)) +
             # (the block runs unconditionally, also for a window past the level's last one: it writes nothing there, and
             # inside a branch the driver compiles it much slower)
             ('\n    barrier();\n    rr2_block();\n    barrier();\n' if G == 1 else
              f'\n    barrier();\n    for (uint w_ = 0u; w_ < {G}u; ++w_) {{\n        rr2_win = w_;\n'
-             '        rr2_block();\n        barrier();\n    }\n') + '    if (rr2_valid) {\n' +
-            '\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n    }\n}\n')
+             '        rr2_block();\n        barrier();\n    }\n') +
+            (tails + '}\n' if tails else
+             '    if (rr2_valid) {\n' + '\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n    }\n}\n'))
     Path(out).write_text(text)
     print(f'{out}: {len(pre)} gather lines, {len(cut)} cut, {len(tail)} tail lines, {len(seen)} A fragment sets')