wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

engine: this GPU's own arithmetic and shared weight tiles as the builders' defaults

Browse Source 
- 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
This commit is contained in:
woutermans committed 2026-10-14 00:52:41 +02:00
1 parent 8eb7a94533
commit f0b55a850a
24 files changed
+1108 -254
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/layer_m.comp
+163 -68
View File
@@ -57,7 +57,17 @@
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#extension GL_EXT_control_flow_attributes : require
// WPW: 4x4-token windows per wave (1, 2 or 4). A workgroup is always one 8x8-token block of four windows; with
// WPW > 1 a wave walks WPW of them and every weight tile is loaded once for all of them (a window reads the whole
// layer's weights, which is a third of the network's time with one window per wave). Values do not change.
#ifndef WPW
#define WPW 1
#endif
#if WPW > 1 && KIND >= 3
#error layer_m.comp: L's input and expansion stages stage their pixels per thread (WPW = 1)
#endif
layout(local_size_x = 128 / WPW) 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;
@@ -144,6 +154,17 @@ const float16_t MCL = float16_t(6.2831855);
// 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).
// M_RELAXED (M_DEFS=-DM_RELAXED for compile_m.py): this GPU's own arithmetic instead. The accumulator stays f32
// across the 32-deep steps and is rounded once where it is used, with no scratch round trip; a layer is then no
// longer bit-identical to NVIDIA's kernel (FP8 codes flip by one step here and there), which two NVIDIA setups
// are not either (engine/rr/README.md).
#ifdef M_RELAXED
#define R16(m) { }
#define K32H(dst, a0, b0, a1, b1) { \
ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(dst))); \
(dst)[0] = float16_t(P_[0]); (dst)[1] = float16_t(P_[1]); (dst)[2] = float16_t(P_[2]); (dst)[3] = float16_t(P_[3]); \
(dst)[4] = float16_t(P_[4]); (dst)[5] = float16_t(P_[5]); (dst)[6] = float16_t(P_[6]); (dst)[7] = float16_t(P_[7]); }
#else
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])); \
@@ -152,6 +173,7 @@ shared f16vec4 RS[128 * 2];
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); }
#endif
// 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
@@ -159,7 +181,8 @@ shared f16vec4 RS[128 * 2];
// 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)
#if TEXL
#if defined(M_RELAXED)
#elif TEXL
// ZLUDA's lowering of m16n8k32 (K32_NO_MIX): the accumulator enters the first product as C, f32 throughout the step,
// one binary16 rounding at its end
#define K32H(dst, a0, b0, a1, b1) { \
@@ -186,11 +209,32 @@ shared f16vec4 RS[128 * 2];
#endif
// 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; }
#ifdef M_RELAXED
// straight from f32 into FP8, two elements per saturating conversion; the activation in f32 as well
#define SAT2_(dst, i, v) { fe4m3vec2 t_; saturatedConvertEXT(t_, (v)); (dst)[i] = t_.x; (dst)[(i) + 1] = t_.y; }
#define MK8_F8(dst, a) { SAT2_(dst, 0, vec2((a)[0], (a)[1])) SAT2_(dst, 2, vec2((a)[2], (a)[3])) SAT2_(dst, 4, vec2((a)[4], (a)[5])) SAT2_(dst, 6, vec2((a)[6], (a)[7])) }
#define GELU2_(dst, i, h) { const vec2 x_ = vec2((h)[i], (h)[(i) + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); \
const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.08108133), vec2(0.41216981)), vec2(0.5)); SAT2_(dst, i, x_ * g_) }
#define GELU8(dst, h) { GELU2_(dst, 0, h) GELU2_(dst, 2, h) GELU2_(dst, 4, h) GELU2_(dst, 6, h) }
#else
#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_); }
#endif
// the normalised tokens into FP8 (inside NM1 / N2: x0..x7 the token, g0 / g1 the gains)
#ifdef M_RELAXED
#define H1N(N, r_) { const float f_ = float(r_); \
SAT2_(h1[N], 0, vec2(x0, x1) * (f_ * vec2(g0.xy))) SAT2_(h1[N], 2, vec2(x2, x3) * (f_ * vec2(g0.zw))) \
SAT2_(h1[N], 4, vec2(x4, x5) * (f_ * vec2(g1.xy))) SAT2_(h1[N], 6, vec2(x6, x7) * (f_ * vec2(g1.zw))) }
#else
#define H1N(N, r_) { \
h1[N][0] = floate4m3_t(CL8(x0 * float16_t(r_ * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(r_ * g0.y))); \
h1[N][2] = floate4m3_t(CL8(x2 * float16_t(r_ * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(r_ * g0.w))); \
h1[N][4] = floate4m3_t(CL8(x4 * float16_t(r_ * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(r_ * g1.y))); \
h1[N][6] = floate4m3_t(CL8(x6 * float16_t(r_ * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(r_ * g1.w))); }
#endif
// 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)
{
@@ -364,18 +408,24 @@ float tanhL(float t)
#endif
void main()
{
const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u;
const uint mt0 = (gl_LocalInvocationID.x >> 5) * uint(WPW), 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;
// per window w of this wave (FORW): its index mt in the block and its token's position (WINV)
#define FORW [[unroll]] for (uint w = 0u; w < uint(WPW); ++w)
#define WINV const uint mt = mt0 + w; \
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
HAC RES_[WPW][NT]; // the residual stream: rows are channels, columns are tokens
B8 H1_[WPW][NT]; // the normalised activation, the K operand of every product
#define res RES_[w]
#define h1 H1_[w]
// ------------------------------------------------------------ patch expand: x0 = q8(E) + skip
#if EXPAND
{
FORW {
WINV
const uint W2 = uint(p.W) / 2u, H2 = uint(p.Hh) / 2u, part = UXC / 8u, c0 = (mt * 2u + hf) * part;
const uint lt = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2));
#if PLANES
@@ -383,7 +433,10 @@ void main()
#else
for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lt * UXC + c];
#endif
SYNC
}
SYNC
FORW {
WINV
// 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);
@@ -401,12 +454,13 @@ void main()
EACH_T(EF_K32)
}
EACH_T(EF_STORE)
SYNC
}
SYNC
#endif
#if KIND == 3
// ------------------------------------------------------------ L's input: features, embedding -> x0
{
FORW {
WINV
const float e1 = ep.expo == 0.0 ? 1.0 : ep.expo, e2 = ep.expoPrev == 0.0 ? 1.0 : ep.expoPrev;
const uint li = gl_LocalInvocationID.x;
if (li < 64u)
@@ -445,6 +499,8 @@ void main()
}
#endif
// ------------------------------------------------------------ norm 1: sum of squares, residual and h1
FORW {
WINV
#if EXPAND
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;
@@ -483,47 +539,74 @@ void main()
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))); }
H1N(N, rn) }
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) { \
// Q of every window, then V of every window: each weight tile is loaded once for the wave's windows
A8 QA_[WPW][2];
{
HAC q[WPW][2];
FORW { q[w][0] = HAC(0.0); q[w][1] = HAC(0.0); }
#define Q_G(A, B) { \
const uint kt = A; \
A8 aq0, aq1; \
LDA8(aq0, oq, kt, 32u, 0u); \
LDA8(aq1, oq, kt + 1u, 32u, 0u); \
FORW { K32H(q[w][0], aq0, h1[A], aq1, h1[B]) } \
LDA8(aq0, oq, kt, 32u, 16u); \
LDA8(aq1, oq, kt + 1u, 32u, 16u); \
FORW { K32H(q[w][1], aq0, h1[A], aq1, h1[B]) } }
EACH_T16(Q_G)
FORW {
for (uint m = 0u; m < 2u; ++m)
{
#ifdef M_RELAXED
MK8_F8(QA_[w][m], q[w][m])
#else
const float16_t q0 = CL8(q[w][m][0]), q1 = CL8(q[w][m][1]), q2 = CL8(q[w][m][2]), q3 = CL8(q[w][m][3]);
const float16_t q4 = CL8(q[w][m][4]), q5 = CL8(q[w][m][5]), q6 = CL8(q[w][m][6]), q7 = CL8(q[w][m][7]);
QA_[w][m][0] = floate4m3_t(q0); QA_[w][m][1] = floate4m3_t(q1);
QA_[w][m][2] = floate4m3_t(q2); QA_[w][m][3] = floate4m3_t(q3);
QA_[w][m][4] = floate4m3_t(q4); QA_[w][m][5] = floate4m3_t(q5);
QA_[w][m][6] = floate4m3_t(q6); QA_[w][m][7] = floate4m3_t(q7);
#endif
}
}
}
{
HAC v[WPW][2];
FORW { v[w][0] = HAC(0.0); v[w][1] = HAC(0.0); }
#define V_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); \
A8 av0, av1; \
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); \
FORW { K32H(v[w][0], av0, h1[A], av1, h1[B]) } \
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];
FORW { K32H(v[w][1], av0, h1[A], av1, h1[B]) } }
EACH_T16(V_G)
// V of every window goes to its shared buffer now, so it does not occupy registers during the scores
FORW {
const uint mt = mt0 + w;
coopMatStore(v[w][0], VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
coopMatStore(v[w][1], VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
}
// 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])
}
// scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias
// (read column-major, so the accumulator holds S^T: this lane holds the keys 8h..8h+7 of one query)
const HAC sbias = bias16(ob, 16u);
B8 OB_[WPW][2];
FORW {
WINV
B8 qb[2];
for (uint m = 0u; m < 2u; ++m) for (int i = 0; i < 8; ++i) qb[m][i] = QA_[w][m][i];
HAC s = sbias;
K32H(s, QA_[w][0], qb[0], QA_[w][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]);
@@ -540,8 +623,6 @@ void main()
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);
@@ -554,20 +635,22 @@ void main()
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]);
MK8_F8(OB_[w][0], o[0]); MK8_F8(OB_[w][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]) }
FORW { K32H(res[N], a0, OB_[w][0], a1, OB_[w][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.
FORW {
WINV
#if TEXL
#define EV(T, I) float16_t(res[T][I]) // natural order: the lane's own elements
#else
@@ -592,50 +675,55 @@ void main()
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))); \
H1N(N, rn2) \
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);
HAC hm[WPW][2];
{ const HAC b0_ = bias16(ob1, 32u), b1_ = bias16(ob1 + 16u, 32u); FORW { hm[w][0] = b0_; hm[w][1] = b1_; } }
#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]) \
FORW { K32H(hm[w][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]) }
FORW { K32H(hm[w][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)); }
B8 G0_[WPW], G1_[WPW];
FORW {
#ifdef M_RELAXED
GELU8(G0_[w], hm[w][0]) GELU8(G1_[w], hm[w][1])
#else
{ const float16_t v0 = gelu1(hm[w][0][0]), v1 = gelu1(hm[w][0][1]), v2 = gelu1(hm[w][0][2]), v3 = gelu1(hm[w][0][3]);
const float16_t v4 = gelu1(hm[w][0][4]), v5 = gelu1(hm[w][0][5]), v6 = gelu1(hm[w][0][6]), v7 = gelu1(hm[w][0][7]);
G0_[w][0] = floate4m3_t(CL8(v0)); G0_[w][1] = floate4m3_t(CL8(v1)); G0_[w][2] = floate4m3_t(CL8(v2)); G0_[w][3] = floate4m3_t(CL8(v3));
G0_[w][4] = floate4m3_t(CL8(v4)); G0_[w][5] = floate4m3_t(CL8(v5)); G0_[w][6] = floate4m3_t(CL8(v6)); G0_[w][7] = floate4m3_t(CL8(v7)); }
{ const float16_t v0 = gelu1(hm[w][1][0]), v1 = gelu1(hm[w][1][1]), v2 = gelu1(hm[w][1][2]), v3 = gelu1(hm[w][1][3]);
const float16_t v4 = gelu1(hm[w][1][4]), v5 = gelu1(hm[w][1][5]), v6 = gelu1(hm[w][1][6]), v7 = gelu1(hm[w][1][7]);
G1_[w][0] = floate4m3_t(CL8(v0)); G1_[w][1] = floate4m3_t(CL8(v1)); G1_[w][2] = floate4m3_t(CL8(v2)); G1_[w][3] = floate4m3_t(CL8(v3));
G1_[w][4] = floate4m3_t(CL8(v4)); G1_[w][5] = floate4m3_t(CL8(v5)); G1_[w][6] = floate4m3_t(CL8(v6)); G1_[w][7] = floate4m3_t(CL8(v7)); }
#endif
}
#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) }
FORW { K32H(res[N], a0, G0_[w], a1, G1_[w]) } }
EACH_T(MLP2_N)
}
#if KIND == 4
// ------------------------------------------------------------ L's expansion head: per-pixel reconstruction inputs
{
FORW {
WINV
const uint hb = (mt * 16u + row) * HS + ch8;
for (uint nt = 0u; nt < NT; ++nt)
for (int i = 0; i < 8; ++i) HO[hb + 16u * nt + uint(i)] = float16_t(res[nt][i]);
@@ -661,6 +749,8 @@ void main()
}
#else
// ------------------------------------------------------------ the block output, FP8
FORW {
WINV
const uint dst = uint(ty * p.W + tx) * UC + ch8;
if (inside)
for (uint nt = 0u; nt < NT; ++nt)
@@ -669,15 +759,20 @@ void main()
OUT8(0) OUT8(1) OUT8(2) OUT8(3) OUT8(4) OUT8(5) OUT8(6) OUT8(7)
#undef OUT8
}
}
#endif
// ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution
#if MERGE
{
FORW {
WINV
// 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
}
SYNC
FORW {
WINV
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;