Files
d4r/engine/cnn_sr/mma_f16.glsl
T
woutermans b9254f7f96 engine: Super Resolution presets E and F, the DLSS 3 CNN, on RDNA2
Presets E (5) and F (6) run on d4r's own Vulkan engine without cooperative
matrices: the input and reconstruction kernels translated from the user's DLL
(engine/compile_sr_e.py, rr/ptx2rr.py), the network engine/sr/e_net.comp with
packed v_pk_fma_f16 accumulation (engine/cnn_sr/mma_f16.glsl, the default
head; --head-accumulation f32 keeps the f32 head). Runtime runtime_e.inc
(EModel/EEngine), game adapter game_e.cpp, replay d4r-e; the shim serves them
from engine\e and engine\f (D4R_ENGINE_MODEL_E_DIR / _F_DIR), LDR or HDR, any
render size, not with AutoExposure. build-engine-models.sh builds both.

Translator fixes that apply to every translated kernel: .relu on f16/f16x2
fma/add/mul was ignored; add.rz was always taken for the roundf idiom (lost
the log1pf correction term); formatted image stores expand at the call site
(SPIR-V image type fix); f16 mma.sync m16n8k8/k16, .const loads, shared
16-bit stores, aligned 32-bit arena words (D4R_WORD_ARENA), unified arena.

Steam Deck gfx1033: 2.55 ms at 640x360 -> 1280x720, 6.20 ms at 1080p
output; Rounded Quality sizes take mid_3x3 (reconstruction 2.44 -> 1.23 ms).
Ray Reconstruction CNN presets A-C (compile_rr3.py, the packed RR flavor)
stay out of this commit; engine/README.md already documents their tuning.
2026-10-18 20:39:26 +02:00

71 lines
3.6 KiB
GLSL

// Extension-free f16 MMA for wave32. PTX m16n8k8 / m16n8k16 row.col fragments:
// A: (row=g+8*(word&1), k=2*t+8*(word/2)+half).
// B: (k=2*t+8*word+half, col=g). C/D: (row=g+8*word, col=2*t+half).
// g=lane/4, t=lane%4. Shuffle sources stay within the current subgroup;
// there is no workgroup scratch or cross-wave synchronization.
// Include engine/round_half.glsl first: RADV's packHalf2x16 truncates instead of rounding to nearest.
#ifndef D4R_MMA_F16_INCLUDED
#define D4R_MMA_F16_INCLUDED
#ifdef D4R_MMA_PACKED
// The SR reconstruction head uses the same packed-half accumulation as e_net.comp.
// Each lane owns two rows and two columns; one packed FMA updates a row's two columns.
// Keep the fragment mapping and K order, rounding after each FMA instead of each MMA.
f16vec4 mma_packed_step8(uint a0, uint a1, uint b0, f16vec4 acc)
{
const uint lane = gl_SubgroupInvocationID, row = lane >> 2u, col = (lane & 3u) * 2u;
[[unroll]] for (uint k = 0u; k < 4u; ++k) {
const f16vec2 x0 = unpackFloat2x16(subgroupShuffle(a0, row * 4u + k));
const f16vec2 x1 = unpackFloat2x16(subgroupShuffle(a1, row * 4u + k));
const f16vec2 y0 = unpackFloat2x16(subgroupShuffle(b0, col * 4u + k));
const f16vec2 y1 = unpackFloat2x16(subgroupShuffle(b0, (col + 1u) * 4u + k));
acc.xy = fma(f16vec2(x0.x), f16vec2(y0.x, y1.x), acc.xy);
acc.xy = fma(f16vec2(x0.y), f16vec2(y0.y, y1.y), acc.xy);
acc.zw = fma(f16vec2(x1.x), f16vec2(y0.x, y1.x), acc.zw);
acc.zw = fma(f16vec2(x1.y), f16vec2(y0.y, y1.y), acc.zw);
}
return acc;
}
void mma_f16_m16n8k8(uint a0, uint a1, uint b0, uint c0, uint c1, out uint d0, out uint d1)
{
const f16vec4 acc = mma_packed_step8(a0, a1, b0, f16vec4(unpackFloat2x16(c0), unpackFloat2x16(c1)));
d0 = packFloat2x16(acc.xy); d1 = packFloat2x16(acc.zw);
}
void mma_f16_m16n8k16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1,
uint c0, uint c1, out uint d0, out uint d1)
{
f16vec4 acc = mma_packed_step8(a0, a1, b0, f16vec4(unpackFloat2x16(c0), unpackFloat2x16(c1)));
acc = mma_packed_step8(a2, a3, b1, acc);
d0 = packFloat2x16(acc.xy); d1 = packFloat2x16(acc.zw);
}
#else
vec4 mma_f16_step8(uint a0, uint a1, uint b0, vec4 acc)
{
const uint lane = gl_SubgroupInvocationID, row = lane >> 2u, col = (lane & 3u) * 2u;
[[unroll]] for (uint k = 0u; k < 4u; ++k) {
const vec2 x0 = unpackHalf2x16(subgroupShuffle(a0, row * 4u + k));
const vec2 x1 = unpackHalf2x16(subgroupShuffle(a1, row * 4u + k));
const vec2 y0 = unpackHalf2x16(subgroupShuffle(b0, col * 4u + k));
const vec2 y1 = unpackHalf2x16(subgroupShuffle(b0, (col + 1u) * 4u + k));
acc.x = fma(x0.x, y0.x, acc.x); acc.x = fma(x0.y, y0.y, acc.x);
acc.y = fma(x0.x, y1.x, acc.y); acc.y = fma(x0.y, y1.y, acc.y);
acc.z = fma(x1.x, y0.x, acc.z); acc.z = fma(x1.y, y0.y, acc.z);
acc.w = fma(x1.x, y1.x, acc.w); acc.w = fma(x1.y, y1.y, acc.w);
}
return acc;
}
void mma_f16_m16n8k8(uint a0, uint a1, uint b0, uint c0, uint c1, out uint d0, out uint d1)
{
const vec4 acc = mma_f16_step8(a0, a1, b0, vec4(unpackHalf2x16(c0), unpackHalf2x16(c1)));
d0 = packHalf2x16(roundHalfValue(acc).xy); d1 = packHalf2x16(roundHalfValue(acc).zw);
}
void mma_f16_m16n8k16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1,
uint c0, uint c1, out uint d0, out uint d1)
{
vec4 acc = mma_f16_step8(a0, a1, b0, vec4(unpackHalf2x16(c0), unpackHalf2x16(c1)));
acc = mma_f16_step8(a2, a3, b1, acc);
d0 = packHalf2x16(roundHalfValue(acc).xy); d1 = packHalf2x16(roundHalfValue(acc).zw);
}
#endif // D4R_MMA_PACKED
#endif