wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
0e851b7f904f1c3b1bb7d3f281d140ee53899422
d4r/engine/layer.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

482 lines
24 KiB
Plaintext
Raw Blame History

#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 View Git Blame Copy Permalink