diff --git a/engine/compile_rr2.py b/engine/compile_rr2.py index eb97e0e..75ae405 100644 --- a/engine/compile_rr2.py +++ b/engine/compile_rr2.py @@ -74,8 +74,11 @@ BUFFERS['downp'] = [72, 168, 176] HYBRID = {'enc0': [], 'dec0_1': ['--dec0_1']} # 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) -SPLIT = os.environ.get('D4R_RR2_SPLIT', '1') != '0' +# 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 +# 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) @@ -191,7 +194,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], diff --git a/engine/replay_rr.cpp b/engine/replay_rr.cpp index 0813568..4e3d605 100644 --- a/engine/replay_rr.cpp +++ b/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 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(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) { diff --git a/engine/rr/README.md b/engine/rr/README.md index daa6855..56c1ad8 100644 --- a/engine/rr/README.md +++ b/engine/rr/README.md @@ -530,6 +530,34 @@ 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. +Frame 4.09 -> 4.00 ms (network 3.39 -> 3.30). 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 diff --git a/engine/rr/gen_rr2_enc0.py b/engine/rr/gen_rr2_enc0.py index 6025983..ac701bc 100644 --- a/engine/rr/gen_rr2_enc0.py +++ b/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,10 +295,11 @@ 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)') - text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' + block + '\n' + - ('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(MSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n' - 'uint RSV(uint row, uint col) { return uint(float16BitsToUint16(RSH[gl_SubgroupID * 2048u + row * 32u + col])); }\n\n' if dec else + wx = (lambda t: t) if one_wave else (lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)')) + 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' diff --git a/engine/rr/rr2_m.comp b/engine/rr/rr2_m.comp index e5f1ed6..bb2b5a5 100644 --- a/engine/rr/rr2_m.comp +++ b/engine/rr/rr2_m.comp @@ -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 -shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels] +#define KVA(o) (o) #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(va, KVS, 512u + c * 256u + 4u * n, 16u, RM); + coopMatStore(ka, KVS, KVA(n * 128u + 4u * c), 8u, 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(k1, KVS, kt * 128u + 4u, 8u, RM); + coopMatLoad(k0, KVS, KVA(kt * 128u), 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(v1, KVS, 512u + c * 256u + 8u * ks + 4u, 16u, RM); + coopMatLoad(v0, KVS, KVA(512u + c * 256u + 8u * ks), 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]; diff --git a/engine/rr/rr2_w.comp b/engine/rr/rr2_w.comp index 2317a76..c359f78 100644 --- a/engine/rr/rr2_w.comp +++ b/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 -// HD), included with RR2_EMBED and RR2_XF_ARENA by the shader engine/rr/gen_rr2_enc0.py --one-wave generates. +// 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 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)) -#error rr2_w.comp: encoder 0's block reads the gather pass's record (gen_rr2_enc0.py --split block --one-wave) +#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's features (gen_rr2_enc0.py --one-wave) #endif #define MT (C / 16) #ifdef RR2_EMBED @@ -56,8 +58,18 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; #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; } 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(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM); + coopMatLoad(xa0, XT, XTA(n * 4u * C + 4u * kt), 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(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ); + coopMatLoad(xb0, XT, XTA(n * 4u * C + 4u * kt), 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; diff --git a/engine/runtime_rr2.inc b/engine/runtime_rr2.inc index 016fabf..85e8fcc 100644 --- a/engine/runtime_rr2.inc +++ b/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);