diff --git a/docs/native-kernels.md b/docs/native-kernels.md index 6ca72e5..56bfe12 100644 --- a/docs/native-kernels.md +++ b/docs/native-kernels.md @@ -57,6 +57,16 @@ The release is built this way, with hashes from DLSS 310.7.0 and 310.9.1. The PT **Motion mitigation.** `[Kernels] NativeSwinEncoders = false` (`D4R_NATIVE_SWIN_ENCODERS=0`) bypasses only `rrlite_enc1_4x4` and `rrlite_enc2_4x4` in the native manifest. L/M then use the original translated PTX for those encoders, retaining native tube, decoder and texture kernels. Keep `PreferAccuracy = true` for the tested conservative path. It reduced the extra building-edge trails in a 64-frame moving Townfall capture at 853×480 → 2560×1440 on gfx1101; it does not recover detail absent from the low-resolution input or establish RTX parity. The default remains `true` for speed. Flat developer sets without manifests fall back entirely when this switch is off; add a manifest to keep selective acceleration. +**L's enc0 and dec0 on RDNA4 native FP8 (`kernels/tex/enc0l_tail.hip`, `dec0l_block.hip`).** The unfolded kernels contain a whole Swin block at pixel resolution (32 channels, 2 heads, 4×4 windows) that M's folded variants do not have: enc0 is the feature code, a two-layer embedding (32 → 32 → 32 with a ReLU), the block, and the 2×2 patch merge to 64 channels; dec0 is the patch expand of the half-resolution tokens, the skip sum, the block, and output code that writes four surfaces. On the `-fp8` build `make_ptx.py` keeps the texture code (enc0's feature staging, dec0's output code) as ZLUDA's compile and replaces the rest with the token-lane block of `swin_tok.h` (`swin_tok_core`), one window after the other in the kernel's single wave per 8×8-pixel block. + +- **enc0:** the cut is the barrier after the feature staging; the tail reads the staged features from shared memory and loads the parameters and the block index itself, so no register of the original code is involved. +- **dec0:** the output code reads the block's f16 result from the D registers of the last 16 MMAs and otherwise depends on nine trivial statements, which `make_ptx.py` finds by following its inputs backwards and keeps. The native function leaves the result in shared memory in that register layout and the edited PTX loads it back. +- **Same bytes as ZLUDA's compile.** These kernels replace code that ZLUDA compiled, so they reproduce its f32 arithmetic where it is not exact (`swin_tex.h`): the two WMMAs of a k32 step hold K 0..15 and K 16..31 in order with C added in the first (`SWIN_TOK_NATURAL`, `K32_NO_MIX`), the f16 P V product has its keys in the K slots of ZLUDA's m16n8k16 lowering (`SWIN_PV_ZLUDA_ORDER`), and the reciprocal square root and reciprocal are `v_rsq_f32` / `v_rcp_f32` (`SWIN_APPROX_RECIP`). The order matters because f32 sums that are not exact round differently in a different order: with M's own slot order the P V product gave other codes than ZLUDA's compile in about one pixel token in 20,000 (each of the two orders then disagrees with `ptxsim.py`'s exact sums in different tokens), and with the single-rounding k32 form one merged code differed in the seventh 3840×2160 frame, where a k32 sum lay 0.0001 ulp from a rounding tie. +- **Validation (RX 9070 XT):** `ptxsim.py` runs the original enc0 from the staging barrier on (texture code stubbed, staged features replaced) and a numpy model built on `swin_model.py` matches its stores on interior and border blocks; the harness's final frame is byte-identical to the `sust_only` builds for all four enc0 flag variants, 1280×720 and 1920×1080 render, 2560×1440 and 3840×2160 output, static and moving scenes, 100 frames, and in the accuracy set. +- **Cost:** enc0 1.69 → 0.63 ms and dec0 1.26 → 0.54 ms at 1280×720 → 3840×2160 ([performance.md](performance.md#rdna4)). The inline assembly of `mix2` is not available in a ZLUDA-compiled kernel (its LLVM has no assembler); `-DD4R_SWIN_K32_MIX` (`D4R_TEX_CFLAGS`) selects the same instructions from `fma(P, one, c)` instead: enc0 0.63 → 0.56 ms and dec0 0.54 → 0.49 ms, but a near-tie sum then rounds as the exact sum does and not as ZLUDA's compile, so it is not byte-identical (first difference in the seventh frame of the test above, 40 dB after 30 frames of temporal feedback). It is not built by default. + +Other targets (RDNA3, and RDNA4 with `NativeFp8` off) keep the `sust_only` builds of these kernels. + Initial validation: a gfx1101 fast L build completed two Ultra Performance harness frames (192×108 → 576×324, DLSS 310.7). Its finite, nonzero RGBA16F output matched the original translated L texture kernels byte for byte with the same native Swin layers. A gfx1201 FP8 accuracy L set compiled successfully; RDNA4 hardware and 4K game output have not been tested. Full packaging rebuilds missing L variants when given an older texture bundle. ## RDNA4 @@ -73,6 +83,7 @@ In that build the GEMM activations live in LDS as e4m3 bytes (`SWIN_A8` in `swin - 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. +- The part between a block's input and its output codes is `swin_tok_core`, which the L texture kernels below reuse. M's layers compile to the same instructions as before the split (a few inverted branches), and M's harness frames are byte-identical. - `-DSWIN_NO_TOKEN_LANES` builds these layers from `swin_block.h` again. enc2 and dec2 (96 channels, four heads) use the same template: byte-identical on the same recorded launches and harness sequences, 0.22 → 0.14 and 0.20 → 0.14 ms at 1280×720 render, 0.75 → 0.47 and 0.69 → 0.47 ms at 2560×1440. - The tube retains upstream's eight weight-image slots plus an overflow slot. Weights and per-channel constants use the same slot; `d4r_prep_identity_at` lets MicroCUDA reuse each image only while its allocation identity and write epoch remain valid. ZLUDA still prepares each tube launch. A 12-frame RX 9070 XT harness sequence matched the token-disabled layers byte-for-byte; MicroCUDA prepared the six tube images once and skipped the following 66 preparations, and repeated feature recreation remained byte-identical. - Tried, no measurable gain: interleaving the WMMAs of independent k32 steps. Skipping windows that lie entirely outside the image was also exact but below the measurement noise (at most 2.4% of the windows of every other tube launch). Not tried: true16 code generation (the half-register conversions would save about 3% of the instructions). diff --git a/docs/performance.md b/docs/performance.md index aae56ff..0e179ab 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -131,3 +131,14 @@ The 16-bit path never re-encoded its operands, so it was already faster than the The shim's whole-evaluation GPU time (`D4R_PROFILE`, no per-kernel synchronisation) went from 5.9–6.0 ms to 2.80 ms with the harness's 200 ms pause between frames, and to 2.39–2.43 ms with frames 0–4 ms apart (2000 frames at 4 ms spacing: median 2.39, maximum 2.56). The difference is the driver's clock governor: with the 200 ms pause the GPU is about 9% busy and runs at about 2790 MHz, with frames 0–4 ms apart it is 80–98% busy and runs at 3180–3300 MHz, as it does under a game's rendering load. Other GPU users on the desktop add occasional slower frames to any of these figures; the previous kernels show the same disturbances. The change is in [native-kernels.md](native-kernels.md#rdna4): the FP8 layers keep their activations as e4m3 bytes, and the downsample kernel is built as wave64 on gfx12. The output image is byte-identical before and after. + +DLSS 4.5 (L) on the same card and harness, with enc0 and dec0 as ZLUDA's compile (`sust_only`) and with their native token-lane parts ([native-kernels.md](native-kernels.md#families)). Median GPU time per frame in milliseconds from `D4R_CUDA_KERNEL_PROFILE` (frames 11–30, moving scene, LDR input unless noted); the final frames are byte-identical in every column pair. + +| | 1280×720 → 4K | | 1280×720 → 4K, HDR flags | | 1920×1080 → 4K | | accuracy set, 1280×720 → 4K | | +|---|---|---|---|---|---|---|---|---| +| Kernel | before | now | before | now | before | now | before | now | +| enc0 | 1.69 | 0.63 | 1.56 | 0.67 | 2.73 | 1.16 | 1.82 | 0.67 | +| dec0 | 1.26 | 0.54 | 1.26 | 0.54 | 2.28 | 1.01 | 1.31 | 0.59 | +| **all kernels** | **5.83** | **4.06** | **5.93** | **4.34** | **8.63** | **5.69** | **5.99** | **4.14** | + +L's post kernel (1.2–1.5 ms) is the largest kernel left; it has no matrix multiplies. The 1920×1080 runs were taken in one of the two clock states the driver's governor settles in at that load (about 13% apart), so compare within the pair, not across resolutions. diff --git a/kernels/build.sh b/kernels/build.sh index 0229fbb..0f141a2 100755 --- a/kernels/build.sh +++ b/kernels/build.sh @@ -112,10 +112,16 @@ if [[ "$WHAT" == all || "$WHAT" == tex || "$WHAT" == l ]]; then fi # L's unfolded variants keep their complete neural arithmetic. The folded M # GEMM heads/tails are incompatible; substitute only the surface stores. - specs+=(rrlite_dec0_4x4:sust_only) + # RDNA4 native FP8: enc0 after its feature staging (embedding, Swin block, patch merge: enc0l_tail.hip) and + # dec0 before its output code (patch expand, Swin block: dec0l_block.hip) are native token-lane code in the + # order of ZLUDA's lowering. Harness frames are byte-identical to the sust_only builds (RX 9070 XT, + # 1280x720 render: enc0 1.71 -> 0.63 ms, dec0 1.26 -> 0.54 ms). + enc0l=sust_only dec0l=sust_only + [[ "$ARCH" == gfx12* && "$FP8" == 1 ]] && enc0l=enc0l_tail dec0l=dec0l_block + specs+=("rrlite_dec0_4x4:$dec0l") for mv in mvhi mvlo; do for range in hdr ldr; do - specs+=("rrlite_enc0_4x4_${mv}_${range}:sust_only") + specs+=("rrlite_enc0_4x4_${mv}_${range}:$enc0l") for v in 3_1 3_2; do specs+=("rrlite_post_${v}_${mv}_${range}:sust_only:w64"); done done done diff --git a/kernels/m/swin_block.h b/kernels/m/swin_block.h index 7b1fea4..41f471d 100644 --- a/kernels/m/swin_block.h +++ b/kernels/m/swin_block.h @@ -9,32 +9,6 @@ #define SWIN_PRAGMA_IMPL(x) _Pragma(#x) #define SWIN_PRAGMA(x) SWIN_PRAGMA_IMPL(x) -template struct SwinLayout -{ - static constexpr int PREB = CIN ? 4 * C * CIN + 8 * C : 0; // patch-expand weights + bias first - static constexpr int NTE = CIN ? 4 * C / NW / 16 : 0; - static constexpr int NPL = C / 32, NT = C / 16, NT8 = C / 8, INNER = 32 * NW, NC = C / 8; - static constexpr int HS = 96 * C + 512; - static constexpr int G1 = PREB, E = PREB + 2 * C + NW * HS, BO = E, G2 = E + 2 * C, M0 = E + 4 * C; - static constexpr int B2 = M0 + 64 * C + 64; - __host__ __device__ static constexpr int head(int h) { return PREB + 2 * C + HS * h; } - __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; - // 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; - __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; - __host__ __device__ static constexpr int pm(int g) { return PM + 8 * C * g; } - static constexpr int PE = PM + (NPM / 32) * 8 * C; - __host__ __device__ static constexpr int pe(int w) { return PE + w * (CIN / 32) * 2 * NTE * 16; } - static constexpr int TOTAL = PE + (CIN ? NW * (CIN / 32) * 2 * NTE * 16 : 0); - 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. diff --git a/kernels/m/swin_common.h b/kernels/m/swin_common.h index 4b8c420..698d764 100644 --- a/kernels/m/swin_common.h +++ b/kernels/m/swin_common.h @@ -6,7 +6,11 @@ // 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 +// (D4R_NO_HIP_RUNTIME: built without the HIP headers, as part of a ZLUDA-compiled texture kernel; see +// ../tex/enc0l_tail.hip, which supplies the few runtime names used here) +#ifndef D4R_NO_HIP_RUNTIME #include +#endif #include #include "../common/wmma_layout.h" #ifdef D4R_ACCURACY @@ -249,9 +253,25 @@ __device__ __forceinline__ half_t f16(float v) return (half_t)v; } +// SWIN_APPROX_RECIP: the reciprocal square root and the reciprocal as the hardware approximations ZLUDA lowers +// rsqrt.approx / rcp.approx to (v_rsq_f32, v_rcp_f32) instead of correctly rounded divisions. The two differ in +// the last f32 bit for some arguments, which changes about one token in 20,000 after the f16 rounding: a kernel +// that has to reproduce ZLUDA's compile of the PTX byte for byte (../tex/enc0l_tail.hip) needs this form. __device__ __forceinline__ half_t rsqrt16(half_t s) { +#ifdef SWIN_APPROX_RECIP + return f16(__builtin_amdgcn_rsqf((float)(s + (half_t)(1.0f / 8192.0f)))); +#else return f16(1.0f / __builtin_sqrtf((float)(s + (half_t)(1.0f / 8192.0f)))); +#endif +} +__device__ __forceinline__ half_t recip16(half_t s) +{ +#ifdef SWIN_APPROX_RECIP + return f16(__builtin_amdgcn_rcpf((float)s)); +#else + return f16(1.0f / (float)s); +#endif } // natural channel c -> pair-order column, and inverse (per 32-channel group) @@ -455,7 +475,12 @@ __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 +#if defined(MIX_ONE) + // Without inline asm (a texture tail is compiled by ZLUDA's LLVM, which has no assembler): the same two + // instructions, selected from fma(P, one, f32(c)) rounded to f16. MIX_ONE must be a float 1.0 the compiler + // cannot fold: with a constant the fma becomes an f32 add followed by a second rounding. + return (hv2){(half_t)__builtin_fmaf(p0, MIX_ONE, (float)c[0]), (half_t)__builtin_fmaf(p1, MIX_ONE, (float)c[1])}; +#elif defined(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)); @@ -636,3 +661,31 @@ __device__ __forceinline__ void expand_weights(const uint8_t* w, wslot* w16, con #endif w16[idx] = v; } + +// ---------------------------------------------------------------- weight layout of a block +// (NVIDIA's weight offsets and the slot layout of the prepared images; used by swin_block.h and swin_tok.h) +template struct SwinLayout +{ + static constexpr int PREB = CIN ? 4 * C * CIN + 8 * C : 0; // patch-expand weights + bias first + static constexpr int NTE = CIN ? 4 * C / NW / 16 : 0; + static constexpr int NPL = C / 32, NT = C / 16, NT8 = C / 8, INNER = 32 * NW, NC = C / 8; + static constexpr int HS = 96 * C + 512; + static constexpr int G1 = PREB, E = PREB + 2 * C + NW * HS, BO = E, G2 = E + 2 * C, M0 = E + 4 * C; + static constexpr int B2 = M0 + 64 * C + 64; + __host__ __device__ static constexpr int head(int h) { return PREB + 2 * C + HS * h; } + __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; + // 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; + __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; + __host__ __device__ static constexpr int pm(int g) { return PM + 8 * C * g; } + static constexpr int PE = PM + (NPM / 32) * 8 * C; + __host__ __device__ static constexpr int pe(int w) { return PE + w * (CIN / 32) * 2 * NTE * 16; } + static constexpr int TOTAL = PE + (CIN ? NW * (CIN / 32) * 2 * NTE * 16 : 0); + static constexpr int NDESC = 2 * NW + 1 + 2 * NC + NPM / 32 + (CIN ? NW : 0); +}; diff --git a/kernels/m/swin_tok.h b/kernels/m/swin_tok.h index 4335be6..ace809f 100644 --- a/kernels/m/swin_tok.h +++ b/kernels/m/swin_tok.h @@ -14,6 +14,13 @@ // // 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. +// +// SWIN_TOK_NATURAL: a half-wave owns the channels whose bit 3 equals the half (dword 4 t + 2 half + g), and +// the two WMMAs of a k32 step hold K 0..15 and K 16..31 in order: a tile's two code dwords are one operand. +// The products are the same; what changes is which of them share a WMMA and in which K slot. That is the +// order of ZLUDA's lowering of NVIDIA's m16n8k32, so a kernel built this way (with K32_NO_MIX, which also adds +// C inside the first WMMA as that lowering does) accumulates in f32 exactly as ZLUDA's compile of the PTX +// does, including the rare sums that are not exact in f32 (../tex/enc0l_tail.hip). #pragma once enum : int { TOK_ID, TOK_PI, TOK_X, TOK_NU, TOK_PE }; @@ -39,7 +46,47 @@ struct TokConst // x layout: natural channel of row r (= 8 half + i) of tile t __device__ __forceinline__ int tok_xnat(int t, int r) { +#ifdef SWIN_TOK_NATURAL + return 16 * t + r; +#else return 16 * t + 8 * ((r >> 2) & 1) + 4 * (r >> 3) + (r & 3); +#endif +} +// logical K (of a 32-value chunk) in operand slot i of WMMA s of a k32 step +__device__ __forceinline__ int tok_kslot(int s, int i) +{ +#ifdef SWIN_TOK_NATURAL + return 16 * s + i; +#else + return kslot(s, i); +#endif +} +// The 4 codes (tile t, g) of half-wave hf within a token's C channel bytes are at tok_half(hf) + tok_word(t, g) +__device__ __forceinline__ int tok_half(int hf) +{ +#ifdef SWIN_TOK_NATURAL + return 8 * hf; +#else + return 4 * hf; +#endif +} +__device__ __forceinline__ int tok_word(int t, int g) +{ +#ifdef SWIN_TOK_NATURAL + return 16 * t + 4 * g; +#else + return 16 * t + 8 * g; +#endif +} +// first of the four values k (of a merged token's 4C inputs) in K slots 8 half + 4 rr .. + 3 of WMMA s of chunk kc +// (without the half-wave's own offset tok_half(hf)) +__device__ __forceinline__ int tok_merge_k(int kc, int s, int rr) +{ +#ifdef SWIN_TOK_NATURAL + return 32 * kc + 16 * s + 4 * rr; +#else + return 32 * kc + 16 * rr + 8 * s; +#endif } template struct TokLayout @@ -120,7 +167,7 @@ __device__ __forceinline__ void expand_tok_weights(const uint8_t* w, wslot* w16, 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)); + n = gperm(tok_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) @@ -138,7 +185,7 @@ __device__ __forceinline__ void expand_tok_weights(const uint8_t* w, wslot* w16, 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); + word |= (uint32_t)w[woff(base, g.Ks, g.Ns, 32 * kc + tok_kslot(s, 4 * j + b), n)] << (8 * b); v[j] = word; } w16[idx] = v; @@ -160,7 +207,7 @@ __device__ __forceinline__ void expand_tok_consts(const uint8_t* w, h8* c16, con if (g.kind == TC_X) off = 2 * gperm(tok_xnat(j, r)); else if (g.kind == TC_PI) - off = 2 * gperm(kslot(j, r)); + off = 2 * gperm(tok_kslot(j, r)); else if (g.kind == TC_NU) off = 2 * gperm(16 * j + r); else @@ -207,6 +254,16 @@ __device__ __forceinline__ T16 tok_t16(h8 c) return t; } +// the two operands (WMMA s = 0, 1) of a k32 step from the code words [tile][g] of the chunk's two tiles +__device__ __forceinline__ gop tok_k32_operand(const uint32_t (*c)[2], int s) +{ +#ifdef SWIN_TOK_NATURAL + return (u2v){c[s][0], c[s][1]}; +#else + return (u2v){c[0][s], c[1][s]}; +#endif +} + // 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) { @@ -238,129 +295,39 @@ template __device__ __forceinline__ half_t tok_rsq(const hv2 (&v)[NT][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 + else if constexpr (NPL == 2) ws[ww] = pp[0] + pp[1]; + else + ws[ww] = pp[0]; } +#ifdef SWIN_TOK_NATURAL + u[g] = ws[0] + ws[1]; + } + // the tree's next partner (dword ^ 1) is this lane's other g, the one after it (dword ^ 2) the other half + const hv2 a = u[0] + u[1]; + const hv2 s = a + tok_other(a); +#else const hv2 a = ws[0] + ws[1]; u[g] = a + tok_other(a); } const hv2 s = u[0] + u[1]; +#endif return rsqrt16(s[0] + s[1]); } -template -__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) +// The block from its input x0 to the e4m3 codes of its output, for the 16 tokens of one 4x4 window (lane l16 = +// token, both half-waves). v: x0 as f16 pairs, [tile][2 g + pair of the dword]; it is overwritten (x1). +// co: the output codes, [tile][g] (see tok_word). ch / wl: this lane's bases into the constant image (c16 + half) +// and the weight image (row l16 of a tile, this half's 8 K slots). x2 (optional): the output before its e4m3 +// encoding, as f16 pairs like v. +template +__device__ __forceinline__ void swin_tok_core(const h8* const ch, const gop* const wl, const int l16, hv2 (&v)[C / 16][4], + uint32_t (&co)[C / 16][2], hv2 (*x2)[4] = nullptr) { using K = TokLayout; 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(ew) + tok_dec2(inw[t][g]); - v[t][2 * g + 1] = tok_dec2(ew) + tok_dec2(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(inw[t][g]); - v[t][2 * g + 1] = tok_dec2(inw[t][g]); - } - } - // residual of the output projection: f16(x0 + b_o) hv2 res[NT][4]; #pragma unroll @@ -391,7 +358,7 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube 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]}; + hop[kc][s] = tok_k32_operand(&c1[2 * kc], s); } // ------------------------------------------------ attention per head: O codes, the K operands of the projection @@ -436,11 +403,25 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube } 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); + const half_t rinv = recip16(sum); hv2 pw[4]; #pragma unroll for (int k = 0; k < 4; ++k) pw[k] = wg[k] * (hv2){rinv, rinv}; +#ifdef SWIN_PV_ZLUDA_ORDER + // P V is the block's only product of two f16 operands: its f32 sum over the 16 keys is not exact, so the + // rounded result depends on which K slot holds which key. Here the slots follow ZLUDA's lowering of + // NVIDIA's m16n8k16 (dword j of half h = key pair (j & 1) * 4 + 2 h + (j >> 1)) instead of this + // layout's own (key pair 4 h + j), which costs two half-wave exchanges per operand. A kernel that has to + // reproduce ZLUDA's compile byte for byte needs it; the two orders differ in about one token in 30,000. + auto zk = [&](hv2 (&w)[4]) { + const bool upper = m_half() != 0; + const hv2 r0 = tok_other(upper ? w[0] : w[2]), r1 = tok_other(upper ? w[1] : w[3]); + const hv2 a = upper ? r0 : w[0], b = upper ? w[2] : r0, c = upper ? r1 : w[1], d = upper ? w[3] : r1; + w[0] = a, w[1] = b, w[2] = c, w[3] = d; + }; + zk(pw); +#endif const sop pb = __builtin_bit_cast(sop, pw); #pragma unroll for (int nt = 0; nt < 2; ++nt) @@ -449,6 +430,9 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube #pragma unroll for (int k = 0; k < 4; ++k) vp[k] = vt[nt].pair(k); +#ifdef SWIN_PV_ZLUDA_ORDER + zk(vp); +#endif 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]))}; } @@ -494,7 +478,7 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube 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]}; + fop[kc][s] = tok_k32_operand(&c2[2 * kc], s); } // ------------------------------------------------ MLP: fc1 transposed, its codes are fc2's operands @@ -559,21 +543,149 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube } // ------------------------------------------------ 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 (x2 != nullptr) + { +#pragma unroll + for (int t = 0; t < NT; ++t) +#pragma unroll + for (int k = 0; k < 4; ++k) + x2[t][k] = x[t].pair(k); + } +} + +template +__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; + 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 + tok_half(hf); +#pragma unroll + for (int j = 0; j < 4; ++j) + inw[2 * pl + (j >> 1)][j & 1] = input_word((const uint32_t*)(b + tok_word(j >> 1, j & 1))); + } + } + + 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 uint8_t* b = p.in + ((size_t)(kc * H2 + MY) * W2 + MX) * 32 + tok_half(hf); +#pragma unroll + for (int s = 0; s < 2; ++s) +#ifdef SWIN_TOK_NATURAL + pa[kc][s] = (u2v){input_word((const uint32_t*)(b + tok_word(s, 0))), input_word((const uint32_t*)(b + tok_word(s, 1)))}; +#else + pa[kc][s] = (u2v){input_word((const uint32_t*)(b + tok_word(0, s))), input_word((const uint32_t*)(b + tok_word(1, s)))}; +#endif + } + 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 + tok_word(t, g) + tok_half(hf)); + v[t][2 * g] = tok_dec2(ew) + tok_dec2(inw[t][g]); + v[t][2 * g + 1] = tok_dec2(ew) + tok_dec2(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(inw[t][g]); + v[t][2 * g + 1] = tok_dec2(inw[t][g]); + } + } + + uint32_t co[NT][2]; + swin_tok_core(ch, wl, l16, v, co); 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); + uint8_t* b = p.out + ((size_t)(pl * p.th + Y) * p.tw + X) * 32 + tok_half(hf); #pragma unroll for (int j = 0; j < 4; ++j) - b[2 * j] = co[2 * pl + (j >> 1)][j & 1]; + *(uint32_t*)(b + tok_word(j >> 1, j & 1)) = co[2 * pl + (j >> 1)][j & 1]; } } @@ -587,7 +699,7 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube 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]; + *(uint32_t*)(U + T * XB + tok_word(t, g) + tok_half(hf)) = co[t][g]; __syncthreads(); if (wv < NTP) { @@ -596,7 +708,7 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube // 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; + 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) @@ -607,7 +719,7 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube #pragma unroll for (int rr = 0; rr < 2; ++rr) { - const int k0 = 32 * kc + 16 * rr + 8 * s, q = k0 / C; + const int k0 = tok_merge_k(kc, s, rr), 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]}; diff --git a/kernels/tex/build_tex.sh b/kernels/tex/build_tex.sh index af71e24..990e3d3 100755 --- a/kernels/tex/build_tex.sh +++ b/kernels/tex/build_tex.sh @@ -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))" \ diff --git a/kernels/tex/dec0l_block.hip b/kernels/tex/dec0l_block.hip new file mode 100644 index 0000000..8509f83 --- /dev/null +++ b/kernels/tex/dec0l_block.hip @@ -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; +constexpr TokTables g_dl_tok = make_tok_tables(); +} // 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(ew) + tok_dec2(sw); + v[t][2 * g + 1] = tok_dec2(ew) + tok_dec2(sw); + } + + uint32_t co[NT][2]; + hv2 x2[NT][4]; + swin_tok_core(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; +} diff --git a/kernels/tex/enc0l_tail.hip b/kernels/tex/enc0l_tail.hip new file mode 100644 index 0000000..643965b --- /dev/null +++ b/kernels/tex/enc0l_tail.hip @@ -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; +// 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 g_el_tok = make_tok_tables(); +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(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; + } + } +} diff --git a/kernels/tex/make_ptx.py b/kernels/tex/make_ptx.py index 4bcd797..6b3d0af 100755 --- a/kernels/tex/make_ptx.py +++ b/kernels/tex/make_ptx.py @@ -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 diff --git a/kernels/tex/swin_tex.h b/kernels/tex/swin_tex.h new file mode 100644 index 0000000..eae6a8f --- /dev/null +++ b/kernels/tex/swin_tex.h @@ -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