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

kernels/rr: staged-image Swin layer engine prototypes

Browse Source 
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 2026-10-12 12:47:53 +02:00
1 parent 798e25611c
commit 0fe3a28cb7
58 files changed
+28330 -5
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/rr/rrswin_dec0.hip
+376
View File
@@ -0,0 +1,376 @@
// Native replacement for the Ray Reconstruction denoiser's first Swin decoder,
// cuda_dldn_engine_swin_dec0_kernel (nvngx_dlssd.dll 310.7).
//
// ---------------------------------------------------------------------------
// STATE: phase 1 (the 48x128x64 patch GEMM, three m-tiles of 16, sixteen n-tiles
// of 8, two k32 steps) is implemented and bit-exact against the debug oracle;
// phases 2-19 are not, so the layer's output is still the known-bad baseline.
// See "Validation" at the bottom for what has been compared byte for byte.
//
// ---------------------------------------------------------------------------
// The kernel is one warp per 8x8 token patch: grid (161, 93, 1), block
// (32, 1, 1) -- one wave, so every %tid.z term is absent from the addressing --
// 14973 blocks, parameter block 152 bytes. The PTX this file replaces is
// kernels/extracted/rr/0023-PREPASS_ENTRYPOINT_NAME.ptx.
//
// kernels/tex/make_ptx.py's recipe (kernels/rr/swin_dec0_recipe.py,
// D4R_RRSWIN_CUT=1) keeps NVIDIA's plane/texture prologue (lines 1-1392) and its
// surface epilogue (46264-end) and replaces the 544 mma of the 19 phases with the
// call below.
//
// 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 DldnSwinDec0Params
{
uint8_t raw[152];
};
static_assert(sizeof(DldnSwinDec0Params) == 152, "parameter block");
enum : int
{
kIn = 8, // the e4m3 input buffer, already in phase 1's A fragment order
kSkip = 24, // the skip input E1 adds to the staged patch
kWeights = 40, // NVIDIA's prepared weight image (arg040), read-only
kPlaneArg = 48, // the output buffer (arg048); make_ptx.py's debug splice writes here
kBlockOff = 80, // the block origin, read as a v2.u16
kScale = 136, // the texture coordinate scale/origin (E19)
};
#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)));
// ---------------------------------------------------------------------------
// 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 offset, from a phase's own base, of the e4m3 weight element (k, Xt) of
// weight tile T, where Xt is the column inside the 512-byte tile (0..15).
//
// A tile is 32 lanes x 16 bytes: lane 4*g+t holds the two fragments of n-tiles
// 2*(T-ks) and 2*(T-ks)+1 at bytes 0..7 and 8..15. The mma's B fragment holds
// the lane's two *columns* 2t and 2t+1, four k values each (B is read as
// [n][k], and D's columns are 2t, 2t+1), so byte b of register r is
// n = 8*(Xt>>3) + (Xt&1), k = <the lane's k pattern> with lane = 4*g'+t.
// That fixes the address from the element's (k, column):
// bit 0 = k&1, bit 1 = (k>>4)&1, bit 2 = column&1, bit 3 = column>>3,
// bits 4-5 = (column>>1)&3, bits 6-8 = g'.
// Phase 1's tile index is T = 2*(nt>>1) + ks (the two k slices of an n-tile pair
// sit in the two 512-byte tiles that pair occupies), and Xt = 8*(nt&1) + 2t + c.
DEVF uint32_t w_byte(int k, int Xt, int T)
{
return 512u * (uint32_t)T + 8u * (uint32_t)(Xt >> 3) + 4u * (uint32_t)(Xt & 1)
+ 16u * (uint32_t)((Xt >> 1) & 3) + 2u * (uint32_t)((k >> 4) & 1)
+ (uint32_t)(k & 1) + 64u * (uint32_t)((k >> 1) & 7);
}
// The same map for the other candidate: byte b holds four *consecutive* k, so
// k = 4*g' + b and the address takes bits 0-1 from k&3 and bits 6-8 from k>>2.
// Both maps are bijections of the 32 k x 16 columns onto the 512-byte tile, so
// only the oracle can say which one NVIDIA's mma uses; D4R_RRSWIN_DEBUG_PROBE
// stores both.
DEVF uint32_t w_byte_c(int k, int Xt, int T)
{
return 512u * (uint32_t)T + 8u * (uint32_t)(Xt >> 3) + 4u * (uint32_t)(Xt & 1)
+ 16u * (uint32_t)((Xt >> 1) & 3) + (uint32_t)(k & 3)
+ 64u * (uint32_t)(k >> 2);
}
// warp butterfly over the four lanes that hold a row's 32 k values (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);
}
// 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);
}
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);
}
DEVF void gemm_p1(const uint32_t* as, const uint32_t* seed, const uint8_t* w, int t,
int g, int aord, int apat, int bpat, uint32_t d[3][2][16])
{
for (int nt = 0; nt < 16; ++nt) {
const uint32_t bias = seed[nt];
for (int mi = 0; mi < 3; ++mi) {
half_t prev[2][2];
#pragma unroll
for (int row = 0; row < 2; ++row) {
for (int c = 0; c < 2; ++c) {
prev[row][c] = __builtin_bit_cast(half_t, (uint16_t)0);
}
}
for (int ks = 0; ks < 2; ++ks) {
const int T = 2 * (nt >> 1) + ks;
float bv[2][32];
for (int c = 0; c < 2; ++c) {
const int Xt = 8 * (nt & 1) + 2 * t + c;
for (int k = 0; k < 32; ++k) {
bv[c][k] = e4m3_to_f32(w[bpat ? w_byte_c(k, Xt, T)
: w_byte(k, Xt, T)]);
}
}
float acc[2][2] = {{0.0f, 0.0f}, {0.0f, 0.0f}};
for (int kg = 0; kg < 2; ++kg) {
for (int row = 0; row < 2; ++row) {
// this lane's word and the same word on lanes t^1, t^2, t^3
const int idx = aord ? 2 * row + kg : 2 * kg + row;
const uint32_t p0 = as[4 * (2 * mi + ks) + idx];
uint32_t p[4];
p[0] = p0;
p[1] = bfly1_bits(p0);
p[2] = bfly2_bits(p0);
p[3] = bfly1_bits(p[2]);
for (int m = 0; m < 4; ++m) {
for (int j = 0; j < 4; ++j) {
const int k = apat
? 16 * kg + 4 * (t ^ m) + j
: 8 * kg + 2 * (t ^ m) + (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];
}
}
}
}
// the chain's C operand: the bias seed on the first k step, the
// first step's f16 result on the second
for (int row = 0; row < 2; ++row) {
for (int c = 0; c < 2; ++c) {
const float base = ks == 0
? h2f(c == 0 ? bias : (bias >> 16))
: (float)prev[row][c];
prev[row][c] = (half_t)(acc[row][c] + base);
}
}
}
for (int row = 0; row < 2; ++row) {
d[mi][row][nt] = pair_of(prev[row][0], prev[row][1]);
}
}
}
}
// ---------------------------------------------------------------------------
// One K=32 GEMM, the shape of phases 2 and 3 (and of every later unfused
// projection): NMI m-tiles of 16 rows x NNT n-tiles of 8 columns, one k32 step,
// no bias. `as` is the phase's A operand in mma operand order (NMI fragments of
// four b32), `d` receives the NMI*NNT tile pairs as d[(mi*NNT + nt)*2 + row],
// the (row g, row g+8) pair of each tile.
//
// The weight image of such a phase is NNT/2 tiles of 512 bytes: sixteen columns
// (two n-tiles) and 32 k values each, at the phase's own byte base -- which is
// not a multiple of 512, so the tile index n>>4 is added to that base rather
// than folded into it.
DEVF void gemm_k32(const uint32_t* as, const uint8_t* w, int base, int nmi, int nnt,
int t, uint32_t* d)
{
for (int nt = 0; nt < nnt; ++nt) {
for (int mi = 0; mi < nmi; ++mi) {
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[base + w_byte(k, n, 0)]);
}
}
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) {
const uint32_t p0 = as[4 * mi + 2 * reg + row];
uint32_t p[4];
p[0] = p0;
p[1] = bfly1_bits(p0);
p[2] = bfly2_bits(p0);
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) {
d[(mi * nnt + nt) * 2 + row] =
pair_of((half_t)acc[row][0], (half_t)acc[row][1]);
}
}
}
}
// ---------------------------------------------------------------------------
// Ground truth for a stage comes from kernels/tex/make_ptx.py's debug recipe:
// with D4R_RRSWIN_DEBUG=<stage> the same splice stores that stage's registers
// into the buffer at arg048 + 46,000,000, which dec0 never writes, and the launch
// dump then holds NVIDIA's 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. The stage table is kernels/rr/swin_dec0_debug_stages.py.
#ifndef D4R_RRSWIN_DEBUG_STAGE
#define D4R_RRSWIN_DEBUG_STAGE 0
#endif
#ifndef D4R_RRSWIN_DEBUG_BASE
// dec0's param_0+48 sits 45,219,840 bytes into the block the launch dump files
// as arg008, so the default 46,000,000 would land past the end of the block and
// every comparison would read ambient memory. 16 MiB keeps the whole
// 64-block payload (64 * 16384) inside it; the oracle splice reads the same
// value from D4R_RRSWIN_DEBUG_BASE.
#define D4R_RRSWIN_DEBUG_BASE 16777216
#endif
DEVF void debug_store_off(const uint8_t* prm, const uint32_t* regs, int n,
uint64_t extra, uint64_t slot)
{
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);
uint32_t* dst = (uint32_t*)(arena + (uint64_t)D4R_RRSWIN_DEBUG_BASE + extra
+ (uint64_t)block * slot
+ (uint64_t)lane_id() * (uint64_t)n * 4ull);
for (int i = 0; i < n; ++i) {
dst[i] = regs[i];
}
}
DEVF void debug_store(const uint8_t* prm, const uint32_t* regs, int n)
{
// 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 (its splice forms block * slot with
// a shl)
const uint64_t want = (uint64_t)n * 128ull > 8192ull ? (uint64_t)n * 128ull : 8192ull;
uint64_t slot = 1;
while (slot < want) {
slot <<= 1;
}
debug_store_off(prm, regs, n, 0, slot);
}
// ---------------------------------------------------------------------------
// The mirror cut. The six 16-row fragments of the staged patch arrive in the
// order the first mma reads them (m-tile, k-slice), then the sixteen bias seeds,
// the weight image and the parameter block. The accumulators the epilogue reads
// cross back through shared memory, not `call.uni`'s return parameters (see
// kernels/rr/rrswin_enc0.hip for why): `scratch` is the base of the entry's
// `smem` array (3200 bytes of .shared, its last access is inside the cut), and
// swin_dec0_recipe.py reads 25 b32 per lane back out of it, 100 bytes apart.
extern "C" __attribute__((device)) void d4r_swin_dec0(
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 c0, uint32_t c1, uint32_t c2, uint32_t c3, uint32_t c4, uint32_t c5,
uint32_t c6, uint32_t c7, uint32_t c8, uint32_t c9, uint32_t c10, uint32_t c11,
uint32_t c12, uint32_t c13, uint32_t c14, uint32_t c15,
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 seed[16] = {c0, c1, c2, c3, c4, c5, c6, c7,
c8, c9, c10, c11, c12, c13, c14, c15};
// Address space 3 (the shared scratch) explicitly: a generic store lowers to
// flat, which faults.
lds_u32* const out = (lds_u32*)(uintptr_t)scratch + 25 * (uint32_t)lane_id();
const int t = (int)(lane_id() & 3); // the lane's two columns of each n-tile
const uint32_t lane = lane_id();
// ---- phase 1 -----------------------------------------------------------
uint32_t d1[3][2][16]; // [m-tile][row half][n-tile]: f16 pair
const int g = (int)(lane_id() >> 2); // the lane's column of each n-tile
// measured against the p1 oracle: see the "Validation" note below
#ifndef D4R_RRSWIN_AORD
#define D4R_RRSWIN_AORD 0
#endif
#ifndef D4R_RRSWIN_APAT
#define D4R_RRSWIN_APAT 0
#endif
gemm_p1(as1, seed, w, t, g, D4R_RRSWIN_AORD, D4R_RRSWIN_APAT, D4R_RRSWIN_APAT, d1);
#if D4R_RRSWIN_DEBUG_STAGE == 1
{
// the oracle's stage p1: the 48 chain-end accumulator pairs, m-tile major,
// n-tile minor
uint32_t regs[96];
int i = 0;
for (int mi = 0; mi < 3; ++mi) {
for (int nt = 0; nt < 16; ++nt) {
regs[i++] = d1[mi][0][nt];
regs[i++] = d1[mi][1][nt];
}
}
debug_store(prm, regs, 96);
}
#endif
#ifdef D4R_RRSWIN_DEBUG_PROBE
// The A register order and the A/B byte->k patterns are not determined by
// the PTX, so the four combinations are computed in one launch and stored a
// mebibyte apart; compare_stage.py re-tiles the oracle's payload to match.
for (int v = 0; v < 4; ++v) {
uint32_t dv[3][2][16];
gemm_p1(as1, seed, w, t, g, v & 1, (v >> 1) & 1, (v >> 1) & 1, dv);
uint32_t regs[96];
int i = 0;
for (int mi = 0; mi < 3; ++mi) {
for (int nt = 0; nt < 16; ++nt) {
regs[i++] = dv[mi][0][nt];
regs[i++] = dv[mi][1][nt];
}
}
debug_store_off(prm, regs, 96, (uint64_t)v * 1048576ull, 16384ull);
}
#endif
// ---- handing the accumulators back ------------------------------------
// E19's chains have not been reproduced yet, so the 24 chain values and E1's
// shuffle default are zero; the two constants are reloaded by the splice.
for (int i = 0; i < 25; ++i) {
out[i] = 0u;
}
}