wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
49439afb818ee4dc74865cf837ffc7f4edc830ca
d4r/engine/rr/rr2_m.comp
T
itchyfeetleech 49439afb81 engine/rr: preset F in less shared memory (encoder 0 in one shader, enc1, dec0_1) 
ACO sizes a shader's registers to the waves its shared memory leaves room
for, so stages that declare less shared memory also get fewer registers and
more waves. Three stages hold buffers whose lifetimes do not overlap:

- Encoder 0 at D4R_RR_RELAXED=3: with the one-wave block the gather no
  longer needs a pass of its own. The generated shader runs gather, block
  and tail in one 32-thread group and hands the features over in shared
  memory that the block's XT and KVS reuse (rr2_w.comp RR2_XF_SHARED). No
  gather pass, no record in the arena (36 MB of it lay past the arena at
  720p). D4R_RR2_SPLIT=1 keeps the pass.
- enc1: the merge input (GT) is in registers before XT or KVS is written:
  17 -> 9 KiB, 240 -> 144 vector registers, 6 -> 10 waves per SIMD.
- dec0_1: r and the MLP sum for the tail are written after the block is done
  with XT and KVS, as binary16 pairs: 14 -> 8 KiB, 192 -> 120 registers,
  8 -> 12 waves per SIMD.

Also D4R_ENGINE_RR_STAGE_STAMPS=1: a timestamp after every network stage and
output group, printed per stage by d4r-rr.

Every image of six temporal frames byte-identical to the previous model
(1280x720 -> 1440p, 1114x626 -> 1080p, 853x480 -> 1440p, robustness2). RX
9070 XT, RADV, 720p -> 1440p: frame 4.09 -> 4.00 ms (enc0 798 -> 772 us,
enc1 286 -> 266, dec0_1 660 -> 623). Other stages' SPIR-V is unchanged.

Signed-off-by: itchyfeetleech <[email protected]>
2026-10-18 20:34:24 +13:00

622 lines
31 KiB
Plaintext
Raw Blame History

#version 460
// Native stage of the DLSS 4.5 Ray Reconstruction network (preset F, the rr2_* kernels of nvngx_dlssd.dll 310.9):
// one shader for enc1 ... enc5 and both kernels of dec4 ... dec1 and dec0's first. The arithmetic is
// engine/rr/rr2_ref.py's (which documents the block); the shape and the weight offsets come from the include
// generated by engine/rr/rr2_pack.py:
// KIND 0 encoder (tokens from the arena), 1 enc1 (first merges encoder 0's tokens 2x2), 2 a decoder's first kernel
// (expands the CL channels one level down and adds the skip), 3 a decoder's second kernel (with the guide blend),
// 4 encoder 0 after its per-pixel feature gather: included (RR2_EMBED) by the shader engine/rr/gen_rr2_enc0.py
// generates around ptx2rr.py's translation of the gather; it embeds the 64 feature slots of XF and leaves
// the guide head in HD for the translated tail; 5 dec0's second kernel, included likewise by the shader
// generated around the translation of its output tail: it leaves r and the MLP sum in RSH / MSH
// C channels (MT = C / 16 tiles), NH heads of 32 channels, NB MLP branches, CM merged channels for the next
// encoder (0: none), GUIDE 1 when the stage gates its guide image (KIND 0 and 1)
// Like engine/rr/swin_m.comp, every product runs transposed (weights x activations), so results stay in registers as
// [channel tile][token] accumulators: one workgroup per 8x8-token window, four waves, wave w owning the token slots
// 16w ... 16w + 15; token and channel slot j holds the natural index inter(j) (NVIDIA's k order, which is also the
// arena's byte order within a 32-channel plane).
#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
#ifndef RR2_EMBED
layout(local_size_x = 128) in;
#endif
#include "rr2_stage.h"
#include "rr_exact.glsl"
#define MT (C / 16)
#ifdef RR2_EMBED
#define PRM(i) param.w[i]
// the generated shader's workgroup holds RR2_GROUP (1 to 4) windows along x (wave w runs the translated code of
// window w, then the block runs once per window): RR2_WIN is the one this call works on
#if RR2_GROUP == 1
#define RR2_WIN 0u
#define RR2_WGX gl_WorkGroupID.x
#else
uint rr2_win;
#define RR2_WIN rr2_win
#define RR2_WGX (uint(RR2_GROUP) * gl_WorkGroupID.x + rr2_win)
#endif
#else
#define RR2_GROUP 1
#define RR2_WIN 0u
#define RR2_WGX gl_WorkGroupID.x
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#define PRM(i) param.p[i]
#endif
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;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CMJ 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 of binary16 values: the reference adds it inside the f32 products and rounds
// RR_FUSE (D4R_RR_FUSE=1): out = in + branch * s and the guide's blend as one fused multiply-add each, which is what
// NVIDIA's compiler makes of those multiply / add pairs on an NVIDIA card; without it the product is rounded first,
// as ZLUDA (the captures) does. The activation table is built to match (enc1_ref.act).
#ifdef RR_FUSE
#define MADD(d, a, b, c) precise float16_t d = fma((a), (b), (c));
#define MADD4(d, a, b, c) precise f16vec4 d = fma((a), (b), (c));
#else
#define MADD(d, a, b, c) precise float16_t d##_p = (a) * (b); precise float16_t d = (c) + d##_p;
#define MADD4(d, a, b, c) precise f16vec4 d##_p = (a) * (b); precise f16vec4 d = d##_p + (c);
#endif
// RR_NATIVE (D4R_RR_RELAXED=3): the residual steps and the MLP activation in f32, straight into FP8 - no binary16 in
// between and no activation table. Not the reference's roundings (about 40 dB from it, like RR_RELAXED).
// (The biases stay binary16: f32 copies were tried for RR_NATIVE and were slower - half the conversions, but one more
// load per accumulator, and a load costs about three arithmetic instructions here.)
#define BV(X, i) RV(w16.w[W_##X / 2u + (i)])
#define BVS(X, stride, n, i) RV(w16.w[(W_##X + (stride) * (n)) / 2u + (i)])
#ifdef RR_NATIVE
#define RVT float
#define E4V(w, b) float(uintBitsToFloate4m3EXT(uint8_t(((w) >> (8u * (b))) & 0xffu)))
#define RESID(d, a, s, c) const float d = fma(float(a), float(s), float(c));
#define Q8R(x) q8s(x)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2(float((h)[i_]), float((h)[i_ + 1])); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
#else
#define RVT float16_t
#define E4V(w, b) e4((w), (b))
#define RESID(d, a, s, c) MADD(d, float16_t(a), s, c)
#define Q8R(x) q8h(x)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
#endif
// once (swin_m.comp explains RR_MIX, the one-rounding variant that is not bit-exact by construction)
#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_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
#define BRANCHB BRANCH
#if KIND == 1
// enc1's merge input is in registers (gb) before XT or KVS is first written, so the three share 8 KiB instead of 16
// (above about 4 KiB per wave this GPU runs the workgroup slower, see KIND 2 below): XT words 0 - 16C-1, KVS after it
#if 16 * C + 2 * 512 > 64 * 32
#error rr2_m.comp: XT and KVS no longer fit in GT
#endif
shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels]; then XT and KVS
#define XT GT
#define KVS GT
#define KVA(o) (16u * C + (o))
#elif KIND == 5 && RR2_GROUP == 1
// dec0's second kernel writes r and the MLP sum for its tail only after the block is done with XT and KVS, so the
// three share 8 KiB instead of 14: XT words 0 - 16C-1, KVS after it; then r (RSW) at word 0 and the MLP sum (MSW) at
// 1024, both as binary16 pairs of natural channels 2k, 2k + 1 per token (the tail reads them in those pairs)
#if 16 * C + 2 * 512 > 2 * 1024
#error rr2_m.comp: XT and KVS no longer fit under RSW and MSW
#endif
shared uint SM[2 * 1024];
#define XT SM
#define KVS SM
#define KVA(o) (16u * C + (o))
#define RSW(w) SM[w]
#define MSW(w) SM[1024u + (w)]
#else
shared uint XT[16 * C]; // the window's tokens: [token slot][C FP8 channels in slot order]; later the merge rows
shared uint KVS[2 * 512]; // the current head's QK ([token slot][32]), then V^T ([32][token slot])
#define KVA(o) (o)
#endif
#if KIND == 2
// Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave (16 KiB for this
// workgroup) the rate falls in proportion (measured: build/rr/f/nv/cost/occ2.py). So the 4 x 4 tokens one level down
// ([row-major][CL FP8 channels]) borrow KVS, which the attention only writes later, and the block's input before its
// FP8 conversion stays in each lane's registers (xr), not in a [token slot][channel slot] array.
#if 4 * CL > 2 * 512
#error rr2_m.comp: LOW no longer fits in KVS
#endif
#define LOW KVS
shared uint EX[16 * C]; // their expansion: [low token][sub-position][C FP8 channels in slot order]
#endif
#if KIND == 4
#ifdef RR2_XF_ARENA
// the features come from the gather pass (engine/rr/gen_rr2_enc0.py --split): the window's record in the arena, as
// the index of its first 32-bit word, set by the stage's main
uint rr2_xf;
#define XF_SRC arena.w
#define XF_AT rr2_xf
#else
shared uint XF[RR2_GROUP * 64 * 16]; // per window, the features from the gather: [token slot][64 FP8 feature slots]
#define XF_SRC XF
#define XF_AT (RR2_WIN * 1024u)
#endif
shared float16_t HD[RR2_GROUP * 64 * 4]; // per window, the guide head per token (raster order), for the translated tail
const uint EBRANCH = W_EWD1 - W_EWD0;
#define EBRANCHB EBRANCH
#endif
#if KIND == 5
#if RR2_GROUP == 1
// r and the MLP sum of token row, natural channel col, for the translated tail (gen_rr2_enc0.py)
float16_t rr2_rsh(uint row, uint col) { return unpackFloat2x16(RSW(row * 16u + (col >> 1u)))[col & 1u]; }
float16_t rr2_msh(uint row, uint col) { return unpackFloat2x16(MSW(row * 16u + (col >> 1u)))[col & 1u]; }
#else
shared float16_t RSH[RR2_GROUP * 64 * 32], MSH[RR2_GROUP * 64 * 32]; // per window, r and the MLP sum per token (raster order) and natural channel
float16_t rr2_rsh(uint row, uint col) { return RSH[gl_SubgroupID * 2048u + row * 32u + col]; }
float16_t rr2_msh(uint row, uint col) { return MSH[gl_SubgroupID * 2048u + row * 32u + col]; }
#endif
#endif
#if GUIDE || KIND == 3
shared float16_t OUTG[64 * 4]; // the guide output per window token (raster order)
#endif
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)); }
// four binary16 values of a guide image at token (x, y) of a level w wide
f16vec4 guideAt(uint off, int x, int y, int w)
{
const uint a = (off + uint(y * w + x) * 8u) >> 2u;
const uint u0 = arena.w[a], u1 = arena.w[a + 1u];
return f16vec4(hb(u0 & 0xffffu), hb(u0 >> 16u), hb(u1 & 0xffffu), hb(u1 >> 16u));
}
#ifdef RR2_EMBED
void rr2_block()
#else
void main()
#endif
{
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(PRM(0)), ey = int(PRM(1)), hx = ex / 2, hy = ey / 2;
const uint inOff = PRM(2), skOff = PRM(4), outOff = PRM(8), mrgOff = PRM(10), shift = PRM(20);
const uint gInOff = PRM(12), gEncOff = PRM(14), gOutOff = PRM(16), gHalfOff = PRM(18);
const float16_t s0 = hb(PRM(23) & 0xffffu), s1 = hb(PRM(23) >> 16u), s2 = hb(PRM(24) & 0xffffu), s3 = hb(PRM(24) >> 16u);
const int x0 = 8 * int(RR2_WGX) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
// this lane's own token (slot 16 wv + col)
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;
// ---------------------------------------------------------------- the block's input x: one token per thread
#if KIND == 0 || KIND == 3 || KIND == 5
if (tid < 64u) {
const uint k = inter(tid);
const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey);
const uint tb = (inOff + uint(yy * ex + xx) * 32u) / 4u;
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) XT[tid * (C / 4u) + wd] = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
}
barrier();
#elif KIND == 4
// encoder 0: x = bu + sum over two branches of Wu . act(Wd . features + bd)
{
B8 fb[4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[kt], XF_SRC, XF_AT + wv * 256u + 4u * kt, 16u, CMJ);
RAC X[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) X[m][i] = BV(EBU, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint br = 0u; br < 2u; ++br) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = BVS(EBD0, EBRANCHB, br, 16u * c + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) {
A8 a0, a1;
LDA(a0, W_EWD0 + EBRANCH * br, 64u, c, kt); LDA(a1, W_EWD0 + EBRANCH * br, 64u, c, kt + 1u);
K32(h, a0, fb[kt], a1, fb[kt + 1u])
}
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, W_EWU0 + EBRANCH * br, 32u, m, 0u); LDA(a1, W_EWU0 + EBRANCH * br, 32u, m, 1u);
K32(X[m], a0, Hb[0], a1, Hb[1])
}
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
B8 q;
Q8(q, X[m])
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
barrier();
#elif KIND == 1
// enc1: the 2x2 group of encoder 0's 32-channel tokens, then x = W1 . group + b1
if (tid < 64u) {
const uint k = inter(tid);
const int xx = mirror(x0 + int(k & 7u), ex), yy = mirror(y0 + int(k >> 3u), ey);
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
const uint tb = (inOff + uint((2 * yy + int(q >> 1u)) * (2 * ex) + 2 * xx + int(q & 1u)) * 32u) / 4u;
[[unroll]] for (uint wd = 0u; wd < 8u; ++wd) GT[tid * 32u + q * 8u + wd] = arena.w[tb + wd];
}
}
barrier();
{
B8 gb[8];
[[unroll]] for (uint kt = 0u; kt < 8u; ++kt) coopMatLoad(gb[kt], GT, wv * 512u + 4u * kt, 32u, CMJ);
barrier(); // every wave has its part of GT in registers: XT takes its place
[[unroll]] for (uint m = 0u; m < MT; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(B1, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < 8u; kt += 2u) {
A8 w0, w1;
LDA(w0, W_W1, 128u, m, kt); LDA(w1, W_W1, 128u, m, kt + 1u);
K32(a, w0, gb[kt], w1, gb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
barrier();
#else
// a decoder's first kernel: wave w expands the 16 low tokens to sub-position w, then x = dq(e) + dq(skip) * s3
float16_t xr[MT][8]; // x of this lane's token (slot 16 wv + col), channel slots 16 m + 8 hl + i
{
const int X0 = x0 >> 1, Y0 = y0 >> 1;
for (uint i = tid; i < 4u * CL; i += 128u) {
const uint r = i / (CL / 4u), wd = i % (CL / 4u);
const int X = mirror(X0 + int(r & 3u), hx), Y = mirror(Y0 + int(r >> 2u), hy);
LOW[i] = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
barrier();
B8 lb[CL / 16];
[[unroll]] for (uint kt = 0u; kt < CL / 16; ++kt) coopMatLoad(lb[kt], LOW, 4u * kt, CL / 4u, CMJ);
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] = BV(BE, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, EX, 4u * m, C, CMJ);
}
barrier();
if (tid < 64u) {
const uint k = inter(tid);
const int kx = x0 + int(k & 7u), ky = y0 + int(k >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) {
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = EX[r * C + s * (C / 4u) + wd];
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
MADD(x, e4(sk, b), s3, e4(e, b))
c[b] = q8h(x);
}
XT[tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
{ // the same x for the residual step: the bytes of the words 4 m + 2 hl and 4 m + 2 hl + 1 of this lane's token
const int kx = x0 + int(t & 7u), ky = y0 + int(t >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (uint h2 = 0u; h2 < 2u; ++h2) {
const uint wd = 4u * m + 2u * hl + h2;
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = EX[r * C + s * (C / 4u) + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
MADD(x, e4(sk, b), s3, e4(e, b))
xr[m][4u * h2 + b] = x;
}
}
}
barrier();
}
#endif
// ---------------------------------------------------------------- attention: the heads' out-projections sum up
RAC AT[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) AT[m][i] = BV(BO, 16u * m + 8u * hl + uint(i));
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's QK / V
for (uint j = 0u; j < 2u; ++j) { // 8 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 < MT; kt += 2u) {
B8 xb0, xb1, wv0, wv1;
A8 xa0, xa1, wk0, wk1;
coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ);
coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ);
coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM);
LDA(wk0, W_WQK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WQK0 + 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, KVS, KVA(n * 128u + 4u * c), 8u, RM);
coopMatStore(va, KVS, KVA(512u + c * 256u + 4u * n), 16u, RM);
}
barrier();
// scores^T (keys x this wave's queries) = QK . QK^T + bias, the cubic softmax down each column
B8 Qb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], KVS, KVA(wv * 128u + 4u * c), 8u, CMJ);
B8 Pb[4];
{
float16_t e[4][8];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
A8 k0, k1;
coopMatLoad(k0, KVS, KVA(kt * 128u), 8u, RM);
coopMatLoad(k1, KVS, KVA(kt * 128u + 4u), 8u, RM);
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
K32H(c, k0, Qb[0], k1, Qb[1])
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
}
// key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1),
// tt = 2hl + (i >> 2), half = i & 1 (rr2_ref.softmax): pairs (e[2q], e[2q + 1]) are the two halves
f16vec2 ep[4][4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = (ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1]);
precise f16vec2 od = (ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
// hl = 0 holds tt 0, 1 and hl = 1 tt 2, 3: the sum is (tt0 + tt1) + (tt2 + tt3) on both halves
precise f16vec2 o = hl == 0u ? loc + unpackFloat2x16(other) : unpackFloat2x16(other) + loc;
precise float16_t rs = o.x + o.y;
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
}
}
// O^T = V^T P^T, then the out-projection
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC o = RAC(RV(float16_t(0.0)));
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) {
A8 v0, v1;
coopMatLoad(v0, KVS, KVA(512u + c * 256u + 8u * ks), 16u, RM);
coopMatLoad(v1, KVS, KVA(512u + c * 256u + 8u * ks + 4u), 16u, RM);
K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u])
}
Q8(Ob[c], o)
}
[[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(AT[m], a0, Ob[0], a1, Ob[1])
}
}
// ---------------------------------------------------------------- r = x + attention * s1; the MLP's input q8(r)
RVT rv[MT][8];
B8 Nb[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m) {
#if KIND != 2
const uint xa = (16u * wv + col) * (C / 4u) + 4u * m + 2u * hl;
const uvec2 xw = uvec2(XT[xa], XT[xa + 1u]);
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 2
const float16_t xb = xr[m][i];
#else
const RVT xb = E4V(xw[i >> 2], uint(i) & 3u);
#endif
RESID(r, AT[m][i], s1, xb)
rv[m][i] = r;
Nb[m][i] = Q8R(r);
}
}
RAC ML[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) ML[m][i] = BV(BU, 16u * m + 8u * hl + uint(i));
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + 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] = BVS(BD0, BRANCHB, br, 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])
}
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u);
K32(ML[m], a0, Hb[0], a1, Hb[1])
}
}
#if KIND == 5
// dec0's second kernel: the translated tail forms y = r + MLP * s2 itself and computes the outputs from it
#if RR2_GROUP == 1
barrier(); // every wave is done with XT and KVS
// (channel slots j and j + 1, j even, hold the natural channels inter(j) and inter(j) + 1)
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; i += 2) {
const uint w = (t * 32u + inter(16u * m + 8u * hl + uint(i))) >> 1u;
RSW(w) = packFloat2x16(f16vec2(float16_t(rv[m][i]), float16_t(rv[m][i + 1])));
MSW(w) = packFloat2x16(f16vec2(float16_t(ML[m][i]), float16_t(ML[m][i + 1])));
}
#else
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint at = RR2_WIN * 2048u + t * 32u + inter(16u * m + 8u * hl + uint(i));
RSH[at] = float16_t(rv[m][i]);
MSH[at] = float16_t(ML[m][i]);
}
#endif
#else
// ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens
B8 Tb[MT];
{
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if CM
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) {
[[unroll]] for (int i = 0; i < 8; ++i) {
RESID(y, ML[m][i], s2, rv[m][i])
Tb[m][i] = Q8R(y);
}
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 CM
const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl;
XT[g] = w0; XT[g + 1u] = w1;
#endif
}
}
#endif
#if GUIDE || KIND == 3 || KIND == 4
{ // the guide head (rows 0-3 on lanes 0-15)
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = BV(BH, 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 == 4
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 4u; ++c) HD[RR2_WIN * 256u + t * 4u + c] = float16_t(hv[c]);
#else
if (hl == 0u) {
f16vec4 o4 = f16vec4(0.0);
if (inside) {
precise f16vec3 gate = f16vec3(sig(float16_t(hv[0])), sig(float16_t(hv[1])), sig(float16_t(hv[2]))) * s0;
#if KIND == 3
// the guide one level down, bilinear (half-pixel centres, edge clamped), every product and sum rounded
const int bx = tx == 0 ? 0 : (tx - 1) >> 1, by = ty == 0 ? 0 : (ty - 1) >> 1;
const float16_t wx = tx == 0 ? float16_t(0.0) : ((tx & 1) != 0 ? float16_t(0.25) : float16_t(0.75));
const float16_t wy = ty == 0 ? float16_t(0.0) : ((ty & 1) != 0 ? float16_t(0.25) : float16_t(0.75));
const int x1 = min(bx + 1, hx - 1), y1 = min(by + 1, hy - 1);
precise float16_t ux = float16_t(1.0) - wx, uy = float16_t(1.0) - wy;
precise f16vec4 p00 = guideAt(gInOff, bx, by, hx) * ux, p10 = guideAt(gInOff, x1, by, hx) * wx;
precise f16vec4 p01 = guideAt(gInOff, bx, y1, hx) * ux, p11 = guideAt(gInOff, x1, y1, hx) * wx;
precise f16vec4 a10 = p10 * uy;
MADD4(u0, p00, f16vec4(uy), a10)
MADD4(u1, p01, f16vec4(wy), u0)
MADD4(up, p11, f16vec4(wy), u1)
const float16_t al = sig(float16_t(hv[3]));
precise f16vec4 gt = up * f16vec4(gate, float16_t(0.0));
precise float16_t ia = float16_t(1.0) - al;
precise f16vec4 pg = gt * al;
MADD4(mx, guideAt(gEncOff, tx, ty, ex), f16vec4(ia), pg)
o4 = min(mx, f16vec4(65504.0));
#else
precise f16vec4 gg = guideAt(gInOff, tx, ty, ex) * f16vec4(gate, float16_t(0.0));
o4 = min(gg, f16vec4(65504.0));
#endif
const uint a = (gOutOff + uint(ty * ex + tx) * 8u) >> 2u;
arena.w[a] = bh(o4.x) | (bh(o4.y) << 16u);
arena.w[a + 1u] = bh(o4.z) | (bh(o4.w) << 16u);
}
[[unroll]] for (uint c = 0u; c < 4u; ++c) OUTG[t * 4u + c] = o4[c];
}
#endif
}
#endif
#if CM
barrier();
{ // the 2x2 merge for the next encoder: merged^T (CM 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 < CM / 16u; m += 4u) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(BM, 16u * m + 8u * hl + uint(i));
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, XT, 8u * ks, C, CMJ);
coopMatLoad(g1, XT, 8u * ks + 4u, C, CMJ);
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 GUIDE
if (wv == 3u && lane < 16u) { // the guide for the next encoder: 2x2 mean, 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;
f16vec4 q4[4];
[[unroll]] for (uint d = 0u; d < 4u; ++d) {
const uint tt = t00 + (d & 1u) + 8u * (d >> 1u);
precise f16vec4 qv = f16vec4(OUTG[tt * 4u], OUTG[tt * 4u + 1u], OUTG[tt * 4u + 2u], OUTG[tt * 4u + 3u]) * float16_t(0.25);
q4[d] = qv;
}
precise f16vec4 h0 = q4[0] + q4[1], h1 = q4[2] + q4[3];
precise f16vec4 hm = h0 + h1;
const uint a = (gHalfOff + uint((y / 2) * hx + x / 2) * 8u) >> 2u;
arena.w[a] = bh(hm.x) | (bh(hm.y) << 16u);
arena.w[a + 1u] = bh(hm.z) | (bh(hm.w) << 16u);
}
}
#endif
}
#endif
}
Reference in new issue View Git Blame Copy Permalink