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

Merge pull request #29 from woutermans/sr-native-layout

Browse Source 
Preset L: native token-lane enc0 and dec0 on RDNA4 native FP8
This commit is contained in:
Woutermans authored and GitHub committed 2026-10-08 02:08:06 +02:00
2 parents 0a12b6ce16 8813ea723c
commit 62738a01cc
11 files changed
+885 -151
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/native-kernels.md
+11
View File
@@ -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. **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. 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 ## 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 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. - 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. - `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. - `-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. - 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). - 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).
docs/performance.md
+11
View File
@@ -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 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. 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.
kernels/build.sh
+8 -2
View File
@@ -112,10 +112,16 @@ if [[ "$WHAT" == all || "$WHAT" == tex || "$WHAT" == l ]]; then
fi fi
# L's unfolded variants keep their complete neural arithmetic. The folded M # L's unfolded variants keep their complete neural arithmetic. The folded M
# GEMM heads/tails are incompatible; substitute only the surface stores. # 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 mv in mvhi mvlo; do
for range in hdr ldr; 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 for v in 3_1 3_2; do specs+=("rrlite_post_${v}_${mv}_${range}:sust_only:w64"); done
done done
done done
kernels/m/swin_block.h
-26
View File
@@ -9,32 +9,6 @@
#define SWIN_PRAGMA_IMPL(x) _Pragma(#x) #define SWIN_PRAGMA_IMPL(x) _Pragma(#x)
#define SWIN_PRAGMA(x) SWIN_PRAGMA_IMPL(x) #define SWIN_PRAGMA(x) SWIN_PRAGMA_IMPL(x)
template <int C, int NW, int NPM, int CIN = 0> 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 // 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 // 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. // constant); this placement puts those 16 reads on 16 distinct 8-byte bank pairs for C = 64, 96 and 128.
kernels/m/swin_common.h
+54 -1
View File
@@ -6,7 +6,11 @@
// D4R_FP8_WMMA (gfx12 layout only): the GEMMs whose operands are both e4m3 values run as native FP8 WMMAs on // 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. // e4m3 bytes (weights prepared as bytes, activations encoded when loaded); P V stays f16.
#pragma once #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 <hip/hip_runtime.h> #include <hip/hip_runtime.h>
#endif
#include <stdint.h> #include <stdint.h>
#include "../common/wmma_layout.h" #include "../common/wmma_layout.h"
#ifdef D4R_ACCURACY #ifdef D4R_ACCURACY
@@ -249,9 +253,25 @@ __device__ __forceinline__ half_t f16(float v)
return (half_t)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) __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)))); 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) // 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) // 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) __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; 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(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)); 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 #endif
w16[idx] = v; 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 <int C, int NW, int NPM, int CIN = 0> 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);
};
kernels/m/swin_tok.h
+233 -121
View File
@@ -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): // 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. // 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 #pragma once
enum : int { TOK_ID, TOK_PI, TOK_X, TOK_NU, TOK_PE }; 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 // x layout: natural channel of row r (= 8 half + i) of tile t
__device__ __forceinline__ int tok_xnat(int t, int r) __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); 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 <int C, int NH, int NPM, int CIN> struct TokLayout template <int C, int NH, int NPM, int CIN> 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; const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
int n = 16 * nt + r, base = g.base; int n = 16 * nt + r, base = g.base;
if (g.perm == TOK_PI) if (g.perm == TOK_PI)
n = gperm(kslot(nt, r)); n = gperm(tok_kslot(nt, r));
else if (g.perm == TOK_X) else if (g.perm == TOK_X)
n = gperm(tok_xnat(nt, r)); n = gperm(tok_xnat(nt, r));
else if (g.perm == TOK_NU || g.perm == TOK_PE) 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; uint32_t word = 0;
#pragma unroll #pragma unroll
for (int b = 0; b < 4; ++b) 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; v[j] = word;
} }
w16[idx] = v; 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) if (g.kind == TC_X)
off = 2 * gperm(tok_xnat(j, r)); off = 2 * gperm(tok_xnat(j, r));
else if (g.kind == TC_PI) 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) else if (g.kind == TC_NU)
off = 2 * gperm(16 * j + r); off = 2 * gperm(16 * j + r);
else else
@@ -207,6 +254,16 @@ __device__ __forceinline__ T16 tok_t16(h8 c)
return t; 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 // 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) __device__ __forceinline__ gop tok_codes(const T16& t)
{ {
@@ -238,129 +295,39 @@ template <int NT> __device__ __forceinline__ half_t tok_rsq(const hv2 (&v)[NT][4
ws[ww] = (pp[0] + pp[1]) + (pp[2] + pp[3]); ws[ww] = (pp[0] + pp[1]) + (pp[2] + pp[3]);
else if constexpr (NPL == 3) else if constexpr (NPL == 3)
ws[ww] = (pp[0] + pp[1]) + pp[2]; ws[ww] = (pp[0] + pp[1]) + pp[2];
else else if constexpr (NPL == 2)
ws[ww] = pp[0] + pp[1]; 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]; const hv2 a = ws[0] + ws[1];
u[g] = a + tok_other(a); u[g] = a + tok_other(a);
} }
const hv2 s = u[0] + u[1]; const hv2 s = u[0] + u[1];
#endif
return rsqrt16(s[0] + s[1]); return rsqrt16(s[0] + s[1]);
} }
template <int C, int NH, int NPM, bool TUBE, int CIN> // The block from its input x0 to the e4m3 codes of its output, for the 16 tokens of one 4x4 window (lane l16 =
__device__ __forceinline__ void swin_tok_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const h8* c16, // token, both half-waves). v: x0 as f16 pairs, [tile][2 g + pair of the dword]; it is overwritten (x1).
const int bx, const int by, const int gx) // 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 <int C, int NH, int NPM, int CIN>
__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<C, NH, NPM, CIN>; using K = TokLayout<C, NH, NPM, CIN>;
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8; 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)]; }; auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
if constexpr (TUBE)
{
if (tp->in_flags != nullptr)
{
if (lane <= 3 && wv == 0)
{
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
{
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
// bounded spin: a missing producer must never hang the GPU
for (int it = 0; it < (1 << 22); ++it)
{
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
break;
__builtin_amdgcn_s_sleep(1);
}
}
}
__syncthreads();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
}
}
// ------------------------------------------------ block input x0 (encoders: the fp8 input; decoders: skip)
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
const int Xm = mirror(X, p.tw), Ym = mirror(Y, p.th);
uint32_t inw[NT][2];
{
const uint8_t* const src = CIN ? p.p32 : p.in;
#pragma unroll
for (int pl = 0; pl < NKC; ++pl)
{
const uint8_t* b = src + ((size_t)(pl * p.th + Ym) * p.tw + Xm) * 32 + 4 * hf;
#pragma unroll
for (int j = 0; j < 4; ++j)
inw[2 * pl + (j >> 1)][j & 1] = input_word((const uint32_t*)(b + 8 * j));
}
}
hv2 v[NT][4]; // x0, then x1, as f16 pairs: [tile][2 g + pair of the dword]
if constexpr (CIN > 0)
{
// patch expand: lanes = the block's 16 low-res tokens, wave = quadrant; the codes of q8(expand) go to
// their child tokens' rows of a byte image, where the windows' waves pick them up
__shared__ __attribute__((aligned(16))) uint8_t E8[64 * C];
const int W2 = p.tw / 2, H2 = p.th / 2;
const int MX = mirror((8 * bx - p.sx) / 2 + (l16 & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (l16 >> 2), H2);
gop pa[CIN / 32][2];
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
{
const uint32_t* b = (const uint32_t*)(p.in + ((size_t)(kc * H2 + MY) * W2 + MX) * 32 + 4 * hf);
#pragma unroll
for (int s = 0; s < 2; ++s)
pa[kc][s] = (u2v){input_word(b + 2 * s), input_word(b + 4 + 2 * s)};
}
const int tx = 2 * (l16 & 3) + (wv & 1), ty = 2 * (l16 >> 2) + (wv >> 1);
const int Tc = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
const int pew = K::pe(0) + wv * (CIN * C / 16);
#pragma unroll
for (int ct = 0; ct < NT; ++ct)
{
T16 e = tok_t16(ch[K::CPEB + 2 * (NT * wv + ct)]);
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
k32(e, wt(pew, NT, kc, 0, ct), pa[kc][0], wt(pew, NT, kc, 1, ct), pa[kc][1]);
*(u2v*)(E8 + Tc * C + 16 * ct + 8 * hf) = tok_codes(e);
}
__syncthreads();
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
{
const uint32_t ew = *(const uint32_t*)(E8 + T * C + 16 * t + 8 * g + 4 * hf);
v[t][2 * g] = tok_dec2<false>(ew) + tok_dec2<false>(inw[t][g]);
v[t][2 * g + 1] = tok_dec2<true>(ew) + tok_dec2<true>(inw[t][g]);
}
}
else
{
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
{
v[t][2 * g] = tok_dec2<false>(inw[t][g]);
v[t][2 * g + 1] = tok_dec2<true>(inw[t][g]);
}
}
// residual of the output projection: f16(x0 + b_o) // residual of the output projection: f16(x0 + b_o)
hv2 res[NT][4]; hv2 res[NT][4];
#pragma unroll #pragma unroll
@@ -391,7 +358,7 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube
for (int kc = 0; kc < NKC; ++kc) for (int kc = 0; kc < NKC; ++kc)
#pragma unroll #pragma unroll
for (int s = 0; s < 2; ++s) 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 // ------------------------------------------------ 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 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 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]; hv2 pw[4];
#pragma unroll #pragma unroll
for (int k = 0; k < 4; ++k) for (int k = 0; k < 4; ++k)
pw[k] = wg[k] * (hv2){rinv, rinv}; 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); const sop pb = __builtin_bit_cast(sop, pw);
#pragma unroll #pragma unroll
for (int nt = 0; nt < 2; ++nt) for (int nt = 0; nt < 2; ++nt)
@@ -449,6 +430,9 @@ __device__ __forceinline__ void swin_tok_block(const CommonParams& p, const Tube
#pragma unroll #pragma unroll
for (int k = 0; k < 4; ++k) for (int k = 0; k < 4; ++k)
vp[k] = vt[nt].pair(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)); 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]))}; 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) for (int kc = 0; kc < NKC; ++kc)
#pragma unroll #pragma unroll
for (int s = 0; s < 2; ++s) 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 // ------------------------------------------------ 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 // ------------------------------------------------ output codes
uint32_t co[NT][2];
#pragma unroll #pragma unroll
for (int t = 0; t < NT; ++t) for (int t = 0; t < NT; ++t)
#pragma unroll #pragma unroll
for (int g = 0; g < 2; ++g) for (int g = 0; g < 2; ++g)
co[t][g] = codes4(x[t].pair(2 * g), x[t].pair(2 * g + 1)); 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 <int C, int NH, int NPM, bool TUBE, int CIN>
__device__ __forceinline__ void swin_tok_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const h8* c16,
const int bx, const int by, const int gx)
{
using K = TokLayout<C, NH, NPM, CIN>;
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8;
static_assert(C % 32 == 0, "32-channel groups");
const int lane = threadIdx.x, wv = threadIdx.z;
const int l16 = lane & 15, hf = lane >> 4;
const int T = 16 * wv + l16; // this lane's token: window wv, position l16
// The lane's bases into the two images, hidden from constant folding: every tile and constant is then
// one load at a fixed offset from a register (folded, each load rebuilt the image's address first).
uintptr_t chv = (uintptr_t)(c16 + hf), wlv = (uintptr_t)((const gop*)w16 + 2 * l16 + hf);
asm("" : "+v"(chv));
asm("" : "+v"(wlv));
const h8* const ch = (const h8*)chv;
// row l16 of a weight tile: this half's 8 K slots
const gop* const wl = (const gop*)wlv;
auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
if constexpr (TUBE)
{
if (tp->in_flags != nullptr)
{
if (lane <= 3 && wv == 0)
{
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
{
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
// bounded spin: a missing producer must never hang the GPU
for (int it = 0; it < (1 << 22); ++it)
{
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
break;
__builtin_amdgcn_s_sleep(1);
}
}
}
__syncthreads();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
}
}
// ------------------------------------------------ block input x0 (encoders: the fp8 input; decoders: skip)
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
const int Xm = mirror(X, p.tw), Ym = mirror(Y, p.th);
uint32_t inw[NT][2];
{
const uint8_t* const src = CIN ? p.p32 : p.in;
#pragma unroll
for (int pl = 0; pl < NKC; ++pl)
{
const uint8_t* b = src + ((size_t)(pl * p.th + Ym) * p.tw + Xm) * 32 + 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<false>(ew) + tok_dec2<false>(inw[t][g]);
v[t][2 * g + 1] = tok_dec2<true>(ew) + tok_dec2<true>(inw[t][g]);
}
}
else
{
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
{
v[t][2 * g] = tok_dec2<false>(inw[t][g]);
v[t][2 * g + 1] = tok_dec2<true>(inw[t][g]);
}
}
uint32_t co[NT][2];
swin_tok_core<C, NH, NPM, CIN>(ch, wl, l16, v, co);
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th) if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
{ {
#pragma unroll #pragma unroll
for (int pl = 0; pl < NKC; ++pl) 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 #pragma unroll
for (int j = 0; j < 4; ++j) 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) for (int t = 0; t < NT; ++t)
#pragma unroll #pragma unroll
for (int g = 0; g < 2; ++g) 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(); __syncthreads();
if (wv < NTP) 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 // 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 // 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. // 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]; gop as[C / 8][2];
#pragma unroll #pragma unroll
for (int kc = 0; kc < C / 8; ++kc) 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 #pragma unroll
for (int rr = 0; rr < 2; ++rr) 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); run4[rr] = *(const uint32_t*)(ub + ((q & 1) + 4 * (q >> 1)) * XB + k0 % C);
} }
as[kc][s] = (u2v){run4[0], run4[1]}; as[kc][s] = (u2v){run4[0], run4[1]};
kernels/tex/build_tex.sh
+1 -1
View File
@@ -27,7 +27,7 @@ read -r -a FLAGS <<< "${D4R_TEX_CFLAGS:-}"
"$LLVM/llvm-dis" "$B/raw.bc" -o "$B/raw.ll" "$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' \ 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" - -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 # 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 \ 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))" \ D4R_ZLUDA_WMMA_FP8_NATIVE="$FP8" D4R_ZLUDA_IGNORE_DENORMAL="$((1 - ACCURACY))" \
kernels/tex/dec0l_block.hip
+142
View File
@@ -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;
}
kernels/tex/enc0l_tail.hip
+193
View File
@@ -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;
}
}
}
kernels/tex/make_ptx.py
+182
View File
@@ -239,6 +239,176 @@ def replace_rz_round(lines):
return out, n 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" ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
# the other flag combinations of the kernels whose only native part is the surface-store rewrite # 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|" 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] return KERNELS[name]
if ENC0_RE.match(name): if ENC0_RE.match(name):
return enc0_variant(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): if SUST_ONLY_RE.match(name):
return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_output")) return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_output"))
sys.exit(f"make_ptx: no recipe for {name}") sys.exit(f"make_ptx: no recipe for {name}")
@@ -334,6 +508,10 @@ def main():
if lines[cut - 1].strip() != k["cut_check"]: if lines[cut - 1].strip() != k["cut_check"]:
sys.exit(f"line {cut} is {lines[cut - 1]!r}, expected {k['cut_check']!r}") sys.exit(f"line {cut} is {lines[cut - 1]!r}, expected {k['cut_check']!r}")
head, tail = lines[:cut], k["tail"].split("\n") 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: elif "delete" in k:
a, b = k["delete"] a, b = k["delete"]
if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]: if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]:
@@ -343,6 +521,10 @@ def main():
head, tail = lines, [] head, tail = lines, []
entry = next(i for i, l in enumerate(head) if l.startswith(".visible .entry")) 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}) 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"): if k.get("rz_round"):
body, nrz = replace_rz_round(body) body, nrz = replace_rz_round(body)
# temporaries for the rewrite, declared after the kernel's own .reg lines # temporaries for the rewrite, declared after the kernel's own .reg lines
kernels/tex/swin_tex.h
+50
View File
@@ -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