kernels/rr: staged-image Swin layer engine prototypes
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).
This commit is contained in:
woutermans committed
1 parent
798e25611c
commit
0fe3a28cb7
58 files changed
+28330
-5
No files matched your search
@@ -0,0 +1,737 @@
|
||||
// Native replacement for the Ray Reconstruction denoiser's second Swin encoder,
|
||||
// cuda_dldn_engine_swin_enc1_kernel (nvngx_dlssd.dll 310.7).
|
||||
//
|
||||
// ---------------------------------------------------------------------------
|
||||
// STATE: the recipe (kernels/rr/swin_enc1_recipe.py) cuts statements s8270-s26139
|
||||
// out of the entry and calls this body instead, handing over the prologue's two
|
||||
// staged A operand sets, the zero register phase 1 uses as its C seed and the
|
||||
// dequantised input tile E6 needs for phase 7's residual. The window implemented
|
||||
// here is phases 1-7: the four 96x32 projections, the 64x96 score, P*V and the
|
||||
// 64x64 output projection. Phase 32's accumulators are still placeholders, so the
|
||||
// epilogue's surface output differs from the translated kernel; every implemented
|
||||
// stage is bit-exact against the debug oracle (see "Ground truth" below).
|
||||
//
|
||||
// The kernel is one warp per 8x8 token patch: grid (81, 47, 1), block (32, 1, 1)
|
||||
// with one wave per block (the prologue also handles tid.y > 0, for waves > 1, but
|
||||
// the shipped manifest launches one), parameter block 144 bytes. The PTX this
|
||||
// file replaces is kernels/extracted/rr/0014-PREPASS_ENTRYPOINT_NAME.ptx.
|
||||
//
|
||||
// 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"
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 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 DldnSwinEnc1Params
|
||||
{
|
||||
uint8_t raw[144];
|
||||
};
|
||||
static_assert(sizeof(DldnSwinEnc1Params) == 144, "parameter block");
|
||||
|
||||
enum : int
|
||||
{
|
||||
kWeights = 40, // NVIDIA's prepared weight image (arg040), read-only
|
||||
kPlaneArg = 48, // plane arena pointer, as make_ptx.py's debug splice reads it
|
||||
kExtents = 0, // ex, ey (v2.u32)
|
||||
};
|
||||
|
||||
#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)));
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// The implemented phases, as byte offsets into arg040 (from the PTX's own loads;
|
||||
// K=32, 512 bytes per group of 16 columns):
|
||||
// P1 weights 128 P2 weights 2176 P3 weights 1152 P4 weights 3200
|
||||
// P5 score bias 4224..16512 (24 tiles of 512)
|
||||
// P6 none (P*V; B is the movmatrix transpose of phase 4's D)
|
||||
// P7 weights 16512..18560, per-column table 36992..37120 (8 words of 16 bytes)
|
||||
|
||||
// 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 for enc0; the B operand layout is
|
||||
// the hardware's, so the same inversion applies here.)
|
||||
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/high f16 of a packed pair (columns 2t / 2t+1)
|
||||
DEVF half_t half_lo(uint32_t pair)
|
||||
{
|
||||
return __builtin_bit_cast(half_t, (uint16_t)pair);
|
||||
}
|
||||
|
||||
DEVF half_t half_hi(uint32_t pair)
|
||||
{
|
||||
return __builtin_bit_cast(half_t, (uint16_t)(pair >> 16));
|
||||
}
|
||||
|
||||
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: `nmi` m-tiles of 16 rows x `nti` n-tiles of 8 columns, K=32, f16 out.
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// The C operand is the 32-bit seed the PTX's mma adds (one packed f16 pair per D
|
||||
// register, indexed [m-tile][row half][n-tile] as `cmat`) or a constant (`fixed`).
|
||||
// d is [m-tile][row half][n-tile], each entry the packed f16 pair 2t, 2t+1.
|
||||
DEVF void gemm(const uint32_t* as, const uint8_t* w, int base, uint32_t fixed,
|
||||
const uint32_t* cmat, int nmi, int nti, int t, uint32_t d[6][2][8])
|
||||
{
|
||||
for (int nt = 0; nt < nti; ++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) {
|
||||
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 seed = (cmat != nullptr) ? cmat[(4 * mi + row) * nti + nt] : fixed;
|
||||
d[mi][row][nt] = pair_of((half_t)(acc[row][0] + h2f(seed)),
|
||||
(half_t)(acc[row][1] + h2f(seed >> 16)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// the f16 the PTX's literals carry: 0x2466, 0xb873, 0x3873, 0x3b6b, 0x3d80
|
||||
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 (or 64x32) block of packed pairs,
|
||||
// col = 8*n-tile + 2t + c
|
||||
DEVF half_t dcell(const uint32_t d[6][2][8], 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
|
||||
// packs tiles 1 and 3 (the mma's own k order)
|
||||
DEVF void row_words(const uint32_t d[6][2][8], 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).
|
||||
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, 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));
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Ground truth for a stage comes from kernels/tex/make_ptx.py's debug recipe:
|
||||
// with D4R_RRSWIN_DEBUG=<stage> the NVIDIA body keeps every mma and stores that
|
||||
// stage's registers into the plane arena (per block 0..63, slot by register count,
|
||||
// 4 bytes per lane), and the launch dump then holds the translated values. 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.
|
||||
//
|
||||
// The offset into the arena is D4R_RRSWIN_DEBUG_BASE on the oracle side (read by
|
||||
// make_ptx.py's swin_debug_block) and D4R_RRSWIN_DEBUG_BASE_NUM here. Two measured
|
||||
// facts fix its value for enc1:
|
||||
//
|
||||
// * enc1's param_0+48 is not dumped as arg048. The launch manifest puts it at byte
|
||||
// 45,219,840 of the 65,945,600-byte block the bridge dumps as arg008
|
||||
// (`pointer 8 0 30146560`, `pointer 48 0 45219840`, `pointer 56 0 60293120`), so
|
||||
// the bridge skips arg048 as a duplicate allocation base. The payload therefore
|
||||
// lives in the arg008 dump at 45,219,840 + base, and enc0's default of 46,000,000
|
||||
// would fall 25 MB past the end of the block.
|
||||
// * enc1's own 44 `st.global` fill that block's whole 20,725,760-byte tail (measured
|
||||
// on an isolated launch: the translated kernel changes every byte of it), so a base
|
||||
// inside the tail is overwritten by the kernel's own later stores and the comparison
|
||||
// silently compares ambient memory. The last byte E31 reaches is arena + 16,957,436,
|
||||
// and the block ends at arena + 20,725,760.
|
||||
//
|
||||
// 17,006,592 is a multiple of the largest slot (16,384), clears E31, and leaves room
|
||||
// for all 64 slots (1,048,576 bytes) inside the block.
|
||||
#ifndef D4R_RRSWIN_DEBUG_STAGE
|
||||
#define D4R_RRSWIN_DEBUG_STAGE 0
|
||||
#endif
|
||||
#ifndef D4R_RRSWIN_DEBUG_BASE_NUM
|
||||
#define D4R_RRSWIN_DEBUG_BASE_NUM 17006592
|
||||
#endif
|
||||
|
||||
DEVF void debug_store(const uint8_t* prm, const uint32_t* regs, int n)
|
||||
{
|
||||
// make_ptx.py's swin_debug_block forms the block index as ctaid.x + 161*ctaid.y
|
||||
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: 32 lanes x n x 4 bytes, at least 8192, rounded up to a power of two
|
||||
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 + (uint64_t)D4R_RRSWIN_DEBUG_BASE_NUM
|
||||
+ (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: 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][8], int nmi, int nti, uint32_t regs[96])
|
||||
{
|
||||
int i = 0;
|
||||
for (int mi = 0; mi < nmi; ++mi) {
|
||||
for (int nt = 0; nt < nti; ++nt) {
|
||||
regs[i++] = d[mi][0][nt];
|
||||
regs[i++] = d[mi][1][nt];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// The native body.
|
||||
//
|
||||
// Live-ins (recipes's order): AS1[24] (phases 1/2/8/9), AS2[24] (phases 3/4/10/11),
|
||||
// the phase-1 C seed (the constant zero register), RES[64] (the dequantised input
|
||||
// tile in (m-tile, n-tile, row half) order), the weight image, the parameter block
|
||||
// and the address of the LDS scratch the accumulators go back through.
|
||||
#define AS(i) a##i
|
||||
extern "C" __attribute__((device)) void d4r_swin_enc1(
|
||||
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 a24, uint32_t a25, uint32_t a26, uint32_t a27,
|
||||
uint32_t a28, uint32_t a29, uint32_t a30, uint32_t a31, uint32_t a32, uint32_t a33, uint32_t a34,
|
||||
uint32_t a35, uint32_t a36, uint32_t a37, uint32_t a38, uint32_t a39, uint32_t a40, uint32_t a41,
|
||||
uint32_t a42, uint32_t a43, uint32_t a44, uint32_t a45, uint32_t a46, uint32_t a47,
|
||||
uint32_t cseed,
|
||||
uint32_t r0, uint32_t r1, uint32_t r2, uint32_t r3, uint32_t r4, uint32_t r5, uint32_t r6,
|
||||
uint32_t r7, uint32_t r8, uint32_t r9, uint32_t r10, uint32_t r11, uint32_t r12, uint32_t r13,
|
||||
uint32_t r14, uint32_t r15, uint32_t r16, uint32_t r17, uint32_t r18, uint32_t r19, uint32_t r20,
|
||||
uint32_t r21, uint32_t r22, uint32_t r23, uint32_t r24, uint32_t r25, uint32_t r26, uint32_t r27,
|
||||
uint32_t r28, uint32_t r29, uint32_t r30, uint32_t r31, uint32_t r32, uint32_t r33, uint32_t r34,
|
||||
uint32_t r35, uint32_t r36, uint32_t r37, uint32_t r38, uint32_t r39, uint32_t r40, uint32_t r41,
|
||||
uint32_t r42, uint32_t r43, uint32_t r44, uint32_t r45, uint32_t r46, uint32_t r47, uint32_t r48,
|
||||
uint32_t r49, uint32_t r50, uint32_t r51, uint32_t r52, uint32_t r53, uint32_t r54, uint32_t r55,
|
||||
uint32_t r56, uint32_t r57, uint32_t r58, uint32_t r59, uint32_t r60, uint32_t r61, uint32_t r62,
|
||||
uint32_t r63,
|
||||
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 as2[24] = {a24, a25, a26, a27, a28, a29, a30, a31, a32, a33, a34, a35,
|
||||
a36, a37, a38, a39, a40, a41, a42, a43, a44, a45, a46, a47};
|
||||
const uint32_t res[64] = {r0, r1, r2, r3, r4, r5, r6, r7, r8, r9, r10, r11, r12, r13, r14, r15,
|
||||
r16, r17, r18, r19, r20, r21, r22, r23, r24, r25, r26, r27, r28, r29,
|
||||
r30, r31, r32, r33, r34, r35, r36, r37, r38, r39, r40, r41, r42, r43,
|
||||
r44, r45, r46, r47, r48, r49, r50, r51, r52, r53, r54, r55, r56, r57,
|
||||
r58, r59, r60, r61, r62, r63};
|
||||
#if D4R_RRSWIN_DEBUG_STAGE
|
||||
// The oracle's payload store is spliced in after E31's last arena write (see
|
||||
// swin_enc1_debug_stages.LINE), because enc1 fills its whole plane-arena tail, so a
|
||||
// store placed beside its own phase is overwritten before the dump reads it back.
|
||||
// This side therefore fills `cap` where each stage's registers are live and hands
|
||||
// the whole thing to debug_store() once, at the end of the body -- after this
|
||||
// body's own E30/E31 stores, which is the same program point.
|
||||
uint32_t cap[128];
|
||||
int capn = 0;
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 10
|
||||
// the oracle's stage a: the two staged A operand sets, the phase-1 C seed and
|
||||
// the three shfl constants the PTX carries (1, 2, -1)
|
||||
for (int i = 0; i < 24; ++i) {
|
||||
cap[i] = as1[i];
|
||||
cap[24 + i] = as2[i];
|
||||
}
|
||||
cap[48] = cseed;
|
||||
cap[49] = 1u;
|
||||
cap[50] = 2u;
|
||||
cap[51] = 0xffffffffu;
|
||||
capn = 52;
|
||||
#endif
|
||||
#endif
|
||||
(void)cseed; // phase 1's C operand is this register's value, the constant zero
|
||||
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
|
||||
|
||||
lds_u32* const slot = (lds_u32*)(uintptr_t)scratch; // the phase 5/6 exchange
|
||||
|
||||
// ---- the four 96x32 projections ----------------------------------------
|
||||
// Phases 1 and 2 have C = 0 (the PTX hands in the zero register); phases 3 and
|
||||
// 4 add phase 1's / phase 2's D back in as their C seed.
|
||||
uint32_t d1[6][2][8];
|
||||
uint32_t d2[6][2][8];
|
||||
gemm(as1, w, 128, 0u, nullptr, 6, 4, t, d1);
|
||||
gemm(as1, w, 2176, 0u, nullptr, 6, 4, t, d2);
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 1
|
||||
stage_regs(d1, 6, 4, cap);
|
||||
capn = 48;
|
||||
#endif
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 2
|
||||
stage_regs(d2, 6, 4, cap);
|
||||
capn = 48;
|
||||
#endif
|
||||
uint32_t d3[6][2][8];
|
||||
uint32_t d4[6][2][8];
|
||||
{
|
||||
uint32_t c1[96], c2[96];
|
||||
for (int mi = 0; mi < 6; ++mi) {
|
||||
for (int nt = 0; nt < 4; ++nt) {
|
||||
for (int rh = 0; rh < 2; ++rh) {
|
||||
c1[(4 * mi + rh) * 4 + nt] = d1[mi][rh][nt];
|
||||
c2[(4 * mi + rh) * 4 + nt] = d2[mi][rh][nt];
|
||||
}
|
||||
}
|
||||
}
|
||||
gemm(as2, w, 1152, 0u, c1, 6, 4, t, d3);
|
||||
gemm(as2, w, 3200, 0u, c2, 6, 4, t, d4);
|
||||
}
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 3
|
||||
stage_regs(d3, 6, 4, cap);
|
||||
capn = 48;
|
||||
#endif
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 4
|
||||
stage_regs(d4, 6, 4, cap);
|
||||
capn = 48;
|
||||
#endif
|
||||
|
||||
// ---- phase 5, the score: Q.K^T over the patch and its halo ---------------
|
||||
// A and B are both phase 3's accumulators. The lane's rows 16*mi + g + 8*rh are the
|
||||
// queries (the 8x8 token patch) and its rows 16*(nb>>1) + 8*(nb&1) + g the keys (the
|
||||
// patch plus a 32-row halo): the PTX builds n-tile nb's B fragment from phase 3's
|
||||
// m-tile nb>>1 and row half nb&1, whose two words are the k-groups {0, 2} and {1, 3}
|
||||
// of the 32 columns. row_words() packs a row that way.
|
||||
//
|
||||
// A lane holds one row group of the queries and one of the keys, so it can only dot
|
||||
// them with each other; every lane therefore publishes its two key words per n-tile
|
||||
// (12 x 2 x 32 lanes = 768 words = 3072 bytes, the whole scratch) and then forms the
|
||||
// dots its own D registers hold.
|
||||
//
|
||||
// The C operand is the 64x96 f16 score bias at weight bytes 4224..16512: 24 tiles of
|
||||
// 512 bytes, the seed of (m-tile, n-tile, row half) at byte
|
||||
// 4224 + 512*(mi + 4*(nb>>1)) + 16*lane + 4*(2*(nb&1) + row).
|
||||
uint32_t d5[4][2][12];
|
||||
{
|
||||
#pragma unroll
|
||||
for (int nb = 0; nb < 12; ++nb) {
|
||||
uint32_t wd[2];
|
||||
row_words(d3, 16 * (nb >> 1) + 8 * (nb & 1) + g, wd);
|
||||
slot[64 * nb + lane_id()] = wd[0];
|
||||
slot[64 * nb + 32 + lane_id()] = wd[1];
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < 4; ++mi) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
uint32_t aw[2], a[2][4];
|
||||
row_words(d3, 16 * mi + g + 8 * row, aw);
|
||||
bfly2x4(aw, a);
|
||||
#pragma unroll
|
||||
for (int nb = 0; nb < 12; ++nb) {
|
||||
float s[2];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c) {
|
||||
uint32_t kb[2][4];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
#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 bias = *(const uint32_t*)(w + 4224 + 512 * (mi + 4 * (nb >> 1))
|
||||
+ 16 * lane_id() + 4 * (2 * (nb & 1) + row));
|
||||
d5[mi][row][nb] = pair_of((half_t)(s[0] + h2f(bias)),
|
||||
(half_t)(s[1] + h2f(bias >> 16)));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 8
|
||||
// the oracle's stage p5o: the 96 C words (the score bias, tile-major: 24 tiles
|
||||
// of four words per lane) and the 24 B words (the two e4m3 key words of each of
|
||||
// the twelve n-tiles) exactly as the mma reads them.
|
||||
#pragma unroll
|
||||
for (int ti = 0; ti < 24; ++ti) {
|
||||
for (int wd = 0; wd < 4; ++wd) {
|
||||
cap[4 * ti + wd] = *(const uint32_t*)(w + 4224 + 512 * ti + 16 * lane_id() + 4 * wd);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nb = 0; nb < 12; ++nb) {
|
||||
cap[96 + 2 * nb] = slot[64 * nb + lane_id()];
|
||||
cap[96 + 2 * nb + 1] = slot[64 * nb + 32 + lane_id()];
|
||||
}
|
||||
capn = 120;
|
||||
#endif
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 5
|
||||
{
|
||||
int i = 0;
|
||||
for (int mi = 0; mi < 4; ++mi) {
|
||||
for (int nb = 0; nb < 12; ++nb) {
|
||||
cap[i++] = d5[mi][0][nb];
|
||||
cap[i++] = d5[mi][1][nb];
|
||||
}
|
||||
}
|
||||
capn = 96;
|
||||
}
|
||||
#endif
|
||||
|
||||
// ---- E5 and phase 6: the softmax and P.V --------------------------------
|
||||
// row_probs() gives the twelve probability pairs of one row; k-step ks packs key
|
||||
// columns {4ks, 4ks+2} 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 scratch as one e4m3 byte per element (96 x 32 = 3072
|
||||
// bytes) and each lane then walks the full k range. Each k-step rounds to f16
|
||||
// before the next one adds to it, which is why the chain has three separate mma.
|
||||
uint32_t d6[4][2][8];
|
||||
{
|
||||
#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;
|
||||
((lds_u8*)(uintptr_t)scratch)[r * 32 + 8 * nt + 2 * t + c] =
|
||||
(uint8_t)q8(dcell(d4, r, 8 * nt + 2 * t + c));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
{
|
||||
const lds_u8* const vs = (const lds_u8*)(uintptr_t)scratch;
|
||||
#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(d5, 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);
|
||||
}
|
||||
#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(vs[r]);
|
||||
acc1 += pv * e4m3_to_f32(vs[r + 1]);
|
||||
}
|
||||
prev0 = (half_t)(acc0 + (float)prev0);
|
||||
prev1 = (half_t)(acc1 + (float)prev1);
|
||||
}
|
||||
d6[mi][rh][nt] = pair_of(prev0, prev1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 6
|
||||
// the oracle's stage p6: the third (last) k-step of each chain, m-tile major,
|
||||
// n-tile minor, the (row g, row g+8) pair
|
||||
{
|
||||
int i = 0;
|
||||
for (int mi = 0; mi < 4; ++mi) {
|
||||
for (int nt = 0; nt < 4; ++nt) {
|
||||
cap[i++] = d6[mi][0][nt];
|
||||
cap[i++] = d6[mi][1][nt];
|
||||
}
|
||||
}
|
||||
capn = 32;
|
||||
}
|
||||
#endif
|
||||
|
||||
// ---- E6 and phase 7, the output projection ------------------------------
|
||||
// A is phase 6's D packed as before; C is the pipeline's residual + the
|
||||
// per-column table at weight bytes 36992..37120 (8 words of two f16 per n-tile,
|
||||
// the lane's own column pair 2t, 2t+1).
|
||||
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 half)][n-tile]
|
||||
for (int mi = 0; mi < 4; ++mi) {
|
||||
for (int nt = 0; nt < 8; ++nt) {
|
||||
const uint32_t tab = *(const uint32_t*)(w + 36992 + 16 * nt + ((lane_id() << 2) & 12));
|
||||
for (int row = 0; row < 2; ++row) {
|
||||
const uint32_t rr = res[(mi * 8 + nt) * 2 + row];
|
||||
c7[(4 * mi + row) * 8 + nt] = pair_of(
|
||||
(half_t)(half_lo(tab) + half_lo(rr)), (half_t)(half_hi(tab) + half_hi(rr)));
|
||||
}
|
||||
}
|
||||
}
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 9
|
||||
// the oracle's stage e6: phase 7's 16 A registers and its 64 C seeds
|
||||
{
|
||||
int i = 0;
|
||||
for (int mi = 0; mi < 4; ++mi) {
|
||||
for (int o = 0; o < 4; ++o) {
|
||||
cap[i++] = as7[4 * mi + o];
|
||||
}
|
||||
}
|
||||
for (int mi = 0; mi < 4; ++mi) {
|
||||
for (int nt = 0; nt < 8; ++nt) {
|
||||
cap[i++] = c7[(4 * mi + 0) * 8 + nt];
|
||||
cap[i++] = c7[(4 * mi + 1) * 8 + nt];
|
||||
}
|
||||
}
|
||||
capn = 80;
|
||||
}
|
||||
#endif
|
||||
uint32_t d7[6][2][8];
|
||||
gemm(as7, w, 16512, 0u, c7, 4, 8, t, d7);
|
||||
#if D4R_RRSWIN_DEBUG_STAGE == 7
|
||||
stage_regs(d7, 4, 8, cap);
|
||||
capn = 64;
|
||||
#endif
|
||||
|
||||
// ---- handing the accumulators back --------------------------------------
|
||||
// Phase 32's eight mma D pairs are not implemented yet. Until they are, fill the
|
||||
// sixteen slots with lane-distinct values so the read-back path is observable
|
||||
// without producing wild magnitudes (the epilogue's surface output changes from
|
||||
// the known-bad baseline). 64 bytes apart per lane, inside the 8 KiB scratch.
|
||||
lds_u32* const out = (lds_u32*)(uintptr_t)scratch + 16 * (uint32_t)lane_id();
|
||||
for (int i = 0; i < 16; ++i) {
|
||||
out[i] = (uint32_t)lane_id() + 32u * (uint32_t)i;
|
||||
}
|
||||
#if D4R_RRSWIN_DEBUG_STAGE
|
||||
// the payload goes out here, after every store this body makes: the oracle's store
|
||||
// is spliced in at the same program point (after E31's last arena write)
|
||||
debug_store(prm, cap, capn);
|
||||
#endif
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Weight image layout of the implemented phases (byte ranges of arg040, from the
|
||||
// PTX's own ld.weak.global loads; 512 bytes per group of 16 columns, K=32):
|
||||
// P1 [128,1152) P2 [2176,3200) P3 [1152,2176) P4 [3200,4224)
|
||||
// P5 bias [4224,16512) (24 tiles of 512, a 64x96 f16 block)
|
||||
// P6 none (P.V: B is the movmatrix transpose of phase 4's D)
|
||||
// P7 [16512,18560) column table [36992,37120)
|
||||
//
|
||||
// Fragment maps (the hardware's, identical to enc0's, VERIFIED against the debug
|
||||
// oracle there; g = lane>>2, t = lane&3):
|
||||
// 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
|
||||
// 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) -> same n, k = 8 + 2t + (j&1) + 16*((j>>1)&1)
|
||||
// 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
|
||||
Reference in new issue
Block a user