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

feat: add d4r inference engines and Windows parity tooling

Browse Source 
Add native K/M/RR inference engines, HIP kernels, replay and validation
harnesses, and Windows DLSS setup and runtime support. Include GPU
profiling tools, build and packaging scripts, platform integration
patches, and implementation documentation.
This commit is contained in:
ayo committed 2026-10-16 08:54:47 -04:00
1 parent 25493b3198
commit 7c5d3c9157
430 files changed
+76191
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/rr/swin_m.comp
+537
View File
@@ -0,0 +1,537 @@
#version 460
// Native Swin stage of the DLSS-D (preset E) denoiser, register-resident, for any stage of engine/rr/swin_ref.STAGES
// (enc2-enc5, dec2-dec4; enc1_m.comp / dec1_m.comp are the hand-written C = 64 instances it generalises). The stage's
// shape and weight offsets come from the include generated by engine/rr/swin_pack.py:
// KIND 0 encoder (tokens from the arena, 2x2 merge to CX channels, enc1's gated filter; CX = 0: the bottleneck,
// ungated filter, no merge), 1 decoder (expand of the CX channels one level down + skip, dec0's filter)
// C channels (KC = C / 32 planes, MT = C / 16 tiles), NH heads of 32 channels, C / 8 MLP branches
// The arithmetic is engine/rr/swin_ref.py's; enc1_m.comp's header explains the transposed products and slot orders.
// One workgroup per 8x8-token window, four waves, wave w owning query slots 16w..16w+15.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#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
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#include "swin_stage.h"
#include "rr_exact.glsl"
#define KC (C / 32)
#define MT (C / 16)
#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;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
layout(set = 0, binding = 3) uniform sampler2D T88;
#if KIND == 1
layout(set = 0, binding = 4) uniform sampler2D T96;
layout(set = 0, binding = 5) uniform sampler2D T104;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112;
#else
layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112;
#if CX
layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128;
#endif
#endif
#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_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#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)
#if KC == 1
#define TREE(v) (v[0])
#elif KC == 2
#define TREE(v) (v[0] + v[1])
#elif KC == 3
#define TREE(v) ((v[0] + v[1]) + v[2])
#elif KC == 4
#define TREE(v) ((v[0] + v[1]) + (v[2] + v[3]))
#else
#define TREE(v) (((v[0] + v[1]) + (v[2] + v[3])) + v[4])
#endif
#if KIND == 1
// one pool, by phase: [EX | LOW] for the expand; the prologue reads EX and writes XT over LOW (dead) and beyond;
// the attention's K and V^T reuse EX (dead after the residual)
#define EX_W (48u * C)
#define POOL_W (EX_W + (12u * CX > 24u * C ? 12u * CX : 24u * C))
shared uint POOL[POOL_W];
#define EX POOL // the expansion: 48 rows x 4 sub-positions x C FP8 channels in slot order
#define LOW_AT EX_W // the low-resolution tokens (6 x 6, padded): CX FP8 channels, arena order
#define XT_AT EX_W // the normalised key tokens: slot-major, C FP8 channels in slot order
#define XTS POOL
#define RR_KS POOL
#define RR_KS0 0u
#define RR_VS POOL
#define RR_VS0 768u
shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions
shared float16_t OUTF[64 * 16]; // the head per window token (raster order)
#else
shared uint XT[96 * C / 4]; // the normalised key tokens: slot-major, C FP8 channels in slot order
#define XT_AT 0u
#define XTS XT
shared uint KVS[2 * 768]; // the current head's K, then V^T
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT // the merge rows (16 x 4C FP8) reuse XT, dead after the projections
#if CX
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
#endif
#endif
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
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));
}
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()
{
const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2;
const uint inOff = param.p[2], skOff = param.p[6], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20];
const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
#if KIND == 1
// ---------------------------------------------------------------- expand: wave w computes sub-position w
const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1;
for (uint i = tid; i < 48u * (CX / 4u); i += 128u) {
const uint r = i / (CX / 4u), wd = i % (CX / 4u);
uint u = 0u;
if (r < 36u) {
const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy);
u = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
POOL[LOW_AT + i] = u;
}
barrier();
[[unroll]] for (uint n = 0u; n < 3u; ++n) {
B8 lb[CX / 16];
[[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;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < CX / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CX, m, kt); LDA(w1, W_WE, CX, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
Q8STB(a, EX, n * 16u * C + 4u * m, C)
}
}
barrier();
#endif
// ---------------------------------------------------------------- prologue: one key token per thread
if (tid < 96u) {
const ivec2 k = keyXY(inter(tid), x0, y0);
const int xx = mirror(k.x, ex), yy = mirror(k.y, ey);
const uint tb = ((KIND == 1 ? skOff : inOff) + uint(yy * ex + xx) * 32u) / 4u;
#if KIND == 1
const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1));
#endif
float16_t v[C];
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
const uint t = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
#if KIND == 1
const uint e = EX[r * C + s * (C / 4u) + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(t, b); v[4u * wd + b] = x; }
#else
[[unroll]] for (uint b = 0u; b < 4u; ++b) v[4u * wd + b] = e4(t, b);
#endif
}
// swin_ref.input_norm: per lane word q (and q + 4) the pair sums of the planes (TREE), the halves, then
// q ^ 1, q ^ 2, the components
precise f16vec2 h[4];
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 hh[2];
[[unroll]] for (uint half_ = 0u; half_ < 2u; ++half_) {
precise f16vec2 pr[KC];
[[unroll]] for (uint p = 0u; p < KC; ++p) {
const uint o = 32u * p + 4u * (q + 4u * half_);
precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]);
precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]);
pr[p] = a + c;
}
precise f16vec2 tr = TREE(pr);
hh[half_] = tr;
}
h[q] = hh[0] + hh[1];
}
precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3];
precise f16vec2 hs = h0 + h2_;
precise float16_t sm = hs.x + hs.y;
precise float16_t se = sm + float16_t(0.0001220703125);
const float16_t rr = rsq16(se);
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
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)];
precise float16_t n = v[j] * g;
c[b] = q8h(n);
}
XTS[XT_AT + tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
barrier();
// ---------------------------------------------------------------- the residual: the window token + table 7
RAC R[MT];
{
const uint tq = inter(16u * wv + col);
const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u);
const uint tb = ((KIND == 1 ? skOff : inOff) + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u;
#if KIND == 1
const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1));
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
const uint ta = tb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl;
const uvec2 t = uvec2(arena.w[ta], arena.w[ta + 1u]);
#if KIND == 1
const uint ea = r * C + s * (C / 4u) + 4u * m + 2u * hl;
const uvec2 e = uvec2(EX[ea], EX[ea + 1u]);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 1
precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(t[i >> 2], uint(i) & 3u);
#else
const float16_t x = e4(t[i >> 2], uint(i) & 3u);
#endif
precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)];
R[m][i] = RV(c);
}
}
}
// ---------------------------------------------------------------- attention, NH heads
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's K / V (and, for the first, the expand's LOW)
for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c)
const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u;
RAC a = RAC(RV(float16_t(0.0))), v = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
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"
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u);
K32(R[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- 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]);
// swin_ref.rowsum_sq: slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) +
// 2((i >> 1) & 1), t = 2hl + (i >> 2), c = i & 1: per parity (m & 1) the tile pairs (i, i + 2) of every
// 32-channel block, TREE over the blocks; the parities; t ^ 1 (local), t ^ 2 (the other lane half)
float16_t h1[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
float16_t ht[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint i0 = 4u * tq + c, i1 = i0 + 2u;
precise float16_t par[2];
[[unroll]] for (uint pa = 0u; pa < 2u; ++pa) {
precise float16_t pr[KC];
[[unroll]] for (uint b = 0u; b < KC; ++b)
pr[b] = rv[2u * b + pa][i0] * rv[2u * b + pa][i0] + rv[2u * b + pa][i1] * rv[2u * b + pa][i1];
precise float16_t tr = TREE(pr);
par[pa] = tr;
}
precise float16_t hv = par[0] + par[1];
ht[tq] = hv;
}
precise float16_t hh = ht[0] + ht[1];
h1[c] = hh;
}
const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u);
precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u);
precise float16_t s = g0 + g1;
precise float16_t se = s + float16_t(0.0001220703125);
const float16_t r = rsq16(se);
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
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;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
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;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(R[m], a0, Hb[0], a1, Hb[1])
}
}
// ---------------------------------------------------------------- outputs
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
B8 Tb[MT];
{
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if KIND == 0 && CX
const uint ty2 = t >> 3u, tx2 = t & 7u;
const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u;
const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u);
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
const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl;
MG[g] = w0; MG[g + 1u] = w1;
#endif
}
}
// the C -> 16 head
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
K32(hv, a0, Tb[kt], a1, Tb[kt + 1u])
}
#if KIND == 0
{ // enc1's kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15)
float16_t l[12];
[[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]);
const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u);
l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u);
const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]);
const float16_t mx = max(max(ma, mb), l[8]);
float e[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); }
precise float s = (e[0] + e[1]) + 0.0;
s = s + (e[2] + e[3]);
s = s + (e[4] + e[5]);
s = s + (e[6] + e[7]);
s = s + e[8];
const float16_t rr = rcp16(s);
float wg[9];
[[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr));
const int x = tx, y = ty;
const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1;
const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1;
const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc),
ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp));
float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0));
[[unroll]] for (uint k = 0u; k < 8u; ++k) {
const vec4 tv = texelFetch(T88, pts[k], 0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); }
}
const vec4 t8 = texelFetch(T88, pts[8], 0);
float o112[3], o120[3];
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float f0 = acc[c][0] + acc[c][1];
precise float f1 = t8[c] * wg[8] + r16(f0);
const float f = r16(f1);
#if CX
precise float a = float(sig(l[9u + c])) * f;
const float ar = r16(a);
precise float b = f - ar;
o120[c] = ar; o112[c] = r16(b);
#else
o112[c] = f; // the bottleneck: the filtered value itself
#endif
}
if (hl == 0u && inside) {
imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0));
#if CX
imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0));
#endif
}
#if CX
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]);
#endif
}
#if CX
barrier();
{ // the 2x2 merge: merged^T (CX x 16 rows), channel tiles m = wv, wv + 4, ...
const uvec2 mxy = mergeXY(col);
const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y);
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CX / 16u; m += 4u) {
RAC a;
[[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;
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) {
const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl;
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;
const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my);
if (x >= 0 && y >= 0 && x < ex && y < ey) {
const uint t00 = (2u * my) * 8u + 2u * mx;
vec4 out_ = vec4(0.0);
[[unroll]] for (uint c = 0u; c < 3u; ++c) {
precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25);
precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25);
precise float16_t b0 = a00 + a01, b1 = a10 + a11;
precise float16_t cc = b0 + b1;
out_[c] = float(cc);
}
imageStore(S128, ivec2(x >> 1, y >> 1), out_);
}
}
}
#endif
#else
[[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]);
barrier();
#define DEC_FILTER
#include "rr_decfilter.glsl"
#endif
}