feat: add d4r inference engines and Windows parity tooling
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
This commit is contained in:
ayo committed
1 parent
25493b3198
commit
7c5d3c9157
430 files changed
+76191
No files matched your search
@@ -0,0 +1,481 @@
|
||||
#version 460
|
||||
// d4r engine, preset K: one Swin block (8x8-token window) written from the layer's mathematics, not from NVIDIA's
|
||||
// kernel. Weights come from the engine's own model file (extract_k.py): plain row-major matrices, with the two
|
||||
// per-channel gains folded into Q/K/V and fc1.
|
||||
// Build-time shape: C, NH, KIND (0 plain, 1 encoder with patch merge, 2 decoder with patch expand + skip),
|
||||
// POSATTN (position-only attention), EMBED (enc0: 16-channel input embedding + ReLU), HEAD (dec0: output head),
|
||||
// XC (merged channels or low-resolution channels), MERGE_N (XC rounded up to 16).
|
||||
// - Four waves per window, one per 16-token tile. The residual stream of a tile stays in registers (one f32
|
||||
// accumulator per 16 channels) from the input to the output.
|
||||
// - Shared memory only stages what has to be read back as a matrix operand or by another wave: one buffer BF and
|
||||
// the Q/K/V regions. It is kept small on purpose: the kernel waits on memory a lot, so its speed depends on how
|
||||
// many windows fit on a compute unit (64 KB of shared memory each).
|
||||
// - Rows are padded: with power-of-two row strides the 16 rows of an operand load land in the same LDS banks.
|
||||
// - Global arrays use f16vec4 units. Shared token tiles are row-major except enc0, which uses padded
|
||||
// column-major scalar-f16 arrays so cooperative-matrix stores can pack adjacent token values.
|
||||
// - FOLD_GAINS=0 keeps the per-channel gains explicit instead of folding them into weights.
|
||||
#extension GL_KHR_cooperative_matrix : require
|
||||
#extension GL_KHR_memory_scope_semantics : require
|
||||
#extension GL_EXT_shader_explicit_arithmetic_types : require
|
||||
#extension GL_KHR_shader_subgroup_basic : require
|
||||
#extension GL_KHR_shader_subgroup_shuffle : require
|
||||
#extension GL_EXT_control_flow_attributes : require
|
||||
#ifndef PADX
|
||||
#define PADX 8
|
||||
#endif
|
||||
#ifndef FOLD_GAINS
|
||||
#define FOLD_GAINS 1
|
||||
#endif
|
||||
#ifndef PADT
|
||||
#define PADT 8
|
||||
#endif
|
||||
#ifndef PADS
|
||||
#define PADS 8
|
||||
#endif
|
||||
#ifndef ABL_HEADS
|
||||
#define ABL_HEADS uint(NH)
|
||||
#endif
|
||||
#ifndef ABL_CHUNKS
|
||||
#define ABL_CHUNKS (UC / 8u)
|
||||
#endif
|
||||
// TPW: 16-token tiles of the 8x8 window per wave (1, 2 or 4). With TPW > 1 a wave walks TPW tiles and loads every
|
||||
// weight tile once for all of them (four waves per window each read the whole layer's weights otherwise, about a
|
||||
// fifth of the network's time). The arithmetic and its order per tile do not change.
|
||||
#ifndef TPW
|
||||
#define TPW 1
|
||||
#endif
|
||||
layout(local_size_x = 128 / TPW) in;
|
||||
layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[108]; } p;
|
||||
layout(set = 0, binding = 1, std430) readonly buffer Wb { f16vec4 w[]; } wb;
|
||||
layout(set = 0, binding = 2, std430) readonly buffer Ib { f16vec4 a[]; } inb;
|
||||
layout(set = 0, binding = 3, std430) readonly buffer Sb { f16vec4 a[]; } skb;
|
||||
layout(set = 0, binding = 4, std430) buffer O0 { f16vec4 a[]; } out0;
|
||||
layout(set = 0, binding = 5, std430) buffer O1 { f16vec4 a[]; } out1;
|
||||
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
||||
#define MA coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
||||
#define MB coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
||||
#define RM gl_CooperativeMatrixLayoutRowMajor
|
||||
#define CMJ gl_CooperativeMatrixLayoutColumnMajor
|
||||
#define LD(m, arr, off, stride) coopMatLoad(m, arr, (off) / 4u, (stride) / 4u, RM)
|
||||
#define ST(m, arr, off, stride) coopMatStore(HACC(m), arr, (off) / 4u, (stride) / 4u, RM)
|
||||
#define SYNC memoryBarrierShared(); barrier();
|
||||
const uint UC = uint(C), NT = uint(C) / 16u, HC = uint(C) / 2u;
|
||||
const uint SX = UC + uint(PADX); // row stride of the token layout (64 rows of C)
|
||||
const uint S4 = 4u * UC + 8u; // row stride of the 16 x 4C layout (patch expand / merge)
|
||||
const uint SQ = 32u + uint(PADT); // row stride of Q, K, V, attention output and MLP hidden (64 rows each)
|
||||
const uint RQ = 64u * SQ; // one such region
|
||||
const uint SLO = uint(XC) + 8u, SEM = 24u, SHD = 56u, SMG = uint(MERGE_N) + 8u; // low-res input, embedding input, head, merged
|
||||
#if POSATTN
|
||||
#define NREG 1
|
||||
#define VOFF 0u
|
||||
#else
|
||||
#define NREG 2
|
||||
#define VOFF RQ
|
||||
#endif
|
||||
// TT must also hold the token layout (decoders: x0), the 16 x 4C layout (encoders: merge input) or the head's output
|
||||
#if KIND == 2
|
||||
#define TT_ALT (16 * (C + PADX))
|
||||
#elif KIND == 1
|
||||
#define TT_ALT (4 * (4 * C + 8))
|
||||
#else
|
||||
#define TT_ALT 0
|
||||
#endif
|
||||
#if HEAD
|
||||
#define TT_HEAD (16 * 56)
|
||||
#else
|
||||
#define TT_HEAD 0
|
||||
#endif
|
||||
#define TT_REG (NREG * 16 * (32 + PADT))
|
||||
#define MAX2(a, b) ((a) > (b) ? (a) : (b))
|
||||
// BF: this tile's rows hold, in turn, the normalised input, Q, the scores / probabilities, the attention output, the
|
||||
// MLP input and the block output; as a whole it is also the 16 x 4C scratch. TT: K | V shared by the four waves; its
|
||||
// first region doubles as MLP hidden and scratch.
|
||||
// enc0's 32-channel, position-only block benefits from column-major shared tiles. All addressing uses
|
||||
// cooperative-matrix layouts rather than hardware fragment indices; other layers retain their existing layout.
|
||||
#ifndef COLUMN_MAJOR_ENC0
|
||||
#define COLUMN_MAJOR_ENC0 (EMBED && C == 32)
|
||||
#endif
|
||||
#if COLUMN_MAJOR_ENC0 && (!EMBED || C != 32 || !POSATTN || KIND != 1 || HEAD)
|
||||
#error Column-major shared tiles require the 32-channel positional enc0 block
|
||||
#endif
|
||||
#if COLUMN_MAJOR_ENC0
|
||||
const uint SP = 72u; // 64 token rows plus padding; patch scratch is transposed with 16 rows
|
||||
shared float16_t BF[32 * 72];
|
||||
shared float16_t TT[32 * 72];
|
||||
#define LOAD_B_TOKEN(m,r,c) coopMatLoad(m,BF,(c)*SP+(r),SP,CMJ)
|
||||
#define STORE_B_TOKEN(m,r,c) coopMatStore(HACC(m),BF,(c)*SP+(r),SP,CMJ)
|
||||
#define LOAD_T_TOKEN(m,r,c,s) coopMatLoad(m,TT,(c)*SP+(r),SP,CMJ)
|
||||
#define STORE_T_TOKEN(m,r,c,s) coopMatStore(HACC(m),TT,(c)*SP+(r),SP,CMJ)
|
||||
#define STORE_B_PATCH(m,r,c,s) coopMatStore(HACC(m),BF,(c)*16u+(r),16u,CMJ)
|
||||
#define LOAD_T_RAW(m,off,s) coopMatLoad(m,TT,off,s,RM)
|
||||
f16vec4 readBToken(uint r,uint c) {uint i=c*SP+r;return f16vec4(BF[i],BF[i+SP],BF[i+2u*SP],BF[i+3u*SP]);}
|
||||
f16vec4 readTToken(uint r,uint c,uint stride) {uint i=c*SP+r;return f16vec4(TT[i],TT[i+SP],TT[i+2u*SP],TT[i+3u*SP]);}
|
||||
f16vec4 readBPatch(uint r,uint c,uint stride) {uint i=c*16u+r;return f16vec4(BF[i],BF[i+16u],BF[i+32u],BF[i+48u]);}
|
||||
void writeBToken(uint r,uint c,f16vec4 v) {uint i=c*SP+r;BF[i]=v.x;BF[i+SP]=v.y;BF[i+2u*SP]=v.z;BF[i+3u*SP]=v.w;}
|
||||
void writeTToken(uint r,uint c,uint stride,f16vec4 v) {uint i=c*SP+r;TT[i]=v.x;TT[i+SP]=v.y;TT[i+2u*SP]=v.z;TT[i+3u*SP]=v.w;}
|
||||
void writeTRaw(uint off,f16vec4 v) {TT[off]=v.x;TT[off+1u]=v.y;TT[off+2u]=v.z;TT[off+3u]=v.w;}
|
||||
#else
|
||||
shared f16vec4 BF[16 * (C + PADX)];
|
||||
shared f16vec4 TT[MAX2(TT_REG, MAX2(TT_ALT, TT_HEAD))];
|
||||
#define LOAD_B_TOKEN(m,r,c) LD(m,BF,(r)*SX+(c),SX)
|
||||
#define STORE_B_TOKEN(m,r,c) ST(m,BF,(r)*SX+(c),SX)
|
||||
#define LOAD_T_TOKEN(m,r,c,s) LD(m,TT,(r)*(s)+(c),s)
|
||||
#define STORE_T_TOKEN(m,r,c,s) ST(m,TT,(r)*(s)+(c),s)
|
||||
#define STORE_B_PATCH(m,r,c,s) ST(m,BF,(r)*(s)+(c),s)
|
||||
#define LOAD_T_RAW(m,off,s) LD(m,TT,off,s)
|
||||
f16vec4 readBToken(uint r,uint c) {return BF[(r*SX+c)/4u];}
|
||||
f16vec4 readTToken(uint r,uint c,uint stride) {return TT[(r*stride+c)/4u];}
|
||||
f16vec4 readBPatch(uint r,uint c,uint stride) {return BF[(r*stride+c)/4u];}
|
||||
void writeBToken(uint r,uint c,f16vec4 v) {BF[(r*SX+c)/4u]=v;}
|
||||
void writeTToken(uint r,uint c,uint stride,f16vec4 v) {TT[(r*stride+c)/4u]=v;}
|
||||
void writeTRaw(uint off,f16vec4 v) {TT[off/4u]=v;}
|
||||
#endif
|
||||
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
|
||||
int tdiv2(int v) { return v < 0 ? -((-v) / 2) : v / 2; }
|
||||
ACC bias_tile(uint off, uint stride) { HACC b; LD(b, wb.w, off, stride); return ACC(b); } // vectors are stored as 16 equal rows
|
||||
f16vec4 softmax_exp(f16vec4 s)
|
||||
{
|
||||
const f16vec4 t = clamp(s, f16vec4(uint16BitsToFloat16(uint16_t(0xCD69))), f16vec4(uint16BitsToFloat16(uint16_t(0x4D69))));
|
||||
const f16vec4 v = fma(t * t, f16vec4(uint16BitsToFloat16(uint16_t(0x91D5))), f16vec4(1.0));
|
||||
return f16vec4(exp2(vec4(fma(t, v, f16vec4(uint16BitsToFloat16(uint16_t(0xC50D)))))));
|
||||
}
|
||||
f16vec4 gelu(f16vec4 x)
|
||||
{
|
||||
const f16vec4 c = clamp(x, f16vec4(-2.0), f16vec4(2.0));
|
||||
// the reference rounds after every binary16 operation (its compiler has fp contraction off): keep the
|
||||
// same steps, or a contracted fma changes the last bit of most outputs
|
||||
precise f16vec4 u = f16vec4(0.41216981) - f16vec4(0.08108133) * abs(c);
|
||||
precise f16vec4 y = f16vec4(0.5) + c * u;
|
||||
return x * y;
|
||||
}
|
||||
void main()
|
||||
{
|
||||
const uint mt0 = (gl_LocalInvocationID.x >> 5) * uint(TPW), lane = gl_LocalInvocationID.x & 31u; // wave = TPW token tiles
|
||||
const uint row = lane & 15u, hf = lane >> 4; // two lanes per token
|
||||
const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy;
|
||||
const uint ch0 = hf * HC; // this lane's first channel
|
||||
// per tile w of this wave (FORW): its index mt in the window and this lane's token (WINV)
|
||||
#define FORW [[unroll]] for (uint w = 0u; w < uint(TPW); ++w)
|
||||
#define WINV const uint mt = mt0 + w, tok = 16u * mt + row; const int ty = Y0 + int(tok >> 3), tx = X0 + int(tok & 7u);
|
||||
ACC ACC_[TPW][C / 16]; // the residual stream of each tile
|
||||
HACC RES_[TPW][C / 16]; // x0, added back at the MLP boundary as one f16 add
|
||||
#define acc ACC_[w]
|
||||
#define residual RES_[w]
|
||||
// ------------------------------------------------------------ input: x0 into acc, x0 * rsqrt(sum x0^2) into BF
|
||||
#if KIND == 2
|
||||
{
|
||||
// low-resolution tokens (one per row index), channels split over the 8 lane groups of the window
|
||||
const int W2 = p.W / 2, H2 = p.Hh / 2;
|
||||
const uint part = uint(XC) / 8u;
|
||||
const uint src = uint(mirror(tdiv2(Y0) + int(row >> 2), H2) * W2 + mirror(tdiv2(X0) + int(row & 3u), W2)) * uint(XC);
|
||||
FORW {
|
||||
WINV
|
||||
const uint c0 = (mt * 2u + hf) * part;
|
||||
for (uint c = c0; c < c0 + part; c += 4u) writeTRaw(row*SLO+c,inb.a[(src+c)/4u]);
|
||||
}
|
||||
SYNC
|
||||
// expand: tile mt computes quarter mt of the 4C columns
|
||||
FORW {
|
||||
WINV
|
||||
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = bias_tile(p.o[105] + mt * UC + nt * 16u, 4u * UC);
|
||||
for (uint kt = 0u; kt < uint(XC) / 16u; ++kt)
|
||||
{
|
||||
MA a; LOAD_T_RAW(a,kt*16u,SLO);
|
||||
MB b[C / 16];
|
||||
for (uint nt = 0u; nt < NT; ++nt) LD(b[nt], wb.w, p.o[104] + kt * 16u * 4u * UC + mt * UC + nt * 16u, 4u * UC);
|
||||
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = coopMatMulAdd(a, b[nt], acc[nt]);
|
||||
}
|
||||
}
|
||||
SYNC // everyone has read the low-resolution tokens
|
||||
FORW { WINV for (uint nt = 0u; nt < NT; ++nt) STORE_B_PATCH(acc[nt],0u,mt*UC+nt*16u,S4); }
|
||||
SYNC
|
||||
FORW {
|
||||
WINV
|
||||
const uint ly = (tok >> 3) >> 1, lx = (tok & 7u) >> 1, q = ((tok >> 3) & 1u) * 2u + (tok & 1u);
|
||||
const uint sk = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0;
|
||||
for (uint c = 0u; c < HC; c += 4u) writeTToken(tok,ch0+c,SX,readBPatch(ly*4u+lx,q*UC+ch0+c,S4)+skb.a[(sk+c)/4u]);
|
||||
}
|
||||
SYNC // BF (the expanded tokens) is free
|
||||
}
|
||||
#define LOAD_X0(m,c) LOAD_T_TOKEN(m,mt*16u,c,SX)
|
||||
#define READ_X0(r,c) readTToken(r,c,SX)
|
||||
#elif EMBED
|
||||
{
|
||||
MA EA_[TPW];
|
||||
FORW {
|
||||
WINV
|
||||
const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * 16u + hf * 8u;
|
||||
for (uint c = 0u; c < 8u; c += 4u) writeTRaw(tok*SEM+hf*8u+c,inb.a[(src+c)/4u]);
|
||||
LOAD_T_RAW(EA_[w],mt*16u*SEM,SEM);
|
||||
}
|
||||
for (uint nt = 0u; nt < NT; ++nt)
|
||||
{
|
||||
MB b; LD(b, wb.w, p.o[106] + nt * 16u, UC);
|
||||
FORW { WINV STORE_B_TOKEN(coopMatMulAdd(EA_[w],b,bias_tile(p.o[107]+nt*16u,UC)),mt*16u,nt*16u); }
|
||||
}
|
||||
FORW { WINV for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,max(readBToken(tok,ch0+c),f16vec4(0.0))); }
|
||||
SYNC // TT is laid out differently from here on
|
||||
}
|
||||
#define LOAD_X0(m,c) LOAD_B_TOKEN(m,mt*16u,c)
|
||||
#define READ_X0(r,c) readBToken(r,c)
|
||||
#else
|
||||
FORW {
|
||||
WINV
|
||||
const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0;
|
||||
for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,inb.a[(src+c)/4u]);
|
||||
}
|
||||
#define LOAD_X0(m,c) LOAD_B_TOKEN(m,mt*16u,c)
|
||||
#define READ_X0(r,c) readBToken(r,c)
|
||||
#endif
|
||||
FORW {
|
||||
WINV
|
||||
// the residual stream starts as x0 + the output projection's bias
|
||||
for (uint nt = 0u; nt < NT; ++nt) {
|
||||
HACC x; LOAD_X0(x,nt*16u);
|
||||
residual[nt] = x;
|
||||
acc[nt] = bias_tile(p.o[1] + nt * 16u, UC);
|
||||
}
|
||||
vec4 s4 = vec4(0.0);
|
||||
f16vec4 x[C / 8];
|
||||
for (uint c = 0u; c < HC; c += 4u) {
|
||||
x[c / 4u] = READ_X0(tok,ch0+c);
|
||||
#if FOLD_GAINS
|
||||
const vec4 f = vec4(x[c / 4u]); s4 += f * f;
|
||||
#else
|
||||
s4 += vec4(x[c / 4u] * x[c / 4u]);
|
||||
#endif
|
||||
}
|
||||
float s = s4.x + s4.y + s4.z + s4.w;
|
||||
s += subgroupShuffleXor(s, 16u);
|
||||
const float16_t r = float16_t(inversesqrt(float(float16_t(s))));
|
||||
for (uint c = 0u; c < HC; c += 4u) {
|
||||
#if FOLD_GAINS
|
||||
writeBToken(tok,ch0+c,x[c/4u]*r);
|
||||
#else
|
||||
writeBToken(tok,ch0+c,x[c/4u]*(r*wb.w[(p.o[0]+ch0+c)/4u]));
|
||||
#endif
|
||||
}
|
||||
}
|
||||
#if KIND == 2
|
||||
SYNC // x0 lived in TT; the heads overwrite it
|
||||
#endif
|
||||
MA XN_[TPW][C / 16]; // the normalised input of each tile, for every head's Q, K, V
|
||||
#define xn XN_[w]
|
||||
FORW { WINV for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(xn[kt],mt*16u,kt*16u); }
|
||||
// ------------------------------------------------------------ attention, one head at a time
|
||||
for (uint h = 0u; h < ABL_HEADS; ++h)
|
||||
{
|
||||
const uint otab = p.o[7u + 5u * h], owo = p.o[8u + 5u * h];
|
||||
// Q, K, V of every tile, two K steps at a time: their weight tiles are requested before the first product,
|
||||
// and each is loaded once for the wave's tiles (one projection at a time: fewer live accumulators)
|
||||
#define PROJ(WOFF, STORE0, STORE1) { \
|
||||
ACC P0_[TPW], P1_[TPW]; \
|
||||
FORW { P0_[w] = ACC(0.0); P1_[w] = ACC(0.0); } \
|
||||
for (uint kt = 0u; kt < NT; kt += 2u) \
|
||||
{ \
|
||||
MB bw[4]; \
|
||||
for (uint i = 0u; i < 2u; ++i) \
|
||||
{ \
|
||||
LD(bw[2u * i], wb.w, (WOFF) + (kt + i) * 512u, 32u); \
|
||||
LD(bw[2u * i + 1u], wb.w, (WOFF) + (kt + i) * 512u + 16u, 32u); \
|
||||
} \
|
||||
FORW { for (uint i = 0u; i < 2u; ++i) { P0_[w] = coopMatMulAdd(xn[kt + i], bw[2u * i], P0_[w]); P1_[w] = coopMatMulAdd(xn[kt + i], bw[2u * i + 1u], P1_[w]); } } \
|
||||
} \
|
||||
FORW { WINV STORE0; STORE1; } }
|
||||
#if !POSATTN
|
||||
PROJ(p.o[4u + 5u * h], STORE_B_TOKEN(P0_[w],mt*16u,0u), STORE_B_TOKEN(P1_[w],mt*16u,16u)) // Q: this tile's own rows
|
||||
PROJ(p.o[5u + 5u * h], STORE_T_TOKEN(P0_[w],mt*16u,0u,SQ), STORE_T_TOKEN(P1_[w],mt*16u,16u,SQ))
|
||||
#endif
|
||||
PROJ(p.o[6u + 5u * h], STORE_T_TOKEN(P0_[w],VOFF/SQ+mt*16u,0u,SQ), STORE_T_TOKEN(P1_[w],VOFF/SQ+mt*16u,16u,SQ))
|
||||
SYNC
|
||||
MA OA0_[TPW], OA1_[TPW];
|
||||
FORW {
|
||||
WINV
|
||||
#if !POSATTN
|
||||
{
|
||||
ACC s[4];
|
||||
for (uint nt = 0u; nt < 4u; ++nt) { HACC tb; LD(tb, wb.w, otab + mt * 1024u + nt * 16u, 64u); s[nt] = ACC(tb); }
|
||||
MA qa[2]; LOAD_B_TOKEN(qa[0],mt*16u,0u); LOAD_B_TOKEN(qa[1],mt*16u,16u);
|
||||
for (uint kt = 0u; kt < 2u; ++kt)
|
||||
{
|
||||
MB b[4];
|
||||
for (uint nt = 0u; nt < 4u; ++nt) coopMatLoad(b[nt], TT, (nt * 16u * SQ + kt * 16u) / 4u, SQ / 4u, CMJ); // K transposed
|
||||
for (uint nt = 0u; nt < 4u; ++nt) s[nt] = coopMatMulAdd(qa[kt], b[nt], s[nt]);
|
||||
}
|
||||
for (uint nt = 0u; nt < 4u; ++nt) STORE_B_TOKEN(s[nt],mt*16u,nt*16u); // over Q, which is in registers
|
||||
// softmax over the token's 64 scores: 32 per lane
|
||||
f16vec4 e[8];
|
||||
vec4 sum = vec4(0.0);
|
||||
for (uint j = 0u; j < 8u; ++j) { e[j] = softmax_exp(readBToken(tok,hf*32u+j*4u)); sum += vec4(e[j]); }
|
||||
float tot = sum.x + sum.y + sum.z + sum.w;
|
||||
tot += subgroupShuffleXor(tot, 16u);
|
||||
const float16_t rc = float16_t(1.0 / float(float16_t(tot)));
|
||||
for (uint j = 0u; j < 8u; ++j) writeBToken(tok,hf*32u+j*4u,e[j]*rc);
|
||||
}
|
||||
#endif
|
||||
{
|
||||
ACC o0 = ACC(0.0), o1 = ACC(0.0);
|
||||
// two K steps at a time: few live operands keep the register count, and with it the occupancy, in check
|
||||
for (uint k2 = 0u; k2 < 4u; k2 += 2u)
|
||||
{
|
||||
MA a[2]; MB b[4];
|
||||
for (uint i = 0u; i < 2u; ++i)
|
||||
{
|
||||
const uint kt = k2 + i;
|
||||
#if POSATTN
|
||||
LD(a[i], wb.w, otab + mt * 1024u + kt * 16u, 64u);
|
||||
#else
|
||||
LOAD_B_TOKEN(a[i],mt*16u,kt*16u);
|
||||
#endif
|
||||
LOAD_T_TOKEN(b[2u*i],VOFF/SQ+kt*16u,0u,SQ); LOAD_T_TOKEN(b[2u*i+1u],VOFF/SQ+kt*16u,16u,SQ);
|
||||
}
|
||||
for (uint i = 0u; i < 2u; ++i) { o0 = coopMatMulAdd(a[i], b[2u * i], o0); o1 = coopMatMulAdd(a[i], b[2u * i + 1u], o1); }
|
||||
}
|
||||
// the two K tiles of this head for the output projection: through memory, to become A operands
|
||||
STORE_B_TOKEN(o0,mt*16u,0u); STORE_B_TOKEN(o1,mt*16u,16u);
|
||||
LOAD_B_TOKEN(OA0_[w],mt*16u,0u); LOAD_B_TOKEN(OA1_[w],mt*16u,16u);
|
||||
}
|
||||
}
|
||||
for (uint n0 = 0u; n0 < NT; n0 += 2u)
|
||||
{
|
||||
MB b0[2], b1[2];
|
||||
for (uint i = 0u; i < 2u; ++i) { LD(b0[i], wb.w, owo + (n0 + i) * 16u, UC); LD(b1[i], wb.w, owo + 16u * UC + (n0 + i) * 16u, UC); }
|
||||
FORW { for (uint i = 0u; i < 2u; ++i) {
|
||||
acc[n0 + i] = coopMatMulAdd(OA1_[w], b1[i], coopMatMulAdd(OA0_[w], b0[i], acc[n0 + i]));
|
||||
} }
|
||||
}
|
||||
SYNC // every wave is done with K and V before the next head overwrites them
|
||||
}
|
||||
// ------------------------------------------------------------ MLP: its input is the stream itself (gain folded into fc1)
|
||||
{
|
||||
MA M_[TPW][C / 16];
|
||||
#define m M_[w]
|
||||
FORW {
|
||||
WINV
|
||||
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = ACC(HACC(acc[nt]) + residual[nt]);
|
||||
for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u);
|
||||
#if !FOLD_GAINS
|
||||
for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,readBToken(tok,ch0+c)*wb.w[(p.o[2]+ch0+c)/4u]);
|
||||
#endif
|
||||
for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(m[kt],mt*16u,kt*16u);
|
||||
for (uint nt = 0u; nt < NT; ++nt) {
|
||||
acc[nt] = ACC(HACC(acc[nt]) + HACC(bias_tile(p.o[3] + nt * 16u, UC)));
|
||||
}
|
||||
}
|
||||
for (uint ch = 0u; ch < ABL_CHUNKS; ++ch)
|
||||
{
|
||||
const uint ow1 = p.o[44u + 3u * ch], ob1 = p.o[45u + 3u * ch], ow2 = p.o[46u + 3u * ch];
|
||||
ACC H0_[TPW], H1_[TPW];
|
||||
FORW { H0_[w] = bias_tile(ob1, 32u); H1_[w] = bias_tile(ob1 + 16u, 32u); }
|
||||
for (uint kt = 0u; kt < NT; kt += 2u)
|
||||
{
|
||||
MB b[4];
|
||||
for (uint i = 0u; i < 2u; ++i) { LD(b[2u * i], wb.w, ow1 + (kt + i) * 512u, 32u); LD(b[2u * i + 1u], wb.w, ow1 + (kt + i) * 512u + 16u, 32u); }
|
||||
FORW { for (uint i = 0u; i < 2u; ++i) { H0_[w] = coopMatMulAdd(m[kt + i], b[2u * i], H0_[w]); H1_[w] = coopMatMulAdd(m[kt + i], b[2u * i + 1u], H1_[w]); } }
|
||||
}
|
||||
FORW { WINV STORE_T_TOKEN(H0_[w],mt*16u,0u,SQ); STORE_T_TOKEN(H1_[w],mt*16u,16u,SQ); }
|
||||
// fc2's first weight tiles are requested before the activation pass, so they arrive while it runs
|
||||
MB w2a[2], w2b[2];
|
||||
for (uint i = 0u; i < 2u; ++i) { LD(w2a[i], wb.w, ow2 + i * 16u, UC); LD(w2b[i], wb.w, ow2 + 16u * UC + i * 16u, UC); }
|
||||
MA GA0_[TPW], GA1_[TPW];
|
||||
FORW {
|
||||
WINV
|
||||
for (uint j = 0u; j < 4u; ++j) writeTToken(tok,hf*16u+j*4u,SQ,gelu(readTToken(tok,hf*16u+j*4u,SQ)));
|
||||
LOAD_T_TOKEN(GA0_[w],mt*16u,0u,SQ); LOAD_T_TOKEN(GA1_[w],mt*16u,16u,SQ);
|
||||
}
|
||||
FORW { for (uint i = 0u; i < 2u; ++i) acc[i] = coopMatMulAdd(GA1_[w], w2b[i], coopMatMulAdd(GA0_[w], w2a[i], acc[i])); }
|
||||
for (uint n0 = 2u; n0 < NT; n0 += 2u)
|
||||
{
|
||||
for (uint i = 0u; i < 2u; ++i) { LD(w2a[i], wb.w, ow2 + (n0 + i) * 16u, UC); LD(w2b[i], wb.w, ow2 + 16u * UC + (n0 + i) * 16u, UC); }
|
||||
FORW { for (uint i = 0u; i < 2u; ++i) acc[n0 + i] = coopMatMulAdd(GA1_[w], w2b[i], coopMatMulAdd(GA0_[w], w2a[i], acc[n0 + i])); }
|
||||
}
|
||||
}
|
||||
FORW { WINV for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u); }
|
||||
#undef m
|
||||
}
|
||||
// ------------------------------------------------------------ outputs (the block output is in BF, token layout)
|
||||
#define OUTV const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W; const uint dst = uint(ty * p.W + tx);
|
||||
#if HEAD
|
||||
SYNC // the head's rows of TT overlap other tiles' MLP rows
|
||||
{
|
||||
ACC HO_[TPW][3];
|
||||
FORW { for (uint nt = 0u; nt < 3u; ++nt) HO_[w][nt] = bias_tile(p.o[107] + nt * 16u, 48u); }
|
||||
for (uint kt = 0u; kt < NT; ++kt)
|
||||
{
|
||||
MB b[3];
|
||||
for (uint nt = 0u; nt < 3u; ++nt) LD(b[nt], wb.w, p.o[106] + kt * 16u * 48u + nt * 16u, 48u);
|
||||
FORW {
|
||||
WINV
|
||||
MA a; LOAD_B_TOKEN(a,mt*16u,kt*16u);
|
||||
for (uint nt = 0u; nt < 3u; ++nt) HO_[w][nt] = coopMatMulAdd(a, b[nt], HO_[w][nt]);
|
||||
}
|
||||
}
|
||||
FORW {
|
||||
WINV OUTV
|
||||
for (uint nt = 0u; nt < 3u; ++nt) STORE_T_TOKEN(HO_[w][nt],mt*16u,nt*16u,SHD);
|
||||
if (inside)
|
||||
for (uint c = hf * 20u; c < hf * 20u + 20u; c += 4u) out0.a[(dst * 40u + c) / 4u] = readTToken(tok,c,SHD);
|
||||
}
|
||||
}
|
||||
#elif KIND == 1
|
||||
FORW {
|
||||
WINV OUTV
|
||||
if (inside)
|
||||
for (uint c = 0u; c < HC; c += 4u) out1.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c);
|
||||
}
|
||||
#else
|
||||
FORW {
|
||||
WINV OUTV
|
||||
if (inside)
|
||||
for (uint c = 0u; c < HC; c += 4u) out0.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c);
|
||||
}
|
||||
#endif
|
||||
#if KIND == 1
|
||||
SYNC
|
||||
{
|
||||
// patch merge: tile mt gathers sub-token mt of the 16 merged tokens into TT, then computes output tiles mt, mt + 4, ...
|
||||
FORW {
|
||||
WINV
|
||||
const uint sourceToken = 8u * (2u * (row >> 2) + (mt >> 1)) + 2u * (row & 3u) + (mt & 1u);
|
||||
for (uint c = 0u; c < HC; c += 4u)
|
||||
writeTRaw(row * S4 + mt * UC + ch0 + c, readBToken(sourceToken, ch0 + c));
|
||||
}
|
||||
SYNC
|
||||
const uint MN = uint(MERGE_N), NI = (MN / 16u + 3u) / 4u;
|
||||
ACC MG_[TPW][(MERGE_N / 16 + 3) / 4];
|
||||
FORW {
|
||||
WINV
|
||||
for (uint i = 0u; i < NI; ++i) MG_[w][i] = bias_tile(p.o[105] + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN);
|
||||
for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u)
|
||||
{
|
||||
MA a[2]; MB b[2 * ((MERGE_N / 16 + 3) / 4)];
|
||||
for (uint j = 0u; j < 2u; ++j)
|
||||
{
|
||||
LOAD_T_RAW(a[j],(kt+j)*16u,S4);
|
||||
for (uint i = 0u; i < NI; ++i) LD(b[j * NI + i], wb.w, p.o[104] + (kt + j) * 16u * MN + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN);
|
||||
}
|
||||
for (uint j = 0u; j < 2u; ++j)
|
||||
for (uint i = 0u; i < NI; ++i) MG_[w][i] = coopMatMulAdd(a[j], b[j * NI + i], MG_[w][i]);
|
||||
}
|
||||
}
|
||||
// BF is free (every wave gathered before the barrier above, and every tile's products are done): the merged
|
||||
// tokens go there
|
||||
FORW {
|
||||
WINV
|
||||
const uint ntn = (MN / 16u + 3u - mt) / 4u;
|
||||
for (uint i = 0u; i < NI; ++i)
|
||||
if (i < ntn) STORE_B_PATCH(MG_[w][i],0u,(mt+4u*i)*16u,SMG);
|
||||
}
|
||||
SYNC
|
||||
const int ym = tdiv2(Y0) + int(row >> 2), xm = tdiv2(X0) + int(row & 3u);
|
||||
const uint part = uint(XC) / 8u;
|
||||
FORW {
|
||||
WINV
|
||||
const uint c0 = (mt * 2u + hf) * part;
|
||||
if (ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2)
|
||||
for (uint c = c0; c < c0 + part; c += 4u) out0.a[(uint(ym * (p.W / 2) + xm) * uint(XC) + c) / 4u] = readBPatch(row,c,SMG);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
Reference in new issue
Block a user