Add RDNA4 (gfx12) support and a 0.1.2 optimization round
RDNA4 (experimental; untested on RDNA4 hardware) - kernels/common/wmma_layout.h picks the WMMA operand/accumulator layout per target; K (pwin_*) and M (swin_*) index through it, so the same sources build for gfx11 and gfx12. D4R_WMMA_LAYOUT=12 on gfx11 is a test shim that runs the gfx12 indexing on gfx11 WMMA. - M can use native e4m3 FP8 WMMA on gfx12 (D4R_NATIVE_FP8=1). Texture tails (enc0_tail, dec0_head) get gfx12 f16 and FP8 paths that match ZLUDA's gfx12 lowering. - ZLUDA patch 0007: gfx12 WMMA helpers, optional native FP8 MMA (D4R_ZLUDA_WMMA_FP8_NATIVE), a layout override for testing, and an architecture argument for d4r_emit. Texture kernels are now compiled offline per target through d4r_emit (D4R_ZLUDA_EMIT) instead of on the local GPU. - New d4r.ini option [Kernels] NativeFp8 (default on), wired through the shim and d4r_config.py. On gfx12 the bridge serves kernels/<arch>-fp8 when it is on. - The bridge picks the KFD GPU with the most SIMDs (so an iGPU is not chosen over a discrete GPU); D4R_GPU_ARCH overrides it. - kernels/build.sh and package_release.sh build gfx1100-gfx1103 and gfx1200/gfx1201 (plus -fp8 variants), and D4R_BUNDLE_TEX takes per-target folders. Performance (RDNA3; harness byte-identical unless noted) - ZLUDA patch 0006: m16n8k8 f16 MMAs on WMMA (E; not bit-exact, 52-55 dB per layer, D4R_ZLUDA_WMMA_K8=0 disables it), per-key weight-image slots for native prep kernels (d4r_prep_key_at, d4r_prep_key_slots), and wave64 builds of MMA-free kernels (D4R_ZLUDA_WAVE64). - M: packed FP8 encoding, conflict-free patch-merge reads, scalar weight loads in enc1, and weight prep once per weight set for the six enc3 tube blocks. - K: CU mode for the six non-wide pwin layers. - Texture kernels: every DLSS flag variant is built natively (24 K output, 8 M post, 4 M downsample, 4 M enc0), not only Townfall's; post and K output run as wave64. - M direct output: rrlite_downsample_kernel writes the shim's output buffer; the shim's default direct-output presets are now 11,13. Tools - dump_runner takes the prep grid from d4r_prep_blocks and refuses texture/surface dumps. - New pwin_check.py, swin_check.py, model_enc3.py and tail_check.hip. Docs: README GPU table and RDNA4 notes, building (patches 0002-0007, d4r_emit), native-kernels (RDNA4 section), performance (new numbers), packaging README. Version 0.1.2.
This commit is contained in:
ayo committed
1 parent
f310d52b69
commit
4334a0c3c4
42 files changed
+3068
-529
No files matched your search
+204
-26
@@ -1,8 +1,29 @@
|
||||
// Shared device code for native RDNA3 (gfx11 wave32 WMMA) rrlite Swin-block kernels (DLSS 310.7.0).
|
||||
// Shared device code for native rrlite Swin-block kernels (DLSS 310.7.0; gfx11 / gfx12 wave32 WMMA).
|
||||
// Semantics: ../ENC3_SPEC.md generalised (numpy reference ../swin_model.py).
|
||||
//
|
||||
// Activations are the WMMA A operand (lane = token row), weights the B operand (lane = channel column), and
|
||||
// the f32 result has lane = channel column, VGPR i = token row wm_acc_row(i) (../common/wmma_layout.h).
|
||||
// D4R_FP8_WMMA (gfx12 layout only): the GEMMs whose operands are both e4m3 values run as native FP8 WMMAs on
|
||||
// e4m3 bytes (weights prepared as bytes, activations encoded when loaded); P V stays f16.
|
||||
#pragma once
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <stdint.h>
|
||||
#include "../common/wmma_layout.h"
|
||||
#if defined(D4R_FP8_WMMA) && D4R_WMMA_LAYOUT != 12 && defined(__AMDGCN__)
|
||||
#error "D4R_FP8_WMMA needs the gfx12 WMMA layout"
|
||||
#endif
|
||||
// FP8 operands must hold e4m3 values, so the FP8 build re-quantises the intermediate activations as NVIDIA's
|
||||
// network does (the gfx11 "fast numerics" SWIN_NO_Q8 keeps them f16, which no gfx12 WMMA can multiply with
|
||||
// e4m3 weights)
|
||||
// (SWIN_FORCE_Q8 does the same on any target: the gfx11 reference of the FP8 build's numerics)
|
||||
#if (defined(D4R_FP8_WMMA) || defined(SWIN_FORCE_Q8) || defined(SWIN_EXACT)) && defined(SWIN_NO_Q8)
|
||||
#undef SWIN_NO_Q8
|
||||
#endif
|
||||
// SWIN_EXACT (tests): NVIDIA's numerics throughout, the semantics of tools/swin_model.py (e4m3 activations and
|
||||
// an f16 accumulator rounded after every k32 step)
|
||||
#if defined(SWIN_EXACT) && defined(SWIN_F32ACC)
|
||||
#undef SWIN_F32ACC
|
||||
#endif
|
||||
|
||||
#pragma clang fp contract(off)
|
||||
|
||||
@@ -82,12 +103,9 @@ typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
|
||||
typedef unsigned short u16x2 __attribute__((ext_vector_type(2)));
|
||||
typedef short i16x2 __attribute__((ext_vector_type(2)));
|
||||
|
||||
// q8 on two halves with packed 16-bit integer / f16 ops (same per-half logic as q8; no carries cross halves)
|
||||
__device__ __forceinline__ hv2 q8x2(hv2 h)
|
||||
// q8_exact on two halves with packed 16-bit integer / f16 ops (same per-half logic; no carries cross halves)
|
||||
__device__ __forceinline__ hv2 q8x2_exact(hv2 h)
|
||||
{
|
||||
#ifdef SWIN_NO_Q8
|
||||
return h;
|
||||
#endif
|
||||
const u16x2 u = __builtin_bit_cast(u16x2, h);
|
||||
const u16x2 m = u & (unsigned short)0x7fff;
|
||||
u16x2 r = (m + (unsigned short)0x3f + ((m >> 7) & (unsigned short)1)) & (unsigned short)0x7f80;
|
||||
@@ -102,6 +120,48 @@ __device__ __forceinline__ hv2 q8x2(hv2 h)
|
||||
return __builtin_bit_cast(hv2, (u16x2)(r | (u & (unsigned short)0x8000)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ hv2 q8x2(hv2 h)
|
||||
{
|
||||
#ifdef SWIN_NO_Q8
|
||||
return h;
|
||||
#endif
|
||||
return q8x2_exact(h);
|
||||
}
|
||||
|
||||
// enc8 of two halves with packed ops: returns the two e4m3 codes in the low bytes of the 16-bit
|
||||
// halves, and q = q8_exact(h) (the values the codes decode to)
|
||||
__device__ __forceinline__ uint32_t codes8x2(hv2 q);
|
||||
__device__ __forceinline__ uint32_t enc8x2(hv2 h, hv2& q)
|
||||
{
|
||||
q = q8x2_exact(h);
|
||||
return codes8x2(q);
|
||||
}
|
||||
|
||||
// the e4m3 codes of two halves that already hold e4m3 values (q8_exact results), in the low bytes of the
|
||||
// 16-bit halves
|
||||
__device__ __forceinline__ uint32_t codes8x2(hv2 q)
|
||||
{
|
||||
const u16x2 qb = __builtin_bit_cast(u16x2, q);
|
||||
const u16x2 m = qb & (unsigned short)0x7fff;
|
||||
// normals: the exponent moves from bias 15 to bias 7
|
||||
const u16x2 norm = (u16x2)(m >> 7) - (unsigned short)64;
|
||||
// subnormals are k * 2^-9 (k = 0..7): f16(k + 1024) has k in its low bits
|
||||
const hv2 mag = __builtin_bit_cast(hv2, m);
|
||||
const hv2 k = (hv2)(mag * (hv2){(half_t)512.0f, (half_t)512.0f}) + (hv2){(half_t)1024.0f, (half_t)1024.0f};
|
||||
const u16x2 sub = __builtin_bit_cast(u16x2, k) - (unsigned short)0x6400;
|
||||
const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
|
||||
u16x2 code = (sub & lt) | (norm & ~lt);
|
||||
const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
|
||||
code = ((u16x2){0x7f, 0x7f} & gt) | (code & ~gt);
|
||||
return __builtin_bit_cast(uint32_t, (u16x2)(((qb >> 8) & (unsigned short)0x80) | code));
|
||||
}
|
||||
|
||||
// four e4m3 codes (a, b from enc8x2) as the bytes of one word in element order
|
||||
__device__ __forceinline__ uint32_t pack_codes(uint32_t a, uint32_t b)
|
||||
{
|
||||
return __builtin_amdgcn_perm(b, a, 0x06040200u);
|
||||
}
|
||||
|
||||
// e4m3 code of h (RNE, satfinite) derived from q8(h): normals shift the exponent, subnormals are k * 2^-9
|
||||
__device__ __forceinline__ uint32_t enc8(half_t h)
|
||||
{
|
||||
@@ -153,9 +213,105 @@ __device__ __forceinline__ int woff(int base, int Ks, int Ns, int k, int n)
|
||||
4 * ((k & 31) >> 4) + (k & 3);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ f8v wmma(h16 a, h16 b, f8v c)
|
||||
// ---------------------------------------------------------------- WMMA operands
|
||||
// half-wave of this lane: blocks are 32 x 1 x waves, so the lane is threadIdx.x (the same value as wm_half(),
|
||||
// but from the source the kernels already use for their lane index)
|
||||
__device__ __forceinline__ int m_half()
|
||||
{
|
||||
return (int)threadIdx.x >> 4;
|
||||
}
|
||||
// token row that accumulator VGPR i holds in this lane
|
||||
__device__ __forceinline__ int mrow(int i)
|
||||
{
|
||||
return wm_acc_row_h(i, m_half());
|
||||
}
|
||||
// sop: an f16 operand (gfx11: the row's 16 K values; gfx12: this half's 8). gop: an operand of the e4m3 GEMMs
|
||||
// (k32 steps): sop, or with D4R_FP8_WMMA the 8 e4m3 bytes of this half's K values.
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
typedef h8 sop;
|
||||
#else
|
||||
typedef h16 sop;
|
||||
#endif
|
||||
typedef uint32_t u2v __attribute__((ext_vector_type(2)));
|
||||
typedef uint32_t u4v_t __attribute__((ext_vector_type(4)));
|
||||
#ifdef D4R_FP8_WMMA
|
||||
typedef u2v gop;
|
||||
typedef u4v_t wslot; // weight image slot: 16 e4m3 bytes (K slots 0..15)
|
||||
#else
|
||||
typedef sop gop;
|
||||
typedef h16 wslot; // weight image slot: 16 f16 values
|
||||
#endif
|
||||
|
||||
__device__ __forceinline__ f8v wmma(sop a, sop b, f8v c)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
return wm_mma(__builtin_bit_cast(wm_op, a), __builtin_bit_cast(wm_op, b), c);
|
||||
#else
|
||||
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
|
||||
#endif
|
||||
}
|
||||
|
||||
// f16 operand from 16 contiguous halves in LDS (gfx12: this half's 8 of them)
|
||||
__device__ __forceinline__ sop lds16(const half_t* p)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
return *(const h8*)(p + 8 * m_half());
|
||||
#else
|
||||
h8 lo = *(const h8*)p, hi = *(const h8*)(p + 8);
|
||||
return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
|
||||
#endif
|
||||
}
|
||||
|
||||
#ifdef D4R_FP8_WMMA
|
||||
// e4m3 bytes of an f16 operand whose values are e4m3 values (K order kept: byte j = element j)
|
||||
__device__ __forceinline__ gop to_fp8(sop v)
|
||||
{
|
||||
uint32_t c[4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
c[j] = codes8x2((hv2){v[2 * j], v[2 * j + 1]});
|
||||
return (u2v){__builtin_amdgcn_perm(c[1], c[0], 0x06040200u), __builtin_amdgcn_perm(c[3], c[2], 0x06040200u)};
|
||||
}
|
||||
// f32 += A B over 16 K for e4m3 operands
|
||||
__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
|
||||
{
|
||||
#if defined(__GFX12__)
|
||||
return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(wm_i2v, a), __builtin_bit_cast(wm_i2v, b), c);
|
||||
#else
|
||||
// layout shim: the e4m3 values widened to f16 (exact), then the f16 WMMA
|
||||
auto widen = [](gop q) {
|
||||
sop r;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
r[j] = e4m3_to_half((q[j >> 2] >> (8 * (j & 3))) & 0xffu);
|
||||
return r;
|
||||
};
|
||||
return wmma(widen(a), widen(b), c);
|
||||
#endif
|
||||
}
|
||||
__device__ __forceinline__ gop lda(const half_t* p)
|
||||
{
|
||||
return to_fp8(lds16(p));
|
||||
}
|
||||
#else
|
||||
__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
|
||||
{
|
||||
return wmma(a, b, c);
|
||||
}
|
||||
// GEMM A operand (activations) from LDS
|
||||
__device__ __forceinline__ gop lda(const half_t* p)
|
||||
{
|
||||
return lds16(p);
|
||||
}
|
||||
#endif
|
||||
// an f16 operand built in registers -> GEMM operand
|
||||
__device__ __forceinline__ gop as_gop(sop v)
|
||||
{
|
||||
#ifdef D4R_FP8_WMMA
|
||||
return to_fp8(v);
|
||||
#else
|
||||
return v;
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ f8v round16(f8v v)
|
||||
@@ -166,7 +322,7 @@ __device__ __forceinline__ f8v round16(f8v v)
|
||||
return v;
|
||||
}
|
||||
|
||||
// A C-fragment tile: element i is row 2i + lane/16. Exact mode stores packed f16 (NVIDIA's f16
|
||||
// A C-fragment tile: element i is row wm_acc_row(i). Exact mode stores packed f16 (NVIDIA's f16
|
||||
// accumulator, rounded after every k32 step); SWIN_F32ACC keeps f32 through whole GEMM chains
|
||||
// (more precise than NVIDIA; values are rounded to f16 where they are consumed).
|
||||
struct T16
|
||||
@@ -217,20 +373,20 @@ __device__ __forceinline__ void mixacc(T16& acc, f8v P)
|
||||
__device__ unsigned d4r_k32_bad[4 + 64 * 4];
|
||||
#endif
|
||||
// f16 accumulate of one k32 step (NVIDIA m16n8k32 semantics: products summed, one rounding with C)
|
||||
__device__ __forceinline__ void k32(T16& acc, h16 a0, h16 b0, h16 a1, h16 b1)
|
||||
__device__ __forceinline__ void k32(T16& acc, gop a0, gop b0, gop a1, gop b1)
|
||||
{
|
||||
#if defined(SWIN_F32ACC)
|
||||
acc.f = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, acc.f));
|
||||
acc.f = wmma8(a1, b1, wmma8(a0, b0, acc.f));
|
||||
#elif defined(K32_MIX)
|
||||
f8v P = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
|
||||
P = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, P);
|
||||
f8v P = wmma8(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
|
||||
P = wmma8(a1, b1, P);
|
||||
mixacc(acc, P);
|
||||
#else // exact: C added inside the f32 WMMA, f16 rounding after the k32 step
|
||||
f8v c;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
c[i] = (float)acc.get(i);
|
||||
c = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, c));
|
||||
c = wmma8(a1, b1, wmma8(a0, b0, c));
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc.set(i, (half_t)c[i]);
|
||||
@@ -273,12 +429,6 @@ __device__ __forceinline__ half_t hload(const uint8_t* w, int off)
|
||||
return __builtin_bit_cast(half_t, *(const uint16_t*)(w + off));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ h16 lds16(const half_t* p)
|
||||
{
|
||||
h8 lo = *(const h8*)p, hi = *(const h8*)(p + 8);
|
||||
return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void wave_sync()
|
||||
{
|
||||
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
|
||||
@@ -296,21 +446,36 @@ __device__ __forceinline__ int tok_y(int T)
|
||||
return ((T >> 2) & 3) + 4 * (T >> 5);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- f16 weight images
|
||||
// A GEMM's B operands expanded to f16 in WMMA lane order: h16 index
|
||||
// dst + ((kc * 2 + s) * NT + nt) * 16 + lane16, holding logical k = 32 kc + kslot(s, i), n = 16 nt + lane16.
|
||||
// ---------------------------------------------------------------- weight images
|
||||
// A GEMM's B operands in WMMA lane order: slot index dst + ((kc * 2 + s) * NT + nt) * 16 + lane16 holds the 16
|
||||
// K slots of logical k = 32 kc + kslot(s, i), n = 16 nt + lane16 (i = 0..15), as f16 values (wslot = h16) or,
|
||||
// with D4R_FP8_WMMA, as the e4m3 bytes. gfx12 lanes read the half of the slot with their 8 K slots.
|
||||
struct GemmDesc
|
||||
{
|
||||
int dst; // in h16 units
|
||||
int dst; // in slots
|
||||
int base, Ks, Ns, KC, NT;
|
||||
};
|
||||
|
||||
__device__ __forceinline__ h16 bload(const h16* w16, int gemm_dst, int NT, int kc, int s, int nt, int l16)
|
||||
// The tile index is wave-uniform (it depends only on the wave, the loop counters and constants), so it
|
||||
// is formed in scalar registers and the load uses a scalar base plus the lane's constant offset.
|
||||
__device__ __forceinline__ gop bload(const wslot* w16, int gemm_dst, int NT, int kc, int s, int nt, int l16)
|
||||
{
|
||||
return w16[gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16];
|
||||
#ifdef SWIN_SCALAR_BLOAD
|
||||
const int tile = __builtin_amdgcn_readfirstlane(gemm_dst + ((kc * 2 + s) * NT + nt) * 16);
|
||||
const wslot* t = w16 + tile;
|
||||
const int at = l16;
|
||||
#else
|
||||
const wslot* t = w16;
|
||||
const int at = gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16;
|
||||
#endif
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
return ((const gop*)t)[2 * at + m_half()];
|
||||
#else
|
||||
return t[at];
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const GemmDesc* descs, int ndesc, int idx)
|
||||
__device__ __forceinline__ void expand_weights(const uint8_t* w, wslot* w16, const GemmDesc* descs, int ndesc, int idx)
|
||||
{
|
||||
int d = 0;
|
||||
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
|
||||
@@ -322,9 +487,22 @@ __device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const
|
||||
const int l16 = local & 15, rest = local >> 4;
|
||||
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
|
||||
const int n = 16 * nt + l16;
|
||||
#ifdef D4R_FP8_WMMA
|
||||
u4v_t v;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
uint32_t word = 0;
|
||||
#pragma unroll
|
||||
for (int b = 0; b < 4; ++b)
|
||||
word |= (uint32_t)w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, 4 * j + b), n)] << (8 * b);
|
||||
v[j] = word;
|
||||
}
|
||||
#else
|
||||
h16 v;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 16; ++i)
|
||||
v[i] = e4m3_to_half(w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, i), n)]);
|
||||
#endif
|
||||
w16[idx] = v;
|
||||
}
|
||||
Reference in new issue
Block a user