kernels/tex: native token-lane parts for preset L's enc0 and dec0
L's unfolded enc0 and dec0 are still ZLUDA's compile of NVIDIA's PTX and cost 2.95 ms of the preset's 5.86 ms per frame (RX 9070 XT, 1280x720 -> 4K). The embedding, the patch expand, the Swin block (32 channels, 2 heads, at pixel resolution) and the patch merge are now native token-lane code; the texture work around them stays ZLUDA's compile. enc0l_tail.hip cuts in at the barrier after the feature staging and dec0l_block.hip before the output code that reads the block's f16 result out of the last MMAs; make_ptx.py splices either in when D4R_TEX_SRC names it and build.sh selects both for the gfx12 native-FP8 builds. swin_tex.h carries the indexing the two tails share; swin_tok.h gains the pixel-order token layout (SWIN_TOK_NATURAL), ZLUDA's attention key order (SWIN_PV_ZLUDA_ORDER) and swin_tok_core, which the encoded tests reuse. swin_common.h takes the block weight layout (moved out of swin_block.h) and the exact and approximate reciprocal / mix forms the tails need without the HIP headers (D4R_NO_HIP_RUNTIME). The f32 sums stay in ZLUDA's order (K slots of the FP8 WMMAs, key order of the f16 attention product) and rsqrt/rcp use the hardware approximations ZLUDA lowers to, so the harness's final frame is byte-identical to the translated kernels: all four enc0 flag variants, 720p and 1080p render, 1440p and 4K output, static and moving scenes, 100 frames, and the accuracy set. With M's own sum order about one token in 20,000 differed. M's layers compile to the same instructions and its frame is byte-identical. enc0 1.69 -> 0.63 ms and dec0 1.26 -> 0.54 ms at 720p->4K, 2.73 -> 1.16 and 2.28 -> 1.01 at 1080p->4K; the preset 5.83 -> 4.06 and 8.63 -> 5.69 ms. -DD4R_SWIN_K32_MIX rounds the f16 products earlier for 0.56 and 0.49 ms and is not byte-identical, so it is not built by default. The repo's tests pass (36, 3 skipped).
This commit is contained in:
woutermans committed
1 parent
0a12b6ce16
commit
8813ea723c
11 files changed
+885
-151
No files matched your search
@@ -27,7 +27,7 @@ read -r -a FLAGS <<< "${D4R_TEX_CFLAGS:-}"
|
||||
"$LLVM/llvm-dis" "$B/raw.bc" -o "$B/raw.ll"
|
||||
sed -E -e '/@llvm.used/d' -e '/wchar_size/d' -e '/llvm.module.flags/d' -e '/__hip_cuid/d' -e 's/optnone//g' \
|
||||
-e "s/\"target-cpu\"=\"$ARCH\"//g" -e 's/"target-features"="[^"]+"//g' "$B/raw.ll" | "$LLVM/llvm-as" -o "$B/extra.bc" -
|
||||
D4R_PREFER_ACCURACY="$ACCURACY" D4R_TEX_FP8="$FP8" python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
|
||||
D4R_PREFER_ACCURACY="$ACCURACY" D4R_TEX_FP8="$FP8" D4R_TEX_SRC="$SRC" python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
|
||||
# the code-generation switches of the runtime (d4r.ini [Kernels]), which ZLUDA's module cache keys on
|
||||
env LD_LIBRARY_PATH="$R/lib" D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 \
|
||||
D4R_ZLUDA_WMMA_FP8_NATIVE="$FP8" D4R_ZLUDA_IGNORE_DENORMAL="$((1 - ACCURACY))" \
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
// Native block of preset L's unfolded rrlite_dec0_4x4, RDNA4 native FP8 only.
|
||||
//
|
||||
// The original kernel is, per 8x8-pixel block (one wave): the patch expand of the block's 16 half-resolution
|
||||
// tokens (K 64 -> 4 x 32 channels), x0 = f16(q8(expand) + skip) with enc0's pixel tokens as the skip input, a
|
||||
// Swin block (C 32, 2 heads, 4x4 windows), and output code that turns the block's f16 result into four surfaces.
|
||||
// make_ptx.py keeps the output code as ZLUDA's compile of the PTX and replaces everything before it with a call
|
||||
// of d4r_dec0l_block. The function leaves the block's result in shared memory in the register layout the
|
||||
// output code reads (the D fragments of NVIDIA's last 16 MMAs), which the edited PTX loads back.
|
||||
// The arithmetic is the token-lane block of ../m/swin_tok.h in the order of ZLUDA's lowering (see
|
||||
// enc0l_tail.hip): harness frames are byte-identical to the original kernel's.
|
||||
#include "tex_common.h"
|
||||
|
||||
#if !defined(D4R_TEX_FP8) || D4R_WMMA_LAYOUT != 12
|
||||
#error "dec0l_block.hip needs the gfx12 native-FP8 build (D4R_TEX_FP8=1)"
|
||||
#endif
|
||||
|
||||
#include "swin_tex.h"
|
||||
|
||||
namespace
|
||||
{
|
||||
using namespace d4r_swin;
|
||||
constexpr int DC = 32, DNH = 2, DCIN = 64, DNT = DC / 16;
|
||||
using DL = TokLayout<DC, DNH, 0, DCIN>;
|
||||
constexpr TokTables<DC, DNH, 0, DCIN> g_dl_tok = make_tok_tables<DC, DNH, 0, DCIN>();
|
||||
} // namespace
|
||||
|
||||
__attribute__((device)) d4r_swin::wslot g_dl_w16[DL::TOTAL];
|
||||
__attribute__((device)) d4r_swin::h8 g_dl_c16[DL::CTOTAL];
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_blocks = ((DL::TOTAL > DL::CTOTAL ? DL::TOTAL : DL::CTOTAL) + 127) / 128;
|
||||
// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
|
||||
|
||||
struct Dec0LParams
|
||||
{
|
||||
uint64_t w;
|
||||
uint8_t rest[160]; // 168-byte parameter block
|
||||
};
|
||||
|
||||
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4x4_prep(Dec0LParams prm)
|
||||
{
|
||||
const int idx = (int)(__builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x());
|
||||
const uint8_t* w = (const uint8_t*)prm.w;
|
||||
if (idx == 0)
|
||||
d4r_swin_one = (float)(prm.w != 0);
|
||||
if (idx < DL::TOTAL)
|
||||
expand_tok_weights(w, g_dl_w16, g_dl_tok.d, DL::NDESC, idx);
|
||||
if (idx < DL::CTOTAL)
|
||||
expand_tok_consts(w, g_dl_c16, g_dl_tok.c, DL::NCONST, idx);
|
||||
}
|
||||
|
||||
// in: e4m3 planes of 32 channels (two) over tw / 2 x th / 2 tokens; skip: one such plane over tw x th tokens
|
||||
// (pixels). (sx, sy): offset of the block grid against the token grid; (bx, by): this block.
|
||||
// Returns the shared-memory address of the result image: register R = 2 (4 m + j) + h of lane L (4 bytes at
|
||||
// 4 (32 R + L)) is D register h of NVIDIA's MMA for token tile m (window) and column tile j, i.e. rows
|
||||
// 16 m + (L >> 2) + 8 h, pair-order columns 8 j + 2 (L & 3), + 1 of the block's f16 output.
|
||||
DEV uint32_t d4r_dec0l_block(uint64_t wp, uint64_t inp, uint64_t skipp, int32_t sx, int32_t sy, int32_t tw, int32_t th,
|
||||
int32_t bx, int32_t by)
|
||||
{
|
||||
using K = DL;
|
||||
constexpr int C = DC, NT = DNT, CIN = DCIN;
|
||||
const int lane = (int)lane_id(), l16 = lane & 15, hf = lane >> 4;
|
||||
(void)wp; // the images were prepared from it
|
||||
uintptr_t chv = (uintptr_t)(g_dl_c16 + hf), wlv = (uintptr_t)((const gop*)g_dl_w16 + 2 * l16 + hf);
|
||||
asm("" : "+v"(chv));
|
||||
asm("" : "+v"(wlv));
|
||||
const h8* const ch = (const h8*)chv;
|
||||
const gop* const wl = (const gop*)wlv;
|
||||
auto wt = [&](int dst, int nt_count, int kc, int st, int nt) { return wl[2 * (dst + ((kc * 2 + st) * nt_count + nt) * 16)]; };
|
||||
|
||||
__shared__ __attribute__((aligned(16))) uint8_t E8[64 * C];
|
||||
__shared__ __attribute__((aligned(16))) uint32_t X2[32 * 32];
|
||||
|
||||
// patch expand: lanes = the block's 16 half-resolution tokens, one sub-position (quadrant) after the other;
|
||||
// the codes of q8(expand) go to their child tokens' rows of E8 (as swin_tok_block, which uses a wave each)
|
||||
{
|
||||
const int W2 = tw / 2, H2 = th / 2;
|
||||
const int MX = mirror((8 * bx - sx) / 2 + (l16 & 3), W2), MY = mirror((8 * by - sy) / 2 + (l16 >> 2), H2);
|
||||
gop pa[CIN / 32][2];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < CIN / 32; ++kc)
|
||||
{
|
||||
const uint8_t* b = (const uint8_t*)inp + ((size_t)(kc * H2 + MY) * W2 + MX) * 32 + tok_half(hf);
|
||||
#pragma unroll
|
||||
for (int st = 0; st < 2; ++st)
|
||||
pa[kc][st] = (u2v){*(const uint32_t*)(b + tok_word(st, 0)), *(const uint32_t*)(b + tok_word(st, 1))};
|
||||
}
|
||||
#pragma unroll 1
|
||||
for (int q = 0; q < 4; ++q)
|
||||
{
|
||||
const int tx = 2 * (l16 & 3) + (q & 1), ty = 2 * (l16 >> 2) + (q >> 1);
|
||||
const int Tc = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
|
||||
const int pew = K::pe(0) + q * (CIN * C / 16);
|
||||
#pragma unroll
|
||||
for (int ct = 0; ct < NT; ++ct)
|
||||
{
|
||||
T16 e = tok_t16(ch[K::CPEB + 2 * (NT * q + 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
wave_sync();
|
||||
|
||||
#pragma unroll 1
|
||||
for (int wv = 0; wv < 4; ++wv)
|
||||
{
|
||||
const int T = 16 * wv + l16;
|
||||
const int Xm = mirror(8 * bx - sx + tok_x(T), tw), Ym = mirror(8 * by - sy + tok_y(T), th);
|
||||
const uint8_t* const sb = (const uint8_t*)skipp + ((size_t)Ym * (size_t)tw + (size_t)Xm) * 32 + tok_half(hf);
|
||||
hv2 v[NT][4];
|
||||
#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 + tok_word(t, g) + tok_half(hf));
|
||||
const uint32_t sw = *(const uint32_t*)(sb + tok_word(t, g));
|
||||
v[t][2 * g] = tok_dec2<false>(ew) + tok_dec2<false>(sw);
|
||||
v[t][2 * g + 1] = tok_dec2<true>(ew) + tok_dec2<true>(sw);
|
||||
}
|
||||
|
||||
uint32_t co[NT][2];
|
||||
hv2 x2[NT][4];
|
||||
swin_tok_core<DC, DNH, 0, DCIN>(ch, wl, l16, v, co, x2);
|
||||
(void)co;
|
||||
|
||||
// pair k of tile t holds natural channels 16 t + 8 half + 2 k, + 1: pair-order columns 8 j + 2 u, + 1 with
|
||||
// j = t + 2 (k & 1), u = 2 half + (k >> 1); token row l16 of tile wv
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
const int R = 2 * (4 * wv + t + 2 * (k & 1)) + (l16 >> 3), L = 4 * (l16 & 7) + 2 * hf + (k >> 1);
|
||||
X2[32 * R + L] = __builtin_bit_cast(uint32_t, x2[t][k]);
|
||||
}
|
||||
}
|
||||
wave_sync();
|
||||
return (uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)X2;
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
// Native tail of preset L's unfolded rrlite_enc0_4x4_* (all flag combinations), RDNA4 native FP8 only.
|
||||
//
|
||||
// The kernel is ZLUDA's compile of NVIDIA's PTX up to the barrier that ends the feature staging: one wave per
|
||||
// 8x8-pixel block has then written 32 e4m3 features per pixel to shared memory. Everything after it is this
|
||||
// tail (make_ptx.py calls it instead of the original code):
|
||||
// embedding x0 = W2 q8(relu(W1 f + b1)) + b2 (K 32 -> 32 -> 32, the layout of an MLP chunk)
|
||||
// Swin block C 32, 2 heads, 4x4 windows (../m/swin_tok.h, one window after the other)
|
||||
// output the block's e4m3 codes at pixel resolution (the skip input of dec0)
|
||||
// patch merge 2x2 pixels (K 128) -> 64 channels, clamped to +-2 pi, at half resolution (the input of enc1)
|
||||
// with NVIDIA's arithmetic (e4m3 activations, an f16 accumulator rounded after every k32 step) in the order of
|
||||
// ZLUDA's lowering of it (swin_tex.h): harness frames are byte-identical to the original kernel's.
|
||||
#include "tex_common.h"
|
||||
|
||||
#if !defined(D4R_TEX_FP8) || D4R_WMMA_LAYOUT != 12
|
||||
#error "enc0l_tail.hip needs the gfx12 native-FP8 build (D4R_TEX_FP8=1)"
|
||||
#endif
|
||||
|
||||
#include "swin_tex.h"
|
||||
|
||||
namespace
|
||||
{
|
||||
using namespace d4r_swin;
|
||||
constexpr int EC = 32, ENH = 2, ENPM = 64, ENT = EC / 16;
|
||||
using EL = TokLayout<EC, ENH, ENPM, 0>;
|
||||
// weights: the embedding (w1 at 4096, b1, w2, b2), then the Swin block and the patch merge
|
||||
constexpr int E_EMB = 4096, E_BLOCK = E_EMB + 2 * (32 * EC + 2 * EC);
|
||||
// images: the block's tiles and constants (bases relative to the block), then the embedding's
|
||||
constexpr int EW1 = EL::TOTAL, EW2 = EW1 + 64, EW_TOTAL = EW2 + 64;
|
||||
constexpr int ECB1 = EL::CTOTAL, ECB2 = ECB1 + 4, EC_TOTAL = ECB2 + 2 * ENT;
|
||||
constexpr TokTables<EC, ENH, ENPM, 0> g_el_tok = make_tok_tables<EC, ENH, ENPM, 0>();
|
||||
constexpr TokDesc g_el_emb[2] = {{EW1, E_EMB, 512, 16 * EC, 1, 2, TOK_PI, 0, 0, 0},
|
||||
{EW2, E_EMB + 32 * EC + 2 * EC, 0, 512, 1, ENT, TOK_X, 0, 0, 0}};
|
||||
constexpr TokConst g_el_embc[2] = {{ECB1, TC_PI, E_EMB + 32 * EC}, {ECB2, TC_X, E_EMB + 2 * (32 * EC) + 2 * EC}};
|
||||
} // namespace
|
||||
|
||||
__attribute__((device)) d4r_swin::wslot g_el_w16[EW_TOTAL];
|
||||
__attribute__((device)) d4r_swin::h8 g_el_c16[EC_TOTAL];
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_blocks = (EW_TOTAL + 127) / 128;
|
||||
// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
|
||||
|
||||
struct Enc0LParams
|
||||
{
|
||||
uint64_t w;
|
||||
uint8_t rest[352];
|
||||
};
|
||||
|
||||
#ifndef D4R_KERNEL_NAME
|
||||
#define D4R_KERNEL_NAME rrlite_enc0_4x4_mvhi_hdr
|
||||
#endif
|
||||
#define D4R_CAT2(a, b) a##b
|
||||
#define D4R_CAT(a, b) D4R_CAT2(a, b)
|
||||
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void D4R_CAT(D4R_KERNEL_NAME, _prep)(Enc0LParams prm)
|
||||
{
|
||||
const int idx = (int)(__builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x());
|
||||
const uint8_t* w = (const uint8_t*)prm.w;
|
||||
if (idx == 0)
|
||||
d4r_swin_one = (float)(prm.w != 0);
|
||||
if (idx < EL::TOTAL)
|
||||
expand_tok_weights(w + E_BLOCK, g_el_w16, g_el_tok.d, EL::NDESC, idx);
|
||||
else if (idx < EW_TOTAL)
|
||||
expand_tok_weights(w, g_el_w16, g_el_emb, 2, idx);
|
||||
if (idx < EL::CTOTAL)
|
||||
expand_tok_consts(w + E_BLOCK, g_el_c16, g_el_tok.c, EL::NCONST, idx);
|
||||
else if (idx < EC_TOTAL)
|
||||
expand_tok_consts(w, g_el_c16, g_el_embc, 2, idx);
|
||||
}
|
||||
|
||||
// e4m3 codes of four halves that are already within +-448 (no saturation clamp)
|
||||
__attribute__((device, always_inline)) static inline uint32_t el_codes4(d4r_swin::hv2 a, d4r_swin::hv2 b)
|
||||
{
|
||||
const uint32_t w = (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32((float)a[0], (float)a[1], 0, false);
|
||||
return (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32((float)b[0], (float)b[1], (int)w, true);
|
||||
}
|
||||
|
||||
// smem: the staged features, byte k of pixel (px, py) of the block at 1024 (k >> 4) + 128 py + 16 px + (k & 15).
|
||||
// out: e4m3 planes of 32 channels over tw x th tokens (pixels); merged: two such planes over tw / 2 x th / 2.
|
||||
// (sx, sy): offset of the block grid against the token grid; (bx, by): this block.
|
||||
DEV void d4r_enc0l_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint64_t mergedp, int32_t sx, int32_t sy, int32_t tw,
|
||||
int32_t th, int32_t bx, int32_t by)
|
||||
{
|
||||
using K = EL;
|
||||
constexpr int C = EC, NT = ENT, XB = C + 4, NTP = ENPM / 16;
|
||||
const int lane = (int)lane_id(), l16 = lane & 15, hf = lane >> 4;
|
||||
const lds_u8* const s = (const lds_u8*)(uintptr_t)smem;
|
||||
(void)wp; // the images were prepared from it
|
||||
// this lane's bases into the two images, as in swin_tok_block
|
||||
uintptr_t chv = (uintptr_t)(g_el_c16 + hf), wlv = (uintptr_t)((const gop*)g_el_w16 + 2 * l16 + hf);
|
||||
asm("" : "+v"(chv));
|
||||
asm("" : "+v"(wlv));
|
||||
const h8* const ch = (const h8*)chv;
|
||||
const gop* const wl = (const gop*)wlv;
|
||||
auto wt = [&](int dst, int nt_count, int kc, int st, int nt) { return wl[2 * (dst + ((kc * 2 + st) * nt_count + nt) * 16)]; };
|
||||
|
||||
// the output codes of the block's 64 tokens, for the patch merge
|
||||
__shared__ __attribute__((aligned(16))) uint8_t U[64 * XB];
|
||||
|
||||
#pragma unroll 1
|
||||
for (int wv = 0; wv < 4; ++wv)
|
||||
{
|
||||
const int T = 16 * wv + l16, px = tok_x(T), py = tok_y(T);
|
||||
// features: K slot 8 half + i of WMMA st is feature 16 st + 8 half + i, 8 contiguous staged bytes
|
||||
gop fa[2];
|
||||
{
|
||||
const lds_u8* f = s + 128 * py + 16 * px + 8 * hf;
|
||||
#pragma unroll
|
||||
for (int st = 0; st < 2; ++st)
|
||||
fa[st] = *(const __attribute__((address_space(3))) u2v*)(f + 1024 * st);
|
||||
}
|
||||
// embedding: the layout of an MLP chunk (32 hidden as two tiles), ReLU instead of the GELU
|
||||
gop gz[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
T16 z = tok_t16(ch[ECB1 + 2 * nt]);
|
||||
k32(z, wt(EW1, 2, 0, 0, nt), fa[0], wt(EW1, 2, 0, 1, nt), fa[1]);
|
||||
hv2 r[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
r[k] = __builtin_elementwise_maximumnum(z.pair(k), (hv2){(half_t)0.0f, (half_t)0.0f});
|
||||
gz[nt] = (u2v){codes4(r[0], r[1]), codes4(r[2], r[3])};
|
||||
}
|
||||
hv2 v[NT][4];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
{
|
||||
T16 x0 = tok_t16(ch[ECB2 + 2 * t]);
|
||||
k32(x0, wt(EW2, NT, 0, 0, t), gz[0], wt(EW2, NT, 0, 1, t), gz[1]);
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
v[t][k] = x0.pair(k);
|
||||
}
|
||||
|
||||
uint32_t co[NT][2];
|
||||
swin_tok_core<EC, ENH, ENPM, 0>(ch, wl, l16, v, co);
|
||||
|
||||
const int X = 8 * bx - sx + px, Y = 8 * by - sy + py;
|
||||
if (X >= 0 && X < tw && Y >= 0 && Y < th)
|
||||
{
|
||||
uint8_t* b = (uint8_t*)outp + ((size_t)Y * (size_t)tw + (size_t)X) * 32 + tok_half(hf);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
*(uint32_t*)(b + tok_word(j >> 1, j & 1)) = co[j >> 1][j & 1];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int t = 0; t < NT; ++t)
|
||||
#pragma unroll
|
||||
for (int g = 0; g < 2; ++g)
|
||||
*(uint32_t*)(U + T * XB + tok_word(t, g) + tok_half(hf)) = co[t][g];
|
||||
}
|
||||
// every lane's codes are in U before any lane reads them
|
||||
wave_sync();
|
||||
|
||||
// patch merge: lanes = the block's 16 merged tokens, the four source pixels' codes as in swin_tok_block
|
||||
{
|
||||
const int mx = l16 & 3, my = l16 >> 2;
|
||||
const uint8_t* const ub = U + ((2 * mx & 3) + 4 * (2 * my & 3) + 16 * (mx >> 1) + 32 * (my >> 1)) * XB + tok_half(hf);
|
||||
gop as[C / 8][2];
|
||||
#pragma unroll
|
||||
for (int kc = 0; kc < C / 8; ++kc)
|
||||
#pragma unroll
|
||||
for (int st = 0; st < 2; ++st)
|
||||
{
|
||||
uint32_t run4[2];
|
||||
#pragma unroll
|
||||
for (int rr = 0; rr < 2; ++rr)
|
||||
{
|
||||
const int k0 = tok_merge_k(kc, st, rr), q = k0 / C;
|
||||
run4[rr] = *(const uint32_t*)(ub + ((q & 1) + 4 * (q >> 1)) * XB + k0 % C);
|
||||
}
|
||||
as[kc][st] = (u2v){run4[0], run4[1]};
|
||||
}
|
||||
const int W2 = tw / 2, H2 = th / 2;
|
||||
const int MX = (8 * bx - sx) / 2 + mx, MY = (8 * by - sy) / 2 + my;
|
||||
const bool inside = MX >= 0 && MX < W2 && MY >= 0 && MY < H2;
|
||||
const hv2 lim = {(half_t)6.2831855f, (half_t)6.2831855f};
|
||||
#pragma unroll 1
|
||||
for (int u = 0; u < NTP; ++u)
|
||||
{
|
||||
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]);
|
||||
hv2 c[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
c[k] = __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(y.pair(k), -lim), lim);
|
||||
const u2v code = {el_codes4(c[0], c[1]), el_codes4(c[2], c[3])};
|
||||
if (inside)
|
||||
*(u2v*)((uint8_t*)mergedp + ((size_t)((u >> 1) * H2 + MY) * (size_t)W2 + (size_t)MX) * 32 + 16 * (u & 1) + 8 * hf) = code;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -239,6 +239,176 @@ def replace_rz_round(lines):
|
||||
return out, n
|
||||
|
||||
|
||||
# Preset L's unfolded enc0 on the gfx12 native-FP8 build (source enc0l_tail): everything after the barrier that
|
||||
# ends the feature staging is the native tail. It needs only the staged features, the parameter block and the
|
||||
# block index, so the call loads those itself and no register of the original code is read.
|
||||
ENC0L_RE = re.compile(r"^rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)$")
|
||||
ENC0L_EXTERN = """.extern .func d4r_enc0l_tail
|
||||
(
|
||||
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b64 d4r_a3, .param .b32 d4r_a4,
|
||||
.param .b32 d4r_a5, .param .b32 d4r_a6, .param .b32 d4r_a7, .param .b32 d4r_a8, .param .b32 d4r_a9
|
||||
)
|
||||
;
|
||||
"""
|
||||
ENC0L_TAIL = """mov.u32 %rel0, {smem};
|
||||
ld.param.u64 %rdel0, [{name}_param_0];
|
||||
ld.param.v2.u32 {{%rel1, %rel2}}, [{name}_param_0+8];
|
||||
ld.param.v2.u32 {{%rel3, %rel4}}, [{name}_param_0+16];
|
||||
ld.param.u64 %rdel1, [{name}_param_0+40];
|
||||
ld.param.u64 %rdel2, [{name}_param_0+48];
|
||||
mov.u32 %rel5, %ctaid.x;
|
||||
mov.u32 %rel6, %ctaid.y;
|
||||
{{
|
||||
.param .b32 q0;
|
||||
st.param.b32 [q0+0], %rel0;
|
||||
.param .b64 q1;
|
||||
st.param.b64 [q1+0], %rdel0;
|
||||
.param .b64 q2;
|
||||
st.param.b64 [q2+0], %rdel1;
|
||||
.param .b64 q3;
|
||||
st.param.b64 [q3+0], %rdel2;
|
||||
.param .b32 q4;
|
||||
st.param.b32 [q4+0], %rel1;
|
||||
.param .b32 q5;
|
||||
st.param.b32 [q5+0], %rel2;
|
||||
.param .b32 q6;
|
||||
st.param.b32 [q6+0], %rel3;
|
||||
.param .b32 q7;
|
||||
st.param.b32 [q7+0], %rel4;
|
||||
.param .b32 q8;
|
||||
st.param.b32 [q8+0], %rel5;
|
||||
.param .b32 q9;
|
||||
st.param.b32 [q9+0], %rel6;
|
||||
call.uni d4r_enc0l_tail, (q0, q1, q2, q3, q4, q5, q6, q7, q8, q9);
|
||||
}}
|
||||
ret;
|
||||
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
def enc0l_variant(name):
|
||||
"""Cut line and tail of an unfolded enc0 variant: three bar.warp.sync (after the feature staging, after
|
||||
the embedding, before the patch merge), 336 MMAs after the first, one 2048-byte shared array."""
|
||||
path = module_file(name, {"file": ""})
|
||||
lines = read_lines(path)
|
||||
syncs = [i for i, l in enumerate(lines) if l.strip() == "bar.warp.sync -1;"]
|
||||
mmas = [i for i, l in enumerate(lines) if "mma.sync.aligned" in l]
|
||||
shared = [m.group(1) for m in (re.match(r"^\s*\.shared \.align \d+ \.b8 (\S+)\[2048\];$", l) for l in lines) if m]
|
||||
if len(syncs) != 3 or len(mmas) != 336 or mmas[0] < syncs[0] or len(shared) != 1:
|
||||
sys.exit(f"enc0l: {name} is not the kernel the tail was written for ({len(syncs)} barriers, {len(mmas)} MMAs, {shared})")
|
||||
return dict(file=path.name, cut_after=syncs[0] + 1, cut_check="bar.warp.sync -1;", sust=True, extern=ENC0L_EXTERN,
|
||||
tail=ENC0L_TAIL.format(name=name, smem=shared[0]), regs=[".reg .b32 %rel<7>;", ".reg .b64 %rdel<3>;"])
|
||||
|
||||
|
||||
# Preset L's unfolded dec0 on the gfx12 native-FP8 build (source dec0l_block): the patch expand, the skip sum and
|
||||
# the Swin block become one native call; the output code after the last MMA stays. That code reads the block's
|
||||
# result from the D registers of the last 16 MMAs (token tile m, column tile j, in that order), which the
|
||||
# native function leaves in shared memory in that layout, and otherwise depends only on a few statements
|
||||
# (block coordinates, constants) that are found by following its inputs backwards and are kept.
|
||||
DEC0L_EXTERN = """.extern .func (.param .b32 d4r_dr) d4r_dec0l_block
|
||||
(
|
||||
.param .b64 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3, .param .b32 d4r_a4,
|
||||
.param .b32 d4r_a5, .param .b32 d4r_a6, .param .b32 d4r_a7, .param .b32 d4r_a8
|
||||
)
|
||||
;
|
||||
"""
|
||||
DEC0L_CALL = """ld.param.u64 %rddl0, [{name}_param_0];
|
||||
ld.param.v2.u32 {{%rdl1, %rdl2}}, [{name}_param_0+8];
|
||||
ld.param.v2.u32 {{%rdl3, %rdl4}}, [{name}_param_0+16];
|
||||
ld.param.u64 %rddl1, [{name}_param_0+24];
|
||||
ld.param.u64 %rddl2, [{name}_param_0+32];
|
||||
mov.u32 %rdl5, %ctaid.x;
|
||||
mov.u32 %rdl6, %ctaid.y;
|
||||
{{
|
||||
.param .b64 q0;
|
||||
st.param.b64 [q0+0], %rddl0;
|
||||
.param .b64 q1;
|
||||
st.param.b64 [q1+0], %rddl1;
|
||||
.param .b64 q2;
|
||||
st.param.b64 [q2+0], %rddl2;
|
||||
.param .b32 q3;
|
||||
st.param.b32 [q3+0], %rdl1;
|
||||
.param .b32 q4;
|
||||
st.param.b32 [q4+0], %rdl2;
|
||||
.param .b32 q5;
|
||||
st.param.b32 [q5+0], %rdl3;
|
||||
.param .b32 q6;
|
||||
st.param.b32 [q6+0], %rdl4;
|
||||
.param .b32 q7;
|
||||
st.param.b32 [q7+0], %rdl5;
|
||||
.param .b32 q8;
|
||||
st.param.b32 [q8+0], %rdl6;
|
||||
.param .b32 d0;
|
||||
call (d0), d4r_dec0l_block, (q0, q1, q2, q3, q4, q5, q6, q7, q8);
|
||||
ld.param.b32 %rdl0, [d0+0];
|
||||
}}
|
||||
mov.u32 %rdl7, %laneid;
|
||||
shl.b32 %rdl7, %rdl7, 2;
|
||||
add.u32 %rdl0, %rdl0, %rdl7;"""
|
||||
PTX_REG_RE = re.compile(r"%(?:rd|rs|r|fd|f|p)\d+")
|
||||
|
||||
|
||||
def ptx_statements(line):
|
||||
"""(destinations, sources, opcode) of each simple statement on a line (inline { ... } blocks included)."""
|
||||
text = re.sub(r"^\s*\{\s*(?=[a-z.@])", "", line.strip())
|
||||
text = re.sub(r";\s*\}\s*$", ";", text)
|
||||
out = []
|
||||
for stmt in text.split(";"):
|
||||
m = re.match(r"\s*(@!?%p\d+\s+)?([a-z][a-z0-9_.]*)\s+(.*)$", stmt)
|
||||
if not m or stmt.strip().startswith("."):
|
||||
continue
|
||||
pred, op, args = m.groups()
|
||||
uses = set(PTX_REG_RE.findall(pred or ""))
|
||||
if op.startswith(("st.", "sust", "bra", "bar", "ret", "call")):
|
||||
out.append((set(), uses | set(PTX_REG_RE.findall(args)), op))
|
||||
continue
|
||||
if args.startswith("{"):
|
||||
end = args.index("}") + 1
|
||||
else:
|
||||
end = len(args.split(",")[0])
|
||||
# (shfl's "%r|%p" destination pair is one operand: both are definitions)
|
||||
out.append((set(PTX_REG_RE.findall(args[:end])), uses | set(PTX_REG_RE.findall(args[end:])), op))
|
||||
return out
|
||||
|
||||
|
||||
def dec0l_spec(name):
|
||||
path = module_file(name, {"file": ""})
|
||||
lines = read_lines(path)
|
||||
mmas = [i for i, l in enumerate(lines) if "mma.sync.aligned" in l]
|
||||
entry = next(i for i, l in enumerate(lines) if l.startswith(".visible .entry"))
|
||||
if len(mmas) != 304 or sum("sust." in l for l in lines[mmas[-1]:]) != 8:
|
||||
sys.exit(f"dec0l: {name} is not the kernel the block was written for ({len(mmas)} MMAs)")
|
||||
x2 = []
|
||||
for i in mmas[-16:]:
|
||||
x2 += [r.strip() for r in re.search(r"\{([^}]*)\}", lines[i]).group(1).split(",")]
|
||||
end = next(i for i in range(mmas[-1], len(lines)) if lines[i].rstrip().endswith(";")) # the MMA's last line
|
||||
body = next(i for i in range(entry, len(lines)) if lines[i].strip() == "{") + 1
|
||||
# registers the output code reads before it defines them
|
||||
need, defined = set(), set()
|
||||
for line in lines[end + 1:]:
|
||||
for d, u, _ in ptx_statements(line):
|
||||
need |= u - defined
|
||||
defined |= d
|
||||
need -= set(x2)
|
||||
keep = set()
|
||||
for i in range(end, body - 1, -1):
|
||||
for d, u, op in ptx_statements(lines[i]):
|
||||
if d & need:
|
||||
if not re.match(r"^(mov|add|sub|shl|shr|and|or|ld\.param)\b", op):
|
||||
sys.exit(f"dec0l: the output code depends on {lines[i].strip()!r}")
|
||||
keep.add(i)
|
||||
need = (need - d) | (u - set(x2))
|
||||
if len(keep) > 24 or len(x2) != 32:
|
||||
sys.exit(f"dec0l: unexpected dependencies of the output code ({len(keep)} statements)")
|
||||
# (the kernel's own declarations come first; declarations further down belong to inline blocks that go)
|
||||
first = next(i for i in range(body, end) if lines[i].strip() and not lines[i].lstrip().startswith((".", "//")))
|
||||
drop = {i for i in range(first, end + 1) if i not in keep}
|
||||
loads = "\n".join(f"ld.shared.u32 {reg}, [%rdl0+{128 * r}];" for r, reg in enumerate(x2))
|
||||
return dict(file=path.name, drop=drop, insert_at=end + 1, insert=DEC0L_CALL.format(name=name) + "\n" + loads,
|
||||
sust=True, extern=DEC0L_EXTERN, regs=[".reg .b32 %rdl<8>;", ".reg .b64 %rddl<3>;"])
|
||||
|
||||
|
||||
ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
|
||||
# the other flag combinations of the kernels whose only native part is the surface-store rewrite
|
||||
SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
|
||||
@@ -307,6 +477,10 @@ def kernel_spec(name):
|
||||
return KERNELS[name]
|
||||
if ENC0_RE.match(name):
|
||||
return enc0_variant(name)
|
||||
if ENC0L_RE.match(name) and os.environ.get("D4R_TEX_SRC") == "enc0l_tail":
|
||||
return enc0l_variant(name)
|
||||
if name == "rrlite_dec0_4x4" and os.environ.get("D4R_TEX_SRC") == "dec0l_block":
|
||||
return dec0l_spec(name)
|
||||
if SUST_ONLY_RE.match(name):
|
||||
return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_output"))
|
||||
sys.exit(f"make_ptx: no recipe for {name}")
|
||||
@@ -334,6 +508,10 @@ def main():
|
||||
if lines[cut - 1].strip() != k["cut_check"]:
|
||||
sys.exit(f"line {cut} is {lines[cut - 1]!r}, expected {k['cut_check']!r}")
|
||||
head, tail = lines[:cut], k["tail"].split("\n")
|
||||
elif "drop" in k:
|
||||
at = k["insert_at"]
|
||||
head = [l for i, l in enumerate(lines[:at]) if i not in k["drop"]] + k["insert"].split("\n") + lines[at:]
|
||||
tail = []
|
||||
elif "delete" in k:
|
||||
a, b = k["delete"]
|
||||
if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]:
|
||||
@@ -343,6 +521,10 @@ def main():
|
||||
head, tail = lines, []
|
||||
entry = next(i for i, l in enumerate(head) if l.startswith(".visible .entry"))
|
||||
body, n = replace_sust(head[entry:]) if k.get("sust") else (head[entry:], {"b16": 0, "p": 0, "b32": 0})
|
||||
if k.get("regs"):
|
||||
# registers of the inserted code, declared after the kernel's own .reg lines
|
||||
at = next(i for i, l in enumerate(body) if l.startswith(".reg"))
|
||||
body = body[:at] + k["regs"] + body[at:]
|
||||
if k.get("rz_round"):
|
||||
body, nrz = replace_rz_round(body)
|
||||
# temporaries for the rewrite, declared after the kernel's own .reg lines
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
// The token-lane Swin block (../m/swin_tok.h) for the native parts of ZLUDA-compiled texture kernels: preset L's
|
||||
// enc0 (enc0l_tail.hip) and dec0 (dec0l_block.hip).
|
||||
#pragma once
|
||||
|
||||
// The Swin sources are written against the HIP headers; a texture tail is built without them. These are the
|
||||
// runtime names they use.
|
||||
#ifndef HIP_INCLUDE_HIP_HIP_RUNTIME_H
|
||||
#define D4R_NO_HIP_RUNTIME
|
||||
#define __device__ __attribute__((device))
|
||||
#define __host__ __attribute__((host))
|
||||
#define __forceinline__ inline __attribute__((always_inline))
|
||||
#define __shared__ __attribute__((shared))
|
||||
typedef __SIZE_TYPE__ size_t;
|
||||
struct d4r_tid_t
|
||||
{
|
||||
struct X { __attribute__((device, always_inline)) operator unsigned() const { return __builtin_amdgcn_workitem_id_x(); } } x;
|
||||
struct Z { __attribute__((device, always_inline)) operator unsigned() const { return __builtin_amdgcn_workitem_id_z(); } } z;
|
||||
};
|
||||
static const d4r_tid_t threadIdx;
|
||||
__attribute__((device, always_inline)) static inline void __syncthreads()
|
||||
{
|
||||
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
|
||||
__builtin_amdgcn_s_barrier();
|
||||
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
|
||||
}
|
||||
#endif
|
||||
|
||||
// 1.0 for the k32 rounding of the D4R_SWIN_K32_MIX variant (see mix2 in ../m/swin_common.h). The prep kernel stores
|
||||
// it, so that it is not a constant to the compiler.
|
||||
__attribute__((device, used)) float d4r_swin_one = 1.0f;
|
||||
#define MIX_ONE d4r_swin_one
|
||||
|
||||
namespace d4r_swin
|
||||
{
|
||||
#define SWIN_EXACT
|
||||
// The original code is ZLUDA's compile of the PTX; these make the f32 arithmetic the same instruction for
|
||||
// instruction where it is not exact (see ../m/swin_tok.h, ../m/swin_common.h): the K order and the C operand of
|
||||
// the FP8 WMMAs, the key order of the f16 P V product, and the hardware reciprocals.
|
||||
#define SWIN_TOK_NATURAL
|
||||
#ifndef D4R_SWIN_K32_MIX
|
||||
#define K32_NO_MIX
|
||||
#endif
|
||||
#define SWIN_PV_ZLUDA_ORDER
|
||||
#define SWIN_APPROX_RECIP
|
||||
#ifndef D4R_FP8_WMMA
|
||||
#define D4R_FP8_WMMA
|
||||
#endif
|
||||
#include "../m/swin_common.h"
|
||||
#include "../m/swin_tok.h"
|
||||
} // namespace d4r_swin
|
||||
Reference in new issue
Block a user