woutermans
0fe3a28cb7
Per-layer staged-image engine for the eleven translated denoiser Swin layers: descriptors (verified offsets per phase), cut recipes that splice an isolated native phase into NVIDIA's PTX, wave templates, epilogue notes and GPU probes. The isolated phase-8 splice produces a byte-identical final 4K frame but is slower than the translated body (14.90 vs 6.59 ms) - a correctness experiment, not a shipping improvement; docs/rr-engine.md records the design and the next direct-converter win (B operands built from transposed accumulators).
1692 lines
72 KiB
Plaintext
1692 lines
72 KiB
Plaintext
// Native replacement for the Ray Reconstruction denoiser's first Swin encoder,
|
|
// cuda_dldn_engine_swin_enc0_kernel (nvngx_dlssd.dll 310.7).
|
|
//
|
|
// ---------------------------------------------------------------------------
|
|
// STATE: every phase of the kernel is implemented, in the PTX's own order. Each one is
|
|
// checked against the debug oracle (see "Ground truth" below) with its own oracle capture,
|
|
// and a stage counts as validated only when its whole payload region matches the
|
|
// translated kernel byte for byte. Phase 1's GEMM, its relu requantisation
|
|
// into phase 2's A operand, phase 2's GEMM, the LayerNorm tree (the sum-of-squares
|
|
// butterfly, the rsqrt and the activation gains) and the fused Q/V projection of phases 3
|
|
// and 4; the first attention window in phases 5 to 7 (the 64x96 score with its position
|
|
// bias, the softmax and P*V, the output projection with the residual and the per-column
|
|
// table); the second window in phases 8 to 12 (two more projections off the same normed
|
|
// tile, its own score table at weight bytes 23744, its P*V, and the output projection that
|
|
// folds phase 7's accumulators in as its C operand); the MLP in phases 13 to 20 (E12's RMS
|
|
// norm, four parallel branches that share its tile and each with the clamped cubic
|
|
// activation, their running residual); and the merge in phases 21 and 22 (E20's plane-arena
|
|
// scatter and eight-pair lane merge, E21's surface scatter, phase 22's projection).
|
|
//
|
|
// The epilogue E22 stays NVIDIA's: it lives outside the cut and reads the eighteen values
|
|
// this body hands back (the parameter block's first word, %tid.x and phase 22's eight mma
|
|
// D pairs).
|
|
//
|
|
// The whole body costs about 6 ms per launch on an RX 9070 XT (gfx1201) with its results
|
|
// kept alive by a debug store; as served, where the optimiser drops them again, it is the
|
|
// kernel's dominant cost.
|
|
//
|
|
// ---------------------------------------------------------------------------
|
|
// The kernel is one warp per 8x8 token patch: grid (161, 93, 1), block (32, 1, 1),
|
|
// 14973 blocks, parameter block 1032 bytes. The PTX this file replaces is
|
|
// kernels/extracted/rr/0013-PREPASS_ENTRYPOINT_NAME.ptx.
|
|
//
|
|
// kernels/tex/make_ptx.py's recipe (D4R_RRSWIN_CUT=1) keeps NVIDIA's texture
|
|
// prologue (lines 1-3963) and surface epilogue (lines 44528-end) and replaces the
|
|
// 536 mma of the 22 phases with the call below. Without that variable the recipe
|
|
// replaces the whole body instead (a served stub).
|
|
//
|
|
// Linked into ZLUDA's compile of that PTX (D4R_ZLUDA_EXTRA_BC), so tex_common.h's
|
|
// helpers and ZLUDA's own device functions are available. Compiled with -nogpuinc,
|
|
// so only plain clang builtins and attributes are used.
|
|
#include "../tex/tex_common.h"
|
|
#ifdef D4R_RRSWIN_NATIVE_GEMM
|
|
#include "rrswin_engine.h"
|
|
#endif
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Parameter block. Fields are addressed by their byte offset in the PTX, so the
|
|
// raw image is kept and read through a typed accessor.
|
|
struct DldnSwinEnc0Params
|
|
{
|
|
uint8_t raw[1032];
|
|
};
|
|
static_assert(sizeof(DldnSwinEnc0Params) == 1032, "parameter block");
|
|
|
|
enum : int
|
|
{
|
|
kOutPlane = 8, // plane arena the encoder writes (arg048)
|
|
kWeights = 40, // NVIDIA's prepared weight image (arg040), read-only
|
|
kSurfaceA = 112, // 1280x736x4 RGBA16F outputs (arg112/arg120)
|
|
kSurfaceB = 120,
|
|
kSurfaceHalf = 128, // 640x368x4 RGBA16F (arg128)
|
|
kSurfaceIn = 304, // 1280x720x4 RGBA16F, read-only (arg304)
|
|
kPlaneArg = 48, // plane arena pointer, as make_ptx.py's debug splice reads it
|
|
};
|
|
|
|
#define DEVF __attribute__((device)) static inline
|
|
typedef float lds_f32 __attribute__((address_space(3)));
|
|
typedef uint32_t lds_u32 __attribute__((address_space(3)));
|
|
typedef uint8_t lds_u8 __attribute__((address_space(3)));
|
|
|
|
#ifdef D4R_RRSWIN_NATIVE_GEMM
|
|
// A token tile (512 bytes), two prepared weight tiles (1024 bytes), and their
|
|
// f16 result (1024 bytes) fit within the existing 3200-byte shared scratch.
|
|
DEVF void gemm_native_unbiased(const uint32_t* as, const uint32_t* bs,
|
|
int nmi, uint32_t scratch, lds_u32* out)
|
|
{
|
|
lds_u8* const shared = (lds_u8*)(uintptr_t)scratch;
|
|
uint8_t* const image = (uint8_t*)shared;
|
|
uint8_t* const weights = (uint8_t*)(shared + 512);
|
|
__attribute__((address_space(3))) half_t* const result =
|
|
(__attribute__((address_space(3))) half_t*)(shared + 1536);
|
|
const uint32_t lane = lane_id();
|
|
#pragma unroll
|
|
for (int tile = 0; tile < 2; ++tile)
|
|
#pragma unroll
|
|
for (int word = 0; word < 4; ++word)
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j)
|
|
{
|
|
const uint32_t n = 16 * tile + (lane >> 2) + 8 * (word >> 1);
|
|
const uint32_t k = 2 * (lane & 3) + (j & 1) + 16 * (j >> 1) + 8 * (word & 1);
|
|
weights[16 * (32 * (k >> 4) + n) + (k & 15)] =
|
|
(uint8_t)(bs[4 * tile + word] >> (8 * j));
|
|
}
|
|
__builtin_amdgcn_wave_barrier();
|
|
#pragma unroll
|
|
for (int mi = 0; mi < nmi; ++mi)
|
|
{
|
|
rrswin::swin_stage_fragment(image, 32, as + 4 * mi, 1);
|
|
__builtin_amdgcn_wave_barrier();
|
|
#pragma unroll
|
|
for (int mt = 0; mt < 2; ++mt)
|
|
{
|
|
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
|
|
c = rrswin::swin_gemm_tile(c, (const kslot*)weights, 16 * mt, 32, image, 32, 0);
|
|
#pragma unroll
|
|
for (int i = 0; i < 8; ++i)
|
|
result[32 * (lane & 15) + 16 * mt + i + 8 * (lane >> 4)] = (half_t)c[i];
|
|
}
|
|
__builtin_amdgcn_wave_barrier();
|
|
#pragma unroll
|
|
for (int rh = 0; rh < 2; ++rh)
|
|
#pragma unroll
|
|
for (int nt = 0; nt < 4; ++nt)
|
|
{
|
|
const int at = 32 * ((lane >> 2) + 8 * rh) + 8 * nt + 2 * (lane & 3);
|
|
out[8 * mi + 2 * nt + rh] = rrswin::pair_of(result[at], result[at + 1]);
|
|
}
|
|
__builtin_amdgcn_wave_barrier();
|
|
}
|
|
}
|
|
#endif
|
|
|
|
#ifdef D4R_RRSWIN_PHASE8
|
|
extern "C" __attribute__((device)) void d4r_swin_enc0_p8(
|
|
uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3,
|
|
uint32_t b0, uint32_t b1, uint32_t b2, uint32_t b3,
|
|
uint32_t b4, uint32_t b5, uint32_t b6, uint32_t b7, uint32_t scratch)
|
|
{
|
|
const uint32_t as[4] = {a0, a1, a2, a3};
|
|
const uint32_t bs[8] = {b0, b1, b2, b3, b4, b5, b6, b7};
|
|
// All WMMA reads finish before the lane-major return overwrites operand
|
|
// storage. The result matrix at +1536 remains disjoint from this return.
|
|
lds_u32* const out = (lds_u32*)(uintptr_t)(scratch + 32 * lane_id());
|
|
gemm_native_unbiased(as, bs, 1, scratch, out);
|
|
}
|
|
#endif
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// The implemented phases, as byte offsets into arg040 (from the PTX's own loads;
|
|
// K=32 and N=32, so each reads two 512-byte weight tiles):
|
|
// P1 weights 4096, bias seeds handed in as b0..b3 (loaded from +5120..+5184)
|
|
// P2 weights 5184, f16 bias at 6208..6272, activation gains at 6272..6336
|
|
// P3 weights 6336, no bias (the fused Q projection)
|
|
// P4 weights 7360, no bias (V)
|
|
|
|
// e4m3 code -> f32, exact. The mma's e4m3 operands reach the FMA path through
|
|
// ZLUDA's exact e4m3 -> f16 widening, so this decode is the same value.
|
|
DEVF float e4m3_to_f32(uint32_t code)
|
|
{
|
|
const uint32_t mag = code & 0x7fu;
|
|
const uint32_t bits = mag >= 8u
|
|
? ((((mag >> 3) + 120u) << 23) | ((mag & 7u) << 20))
|
|
: __builtin_bit_cast(uint32_t, (float)mag * 0x1p-9f);
|
|
return __builtin_bit_cast(float, bits | ((code & 0x80u) << 24));
|
|
}
|
|
|
|
// byte j of a lane's 16-byte weight tile holds (k, n) at
|
|
// addr(n, k) = base + 512*(n>>4) + 64*(n&7) + 16*((k>>1)&3) + 8*((n>>3)&1)
|
|
// + 4*((k>>3)&1) + 2*((k>>4)&1) + (k&1)
|
|
// (VERIFIED on all 1024 (n, k) pairs of a tile; the 64*(n&7) term is what keeps
|
|
// columns 16..31 from silently reading as zero.)
|
|
DEVF uint32_t waddr(int base, int k, int n)
|
|
{
|
|
return (uint32_t)base + 512u * (uint32_t)(n >> 4) + 64u * (uint32_t)(n & 7)
|
|
+ 16u * (uint32_t)((k >> 1) & 3) + 8u * (uint32_t)((n >> 3) & 1)
|
|
+ 4u * (uint32_t)((k >> 3) & 1) + 2u * (uint32_t)((k >> 4) & 1)
|
|
+ (uint32_t)(k & 1);
|
|
}
|
|
|
|
// warp butterfly over the four lanes that hold a row's 32 columns (laneid ^ 1 and
|
|
// laneid ^ 2): the two shfl.sync.bfly steps the PTX's reductions perform.
|
|
DEVF uint32_t bfly1_bits(uint32_t v)
|
|
{
|
|
return (uint32_t)__builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xf, 0xf, false);
|
|
}
|
|
|
|
DEVF uint32_t bfly2_bits(uint32_t v)
|
|
{
|
|
return (uint32_t)__builtin_amdgcn_update_dpp(0, (int)v, 0x4E, 0xf, 0xf, false);
|
|
}
|
|
|
|
DEVF float bfly1(float v)
|
|
{
|
|
return __builtin_bit_cast(float, bfly1_bits(__builtin_bit_cast(uint32_t, v)));
|
|
}
|
|
|
|
DEVF float bfly2(float v)
|
|
{
|
|
return __builtin_bit_cast(float, bfly2_bits(__builtin_bit_cast(uint32_t, v)));
|
|
}
|
|
|
|
// the low f16 of a packed pair (column 2t)
|
|
DEVF half_t half_lo(uint32_t pair)
|
|
{
|
|
return __builtin_bit_cast(half_t, (uint16_t)pair);
|
|
}
|
|
|
|
// the high f16 of a packed pair (column 2t+1)
|
|
DEVF half_t half_hi(uint32_t pair)
|
|
{
|
|
return __builtin_bit_cast(half_t, (uint16_t)(pair >> 16));
|
|
}
|
|
|
|
// f16 bits -> f32
|
|
DEVF float h2f(uint32_t bits)
|
|
{
|
|
return (float)__builtin_bit_cast(half_t, (uint16_t)bits);
|
|
}
|
|
|
|
// add.f16x2 on a packed pair (the PTX's reduction and requantisation steps)
|
|
DEVF uint32_t pair_add(uint32_t a, uint32_t b)
|
|
{
|
|
return __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, a) + __builtin_bit_cast(hv2, b));
|
|
}
|
|
|
|
DEVF uint32_t pair_of(half_t a, half_t b)
|
|
{
|
|
return (uint32_t)__builtin_bit_cast(uint16_t, a)
|
|
| ((uint32_t)__builtin_bit_cast(uint16_t, b) << 16);
|
|
}
|
|
|
|
// ZLUDA's relu_keep_nan_halves: a negative (or -0) half becomes +0, NaN is kept.
|
|
DEVF uint32_t relu_pair(uint32_t p)
|
|
{
|
|
if ((p & 0x8000u) != 0 && (p & 0x7fffu) <= 0x7c00u) {
|
|
p &= 0xffff0000u;
|
|
}
|
|
if ((p & 0x80000000u) != 0 && ((p >> 16) & 0x7fffu) <= 0x7c00u) {
|
|
p &= 0x0000ffffu;
|
|
}
|
|
return p;
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// One GEMM: 96 rows (six m-tiles of 16) x 32 columns, K=32, f16 output.
|
|
//
|
|
// A lane's A registers hold, for each m-tile and each of its two rows, the eight
|
|
// e4m3 values of k = 2t + (j&1) + 16*((j>>1)&1) (byte j = 0..3) and of k + 8; the
|
|
// four lanes that differ in t hold the four disjoint eighths of the 32 k's, so two
|
|
// butterfly steps broadcast the other three eighths in. The B operand is read from
|
|
// the weight image (waddr) for the lane's own two columns of each n-tile. The dot
|
|
// product runs in f32 and is rounded to f16 once, which is what the translated mma
|
|
// does: e4m3 x e4m3 products are exact and 32 of them sum exactly in f32, so this
|
|
// reproduces the translated kernel bit for bit.
|
|
//
|
|
// Exactly one of the three C operands carries the bias: `wbias` (a byte pointer at the phase's
|
|
// bias offset, 16 bytes per n-tile), `seeds` (four already lane-indexed packed pairs, phase 1),
|
|
// or `cmat` (32 lane-indexed pairs indexed [4*m-tile + row][n-tile], phase 7, whose seed is
|
|
// or `cmat` (32 lane-indexed pairs indexed [4*m-tile + row][n-tile], phases 7 and 12-20,
|
|
// whose seed is the residual plus the per-column table and so differs per row).
|
|
// d is [m-tile][row half][n-tile], each entry the packed f16 pair for columns 2t, 2t+1;
|
|
// `nnt` is the number of n-tiles (four everywhere but phase 22's two).
|
|
DEVF void gemm(const uint32_t* as, const uint8_t* w, int base, const uint8_t* wbias,
|
|
const uint32_t* seeds, const uint32_t* cmat, int nmi, int t,
|
|
uint32_t d[6][2][4], int nnt = 4)
|
|
{
|
|
for (int nt = 0; nt < nnt; ++nt) {
|
|
float bv[2][32];
|
|
for (int c = 0; c < 2; ++c) {
|
|
const int n = 8 * nt + 2 * t + c;
|
|
for (int k = 0; k < 32; ++k) {
|
|
bv[c][k] = e4m3_to_f32(w[waddr(base, k, n)]);
|
|
}
|
|
}
|
|
for (int mi = 0; mi < nmi; ++mi) {
|
|
uint32_t fixed = 0;
|
|
if (wbias != nullptr) {
|
|
fixed = *(const uint32_t*)(wbias + 16 * nt + 4 * t);
|
|
} else if (seeds != nullptr) {
|
|
fixed = seeds[nt];
|
|
}
|
|
float acc[2][2] = {{0.0f, 0.0f}, {0.0f, 0.0f}};
|
|
for (int reg = 0; reg < 2; ++reg) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
// the register's value on this lane and on lanes t^1, t^2, t^3
|
|
uint32_t p[4];
|
|
p[0] = as[4 * mi + 2 * reg + row];
|
|
p[1] = bfly1_bits(p[0]);
|
|
p[2] = bfly2_bits(p[0]);
|
|
p[3] = bfly1_bits(p[2]);
|
|
for (int m = 0; m < 4; ++m) {
|
|
const int kk = 2 * (t ^ m);
|
|
for (int j = 0; j < 4; ++j) {
|
|
const int k = 8 * reg + kk + (j & 1) + 16 * ((j >> 1) & 1);
|
|
const float av = e4m3_to_f32((p[m] >> (8 * j)) & 0xffu);
|
|
acc[row][0] += av * bv[0][k];
|
|
acc[row][1] += av * bv[1][k];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
for (int row = 0; row < 2; ++row) {
|
|
const uint32_t bias_pair = (cmat != nullptr) ? cmat[(4 * mi + row) * 4 + nt] : fixed;
|
|
d[mi][row][nt] = pair_of((half_t)(acc[row][0] + h2f(bias_pair)),
|
|
(half_t)(acc[row][1] + h2f(bias_pair >> 16)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Phases 5-7: the window's self-attention and the output projection. Both operands of
|
|
// phase 5 come from phase 3's accumulators (rows 0..63 are the queries, rows 0..95 the keys
|
|
// -- the 8x8 token patch plus its 32-row halo), so A and B are built exactly the way the
|
|
// PTX's 24 packs build them: word 0 of a row packs n-tiles 0 and 2, word 1 packs 1 and 3,
|
|
// and the four lanes t=0..3 of a group hold the four disjoint eighths of the 32 k's, which
|
|
// the two butterfly steps bring together. Every e4m3 product is exact and a 32-term k
|
|
// slice sums exactly in f32, so the dot product is plain f32 accumulation rounded to f16
|
|
// once -- what the translated mma does.
|
|
|
|
// the f16 the PTX's literals carry: 0f3C8CCB50, 0dBFE1CC0000000000, 0f3F6D6000, 0f3FB00000
|
|
DEVF half_t hf16(uint16_t bits)
|
|
{
|
|
return __builtin_bit_cast(half_t, bits);
|
|
}
|
|
|
|
DEVF hv2 hv2_of(uint32_t p)
|
|
{
|
|
return __builtin_bit_cast(hv2, p);
|
|
}
|
|
|
|
DEVF uint32_t u32_of(hv2 v)
|
|
{
|
|
return __builtin_bit_cast(uint32_t, v);
|
|
}
|
|
|
|
// ZLUDA's own rcp.approx.ftz.f32 helper: calling it is what makes the reciprocal bit-exact.
|
|
extern "C" __attribute__((device)) float __zluda_ptx_impl_rcp_approx_f32(float);
|
|
|
|
// the e4m3 code of one f16 (cvt.rn.satfinite.e4m3x2.f16x2 of a pair of equal halves)
|
|
DEVF uint32_t q8(half_t v)
|
|
{
|
|
return e4m3x2(pair_of(v, v));
|
|
}
|
|
|
|
// the f16 at (row, col) of a 96x32 block of packed pairs, col = 8*n-tile + 2t + c
|
|
DEVF half_t dcell(const uint32_t d[6][2][4], int row, int col)
|
|
{
|
|
const uint32_t p = d[row >> 4][(row >> 3) & 1][col >> 3];
|
|
return (col & 1) ? half_hi(p) : half_lo(p);
|
|
}
|
|
|
|
// the lane's two e4m3-packed words for row r: word 0 packs n-tiles 0 and 2, word 1 tiles 1 and 3
|
|
DEVF void row_words(const uint32_t d[6][2][4], int r, uint32_t w[2])
|
|
{
|
|
w[0] = e4m3x2(pair_of(dcell(d, r, 0), dcell(d, r, 1)))
|
|
| (e4m3x2(pair_of(dcell(d, r, 16), dcell(d, r, 17))) << 16);
|
|
w[1] = e4m3x2(pair_of(dcell(d, r, 8), dcell(d, r, 9)))
|
|
| (e4m3x2(pair_of(dcell(d, r, 24), dcell(d, r, 25))) << 16);
|
|
}
|
|
|
|
// out[i][m]: word i as lane t^m holds it, which is what the mma pairs across the group
|
|
DEVF void bfly2x4(const uint32_t w[2], uint32_t out[2][4])
|
|
{
|
|
#pragma unroll
|
|
for (int i = 0; i < 2; ++i) {
|
|
out[i][0] = w[i];
|
|
out[i][1] = bfly1_bits(w[i]);
|
|
out[i][2] = bfly2_bits(w[i]);
|
|
out[i][3] = bfly1_bits(out[i][2]);
|
|
}
|
|
}
|
|
|
|
|
|
// word i at butterfly index m (0..3), without the dynamic index into the pair
|
|
DEVF uint32_t word_at(const uint32_t out[2][4], int i, int m)
|
|
{
|
|
return m == 0 ? out[i][0] : (m == 1 ? out[i][1] : (m == 2 ? out[i][2] : out[i][3]));
|
|
}
|
|
|
|
// the 32-term dot product of two packed operand word pairs, in the mma's own k order
|
|
DEVF float dot32(const uint32_t a[2][4], const uint32_t b[2][4])
|
|
{
|
|
float acc = 0.0f;
|
|
#pragma unroll
|
|
for (int i = 0; i < 2; ++i) {
|
|
#pragma unroll
|
|
for (int m = 0; m < 4; ++m) {
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j) {
|
|
acc += e4m3_to_f32((a[i][m] >> (8 * j)) & 0xffu)
|
|
* e4m3_to_f32((b[i][m] >> (8 * j)) & 0xffu);
|
|
}
|
|
}
|
|
}
|
|
return acc;
|
|
}
|
|
|
|
// E5(a), one element: the scale, the clamp, the two clamped-cubic fma pairs and the trick
|
|
// that moves the mantissa into the exponent field ((u << 5) & 0x7FE07FE0 over the pair).
|
|
DEVF uint32_t expval_pair(uint32_t score)
|
|
{
|
|
hv2 m = hv2_of(score) * __builtin_bit_cast(hv2, pair_of(hf16(0x2466), hf16(0x2466)));
|
|
m = __builtin_elementwise_max(m, __builtin_bit_cast(hv2, pair_of(hf16(0xb873), hf16(0xb873))));
|
|
m = __builtin_elementwise_min(m, __builtin_bit_cast(hv2, pair_of(hf16(0x3873), hf16(0x3873))));
|
|
const hv2 t = __builtin_elementwise_fma(m, -m,
|
|
__builtin_bit_cast(hv2, pair_of(hf16(0x3b6b), hf16(0x3b6b))));
|
|
const uint32_t u = u32_of(__builtin_elementwise_fma(m, t,
|
|
__builtin_bit_cast(hv2, pair_of(hf16(0x3d80), hf16(0x3d80)))));
|
|
return (u << 5) & 0x7fe07fe0u;
|
|
}
|
|
|
|
// E5(a)-(e), one row: the twelve exponent pairs, the in-lane sum tree
|
|
// (((e0+e2)+(e4+e6))+(e8+e10)) + (((e1+e3)+(e5+e7))+(e9+e11)), the two butterfly steps that
|
|
// finish the row's 96 columns, the two halves' add, the reciprocal and the twelve products.
|
|
DEVF void row_probs(const uint32_t d5[4][2][12], int mi, int rh, uint32_t prob[12])
|
|
{
|
|
hv2 ev[6], eo[6];
|
|
#pragma unroll
|
|
for (int i = 0; i < 6; ++i) {
|
|
ev[i] = hv2_of(expval_pair(d5[mi][rh][2 * i]));
|
|
eo[i] = hv2_of(expval_pair(d5[mi][rh][2 * i + 1]));
|
|
}
|
|
const hv2 even = ((ev[0] + ev[1]) + (ev[2] + ev[3])) + (ev[4] + ev[5]);
|
|
const hv2 odd = ((eo[0] + eo[1]) + (eo[2] + eo[3])) + (eo[4] + eo[5]);
|
|
hv2 s = even + odd;
|
|
s = s + hv2_of(bfly1_bits(u32_of(s)));
|
|
s = s + hv2_of(bfly2_bits(u32_of(s)));
|
|
const uint32_t p = u32_of(s);
|
|
// the PTX broadcasts the f16 sum of the two halves (add.f16) and takes cvt.f32.f16 of it:
|
|
// the value, not its bits
|
|
const half_t r = (half_t)__zluda_ptx_impl_rcp_approx_f32((float)(half_lo(p) + half_hi(p)));
|
|
const uint32_t rcp = pair_of(r, r);
|
|
#pragma unroll
|
|
for (int i = 0; i < 6; ++i) {
|
|
prob[2 * i] = u32_of(ev[i] * hv2_of(rcp));
|
|
prob[2 * i + 1] = u32_of(eo[i] * hv2_of(rcp));
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// The two GEMMs the two attention windows share: the 64x96 score (phases 5 and 10) and
|
|
// the softmax plus P·V that follow it (phases 6 and 11). The second window repeats the
|
|
// first one's code with another weight base and another source accumulator, so each is
|
|
// one function.
|
|
|
|
// The score block. A and B are both the one 96x32 accumulator `dq`: the lane's rows
|
|
// 16*mi + g + 8*row are the queries (the 8x8 token patch, 64 rows) and its rows
|
|
// 16*(nb>>1) + 8*(nb&1) + g the keys (the patch plus a 32-row halo, 96 rows), which is
|
|
// how the PTX's 24 packs build the two operands out of the same D. Every lane publishes
|
|
// its two key words per n-tile into `slot` (twelve n-tiles x 64 words = 3072 bytes, inside
|
|
// the dead input_tensor the cut hands over) and then forms the dots its own D registers
|
|
// hold: query row 16*mi + g + 8*row (the A side keeps the lane's own row group g) against
|
|
// key rows 16*(nb>>1) + 8*(nb&1) + 2t + c, published by the lanes 4*(2t + c) + 0..3. The
|
|
// C operand is the phase's 64x96 f16 position-bias table at `bias` (weight bytes 8384 or
|
|
// 23744): 24 tiles of 512 bytes and 16 per lane, the seed of (m-tile, n-tile) in tile
|
|
// mi + 4*(nb>>1) at word 2*(nb&1) + row of the lane's eight f16 (VERIFIED against the
|
|
// oracle's own C registers, debug stage 11 for phase 5 and 204 for phase 10).
|
|
// Every e4m3 product is exact and a 32-term k slice sums exactly in f32, so the dot
|
|
// product is plain f32 accumulation rounded to f16 once -- what the translated mma does.
|
|
//
|
|
// `cap`, when the debug stage asks for the published key words (203) or the C and B
|
|
// operand registers (204), receives them in the oracle's order.
|
|
DEVF void score_gemm(const uint32_t dq[6][2][4], const uint8_t* w, int bias, int t, int g,
|
|
int lane, lds_u32* slot, uint32_t* cap, uint32_t d[4][2][12])
|
|
{
|
|
#pragma unroll
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
uint32_t kw[2];
|
|
row_words(dq, 16 * (nb >> 1) + 8 * (nb & 1) + g, kw);
|
|
slot[64 * nb + lane] = kw[0];
|
|
slot[64 * nb + 32 + lane] = kw[1];
|
|
}
|
|
__builtin_amdgcn_wave_barrier();
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 203
|
|
if (cap != nullptr) {
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
cap[2 * nb] = slot[64 * nb + 4 * (2 * t)];
|
|
cap[2 * nb + 1] = slot[64 * nb + 32 + 4 * (2 * t + 1)];
|
|
}
|
|
}
|
|
#endif
|
|
#pragma unroll
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 204
|
|
if (cap != nullptr) {
|
|
// the mma's A fragment: rows g and g+8 of m-tile mi, word 0 (k = 2t + (j&1) +
|
|
// 16*((j>>1)&1), i.e. n-tiles 0 and 2) then word 1 (n-tiles 1 and 3)
|
|
uint32_t q0[2], q1[2];
|
|
row_words(dq, 16 * mi + g, q0);
|
|
row_words(dq, 16 * mi + g + 8, q1);
|
|
cap[4 * mi + 0] = q0[0];
|
|
cap[4 * mi + 1] = q1[0];
|
|
cap[4 * mi + 2] = q0[1];
|
|
cap[4 * mi + 3] = q1[1];
|
|
}
|
|
#endif
|
|
for (int row = 0; row < 2; ++row) {
|
|
uint32_t aw[2], a[2][4];
|
|
row_words(dq, 16 * mi + g + 8 * row, aw);
|
|
bfly2x4(aw, a);
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
float s[2];
|
|
for (int c = 0; c < 2; ++c) {
|
|
uint32_t kb[2][4];
|
|
#pragma unroll
|
|
for (int i = 0; i < 2; ++i) {
|
|
// The mma's eight n-index rows are the key rows
|
|
// 16*(nb>>1) + 8*(nb&1) + n, and the lane's two D columns c = 0, 1
|
|
// are n = 2*t + c: each half dots a *different* key row, not the two
|
|
// halves of one. Its B words are the four lanes 4*(2*t + c) + (t^m)
|
|
// of that row's group (the same butterfly the A side walks, on the
|
|
// row group the D column selects).
|
|
#pragma unroll
|
|
for (int m = 0; m < 4; ++m) {
|
|
kb[i][m] = slot[64 * nb + 32 * i + 4 * (2 * t + c) + (t ^ m)];
|
|
}
|
|
}
|
|
s[c] = dot32(a, kb);
|
|
}
|
|
const uint32_t b =
|
|
*(const uint32_t*)(w + bias + 512 * (mi + 4 * (nb >> 1))
|
|
+ 16 * lane + 8 * (nb & 1) + 4 * row);
|
|
d[mi][row][nb] = pair_of((half_t)(s[0] + h2f(b)),
|
|
(half_t)(s[1] + h2f(b >> 16)));
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 204
|
|
if (cap != nullptr) {
|
|
// the position bias, the 24 tiles of four words after the A and B
|
|
cap[40 + 4 * (mi + 4 * (nb >> 1)) + 2 * (nb & 1) + row] = b;
|
|
}
|
|
#elif D4R_RRSWIN_DEBUG_STAGE == 211
|
|
if (cap != nullptr) {
|
|
cap[4 * (mi + 4 * (nb >> 1)) + 2 * (nb & 1) + row] = b;
|
|
}
|
|
#endif
|
|
}
|
|
}
|
|
}
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 204
|
|
if (cap != nullptr) {
|
|
// the twelve B fragment pairs, as the mma reads them: the lane's two packed words
|
|
// of the key row 16*(nb>>1) + 8*(nb&1) + g, which is what it published in `slot`
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
cap[16 + 2 * nb] = slot[64 * nb + lane];
|
|
cap[16 + 2 * nb + 1] = slot[64 * nb + 32 + lane];
|
|
}
|
|
}
|
|
#elif D4R_RRSWIN_DEBUG_STAGE == 211
|
|
if (cap != nullptr) {
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
cap[96 + 2 * nb] = slot[64 * nb + lane];
|
|
cap[96 + 2 * nb + 1] = slot[64 * nb + 32 + lane];
|
|
}
|
|
}
|
|
#endif
|
|
}
|
|
|
|
// E5 and E10 with the P·V of phases 6 and 11. row_probs() gives the twelve probability
|
|
// pairs of one row; k-step ks packs key columns {4ks, 4ks+2} -- the columns of prob[4ks]
|
|
// and prob[4ks+2], low half first -- into the fragment's first word and {4ks+1, 4ks+3}
|
|
// into its second. The lane's A fragment is its *own* two rows (16*mi + g and
|
|
// 16*mi + g + 8), so the D row it accumulates is its own g; what it lacks is V. The key
|
|
// rows 32*ks + kk live in the lanes whose g is kk % 8, and every lane needs all 96 of
|
|
// them, so V goes through the dead input_tensor as one e4m3 byte per element
|
|
// (96 x 32 = 3072 bytes) and each lane then walks the full k range. The P byte for key
|
|
// column kk is the one the lane with t = (kk%8)>>1 packs: word (kk>>3)&1, byte
|
|
// 2*(kk>>4) + (kk&1) -- the two butterfly steps of bfly2x4 bring the other three eighths
|
|
// in. Each k-step rounds to f16 before the next one adds to it, which is why the chain
|
|
// has three separate mma.
|
|
//
|
|
// `cap`, when the debug stage asks for the A fragments (201), receives the twelve packed
|
|
// word pairs in the oracle's order.
|
|
DEVF void pv_gemm(const uint32_t ds[4][2][12], const uint32_t dv[6][2][4], int t, int g,
|
|
lds_u8* vs, uint32_t* cap, uint32_t d[4][2][4])
|
|
{
|
|
__builtin_amdgcn_wave_barrier(); // the score GEMM's reads of `slot` are done
|
|
#pragma unroll
|
|
for (int mi4 = 0; mi4 < 6; ++mi4) {
|
|
for (int rh4 = 0; rh4 < 2; ++rh4) {
|
|
#pragma unroll
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
for (int c = 0; c < 2; ++c) {
|
|
const int r = 16 * mi4 + 8 * rh4 + g;
|
|
vs[r * 32 + 8 * nt + 2 * t + c] =
|
|
(uint8_t)q8(dcell(dv, r, 8 * nt + 2 * t + c));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
__builtin_amdgcn_wave_barrier();
|
|
const lds_u8* const vp = vs;
|
|
#pragma unroll
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
#pragma unroll
|
|
for (int rh = 0; rh < 2; ++rh) {
|
|
uint32_t prob[12], aw[3][2];
|
|
row_probs(ds, mi, rh, prob);
|
|
#pragma unroll
|
|
for (int ks = 0; ks < 3; ++ks) {
|
|
aw[ks][0] = e4m3x2(prob[4 * ks]) | (e4m3x2(prob[4 * ks + 2]) << 16);
|
|
aw[ks][1] = e4m3x2(prob[4 * ks + 1]) | (e4m3x2(prob[4 * ks + 3]) << 16);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 201
|
|
if (cap != nullptr) {
|
|
cap[4 * (mi * 3 + ks) + rh] = aw[ks][0];
|
|
cap[4 * (mi * 3 + ks) + 2 + rh] = aw[ks][1];
|
|
}
|
|
#endif
|
|
}
|
|
#pragma unroll
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
half_t prev0 = hf16(0x0000), prev1 = hf16(0x0000);
|
|
#pragma unroll
|
|
for (int ks = 0; ks < 3; ++ks) {
|
|
uint32_t a[2][4];
|
|
bfly2x4(aw[ks], a);
|
|
float acc0 = 0.0f, acc1 = 0.0f;
|
|
#pragma unroll
|
|
for (int kk = 0; kk < 32; ++kk) {
|
|
const float pv = e4m3_to_f32(
|
|
(word_at(a, (kk >> 3) & 1, t ^ ((kk >> 1) & 3))
|
|
>> (8 * (2 * (kk >> 4) + (kk & 1)))) & 0xffu);
|
|
const int r = (32 * ks + kk) * 32 + 8 * nt + 2 * t;
|
|
acc0 += pv * e4m3_to_f32(vp[r]);
|
|
acc1 += pv * e4m3_to_f32(vp[r + 1]);
|
|
}
|
|
prev0 = (half_t)(acc0 + (float)prev0);
|
|
prev1 = (half_t)(acc1 + (float)prev1);
|
|
}
|
|
d[mi][rh][nt] = pair_of(prev0, prev1);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// The MLP's epilogues and the two operand helpers every mma in the entry shares.
|
|
|
|
// The sum-of-squares tree: the LayerNorm before phases 3 and 4 (PTX stmts 2399-3701) and
|
|
// E12's RMS norm before phase 13 (s12288-s13110), which are the same code with another
|
|
// gain table. Per row of the lane: the four n-tiles are squared and added as (n0+n2) then
|
|
// (n1+n3), both halves in f16; the butterfly finishes the sum over the row's 32 columns
|
|
// and the two halves are added with add.f16; +eps, rsqrt.approx.ftz.f32 and back to f16,
|
|
// times the activation gain of the column. The row's value times that factor is the A
|
|
// operand of the phase (cvt.rn.satfinite.e4m3x2.f16x2, no relu), packed as {n-tile 0,
|
|
// n-tile 2} into the k-group 0 register and {n-tile 1, n-tile 3} into k-group 1. There is
|
|
// no mean subtraction and no 1/N: the quantity is the raw sum of the 32 squares.
|
|
DEVF void norm_as(const uint32_t (*d)[2][4], const uint8_t* w, int gain, half_t eps,
|
|
int nmi, uint32_t as[24])
|
|
{
|
|
#pragma unroll
|
|
for (int mi = 0; mi < nmi; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
// the row's partial sums, packed for the two columns 2t and 2t+1
|
|
uint32_t x = 0;
|
|
for (int c = 0; c < 2; ++c) {
|
|
half_t sq[4];
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
const half_t v =
|
|
c == 0 ? half_lo(d[mi][row][nt]) : half_hi(d[mi][row][nt]);
|
|
sq[nt] = v * v;
|
|
}
|
|
const half_t s = (sq[0] + sq[2]) + (sq[1] + sq[3]);
|
|
x |= (uint32_t)__builtin_bit_cast(uint16_t, s) << (16 * c);
|
|
}
|
|
// the PTX's two shfl.sync.bfly steps over the four lanes that hold
|
|
// this row's 32 columns
|
|
x = pair_add(x, bfly1_bits(x));
|
|
x = pair_add(x, bfly2_bits(x));
|
|
const half_t sum = (half_t)(half_lo(x) + half_hi(x));
|
|
const half_t r = (half_t)__builtin_amdgcn_rsqf((float)(sum + eps));
|
|
const uint32_t* dr = d[mi][row];
|
|
uint32_t p0 = 0;
|
|
uint32_t p1 = 0;
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
uint32_t prod = 0;
|
|
for (int c = 0; c < 2; ++c) {
|
|
const half_t g = __builtin_bit_cast(half_t,
|
|
*(const uint16_t*)(w + gain + 16 * nt + 4 * (lane_id() & 3) + 2 * c));
|
|
const half_t v = (c == 0 ? half_lo(dr[nt]) : half_hi(dr[nt])) * (r * g);
|
|
prod |= (uint32_t)__builtin_bit_cast(uint16_t, v) << (16 * c);
|
|
}
|
|
const uint32_t code = e4m3x2(prod);
|
|
(nt & 1 ? p1 : p0) |= code << 16 * (nt >> 1);
|
|
}
|
|
as[4 * mi + row] = p0;
|
|
as[4 * mi + 2 + row] = p1;
|
|
}
|
|
}
|
|
}
|
|
|
|
// E13/E15/E17/E19, one packed pair: the clamped cubic activation, in f16 throughout. The
|
|
// PTX's own sequence (s13997.. / s33795..) is
|
|
// y = max(x, -2); y = min(y, +2); a = abs(y)
|
|
// g = mid - slope * a; g = y * g; g = 0.5 + g; out = x * g
|
|
// so the *signed* clamped y is the multiplier and only the inner `slope * a` term takes the
|
|
// absolute value; using |y| throughout would make every negative input's gain 0.5 + |y|(...)
|
|
// instead of 0.5 - |y|(...) and shift all four branches' activations.
|
|
DEVF uint32_t act_pair(uint32_t x)
|
|
{
|
|
const hv2 two = hv2_of(pair_of(hf16(0x4000), hf16(0x4000)));
|
|
const hv2 slope = hv2_of(pair_of(hf16(0x2d30), hf16(0x2d30)));
|
|
const hv2 mid = hv2_of(pair_of(hf16(0x3698), hf16(0x3698)));
|
|
const hv2 half0 = hv2_of(pair_of(hf16(0x3800), hf16(0x3800)));
|
|
hv2 y = __builtin_elementwise_max(hv2_of(x), -two);
|
|
y = __builtin_elementwise_min(y, two);
|
|
const hv2 a = __builtin_elementwise_max(y, -y); // abs.f16x2
|
|
const hv2 g = half0 + y * (mid - slope * a);
|
|
return u32_of(hv2_of(x) * g);
|
|
}
|
|
|
|
// E13/E15/E17/E19's pack: the activation of `d` requantised into the next phase's A operand.
|
|
// The A operand is four m-tiles of four registers: register 4*mi+row holds n-tiles 0 and 2,
|
|
// register 4*mi+2+row holds n-tiles 1 and 3, each as one e4m3x2 of the pair the mma reads.
|
|
DEVF void act_pack(const uint32_t d[4][2][4], uint32_t as[16])
|
|
{
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
const uint32_t a0 = e4m3x2(act_pair(d[mi][row][0]));
|
|
const uint32_t a1 = e4m3x2(act_pair(d[mi][row][1]));
|
|
const uint32_t a2 = e4m3x2(act_pair(d[mi][row][2]));
|
|
const uint32_t a3 = e4m3x2(act_pair(d[mi][row][3]));
|
|
as[4 * mi + row] = a0 | (a2 << 16);
|
|
as[4 * mi + 2 + row] = a1 | (a3 << 16);
|
|
}
|
|
}
|
|
}
|
|
|
|
// The next phase's C seeds: the block's residual plus the per-column table at `tab`
|
|
// (eight f16 per lane, the lane's 8 columns of all four n-tiles), added in f16.
|
|
// Phase 14's residual is *phase 12's* accumulators -- the MLP adds its input block back, so
|
|
// the running residual enters phase 14, not phase 13's own first GEMM (VERIFIED against the
|
|
// oracle's own C registers, debug stage p14o: 1024/1024 over four blocks).
|
|
DEVF void tab_seed(const uint32_t d[4][2][4], const uint8_t* tab, int lane, uint32_t c[64])
|
|
{
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
const uint32_t r = d[mi][row][nt];
|
|
const uint32_t tb = *(const uint32_t*)(tab + 16 * nt + ((lane << 2) & 12));
|
|
c[(4 * mi + row) * 4 + nt] = pair_of(
|
|
(half_t)(half_lo(tb) + half_lo(r)),
|
|
(half_t)(half_hi(tb) + half_hi(r)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// The four B fragment pairs of a phase (a weight-image B), as the mma reads them: the
|
|
// lane's own two words of tile nt>>1 at byte offset 8*(nt&1), which is its column
|
|
// 8*nt + g (VERIFIED against the oracle's own B registers, debug stage p8o).
|
|
DEVF void bpair_regs(const uint8_t* w, int base, int lane, uint32_t regs[8])
|
|
{
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
const uint8_t* const p = w + base + 512 * (nt >> 1) + 16 * lane + 8 * (nt & 1);
|
|
regs[2 * nt] = *(const uint32_t*)p;
|
|
regs[2 * nt + 1] = *(const uint32_t*)(p + 4);
|
|
}
|
|
}
|
|
|
|
// The four bias words of a phase, one per n-tile: the f16 pair of the lane's columns 2t
|
|
// and 2t+1, each replicated over the 16 rows of its tile and used by all four m-tiles.
|
|
DEVF void wbias_regs(const uint8_t* w, int base, int lane, uint32_t regs[4])
|
|
{
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[nt] = *(const uint32_t*)(w + base + 16 * nt + ((lane << 2) & 12));
|
|
}
|
|
}
|
|
|
|
// A phase's C seeds in the oracle's register order: the mma run visits (m-tile, n-tile)
|
|
// with n-tile fastest, and each of its C pairs is the (row g, row g+8) pair.
|
|
DEVF void cseed_regs(const uint32_t c[64], uint32_t regs[32])
|
|
{
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
regs[2 * (4 * mi + nt) + row] = c[(4 * mi + row) * 4 + nt];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// An accumulator block as a C operand: the mma's C pair for m-tile mi and n-tile nt is
|
|
// d[mi][row][nt] (VERIFIED for phase 12 against the oracle's own C registers, stage p12o).
|
|
DEVF void cres_regs(const uint32_t d[4][2][4], uint32_t c[64])
|
|
{
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
c[(4 * mi + row) * 4 + nt] = d[mi][row][nt];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Ground truth for a stage comes from kernels/tex/make_ptx.py's debug recipe:
|
|
// with D4R_RRSWIN_DEBUG=p1 (p2, p3, p4, or a) the same splice stores that stage's
|
|
// registers into the plane arena at +46,000,000, which the kernel never otherwise
|
|
// writes, and the launch dump then holds the translated values for a block/lane.
|
|
// Built with -DD4R_RRSWIN_DEBUG_STAGE=<n> (kernels/build.sh rrswin,
|
|
// D4R_TEX_CFLAGS), this body stores the same registers of its own phase n in the
|
|
// same order, so the two dumps compare byte for byte.
|
|
#ifndef D4R_RRSWIN_DEBUG_STAGE
|
|
#define D4R_RRSWIN_DEBUG_STAGE 0
|
|
#endif
|
|
|
|
DEVF void debug_store(const uint8_t* prm, const uint32_t* regs, int n)
|
|
{
|
|
const uint32_t block = __builtin_amdgcn_workgroup_id_x()
|
|
+ 161u * __builtin_amdgcn_workgroup_id_y();
|
|
if (block >= 64u) {
|
|
return; // the oracle's per-block slots only cover the first 64 blocks
|
|
}
|
|
const uint64_t arena = *(const uint64_t*)(prm + kPlaneArg);
|
|
// the oracle's slot, make_ptx.py's swin_debug_slot: 32 lanes x n x 4 bytes, at least 8192,
|
|
// rounded up to a power of two (the oracle's splice forms block*slot with a shift)
|
|
const uint64_t want = (uint64_t)n * 128ull > 8192ull ? (uint64_t)n * 128ull : 8192ull;
|
|
uint64_t slot = 1;
|
|
while (slot < want) {
|
|
slot <<= 1;
|
|
}
|
|
uint32_t* dst = (uint32_t*)(arena + 46000000ull + (uint64_t)block * slot
|
|
+ (uint64_t)lane_id() * (uint64_t)n * 4ull);
|
|
for (int i = 0; i < n; ++i) {
|
|
dst[i] = regs[i];
|
|
}
|
|
}
|
|
|
|
// stage registers in the oracle's order: 24 mma, m-tile major, n-tile minor, each
|
|
// as the (row g, row g+8) pair the mma writes.
|
|
DEVF void stage_regs(const uint32_t d[6][2][4], uint32_t regs[48])
|
|
{
|
|
int i = 0;
|
|
for (int mi = 0; mi < 6; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = d[mi][0][nt];
|
|
regs[i++] = d[mi][1][nt];
|
|
}
|
|
}
|
|
}
|
|
|
|
// the same for the sixteen-mma phases: m-tile major, n-tile minor, the (row g, row g+8)
|
|
// pair each mma writes.
|
|
DEVF void stage_regs4(const uint32_t d[4][2][4], uint32_t regs[32])
|
|
{
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = d[mi][0][nt];
|
|
regs[i++] = d[mi][1][nt];
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// The merge: E20's arena scatter and its eight shfl pairs, then phases 21 and 22.
|
|
|
|
// two u16 into one b32 (the PTX's mov.b32 of two cvt results)
|
|
DEVF uint32_t pair16(uint16_t a, uint16_t b)
|
|
{
|
|
return (uint32_t)a | ((uint32_t)b << 16);
|
|
}
|
|
|
|
// E20(b): the sixteen guarded stores of phase 20's requantised output into the plane
|
|
// arena. The index is a word index, idx = 8*ex*y + 8*x + d, where (ex, ey) are the
|
|
// extents of the parameter block's first word pair, x0 = 8*ctaid.x - param[80] and
|
|
// y0 = 8*(ctaid.y + tid.y) - param[82] are the block origins (the PTX's s6 and s16), and
|
|
// the lane's cell is x = x0 + (g & 7) with g = laneid>>2 and t = laneid&3: the eight
|
|
// stores of one half walk the eight rows y0 + ((g + 8*j)>>3) = y0 + j and the two words
|
|
// d = t and t|4 of them. A store whose y or x falls outside the arena is skipped (the
|
|
// PTX turns its index into -1 and branches).
|
|
DEVF void arena_scatter(const uint8_t* prm, const uint32_t r[16], int lane)
|
|
{
|
|
const uint32_t ex = *(const uint32_t*)prm;
|
|
const uint32_t ey = *(const uint32_t*)(prm + 4);
|
|
const int x0 = 8 * (int)__builtin_amdgcn_workgroup_id_x()
|
|
- (int)*(const uint16_t*)(prm + 80);
|
|
const int y0 = 8 * (int)(__builtin_amdgcn_workgroup_id_y()
|
|
+ __builtin_amdgcn_workitem_id_y())
|
|
- (int)*(const uint16_t*)(prm + 82);
|
|
uint8_t* const s = (uint8_t*)(uintptr_t)*(const uint64_t*)(prm + kPlaneArg);
|
|
const int x = x0 + ((lane >> 2) & 7);
|
|
const int t = lane & 3;
|
|
if ((unsigned)x >= ex) {
|
|
return;
|
|
}
|
|
const int stride = 8 * (int)ex;
|
|
for (int j = 0; j < 8; ++j) {
|
|
const int y = y0 + (((lane >> 2) + 8 * j) >> 3);
|
|
if ((unsigned)y >= ey) {
|
|
continue;
|
|
}
|
|
const int idx = stride * y + 8 * x;
|
|
*(uint32_t*)(s + 4 * (uint32_t)(idx + t)) = r[2 * j];
|
|
*(uint32_t*)(s + 4 * (uint32_t)(idx + (t | 4))) = r[2 * j + 1];
|
|
}
|
|
}
|
|
|
|
// E20(c): the eight shfl pairs, all with delta 4, that re-pair the eight lanes of every
|
|
// group of eight. For a source pair (lo, hi) the lane keeps its own `lo` and its
|
|
// partner's `hi` when it is in the lower group of four (laneid & 4 == 0) and the other way
|
|
// round when it is in the upper one, so each destination ends up holding both members:
|
|
// d1 = (lane & 4) == 0 ? lo[lane] : hi[lane - 4]
|
|
// d2 = (lane & 4) == 0 ? lo[lane + 4] : hi[lane]
|
|
// The eight pairs are (%r262,%r266) (%r264,%r268) (%r263,%r267) (%r265,%r269)
|
|
// (%r270,%r274) (%r272,%r276) (%r271,%r275) (%r273,%r277) and their sixteen results are
|
|
// phase 21's four A fragments: %r17490..%r17493, %r17510..%r17513, %r17530..%r17533 and
|
|
// %r17550..%r17553, in that order.
|
|
// The partner values go through the shared scratch (sixteen words per lane, 512 of the
|
|
// 3200 bytes the cut hands over and dead by now) rather than through two more warp
|
|
// shuffles: the exchange is the same read of the same lane's value at the same point.
|
|
DEVF void merge_frag(const uint32_t r[16], uint32_t scratch, int lane, uint32_t a[16])
|
|
{
|
|
const int lo[8] = {0, 2, 1, 3, 8, 10, 9, 11};
|
|
const int hi[8] = {4, 6, 5, 7, 12, 14, 13, 15};
|
|
// where each pair's first and second result land, as indices into a[]
|
|
const int d1[8] = {0, 8, 2, 10, 1, 9, 3, 11};
|
|
const int d2[8] = {4, 12, 6, 14, 5, 13, 7, 15};
|
|
lds_u32* const s = (lds_u32*)(uintptr_t)scratch;
|
|
#pragma unroll
|
|
for (int i = 0; i < 16; ++i) {
|
|
s[16 * lane + i] = r[i];
|
|
}
|
|
__builtin_amdgcn_wave_barrier();
|
|
const int partner = lane ^ 4;
|
|
const bool up = (lane & 4) != 0;
|
|
#pragma unroll
|
|
for (int i = 0; i < 8; ++i) {
|
|
a[d1[i]] = up ? s[16 * partner + hi[i]] : r[lo[i]];
|
|
a[d2[i]] = up ? s[16 * lane + hi[i]] : s[16 * partner + lo[i]];
|
|
}
|
|
}
|
|
|
|
// Phase 21: 16 rows x 64 columns contracting over K = 128 in four 32-wide steps. All 32
|
|
// mma read the same four A fragments (one per k-step, the merged `a`), eight n-tiles' B
|
|
// fragment pairs from the weight image at 45696, and eight C words at 53888 (one per
|
|
// n-tile, replicated over the tile's 16 rows). The A operand's row is the lane's own g
|
|
// and its column index 2*(lane & 3), so the four eighths of every k slice come from the
|
|
// four lanes of the row group and the two butterfly steps bring them in; the B operand's
|
|
// column index is the lane's own g, so the output column 2t + c is the B column of the
|
|
// lanes whose g is 2t + c -- which the weight image already holds, and each k-step rounds
|
|
// to f16 before the next one adds to it.
|
|
DEVF void band21(const uint32_t a[16], const uint8_t* w, const uint8_t* cbias, int lane,
|
|
uint32_t d[8][2])
|
|
{
|
|
const int t = lane & 3;
|
|
#pragma unroll
|
|
for (int nt = 0; nt < 8; ++nt) {
|
|
half_t prev[2][2] = {{hf16(0x0000), hf16(0x0000)}, {hf16(0x0000), hf16(0x0000)}};
|
|
const uint32_t cb = *(const uint32_t*)(cbias + 16 * nt + ((lane << 2) & 12));
|
|
#pragma unroll
|
|
for (int ks = 0; ks < 4; ++ks) {
|
|
// the lane's own A fragment of this k-step, as the four lanes of its row group
|
|
// hold it: register o is row o&1, k slice o>>1
|
|
uint32_t ar[2][2][4]; // [row][k slice][lane t^m]
|
|
#pragma unroll
|
|
for (int slice = 0; slice < 2; ++slice) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
const uint32_t v = a[4 * ks + 2 * slice + row];
|
|
ar[row][slice][0] = v;
|
|
ar[row][slice][1] = bfly1_bits(v);
|
|
ar[row][slice][2] = bfly2_bits(v);
|
|
ar[row][slice][3] = bfly1_bits(ar[row][slice][2]);
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int c = 0; c < 2; ++c) {
|
|
// the B fragment pairs of the four lanes whose g is the D column 2t + c
|
|
const int gp = 2 * t + c;
|
|
uint32_t br[2][4]; // [k slice][lane m]
|
|
#pragma unroll
|
|
for (int m = 0; m < 4; ++m) {
|
|
const uint8_t* const p = w + 45696 + 512 * (ks + 4 * (nt >> 1))
|
|
+ 16 * (4 * gp + m) + 8 * (nt & 1);
|
|
br[0][m] = *(const uint32_t*)p;
|
|
br[1][m] = *(const uint32_t*)(p + 4);
|
|
}
|
|
#pragma unroll
|
|
for (int row = 0; row < 2; ++row) {
|
|
float s2 = 0.0f;
|
|
#pragma unroll
|
|
for (int m = 0; m < 4; ++m) {
|
|
#pragma unroll
|
|
for (int slice = 0; slice < 2; ++slice) {
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j) {
|
|
s2 += e4m3_to_f32((ar[row][slice][m] >> (8 * j)) & 0xffu)
|
|
* e4m3_to_f32((br[slice][m] >> (8 * j)) & 0xffu);
|
|
}
|
|
}
|
|
}
|
|
prev[row][c] = (half_t)(s2 + (ks == 0 ? h2f(cb) : 0.0f)
|
|
+ (float)prev[row][c]);
|
|
}
|
|
}
|
|
}
|
|
d[nt][0] = pair_of(prev[0][0], prev[0][1]);
|
|
d[nt][1] = pair_of(prev[1][0], prev[1][1]);
|
|
}
|
|
}
|
|
|
|
// E21: phase 21's requantisation and its eight guarded stores into the buffer the
|
|
// parameter block holds at +56 -- not the plane arena. The extents come from the first
|
|
// word pair halved, the origin from %r1 and %r6 halved, and the two index bases are
|
|
// base1 = 8*(ex>>1)*y1 + 8*x1 with y1 = (%r6>>1) + ((g&1)<<1),
|
|
// x1 = (%r1>>1) + ((laneid>>3) & 3)
|
|
// base2 = 8*(ex>>1)*y2 + 8*x2 with y2 = y1 + 1, x2 = (%r1>>1) + ((g>>1) & 3)
|
|
// (the laneid>>5 term of y1 and the tile-bit test of the guard are both zero for a
|
|
// 32-lane block), each store landing at base + t or base + (t|4) words, the last four
|
|
// of them one row-stride further on. The sixteen payloads are phase 21's D registers in
|
|
// mma order, one e4m3x2 each: the first and the last two travel as v2.u16 and the middle
|
|
// twelve as six u32.
|
|
DEVF void surface_scatter21(const uint8_t* prm, const uint32_t d[8][2], int lane)
|
|
{
|
|
const int r1 = 8 * (int)__builtin_amdgcn_workgroup_id_x()
|
|
- (int)*(const uint16_t*)(prm + 80);
|
|
const int r6 = 8 * (int)(__builtin_amdgcn_workgroup_id_y()
|
|
+ __builtin_amdgcn_workitem_id_y())
|
|
- (int)*(const uint16_t*)(prm + 82);
|
|
const int ex = (int)(*(const uint32_t*)prm >> 1);
|
|
const int ey = (int)(*(const uint32_t*)(prm + 4) >> 1);
|
|
const int g = lane >> 2;
|
|
const int t = lane & 3;
|
|
const int sx = 8 * ex;
|
|
const int y1 = (r6 >> 1) + ((g & 1) << 1);
|
|
const int y2 = y1 + 1;
|
|
const int x1 = (r1 >> 1) + ((lane >> 3) & 3);
|
|
const int x2 = (r1 >> 1) + ((g >> 1) & 3);
|
|
uint8_t* const dst = (uint8_t*)(uintptr_t)*(const uint64_t*)(prm + 56);
|
|
uint16_t q[16];
|
|
#pragma unroll
|
|
for (int nt = 0; nt < 8; ++nt) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
q[2 * nt + row] = (uint16_t)e4m3x2(d[nt][row]);
|
|
}
|
|
}
|
|
const int b1 = sx * y1 + 8 * x1;
|
|
const int b2 = sx * y2 + 8 * x2;
|
|
const int st = sx * ey;
|
|
const bool ok1 = (unsigned)y1 < (unsigned)ey && (unsigned)x1 < (unsigned)ex;
|
|
const bool ok2 = (unsigned)y2 < (unsigned)ey && (unsigned)x2 < (unsigned)ex;
|
|
if (ok1) {
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b1 + t)) = pair16(q[0], q[1]);
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b1 + (t | 4))) = pair16(q[2], q[3]);
|
|
}
|
|
if (ok2) {
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b2 + t)) = pair16(q[4], q[5]);
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b2 + (t | 4))) = pair16(q[6], q[7]);
|
|
}
|
|
if (ok1) {
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b1 + st + t)) = pair16(q[8], q[9]);
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b1 + st + (t | 4))) = pair16(q[10], q[11]);
|
|
}
|
|
if (ok2) {
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b2 + st + t)) = pair16(q[12], q[13]);
|
|
*(uint32_t*)(dst + 4 * (uint32_t)(b2 + st + (t | 4))) = pair16(q[14], q[15]);
|
|
}
|
|
}
|
|
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Mirror cut. The six 16-row fragments of the staged block input arrive in this
|
|
// order (the first mma's operand order, .f16.e4m3 fragments of 4 b32 each):
|
|
// f0 %r120 %r136 %r128 %r144 f3 %r216 %r232 %r224 %r240
|
|
// f1 %r152 %r168 %r160 %r176 f4 %r246 %r257 %r253 %r1042
|
|
// f2 %r184 %r200 %r192 %r208 f5 %r1079 %r1080 %r1081 %r1082
|
|
// then phase 1's four bias C seeds (%r1086 %r1096 %r1106 %r1116, read from
|
|
// +5120/+5136/+5152/+5168), the weight image and the parameter block.
|
|
//
|
|
// The eighteen accumulators phase 22's epilogue reads cross the cut through shared memory, not
|
|
// `call.uni`'s return parameters: ZLUDA lowers those to the function's LLVM return type (emit.rs's
|
|
// emit_call makes one value, or an get_or_create_struct_type struct for several, and distributes
|
|
// them with LLVMBuildExtractValue), so there is no .param-space object and no pointer is formed --
|
|
// the callee reads 0x0 in all eighteen slots, and an 18 x i32 C++ struct return is an sret pointer,
|
|
// a different LLVM signature that will not link. `scratch` is the base of the entry's `input_tensor`
|
|
// (.shared 3200 bytes, dead after the cut), 72 bytes apart per lane.
|
|
extern "C" __attribute__((device)) void d4r_swin_enc0(
|
|
uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t a4, uint32_t a5, uint32_t a6,
|
|
uint32_t a7, uint32_t a8, uint32_t a9, uint32_t a10, uint32_t a11, uint32_t a12, uint32_t a13,
|
|
uint32_t a14, uint32_t a15, uint32_t a16, uint32_t a17, uint32_t a18, uint32_t a19, uint32_t a20,
|
|
uint32_t a21, uint32_t a22, uint32_t a23, uint32_t b0, uint32_t b1, uint32_t b2, uint32_t b3,
|
|
const uint8_t* w, const uint8_t* prm, uint32_t scratch)
|
|
{
|
|
const uint32_t as1[24] = {a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11,
|
|
a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23};
|
|
const uint32_t seeds[4] = {b0, b1, b2, b3};
|
|
// The epilogue's eighteen registers: %r17597, %r18203 (%tid.x) and phase 22's eight mma D pairs.
|
|
// Address space 3 (the shared scratch) explicitly: a generic store lowers to flat, which faults.
|
|
__attribute__((address_space(3))) uint32_t* out =
|
|
((__attribute__((address_space(3))) uint32_t*)(uintptr_t)scratch) + 18 * (uint32_t)lane_id();
|
|
const int t = (int)(lane_id() & 3); // the lane's two columns of each n-tile
|
|
const int g = (int)(lane_id() >> 2); // ... and of its sixteen rows
|
|
const uint32_t lane = lane_id();
|
|
const half_t eps = (half_t)0.0009765625f; // the tree's sum-of-squares epsilon
|
|
|
|
// ---- phase 1 -----------------------------------------------------------
|
|
uint32_t d1[6][2][4]; // [m-tile][row half][n-tile]: f16 pair
|
|
gemm(as1, w, 4096, nullptr, seeds, nullptr, 6, t, d1);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 1
|
|
{
|
|
uint32_t regs[48];
|
|
stage_regs(d1, regs);
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
|
|
// ---- phase 2, on phase 1's relu requantisation --------------------------
|
|
// cvt.rn.satfinite.relu.e4m3x2.f16x2, packed by the epilogue as {n-tile 0,
|
|
// n-tile 2} into the k-group 0 register and {n-tile 1, n-tile 3} into k-group 1.
|
|
uint32_t as2[24];
|
|
for (int mi = 0; mi < 6; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
const uint32_t* d = d1[mi][row];
|
|
const uint32_t c0 = e4m3x2(relu_pair(d[0])) | (e4m3x2(relu_pair(d[2])) << 16);
|
|
const uint32_t c1 = e4m3x2(relu_pair(d[1])) | (e4m3x2(relu_pair(d[3])) << 16);
|
|
as2[4 * mi + row] = c0;
|
|
as2[4 * mi + 2 + row] = c1;
|
|
}
|
|
}
|
|
uint32_t d2[6][2][4];
|
|
gemm(as2, w, 5184, w + 6208, nullptr, nullptr, 6, t, d2);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 2
|
|
{
|
|
uint32_t regs[48];
|
|
stage_regs(d2, regs);
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
|
|
// ---- the LayerNorm tree (PTX stmts 2399-3701) --------------------------
|
|
// Per row: the four n-tiles are squared and added as (n0+n2) then (n1+n3),
|
|
// both halves in f16; the butterfly finishes the sum over the row's 32
|
|
// columns and the two halves are added with add.f16; +2^-13, then
|
|
// rsqrt.approx.ftz.f32 and back to f16, times the activation gain of the
|
|
// column. The row's value times that factor is the A operand of phases 3 and
|
|
// 4 (cvt.rn.satfinite.e4m3x2.f16x2, no relu).
|
|
uint32_t as3[24];
|
|
norm_as(d2, w, 6272, eps, 6, as3);
|
|
|
|
// ---- phases 3 and 4, the fused Q/V projection, on that A operand -------
|
|
uint32_t d3[6][2][4];
|
|
uint32_t d4[6][2][4];
|
|
gemm(as3, w, 6336, nullptr, nullptr, nullptr, 6, t, d3);
|
|
gemm(as3, w, 7360, nullptr, nullptr, nullptr, 6, t, d4);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 3
|
|
{
|
|
uint32_t regs[48];
|
|
stage_regs(d3, regs);
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 4
|
|
{
|
|
uint32_t regs[48];
|
|
stage_regs(d4, regs);
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
|
|
// ---- phases 5 to 7: the first window --------------------------------------
|
|
// score_gemm() and pv_gemm() are the two GEMMs both windows run; see their comments
|
|
// for the fragment bookkeeping. The first window scores phase 3's accumulator against
|
|
// itself (A and B are both D3 -- the PTX's 24 B fragments are %r2932+10*tt, m-tile
|
|
// tt>>1, row half tt&1, VERIFIED against the oracle's own B registers) with the bias
|
|
// at weight byte 8384, and takes its V from phase 4's accumulator.
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 201 || D4R_RRSWIN_DEBUG_STAGE == 203 \
|
|
|| D4R_RRSWIN_DEBUG_STAGE == 204 || D4R_RRSWIN_DEBUG_STAGE == 211
|
|
uint32_t cap[136];
|
|
uint32_t* const cp = cap;
|
|
#else
|
|
uint32_t* const cp = nullptr;
|
|
#endif
|
|
lds_u32* const slot = (lds_u32*)(uintptr_t)scratch; // twelve n-tiles x 64 words
|
|
uint32_t d5[4][2][12]; // [m-tile][row half][n-tile]: f16 pair
|
|
score_gemm(d3, w, 8384, t, g, (int)lane, slot, cp, d5);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 5
|
|
{
|
|
uint32_t regs[96];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
regs[i++] = d5[mi][0][nb];
|
|
regs[i++] = d5[mi][1][nb];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 96);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 203
|
|
debug_store(prm, cp, 24);
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 204
|
|
debug_store(prm, cp, 136);
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 211
|
|
debug_store(prm, cp, 120);
|
|
#endif
|
|
|
|
uint32_t d6[4][2][4];
|
|
pv_gemm(d5, d4, t, g, (lds_u8*)(uintptr_t)scratch, cp, d6);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 6
|
|
{
|
|
uint32_t regs[32];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = d6[mi][0][nt];
|
|
regs[i++] = d6[mi][1][nt];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 201
|
|
debug_store(prm, cp, 48);
|
|
#endif
|
|
|
|
// ---- E6 and phase 7, the first window's output projection ----------------
|
|
// A is phase 6's D packed as before; C is phase 2's D plus the per-column table at
|
|
// weight bytes 37056..37120 (four words of eight f16, indexed by the 64-byte record
|
|
// each group of four lanes shares), taken from the d2 m-tile and n-tile the register
|
|
// order names: the mma's C pair for (m-tile mi, n-tile nt) is d2[mi][row][nt] -- the
|
|
// identity: phase 7's first four m-tiles are phase 2's rows 0..63 (VERIFIED, stage e6).
|
|
// Its output is also phase 12's C operand, so it has to exist before the second window.
|
|
uint32_t as7[16];
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
as7[4 * mi + row] = e4m3x2(d6[mi][row][0]) | (e4m3x2(d6[mi][row][2]) << 16);
|
|
as7[4 * mi + 2 + row] = e4m3x2(d6[mi][row][1]) | (e4m3x2(d6[mi][row][3]) << 16);
|
|
}
|
|
}
|
|
uint32_t c7[64]; // [(4*m-tile + row)][n-tile]
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
const uint32_t tab = *(const uint32_t*)(w + 37056 + 16 * nt + ((lane << 2) & 12));
|
|
for (int row = 0; row < 2; ++row) {
|
|
const uint32_t r = d2[mi][row][nt];
|
|
c7[(4 * mi + row) * 4 + nt] = pair_of(
|
|
(half_t)(half_lo(tab) + half_lo(r)), (half_t)(half_hi(tab) + half_hi(r)));
|
|
}
|
|
}
|
|
}
|
|
uint32_t d7[4][2][4];
|
|
gemm(as7, w, 20672, nullptr, nullptr, c7, 4, t, d7);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 202
|
|
{
|
|
uint32_t regs[48];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int o = 0; o < 4; ++o) {
|
|
regs[i++] = as7[4 * mi + o];
|
|
}
|
|
}
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = c7[(4 * mi + 0) * 4 + nt];
|
|
regs[i++] = c7[(4 * mi + 1) * 4 + nt];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 7
|
|
{
|
|
uint32_t regs[32];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = d7[mi][0][nt];
|
|
regs[i++] = d7[mi][1][nt];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
|
|
// ---- phases 8 and 9: the second window's two projections -----------------
|
|
// The same LayerNorm tile as the A operand of phases 3 and 4 (all four mma read the
|
|
// registers %r6784+40*mi -- VERIFIED against the oracle's own A fragment registers,
|
|
// debug stage p8o) and no bias: the second window scores D8 against itself and takes
|
|
// its V from D9. E7 and E8 between them only load the two weight tiles this reads.
|
|
uint32_t d8[6][2][4];
|
|
uint32_t d9[6][2][4];
|
|
gemm(as3, w, 21696, nullptr, nullptr, nullptr, 6, t, d8);
|
|
gemm(as3, w, 22720, nullptr, nullptr, nullptr, 6, t, d9);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 8
|
|
{
|
|
uint32_t regs[48];
|
|
stage_regs(d8, regs);
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 9
|
|
{
|
|
uint32_t regs[48];
|
|
stage_regs(d9, regs);
|
|
debug_store(prm, regs, 48);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 108
|
|
{
|
|
// phase 8's mma operands: the six A fragments (the LayerNorm tile) and the four
|
|
// B fragment pairs, which are the lane's own two words of tile nt>>1 at byte
|
|
// offset 8*(nt&1) -- its column 8*nt + g
|
|
uint32_t regs[32];
|
|
for (int o = 0; o < 24; ++o) {
|
|
regs[o] = as3[o];
|
|
}
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
const uint8_t* const p = w + 21696 + 512 * (nt >> 1) + 16 * lane + 8 * (nt & 1);
|
|
regs[24 + 2 * nt] = *(const uint32_t*)p;
|
|
regs[24 + 2 * nt + 1] = *(const uint32_t*)(p + 4);
|
|
}
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
|
|
// ---- phases 10 and 11, and E10: the second window's score and P·V --------
|
|
// Identical to phases 5 and 6 with phase 8's accumulator in place of phase 3's and the
|
|
// bias at weight bytes 23744..35532 (the second 12288-byte table, same layout).
|
|
uint32_t d10[4][2][12];
|
|
score_gemm(d8, w, 23744, t, g, (int)lane, slot, cp, d10);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 10
|
|
{
|
|
uint32_t regs[96];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nb = 0; nb < 12; ++nb) {
|
|
regs[i++] = d10[mi][0][nb];
|
|
regs[i++] = d10[mi][1][nb];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 96);
|
|
}
|
|
#endif
|
|
uint32_t d11[4][2][4];
|
|
pv_gemm(d10, d9, t, g, (lds_u8*)(uintptr_t)scratch, cp, d11);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 11
|
|
{
|
|
uint32_t regs[32];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = d11[mi][0][nt];
|
|
regs[i++] = d11[mi][1][nt];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
|
|
// ---- E11 and phase 12: the second window's output projection --------------
|
|
// Unlike E6 there is no table: phase 11's output is requantised only, because the
|
|
// residual of the first window is phase 12's C operand (the mma's C pair for m-tile
|
|
// mi and n-tile nt is phase 7's D register pair %r6334+10*(4*mi+nt) -- VERIFIED
|
|
// against the oracle's own C registers, debug stage p12o).
|
|
uint32_t as12[16];
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
as12[4 * mi + row] = e4m3x2(d11[mi][row][0]) | (e4m3x2(d11[mi][row][2]) << 16);
|
|
as12[4 * mi + 2 + row] = e4m3x2(d11[mi][row][1]) | (e4m3x2(d11[mi][row][3]) << 16);
|
|
}
|
|
}
|
|
uint32_t c12[64]; // [(4*m-tile + row)][n-tile]: phase 7's accumulators
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
c12[(4 * mi + row) * 4 + nt] = d7[mi][row][nt];
|
|
}
|
|
}
|
|
}
|
|
uint32_t d12[4][2][4];
|
|
gemm(as12, w, 36032, nullptr, nullptr, c12, 4, t, d12);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 12
|
|
{
|
|
uint32_t regs[32];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 4; ++nt) {
|
|
regs[i++] = d12[mi][0][nt];
|
|
regs[i++] = d12[mi][1][nt];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 112
|
|
{
|
|
// phase 12's mma operands: the four A fragments, the four B fragment pairs and
|
|
// the 32 C seeds in the oracle's (m-tile, n-tile) pair order
|
|
uint32_t regs[56];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = as12[o];
|
|
}
|
|
bpair_regs(w, 36032, (int)lane, regs + 16);
|
|
cseed_regs(c12, regs + 24);
|
|
debug_store(prm, regs, 56);
|
|
}
|
|
#endif
|
|
|
|
// ---- E12 and phases 13 to 20: the MLP ------------------------------------
|
|
// E12's RMS norm is the same sum-of-squares tree as the LayerNorm above, on phase 12's
|
|
// accumulators and with the gain at weight bytes 37120..37184; its epsilon is E12's own
|
|
// add (%fd770 = 2^-13) and there is no mean subtraction and no 1/N. Its output is the
|
|
// shared A operand of the four branches: phases 13, 15, 17 and 19 all read the same
|
|
// sixteen registers (%r15389..%r15420, re-packed by E14, E16 and E18).
|
|
uint32_t asN[16];
|
|
norm_as(d12, w, 37120, (half_t)0.0001220703125f, 4, asN);
|
|
|
|
// ---- phase 13 and E13, the first branch -----------------------------------
|
|
// A = the normed tile, B = the weights at 37184, C = the bias at 38208 (one word per
|
|
// n-tile, replicated over the 16 rows of its tile). E13 activates the result and adds
|
|
// the per-column table at 39296 to the activation's *input*, which is phase 14's C.
|
|
uint32_t d13[4][2][4];
|
|
gemm(asN, w, 37184, w + 38208, nullptr, nullptr, 4, t, d13);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 13
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d13, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 113
|
|
{
|
|
uint32_t regs[28];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = asN[o];
|
|
}
|
|
bpair_regs(w, 37184, (int)lane, regs + 16);
|
|
wbias_regs(w, 38208, (int)lane, regs + 24);
|
|
debug_store(prm, regs, 28);
|
|
}
|
|
#endif
|
|
uint32_t as14[16];
|
|
act_pack(d13, as14);
|
|
uint32_t c14[64];
|
|
tab_seed(d12, w + 39296, (int)lane, c14);
|
|
uint32_t d14[4][2][4];
|
|
gemm(as14, w, 38272, nullptr, nullptr, c14, 4, t, d14);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 14
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d14, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 114
|
|
{
|
|
uint32_t regs[56];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = as14[o];
|
|
}
|
|
bpair_regs(w, 38272, (int)lane, regs + 16);
|
|
cseed_regs(c14, regs + 24);
|
|
debug_store(prm, regs, 56);
|
|
}
|
|
#endif
|
|
|
|
// ---- phases 15 and 16, the second branch ---------------------------------
|
|
// The same pair with the weights at 39360 and 40448 and the bias at 40384: no table
|
|
// add, and phase 16's C is branch one's accumulators.
|
|
uint32_t d15[4][2][4];
|
|
gemm(asN, w, 39360, w + 40384, nullptr, nullptr, 4, t, d15);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 15
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d15, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 115
|
|
{
|
|
uint32_t regs[28];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = asN[o];
|
|
}
|
|
bpair_regs(w, 39360, (int)lane, regs + 16);
|
|
wbias_regs(w, 40384, (int)lane, regs + 24);
|
|
debug_store(prm, regs, 28);
|
|
}
|
|
#endif
|
|
uint32_t as16[16];
|
|
act_pack(d15, as16);
|
|
uint32_t c16[64];
|
|
cres_regs(d14, c16);
|
|
uint32_t d16[4][2][4];
|
|
gemm(as16, w, 40448, nullptr, nullptr, c16, 4, t, d16);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 16
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d16, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 116
|
|
{
|
|
uint32_t regs[56];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = as16[o];
|
|
}
|
|
bpair_regs(w, 40448, (int)lane, regs + 16);
|
|
cseed_regs(c16, regs + 24);
|
|
debug_store(prm, regs, 56);
|
|
}
|
|
#endif
|
|
|
|
// ---- phases 17 and 18, the third branch ----------------------------------
|
|
uint32_t d17[4][2][4];
|
|
gemm(asN, w, 41472, w + 42496, nullptr, nullptr, 4, t, d17);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 17
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d17, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 117
|
|
{
|
|
uint32_t regs[28];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = asN[o];
|
|
}
|
|
bpair_regs(w, 41472, (int)lane, regs + 16);
|
|
wbias_regs(w, 42496, (int)lane, regs + 24);
|
|
debug_store(prm, regs, 28);
|
|
}
|
|
#endif
|
|
uint32_t as18[16];
|
|
act_pack(d17, as18);
|
|
uint32_t c18[64];
|
|
cres_regs(d16, c18);
|
|
uint32_t d18[4][2][4];
|
|
gemm(as18, w, 42560, nullptr, nullptr, c18, 4, t, d18);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 18
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d18, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 118
|
|
{
|
|
uint32_t regs[56];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = as18[o];
|
|
}
|
|
bpair_regs(w, 42560, (int)lane, regs + 16);
|
|
cseed_regs(c18, regs + 24);
|
|
debug_store(prm, regs, 56);
|
|
}
|
|
#endif
|
|
|
|
// ---- phases 19 and 20, the fourth branch ---------------------------------
|
|
uint32_t d19[4][2][4];
|
|
gemm(asN, w, 43584, w + 44608, nullptr, nullptr, 4, t, d19);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 19
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d19, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 119
|
|
{
|
|
uint32_t regs[28];
|
|
for (int o = 0; o < 16; ++o) {
|
|
regs[o] = asN[o];
|
|
}
|
|
bpair_regs(w, 43584, (int)lane, regs + 16);
|
|
wbias_regs(w, 44608, (int)lane, regs + 24);
|
|
debug_store(prm, regs, 28);
|
|
}
|
|
#endif
|
|
uint32_t as20[16];
|
|
act_pack(d19, as20);
|
|
uint32_t c20[64];
|
|
cres_regs(d18, c20);
|
|
uint32_t d20[4][2][4];
|
|
gemm(as20, w, 44672, nullptr, nullptr, c20, 4, t, d20);
|
|
|
|
// ---- E20 and phases 21 and 22: the merge ---------------------------------
|
|
// E20 requantises phase 20 into sixteen packed words (r20, the four A fragments of
|
|
// phase 22), scatters them into the plane arena, and merges the lanes of each group of
|
|
// eight back into phase 21's four A fragments.
|
|
uint32_t r20[16];
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int row = 0; row < 2; ++row) {
|
|
const uint32_t a0 = e4m3x2(d20[mi][row][0]) | (e4m3x2(d20[mi][row][2]) << 16);
|
|
const uint32_t a1 = e4m3x2(d20[mi][row][1]) | (e4m3x2(d20[mi][row][3]) << 16);
|
|
r20[4 * mi + row] = a0;
|
|
r20[4 * mi + 2 + row] = a1;
|
|
}
|
|
}
|
|
arena_scatter(prm, r20, (int)lane);
|
|
uint32_t m21[16];
|
|
merge_frag(r20, scratch, (int)lane, m21);
|
|
uint32_t d21[8][2];
|
|
band21(m21, w, w + 53888, (int)lane, d21);
|
|
|
|
// ---- E21 and phase 22 ---------------------------------------------------
|
|
// E21 requantises phase 21 and scatters it into the buffer at parameter block +56;
|
|
// phase 22 takes the sixteen words E20 packed as its A operand (four m-tiles, K=32)
|
|
// with the weights at 54016 and the two bias words at 54528.
|
|
surface_scatter21(prm, d21, (int)lane);
|
|
uint32_t d22[4][2][4];
|
|
gemm(r20, w, 54016, w + 54528, nullptr, nullptr, 4, t, d22, 2);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 20
|
|
{
|
|
uint32_t regs[32];
|
|
stage_regs4(d20, regs);
|
|
debug_store(prm, regs, 32);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 21
|
|
{
|
|
uint32_t regs[16];
|
|
for (int nt = 0; nt < 8; ++nt) {
|
|
regs[2 * nt] = d21[nt][0];
|
|
regs[2 * nt + 1] = d21[nt][1];
|
|
}
|
|
debug_store(prm, regs, 16);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 22
|
|
{
|
|
uint32_t regs[16];
|
|
int i = 0;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 2; ++nt) {
|
|
regs[i++] = d22[mi][0][nt];
|
|
regs[i++] = d22[mi][1][nt];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 16);
|
|
}
|
|
#endif
|
|
|
|
// ---- handing the accumulators back ------------------------------------
|
|
// %r17597 is param_0[0] (ld.param.v2.u32 at line 44316) and %r18203 is %tid.x
|
|
// (line 44456); both sit inside the cut, so the epilogue cannot recompute them.
|
|
// The other sixteen are phase 22's eight mma D pairs in mma order.
|
|
out[0] = *(const uint32_t*)prm;
|
|
out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x();
|
|
int i = 2;
|
|
for (int mi = 0; mi < 4; ++mi) {
|
|
for (int nt = 0; nt < 2; ++nt) {
|
|
out[i++] = d22[mi][0][nt];
|
|
out[i++] = d22[mi][1][nt];
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Weight image layout of the 22 phases (byte ranges of arg040, from the PTX's own
|
|
// ld.weak.global loads; 512 bytes per B-fragment pair, K=32, N=32):
|
|
// P1 [4096,5120) bias 5120..5184 P12 [36032,37056)
|
|
// P2 [5184,6208) bias 6208..6272 P13 [37184,38208) bias 38208..38272
|
|
// P3 [6336,7360) P14 [38272,39296)
|
|
// P4 [7360,8384) P15 [39360,40384) bias 40384..40448
|
|
// P5 bias [8384,20672) (64x96 f16) P16 [40448,41472)
|
|
// P6 none (P V, B from mma output) P17 [41472,42496) bias 42496..42560
|
|
// P7 [20672,21696) P18 [42560,43584)
|
|
// P8 [21696,22720) P19 [43584,44608) bias 44608..44672
|
|
// P9 [22720,23744) P20 [44672,45696)
|
|
// P10 bias [23744,36032) (64x96 f16) P21 [45696,53888) bias 53888..54016
|
|
// P11 none (P V) P22 [54016,54544) bias 54528..54560
|
|
//
|
|
// Fragment maps (g = lane>>2, t = lane&3). The corrected full-lane phase-8
|
|
// model matches all 3072 f16 outputs of captured block 0. The checked native
|
|
// pilot matches 196608 f16 outputs over 64 captured blocks. Its former
|
|
// mismatch came from invalid accumulator bounds and device/host block strides,
|
|
// not from the fragment map. The software fallback's different fragment
|
|
// convention is not the active gfx12 FP8 WMMA reference.
|
|
// The hardware order of an m16n8k32 .e4m3 fragment pair is *not* the order the
|
|
// PTX ISA figures suggest. NVIDIA's own epilogue (stmts 2256-2327) settles it: it
|
|
// builds phase 2's A operand as
|
|
// %r1207 = {cvt(%r877), cvt(%r897)}
|
|
// where %r877/%r897 are the D halves of mma rows 2231/2233, i.e. (row g, n-tile 0,
|
|
// cols 2t, 2t+1) and (row g, n-tile 2, cols 2t, 2t+1). So with the four bytes of
|
|
// an A register in memory order:
|
|
// A reg 0 -> row g, k = 2t + (j&1) + 16*((j>>1)&1)
|
|
// A reg 1 -> row g+8, same k
|
|
// A reg 2 -> row g, k = 8 + 2t + (j&1) + 16*((j>>1)&1)
|
|
// A reg 3 -> row g+8, same k
|
|
// and the B half that must pair with it (a0/a1 with b0, a2/a3 with b1) is
|
|
// B byte j (j<4) -> n = 16*tt + 8*(b>>3) + g, k = 2t + (j&1) + 16*((j>>1)&1)
|
|
// B byte j (j>=4) -> n = 16*tt + 8*(b>>3) + g, k = 8 + 2t + (j&1) + 16*((j>>1)&1)
|
|
// with b the byte index within the lane's 16 and tt the 512-byte tile. Inverted,
|
|
// that is waddr() above. C/D:
|
|
// C/D reg 0 -> row g, cols t*2, t*2+1
|
|
// C/D reg 1 -> row g+8, cols t*2, t*2+1
|
|
//
|
|
// A plain f32 dot product over K plus C, rounded to f16 once, reproduces the
|
|
// translated kernel's phase-1 output exactly (96x32 elements, 4 blocks x 32 lanes,
|
|
// 0 mismatches out of 12288), so the phases need no operand-layout gymnastics and
|
|
// no WMMA intrinsics: every GEMM here is exact in f32 before the f16 rounding.
|
|
// The textbook order (reg 0 = row g, k = 4t..4t+3) also reproduces phase 1, but
|
|
// only because it is a self-consistent relabeling of k on *both* sides; it cannot
|
|
// be used for A operands that an epilogue produced, which arrive in the hardware
|
|
// order above.
|