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_dec1.hip
+336
View File
@@ -0,0 +1,336 @@
// Native replacement for the Ray Reconstruction denoiser's third decoder,
// cuda_dldn_engine_swin_dec1_kernel (nvngx_dlssd.dll 310.7).
//
// ---------------------------------------------------------------------------
// STATE: phase 1 (the 48x384 patch-expand GEMM, 144 mma) is implemented and
// bit-exact against the GPU oracle -- debug stage p1, 4,194,304 payload bytes
// over 64 blocks x 32 lanes x 288 registers, 0 differing bytes, identical md5
// (ae431ae0baff). Stage a1 (the 52 staged input registers) is likewise
// 100.0000%. E1 and phases 2-26 are not implemented yet, so the eight
// accumulators the epilogue reads back are placeholders: the served frame is
// wrong (compare_launch_dump.py: 14 arguments compared, 4 differ -- arg[8] the
// feature plane E25 writes and arg[112] the output surface), and the layer
// costs 0.036 ms/frame instead of the translated 1-2 ms, because phase 1's
// accumulators are dead after the read-back.
//
// The next piece is E1 (s547-s7095): the 48 packed patch cells per lane into a
// 10x10x4 shared grid, the 3x3-ish gather from param_0+24, the per-column RMS
// norm (raw sum of squares + 2^-13, rsqrt.approx.ftz, gain at weight +25088) and
// the restage into phases 2-5's A operands. Its lane -> (sub-array, row, col)
// map is 438 basic blocks in the PTX and is *not* characterised in
// rrswin_dec1_epilogues.md (its U4), so it has to be lifted out of the index
// arithmetic rather than guessed.
//
// The mirror-cut splice is kernels/rr/swin_dec1_recipe.py: it keeps NVIDIA's
// texture prologue (which stages phase 1's A fragments and bias seeds) and its
// surface epilogue (the shared staging, the softmax, the three texture reads
// and the surface store), and replaces physical lines 1486-38472 -- the whole
// 26-phase mma region, 632 mma -- with one call into d4r_swin_dec1.
//
// 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.
//
// ---------------------------------------------------------------------------
// Launch: grid 81x47, block 32x1x2 (two waves per block, which is why the weight
// addressing carries a per-tid.z stride), parameter block 144 bytes. The PTX this
// file replaces is ~/.cache/d4r-rr-corpus/0022-PREPASS_ENTRYPOINT_NAME.ptx.
#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 DldnSwinDec1Params
{
uint8_t raw[144];
};
static_assert(sizeof(DldnSwinDec1Params) == 144, "parameter block");
enum : int
{
kExtents = 0, // ex, ey: two u32 (s14)
kInPlane = 8, // input feature plane arena, read by phase 1's A loads
kPatchGather = 24, // the skip input E1 gathers the 10x10 patch from
kWeights = 40, // NVIDIA's prepared weight image (arg040), read-only
kOutPlane = 48, // plane arena E25 scatters phase 25's output into
};
#define DEVF __attribute__((device)) static inline
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 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)
// (the settled formula, VERIFIED on enc0's 1024 (n, k) pairs per tile: the
// 512-byte block is 32 lanes of 16 contiguous bytes, and 64*(n&7) + 16*((k>>1)&3)
// is exactly 16*laneid of the lane that owns the byte, so the remaining four terms
// are that byte's index inside the lane's slot).
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);
}
// 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);
}
// ---------------------------------------------------------------------------
// One m16n8k32 e4m3 mma, f16 accumulate, done the way the hardware does it.
//
// The lane's four A registers hold, for each of its two rows (row g and row g+8),
// the eight e4m3 values of k = 2t + (j&1) + 16*((j>>1)&1) (byte j = 0..3) and of
// k + 8; the four lanes of the group hold the four disjoint eighths of the 32 k's,
// so two butterfly steps bring the other three in. The B operand is read from the
// weight image at waddr() for the lane's own two D columns (D register 0 is row g,
// columns 2t and 2t+1; register 1 is row g+8). Every e4m3 product is exact and 32
// of them sum exactly in f32, so a plain f32 accumulation rounded to f16 once is
// what the translated mma computes.
//
// a[0..3] the A fragment as the mma reads it
// w the weight image
// base byte offset of the B tile's first 512-byte block
// n0 the mma's first output column (a multiple of 8)
// c0, c1 the C operand's two registers (row g, row g+8)
// t laneid & 3
// d[0], d[1] the two D registers
DEVF void mma_tile(const uint32_t* a, const uint8_t* w, int base, int n0,
uint32_t c0, uint32_t c1, int t, uint32_t d[2])
{
float bv[2][32];
#pragma unroll
for (int c = 0; c < 2; ++c) {
const int n = n0 + 2 * t + c;
#pragma unroll
for (int k = 0; k < 32; ++k) {
bv[c][k] = e4m3_to_f32(w[waddr(base, k, n)]);
}
}
#pragma unroll
for (int row = 0; row < 2; ++row) {
float acc[2] = {0.0f, 0.0f};
#pragma unroll
for (int reg = 0; reg < 2; ++reg) {
const uint32_t q = a[2 * reg + row];
uint32_t p[4];
p[0] = q;
p[1] = bfly1_bits(q);
p[2] = bfly2_bits(q);
p[3] = bfly1_bits(p[2]);
#pragma unroll
for (int m = 0; m < 4; ++m) {
const int kk = 2 * (t ^ m);
#pragma unroll
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[0] += av * bv[0][k];
acc[1] += av * bv[1][k];
}
}
}
const uint32_t cpair = row ? c1 : c0;
d[row] = pair_of((half_t)(acc[0] + h2f(cpair)),
(half_t)(acc[1] + h2f(cpair >> 16)));
}
}
// ---------------------------------------------------------------------------
// Ground truth for a stage comes from kernels/rr/swin_dec1_recipe.py's debug
// stages: with D4R_RRSWIN_DEBUG=a1 (or p1) the same splice stores that stage's
// registers into the param_0+8 feature-plane buffer at +46,000,000, which the
// entry has finished reading by then (the prologue's last read is s324, before
// the cut at line 1486), and the launch dump then holds the translated values
// for a block/lane. Built with -DD4R_RRSWIN_DEBUG_STAGE=<n>
// (kernels/tex/build_tex.sh, D4R_TEX_CFLAGS) this body stores the same
// registers of its own stage 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()
+ 81u * __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 + kInPlane);
// the oracle's slot, swin_dec1_recipe.py's 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 + (uint64_t)block * slot
+ (uint64_t)lane_id() * (uint64_t)n * 4ull);
for (int i = 0; i < n; ++i) {
dst[i] = regs[i];
}
}
// ---------------------------------------------------------------------------
// Mirror cut. The nine 16-row A fragments of the staged block input arrive in
// this order (the first mma's operand order, .f16.e4m3 fragments of four b32
// each):
// f0 %r1084.. f1 %r1104.. f2 %r1124.. (m-group 0's three steps)
// f3 %r1564.. f4 %r1584.. f5 %r1604.. (m-group 1)
// f6 %r2044.. f7 %r2064.. f8 %r2084.. (m-group 2)
// then phase 1's sixteen bias seeds (%r1621 %r1631 ... %r2041 %r2051, read from
// W + 256*tid.z + 4*(laneid&3) + 24576..24816), the weight image and the parameter
// block. No other C register is passed: every other mma's C operand is another
// mma's D register, which this body computes itself.
//
// The ten values the epilogue reads out of the deleted region cross the cut
// through the shared array, not `call.uni`'s return parameters: ZLUDA lowers
// those to the function's LLVM return type, so there is no .param-space object
// and no pointer is formed. `scratch` is the base of the entry's 6400-byte
// .shared array, 40 bytes apart per lane.
extern "C" __attribute__((device)) void d4r_swin_dec1(
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 b0, uint32_t b1, uint32_t b2, uint32_t b3, uint32_t b4, uint32_t b5,
uint32_t b6, uint32_t b7, uint32_t b8, uint32_t b9, uint32_t b10, uint32_t b11,
uint32_t b12, uint32_t b13, uint32_t b14, uint32_t b15,
const uint8_t* w, const uint8_t* prm, uint32_t scratch)
{
const uint32_t as1[36] = {a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11,
a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23,
a24, a25, a26, a27, a28, a29, a30, a31, a32, a33, a34, a35};
const uint32_t seeds[16] = {b0, b1, b2, b3, b4, b5, b6, b7,
b8, b9, b10, b11, b12, b13, b14, b15};
// The epilogue's ten registers: phase 26's eight accumulator halves, then
// %r18991 and %r18993. Address space 3 (the shared scratch) explicitly: a
// generic store lowers to flat, which faults.
lds_u32* const out = (lds_u32*)(uintptr_t)scratch + 10 * (uint32_t)lane_id();
const uint32_t lane = lane_id();
const int t = (int)(lane & 3); // the lane's two columns of each n-tile
#if D4R_RRSWIN_DEBUG_STAGE == 1
{
uint32_t regs[52];
#pragma unroll
for (int i = 0; i < 36; ++i) {
regs[i] = as1[i];
}
#pragma unroll
for (int i = 0; i < 16; ++i) {
regs[36 + i] = seeds[i];
}
debug_store(prm, regs, 52);
}
#endif
// ---- phase 1 -----------------------------------------------------------
// The 144 mma of s403-s546 are three m-groups of 48 tiles. Within a group the
// 48 tiles are eight chains of three: chain gp starts from the bias seed
// 2*gp + h (broadcast into both C registers), takes its A operand from
// fragment 3*mgrp + j at step j, its B operand from the 512-byte block
// p = 3*gp + j at the output columns 16*p + 8*h, and adds to the previous
// step's D -- so the chain is a running sum along j, not a k-accumulation.
// The tile index is i = 48*mgrp + 2*p + h; the output is 48 rows by 384
// columns and every element is written by exactly one tile (m-tile i/48,
// n-tile i%48). The weight image is per wave: the B blocks at
// W + 12288*tid.z, the seeds at W + 256*tid.z + 4*t + 24576 + 16*seed.
const int base1 = 12288 * (int)__builtin_amdgcn_workitem_id_z();
uint32_t d1[3][8][3][2][2]; // [m-group][chain][step][n-tile half][row half]
#pragma unroll
for (int mgrp = 0; mgrp < 3; ++mgrp) {
#pragma unroll
for (int gp = 0; gp < 8; ++gp) {
#pragma unroll
for (int j = 0; j < 3; ++j) {
#pragma unroll
for (int h = 0; h < 2; ++h) {
const int p = 3 * gp + j;
const uint32_t seed = seeds[2 * gp + h];
const uint32_t c0 = j == 0 ? seed : d1[mgrp][gp][j - 1][h][0];
const uint32_t c1 = j == 0 ? seed : d1[mgrp][gp][j - 1][h][1];
mma_tile(as1 + 4 * (3 * mgrp + j), w, base1, 16 * p + 8 * h,
c0, c1, t, d1[mgrp][gp][j][h]);
}
}
}
}
#if D4R_RRSWIN_DEBUG_STAGE == 2
{
// the oracle's stage p1: the 144 mma D pairs in mma order
uint32_t regs[288];
int i = 0;
#pragma unroll
for (int mgrp = 0; mgrp < 3; ++mgrp) {
#pragma unroll
for (int gp = 0; gp < 8; ++gp) {
#pragma unroll
for (int j = 0; j < 3; ++j) {
#pragma unroll
for (int h = 0; h < 2; ++h) {
regs[i++] = d1[mgrp][gp][j][h][0];
regs[i++] = d1[mgrp][gp][j][h][1];
}
}
}
}
debug_store(prm, regs, 288);
}
#endif
// ---- handing the accumulators back ------------------------------------
// %r18991 = 32*tid.z (PTX line 38396) and %r18993 = the shared base (38394);
// both sit inside the cut, so the epilogue cannot recompute them. Phase 26's
// eight accumulator halves (%r18794 %r18795 %r18804 %r18805 %r18834 %r18835
// %r18844 %r18845) are not implemented yet: lane-distinct placeholders keep
// the read-back path observable without producing wild magnitudes.
out[0] = lane;
out[1] = lane + 1u;
out[2] = lane + 2u;
out[3] = lane + 3u;
out[4] = lane + 4u;
out[5] = lane + 5u;
out[6] = lane + 6u;
out[7] = lane + 7u;
out[8] = 32u * __builtin_amdgcn_workitem_id_z();
out[9] = scratch;
}