wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add RDNA4 (gfx12) support and a 0.1.2 optimization round

Browse Source 
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 2026-09-29 20:36:10 -04:00
1 parent f310d52b69
commit 4334a0c3c4
42 files changed
+3068 -529
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/m/rrlite_enc1_4x4.hip
+1
View File
@@ -3,5 +3,6 @@
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 1
#define SWIN_SCALAR_BLOAD // scalar-base weight loads: faster here (no spills), slower in the 8-wave layers
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc1_4x4, 64, 2, 4, 96, false, CommonParams, 0)
kernels/m/rrlite_enc3_tube_4x4.hip
+1
View File
@@ -3,5 +3,6 @@
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 2
#define SWIN_PREP_SLOTS 8 // one weight image per tube block (six per frame)
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, 8, 0, true, TubeParams, 0)
kernels/m/swin_block.h
+144 -90
View File
@@ -19,7 +19,7 @@ template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
__host__ __device__ static constexpr int b1(int c) { return c == 0 ? M0 : M0 + 66 * C + 64 + (64 * C + 64) * (c - 1); }
static constexpr int PM0 = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1);
static constexpr int PMB = PM0 + 4 * C * NPM;
// f16 weight image (h16 units)
// weight image (slots of 16 K values, see expand_weights)
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
static constexpr int WO = 4 * C * NW;
static constexpr int W1 = WO + 2 * NW * C;
@@ -33,6 +33,14 @@ template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
static constexpr int NDESC = 2 * NW + 1 + 2 * NC + NPM / 32 + (CIN ? NW : 0);
};
// Start (in halves) of token T's row in the output staging buffer: stride C + 4 plus a shift by bits 4 and 5
// of T. The patch merge's lanes read 8 bytes each from 16 tokens that differ in bits 1, 3, 4 and 5 (plus a
// constant); this placement puts those 16 reads on 16 distinct 8-byte bank pairs for C = 64, 96 and 128.
template <int C> __host__ __device__ constexpr int xo_row(int T)
{
return (C + 4) * T + 4 * ((T >> 4) & 1) + 16 * (T >> 5);
}
template <int C, int NW, int NPM, int CIN> struct SwinDescs
{
GemmDesc d[SwinLayout<C, NW, NPM, CIN>::NDESC];
@@ -61,7 +69,7 @@ template <int C, int NW, int NPM, int CIN> constexpr SwinDescs<C, NW, NPM, CIN>
// NH heads (NVIDIA's warps per block) run on NW waves (NW = NH, or 2 NH to halve LDS per wave)
template <int C, int NH, int NW, int NPM, bool TUBE, int CIN>
__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const h16* w16, const int bx,
__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const int bx,
const int by, const int gx)
{
using L = SwinLayout<C, NH, NPM, CIN>;
@@ -77,8 +85,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
static_assert(NTW * NW == NT * MG, "column split");
constexpr int PAST = CIN + 8;
constexpr int SCR = NW * 16 * QST * 2 > (CIN ? 16 * PAST * 2 : 0) ? NW * 16 * QST * 2 : 16 * PAST * 2;
constexpr int U_ST4 = R * 64 * GST * 2, U_OB = NPM ? 64 * C : 0;
constexpr int USIZE = SCR > U_ST4 ? (SCR > U_OB ? SCR : U_OB) : (U_ST4 > U_OB ? U_ST4 : U_OB);
constexpr int U_ST4 = R * 64 * GST * 2;
constexpr int USIZE = SCR > U_ST4 ? SCR : U_ST4;
__shared__ __attribute__((aligned(16))) half_t A[64 * AST];
__shared__ __attribute__((aligned(16))) uint8_t U[USIZE];
@@ -118,11 +126,11 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
// block input words (encoders: fp8 input, decoders: fp8 skip), issued first so their latency
// overlaps the patch expand; word k of thread tid is idx = tid + 32 NW k
constexpr int NSW = CIN ? 1 : 64 * C / 4 / (32 * NW);
static_assert(CIN || NSW * 32 * NW == 64 * C / 4, "input words per thread");
constexpr int NSW = 64 * C / 4 / (32 * NW);
static_assert(NSW * 32 * NW == 64 * C / 4, "input words per thread");
uint32_t inw[NSW];
#pragma unroll
for (int k = 0; k < (CIN ? 0 : NSW); ++k)
for (int k = 0; k < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
@@ -130,26 +138,24 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const uint8_t* src = CIN ? p.p32 : p.in;
inw[k] = *(const uint32_t*)(src + ((size_t)((c4 >> 5) * p.th + Y) * p.tw + X) * 32 + (c4 & 31));
}
// the input words as f16 into X0 (natural channel order); decoders add the patch expand to them later
auto stage_input_words = [&]() {
#pragma unroll
for (int k = 0; k < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
const uint32_t bytes = inw[k];
#pragma unroll
for (int b = 0; b < 4; ++b)
X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
};
// ------------------------------------------------ stage 0: block input x0 (f16) into A
// encoders: the fp8 input; decoders: q8(patch expand(low-res)) + skip
if constexpr (CIN > 0)
{
constexpr int NTEW_ = 4 * C / NW / 16;
uint32_t sk[NTEW_][8];
#pragma unroll
for (int j = 0; j < NTEW_; ++j)
{
const int Gc = 16 * (NTEW_ * wv + j) + l16, q = Gc / C, ch = ginv(Gc % C);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 2 * i + hi;
const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
const int X = mirror(8 * bx - p.sx + tx, p.tw), Y = mirror(8 * by - p.sy + ty, p.th);
sk[j][i] = p.p32[((size_t)((ch >> 5) * p.th + Y) * p.tw + X) * 32 + (ch & 31)];
}
}
const int W2 = p.tw / 2, H2 = p.th / 2;
for (int idx = tid; idx < 16 * CIN / 4; idx += 32 * NW)
{
@@ -161,6 +167,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int b = 0; b < 4; ++b)
PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
stage_input_words(); // the skip input
__syncthreads();
constexpr int NTEW = 4 * C / NW / 16; // expand tiles per wave (L::NTE per weight block)
T16 e[NTEW];
@@ -170,7 +177,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
{
const h16 a0 = lds16(PA + l16 * PAST + 32 * kc), a1 = lds16(PA + l16 * PAST + 32 * kc + 16);
const gop a0 = lda(PA + l16 * PAST + 32 * kc), a1 = lda(PA + l16 * PAST + 32 * kc + 16);
#pragma unroll
for (int j = 0; j < NTEW; ++j)
{
@@ -185,28 +192,16 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 2 * i + hi;
const int r = mrow(i);
const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + e4m3_to_half(sk[j][i]));
// e4m3_to_half is exact in f16, so adding the staged skip value is the same sum
X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + X0[T * AST + ch]);
}
}
}
else
{
#pragma unroll
for (int k = 0; k < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
const uint32_t bytes = inw[k];
#pragma unroll
for (int b = 0; b < 4; ++b)
{
X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
}
}
stage_input_words();
__syncthreads();
// residual accumulators of the output projection: f16(x0 + b_o) in the C-fragment layout
@@ -223,8 +218,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const int ch = ginv(n);
#pragma unroll
for (int i = 0; i < 8; i += 2)
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + 2 * i + hi) * AST + ch],
X0[(16 * (mg * MT + mm) + 2 * i + 2 + hi) * AST + ch]} + (hv2){bo, bo});
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + mrow(i)) * AST + ch],
X0[(16 * (mg * MT + mm) + mrow(i + 1)) * AST + ch]} + (hv2){bo, bo});
}
__syncthreads(); // stage 1 overwrites x0 rows with h1
@@ -308,7 +303,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < C / 32; ++kc)
{
h16 b[2][4];
gop b[2][4];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -317,8 +312,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int m = 0; m < MW; ++m)
{
const h16 a0 = lds16(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
const h16 a1 = lds16(A + (16 * (m0 + m) + l16) * AST + 32 * kc + 16);
const gop a0 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
const gop a1 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc + 16);
#pragma unroll
for (int q = 0; q < 4; ++q)
k32(acc[m][q], a0, b[0][q], a1, b[1][q]);
@@ -344,7 +339,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int ri = 2 * i + hi, cj = l16;
const int ri = mrow(i), cj = l16;
const int bl = 4 * (ri & 7) + ((cj & 7) >> 1), word = (ri >> 3) + 2 * (cj >> 3);
bias.set(i, hload(W, bbase + 2 * (8 * bl + 2 * word + (cj & 1))));
}
@@ -359,14 +354,14 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
Qs[(2 * i + hi) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
Qs[(2 * i + 2 + hi) * QST + 16 * q + l16] = qp[m][q][i / 2][1];
Qs[mrow(i) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
Qs[mrow(i + 1) * QST + 16 * q + l16] = qp[m][q][i / 2][1];
}
wave_sync();
const h16 a0 = lds16(Qs + l16 * QST), a1 = lds16(Qs + l16 * QST + 16);
const gop a0 = lda(Qs + l16 * QST), a1 = lda(Qs + l16 * QST + 16);
T16 S = bias;
k32(S, a0, a0, a1, a1);
// rows 2k+hi and 2k+2+hi of this lane's column, packed: scale, clamp and the cubic
// rows mrow(2k), mrow(2k+1) of this lane's column, packed: scale, clamp and the cubic
// polynomial as f16x2, the exponent trick per row on the (even, odd) column word
hv2 wg[4];
const bool oddc = l16 & 1;
@@ -398,15 +393,19 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const hv2 sum = v3 + xr(v3, std::integral_constant<int, 1>{});
const hv2 rinv = {f16(1.0f / (float)sum[0]), f16(1.0f / (float)sum[1])};
const hv2 pw = wg[k] * rinv;
Qs[(4 * k + hi) * PST + l16] = pw[0];
Qs[(4 * k + 2 + hi) * PST + l16] = pw[1];
Qs[mrow(2 * k) * PST + l16] = pw[0];
Qs[mrow(2 * k + 1) * PST + l16] = pw[1];
}
wave_sync();
const h16 pa = lds16(Qs + l16 * PST);
const sop pa = lds16(Qs + l16 * PST);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
// V as B operand: lane holds column l16 over the 16 tokens (rows); own rows 2i+hi
// V as B operand: lane holds column l16 over the 16 tokens (rows, the K of P V)
#if D4R_WMMA_LAYOUT == 12
// own pair k holds rows (2k + 8hi, 2k + 1 + 8hi): this half's K values in order
const sop vb = __builtin_bit_cast(sop, vp[m][nt]);
#else
// own pair k holds rows (4k+hi, 4k+2+hi); the other half-wave has (4k+1-hi, 4k+3-hi)
uint32_t vw[8];
const uint32_t s_lo = hi ? 0x01000504u : 0x05040100u, s_hi = hi ? 0x03020706u : 0x07060302u;
@@ -419,14 +418,15 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
vw[2 * k + 1] = __builtin_amdgcn_perm(other, own, s_hi); // rows 4k+2, 4k+3
}
const h16 vb = __builtin_bit_cast(h16, vw);
#endif
const f8v o = wmma(pa, vb, splat(0.0f));
const int ocol = apos(32 * h + ginv(16 * nt + l16));
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
const hv2 oq = q8x2((hv2){(half_t)o[i], (half_t)o[i + 1]});
A[(16 * (m0 + m) + 2 * i + hi) * AST + ocol] = oq[0];
A[(16 * (m0 + m) + 2 * i + 2 + hi) * AST + ocol] = oq[1];
A[(16 * (m0 + m) + mrow(i)) * AST + ocol] = oq[0];
A[(16 * (m0 + m) + mrow(i + 1)) * AST + ocol] = oq[1];
}
}
wave_sync();
@@ -438,7 +438,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < NH; ++kc)
{
h16 b[2][NTW];
gop b[2][NTW];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -447,8 +447,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
{
const h16 a0 = lds16(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
const h16 a1 = lds16(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc + 16);
const gop a0 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
const gop a1 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc + 16);
#pragma unroll
for (int j = 0; j < NTW; ++j)
k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
@@ -463,7 +463,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int j = 0; j < NTW; ++j)
#pragma unroll
for (int i = 0; i < 8; ++i)
A[(16 * (mg * MT + mm) + 2 * i + hi) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
A[(16 * (mg * MT + mm) + mrow(i)) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
__syncthreads();
if (tid < 64)
{
@@ -514,12 +514,12 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
const int row = 16 * (mg * MT + mm) + 2 * i + hi;
const int row = 16 * (mg * MT + mm) + mrow(i), row2 = 16 * (mg * MT + mm) + mrow(i + 1);
const hv2 v = x[mm][j].pair(i / 2);
const hv2 r2 = {R2[row], R2[row + 2]};
const hv2 r2 = {R2[row], R2[row2]};
const hv2 hq = q8x2(v * (hv2)(r2 * (hv2){g2, g2}));
A[row * AST + col] = hq[0];
A[(row + 2) * AST + col] = hq[1];
A[row2 * AST + col] = hq[1];
x[mm][j].set_pair(i / 2, v + (hv2){b2, b2});
}
}
@@ -546,7 +546,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < C / 32; ++kc)
{
h16 b[2][2];
gop b[2][2];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -555,8 +555,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int mm = 0; mm < MPW; ++mm)
{
const h16 a0 = lds16(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
const h16 a1 = lds16(A + (16 * (mb + mm) + l16) * AST + 32 * kc + 16);
const gop a0 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
const gop a1 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc + 16);
k32(z[mm][0], a0, b[0][0], a1, b[1][0]);
k32(z[mm][1], a0, b[0][1], a1, b[1][1]);
}
@@ -579,8 +579,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
const hv2 gl = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
const hv2 gq = q8x2(gl);
Gc[(16 * (mb + mm) + 2 * i + hi) * GST + col] = gq[0];
Gc[(16 * (mb + mm) + 2 * i + 2 + hi) * GST + col] = gq[1];
Gc[(16 * (mb + mm) + mrow(i)) * GST + col] = gq[0];
Gc[(16 * (mb + mm) + mrow(i + 1)) * GST + col] = gq[1];
}
}
__syncthreads();
@@ -588,7 +588,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int cc = 0; cc < R; ++cc)
{
const int c2 = R * r + cc;
h16 b[2][NTW];
gop b[2][NTW];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -597,8 +597,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
{
const h16 a0 = lds16(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
const h16 a1 = lds16(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST + 16);
const gop a0 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
const gop a1 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST + 16);
#pragma unroll
for (int j = 0; j < NTW; ++j)
k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
@@ -609,7 +609,10 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
}
// ------------------------------------------------ output: f16 staged in LDS, 16-value segments encoded by a loop
half_t* XO = A; // [64 tokens][C channels] natural order
// XO: [64 tokens][C channels] natural order; the row padding keeps the patch merge's 8-byte reads of 16
// tokens on distinct LDS banks (see xo_row)
half_t* XO = A;
static_assert(xo_row<C>(63) + C <= 64 * AST, "XO fits in A");
#pragma unroll
for (int j = 0; j < NTW; ++j)
{
@@ -618,26 +621,28 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int i = 0; i < 8; ++i)
XO[(16 * (mg * MT + mm) + 2 * i + hi) * C + ch] = x[mm][j].get(i);
XO[xo_row<C>(16 * (mg * MT + mm) + mrow(i)) + ch] = x[mm][j].get(i);
}
__syncthreads();
uint8_t* OB = U; // e4m3 codes [64][C] for the patch merge
#pragma unroll 1
for (int seg = tid; seg < 64 * C / 16; seg += 32 * NW)
{
const int T = seg / (C / 16), c16 = 16 * (seg % (C / 16));
// rows are 8-byte aligned (xo_row), so the segment moves as four 8-byte words of two f16 pairs
uint2* xs = (uint2*)(XO + xo_row<C>(T) + c16);
uint32_t w[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
{
w[k] = 0;
#pragma unroll
for (int b = 0; b < 4; ++b)
w[k] |= enc8(XO[T * C + c16 + 4 * k + b]) << (8 * b);
const uint2 pr = xs[k];
hv2 qa, qb;
const uint32_t ca = enc8x2(__builtin_bit_cast(hv2, pr.x), qa), cb = enc8x2(__builtin_bit_cast(hv2, pr.y), qb);
w[k] = pack_codes(ca, cb);
// the patch merge consumes decode(code) = q8_exact(x), written back over x (this thread's own values)
if constexpr (NPM > 0)
xs[k] = (uint2){__builtin_bit_cast(uint32_t, qa), __builtin_bit_cast(uint32_t, qb)};
}
const uint4 v = {w[0], w[1], w[2], w[3]};
if constexpr (NPM > 0)
*(uint4*)(OB + T * C + c16) = v;
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
*(uint4*)(p.out + ((size_t)((c16 >> 5) * p.th + Y) * p.tw + X) * 32 + (c16 & 31)) = v;
@@ -663,23 +668,28 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < C / 8; ++kc)
{
h16 as[2];
gop as[2];
#pragma unroll
for (int s = 0; s < 2; ++s)
{
h16& a = as[s];
sop a;
// the operand's K slots in runs of 4 (gfx12: this half's slots 8 hi .. 8 hi + 7)
constexpr int RUNS = sizeof(sop) / sizeof(half_t) / 4;
#pragma unroll
for (int run = 0; run < 4; ++run)
for (int rr = 0; rr < RUNS; ++rr)
{
const int run = RUNS == 4 ? rr : 2 * hi + rr;
const int k = 32 * kc + kslot(s, 4 * run);
const int q = k / C, ch = k % C;
const int tx = 2 * mx + (q & 1), ty = 2 * my + (q >> 1);
const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
const uint32_t bytes = *(const uint32_t*)(OB + T * C + ch);
typedef _Float16 h4 __attribute__((ext_vector_type(4)));
const h4 v = *(const h4*)(XO + xo_row<C>(T) + ch);
#pragma unroll
for (int b = 0; b < 4; ++b)
a[4 * run + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
a[4 * rr + b] = v[b];
}
as[s] = as_gop(a);
}
#pragma unroll
for (int j = 0; j < NTP; ++j)
@@ -700,7 +710,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 2 * i + hi;
const int r = mrow(i);
const int MX = (8 * bx - p.sx) / 2 + (r & 3), MY = (8 * by - p.sy) / 2 + (r >> 2);
if (MX < 0 || MX >= W2 || MY < 0 || MY >= H2)
continue;
@@ -734,35 +744,79 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#ifndef SWIN_PERSIST
#define SWIN_PERSIST 0 // persistent workgroups in x (0: one workgroup per original block)
#endif
// Module boilerplate: f16 weight image, prep kernel and its grid for the ZLUDA override hook.
#ifndef SWIN_PREP_SLOTS
#define SWIN_PREP_SLOTS 0 // weight images kept per weights pointer (0: one image, re-prepared when the pointer changes)
#endif
// Weight-image slots for kernels launched with several weight sets (the six enc3 tube blocks). The ZLUDA
// hook runs the prep only for a weights pointer it has not seen (d4r_prep_key_at / d4r_prep_key_slots);
// the prep claims the first free slot for that pointer, and pointers beyond S share the overflow slot S,
// which the hook re-prepares on every launch. Keys are written by one thread of the prep; any prep
// thread that reads before that write finds the same (first free) slot.
template <int S> __device__ __forceinline__ int prep_slot_find(const uint64_t* keys, uint64_t w)
{
for (int i = 0; i < S; ++i)
if (__hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) == w)
return i;
return S;
}
template <int S> __device__ __forceinline__ int prep_slot_claim(uint64_t* keys, uint64_t w, bool writer)
{
// one pass: the first slot holding w or nothing. Keys only ever go from 0 to a pointer, so every
// prep thread picks the same slot whether it reads that slot before or after the writer's store.
for (int i = 0; i < S; ++i)
{
const uint64_t k = __hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
if (k == w)
return i;
if (k == 0)
{
if (writer)
__hip_atomic_store(keys + i, w, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
return i;
}
}
return S;
}
// Module boilerplate: f16 weight image(s), prep kernel and its grid for the ZLUDA override hook.
#define SWIN_MODULE(NAME, C, NW, NPM, TUBE, PARAMS, CIN) SWIN_MODULE_W(NAME, C, NW, NW, NPM, TUBE, PARAMS, CIN)
// NH heads on NWAVES waves; d4r_block_z tells the ZLUDA hook the block z size to launch with
#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN) \
using NAME##_L = SwinLayout<C, NH, NPM, CIN>; \
__device__ h16 g_w16[NAME##_L::TOTAL]; \
constexpr int NAME##_S = SWIN_PREP_SLOTS; \
__device__ wslot g_w16[(NAME##_S + 1) * NAME##_L::TOTAL]; \
__device__ uint64_t g_prep_keys[NAME##_S + 1]; \
__constant__ SwinDescs<C, NH, NPM, CIN> g_descs = make_descs<C, NH, NPM, CIN>(); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = NWAVES; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::TOTAL + 127) / 128; \
/* the prep reads only p.w (offset 0): skip it while the image for that pointer is resident */ \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = NAME##_S; \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PARAMS p) \
{ \
const int idx = blockIdx.x * 128 + threadIdx.x; \
const uint8_t* w = ((const CommonParams*)&p)->w; \
const int slot = NAME##_S ? prep_slot_claim<NAME##_S>(g_prep_keys, (uint64_t)w, idx == 0) : 0; \
if (idx < NAME##_L::TOTAL) \
expand_weights(((const CommonParams*)&p)->w, g_w16, g_descs.d, NAME##_L::NDESC, idx); \
expand_weights(w, g_w16 + slot * NAME##_L::TOTAL, g_descs.d, NAME##_L::NDESC, idx); \
} \
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = SWIN_PERSIST; \
extern "C" __global__ void __launch_bounds__(32 * NWAVES) SWIN_VGPR_ATTR NAME(PARAMS p) \
{ \
const CommonParams& cp = *(const CommonParams*)&p; \
const int slot = NAME##_S ? prep_slot_find<NAME##_S>(g_prep_keys, (uint64_t)cp.w) : 0; \
const wslot* w16 = g_w16 + slot * NAME##_L::TOTAL; \
/* the original grid: 8x8-token blocks covering the shifted token grid */ \
const int gx = (cp.tw + cp.sx + 7) / 8, gy = (cp.th + cp.sy + 7) / 8; \
if (SWIN_PERSIST == 0) \
{ \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, blockIdx.x, blockIdx.y, gx); \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, blockIdx.x, blockIdx.y, gx); \
return; \
} \
_Pragma("unroll 1") for (int b = blockIdx.x; b < gx * gy; b += gridDim.x) \
{ \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, b % gx, b / gx, gx); \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, b % gx, b / gx, gx); \
__syncthreads(); \
} \
}
kernels/m/swin_common.h
+204 -26
View File
@@ -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;
}