Port K kernel tuning and M gfx1101 flags

Preset K, byte-identical to the previous output in all captured buffers
and the harness image, measured on an RX 7700 XT:

- LDS row padding so WMMA operand loads hit distinct banks

- WGP mode for enc1/enc2/dec1/dec2, ds_bpermute half-wave exchange for
  enc0/dec0/dec3, more waves per window in enc4/dec4/dec5

- cheaper operand plumbing in common/wmma_layout.h (packed f16
  conversions, no dynamic vector indexing), raw buffer weight loads,
  accumulator-init vectors prepared once

- enc1 as a single-buffer recompute layer, dec1 Q/K/V chains split

- dec5 as four dependent phase kernels
This commit is contained in:
ayo committed 2026-10-05 14:48:37 -04:00
1 parent 8ce8618181
commit 22026ede12
45 files changed
+4514 -354

No files matched your search

+97 -3
View File
@@ -56,7 +56,21 @@ WM_FN uint32_t wm_half()
// value of the lane with the same index in the other half of the wave
WM_FN uint32_t wm_other_half(uint32_t v)
{
#ifdef WM_HALF_SWIZZLE
// Bit-mode AND=31, OR=0, XOR=16: the same half-wave partner, no address VGPR.
return (uint32_t)__builtin_amdgcn_ds_swizzle((int)v, 0x401f);
#elif defined(WM_HALF_BPERMUTE)
// through the LDS crossbar (no VALU issue; the same value)
return (uint32_t)__builtin_amdgcn_ds_bpermute((int)((__builtin_amdgcn_mbcnt_lo(~0u, 0u) ^ 16u) << 2), (int)v);
#elif defined(WM_PERMLANE_ASM)
// the builtin ties the destination to an "old" input (unused here: all lanes active, every select valid), which
// costs a v_mov per exchange; a write-only asm output avoids it. Same instruction, same result.
uint32_t r;
__asm__("v_permlanex16_b32 %0, %1, %2, %3" : "=&v"(r) : "v"(v), "s"(0x76543210u), "s"(0xfedcba98u));
return r;
#else
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
#endif
}
WM_FN float wm_other_half_f(float v)
{
@@ -134,8 +148,9 @@ WM_FN wm_f16 wm_op_acc_elem(const wm_op& v, int i)
#if D4R_WMMA_LAYOUT == 12
return (i & 1) ? wm_hi16(v[i >> 1]) : wm_lo16(v[i >> 1]);
#else
const int c = 2 * i + (int)wm_half();
return (c & 1) ? wm_hi16(v[c >> 1]) : wm_lo16(v[c >> 1]);
// element 2 i + h is half h of dword i: one shift by 16 h (written as v[(2 i + h) >> 1] the compiler could not
// see that the dword index is the constant i and indexed the operand dynamically, ~9 VALU per element)
return __builtin_bit_cast(wm_f16, (uint16_t)(v[i] >> (wm_half() << 4)));
#endif
}
@@ -146,6 +161,24 @@ WM_FN wm_op wm_op_from_acc(const wm_f16 own[8])
#if D4R_WMMA_LAYOUT == 12
return (wm_u4v){wm_pack2(own[0], own[1]), wm_pack2(own[2], own[3]), wm_pack2(own[4], own[5]), wm_pack2(own[6], own[7])};
#else
#ifdef WM_OP_VIA_LDS
// Through a per-wave LDS row buffer (WM_OP_VIA_LDS = waves per block, wave index = work-item z): every lane
// stores its 8 rows of its token's row, then loads the whole 16-value row. 8 VALU conversions + 10 LDS
// instructions instead of ~20 VALU (pack, exchange, perm). The same values. Rows 48 bytes apart (16 mod 32).
// A wave's LDS instructions complete in order, so the loads see all lanes' stores and the next call's stores
// cannot overtake these loads.
__shared__ __attribute__((aligned(16))) wm_f16 wm_xbuf[WM_OP_VIA_LDS][16][24];
{
const uint32_t ln = __builtin_amdgcn_mbcnt_lo(~0u, 0u);
wm_f16* row = wm_xbuf[__builtin_amdgcn_workitem_id_z()][ln & 15];
const uint32_t h = ln >> 4;
#pragma unroll
for (int i = 0; i < 8; ++i)
row[2 * i + h] = own[i];
const wm_u4v a = *(const wm_u4v*)row, b = *(const wm_u4v*)(row + 8);
return (wm_u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
}
#endif
const uint32_t hf = wm_half();
uint32_t mine[4], theirs[4];
#pragma unroll
@@ -177,6 +210,66 @@ WM_FN wm_op wm_op_from_acc(const wm_f16 own[8])
#endif
}
// two f32 -> f16 conversions (round to nearest even) packed into one register (low = a, high = b). WM_CVT_ASM: the
// same v_cvt_f16_f32 instructions as inline asm writing %0.l then %0.h; the compiler otherwise re-packs the halves
// with v_pack_b32_f16 copies
WM_FN uint32_t wm_cvt2(float a, float b)
{
#ifdef WM_CVT_ASM
uint32_t r;
asm("v_cvt_f16_f32 %0.l, %1" : "=v"(r) : "v"(a)); // r may reuse a's register (a is read first)
asm("v_cvt_f16_f32 %0.h, %1" : "+v"(r) : "v"(b)); // keeps the low half
return r;
#else
return wm_pack2((wm_f16)a, (wm_f16)b);
#endif
}
#if D4R_WMMA_LAYOUT == 12
// gfx12: an operand is this half's 8 K values in 4 dwords, accumulator VGPR i = row (K) 8 h + i: the packed pairs
// ARE the operand (no exchange, no byte permutation)
WM_FN wm_op wm_op_from_mine(const uint32_t mine[4])
{
return (wm_u4v){mine[0], mine[1], mine[2], mine[3]};
}
// the pair (element 2 j, 2 j + 1) of wm_op_acc_elem(v, .): operand dword j
WM_FN uint32_t wm_op_acc_pair(const wm_op& v, int j)
{
return v[j];
}
#else
// operand from this lane's 8 accumulator-row halves already packed in pairs (mine[j] = own[2 j], own[2 j + 1]);
// the same exchange + byte permutation as wm_op_from_acc
WM_FN wm_op wm_op_from_mine(const uint32_t mine[4])
{
const uint32_t hf = wm_half();
uint32_t theirs[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
theirs[j] = wm_other_half(mine[j]);
const uint32_t sel_lo = hf ? 0x01000504u : 0x05040100u, sel_hi = hf ? 0x03020706u : 0x07060302u;
wm_u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
r[j] = __builtin_amdgcn_perm(theirs[j >> 1], mine[j >> 1], (j & 1) ? sel_hi : sel_lo);
return r;
}
// the pair (element 2 j, 2 j + 1) of wm_op_acc_elem(v, .) as one packed register (one byte permutation)
WM_FN uint32_t wm_op_acc_pair(const wm_op& v, int j)
{
return __builtin_amdgcn_perm(v[2 * j + 1], v[2 * j], wm_half() ? 0x07060302u : 0x05040100u);
}
#endif
#if defined(WM_CVT_ASM) && (D4R_WMMA_LAYOUT == 12 || (!defined(WM_OP_VIA_LDS) && !defined(PWIN_NO_PERM)))
#define WM_HAVE_FROM_MINE 1
WM_FN wm_op wm_op_from_f8(wm_f8v d)
{
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
mine[j] = wm_cvt2(d[2 * j], d[2 * j + 1]);
return wm_op_from_mine(mine);
}
#else
WM_FN wm_op wm_op_from_f8(wm_f8v d)
{
wm_f16 h[8];
@@ -185,6 +278,7 @@ WM_FN wm_op wm_op_from_f8(wm_f8v d)
h[i] = (wm_f16)d[i];
return wm_op_from_acc(h);
}
#endif
// Stores an operand back as a row of 16 halves (K 0..15). In gfx11 both halves hold the whole row and the
// lanes of half 0 store it; in gfx12 each half stores its 8 values. `count` (8 or 16) limits the row to its
@@ -219,7 +313,7 @@ WM_FN void wm_acc_vec_load(const wm_f16* p, wm_f16 out[8])
const wm_u8v v = (wm_u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
#pragma unroll
for (int i = 0; i < 8; ++i)
out[i] = hf ? wm_hi16(v[i]) : wm_lo16(v[i]);
out[i] = __builtin_bit_cast(wm_f16, (uint16_t)(v[i] >> (hf << 4)));
#endif
}