#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 #extension GL_EXT_float_e4m3 : require #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; layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; 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 #define HAC coopmat #define A8 coopmat #define B8 coopmat #define RM gl_CooperativeMatrixLayoutRowMajor #define CM gl_CooperativeMatrixLayoutColumnMajor #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 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); } 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))); } 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) coopMatLoad(lb[kt], POOL, LOW_AT + n * 4u * CX + 4u * kt, CX / 4u, CM); 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]) } B8 q; Q8(q, a) coopMatStore(q, EX, n * 16u * C + 4u * m, C, CM); } } 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) { floate4m3_t 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; 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); 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) } 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); } 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]; { 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; Nb[m][i] = q8h(n); precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k]; R[m][i] = RV(a); } } [[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]) } [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i])); } [[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) { 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]); 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; coopMatLoad(g0, MG, 8u * ks, C, CM); coopMatLoad(g1, MG, 8u * ks + 4u, C, CM); 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) } 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]); } } 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 }