Merge remote-tracking branch 'origin/linux' into linux
# Conflicts: # engine/rr/ptx2rr.py
This commit is contained in:
woutermans committed
commit
b12f43b913
15 files changed
+758
-69
No files matched your search
+14
-4
@@ -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)
|
||||
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)
|
||||
@@ -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:
|
||||
|
||||
+2
-1
@@ -16,7 +16,8 @@ bool gameColorFormat(VkFormat format) {
|
||||
}
|
||||
}
|
||||
bool gameMotionFormat(VkFormat format) {
|
||||
return format == VK_FORMAT_R16G16_SFLOAT || format == VK_FORMAT_R32G32_SFLOAT ||
|
||||
// R16G16_SNORM (Enshrouded): sampled as floats in [-1, 1]; MV.Scale carries them to pixels as for float vectors.
|
||||
return format == VK_FORMAT_R16G16_SFLOAT || format == VK_FORMAT_R32G32_SFLOAT || format == VK_FORMAT_R16G16_SNORM ||
|
||||
format == VK_FORMAT_R16G16B16A16_SFLOAT || format == VK_FORMAT_R32G32B32A32_SFLOAT;
|
||||
}
|
||||
bool gameDepthFormat(VkFormat format, VkImageAspectFlags aspect) {
|
||||
|
||||
+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,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
|
||||
|
||||
@@ -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')
|
||||
|
||||
|
||||
+15
-2
@@ -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)}
|
||||
|
||||
+53
-10
@@ -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];
|
||||
|
||||
+34
-12
@@ -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
|
||||
|
||||
+16
-3
@@ -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
|
||||
|
||||
@@ -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