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,7 @@
|
||||
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
|
||||
#define SWIN_NO_Q8
|
||||
#define SWIN_F32ACC
|
||||
#define SWIN_MAX_VGPR 128
|
||||
#define SWIN_MLP_R(NW) 2
|
||||
#include "swin_block.h"
|
||||
SWIN_MODULE_W(rrlite_dec1_4x4, 64, 2, 8, 0, false, CommonParams, 96)
|
||||
@@ -0,0 +1,7 @@
|
||||
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
|
||||
#define SWIN_NO_Q8
|
||||
#define SWIN_F32ACC
|
||||
#define SWIN_MAX_VGPR 192
|
||||
#define SWIN_MLP_R(NW) 2
|
||||
#include "swin_block.h"
|
||||
SWIN_MODULE_W(rrlite_dec2_4x4, 96, 4, 8, 0, false, CommonParams, 128)
|
||||
@@ -0,0 +1,8 @@
|
||||
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
|
||||
#define SWIN_NO_Q8
|
||||
#define SWIN_F32ACC
|
||||
#define SWIN_MAX_VGPR 192
|
||||
#define SWIN_MLP_R(NW) 1
|
||||
#define SWIN_SCALAR_BLOAD // scalar-base weight loads: faster here (no spills), slower in the 8-wave layers
|
||||
#include "swin_block.h"
|
||||
SWIN_MODULE_W(rrlite_enc1_4x4, 64, 2, 4, 96, false, CommonParams, 0)
|
||||
@@ -0,0 +1,7 @@
|
||||
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
|
||||
#define SWIN_NO_Q8
|
||||
#define SWIN_F32ACC
|
||||
#define SWIN_MAX_VGPR 192
|
||||
#define SWIN_MLP_R(NW) 2
|
||||
#include "swin_block.h"
|
||||
SWIN_MODULE_W(rrlite_enc2_4x4, 96, 4, 8, 128, false, CommonParams, 0)
|
||||
@@ -0,0 +1,8 @@
|
||||
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
|
||||
#define SWIN_NO_Q8
|
||||
#define SWIN_F32ACC
|
||||
#define SWIN_MAX_VGPR 192
|
||||
#define SWIN_MLP_R(NW) 2
|
||||
#define SWIN_PREP_SLOTS 8 // one weight image per tube block (six per frame)
|
||||
#include "swin_block.h"
|
||||
SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, 8, 0, true, TubeParams, 0)
|
||||
@@ -0,0 +1,897 @@
|
||||
// Templated rrlite Swin block (64 tokens = 4 windows of 4x4 per block, NW waves = heads of 32 dims).
|
||||
// C channels (64 / 96 / 128), NW waves (2 or 4)
|
||||
// NPM patch-merge output channels (0 = none): 2x2 merge of the block output, 4C -> NPM
|
||||
// TUBE 88-byte parameter block with the cross-launch flag protocol
|
||||
// CIN patch-expand input channels (0 = none): decoder input = f16(q8(expand(low-res CIN)) + skip)
|
||||
#pragma once
|
||||
#include "swin_common.h"
|
||||
#ifdef D4R_SWIN_DIAGNOSTICS
|
||||
extern "C" {
|
||||
__device__ int d4r_swin_debug_block[2] = {0, 4};
|
||||
__device__ half_t d4r_swin_debug_values[4 * 64 * 128];
|
||||
}
|
||||
template <int C, int AST> __device__ __forceinline__ void swin_snapshot(
|
||||
const half_t* activation, int stage, int bx, int by, int tid, bool pair_order) {
|
||||
if (bx == d4r_swin_debug_block[0] && by == d4r_swin_debug_block[1] && tid < 64)
|
||||
for (int c = 0; c < C; ++c)
|
||||
d4r_swin_debug_values[(stage * 64 + tid) * 128 + c] =
|
||||
activation[tid * AST + (pair_order ? gperm(c) : apos(c))];
|
||||
__syncthreads();
|
||||
}
|
||||
#endif
|
||||
|
||||
// Immediate lane masks need only a compile-time integer, not the host C++
|
||||
// standard library (which is absent from device-only Windows HIP builds).
|
||||
template <int Mask> struct SwinLaneMask { static constexpr int value = Mask; };
|
||||
|
||||
template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
|
||||
{
|
||||
static constexpr int PREB = CIN ? 4 * C * CIN + 8 * C : 0; // patch-expand weights + bias first
|
||||
static constexpr int NTE = CIN ? 4 * C / NW / 16 : 0;
|
||||
static constexpr int NPL = C / 32, NT = C / 16, NT8 = C / 8, INNER = 32 * NW, NC = C / 8;
|
||||
static constexpr int HS = 96 * C + 512;
|
||||
static constexpr int G1 = PREB, E = PREB + 2 * C + NW * HS, BO = E, G2 = E + 2 * C, M0 = E + 4 * C;
|
||||
static constexpr int B2 = M0 + 64 * C + 64;
|
||||
__host__ __device__ static constexpr int head(int h) { return PREB + 2 * C + HS * h; }
|
||||
__host__ __device__ static constexpr int b1(int c) { return c == 0 ? M0 : M0 + 66 * C + 64 + (64 * C + 64) * (c - 1); }
|
||||
static constexpr int PM0 = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1);
|
||||
static constexpr int PMB = PM0 + 4 * C * NPM;
|
||||
// weight image (slots of 16 K values, see expand_weights)
|
||||
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
|
||||
static constexpr int WO = 4 * C * NW;
|
||||
static constexpr int W1 = WO + 2 * NW * C;
|
||||
__host__ __device__ static constexpr int w1(int c) { return W1 + 4 * C * c; }
|
||||
__host__ __device__ static constexpr int w2(int c) { return W1 + 4 * C * c + 2 * C; }
|
||||
static constexpr int PM = W1 + 4 * C * NC;
|
||||
__host__ __device__ static constexpr int pm(int g) { return PM + 8 * C * g; }
|
||||
static constexpr int PE = PM + (NPM / 32) * 8 * C;
|
||||
__host__ __device__ static constexpr int pe(int w) { return PE + w * (CIN / 32) * 2 * NTE * 16; }
|
||||
static constexpr int TOTAL = PE + (CIN ? NW * (CIN / 32) * 2 * NTE * 16 : 0);
|
||||
static constexpr int NDESC = 2 * NW + 1 + 2 * NC + NPM / 32 + (CIN ? NW : 0);
|
||||
};
|
||||
|
||||
// Start (in halves) of token T's row in the output staging buffer: stride C + 4 plus a shift by bits 4 and 5
|
||||
// of T. The patch merge's lanes read 8 bytes each from 16 tokens that differ in bits 1, 3, 4 and 5 (plus a
|
||||
// constant); this placement puts those 16 reads on 16 distinct 8-byte bank pairs for C = 64, 96 and 128.
|
||||
template <int C> __host__ __device__ constexpr int xo_row(int T)
|
||||
{
|
||||
return (C + 4) * T + 4 * ((T >> 4) & 1) + 16 * (T >> 5);
|
||||
}
|
||||
|
||||
template <int C, int NW, int NPM, int CIN> struct SwinDescs
|
||||
{
|
||||
GemmDesc d[SwinLayout<C, NW, NPM, CIN>::NDESC];
|
||||
};
|
||||
|
||||
template <int C, int NW, int NPM, int CIN> constexpr SwinDescs<C, NW, NPM, CIN> make_descs()
|
||||
{
|
||||
using L = SwinLayout<C, NW, NPM, CIN>;
|
||||
SwinDescs<C, NW, NPM, CIN> r{};
|
||||
int n = 0;
|
||||
for (int h = 0; h < NW; ++h)
|
||||
for (int q = 0; q < 2; ++q)
|
||||
r.d[n++] = GemmDesc{L::qv(h, q), L::head(h) + (q ? 32 * C : 0), 1024, 512, C / 32, 2};
|
||||
r.d[n++] = GemmDesc{L::WO, L::head(0) + 64 * C + 512, L::HS, 512, NW, C / 16};
|
||||
for (int c = 0; c < L::NC; ++c)
|
||||
{
|
||||
r.d[n++] = GemmDesc{L::w1(c), L::b1(c), 512, 16 * C, C / 32, 2};
|
||||
r.d[n++] = GemmDesc{L::w2(c), L::b1(c) + 32 * C + 64, 0, 512, 1, C / 16};
|
||||
}
|
||||
for (int g = 0; g < NPM / 32; ++g)
|
||||
r.d[n++] = GemmDesc{L::pm(g), L::PM0 + g * 128 * C, 512, 64 * C, C / 8, 2};
|
||||
for (int w = 0; w < (CIN ? NW : 0); ++w)
|
||||
r.d[n++] = GemmDesc{L::pe(w), w * CIN * 4 * C / NW, 512, 16 * CIN, CIN / 32, L::NTE};
|
||||
return r;
|
||||
}
|
||||
|
||||
// NH heads (NVIDIA's warps per block) run on NW waves (NW = NH, or 2 NH to halve LDS per wave)
|
||||
template <int C, int NH, int NW, int NPM, bool TUBE, int CIN>
|
||||
__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const int bx,
|
||||
const int by, const int gx)
|
||||
{
|
||||
using L = SwinLayout<C, NH, NPM, CIN>;
|
||||
static_assert(NW % NH == 0 && NW / NH <= 4, "waves per head");
|
||||
constexpr int NPL = L::NPL, NT = L::NT, INNER = L::INNER;
|
||||
constexpr int AST = (C > INNER ? C : INNER) + 8;
|
||||
constexpr int GST = 40, PST = 24;
|
||||
#ifdef SWIN_EXACT_PV
|
||||
constexpr int QST = 56;
|
||||
#else
|
||||
constexpr int QST = 40;
|
||||
#endif
|
||||
#ifndef SWIN_MLP_R
|
||||
#define SWIN_MLP_R(NW) ((NW) / 2)
|
||||
#endif
|
||||
constexpr int R = SWIN_MLP_R(NW); // MLP chunks per round
|
||||
constexpr int MG = (NT % NW == 0) ? 1 : ((2 * NT) % NW == 0 ? 2 : 4), MT = 4 / MG, NTW = NT * MG / NW;
|
||||
static_assert(NTW * NW == NT * MG, "column split");
|
||||
constexpr int PAST = CIN + 8;
|
||||
constexpr int SCR = NW * 16 * QST * 2 > (CIN ? 16 * PAST * 2 : 0) ? NW * 16 * QST * 2 : 16 * PAST * 2;
|
||||
constexpr int U_ST4 = R * 64 * GST * 2;
|
||||
constexpr int USIZE = SCR > U_ST4 ? SCR : U_ST4;
|
||||
|
||||
__shared__ __attribute__((aligned(16))) half_t A[64 * AST];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t U[USIZE];
|
||||
__shared__ half_t R2[64];
|
||||
half_t* X0 = A; // stages 0-1: block input f16 [64][AST], natural order
|
||||
half_t* QP = (half_t*)U; // stage 2: per-wave Q / P scratch [NW][16][QST]
|
||||
half_t* PA = QP; // stage 0: low-res tokens f16 [16][PAST] (apos order)
|
||||
half_t* G = (half_t*)U; // stage 4: fc1 outputs [R][64][GST]
|
||||
|
||||
const int lane = threadIdx.x, wv = threadIdx.z, tid = lane + 32 * wv;
|
||||
const int l16 = lane & 15, hi = lane >> 4;
|
||||
const uint8_t* W = p.w;
|
||||
|
||||
if constexpr (TUBE)
|
||||
{
|
||||
if (tp->in_flags != nullptr)
|
||||
{
|
||||
if (lane <= 3 && wv == 0)
|
||||
{
|
||||
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
|
||||
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
|
||||
{
|
||||
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
|
||||
// bounded spin: a missing producer must never hang the GPU
|
||||
for (int it = 0; it < (1 << 22); ++it)
|
||||
{
|
||||
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
|
||||
break;
|
||||
__builtin_amdgcn_s_sleep(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
|
||||
}
|
||||
}
|
||||
|
||||
// block input words (encoders: fp8 input, decoders: fp8 skip), issued first so their latency
|
||||
// overlaps the patch expand; word k of thread tid is idx = tid + 32 NW k
|
||||
constexpr int NSW = 64 * C / 4 / (32 * NW);
|
||||
static_assert(NSW * 32 * NW == 64 * C / 4, "input words per thread");
|
||||
uint32_t inw[NSW];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < NSW; ++k)
|
||||
{
|
||||
const int idx = tid + 32 * NW * k;
|
||||
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
|
||||
const int X = mirror(8 * bx - p.sx + tok_x(T), p.tw), Y = mirror(8 * by - p.sy + tok_y(T), p.th);
|
||||
const uint8_t* src = CIN ? p.p32 : p.in;
|
||||
inw[k] = *(const uint32_t*)(src + ((size_t)((c4 >> 5) * p.th + Y) * p.tw + X) * 32 + (c4 & 31));
|
||||
}
|
||||
// the input words as f16 into X0 (natural channel order); decoders add the patch expand to them later
|
||||
auto stage_input_words = [&]() {
|
||||
#pragma unroll
|
||||
for (int k = 0; k < NSW; ++k)
|
||||
{
|
||||
const int idx = tid + 32 * NW * k;
|
||||
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
|
||||
const uint32_t bytes = inw[k];
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
|
||||
}
|
||||
};
|
||||
|
||||
// ------------------------------------------------ stage 0: block input x0 (f16) into A
|
||||
// encoders: the fp8 input; decoders: q8(patch expand(low-res)) + skip
|
||||
if constexpr (CIN > 0)
|
||||
{
|
||||
const int W2 = p.tw / 2, H2 = p.th / 2;
|
||||
for (int idx = tid; idx < 16 * CIN / 4; idx += 32 * NW)
|
||||
{
|
||||
const int r = idx / (CIN / 4), c4 = 4 * (idx % (CIN / 4));
|
||||
const int MX = mirror((8 * bx - p.sx) / 2 + (r & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (r >> 2), H2);
|
||||
const uint32_t bytes =
|
||||
*(const uint32_t*)(p.in + ((size_t)((c4 >> 5) * H2 + MY) * W2 + MX) * 32 + (c4 & 31));
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
|
||||
}
|
||||
stage_input_words(); // the skip input
|
||||
__syncthreads();
|
||||
constexpr int NTEW = 4 * C / NW / 16; // expand tiles per wave (L::NTE per weight block)
|
||||
T16 e[NTEW];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTEW; ++j)
|
||||
e[j] = t16_splat(hload(W, 4 * C * CIN + 2 * (16 * (NTEW * wv + j) + l16)));
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < CIN / 32; ++kc)
|
||||
{
|
||||
const gop a0 = lda(PA + l16 * PAST + 32 * kc), a1 = lda(PA + l16 * PAST + 32 * kc + 16);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTEW; ++j)
|
||||
{
|
||||
const int t = NTEW * wv + j, blk = t / L::NTE, lt = t % L::NTE;
|
||||
k32(e[j], a0, bload(w16, L::pe(blk), L::NTE, kc, 0, lt, l16), a1, bload(w16, L::pe(blk), L::NTE, kc, 1, lt, l16));
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTEW; ++j)
|
||||
{
|
||||
const int Gc = 16 * (NTEW * wv + j) + l16, q = Gc / C, ch = ginv(Gc % C);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int r = mrow(i);
|
||||
const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
|
||||
const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
|
||||
// e4m3_to_half is exact in f16, so adding the staged skip value is the same sum
|
||||
X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + X0[T * AST + ch]);
|
||||
}
|
||||
}
|
||||
}
|
||||
else
|
||||
stage_input_words();
|
||||
__syncthreads();
|
||||
|
||||
// residual accumulators of the output projection: f16(x0 + b_o) in the C-fragment layout
|
||||
// (wave owns pair columns of n tiles ng*NTW.. for m tiles mg*MT..)
|
||||
const int ng = wv % (NW / MG), mg = wv / (NW / MG);
|
||||
T16 x[MT][NTW]; // packed f16 accumulators: residual f16(x0 + b_o), live through stage 2
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MT; ++mm)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
{
|
||||
const int n = 16 * (ng * NTW + j) + l16;
|
||||
const half_t bo = hload(W, L::BO + 2 * n);
|
||||
const int ch = ginv(n);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i += 2)
|
||||
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + mrow(i)) * AST + ch],
|
||||
X0[(16 * (mg * MT + mm) + mrow(i + 1)) * AST + ch]} + (hv2){bo, bo});
|
||||
}
|
||||
__syncthreads(); // stage 1 overwrites x0 rows with h1
|
||||
|
||||
// ------------------------------------------------ stage 1: norm1 (kernel-exact tree)
|
||||
{
|
||||
const int g = lane >> 2, t = lane & 3;
|
||||
#pragma unroll
|
||||
for (int r = 0; r < 8 / NW; ++r)
|
||||
{
|
||||
const int T = (64 / NW) * wv + g + 8 * r;
|
||||
half_t v[NPL][2][4];
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NPL; ++pl)
|
||||
#pragma unroll
|
||||
for (int ww = 0; ww < 2; ++ww)
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
v[pl][ww][b] = X0[T * AST + 32 * pl + 4 * (t + 4 * ww) + b];
|
||||
// the 4 lanes of a token read the whole row before any of them writes h1 over it
|
||||
wave_sync();
|
||||
half_t ws[2][2];
|
||||
#pragma unroll
|
||||
for (int ww = 0; ww < 2; ++ww)
|
||||
{
|
||||
half_t pp[NPL][2];
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NPL; ++pl)
|
||||
#pragma unroll
|
||||
for (int e = 0; e < 2; ++e)
|
||||
{
|
||||
half_t s0 = v[pl][ww][e] * v[pl][ww][e];
|
||||
half_t s1 = v[pl][ww][2 + e] * v[pl][ww][2 + e];
|
||||
pp[pl][e] = s0 + s1;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int e = 0; e < 2; ++e)
|
||||
{
|
||||
if constexpr (NPL == 4)
|
||||
ws[ww][e] = (pp[0][e] + pp[1][e]) + (pp[2][e] + pp[3][e]);
|
||||
else if constexpr (NPL == 3)
|
||||
ws[ww][e] = (pp[0][e] + pp[1][e]) + pp[2][e];
|
||||
else
|
||||
ws[ww][e] = pp[0][e] + pp[1][e];
|
||||
}
|
||||
}
|
||||
half_t a0 = ws[0][0] + ws[1][0], a1 = ws[0][1] + ws[1][1];
|
||||
a0 = a0 + hxor<1>(a0);
|
||||
a1 = a1 + hxor<1>(a1);
|
||||
a0 = a0 + hxor<2>(a0);
|
||||
a1 = a1 + hxor<2>(a1);
|
||||
const half_t rn = rsqrt16(a0 + a1);
|
||||
const hv2 rn2 = {rn, rn};
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NPL; ++pl)
|
||||
#pragma unroll
|
||||
for (int ww = 0; ww < 2; ++ww)
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; b += 2)
|
||||
{
|
||||
const int c = 32 * pl + 4 * (t + 4 * ww) + b;
|
||||
// gperm(c + 1) = gperm(c) + 1 for even c: one 32-bit gamma load per pair
|
||||
const hv2 gm = __builtin_bit_cast(hv2, *(const uint32_t*)(W + L::G1 + 2 * gperm(c)));
|
||||
const hv2 hq = q8x2((hv2){v[pl][ww][b], v[pl][ww][b + 1]} * (hv2)(rn2 * gm));
|
||||
// apos(c + 1) = apos(c) + 1 for even c
|
||||
*(hv2*)(A + T * AST + apos(c)) = hq;
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ------------------------------------------------ stage 2: Q/V projection + attention (wave = head)
|
||||
#ifdef D4R_SWIN_DIAGNOSTICS
|
||||
swin_snapshot<C, AST>(A, 0, bx, by, tid, false);
|
||||
#endif
|
||||
{
|
||||
constexpr int MW = 4 / (NW / NH); // windows per wave
|
||||
const int h = NW == NH ? wv : wv % NH, m0 = NW == NH ? 0 : MW * (wv / NH);
|
||||
T16 acc[MW][4]; // [window][Q nt0, Q nt1, V nt0, V nt1]
|
||||
#pragma unroll
|
||||
for (int m = 0; m < MW; ++m)
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q)
|
||||
acc[m][q] = t16_splat((half_t)0.0f);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 32; ++kc)
|
||||
{
|
||||
gop b[2][4];
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q)
|
||||
b[s][q] = bload(w16, L::qv(h, q >> 1), 2, kc, s, q & 1, l16);
|
||||
#pragma unroll
|
||||
for (int m = 0; m < MW; ++m)
|
||||
{
|
||||
const gop a0 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
|
||||
const gop a1 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc + 16);
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 4; ++q)
|
||||
k32(acc[m][q], a0, b[0][q], a1, b[1][q]);
|
||||
}
|
||||
}
|
||||
// pack: Q as q8 pairs, V as f16 pairs (all that attention needs), halving the live registers
|
||||
hv2 qp[MW][2][4], vp[MW][2][4];
|
||||
#pragma unroll
|
||||
for (int m = 0; m < MW; ++m)
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 2; ++q)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i += 2)
|
||||
{
|
||||
qp[m][q][i / 2] = q8x2(acc[m][q].pair(i / 2));
|
||||
vp[m][q][i / 2] = acc[m][2 + q].pair(i / 2);
|
||||
}
|
||||
__syncthreads(); // every wave is done reading h1; O may overwrite A
|
||||
|
||||
half_t* Qs = QP + wv * 16 * QST;
|
||||
const int bbase = L::head(h) + 64 * C;
|
||||
T16 bias;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int ri = mrow(i), cj = l16;
|
||||
const int bl = 4 * (ri & 7) + ((cj & 7) >> 1), word = (ri >> 3) + 2 * (cj >> 3);
|
||||
bias.set(i, hload(W, bbase + 2 * (8 * bl + 2 * word + (cj & 1))));
|
||||
}
|
||||
const half_t cs = f16(0.0972222164273262f / 5.656854152679443f);
|
||||
const half_t cl = (half_t)0.55615234375f;
|
||||
const half_t p1 = f16(0.92730712890625f), p0 = (half_t)1.375f;
|
||||
#pragma unroll
|
||||
for (int m = 0; m < MW; ++m)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int q = 0; q < 2; ++q)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i += 2)
|
||||
{
|
||||
Qs[mrow(i) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
|
||||
Qs[mrow(i + 1) * QST + 16 * q + l16] = qp[m][q][i / 2][1];
|
||||
}
|
||||
wave_sync();
|
||||
const gop a0 = lda(Qs + l16 * QST), a1 = lda(Qs + l16 * QST + 16);
|
||||
T16 S = bias;
|
||||
k32(S, a0, a0, a1, a1);
|
||||
// rows mrow(2k), mrow(2k+1) of this lane's column, packed: scale, clamp and the cubic
|
||||
// polynomial as f16x2, the exponent trick per row on the (even, odd) column word
|
||||
hv2 wg[4];
|
||||
const bool oddc = l16 & 1;
|
||||
const uint32_t sel_a = oddc ? 0x01000504u : 0x05040100u, sel_b = oddc ? 0x03020706u : 0x07060302u;
|
||||
const uint32_t sel_r = oddc ? 0x07060302u : 0x05040100u;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
hv2 tv = S.pair(k) * (hv2){cs, cs};
|
||||
tv = __builtin_elementwise_max(__builtin_elementwise_min(tv, (hv2){cl, cl}), (hv2){-cl, -cl});
|
||||
const hv2 e1 = __builtin_elementwise_fma(tv, -tv, (hv2){p1, p1});
|
||||
const hv2 poly = __builtin_elementwise_fma(tv, e1, (hv2){p0, p0});
|
||||
const uint32_t mine = __builtin_bit_cast(uint32_t, poly);
|
||||
const uint32_t other = xor_lane<1>(mine);
|
||||
const uint32_t ua = (__builtin_amdgcn_perm(other, mine, sel_a) << 5) + 0x7FF88000u;
|
||||
const uint32_t ub = (__builtin_amdgcn_perm(other, mine, sel_b) << 5) + 0x7FF88000u;
|
||||
wg[k] = __builtin_bit_cast(hv2, __builtin_amdgcn_perm(ub, ua, sel_r));
|
||||
}
|
||||
wave_sync(); // Q scratch is dead: P overwrites it
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
auto xr = [](hv2 v, auto mask) {
|
||||
return __builtin_bit_cast(hv2, xor_lane<decltype(mask)::value>(__builtin_bit_cast(uint32_t, v)));
|
||||
};
|
||||
const hv2 v1 = wg[k] + xr(wg[k], SwinLaneMask<8>{});
|
||||
const hv2 v2 = v1 + xr(v1, SwinLaneMask<2>{});
|
||||
const hv2 v3 = v2 + xr(v2, SwinLaneMask<4>{});
|
||||
const hv2 sum = v3 + xr(v3, SwinLaneMask<1>{});
|
||||
const hv2 rinv = {f16(1.0f / (float)sum[0]), f16(1.0f / (float)sum[1])};
|
||||
const hv2 pw = wg[k] * rinv;
|
||||
Qs[mrow(2 * k) * PST + l16] = pw[0];
|
||||
Qs[mrow(2 * k + 1) * PST + l16] = pw[1];
|
||||
}
|
||||
wave_sync();
|
||||
#ifdef SWIN_EXACT_PV
|
||||
half_t* Vs = Qs + 16 * PST;
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
for (int i = 0; i < 8; ++i)
|
||||
Vs[mrow(i) * 32 + 16 * nt + l16] = vp[m][nt][i / 2][i & 1];
|
||||
wave_sync();
|
||||
#endif
|
||||
const sop pa = lds16(Qs + l16 * PST);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
// V as B operand: lane holds column l16 over the 16 tokens (rows, the K of P V)
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
// own pair k holds rows (2k + 8hi, 2k + 1 + 8hi): this half's K values in order
|
||||
const sop vb = __builtin_bit_cast(sop, vp[m][nt]);
|
||||
#else
|
||||
// own pair k holds rows (4k+hi, 4k+2+hi); the other half-wave has (4k+1-hi, 4k+3-hi)
|
||||
uint32_t vw[8];
|
||||
const uint32_t s_lo = hi ? 0x01000504u : 0x05040100u, s_hi = hi ? 0x03020706u : 0x07060302u;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
const uint32_t own = __builtin_bit_cast(uint32_t, vp[m][nt][k]);
|
||||
const uint32_t other = xor_lane<16>(own);
|
||||
vw[2 * k] = __builtin_amdgcn_perm(other, own, s_lo); // rows 4k, 4k+1
|
||||
vw[2 * k + 1] = __builtin_amdgcn_perm(other, own, s_hi); // rows 4k+2, 4k+3
|
||||
}
|
||||
const h16 vb = __builtin_bit_cast(h16, vw);
|
||||
#endif
|
||||
f8v o = wmma(pa, vb, splat(0.0f));
|
||||
#ifdef SWIN_EXACT_PV
|
||||
// FP16 P/V products need more than 24 significant bits in a
|
||||
// cancellation-heavy dot. An f32 WMMA can round to the wrong
|
||||
// FP16 neighbour at a midpoint, then flip the FP8 output code.
|
||||
// This baseline sums the 16 FP16 products before one FP16 round.
|
||||
for (int i = 0; i < 8; ++i) {
|
||||
double dot = 0.0;
|
||||
for (int k = 0; k < 16; ++k)
|
||||
dot += (double)Qs[mrow(i) * PST + k] * (double)Vs[k * 32 + 16 * nt + l16];
|
||||
o[i] = (float)(half_t)dot;
|
||||
}
|
||||
#endif
|
||||
const int ocol = apos(32 * h + ginv(16 * nt + l16));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i += 2)
|
||||
{
|
||||
const hv2 oq = q8x2((hv2){(half_t)o[i], (half_t)o[i + 1]});
|
||||
A[(16 * (m0 + m) + mrow(i)) * AST + ocol] = oq[0];
|
||||
A[(16 * (m0 + m) + mrow(i + 1)) * AST + ocol] = oq[1];
|
||||
}
|
||||
}
|
||||
wave_sync();
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ------------------------------------------------ stage 3: output projection + residual
|
||||
#ifdef D4R_SWIN_DIAGNOSTICS
|
||||
swin_snapshot<INNER, AST>(A, 3, bx, by, tid, false);
|
||||
#endif
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NH; ++kc)
|
||||
{
|
||||
gop b[2][NTW];
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
b[s][j] = bload(w16, L::WO, NT, kc, s, ng * NTW + j, l16);
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MT; ++mm)
|
||||
{
|
||||
const gop a0 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
|
||||
const gop a1 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc + 16);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
|
||||
}
|
||||
}
|
||||
__syncthreads(); // X0b and O are dead
|
||||
|
||||
// ------------------------------------------------ norm2 (kernel-exact tree), x1 staged in A
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MT; ++mm)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
A[(16 * (mg * MT + mm) + mrow(i)) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
|
||||
__syncthreads();
|
||||
if (tid < 64)
|
||||
{
|
||||
// the row as 128-bit LDS loads: element 8 j + c of n8 tile j
|
||||
#ifdef D4R_SWIN_DIAGNOSTICS
|
||||
if (bx == d4r_swin_debug_block[0] && by == d4r_swin_debug_block[1])
|
||||
for (int c = 0; c < C; ++c)
|
||||
d4r_swin_debug_values[(64 + tid) * 128 + c] = A[tid * AST + gperm(c)];
|
||||
#endif
|
||||
h8 row[L::NT8];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < L::NT8; ++j)
|
||||
row[j] = *(const h8*)(A + tid * AST + 8 * j);
|
||||
half_t L8[8];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 8; ++c)
|
||||
{
|
||||
half_t sq[L::NT8];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < L::NT8; ++j)
|
||||
sq[j] = row[j][c] * row[j][c];
|
||||
half_t even, odd;
|
||||
if constexpr (L::NT8 == 16)
|
||||
{
|
||||
even = ((sq[0] + sq[2]) + (sq[4] + sq[6])) + ((sq[8] + sq[10]) + (sq[12] + sq[14]));
|
||||
odd = ((sq[1] + sq[3]) + (sq[5] + sq[7])) + ((sq[9] + sq[11]) + (sq[13] + sq[15]));
|
||||
}
|
||||
else if constexpr (L::NT8 == 12)
|
||||
{
|
||||
even = ((sq[0] + sq[2]) + (sq[4] + sq[6])) + (sq[8] + sq[10]);
|
||||
odd = ((sq[1] + sq[3]) + (sq[5] + sq[7])) + (sq[9] + sq[11]);
|
||||
}
|
||||
else
|
||||
{
|
||||
even = (sq[0] + sq[2]) + (sq[4] + sq[6]);
|
||||
odd = (sq[1] + sq[3]) + (sq[5] + sq[7]);
|
||||
}
|
||||
L8[c] = even + odd;
|
||||
}
|
||||
const half_t b0 = (L8[0] + L8[2]) + (L8[4] + L8[6]);
|
||||
const half_t b1 = (L8[1] + L8[3]) + (L8[5] + L8[7]);
|
||||
R2[tid] = rsqrt16(b0 + b1);
|
||||
}
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
{
|
||||
const int n = 16 * (ng * NTW + j) + l16;
|
||||
const half_t g2 = hload(W, L::G2 + 2 * n), b2 = hload(W, L::B2 + 2 * n);
|
||||
const int col = apos(ginv(n));
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MT; ++mm)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i += 2)
|
||||
{
|
||||
const int row = 16 * (mg * MT + mm) + mrow(i), row2 = 16 * (mg * MT + mm) + mrow(i + 1);
|
||||
const hv2 v = x[mm][j].pair(i / 2);
|
||||
const hv2 r2 = {R2[row], R2[row2]};
|
||||
const hv2 hq = q8x2(v * (hv2)(r2 * (hv2){g2, g2}));
|
||||
A[row * AST + col] = hq[0];
|
||||
A[row2 * AST + col] = hq[1];
|
||||
x[mm][j].set_pair(i / 2, v + (hv2){b2, b2});
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// ------------------------------------------------ stage 4: MLP, R chunks per round
|
||||
#ifdef D4R_SWIN_DIAGNOSTICS
|
||||
swin_snapshot<C, AST>(A, 2, bx, by, tid, false);
|
||||
#endif
|
||||
{
|
||||
constexpr int WPC = NW / R, MPW = 4 / WPC; // waves per chunk, fc1 m tiles per wave
|
||||
const int cw = wv / WPC, mb = MPW * (wv % WPC);
|
||||
#pragma unroll 1
|
||||
for (int r = 0; r < L::NC / R; ++r)
|
||||
{
|
||||
const int c = R * r + cw;
|
||||
const int B1 = L::b1(c);
|
||||
T16 z[MPW][2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const half_t b1 = hload(W, B1 + 32 * C + 2 * (16 * nt + l16));
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MPW; ++mm)
|
||||
z[mm][nt] = t16_splat(b1);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 32; ++kc)
|
||||
{
|
||||
gop b[2][2];
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
b[s][nt] = bload(w16, L::w1(c), 2, kc, s, nt, l16);
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MPW; ++mm)
|
||||
{
|
||||
const gop a0 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
|
||||
const gop a1 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc + 16);
|
||||
k32(z[mm][0], a0, b[0][0], a1, b[1][0]);
|
||||
k32(z[mm][1], a0, b[0][1], a1, b[1][1]);
|
||||
}
|
||||
}
|
||||
half_t* Gc = G + cw * 64 * GST;
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MPW; ++mm)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const int col = apos(ginv(16 * nt + l16));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; i += 2)
|
||||
{
|
||||
const hv2 zv = z[mm][nt].pair(i / 2);
|
||||
const hv2 cz = __builtin_elementwise_max(__builtin_elementwise_min(zv, (hv2){2.0f16, 2.0f16}),
|
||||
(hv2){-2.0f16, -2.0f16});
|
||||
const hv2 az = __builtin_bit_cast(hv2, (u16x2)(__builtin_bit_cast(u16x2, cz) & (unsigned short)0x7fff));
|
||||
const hv2 inner = (hv2){f16(0.41216981f), f16(0.41216981f)} -
|
||||
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
|
||||
const hv2 gl = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
|
||||
const hv2 gq = q8x2(gl);
|
||||
Gc[(16 * (mb + mm) + mrow(i)) * GST + col] = gq[0];
|
||||
Gc[(16 * (mb + mm) + mrow(i + 1)) * GST + col] = gq[1];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int cc = 0; cc < R; ++cc)
|
||||
{
|
||||
const int c2 = R * r + cc;
|
||||
gop b[2][NTW];
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
b[s][j] = bload(w16, L::w2(c2), NT, 0, s, ng * NTW + j, l16);
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MT; ++mm)
|
||||
{
|
||||
const gop a0 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
|
||||
const gop a1 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST + 16);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ output: f16 staged in LDS, 16-value segments encoded by a loop
|
||||
// XO: [64 tokens][C channels] natural order; the row padding keeps the patch merge's 8-byte reads of 16
|
||||
// tokens on distinct LDS banks (see xo_row)
|
||||
half_t* XO = A;
|
||||
static_assert(xo_row<C>(63) + C <= 64 * AST, "XO fits in A");
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTW; ++j)
|
||||
{
|
||||
const int ch = ginv(16 * (ng * NTW + j) + l16);
|
||||
#pragma unroll
|
||||
for (int mm = 0; mm < MT; ++mm)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
XO[xo_row<C>(16 * (mg * MT + mm) + mrow(i)) + ch] = x[mm][j].get(i);
|
||||
}
|
||||
__syncthreads();
|
||||
#pragma unroll 1
|
||||
for (int seg = tid; seg < 64 * C / 16; seg += 32 * NW)
|
||||
{
|
||||
const int T = seg / (C / 16), c16 = 16 * (seg % (C / 16));
|
||||
// rows are 8-byte aligned (xo_row), so the segment moves as four 8-byte words of two f16 pairs
|
||||
uint2* xs = (uint2*)(XO + xo_row<C>(T) + c16);
|
||||
uint32_t w[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
const uint2 pr = xs[k];
|
||||
hv2 qa, qb;
|
||||
const uint32_t ca = enc8x2(__builtin_bit_cast(hv2, pr.x), qa), cb = enc8x2(__builtin_bit_cast(hv2, pr.y), qb);
|
||||
w[k] = pack_codes(ca, cb);
|
||||
// the patch merge consumes decode(code) = q8_exact(x), written back over x (this thread's own values)
|
||||
if constexpr (NPM > 0)
|
||||
xs[k] = (uint2){__builtin_bit_cast(uint32_t, qa), __builtin_bit_cast(uint32_t, qb)};
|
||||
}
|
||||
const uint4 v = {w[0], w[1], w[2], w[3]};
|
||||
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
|
||||
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
|
||||
*(uint4*)(p.out + ((size_t)((c16 >> 5) * p.th + Y) * p.tw + X) * 32 + (c16 & 31)) = v;
|
||||
}
|
||||
if constexpr (NPM > 0)
|
||||
__syncthreads();
|
||||
|
||||
// ------------------------------------------------ patch merge: 2x2 tokens (4C) -> NPM, half resolution
|
||||
if constexpr (NPM > 0)
|
||||
{
|
||||
// merge tiles split evenly over the largest wave count dividing them (remaining waves idle)
|
||||
constexpr int PT = NPM / 16 ? NPM / 16 : 1;
|
||||
constexpr int PMW = PT % NW == 0 ? NW : (PT % 4 == 0 && NW >= 4 ? 4 : (PT % 3 == 0 && NW >= 3 ? 3 : (PT % 2 == 0 && NW >= 2 ? 2 : 1)));
|
||||
constexpr int NTP = PT / PMW;
|
||||
constexpr bool PM_EVEN = true;
|
||||
if (wv < PMW)
|
||||
{
|
||||
T16 y[NTP];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTP; ++j)
|
||||
y[j] = t16_splat(hload(W, L::PMB + 2 * (16 * (PM_EVEN ? wv * NTP + j : (wv * NTP + j) % (NPM / 16)) + l16)));
|
||||
const int mx = l16 & 3, my = l16 >> 2;
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 8; ++kc)
|
||||
{
|
||||
gop as[2];
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
{
|
||||
sop a;
|
||||
// the operand's K slots in runs of 4 (gfx12: this half's slots 8 hi .. 8 hi + 7)
|
||||
constexpr int RUNS = sizeof(sop) / sizeof(half_t) / 4;
|
||||
#pragma unroll
|
||||
for (int rr = 0; rr < RUNS; ++rr)
|
||||
{
|
||||
const int run = RUNS == 4 ? rr : 2 * hi + rr;
|
||||
const int k = 32 * kc + kslot(s, 4 * run);
|
||||
const int q = k / C, ch = k % C;
|
||||
const int tx = 2 * mx + (q & 1), ty = 2 * my + (q >> 1);
|
||||
const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
|
||||
typedef _Float16 h4 __attribute__((ext_vector_type(4)));
|
||||
const h4 v = *(const h4*)(XO + xo_row<C>(T) + ch);
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
a[4 * rr + b] = v[b];
|
||||
}
|
||||
as[s] = as_gop(a);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTP; ++j)
|
||||
{
|
||||
// out-of-range tiles (uneven split) recompute a valid one and are not stored
|
||||
const int t16 = PM_EVEN ? wv * NTP + j : (wv * NTP + j) % (NPM / 16);
|
||||
k32(y[j], as[0], bload(w16, L::pm(t16 >> 1), 2, kc, 0, t16 & 1, l16), as[1],
|
||||
bload(w16, L::pm(t16 >> 1), 2, kc, 1, t16 & 1, l16));
|
||||
}
|
||||
}
|
||||
const int W2 = p.tw / 2, H2 = p.th / 2;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTP; ++j)
|
||||
{
|
||||
if (!PM_EVEN && wv * NTP + j >= NPM / 16)
|
||||
continue;
|
||||
const int ch = ginv(16 * (wv * NTP + j) + l16);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int r = mrow(i);
|
||||
const int MX = (8 * bx - p.sx) / 2 + (r & 3), MY = (8 * by - p.sy) / 2 + (r >> 2);
|
||||
if (MX < 0 || MX >= W2 || MY < 0 || MY >= H2)
|
||||
continue;
|
||||
p.p48[((size_t)((ch >> 5) * H2 + MY) * W2 + MX) * 32 + (ch & 31)] =
|
||||
(uint8_t)enc8(
|
||||
#ifdef SWIN_EXACT
|
||||
// The audited 310.9 patch-merge path bounds the f16
|
||||
// result to +/- 2*pi before e4m3 quantisation.
|
||||
__builtin_elementwise_max(__builtin_elementwise_min(y[j].get(i), (half_t)6.283185307179586f),
|
||||
(half_t)-6.283185307179586f)
|
||||
#else
|
||||
y[j].get(i)
|
||||
#endif
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (TUBE)
|
||||
{
|
||||
if (tp->out_flags != nullptr)
|
||||
{
|
||||
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "agent");
|
||||
__syncthreads();
|
||||
if (tid == 0)
|
||||
__hip_atomic_store(tp->out_flags + (size_t)(by * gx + bx) * 32, (uint8_t)1, __ATOMIC_RELEASE,
|
||||
__HIP_MEMORY_SCOPE_AGENT);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(SWIN_WAVES_PER_EU)
|
||||
#define SWIN_VGPR_ATTR __attribute__((amdgpu_waves_per_eu(SWIN_WAVES_PER_EU)))
|
||||
#elif defined(SWIN_MAX_VGPR)
|
||||
#define SWIN_VGPR_ATTR __attribute__((amdgpu_num_vgpr(SWIN_MAX_VGPR)))
|
||||
#else
|
||||
#define SWIN_VGPR_ATTR
|
||||
#endif
|
||||
#ifndef SWIN_PERSIST
|
||||
#define SWIN_PERSIST 0 // persistent workgroups in x (0: one workgroup per original block)
|
||||
#endif
|
||||
#ifndef SWIN_PREP_SLOTS
|
||||
#define SWIN_PREP_SLOTS 0 // weight images kept per weights pointer (0: one image, re-prepared when the pointer changes)
|
||||
#endif
|
||||
#ifndef SWIN_PREP_KEY_AT
|
||||
#define SWIN_PREP_KEY_AT 1
|
||||
#endif
|
||||
#ifndef SWIN_PREP_IDENTITY_AT
|
||||
#define SWIN_PREP_IDENTITY_AT 0
|
||||
#endif
|
||||
|
||||
// Weight-image slots for kernels launched with several weight sets (the six enc3 tube blocks). The ZLUDA
|
||||
// hook runs the prep only for a weights pointer it has not seen (d4r_prep_key_at / d4r_prep_key_slots);
|
||||
// the prep claims the first free slot for that pointer, and pointers beyond S share the overflow slot S,
|
||||
// which the hook re-prepares on every launch. Keys are written by one thread of the prep; any prep
|
||||
// thread that reads before that write finds the same (first free) slot.
|
||||
template <int S> __device__ __forceinline__ int prep_slot_find(const uint64_t* keys, uint64_t w)
|
||||
{
|
||||
for (int i = 0; i < S; ++i)
|
||||
if (__hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) == w)
|
||||
return i;
|
||||
return S;
|
||||
}
|
||||
template <int S> __device__ __forceinline__ int prep_slot_claim(uint64_t* keys, uint64_t w, bool writer)
|
||||
{
|
||||
// one pass: the first slot holding w or nothing. Keys only ever go from 0 to a pointer, so every
|
||||
// prep thread picks the same slot whether it reads that slot before or after the writer's store.
|
||||
for (int i = 0; i < S; ++i)
|
||||
{
|
||||
const uint64_t k = __hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
|
||||
if (k == w)
|
||||
return i;
|
||||
if (k == 0)
|
||||
{
|
||||
if (writer)
|
||||
__hip_atomic_store(keys + i, w, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
|
||||
return i;
|
||||
}
|
||||
}
|
||||
return S;
|
||||
}
|
||||
|
||||
// Module boilerplate: f16 weight image(s), prep kernel and its grid for the ZLUDA override hook.
|
||||
#define SWIN_MODULE(NAME, C, NW, NPM, TUBE, PARAMS, CIN) SWIN_MODULE_W(NAME, C, NW, NW, NPM, TUBE, PARAMS, CIN)
|
||||
// NH heads on NWAVES waves; d4r_block_z tells the ZLUDA hook the block z size to launch with
|
||||
#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN) \
|
||||
using NAME##_L = SwinLayout<C, NH, NPM, CIN>; \
|
||||
constexpr int NAME##_S = SWIN_PREP_SLOTS; \
|
||||
__device__ wslot g_w16[(NAME##_S + 1) * NAME##_L::TOTAL]; \
|
||||
__device__ uint64_t g_prep_keys[NAME##_S + 1]; \
|
||||
__constant__ SwinDescs<C, NH, NPM, CIN> g_descs = make_descs<C, NH, NPM, CIN>(); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = NWAVES; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::TOTAL + 127) / 128; \
|
||||
/* the prep reads only p.w (offset 0): skip it while the image for that pointer is resident */ \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = NAME##_S; \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PARAMS p) \
|
||||
{ \
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x; \
|
||||
const uint8_t* w = ((const CommonParams*)&p)->w; \
|
||||
const int slot = NAME##_S ? prep_slot_claim<NAME##_S>(g_prep_keys, (uint64_t)w, idx == 0) : 0; \
|
||||
if (idx < NAME##_L::TOTAL) \
|
||||
expand_weights(w, g_w16 + slot * NAME##_L::TOTAL, g_descs.d, NAME##_L::NDESC, idx); \
|
||||
} \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = SWIN_PERSIST; \
|
||||
extern "C" __global__ void __launch_bounds__(32 * NWAVES) SWIN_VGPR_ATTR NAME(PARAMS p) \
|
||||
{ \
|
||||
const CommonParams& cp = *(const CommonParams*)&p; \
|
||||
const int slot = NAME##_S ? prep_slot_find<NAME##_S>(g_prep_keys, (uint64_t)cp.w) : 0; \
|
||||
const wslot* w16 = g_w16 + slot * NAME##_L::TOTAL; \
|
||||
/* the original grid: 8x8-token blocks covering the shifted token grid */ \
|
||||
const int gx = (cp.tw + cp.sx + 7) / 8, gy = (cp.th + cp.sy + 7) / 8; \
|
||||
if (SWIN_PERSIST == 0) \
|
||||
{ \
|
||||
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, blockIdx.x, blockIdx.y, gx); \
|
||||
return; \
|
||||
} \
|
||||
_Pragma("unroll 1") for (int b = blockIdx.x; b < gx * gy; b += gridDim.x) \
|
||||
{ \
|
||||
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, b % gx, b / gx, gx); \
|
||||
__syncthreads(); \
|
||||
} \
|
||||
}
|
||||
#if defined(SWIN_A8) && defined(SWIN_TOKEN_LANES)
|
||||
#include "swin_tok.h"
|
||||
#undef SWIN_MODULE_W
|
||||
#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN) SWIN_TOK_MODULE(NAME, C, NH, NPM, TUBE, PARAMS, CIN)
|
||||
#endif
|
||||
@@ -0,0 +1,573 @@
|
||||
// Shared device code for native rrlite Swin-block kernels (DLSS 310.7.0; gfx11 / gfx12 wave32 WMMA).
|
||||
// Semantics: ../ENC3_SPEC.md generalised (numpy reference ../swin_model.py).
|
||||
//
|
||||
// Activations are the WMMA A operand (lane = token row), weights the B operand (lane = channel column), and
|
||||
// the f32 result has lane = channel column, VGPR i = token row wm_acc_row(i) (../common/wmma_layout.h).
|
||||
// D4R_FP8_WMMA (gfx12 layout only): the GEMMs whose operands are both e4m3 values run as native FP8 WMMAs on
|
||||
// e4m3 bytes (weights prepared as bytes, activations encoded when loaded); P V stays f16.
|
||||
#pragma once
|
||||
#ifdef D4R_DEVICE_ONLY_MINIMAL
|
||||
#include "../common/hip_device_minimal.h"
|
||||
#else
|
||||
#include <hip/hip_runtime.h>
|
||||
#endif
|
||||
#include <stdint.h>
|
||||
#include "../common/wmma_layout.h"
|
||||
#ifdef D4R_ACCURACY
|
||||
#define SWIN_EXACT
|
||||
#endif
|
||||
#if defined(D4R_FP8_WMMA) && D4R_WMMA_LAYOUT != 12 && defined(__AMDGCN__)
|
||||
#error "D4R_FP8_WMMA needs the gfx12 WMMA layout"
|
||||
#endif
|
||||
// FP8 operands must hold e4m3 values, so the FP8 build re-quantises the intermediate activations as NVIDIA's
|
||||
// network does (the gfx11 "fast numerics" SWIN_NO_Q8 keeps them f16, which no gfx12 WMMA can multiply with
|
||||
// e4m3 weights)
|
||||
// (SWIN_FORCE_Q8 does the same on any target: the gfx11 reference of the FP8 build's numerics)
|
||||
#if (defined(D4R_FP8_WMMA) || defined(SWIN_FORCE_Q8) || defined(SWIN_EXACT)) && defined(SWIN_NO_Q8)
|
||||
#undef SWIN_NO_Q8
|
||||
#endif
|
||||
// SWIN_EXACT (tests): NVIDIA's numerics throughout, the semantics of tools/swin_model.py (e4m3 activations and
|
||||
// an f16 accumulator rounded after every k32 step)
|
||||
#if defined(SWIN_EXACT) && defined(SWIN_F32ACC)
|
||||
#undef SWIN_F32ACC
|
||||
#endif
|
||||
|
||||
// Experimental RDNA4 byte-activation path. Keep disabled unless the caller
|
||||
// explicitly builds native FP8 WMMA for gfx12.
|
||||
#if defined(D4R_FP8_WMMA) && defined(__GFX12__) && !defined(SWIN_NO_A8)
|
||||
#define SWIN_A8 1
|
||||
#endif
|
||||
#if defined(SWIN_A8) && defined(SWIN_EXACT) && !defined(K32_NO_MIX) && !defined(K32_MIX)
|
||||
#define K32_MIX
|
||||
#endif
|
||||
|
||||
#pragma clang fp contract(off)
|
||||
|
||||
typedef _Float16 half_t;
|
||||
typedef _Float16 h16 __attribute__((ext_vector_type(16)));
|
||||
typedef _Float16 h8 __attribute__((ext_vector_type(8)));
|
||||
typedef float f8v __attribute__((ext_vector_type(8)));
|
||||
|
||||
struct CommonParams
|
||||
{
|
||||
const uint8_t* w; // 0
|
||||
int sx, sy; // 8
|
||||
int tw, th; // 16
|
||||
const uint8_t* in; // 24
|
||||
const uint8_t* p32; // 32 (decoders: skip input)
|
||||
uint8_t* out; // 40
|
||||
uint8_t* p48; // 48 (encoders with patch merge: merged output)
|
||||
};
|
||||
static_assert(sizeof(CommonParams) == 56, "param block");
|
||||
|
||||
struct TubeParams
|
||||
{
|
||||
CommonParams c;
|
||||
int r69, r70;
|
||||
int r71, r72;
|
||||
const uint8_t* in_flags;
|
||||
uint8_t* out_flags;
|
||||
};
|
||||
static_assert(sizeof(TubeParams) == 88, "param block");
|
||||
|
||||
// ---------------------------------------------------------------- numerics
|
||||
__device__ __forceinline__ half_t e4m3_to_half(uint32_t code)
|
||||
{
|
||||
uint16_t bits = (uint16_t)(((code & 0x7fu) << 7) | ((code & 0x80u) << 8));
|
||||
return __builtin_bit_cast(half_t, bits) * (half_t)256.0f;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t half_to_e4m3(half_t h)
|
||||
{
|
||||
uint32_t half = __builtin_bit_cast(uint16_t, h);
|
||||
uint32_t magnitude = half & 0x7fffu;
|
||||
uint32_t normal = (magnitude - 0x2000u + 0x3fu + ((magnitude >> 7) & 1u)) >> 7;
|
||||
uint32_t exponent = magnitude >> 10;
|
||||
uint32_t significand = (magnitude & 0x3ffu) | (exponent != 0 ? 0x400u : 0u);
|
||||
uint32_t clamped = exponent < 1u ? 1u : (exponent > 8u ? 8u : exponent);
|
||||
uint32_t shift = 16u - clamped;
|
||||
uint32_t subnormal = (significand + (1u << (shift - 1u)) - 1u + ((significand >> shift) & 1u)) >> shift;
|
||||
uint32_t code = magnitude >= 0x2400u ? normal : subnormal;
|
||||
code = magnitude > 0x5f00u ? 0x7eu : code;
|
||||
code = magnitude > 0x7c00u ? 0x7fu : code;
|
||||
return ((half >> 8) & 0x80u) | code;
|
||||
}
|
||||
|
||||
// e4m3 RNE satfinite round trip evaluated in the f16 domain (exhaustively equal to decode(encode(h)))
|
||||
__device__ __forceinline__ half_t q8_exact(half_t h);
|
||||
__device__ __forceinline__ half_t q8(half_t h)
|
||||
{
|
||||
#ifdef SWIN_NO_Q8
|
||||
return h; // fast mode: intermediate activations stay f16 (stored outputs are still FP8)
|
||||
#endif
|
||||
return q8_exact(h);
|
||||
}
|
||||
__device__ __forceinline__ half_t q8_exact(half_t h)
|
||||
{
|
||||
const uint32_t u = __builtin_bit_cast(uint16_t, h);
|
||||
const uint32_t m = u & 0x7fffu;
|
||||
uint32_t r = (m + 0x3fu + ((m >> 7) & 1u)) & 0x7f80u;
|
||||
r = r < 0x5f00u ? r : 0x5f00u;
|
||||
const half_t mag = __builtin_bit_cast(half_t, (uint16_t)m);
|
||||
const uint32_t sub = __builtin_bit_cast(uint16_t, (half_t)((half_t)(mag + (half_t)2.0f) - (half_t)2.0f));
|
||||
r = m < 0x2400u ? sub : r;
|
||||
r = m > 0x7c00u ? m : r;
|
||||
return __builtin_bit_cast(half_t, (uint16_t)((u & 0x8000u) | r));
|
||||
}
|
||||
|
||||
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
|
||||
typedef unsigned short u16x2 __attribute__((ext_vector_type(2)));
|
||||
typedef short i16x2 __attribute__((ext_vector_type(2)));
|
||||
|
||||
// q8_exact on two halves with packed 16-bit integer / f16 ops (same per-half logic; no carries cross halves)
|
||||
__device__ __forceinline__ hv2 q8x2_exact(hv2 h)
|
||||
{
|
||||
const u16x2 u = __builtin_bit_cast(u16x2, h);
|
||||
const u16x2 m = u & (unsigned short)0x7fff;
|
||||
u16x2 r = (m + (unsigned short)0x3f + ((m >> 7) & (unsigned short)1)) & (unsigned short)0x7f80;
|
||||
r = __builtin_elementwise_min(r, (u16x2){0x5f00, 0x5f00});
|
||||
const hv2 mag = __builtin_bit_cast(hv2, m);
|
||||
const hv2 two = {(half_t)2.0f, (half_t)2.0f};
|
||||
const u16x2 sub = __builtin_bit_cast(u16x2, (hv2)((hv2)(mag + two) - two));
|
||||
const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
|
||||
r = (sub & lt) | (r & ~lt);
|
||||
const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
|
||||
r = (m & gt) | (r & ~gt);
|
||||
return __builtin_bit_cast(hv2, (u16x2)(r | (u & (unsigned short)0x8000)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ hv2 q8x2(hv2 h)
|
||||
{
|
||||
#ifdef SWIN_NO_Q8
|
||||
return h;
|
||||
#endif
|
||||
return q8x2_exact(h);
|
||||
}
|
||||
|
||||
// enc8 of two halves with packed ops: returns the two e4m3 codes in the low bytes of the 16-bit
|
||||
// halves, and q = q8_exact(h) (the values the codes decode to)
|
||||
__device__ __forceinline__ uint32_t codes8x2(hv2 q);
|
||||
__device__ __forceinline__ uint32_t enc8x2(hv2 h, hv2& q)
|
||||
{
|
||||
q = q8x2_exact(h);
|
||||
return codes8x2(q);
|
||||
}
|
||||
|
||||
// e4m3 byte codes of four f16 values; caller provides exact e4m3 values.
|
||||
#ifdef SWIN_A8
|
||||
__device__ __forceinline__ uint32_t codes4(hv2 a, hv2 b)
|
||||
{
|
||||
const hv2 lim = {(half_t)448.0f, (half_t)448.0f};
|
||||
a = __builtin_elementwise_maximum(__builtin_elementwise_minimum(a, lim), -lim);
|
||||
b = __builtin_elementwise_maximum(__builtin_elementwise_minimum(b, lim), -lim);
|
||||
const uint32_t w = (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32((float)a[0], (float)a[1], 0, false);
|
||||
return (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32((float)b[0], (float)b[1], (int)w, true);
|
||||
}
|
||||
#endif
|
||||
|
||||
// the e4m3 codes of two halves that already hold e4m3 values (q8_exact results), in the low bytes of the
|
||||
// 16-bit halves
|
||||
__device__ __forceinline__ uint32_t codes8x2(hv2 q)
|
||||
{
|
||||
const u16x2 qb = __builtin_bit_cast(u16x2, q);
|
||||
const u16x2 m = qb & (unsigned short)0x7fff;
|
||||
// normals: the exponent moves from bias 15 to bias 7
|
||||
const u16x2 norm = (u16x2)(m >> 7) - (unsigned short)64;
|
||||
// subnormals are k * 2^-9 (k = 0..7): f16(k + 1024) has k in its low bits
|
||||
const hv2 mag = __builtin_bit_cast(hv2, m);
|
||||
const hv2 k = (hv2)(mag * (hv2){(half_t)512.0f, (half_t)512.0f}) + (hv2){(half_t)1024.0f, (half_t)1024.0f};
|
||||
const u16x2 sub = __builtin_bit_cast(u16x2, k) - (unsigned short)0x6400;
|
||||
const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
|
||||
u16x2 code = (sub & lt) | (norm & ~lt);
|
||||
const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
|
||||
code = ((u16x2){0x7f, 0x7f} & gt) | (code & ~gt);
|
||||
return __builtin_bit_cast(uint32_t, (u16x2)(((qb >> 8) & (unsigned short)0x80) | code));
|
||||
}
|
||||
|
||||
// four e4m3 codes (a, b from enc8x2) as the bytes of one word in element order
|
||||
__device__ __forceinline__ uint32_t pack_codes(uint32_t a, uint32_t b)
|
||||
{
|
||||
return __builtin_amdgcn_perm(b, a, 0x06040200u);
|
||||
}
|
||||
|
||||
// e4m3 code of h (RNE, satfinite) derived from q8(h): normals shift the exponent, subnormals are k * 2^-9
|
||||
__device__ __forceinline__ uint32_t enc8(half_t h)
|
||||
{
|
||||
const uint32_t qb = __builtin_bit_cast(uint16_t, q8_exact(h));
|
||||
const uint32_t m = qb & 0x7fffu;
|
||||
const half_t mag = __builtin_bit_cast(half_t, (uint16_t)m);
|
||||
uint32_t code = m >= 0x2400u ? (m >> 7) - 64u : (uint32_t)(uint16_t)(mag * (half_t)512.0f);
|
||||
code = m > 0x7c00u ? 0x7fu : code;
|
||||
return ((qb >> 8) & 0x80u) | code;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t f16(float v)
|
||||
{
|
||||
return (half_t)v;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t rsqrt16(half_t s)
|
||||
{
|
||||
return f16(1.0f / __builtin_sqrtf((float)(s + (half_t)(1.0f / 8192.0f))));
|
||||
}
|
||||
__device__ __forceinline__ half_t recip16(half_t s)
|
||||
{
|
||||
return f16(1.0f / (float)s);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ hv2 merge_clamp(hv2 v)
|
||||
{
|
||||
const hv2 lim = {(half_t)6.2831855f, (half_t)6.2831855f};
|
||||
return __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, -lim), lim);
|
||||
}
|
||||
__device__ __forceinline__ half_t merge_clamp(half_t v)
|
||||
{
|
||||
return merge_clamp((hv2){v, v})[0];
|
||||
}
|
||||
|
||||
// natural channel c -> pair-order column, and inverse (per 32-channel group)
|
||||
__device__ __forceinline__ int gperm(int c)
|
||||
{
|
||||
return 32 * (c >> 5) + 2 * ((c & 31) >> 2) + (c & 1) + 16 * ((c & 3) >> 1);
|
||||
}
|
||||
__device__ __forceinline__ int ginv(int n)
|
||||
{
|
||||
int m = n & 31, e = m >> 4, a = (m & 15) >> 1, b = m & 1;
|
||||
return (n & ~31) + 4 * a + 2 * e + b;
|
||||
}
|
||||
|
||||
// WMMA k16 step s of a k32 chunk: operand slot i holds logical k = kslot(s, i) (matches the weight
|
||||
// memory runs). Activations live in LDS at apos(k), so an A operand is 16 contiguous halves.
|
||||
__device__ __forceinline__ int kslot(int s, int i)
|
||||
{
|
||||
int a = 2 * s + (i >> 3), r = i & 7;
|
||||
return (r < 4) ? (4 * a + r) : (16 + 4 * a + (r - 4));
|
||||
}
|
||||
__device__ __forceinline__ int apos(int k)
|
||||
{
|
||||
return (k & ~31) + 8 * ((k & 15) >> 2) + 4 * ((k >> 4) & 1) + (k & 3);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int woff(int base, int Ks, int Ns, int k, int n)
|
||||
{
|
||||
int j = n >> 3;
|
||||
return base + Ks * (k >> 5) + Ns * (j >> 1) + 64 * (n & 7) + 16 * ((k & 15) >> 2) + 8 * (j & 1) +
|
||||
4 * ((k & 31) >> 4) + (k & 3);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- WMMA operands
|
||||
// half-wave of this lane: blocks are 32 x 1 x waves, so the lane is threadIdx.x (the same value as wm_half(),
|
||||
// but from the source the kernels already use for their lane index)
|
||||
__device__ __forceinline__ int m_half()
|
||||
{
|
||||
return (int)threadIdx.x >> 4;
|
||||
}
|
||||
// token row that accumulator VGPR i holds in this lane
|
||||
__device__ __forceinline__ int mrow(int i)
|
||||
{
|
||||
return wm_acc_row_h(i, m_half());
|
||||
}
|
||||
// sop: an f16 operand (gfx11: the row's 16 K values; gfx12: this half's 8). gop: an operand of the e4m3 GEMMs
|
||||
// (k32 steps): sop, or with D4R_FP8_WMMA the 8 e4m3 bytes of this half's K values.
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
typedef h8 sop;
|
||||
#else
|
||||
typedef h16 sop;
|
||||
#endif
|
||||
typedef uint32_t u2v __attribute__((ext_vector_type(2)));
|
||||
typedef uint32_t u4v_t __attribute__((ext_vector_type(4)));
|
||||
#ifdef D4R_FP8_WMMA
|
||||
typedef u2v gop;
|
||||
typedef u4v_t wslot; // weight image slot: 16 e4m3 bytes (K slots 0..15)
|
||||
#else
|
||||
typedef sop gop;
|
||||
typedef h16 wslot; // weight image slot: 16 f16 values
|
||||
#endif
|
||||
|
||||
__device__ __forceinline__ f8v wmma(sop a, sop b, f8v c)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
return wm_mma(__builtin_bit_cast(wm_op, a), __builtin_bit_cast(wm_op, b), c);
|
||||
#else
|
||||
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
|
||||
#endif
|
||||
}
|
||||
|
||||
// f16 operand from 16 contiguous halves in LDS (gfx12: this half's 8 of them)
|
||||
__device__ __forceinline__ sop lds16(const half_t* p)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
return *(const h8*)(p + 8 * m_half());
|
||||
#else
|
||||
h8 lo = *(const h8*)p, hi = *(const h8*)(p + 8);
|
||||
return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef D4R_FP8_WMMA
|
||||
// e4m3 bytes of an f16 operand whose values are e4m3 values (K order kept: byte j = element j)
|
||||
__device__ __forceinline__ gop to_fp8(sop v)
|
||||
{
|
||||
#if (defined(D4R_FP8_HW_PACK) || defined(SWIN_FP8_HWCVT)) && defined(__GFX12__)
|
||||
// Inputs already contain exact e4m3 values. gfx12's native packed convert
|
||||
// avoids the software subnormal/sign/exponent reconstruction below. Keep
|
||||
// opt-in until all finite encodings and layer/network outputs pass replay.
|
||||
gop result;
|
||||
#pragma unroll
|
||||
for (int word = 0; word < 2; ++word) {
|
||||
int packed = __builtin_amdgcn_cvt_pk_fp8_f32((float)v[4 * word], (float)v[4 * word + 1], 0, false);
|
||||
packed = __builtin_amdgcn_cvt_pk_fp8_f32((float)v[4 * word + 2], (float)v[4 * word + 3], packed, true);
|
||||
result[word] = (uint32_t)packed;
|
||||
}
|
||||
return result;
|
||||
#else
|
||||
uint32_t c[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
c[j] = codes8x2((hv2){v[2 * j], v[2 * j + 1]});
|
||||
return (u2v){__builtin_amdgcn_perm(c[1], c[0], 0x06040200u), __builtin_amdgcn_perm(c[3], c[2], 0x06040200u)};
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t input_word(const uint32_t* p)
|
||||
{
|
||||
#if defined(SWIN_NT_INPUT)
|
||||
return __builtin_nontemporal_load(p);
|
||||
#else
|
||||
return *p;
|
||||
#endif
|
||||
}
|
||||
// f32 += A B over 16 K for e4m3 operands
|
||||
__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
|
||||
{
|
||||
#if defined(__GFX12__)
|
||||
return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(wm_i2v, a), __builtin_bit_cast(wm_i2v, b), c);
|
||||
#else
|
||||
// layout shim: the e4m3 values widened to f16 (exact), then the f16 WMMA
|
||||
auto widen = [](gop q) {
|
||||
sop r;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
r[j] = e4m3_to_half((q[j >> 2] >> (8 * (j & 3))) & 0xffu);
|
||||
return r;
|
||||
};
|
||||
return wmma(widen(a), widen(b), c);
|
||||
#endif
|
||||
}
|
||||
__device__ __forceinline__ gop lda(const half_t* p)
|
||||
{
|
||||
return to_fp8(lds16(p));
|
||||
}
|
||||
#else
|
||||
__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
|
||||
{
|
||||
return wmma(a, b, c);
|
||||
}
|
||||
// GEMM A operand (activations) from LDS
|
||||
__device__ __forceinline__ gop lda(const half_t* p)
|
||||
{
|
||||
return lds16(p);
|
||||
}
|
||||
#endif
|
||||
// an f16 operand built in registers -> GEMM operand
|
||||
__device__ __forceinline__ gop as_gop(sop v)
|
||||
{
|
||||
#ifdef D4R_FP8_WMMA
|
||||
return to_fp8(v);
|
||||
#else
|
||||
return v;
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ f8v round16(f8v v)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
v[i] = (float)(half_t)v[i];
|
||||
return v;
|
||||
}
|
||||
|
||||
// A C-fragment tile: element i is row wm_acc_row(i). Exact mode stores packed f16 (NVIDIA's f16
|
||||
// accumulator, rounded after every k32 step); SWIN_F32ACC keeps f32 through whole GEMM chains
|
||||
// (more precise than NVIDIA; values are rounded to f16 where they are consumed).
|
||||
struct T16
|
||||
{
|
||||
#ifdef SWIN_F32ACC
|
||||
f8v f;
|
||||
__device__ hv2 pair(int k) const { return (hv2){(half_t)f[2 * k], (half_t)f[2 * k + 1]}; }
|
||||
__device__ void set_pair(int k, hv2 p) { f[2 * k] = (float)p[0]; f[2 * k + 1] = (float)p[1]; }
|
||||
__device__ half_t get(int i) const { return (half_t)f[i]; }
|
||||
__device__ void set(int i, half_t h) { f[i] = (float)h; }
|
||||
#else
|
||||
hv2 v[4];
|
||||
__device__ hv2 pair(int k) const { return v[k]; }
|
||||
__device__ void set_pair(int k, hv2 p) { v[k] = p; }
|
||||
__device__ half_t get(int i) const { return v[i / 2][i & 1]; }
|
||||
__device__ void set(int i, half_t h) { v[i / 2][i & 1] = h; }
|
||||
#endif
|
||||
};
|
||||
|
||||
__device__ __forceinline__ T16 t16_splat(half_t h)
|
||||
{
|
||||
T16 t;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
t.set_pair(k, (hv2){h, h});
|
||||
return t;
|
||||
}
|
||||
|
||||
#ifndef MIX_PAD
|
||||
#define MIX_PAD ""
|
||||
#endif
|
||||
// acc = f16(P + acc), rounded once (v_fma_mix: verified == f16(exact sum) on 8.4M cases)
|
||||
__device__ __forceinline__ hv2 mix2(float p0, float p1, hv2 c)
|
||||
{
|
||||
hv2 r = c;
|
||||
asm(MIX_PAD "v_fma_mixlo_f16 %0, %1, 1.0, %2 op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p0), "v"(c));
|
||||
asm("v_fma_mixhi_f16 %0, %1, 1.0, %2 op_sel:[0,0,1] op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p1), "v"(c));
|
||||
return r;
|
||||
}
|
||||
__device__ __forceinline__ void mixacc(T16& acc, f8v P)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
acc.set_pair(k, mix2(P[2 * k], P[2 * k + 1], acc.pair(k)));
|
||||
}
|
||||
|
||||
#ifdef K32_CHECK
|
||||
__device__ unsigned d4r_k32_bad[4 + 64 * 4];
|
||||
#endif
|
||||
// f16 accumulate of one k32 step (NVIDIA m16n8k32 semantics: products summed, one rounding with C)
|
||||
__device__ __forceinline__ void k32(T16& acc, gop a0, gop b0, gop a1, gop b1)
|
||||
{
|
||||
#if defined(SWIN_F32ACC)
|
||||
acc.f = wmma8(a1, b1, wmma8(a0, b0, acc.f));
|
||||
#elif defined(K32_MIX)
|
||||
f8v P = wmma8(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
|
||||
P = wmma8(a1, b1, P);
|
||||
mixacc(acc, P);
|
||||
#else // exact: C added inside the f32 WMMA, f16 rounding after the k32 step
|
||||
f8v c;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
c[i] = (float)acc.get(i);
|
||||
c = wmma8(a1, b1, wmma8(a0, b0, c));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc.set(i, (half_t)c[i]);
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ f8v splat(float v)
|
||||
{
|
||||
return (f8v){v, v, v, v, v, v, v, v};
|
||||
}
|
||||
|
||||
// lane xor shuffles without LDS: DPP within rows of 16, permlanex16 across the two rows
|
||||
template <int MASK> __device__ __forceinline__ uint32_t xor_lane(uint32_t v)
|
||||
{
|
||||
if constexpr (MASK == 1)
|
||||
return __builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xf, 0xf, false);
|
||||
else if constexpr (MASK == 2)
|
||||
return __builtin_amdgcn_update_dpp(0, (int)v, 0x4E, 0xf, 0xf, false);
|
||||
else if constexpr (MASK == 4)
|
||||
return __builtin_amdgcn_update_dpp(0, (int)v, 0x164, 0xf, 0xf, false);
|
||||
else if constexpr (MASK == 8)
|
||||
return __builtin_amdgcn_update_dpp(0, (int)v, 0x168, 0xf, 0xf, false);
|
||||
else
|
||||
return __builtin_amdgcn_permlanex16((int)v, (int)v, 0x76543210, 0xfedcba98, false, false);
|
||||
}
|
||||
template <int MASK> __device__ __forceinline__ half_t hxor(half_t v)
|
||||
{
|
||||
return __builtin_bit_cast(half_t, (uint16_t)xor_lane<MASK>(__builtin_bit_cast(uint16_t, v)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int mirror(int v, int n)
|
||||
{
|
||||
v = v < 0 ? -v : v;
|
||||
int o = 2 * n - 2 - v;
|
||||
return v < o ? v : o;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t hload(const uint8_t* w, int off)
|
||||
{
|
||||
return __builtin_bit_cast(half_t, *(const uint16_t*)(w + off));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void wave_sync()
|
||||
{
|
||||
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");
|
||||
}
|
||||
|
||||
// Token T (0..63) of a block: 8x8 tokens, windows of 4x4 (T >> 4)
|
||||
__device__ __forceinline__ int tok_x(int T)
|
||||
{
|
||||
return (T & 3) + 4 * ((T >> 4) & 1);
|
||||
}
|
||||
__device__ __forceinline__ int tok_y(int T)
|
||||
{
|
||||
return ((T >> 2) & 3) + 4 * (T >> 5);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- weight images
|
||||
// A GEMM's B operands in WMMA lane order: slot index dst + ((kc * 2 + s) * NT + nt) * 16 + lane16 holds the 16
|
||||
// K slots of logical k = 32 kc + kslot(s, i), n = 16 nt + lane16 (i = 0..15), as f16 values (wslot = h16) or,
|
||||
// with D4R_FP8_WMMA, as the e4m3 bytes. gfx12 lanes read the half of the slot with their 8 K slots.
|
||||
struct GemmDesc
|
||||
{
|
||||
int dst; // in slots
|
||||
int base, Ks, Ns, KC, NT;
|
||||
};
|
||||
|
||||
// The tile index is wave-uniform (it depends only on the wave, the loop counters and constants), so it
|
||||
// is formed in scalar registers and the load uses a scalar base plus the lane's constant offset.
|
||||
__device__ __forceinline__ gop bload(const wslot* w16, int gemm_dst, int NT, int kc, int s, int nt, int l16)
|
||||
{
|
||||
#ifdef SWIN_SCALAR_BLOAD
|
||||
const int tile = __builtin_amdgcn_readfirstlane(gemm_dst + ((kc * 2 + s) * NT + nt) * 16);
|
||||
const wslot* t = w16 + tile;
|
||||
const int at = l16;
|
||||
#else
|
||||
const wslot* t = w16;
|
||||
const int at = gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16;
|
||||
#endif
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
return ((const gop*)t)[2 * at + m_half()];
|
||||
#else
|
||||
return t[at];
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void expand_weights(const uint8_t* w, wslot* w16, const GemmDesc* descs, int ndesc, int idx)
|
||||
{
|
||||
int d = 0;
|
||||
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
|
||||
++d;
|
||||
const GemmDesc g = descs[d];
|
||||
const int local = idx - g.dst;
|
||||
if (local >= g.KC * 2 * g.NT * 16)
|
||||
return;
|
||||
const int l16 = local & 15, rest = local >> 4;
|
||||
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
|
||||
const int n = 16 * nt + l16;
|
||||
#ifdef D4R_FP8_WMMA
|
||||
u4v_t v;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
uint32_t word = 0;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
word |= (uint32_t)w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, 4 * j + b), n)] << (8 * b);
|
||||
v[j] = word;
|
||||
}
|
||||
#else
|
||||
h16 v;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i)
|
||||
v[i] = e4m3_to_half(w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, i), n)]);
|
||||
#endif
|
||||
w16[idx] = v;
|
||||
}
|
||||
@@ -0,0 +1,801 @@
|
||||
// Token-lane rrlite Swin block for RDNA4 native FP8 (SWIN_A8), selected with SWIN_TOKEN_LANES.
|
||||
//
|
||||
// swin_block.h computes every GEMM as activations x weights: the result has lane = channel column and
|
||||
// VGPR = token row, so each result is transposed through LDS (byte stores, barriers) before it can be the
|
||||
// next GEMM's operand. Here the weights are the first WMMA operand and the activations the second, which
|
||||
// gives the transposed result: lane = token, VGPR i = row 8 half + i of the output. A gfx12 operand is this
|
||||
// half's 8 K values, so an encoded result IS the next GEMM's operand and stays in registers. The products
|
||||
// and the K slot of every product are the same as in swin_block.h, so the values are the same (operand
|
||||
// roles swapped: bit-identical on an RX 9070 XT for the FP8 and the f16 WMMA).
|
||||
//
|
||||
// One wave owns one 4x4 window (16 tokens) for the whole block: norms, attention and the MLP need no LDS
|
||||
// and no barriers. Only the patch expand (decoders) and the patch merge (encoders) cross windows and go
|
||||
// through an e4m3 byte image in LDS.
|
||||
//
|
||||
// A half-wave owns the channels whose bit 2 equals the half (dword 4 t + 2 g + half of tile t, g = 0, 1):
|
||||
// both halves of a k32 operand's K slots are then two of the lane's own code dwords.
|
||||
//
|
||||
// SWIN_TOK_NATURAL: a half-wave owns the channels whose bit 3 equals the half (dword 4 t + 2 half + g), and
|
||||
// the two WMMAs of a k32 step hold K 0..15 and K 16..31 in order: a tile's two code dwords are one operand.
|
||||
// The products are the same; what changes is which of them share a WMMA and in which K slot. That is the
|
||||
// order of ZLUDA's lowering of NVIDIA's m16n8k32, so a kernel built this way (with K32_NO_MIX, which also adds
|
||||
// C inside the first WMMA as that lowering does) accumulates in f32 exactly as ZLUDA's compile of the PTX
|
||||
// does, including the rare sums that are not exact in f32 (../tex/enc0l_tail.hip).
|
||||
#pragma once
|
||||
|
||||
enum : int { TOK_ID, TOK_PI, TOK_X, TOK_NU, TOK_PE };
|
||||
// A GEMM's weight tiles as in GemmDesc (slot dst + ((kc * 2 + s) * NT + nt) * 16 + r holds the 16 K slots of
|
||||
// row r of tile nt), with the raw weight column of (nt, r) chosen by perm:
|
||||
// TOK_ID 16 nt + r (Q: the attention's K order)
|
||||
// TOK_PI gperm(kslot(nt, r)) (V, fc1: row r is K slot r of step nt of the GEMM that consumes it)
|
||||
// TOK_X gperm(tok_xnat(nt, r)) (output projection, fc2: the x layout)
|
||||
// TOK_NU gperm(16 nt + r) (patch merge: natural channel order, 8 contiguous output bytes)
|
||||
// TOK_PE col0 + gperm(16 nt + r) over weight blocks of blkcols columns / blkbytes bytes (patch expand)
|
||||
struct TokDesc
|
||||
{
|
||||
int dst, base, Ks, Ns, KC, NT, perm, col0, blkcols, blkbytes;
|
||||
};
|
||||
enum : int { TC_X, TC_PI, TC_NU, TC_ATT };
|
||||
// Per-channel f16 vectors as 8-value slots: slot dst + 2 j + half holds rows 8 half + i of tile j in the
|
||||
// layout of kind (TC_ATT: j = query token, row = key token, the transposed attention bias)
|
||||
struct TokConst
|
||||
{
|
||||
int dst, kind, base;
|
||||
};
|
||||
|
||||
// x layout: natural channel of row r (= 8 half + i) of tile t
|
||||
__device__ __forceinline__ int tok_xnat(int t, int r)
|
||||
{
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
return 16 * t + r;
|
||||
#else
|
||||
return 16 * t + 8 * ((r >> 2) & 1) + 4 * (r >> 3) + (r & 3);
|
||||
#endif
|
||||
}
|
||||
// logical K (of a 32-value chunk) in operand slot i of WMMA s of a k32 step
|
||||
__device__ __forceinline__ int tok_kslot(int s, int i)
|
||||
{
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
return 16 * s + i;
|
||||
#else
|
||||
return kslot(s, i);
|
||||
#endif
|
||||
}
|
||||
// The 4 codes (tile t, g) of half-wave hf within a token's C channel bytes are at tok_half(hf) + tok_word(t, g)
|
||||
__device__ __forceinline__ int tok_half(int hf)
|
||||
{
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
return 8 * hf;
|
||||
#else
|
||||
return 4 * hf;
|
||||
#endif
|
||||
}
|
||||
__device__ __forceinline__ int tok_word(int t, int g)
|
||||
{
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
return 16 * t + 4 * g;
|
||||
#else
|
||||
return 16 * t + 8 * g;
|
||||
#endif
|
||||
}
|
||||
// first of the four values k (of a merged token's 4C inputs) in K slots 8 half + 4 rr .. + 3 of WMMA s of chunk kc
|
||||
// (without the half-wave's own offset tok_half(hf))
|
||||
__device__ __forceinline__ int tok_merge_k(int kc, int s, int rr)
|
||||
{
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
return 32 * kc + 16 * s + 4 * rr;
|
||||
#else
|
||||
return 32 * kc + 16 * rr + 8 * s;
|
||||
#endif
|
||||
}
|
||||
|
||||
template <int C, int NH, int NPM, int CIN> struct TokLayout
|
||||
{
|
||||
using L = SwinLayout<C, NH, NPM, CIN>;
|
||||
static constexpr int NT = C / 16, NC = C / 8, NTP = NPM / 16;
|
||||
// weight image (slots)
|
||||
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
|
||||
static constexpr int WO = 4 * C * NH;
|
||||
static constexpr int W1 = WO + 2 * NH * C;
|
||||
__host__ __device__ static constexpr int w1(int c) { return W1 + 4 * C * c; }
|
||||
__host__ __device__ static constexpr int w2(int c) { return W1 + 4 * C * c + 2 * C; }
|
||||
static constexpr int PM = W1 + 4 * C * NC;
|
||||
static constexpr int PE = PM + NPM * C / 4;
|
||||
__host__ __device__ static constexpr int pe(int q) { return PE + q * (CIN * C / 16); }
|
||||
static constexpr int TOTAL = PE + CIN * C / 4;
|
||||
static constexpr int NDESC = 2 * NH + 1 + 2 * NC + (NPM ? 1 : 0) + (CIN ? 4 : 0);
|
||||
// constant image (slots of 8 f16)
|
||||
static constexpr int CBO = 0, CG1 = 2 * NT, CG2 = 4 * NT, CB2 = 6 * NT, CB1 = 8 * NT;
|
||||
static constexpr int CATT = CB1 + 4 * NC, CPMB = CATT + 32 * NH, CPEB = CPMB + 2 * NTP;
|
||||
static constexpr int CTOTAL = CPEB + (CIN ? 8 * NT : 0);
|
||||
static constexpr int NCONST = 4 + NC + NH + (NPM ? 1 : 0) + (CIN ? 4 : 0);
|
||||
};
|
||||
|
||||
template <int C, int NH, int NPM, int CIN> struct TokTables
|
||||
{
|
||||
TokDesc d[TokLayout<C, NH, NPM, CIN>::NDESC];
|
||||
TokConst c[TokLayout<C, NH, NPM, CIN>::NCONST];
|
||||
};
|
||||
|
||||
template <int C, int NH, int NPM, int CIN> constexpr TokTables<C, NH, NPM, CIN> make_tok_tables()
|
||||
{
|
||||
using L = SwinLayout<C, NH, NPM, CIN>;
|
||||
using K = TokLayout<C, NH, NPM, CIN>;
|
||||
TokTables<C, NH, NPM, CIN> r{};
|
||||
int n = 0;
|
||||
for (int h = 0; h < NH; ++h)
|
||||
{
|
||||
r.d[n++] = TokDesc{K::qv(h, 0), L::head(h), 1024, 512, C / 32, 2, TOK_ID, 0, 0, 0};
|
||||
r.d[n++] = TokDesc{K::qv(h, 1), L::head(h) + 32 * C, 1024, 512, C / 32, 2, TOK_PI, 0, 0, 0};
|
||||
}
|
||||
r.d[n++] = TokDesc{K::WO, L::head(0) + 64 * C + 512, L::HS, 512, NH, C / 16, TOK_X, 0, 0, 0};
|
||||
for (int c = 0; c < K::NC; ++c)
|
||||
{
|
||||
r.d[n++] = TokDesc{K::w1(c), L::b1(c), 512, 16 * C, C / 32, 2, TOK_PI, 0, 0, 0};
|
||||
r.d[n++] = TokDesc{K::w2(c), L::b1(c) + 32 * C + 64, 0, 512, 1, C / 16, TOK_X, 0, 0, 0};
|
||||
}
|
||||
if (NPM)
|
||||
r.d[n++] = TokDesc{K::PM, L::PM0, 512, 64 * C, C / 8, NPM / 16, TOK_NU, 0, 0, 0};
|
||||
for (int q = 0; q < (CIN ? 4 : 0); ++q)
|
||||
r.d[n++] = TokDesc{K::pe(q), 0, 512, 16 * CIN, CIN / 32, C / 16, TOK_PE, q * C, 4 * C / NH, CIN * 4 * C / NH};
|
||||
int m = 0;
|
||||
r.c[m++] = TokConst{K::CBO, TC_X, L::BO};
|
||||
r.c[m++] = TokConst{K::CG1, TC_X, L::G1};
|
||||
r.c[m++] = TokConst{K::CG2, TC_X, L::G2};
|
||||
r.c[m++] = TokConst{K::CB2, TC_X, L::B2};
|
||||
for (int c = 0; c < K::NC; ++c)
|
||||
r.c[m++] = TokConst{K::CB1 + 4 * c, TC_PI, L::b1(c) + 32 * C};
|
||||
for (int h = 0; h < NH; ++h)
|
||||
r.c[m++] = TokConst{K::CATT + 32 * h, TC_ATT, L::head(h) + 64 * C};
|
||||
if (NPM)
|
||||
r.c[m++] = TokConst{K::CPMB, TC_NU, L::PMB};
|
||||
for (int q = 0; q < (CIN ? 4 : 0); ++q)
|
||||
r.c[m++] = TokConst{K::CPEB + 2 * K::NT * q, TC_NU, 4 * C * CIN + 2 * q * C};
|
||||
return r;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void expand_tok_weights(const uint8_t* w, wslot* w16, const TokDesc* descs, int ndesc, int idx)
|
||||
{
|
||||
int d = 0;
|
||||
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
|
||||
++d;
|
||||
const TokDesc g = descs[d];
|
||||
const int local = idx - g.dst;
|
||||
if (local >= g.KC * 2 * g.NT * 16)
|
||||
return;
|
||||
const int r = local & 15, rest = local >> 4;
|
||||
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
|
||||
int n = 16 * nt + r, base = g.base;
|
||||
if (g.perm == TOK_PI)
|
||||
n = gperm(tok_kslot(nt, r));
|
||||
else if (g.perm == TOK_X)
|
||||
n = gperm(tok_xnat(nt, r));
|
||||
else if (g.perm == TOK_NU || g.perm == TOK_PE)
|
||||
n = gperm(16 * nt + r);
|
||||
if (g.perm == TOK_PE)
|
||||
{
|
||||
const int gc = g.col0 + n;
|
||||
base += (gc / g.blkcols) * g.blkbytes;
|
||||
n = gc % g.blkcols;
|
||||
}
|
||||
u4v_t v;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
uint32_t word = 0;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
word |= (uint32_t)w[woff(base, g.Ks, g.Ns, 32 * kc + tok_kslot(s, 4 * j + b), n)] << (8 * b);
|
||||
v[j] = word;
|
||||
}
|
||||
w16[idx] = v;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void expand_tok_consts(const uint8_t* w, h8* c16, const TokConst* cs, int ncs, int idx)
|
||||
{
|
||||
int d = 0;
|
||||
while (d + 1 < ncs && idx >= cs[d + 1].dst)
|
||||
++d;
|
||||
const TokConst g = cs[d];
|
||||
const int local = idx - g.dst, hf = local & 1, j = local >> 1;
|
||||
h8 v;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int r = 8 * hf + i;
|
||||
int off;
|
||||
if (g.kind == TC_X)
|
||||
off = 2 * gperm(tok_xnat(j, r));
|
||||
else if (g.kind == TC_PI)
|
||||
off = 2 * gperm(tok_kslot(j, r));
|
||||
else if (g.kind == TC_NU)
|
||||
off = 2 * gperm(16 * j + r);
|
||||
else
|
||||
{
|
||||
// the attention bias of (query j, key r)
|
||||
const int bl = 4 * (j & 7) + ((r & 7) >> 1), word = (j >> 3) + 2 * (r >> 3);
|
||||
off = 2 * (8 * bl + 2 * word + (r & 1));
|
||||
}
|
||||
v[i] = hload(w, g.base + off);
|
||||
}
|
||||
c16[idx] = v;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- register helpers
|
||||
typedef float tok_f2v __attribute__((ext_vector_type(2)));
|
||||
|
||||
// two f32 -> f16 (round to nearest even) packed into one register
|
||||
__device__ __forceinline__ hv2 tok_pk16(float a, float b)
|
||||
{
|
||||
return (hv2){(half_t)a, (half_t)b};
|
||||
}
|
||||
|
||||
// two e4m3 codes of a word (bytes 0, 1 or 2, 3) as f16. Equal to e4m3_to_half for every code except the NaN
|
||||
// codes 0x7f / 0xff, which the conversion decodes as NaN (e4m3_to_half: +-480). The FP8 WMMAs read the
|
||||
// same bytes with the hardware's decoding, and no encoder here stores those codes for a non-NaN value.
|
||||
template <bool UPPER> __device__ __forceinline__ hv2 tok_dec2(uint32_t word)
|
||||
{
|
||||
const tok_f2v f = __builtin_amdgcn_cvt_pk_f32_fp8((int)word, UPPER);
|
||||
return tok_pk16(f[0], f[1]);
|
||||
}
|
||||
|
||||
// the same lane of the other half-wave
|
||||
__device__ __forceinline__ hv2 tok_other(hv2 v)
|
||||
{
|
||||
return __builtin_bit_cast(hv2, xor_lane<16>(__builtin_bit_cast(uint32_t, v)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ T16 tok_t16(h8 c)
|
||||
{
|
||||
T16 t;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
t.set_pair(k, (hv2){c[2 * k], c[2 * k + 1]});
|
||||
return t;
|
||||
}
|
||||
|
||||
// the two operands (WMMA s = 0, 1) of a k32 step from the code words [tile][g] of the chunk's two tiles
|
||||
__device__ __forceinline__ gop tok_k32_operand(const uint32_t (*c)[2], int s)
|
||||
{
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
return (u2v){c[s][0], c[s][1]};
|
||||
#else
|
||||
return (u2v){c[0][s], c[1][s]};
|
||||
#endif
|
||||
}
|
||||
|
||||
// e4m3 codes of a tile's 8 rows: this half's 8 K slots of an operand
|
||||
__device__ __forceinline__ gop tok_codes(const T16& t)
|
||||
{
|
||||
return (u2v){codes4(t.pair(0), t.pair(1)), codes4(t.pair(2), t.pair(3))};
|
||||
}
|
||||
|
||||
// 1 / sqrt(sum of squares + eps) of a token's C values, with the summation tree of norm1 and norm2 of
|
||||
// swin_block.h (both are this tree in the x layout): per (g, element) the squares of a dword's two pairs,
|
||||
// the 32-channel groups, the two 16-channel halves of a group, then the half-waves, g and the elements
|
||||
template <int NT> __device__ __forceinline__ half_t tok_rsq(const hv2 (&v)[NT][4])
|
||||
{
|
||||
constexpr int NPL = NT / 2;
|
||||
hv2 u[2];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
{
|
||||
hv2 ws[2];
|
||||
#pragma unroll
|
||||
for (int ww = 0; ww < 2; ++ww)
|
||||
{
|
||||
hv2 pp[NPL];
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NPL; ++pl)
|
||||
{
|
||||
const hv2 a = v[2 * pl + ww][2 * g], b = v[2 * pl + ww][2 * g + 1];
|
||||
pp[pl] = (hv2)(a * a) + (hv2)(b * b);
|
||||
}
|
||||
if constexpr (NPL == 4)
|
||||
ws[ww] = (pp[0] + pp[1]) + (pp[2] + pp[3]);
|
||||
else if constexpr (NPL == 3)
|
||||
ws[ww] = (pp[0] + pp[1]) + pp[2];
|
||||
else if constexpr (NPL == 2)
|
||||
ws[ww] = pp[0] + pp[1];
|
||||
else
|
||||
ws[ww] = pp[0];
|
||||
}
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
u[g] = ws[0] + ws[1];
|
||||
}
|
||||
// the tree's next partner (dword ^ 1) is this lane's other g, the one after it (dword ^ 2) the other half
|
||||
const hv2 a = u[0] + u[1];
|
||||
const hv2 s = a + tok_other(a);
|
||||
#else
|
||||
const hv2 a = ws[0] + ws[1];
|
||||
u[g] = a + tok_other(a);
|
||||
}
|
||||
const hv2 s = u[0] + u[1];
|
||||
#endif
|
||||
return rsqrt16(s[0] + s[1]);
|
||||
}
|
||||
|
||||
// The block from its input x0 to the e4m3 codes of its output, for the 16 tokens of one 4x4 window (lane l16 =
|
||||
// token, both half-waves). v: x0 as f16 pairs, [tile][2 g + pair of the dword]; it is overwritten (x1).
|
||||
// co: the output codes, [tile][g] (see tok_word). ch / wl: this lane's bases into the constant image (c16 + half)
|
||||
// and the weight image (row l16 of a tile, this half's 8 K slots). x2 (optional): the output before its e4m3
|
||||
// encoding, as f16 pairs like v.
|
||||
template <int C, int NH, int NPM, int CIN>
|
||||
__device__ __forceinline__ void swin_tok_core(const h8* const ch, const gop* const wl, const int l16, hv2 (&v)[C / 16][4],
|
||||
uint32_t (&co)[C / 16][2], hv2 (*x2)[4] = nullptr)
|
||||
{
|
||||
using K = TokLayout<C, NH, NPM, CIN>;
|
||||
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8;
|
||||
auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
|
||||
|
||||
// residual of the output projection: f16(x0 + b_o)
|
||||
hv2 res[NT][4];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
const h8 bo = ch[K::CBO + 2 * t];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
res[t][k] = v[t][k] + (hv2){bo[2 * k], bo[2 * k + 1]};
|
||||
}
|
||||
|
||||
// ------------------------------------------------ norm1: h1 codes, the K operands of Q and V
|
||||
gop hop[NKC][2];
|
||||
{
|
||||
const half_t rn = tok_rsq<NT>(v);
|
||||
const hv2 rn2 = {rn, rn};
|
||||
uint32_t c1[NT][2];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
const h8 g1 = ch[K::CG1 + 2 * t];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
c1[t][g] = codes4(v[t][2 * g] * (hv2)(rn2 * (hv2){g1[4 * g], g1[4 * g + 1]}),
|
||||
v[t][2 * g + 1] * (hv2)(rn2 * (hv2){g1[4 * g + 2], g1[4 * g + 3]}));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
hop[kc][s] = tok_k32_operand(&c1[2 * kc], s);
|
||||
}
|
||||
|
||||
// ------------------------------------------------ attention per head: O codes, the K operands of the projection
|
||||
gop oc[NH][2];
|
||||
#pragma unroll
|
||||
for (int hd = 0; hd < NH; ++hd)
|
||||
{
|
||||
// Q transposed (lane = token, rows = Q's 32 columns), V as in swin_block.h (lane = column, rows = tokens)
|
||||
T16 qt[2], vt[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
qt[nt] = t16_splat((half_t)0.0f);
|
||||
vt[nt] = t16_splat((half_t)0.0f);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
k32(qt[nt], wt(K::qv(hd, 0), 2, kc, 0, nt), hop[kc][0], wt(K::qv(hd, 0), 2, kc, 1, nt), hop[kc][1]);
|
||||
k32(vt[nt], hop[kc][0], wt(K::qv(hd, 1), 2, kc, 0, nt), hop[kc][1], wt(K::qv(hd, 1), 2, kc, 1, nt));
|
||||
}
|
||||
const gop q0 = tok_codes(qt[0]), q1 = tok_codes(qt[1]);
|
||||
// S transposed: lane = query token, rows = key tokens
|
||||
T16 S = tok_t16(ch[K::CATT + 32 * hd + 2 * l16]);
|
||||
k32(S, q0, q0, q1, q1);
|
||||
const half_t cs = f16(0.0972222164273262f / 5.656854152679443f);
|
||||
const half_t cl = (half_t)0.55615234375f;
|
||||
const half_t p1 = f16(0.92730712890625f), p0 = (half_t)1.375f;
|
||||
hv2 wg[4], v1[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
hv2 tv = S.pair(k) * (hv2){cs, cs};
|
||||
tv = __builtin_elementwise_max(__builtin_elementwise_min(tv, (hv2){cl, cl}), (hv2){-cl, -cl});
|
||||
const hv2 e1 = __builtin_elementwise_fma(tv, -tv, (hv2){p1, p1});
|
||||
const hv2 poly = __builtin_elementwise_fma(tv, e1, (hv2){p0, p0});
|
||||
// the pair is the (even, odd) key word of the exponent trick
|
||||
wg[k] = __builtin_bit_cast(hv2, (__builtin_bit_cast(uint32_t, poly) << 5) + 0x7FF88000u);
|
||||
v1[k] = wg[k] + tok_other(wg[k]); // keys c and c ^ 8
|
||||
}
|
||||
const hv2 v3 = (v1[0] + v1[1]) + (v1[2] + v1[3]); // c ^ 2, then c ^ 4
|
||||
const half_t sum = v3[0] + v3[1]; // c ^ 1
|
||||
const half_t rinv = recip16(sum);
|
||||
hv2 pw[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
pw[k] = wg[k] * (hv2){rinv, rinv};
|
||||
#ifdef SWIN_PV_ZLUDA_ORDER
|
||||
// P V is the block's only product of two f16 operands: its f32 sum over the 16 keys is not exact, so the
|
||||
// rounded result depends on which K slot holds which key. Here the slots follow ZLUDA's lowering of
|
||||
// NVIDIA's m16n8k16 (dword j of half h = key pair (j & 1) * 4 + 2 h + (j >> 1)) instead of this
|
||||
// layout's own (key pair 4 h + j), which costs two half-wave exchanges per operand. A kernel that has to
|
||||
// reproduce ZLUDA's compile byte for byte needs it; the two orders differ in about one token in 30,000.
|
||||
auto zk = [&](hv2 (&w)[4]) {
|
||||
const bool upper = m_half() != 0;
|
||||
const hv2 r0 = tok_other(upper ? w[0] : w[2]), r1 = tok_other(upper ? w[1] : w[3]);
|
||||
const hv2 a = upper ? r0 : w[0], b = upper ? w[2] : r0, c = upper ? r1 : w[1], d = upper ? w[3] : r1;
|
||||
w[0] = a, w[1] = b, w[2] = c, w[3] = d;
|
||||
};
|
||||
zk(pw);
|
||||
#endif
|
||||
const sop pb = __builtin_bit_cast(sop, pw);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
hv2 vp[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
vp[k] = vt[nt].pair(k);
|
||||
#ifdef SWIN_PV_ZLUDA_ORDER
|
||||
zk(vp);
|
||||
#endif
|
||||
const f8v o = wmma(__builtin_bit_cast(sop, vp), pb, splat(0.0f));
|
||||
oc[hd][nt] = (u2v){codes4(tok_pk16(o[0], o[1]), tok_pk16(o[2], o[3])), codes4(tok_pk16(o[4], o[5]), tok_pk16(o[6], o[7]))};
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ output projection + residual
|
||||
T16 x[NT];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
x[t].set_pair(k, res[t][k]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NH; ++kc)
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
k32(x[t], wt(K::WO, NT, kc, 0, t), oc[kc][0], wt(K::WO, NT, kc, 1, t), oc[kc][1]);
|
||||
|
||||
// ------------------------------------------------ norm2: h2 codes, the K operands of fc1
|
||||
gop fop[NKC][2];
|
||||
{
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
v[t][k] = x[t].pair(k);
|
||||
const half_t rn = tok_rsq<NT>(v);
|
||||
const hv2 rn2 = {rn, rn};
|
||||
uint32_t c2[NT][2];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
const h8 g2 = ch[K::CG2 + 2 * t], b2 = ch[K::CB2 + 2 * t];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
c2[t][g] = codes4(v[t][2 * g] * (hv2)(rn2 * (hv2){g2[4 * g], g2[4 * g + 1]}),
|
||||
v[t][2 * g + 1] * (hv2)(rn2 * (hv2){g2[4 * g + 2], g2[4 * g + 3]}));
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
x[t].set_pair(k, v[t][k] + (hv2){b2[2 * k], b2[2 * k + 1]});
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
fop[kc][s] = tok_k32_operand(&c2[2 * kc], s);
|
||||
}
|
||||
|
||||
// ------------------------------------------------ MLP: fc1 transposed, its codes are fc2's operands
|
||||
#pragma unroll 1
|
||||
for (int c = 0; c < NC; ++c)
|
||||
{
|
||||
const int w1 = K::W1 + 4 * C * c;
|
||||
// the round's operands first: 2 bias slots, fc1's and fc2's weight tiles
|
||||
h8 zb[2];
|
||||
gop wa[NKC][2][2], wb[2][NT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
zb[nt] = ch[K::CB1 + 4 * c + 2 * nt];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
wa[kc][s][nt] = wt(w1, 2, kc, s, nt);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
wb[s][t] = wt(w1 + 2 * C, NT, 0, s, t);
|
||||
#ifndef SWIN_TOK_NO_LOAD_SYNC
|
||||
// one wait for the whole batch here instead of one before each WMMA
|
||||
asm("" : "+v"(wb[1][NT - 1]));
|
||||
#endif
|
||||
T16 z[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
z[nt] = tok_t16(zb[nt]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
k32(z[nt], wa[kc][0][nt], fop[kc][0], wa[kc][1][nt], fop[kc][1]);
|
||||
gop gz[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
hv2 gl[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
const hv2 zv = z[nt].pair(k);
|
||||
// minimum / maximum instead of swin_block.h's minnum / maxnum: they differ only for a NaN z, and
|
||||
// then the product with z below is NaN either way (no canonicalising instruction per pair)
|
||||
const hv2 cz = __builtin_elementwise_maximum(__builtin_elementwise_minimum(zv, (hv2){2.0f16, 2.0f16}),
|
||||
(hv2){-2.0f16, -2.0f16});
|
||||
const hv2 az = __builtin_bit_cast(hv2, (u16x2)(__builtin_bit_cast(u16x2, cz) & (unsigned short)0x7fff));
|
||||
const hv2 inner = (hv2){f16(0.41216981f), f16(0.41216981f)} -
|
||||
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
|
||||
gl[k] = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
|
||||
}
|
||||
gz[nt] = (u2v){codes4(gl[0], gl[1]), codes4(gl[2], gl[3])};
|
||||
}
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
k32(x[t], wb[0][t], gz[0], wb[1][t], gz[1]);
|
||||
}
|
||||
|
||||
// ------------------------------------------------ output codes
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
co[t][g] = codes4(x[t].pair(2 * g), x[t].pair(2 * g + 1));
|
||||
if (x2 != nullptr)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
x2[t][k] = x[t].pair(k);
|
||||
}
|
||||
}
|
||||
|
||||
template <int C, int NH, int NPM, bool TUBE, int CIN>
|
||||
__device__ __forceinline__ void swin_tok_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const h8* c16,
|
||||
const int bx, const int by, const int gx)
|
||||
{
|
||||
using K = TokLayout<C, NH, NPM, CIN>;
|
||||
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8;
|
||||
static_assert(C % 32 == 0, "32-channel groups");
|
||||
|
||||
const int lane = threadIdx.x, wv = threadIdx.z;
|
||||
const int l16 = lane & 15, hf = lane >> 4;
|
||||
const int T = 16 * wv + l16; // this lane's token: window wv, position l16
|
||||
// The lane's bases into the two images, hidden from constant folding: every tile and constant is then
|
||||
// one load at a fixed offset from a register (folded, each load rebuilt the image's address first).
|
||||
uintptr_t chv = (uintptr_t)(c16 + hf), wlv = (uintptr_t)((const gop*)w16 + 2 * l16 + hf);
|
||||
asm("" : "+v"(chv));
|
||||
asm("" : "+v"(wlv));
|
||||
const h8* const ch = (const h8*)chv;
|
||||
// row l16 of a weight tile: this half's 8 K slots
|
||||
const gop* const wl = (const gop*)wlv;
|
||||
auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
|
||||
|
||||
if constexpr (TUBE)
|
||||
{
|
||||
if (tp->in_flags != nullptr)
|
||||
{
|
||||
if (lane <= 3 && wv == 0)
|
||||
{
|
||||
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
|
||||
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
|
||||
{
|
||||
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
|
||||
// bounded spin: a missing producer must never hang the GPU
|
||||
for (int it = 0; it < (1 << 22); ++it)
|
||||
{
|
||||
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
|
||||
break;
|
||||
__builtin_amdgcn_s_sleep(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ block input x0 (encoders: the fp8 input; decoders: skip)
|
||||
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
|
||||
const int Xm = mirror(X, p.tw), Ym = mirror(Y, p.th);
|
||||
uint32_t inw[NT][2];
|
||||
{
|
||||
const uint8_t* const src = CIN ? p.p32 : p.in;
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NKC; ++pl)
|
||||
{
|
||||
const uint8_t* b = src + ((size_t)(pl * p.th + Ym) * p.tw + Xm) * 32 + tok_half(hf);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
inw[2 * pl + (j >> 1)][j & 1] = input_word((const uint32_t*)(b + tok_word(j >> 1, j & 1)));
|
||||
}
|
||||
}
|
||||
|
||||
hv2 v[NT][4]; // x0, then x1, as f16 pairs: [tile][2 g + pair of the dword]
|
||||
if constexpr (CIN > 0)
|
||||
{
|
||||
// patch expand: lanes = the block's 16 low-res tokens, wave = quadrant; the codes of q8(expand) go to
|
||||
// their child tokens' rows of a byte image, where the windows' waves pick them up
|
||||
__shared__ __attribute__((aligned(16))) uint8_t E8[64 * C];
|
||||
const int W2 = p.tw / 2, H2 = p.th / 2;
|
||||
const int MX = mirror((8 * bx - p.sx) / 2 + (l16 & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (l16 >> 2), H2);
|
||||
gop pa[CIN / 32][2];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < CIN / 32; ++kc)
|
||||
{
|
||||
const uint8_t* b = p.in + ((size_t)(kc * H2 + MY) * W2 + MX) * 32 + tok_half(hf);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#ifdef SWIN_TOK_NATURAL
|
||||
pa[kc][s] = (u2v){input_word((const uint32_t*)(b + tok_word(s, 0))), input_word((const uint32_t*)(b + tok_word(s, 1)))};
|
||||
#else
|
||||
pa[kc][s] = (u2v){input_word((const uint32_t*)(b + tok_word(0, s))), input_word((const uint32_t*)(b + tok_word(1, s)))};
|
||||
#endif
|
||||
}
|
||||
const int tx = 2 * (l16 & 3) + (wv & 1), ty = 2 * (l16 >> 2) + (wv >> 1);
|
||||
const int Tc = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
|
||||
const int pew = K::pe(0) + wv * (CIN * C / 16);
|
||||
#pragma unroll
|
||||
for (int ct = 0; ct < NT; ++ct)
|
||||
{
|
||||
T16 e = tok_t16(ch[K::CPEB + 2 * (NT * wv + ct)]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < CIN / 32; ++kc)
|
||||
k32(e, wt(pew, NT, kc, 0, ct), pa[kc][0], wt(pew, NT, kc, 1, ct), pa[kc][1]);
|
||||
*(u2v*)(E8 + Tc * C + 16 * ct + 8 * hf) = tok_codes(e);
|
||||
}
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
{
|
||||
const uint32_t ew = *(const uint32_t*)(E8 + T * C + tok_word(t, g) + tok_half(hf));
|
||||
v[t][2 * g] = tok_dec2<false>(ew) + tok_dec2<false>(inw[t][g]);
|
||||
v[t][2 * g + 1] = tok_dec2<true>(ew) + tok_dec2<true>(inw[t][g]);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
{
|
||||
v[t][2 * g] = tok_dec2<false>(inw[t][g]);
|
||||
v[t][2 * g + 1] = tok_dec2<true>(inw[t][g]);
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t co[NT][2];
|
||||
swin_tok_core<C, NH, NPM, CIN>(ch, wl, l16, v, co);
|
||||
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NKC; ++pl)
|
||||
{
|
||||
uint8_t* b = p.out + ((size_t)(pl * p.th + Y) * p.tw + X) * 32 + tok_half(hf);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
*(uint32_t*)(b + tok_word(j >> 1, j & 1)) = co[2 * pl + (j >> 1)][j & 1];
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ patch merge: 2x2 tokens (4C) -> NPM, half resolution
|
||||
if constexpr (NPM > 0)
|
||||
{
|
||||
// lanes = the block's 16 merged tokens; the merge reads the output codes from a byte image
|
||||
constexpr int XB = C + 4, NTP = NPM / 16;
|
||||
__shared__ __attribute__((aligned(16))) uint8_t U[64 * XB];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
*(uint32_t*)(U + T * XB + tok_word(t, g) + tok_half(hf)) = co[t][g];
|
||||
__syncthreads();
|
||||
if (wv < NTP)
|
||||
{
|
||||
const int mx = l16 & 3, my = l16 >> 2;
|
||||
// K slot 8 half + 4 rr + b of step (kc, s) is value k = 32 kc + 16 rr + 8 s + 4 half + b of the merged
|
||||
// token's 4C inputs: channel k % C of its source token k / C. The four source tokens are consecutive
|
||||
// in x and 4 apart in y, and 4 half never crosses a token, so every run of four codes is at a fixed
|
||||
// offset from the lane's first source token.
|
||||
const uint8_t* const ub = U + ((2 * mx & 3) + 4 * (2 * my & 3) + 16 * (mx >> 1) + 32 * (my >> 1)) * XB + tok_half(hf);
|
||||
gop as[C / 8][2];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 8; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
{
|
||||
uint32_t run4[2];
|
||||
#pragma unroll
|
||||
for (int rr = 0; rr < 2; ++rr)
|
||||
{
|
||||
const int k0 = tok_merge_k(kc, s, rr), q = k0 / C;
|
||||
run4[rr] = *(const uint32_t*)(ub + ((q & 1) + 4 * (q >> 1)) * XB + k0 % C);
|
||||
}
|
||||
as[kc][s] = (u2v){run4[0], run4[1]};
|
||||
}
|
||||
const int W2 = p.tw / 2, H2 = p.th / 2;
|
||||
const int MX = (8 * bx - p.sx) / 2 + mx, MY = (8 * by - p.sy) / 2 + my;
|
||||
const bool inside = MX >= 0 && MX < W2 && MY >= 0 && MY < H2;
|
||||
// output tiles u = wv, wv + 4, ...: rows = 16 natural channels, 8 contiguous bytes per half-wave
|
||||
#pragma unroll
|
||||
for (int j = 0; j < (NTP + 3) / 4; ++j)
|
||||
{
|
||||
const int u = wv + 4 * j;
|
||||
if (u < NTP)
|
||||
{
|
||||
T16 y = tok_t16(ch[K::CPMB + 2 * u]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 8; ++kc)
|
||||
k32(y, wt(K::PM, NTP, kc, 0, u), as[kc][0], wt(K::PM, NTP, kc, 1, u), as[kc][1]);
|
||||
const u2v code = {codes4(merge_clamp(y.pair(0)), merge_clamp(y.pair(1))),
|
||||
codes4(merge_clamp(y.pair(2)), merge_clamp(y.pair(3)))};
|
||||
if (inside)
|
||||
*(u2v*)(p.p48 + ((size_t)((u >> 1) * H2 + MY) * W2 + MX) * 32 + 16 * (u & 1) + 8 * hf) = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (TUBE)
|
||||
{
|
||||
if (tp->out_flags != nullptr)
|
||||
{
|
||||
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "agent");
|
||||
__syncthreads();
|
||||
if (lane == 0 && wv == 0)
|
||||
__hip_atomic_store(tp->out_flags + (size_t)(by * gx + bx) * 32, (uint8_t)1, __ATOMIC_RELEASE,
|
||||
__HIP_MEMORY_SCOPE_AGENT);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(SWIN_TOK_MAX_VGPR)
|
||||
#define SWIN_TOK_VGPR_ATTR __attribute__((amdgpu_num_vgpr(SWIN_TOK_MAX_VGPR)))
|
||||
#else
|
||||
#define SWIN_TOK_VGPR_ATTR
|
||||
#endif
|
||||
|
||||
// Module boilerplate: the weight and constant images, their prep kernel, and the 4-wave main kernel
|
||||
#define SWIN_TOK_MODULE(NAME, C, NH, NPM, TUBE, PARAMS, CIN) \
|
||||
using NAME##_K = TokLayout<C, NH, NPM, CIN>; \
|
||||
constexpr int NAME##_S = SWIN_PREP_SLOTS; \
|
||||
__device__ wslot g_w16[(NAME##_S + 1) * NAME##_K::TOTAL]; \
|
||||
__device__ h8 g_c16[(NAME##_S + 1) * NAME##_K::CTOTAL]; \
|
||||
extern "C" __device__ __attribute__((used)) uint64_t g_prep_keys[NAME##_S + 1] = {}; \
|
||||
__constant__ TokTables<C, NH, NPM, CIN> g_tok = make_tok_tables<C, NH, NPM, CIN>(); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = \
|
||||
((NAME##_K::TOTAL > NAME##_K::CTOTAL ? NAME##_K::TOTAL : NAME##_K::CTOTAL) + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = SWIN_PREP_KEY_AT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = NAME##_S; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_identity_at = SWIN_PREP_IDENTITY_AT; \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PARAMS p) \
|
||||
{ \
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x; \
|
||||
const uint8_t* w = ((const CommonParams*)&p)->w; \
|
||||
const int slot = NAME##_S ? prep_slot_claim<NAME##_S>(g_prep_keys, (uint64_t)w, idx == 0) : 0; \
|
||||
if (idx < NAME##_K::TOTAL) \
|
||||
expand_tok_weights(w, g_w16 + slot * NAME##_K::TOTAL, g_tok.d, NAME##_K::NDESC, idx); \
|
||||
if (idx < NAME##_K::CTOTAL) \
|
||||
expand_tok_consts(w, g_c16 + slot * NAME##_K::CTOTAL, g_tok.c, NAME##_K::NCONST, idx); \
|
||||
} \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = 0; \
|
||||
extern "C" __global__ void __launch_bounds__(128) SWIN_TOK_VGPR_ATTR NAME(PARAMS p) \
|
||||
{ \
|
||||
const CommonParams& cp = *(const CommonParams*)&p; \
|
||||
const int slot = NAME##_S ? prep_slot_find<NAME##_S>(g_prep_keys, (uint64_t)cp.w) : 0; \
|
||||
const int gx = (cp.tw + cp.sx + 7) / 8; \
|
||||
swin_tok_block<C, NH, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16 + slot * NAME##_K::TOTAL, \
|
||||
g_c16 + slot * NAME##_K::CTOTAL, blockIdx.x, blockIdx.y, gx); \
|
||||
}
|
||||
Reference in new issue
Block a user