Engine: native presets K and M in the D3D12 shim, wider game coverage
A standalone Vulkan inference engine (engine/) for DLSS presets K and M that runs from weights extracted from the user's own nvngx_dlss.dll, with d4r's own input, network and reconstruction shaders or the native HIP network layers, recorded into the game's command list by the shim. Opt-in ([Engine] Enabled / D4R_ENGINE=1); anything uncovered falls back to the CUDA/ZLUDA path with the reason logged. Byte-identical to NGX's CUDA backend under the phase-capable reference (1706x960 -> 2560x1440, native HIP network, eleven frames, RX 9070 XT): - K render- and display-resolution motion vectors, HDR and LDR - K game-supplied exposure and automatic exposure with any pre-exposure and exposure scale (one residual: effective exposure exactly 2 or 4) - K and M with the DLSS indicator's inverted-axis parameters - M with any exposure settings (NGX ignores them for M) K packages need the current compile_k.py (marker D4RO0003). Also: HIP shared buffers are pooled across feature recreation (ROCm leaks destroyed imports), hipnet.bin must match the preset, M-only games do not load the K model, install_d4r_runtime.sh stages k, k-ldr, m, m-ldr with hip/, the activation table has a generator (gen_activation_m.py), and package_release.sh can bundle the engine's shaders (never its weights). Signed-off-by: woutermans <[email protected]>
This commit is contained in:
woutermans committed
1 parent
16cfc1f063
commit
e0d1549a21
82 files changed
+13335
-96
No files matched your search
@@ -0,0 +1,446 @@
|
||||
#version 460
|
||||
// d4r engine, preset M (rrlite): one Swin block over 4x4-token windows, register-resident.
|
||||
//
|
||||
// The cooperative matrix is opaque in GLSL, but RADV's gfx1201 lowering has a fixed lane/slot <-> element
|
||||
// mapping, the same for every 16x16x16 use. It was proven on this GPU with the probes in
|
||||
// build/engine-quality/rm-probe (each lane writes its own slot values into an operand and the resulting image is
|
||||
// read back):
|
||||
//
|
||||
// RowMajor sources and destinations are indexed [row][col]
|
||||
// ColumnMajor sources and destinations are indexed [col][row]
|
||||
// accumulator and B operand: element i of lane l is (row = 8*(l>>4) + i, col = l & 15)
|
||||
// A operand: element i of lane l is (row = l & 15, col = 8*(l>>4) + i)
|
||||
//
|
||||
// Every product here is therefore weights x activations (A = weights, B = activations): a product's rows are the
|
||||
// output channels and its columns the tokens, and the per-lane elements of a result ARE the per-lane elements of
|
||||
// the next product's B operand. Cooperative-matrix subscripts read and write those elements in place, so the
|
||||
// attention and the MLP chain never leave the registers. Only the patch expand (decoders) and the patch merge
|
||||
// (encoders) move data between waves, and only those use shared memory and a workgroup barrier.
|
||||
//
|
||||
// NVIDIA's arithmetic (the semantics of kernels/m/swin_common.h K32_MIX, which the captured launches agree with)
|
||||
// rounds the accumulator to binary16 after every logical 32 products; the 16x16x16 cells have no K32 form, so a
|
||||
// logical k32 is two K16 steps summed in f32 and rounded once. RADV folds a register-only f32 -> f16 -> f32
|
||||
// conversion away, so the rounding of P + the previous accumulator passes through a per-invocation binary16
|
||||
// scratch slot (no cross-lane traffic: a subgroup memory barrier is enough). The tile loops are written out one
|
||||
// tile at a time because a cooperative matrix indexed by a loop variable is not a register any more.
|
||||
//
|
||||
// The two norms do NOT use the natural element order of this shader's accumulators: the native register element
|
||||
// (tile T, half hf, element I) holds the natural channel tok_xnat(T, 8hf+I) = 16T + 8*((I>>2)&1) + 4hf + (I&3),
|
||||
// and tok_rsq's summation tree groups the lanes' squares in that order. Norm 1 therefore reads the input at that
|
||||
// channel (XVALN, the absolute one, not the lane's own half) and norm 2 rebuilds the native order from the
|
||||
// accumulator registers with a pair-aligned half-wave exchange. Without this the ten layers reach only 41-46 dB
|
||||
// against the capture even though every product, weight map and bias is already the native one.
|
||||
//
|
||||
// Build-time shape: C, NH, KIND (0 plain, 1 encoder with 2x2 patch merge, 2 decoder with patch expand + skip),
|
||||
// XC (merged channels or low-resolution channels). One workgroup is an 8x8-token block: four waves, one per
|
||||
// 4x4-token window. Bindings and the parameter block are unchanged; the FP8 weight matrices are the packed
|
||||
// fixed-16 tiles of test_m.pack8 (the f16 vectors stay row-major).
|
||||
#extension GL_KHR_cooperative_matrix : require
|
||||
#extension GL_KHR_memory_scope_semantics : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types : require
|
||||
#extension GL_EXT_float_e4m3 : require
|
||||
#extension GL_KHR_shader_subgroup_basic : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle : require
|
||||
layout(local_size_x = 128) in;
|
||||
// o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8);
|
||||
// per MLP chunk c at 20 + 3c: w1 (FP8), b1 (f16), w2 (FP8); 68, 69: merge or expand weights (FP8) and bias (f16)
|
||||
layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[72]; } p;
|
||||
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
|
||||
layout(set = 0, binding = 2, std430) readonly buffer W16 { f16vec4 w[]; } w16;
|
||||
layout(set = 0, binding = 3, std430) readonly buffer Ib { floate4m3_t a[]; } inb; // [H][W][C]; decoders: [H/2][W/2][XC]
|
||||
layout(set = 0, binding = 4, std430) readonly buffer Sb { floate4m3_t a[]; } skb; // decoders: skip [H][W][C]
|
||||
layout(set = 0, binding = 5, std430) buffer O0 { floate4m3_t a[]; } out0; // [H][W][C]
|
||||
layout(set = 0, binding = 6, std430) buffer O1 { floate4m3_t a[]; } out1; // encoders: merged [H/2][W/2][XC]
|
||||
#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 A16 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define B16 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#define RM gl_CooperativeMatrixLayoutRowMajor
|
||||
#define CM gl_CooperativeMatrixLayoutColumnMajor
|
||||
const uint UC = uint(C), NT = uint(C) / 16u, UXC = uint(XC);
|
||||
const uint SLO = UXC + 8u; // row stride of the low-resolution token image
|
||||
const uint S4 = 4u * UC + 8u; // row stride of the 4C rows (patch expand / merge)
|
||||
const uint VWD = 16u; // row width of the per-wave V transpose buffer, in f16
|
||||
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
|
||||
#define Q16(v) float16_t(v)
|
||||
#define CL8(v) clamp(v, float16_t(-448.0), float16_t(448.0))
|
||||
// NVIDIA's accumulator rounding: one binary16 rounding of the whole tile after every logical 32 products.
|
||||
// RADV folds a register-only f32->f16->f32 conversion away (measured: half of all values and every tie
|
||||
// midpoint unchanged), so the rounding has to pass through binary16 memory. Each invocation owns 8 halves of
|
||||
// its own scratch: no cross-lane traffic, so a subgroup fence is enough (no workgroup barrier).
|
||||
shared f16vec4 RS[128 * 2];
|
||||
#define R16(m) { \
|
||||
RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4(float16_t((m)[0]), float16_t((m)[1]), float16_t((m)[2]), float16_t((m)[3])); \
|
||||
RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4(float16_t((m)[4]), float16_t((m)[5]), float16_t((m)[6]), float16_t((m)[7])); \
|
||||
subgroupMemoryBarrierShared(); \
|
||||
const f16vec4 r0_ = RS[gl_LocalInvocationID.x * 2u + 0u], r1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
|
||||
(m)[0] = float(r0_.x); (m)[1] = float(r0_.y); (m)[2] = float(r0_.z); (m)[3] = float(r0_.w); \
|
||||
(m)[4] = float(r1_.x); (m)[5] = float(r1_.y); (m)[6] = float(r1_.z); (m)[7] = float(r1_.w); }
|
||||
// one logical k32, exactly as the native K32_MIX: P = A0*B0 + A1*B1 from a zero seed, then one binary16
|
||||
// rounding of P + the previous accumulator (the native v_fma_mix), through the scratch so it is not folded.
|
||||
// Packed fixed-16 weight tiles (test_m.pack8): matrix W[K][N] as [K/16][N/16][N16][K16], so one 16x16 A tile
|
||||
// is 256 contiguous bytes at base + 256*(kTile*(N/16) + nOffset/16), read RowMajor with stride 16. Same A
|
||||
// operand values as the old plain W[K][N] ColumnMajor stride-N load, but a plain (non-transposed) load.
|
||||
#define LDA8(dst, base, ks, stride, moff) coopMatLoad((dst), w8.w, (base) + 256u * (uint(ks) * (uint(stride) / 16u) + (uint(moff) / 16u)), 16u, RM)
|
||||
|
||||
#define K32H(dst, a0, b0, a1, b1) { \
|
||||
ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); \
|
||||
RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4( \
|
||||
float16_t(P_[0] + float((dst)[0])), float16_t(P_[1] + float((dst)[1])), \
|
||||
float16_t(P_[2] + float((dst)[2])), float16_t(P_[3] + float((dst)[3]))); \
|
||||
RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4( \
|
||||
float16_t(P_[4] + float((dst)[4])), float16_t(P_[5] + float((dst)[5])), \
|
||||
float16_t(P_[6] + float((dst)[6])), float16_t(P_[7] + float((dst)[7]))); \
|
||||
subgroupMemoryBarrierShared(); \
|
||||
const f16vec4 u0_ = RS[gl_LocalInvocationID.x * 2u + 0u], u1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
|
||||
(dst)[0] = u0_.x; (dst)[1] = u0_.y; (dst)[2] = u0_.z; (dst)[3] = u0_.w; \
|
||||
(dst)[4] = u1_.x; (dst)[5] = u1_.y; (dst)[6] = u1_.z; (dst)[7] = u1_.w; }
|
||||
// a tiled f16 vector (16 equal rows) as the initial value of an accumulator whose rows are its channels
|
||||
HAC bias16(uint off, uint stride) { HAC b; coopMatLoad(b, w16.w, off / 4u, stride / 4u, CM); return b; }
|
||||
#define MK8_F8(dst, a) { \
|
||||
const float16_t q0_ = CL8(Q16((a)[0])), q1_ = CL8(Q16((a)[1])), q2_ = CL8(Q16((a)[2])), q3_ = CL8(Q16((a)[3])); \
|
||||
const float16_t q4_ = CL8(Q16((a)[4])), q5_ = CL8(Q16((a)[5])), q6_ = CL8(Q16((a)[6])), q7_ = CL8(Q16((a)[7])); \
|
||||
(dst)[0] = floate4m3_t(q0_); (dst)[1] = floate4m3_t(q1_); (dst)[2] = floate4m3_t(q2_); (dst)[3] = floate4m3_t(q3_); \
|
||||
(dst)[4] = floate4m3_t(q4_); (dst)[5] = floate4m3_t(q5_); (dst)[6] = floate4m3_t(q6_); (dst)[7] = floate4m3_t(q7_); }
|
||||
// NVIDIA's softmax numerator, two neighbouring columns in one 32-bit word (the even column's carry reaches the odd)
|
||||
void softmax2(float16_t s0, float16_t s1, out float16_t e0, out float16_t e1)
|
||||
{
|
||||
const float16_t lim = float16_t(0.55615234375);
|
||||
const float16_t cs = float16_t(0.0972222164273262 / 5.656854152679443);
|
||||
const float16_t t0 = clamp(s0 * cs, -lim, lim), t1 = clamp(s1 * cs, -lim, lim);
|
||||
const float16_t p0 = fma(t0, fma(t0, -t0, float16_t(0.92730712890625)), float16_t(1.375));
|
||||
const float16_t p1 = fma(t1, fma(t1, -t1, float16_t(0.92730712890625)), float16_t(1.375));
|
||||
const uint w = ((uint(float16BitsToUint16(p0)) | (uint(float16BitsToUint16(p1)) << 16)) << 5) + 0x7FF88000u;
|
||||
e0 = uint16BitsToFloat16(uint16_t(w & 0xFFFFu)); e1 = uint16BitsToFloat16(uint16_t(w >> 16));
|
||||
}
|
||||
float16_t gelu1(float16_t x)
|
||||
{
|
||||
const float16_t c = clamp(x, float16_t(-2.0), float16_t(2.0));
|
||||
const float16_t ac = abs(c);
|
||||
precise float16_t inner = float16_t(0.41216981) - float16_t(0.08108133) * ac;
|
||||
precise float16_t outer = float16_t(0.5) + c * inner;
|
||||
return x * outer;
|
||||
}
|
||||
#if C == 64
|
||||
#define EACH_T(X) X(0) X(1) X(2) X(3)
|
||||
#define EACH_T16(X) X(0, 1) X(2, 3)
|
||||
#elif C == 96
|
||||
#define EACH_T(X) X(0) X(1) X(2) X(3) X(4) X(5)
|
||||
#define EACH_T16(X) X(0, 1) X(2, 3) X(4, 5)
|
||||
#else
|
||||
#define EACH_T(X) X(0) X(1) X(2) X(3) X(4) X(5) X(6) X(7)
|
||||
#define EACH_T16(X) X(0, 1) X(2, 3) X(4, 5) X(6, 7)
|
||||
#endif
|
||||
|
||||
// ------------------------------------------------------------------ native arithmetic order
|
||||
// The accuracy build (kernels/build.sh D4R_PREFER_ACCURACY=1 -> -DD4R_ACCURACY -> SWIN_EXACT, and
|
||||
// D4R_NATIVE_FP8=1 -> -DD4R_FP8_WMMA -> SWIN_A8/K32_MIX) runs the token-lane block of kernels/m/swin_tok.h.
|
||||
// Its two non-obvious orderings are reproduced here in binary16:
|
||||
// * tok_rsq: per tile pp = a*a + b*b over the pairs (0,1)-(2,3) and (4,5)-(6,7), then a tree that splits the
|
||||
// tiles by parity, then a + other(a) across the half-waves, then u[0] + u[1], then the two components.
|
||||
// * the softmax: v1[k] = wg[k] + other(wg[k]) elementwise first, then (v1[0]+v1[1]) + (v1[2]+v1[3]).
|
||||
// f16 arithmetic is marked precise so nothing is contracted into an FMA (the native kernel builds with
|
||||
// #pragma clang fp contract(off)).
|
||||
#define SH16(v) subgroupShuffleXor(packHalf2x16(vec2(float((v).x), float((v).y))), 16u)
|
||||
#define UN16(p) f16vec2(float16_t(unpackHalf2x16(p).x), float16_t(unpackHalf2x16(p).y))
|
||||
#if C == 64
|
||||
#define RSQ_TREE(OUT) { \
|
||||
f16vec2 P0_[NT], P1_[NT]; \
|
||||
EACH_T(RSQ_TILE) \
|
||||
precise f16vec2 W00_ = P0_[0] + P0_[2], W01_ = P1_[0] + P1_[2]; \
|
||||
precise f16vec2 W10_ = P0_[1] + P0_[3], W11_ = P1_[1] + P1_[3]; \
|
||||
precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
|
||||
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
|
||||
precise f16vec2 S_ = U0_ + U1_; \
|
||||
OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
|
||||
#elif C == 96
|
||||
#define RSQ_TREE(OUT) { \
|
||||
f16vec2 P0_[NT], P1_[NT]; \
|
||||
EACH_T(RSQ_TILE) \
|
||||
precise f16vec2 W00_ = (P0_[0] + P0_[2]) + P0_[4], W01_ = (P1_[0] + P1_[2]) + P1_[4]; \
|
||||
precise f16vec2 W10_ = (P0_[1] + P0_[3]) + P0_[5], W11_ = (P1_[1] + P1_[3]) + P1_[5]; \
|
||||
precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
|
||||
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
|
||||
precise f16vec2 S_ = U0_ + U1_; \
|
||||
OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
|
||||
#else
|
||||
#define RSQ_TREE(OUT) { \
|
||||
f16vec2 P0_[NT], P1_[NT]; \
|
||||
EACH_T(RSQ_TILE) \
|
||||
precise f16vec2 W00_ = (P0_[0] + P0_[2]) + (P0_[4] + P0_[6]), W01_ = (P1_[0] + P1_[2]) + (P1_[4] + P1_[6]); \
|
||||
precise f16vec2 W10_ = (P0_[1] + P0_[3]) + (P0_[5] + P0_[7]), W11_ = (P1_[1] + P1_[3]) + (P1_[5] + P1_[7]); \
|
||||
precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
|
||||
precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
|
||||
precise f16vec2 S_ = U0_ + U1_; \
|
||||
OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
|
||||
#endif
|
||||
#define RSQ_TILE(T) { \
|
||||
precise f16vec2 A_ = f16vec2(EV(T, 0), EV(T, 1)), B_ = f16vec2(EV(T, 2), EV(T, 3)); \
|
||||
precise f16vec2 C_ = f16vec2(EV(T, 4), EV(T, 5)), D_ = f16vec2(EV(T, 6), EV(T, 7)); \
|
||||
P0_[T] = (A_ * A_) + (B_ * B_); P1_[T] = (C_ * C_) + (D_ * D_); }
|
||||
|
||||
shared f16vec4 VSH[4 * 2 * 16 * VWD / 4]; // per wave and head: the V tiles, transposed through shared memory
|
||||
#if KIND == 1
|
||||
shared floate4m3_t NB[32 * S4]; // 16 merged rows, preceded by 16 rows of slack for the gather
|
||||
#endif
|
||||
#if KIND == 2
|
||||
shared floate4m3_t NLI[16 * SLO]; // the block's 16 low-resolution tokens
|
||||
shared floate4m3_t NEX[16 * S4]; // their expanded 4C columns
|
||||
#endif
|
||||
#define SYNC memoryBarrierShared(); barrier();
|
||||
void main()
|
||||
{
|
||||
const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u;
|
||||
const uint row = lane & 15u, hf = lane >> 4, ch8 = 8u * hf; // one token per lane, its channels 8h..8h+7
|
||||
const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy;
|
||||
const uint bxr = 4u * (mt & 1u) + (row & 3u), byr = 4u * (mt >> 1) + (row >> 2);
|
||||
const int tx = X0 + int(bxr), ty = Y0 + int(byr);
|
||||
const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W;
|
||||
const uint iy = uint(mirror(ty, p.Hh)), ix = uint(mirror(tx, p.W));
|
||||
HAC res[NT]; // the residual stream: rows are channels, columns are tokens
|
||||
B8 h1[NT]; // the normalised activation, the K operand of every product
|
||||
// ------------------------------------------------------------ patch expand: x0 = q8(E) + skip
|
||||
#if KIND == 2
|
||||
{
|
||||
const uint W2 = uint(p.W) / 2u, H2 = uint(p.Hh) / 2u, part = UXC / 8u, c0 = (mt * 2u + hf) * part;
|
||||
const uint lo = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2)) * UXC;
|
||||
for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lo + c];
|
||||
SYNC
|
||||
// this wave computes sub-position mt of every low-resolution token's 4C columns
|
||||
HAC ef[NT];
|
||||
#define EF_INIT(N) ef[N] = bias16(p.o[69] + 16u * (mt * NT + (N)), 4u * UC);
|
||||
#define EF_K32(N) { A8 xa0, xa1; \
|
||||
LDA8(xa0, p.o[68], kt, 4u * UC, 16u * (mt * NT + (N))); \
|
||||
LDA8(xa1, p.o[68], kt + 1u, 4u * UC, 16u * (mt * NT + (N))); \
|
||||
K32H(ef[N], xa0, b0, xa1, b1) }
|
||||
#define EF_STORE(N) { A8 ey; MK8_F8(ey, ef[N]); coopMatStore(ey, NEX, mt * UC + 16u * (N), S4, RM); }
|
||||
EACH_T(EF_INIT)
|
||||
for (uint kt = 0u; kt < UXC / 16u; kt += 2u)
|
||||
{
|
||||
B8 b0, b1;
|
||||
coopMatLoad(b0, NLI, kt * 16u, SLO, CM);
|
||||
coopMatLoad(b1, NLI, (kt + 1u) * 16u, SLO, CM);
|
||||
EACH_T(EF_K32)
|
||||
}
|
||||
EACH_T(EF_STORE)
|
||||
SYNC
|
||||
}
|
||||
#endif
|
||||
// ------------------------------------------------------------ norm 1: sum of squares, residual and h1
|
||||
#if KIND == 2
|
||||
const uint lb0 = ((byr >> 1) * 4u + (bxr >> 1)) * S4 + ((byr & 1u) * 2u + (bxr & 1u)) * UC + ch8;
|
||||
const uint sb0 = (iy * uint(p.W) + ix) * UC + ch8;
|
||||
#define XVAL(C) (float16_t(NEX[lb0 + (C)]) + float16_t(skb.a[sb0 + (C)]))
|
||||
#else
|
||||
const uint ib0 = (iy * uint(p.W) + ix) * UC + ch8;
|
||||
#define XVAL(C) float16_t(inb.a[ib0 + (C)])
|
||||
#endif
|
||||
// The native's register element (tile T, half hf, element I) holds the natural channel tok_xnat(T, 8hf+I) =
|
||||
// 16T + 8*((I>>2)&1) + 4hf + (I&3) (kernels/m/swin_tok.h tok_rsq on the register layout the input plane implies).
|
||||
// This shader's own accumulators and A rows keep natural channels, so h1 and the residual use XVAL (the lane's
|
||||
// own half); the norm tree, however, groups per the native element order, so it reads the absolute channel.
|
||||
#if KIND == 2
|
||||
#define XVALN(C) (float16_t(NEX[(lb0 - ch8) + (C)]) + float16_t(skb.a[(sb0 - ch8) + (C)]))
|
||||
#else
|
||||
#define XVALN(C) float16_t(inb.a[(iy * uint(p.W) + ix) * UC + (C)])
|
||||
#endif
|
||||
#define EV(T, I) XVALN(16u * (T) + 8u * ((uint(I) >> 2) & 1u) + 4u * hf + (uint(I) & 3u))
|
||||
float16_t rn; RSQ_TREE(rn)
|
||||
#undef EV
|
||||
#define NM1(N) { \
|
||||
const uint wb = (p.o[0] + 16u * (N) + ch8) / 4u, bb = (p.o[1] + 16u * (N) + ch8) / 4u; \
|
||||
const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \
|
||||
const float16_t x0 = XVAL(16u * (N) + 0u), x1 = XVAL(16u * (N) + 1u), x2 = XVAL(16u * (N) + 2u), x3 = XVAL(16u * (N) + 3u); \
|
||||
const float16_t x4 = XVAL(16u * (N) + 4u), x5 = XVAL(16u * (N) + 5u), x6 = XVAL(16u * (N) + 6u), x7 = XVAL(16u * (N) + 7u); \
|
||||
res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \
|
||||
res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \
|
||||
res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \
|
||||
res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); \
|
||||
h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn * g0.y))); \
|
||||
h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn * g0.w))); \
|
||||
h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn * g1.y))); \
|
||||
h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn * g1.w))); }
|
||||
EACH_T(NM1)
|
||||
// ------------------------------------------------------------ attention, one head at a time
|
||||
for (uint hd = 0u; hd < uint(NH); ++hd)
|
||||
{
|
||||
const uint oq = p.o[4u + 4u * hd], ov = p.o[5u + 4u * hd], ob = p.o[6u + 4u * hd], ow = p.o[7u + 4u * hd];
|
||||
HAC q[2], v[2];
|
||||
q[0] = HAC(0.0); q[1] = HAC(0.0); v[0] = HAC(0.0); v[1] = HAC(0.0);
|
||||
#define QV_G(A, B) { \
|
||||
const uint kt = A; \
|
||||
A8 aq0, aq1, av0, av1; \
|
||||
LDA8(aq0, oq, kt, 32u, 0u); \
|
||||
LDA8(aq1, oq, kt + 1u, 32u, 0u); \
|
||||
LDA8(av0, ov, kt, 32u, 0u); \
|
||||
LDA8(av1, ov, kt + 1u, 32u, 0u); \
|
||||
K32H(q[0], aq0, h1[A], aq1, h1[B]) \
|
||||
K32H(v[0], av0, h1[A], av1, h1[B]) \
|
||||
LDA8(aq0, oq, kt, 32u, 16u); \
|
||||
LDA8(aq1, oq, kt + 1u, 32u, 16u); \
|
||||
LDA8(av0, ov, kt, 32u, 16u); \
|
||||
LDA8(av1, ov, kt + 1u, 32u, 16u); \
|
||||
K32H(q[1], aq0, h1[A], aq1, h1[B]) \
|
||||
K32H(v[1], av0, h1[A], av1, h1[B]) }
|
||||
EACH_T16(QV_G)
|
||||
// scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias
|
||||
A8 qa[2]; B8 qb[2];
|
||||
for (uint m = 0u; m < 2u; ++m)
|
||||
{
|
||||
const float16_t q0 = CL8(q[m][0]), q1 = CL8(q[m][1]), q2 = CL8(q[m][2]), q3 = CL8(q[m][3]);
|
||||
const float16_t q4 = CL8(q[m][4]), q5 = CL8(q[m][5]), q6 = CL8(q[m][6]), q7 = CL8(q[m][7]);
|
||||
qa[m][0] = floate4m3_t(q0); qb[m][0] = qa[m][0]; qa[m][1] = floate4m3_t(q1); qb[m][1] = qa[m][1];
|
||||
qa[m][2] = floate4m3_t(q2); qb[m][2] = qa[m][2]; qa[m][3] = floate4m3_t(q3); qb[m][3] = qa[m][3];
|
||||
qa[m][4] = floate4m3_t(q4); qb[m][4] = qa[m][4]; qa[m][5] = floate4m3_t(q5); qb[m][5] = qa[m][5];
|
||||
qa[m][6] = floate4m3_t(q6); qb[m][6] = qa[m][6]; qa[m][7] = floate4m3_t(q7); qb[m][7] = qa[m][7];
|
||||
}
|
||||
// the bias is read column-major, so the accumulator holds S^T: this lane holds the keys 8h..8h+7 of one query
|
||||
HAC s = bias16(ob, 16u);
|
||||
K32H(s, qa[0], qb[0], qa[1], qb[1])
|
||||
float16_t e[8];
|
||||
softmax2(float16_t(s[0]), float16_t(s[1]), e[0], e[1]);
|
||||
softmax2(float16_t(s[2]), float16_t(s[3]), e[2], e[3]);
|
||||
softmax2(float16_t(s[4]), float16_t(s[5]), e[4], e[5]);
|
||||
softmax2(float16_t(s[6]), float16_t(s[7]), e[6], e[7]);
|
||||
precise f16vec2 G0_ = f16vec2(e[0], e[1]), G1_ = f16vec2(e[2], e[3]);
|
||||
precise f16vec2 G2_ = f16vec2(e[4], e[5]), G3_ = f16vec2(e[6], e[7]);
|
||||
precise f16vec2 V0_ = G0_ + UN16(SH16(G0_)), V1_ = G1_ + UN16(SH16(G1_));
|
||||
precise f16vec2 V2_ = G2_ + UN16(SH16(G2_)), V3_ = G3_ + UN16(SH16(G3_));
|
||||
precise f16vec2 VV_ = (V0_ + V1_) + (V2_ + V3_);
|
||||
const float16_t rc = float16_t(1.0 / float(VV_.x + VV_.y));
|
||||
B16 pb;
|
||||
pb[0] = G0_.x * rc; pb[1] = G0_.y * rc; pb[2] = G1_.x * rc; pb[3] = G1_.y * rc;
|
||||
pb[4] = G2_.x * rc; pb[5] = G2_.y * rc; pb[6] = G3_.x * rc; pb[7] = G3_.y * rc;
|
||||
// V through this wave's shared buffer, then O^T = V^T x P^T: rows are the head channels, columns the tokens
|
||||
ACC o[2];
|
||||
coopMatStore(v[0], VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
|
||||
coopMatStore(v[1], VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
|
||||
A16 vt0, vt1;
|
||||
coopMatLoad(vt0, VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
|
||||
coopMatLoad(vt1, VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
|
||||
o[0] = coopMatMulAdd(vt0, pb, ACC(0.0));
|
||||
o[1] = coopMatMulAdd(vt1, pb, ACC(0.0));
|
||||
R16(o[0]); R16(o[1]);
|
||||
B8 ob8[2];
|
||||
MK8_F8(ob8[0], o[0]); MK8_F8(ob8[1], o[1]);
|
||||
// the output projection of this head's 32 attention channels, straight into the residual stream
|
||||
#define WO_N(N) { \
|
||||
A8 a0, a1; \
|
||||
LDA8(a0, ow, 0u, UC, 16u * (N)); \
|
||||
LDA8(a1, ow, 1u, UC, 16u * (N)); \
|
||||
K32H(res[N], a0, ob8[0], a1, ob8[1]) }
|
||||
EACH_T(WO_N)
|
||||
}
|
||||
// ------------------------------------------------------------ norm 2
|
||||
// the native element (T, 8h+i) holds the natural channel tok_xnat(T, 8h+i) = 16T + 4h + (i&3) + 8*((i>>2)&1),
|
||||
// which lives in my natural-order accumulator at the half ((i>>2)&1) and the element 4h+(i&3). The tree needs
|
||||
// the pairs in that order, so fetch the (pair 2h, pair 2h+1) of the half ((i>>2)&1) and select by this lane's half.
|
||||
f16vec2 XA_[NT], XB_[NT], XC_[NT], XD_[NT];
|
||||
for (uint t_ = 0u; t_ < NT; ++t_)
|
||||
{
|
||||
const f16vec2 p0_ = f16vec2(float16_t(res[t_][0]), float16_t(res[t_][1]));
|
||||
const f16vec2 p1_ = f16vec2(float16_t(res[t_][2]), float16_t(res[t_][3]));
|
||||
const f16vec2 p2_ = f16vec2(float16_t(res[t_][4]), float16_t(res[t_][5]));
|
||||
const f16vec2 p3_ = f16vec2(float16_t(res[t_][6]), float16_t(res[t_][7]));
|
||||
const f16vec2 o0_ = UN16(SH16(p0_)), o1_ = UN16(SH16(p1_)), o2_ = UN16(SH16(p2_)), o3_ = UN16(SH16(p3_));
|
||||
if (hf == 0u) { XA_[t_] = p0_; XB_[t_] = p1_; XC_[t_] = o0_; XD_[t_] = o1_; }
|
||||
else { XA_[t_] = o2_; XB_[t_] = o3_; XC_[t_] = p2_; XD_[t_] = p3_; }
|
||||
}
|
||||
#define EV(T, I) ((uint(I) < 2u) ? XA_[T][uint(I) & 1u] : (uint(I) < 4u) ? XB_[T][uint(I) & 1u] : \
|
||||
(uint(I) < 6u) ? XC_[T][uint(I) & 1u] : XD_[T][uint(I) & 1u])
|
||||
float16_t rn2; RSQ_TREE(rn2)
|
||||
#undef EV
|
||||
#define N2(N) { \
|
||||
const uint wb = (p.o[2] + 16u * (N) + ch8) / 4u, bb = (p.o[3] + 16u * (N) + ch8) / 4u; \
|
||||
const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \
|
||||
const float16_t x0 = res[N][0], x1 = res[N][1], x2 = res[N][2], x3 = res[N][3]; \
|
||||
const float16_t x4 = res[N][4], x5 = res[N][5], x6 = res[N][6], x7 = res[N][7]; \
|
||||
h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn2 * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn2 * g0.y))); \
|
||||
h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn2 * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn2 * g0.w))); \
|
||||
h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn2 * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn2 * g1.y))); \
|
||||
h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn2 * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn2 * g1.w))); \
|
||||
res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \
|
||||
res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \
|
||||
res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \
|
||||
res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); }
|
||||
EACH_T(N2)
|
||||
// ------------------------------------------------------------ the MLP, 32 hidden channels at a time
|
||||
for (uint c = 0u; c < UC / 8u; ++c)
|
||||
{
|
||||
const uint ow1 = p.o[20u + 3u * c], ob1 = p.o[21u + 3u * c], ow2 = p.o[22u + 3u * c];
|
||||
HAC hm[2];
|
||||
hm[0] = bias16(ob1, 32u); hm[1] = bias16(ob1 + 16u, 32u);
|
||||
#define MLP_G(A, B) { \
|
||||
const uint kt = A; \
|
||||
A8 a0, a1; \
|
||||
LDA8(a0, ow1, kt, 32u, 0u); \
|
||||
LDA8(a1, ow1, kt + 1u, 32u, 0u); \
|
||||
K32H(hm[0], a0, h1[A], a1, h1[B]) \
|
||||
LDA8(a0, ow1, kt, 32u, 16u); \
|
||||
LDA8(a1, ow1, kt + 1u, 32u, 16u); \
|
||||
K32H(hm[1], a0, h1[A], a1, h1[B]) }
|
||||
EACH_T16(MLP_G)
|
||||
B8 g0, g1;
|
||||
{ const float16_t v0 = gelu1(hm[0][0]), v1 = gelu1(hm[0][1]), v2 = gelu1(hm[0][2]), v3 = gelu1(hm[0][3]);
|
||||
const float16_t v4 = gelu1(hm[0][4]), v5 = gelu1(hm[0][5]), v6 = gelu1(hm[0][6]), v7 = gelu1(hm[0][7]);
|
||||
g0[0] = floate4m3_t(CL8(v0)); g0[1] = floate4m3_t(CL8(v1)); g0[2] = floate4m3_t(CL8(v2)); g0[3] = floate4m3_t(CL8(v3));
|
||||
g0[4] = floate4m3_t(CL8(v4)); g0[5] = floate4m3_t(CL8(v5)); g0[6] = floate4m3_t(CL8(v6)); g0[7] = floate4m3_t(CL8(v7)); }
|
||||
{ const float16_t v0 = gelu1(hm[1][0]), v1 = gelu1(hm[1][1]), v2 = gelu1(hm[1][2]), v3 = gelu1(hm[1][3]);
|
||||
const float16_t v4 = gelu1(hm[1][4]), v5 = gelu1(hm[1][5]), v6 = gelu1(hm[1][6]), v7 = gelu1(hm[1][7]);
|
||||
g1[0] = floate4m3_t(CL8(v0)); g1[1] = floate4m3_t(CL8(v1)); g1[2] = floate4m3_t(CL8(v2)); g1[3] = floate4m3_t(CL8(v3));
|
||||
g1[4] = floate4m3_t(CL8(v4)); g1[5] = floate4m3_t(CL8(v5)); g1[6] = floate4m3_t(CL8(v6)); g1[7] = floate4m3_t(CL8(v7)); }
|
||||
#define MLP2_N(N) { \
|
||||
A8 a0, a1; \
|
||||
LDA8(a0, ow2, 0u, UC, 16u * (N)); \
|
||||
LDA8(a1, ow2, 1u, UC, 16u * (N)); \
|
||||
K32H(res[N], a0, g0, a1, g1) }
|
||||
EACH_T(MLP2_N)
|
||||
}
|
||||
// ------------------------------------------------------------ the block output, FP8
|
||||
const uint dst = uint(ty * p.W + tx) * UC + ch8;
|
||||
if (inside)
|
||||
for (uint nt = 0u; nt < NT; ++nt)
|
||||
{
|
||||
#define OUT8(I) out0.a[dst + 16u * nt + uint(I)] = floate4m3_t(CL8(res[nt][I]));
|
||||
OUT8(0) OUT8(1) OUT8(2) OUT8(3) OUT8(4) OUT8(5) OUT8(6) OUT8(7)
|
||||
#undef OUT8
|
||||
}
|
||||
// ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution
|
||||
#if KIND == 1
|
||||
{
|
||||
// the 16 merged rows are gathered from the four waves: the row is the merged token of the lane's pixel
|
||||
const uint mrow = (byr >> 1) * 4u + (bxr >> 1), qq = (byr & 1u) * 2u + (bxr & 1u);
|
||||
#define MG_W(N) { A8 y; MK8_F8(y, res[N]); coopMatStore(y, NB, 16u * S4 + mrow * S4 + qq * UC + 16u * (N) - row * S4, S4, RM); }
|
||||
EACH_T(MG_W)
|
||||
SYNC
|
||||
const uint MT = UXC / 16u, NI = (MT + 3u) / 4u;
|
||||
const int ym = (Y0 >> 1) + int(row >> 2), xm = (X0 >> 1) + int(row & 3u);
|
||||
const bool mok = ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2;
|
||||
const uint md = uint(ym * (p.W / 2) + xm) * UXC + ch8;
|
||||
for (uint i = 0u; i < NI; ++i)
|
||||
{
|
||||
const uint u = mt + 4u * i;
|
||||
if (u < MT)
|
||||
{
|
||||
HAC y = bias16(p.o[69] + 16u * u, UXC);
|
||||
for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u)
|
||||
{
|
||||
B8 b0, b1; coopMatLoad(b0, NB, 16u * S4 + kt * 16u, S4, CM);
|
||||
coopMatLoad(b1, NB, 16u * S4 + (kt + 1u) * 16u, S4, CM);
|
||||
A8 a0, a1; LDA8(a0, p.o[68], kt, UXC, 16u * u);
|
||||
LDA8(a1, p.o[68], kt + 1u, UXC, 16u * u);
|
||||
K32H(y, a0, b0, a1, b1)
|
||||
}
|
||||
if (mok)
|
||||
{
|
||||
#define MGO(I) out1.a[md + 16u * u + uint(I)] = floate4m3_t(CL8(y[I]));
|
||||
MGO(0) MGO(1) MGO(2) MGO(3) MGO(4) MGO(5) MGO(6) MGO(7)
|
||||
#undef MGO
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
Reference in new issue
Block a user