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/compile_rr2.py

@@ -71,11 +71,14 @@ KERNEL['downp'] = 'rr2_downsample_kernel_static_hdr'
 BLOCK['downp'] = '16,16,1'
 BUFFERS['downp'] = [72, 168, 176]
 # with --native: the translation of the per-pixel code around the native block (engine/rr/gen_rr2_enc0.py)
-HYBRID = {'enc0': [], 'dec0_1': ['--dec0_1']}
+HYBRID = {'enc0': [], 'dec0_1': ['--dec0_1', '--tail-waves', os.environ.get('D4R_RR2_TAIL_WAVES', '2')]}
 # windows per workgroup of the hybrid stages (gen_rr2_enc0.py --group; D4R_RR2_GROUP to experiment)
 GROUP = int(os.environ.get('D4R_RR2_GROUP', '1'))
-# encoder 0's gather as a pass of its own (D4R_RR2_SPLIT=0: one shader, the gather on one wave of the block's four)
+# encoder 0's gather as a pass of its own, which pays with the four-wave block (one shader would run the gather on one
-SPLIT = os.environ.get('D4R_RR2_SPLIT', '1') != '0'
+# wave of the four). With the one-wave block (level 3) the gather runs in the block's own shader and hands its
+# features over in shared memory instead (rr2_w.comp RR2_XF_SHARED): no record in the arena, one dispatch less.
+# D4R_RR2_SPLIT=1 / 0 forces the gather pass / the single shader.
+SPLIT = os.environ.get('D4R_RR2_SPLIT', '0' if RELAXED >= 3 and 'enc0' in ONE_WAVE else '1') != '0'
 # The pooled histories as RGBA32F holding q / 1023, as ZLUDA backs them (16 bytes a pixel). D4R_RR2_POOLED=rgb10_a2
 # makes them real UNORM 10:10:10:2 images (0.1 ms faster at 4K), but the hardware's bilinear filter then rounds
 # differently: identical at 3x and 2x scaling, 44 - 51 dB on temporal frames at 1.5x (capFQ, 2026-10-13)
@@ -165,6 +168,10 @@ def main():
         cmd += [f'--buffer={off}' for off in BUFFERS.get(stage, [])]
         cmd += [f'--snorm16={off}' for off in SNORM16.get(stage, [])]
         cmd += [f'--surface={off}={fmt}' for off, fmt in SURFACES.get(stage, {}).items()]
+        # the hybrid stages' and the output passes' modules go through spv_exact_convert.py, so their conversions need
+        # only the lighter guard (ptx2rr.py --exact-convert)
+        exact = (bool(a.native) and stage in HYBRID) or not network
+        cmd += ['--exact-convert'] if exact else []
         subprocess.run(cmd, check=True, stderr=subprocess.DEVNULL)
         info = json.loads(generated.with_suffix('.json').read_text())
         weights = '-'
@@ -191,7 +198,7 @@ def main():
                 subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / 'enc0g.spv')],
                                check=True, stdout=subprocess.DEVNULL)
                 record = int((work / 'enc0_g.record').read_text())
-            if RELAXED >= 3 and stage in ONE_WAVE and split:
+            if RELAXED >= 3 and stage in ONE_WAVE and (split or GROUP == 1):
                 how = how + ['--one-wave']
             subprocess.run([sys.executable, str(HERE / 'rr' / 'gen_rr2_enc0.py'), str(find_module(a.ptx_dir, kernel)),
                             str(generated), str(work / f'{stage}_m.comp'), *HYBRID[stage], *how],
@@ -204,6 +211,9 @@ def main():
         else:
             subprocess.run(capped + ['glslangValidator', '-V', '--target-env', 'vulkan1.3', str(generated), '-o',
                                      str(out / f'{stage}.spv')], check=True, stdout=subprocess.DEVNULL)
+            if exact:
+                subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}.spv')],
+                               check=True, stdout=subprocess.DEVNULL)
         if hybrid:
             pass
         elif network:

engine/replay_rr.cpp

@@ -185,7 +185,9 @@ int main(int argc, char** argv) try {
     HostBuffer staging(c, VkDeviceSize(ow) * oh * 8 * 4);
     // GPU time of every frame: REngine's four stamps (start, after exposure, after the network, after the output)
     VkQueryPoolCreateInfo qp{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO};
-    qp.queryType = VK_QUERY_TYPE_TIMESTAMP; qp.queryCount = 4;
+    qp.queryType = VK_QUERY_TYPE_TIMESTAMP; qp.queryCount = 32;
+    // D4R_ENGINE_RR_STAGE_STAMPS=1: the time of every network stage and output group as well (best frame of the repeats)
+    const bool stageTimes = std::getenv("D4R_ENGINE_RR_STAGE_STAMPS") != nullptr;
     VkQueryPool stamps; CK(vkCreateQueryPool(c.device, &qp, nullptr, &stamps));
     VkPhysicalDeviceProperties props; vkGetPhysicalDeviceProperties(c.physical, &props);
     const int repeat = std::getenv("D4R_RR_TIME_REPEAT") ? std::atoi(std::getenv("D4R_RR_TIME_REPEAT")) : 1;
@@ -233,21 +235,36 @@ int main(int argc, char** argv) try {
         engine.setFrame(images, s);
         if (first) c.submit([&](VkCommandBuffer cb) { model.recordUpload(cb); engine.recordInitialize(cb); first = false; });
         double best = 1e30, network = 0, output = 0, exposure = 0;
+        std::vector<double> stageBest(28, 1e30);
         for (int r = 0; r < repeat; ++r) {
             if (r) engine.setFrame(images, s);
             c.submit([&](VkCommandBuffer cb) {
-                vkCmdResetQueryPool(cb, stamps, 0, 4);
+                vkCmdResetQueryPool(cb, stamps, 0, 32);
                 engine.recordFrame(cb, stamps, 0);
             });
             uint64_t t[4];
             CK(vkGetQueryPoolResults(c.device, stamps, 0, 4, sizeof(t), t, 8, VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT));
             const double ms = props.limits.timestampPeriod * 1e-6;
+            if (stageTimes) {   // per stage: the minimum over the repeats (stamps 1, then 4 ... as written)
+                uint64_t all[32][2];
+                vkGetQueryPoolResults(c.device, stamps, 0, 32, sizeof(all), all, 16, VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WITH_AVAILABILITY_BIT);
+                uint64_t prev = all[1][0];
+                for (uint32_t q = 4; q < 32 && all[q][1]; ++q) {
+                    stageBest[q - 4] = std::min(stageBest[q - 4], (all[q][0] - prev) * ms);
+                    prev = all[q][0];
+                }
+            }
             if ((t[3] - t[0]) * ms < best) {
                 best = (t[3] - t[0]) * ms; exposure = (t[1] - t[0]) * ms; network = (t[2] - t[1]) * ms; output = (t[3] - t[2]) * ms;
             }
         }
         std::cout << "frame " << f << ": " << best << " ms (exposure " << exposure << ", network " << network
                   << ", output " << output << ")\n";
+        if (stageTimes) {
+            std::cout << "stages (us):";
+            for (double t : stageBest) if (t < 1e29) std::cout << " " << static_cast<int>(t * 1000 + 0.5);
+            std::cout << "\n";
+        }
         for (const auto& d : engine.debugImages()) {
             const size_t bytes = size_t(d.width) * d.height * texelBytes(d.format);
             c.submit([&](VkCommandBuffer cb) {

engine/rr/README.md

@@ -530,6 +530,65 @@ only; the other stages still use `rr2_m.comp`.
 well 4.98 / 5.59. Decoder 0's second kernel is not converted: it hands r and the MLP sum (8 KiB of binary16) to its
 translated tail through shared memory, which a single wave cannot hold without losing the 4 KiB rate.
 
+Shared memory, later (2026-10-18, RX 9070 XT, RADV, 1280x720 -> 2560x1440, `d4r-rr` with `D4R_RR_TIME_REPEAT=40`
+and `D4R_ENGINE_RR_STAGE_STAMPS=1`, five interleaved runs; every image of six temporal frames byte-identical to the
+previous model, also at 1114x626 -> 1080p, 853x480 -> 1440p and with `D4R_RR_TEST_ROBUST2=1`). ACO sizes a shader's
+registers to the waves its shared memory leaves room for (enc1 needed 97 vector registers before scheduling and got
+240, dec0_1 79 and 192), so a stage that declares less shared memory also gets fewer registers and more waves:
+- **Encoder 0 in one shader again at level 3.** With the one-wave block nothing waits while the gather runs, so
+  `gen_rr2_enc0.py --one-wave` without `--split` runs the gather, the block and the tail in one 32-thread group and
+  hands the features over in shared memory: XF shares its 4 KiB with the block's XT and KVS (`rr2_w.comp`
+  `RR2_XF_SHARED`; every feature tile is in registers before the block first writes either). No gather pass, no
+  record in the arena (76 bytes a render pixel, of which 36 MB lay past the arena at 720p). The gather now runs at the
+  block's 6 waves per SIMD instead of 16 as a pass of its own, which takes back most of the saved traffic:
+  enc0g + enc0 798 -> 772 us. `D4R_RR2_SPLIT=1` builds the gather pass.
+- **enc1's merge input shares its 8 KiB with XT and KVS** (`rr2_m.comp` KIND 1): 17 -> 9 KiB of shared memory, 240 ->
+  144 vector registers, 6 -> 10 waves per SIMD: 286 -> 266 us.
+- **dec0_1 hands r and the MLP sum to its tail in XT and KVS' place** (KIND 5, one window per workgroup): written as
+  binary16 pairs of natural channels after a barrier, read in those pairs by the tail; 14 -> 8 KiB, 192 -> 120 vector
+  registers, 8 -> 12 waves per SIMD: 660 -> 623 us.
+- **Lighter conversion guards in the hybrid stages** (enc0, dec0_1; `ptx2rr.py --exact-convert`): their modules go
+  through `spv_exact_convert.py`, whose NoContraction on every conversion already stops the fold of
+  `float(float16_t(x))`, so a `cvt.rn.f16.f32` keeps only the opaque zero on its f32 bits (against the fused
+  multiply-convert, which `precise` alone does not stop: without it ACO emits `v_fma_mixlo_f16` there). In the ISA the
+  same number of `v_cvt_f16_f32` / `v_cvt_f32_f16` as before and no `v_fma_mix*_f16`; 224 / 153 fewer instructions:
+  enc0 747 -> 729 us, dec0_1 598 -> 592. The output passes (post, down, downp, pooled) now go through
+  `spv_exact_convert.py` too and take the same guard. Their ISA already uses `v_fma_mix*_f16` for plain conversions,
+  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.
+- **Hardware saturation for the FP8 conversions of f32 values** (`rr_exact.glsl` `FINITE`). RADV lowers a saturating
+  conversion (`saturatedConvertEXT`) as a NaN-preserving clamp per element (`v_minimummaximum_f32`) and then the
+  conversion, because the hardware's FP16_OVFL saturation turns infinities into NaN instead of +-448 (ACO's comment;
+  confirmed on all 2^32 f32 inputs, RX 9070 XT: FP16_OVFL is identical to the clamp for every finite value and NaN,
+  only +-Inf differ). The values converted here never are infinite: sums of FP8 products (each at most 448^2) and
+  binary16 biases stay far below 3.4e38. Written as "an infinity becomes NaN, then saturate", the conversion is the
+  pattern `nir_opt_algebraic` turns into `f2e4m3fn_satfn` (vkd3d-proton's; Mesa >= 25.2, `has_f2e4m3fn_satfn`):
+  the hardware conversion with FP16_OVFL and no clamp. 168 - 232 clamps fewer per network stage (enc0 960 -> 128);
+  frame 3.92 -> 3.79 ms, every stage faster (dec0_1 607 -> 576 us, enc0 752 -> 727, enc1 262 -> 247). The binary16
+  conversions (the decoders' first kernels' input) keep the clamp: NIR moves the infinity test through the widening
+  and the pattern no longer matches. `RR_SAT_CLAMP=1` builds the old form.
+Frame 4.09 -> 3.84 ms. Tried without gain: the decoders' first kernels' expansion in KVS (16 /
+12 / 8 KiB, but the deep levels have too few windows for the waves to matter and the extra barrier cost 4 - 5 us
+each), all twelve tiles of a 1440p frame in one group (sixteen tile slots: the same), the output passes on 32-wide
+waves (0.67 -> 0.82 ms) and unrolled head / branch loops in the 160-channel stages (+10 - 20 us each). Per stage
+(us): enc0 772, enc1 266, enc2 137, dec1_0 246, dec1_1 208, dec0_0 405, dec0_1 623, levels 3 - 5 (eight stages) 395,
+output 670 (post 490, down 180). Levels 4 and 5 have 15 - 77 windows for 64 compute units.
+Eight waves per window for those levels (each token tile's channel tiles split over two waves, the hidden layer's
+two tiles one each, the halves meeting in shared memory; byte-identical) does not pay: enc4 / dec4_0 / dec4_1 46 / 45 /
+36 -> 49 / 45 / 40 us, enc5 35 -> 31, level 3 64 - 66 -> 78 - 86. A workgroup stays on one compute unit, so the idle
+units stay idle and the extra waves only add the shared work and barriers.
+
 ### Presets D and E against a real NVIDIA card, and level 3 for them (2026-10-14)
 
 NVIDIA's own kernels for presets D and E were replayed on an RTX PRO 4500 Blackwell (drivers 595.91 and 580.178: the

engine/rr/gen_rr2_enc0.py

@@ -15,6 +15,10 @@ out:
 With --group N (2 to 4) a workgroup holds N windows along x: waves 0 ... N - 1 run the translated code of one window
 each and the block then runs once per window; the stage's dispatch is 1 / N as wide ("xN" in rr.txt). Measured on
 an RX 9070 XT (2026-10-13) that is slower than one window per workgroup (network +0.4 ms at 1280x720 with N = 4).
+With --one-wave (D4R_RR_RELAXED=3) the block is engine/rr/rr2_w.comp in NVIDIA's own 32-thread group: the gather, the
+block and the tail all run on the one wave, and XF shares its 4 KiB with the block's XT and KVS (rr2_w.comp
+RR2_XF_SHARED), so nothing goes through memory. (The gather was made a pass of its own because on the four-wave block
+it ran on one wave while three waited; on one wave there is nothing to wait for.)
 With --dec0_1 the same for dec0's second kernel (rr2_m.comp KIND 5): nothing before the block (it reads its tokens
 from the arena); after it the last product's D fragments (the MLP sum) and the registers holding r are loaded from
 MSH / RSH, and the translated tail forms y = r + MLP * s2 and computes the kernel's outputs from it.
@@ -177,6 +181,8 @@ def main():
     m = next(i for i, l in enumerate(g) if l.startswith('void main()'))
     # --one-wave (D4R_RR_RELAXED=3): the block as rr2_w.comp, in NVIDIA's own 32-thread group, the tail on the same wave
     one_wave = '--one-wave' in sys.argv
+    if one_wave and (dec or G != 1):
+        raise SystemExit('--one-wave is for encoder 0 with one window per workgroup')
     head = '\n'.join(g[:m])
     if not one_wave:
         head = head.replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
@@ -289,23 +295,43 @@ 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.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.
-    text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' + block + '\n' +
+    # Each wave runs its own copy of the tail with only the N-th part of its stores (in order: the first pixel's,
-            ('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(MSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n'
+    # then the second's); the compiler drops what a copy's stores do not need. Every stored value comes from the
-             'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(RSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n\n' if dec else
+    # 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'
+             'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(rr2_rsh(row, col))); }\n\n' if dec else
              'uint HDV(uint row, uint col) { return col < 4u ? uint(float16BitsToUint16(HD[gl_SubgroupID * 256u + row * 4u + col])) : 0u; }\n\n') +
             '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' +
+             '        rr2_block();\n        barrier();\n    }\n') +
-            '\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n    }\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')
 

engine/rr/ptx2rr.py

@@ -59,6 +59,9 @@ def mem_shape(o):
     return size, count
 
 
+# --exact-convert (the module is passed through engine/rr/spv_exact_convert.py, which keeps every conversion): a
+# cvt.rn.f16.f32 only needs the opaque zero on its f32 bits, against the fused multiply-convert (see cvt)
+CVT_GUARD = '--exact-convert' not in sys.argv
 RR_RUNTIME = r'''
 float16_t H(uint u) { return uint16BitsToFloat16(uint16_t(u)); }
 uint HU(float16_t h) { return uint(float16BitsToUint16(h)); }
@@ -67,7 +70,8 @@ uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); }
 uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); }
 int s16(uint u) { return (int(u) << 16) >> 16; }
 float rcp_a(float x) { return 1.0 / x; }
-vec4 fetch0(sampler2D t, ivec2 c) { const ivec2 n = textureSize(t, 0); return (c.x < 0 || c.y < 0 || c.x >= n.x || c.y >= n.y) ? vec4(0.0) : texelFetch(t, c, 0); }
+// (one unsigned comparison per axis: a negative coordinate is above every size)
+vec4 fetch0(sampler2D t, ivec2 c) { const ivec2 n = textureSize(t, 0); return any(greaterThanEqual(uvec2(c), uvec2(n))) ? vec4(0.0) : texelFetch(t, c, 0); }
 // x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly
 float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; }
 uint PARAM32(uint o) { return param.w[o >> 2]; }
@@ -716,7 +720,7 @@ class RRT(base.T):
         for i, d in enumerate(dst):
             # a binary16 result is the texel rounded to binary16: through H(HU(...) | SUST_Z) so that the driver cannot
             # fold the rounding away when the value is widened again (RR_RUNTIME's note on sust_b)
-            self.A(d, f'H(HU(float16_t(t_[{i}])) | SUST_Z)' if f16 else f't_[{i}]', 'f16' if f16 else 'float')
+            self.A(d, (f'H(HU(float16_t(t_[{i}])) | SUST_Z)' if CVT_GUARD else f'float16_t(t_[{i}])') if f16 else f't_[{i}]', 'f16' if f16 else 'float')
         self.E('}')
 
     def sust(self, op, a):
@@ -777,6 +781,11 @@ class RRT(base.T):
             # conversion into v_fma_mix, one rounding of the exact value, where the kernel rounds to f32 first: an f32
             # result that is exactly a binary16 midpoint then rounds the other way (the output stages' rare one-step
             # differences, e.g. log2(x) * c1 * c2; found 2026-10-13): the OR on the f32 bits keeps the f32 value.
+            # With --exact-convert the conversion itself is decorated NoContraction (spv_exact_convert.py), which keeps the
+            # first: only the f32 OR is needed (checked in the ISA: the same v_cvt_f16_f32 / v_cvt_f32_f16 counts, no
+            # v_fma_mix*_f16; without the f32 OR, ACO fuses some products into v_fma_mixlo_f16 despite `precise`)
+            if not CVT_GUARD:
+                return self.A(a[0], f'float16_t(uintBitsToFloat(floatBitsToUint({self.val(a[1], "f32")}) | SUST_Z))', 'f16')
             return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint({self.val(a[1], "f32")}) | SUST_Z))) | SUST_Z)', 'f16')
         return super().cvt(op, o, a)
 
@@ -859,6 +868,8 @@ class RRT(base.T):
             return self.A(a[0], f'tanh_ocml({self.val(a[1], "f32")})', 'float')
         if op == 'ex2.approx.f16':
             # v_exp_f32 of the widened value, rounded to binary16; the ORs keep that rounding (see cvt)
+            if not CVT_GUARD:
+                return self.A(a[0], f'float16_t(uintBitsToFloat(floatBitsToUint(exp2(float({self.val(a[1], "f16")}))) | SUST_Z))', 'f16')
             return self.A(a[0], f'H(HU(float16_t(uintBitsToFloat(floatBitsToUint(exp2(float({self.val(a[1], "f16")}))) | SUST_Z))) | SUST_Z)', 'f16')
         if op == 'neg.f16':
             return self.A(a[0], f'-{self.val(a[1], "f16")}', 'f16')
@@ -945,6 +956,8 @@ def main():
                          'execution) for workgroup --probe-block, into a debug buffer bound after the surfaces')
     ap.add_argument('--probe-block', default='0,0')
     ap.add_argument('--oob-report', action='store_true', help='debugging: count out-of-range buffer accesses')
+    ap.add_argument('--exact-convert', action='store_true',
+                    help='the module goes through spv_exact_convert.py: lighter guards on binary16 conversions (CVT_GUARD)')
     ns = ap.parse_args()
     block = tuple(int(x) for x in ns.block.split(',')) if ns.block else default_block(ns.kernel)
     formats = {int(k): v for k, v in (x.split('=') for x in ns.surface)}

engine/rr/rr2_m.comp

@@ -90,7 +90,7 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
 #define E4V(w, b) float(uintBitsToFloate4m3EXT(uint8_t(((w) >> (8u * (b))) & 0xffu)))
 #define RESID(d, a, s, c) const float d = fma(float(a), float(s), float(c));
 #define Q8R(x) q8s(x)
-#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2(float((h)[i_]), float((h)[i_ + 1])); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
+#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2(float((h)[i_]), float((h)[i_ + 1])); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, FINITE2(x_ * g_)); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
 #else
 #define RVT float16_t
 #define E4V(w, b) e4((w), (b))
@@ -113,10 +113,33 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
 const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
 #define BRANCHB BRANCH
 
+#if KIND == 1
+// enc1's merge input is in registers (gb) before XT or KVS is first written, so the three share 8 KiB instead of 16
+// (above about 4 KiB per wave this GPU runs the workgroup slower, see KIND 2 below): XT words 0 - 16C-1, KVS after it
+#if 16 * C + 2 * 512 > 64 * 32
+#error rr2_m.comp: XT and KVS no longer fit in GT
+#endif
+shared uint GT[64 * 32];         // the merge input: [token slot][4 sub-positions x 32 FP8 channels]; then XT and KVS
+#define XT GT
+#define KVS GT
+#define KVA(o) (16u * C + (o))
+#elif KIND == 5 && RR2_GROUP == 1
+// dec0's second kernel writes r and the MLP sum for its tail only after the block is done with XT and KVS, so the
+// three share 8 KiB instead of 14: XT words 0 - 16C-1, KVS after it; then r (RSW) at word 0 and the MLP sum (MSW) at
+// 1024, both as binary16 pairs of natural channels 2k, 2k + 1 per token (the tail reads them in those pairs)
+#if 16 * C + 2 * 512 > 2 * 1024
+#error rr2_m.comp: XT and KVS no longer fit under RSW and MSW
+#endif
+shared uint SM[2 * 1024];
+#define XT SM
+#define KVS SM
+#define KVA(o) (16u * C + (o))
+#define RSW(w) SM[w]
+#define MSW(w) SM[1024u + (w)]
+#else
 shared uint XT[16 * C];          // the window's tokens: [token slot][C FP8 channels in slot order]; later the merge rows
 shared uint KVS[2 * 512];        // the current head's QK ([token slot][32]), then V^T ([32][token slot])
-#if KIND == 1
+#define KVA(o) (o)
-shared uint GT[64 * 32];         // the merge input: [token slot][4 sub-positions x 32 FP8 channels]
 #endif
 #if KIND == 2
 // Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave (16 KiB for this
@@ -146,7 +169,15 @@ const uint EBRANCH = W_EWD1 - W_EWD0;
 #define EBRANCHB EBRANCH
 #endif
 #if KIND == 5
+#if RR2_GROUP == 1
+// r and the MLP sum of token row, natural channel col, for the translated tail (gen_rr2_enc0.py)
+float16_t rr2_rsh(uint row, uint col) { return unpackFloat2x16(RSW(row * 16u + (col >> 1u)))[col & 1u]; }
+float16_t rr2_msh(uint row, uint col) { return unpackFloat2x16(MSW(row * 16u + (col >> 1u)))[col & 1u]; }
+#else
 shared float16_t RSH[RR2_GROUP * 64 * 32], MSH[RR2_GROUP * 64 * 32];   // per window, r and the MLP sum per token (raster order) and natural channel
+float16_t rr2_rsh(uint row, uint col) { return RSH[gl_SubgroupID * 2048u + row * 32u + col]; }
+float16_t rr2_msh(uint row, uint col) { return MSH[gl_SubgroupID * 2048u + row * 32u + col]; }
+#endif
 #endif
 #if GUIDE || KIND == 3
 shared float16_t OUTG[64 * 4];   // the guide output per window token (raster order)
@@ -242,6 +273,7 @@ void main()
     {
         B8 gb[8];
         [[unroll]] for (uint kt = 0u; kt < 8u; ++kt) coopMatLoad(gb[kt], GT, wv * 512u + 4u * kt, 32u, CMJ);
+        barrier();   // every wave has its part of GT in registers: XT takes its place
         [[unroll]] for (uint m = 0u; m < MT; ++m) {
             RAC a;
             [[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(B1, 16u * m + 8u * hl + uint(i));
@@ -342,20 +374,20 @@ void main()
             A8 ka, va;
             Q8(ka, a)
             Q8(va, v)
-            coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
+            coopMatStore(ka, KVS, KVA(n * 128u + 4u * c), 8u, RM);
-            coopMatStore(va, KVS, 512u + c * 256u + 4u * n, 16u, RM);
+            coopMatStore(va, KVS, KVA(512u + c * 256u + 4u * n), 16u, RM);
         }
         barrier();
         // scores^T (keys x this wave's queries) = QK . QK^T + bias, the cubic softmax down each column
         B8 Qb[2];
-        [[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], KVS, wv * 128u + 4u * c, 8u, CMJ);
+        [[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], KVS, KVA(wv * 128u + 4u * c), 8u, CMJ);
         B8 Pb[4];
         {
             float16_t e[4][8];
             [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
                 A8 k0, k1;
-                coopMatLoad(k0, KVS, kt * 128u, 8u, RM);
+                coopMatLoad(k0, KVS, KVA(kt * 128u), 8u, RM);
-                coopMatLoad(k1, KVS, kt * 128u + 4u, 8u, RM);
+                coopMatLoad(k1, KVS, KVA(kt * 128u + 4u), 8u, RM);
                 HAC c;
                 coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
                 K32H(c, k0, Qb[0], k1, Qb[1])
@@ -392,8 +424,8 @@ void main()
             RAC o = RAC(RV(float16_t(0.0)));
             [[unroll]] for (uint ks = 0u; ks < 2u; ++ks) {
                 A8 v0, v1;
-                coopMatLoad(v0, KVS, 512u + c * 256u + 8u * ks, 16u, RM);
+                coopMatLoad(v0, KVS, KVA(512u + c * 256u + 8u * ks), 16u, RM);
-                coopMatLoad(v1, KVS, 512u + c * 256u + 8u * ks + 4u, 16u, RM);
+                coopMatLoad(v1, KVS, KVA(512u + c * 256u + 8u * ks + 4u), 16u, RM);
                 K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
             }
             Q8(Ob[c], o)
@@ -449,12 +481,23 @@ void main()
 
 #if KIND == 5
     // dec0's second kernel: the translated tail forms y = r + MLP * s2 itself and computes the outputs from it
+#if RR2_GROUP == 1
+    barrier();   // every wave is done with XT and KVS
+    // (channel slots j and j + 1, j even, hold the natural channels inter(j) and inter(j) + 1)
+    [[unroll]] for (uint m = 0u; m < MT; ++m)
+        [[unroll]] for (int i = 0; i < 8; i += 2) {
+            const uint w = (t * 32u + inter(16u * m + 8u * hl + uint(i))) >> 1u;
+            RSW(w) = packFloat2x16(f16vec2(float16_t(rv[m][i]), float16_t(rv[m][i + 1])));
+            MSW(w) = packFloat2x16(f16vec2(float16_t(ML[m][i]), float16_t(ML[m][i + 1])));
+        }
+#else
     [[unroll]] for (uint m = 0u; m < MT; ++m)
         [[unroll]] for (int i = 0; i < 8; ++i) {
             const uint at = RR2_WIN * 2048u + t * 32u + inter(16u * m + 8u * hl + uint(i));
             RSH[at] = float16_t(rv[m][i]);
             MSH[at] = float16_t(ML[m][i]);
         }
+#endif
 #else
     // ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens
     B8 Tb[MT];

engine/rr/rr2_w.comp

@@ -6,8 +6,10 @@
 // Shared memory stays at 4 KiB, above which this GPU runs fewer waves (engine/rr/README.md). The arithmetic is
 // rr2_m.comp's with -DRR_RELAXED=2 -DRR_NATIVE and the arena comes out byte-identical to that build
 // (build/rr/f/nv/invest-20261013/REPORT.md: 453 us against 581 us at 1280x720). compile_rr2.py uses it for dec0_0.
-// KIND 4 is encoder 0's block the same way (the embedding of the gather pass's record, the tokens, the guide head to
+// KIND 4 is encoder 0's block the same way (the embedding of the gather's features, the tokens, the guide head to
-// HD), included with RR2_EMBED and RR2_XF_ARENA by the shader engine/rr/gen_rr2_enc0.py --one-wave generates.
+// HD), included with RR2_EMBED by the shader engine/rr/gen_rr2_enc0.py --one-wave generates: with RR2_XF_SHARED the
+// gather runs in the same shader and leaves the features in XF (the default), with RR2_XF_ARENA they come from the
+// gather pass's record in the arena (--split block).
 #extension GL_KHR_cooperative_matrix : require
 #extension GL_KHR_memory_scope_semantics : require
 #extension GL_EXT_shader_explicit_arithmetic_types : require
@@ -26,8 +28,8 @@ layout(local_size_x = 32) in;
 #if !(KIND == 2 || KIND == 4) || CM != 0 || GUIDE != 0 || C != 32
 #error rr2_w.comp: a 32-channel decoder's first kernel, or encoder 0's block
 #endif
-#if KIND == 4 && !(defined(RR2_EMBED) && defined(RR2_XF_ARENA))
+#if KIND == 4 && !(defined(RR2_EMBED) && (defined(RR2_XF_ARENA) || defined(RR2_XF_SHARED)))
-#error rr2_w.comp: encoder 0's block reads the gather pass's record (gen_rr2_enc0.py --split block --one-wave)
+#error rr2_w.comp: encoder 0's block reads the gather's features (gen_rr2_enc0.py --one-wave)
 #endif
 #define MT (C / 16)
 #ifdef RR2_EMBED
@@ -53,11 +55,21 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
 #define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
 #define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); }
 #define Q8(dst, src) Q8PAIRS(dst, src)
-#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
+#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, FINITE2(x_ * g_)); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
 const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
 
+#if KIND == 4 && defined(RR2_XF_SHARED)
+// the gather's features (XF, [token slot][64 FP8 feature slots], 4 KiB) are all in registers before the block first
+// writes XT or KVS, so the three share 4 KiB: KVS words 0 - 511, XT words 512 - 1023
+shared uint XF[64 * 16];
+#define KVS XF
+#define XT XF
+#define XTA(o) (512u + (o))
+#else
 shared uint XT[16 * C];
 shared uint KVS[512];          // the expansion (input phase), then per head V^T, then QK
+#define XTA(o) (o)
+#endif
 // Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave the rate falls in
 // proportion (build/rr/f/nv/cost/occ2.py). So the input phase borrows KVS, which the attention only needs later
 // (LOW at word 0, EX at word 512), and the block's unquantised input stays in registers (xr), not in memory.
@@ -68,7 +80,14 @@ shared uint KVS[512];          // the expansion (input phase), then per head V^T
 #endif
 #define EX_AT 0u
 #if KIND == 4
+#ifdef RR2_XF_ARENA
 uint rr2_xf;                     // the window's record from the gather pass: the index of its first arena word
+#define XF_SRC arena.w
+#define XF_AT rr2_xf
+#else
+#define XF_SRC XF
+#define XF_AT 0u
+#endif
 shared float16_t HD[64 * 4];     // the guide head per token (raster order), for the translated tail
 const uint EBRANCH = W_EWD1 - W_EWD0;
 #endif
@@ -103,10 +122,13 @@ void main()
         B8 fb[4][4];
         ACC X[4][MT];
         [[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
-            [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[wv][kt], arena.w, rr2_xf + wv * 256u + 4u * kt, 16u, CMJ);
+            [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[wv][kt], XF_SRC, XF_AT + wv * 256u + 4u * kt, 16u, CMJ);
             [[unroll]] for (uint m = 0u; m < MT; ++m)
                 [[unroll]] for (int i = 0; i < 8; ++i) X[wv][m][i] = float(w16.w[W_EBU / 2u + 16u * m + 8u * hl + uint(i)]);
         }
+#ifdef RR2_XF_SHARED
+        barrier();   // XF is read: XT and KVS take its place
+#endif
         [[unroll]] for (uint br = 0u; br < 2u; ++br) {
             A8 ewd[2][4], ewu0[MT], ewu1[MT];
             [[unroll]] for (uint c = 0u; c < 2u; ++c)
@@ -128,7 +150,7 @@ void main()
         [[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
             [[unroll]] for (uint m = 0u; m < MT; ++m) {
                 Q8(qx[wv][m], X[wv][m])
-                coopMatStore(qx[wv][m], XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
+                coopMatStore(qx[wv][m], XT, XTA(wv * 4u * C + 4u * m), C / 4u, CMJ);
             }
     }
 #else
@@ -198,8 +220,8 @@ void main()
                 ACC v = ACC(0.0);
                 [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
                     A8 xa0, xa1;
-                    coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM);
+                    coopMatLoad(xa0, XT, XTA(n * 4u * C + 4u * kt), C / 4u, RM);
-                    coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM);
+                    coopMatLoad(xa1, XT, XTA(n * 4u * C + 4u * kt + 4u), C / 4u, RM);
                     K32(v, xa0, wvt[kt], xa1, wvt[kt + 1u])
                 }
                 A8 va;
@@ -219,8 +241,8 @@ void main()
                 ACC a = ACC(0.0);
                 [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
                     B8 xb0, xb1;
-                    coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ);
+                    coopMatLoad(xb0, XT, XTA(n * 4u * C + 4u * kt), C / 4u, CMJ);
-                    coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ);
+                    coopMatLoad(xb1, XT, XTA(n * 4u * C + 4u * kt + 4u), C / 4u, CMJ);
                     K32(a, wk[kt], xb0, wk[kt + 1u], xb1)
                 }
                 A8 ka;
@@ -348,7 +370,7 @@ void main()
         [[unroll]] for (uint m = 0u; m < MT; ++m) {
             floate4m3_t Tb[8];
             [[unroll]] for (int i = 0; i < 8; i += 2) {
-                fe4m3vec2 q; saturatedConvertEXT(q, vec2(fma(ML[wv][m][i], f2, rv[wv][m][i]), fma(ML[wv][m][i+1], f2, rv[wv][m][i+1])));
+                fe4m3vec2 q; saturatedConvertEXT(q, FINITE2(vec2(fma(ML[wv][m][i], f2, rv[wv][m][i]), fma(ML[wv][m][i+1], f2, rv[wv][m][i+1]))));
                 Tb[i] = q.x; Tb[i+1] = q.y;
             }
 #if KIND == 4

engine/rr/rr_exact.glsl

@@ -31,6 +31,19 @@ uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) |
 uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); }
 int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
 
+// FINITE(v): v, but an infinity (which these f32 values never are: sums of FP8 products and binary16 biases stay far
+// below 3.4e38) becomes NaN. Written so, a saturating FP8 conversion is RADV's f2e4m3fn_satfn (nir_opt_algebraic's
+// vkd3d-proton pattern): the hardware conversion with FP16_OVFL, which saturates every finite value the same way,
+// instead of a NaN-preserving clamp per element and then the conversion. RR_SAT_CLAMP=1 keeps the clamp, and so does
+// RR_MIX, whose binary16 accumulators can overflow.
+#if defined(RR_SAT_CLAMP) || defined(RR_MIX)
+#define FINITE(v) (v)
+#define FINITE2(v) (v)
+#else
+#define FINITE(v) (isinf(v) ? uintBitsToFloat(0x7fc00000u) : (v))
+#define FINITE2(v) vec2(FINITE((v).x), FINITE((v).y))
+#endif
+
 #ifdef RR_GFX11
 #include "rr_gfx11.glsl"
 #define F8 uint
@@ -52,7 +65,7 @@ floate4m3_t q8(float16_t h)
 }
 // the same from f32 / f16 values in one saturating hardware conversion (satfinite; for an f32 value that is a binary16
 // number the result equals the f16 conversion's)
-floate4m3_t q8s(float x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; }
+floate4m3_t q8s(float x) { floate4m3_t r; saturatedConvertEXT(r, FINITE(x)); return r; }
 floate4m3_t q8h(float16_t x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; }
 
 uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); }
@@ -124,13 +137,13 @@ float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid)
 }
 #ifndef RR_GFX11
 // The saturating FP8 conversion of a whole accumulator, two elements per conversion instruction
-#define Q8PAIRS(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { fe4m3vec2 t_; saturatedConvertEXT(t_, vec2((src)[i_], (src)[i_ + 1])); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
+#define Q8PAIRS(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { fe4m3vec2 t_; saturatedConvertEXT(t_, FINITE2(vec2((src)[i_], (src)[i_ + 1]))); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
 // q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
 #define ACTQ(h) uintBitsToFloate4m3EXT(w8u.w[W_ACT + bh(h)])
 // the activation of a whole accumulator into FP8. RR_NATIVE (D4R_RR_RELAXED=3): in f32, two elements per
 // instruction, straight into FP8 - no binary16 in between and no table (not the reference's roundings)
 #ifdef RR_NATIVE
-#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
+#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, FINITE2(x_ * g_)); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
 #else
 #define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
 #endif

engine/runtime_rr2.inc

@@ -638,6 +638,11 @@ struct R2Engine {
         const auto& b = *bound;
         const uint32_t f = frame, cur = f & 1u;
         auto stamp = [&](uint32_t n) { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, stamps, first + n); };
+        // timing only: D4R_ENGINE_RR_STAGE_STAMPS=1 also stamps after every network stage and output group, from query
+        // first + 4 on (the caller's pool needs 32 queries)
+        static const bool stageStamps = std::getenv("D4R_ENGINE_RR_STAGE_STAMPS") != nullptr;
+        uint32_t extra = 4;
+        auto stageStamp = [&] { if (stageStamps) stamp(extra++); };
         auto step = [&] {
             memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
                           VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT);
@@ -685,6 +690,7 @@ struct R2Engine {
             }
             dispatch(cb, model.stages[stage].pipeline, b.network[stage][cur], (gx + model.stages[stage].group - 1) / model.stages[stage].group, (lh(l) + sy + 7) / 8);
             step();
+            stageStamp();
         }
         stamp(2);
         // timing only: D4R_ENGINE_RR_SKIP=post,down,pooled leaves those passes out (the picture is wrong then)
@@ -706,6 +712,7 @@ struct R2Engine {
                 if (serial) step();
             }
             if (g < groups) step();
+            stageStamp();
         }
         // (the pooling pass reads what the post passes wrote: the last group's down passes do not touch it)
         if (!noPooled) dispatch(cb, model.stages[R2Pooled].pipeline, b.pooled, (qw + 15) / 16, (qh + 15) / 16);