Token-lane block for the RDNA4 native-FP8 layers: enc1, dec1 and the enc3 tube
swin_tok.h passes the weights as the first WMMA operand and the activations as the second, so each GEMM's encoded result is already the next one's operand and stays in registers (swin_block.h transposes every result through LDS). One wave owns one 4x4 window for the whole block: norms, attention and the MLP need no LDS and no barriers; only the patch expand and the patch merge go through an e4m3 image in LDS. Outputs are bit-identical to swin_block.h (swapped WMMA operand roles, verified on an RX 9070 XT and on recorded launches). mix2 rounds in place (accumulator is addend and destination, no copy per pair). -DSWIN_NO_TOKEN_LANES builds these layers from swin_block.h again; enc2 and dec2 are unchanged. RX 9070 XT accuracy set, 2560x1440 render: whole evaluation 8.56 -> 7.18 ms.
This commit is contained in:
woutermans committed
1 parent
7c60d53cd0
commit
a7d565d1cd
7 files changed
+730
No files matched your search
@@ -67,6 +67,25 @@ The gfx12 M variant can use native e4m3 FP8 WMMA with `D4R_NATIVE_FP8=1`. Its ac
|
||||
|
||||
In that build the GEMM activations live in LDS as e4m3 bytes (`SWIN_A8` in `swin_common.h`). Every value is encoded once where it is produced, four at a time with `v_cvt_pk_fp8_f32`, and an operand load is one 8-byte LDS read; the earlier build kept f16 values and encoded them again at every operand load (up to 16 times per value in the MLP). The instruction rounds like NVIDIA's conversion but encodes overflow as `0x7f`, so values are clamped to ±448 first with the NaN-propagating packed minimum and maximum (`codes4`). It equals `half_to_e4m3` for every non-NaN f16 input on an RX 9070 XT; a NaN encodes as the NaN code `0xff`, where the reference keeps its sign (`0x7f` or `0xff`). The accuracy build of these layers rounds each k32 step with `v_fma_mix` (`K32_MIX`; `-DK32_NO_MIX` restores the convert-and-round form) and its stored codes still equal `swin_model.py`. The patch merge reads the stored output codes from a byte image. `-DSWIN_NO_A8` restores the previous operand path. Builds without `D4R_FP8_WMMA` (RDNA3, and RDNA4 with `NativeFp8` off) compile to the same code as before.
|
||||
|
||||
**Token-lane block (`kernels/m/swin_tok.h`, enc1, dec1 and the enc3 tube on the native-FP8 path).** `swin_block.h` computes every GEMM as activations × weights, so a result has lane = channel and VGPR = token and goes through LDS (byte stores and a barrier) before it can be the next GEMM's operand. These three layers pass the weights as the first WMMA operand and the activations as the second. The result is the transpose, lane = token with a half-wave's 8 rows in its VGPRs, which is the layout of a gfx12 operand: the encoded result of one GEMM is the next one's operand and stays in registers. One wave owns one 4×4 window for the whole block, so norms, attention and the MLP use no LDS and no barriers; only the patch expand and the patch merge cross windows, through a byte image. The weight image keeps its slot format with the rows permuted per GEMM (`TokDesc`), and the per-channel vectors come from a second prepared image (`TokConst`).
|
||||
|
||||
- The products and the K slot of every product are unchanged, and swapping the operand roles of the FP8 and the f16 WMMA is bit-identical on an RX 9070 XT (1M random elements each). Both sets' outputs are byte-identical to the previous build on recorded launches (2560×1440 and 1280×720 render, shifted and unshifted tube launches) and over a 16-frame moving harness sequence.
|
||||
- The input is decoded with `v_cvt_pk_f32_fp8`. It equals `e4m3_to_half` for every code except the NaN codes `0x7f` / `0xff` (NaN instead of ±480); the FP8 WMMAs already read those bytes that way.
|
||||
- The MLP clamp uses the NaN-propagating minimum / maximum; it differs from minnum / maxnum only for a NaN input, whose product is NaN either way.
|
||||
- `mix2` (all exact FP8 layers, enc2 and dec2 included) rounds in place: the accumulator is both addend and destination, which removes a register copy per pair.
|
||||
- `-DSWIN_NO_TOKEN_LANES` builds these layers from `swin_block.h` again. enc2 and dec2 still use `swin_block.h`.
|
||||
- Tried, no measurable gain: interleaving the WMMAs of independent k32 steps. Not tried: skipping windows that lie entirely outside the image, true16 code generation (the half-register conversions would save about 3% of the instructions).
|
||||
|
||||
RX 9070 XT, accuracy set, D3D12 harness at 3840×2160 output, median GPU ms per frame (`D4R_CUDA_KERNEL_PROFILE`, frames 81–200), previous build → this one:
|
||||
|
||||
| Render | enc1 | dec1 | enc3 tube (6 launches) | whole evaluation |
|
||||
|---|---|---|---|---|
|
||||
| 1280×720 | 0.36 → 0.25 | 0.36 → 0.25 | 0.56 → 0.49 | 3.47 → 3.15 |
|
||||
| 1920×1080 | 0.71 → 0.50 | 0.74 → 0.50 | 0.80 → 0.64 | 5.15 → 4.54 |
|
||||
| 2227×1253 | 0.98 → 0.70 | 1.08 → 0.71 | 1.07 → 0.84 | 6.70 → 5.82 |
|
||||
| 2560×1440 | 1.42 → 0.89 | 1.49 → 0.91 | 1.34 → 1.05 | 8.56 → 7.18 |
|
||||
|
||||
|
||||
RX 9070 XT (gfx1201), 1280×720 → 3840×2160, D3D12 harness, DLSS 310.7: the five Swin layers went from 3.99 ms to 1.37 ms per frame in replay (enc1 0.97 → 0.28, enc2 0.42 → 0.17, enc3 6 × 0.22 → 6 × 0.08, dec2 0.41 → 0.15, dec1 0.89 → 0.31), and the whole evaluation from 5.9 ms to 2.8 ms of GPU time (2.4 ms when the GPU is not left idle between frames) ([performance.md](performance.md#rdna4)). The layer outputs are byte-identical to the previous build on recorded launches, and the harness output is byte-identical over a 12-frame temporal sequence. The accuracy build's stored codes equal `swin_model.py` on sampled blocks of all five layers. The complete accuracy set costs 3.06 ms per evaluation in the same harness (2.59 ms under load) against 2.77 ms (2.38 ms) for the fast set; with the previous FP8 layers it cost 6.49 ms (6.32 ms). On gfx12 the accuracy builds of M's post and downsample kernels are wave64 with denormal handling kept; the image was identical to the wave32 accuracy set over an 8-frame temporal sequence.
|
||||
|
||||
The gfx11 refactor and gfx12-layout shim matched the recorded gfx11 replays and 88-frame K/M harness images byte for byte on the RX 7700 XT. rocjitsu gfx1201 replay comparisons put the native K layers within 96–109 dB PSNR of gfx11, and M differed in 0.001–0.3% of e4m3 bytes. These are emulator results, not RDNA4 hardware results. rocjitsu's FP8 arithmetic is not hardware-checked; bit-exactness, speed and driver stability on an RDNA4 GPU remain unverified.
|
||||
|
||||
@@ -9,6 +9,10 @@
|
||||
#ifndef SWIN_MLP_R
|
||||
#define SWIN_MLP_R(NW) 2
|
||||
#endif
|
||||
// RDNA4 native FP8 builds use the token-lane block (swin_tok.h); -DSWIN_NO_TOKEN_LANES keeps swin_block.h
|
||||
#ifndef SWIN_NO_TOKEN_LANES
|
||||
#define SWIN_TOKEN_LANES
|
||||
#endif
|
||||
#include "swin_block.h"
|
||||
#ifndef SWIN_NW
|
||||
#define SWIN_NW 8
|
||||
|
||||
@@ -8,6 +8,10 @@
|
||||
#define SWIN_MLP_R(NW) 1
|
||||
#endif
|
||||
#define SWIN_SCALAR_BLOAD // scalar-base weight loads: faster here (no spills), slower in the 8-wave layers
|
||||
// RDNA4 native FP8 builds use the token-lane block (swin_tok.h); -DSWIN_NO_TOKEN_LANES keeps swin_block.h
|
||||
#ifndef SWIN_NO_TOKEN_LANES
|
||||
#define SWIN_TOKEN_LANES
|
||||
#endif
|
||||
#include "swin_block.h"
|
||||
#ifndef SWIN_NW
|
||||
#define SWIN_NW 4
|
||||
|
||||
@@ -11,6 +11,10 @@
|
||||
// recreation. Address-only caching then serves stale packed weights. Prepare
|
||||
// each launch; the default single image is sufficient on the launch stream.
|
||||
#define SWIN_PREP_KEY_AT 0
|
||||
// RDNA4 native FP8 builds use the token-lane block (swin_tok.h); -DSWIN_NO_TOKEN_LANES keeps swin_block.h
|
||||
#ifndef SWIN_NO_TOKEN_LANES
|
||||
#define SWIN_TOKEN_LANES
|
||||
#endif
|
||||
#include "swin_block.h"
|
||||
#ifndef SWIN_NW
|
||||
#define SWIN_NW 8
|
||||
|
||||
@@ -1124,3 +1124,11 @@ template <int S> __device__ __forceinline__ int prep_slot_claim(uint64_t* keys,
|
||||
__syncthreads(); \
|
||||
} \
|
||||
}
|
||||
|
||||
// RDNA4 native FP8: a module that defines SWIN_TOKEN_LANES is built from the token-lane block instead
|
||||
// (swin_tok.h: one wave per window, transposed GEMMs, results stay in registers)
|
||||
#if defined(SWIN_A8) && defined(SWIN_TOKEN_LANES)
|
||||
#include "swin_tok.h"
|
||||
#undef SWIN_MODULE_W
|
||||
#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN) SWIN_TOK_MODULE(NAME, C, NH, NPM, TUBE, PARAMS, CIN)
|
||||
#endif
|
||||
@@ -441,10 +441,18 @@ __device__ __forceinline__ T16 t16_splat(half_t h)
|
||||
// acc = f16(P + acc), rounded once (v_fma_mix: verified == f16(exact sum) on 8.4M cases)
|
||||
__device__ __forceinline__ hv2 mix2(float p0, float p1, hv2 c)
|
||||
{
|
||||
#ifdef MIX_COPY
|
||||
hv2 r = c;
|
||||
asm(MIX_PAD "v_fma_mixlo_f16 %0, %1, 1.0, %2 op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p0), "v"(c));
|
||||
asm("v_fma_mixhi_f16 %0, %1, 1.0, %2 op_sel:[0,0,1] op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p1), "v"(c));
|
||||
return r;
|
||||
#else
|
||||
// in place: mixlo writes only the low half and mixhi reads only the high half of the addend, so the
|
||||
// accumulator is both the addend and the destination (no copy of it per pair)
|
||||
asm(MIX_PAD "v_fma_mixlo_f16 %0, %1, 1.0, %0 op_sel_hi:[0,0,1]" : "+v"(c) : "v"(p0));
|
||||
asm("v_fma_mixhi_f16 %0, %1, 1.0, %0 op_sel:[0,0,1] op_sel_hi:[0,0,1]" : "+v"(c) : "v"(p1));
|
||||
return c;
|
||||
#endif
|
||||
}
|
||||
__device__ __forceinline__ void mixacc(T16& acc, f8v P)
|
||||
{
|
||||
|
||||
@@ -0,0 +1,683 @@
|
||||
// Token-lane rrlite Swin block for RDNA4 native FP8 (SWIN_A8), selected with SWIN_TOKEN_LANES.
|
||||
//
|
||||
// swin_block.h computes every GEMM as activations x weights: the result has lane = channel column and
|
||||
// VGPR = token row, so each result is transposed through LDS (byte stores, barriers) before it can be the
|
||||
// next GEMM's operand. Here the weights are the first WMMA operand and the activations the second, which
|
||||
// gives the transposed result: lane = token, VGPR i = row 8 half + i of the output. A gfx12 operand is this
|
||||
// half's 8 K values, so an encoded result IS the next GEMM's operand and stays in registers. The products
|
||||
// and the K slot of every product are the same as in swin_block.h, so the values are the same (operand
|
||||
// roles swapped: bit-identical on an RX 9070 XT for the FP8 and the f16 WMMA).
|
||||
//
|
||||
// One wave owns one 4x4 window (16 tokens) for the whole block: norms, attention and the MLP need no LDS
|
||||
// and no barriers. Only the patch expand (decoders) and the patch merge (encoders) cross windows and go
|
||||
// through an e4m3 byte image in LDS.
|
||||
//
|
||||
// A half-wave owns the channels whose bit 2 equals the half (dword 4 t + 2 g + half of tile t, g = 0, 1):
|
||||
// both halves of a k32 operand's K slots are then two of the lane's own code dwords.
|
||||
#pragma once
|
||||
|
||||
enum : int { TOK_ID, TOK_PI, TOK_X, TOK_NU, TOK_PE };
|
||||
// A GEMM's weight tiles as in GemmDesc (slot dst + ((kc * 2 + s) * NT + nt) * 16 + r holds the 16 K slots of
|
||||
// row r of tile nt), with the raw weight column of (nt, r) chosen by perm:
|
||||
// TOK_ID 16 nt + r (Q: the attention's K order)
|
||||
// TOK_PI gperm(kslot(nt, r)) (V, fc1: row r is K slot r of step nt of the GEMM that consumes it)
|
||||
// TOK_X gperm(tok_xnat(nt, r)) (output projection, fc2: the x layout)
|
||||
// TOK_NU gperm(16 nt + r) (patch merge: natural channel order, 8 contiguous output bytes)
|
||||
// TOK_PE col0 + gperm(16 nt + r) over weight blocks of blkcols columns / blkbytes bytes (patch expand)
|
||||
struct TokDesc
|
||||
{
|
||||
int dst, base, Ks, Ns, KC, NT, perm, col0, blkcols, blkbytes;
|
||||
};
|
||||
enum : int { TC_X, TC_PI, TC_NU, TC_ATT };
|
||||
// Per-channel f16 vectors as 8-value slots: slot dst + 2 j + half holds rows 8 half + i of tile j in the
|
||||
// layout of kind (TC_ATT: j = query token, row = key token, the transposed attention bias)
|
||||
struct TokConst
|
||||
{
|
||||
int dst, kind, base;
|
||||
};
|
||||
|
||||
// x layout: natural channel of row r (= 8 half + i) of tile t
|
||||
__device__ __forceinline__ int tok_xnat(int t, int r)
|
||||
{
|
||||
return 16 * t + 8 * ((r >> 2) & 1) + 4 * (r >> 3) + (r & 3);
|
||||
}
|
||||
|
||||
template <int C, int NH, int NPM, int CIN> struct TokLayout
|
||||
{
|
||||
using L = SwinLayout<C, NH, NPM, CIN>;
|
||||
static constexpr int NT = C / 16, NC = C / 8, NTP = NPM / 16;
|
||||
// weight image (slots)
|
||||
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
|
||||
static constexpr int WO = 4 * C * NH;
|
||||
static constexpr int W1 = WO + 2 * NH * C;
|
||||
__host__ __device__ static constexpr int w1(int c) { return W1 + 4 * C * c; }
|
||||
__host__ __device__ static constexpr int w2(int c) { return W1 + 4 * C * c + 2 * C; }
|
||||
static constexpr int PM = W1 + 4 * C * NC;
|
||||
static constexpr int PE = PM + NPM * C / 4;
|
||||
__host__ __device__ static constexpr int pe(int q) { return PE + q * (CIN * C / 16); }
|
||||
static constexpr int TOTAL = PE + CIN * C / 4;
|
||||
static constexpr int NDESC = 2 * NH + 1 + 2 * NC + (NPM ? 1 : 0) + (CIN ? 4 : 0);
|
||||
// constant image (slots of 8 f16)
|
||||
static constexpr int CBO = 0, CG1 = 2 * NT, CG2 = 4 * NT, CB2 = 6 * NT, CB1 = 8 * NT;
|
||||
static constexpr int CATT = CB1 + 4 * NC, CPMB = CATT + 32 * NH, CPEB = CPMB + 2 * NTP;
|
||||
static constexpr int CTOTAL = CPEB + (CIN ? 8 * NT : 0);
|
||||
static constexpr int NCONST = 4 + NC + NH + (NPM ? 1 : 0) + (CIN ? 4 : 0);
|
||||
};
|
||||
|
||||
template <int C, int NH, int NPM, int CIN> struct TokTables
|
||||
{
|
||||
TokDesc d[TokLayout<C, NH, NPM, CIN>::NDESC];
|
||||
TokConst c[TokLayout<C, NH, NPM, CIN>::NCONST];
|
||||
};
|
||||
|
||||
template <int C, int NH, int NPM, int CIN> constexpr TokTables<C, NH, NPM, CIN> make_tok_tables()
|
||||
{
|
||||
using L = SwinLayout<C, NH, NPM, CIN>;
|
||||
using K = TokLayout<C, NH, NPM, CIN>;
|
||||
TokTables<C, NH, NPM, CIN> r{};
|
||||
int n = 0;
|
||||
for (int h = 0; h < NH; ++h)
|
||||
{
|
||||
r.d[n++] = TokDesc{K::qv(h, 0), L::head(h), 1024, 512, C / 32, 2, TOK_ID, 0, 0, 0};
|
||||
r.d[n++] = TokDesc{K::qv(h, 1), L::head(h) + 32 * C, 1024, 512, C / 32, 2, TOK_PI, 0, 0, 0};
|
||||
}
|
||||
r.d[n++] = TokDesc{K::WO, L::head(0) + 64 * C + 512, L::HS, 512, NH, C / 16, TOK_X, 0, 0, 0};
|
||||
for (int c = 0; c < K::NC; ++c)
|
||||
{
|
||||
r.d[n++] = TokDesc{K::w1(c), L::b1(c), 512, 16 * C, C / 32, 2, TOK_PI, 0, 0, 0};
|
||||
r.d[n++] = TokDesc{K::w2(c), L::b1(c) + 32 * C + 64, 0, 512, 1, C / 16, TOK_X, 0, 0, 0};
|
||||
}
|
||||
if (NPM)
|
||||
r.d[n++] = TokDesc{K::PM, L::PM0, 512, 64 * C, C / 8, NPM / 16, TOK_NU, 0, 0, 0};
|
||||
for (int q = 0; q < (CIN ? 4 : 0); ++q)
|
||||
r.d[n++] = TokDesc{K::pe(q), 0, 512, 16 * CIN, CIN / 32, C / 16, TOK_PE, q * C, 4 * C / NH, CIN * 4 * C / NH};
|
||||
int m = 0;
|
||||
r.c[m++] = TokConst{K::CBO, TC_X, L::BO};
|
||||
r.c[m++] = TokConst{K::CG1, TC_X, L::G1};
|
||||
r.c[m++] = TokConst{K::CG2, TC_X, L::G2};
|
||||
r.c[m++] = TokConst{K::CB2, TC_X, L::B2};
|
||||
for (int c = 0; c < K::NC; ++c)
|
||||
r.c[m++] = TokConst{K::CB1 + 4 * c, TC_PI, L::b1(c) + 32 * C};
|
||||
for (int h = 0; h < NH; ++h)
|
||||
r.c[m++] = TokConst{K::CATT + 32 * h, TC_ATT, L::head(h) + 64 * C};
|
||||
if (NPM)
|
||||
r.c[m++] = TokConst{K::CPMB, TC_NU, L::PMB};
|
||||
for (int q = 0; q < (CIN ? 4 : 0); ++q)
|
||||
r.c[m++] = TokConst{K::CPEB + 2 * K::NT * q, TC_NU, 4 * C * CIN + 2 * q * C};
|
||||
return r;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void expand_tok_weights(const uint8_t* w, wslot* w16, const TokDesc* descs, int ndesc, int idx)
|
||||
{
|
||||
int d = 0;
|
||||
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
|
||||
++d;
|
||||
const TokDesc g = descs[d];
|
||||
const int local = idx - g.dst;
|
||||
if (local >= g.KC * 2 * g.NT * 16)
|
||||
return;
|
||||
const int r = local & 15, rest = local >> 4;
|
||||
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
|
||||
int n = 16 * nt + r, base = g.base;
|
||||
if (g.perm == TOK_PI)
|
||||
n = gperm(kslot(nt, r));
|
||||
else if (g.perm == TOK_X)
|
||||
n = gperm(tok_xnat(nt, r));
|
||||
else if (g.perm == TOK_NU || g.perm == TOK_PE)
|
||||
n = gperm(16 * nt + r);
|
||||
if (g.perm == TOK_PE)
|
||||
{
|
||||
const int gc = g.col0 + n;
|
||||
base += (gc / g.blkcols) * g.blkbytes;
|
||||
n = gc % g.blkcols;
|
||||
}
|
||||
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(base, g.Ks, g.Ns, 32 * kc + kslot(s, 4 * j + b), n)] << (8 * b);
|
||||
v[j] = word;
|
||||
}
|
||||
w16[idx] = v;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void expand_tok_consts(const uint8_t* w, h8* c16, const TokConst* cs, int ncs, int idx)
|
||||
{
|
||||
int d = 0;
|
||||
while (d + 1 < ncs && idx >= cs[d + 1].dst)
|
||||
++d;
|
||||
const TokConst g = cs[d];
|
||||
const int local = idx - g.dst, hf = local & 1, j = local >> 1;
|
||||
h8 v;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int r = 8 * hf + i;
|
||||
int off;
|
||||
if (g.kind == TC_X)
|
||||
off = 2 * gperm(tok_xnat(j, r));
|
||||
else if (g.kind == TC_PI)
|
||||
off = 2 * gperm(kslot(j, r));
|
||||
else if (g.kind == TC_NU)
|
||||
off = 2 * gperm(16 * j + r);
|
||||
else
|
||||
{
|
||||
// the attention bias of (query j, key r)
|
||||
const int bl = 4 * (j & 7) + ((r & 7) >> 1), word = (j >> 3) + 2 * (r >> 3);
|
||||
off = 2 * (8 * bl + 2 * word + (r & 1));
|
||||
}
|
||||
v[i] = hload(w, g.base + off);
|
||||
}
|
||||
c16[idx] = v;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- register helpers
|
||||
typedef float tok_f2v __attribute__((ext_vector_type(2)));
|
||||
|
||||
// two f32 -> f16 (round to nearest even) packed into one register
|
||||
__device__ __forceinline__ hv2 tok_pk16(float a, float b)
|
||||
{
|
||||
return (hv2){(half_t)a, (half_t)b};
|
||||
}
|
||||
|
||||
// two e4m3 codes of a word (bytes 0, 1 or 2, 3) as f16. Equal to e4m3_to_half for every code except the NaN
|
||||
// codes 0x7f / 0xff, which the conversion decodes as NaN (e4m3_to_half: +-480). The FP8 WMMAs read the
|
||||
// same bytes with the hardware's decoding, and no encoder here stores those codes for a non-NaN value.
|
||||
template <bool UPPER> __device__ __forceinline__ hv2 tok_dec2(uint32_t word)
|
||||
{
|
||||
const tok_f2v f = __builtin_amdgcn_cvt_pk_f32_fp8((int)word, UPPER);
|
||||
return tok_pk16(f[0], f[1]);
|
||||
}
|
||||
|
||||
// the same lane of the other half-wave
|
||||
__device__ __forceinline__ hv2 tok_other(hv2 v)
|
||||
{
|
||||
return __builtin_bit_cast(hv2, xor_lane<16>(__builtin_bit_cast(uint32_t, v)));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ T16 tok_t16(h8 c)
|
||||
{
|
||||
T16 t;
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
t.set_pair(k, (hv2){c[2 * k], c[2 * k + 1]});
|
||||
return t;
|
||||
}
|
||||
|
||||
// e4m3 codes of a tile's 8 rows: this half's 8 K slots of an operand
|
||||
__device__ __forceinline__ gop tok_codes(const T16& t)
|
||||
{
|
||||
return (u2v){codes4(t.pair(0), t.pair(1)), codes4(t.pair(2), t.pair(3))};
|
||||
}
|
||||
|
||||
// 1 / sqrt(sum of squares + eps) of a token's C values, with the summation tree of norm1 and norm2 of
|
||||
// swin_block.h (both are this tree in the x layout): per (g, element) the squares of a dword's two pairs,
|
||||
// the 32-channel groups, the two 16-channel halves of a group, then the half-waves, g and the elements
|
||||
template <int NT> __device__ __forceinline__ half_t tok_rsq(const hv2 (&v)[NT][4])
|
||||
{
|
||||
constexpr int NPL = NT / 2;
|
||||
hv2 u[2];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
{
|
||||
hv2 ws[2];
|
||||
#pragma unroll
|
||||
for (int ww = 0; ww < 2; ++ww)
|
||||
{
|
||||
hv2 pp[NPL];
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NPL; ++pl)
|
||||
{
|
||||
const hv2 a = v[2 * pl + ww][2 * g], b = v[2 * pl + ww][2 * g + 1];
|
||||
pp[pl] = (hv2)(a * a) + (hv2)(b * b);
|
||||
}
|
||||
if constexpr (NPL == 4)
|
||||
ws[ww] = (pp[0] + pp[1]) + (pp[2] + pp[3]);
|
||||
else if constexpr (NPL == 3)
|
||||
ws[ww] = (pp[0] + pp[1]) + pp[2];
|
||||
else
|
||||
ws[ww] = pp[0] + pp[1];
|
||||
}
|
||||
const hv2 a = ws[0] + ws[1];
|
||||
u[g] = a + tok_other(a);
|
||||
}
|
||||
const hv2 s = u[0] + u[1];
|
||||
return rsqrt16(s[0] + s[1]);
|
||||
}
|
||||
|
||||
template <int C, int NH, int NPM, bool TUBE, int CIN>
|
||||
__device__ __forceinline__ void swin_tok_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const h8* c16,
|
||||
const int bx, const int by, const int gx)
|
||||
{
|
||||
using K = TokLayout<C, NH, NPM, CIN>;
|
||||
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8;
|
||||
static_assert(C % 32 == 0, "32-channel groups");
|
||||
|
||||
const int lane = threadIdx.x, wv = threadIdx.z;
|
||||
const int l16 = lane & 15, hf = lane >> 4;
|
||||
const int T = 16 * wv + l16; // this lane's token: window wv, position l16
|
||||
// The lane's bases into the two images, hidden from constant folding: every tile and constant is then
|
||||
// one load at a fixed offset from a register (folded, each load rebuilt the image's address first).
|
||||
uintptr_t chv = (uintptr_t)(c16 + hf), wlv = (uintptr_t)((const gop*)w16 + 2 * l16 + hf);
|
||||
asm("" : "+v"(chv));
|
||||
asm("" : "+v"(wlv));
|
||||
const h8* const ch = (const h8*)chv;
|
||||
// row l16 of a weight tile: this half's 8 K slots
|
||||
const gop* const wl = (const gop*)wlv;
|
||||
auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
|
||||
|
||||
if constexpr (TUBE)
|
||||
{
|
||||
if (tp->in_flags != nullptr)
|
||||
{
|
||||
if (lane <= 3 && wv == 0)
|
||||
{
|
||||
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
|
||||
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
|
||||
{
|
||||
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
|
||||
// bounded spin: a missing producer must never hang the GPU
|
||||
for (int it = 0; it < (1 << 22); ++it)
|
||||
{
|
||||
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
|
||||
break;
|
||||
__builtin_amdgcn_s_sleep(1);
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ block input x0 (encoders: the fp8 input; decoders: skip)
|
||||
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
|
||||
const int Xm = mirror(X, p.tw), Ym = mirror(Y, p.th);
|
||||
uint32_t inw[NT][2];
|
||||
{
|
||||
const uint8_t* const src = CIN ? p.p32 : p.in;
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NKC; ++pl)
|
||||
{
|
||||
const uint8_t* b = src + ((size_t)(pl * p.th + Ym) * p.tw + Xm) * 32 + 4 * hf;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
inw[2 * pl + (j >> 1)][j & 1] = input_word((const uint32_t*)(b + 8 * j));
|
||||
}
|
||||
}
|
||||
|
||||
hv2 v[NT][4]; // x0, then x1, as f16 pairs: [tile][2 g + pair of the dword]
|
||||
if constexpr (CIN > 0)
|
||||
{
|
||||
// patch expand: lanes = the block's 16 low-res tokens, wave = quadrant; the codes of q8(expand) go to
|
||||
// their child tokens' rows of a byte image, where the windows' waves pick them up
|
||||
__shared__ __attribute__((aligned(16))) uint8_t E8[64 * C];
|
||||
const int W2 = p.tw / 2, H2 = p.th / 2;
|
||||
const int MX = mirror((8 * bx - p.sx) / 2 + (l16 & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (l16 >> 2), H2);
|
||||
gop pa[CIN / 32][2];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < CIN / 32; ++kc)
|
||||
{
|
||||
const uint32_t* b = (const uint32_t*)(p.in + ((size_t)(kc * H2 + MY) * W2 + MX) * 32 + 4 * hf);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
pa[kc][s] = (u2v){input_word(b + 2 * s), input_word(b + 4 + 2 * s)};
|
||||
}
|
||||
const int tx = 2 * (l16 & 3) + (wv & 1), ty = 2 * (l16 >> 2) + (wv >> 1);
|
||||
const int Tc = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
|
||||
const int pew = K::pe(0) + wv * (CIN * C / 16);
|
||||
#pragma unroll
|
||||
for (int ct = 0; ct < NT; ++ct)
|
||||
{
|
||||
T16 e = tok_t16(ch[K::CPEB + 2 * (NT * wv + ct)]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < CIN / 32; ++kc)
|
||||
k32(e, wt(pew, NT, kc, 0, ct), pa[kc][0], wt(pew, NT, kc, 1, ct), pa[kc][1]);
|
||||
*(u2v*)(E8 + Tc * C + 16 * ct + 8 * hf) = tok_codes(e);
|
||||
}
|
||||
__syncthreads();
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
{
|
||||
const uint32_t ew = *(const uint32_t*)(E8 + T * C + 16 * t + 8 * g + 4 * hf);
|
||||
v[t][2 * g] = tok_dec2<false>(ew) + tok_dec2<false>(inw[t][g]);
|
||||
v[t][2 * g + 1] = tok_dec2<true>(ew) + tok_dec2<true>(inw[t][g]);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
{
|
||||
v[t][2 * g] = tok_dec2<false>(inw[t][g]);
|
||||
v[t][2 * g + 1] = tok_dec2<true>(inw[t][g]);
|
||||
}
|
||||
}
|
||||
|
||||
// residual of the output projection: f16(x0 + b_o)
|
||||
hv2 res[NT][4];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
const h8 bo = ch[K::CBO + 2 * t];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
res[t][k] = v[t][k] + (hv2){bo[2 * k], bo[2 * k + 1]};
|
||||
}
|
||||
|
||||
// ------------------------------------------------ norm1: h1 codes, the K operands of Q and V
|
||||
gop hop[NKC][2];
|
||||
{
|
||||
const half_t rn = tok_rsq<NT>(v);
|
||||
const hv2 rn2 = {rn, rn};
|
||||
uint32_t c1[NT][2];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
const h8 g1 = ch[K::CG1 + 2 * t];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
c1[t][g] = codes4(v[t][2 * g] * (hv2)(rn2 * (hv2){g1[4 * g], g1[4 * g + 1]}),
|
||||
v[t][2 * g + 1] * (hv2)(rn2 * (hv2){g1[4 * g + 2], g1[4 * g + 3]}));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
hop[kc][s] = (u2v){c1[2 * kc][s], c1[2 * kc + 1][s]};
|
||||
}
|
||||
|
||||
// ------------------------------------------------ attention per head: O codes, the K operands of the projection
|
||||
gop oc[NH][2];
|
||||
#pragma unroll
|
||||
for (int hd = 0; hd < NH; ++hd)
|
||||
{
|
||||
// Q transposed (lane = token, rows = Q's 32 columns), V as in swin_block.h (lane = column, rows = tokens)
|
||||
T16 qt[2], vt[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
qt[nt] = t16_splat((half_t)0.0f);
|
||||
vt[nt] = t16_splat((half_t)0.0f);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
k32(qt[nt], wt(K::qv(hd, 0), 2, kc, 0, nt), hop[kc][0], wt(K::qv(hd, 0), 2, kc, 1, nt), hop[kc][1]);
|
||||
k32(vt[nt], hop[kc][0], wt(K::qv(hd, 1), 2, kc, 0, nt), hop[kc][1], wt(K::qv(hd, 1), 2, kc, 1, nt));
|
||||
}
|
||||
const gop q0 = tok_codes(qt[0]), q1 = tok_codes(qt[1]);
|
||||
// S transposed: lane = query token, rows = key tokens
|
||||
T16 S = tok_t16(ch[K::CATT + 32 * hd + 2 * l16]);
|
||||
k32(S, q0, q0, q1, q1);
|
||||
const half_t cs = f16(0.0972222164273262f / 5.656854152679443f);
|
||||
const half_t cl = (half_t)0.55615234375f;
|
||||
const half_t p1 = f16(0.92730712890625f), p0 = (half_t)1.375f;
|
||||
hv2 wg[4], v1[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
hv2 tv = S.pair(k) * (hv2){cs, cs};
|
||||
tv = __builtin_elementwise_max(__builtin_elementwise_min(tv, (hv2){cl, cl}), (hv2){-cl, -cl});
|
||||
const hv2 e1 = __builtin_elementwise_fma(tv, -tv, (hv2){p1, p1});
|
||||
const hv2 poly = __builtin_elementwise_fma(tv, e1, (hv2){p0, p0});
|
||||
// the pair is the (even, odd) key word of the exponent trick
|
||||
wg[k] = __builtin_bit_cast(hv2, (__builtin_bit_cast(uint32_t, poly) << 5) + 0x7FF88000u);
|
||||
v1[k] = wg[k] + tok_other(wg[k]); // keys c and c ^ 8
|
||||
}
|
||||
const hv2 v3 = (v1[0] + v1[1]) + (v1[2] + v1[3]); // c ^ 2, then c ^ 4
|
||||
const half_t sum = v3[0] + v3[1]; // c ^ 1
|
||||
const half_t rinv = f16(1.0f / (float)sum);
|
||||
hv2 pw[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
pw[k] = wg[k] * (hv2){rinv, rinv};
|
||||
const sop pb = __builtin_bit_cast(sop, pw);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
hv2 vp[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
vp[k] = vt[nt].pair(k);
|
||||
const f8v o = wmma(__builtin_bit_cast(sop, vp), pb, splat(0.0f));
|
||||
oc[hd][nt] = (u2v){codes4(tok_pk16(o[0], o[1]), tok_pk16(o[2], o[3])), codes4(tok_pk16(o[4], o[5]), tok_pk16(o[6], o[7]))};
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ output projection + residual
|
||||
T16 x[NT];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
x[t].set_pair(k, res[t][k]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NH; ++kc)
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
k32(x[t], wt(K::WO, NT, kc, 0, t), oc[kc][0], wt(K::WO, NT, kc, 1, t), oc[kc][1]);
|
||||
|
||||
// ------------------------------------------------ norm2: h2 codes, the K operands of fc1
|
||||
gop fop[NKC][2];
|
||||
{
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
v[t][k] = x[t].pair(k);
|
||||
const half_t rn = tok_rsq<NT>(v);
|
||||
const hv2 rn2 = {rn, rn};
|
||||
uint32_t c2[NT][2];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
const h8 g2 = ch[K::CG2 + 2 * t], b2 = ch[K::CB2 + 2 * t];
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
c2[t][g] = codes4(v[t][2 * g] * (hv2)(rn2 * (hv2){g2[4 * g], g2[4 * g + 1]}),
|
||||
v[t][2 * g + 1] * (hv2)(rn2 * (hv2){g2[4 * g + 2], g2[4 * g + 3]}));
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
x[t].set_pair(k, v[t][k] + (hv2){b2[2 * k], b2[2 * k + 1]});
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
fop[kc][s] = (u2v){c2[2 * kc][s], c2[2 * kc + 1][s]};
|
||||
}
|
||||
|
||||
// ------------------------------------------------ MLP: fc1 transposed, its codes are fc2's operands
|
||||
#pragma unroll 1
|
||||
for (int c = 0; c < NC; ++c)
|
||||
{
|
||||
const int w1 = K::W1 + 4 * C * c;
|
||||
// the round's operands first: 2 bias slots, fc1's and fc2's weight tiles
|
||||
h8 zb[2];
|
||||
gop wa[NKC][2][2], wb[2][NT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
zb[nt] = ch[K::CB1 + 4 * c + 2 * nt];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
wa[kc][s][nt] = wt(w1, 2, kc, s, nt);
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
wb[s][t] = wt(w1 + 2 * C, NT, 0, s, t);
|
||||
#ifndef SWIN_TOK_NO_LOAD_SYNC
|
||||
// one wait for the whole batch here instead of one before each WMMA
|
||||
asm("" : "+v"(wb[1][NT - 1]));
|
||||
#endif
|
||||
T16 z[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
z[nt] = tok_t16(zb[nt]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < NKC; ++kc)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
k32(z[nt], wa[kc][0][nt], fop[kc][0], wa[kc][1][nt], fop[kc][1]);
|
||||
gop gz[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
hv2 gl[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
const hv2 zv = z[nt].pair(k);
|
||||
// minimum / maximum instead of swin_block.h's minnum / maxnum: they differ only for a NaN z, and
|
||||
// then the product with z below is NaN either way (no canonicalising instruction per pair)
|
||||
const hv2 cz = __builtin_elementwise_maximum(__builtin_elementwise_minimum(zv, (hv2){2.0f16, 2.0f16}),
|
||||
(hv2){-2.0f16, -2.0f16});
|
||||
const hv2 az = __builtin_bit_cast(hv2, (u16x2)(__builtin_bit_cast(u16x2, cz) & (unsigned short)0x7fff));
|
||||
const hv2 inner = (hv2){f16(0.41216981f), f16(0.41216981f)} -
|
||||
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
|
||||
gl[k] = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
|
||||
}
|
||||
gz[nt] = (u2v){codes4(gl[0], gl[1]), codes4(gl[2], gl[3])};
|
||||
}
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
k32(x[t], wb[0][t], gz[0], wb[1][t], gz[1]);
|
||||
}
|
||||
|
||||
// ------------------------------------------------ output codes
|
||||
uint32_t co[NT][2];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
co[t][g] = codes4(x[t].pair(2 * g), x[t].pair(2 * g + 1));
|
||||
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < NKC; ++pl)
|
||||
{
|
||||
uint32_t* b = (uint32_t*)(p.out + ((size_t)(pl * p.th + Y) * p.tw + X) * 32 + 4 * hf);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
b[2 * j] = co[2 * pl + (j >> 1)][j & 1];
|
||||
}
|
||||
}
|
||||
|
||||
// ------------------------------------------------ patch merge: 2x2 tokens (4C) -> NPM, half resolution
|
||||
if constexpr (NPM > 0)
|
||||
{
|
||||
// lanes = the block's 16 merged tokens; the merge reads the output codes from a byte image
|
||||
constexpr int XB = C + 4, NTP = NPM / 16;
|
||||
__shared__ __attribute__((aligned(16))) uint8_t U[64 * XB];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
*(uint32_t*)(U + T * XB + 16 * t + 8 * g + 4 * hf) = co[t][g];
|
||||
__syncthreads();
|
||||
if (wv < NTP)
|
||||
{
|
||||
const int mx = l16 & 3, my = l16 >> 2;
|
||||
// K slot 8 half + 4 rr + b of step (kc, s) is value k = 32 kc + 16 rr + 8 s + 4 half + b of the merged
|
||||
// token's 4C inputs: channel k % C of its source token k / C. The four source tokens are consecutive
|
||||
// in x and 4 apart in y, and 4 half never crosses a token, so every run of four codes is at a fixed
|
||||
// offset from the lane's first source token.
|
||||
const uint8_t* const ub = U + ((2 * mx & 3) + 4 * (2 * my & 3) + 16 * (mx >> 1) + 32 * (my >> 1)) * XB + 4 * hf;
|
||||
gop as[C / 8][2];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 8; ++kc)
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 2; ++s)
|
||||
{
|
||||
uint32_t run4[2];
|
||||
#pragma unroll
|
||||
for (int rr = 0; rr < 2; ++rr)
|
||||
{
|
||||
const int k0 = 32 * kc + 16 * rr + 8 * s, q = k0 / C;
|
||||
run4[rr] = *(const uint32_t*)(ub + ((q & 1) + 4 * (q >> 1)) * XB + k0 % C);
|
||||
}
|
||||
as[kc][s] = (u2v){run4[0], run4[1]};
|
||||
}
|
||||
const int W2 = p.tw / 2, H2 = p.th / 2;
|
||||
const int MX = (8 * bx - p.sx) / 2 + mx, MY = (8 * by - p.sy) / 2 + my;
|
||||
const bool inside = MX >= 0 && MX < W2 && MY >= 0 && MY < H2;
|
||||
// output tiles u = wv, wv + 4, ...: rows = 16 natural channels, 8 contiguous bytes per half-wave
|
||||
#pragma unroll
|
||||
for (int j = 0; j < (NTP + 3) / 4; ++j)
|
||||
{
|
||||
const int u = wv + 4 * j;
|
||||
if (u < NTP)
|
||||
{
|
||||
T16 y = tok_t16(ch[K::CPMB + 2 * u]);
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 8; ++kc)
|
||||
k32(y, wt(K::PM, NTP, kc, 0, u), as[kc][0], wt(K::PM, NTP, kc, 1, u), as[kc][1]);
|
||||
const u2v code = tok_codes(y);
|
||||
if (inside)
|
||||
*(u2v*)(p.p48 + ((size_t)((u >> 1) * H2 + MY) * W2 + MX) * 32 + 16 * (u & 1) + 8 * hf) = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (TUBE)
|
||||
{
|
||||
if (tp->out_flags != nullptr)
|
||||
{
|
||||
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "agent");
|
||||
__syncthreads();
|
||||
if (lane == 0 && wv == 0)
|
||||
__hip_atomic_store(tp->out_flags + (size_t)(by * gx + bx) * 32, (uint8_t)1, __ATOMIC_RELEASE,
|
||||
__HIP_MEMORY_SCOPE_AGENT);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#if defined(SWIN_TOK_MAX_VGPR)
|
||||
#define SWIN_TOK_VGPR_ATTR __attribute__((amdgpu_num_vgpr(SWIN_TOK_MAX_VGPR)))
|
||||
#else
|
||||
#define SWIN_TOK_VGPR_ATTR
|
||||
#endif
|
||||
|
||||
// Module boilerplate: the weight and constant images, their prep kernel, and the 4-wave main kernel
|
||||
#define SWIN_TOK_MODULE(NAME, C, NH, NPM, TUBE, PARAMS, CIN) \
|
||||
using NAME##_K = TokLayout<C, NH, NPM, CIN>; \
|
||||
static_assert(SWIN_PREP_SLOTS == 0, "one prepared image"); \
|
||||
__device__ wslot g_w16[NAME##_K::TOTAL]; \
|
||||
__device__ h8 g_c16[NAME##_K::CTOTAL]; \
|
||||
__constant__ TokTables<C, NH, NPM, CIN> g_tok = make_tok_tables<C, NH, NPM, CIN>(); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = \
|
||||
((NAME##_K::TOTAL > NAME##_K::CTOTAL ? NAME##_K::TOTAL : NAME##_K::CTOTAL) + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = SWIN_PREP_KEY_AT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = 0; \
|
||||
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; \
|
||||
if (idx < NAME##_K::TOTAL) \
|
||||
expand_tok_weights(w, g_w16, g_tok.d, NAME##_K::NDESC, idx); \
|
||||
if (idx < NAME##_K::CTOTAL) \
|
||||
expand_tok_consts(w, g_c16, g_tok.c, NAME##_K::NCONST, idx); \
|
||||
} \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = 0; \
|
||||
extern "C" __global__ void __launch_bounds__(128) SWIN_TOK_VGPR_ATTR NAME(PARAMS p) \
|
||||
{ \
|
||||
const CommonParams& cp = *(const CommonParams*)&p; \
|
||||
const int gx = (cp.tw + cp.sx + 7) / 8; \
|
||||
swin_tok_block<C, NH, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, g_c16, blockIdx.x, blockIdx.y, gx); \
|
||||
}
|
||||
Reference in new issue
Block a user