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/README.md

@@ -557,6 +557,13 @@ registers to the waves its shared memory leaves room for (enc1 needed 97 vector
   so the check there is the whole instruction mix: only `v_and_or_b32` / `v_or` / `v_perm` / shifts go (post -16
   VALU, down -9, downp -73), every `v_fma_f32` / `v_mul_f32` / `v_add_f32` / conversion / `v_exp` / `v_log` count is
   unchanged (post moves 5 of its 57 mixes from the low to the high half of their registers). Output 673 -> 665 us.
+- **dec0_1's translated tail on two waves** (`gen_rr2_enc0.py --tail-waves 2`). The tail ran on wave 0 while the
+  block's other three waves had nothing left to do; it is a prologue (y = r + MLP * s2 and the shuffles that give each
+  lane its two pixels) and then the same per-pixel code twice. Each of two waves now runs its own copy of the tail with
+  only one pixel's stores (the first six, the last six), and the compiler drops what a copy's stores do not need:
+  every stored value comes from the same instructions on the same values. 120 -> 96 vector registers, 12 -> 16 waves
+  per SIMD; dec0_1 565 -> 474 us. Four waves (each pixel's stores split in two) is slower, 529 us: the per-pixel work
+  the parts share is then done twice. `D4R_RR2_TAIL_WAVES=1` builds the single-wave tail.
 - **One unsigned comparison per axis in the translation's bounds-checked fetch** (`ptx2rr.py` `fetch0`: a negative
   coordinate is above every size): 231 instructions fewer in enc0, whose gather fetches 102 times, its scalar spills
   200 -> 176; enc0 745 -> 738 us.

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