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]>
This commit is contained in:
itchyfeetleech committed
1 parent
0e851b7f90
commit
49439afb81
7 files changed
+155
-28
No files matched your search
@@ -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],
|
||||
|
||||
+19
-2
@@ -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) {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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'
|
||||
|
||||
+52
-9
@@ -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];
|
||||
|
||||
+32
-10
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in new issue
Block a user