wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
3f66081add6ca4bdf4880fcdac5b92458164128c
d4r/engine/rr/rr2_m.comp
T
woutermans f0b55a850a engine: this GPU's own arithmetic and shared weight tiles as the builders' defaults 
- layer_m.comp builds relaxed by default (M_RELAXED: f32 accumulators, paired FP8
  conversions, no activation table) with two windows per wave sharing weight tiles
  (WPW=2); M_EXACT=1 keeps NVIDIA's roundings, bit-identical to the CUDA path
- layer.comp reuses weight tiles across its waves (TPW=2 in the 32-channel layers)
- engine/rr/rr2_w.comp: one wave per window for preset F's dec0_0 and enc0, used by
  compile_rr2.py at the new default D4R_RR_RELAXED=3 (build_native_model.sh too);
  D4R_RR_FUSE=1 matches the reference's fused multiply-adds, ptx2rr.py --snorm16
- [Engine] Network defaults to vulkan; the HIP network is opt-in
2026-10-14 00:52:41 +02:00

579 lines
29 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
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])
#if KIND == 1
shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels]
#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
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
#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);
[[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, n * 128u + 4u * c, 8u, RM);
coopMatStore(va, KVS, 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, 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, kt * 128u, 8u, RM);
coopMatLoad(k1, KVS, 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, 512u + c * 256u + 8u * ks, 16u, RM);
coopMatLoad(v1, KVS, 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
[[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]);
}
#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