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).
387 lines
20 KiB
Plaintext
387 lines
20 KiB
Plaintext
// Native replacement for the Ray Reconstruction denoiser's third Swin decoder,
|
|
// cuda_dldn_engine_swin_dec3_kernel (nvngx_dlssd.dll 310.7).
|
|
//
|
|
// ---------------------------------------------------------------------------
|
|
// STATE: phases 1 is implemented; phases 2-48 are not. Phase 1 does NOT yet reproduce
|
|
// the oracle's payload: against the dec3 `p1` stage the frame-1 payload differs in
|
|
// every one of the 96 accumulators of every lane of all 64 slots, so this layer's
|
|
// image does not match the translated reference (arg112 differs, arg120 is still ok).
|
|
// The store plumbing itself is proven correct -- the `a` stage, which carries the same
|
|
// 60 A fragments and 16 seeds the cut hands over, is byte-identical (md5
|
|
// 6711adc5517ea8d6ce16d908946189e4 both sides). What phase 1 still gets wrong is the
|
|
// mma's cross-lane structure: an m16n8k32 lane's B fragment holds one column's k
|
|
// slice, so D[row g][cols 2t,2t+1] needs a reduction across the eight column groups,
|
|
// not a butterfly within the four-lane quad. ZLUDA's fallback_mma_sync_aligned
|
|
// (zluda_ptx_impl.cpp:1226) does that with dpp and ds_bpermute; see the comment above
|
|
// phase1() for what has been established so far. Phases 2-48's accumulators are
|
|
// written back as lane-distinct placeholders, so the kernel runs and is served.
|
|
//
|
|
// The kernel is four waves per block (blockDim (32, 1, 4)); %tid.z is the slab index
|
|
// that scales the input stride, the weight-image stride and the bias stride. The PTX
|
|
// this file replaces is ~/.cache/d4r-rr-corpus/0020-PREPASS_ENTRYPOINT_NAME.ptx, whose
|
|
// entry opens at file line 1016 and whose mma region is statements s543-s19381 (file
|
|
// lines 1658-46565).
|
|
//
|
|
// kernels/rr/swin_dec3_recipe.py (D4R_RRSWIN_CUT, the default) keeps NVIDIA's texture
|
|
// prologue (file lines 1-1657) and surface epilogue (46566-end) and replaces the 1112
|
|
// mma of the 48 phases with the call below. The accumulators cross the cut through
|
|
// shared memory, because ZLUDA lowers a multi-value C++ return to LLVM's function return
|
|
// type and never forms a callee-side .param-space object.
|
|
//
|
|
// 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 DldnSwinDec3Params
|
|
{
|
|
uint8_t raw[144];
|
|
};
|
|
static_assert(sizeof(DldnSwinDec3Params) == 144, "parameter block");
|
|
|
|
enum : int
|
|
{
|
|
kWeights = 40, // NVIDIA's prepared weight image (arg040), read-only
|
|
kPlaneArg = 48, // plane arena pointer (arg048)
|
|
};
|
|
|
|
#define DEVF __attribute__((device)) static inline
|
|
typedef uint32_t lds_u32 __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));
|
|
}
|
|
|
|
// 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);
|
|
}
|
|
|
|
// word v as the lane t^m of its group of four holds it, m = 0..3
|
|
DEVF uint32_t bfly_m(uint32_t v, int m)
|
|
{
|
|
return m == 0 ? v : (m == 1 ? bfly1_bits(v) : (m == 2 ? bfly2_bits(v) : bfly1_bits(bfly2_bits(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);
|
|
}
|
|
|
|
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);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// The mma's fragment layouts, which are fixed by the instruction (they are not
|
|
// NVIDIA's choice, so they are the same for every layer):
|
|
// A reg q, byte j -> row g + 8*(q&1), k = 16*(q>>1) + 2t + (j&1) + 8*((j>>1)&1)
|
|
// B byte b -> column g, k = 16*(b>>3) + 2t + (b&1) + 8*((b>>1)&1)
|
|
// C/D reg 0 -> row g, columns 2t, 2t+1; reg 1 -> row g+8, same columns
|
|
// with g = laneid >> 2 and t = laneid & 3. A lane therefore holds 8 of the 32 k's of
|
|
// each of two rows, and 8 of the 32 k's of one column of the B fragment; the four
|
|
// lanes t = 0..3 hold the four disjoint eighths, so the two butterfly steps bring the
|
|
// other three in. Because both operands' k index depends on t in the same way, one
|
|
// butterfly index pairs an A byte with the B byte that sits at the same k.
|
|
//
|
|
// The 16 and the 8 are not interchangeable. ZLUDA splits NVIDIA's m16n8k32 e4m3 mma
|
|
// into two m16n8k16 f16 mmas (`e4m3_a_k16_fragment`, zluda_ptx_impl.cpp:1390): the
|
|
// first takes A registers {x, y} and B register x, the second takes {z, w} and B
|
|
// register y. So register q covers the k range 16*(q>>1), and inside a register the
|
|
// byte pair j >> 1 selects the +8 block. An earlier revision of this comment had the
|
|
// two swapped, which permuted the k's of both operands.
|
|
//
|
|
// That split is also why the rounding is not a single one at the end. The two halves
|
|
// are separate mma instructions, so `pack_f16_accumulator` rounds to f16 after the
|
|
// first half and again after the second:
|
|
//
|
|
// t = f16(dot(A[0..1], B[0..3]) + C)
|
|
// D = f16(dot(A[2..3], B[4..7]) + t)
|
|
//
|
|
// Each dot runs in f32 (ZLUDA's `dot_product<float, f16x2>`), and every e4m3 product
|
|
// is exact in f32, so a plain f32 accumulation of one half's k's reproduces it.
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Phase 1 (file lines 1658-3334, 240 mma).
|
|
//
|
|
// The 240 mma are 24 groups of ten. Group g uses A fragments 5*(g/8) + i and B
|
|
// fragments 10*(g%8) + 2i + c, i = 0..4 (the k-slices, K = 160) and c in {0,1} (the
|
|
// two n-tiles of the group's B block), chained through C: the first step of each chain
|
|
// is seeded with bias[g%8][c], every later step with the previous step's f16 result.
|
|
// So the phase is 48 output tiles of 16x8 (3 m-blocks x 16 n-tiles) over K = 160.
|
|
//
|
|
// The weight image holds phase 1's B block at [0, 20480) relative to W + z*20480:
|
|
// forty 512-byte tiles, tile t = 5*(g%8) + i covering sixteen columns (n = 0..7 for
|
|
// c = 0, n = 8..15 for c = 1) and k = 0..31, one lane's 16 bytes at +16*laneid. That
|
|
// is the `ld.weak.global.ca.v4.u32 {..}, [rd9 + 512*t]` chain of the prologue with
|
|
// rd9 = W + z*20480 + 16*laneid (file line 1471).
|
|
|
|
// one output tile: the two f16x2 pairs the mma writes, [m-block][n-tile][row half]
|
|
DEVF void phase1(const uint32_t as[60], const uint8_t* w, const uint32_t seeds[16],
|
|
int z, uint32_t lane, uint32_t d[3][16][2])
|
|
{
|
|
#pragma unroll
|
|
for (int nb = 0; nb < 8; ++nb) {
|
|
#pragma unroll
|
|
for (int c = 0; c < 2; ++c) {
|
|
const uint32_t seed = seeds[2 * nb + c];
|
|
#pragma unroll
|
|
for (int mt = 0; mt < 3; ++mt) {
|
|
// C of the first k-step: the seed u32 duplicated into both f32 halves,
|
|
// i.e. the lane's two columns of both rows start at the same f16 pair
|
|
half_t prev[2][2];
|
|
for (int col = 0; col < 2; ++col) {
|
|
const half_t v = col == 0 ? half_lo(seed) : half_hi(seed);
|
|
prev[0][col] = v;
|
|
prev[1][col] = v;
|
|
}
|
|
#pragma unroll
|
|
for (int i = 0; i < 5; ++i) {
|
|
const int f = 10 * nb + 2 * i + c; // B fragment index
|
|
const uint32_t* a = as + 4 * (5 * mt + i);
|
|
// tile f>>1, half f&1 (bytes 8..15 of the lane's 16), at +16*laneid
|
|
const uint8_t* tile = w + (uint32_t)z * 20480u + 512u * (uint32_t)(f >> 1)
|
|
+ 16u * lane + 8u * (uint32_t)(f & 1);
|
|
const uint32_t b0 = *(const uint32_t*)(tile);
|
|
const uint32_t b1 = *(const uint32_t*)(tile + 4);
|
|
// ZLUDA runs this mma as two m16n8k16 f16 mmas, so each half's
|
|
// dot product is accumulated and then rounded to f16 on its own.
|
|
// A register q covers k + 16*(q>>1) and byte j within it selects
|
|
// the +8 block, so half kg uses A registers 2*kg+0 and 2*kg+1 and
|
|
// B bytes 4*kg..4*kg+3.
|
|
float acc[2][2];
|
|
#pragma unroll
|
|
for (int col = 0; col < 2; ++col) {
|
|
acc[0][col] = (float)prev[0][col];
|
|
acc[1][col] = (float)prev[1][col];
|
|
}
|
|
#pragma unroll
|
|
for (int kg = 0; kg < 2; ++kg) {
|
|
float dot[2][2] = {{0.0f, 0.0f}, {0.0f, 0.0f}};
|
|
// butterfly index m brings in the lane whose t is t^m, which
|
|
// holds the k's 2*(t^m), 2*(t^m)+1, 8+2*(t^m), 8+2*(t^m)+1
|
|
#pragma unroll
|
|
for (int m = 0; m < 4; ++m) {
|
|
const uint32_t a0 = bfly_m(a[2 * kg + 0], m); // row g
|
|
const uint32_t a1 = bfly_m(a[2 * kg + 1], m); // row g+8
|
|
const uint32_t bb = bfly_m(kg == 0 ? b0 : b1, m);
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; ++j) {
|
|
const float bv = e4m3_to_f32((bb >> (8 * j)) & 0xffu);
|
|
dot[0][0] += e4m3_to_f32((a0 >> (8 * j)) & 0xffu) * bv;
|
|
dot[0][1] += e4m3_to_f32((a1 >> (8 * j)) & 0xffu) * bv;
|
|
}
|
|
}
|
|
// Both k16 halves accumulate into the same f32 accumulator and
|
|
// the pair rounds to f16 once. ZLUDA's scalar fallback rounds
|
|
// between the halves, but the gfx12 native-FP8 path the runtime
|
|
// actually takes does not -- it is two chained f32-accumulating
|
|
// WMMAs, exactly what the family's k32_e4m3() in tex_common.h
|
|
// reproduces.
|
|
#pragma unroll
|
|
for (int col = 0; col < 2; ++col) {
|
|
acc[0][col] += dot[0][col];
|
|
acc[1][col] += dot[1][col];
|
|
}
|
|
}
|
|
#pragma unroll
|
|
for (int col = 0; col < 2; ++col) {
|
|
prev[0][col] = (half_t)acc[0][col];
|
|
prev[1][col] = (half_t)acc[1][col];
|
|
}
|
|
}
|
|
d[mt][2 * nb + c][0] = pair_of(prev[0][0], prev[0][1]);
|
|
d[mt][2 * nb + c][1] = pair_of(prev[1][0], prev[1][1]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Ground truth for a stage comes from kernels/rr/swin_dec3_recipe.py's oracle recipe:
|
|
// with D4R_RRSWIN_DEBUG_DEC3=p1 the same splice keeps NVIDIA's whole body and stores
|
|
// that stage's registers into the plane arena at +D4R_RRSWIN_DEBUG_BASE, 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) and the same -DD4R_RRSWIN_DEBUG_BASE the oracle build used, 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
|
|
|
|
#ifndef D4R_RRSWIN_DEBUG_BASE
|
|
#define D4R_RRSWIN_DEBUG_BASE 27000000
|
|
#endif
|
|
|
|
DEVF void debug_store(const uint8_t* prm, const uint32_t* regs, int n)
|
|
{
|
|
const uint32_t cx = __builtin_amdgcn_workgroup_id_x();
|
|
const uint32_t cy = __builtin_amdgcn_workgroup_id_y();
|
|
const uint32_t cz = __builtin_amdgcn_workitem_id_z();
|
|
// the oracle's slots cover ctaid.x < 8, ctaid.y < 2 and all four waves of the block
|
|
if (cx >= 8u || cy >= 2u || cz >= 4u) {
|
|
return;
|
|
}
|
|
const uint64_t arena = *(const uint64_t*)(prm + kPlaneArg);
|
|
// the oracle's slot, swin_dec3_recipe.py's 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 shift)
|
|
const uint64_t want = (uint64_t)n * 128ull > 8192ull ? (uint64_t)n * 128ull : 8192ull;
|
|
uint64_t slot = 1;
|
|
while (slot < want) {
|
|
slot <<= 1;
|
|
}
|
|
const uint64_t block = ((uint64_t)cy * 8ull + cx) * 4ull + cz;
|
|
// swin_dec3_recipe.py's DEBUG_BASE, which reads D4R_RRSWIN_DEBUG_BASE and defaults
|
|
// to the tail described there; pass the same value to both builds or the two sides
|
|
// store to different places. dec3's plane arena is the tail of the 65,945,600-byte
|
|
// scratch block all eleven denoiser layers share: arg048 is at block offset
|
|
// 35,799,040, so the arena is [35,799,040, 65,945,600) and the family default of
|
|
// 46,000,000 falls outside the block entirely. A payload must also clear the other
|
|
// layers' sub-buffers (the replay manifests list block offsets 0, 30146560,
|
|
// 35799040, 36387840, 36535040, 37683200, 39567360, 45219840 and 60293120) as well
|
|
// as dec3's own inputs at 37683200 and 39567360.
|
|
const uint64_t debug_base = D4R_RRSWIN_DEBUG_BASE;
|
|
uint32_t* dst = (uint32_t*)(arena + debug_base + block * slot
|
|
+ (uint64_t)lane_id() * (uint64_t)n * 4ull);
|
|
for (int i = 0; i < n; ++i) {
|
|
dst[i] = regs[i];
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Mirror cut. The fifteen 16-row A fragments arrive in the order the mma reads them
|
|
// (fragment f is the four consecutive registers %r<base>+0..3 of A_BASES in
|
|
// swin_dec3_recipe.py), then phase 1's sixteen bias seeds (%r2453 %r2463 ... %r3163,
|
|
// read from W + z*256 + ((laneid<<2)&12) + 81920 + 16*i at file lines 1633-1648), the
|
|
// weight image and the parameter block.
|
|
//
|
|
// The four accumulators phase 48's epilogue reads cross the cut through the shared
|
|
// scratch, not `call.uni`'s return parameters: `scratch` is the base of the entry's
|
|
// only shared array (12800 bytes, dead across the cut -- E48 re-materialises it), and
|
|
// the four values sit 16 bytes apart per lane.
|
|
//
|
|
// Address space 3 (the shared scratch) explicitly: a generic store lowers to flat,
|
|
// which faults.
|
|
extern "C" __attribute__((device)) void d4r_swin_dec3(
|
|
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 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[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[16] = {c0, c1, c2, c3, c4, c5, c6, c7,
|
|
c8, c9, c10, c11, c12, c13, c14, c15};
|
|
// The epilogue's four accumulators: %r26071, %r26072, %r26081 and %r26082 -- phase
|
|
// 48's two mma D pairs. Address space 3 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) + 4 * (uint32_t)lane_id();
|
|
const uint32_t lane = lane_id();
|
|
const int z = (int)__builtin_amdgcn_workitem_id_z();
|
|
|
|
// ---- phase 1 -----------------------------------------------------------
|
|
uint32_t d1[3][16][2]; // [m-block][n-tile][row half]: f16 pair
|
|
phase1(as1, w, seeds, z, lane, d1);
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 1
|
|
{
|
|
// the oracle's stage p1: the 48 chain finals in mma order (group major, the two
|
|
// n-tiles of the group minor), each the (row g, row g+8) pair the mma writes
|
|
uint32_t regs[96];
|
|
int i = 0;
|
|
for (int g = 0; g < 24; ++g) {
|
|
for (int c = 0; c < 2; ++c) {
|
|
regs[i++] = d1[g / 8][2 * (g % 8) + c][0];
|
|
regs[i++] = d1[g / 8][2 * (g % 8) + c][1];
|
|
}
|
|
}
|
|
debug_store(prm, regs, 96);
|
|
}
|
|
#endif
|
|
#if D4R_RRSWIN_DEBUG_STAGE == 100
|
|
{
|
|
// the oracle's stage a: the A fragments and C seeds as the mma reads them
|
|
uint32_t regs[76];
|
|
for (int i = 0; i < 60; ++i) {
|
|
regs[i] = as1[i];
|
|
}
|
|
for (int i = 0; i < 16; ++i) {
|
|
regs[60 + i] = seeds[i];
|
|
}
|
|
debug_store(prm, regs, 76);
|
|
}
|
|
#endif
|
|
|
|
// ---- handing the accumulators back ------------------------------------
|
|
// %r26071/%r26072 and %r26081/%r26082 are phase 48's two mma D pairs (E48 stores
|
|
// them into the shared array at (48*y + 6*x) and softmaxes five taps of it).
|
|
// Phase 48 is not implemented yet; until it is, fill the four slots with
|
|
// lane-distinct values so the read-back path is observable without wild magnitudes.
|
|
for (int i = 0; i < 4; ++i) {
|
|
out[i] = (uint32_t)lane + 32u * (uint32_t)i;
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Weight image layout of phase 1 (byte offsets from the PTX's own loads; the phase
|
|
// reads [0, 20480) of W + z*20480, forty 512-byte tiles):
|
|
// tile t = 5*(g%8) + i, half 0 = n 0..7 (c = 0), half 1 = n 8..15 (c = 1), k = 0..31
|
|
// byte j of a lane's 16-byte tile -> (k, n) at
|
|
// 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)
|
|
// with base = W + z*20480 + 512*t + 16*laneid + 8*(f&1). phase1() reads that layout
|
|
// directly through the lane's own two registers instead of recomputing the address.
|
|
//
|
|
// Fragment maps (fixed by the mma instruction, see the comment above phase 1):
|
|
// A reg r, byte j -> row g + 8*(r&1), k = 8*(r>>1) + 2t + (j&1) + 16*((j>>1)&1)
|
|
// B byte b -> column g, k = 8*(b>>3) + 2t + (b&1) + 16*((b>>1)&1)
|
|
// C/D reg 0 -> row g, cols 2t, 2t+1; reg 1 -> row g+8, same cols
|
|
// A plain f32 dot product over K plus C, rounded to f16 once, reproduces the
|
|
// translated kernel's output exactly, so the phases need no operand-layout gymnastics
|
|
// and no WMMA intrinsics: every GEMM here is exact in f32 before the f16 rounding. |