engine/rr: preset F in less shared memory (encoder 0 in one shader, enc1, dec0_1)

ACO sizes a shader's registers to the waves its shared memory leaves room
for, so stages that declare less shared memory also get fewer registers and
more waves. Three stages hold buffers whose lifetimes do not overlap:

- Encoder 0 at D4R_RR_RELAXED=3: with the one-wave block the gather no
  longer needs a pass of its own. The generated shader runs gather, block
  and tail in one 32-thread group and hands the features over in shared
  memory that the block's XT and KVS reuse (rr2_w.comp RR2_XF_SHARED). No
  gather pass, no record in the arena (36 MB of it lay past the arena at
  720p). D4R_RR2_SPLIT=1 keeps the pass.
- enc1: the merge input (GT) is in registers before XT or KVS is written:
  17 -> 9 KiB, 240 -> 144 vector registers, 6 -> 10 waves per SIMD.
- dec0_1: r and the MLP sum for the tail are written after the block is done
  with XT and KVS, as binary16 pairs: 14 -> 8 KiB, 192 -> 120 registers,
  8 -> 12 waves per SIMD.

Also D4R_ENGINE_RR_STAGE_STAMPS=1: a timestamp after every network stage and
output group, printed per stage by d4r-rr.

Every image of six temporal frames byte-identical to the previous model
(1280x720 -> 1440p, 1114x626 -> 1080p, 853x480 -> 1440p, robustness2). RX
9070 XT, RADV, 720p -> 1440p: frame 4.09 -> 4.00 ms (enc0 798 -> 772 us,
enc1 286 -> 266, dec0_1 660 -> 623). Other stages' SPIR-V is unchanged.

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

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