wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/linux' into linux

Browse Source 
# Conflicts:
#	engine/rr/ptx2rr.py
This commit is contained in:
woutermans committed 2026-10-19 11:33:30 +02:00
2 parents 3f66081add 689ac6f270
commit b12f43b913
15 files changed
+758 -69
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/rr/README.md
+59
View File
@@ -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
+33 -7
View File
@@ -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)')
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
# --tail-waves N (dec0_1, one window per workgroup): the translated tail on N of the block's waves instead of one.
# Each wave runs its own copy of the tail with only the N-th part of its stores (in order: the first pixel's,
# then the second's); the compiler drops what a copy's stores do not need. Every stored value comes from the
# same instructions on the same values as on one wave, so the outputs are the same.
TW = int(sys.argv[sys.argv.index('--tail-waves') + 1]) if '--tail-waves' in sys.argv else 1
if TW > 1 and not (dec and G == 1):
raise SystemExit('--tail-waves is for dec0_1 with one window per workgroup')
wx = (lambda t: t) if one_wave or TW > 1 else (lambda t: t.replace('gl_WorkGroupID.x', f'({G}u * gl_WorkGroupID.x + gl_SubgroupID)'))
STORE = re.compile(r'^\s*(?!if\b|for\b|while\b|switch\b|return\b)[A-Za-z_]\w*\(')
stores = [i for i, l in enumerate(tail) if STORE.match(l)]
if TW > 1 and len(stores) % TW:
raise SystemExit(f'--tail-waves {TW}: {len(stores)} stores do not split evenly')
def tail_part(w):
mine = set(stores[w * len(stores) // TW:(w + 1) * len(stores) // TW])
return [l for i, l in enumerate(tail) if i not in stores or i in mine]
tails = ''.join(f' if (rr2_valid && gl_SubgroupID == {w}u) {{\n' + '\n'.join(load + keep) + '\n' + '\n'.join(tail_part(w)) + '\n }\n'
for w in range(TW)) if TW > 1 else None
valid = (f'gl_SubgroupID < {TW}u && gl_WorkGroupID.x < rr2_gx' if TW > 1 else
f'gl_SubgroupID < {G}u && {G}u * gl_WorkGroupID.x + gl_SubgroupID < rr2_gx')
text = (head + f'\n#extension GL_GOOGLE_include_directive : require\n#define RR2_EMBED 1\n#define RR2_GROUP {G}\n' +
('#define RR2_XF_SHARED 1\n' if one_wave else '') + block + '\n' +
('uint HDV(uint row, uint col) { return uint(float16BitsToUint16(rr2_msh(row, col))); }\n'
'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' +
'\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n }\n}\n')
' rr2_block();\n barrier();\n }\n') +
(tails + '}\n' if tails else
' if (rr2_valid) {\n' + '\n'.join(load + keep) + '\n' + wx('\n'.join(tail)) + '\n }\n}\n'))
Path(out).write_text(text)
print(f'{out}: {len(pre)} gather lines, {len(cut)} cut, {len(tail)} tail lines, {len(seen)} A fragment sets')
engine/rr/ptx2rr.py
+15 -2
View File
@@ -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)); }
@@ -75,7 +78,8 @@ float lg2_ftz(float x) { return log2(abs(x) < 1.17549435e-38 ? 0.0 : x); }
// PTX .relu: negative results become +0, NaN stays NaN
f16vec2 relu2(f16vec2 v) { return f16vec2(v.x < 0.0hf ? 0.0hf : v.x, v.y < 0.0hf ? 0.0hf : v.y); }
float16_t relu1(float16_t v) { return v < 0.0hf ? 0.0hf : v; }
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]; }
@@ -753,7 +757,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):
@@ -814,6 +818,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)
@@ -910,6 +919,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')
@@ -1048,6 +1059,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
+53 -10
View File
@@ -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
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];
engine/rr/rr2_w.comp
+34 -12
View File
@@ -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
@@ -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(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;
@@ -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
+16 -3
View File
@@ -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