wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
78a3880c4475779503fab0d7616c1508ee166ba9
d4r/kernels/m/swin_block.h
T
ayo ac4ee6c225 Expand Windows GPU support, add native FSR RR, and improve SR interop 
Add a native D3D12 FSR Ray Regeneration route for gfx1200/gfx1201,
using the game's device and queue with embedded HLSL guide preparation
and output composition. Preserve the vkd3d compatibility route for
RDNA3 and expose explicit route overrides and diagnostics.

Improve Windows runtime compatibility:
- Retry rejected application-ID initialization with a project identity.
- Defer expensive RR backend creation until feature creation.
- Recover from transient submission failures without disabling later frames.
- Add shared-buffer export fallbacks through explicit heaps and a private
  device on the same adapter.
- Safely evict old RR descriptor sets when input views change.
- Treat requested RR presets as hints when a configured model is available.
- Expand SR/RR input formats and handle diffuse/specular hit distances,
  packed normals, depth-stencil resources, and unused FSR guides.

Reduce SR overhead by batching output downloads and format conversion
under one completion wait. Optimize strict M enc2/tube attention with
bounded FP32 accumulation and the original FP64 fallback for ambiguous
rounding cases.

Extend RDNA4 tooling with gfx1200/gfx1201 native FP8 RR model builds,
capability checks, architecture-aware deployment, and experimental K/M
kernel paths. Require matching model and runtime arithmetic families.

Add optional embedded FSR shader optimizations keyed by shader hash and
exact original SPIR-V contents. Preserve reflected metadata and record
optimizer provenance, artifact hashes, and cache identity.

Simplify Windows installation with automatic backend dependency downloads,
download progress and ETA, backed-up absolute OptiScaler library paths,
and recovery from interrupted RR model builds.

Add build-kernels.bat and a kernels-only build matrix for all eleven
registered RDNA targets, with code-object target validation and SHA256
manifests. Expand conversion probes, numerical stress tests, temporal
benchmark tooling, installer/deployment tests, and setup documentation.

Document synthetic regression and timing results, experimental build
options, and remaining RDNA4 hardware validation requirements.
2026-10-17 07:40:34 -04:00

895 lines
41 KiB
C++
Raw Blame History

// 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) {
o[i] = (float)exact_pv_dot(Qs + mrow(i) * PST, Vs + 16 * nt + l16);
}
#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
Reference in new issue View Git Blame Copy Permalink