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).
411 lines
20 KiB
Plaintext
411 lines
20 KiB
Plaintext
// Native replacement for the Ray Reconstruction denoiser's last Swin decoder,
|
|
// cuda_dldn_engine_swin_dec4_kernel (nvngx_dlssd.dll 310.7).
|
|
//
|
|
// ---------------------------------------------------------------------------
|
|
// STATE: the mirror-cut frame is in place, builds, and *runs* (measured exit 0,
|
|
// after carrying the epilogue's six live registers across the cut -- see
|
|
// swin_dec4_recipe.py's RET_REGS). Phase 1 (the patch-expand GEMM, 300 mma in
|
|
// sixty independent five-step chains) is implemented, but it is NOT yet validated:
|
|
// the debug payload cannot be extracted from the launch dumps yet, because the
|
|
// bridge deduplicates param_0+48 against arg008's allocation and never writes that
|
|
// allocation out, so the oracle's store lands outside every dumped file. Phases
|
|
// 2-62 are unimplemented, so the six registers that cross the cut back are still
|
|
// lane-distinct placeholders and the layer's surface output does not match the
|
|
// translated kernel (compare_launch_dump: arg112 DIFFERS, arg120 ok).
|
|
//
|
|
// dec4 is the widest of the eleven layers: the attention/MLP width is 160 (dec3's
|
|
// is 128), so no count here is interchangeable with another layer's. It has 62
|
|
// phases and 1606 mma against enc0's 22 and 536, and its E1 alone is 44% of the
|
|
// entry.
|
|
//
|
|
// The kernel is one warp per 8x8 token patch with eight waves per block, block z
|
|
// selecting the slab (%r9 = %tid.z, guarded to 0..3). The PTX this file replaces
|
|
// is kernels/extracted/rr/0019-PREPASS_ENTRYPOINT_NAME.ptx.
|
|
//
|
|
// kernels/tex/make_ptx.py's recipe (kernels/rr/swin_dec4_recipe.py) keeps
|
|
// NVIDIA's texture prologue (file lines 1-1706) and surface epilogue (from line
|
|
// 58080) and replaces the 1606 mma of the 62 phases (lines 1707-58079) 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. dec4's block is 144 bytes.
|
|
struct DldnSwinDec4Params
|
|
{
|
|
uint8_t raw[144];
|
|
};
|
|
static_assert(sizeof(DldnSwinDec4Params) == 144, "parameter block");
|
|
|
|
enum : int
|
|
{
|
|
kPlaneArg = 48, // plane arena the decoder writes (arg048), also the debug base's arena
|
|
};
|
|
|
|
// The shared scratch the six registers the surviving epilogue reads cross the cut
|
|
// through: the top 768 bytes of `SM` (16384 bytes), twenty-four bytes per lane, at
|
|
// 15616. See swin_dec4_recipe.py (RET_BASE, RET_REGS) for why the top of the array is
|
|
// free -- SM itself is *not* dead after the cut, E62 reads it back at line 58152, but
|
|
// its deepest reach is 6328. A base past 15616 is out of bounds of the array, and the
|
|
// compiler then proves the store undefined and drops the whole body, so this constant
|
|
// has to match the recipe's.
|
|
enum : int
|
|
{
|
|
kRetBase = 15616,
|
|
kRetStride = 24,
|
|
};
|
|
|
|
#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 weight image is prepared: `NAME_prep` may permute it once, so the native side
|
|
// consumes it in the fragment order its own gemm wants rather than NVIDIA's. What
|
|
// follows is the mapping the *splice's* own reading of the image implies.
|
|
//
|
|
// Phase 1's B block is 25600 bytes: fifty 512-byte tiles at
|
|
// `W + z*25600 + 16*laneid + 512*i` (file lines 1478-1674), i.e. one tile per two
|
|
// B fragments, read as `ld.weak.global.ca.v4.u32` four b32 at a time. The hundred
|
|
// B fragments (%r1017..%r1216, two b32 each) are the mma's B operands in order.
|
|
//
|
|
// Within one 512-byte tile, byte j of a lane's sixteen holds the e4m3 for the
|
|
// (k, n) pair
|
|
// byte = 64*(n & 7) + 8*((n >> 3) & 1) + 16*((k >> 1) & 3)
|
|
// + 4*((k >> 3) & 1) + 2*((k >> 4) & 1) + (k & 1)
|
|
// and `16*laneid` is the tile's lane stride. VERIFIED: this is a bijection from the
|
|
// 512 (k, n) pairs of a tile onto the tile's 512 bytes, and it agrees with the
|
|
// fragment layout enc0's waddr() encodes for the same mma shape (checked over all
|
|
// 1024 pairs of a two-tile block), so the two layers' readers see one image.
|
|
//
|
|
// n here is the *tile-local* column: tile t covers n = 16*t .. 16*t+15, so a phase
|
|
// whose B block starts at `base` addresses column `n` of the phase at
|
|
// waddr(base, k, n) with n already phase-relative.
|
|
DEVF uint32_t waddr(int base, int k, int n)
|
|
{
|
|
return (uint32_t)base + 64u * (uint32_t)(n & 7) + 8u * (uint32_t)((n >> 3) & 1)
|
|
+ 16u * (uint32_t)((k >> 1) & 3) + 4u * (uint32_t)((k >> 3) & 1)
|
|
+ 2u * (uint32_t)((k >> 4) & 1) + (uint32_t)(k & 1);
|
|
}
|
|
|
|
// 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));
|
|
}
|
|
|
|
// 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 mma's k-split needs.
|
|
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
|
|
DEVF half_t half_lo(uint32_t pair)
|
|
{
|
|
return __builtin_bit_cast(half_t, (uint16_t)pair);
|
|
}
|
|
|
|
// the high f16 of a packed pair
|
|
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);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// One mma.m16n8k32.row.col.f16.e4m3.e4m3.f16 step, the shape every phase of dec4
|
|
// uses: sixteen rows, eight columns, K=32, e4m3 in, f16 out.
|
|
//
|
|
// `a` is the A fragment (four b32) in the hardware order enc0's comment documents:
|
|
// A reg 0 -> row g, k = 2t + (j&1) + 16*((j>>1)&1), byte j = 0..3
|
|
// 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 pairs with a0/a1 is b0, with a2/a3 is b1:
|
|
// 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)
|
|
// with b the byte index within the lane's sixteen and tt the 512-byte tile.
|
|
//
|
|
// `bv[2][32]` is therefore the lane's own two columns of the eight, over all 32 k,
|
|
// read straight from the weight image; the four lanes that differ in t hold the four
|
|
// disjoint eighths of each row's k, so the A side needs both butterfly steps to see
|
|
// all 32 k for its own t.
|
|
//
|
|
// `d` receives the two b32 the mma writes: d[0] is (row g, cols 2t, 2t+1) and d[1]
|
|
// is (row g+8, the same two columns), matching the PTX's D-register order.
|
|
//
|
|
// `c0`/`c1` are the C operand's two b32 (or zero for an unseeded step). e4m3 x e4m3
|
|
// products are exact and 32 of them sum exactly in f32, so the dot product is a plain
|
|
// f32 accumulation rounded to f16 once -- what the translated mma does.
|
|
DEVF void mma_step(const uint32_t* a, const float bv[2][32], uint32_t c0, uint32_t c1,
|
|
uint32_t d[2])
|
|
{
|
|
const int t = (int)(lane_id() & 3);
|
|
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
|
|
const uint32_t p0 = a[2 * reg + row];
|
|
const uint32_t p[4] = {p0, bfly1_bits(p0), bfly2_bits(p0), bfly1_bits(bfly2_bits(p0))};
|
|
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];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
const uint32_t c[2] = {c0, c1};
|
|
for (int row = 0; row < 2; ++row) {
|
|
d[row] = pair_of((half_t)(acc[row][0] + h2f(c[row])),
|
|
(half_t)(acc[row][1] + h2f(c[row] >> 16)));
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Ground truth for a stage comes from kernels/tex/make_ptx.py's debug recipe: with
|
|
// D4R_RRSWIN_DEBUG=<stage> (resolved from kernels/rr/swin_dec4_debug_stages.py) 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.
|
|
//
|
|
// The block index, the per-block slot and the lane stride all mirror
|
|
// make_ptx.py's swin_debug_block() exactly, including its 161-wide grid pitch and
|
|
// its skip of everything past the first 64 blocks.
|
|
#ifndef D4R_RRSWIN_DEBUG_STAGE
|
|
#define D4R_RRSWIN_DEBUG_STAGE 0
|
|
#endif
|
|
|
|
// The debug payload's byte offset into the plane arena. It must agree with
|
|
// D4R_RRSWIN_DEBUG_BASE, which make_ptx reads for the oracle splice. The default
|
|
// 46000000 is wrong for dec4: VERIFIED by storing param_0+48's value into a dumped
|
|
// argument, the arena is the 65,945,600-byte block at offset 39,567,360, so
|
|
// 39,567,360 + 46,000,000 runs 19.6 MB past the end of the block and the payload
|
|
// lands in memory no launch dump contains -- which is why every dec4 stage
|
|
// comparison first looked like a MATCH while comparing stale buffer content
|
|
// (the same window is byte-identical in the reference run, which stores nothing).
|
|
// 16 MiB leaves 64*slot bytes of room inside the block (the tail after the arena is
|
|
// 26,378,240) and clears arg008's own window at offset 2,282,752.
|
|
#ifndef D4R_RRSWIN_DEBUG_BASE
|
|
#define D4R_RRSWIN_DEBUG_BASE 16777216
|
|
#endif
|
|
|
|
enum : int
|
|
{
|
|
kDebugBase = D4R_RRSWIN_DEBUG_BASE,
|
|
kDebugBlocks = 64,
|
|
kDebugGridPitch = 161,
|
|
};
|
|
|
|
DEVF uint32_t debug_slot(int nregs)
|
|
{
|
|
const uint32_t want = (uint32_t)nregs * 128u > 8192u ? (uint32_t)nregs * 128u : 8192u;
|
|
uint32_t slot = 1;
|
|
while (slot < want) {
|
|
slot <<= 1;
|
|
}
|
|
return slot;
|
|
}
|
|
|
|
DEVF void debug_store(const uint8_t* prm, const uint32_t* regs, int n)
|
|
{
|
|
const uint32_t block = __builtin_amdgcn_workgroup_id_x()
|
|
+ kDebugGridPitch * __builtin_amdgcn_workgroup_id_y();
|
|
if (block >= (uint32_t)kDebugBlocks) {
|
|
return; // the oracle's per-block slots only cover the first 64 blocks
|
|
}
|
|
const uint64_t arena = *(const uint64_t*)(prm + kPlaneArg);
|
|
const uint32_t slot = debug_slot(n);
|
|
uint32_t* dst = (uint32_t*)(arena + (uint64_t)kDebugBase
|
|
+ ((uint64_t)block * slot)
|
|
+ (uint64_t)lane_id() * (uint64_t)n * 4ull);
|
|
for (int i = 0; i < n; ++i) {
|
|
dst[i] = regs[i];
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Phase 1: the patch-expand GEMM, the entry's first 300 mma.
|
|
//
|
|
// Shape: M = 240 (fifteen 16-row A fragments), N = 800 (a hundred 8-column B
|
|
// fragments), K = 160 (five k-steps), and 300 mma rather than the 7500 a dense
|
|
// 240x800x160 would need -- so this is a band, not a dense GEMM, and the tool's
|
|
// M/N columns are its "distinct fragment" counts rather than a real matrix shape.
|
|
//
|
|
// The mma order (VERIFIED by walking all 300 statements) is: sixty five-step chains,
|
|
// each seeded once and then accumulating through C. Statement i uses
|
|
// A fragment 5*(i/100) + (i/2)%5
|
|
// B fragment i % 100
|
|
// C the D pair of statement i-2, or the chain's seed when i%10 is 0 or 1
|
|
// so chain (g, q, p) -- g = 0..2, q = 0..9, p = 0,1 -- runs the statements
|
|
// 100g + 10q + p, +2, +4, +6, +8, carries A fragment 5*g + step, and is seeded with
|
|
// the seed register at index 10*q + p of the twenty (3227, 3237, 3327, 3337, ...),
|
|
// three chains per seed, one per A group. The sixty chains are disjoint: their
|
|
// accumulators are %r(1218 + 10*i) with i = 100g + 10q + p + 8, so the chain's own
|
|
// statement order is also its accumulator order (g major, then q, then p) -- which is
|
|
// the order the oracle's p1 stage lists the sixty pairs in.
|
|
//
|
|
// The B block is [0, 25600) per slab: `W + z*25600 + 16*laneid + 512*i` for the
|
|
// hundred fragments, fifty 512-byte tiles, two fragments per tile. Fragment f lives
|
|
// at tile f/2 with `8*(f&1)` selecting which half, so its eight columns are
|
|
// n = 8*f .. 8*f+7 phase-relative -- which is exactly waddr()'s phase-relative n.
|
|
//
|
|
// The lane's two columns of fragment f are 2t and 2t+1, so bv[c][k] is read at
|
|
// n = 8*f + 2*t + c.
|
|
DEVF void phase1(const uint32_t* as, const uint8_t* w, const uint32_t* seeds, int z,
|
|
const uint8_t* prm)
|
|
{
|
|
const int t = (int)(lane_id() & 3);
|
|
const int base = z * 25600;
|
|
uint32_t acc[60][2];
|
|
|
|
for (int chain = 0; chain < 60; ++chain) {
|
|
const int g = chain / 20; // the A group: statements 100g..100g+99
|
|
const int sp = chain % 20; // 10*q + p: the seed index and the B fragment base
|
|
uint32_t c0 = seeds[sp];
|
|
uint32_t c1 = c0; // PTX: {%r3227, %r3227} -- one u32 in both halves
|
|
for (int step = 0; step < 5; ++step) {
|
|
const int af = 5 * g + step;
|
|
const int bf = sp + 2 * step;
|
|
float bv[2][32];
|
|
for (int c = 0; c < 2; ++c) {
|
|
const int n = 8 * bf + 2 * t + c;
|
|
for (int k = 0; k < 32; ++k) {
|
|
bv[c][k] = e4m3_to_f32(w[waddr(base, k, n)]);
|
|
}
|
|
}
|
|
mma_step(&as[4 * af], bv, c0, c1, acc[chain]);
|
|
c0 = acc[chain][0];
|
|
c1 = acc[chain][1];
|
|
}
|
|
}
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 1
|
|
{
|
|
// the oracle's p1 stage: the sixty chain-final D pairs in mma statement
|
|
// order, each as the (row g, row g+8) pair its mma wrote
|
|
uint32_t regs[120];
|
|
for (int c = 0; c < 60; ++c) {
|
|
regs[2 * c] = acc[c][0];
|
|
regs[2 * c + 1] = acc[c][1];
|
|
}
|
|
debug_store(prm, regs, 120);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 2
|
|
{
|
|
// the oracle's p1a stage: the fifteen A fragments, four b32 each, in the
|
|
// first mma's operand order -- the sixty registers the splice hands over
|
|
uint32_t regs[60];
|
|
for (int i = 0; i < 60; ++i) {
|
|
regs[i] = as[i];
|
|
}
|
|
debug_store(prm, regs, 60);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 3
|
|
{
|
|
// the oracle's p1s stage: the twenty bias seeds in chain-root order
|
|
uint32_t regs[20];
|
|
for (int i = 0; i < 20; ++i) {
|
|
regs[i] = seeds[i];
|
|
}
|
|
debug_store(prm, regs, 20);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Mirror cut. dec4's first phase reads its A operands from the e4m3 feature buffer
|
|
// in the prologue (param+8, a clamped halo gather, file lines 1063-1468), so its
|
|
// fifteen 16-row fragments arrive here already in registers: %r2130..%r2213,
|
|
// %r3130..%r3213 and %r4130..%r4213, four b32 each in the first mma's operand
|
|
// order. Then phase 1's twenty bias seeds, the weight image (arg040) and the
|
|
// parameter block; see kernels/rr/swin_dec4_recipe.py for the register lists and
|
|
// the file-line evidence behind each.
|
|
//
|
|
// The four accumulators phase 62 leaves live across the cut (%r34761, %r34762,
|
|
// %r34771, %r34772, read back out of `SM` by E62's softmax at line 58152) cross it
|
|
// 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 a 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 four
|
|
// slots. `scratch` is the base of `SM` (16384 bytes), whose top 128 bytes no
|
|
// ld.shared/st.shared in the entry can reach.
|
|
extern "C" __attribute__((device)) void d4r_swin_dec4(
|
|
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 a48,
|
|
uint32_t a49, uint32_t a50, uint32_t a51, uint32_t a52, uint32_t a53, uint32_t a54, uint32_t a55,
|
|
uint32_t a56, uint32_t a57, uint32_t a58, uint32_t a59,
|
|
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, uint32_t b16, uint32_t b17, uint32_t b18, uint32_t b19,
|
|
const uint8_t* w, const uint8_t* prm, uint32_t scratch)
|
|
{
|
|
const uint32_t as1[60] = {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,
|
|
a36, a37, a38, a39, a40, a41, a42, a43, a44, a45, a46, a47,
|
|
a48, a49, a50, a51, a52, a53, a54, a55, a56, a57, a58, a59};
|
|
const uint32_t seeds[20] = {b0, b1, b2, b3, b4, b5, b6, b7, b8, b9,
|
|
b10, b11, b12, b13, b14, b15, b16, b17, b18, b19};
|
|
// the slab index: %r9 = %tid.z, guarded to 0..3 by the prologue's `setp.gt.u32`
|
|
const int z = (int)__builtin_amdgcn_workitem_id_z();
|
|
// Address space 3 (the shared scratch) explicitly: a generic store lowers to flat,
|
|
// which faults. The splice reads the same six words back at SM + 15616 + 24*laneid
|
|
// (swin_dec4_recipe.py's RET_BASE/RET_STRIDE).
|
|
__attribute__((address_space(3))) uint32_t* out =
|
|
((__attribute__((address_space(3))) uint32_t*)(uintptr_t)scratch)
|
|
+ kRetBase / 4 + (uint32_t)lane_id() * (kRetStride / 4);
|
|
|
|
// ---- phase 1 -----------------------------------------------------------
|
|
phase1(as1, w, seeds, z, prm);
|
|
|
|
// ---- phases 2-62: not implemented yet ----------------------------------
|
|
// The three D pairs the epilogue reads (%r34761/62, %r34771/72, %r34888/89) are
|
|
// still placeholders. Fill them with lane-distinct values so the read-back path is
|
|
// observable without producing wild magnitudes. Every one of the six must be
|
|
// written: a register PTX leaves undefined is poison to ZLUDA's optimizer, and one
|
|
// such survivor folds the whole entry to `unreachable` (measured: HSA_STATUS_ERROR_ILLEGAL_INSTRUCTION).
|
|
for (int i = 0; i < 6; ++i) {
|
|
out[i] = (uint32_t)lane_id() + 32u * (uint32_t)i;
|
|
}
|
|
}
|