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

rr: add FSR Ray Regeneration, exact RDNA3 inference and motion experiments

Browse Source 
Add an experimental FSR Ray Regeneration backend for D3D12 RR requests:
- Convert depth, motion, normals and albedos into AMD denoiser inputs.
- Split and recompose diffuse/specular signals, then run Super Sampling.
- Add stability/detail controls, demodulation guards and halo limiting.
- Handle AMD provider discovery under NVIDIA GPU spoofing.
- Add runtime probes, synthetic reference scenes and per-dispatch profiling.

Add exact gfx11 DLSS-RR inference using f16 cooperative matrices:
- Reproduce FP8 conversions, accumulation and camera reprojection arithmetic.
- Widen model weights and avoid requiring VK_EXT_shader_float8.
- Extend model building, replay tools and descriptor metadata.

Expand motion preprocessing and experimental AW2 native-motion support:
- Add tiled LDS dilation, padded allocation handling and thin-only sky fill.
- Order dilation with input waits and NGX on the selected stream.
- Add jitter-aligned native composition, hybrid routing and wire history repair.
- Add prepared masks, sparse repair and independent motion/repair controls.

Add optional fused R11G11B10 input/output conversion, sampled stage timings
and profile summarization. Extend configuration, packaging, documentation,
FidelityFX headers and CPU/GPU regression coverage.
This commit is contained in:
ayo committed 2026-10-16 04:55:25 -04:00
1 parent 17c1828f97
commit 34e6b0fe71
83 files changed
+11300 -195
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/rr/README.md
+31
View File
@@ -15,6 +15,37 @@ NVIDIA kernel of the frame; the engine records them in NVIDIA's order:
The D variants of enc0, dec0 and the output kernel are different code; the other nine kernels are shared and only
their weights differ between presets.
## RDNA3 exact preset D (2026-10-14)
Build with `D4R_RR_GFX11=1 D4R_RR_RELAXED=0 engine/rr/build_native_model.sh CAPTURE OUT DLL`.
`rr_gfx11.glsl` reproduces ZLUDA's gfx11 FP8 widening, K positions, scaled f32 accumulators and final binary16
rounding using f16 cooperative matrices. The weight files include a widened copy. Relaxed and mix arithmetic are
unavailable in this build; `VK_EXT_shader_float8` is unnecessary. `kvk`, `evk` and `d4r-rr` enable float8 only when
the driver advertises the extension, because RADV gfx11 may report its feature bits without the extension.
`arithmetic gfx11` in `rr.txt` also selects NGX 310.7's CPU camera operation order. Its inversions, relative camera
translation and intermediate reprojection products leave small residues that affect specular motion history;
double-precision equivalents were insufficient for temporal equality. Other model builds keep their existing path.
RX 7700 XT (gfx1101), RADV, preset D, AW2 town at 853×480 → 2560×1440: enc0 and hkpn are byte-identical on all three
captured launches, joining the ten previously verified stages. Encoder surface 112's undefined alpha is excluded
from the isolated stage check. The complete presented RGBA16F output, including alpha, is byte-identical over all
24 replay frames in two engine runs versus the deterministic CUDA control. Engine GPU time is about 11.4 ms versus
27.9 ms for the current control; warmed samples varied by about 0.1 ms within each engine run. Older CUDA runs
included timing outliers, so their lower 26.7 ms figure is historical, not the paired control. No frame-age or
resolution change contributes to the result. Other RDNA3 GPUs, preset E, and other games require validation.
AW2 live at the same setting: CUDA / engine denoiser GPU medians 27.3 / 12.0 ms; frame intervals 42.6 / 24.4 ms
in the faster sessions and 53.4 / 36.2 ms in the slower sessions. These are separate session states, one launch per
backend in each state; the engine does not resolve the pre-existing slower game state. Frame intervals use the
frame-count/time slope over the final 30 seconds, with sampled instantaneous spreads of about 1 ms for CUDA and
3-4 ms for the engine. The same-frame engine is opt-in; configuration was restored after testing.
Put the local model under `d4r/engine/rr-d`, or set `D4R_ENGINE_MODEL_RR_D_DIR` to its Windows path. The shim runs it
with `[RayReconstruction] Backend = dlss` and `[Engine] Enabled = true`. `D4R_ENGINE_RR_PROFILE` log lines give
game-queue GPU time around the complete engine evaluation; `gpu_frame` identifies the earlier frame whose query
has completed, while `game_interval` is the current call interval. Query readback does not delay the image.
## Stage contract
Every stage shader reads the same data NVIDIA's kernel reads, in the same layouts, so each one can be checked on
engine/rr/build_native_model.sh
+33 -6
View File
@@ -74,10 +74,23 @@ if (( D4R_RR_RELAXED > 0 )); then
fi
# D4R_RR_MIX=1: the 32-deep steps round once through v_fma_mix (the shaders' K32): ~5% faster, not provably bit-exact
[[ "${D4R_RR_MIX:-0}" == 1 ]] && DEFS+=(-DRR_MIX)
glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR...
# D4R_RR_GFX11=1: RDNA3 (no FP8 matrices, no VK_EXT_shader_float8): the stages reproduce ZLUDA's gfx11 lowering bit for
# bit (engine/rr/rr_gfx11.glsl); exact arithmetic only, and each native weight file carries a widened copy (W_WIDE)
GFX11=0
if [[ "${D4R_RR_GFX11:-0}" == 1 ]]; then
GFX11=1; DEFS=(-DRR_EXACT_CONVERT -DRR_GFX11); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D)
(( D4R_RR_RELAXED == 0 )) && [[ "${D4R_RR_MIX:-0}" != 1 ]] ||
echo "D4R_RR_GFX11: relaxed / mix builds do not exist for RDNA3; building the exact model" >&2
D4R_RR_RELAXED=0
fi
# The gfx11 runtime also reproduces NGX's CPU camera arithmetic, whose residues feed temporal specular motion.
(( GFX11 == 0 )) || sed -i '/^preset /a arithmetic gfx11' "$OUT/rr.txt"
# widen WEIGHTS: on RDNA3, append the widened copy and print -DW_WIDE=...; nothing otherwise
widen() { (( GFX11 )) && printf -- '-DW_WIDE=%su' "$(python3 "$RR/widen_weights.py" "$1")" || true; }
glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR... [-DEXTRA]
local out="$1" src="$2"; shift 2
local dir inc=(); for dir in "$@"; do inc+=("-I$dir"); done
"${CAPPED[@]}" glslangValidator --target-env vulkan1.3 -V "${DEFS[@]}" "${inc[@]}" -o "$out" "$src" > "$WORK/glsl.log" 2>&1 ||
local dir inc=() extra=(); for dir in "$@"; do [[ "$dir" == -D* ]] && extra+=("$dir") || inc+=("-I$dir"); done
"${CAPPED[@]}" glslangValidator --target-env vulkan1.3 -V "${DEFS[@]}" "${extra[@]}" "${inc[@]}" -o "$out" "$src" > "$WORK/glsl.log" 2>&1 ||
{ grep -E 'ERROR' "$WORK/glsl.log" | head -5; exit 1; }
# from level 2 on the conversions are left to the driver (as compile_rr2.py does for preset F)
(( D4R_RR_RELAXED >= 2 )) || python3 "$RR/spv_exact_convert.py" "$out" > /dev/null
@@ -94,7 +107,7 @@ for s in enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1; do
cp "$WORK/dll/$s.npz" "$d/weights.npz"
[[ -z "$CAP" ]] || python3 "$RR/swin_ref.py" "$CAP" "$s" "$FRAME" "$d/weights.npz" --blocks "$b" 0,0 | tail -2 | sed "s/^/$s: /"
python3 "$RR/swin_pack.py" "$s" "$d/weights.npz" "$OUT/weights-${s}m.bin" "$d/swin_stage.h" > /dev/null
glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR"
glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR" $(widen "$OUT/weights-${s}m.bin")
native "$s" "weights-${s}m.bin"
done
@@ -103,7 +116,7 @@ d="$WORK/dec0"; mkdir -p "$d"
cp "$WORK/dll/dec0.npz" "$d/weights.npz"
[[ -z "$CAP" ]] || python3 "$RR/dec0_ref.py" "$CAP" "$FRAME" "$d/weights.npz" --blocks 5,5 0,0 | tail -2 | sed 's/^/dec0: /'
python3 "$RR/dec0_pack.py" "$d/weights.npz" "$OUT/weights-dec0m.bin" "$d/dec0_m_w.h" > /dev/null
glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR"
glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR" $(widen "$OUT/weights-dec0m.bin")
native dec0 weights-dec0m.bin
# encoder 0: the translated feature gather + the native rest
@@ -119,7 +132,7 @@ else
cp "$d/enc0_m.bin" "$OUT/weights-enc0m.bin"
fi
python3 "$RR/gen_enc0_m.py" "$OUT/enc0.generated.comp" "$d/enc0_m.comp" > /dev/null
glsl "$OUT/enc0_m.spv" "$d/enc0_m.comp" "$d" "$RR"
glsl "$OUT/enc0_m.spv" "$d/enc0_m.comp" "$d" "$RR" $(widen "$OUT/weights-enc0m.bin")
native enc0 weights-enc0m.bin
# the translated output stage (preset D; preset E's when engine/rr/hkpn.comp is not used) has no subgroup operations
@@ -127,5 +140,19 @@ if grep -q '"shader": "generated from the DLL"' <(python3 -c "import json,sys; p
sed -i -E 's/^stage hkpn hkpn\.spv 32 /stage hkpn hkpn.spv 64 /' "$OUT/rr.txt"
fi
for s in enc0 enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1 dec0; do rm -f "$OUT/weights-$s.bin"; done
# test_layer.py must bind optional images as well as the ones present in a capture. Keep a sidecar for every used
# shader, including native stages, so an absent input never shifts the shader's descriptor bindings.
python3 - "$OUT" <<'PY'
import json, sys
from pathlib import Path
out = Path(sys.argv[1])
for line in (out / 'rr.txt').read_text().splitlines():
fields = line.split()
if not fields or fields[0] != 'stage':
continue
textures = [] if fields[5] == '-' else list(map(int, fields[5].split(',')))
surfaces = [] if fields[6] == '-' else [[int(s.split(':')[0]), s.split(':')[1]] for s in fields[6].split(',')]
(out / fields[2]).with_suffix('.json').write_text(json.dumps({'textures': textures, 'surfaces': surfaces}) + '\n')
PY
printf 'Native Ray Reconstruction model (preset %s) in %s\n' "$([[ "$PRESET" == 4 ]] && echo D || echo E)" "$OUT"
grep '^stage' "$OUT/rr.txt" | awk '{ printf " %s %s (%s-wide)\n", $2, $3, $4 }'
engine/rr/dec0_m.comp
+39 -13
View File
@@ -17,7 +17,9 @@
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#ifndef RR_GFX11
#extension GL_EXT_float_e4m3 : require
#endif
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
@@ -27,7 +29,9 @@ layout(local_size_x = 128) in;
#include "rr_exact.glsl"
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
@@ -40,10 +44,15 @@ layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
@@ -65,6 +74,7 @@ layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144;
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// shared memory, aliased by phase: expand and prologue [LOW | EX], then the attention [K | V^T]
@@ -124,16 +134,14 @@ void main()
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDA(we[mm][kt], W_WE, 64u, 2u * wv + mm, kt);
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(lb[kt], LOW, n * 256u + 4u * kt, 16u, CM);
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDS8B(lb[kt], LOW, n * 256u + 4u * kt, 16u);
[[unroll]] for (uint mm = 0u; mm < 2u; ++mm) {
const uint m = 2u * wv + mm;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
K32(a, we[mm][0], lb[0], we[mm][1], lb[1])
K32(a, we[mm][2], lb[2], we[mm][3], lb[3])
B8 q;
Q8(q, a)
coopMatStore(q, POOL, EX_AT + n * 512u + 4u * m, 32u, CM);
Q8STB(a, POOL, EX_AT + n * 512u + 4u * m, 32u)
}
}
}
@@ -149,8 +157,12 @@ void main()
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) {
const uint e = POOL[EX_AT + r * 32u + 8u * s + wd], sk = arena.w[sb + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
#ifdef RR_GFX11
precise float16_t x = g11_e4h((e >> (8u * b)) & 0xffu) + g11_e4h((sk >> (8u * b)) & 0xffu);
#else
precise float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((e >> (8u * b)) & 0xffu)))
+ float16_t(uintBitsToFloate4m3EXT(uint8_t((sk >> (8u * b)) & 0xffu)));
#endif
v[4u * wd + b] = x;
}
}
@@ -169,7 +181,7 @@ void main()
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) {
floate4m3_t c[4];
F8 c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
@@ -186,7 +198,11 @@ void main()
RAC R[2];
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
HAC x;
#ifdef RR_GFX11
[[unroll]] for (int i = 0; i < 8; ++i) x[i] = XR[wv * 512u + 16u * m + col * 32u + 8u * hl + uint(i)];
#else
coopMatLoad(x, XR, wv * 512u + 16u * m, 32u, CM);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t c = float16_t(x[i]) + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = RV(c);
@@ -201,21 +217,17 @@ void main()
A8 wk[2], xa[2];
B8 wvb[2], xb[2];
[[unroll]] for (uint kt = 0u; kt < 2u; ++kt) {
coopMatLoad(xb[kt], XT, n * 128u + 4u * kt, 8u, CM);
coopMatLoad(xa[kt], XT, n * 128u + 4u * kt, 8u, RM);
LDS8B(xb[kt], XT, n * 128u + 4u * kt, 8u);
LDS8A(xa[kt], XT, n * 128u + 4u * kt, 8u);
LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
RAC a = RAC(RV(float16_t(0.0)));
K32(a, wk[0], xb[0], wk[1], xb[1])
A8 ka;
Q8(ka, a)
coopMatStore(ka, POOL, n * 128u + 4u * c, 8u, RM);
Q8STA(a, POOL, n * 128u + 4u * c, 8u)
RAC v = RAC(RV(float16_t(0.0)));
K32(v, xa[0], wvb[0], xa[1], wvb[1])
A8 va;
Q8(va, v)
coopMatStore(va, POOL, 768u + c * 384u + 4u * n, 24u, RM);
Q8STA(v, POOL, 768u + c * 384u + 4u * n, 24u)
}
barrier();
#include "rr_attn.glsl"
@@ -228,6 +240,9 @@ void main()
// ---------------------------------------------------------------- MLP
B8 Nb[2];
#ifdef RR_GFX11
F8V NbC[2] = F8V[2](uvec2(0u), uvec2(0u));
#endif
{
float16_t rv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
@@ -255,10 +270,17 @@ void main()
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN1 / 2u + k];
precise float16_t n = rv[m][i] * g;
#ifdef RR_GFX11
NbC[m][i >> 2] |= q8h(n) << (8u * uint(i & 3));
#else
Nb[m][i] = q8h(n);
#endif
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = RV(a);
}
#ifdef RR_GFX11
[[unroll]] for (uint m = 0u; m < 2u; ++m) G11_B(Nb[m], NbC[m])
#endif
}
for (uint br = 0u; br < 4u; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
@@ -269,7 +291,11 @@ void main()
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
#ifdef RR_GFX11
{ F8V hc_; ACTQ8C(hc_, h) G11_B(Hb[c], hc_) }
#else
ACTQ8(Hb[c], h)
#endif
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
engine/rr/enc0_m_decl.glsl
+10
View File
@@ -7,17 +7,24 @@
#extension GL_GOOGLE_include_directive : require
#include "enc0_m_w.h"
#include "rr_exact.glsl"
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
@@ -39,6 +46,7 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features
#define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major
@@ -55,10 +63,12 @@ ivec2 keyXY(uint t, int x0, int y0)
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
#ifndef RR_GFX11
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
#endif
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// the translated prologue's shared buffer sh: input_noise (10 x 10 pixels, 4 f16) at byte 0, input_tensor (two planes
// of 10 x 10 pixels x 16 FP8 features) at byte 800
engine/rr/enc0_m_tail.glsl
+56 -12
View File
@@ -23,15 +23,20 @@
RAC R[2];
for (uint n = wv; n < 6u; n += 4u) {
B8 fb0, fb1, hb_[2];
coopMatLoad(fb0, XF, n * 128u, 8u, CM);
coopMatLoad(fb1, XF, n * 128u + 4u, 8u, CM);
LDS8B(fb0, XF, n * 128u, 8u);
LDS8B(fb1, XF, n * 128u + 4u, 8u);
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]);
A8 a0, a1;
LDA(a0, W_W1, 32u, m, 0u); LDA(a1, W_W1, 32u, m, 1u);
K32(a, a0, fb0, a1, fb1)
#ifdef RR_GFX11
{ F8V hc_ = uvec2(0u); [[unroll]] for (int i = 0; i < 8; ++i) hc_[i >> 2] |= q8s(max(float(a[i]), 0.0)) << (8u * uint(i & 3));
G11_B(hb_[m], hc_) }
#else
[[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(float(a[i]), 0.0));
#endif
}
float16_t xv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
@@ -67,13 +72,25 @@
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
#ifdef RR_GFX11
F8V xq = uvec2(0u);
#else
B8 xq;
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(16u * m + 8u * hl + uint(i))];
precise float16_t nn = xv[m][i] * g;
#ifdef RR_GFX11
xq[i >> 2] |= q8h(nn) << (8u * uint(i & 3));
#else
xq[i] = q8h(nn);
#endif
}
#ifdef RR_GFX11
G11_ST(xq, XT, n * 128u + 4u * m, 8u)
#else
coopMatStore(xq, XT, n * 128u + 4u * m, 8u, CM);
#endif
}
}
barrier(); // XT complete; XF is read (the attention reuses its memory)
@@ -85,21 +102,17 @@
A8 wk[2], xa[2];
B8 wvb[2], xb[2];
[[unroll]] for (uint kt = 0u; kt < 2u; ++kt) {
coopMatLoad(xb[kt], XT, n * 128u + 4u * kt, 8u, CM);
coopMatLoad(xa[kt], XT, n * 128u + 4u * kt, 8u, RM);
LDS8B(xb[kt], XT, n * 128u + 4u * kt, 8u);
LDS8A(xa[kt], XT, n * 128u + 4u * kt, 8u);
LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt);
LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c);
}
RAC a = RAC(RV(float16_t(0.0)));
K32(a, wk[0], xb[0], wk[1], xb[1])
A8 ka;
Q8(ka, a)
coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
Q8STA(a, KVS, n * 128u + 4u * c, 8u)
RAC vv = RAC(RV(float16_t(0.0)));
K32(vv, xa[0], wvb[0], xa[1], wvb[1])
A8 va;
Q8(va, vv)
coopMatStore(va, KVS, 768u + c * 384u + 4u * n, 24u, RM);
Q8STA(vv, KVS, 768u + c * 384u + 4u * n, 24u)
}
barrier();
#include "rr_attn.glsl"
@@ -111,6 +124,9 @@
}
// MLP (dec0_m.comp's)
B8 Nb[2];
#ifdef RR_GFX11
F8V NbC[2] = F8V[2](uvec2(0u), uvec2(0u));
#endif
{
float16_t rv[2][8];
[[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
@@ -137,10 +153,17 @@
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN1 / 2u + k];
precise float16_t nn = rv[m][i] * g;
#ifdef RR_GFX11
NbC[m][i >> 2] |= q8h(nn) << (8u * uint(i & 3));
#else
Nb[m][i] = q8h(nn);
#endif
precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k];
R[m][i] = RV(a);
}
#ifdef RR_GFX11
[[unroll]] for (uint m = 0u; m < 2u; ++m) G11_B(Nb[m], NbC[m])
#endif
}
for (uint br = 0u; br < 4u; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
@@ -151,7 +174,11 @@
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
#ifdef RR_GFX11
{ F8V hc_; ACTQ8C(hc_, h) G11_B(Hb[c], hc_) }
#else
ACTQ8(Hb[c], h)
#endif
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
@@ -171,8 +198,15 @@
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
barrier(); // XT is read; the merge rows reuse it
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
#ifdef RR_GFX11
F8V tc_;
G11_CODES(tc_, R[m])
G11_B(Tb[m], tc_)
const uint q0 = tc_.x, q1 = tc_.y;
#else
Q8(Tb[m], R[m])
const uint q0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), q1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
#endif
const uint at = 4u * m + 2u * hl;
if (inside) { arena.w[o + at] = q0; arena.w[o + at + 1u] = q1; }
const uint gm = mr * 32u + blk * 8u + 4u * m + 2u * hl;
@@ -232,11 +266,20 @@
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint ks = 0u; ks < 4u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, 32u, CM);
coopMatLoad(g1, MG, 8u * ks + 4u, 32u, CM);
LDS8B(g0, MG, 8u * ks, 32u);
LDS8B(g1, MG, 8u * ks + 4u, 32u);
LDA(w0, W_WM, 128u, m, 2u * ks); LDA(w1, W_WM, 128u, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
#ifdef RR_GFX11
F8V q;
G11_CODES(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = (mrgOff + (m >> 1u) * lplane + uint(Y * hx + X) * 32u) / 4u + 4u * (m & 1u) + 2u * hl;
arena.w[at] = q.x;
arena.w[at + 1u] = q.y;
}
#else
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
@@ -244,6 +287,7 @@
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
#endif
}
if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
engine/rr/ptx2rr.py
+90 -4
View File
@@ -256,7 +256,17 @@ uint e4m3_to_h(uint b)
float f = e == 0u ? ldexp(float(m), -9) : ldexp(float(8u + m), int(max(e, 1u)) - 10);
return s | HU(float16_t(f));
}
#ifdef SOFTWARE_E4M3
#if defined(RR_GFX11)
// RDNA3 (engine/rr/rr_gfx11.glsl): ZLUDA's gfx11 conversions exactly. e4m3 -> f16: NaN -> 0x7fff with the sign
uint g11_e4m3_to_h(uint c)
{
const uint m = c & 0x7fu;
uint h = m >= 8u ? (m << 7u) + 0x2000u : uint(float16BitsToUint16(float16_t(float(m & 7u) * 0.001953125)));
h = m == 0x7fu ? 0x7fffu : h;
return ((c & 0x80u) << 8u) | h;
}
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(g11_e4m3_to_h(v & 0xffu) | (g11_e4m3_to_h((v >> 8) & 0xffu) << 16)); }
#elif defined(SOFTWARE_E4M3)
f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h(v & 0xffu) | (e4m3_to_h((v >> 8) & 0xffu) << 16)); }
#else
// the hardware conversion: exact for every non-NaN code (2026-10-12); NaN keeps e4m3_to_h's sign | 0x7e00
@@ -284,7 +294,31 @@ uint h2e4m3(float f, bool relu)
else { int p = findMSB(n); bits = (uint(qe + p + 7) << 3) | ((n >> uint(p - 3)) & 7u); }
return s | min(bits, 0x7eu);
}
#ifdef SOFTWARE_E4M3
#if defined(RR_GFX11)
// f16 -> e4m3 as ZLUDA computes it on gfx11 (HIP's satfinite conversion, infinities clamped first): nearest even,
// |x| > 448 -> 0x7e, NaN -> 0x7f, the sign kept; .relu first zeroes negative non-NaN halves
uint g11_h2e4m3(uint h)
{
const uint s = (h >> 8u) & 0x80u, a = h & 0x7fffu;
if (a > 0x7c00u) return s | 0x7fu;
uint c;
if (a >= 0x2400u) {
const uint t = a - 0x2000u;
c = min((t + 0x3fu + ((t >> 7u) & 1u)) >> 7u, 0x7eu);
} else {
c = uint(roundEven(float(uint16BitsToFloat16(uint16_t(a))) * 512.0));
}
return s | c;
}
#define G11_HAVE_H2E4M3
uint g11_relu(uint h) { return (h & 0x8000u) != 0u && (h & 0x7fffu) <= 0x7c00u ? 0u : h; }
uint f16x2_to_e4m3x2(f16vec2 h, bool relu)
{
uint x = HU(h.x), y = HU(h.y);
if (relu) { x = g11_relu(x); y = g11_relu(y); }
return g11_h2e4m3(x) | (g11_h2e4m3(y) << 8);
}
#elif defined(SOFTWARE_E4M3)
uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3(float(h.x), relu) | (h2e4m3(float(h.y), relu) << 8); }
#else
// cvt.rn.satfinite(.relu).e4m3x2.f16x2 on the hardware conversion: round to nearest even after clamping to +-448
@@ -305,10 +339,38 @@ uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3_hw(h.x, relu) | (h2e4
// place (A row-major 16x32, B column-major 32x8 in a 16-column tile, C/D row-major f16 16x16), then loaded as
// whole matrices. Lane l = 4g + t holds A rows g and g+8 (k 4t.. and 16+4t..), B column g (same k), C/D row g
// and g+8, columns 2t and 2t+1.
#define MMA coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MMB coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define MMF coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define MMH coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifdef RR_GFX11
// RDNA3: f16 operands holding ZLUDA's widening (value * 2^-8, k positions 2q + j <- 4q + j, 8 + 2q + j <- 4q + 2 + j
// in each K half), C scaled by 2^-16 into the two products and back (engine/rr/rr_gfx11.glsl)
#define MMA coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MMB coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
uvec2 g11_sc4(uint packed)
{
uint mag = packed & 0x7f7f7f7fu;
mag |= (mag + 0x01010101u) & 0x80808080u;
const uint sign = packed & 0x80808080u;
return uvec2(((mag & 0xffu) << 7u) | ((mag & 0xff00u) << 15u) | ((sign & 0x80u) << 8u) | ((sign & 0x8000u) << 16u),
(((mag >> 16u) & 0xffu) << 7u) | ((mag >> 24u) << 23u) | ((sign & 0x800000u) >> 8u) | (sign & 0x80000000u));
}
#define G11_FILL(m, w0, w1, w2, w3) { \
const uvec2 s0_ = g11_sc4(w0), s1_ = g11_sc4(w1), s2_ = g11_sc4(w2), s3_ = g11_sc4(w3); f16vec2 l_; \
l_ = unpackFloat2x16(s0_.x); (m)[0] = l_.x; (m)[1] = l_.y; l_ = unpackFloat2x16(s0_.y); (m)[8] = l_.x; (m)[9] = l_.y; \
l_ = unpackFloat2x16(s1_.x); (m)[2] = l_.x; (m)[3] = l_.y; l_ = unpackFloat2x16(s1_.y); (m)[10] = l_.x; (m)[11] = l_.y; \
l_ = unpackFloat2x16(s2_.x); (m)[4] = l_.x; (m)[5] = l_.y; l_ = unpackFloat2x16(s2_.y); (m)[12] = l_.x; (m)[13] = l_.y; \
l_ = unpackFloat2x16(s3_.x); (m)[6] = l_.x; (m)[7] = l_.y; l_ = unpackFloat2x16(s3_.y); (m)[14] = l_.x; (m)[15] = l_.y; }
#define G11_HAVE_FILL
#define coopMatLoad8(m, arr, off, stride) { const uint b_ = (off) + (gl_SubgroupInvocationID & 15u) * (stride); \
G11_FILL(m, arr[b_], arr[b_ + 1u], arr[b_ + 2u], arr[b_ + 3u]) }
#define MM8(x1, y1, x0, y0, c) MMF acc_ = MMF(c); for (int i_ = 0; i_ < acc_.length(); ++i_) acc_[i_] *= 1.52587890625e-05; \
acc_ = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, acc_)); for (int i_ = 0; i_ < acc_.length(); ++i_) acc_[i_] *= 65536.0;
#else
#define MMA coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MMB coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define coopMatLoad8(m, arr, off, stride) coopMatLoad(m, arr, off, stride, gl_CooperativeMatrixLayoutRowMajor)
#define MM8(x1, y1, x0, y0, c) MMF acc_ = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
#endif
void mma_e4m3_f16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1, out uint d0, out uint d1)
{
const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u;
@@ -323,6 +385,16 @@ void mma_e4m3_f16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0,
mms[C + g * 8u + 4u + t] = 0u; mms[C + (g + 8u) * 8u + 4u + t] = 0u;
memoryBarrierShared(); subgroupBarrier();
MMA x0, x1; MMB y0, y1; MMH c;
#ifdef RR_GFX11
coopMatLoad8(x0, mms, A, 8u);
coopMatLoad8(x1, mms, A + 4u, 8u);
coopMatLoad8(y0, mms, B, 8u);
coopMatLoad8(y1, mms, B + 4u, 8u);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MM8(x1, y1, x0, y0, c)
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc_), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#else
coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor);
@@ -331,6 +403,7 @@ void mma_e4m3_f16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0,
MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#endif
memoryBarrierShared(); subgroupBarrier();
d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t];
}
@@ -350,6 +423,16 @@ void mma_e4m3_f16_pair(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uin
mms[C + g * 8u + 4u + t] = f0; mms[C + (g + 8u) * 8u + 4u + t] = f1;
memoryBarrierShared(); subgroupBarrier();
MMA x0, x1; MMB y0, y1; MMH c;
#ifdef RR_GFX11
coopMatLoad8(x0, mms, A, 8u);
coopMatLoad8(x1, mms, A + 4u, 8u);
coopMatLoad8(y0, mms, B, 8u);
coopMatLoad8(y1, mms, B + 4u, 8u);
coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
MM8(x1, y1, x0, y0, c)
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc_), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#else
coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor);
@@ -358,6 +441,7 @@ void mma_e4m3_f16_pair(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uin
MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c)));
memoryBarrierShared(); subgroupBarrier();
coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor);
#endif
memoryBarrierShared(); subgroupBarrier();
d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t];
h0 = mms[C + g * 8u + 4u + t]; h1 = mms[C + (g + 8u) * 8u + 4u + t];
@@ -803,7 +887,9 @@ class RRT(base.T):
w('#extension GL_KHR_shader_subgroup_ballot : require')
w('#extension GL_KHR_cooperative_matrix : require')
w('#extension GL_KHR_memory_scope_semantics : require')
w('#ifndef RR_GFX11')
w('#extension GL_EXT_float_e4m3 : require')
w('#endif')
w('#extension GL_EXT_spirv_intrinsics : require')
w(f'// Generated by engine/rr/ptx2rr.py from {self.kernel}. Model data derived from the user\'s DLL; do not commit.')
w(f'layout(local_size_x = {self.bx * self.by * self.bz}, local_size_y = 1, local_size_z = 1) in;')
engine/rr/rr_attn.glsl
+27 -6
View File
@@ -13,17 +13,23 @@
#define RR_VS0 0u
#endif
B8 Qb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], RR_KS, RR_KS0 + wv * 128u + 4u * c, 8u, CM);
[[unroll]] for (uint c = 0u; c < 2u; ++c) LDS8B(Qb[c], RR_KS, RR_KS0 + wv * 128u + 4u * c, 8u);
// scores S^T (keys x this wave's queries) + bias, then the cubic softmax down each column
B8 Pb[6];
{
float16_t e[6][8];
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) {
HAC c;
#ifdef RR_GFX11
// the accumulator layout the stages are written for (rr_gfx11.glsl): element i is row i + 8 hl
[[unroll]] for (int i = 0; i < 8; ++i)
c[i] = w16.w[(W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv + (8u * hl + uint(i)) * 64u + col];
#else
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
#endif
A8 k0, k1;
coopMatLoad(k0, RR_KS, RR_KS0 + kt * 128u, 8u, RM);
coopMatLoad(k1, RR_KS, RR_KS0 + kt * 128u + 4u, 8u, RM);
LDS8A(k0, RR_KS, RR_KS0 + kt * 128u, 8u);
LDS8A(k1, RR_KS, RR_KS0 + kt * 128u + 4u, 8u);
K32H(c, k0, Qb[0], k1, Qb[1])
// the cubic numerator from enc1_pack.cubic_table, indexed by the score's binary16 bits
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
@@ -49,12 +55,23 @@
// reciprocal (within an ulp) already rounds correctly
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt)
[[unroll]] for (uint kt = 0u; kt < 6u; ++kt) {
#ifdef RR_GFX11
F8V pc_ = uvec2(0u);
#endif
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
#ifdef RR_GFX11
pc_[q >> 1u] |= (q8h(pv.x) | (q8h(pv.y) << 8u)) << (16u * (q & 1u));
#else
// a probability is in [0, 1]: the plain conversion, no saturation clamp, both in one instruction
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
#endif
}
#ifdef RR_GFX11
G11_B(Pb[kt], pc_)
#endif
}
}
// O^T = V^T P^T, then the out-projection onto the residual
B8 Ob[2];
@@ -62,9 +79,13 @@
RAC o = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint ks = 0u; ks < 3u; ++ks) {
A8 v0, v1;
coopMatLoad(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u, RM);
coopMatLoad(v1, RR_VS, RR_VS0 + c * 384u + 8u * ks + 4u, 24u, RM);
LDS8A(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u);
LDS8A(v1, RR_VS, RR_VS0 + c * 384u + 8u * ks + 4u, 24u);
K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
}
#ifdef RR_GFX11
{ F8V oc_; G11_CODES(oc_, o) G11_B(Ob[c], oc_) }
#else
Q8(Ob[c], o)
#endif
}
engine/rr/rr_exact.glsl
+32 -1
View File
@@ -1,5 +1,6 @@
// Exact NVIDIA arithmetic helpers for the native Ray Reconstruction stages (engine/rr/*_m.comp, enc1_native.comp).
// Requires GL_EXT_shader_explicit_arithmetic_types, GL_EXT_float_e4m3.
// Requires GL_EXT_shader_explicit_arithmetic_types, GL_EXT_float_e4m3 (RR_GFX11: GL_KHR_shader_subgroup_shuffle instead;
// FP8 values are then codes, rr_gfx11.glsl).
//
// RADV folds a register-only float(float16_t(x)) back to x, so every binary16 rounding whose result is used again goes
// through a conversion with an opaque zero (RR_ZERO16, set in main from something the compiler cannot evaluate, e.g.
@@ -30,6 +31,20 @@ 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); }
#ifdef RR_GFX11
#include "rr_gfx11.glsl"
#define F8 uint
uint q8h(float16_t x) { return g11_code(x); }
uint q8s(float x) { return g11_codef(x); }
uint code(uint c) { return c; }
float16_t e4h(uint c) { return g11_e4h(c); }
// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
#define ACTQ(h) uint(w8u.w[W_ACT + bh(h)])
// an accumulator's activation codes (accumulator layout)
#define ACTQ8C(dst, h) { (dst) = uvec2(0u); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_ >> 2] |= ACTQ(float16_t((h)[i_])) << (8u * uint(i_ & 3)); }
#else
#define F8 floate4m3_t
float16_t e4h(floate4m3_t v) { return float16_t(v); }
// cvt.rn.satfinite.e4m3 (NaN -> 0x7f)
floate4m3_t q8(float16_t h)
{
@@ -41,6 +56,7 @@ floate4m3_t q8s(float x) { floate4m3_t r; saturatedConvertEXT(r, 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)); }
#endif
float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); }
uint bh(float16_t h) { return uint(float16BitsToUint16(h)); }
@@ -106,6 +122,7 @@ float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid)
return hb(bh(float16_t(s)) | uint(RR_ZERO16));
#endif
}
#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; }
// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
@@ -117,6 +134,7 @@ float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid)
#else
#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
#endif
#endif
float16_t act(float16_t x)
{
precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0));
@@ -152,3 +170,16 @@ float16_t cubic_e(float16_t x)
const float16_t u = fma16(m, t, float16_t(1.375));
return hb((bh(u) << 5u) & 0x7FE0u);
}
// FP8 operands from a shared byte image (word offset and stride): an A row-major, a B column-major; and an accumulator
// quantized and stored as the transposed A row-major / as a B column-major (the same bytes)
#ifdef RR_GFX11
#define LDS8A(m, ARR, off, stride) LD8A(m, ARR, off, stride)
#define LDS8B(m, ARR, off, stride) LD8B(m, ARR, off, stride)
#define Q8STA(src, ARR, off, stride) { F8V c_; G11_CODES(c_, src) G11_ST(c_, ARR, off, stride) }
#define Q8STB(src, ARR, off, stride) Q8STA(src, ARR, off, stride)
#else
#define LDS8A(m, ARR, off, stride) coopMatLoad(m, ARR, off, stride, gl_CooperativeMatrixLayoutRowMajor)
#define LDS8B(m, ARR, off, stride) coopMatLoad(m, ARR, off, stride, gl_CooperativeMatrixLayoutColumnMajor)
#define Q8STA(src, ARR, off, stride) { A8 q; Q8(q, src) coopMatStore(q, ARR, off, stride, gl_CooperativeMatrixLayoutRowMajor); }
#define Q8STB(src, ARR, off, stride) { B8 q; Q8(q, src) coopMatStore(q, ARR, off, stride, gl_CooperativeMatrixLayoutColumnMajor); }
#endif
engine/rr/rr_gfx11.glsl
+95
View File
@@ -0,0 +1,95 @@
// RDNA3 (gfx11, RR_GFX11) forms of the native Ray Reconstruction stages' FP8 operations, included by rr_exact.glsl.
// gfx11 has no FP8 matrix instruction and RADV no VK_EXT_shader_float8, so the reference is ZLUDA's gfx11 lowering of
// mma.m16n8k32.f16.e4m3.e4m3.f16 (D4R_ZLUDA_WMMA_FP8, patches/zluda 0004), which these reproduce bit for bit:
// - each e4m3 byte is widened on its bits to the binary16 of value * 2^-8 (magnitude at bits 7..14, the sign at 15,
// the NaN code to a NaN), so the products carry 2^-16;
// - the accumulator enters scaled by 2^-16, the two K halves (k 0..15, 16..31) run as two chained
// v_wmma_f32_16x16x16_f16, the sum is scaled back and rounded to binary16 once;
// - inside each K half the operand positions hold k in ZLUDA's order: position 2q + j holds k 4q + j and position
// 8 + 2q + j holds k 4q + 2 + j (q < 4, j < 2). The instruction's f32 sum depends on that order (tests on an
// RX 7700 XT: 29% of the outputs change when the same products are permuted), so it is part of the arithmetic.
// FP8 values in registers are their codes (uint); an operand matrix is coopmat<float16_t> holding the widened values.
// RADV's gfx11 layouts: an A or B operand lane l holds row / column l & 15, element i = k i (both half-waves alike);
// an accumulator lane holds column l & 15, element i = row 2i + (l >> 4). The stages' elementwise code is written for
// gfx12's accumulator layout (element i = row i + 8 (l >> 4)), so every A operand is loaded with its rows permuted:
// lane l takes row g11_row(l & 15), and a product's physical row 2i + h is then row i + 8h, i.e. the gfx12 layout.
// That only relabels the outputs (each is its own dot product), so it is exact. Accumulators loaded or built
// elementwise therefore use the gfx12 positions too.
#define A8 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
// cvt.rn.satfinite.e4m3 of binary16 bits as ZLUDA computes it on gfx11 (HIP's conversion after clamping infinities
// to 448): nearest even, |x| > 448 -> 0x7e, NaN -> 0x7f, the sign kept. Equal to it for all 65536 inputs.
#ifndef G11_HAVE_H2E4M3 // ptx2rr.py's preamble (a translated head before this file) defines the same
uint g11_h2e4m3(uint h)
{
const uint s = (h >> 8u) & 0x80u, a = h & 0x7fffu;
if (a > 0x7c00u) return s | 0x7fu;
uint c;
if (a >= 0x2400u) { // |x| >= 2^-6: an e4m3 normal; rebias 15 -> 7 and round 10 -> 3 mantissa bits
const uint t = a - 0x2000u;
c = min((t + 0x3fu + ((t >> 7u) & 1u)) >> 7u, 0x7eu);
} else { // e4m3 subnormals: steps of 2^-9 (x * 2^9 is exact in f32)
c = uint(roundEven(float(uint16BitsToFloat16(uint16_t(a))) * 512.0));
}
return s | c;
}
#endif
uint g11_code(float16_t h) { return g11_h2e4m3(uint(float16BitsToUint16(h))); }
// an f32 that holds a binary16 value
uint g11_codef(float x) { return g11_code(float16_t(x)); }
// the exact binary16 value of a code (ZLUDA's e4m3_to_f16_bits; NaN -> 0x7fff with the sign)
float16_t g11_e4h(uint c)
{
const uint m = c & 0x7fu;
uint h = m >= 8u ? (m << 7u) + 0x2000u : uint(float16BitsToUint16(float16_t(float(m & 7u) * 0.001953125)));
h = m == 0x7fu ? 0x7fffu : h;
return uint16BitsToFloat16(uint16_t(((c & 0x80u) << 8u) | h));
}
#ifndef G11_HAVE_FILL
// four codes -> the scaled binary16 pairs (bytes 0, 1) and (bytes 2, 3), ZLUDA's e4m3x4_to_scaled_f16x2x2
uvec2 g11_sc4(uint packed)
{
uint mag = packed & 0x7f7f7f7fu;
mag |= (mag + 0x01010101u) & 0x80808080u;
const uint sign = packed & 0x80808080u;
const uint lo = ((mag & 0xffu) << 7u) | ((mag & 0xff00u) << 15u) | ((sign & 0x80u) << 8u) | ((sign & 0x8000u) << 16u);
const uint hi = (((mag >> 16u) & 0xffu) << 7u) | ((mag >> 24u) << 23u) | ((sign & 0x800000u) >> 8u) | (sign & 0x80000000u);
return uvec2(lo, hi);
}
// an operand from its lane's 16 codes in k order (four words): positions in ZLUDA's order
#define G11_FILL(m, w0, w1, w2, w3) { \
const uvec2 s0_ = g11_sc4(w0), s1_ = g11_sc4(w1), s2_ = g11_sc4(w2), s3_ = g11_sc4(w3); f16vec2 l_; \
l_ = unpackFloat2x16(s0_.x); (m)[0] = l_.x; (m)[1] = l_.y; l_ = unpackFloat2x16(s0_.y); (m)[8] = l_.x; (m)[9] = l_.y; \
l_ = unpackFloat2x16(s1_.x); (m)[2] = l_.x; (m)[3] = l_.y; l_ = unpackFloat2x16(s1_.y); (m)[10] = l_.x; (m)[11] = l_.y; \
l_ = unpackFloat2x16(s2_.x); (m)[4] = l_.x; (m)[5] = l_.y; l_ = unpackFloat2x16(s2_.y); (m)[12] = l_.x; (m)[13] = l_.y; \
l_ = unpackFloat2x16(s3_.x); (m)[6] = l_.x; (m)[7] = l_.y; l_ = unpackFloat2x16(s3_.y); (m)[14] = l_.x; (m)[15] = l_.y; }
#endif
// the A operand row a lane holds: physical row r carries row (r >> 1) + 8 (r & 1)
uint g11_row(uint r) { return (r >> 1u) | ((r & 1u) << 3u); }
// an A operand row-major (rows permuted) or a B operand column-major from a shared byte image (word offset and
// stride, as coopMatLoad takes them for an FP8 matrix): the lane's row / column is 16 consecutive bytes
#define LD8A(m, ARR, off, stride) { const uint b_ = (off) + g11_row(gl_SubgroupInvocationID & 15u) * (stride); \
G11_FILL(m, ARR[b_], ARR[b_ + 1u], ARR[b_ + 2u], ARR[b_ + 3u]) }
#define LD8B(m, ARR, off, stride) { const uint b_ = (off) + (gl_SubgroupInvocationID & 15u) * (stride); \
G11_FILL(m, ARR[b_], ARR[b_ + 1u], ARR[b_ + 2u], ARR[b_ + 3u]) }
// natural k order (a matrix loaded from the widened weights) -> ZLUDA's positions
#define G11_PERM(d, s) { (d)[0] = (s)[0]; (d)[1] = (s)[1]; (d)[2] = (s)[4]; (d)[3] = (s)[5]; (d)[4] = (s)[8]; \
(d)[5] = (s)[9]; (d)[6] = (s)[12]; (d)[7] = (s)[13]; (d)[8] = (s)[2]; (d)[9] = (s)[3]; (d)[10] = (s)[6]; \
(d)[11] = (s)[7]; (d)[12] = (s)[10]; (d)[13] = (s)[11]; (d)[14] = (s)[14]; (d)[15] = (s)[15]; }
// FP8 vectors in the accumulator layout: element i of a lane (row i + 8 half) as byte i & 3 of word i >> 2
#define F8V uvec2
// the codes of an accumulator that holds binary16 values
#define G11_CODES(dst, src) { (dst) = uvec2(0u); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) \
(dst)[i_ >> 2] |= g11_codef((src)[i_]) << (8u * uint(i_ & 3)); }
// the B operand that is the accumulator-layout codes c (the gfx12 elementwise conversion of an accumulator into B):
// rows 0..7 are half-wave 0's codes, 8..15 half-wave 1's
#define G11_B(dst, c) { const uvec2 o_ = uvec2(subgroupShuffleXor((c).x, 16u), subgroupShuffleXor((c).y, 16u)); \
const bool h_ = gl_SubgroupInvocationID >= 16u; const uvec2 e0_ = h_ ? o_ : (c), e1_ = h_ ? (c) : o_; \
G11_FILL(dst, e0_.x, e0_.y, e1_.x, e1_.y) }
// store the codes c column-major as bytes (column j at word off + j * stride): what an FP8 B stored column-major or
// an FP8 A (the accumulator's transpose) stored row-major leaves; each half-wave writes its 8 rows
#define G11_ST(c, ARR, off, stride) { const uint b_ = (off) + (gl_SubgroupInvocationID & 15u) * (stride) + \
(gl_SubgroupInvocationID >= 16u ? 2u : 0u); ARR[b_] = (c).x; ARR[b_ + 1u] = (c).y; }
engine/rr/rr_gfx11_k32.glsl
+29
View File
@@ -0,0 +1,29 @@
// RR_GFX11: the weight loads and the 32-deep step of the native stages (rr_gfx11.glsl explains the arithmetic).
// Included after ACC / HAC / RM / CM and the binding-1 views. W_WIDE (a -D of the build) is the f16 element at which
// the weight file carries its widened copy: element W_WIDE + j is byte j of the image widened like an FP8 operand
// (value * 2^-8), so a stage's byte offsets address it unchanged.
layout(set = 0, binding = 1, std430) readonly buffer W128 { uvec4 w[]; } w128;
#ifndef W_WIDE
#error RR_GFX11 needs W_WIDE (the widened weight copy, engine/rr/widen_weights.py)
#endif
// A: the lane's (permuted) row is 16 consecutive f16, two 16-byte loads (base and tiles are multiples of 16 bytes)
#define LDA(dst, base, k, mt, kt) { const uint e_ = W_WIDE + (base) + 256u * ((mt) * ((k) / 16u) + (kt)) + \
16u * g11_row(gl_SubgroupInvocationID & 15u); const uvec4 u0_ = w128.w[e_ >> 3u], u1_ = w128.w[(e_ >> 3u) + 1u]; \
f16vec2 l_; l_ = unpackFloat2x16(u0_.x); (dst)[0] = l_.x; (dst)[1] = l_.y; l_ = unpackFloat2x16(u0_.y); (dst)[8] = l_.x; (dst)[9] = l_.y; \
l_ = unpackFloat2x16(u0_.z); (dst)[2] = l_.x; (dst)[3] = l_.y; l_ = unpackFloat2x16(u0_.w); (dst)[10] = l_.x; (dst)[11] = l_.y; \
l_ = unpackFloat2x16(u1_.x); (dst)[4] = l_.x; (dst)[5] = l_.y; l_ = unpackFloat2x16(u1_.y); (dst)[12] = l_.x; (dst)[13] = l_.y; \
l_ = unpackFloat2x16(u1_.z); (dst)[6] = l_.x; (dst)[7] = l_.y; l_ = unpackFloat2x16(u1_.w); (dst)[14] = l_.x; (dst)[15] = l_.y; }
#define LDB(dst, base, n, kt, nt) { B8 t_; coopMatLoad(t_, w16.w, W_WIDE + (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM); G11_PERM(dst, t_) }
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#if defined(RR_MIX) || defined(RR_RELAXED)
#error RR_GFX11 builds the exact arithmetic only
#endif
#define RAC ACC
#define RV(x) float(x)
// C * 2^-16 (exact), the two halves, * 2^16 (exact), one rounding to binary16
#define K32(acc, a0, b0, a1, b1) { [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] *= 1.52587890625e-05; \
(acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); \
[[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = rnd16((acc)[i_] * 65536.0); }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = ACC(acc); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) p_[i_] *= 1.52587890625e-05; \
p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), p_)); \
[[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] * 65536.0); }
engine/rr/swin_m.comp
+61 -11
View File
@@ -12,7 +12,9 @@
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#ifndef RR_GFX11
#extension GL_EXT_float_e4m3 : require
#endif
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
@@ -26,7 +28,9 @@ layout(local_size_x = 128) in;
#define NB (C / 8)
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
@@ -46,10 +50,15 @@ layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
@@ -71,6 +80,7 @@ layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
// NVIDIA's reduction over a short list (swin_ref.pairs_tree)
@@ -124,11 +134,16 @@ ivec2 keyXY(uint t, int x0, int y0)
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); }
#ifdef RR_GFX11
uint pack4(uint a, uint b, uint c, uint d) { return a | (b << 8u) | (c << 16u) | (d << 24u); }
float16_t e4(uint w, uint b) { return g11_e4h((w >> (8u * b)) & 0xffu); }
#else
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
#endif
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
void main()
@@ -156,7 +171,7 @@ void main()
barrier();
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[CX / 16];
[[unroll]] for (uint kt = 0u; kt < CX / 16; ++kt) coopMatLoad(lb[kt], POOL, LOW_AT + n * 4u * CX + 4u * kt, CX / 4u, CM);
[[unroll]] for (uint kt = 0u; kt < CX / 16; ++kt) LDS8B(lb[kt], POOL, LOW_AT + n * 4u * CX + 4u * kt, CX / 4u);
for (uint mm = 0u; mm < MT; ++mm) {
const uint m = MT * wv + mm;
RAC a;
@@ -166,9 +181,7 @@ void main()
LDA(w0, W_WE, CX, m, kt); LDA(w1, W_WE, CX, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, EX, n * 16u * C + 4u * m, C, CM);
Q8STB(a, EX, n * 16u * C + 4u * m, C)
}
}
barrier();
@@ -216,7 +229,7 @@ void main()
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
floate4m3_t c[4];
F8 c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
const uint j = 4u * wd + b;
precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)];
@@ -265,20 +278,25 @@ void main()
[[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
coopMatLoad(xb0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u, CM);
coopMatLoad(xb1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u, CM);
coopMatLoad(xa0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u, RM);
LDS8B(xb0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u);
LDS8B(xb1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u);
LDS8A(xa0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u);
LDS8A(xa1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u);
LDA(wk0, W_WK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WK0 + HEAD * hd, C, c, kt + 1u);
LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c);
K32(a, wk0, xb0, wk1, xb1)
K32(v, xa0, wv0, xa1, wv1)
}
#ifdef RR_GFX11
Q8STA(a, RR_KS, RR_KS0 + n * 128u + 4u * c, 8u)
Q8STA(v, RR_VS, RR_VS0 + c * 384u + 4u * n, 24u)
#else
A8 ka, va;
Q8(ka, a)
Q8(va, v)
coopMatStore(ka, RR_KS, RR_KS0 + n * 128u + 4u * c, 8u, RM);
coopMatStore(va, RR_VS, RR_VS0 + c * 384u + 4u * n, 24u, RM);
#endif
}
barrier();
#include "rr_attn.glsl"
@@ -291,6 +309,10 @@ void main()
// ---------------------------------------------------------------- MLP
B8 Nb[MT];
#ifdef RR_GFX11
F8V NbC[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m) NbC[m] = uvec2(0u);
#endif
{
float16_t rv[MT][8];
[[unroll]] for (uint m = 0u; m < MT; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]);
@@ -326,10 +348,17 @@ void main()
const uint k = 16u * m + 8u * hl + uint(i);
precise float16_t g = r * w16.w[W_GAIN14 / 2u + k];
precise float16_t n = rv[m][i] * g;
#ifdef RR_GFX11
NbC[m][i >> 2] |= q8h(n) << (8u * uint(i & 3));
#else
Nb[m][i] = q8h(n);
#endif
precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k];
R[m][i] = RV(a);
}
#ifdef RR_GFX11
[[unroll]] for (uint m = 0u; m < MT; ++m) G11_B(Nb[m], NbC[m])
#endif
}
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br;
@@ -342,7 +371,11 @@ void main()
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
K32(h, a0, Nb[kt], a1, Nb[kt + 1u])
}
#ifdef RR_GFX11
{ F8V hc_; ACTQ8C(hc_, h) G11_B(Hb[c], hc_) }
#else
ACTQ8(Hb[c], h)
#endif
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
@@ -365,8 +398,15 @@ void main()
barrier(); // every wave is done with XT (the merge rows reuse it)
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
#ifdef RR_GFX11
F8V tc_;
G11_CODES(tc_, R[m])
G11_B(Tb[m], tc_)
const uint w0 = tc_.x, w1 = tc_.y;
#else
Q8(Tb[m], R[m])
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
#endif
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
#if KIND == 0 && CX
@@ -447,11 +487,20 @@ void main()
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, MG, 8u * ks, C, CM);
coopMatLoad(g1, MG, 8u * ks + 4u, C, CM);
LDS8B(g0, MG, 8u * ks, C);
LDS8B(g1, MG, 8u * ks + 4u, C);
LDA(w0, W_WM, 4u * C, m, 2u * ks); LDA(w1, W_WM, 4u * C, m, 2u * ks + 1u);
K32(a, w0, g0, w1, g1)
}
#ifdef RR_GFX11
F8V q;
G11_CODES(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl;
arena.w[at] = q.x;
arena.w[at + 1u] = q.y;
}
#else
B8 q;
Q8(q, a)
if (X >= 0 && Y >= 0 && X < hx && Y < hy) {
@@ -459,6 +508,7 @@ void main()
arena.w[at] = pack4(q[0], q[1], q[2], q[3]);
arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]);
}
#endif
}
if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards
const uint mx = lane & 3u, my = lane >> 2u;
engine/rr/widen_weights.py
+33
View File
@@ -0,0 +1,33 @@
#!/usr/bin/env python3
"""Append the RDNA3 (RR_GFX11) widened copy of a native stage's weight file and print its f16 element offset (W_WIDE).
usage: widen_weights.py WEIGHTS.bin [OUT.bin] (default: in place)
Element W_WIDE + j of the file's f16 view is byte j of the original image as ZLUDA widens an FP8 operand on gfx11
(rr_gfx11.glsl): the magnitude at bits 7..14, the sign at bit 15, the NaN code with bit 14 set too, i.e. the binary16
of value * 2^-8. Bytes that are not FP8 (tables, binary16 data) are widened as well and never read. W_WIDE is the
original size rounded up to 512 bytes, halved.
"""
import sys
from pathlib import Path
import numpy as np
def widen(image):
b = np.frombuffer(image, np.uint8).astype(np.uint16)
mag = b & 0x7f
mag = mag | ((mag + 1) & 0x80)
return ((mag << 7) | ((b & 0x80) << 8)).astype('<u2').tobytes()
def main():
src = Path(sys.argv[1]); dst = Path(sys.argv[2]) if len(sys.argv) > 2 else src
image = src.read_bytes()
base = (len(image) + 511) // 512 * 512
dst.write_bytes(image + bytes(base - len(image)) + widen(image))
print(base // 2)
if __name__ == '__main__':
main()