d4r: official NVIDIA DLSS on AMD GPUs under Linux and Proton

kernels/build.sh

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Builds the native RDNA3 replacements for DLSS kernels into one directory that ZLUDA serves them from
+# (D4R_ZLUDA_NATIVE_DIR, or NativeKernelDirFast in d4r.ini). See docs/native-kernels.md.
+#
+# usage: kernels/build.sh [all|k|m|tex] [OUT_DIR]      (default: all, kernels/out/native)
+#
+#   D4R_ROCM_DIR    ROCm installation with clang and the HIP device libraries (default /opt/rocm)
+#   D4R_GPU_ARCH    target GPU (default gfx1101; the kernels need RDNA3 / gfx11 WMMA)
+# Texture kernels (tex) are NVIDIA's PTX with parts replaced, so they additionally need:
+#   D4R_DLSS_DLL    nvngx_dlss.dll 310.7 (its PTX is extracted into kernels/extracted/, never committed)
+#   D4R_ZLUDA_BUILD a ZLUDA build with patches/zluda applied (directory holding libcuda.so)
+set -euo pipefail
+
+HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+WHAT="${1:-all}"
+OUT="$(realpath -m "${2:-$HERE/out/native}")"
+ROCM="${D4R_ROCM_DIR:-/opt/rocm}"
+ARCH="${D4R_GPU_ARCH:-gfx1101}"
+CLANG="$ROCM/lib/llvm/bin/clang++"
+[[ -x "$CLANG" ]] || { echo "clang++ not found in $ROCM/lib/llvm/bin (set D4R_ROCM_DIR)" >&2; exit 2; }
+mkdir -p "$OUT"
+
+# HIP source -> raw code object (the kernel name inside matches the DLSS kernel it replaces)
+build_hip() {
+    local src="$1" name
+    name="$(basename "$src" .hip)"
+    local tmp
+    tmp="$(mktemp -d)"
+    (cd "$tmp" && "$CLANG" -x hip --offload-arch="$ARCH" --offload-device-only -O3 -I"$ROCM/include" \
+        --rocm-path="$ROCM" --rocm-device-lib-path="$ROCM/amdgcn/bitcode" -I"$(dirname "$src")" \
+        -o "$OUT/$name.hsaco" "$src" -save-temps=cwd --no-gpu-bundle-output)
+    printf '%-52s %s\n' "$name" "$(grep -hE '^; (NumVgprs|ScratchSize|Occupancy)' "$tmp"/*.s | tail -3 | paste -sd' ')"
+    rm -rf "$tmp"
+}
+
+if [[ "$WHAT" == all || "$WHAT" == k ]]; then
+    echo "== DLSS 4 (preset K) transformer layers"
+    for src in "$HERE"/k/dltss_pwin_*.hip; do build_hip "$src"; done
+fi
+if [[ "$WHAT" == all || "$WHAT" == m ]]; then
+    echo "== DLSS 4.5 (preset M) Swin layers"
+    for src in "$HERE"/m/rrlite_*.hip; do build_hip "$src"; done
+fi
+if [[ "$WHAT" == all || "$WHAT" == tex ]]; then
+    echo "== texture kernels (NVIDIA PTX with native parts)"
+    : "${D4R_DLSS_DLL:?set D4R_DLSS_DLL to nvngx_dlss.dll (310.7)}"
+    : "${D4R_ZLUDA_BUILD:?set D4R_ZLUDA_BUILD to a ZLUDA build with patches/zluda}"
+    PTX_DIR="$HERE/extracted/ptx"
+    python3 "$HERE/tools/extract_dlss_ptx.py" "$D4R_DLSS_DLL" "$PTX_DIR"
+    PTXLOAD="$HERE/out/tools/ptxload"
+    mkdir -p "$(dirname "$PTXLOAD")"
+    cc -O2 -o "$PTXLOAD" "$HERE/tools/ptxload.c" -ldl
+    # kernel : source of its native part
+    for pair in hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel:sust_only rrlite_post_3_2_mvhi_hdr_folded:sust_only \
+        rrlite_downsample_kernel_static_hdr:sust_only rrlite_enc0_4x4_mvhi_hdr_folded:enc0_tail \
+        rrlite_dec0_4x4_folded:dec0_head; do
+        D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" D4R_DLSS_PTX_DIR="$PTX_DIR" D4R_PTXLOAD="$PTXLOAD" \
+            "$HERE/tex/build_tex.sh" "${pair%%:*}" "${pair#*:}" "$OUT"
+    done
+fi
+echo "native kernels in $OUT"

kernels/k/dltss_pwin_dec0_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_pos.h"
+POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)

kernels/k/dltss_pwin_dec1_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)

kernels/k/dltss_pwin_dec2_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_DECODER(dltss_pwin_dec2_layer, 2, 64, 96)

kernels/k/dltss_pwin_dec3_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_DECODER(dltss_pwin_dec3_layer, 4, 96, 128, 2)

kernels/k/dltss_pwin_dec4_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 4)

kernels/k/dltss_pwin_dec5_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 4)

kernels/k/dltss_pwin_enc0_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_pos.h"
+POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)

kernels/k/dltss_pwin_enc1_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)

kernels/k/dltss_pwin_enc2_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)

kernels/k/dltss_pwin_enc3_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_ENCODER(dltss_pwin_enc3_layer, 4, 96, 128, 2)

kernels/k/dltss_pwin_enc4_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 4)

kernels/k/pwin_common.h

@@ -0,0 +1,188 @@
+// Native RDNA3 (gfx11 wave32 WMMA) kernels for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
+// Semantics: kernels/tools/pwin_model.py (verified against the PTX interpreter, kernels/tools/ptxsim.py).
+//
+// Layout conventions (transposed WMMA, as in the preset M kernels): a GEMM Y = X W is computed as
+// Y^T = W^T X^T with the weights as the WMMA A operand (lane l: output channel l & 15, 16 K values)
+// and the activations as the B operand (lane l: token l & 15, 16 K values). The f32 result then has
+// lane l = token l & 15 and VGPR i = output channel 2i + (l >> 4) of the 16-channel tile ("D^T").
+#pragma once
+#include <hip/hip_runtime.h>
+#include <stdint.h>
+
+#pragma clang fp contract(off)
+
+typedef _Float16 half_t;
+typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
+typedef _Float16 h16 __attribute__((ext_vector_type(16)));
+typedef float f8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u4v __attribute__((ext_vector_type(4)));
+
+__device__ __forceinline__ uint32_t lane_id()
+{
+    return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
+}
+
+__device__ __forceinline__ f8v wmma(u8v a, u8v b, f8v c)
+{
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
+}
+
+// one k16 step with the f16 accumulator of NVIDIA's f16 wmma (rounded after the step)
+// PWIN_F32ACC: keep the accumulator in f32 through the chain (rounded to f16 where the values are used)
+__device__ __forceinline__ f8v mma16(u8v a, u8v b, f8v c)
+{
+    f8v d = wmma(a, b, c);
+#ifndef PWIN_F32ACC
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        d[i] = (float)(half_t)d[i];
+#endif
+    return d;
+}
+
+__device__ __forceinline__ f8v splat8(float v)
+{
+    return (f8v){v, v, v, v, v, v, v, v};
+}
+
+__device__ __forceinline__ uint32_t other_half(uint32_t v)
+{
+    return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
+}
+
+__device__ __forceinline__ uint32_t pack2(half_t lo, half_t hi)
+{
+    return __builtin_bit_cast(uint32_t, (hv2){lo, hi});
+}
+
+__device__ __forceinline__ half_t lo16(uint32_t v)
+{
+    return __builtin_bit_cast(hv2, v)[0];
+}
+__device__ __forceinline__ half_t hi16(uint32_t v)
+{
+    return __builtin_bit_cast(hv2, v)[1];
+}
+
+// D^T f16 values (own[i] = channel 2i + hf of token l & 15) -> B operand of the same token (all 16
+// channels in order): the other parity comes from lane l ^ 16
+__device__ __forceinline__ u8v operand_from_dt(const half_t own[8])
+{
+    const uint32_t hf = lane_id() >> 4;
+    uint32_t mine[4], theirs[4];
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        mine[j] = pack2(own[2 * j], own[2 * j + 1]);
+        theirs[j] = other_half(mine[j]);
+    }
+#ifndef PWIN_NO_PERM
+    // v_perm_b32(theirs, mine): bytes 0-3 = mine, 4-7 = theirs. Even j: the low halves, odd j: the high
+    // halves, in the order (mine, theirs) for hf = 0 and (theirs, mine) for hf = 1.
+    const uint32_t sel_lo = hf ? 0x01000504u : 0x05040100u, sel_hi = hf ? 0x03020706u : 0x07060302u;
+    u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+        r[j] = __builtin_amdgcn_perm(theirs[j >> 1], mine[j >> 1], (j & 1) ? sel_hi : sel_lo);
+    return r;
+#else
+    u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+    {
+        const half_t a = (j & 1) ? hi16(mine[j >> 1]) : lo16(mine[j >> 1]);
+        const half_t b = (j & 1) ? hi16(theirs[j >> 1]) : lo16(theirs[j >> 1]);
+        r[j] = hf ? pack2(b, a) : pack2(a, b);
+    }
+    return r;
+#endif
+}
+
+__device__ __forceinline__ u8v operand_from_f8(f8v d)
+{
+    half_t h[8];
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        h[i] = (half_t)d[i];
+    return operand_from_dt(h);
+}
+
+// element c (0..15) of a B operand (16 halves)
+__device__ __forceinline__ half_t op_get(const u8v& v, int c)
+{
+    return (c & 1) ? hi16(v[c >> 1]) : lo16(v[c >> 1]);
+}
+
+__device__ __forceinline__ u8v lds_row16(const half_t* p)
+{
+    const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
+    return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+}
+
+// Global activation traffic (tokens in, rows out): PWIN_NT_LOAD / PWIN_NT_STORE mark it non-temporal so
+// it does not evict the weight images from the caches. Same data either way.
+__device__ __forceinline__ u8v gload_row16(const half_t* p)
+{
+#ifdef PWIN_NT_LOAD
+    const u4v a = __builtin_nontemporal_load((const u4v*)p), b = __builtin_nontemporal_load((const u4v*)(p + 8));
+#else
+    const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
+#endif
+    return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+}
+__device__ __forceinline__ void gstore16(half_t* p, u4v v)
+{
+#ifdef PWIN_NT_STORE
+    __builtin_nontemporal_store(v, (u4v*)p);
+#else
+    *(u4v*)p = v;
+#endif
+}
+__device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
+{
+    gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
+    gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
+}
+
+__device__ __forceinline__ void block_sync()
+{
+    __syncthreads();
+}
+
+__device__ __forceinline__ int mirror(int v, int n)
+{
+    v = v < 0 ? -v : v;
+    return v < 2 * n - 2 - v ? v : 2 * n - 2 - v;
+}
+
+// ------------------------------------------------ pre-swizzled f16 weight fragments (NVIDIA)
+// logical B[k][n] of one 16x16 fragment (512 bytes)
+__device__ __forceinline__ int frag_offset(int k, int n)
+{
+    return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1);
+}
+
+// A-operand image of a K x N weight: tile (kt, nt) = 16 lanes x 8 dwords (lane = output channel)
+struct WDesc
+{
+    int src;       // byte offset of the fragments in the weight buffer
+    int kt_stride; // bytes between k tiles
+    int nt_stride; // bytes between n tiles
+    int KT, NT;    // tiles
+    int dst;       // u8v index of the image in the prep buffer
+};
+
+// weight operand tile `tile` of an A-operand image (16 lanes x 32 bytes): lanes l and l + 16 need the same
+// row, so each half-wave loads one 16-byte half and the halves are swapped across (halves L2 traffic)
+__device__ __forceinline__ u8v wtile(const u8v* __restrict__ img, int tile)
+{
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const u4v mine = ((const u4v*)&img[tile * 16 + m])[hf];
+    u4v other;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+        other[j] = other_half(mine[j]);
+    const u4v lo = hf ? other : mine, hi = hf ? mine : other;
+    return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
+}

kernels/k/pwin_layer.h

@@ -0,0 +1,837 @@
+// Native DltssPaddedWinLayer (DLSS 4, preset K) template: 8x8 windows, H heads x 32, C channels.
+// One block per window, 4 waves (one 16-token tile each). Semantics: kernels/tools/pwin_model.py (swin_core,
+// pwin_encoder); weight layout: pwin_model.Layout.
+#pragma once
+#include "pwin_common.h"
+#ifndef PWIN_VGPR
+#define PWIN_VGPR
+#endif
+
+struct PwinParams
+{
+    int W, H;              // 0: token grid
+    const uint8_t* in;     // 8
+    const uint8_t* skip;   // 16 (decoders)
+    uint8_t* out24;        // 24: encoders: merged output; decoders: output
+    uint8_t* out32;        // 32: encoders: full-resolution output
+    uint64_t pad40, pad48; // 40
+    int sx, sy;            // 56
+    const uint8_t* w;      // 64
+    uint8_t rest[176 - 72];
+};
+static_assert(sizeof(PwinParams) == 176, "param block");
+
+template <int H_, int C_, int COUT_, bool POS_ = false> struct PwinLayout
+{
+    static constexpr bool POS = POS_; // position-only attention (trait flag 2): the table is P itself
+    static constexpr int H = H_, C = C_, COUT = COUT_, KT = C / 16, NCH = C / 32, NMLP = C / 8;
+    static constexpr int G1 = 0, QKV = 2 * C, HEAD = 192 * C, BIAS = QKV + HEAD * H, BO = BIAS + 8192 * H,
+                         WO = BO + 2 * C, G2 = WO + 64 * C * H, B2 = G2 + 2 * C, W1 = B2 + 2 * C, B1 = W1 + 8 * C * C,
+                         W2 = B1 + 8 * C, PM = W2 + 8 * C * C;
+    // patch merge: NVIDIA's warps own COUT / H channels each, allocated as whole 16-column tiles
+    static constexpr int NPW = COUT / H, NPA = (NPW + 15) / 16 * 16, PMT = H * NPA / 16, PMB = PM + 8 * C * NPA * H;
+    // prep image tiles (16 lanes of u8v each)
+    static constexpr int T_QKV = 0;                           // [h][j][kt KT][nt 2]
+    static constexpr int T_WO = T_QKV + H * 3 * KT * 2;        // [ks 2H][nt KT]
+    static constexpr int T_W1 = T_WO + 2 * H * KT;             // [c][kt KT][nt 2]
+    static constexpr int T_W2 = T_W1 + NMLP * KT * 2;          // [c][ks 2][nt KT]
+    static constexpr int T_PM = T_W2 + NMLP * 2 * KT;          // [ks C/4][nt COUT/16]
+    static constexpr int T_END = T_PM + (C / 4) * PMT;
+    static constexpr int NBIAS = H * 4 * 4 * 32;               // u4v per lane
+    static constexpr int PREP_ITEMS = T_END * 16 + NBIAS;
+
+    __device__ static int src_of(int tile)
+    {
+        if (tile < T_WO)
+        {
+            const int nt = tile & 1, kt = (tile >> 1) % KT, j = (tile / (2 * KT)) % 3, h = tile / (6 * KT);
+            return QKV + HEAD * h + 2048 * (3 * (kt >> 1) + j) + 512 * (kt & 1) + 1024 * nt;
+        }
+        if (tile < T_W1)
+        {
+            const int t = tile - T_WO, nt = t % KT, ks = t / KT;
+            return WO + 64 * C * (ks >> 1) + 512 * (ks & 1) + 1024 * nt;
+        }
+        if (tile < T_W2)
+        {
+            const int t = tile - T_W1, nt = t & 1, kt = (t >> 1) % KT, c = t / (2 * KT);
+            return W1 + 64 * C * c + 512 * kt + 32 * C * nt;
+        }
+        if (tile < T_PM)
+        {
+            const int t = tile - T_W2, nt = t % KT, ks = (t / KT) & 1, c = t / (2 * KT);
+            return W2 + 64 * C * c + 512 * ks + 1024 * nt;
+        }
+        const int t = tile - T_PM, nt = t % PMT, ks = t / PMT;
+        return PM + 512 * ks + 128 * C * nt; // allocated tiles are contiguous
+    }
+};
+
+__device__ __forceinline__ half_t wload(const uint8_t* w, int byte)
+{
+    return *(const half_t*)(w + byte);
+}
+
+// channels 16 nt + 2i + hf (i = 0..7) of an f16 vector at byte offset `byte` (16 nt already added)
+__device__ __forceinline__ void dvec8(const uint8_t* w, int byte, uint32_t hf, half_t out[8])
+{
+    const u4v a = *(const u4v*)(w + byte), b = *(const u4v*)(w + byte + 16);
+    const u8v v = (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        out[i] = hf ? hi16(v[i]) : lo16(v[i]);
+}
+
+__device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
+{
+    u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+        r[j] = pack2(wload(w, src + frag_offset(2 * j, n)), wload(w, src + frag_offset(2 * j + 1, n)));
+    return r;
+}
+
+template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx)
+{
+    const uint8_t* w = p.w;
+    if (idx < L::T_END * 16)
+    {
+        img[idx] = img_tile(w, L::src_of(idx >> 4), idx & 15);
+        return;
+    }
+    const int b = idx - L::T_END * 16;
+    if (b >= L::NBIAS)
+        return;
+    if constexpr (L::POS)
+    {
+        // B operand of P V: lane = query 16 qt + l16, 16 keys 16 kt .. 16 kt + 15 (u8v per lane; 16 lanes)
+        if (b >= L::H * 4 * 4 * 16)
+            return;
+        const int l16 = b & 15, kt = (b >> 4) & 3, qt = (b >> 6) & 3, h = b >> 8, r = l16;
+        u8v v;
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+        {
+            half_t e[2];
+#pragma unroll
+            for (int u = 0; u < 2; ++u)
+            {
+                const int c = 2 * j + u;
+                const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
+                e[u] = wload(w, L::BIAS + 8192 * h + off);
+            }
+            v[j] = pack2(e[0], e[1]);
+        }
+        ((u8v*)bias)[b] = v;
+        return;
+    }
+    const int lane = b & 31, kt = (b >> 5) & 3, qt = (b >> 7) & 3, h = b >> 9;
+    const int q = 16 * qt + (lane & 15);
+    half_t v[8];
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+    {
+        const int key = 16 * kt + 2 * i + (lane >> 4), r = q & 15, c = key & 15;
+        const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
+        v[i] = wload(w, L::BIAS + 8192 * h + off);
+    }
+    bias[b] = (u4v){pack2(v[0], v[1]), pack2(v[2], v[3]), pack2(v[4], v[5]), pack2(v[6], v[7])};
+}
+
+template <class L> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
+{
+    pwin_prep_at<L>(p, img, bias, blockIdx.x * 128 + threadIdx.x);
+}
+
+__device__ __forceinline__ half_t hmin(half_t a, half_t b) { return a < b ? a : b; }
+__device__ __forceinline__ half_t hmax(half_t a, half_t b) { return a > b ? a : b; }
+
+__device__ __forceinline__ half_t softmax_e(half_t s)
+{
+    const half_t hi = __builtin_bit_cast(half_t, (uint16_t)0x4D69), lo = __builtin_bit_cast(half_t, (uint16_t)0xCD69);
+    const half_t a = __builtin_bit_cast(half_t, (uint16_t)0x91D5), c = __builtin_bit_cast(half_t, (uint16_t)0xC50D);
+    const half_t t = hmax(hmin(s, hi), lo);
+    const half_t u = t * t;
+    const half_t v = __builtin_fmaf16(u, a, (half_t)1.0f);
+    const half_t x = __builtin_fmaf16(t, v, c);
+    return (half_t)__builtin_amdgcn_exp2f((float)x);
+}
+
+__device__ __forceinline__ half_t gelu(half_t x)
+{
+    const half_t k1 = (half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
+    const half_t c = hmin(hmax(x, (half_t)-2.0f), (half_t)2.0f);
+    const half_t ac = c < (half_t)0.0f ? -c : c;
+    const half_t u = k1 - k2 * ac;
+    return x * ((half_t)0.5f + c * u);
+}
+
+// gelu of 8 accumulator values as 4 packed f16 pairs (the same f16 operations as gelu(), element-wise)
+__device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
+{
+    const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        const hv2 x = (hv2){(half_t)d[2 * j], (half_t)d[2 * j + 1]};
+        const hv2 c = __builtin_elementwise_min(__builtin_elementwise_max(x, (hv2)(half_t)-2.0f), (hv2)(half_t)2.0f);
+        const hv2 u = k1 - k2 * __builtin_elementwise_abs(c);
+        const hv2 r = x * ((hv2)(half_t)0.5f + c * u);
+        g[2 * j] = r[0];
+        g[2 * j + 1] = r[1];
+    }
+}
+
+// store the D^T tile of token m as 16 contiguous channels (lanes of one half-wave write)
+__device__ __forceinline__ void store_row16(half_t* dst, const u8v& v)
+{
+    *(u4v*)dst = (u4v){v[0], v[1], v[2], v[3]};
+    *(u4v*)(dst + 8) = (u4v){v[4], v[5], v[6], v[7]};
+}
+
+#ifdef PWIN_STAMP
+// per-block stage timestamps of wave 0 (s_memtime), slot k
+__device__ uint64_t g_stamp[16384][16];
+#define STAMP(k)                                                                                                   \
+    do                                                                                                             \
+    {                                                                                                              \
+        if (threadIdx.z == 0 && lane_id() == 0)                                                                    \
+        {                                                                                                          \
+            const int bid = blockIdx.y * gridDim.x + blockIdx.x;                                                   \
+            if (bid < 16384)                                                                                       \
+                g_stamp[bid][k] = __builtin_readcyclecounter();                                                    \
+        }                                                                                                          \
+    } while (0)
+#else
+#define STAMP(k) do {} while (0)
+#endif
+#ifdef PWIN_DBG
+__device__ float g_dbg[12][64][160];
+#define DBG_ON (blockIdx.x == DBX && blockIdx.y == DBY)
+#endif
+// The Swin core for one window: x0 rows of this wave's 16 tokens are in act (f16, C per token).
+// On return act holds the block output y. hb: normalised input rows; K/V of one head at a time in kl/vt.
+template <class L>
+__device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
+                          half_t (*act)[L::C], half_t (*hb)[L::C], half_t (*kl)[32], half_t (*vt)[64])
+{
+    constexpr int C = L::C, KT = L::KT, H = L::H;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z;
+    const uint8_t* w = p.w;
+    half_t* arow = act[16 * wv + m];
+    half_t* hrow = hb[16 * wv + m];
+
+    STAMP(1);
+    // ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb
+    float ss = 0.0f;
+    for (int kt = 0; kt < KT; ++kt)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+#pragma unroll
+        for (int c = 0; c < 16; ++c)
+        {
+            const half_t v = op_get(xv, c);
+            ss += (float)(half_t)(v * v);
+        }
+    }
+    const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
+    for (int kt = (int)hf; kt < KT; kt += 2)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+        u8v hv;
+        const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt); // 16 contiguous gammas
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+            hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
+                          op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
+        store_row16(hrow + 16 * kt, hv);
+    }
+    __builtin_amdgcn_wave_barrier();
+#ifdef PWIN_DBG
+    if (DBG_ON && !hf)
+        for (int c = 0; c < C; ++c)
+            g_dbg[0][16 * wv + m][c] = (float)hrow[c];
+#endif
+
+    if constexpr (L::POS)
+        block_sync(); // V of all 64 tokens is computed by every wave from hb
+    STAMP(2);
+    // ---- attention, one head at a time; the output projection accumulates head by head
+    f8v acc[KT];
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+{
+
+    half_t vv_[8];
+
+    dvec8(w, L::BO + 32 * nt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        acc[nt][i] = (float)vv_[i];
+
+}
+    for (int h = 0; h < H; ++h)
+    {
+        // position-only attention: V^T A operands straight from D = h1 Wv (lane = dim, 16 keys per tile)
+        u8v vtop[2][4];
+        if constexpr (L::POS)
+        {
+            // each wave computes the V^T operands of its own key tile; shared through LDS, one slot per head
+            __shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                f8v d = splat8(0.0f);
+                for (int ks = 0; ks < KT; ++ks)
+                    d = mma16(lds_row16(&hb[16 * wv + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
+                const u8v vo = operand_from_f8(d);
+                if (!hf)
+                    vsh[h][nt][wv][m] = vo;
+            }
+            block_sync();
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+                for (int kt = 0; kt < 4; ++kt)
+                    vtop[nt][kt] = vsh[h][nt][kt][m];
+        }
+        u8v pop[4];
+        if constexpr (L::POS)
+        {
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+                pop[kt] = ((const u8v*)bias)[((h * 4 + wv) * 4 + kt) * 16 + m];
+        }
+        else
+        {
+        u8v qop[2];
+        // six independent chains (Q, K, V x 2 n-tiles), each over kt in order
+        f8v dq[6];
+#pragma unroll
+        for (int u = 0; u < 6; ++u)
+            dq[u] = splat8(0.0f);
+        for (int kt = 0; kt < KT; ++kt)
+        {
+            const u8v row = lds_row16(hrow + 16 * kt);
+#pragma unroll
+            for (int u = 0; u < 6; ++u)
+                dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[u]);
+        }
+#pragma unroll
+        for (int j = 0; j < 3; ++j)
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                const f8v d = dq[2 * j + nt];
+                if (j == 0)
+                    qop[nt] = operand_from_f8(d);
+                else if (j == 1)
+                {
+                    const u8v kop = operand_from_f8(d);
+                    if (!hf)
+                        store_row16(&kl[16 * wv + m][16 * nt], kop);
+                }
+                else
+                {
+#pragma unroll
+                    for (int i = 0; i < 8; ++i)
+                        vt[16 * nt + 2 * i + hf][16 * wv + m] = (half_t)d[i];
+                }
+            }
+        block_sync();
+        half_t e[4][8];
+        float rs = 0.0f;
+#pragma unroll
+        for (int kt = 0; kt < 4; ++kt)
+        {
+            const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
+            f8v d;
+#pragma unroll
+            for (int i = 0; i < 4; ++i)
+            {
+                d[2 * i] = (float)lo16(bi[i]);
+                d[2 * i + 1] = (float)hi16(bi[i]);
+            }
+#pragma unroll
+            for (int ds = 0; ds < 2; ++ds)
+                d = mma16(lds_row16(&kl[16 * kt + m][16 * ds]), qop[ds], d);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                e[kt][i] = softmax_e((half_t)d[i]);
+                rs += (float)e[kt][i];
+            }
+        }
+        rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+        const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
+#pragma unroll
+        for (int kt = 0; kt < 4; ++kt)
+        {
+            half_t pv[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                pv[i] = e[kt][i] * rc;
+            pop[kt] = operand_from_dt(pv);
+        }
+        }
+        u8v oop[2];
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+        {
+            f8v d = splat8(0.0f);
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                if constexpr (L::POS)
+                    d = mma16(vtop[nt][kt], pop[kt], d);
+                else
+                    d = mma16(lds_row16(&vt[16 * nt + m][16 * kt]), pop[kt], d);
+            }
+            oop[nt] = operand_from_f8(d);
+#ifdef PWIN_DBG
+            if (DBG_ON)
+                for (int i = 0; i < 8; ++i)
+                    g_dbg[2 + h][16 * wv + m][16 * nt + 2 * i + hf] = (float)(half_t)d[i];
+#endif
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+                acc[nt] = mma16(img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m], oop[ks], acc[nt]);
+        if constexpr (!L::POS)
+            block_sync(); // kl / vt are reused by the next head
+    }
+
+#ifdef PWIN_DBG
+    if (DBG_ON)
+        for (int nt = 0; nt < KT; ++nt)
+            for (int i = 0; i < 8; ++i)
+                g_dbg[1][16 * wv + m][16 * nt + 2 * i + hf] = acc[nt][i];
+#endif
+    STAMP(3);
+    // ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        const u8v xv = lds_row16(arow + 16 * nt);
+        half_t mv[8], g2[8], b2[8];
+        dvec8(w, L::G2 + 32 * nt, hf, g2);
+        dvec8(w, L::B2 + 32 * nt, hf, b2);
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+        {
+            const half_t x1 = (half_t)acc[nt][i] + op_get(xv, 2 * i + hf);
+            mv[i] = x1 * g2[i];
+            acc[nt][i] = (float)(half_t)(x1 + b2[i]);
+        }
+        const u8v mo = operand_from_dt(mv);
+        if (!hf)
+            store_row16(hrow + 16 * nt, mo);
+    }
+    __builtin_amdgcn_wave_barrier();
+
+    // ---- MLP: C/8 chunks of 32 hidden
+    for (int c = 0; c < L::NMLP; ++c)
+    {
+        u8v gop[2];
+#pragma unroll
+        for (int hn = 0; hn < 2; ++hn)
+        {
+            f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+}
+            for (int kt = 0; kt < KT; ++kt)
+                d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
+            half_t g[8];
+            gelu8(d, g);
+            gop[hn] = operand_from_dt(g);
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+                acc[nt] = mma16(img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m], gop[ks], acc[nt]);
+    }
+    STAMP(4);
+    // y -> act (rows of this wave)
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        const u8v yo = operand_from_f8(acc[nt]);
+        if (!hf)
+            store_row16(arow + 16 * nt, yo);
+    }
+    __builtin_amdgcn_wave_barrier();
+}
+
+// first layer: token embedding x0 = relu(in (CIN) W_e + b_e) in front of the core (core weights at CB)
+template <int H_, int C_, int COUT_, bool POS_, int CIN_> struct PwinEmbLayout : PwinLayout<H_, C_, COUT_, POS_>
+{
+    using B = PwinLayout<H_, C_, COUT_, POS_>;
+    static constexpr int CIN = CIN_, KE = CIN / 16, CB = CIN ? 2 * C_ + 2 * CIN * C_ : 0;
+    static constexpr int T_EMB = B::T_END, T_END = T_EMB + KE * (C_ / 16);
+    static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
+};
+
+template <class L> __device__ void pwin_emb_prep(const PwinParams& p, u8v* img, u4v* bias)
+{
+    const int idx = blockIdx.x * 128 + threadIdx.x;
+    if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
+    {
+        const int t = (idx >> 4) - L::T_EMB, nt = t % (L::C / 16), ks = t / (L::C / 16);
+        img[idx] = img_tile(p.w, 2 * L::C + 512 * ks + 512 * L::KE * nt, idx & 15);
+        return;
+    }
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    // core tiles and the bias / P images (indices after T_END are shifted down to the base layout's)
+    const int base_idx = idx < L::T_EMB * 16 ? idx : idx - (L::T_END - L::T_EMB) * 16;
+    if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
+        return;
+    pwin_prep_at<typename L::B>(q, img, bias, base_idx);
+}
+
+// encoder: core + full-resolution output + 2x2 patch merge
+template <class L>
+__device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    PwinParams p = p0;
+    if constexpr (L::CB != 0)
+        p.w = p0.w + L::CB;
+    constexpr int C = L::C, COUT = L::COUT;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) half_t kl[64][32];
+    __shared__ __attribute__((aligned(16))) half_t vt[32][64];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
+    const int bx = blockIdx.x, by = blockIdx.y;
+
+    STAMP(0);
+    // input rows (mirror padded) -> act (first layer: embedding + ReLU)
+    {
+        const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
+        if constexpr (L::CIN != 0)
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
+            u8v xe[L::KE];
+#pragma unroll
+            for (int ks = 0; ks < L::KE; ++ks)
+                xe[ks] = gload_row16(xin + 16 * ks);
+#pragma unroll
+            for (int nt = 0; nt < C / 16; ++nt)
+            {
+                f8v d;
+                {
+                    half_t vv_[8];
+                    dvec8(p0.w, 32 * nt, hf, vv_);
+#pragma unroll
+                    for (int i = 0; i < 8; ++i)
+                        d[i] = (float)vv_[i];
+                }
+#pragma unroll
+                for (int ks = 0; ks < L::KE; ++ks)
+                    d = mma16(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
+                half_t xo[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    xo[i] = hmax((half_t)d[i], (half_t)0.0f);
+                const u8v xv = operand_from_dt(xo);
+                if (!hf)
+                    store_row16(&act[tok][16 * nt], xv);
+            }
+        }
+        else
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+            for (int kt = (int)hf; kt < L::KT; kt += 2)
+                store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+        }
+    }
+    __builtin_amdgcn_wave_barrier();
+    pwin_core<L>(p, img, bias, act, hb, kl, vt);
+
+    STAMP(5);
+    // full-resolution output
+    const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
+    if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+    {
+        half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+    }
+    block_sync();
+
+    STAMP(6);
+    // patch merge: 16 merged tokens, K = 4 sub-tokens x C, COUT/16 n-tiles over the 4 waves
+    const int my = m >> 2, mx = m & 3;
+    const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
+    const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
+    half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
+    const uint8_t* w = p.w;
+    for (int vt = wv; vt < L::PMT; vt += 4)
+    {
+        f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(w, L::PMB + 32 * vt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+    STAMP(7);
+}
+        for (int ks = 0; ks < C / 4; ++ks)
+        {
+            const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
+            const int src = 8 * (2 * my + dy) + 2 * mx + dx;
+            d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
+        }
+        const u8v mo = operand_from_f8(d);
+        // allocated column 16 vt + c belongs to original warp g, local column 16 u + c
+        const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
+        half_t* dst = mout + L::NPW * g + 16 * u;
+        if (!hf && minb)
+        {
+            gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
+            if (valid >= 16)
+                gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
+        }
+    }
+}
+
+#define PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, CIN)                                                               \
+    using NAME##_L = PwinEmbLayout<H, C, COUT, POS, CIN>;                                                                       \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4;                                          \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_emb_prep<NAME##_L>(p, g_img, g_bias);                                                                     \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pwin_encoder<NAME##_L>(p, g_img, g_bias);                                                                  \
+    }
+#define PWIN_ENCODER(NAME, H, C, COUT) PWIN_ENCODER_EMB(NAME, H, C, COUT, false, 0)
+#define PWIN_ENCODER_POS(NAME, H, C, COUT, POS) PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, 0)
+
+// bottleneck (dec5): core only, output at the same resolution to +24
+template <class L>
+__device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) half_t kl[64][32];
+    __shared__ __attribute__((aligned(16))) half_t vt[32][64];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
+    const int bx = blockIdx.x, by = blockIdx.y;
+    {
+        const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
+        const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+    }
+    __builtin_amdgcn_wave_barrier();
+    pwin_core<L>(p, img, bias, act, hb, kl, vt);
+    const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
+    if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+    {
+        half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+    }
+}
+
+#define PWIN_KERNEL(NAME, H, C, COUT, BODY)                                                                        \
+    using NAME##_L = PwinLayout<H, C, COUT>;                                                                       \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4;                                          \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_prep<NAME##_L>(p, g_img, g_bias);                                                                     \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p)                                           \
+    {                                                                                                              \
+        BODY<NAME##_L>(p, g_img, g_bias);                                                                          \
+    }
+
+// ---------------------------------------------------------------- decoders (patch expand + skip)
+template <int H_, int C_, int CL_, bool POS_ = false, int NOUT_ = 0, int NOUTA_ = 0>
+struct PwinDecLayout : PwinLayout<H_, C_, 16, POS_>
+{
+    using B = PwinLayout<H_, C_, 16, POS_>;
+    // last layer: output head out = y W_h + b_h (NOUTA computed, NOUT stored); b_h at B::PM, W_h after it
+    static constexpr int NOUT = NOUT_, NOUTA = NOUTA_, HB = B::PM, HW = B::PM + 2 * NOUTA_;
+    static constexpr int CL = CL_, KL = CL / 16, EXPN = 4 * C_ / 16; // expand: K CL, N 4C
+    static constexpr int EXPB = 2 * CL * 4 * C_, CB = EXPB + 8 * C_;  // expand bias, core weights base
+    static constexpr int T_EXP = B::T_PM;                              // expand tiles replace the merge tiles
+    static constexpr int T_HEAD = T_EXP + KL * EXPN;
+    static constexpr int T_END = T_HEAD + B::KT * (NOUTA_ / 16);
+    static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
+};
+
+template <class L> __device__ void pwin_dec_prep(const PwinParams& p, u8v* img, u4v* bias)
+{
+    const int idx = blockIdx.x * 128 + threadIdx.x;
+    if (idx >= L::T_EXP * 16 && idx < L::T_HEAD * 16)
+    {
+        const int t = (idx >> 4) - L::T_EXP, nt = t % L::EXPN, ks = t / L::EXPN;
+        img[idx] = img_tile(p.w, 512 * ks + 32 * L::CL * nt, idx & 15);
+        return;
+    }
+    if (idx >= L::T_HEAD * 16 && idx < L::T_END * 16)
+    {
+        const int t = (idx >> 4) - L::T_HEAD, nt = t % (L::NOUTA / 16), ks = t / (L::NOUTA / 16);
+        img[idx] = img_tile(p.w + L::CB, L::HW + 512 * ks + 512 * L::KT * nt, idx & 15);
+        return;
+    }
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    // core tiles lie below T_EXP; bias / P images follow the base layout's (unused) merge tiles
+    pwin_prep_at<typename L::B>(q, img, bias, idx < L::T_EXP * 16 ? idx : L::B::T_END * 16 + (idx - L::T_END * 16));
+}
+
+template <class L>
+__device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : 32;
+    __shared__ __attribute__((aligned(16))) half_t kl[KVR][32];
+    __shared__ __attribute__((aligned(16))) half_t vt[KVC][64];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
+
+    // ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens
+    {
+        const int W2 = p.W / 2, H2 = p.H / 2;
+        const int ly = m >> 2, lx = m & 3;
+        const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
+        const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
+        const int dy = wv >> 1, dx = wv & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
+        const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
+        const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
+        for (int nt = 0; nt < C / 16; ++nt)
+        {
+            const int gnt = wv * (C / 16) + nt; // expand output column tile
+            f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(p.w, L::EXPB + 32 * gnt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+}
+            for (int ks = 0; ks < L::KL; ++ks)
+                d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
+            const u8v sv = gload_row16(sk + 16 * nt);
+            half_t xo[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
+            const u8v xv = operand_from_dt(xo);
+            if (!hf)
+                store_row16(&act[t][16 * nt], xv);
+        }
+    }
+    block_sync();
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    pwin_core<L>(q, img, bias, act, hb, kl, vt);
+
+    const int tok = 16 * wv + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
+    const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
+    if constexpr (L::NOUTA != 0)
+    {
+        half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
+#pragma unroll
+        for (int nt = 0; nt < L::NOUTA / 16; ++nt)
+        {
+            f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(q.w, L::HB + 32 * nt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+}
+#pragma unroll
+            for (int kt = 0; kt < L::KT; ++kt)
+                d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
+            const u8v ho = operand_from_f8(d);
+            const int valid = L::NOUT - 16 * nt;
+            if (!hf && inb)
+            {
+                gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
+                if (valid >= 16)
+                    gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
+            }
+        }
+    }
+    else if (inb)
+    {
+        half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+    }
+}
+
+#define PWIN_DECODER_EX(NAME, H, C, CL, POS, NOUT, NOUTA)                                                          \
+    using NAME##_L = PwinDecLayout<H, C, CL, POS, NOUT, NOUTA>;                                                                      \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4;                                          \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_dec_prep<NAME##_L>(p, g_img, g_bias);                                                                 \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pwin_decoder<NAME##_L>(p, g_img, g_bias);                                                                  \
+    }
+
+#define PWIN_DECODER(NAME, H, C, CL) PWIN_DECODER_EX(NAME, H, C, CL, false, 0, 0)

kernels/k/pwin_pos.h

@@ -0,0 +1,433 @@
+// Position-only-attention layers (trait flag 2: enc0, enc1, dec0) with MT token tiles per wave, so every
+// weight operand loaded from L2 is used MT times (4 / MT waves per window). Same math as pwin_core.
+#pragma once
+#include "pwin_layer.h"
+
+template <class L, int MT>
+__device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
+                         half_t (*act)[L::C], half_t (*hb)[L::C])
+{
+    constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
+    static_assert(L::POS, "position-only attention layers");
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z;
+    const uint8_t* w = p.w;
+
+    // ---- L2 norm of every tile of this wave -> hb
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const half_t* arow = act[16 * (wv * MT + mi) + m];
+        half_t* hrow = hb[16 * (wv * MT + mi) + m];
+        float ss = 0.0f;
+#pragma unroll
+        for (int kt = 0; kt < KT; ++kt)
+        {
+            const u8v xv = lds_row16(arow + 16 * kt);
+#pragma unroll
+            for (int c = 0; c < 16; ++c)
+            {
+                const half_t v = op_get(xv, c);
+                ss += (float)(half_t)(v * v);
+            }
+        }
+        const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
+        for (int kt = (int)hf; kt < KT; kt += 2)
+        {
+            const u8v xv = lds_row16(arow + 16 * kt);
+            const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
+            u8v hv;
+#pragma unroll
+            for (int j = 0; j < 8; ++j)
+                hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
+                              op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
+            store_row16(hrow + 16 * kt, hv);
+        }
+    }
+    block_sync(); // V of all 64 tokens is computed from hb
+    STAMP(2);
+
+    // ---- attention: O = P V per head (P from the table), output projection accumulated head by head
+    f8v acc[MT][KT];
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        half_t vv[8];
+        dvec8(w, L::BO + 32 * nt, hf, vv);
+#pragma unroll
+        for (int mi = 0; mi < MT; ++mi)
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[mi][nt][i] = (float)vv[i];
+    }
+#ifdef POS_VSHARE
+    // V^T operands of every head: each wave computes those of its own key tiles, shared through LDS
+    // (one slot per head, so a head's slot is never rewritten while it is read)
+    __shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
+#endif
+    for (int h = 0; h < H; ++h)
+    {
+        u8v vtop[2][4];
+#ifdef POS_VSHARE
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+            for (int mi = 0; mi < MT; ++mi)
+            {
+                const int kt = wv * MT + mi;
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int ks = 0; ks < KT; ++ks)
+                    d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
+                const u8v vo = operand_from_f8(d);
+                if (!hf)
+                    vsh[h][nt][kt][m] = vo;
+            }
+        block_sync();
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+                vtop[nt][kt] = vsh[h][nt][kt][m];
+#else
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int ks = 0; ks < KT; ++ks)
+                    d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
+                vtop[nt][kt] = operand_from_f8(d);
+            }
+#endif
+        u8v oop[MT][2];
+#pragma unroll
+        for (int mi = 0; mi < MT; ++mi)
+        {
+            const int qt = wv * MT + mi;
+            u8v pop[4];
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+                pop[kt] = ((const u8v*)bias)[((h * 4 + qt) * 4 + kt) * 16 + m];
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int kt = 0; kt < 4; ++kt)
+                    d = mma16(vtop[nt][kt], pop[kt], d);
+                oop[mi][nt] = operand_from_f8(d);
+            }
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+            {
+                const u8v wo = img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m];
+#pragma unroll
+                for (int mi = 0; mi < MT; ++mi)
+                    acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
+            }
+    }
+
+    STAMP(3);
+    // ---- residual, MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        half_t g2[8], b2[8];
+        dvec8(w, L::G2 + 32 * nt, hf, g2);
+        dvec8(w, L::B2 + 32 * nt, hf, b2);
+#pragma unroll
+        for (int mi = 0; mi < MT; ++mi)
+        {
+            const int tok = 16 * (wv * MT + mi) + m;
+            const u8v xv = lds_row16(&act[tok][16 * nt]);
+            half_t mv[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const half_t x1 = (half_t)acc[mi][nt][i] + op_get(xv, 2 * i + hf);
+                mv[i] = x1 * g2[i];
+                acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
+            }
+            const u8v mo = operand_from_dt(mv);
+            if (!hf)
+                store_row16(&hb[tok][16 * nt], mo);
+        }
+    }
+    __builtin_amdgcn_wave_barrier();
+
+    STAMP(4);
+    // ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
+    for (int c = 0; c < L::NMLP; ++c)
+    {
+        u8v gop[MT][2];
+#pragma unroll
+        for (int hn = 0; hn < 2; ++hn)
+        {
+            half_t bb[8];
+            dvec8(w, L::B1 + 64 * c + 32 * hn, hf, bb);
+            f8v d[MT];
+#pragma unroll
+            for (int mi = 0; mi < MT; ++mi)
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[mi][i] = (float)bb[i];
+#pragma unroll
+            for (int kt = 0; kt < KT; ++kt)
+            {
+                const u8v w1 = img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m];
+#pragma unroll
+                for (int mi = 0; mi < MT; ++mi)
+                    d[mi] = mma16(w1, lds_row16(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
+            }
+#pragma unroll
+            for (int mi = 0; mi < MT; ++mi)
+            {
+                half_t g[8];
+                gelu8(d[mi], g);
+                gop[mi][hn] = operand_from_dt(g);
+            }
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+            {
+                const u8v w2 = img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m];
+#pragma unroll
+                for (int mi = 0; mi < MT; ++mi)
+                    acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
+            }
+    }
+    STAMP(5);
+    // y -> act
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+        {
+            const u8v yo = operand_from_f8(acc[mi][nt]);
+            if (!hf)
+                store_row16(&act[16 * (wv * MT + mi) + m][16 * nt], yo);
+        }
+    block_sync();
+}
+
+// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
+template <class L, int MT>
+__device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C, NW = 4 / MT;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    PwinParams p = p0;
+    if constexpr (L::CB != 0)
+        p.w = p0.w + L::CB;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
+    STAMP(0);
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const int tok = 16 * (wv * MT + mi) + m, ty = tok >> 3, tx = tok & 7;
+        const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
+        if constexpr (L::CIN != 0)
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
+            u8v xe[L::KE];
+#pragma unroll
+            for (int ks = 0; ks < L::KE; ++ks)
+                xe[ks] = gload_row16(xin + 16 * ks);
+#pragma unroll
+            for (int nt = 0; nt < C / 16; ++nt)
+            {
+                half_t vv[8];
+                dvec8(p0.w, 32 * nt, hf, vv);
+                f8v d;
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+#pragma unroll
+                for (int ks = 0; ks < L::KE; ++ks)
+                    d = mma16(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
+                half_t xo[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    xo[i] = hmax((half_t)d[i], (half_t)0.0f);
+                const u8v xv = operand_from_dt(xo);
+                if (!hf)
+                    store_row16(&act[tok][16 * nt], xv);
+            }
+        }
+        else
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+            for (int kt = (int)hf; kt < L::KT; kt += 2)
+                store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+        }
+    }
+    __builtin_amdgcn_wave_barrier();
+    STAMP(1);
+    pos_core<L, MT>(p, img, bias, act, hb);
+    STAMP(6);
+
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const int tok = 16 * (wv * MT + mi) + m;
+        const int Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
+        if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+        {
+            half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
+            for (int kt = (int)hf; kt < L::KT; kt += 2)
+                gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+        }
+    }
+    STAMP(7);
+    const int my = m >> 2, mx = m & 3;
+    const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
+    const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
+    half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
+    for (int vt = wv; vt < L::PMT; vt += NW)
+    {
+        half_t vv[8];
+        dvec8(p.w, L::PMB + 32 * vt, hf, vv);
+        f8v d;
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+            d[i] = (float)vv[i];
+        for (int ks = 0; ks < C / 4; ++ks)
+        {
+            const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
+            const int src = 8 * (2 * my + dy) + 2 * mx + dx;
+            d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
+        }
+        const u8v mo = operand_from_f8(d);
+        const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
+        half_t* dst = mout + L::NPW * g + 16 * u;
+        if (!hf && minb)
+        {
+            gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
+            if (valid >= 16)
+                gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
+        }
+    }
+    STAMP(8);
+}
+
+// dec0: patch expand + skip, core, output head
+template <class L, int MT>
+__device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C, NW = 4 / MT;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
+    {
+        const int W2 = p.W / 2, H2 = p.H / 2, ly = m >> 2, lx = m & 3;
+        const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
+        const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
+        u8v xls[L::KL];
+#pragma unroll
+        for (int ks = 0; ks < L::KL; ++ks)
+            xls[ks] = gload_row16(xl + 16 * ks);
+        for (int q = wv; q < 4; q += NW)
+        {
+            const int dy = q >> 1, dx = q & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
+            const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
+            const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
+#pragma unroll
+            for (int nt = 0; nt < C / 16; ++nt)
+            {
+                const int gnt = q * (C / 16) + nt;
+                half_t vv[8];
+                dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
+                f8v d;
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+#pragma unroll
+                for (int ks = 0; ks < L::KL; ++ks)
+                    d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], xls[ks], d);
+                const u8v sv = gload_row16(sk + 16 * nt);
+                half_t xo[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
+                const u8v xv = operand_from_dt(xo);
+                if (!hf)
+                    store_row16(&act[t][16 * nt], xv);
+            }
+        }
+    }
+    block_sync();
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    pos_core<L, MT>(q, img, bias, act, hb);
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const int tok = 16 * (wv * MT + mi) + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
+        const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
+        half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
+#pragma unroll
+        for (int nt = 0; nt < L::NOUTA / 16; ++nt)
+        {
+            half_t vv[8];
+            dvec8(q.w, L::HB + 32 * nt, hf, vv);
+            f8v d;
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                d[i] = (float)vv[i];
+#pragma unroll
+            for (int kt = 0; kt < L::KT; ++kt)
+                d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
+            const u8v ho = operand_from_f8(d);
+            const int valid = L::NOUT - 16 * nt;
+            if (!hf && inb)
+            {
+                gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
+                if (valid >= 16)
+                    gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
+            }
+        }
+    }
+}
+
+#define POS_ENCODER(NAME, H, C, COUT, CIN, MT)                                                                     \
+    using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>;                                                         \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT;                                     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_emb_prep<NAME##_L>(p, g_img, g_bias);                                                                 \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pos_encoder<NAME##_L, MT>(p, g_img, g_bias);                                                               \
+    }
+
+#define POS_DECODER(NAME, H, C, CL, NOUT, NOUTA, MT)                                                               \
+    using NAME##_L = PwinDecLayout<H, C, CL, true, NOUT, NOUTA>;                                                   \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT;                                     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_dec_prep<NAME##_L>(p, g_img, g_bias);                                                                 \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pos_decoder<NAME##_L, MT>(p, g_img, g_bias);                                                               \
+    }

kernels/k/pwin_wide.h

@@ -0,0 +1,455 @@
+// Wide variant of the DltssPaddedWinLayer template for the deep, latency-bound levels (few windows):
+// 4 * NG waves per window. Wave wv = threadIdx.z works on token tile t = wv & 3 and channel group
+// g = wv >> 2: heads g, g + NG, ... (one head per group and round), output tiles nt = g + NG j of the
+// output projection and MLP, and MLP hidden chunks c = NG rr + g. Every output element goes through
+// the same f16 operations in the same order as pwin_layer.h (bit-identical results).
+//
+// Shared memory: hb (64 x C) = normalised input, then MLP input, then the block output y; R = union of
+// x0 (64 x C), the K / V images of the NG heads of an attention round, their outputs O, and the double
+// buffered MLP hidden chunks of a round.
+#pragma once
+#include "pwin_layer.h"
+
+template <class L, int NG> struct WideDims
+{
+    // NA: heads per attention round (their K / V images share R)
+    static constexpr int NA = NG < 4 ? NG : 4;
+    static constexpr int C = L::C, KT = L::KT, H = L::H, NTL = (KT + NG - 1) / NG, ROUNDS = (H + NA - 1) / NA;
+    static constexpr int X0_BYTES = 64 * C * 2, KV_BYTES = NA * 2 * 64 * 32 * 2;
+    // MLP hidden chunks: one per group and round, double-buffered while that fits
+    static constexpr int NB = NG <= 4 ? 2 : 1, MROUNDS = (L::NMLP + NG - 1) / NG, HID_BYTES = NB * NG * 64 * 32 * 2;
+    static constexpr int R_BYTES = X0_BYTES > KV_BYTES ? (X0_BYTES > HID_BYTES ? X0_BYTES : HID_BYTES)
+                                                       : (KV_BYTES > HID_BYTES ? KV_BYTES : HID_BYTES);
+    static_assert(!L::POS, "position-only attention layers use pwin_pos.h");
+};
+
+// Core for one window. On entry R holds x0 (64 x C); on return hb holds the block output y.
+template <class L, int NG>
+__device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
+                          half_t (*hb)[L::C], uint8_t* R)
+{
+    using D = WideDims<L, NG>;
+    constexpr int C = L::C, KT = L::KT, H = L::H, NTL = D::NTL;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
+    const uint8_t* w = p.w;
+    half_t(*act)[C] = (half_t(*)[C])R;
+    const half_t* arow = act[tok];
+    half_t* hrow = hb[tok];
+
+    // ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb (all groups of a token tile compute the sum)
+    float ss = 0.0f;
+    for (int kt = 0; kt < KT; ++kt)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+#pragma unroll
+        for (int c = 0; c < 16; ++c)
+        {
+            const half_t v = op_get(xv, c);
+            ss += (float)(half_t)(v * v);
+        }
+    }
+    const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
+    for (int kt = 2 * g + (int)hf; kt < KT; kt += 2 * NG)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+        u8v hv;
+        const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+            hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
+                          op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
+        store_row16(hrow + 16 * kt, hv);
+    }
+    // residual input of the owned output tiles (D^T order: channel 16 nt + 2 i + hf)
+    half_t x0[NTL][8];
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            const u8v xv = lds_row16(arow + 16 * nt);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                x0[j][i] = op_get(xv, 2 * i + hf);
+        }
+    }
+    block_sync(); // hb complete, R free
+
+    // ---- attention: NG heads per round; the output projection accumulates head by head
+    f8v acc[NTL];
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            half_t vv[8];
+            dvec8(w, L::BO + 32 * nt, hf, vv);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[j][i] = (float)vv[i];
+        }
+    }
+    half_t(*kl)[64][32] = (half_t(*)[64][32])R;                   // [group][token][dim]
+    constexpr int NA = D::NA;
+    half_t(*vt)[32][64] = (half_t(*)[32][64])(R + NA * 64 * 32 * 2); // [group][dim][token]
+    half_t(*ob)[64][32] = (half_t(*)[64][32])R;                   // [group][token][dim] (after K / V)
+    for (int r = 0; r < D::ROUNDS; ++r)
+    {
+        const int h = g < NA ? r * NA + g : H;
+        u8v qop[2], oop[2];
+        if (h < H)
+        {
+            // six independent chains (Q, K, V x 2 n-tiles), each over kt in order
+            f8v dq[6];
+#pragma unroll
+            for (int u = 0; u < 6; ++u)
+                dq[u] = splat8(0.0f);
+            for (int kt = 0; kt < KT; ++kt)
+            {
+                const u8v row = lds_row16(hrow + 16 * kt);
+#pragma unroll
+                for (int u = 0; u < 6; ++u)
+                    dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[u]);
+            }
+#pragma unroll
+            for (int j = 0; j < 3; ++j)
+#pragma unroll
+                for (int nt = 0; nt < 2; ++nt)
+                {
+                    const f8v d = dq[2 * j + nt];
+                    if (j == 0)
+                        qop[nt] = operand_from_f8(d);
+                    else if (j == 1)
+                    {
+                        const u8v kop = operand_from_f8(d);
+                        if (!hf)
+                            store_row16(&kl[g][tok][16 * nt], kop);
+                    }
+                    else
+                    {
+#pragma unroll
+                        for (int i = 0; i < 8; ++i)
+                            vt[g][16 * nt + 2 * i + hf][tok] = (half_t)d[i];
+                    }
+                }
+        }
+        block_sync();
+        if (h < H)
+        {
+            half_t e[4][8];
+            float rs = 0.0f;
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
+                f8v d;
+#pragma unroll
+                for (int i = 0; i < 4; ++i)
+                {
+                    d[2 * i] = (float)lo16(bi[i]);
+                    d[2 * i + 1] = (float)hi16(bi[i]);
+                }
+#pragma unroll
+                for (int ds = 0; ds < 2; ++ds)
+                    d = mma16(lds_row16(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                {
+                    e[kt][i] = softmax_e((half_t)d[i]);
+                    rs += (float)e[kt][i];
+                }
+            }
+            rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+            const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
+            u8v pop[4];
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                half_t pv[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    pv[i] = e[kt][i] * rc;
+                pop[kt] = operand_from_dt(pv);
+            }
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int kt = 0; kt < 4; ++kt)
+                    d = mma16(lds_row16(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
+                oop[nt] = operand_from_f8(d);
+            }
+        }
+        block_sync(); // K / V of the round consumed
+        if (h < H && !hf)
+        {
+            store_row16(&ob[g][tok][0], oop[0]);
+            store_row16(&ob[g][tok][16], oop[1]);
+        }
+        block_sync();
+        for (int jh = 0; jh < NA; ++jh)
+        {
+            const int hh = r * NA + jh;
+            if (hh >= H)
+                break;
+            const u8v o0 = lds_row16(&ob[jh][tok][0]), o1 = lds_row16(&ob[jh][tok][16]);
+#pragma unroll
+            for (int j = 0; j < NTL; ++j)
+            {
+                const int nt = g + NG * j;
+                if (nt < KT)
+                {
+                    acc[j] = mma16(img[(L::T_WO + (2 * hh) * KT + nt) * 16 + m], o0, acc[j]);
+                    acc[j] = mma16(img[(L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m], o1, acc[j]);
+                }
+            }
+        }
+        if (r + 1 < D::ROUNDS)
+            block_sync(); // ob is overwritten by the next round's K / V
+    }
+
+    // ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            half_t mv[8], g2[8], b2[8];
+            dvec8(w, L::G2 + 32 * nt, hf, g2);
+            dvec8(w, L::B2 + 32 * nt, hf, b2);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const half_t x1 = (half_t)acc[j][i] + x0[j][i];
+                mv[i] = x1 * g2[i];
+                acc[j][i] = (float)(half_t)(x1 + b2[i]);
+            }
+            const u8v mo = operand_from_dt(mv);
+            if (!hf)
+                store_row16(hrow + 16 * nt, mo);
+        }
+    }
+    block_sync();
+
+    // ---- MLP: rounds of NG chunks of 32 hidden (double-buffered in R)
+    half_t(*hid)[NG][64][32] = (half_t(*)[NG][64][32])R;
+    for (int rr = 0; rr < D::MROUNDS; ++rr)
+    {
+        const int c = rr * NG + g;
+#pragma unroll
+        for (int hn = 0; hn < 2; ++hn)
+        {
+            if (c >= L::NMLP)
+                break;
+            f8v d;
+            {
+                half_t vv[8];
+                dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv);
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+            }
+            for (int kt = 0; kt < KT; ++kt)
+                d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
+            half_t gg[8];
+            gelu8(d, gg);
+            const u8v gop = operand_from_dt(gg);
+            if (!hf)
+                store_row16(&hid[rr % D::NB][g][tok][16 * hn], gop);
+        }
+        block_sync();
+        for (int jc = 0; jc < NG; ++jc)
+        {
+            const int cc = rr * NG + jc;
+            if (cc >= L::NMLP)
+                break;
+            const u8v h0 = lds_row16(&hid[rr % D::NB][jc][tok][0]), h1 = lds_row16(&hid[rr % D::NB][jc][tok][16]);
+#pragma unroll
+            for (int j = 0; j < NTL; ++j)
+            {
+                const int nt = g + NG * j;
+                if (nt < KT)
+                {
+                    acc[j] = mma16(img[(L::T_W2 + (cc * 2) * KT + nt) * 16 + m], h0, acc[j]);
+                    acc[j] = mma16(img[(L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m], h1, acc[j]);
+                }
+            }
+        }
+        if (D::NB == 1 && rr + 1 < D::MROUNDS)
+            block_sync();
+    }
+    // y -> hb (the last reads of hb were the W1 operands before the final barrier)
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            const u8v yo = operand_from_f8(acc[j]);
+            if (!hf)
+                store_row16(hrow + 16 * nt, yo);
+        }
+    }
+    block_sync();
+}
+
+// x0 rows of a window (mirror padded) -> R
+template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uint8_t* R)
+{
+    constexpr int C = L::C;
+    half_t(*act)[C] = (half_t(*)[C])R;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
+    const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W), Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
+    const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+    for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
+        store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+}
+
+// y rows of the window's own tokens -> dst (C per token)
+template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[L::C], uint8_t* dst)
+{
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
+    const int Y0 = 8 * (int)blockIdx.y - p.sy + (tok >> 3), X0 = 8 * (int)blockIdx.x - p.sx + (tok & 7);
+    if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+    {
+        half_t* yout = (half_t*)dst + (size_t)(Y0 * p.W + X0) * L::C;
+        for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
+            gstore_row16(yout + 16 * kt, lds_row16(&hb[tok][16 * kt]));
+    }
+}
+
+template <class L, int NG>
+__device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C, COUT = L::COUT;
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
+    wide_load_x0<L, NG>(p, R);
+    block_sync();
+    wide_core<L, NG>(p, img, bias, hb, R);
+    wide_store_y<L, NG>(p, hb, p.out32);
+
+    // patch merge: 16 merged tokens, K = 4 sub-tokens x C, the allocated n-tiles over the waves
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z;
+    const int my = m >> 2, mx = m & 3;
+    const int Ym = (8 * (int)blockIdx.y - p.sy) / 2 + my, Xm = (8 * (int)blockIdx.x - p.sx) / 2 + mx;
+    const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
+    half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
+    const uint8_t* w = p.w;
+    for (int vt = wv; vt < L::PMT; vt += 4 * NG)
+    {
+        f8v d;
+        {
+            half_t vv[8];
+            dvec8(w, L::PMB + 32 * vt, hf, vv);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                d[i] = (float)vv[i];
+        }
+        for (int ks = 0; ks < C / 4; ++ks)
+        {
+            const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
+            const int src = 8 * (2 * my + dy) + 2 * mx + dx;
+            d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&hb[src][16 * (ks % L::KT)]), d);
+        }
+        const u8v mo = operand_from_f8(d);
+        const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
+        half_t* dst = mout + L::NPW * gw + 16 * u;
+        if (!hf && minb)
+        {
+            gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
+            if (valid >= 16)
+                gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
+        }
+    }
+}
+
+// bottleneck (dec5): core only, output at the same resolution to +24
+template <class L, int NG>
+__device__ void wide_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    __shared__ __attribute__((aligned(16))) half_t hb[64][L::C];
+    __shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
+    wide_load_x0<L, NG>(p, R);
+    block_sync();
+    wide_core<L, NG>(p, img, bias, hb, R);
+    wide_store_y<L, NG>(p, hb, p.out24);
+}
+
+template <class L, int NG>
+__device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C;
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
+    half_t(*act)[C] = (half_t(*)[C])R;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
+
+    // ---- patch expand: waves of sub-token q = (dy, dx) produce its tiles nt = g + NG j for the 16 low-res tokens
+    {
+        const int W2 = p.W / 2, H2 = p.H / 2;
+        const int ly = m >> 2, lx = m & 3;
+        const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
+        const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
+        const int dy = q >> 1, dx = q & 1, tt = 8 * (2 * ly + dy) + 2 * lx + dx;
+        const int SX = mirror(8 * bx - p.sx + (tt & 7), p.W), SY = mirror(8 * by - p.sy + (tt >> 3), p.H);
+        const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
+        for (int nt = g; nt < C / 16; nt += NG)
+        {
+            const int gnt = q * (C / 16) + nt;
+            f8v d;
+            {
+                half_t vv[8];
+                dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+            }
+            for (int ks = 0; ks < L::KL; ++ks)
+                d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
+            const u8v sv = gload_row16(sk + 16 * nt);
+            half_t xo[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
+            const u8v xv = operand_from_dt(xo);
+            if (!hf)
+                store_row16(&act[tt][16 * nt], xv);
+        }
+    }
+    block_sync();
+    PwinParams qp = p;
+    qp.w = p.w + L::CB;
+    wide_core<L, NG>(qp, img, bias, hb, R);
+    static_assert(L::NOUTA == 0, "output head: use pwin_layer.h");
+    wide_store_y<L, NG>(p, hb, p.out24);
+}
+
+#define WIDE_KERNEL(NAME, PREP, BODY, NG, ...)                                                                  \
+    using NAME##_L = __VA_ARGS__;                                                                                     \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 * (NG);                                   \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        PREP<NAME##_L>(p, g_img, g_bias);                                                                          \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128 * (NG)) PWIN_VGPR NAME(PwinParams p)                          \
+    {                                                                                                              \
+        BODY<NAME##_L, NG>(p, g_img, g_bias);                                                                      \
+    }
+
+#define WIDE_ENCODER(NAME, H, C, COUT, NG) WIDE_KERNEL(NAME, pwin_prep, wide_encoder, NG, PwinLayout<H, C, COUT>)
+#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep, wide_plain, NG, PwinLayout<H, C, 16>)
+#define WIDE_DECODER(NAME, H, C, CL, NG) WIDE_KERNEL(NAME, pwin_dec_prep, wide_decoder, NG, PwinDecLayout<H, C, CL>)

kernels/m/rrlite_dec1_4x4.hip

@@ -0,0 +1,7 @@
+// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
+#define SWIN_NO_Q8
+#define SWIN_F32ACC
+#define SWIN_MAX_VGPR 128
+#define SWIN_MLP_R(NW) 2
+#include "swin_block.h"
+SWIN_MODULE_W(rrlite_dec1_4x4, 64, 2, 8, 0, false, CommonParams, 96)

kernels/m/rrlite_dec2_4x4.hip

@@ -0,0 +1,7 @@
+// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
+#define SWIN_NO_Q8
+#define SWIN_F32ACC
+#define SWIN_MAX_VGPR 192
+#define SWIN_MLP_R(NW) 2
+#include "swin_block.h"
+SWIN_MODULE_W(rrlite_dec2_4x4, 96, 4, 8, 0, false, CommonParams, 128)

kernels/m/rrlite_enc1_4x4.hip

@@ -0,0 +1,7 @@
+// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
+#define SWIN_NO_Q8
+#define SWIN_F32ACC
+#define SWIN_MAX_VGPR 192
+#define SWIN_MLP_R(NW) 1
+#include "swin_block.h"
+SWIN_MODULE_W(rrlite_enc1_4x4, 64, 2, 4, 96, false, CommonParams, 0)

kernels/m/rrlite_enc2_4x4.hip

@@ -0,0 +1,7 @@
+// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
+#define SWIN_NO_Q8
+#define SWIN_F32ACC
+#define SWIN_MAX_VGPR 192
+#define SWIN_MLP_R(NW) 2
+#include "swin_block.h"
+SWIN_MODULE_W(rrlite_enc2_4x4, 96, 4, 8, 128, false, CommonParams, 0)

kernels/m/rrlite_enc3_tube_4x4.hip

@@ -0,0 +1,7 @@
+// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
+#define SWIN_NO_Q8
+#define SWIN_F32ACC
+#define SWIN_MAX_VGPR 192
+#define SWIN_MLP_R(NW) 2
+#include "swin_block.h"
+SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, 8, 0, true, TubeParams, 0)

kernels/m/swin_block.h

@@ -0,0 +1,768 @@
+// Templated rrlite Swin block (64 tokens = 4 windows of 4x4 per block, NW waves = heads of 32 dims).
+//   C    channels (64 / 96 / 128), NW waves (2 or 4)
+//   NPM  patch-merge output channels (0 = none): 2x2 merge of the block output, 4C -> NPM
+//   TUBE 88-byte parameter block with the cross-launch flag protocol
+//   CIN  patch-expand input channels (0 = none): decoder input = f16(q8(expand(low-res CIN)) + skip)
+#pragma once
+#include "swin_common.h"
+#include <type_traits>
+
+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;
+    // f16 weight image (h16 units)
+    __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);
+};
+
+template <int C, int NW, int NPM, int CIN> struct SwinDescs
+{
+    GemmDesc d[SwinLayout<C, NW, NPM, CIN>::NDESC];
+};
+
+template <int C, int NW, int NPM, int CIN> constexpr SwinDescs<C, NW, NPM, CIN> make_descs()
+{
+    using L = SwinLayout<C, NW, NPM, CIN>;
+    SwinDescs<C, NW, NPM, CIN> r{};
+    int n = 0;
+    for (int h = 0; h < NW; ++h)
+        for (int q = 0; q < 2; ++q)
+            r.d[n++] = GemmDesc{L::qv(h, q), L::head(h) + (q ? 32 * C : 0), 1024, 512, C / 32, 2};
+    r.d[n++] = GemmDesc{L::WO, L::head(0) + 64 * C + 512, L::HS, 512, NW, C / 16};
+    for (int c = 0; c < L::NC; ++c)
+    {
+        r.d[n++] = GemmDesc{L::w1(c), L::b1(c), 512, 16 * C, C / 32, 2};
+        r.d[n++] = GemmDesc{L::w2(c), L::b1(c) + 32 * C + 64, 0, 512, 1, C / 16};
+    }
+    for (int g = 0; g < NPM / 32; ++g)
+        r.d[n++] = GemmDesc{L::pm(g), L::PM0 + g * 128 * C, 512, 64 * C, C / 8, 2};
+    for (int w = 0; w < (CIN ? NW : 0); ++w)
+        r.d[n++] = GemmDesc{L::pe(w), w * CIN * 4 * C / NW, 512, 16 * CIN, CIN / 32, L::NTE};
+    return r;
+}
+
+// NH heads (NVIDIA's warps per block) run on NW waves (NW = NH, or 2 NH to halve LDS per wave)
+template <int C, int NH, int NW, int NPM, bool TUBE, int CIN>
+__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const h16* w16, const int bx,
+                                           const int by, const int gx)
+{
+    using L = SwinLayout<C, NH, NPM, CIN>;
+    static_assert(NW % NH == 0 && NW / NH <= 4, "waves per head");
+    constexpr int NPL = L::NPL, NT = L::NT, INNER = L::INNER;
+    constexpr int AST = (C > INNER ? C : INNER) + 8;
+    constexpr int GST = 40, QST = 40, PST = 24;
+#ifndef SWIN_MLP_R
+#define SWIN_MLP_R(NW) ((NW) / 2)
+#endif
+    constexpr int R = SWIN_MLP_R(NW); // MLP chunks per round
+    constexpr int MG = (NT % NW == 0) ? 1 : ((2 * NT) % NW == 0 ? 2 : 4), MT = 4 / MG, NTW = NT * MG / NW;
+    static_assert(NTW * NW == NT * MG, "column split");
+    constexpr int PAST = CIN + 8;
+    constexpr int SCR = NW * 16 * QST * 2 > (CIN ? 16 * PAST * 2 : 0) ? NW * 16 * QST * 2 : 16 * PAST * 2;
+    constexpr int U_ST4 = R * 64 * GST * 2, U_OB = NPM ? 64 * C : 0;
+    constexpr int USIZE = SCR > U_ST4 ? (SCR > U_OB ? SCR : U_OB) : (U_ST4 > U_OB ? U_ST4 : U_OB);
+
+    __shared__ __attribute__((aligned(16))) half_t A[64 * AST];
+    __shared__ __attribute__((aligned(16))) uint8_t U[USIZE];
+    __shared__ half_t R2[64];
+    half_t* X0 = A;                            // stages 0-1: block input f16 [64][AST], natural order
+    half_t* QP = (half_t*)U;                   // stage 2: per-wave Q / P scratch [NW][16][QST]
+    half_t* PA = QP;                           // stage 0: low-res tokens f16 [16][PAST] (apos order)
+    half_t* G = (half_t*)U;                    // stage 4: fc1 outputs [R][64][GST]
+
+    const int lane = threadIdx.x, wv = threadIdx.z, tid = lane + 32 * wv;
+    const int l16 = lane & 15, hi = lane >> 4;
+    const uint8_t* W = p.w;
+
+    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 words (encoders: fp8 input, decoders: fp8 skip), issued first so their latency
+    // overlaps the patch expand; word k of thread tid is idx = tid + 32 NW k
+    constexpr int NSW = CIN ? 1 : 64 * C / 4 / (32 * NW);
+    static_assert(CIN || NSW * 32 * NW == 64 * C / 4, "input words per thread");
+    uint32_t inw[NSW];
+#pragma unroll
+    for (int k = 0; k < (CIN ? 0 : NSW); ++k)
+    {
+        const int idx = tid + 32 * NW * k;
+        const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
+        const int X = mirror(8 * bx - p.sx + tok_x(T), p.tw), Y = mirror(8 * by - p.sy + tok_y(T), p.th);
+        const uint8_t* src = CIN ? p.p32 : p.in;
+        inw[k] = *(const uint32_t*)(src + ((size_t)((c4 >> 5) * p.th + Y) * p.tw + X) * 32 + (c4 & 31));
+    }
+
+    // ------------------------------------------------ stage 0: block input x0 (f16) into A
+    // encoders: the fp8 input; decoders: q8(patch expand(low-res)) + skip
+    if constexpr (CIN > 0)
+    {
+        constexpr int NTEW_ = 4 * C / NW / 16;
+        uint32_t sk[NTEW_][8];
+#pragma unroll
+        for (int j = 0; j < NTEW_; ++j)
+        {
+            const int Gc = 16 * (NTEW_ * wv + j) + l16, q = Gc / C, ch = ginv(Gc % C);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const int r = 2 * i + hi;
+                const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
+                const int X = mirror(8 * bx - p.sx + tx, p.tw), Y = mirror(8 * by - p.sy + ty, p.th);
+                sk[j][i] = p.p32[((size_t)((ch >> 5) * p.th + Y) * p.tw + X) * 32 + (ch & 31)];
+            }
+        }
+        const int W2 = p.tw / 2, H2 = p.th / 2;
+        for (int idx = tid; idx < 16 * CIN / 4; idx += 32 * NW)
+        {
+            const int r = idx / (CIN / 4), c4 = 4 * (idx % (CIN / 4));
+            const int MX = mirror((8 * bx - p.sx) / 2 + (r & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (r >> 2), H2);
+            const uint32_t bytes =
+                *(const uint32_t*)(p.in + ((size_t)((c4 >> 5) * H2 + MY) * W2 + MX) * 32 + (c4 & 31));
+#pragma unroll
+            for (int b = 0; b < 4; ++b)
+                PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
+        }
+        __syncthreads();
+        constexpr int NTEW = 4 * C / NW / 16; // expand tiles per wave (L::NTE per weight block)
+        T16 e[NTEW];
+#pragma unroll
+        for (int j = 0; j < NTEW; ++j)
+            e[j] = t16_splat(hload(W, 4 * C * CIN + 2 * (16 * (NTEW * wv + j) + l16)));
+#pragma unroll
+        for (int kc = 0; kc < CIN / 32; ++kc)
+        {
+            const h16 a0 = lds16(PA + l16 * PAST + 32 * kc), a1 = lds16(PA + l16 * PAST + 32 * kc + 16);
+#pragma unroll
+            for (int j = 0; j < NTEW; ++j)
+            {
+                const int t = NTEW * wv + j, blk = t / L::NTE, lt = t % L::NTE;
+                k32(e[j], a0, bload(w16, L::pe(blk), L::NTE, kc, 0, lt, l16), a1, bload(w16, L::pe(blk), L::NTE, kc, 1, lt, l16));
+            }
+        }
+#pragma unroll
+        for (int j = 0; j < NTEW; ++j)
+        {
+            const int Gc = 16 * (NTEW * wv + j) + l16, q = Gc / C, ch = ginv(Gc % C);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const int r = 2 * i + hi;
+                const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
+                const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
+                X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + e4m3_to_half(sk[j][i]));
+            }
+        }
+    }
+    else
+    {
+#pragma unroll
+    for (int k = 0; k < NSW; ++k)
+    {
+        const int idx = tid + 32 * NW * k;
+        const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
+        const uint32_t bytes = inw[k];
+#pragma unroll
+        for (int b = 0; b < 4; ++b)
+        {
+            X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
+        }
+    }
+    }
+    __syncthreads();
+
+    // residual accumulators of the output projection: f16(x0 + b_o) in the C-fragment layout
+    // (wave owns pair columns of n tiles ng*NTW.. for m tiles mg*MT..)
+    const int ng = wv % (NW / MG), mg = wv / (NW / MG);
+    T16 x[MT][NTW]; // packed f16 accumulators: residual f16(x0 + b_o), live through stage 2
+#pragma unroll
+    for (int mm = 0; mm < MT; ++mm)
+#pragma unroll
+        for (int j = 0; j < NTW; ++j)
+        {
+            const int n = 16 * (ng * NTW + j) + l16;
+            const half_t bo = hload(W, L::BO + 2 * n);
+            const int ch = ginv(n);
+#pragma unroll
+            for (int i = 0; i < 8; i += 2)
+                x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + 2 * i + hi) * AST + ch],
+                                         X0[(16 * (mg * MT + mm) + 2 * i + 2 + hi) * AST + ch]} + (hv2){bo, bo});
+        }
+    __syncthreads(); // stage 1 overwrites x0 rows with h1
+
+    // ------------------------------------------------ stage 1: norm1 (kernel-exact tree)
+    {
+        const int g = lane >> 2, t = lane & 3;
+#pragma unroll
+        for (int r = 0; r < 8 / NW; ++r)
+        {
+            const int T = (64 / NW) * wv + g + 8 * r;
+            half_t v[NPL][2][4];
+#pragma unroll
+            for (int pl = 0; pl < NPL; ++pl)
+#pragma unroll
+                for (int ww = 0; ww < 2; ++ww)
+#pragma unroll
+                    for (int b = 0; b < 4; ++b)
+                        v[pl][ww][b] = X0[T * AST + 32 * pl + 4 * (t + 4 * ww) + b];
+            // the 4 lanes of a token read the whole row before any of them writes h1 over it
+            wave_sync();
+            half_t ws[2][2];
+#pragma unroll
+            for (int ww = 0; ww < 2; ++ww)
+            {
+                half_t pp[NPL][2];
+#pragma unroll
+                for (int pl = 0; pl < NPL; ++pl)
+#pragma unroll
+                    for (int e = 0; e < 2; ++e)
+                    {
+                        half_t s0 = v[pl][ww][e] * v[pl][ww][e];
+                        half_t s1 = v[pl][ww][2 + e] * v[pl][ww][2 + e];
+                        pp[pl][e] = s0 + s1;
+                    }
+#pragma unroll
+                for (int e = 0; e < 2; ++e)
+                {
+                    if constexpr (NPL == 4)
+                        ws[ww][e] = (pp[0][e] + pp[1][e]) + (pp[2][e] + pp[3][e]);
+                    else if constexpr (NPL == 3)
+                        ws[ww][e] = (pp[0][e] + pp[1][e]) + pp[2][e];
+                    else
+                        ws[ww][e] = pp[0][e] + pp[1][e];
+                }
+            }
+            half_t a0 = ws[0][0] + ws[1][0], a1 = ws[0][1] + ws[1][1];
+            a0 = a0 + hxor<1>(a0);
+            a1 = a1 + hxor<1>(a1);
+            a0 = a0 + hxor<2>(a0);
+            a1 = a1 + hxor<2>(a1);
+            const half_t rn = rsqrt16(a0 + a1);
+            const hv2 rn2 = {rn, rn};
+#pragma unroll
+            for (int pl = 0; pl < NPL; ++pl)
+#pragma unroll
+                for (int ww = 0; ww < 2; ++ww)
+#pragma unroll
+                    for (int b = 0; b < 4; b += 2)
+                    {
+                        const int c = 32 * pl + 4 * (t + 4 * ww) + b;
+                        // gperm(c + 1) = gperm(c) + 1 for even c: one 32-bit gamma load per pair
+                        const hv2 gm = __builtin_bit_cast(hv2, *(const uint32_t*)(W + L::G1 + 2 * gperm(c)));
+                        const hv2 hq = q8x2((hv2){v[pl][ww][b], v[pl][ww][b + 1]} * (hv2)(rn2 * gm));
+                        // apos(c + 1) = apos(c) + 1 for even c
+                        *(hv2*)(A + T * AST + apos(c)) = hq;
+                    }
+        }
+    }
+    __syncthreads();
+
+    // ------------------------------------------------ stage 2: Q/V projection + attention (wave = head)
+    {
+        constexpr int MW = 4 / (NW / NH); // windows per wave
+        const int h = NW == NH ? wv : wv % NH, m0 = NW == NH ? 0 : MW * (wv / NH);
+        T16 acc[MW][4]; // [window][Q nt0, Q nt1, V nt0, V nt1]
+#pragma unroll
+        for (int m = 0; m < MW; ++m)
+#pragma unroll
+            for (int q = 0; q < 4; ++q)
+                acc[m][q] = t16_splat((half_t)0.0f);
+#pragma unroll
+        for (int kc = 0; kc < C / 32; ++kc)
+        {
+            h16 b[2][4];
+#pragma unroll
+            for (int s = 0; s < 2; ++s)
+#pragma unroll
+                for (int q = 0; q < 4; ++q)
+                    b[s][q] = bload(w16, L::qv(h, q >> 1), 2, kc, s, q & 1, l16);
+#pragma unroll
+            for (int m = 0; m < MW; ++m)
+            {
+                const h16 a0 = lds16(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
+                const h16 a1 = lds16(A + (16 * (m0 + m) + l16) * AST + 32 * kc + 16);
+#pragma unroll
+                for (int q = 0; q < 4; ++q)
+                    k32(acc[m][q], a0, b[0][q], a1, b[1][q]);
+            }
+        }
+        // pack: Q as q8 pairs, V as f16 pairs (all that attention needs), halving the live registers
+        hv2 qp[MW][2][4], vp[MW][2][4];
+#pragma unroll
+        for (int m = 0; m < MW; ++m)
+#pragma unroll
+            for (int q = 0; q < 2; ++q)
+#pragma unroll
+                for (int i = 0; i < 8; i += 2)
+                {
+                    qp[m][q][i / 2] = q8x2(acc[m][q].pair(i / 2));
+                    vp[m][q][i / 2] = acc[m][2 + q].pair(i / 2);
+                }
+        __syncthreads(); // every wave is done reading h1; O may overwrite A
+
+        half_t* Qs = QP + wv * 16 * QST;
+        const int bbase = L::head(h) + 64 * C;
+        T16 bias;
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+        {
+            const int ri = 2 * i + hi, cj = l16;
+            const int bl = 4 * (ri & 7) + ((cj & 7) >> 1), word = (ri >> 3) + 2 * (cj >> 3);
+            bias.set(i, hload(W, bbase + 2 * (8 * bl + 2 * word + (cj & 1))));
+        }
+        const half_t cs = f16(0.0972222164273262f / 5.656854152679443f);
+        const half_t cl = (half_t)0.55615234375f;
+        const half_t p1 = f16(0.92730712890625f), p0 = (half_t)1.375f;
+#pragma unroll
+        for (int m = 0; m < MW; ++m)
+        {
+#pragma unroll
+            for (int q = 0; q < 2; ++q)
+#pragma unroll
+                for (int i = 0; i < 8; i += 2)
+                {
+                    Qs[(2 * i + hi) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
+                    Qs[(2 * i + 2 + hi) * QST + 16 * q + l16] = qp[m][q][i / 2][1];
+                }
+            wave_sync();
+            const h16 a0 = lds16(Qs + l16 * QST), a1 = lds16(Qs + l16 * QST + 16);
+            T16 S = bias;
+            k32(S, a0, a0, a1, a1);
+            // rows 2k+hi and 2k+2+hi of this lane's column, packed: scale, clamp and the cubic
+            // polynomial as f16x2, the exponent trick per row on the (even, odd) column word
+            hv2 wg[4];
+            const bool oddc = l16 & 1;
+            const uint32_t sel_a = oddc ? 0x01000504u : 0x05040100u, sel_b = oddc ? 0x03020706u : 0x07060302u;
+            const uint32_t sel_r = oddc ? 0x07060302u : 0x05040100u;
+#pragma unroll
+            for (int k = 0; k < 4; ++k)
+            {
+                hv2 tv = S.pair(k) * (hv2){cs, cs};
+                tv = __builtin_elementwise_max(__builtin_elementwise_min(tv, (hv2){cl, cl}), (hv2){-cl, -cl});
+                const hv2 e1 = __builtin_elementwise_fma(tv, -tv, (hv2){p1, p1});
+                const hv2 poly = __builtin_elementwise_fma(tv, e1, (hv2){p0, p0});
+                const uint32_t mine = __builtin_bit_cast(uint32_t, poly);
+                const uint32_t other = xor_lane<1>(mine);
+                const uint32_t ua = (__builtin_amdgcn_perm(other, mine, sel_a) << 5) + 0x7FF88000u;
+                const uint32_t ub = (__builtin_amdgcn_perm(other, mine, sel_b) << 5) + 0x7FF88000u;
+                wg[k] = __builtin_bit_cast(hv2, __builtin_amdgcn_perm(ub, ua, sel_r));
+            }
+            wave_sync(); // Q scratch is dead: P overwrites it
+#pragma unroll
+            for (int k = 0; k < 4; ++k)
+            {
+                auto xr = [](hv2 v, auto mask) {
+                    return __builtin_bit_cast(hv2, xor_lane<decltype(mask)::value>(__builtin_bit_cast(uint32_t, v)));
+                };
+                const hv2 v1 = wg[k] + xr(wg[k], std::integral_constant<int, 8>{});
+                const hv2 v2 = v1 + xr(v1, std::integral_constant<int, 2>{});
+                const hv2 v3 = v2 + xr(v2, std::integral_constant<int, 4>{});
+                const hv2 sum = v3 + xr(v3, std::integral_constant<int, 1>{});
+                const hv2 rinv = {f16(1.0f / (float)sum[0]), f16(1.0f / (float)sum[1])};
+                const hv2 pw = wg[k] * rinv;
+                Qs[(4 * k + hi) * PST + l16] = pw[0];
+                Qs[(4 * k + 2 + hi) * PST + l16] = pw[1];
+            }
+            wave_sync();
+            const h16 pa = lds16(Qs + l16 * PST);
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                // V as B operand: lane holds column l16 over the 16 tokens (rows); own rows 2i+hi
+                // own pair k holds rows (4k+hi, 4k+2+hi); the other half-wave has (4k+1-hi, 4k+3-hi)
+                uint32_t vw[8];
+                const uint32_t s_lo = hi ? 0x01000504u : 0x05040100u, s_hi = hi ? 0x03020706u : 0x07060302u;
+#pragma unroll
+                for (int k = 0; k < 4; ++k)
+                {
+                    const uint32_t own = __builtin_bit_cast(uint32_t, vp[m][nt][k]);
+                    const uint32_t other = xor_lane<16>(own);
+                    vw[2 * k] = __builtin_amdgcn_perm(other, own, s_lo);     // rows 4k, 4k+1
+                    vw[2 * k + 1] = __builtin_amdgcn_perm(other, own, s_hi); // rows 4k+2, 4k+3
+                }
+                const h16 vb = __builtin_bit_cast(h16, vw);
+                const f8v o = wmma(pa, vb, splat(0.0f));
+                const int ocol = apos(32 * h + ginv(16 * nt + l16));
+#pragma unroll
+                for (int i = 0; i < 8; i += 2)
+                {
+                    const hv2 oq = q8x2((hv2){(half_t)o[i], (half_t)o[i + 1]});
+                    A[(16 * (m0 + m) + 2 * i + hi) * AST + ocol] = oq[0];
+                    A[(16 * (m0 + m) + 2 * i + 2 + hi) * AST + ocol] = oq[1];
+                }
+            }
+            wave_sync();
+        }
+    }
+    __syncthreads();
+
+    // ------------------------------------------------ stage 3: output projection + residual
+#pragma unroll
+    for (int kc = 0; kc < NH; ++kc)
+    {
+        h16 b[2][NTW];
+#pragma unroll
+        for (int s = 0; s < 2; ++s)
+#pragma unroll
+            for (int j = 0; j < NTW; ++j)
+                b[s][j] = bload(w16, L::WO, NT, kc, s, ng * NTW + j, l16);
+#pragma unroll
+        for (int mm = 0; mm < MT; ++mm)
+        {
+            const h16 a0 = lds16(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
+            const h16 a1 = lds16(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc + 16);
+#pragma unroll
+            for (int j = 0; j < NTW; ++j)
+                k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
+        }
+    }
+    __syncthreads(); // X0b and O are dead
+
+    // ------------------------------------------------ norm2 (kernel-exact tree), x1 staged in A
+#pragma unroll
+    for (int mm = 0; mm < MT; ++mm)
+#pragma unroll
+        for (int j = 0; j < NTW; ++j)
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                A[(16 * (mg * MT + mm) + 2 * i + hi) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
+    __syncthreads();
+    if (tid < 64)
+    {
+        // the row as 128-bit LDS loads: element 8 j + c of n8 tile j
+        h8 row[L::NT8];
+#pragma unroll
+        for (int j = 0; j < L::NT8; ++j)
+            row[j] = *(const h8*)(A + tid * AST + 8 * j);
+        half_t L8[8];
+#pragma unroll
+        for (int c = 0; c < 8; ++c)
+        {
+            half_t sq[L::NT8];
+#pragma unroll
+            for (int j = 0; j < L::NT8; ++j)
+                sq[j] = row[j][c] * row[j][c];
+            half_t even, odd;
+            if constexpr (L::NT8 == 16)
+            {
+                even = ((sq[0] + sq[2]) + (sq[4] + sq[6])) + ((sq[8] + sq[10]) + (sq[12] + sq[14]));
+                odd = ((sq[1] + sq[3]) + (sq[5] + sq[7])) + ((sq[9] + sq[11]) + (sq[13] + sq[15]));
+            }
+            else if constexpr (L::NT8 == 12)
+            {
+                even = ((sq[0] + sq[2]) + (sq[4] + sq[6])) + (sq[8] + sq[10]);
+                odd = ((sq[1] + sq[3]) + (sq[5] + sq[7])) + (sq[9] + sq[11]);
+            }
+            else
+            {
+                even = (sq[0] + sq[2]) + (sq[4] + sq[6]);
+                odd = (sq[1] + sq[3]) + (sq[5] + sq[7]);
+            }
+            L8[c] = even + odd;
+        }
+        const half_t b0 = (L8[0] + L8[2]) + (L8[4] + L8[6]);
+        const half_t b1 = (L8[1] + L8[3]) + (L8[5] + L8[7]);
+        R2[tid] = rsqrt16(b0 + b1);
+    }
+    __syncthreads();
+#pragma unroll
+    for (int j = 0; j < NTW; ++j)
+    {
+        const int n = 16 * (ng * NTW + j) + l16;
+        const half_t g2 = hload(W, L::G2 + 2 * n), b2 = hload(W, L::B2 + 2 * n);
+        const int col = apos(ginv(n));
+#pragma unroll
+        for (int mm = 0; mm < MT; ++mm)
+#pragma unroll
+            for (int i = 0; i < 8; i += 2)
+            {
+                const int row = 16 * (mg * MT + mm) + 2 * i + hi;
+                const hv2 v = x[mm][j].pair(i / 2);
+                const hv2 r2 = {R2[row], R2[row + 2]};
+                const hv2 hq = q8x2(v * (hv2)(r2 * (hv2){g2, g2}));
+                A[row * AST + col] = hq[0];
+                A[(row + 2) * AST + col] = hq[1];
+                x[mm][j].set_pair(i / 2, v + (hv2){b2, b2});
+            }
+    }
+    __syncthreads();
+
+    // ------------------------------------------------ stage 4: MLP, R chunks per round
+    {
+        constexpr int WPC = NW / R, MPW = 4 / WPC; // waves per chunk, fc1 m tiles per wave
+        const int cw = wv / WPC, mb = MPW * (wv % WPC);
+#pragma unroll 1
+        for (int r = 0; r < L::NC / R; ++r)
+        {
+            const int c = R * r + cw;
+            const int B1 = L::b1(c);
+            T16 z[MPW][2];
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                const half_t b1 = hload(W, B1 + 32 * C + 2 * (16 * nt + l16));
+#pragma unroll
+                for (int mm = 0; mm < MPW; ++mm)
+                    z[mm][nt] = t16_splat(b1);
+            }
+#pragma unroll
+            for (int kc = 0; kc < C / 32; ++kc)
+            {
+                h16 b[2][2];
+#pragma unroll
+                for (int s = 0; s < 2; ++s)
+#pragma unroll
+                    for (int nt = 0; nt < 2; ++nt)
+                        b[s][nt] = bload(w16, L::w1(c), 2, kc, s, nt, l16);
+#pragma unroll
+                for (int mm = 0; mm < MPW; ++mm)
+                {
+                    const h16 a0 = lds16(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
+                    const h16 a1 = lds16(A + (16 * (mb + mm) + l16) * AST + 32 * kc + 16);
+                    k32(z[mm][0], a0, b[0][0], a1, b[1][0]);
+                    k32(z[mm][1], a0, b[0][1], a1, b[1][1]);
+                }
+            }
+            half_t* Gc = G + cw * 64 * GST;
+#pragma unroll
+            for (int mm = 0; mm < MPW; ++mm)
+#pragma unroll
+                for (int nt = 0; nt < 2; ++nt)
+                {
+                    const int col = apos(ginv(16 * nt + l16));
+#pragma unroll
+                    for (int i = 0; i < 8; i += 2)
+                    {
+                        const hv2 zv = z[mm][nt].pair(i / 2);
+                        const hv2 cz = __builtin_elementwise_max(__builtin_elementwise_min(zv, (hv2){2.0f16, 2.0f16}),
+                                                                 (hv2){-2.0f16, -2.0f16});
+                        const hv2 az = __builtin_bit_cast(hv2, (u16x2)(__builtin_bit_cast(u16x2, cz) & (unsigned short)0x7fff));
+                        const hv2 inner = (hv2){f16(0.41216981f), f16(0.41216981f)} -
+                                          (hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
+                        const hv2 gl = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
+                        const hv2 gq = q8x2(gl);
+                        Gc[(16 * (mb + mm) + 2 * i + hi) * GST + col] = gq[0];
+                        Gc[(16 * (mb + mm) + 2 * i + 2 + hi) * GST + col] = gq[1];
+                    }
+                }
+            __syncthreads();
+#pragma unroll
+            for (int cc = 0; cc < R; ++cc)
+            {
+                const int c2 = R * r + cc;
+                h16 b[2][NTW];
+#pragma unroll
+                for (int s = 0; s < 2; ++s)
+#pragma unroll
+                    for (int j = 0; j < NTW; ++j)
+                        b[s][j] = bload(w16, L::w2(c2), NT, 0, s, ng * NTW + j, l16);
+#pragma unroll
+                for (int mm = 0; mm < MT; ++mm)
+                {
+                    const h16 a0 = lds16(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
+                    const h16 a1 = lds16(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST + 16);
+#pragma unroll
+                    for (int j = 0; j < NTW; ++j)
+                        k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
+                }
+            }
+            __syncthreads();
+        }
+    }
+
+    // ------------------------------------------------ output: f16 staged in LDS, 16-value segments encoded by a loop
+    half_t* XO = A; // [64 tokens][C channels] natural order
+#pragma unroll
+    for (int j = 0; j < NTW; ++j)
+    {
+        const int ch = ginv(16 * (ng * NTW + j) + l16);
+#pragma unroll
+        for (int mm = 0; mm < MT; ++mm)
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                XO[(16 * (mg * MT + mm) + 2 * i + hi) * C + ch] = x[mm][j].get(i);
+    }
+    __syncthreads();
+    uint8_t* OB = U; // e4m3 codes [64][C] for the patch merge
+#pragma unroll 1
+    for (int seg = tid; seg < 64 * C / 16; seg += 32 * NW)
+    {
+        const int T = seg / (C / 16), c16 = 16 * (seg % (C / 16));
+        uint32_t w[4];
+#pragma unroll
+        for (int k = 0; k < 4; ++k)
+        {
+            w[k] = 0;
+#pragma unroll
+            for (int b = 0; b < 4; ++b)
+                w[k] |= enc8(XO[T * C + c16 + 4 * k + b]) << (8 * b);
+        }
+        const uint4 v = {w[0], w[1], w[2], w[3]};
+        if constexpr (NPM > 0)
+            *(uint4*)(OB + T * C + c16) = v;
+        const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
+        if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
+            *(uint4*)(p.out + ((size_t)((c16 >> 5) * p.th + Y) * p.tw + X) * 32 + (c16 & 31)) = v;
+    }
+    if constexpr (NPM > 0)
+        __syncthreads();
+
+    // ------------------------------------------------ patch merge: 2x2 tokens (4C) -> NPM, half resolution
+    if constexpr (NPM > 0)
+    {
+        // merge tiles split evenly over the largest wave count dividing them (remaining waves idle)
+        constexpr int PT = NPM / 16 ? NPM / 16 : 1;
+        constexpr int PMW = PT % NW == 0 ? NW : (PT % 4 == 0 && NW >= 4 ? 4 : (PT % 3 == 0 && NW >= 3 ? 3 : (PT % 2 == 0 && NW >= 2 ? 2 : 1)));
+        constexpr int NTP = PT / PMW;
+        constexpr bool PM_EVEN = true;
+        if (wv < PMW)
+        {
+        T16 y[NTP];
+#pragma unroll
+        for (int j = 0; j < NTP; ++j)
+            y[j] = t16_splat(hload(W, L::PMB + 2 * (16 * (PM_EVEN ? wv * NTP + j : (wv * NTP + j) % (NPM / 16)) + l16)));
+        const int mx = l16 & 3, my = l16 >> 2;
+#pragma unroll
+        for (int kc = 0; kc < C / 8; ++kc)
+        {
+            h16 as[2];
+#pragma unroll
+            for (int s = 0; s < 2; ++s)
+            {
+                h16& a = as[s];
+#pragma unroll
+                for (int run = 0; run < 4; ++run)
+                {
+                    const int k = 32 * kc + kslot(s, 4 * run);
+                    const int q = k / C, ch = k % C;
+                    const int tx = 2 * mx + (q & 1), ty = 2 * my + (q >> 1);
+                    const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
+                    const uint32_t bytes = *(const uint32_t*)(OB + T * C + ch);
+#pragma unroll
+                    for (int b = 0; b < 4; ++b)
+                        a[4 * run + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
+                }
+            }
+#pragma unroll
+            for (int j = 0; j < NTP; ++j)
+            {
+                // out-of-range tiles (uneven split) recompute a valid one and are not stored
+                const int t16 = PM_EVEN ? wv * NTP + j : (wv * NTP + j) % (NPM / 16);
+                k32(y[j], as[0], bload(w16, L::pm(t16 >> 1), 2, kc, 0, t16 & 1, l16), as[1],
+                    bload(w16, L::pm(t16 >> 1), 2, kc, 1, t16 & 1, l16));
+            }
+        }
+        const int W2 = p.tw / 2, H2 = p.th / 2;
+#pragma unroll
+        for (int j = 0; j < NTP; ++j)
+        {
+            if (!PM_EVEN && wv * NTP + j >= NPM / 16)
+                continue;
+            const int ch = ginv(16 * (wv * NTP + j) + l16);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const int r = 2 * i + hi;
+                const int MX = (8 * bx - p.sx) / 2 + (r & 3), MY = (8 * by - p.sy) / 2 + (r >> 2);
+                if (MX < 0 || MX >= W2 || MY < 0 || MY >= H2)
+                    continue;
+                p.p48[((size_t)((ch >> 5) * H2 + MY) * W2 + MX) * 32 + (ch & 31)] =
+                    (uint8_t)enc8(y[j].get(i));
+            }
+        }
+    }
+    }
+
+    if constexpr (TUBE)
+    {
+        if (tp->out_flags != nullptr)
+        {
+            __builtin_amdgcn_fence(__ATOMIC_RELEASE, "agent");
+            __syncthreads();
+            if (tid == 0)
+                __hip_atomic_store(tp->out_flags + (size_t)(by * gx + bx) * 32, (uint8_t)1, __ATOMIC_RELEASE,
+                                   __HIP_MEMORY_SCOPE_AGENT);
+        }
+    }
+}
+
+#if defined(SWIN_WAVES_PER_EU)
+#define SWIN_VGPR_ATTR __attribute__((amdgpu_waves_per_eu(SWIN_WAVES_PER_EU)))
+#elif defined(SWIN_MAX_VGPR)
+#define SWIN_VGPR_ATTR __attribute__((amdgpu_num_vgpr(SWIN_MAX_VGPR)))
+#else
+#define SWIN_VGPR_ATTR
+#endif
+#ifndef SWIN_PERSIST
+#define SWIN_PERSIST 0 // persistent workgroups in x (0: one workgroup per original block)
+#endif
+// Module boilerplate: f16 weight image, prep kernel and its grid for the ZLUDA override hook.
+#define SWIN_MODULE(NAME, C, NW, NPM, TUBE, PARAMS, CIN) SWIN_MODULE_W(NAME, C, NW, NW, NPM, TUBE, PARAMS, CIN)
+// NH heads on NWAVES waves; d4r_block_z tells the ZLUDA hook the block z size to launch with
+#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN)                                                      \
+    using NAME##_L = SwinLayout<C, NH, NPM, CIN>;                                                                       \
+    __device__ h16 g_w16[NAME##_L::TOTAL];                                                                              \
+    __constant__ SwinDescs<C, NH, NPM, CIN> g_descs = make_descs<C, NH, NPM, CIN>();                                    \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = NWAVES;                                         \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::TOTAL + 127) / 128;              \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PARAMS p)                                             \
+    {                                                                                                                   \
+        const int idx = blockIdx.x * 128 + threadIdx.x;                                                                 \
+        if (idx < NAME##_L::TOTAL)                                                                                      \
+            expand_weights(((const CommonParams*)&p)->w, g_w16, g_descs.d, NAME##_L::NDESC, idx);                      \
+    }                                                                                                                   \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = SWIN_PERSIST;                                     \
+    extern "C" __global__ void __launch_bounds__(32 * NWAVES) SWIN_VGPR_ATTR NAME(PARAMS p)                             \
+    {                                                                                                                   \
+        const CommonParams& cp = *(const CommonParams*)&p;                                                              \
+        /* the original grid: 8x8-token blocks covering the shifted token grid */                                        \
+        const int gx = (cp.tw + cp.sx + 7) / 8, gy = (cp.th + cp.sy + 7) / 8;                                           \
+        if (SWIN_PERSIST == 0)                                                                                          \
+        {                                                                                                               \
+            swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, blockIdx.x, blockIdx.y, gx);    \
+            return;                                                                                                     \
+        }                                                                                                               \
+        _Pragma("unroll 1") for (int b = blockIdx.x; b < gx * gy; b += gridDim.x)                                      \
+        {                                                                                                               \
+            swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, b % gx, b / gx, gx);            \
+            __syncthreads();                                                                                            \
+        }                                                                                                               \
+    }

kernels/m/swin_common.h

@@ -0,0 +1,330 @@
+// Shared device code for native RDNA3 (gfx11 wave32 WMMA) rrlite Swin-block kernels (DLSS 310.7.0).
+// Semantics: ../ENC3_SPEC.md generalised (numpy reference ../swin_model.py).
+#pragma once
+#include <hip/hip_runtime.h>
+#include <stdint.h>
+
+#pragma clang fp contract(off)
+
+typedef _Float16 half_t;
+typedef _Float16 h16 __attribute__((ext_vector_type(16)));
+typedef _Float16 h8 __attribute__((ext_vector_type(8)));
+typedef float f8v __attribute__((ext_vector_type(8)));
+
+struct CommonParams
+{
+    const uint8_t* w; // 0
+    int sx, sy;       // 8
+    int tw, th;       // 16
+    const uint8_t* in;  // 24
+    const uint8_t* p32; // 32 (decoders: skip input)
+    uint8_t* out;       // 40
+    uint8_t* p48;       // 48 (encoders with patch merge: merged output)
+};
+static_assert(sizeof(CommonParams) == 56, "param block");
+
+struct TubeParams
+{
+    CommonParams c;
+    int r69, r70;
+    int r71, r72;
+    const uint8_t* in_flags;
+    uint8_t* out_flags;
+};
+static_assert(sizeof(TubeParams) == 88, "param block");
+
+// ---------------------------------------------------------------- numerics
+__device__ __forceinline__ half_t e4m3_to_half(uint32_t code)
+{
+    uint16_t bits = (uint16_t)(((code & 0x7fu) << 7) | ((code & 0x80u) << 8));
+    return __builtin_bit_cast(half_t, bits) * (half_t)256.0f;
+}
+
+__device__ __forceinline__ uint32_t half_to_e4m3(half_t h)
+{
+    uint32_t half = __builtin_bit_cast(uint16_t, h);
+    uint32_t magnitude = half & 0x7fffu;
+    uint32_t normal = (magnitude - 0x2000u + 0x3fu + ((magnitude >> 7) & 1u)) >> 7;
+    uint32_t exponent = magnitude >> 10;
+    uint32_t significand = (magnitude & 0x3ffu) | (exponent != 0 ? 0x400u : 0u);
+    uint32_t clamped = exponent < 1u ? 1u : (exponent > 8u ? 8u : exponent);
+    uint32_t shift = 16u - clamped;
+    uint32_t subnormal = (significand + (1u << (shift - 1u)) - 1u + ((significand >> shift) & 1u)) >> shift;
+    uint32_t code = magnitude >= 0x2400u ? normal : subnormal;
+    code = magnitude > 0x5f00u ? 0x7eu : code;
+    code = magnitude > 0x7c00u ? 0x7fu : code;
+    return ((half >> 8) & 0x80u) | code;
+}
+
+// e4m3 RNE satfinite round trip evaluated in the f16 domain (exhaustively equal to decode(encode(h)))
+__device__ __forceinline__ half_t q8_exact(half_t h);
+__device__ __forceinline__ half_t q8(half_t h)
+{
+#ifdef SWIN_NO_Q8
+    return h; // fast mode: intermediate activations stay f16 (stored outputs are still FP8)
+#endif
+    return q8_exact(h);
+}
+__device__ __forceinline__ half_t q8_exact(half_t h)
+{
+    const uint32_t u = __builtin_bit_cast(uint16_t, h);
+    const uint32_t m = u & 0x7fffu;
+    uint32_t r = (m + 0x3fu + ((m >> 7) & 1u)) & 0x7f80u;
+    r = r < 0x5f00u ? r : 0x5f00u;
+    const half_t mag = __builtin_bit_cast(half_t, (uint16_t)m);
+    const uint32_t sub = __builtin_bit_cast(uint16_t, (half_t)((half_t)(mag + (half_t)2.0f) - (half_t)2.0f));
+    r = m < 0x2400u ? sub : r;
+    r = m > 0x7c00u ? m : r;
+    return __builtin_bit_cast(half_t, (uint16_t)((u & 0x8000u) | r));
+}
+
+typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
+typedef unsigned short u16x2 __attribute__((ext_vector_type(2)));
+typedef short i16x2 __attribute__((ext_vector_type(2)));
+
+// q8 on two halves with packed 16-bit integer / f16 ops (same per-half logic as q8; no carries cross halves)
+__device__ __forceinline__ hv2 q8x2(hv2 h)
+{
+#ifdef SWIN_NO_Q8
+    return h;
+#endif
+    const u16x2 u = __builtin_bit_cast(u16x2, h);
+    const u16x2 m = u & (unsigned short)0x7fff;
+    u16x2 r = (m + (unsigned short)0x3f + ((m >> 7) & (unsigned short)1)) & (unsigned short)0x7f80;
+    r = __builtin_elementwise_min(r, (u16x2){0x5f00, 0x5f00});
+    const hv2 mag = __builtin_bit_cast(hv2, m);
+    const hv2 two = {(half_t)2.0f, (half_t)2.0f};
+    const u16x2 sub = __builtin_bit_cast(u16x2, (hv2)((hv2)(mag + two) - two));
+    const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
+    r = (sub & lt) | (r & ~lt);
+    const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
+    r = (m & gt) | (r & ~gt);
+    return __builtin_bit_cast(hv2, (u16x2)(r | (u & (unsigned short)0x8000)));
+}
+
+// e4m3 code of h (RNE, satfinite) derived from q8(h): normals shift the exponent, subnormals are k * 2^-9
+__device__ __forceinline__ uint32_t enc8(half_t h)
+{
+    const uint32_t qb = __builtin_bit_cast(uint16_t, q8_exact(h));
+    const uint32_t m = qb & 0x7fffu;
+    const half_t mag = __builtin_bit_cast(half_t, (uint16_t)m);
+    uint32_t code = m >= 0x2400u ? (m >> 7) - 64u : (uint32_t)(uint16_t)(mag * (half_t)512.0f);
+    code = m > 0x7c00u ? 0x7fu : code;
+    return ((qb >> 8) & 0x80u) | code;
+}
+
+__device__ __forceinline__ half_t f16(float v)
+{
+    return (half_t)v;
+}
+
+__device__ __forceinline__ half_t rsqrt16(half_t s)
+{
+    return f16(1.0f / __builtin_sqrtf((float)(s + (half_t)(1.0f / 8192.0f))));
+}
+
+// natural channel c -> pair-order column, and inverse (per 32-channel group)
+__device__ __forceinline__ int gperm(int c)
+{
+    return 32 * (c >> 5) + 2 * ((c & 31) >> 2) + (c & 1) + 16 * ((c & 3) >> 1);
+}
+__device__ __forceinline__ int ginv(int n)
+{
+    int m = n & 31, e = m >> 4, a = (m & 15) >> 1, b = m & 1;
+    return (n & ~31) + 4 * a + 2 * e + b;
+}
+
+// WMMA k16 step s of a k32 chunk: operand slot i holds logical k = kslot(s, i) (matches the weight
+// memory runs). Activations live in LDS at apos(k), so an A operand is 16 contiguous halves.
+__device__ __forceinline__ int kslot(int s, int i)
+{
+    int a = 2 * s + (i >> 3), r = i & 7;
+    return (r < 4) ? (4 * a + r) : (16 + 4 * a + (r - 4));
+}
+__device__ __forceinline__ int apos(int k)
+{
+    return (k & ~31) + 8 * ((k & 15) >> 2) + 4 * ((k >> 4) & 1) + (k & 3);
+}
+
+__device__ __forceinline__ int woff(int base, int Ks, int Ns, int k, int n)
+{
+    int j = n >> 3;
+    return base + Ks * (k >> 5) + Ns * (j >> 1) + 64 * (n & 7) + 16 * ((k & 15) >> 2) + 8 * (j & 1) +
+           4 * ((k & 31) >> 4) + (k & 3);
+}
+
+__device__ __forceinline__ f8v wmma(h16 a, h16 b, f8v c)
+{
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
+}
+
+__device__ __forceinline__ f8v round16(f8v v)
+{
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        v[i] = (float)(half_t)v[i];
+    return v;
+}
+
+// A C-fragment tile: element i is row 2i + lane/16. Exact mode stores packed f16 (NVIDIA's f16
+// accumulator, rounded after every k32 step); SWIN_F32ACC keeps f32 through whole GEMM chains
+// (more precise than NVIDIA; values are rounded to f16 where they are consumed).
+struct T16
+{
+#ifdef SWIN_F32ACC
+    f8v f;
+    __device__ hv2 pair(int k) const { return (hv2){(half_t)f[2 * k], (half_t)f[2 * k + 1]}; }
+    __device__ void set_pair(int k, hv2 p) { f[2 * k] = (float)p[0]; f[2 * k + 1] = (float)p[1]; }
+    __device__ half_t get(int i) const { return (half_t)f[i]; }
+    __device__ void set(int i, half_t h) { f[i] = (float)h; }
+#else
+    hv2 v[4];
+    __device__ hv2 pair(int k) const { return v[k]; }
+    __device__ void set_pair(int k, hv2 p) { v[k] = p; }
+    __device__ half_t get(int i) const { return v[i / 2][i & 1]; }
+    __device__ void set(int i, half_t h) { v[i / 2][i & 1] = h; }
+#endif
+};
+
+__device__ __forceinline__ T16 t16_splat(half_t h)
+{
+    T16 t;
+#pragma unroll
+    for (int k = 0; k < 4; ++k)
+        t.set_pair(k, (hv2){h, h});
+    return t;
+}
+
+#ifndef MIX_PAD
+#define MIX_PAD ""
+#endif
+// 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)
+{
+    hv2 r = c;
+    asm(MIX_PAD "v_fma_mixlo_f16 %0, %1, 1.0, %2 op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p0), "v"(c));
+    asm("v_fma_mixhi_f16 %0, %1, 1.0, %2 op_sel:[0,0,1] op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p1), "v"(c));
+    return r;
+}
+__device__ __forceinline__ void mixacc(T16& acc, f8v P)
+{
+#pragma unroll
+    for (int k = 0; k < 4; ++k)
+        acc.set_pair(k, mix2(P[2 * k], P[2 * k + 1], acc.pair(k)));
+}
+
+#ifdef K32_CHECK
+__device__ unsigned d4r_k32_bad[4 + 64 * 4];
+#endif
+// f16 accumulate of one k32 step (NVIDIA m16n8k32 semantics: products summed, one rounding with C)
+__device__ __forceinline__ void k32(T16& acc, h16 a0, h16 b0, h16 a1, h16 b1)
+{
+#if defined(SWIN_F32ACC)
+    acc.f = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, acc.f));
+#elif defined(K32_MIX)
+    f8v P = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
+    P = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, P);
+    mixacc(acc, P);
+#else // exact: C added inside the f32 WMMA, f16 rounding after the k32 step
+    f8v c;
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        c[i] = (float)acc.get(i);
+    c = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, c));
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        acc.set(i, (half_t)c[i]);
+#endif
+}
+
+__device__ __forceinline__ f8v splat(float v)
+{
+    return (f8v){v, v, v, v, v, v, v, v};
+}
+
+// lane xor shuffles without LDS: DPP within rows of 16, permlanex16 across the two rows
+template <int MASK> __device__ __forceinline__ uint32_t xor_lane(uint32_t v)
+{
+    if constexpr (MASK == 1)
+        return __builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xf, 0xf, false);
+    else if constexpr (MASK == 2)
+        return __builtin_amdgcn_update_dpp(0, (int)v, 0x4E, 0xf, 0xf, false);
+    else if constexpr (MASK == 4)
+        return __builtin_amdgcn_update_dpp(0, (int)v, 0x164, 0xf, 0xf, false);
+    else if constexpr (MASK == 8)
+        return __builtin_amdgcn_update_dpp(0, (int)v, 0x168, 0xf, 0xf, false);
+    else
+        return __builtin_amdgcn_permlanex16((int)v, (int)v, 0x76543210, 0xfedcba98, false, false);
+}
+template <int MASK> __device__ __forceinline__ half_t hxor(half_t v)
+{
+    return __builtin_bit_cast(half_t, (uint16_t)xor_lane<MASK>(__builtin_bit_cast(uint16_t, v)));
+}
+
+__device__ __forceinline__ int mirror(int v, int n)
+{
+    v = v < 0 ? -v : v;
+    int o = 2 * n - 2 - v;
+    return v < o ? v : o;
+}
+
+__device__ __forceinline__ half_t hload(const uint8_t* w, int off)
+{
+    return __builtin_bit_cast(half_t, *(const uint16_t*)(w + off));
+}
+
+__device__ __forceinline__ h16 lds16(const half_t* p)
+{
+    h8 lo = *(const h8*)p, hi = *(const h8*)(p + 8);
+    return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
+}
+
+__device__ __forceinline__ void wave_sync()
+{
+    __builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
+    __builtin_amdgcn_wave_barrier();
+    __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");
+}
+
+// Token T (0..63) of a block: 8x8 tokens, windows of 4x4 (T >> 4)
+__device__ __forceinline__ int tok_x(int T)
+{
+    return (T & 3) + 4 * ((T >> 4) & 1);
+}
+__device__ __forceinline__ int tok_y(int T)
+{
+    return ((T >> 2) & 3) + 4 * (T >> 5);
+}
+
+// ---------------------------------------------------------------- f16 weight images
+// A GEMM's B operands expanded to f16 in WMMA lane order: h16 index
+// dst + ((kc * 2 + s) * NT + nt) * 16 + lane16, holding logical k = 32 kc + kslot(s, i), n = 16 nt + lane16.
+struct GemmDesc
+{
+    int dst; // in h16 units
+    int base, Ks, Ns, KC, NT;
+};
+
+__device__ __forceinline__ h16 bload(const h16* w16, int gemm_dst, int NT, int kc, int s, int nt, int l16)
+{
+    return w16[gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16];
+}
+
+__device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const GemmDesc* descs, int ndesc, int idx)
+{
+    int d = 0;
+    while (d + 1 < ndesc && idx >= descs[d + 1].dst)
+        ++d;
+    const GemmDesc g = descs[d];
+    const int local = idx - g.dst;
+    if (local >= g.KC * 2 * g.NT * 16)
+        return;
+    const int l16 = local & 15, rest = local >> 4;
+    const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
+    const int n = 16 * nt + l16;
+    h16 v;
+#pragma unroll
+    for (int i = 0; i < 16; ++i)
+        v[i] = e4m3_to_half(w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, i), n)]);
+    w16[idx] = v;
+}

kernels/tex/build_tex.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+# build_tex.sh KERNEL SRC OUT_DIR: one texture kernel. NVIDIA's PTX for KERNEL, edited by make_ptx.py so
+# part of it calls functions from SRC.hip, is compiled by ZLUDA with SRC's bitcode linked in
+# (D4R_ZLUDA_EXTRA_BC) and the resulting code object is saved as OUT_DIR/KERNEL.hsaco.
+# Called by kernels/build.sh, which sets D4R_ROCM_DIR, D4R_GPU_ARCH, D4R_DLSS_PTX_DIR, D4R_PTXLOAD and
+# needs D4R_ZLUDA_BUILD.
+set -euo pipefail
+K=$1 SRC=$2 OUT=$3
+D="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+R="${D4R_ROCM_DIR:-/opt/rocm}"
+ARCH="${D4R_GPU_ARCH:-gfx1101}"
+LLVM="$R/lib/llvm/bin"
+B="$(mktemp -d)"
+trap 'rm -rf "$B"' EXIT
+# SRC.hip -> bitcode without target attributes, so ZLUDA's own target settings apply after linking
+"$LLVM/clang++" -x hip -std=c++20 -nogpuinc -nogpulib -O3 -mno-wavefrontsize64 --offload-device-only \
+    --offload-arch="$ARCH" -fgpu-rdc -emit-llvm -c -Xclang -fdenormal-fp-math=dynamic -o "$B/raw.bc" "$D/$SRC.hip"
+"$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" -
+python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
+env LD_LIBRARY_PATH="$D4R_ZLUDA_BUILD:$R/lib" XDG_CACHE_HOME="$B/jit" D4R_ZLUDA_DUMP_DIR="$B/dump" \
+    D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 D4R_ZLUDA_IGNORE_DENORMAL=1 \
+    D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 "$D4R_PTXLOAD" "$B/$K.ptx" >/dev/null
+cp "$B/dump/module.hsaco" "$OUT/$K.hsaco"
+printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"

kernels/tex/dec0_head.hip

@@ -0,0 +1,112 @@
+// Native head of rrlite_dec0_4x4_folded: the token GEMM (16 half-resolution tokens x K 64 (two
+// e4m3 planes) x N 128, f16 bias, f16 accumulation per k32 step) and the e4m3 staging of its result
+// in shared memory, which the rest of the kernel (ZLUDA's compile of NVIDIA's PTX, dec0_head.ptx)
+// reads. Arithmetic as ZLUDA's m16n8k32 lowering (see enc0_tail.hip); tokens are mirror padded.
+#include "tex_common.h"
+
+// weights as WMMA A images: index ((p * 2 + c) * 2 + h) * 16 + n16 for channel pair p (16p..16p+15),
+// k32 chunk c (input plane), k16 half h
+__attribute__((device)) u8v g_dec0_w16[8 * 2 * 2 * 16];
+DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
+
+struct Dec0Params
+{
+    uint64_t w;
+    uint8_t rest[160]; // 168-byte parameter block
+};
+
+KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4x4_folded_prep(Dec0Params prm)
+{
+    const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
+    if (idx >= 512)
+        return;
+    const uint8_t* w = (const uint8_t*)prm.w;
+    const int n16 = idx & 15, h = (idx >> 4) & 1, c = (idx >> 5) & 1, p = idx >> 6;
+    const int g = n16 & 7, second = n16 >> 3;
+    u4v d;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+        d[j] = *(const uint32_t*)(w + 512 * (2 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
+    g_dec0_w16[idx] = operand16(d);
+}
+
+__attribute__((device)) static inline uint32_t mirror_clamp(int32_t v, uint32_t full, uint32_t half)
+{
+    const int32_t a = v < 0 ? -v : v;
+    const int32_t m = __builtin_elementwise_min(a, (int32_t)(full & ~1u) - a - 2);
+    return __builtin_elementwise_min(half - 1u, (uint32_t)m);
+}
+
+DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx, uint32_t dimy, int32_t ox, int32_t oy)
+{
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const uint8_t* w = (const uint8_t*)wp;
+    lds_u8* s = (lds_u8*)(uintptr_t)smem;
+
+    const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
+    const uint32_t Y = mirror_clamp(oy / 2 + (int32_t)(m >> 2), dimy, H2);
+    const uint32_t Xc = mirror_clamp(ox / 2 + (int32_t)(m & 3), dimx, W2);
+    const uint32_t* in = (const uint32_t*)inp + rowstride * Y + 8 * Xc;
+    u8v X[2][2];
+#pragma unroll
+    for (int c = 0; c < 2; ++c)
+#pragma unroll
+        for (int h = 0; h < 2; ++h)
+            X[c][h] = operand16(*(const u4v*)(in + plane * c + 4 * h));
+
+    uint32_t E[8][2];
+#pragma unroll
+    for (int p = 0; p < 8; ++p)
+    {
+        half_t acc[8];
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+            acc[i] = *(const half_t*)(w + 8192 + 2 * (16 * p + 2 * i + hf));
+#pragma unroll
+        for (int c = 0; c < 2; ++c)
+        {
+            f8v cf;
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                cf[i] = (float)acc[i];
+            cf *= 0x1p-16f;
+            const u8v* wv = &g_dec0_w16[((p * 2 + c) * 2) * 16 + m];
+            cf = wmma(wv[0], X[c][0], cf);
+            cf = wmma(wv[16], X[c][1], cf);
+            cf = cf * 0x1p16f;
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[i] = (half_t)cf[i];
+        }
+        uint32_t codes[4];
+#pragma unroll
+        for (int k = 0; k < 4; ++k)
+            codes[k] = e4m3x2(pack2(acc[2 * k], acc[2 * k + 1]));
+        E[p][0] = codes[0] | (codes[1] << 16);
+        E[p][1] = codes[2] | (codes[3] << 16);
+    }
+
+    // staging: word t of token m at 1024 * H + 32 * (m & 3) + 256 * (m >> 2) + {0, 16, 128, 144}[G] + 4t
+    // holds columns 2t, 2t + 1 of tiles 4G + H and 4G + 2 + H; lanes 0..15 write H = 0, 16..31 H = 1
+    uint32_t EV[8], OD[8];
+#pragma unroll
+    for (int p = 0; p < 8; ++p)
+    {
+        const uint32_t theirs = other_half(hf ? E[p][0] : E[p][1]);
+        EV[p] = hf ? theirs : E[p][0];
+        OD[p] = hf ? E[p][1] : theirs;
+    }
+    const uint32_t base = 1024 * hf + 32 * (m & 3) + 256 * (m >> 2);
+    const uint32_t goff[4] = {0, 16, 128, 144};
+#pragma unroll
+    for (int G = 0; G < 4; ++G)
+    {
+        const uint32_t a01 = __builtin_amdgcn_perm(OD[2 * G], EV[2 * G], 0x05010400u);
+        const uint32_t a23 = __builtin_amdgcn_perm(OD[2 * G], EV[2 * G], 0x07030602u);
+        const uint32_t b01 = __builtin_amdgcn_perm(OD[2 * G + 1], EV[2 * G + 1], 0x05010400u);
+        const uint32_t b23 = __builtin_amdgcn_perm(OD[2 * G + 1], EV[2 * G + 1], 0x07030602u);
+        u4v words = {__builtin_amdgcn_perm(b01, a01, 0x05040100u), __builtin_amdgcn_perm(b01, a01, 0x07060302u),
+                     __builtin_amdgcn_perm(b23, a23, 0x05040100u), __builtin_amdgcn_perm(b23, a23, 0x07060302u)};
+        *(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
+    }
+}

kernels/tex/enc0_tail.hip

@@ -0,0 +1,124 @@
+// Native tail of rrlite_enc0_4x4_mvhi_hdr_folded: the 2x2-pixel patch embedding GEMM (16 tokens x
+// K 128 x N 64, e4m3 inputs, f16 accumulation per k32 step, f16 bias), clamp to +-2pi and e4m3 store.
+// The feature code before it is ZLUDA's compile of the original PTX (enc0_tail.ptx calls this after
+// the kernel's own shared-memory staging); the arithmetic reproduces ZLUDA's m16n8k32 lowering
+// (scaled e4m3 widening, same k slot order, C * 2^-16 inside the f32 WMMA, f16 rounding per k32).
+#include "tex_common.h"
+
+// B operands (weights) as WMMA A images: index ((p * 4 + c) * 2 + h) * 16 + n16 for channel pair p
+// (channels 16p..16p+15), k32 chunk c, k16 half h.
+__attribute__((device)) u8v g_enc0_w16[4 * 4 * 2 * 16];
+DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
+
+struct Enc0Params
+{
+    uint64_t w;
+    uint8_t rest[352];
+};
+
+KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4x4_mvhi_hdr_folded_prep(Enc0Params prm)
+{
+    const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
+    if (idx >= 512)
+        return;
+    const uint8_t* w = (const uint8_t*)prm.w;
+    const int n16 = idx & 15, h = (idx >> 4) & 1, c = (idx >> 5) & 3, p = idx >> 7;
+    // NVIDIA B fragments: v4 load 4p + c of lane 4 * column + j holds {b0, b1} of tile 2p then 2p + 1
+    const int g = n16 & 7, second = n16 >> 3;
+    u4v d;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+        d[j] = *(const uint32_t*)(w + 4096 + 512 * (4 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
+    g_enc0_w16[idx] = operand16(d);
+}
+
+DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx, uint32_t dimy, int32_t ox, int32_t oy)
+{
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const uint8_t* w = (const uint8_t*)wp;
+    lds_u8* s = (lds_u8*)(uintptr_t)smem;
+
+    // activations: token m's 128 e4m3 features; A[m][k] of chunk c at
+    // 1024 * (k >> 4) + 32 * (m & 3) + 256 * (m >> 2) + (k & 15) + {0, 16, 128, 144}[c]
+    const uint32_t rowoff = 32 * (m & 3) + 256 * (m >> 2);
+    const uint32_t choff[4] = {0, 16, 128, 144};
+    u8v X[4][2];
+#pragma unroll
+    for (int c = 0; c < 4; ++c)
+#pragma unroll
+        for (int h = 0; h < 2; ++h)
+            X[c][h] = operand16(*(const __attribute__((address_space(3))) u4v*)(s + 1024 * h + rowoff + choff[c]));
+
+    const hv2 lo = {(half_t)-6.2831855f, (half_t)-6.2831855f}, hi = {(half_t)6.2831855f, (half_t)6.2831855f};
+    uint32_t E[4][2];
+#pragma unroll
+    for (int p = 0; p < 4; ++p)
+    {
+        // D^T layout: element i is channel 16p + 2i + hf of token m
+        half_t acc[8];
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+            acc[i] = *(const half_t*)(w + 12288 + 2 * (16 * p + 2 * i + hf));
+#pragma unroll
+        for (int c = 0; c < 4; ++c)
+        {
+            f8v cf;
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                cf[i] = (float)acc[i];
+            cf *= 0x1p-16f;
+            const u8v* wv = &g_enc0_w16[((p * 4 + c) * 2) * 16 + m];
+            cf = wmma(wv[0], X[c][0], cf);
+            cf = wmma(wv[16], X[c][1], cf);
+            cf = cf * 0x1p16f;
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[i] = (half_t)cf[i];
+        }
+        uint32_t codes[4];
+#pragma unroll
+        for (int k = 0; k < 4; ++k)
+        {
+            hv2 v = {acc[2 * k], acc[2 * k + 1]};
+            v = __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, lo), hi);
+            codes[k] = e4m3x2(__builtin_bit_cast(uint32_t, v));
+        }
+        // byte i of E[p][0] (E[p][1]) is column 2i + hf (2(i - 4) + hf) of tile 2p (2p + 1)
+        E[p][0] = codes[0] | (codes[1] << 16);
+        E[p][1] = codes[2] | (codes[3] << 16);
+    }
+
+    // lanes 0..15 write plane 0 (tiles 0..3), lanes 16..31 plane 1 (tiles 4..7): swap the other
+    // plane's codes with the partner lane holding the other column parity
+    uint32_t mine[2][2], theirs[2][2];
+#pragma unroll
+    for (int a = 0; a < 2; ++a)
+#pragma unroll
+        for (int j = 0; j < 2; ++j)
+        {
+            const uint32_t send = hf ? E[a][j] : E[2 + a][j];
+            theirs[a][j] = other_half(send);
+            mine[a][j] = hf ? E[2 + a][j] : E[a][j];
+        }
+    u8v words;
+#pragma unroll
+    for (int j = 0; j < 2; ++j)
+    {
+        const uint32_t ev0 = hf ? theirs[0][j] : mine[0][j], od0 = hf ? mine[0][j] : theirs[0][j];
+        const uint32_t ev1 = hf ? theirs[1][j] : mine[1][j], od1 = hf ? mine[1][j] : theirs[1][j];
+        const uint32_t a01 = __builtin_amdgcn_perm(od0, ev0, 0x05010400u), a23 = __builtin_amdgcn_perm(od0, ev0, 0x07030602u);
+        const uint32_t b01 = __builtin_amdgcn_perm(od1, ev1, 0x05010400u), b23 = __builtin_amdgcn_perm(od1, ev1, 0x07030602u);
+        words[4 * j + 0] = __builtin_amdgcn_perm(b01, a01, 0x05040100u);
+        words[4 * j + 1] = __builtin_amdgcn_perm(b01, a01, 0x07060302u);
+        words[4 * j + 2] = __builtin_amdgcn_perm(b23, a23, 0x05040100u);
+        words[4 * j + 3] = __builtin_amdgcn_perm(b23, a23, 0x07060302u);
+    }
+
+    const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
+    const int32_t Y = oy / 2 + (int32_t)(m >> 2), Xc = ox / 2 + (int32_t)(m & 3);
+    if ((uint32_t)Y < H2 && (uint32_t)Xc < W2)
+    {
+        uint32_t* out = (uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * hf;
+        *(u8v*)out = words;
+    }
+}

kernels/tex/make_ptx.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""Hand-edited PTX for the texture kernels: keep NVIDIA's code up to a cut line, then call a native tail.
+
+usage: make_ptx.py KERNEL OUT.ptx
+"""
+import os
+import re
+import sys
+from pathlib import Path
+
+# PTX extracted from nvngx_dlss.dll by kernels/tools/extract_dlss_ptx.py (kernels/build.sh does this)
+PTX_DIR = Path(os.environ.get("D4R_DLSS_PTX_DIR", Path(__file__).resolve().parent.parent / "extracted" / "ptx"))
+
+KERNELS = {
+    # enc0: cut after the bar.warp.sync that ends the smem staging of the GEMM A operand
+    "rrlite_enc0_4x4_mvhi_hdr_folded": dict(
+        file="dlss-0185-00.ptx",
+        cut_after=2548,
+        cut_check="bar.warp.sync -1;",
+        sust=True,
+        extern=""".extern .func d4r_enc0_tail
+(
+	.param .b32 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
+)
+;
+""",
+        tail="""{
+	.param .b32 q0;
+	st.param.b32 [q0+0], %r259;
+	.param .b64 q1;
+	st.param.b64 [q1+0], %rd106;
+	.param .b64 q2;
+	st.param.b64 [q2+0], %rd78;
+	.param .b32 q3;
+	st.param.b32 [q3+0], %r1066;
+	.param .b32 q4;
+	st.param.b32 [q4+0], %r1067;
+	.param .b32 q5;
+	st.param.b32 [q5+0], %r1230;
+	.param .b32 q6;
+	st.param.b32 [q6+0], %r1233;
+	call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
+}
+ret;
+
+}
+""",
+    ),
+    # dec0: the GEMM (from the B fragment loads to the e4m3 staging stores) becomes a native call; the
+    # A loads before it are left dead. %r560 (shared base) is the only register of the range used later.
+    "rrlite_dec0_4x4_folded": dict(
+        file="dlss-0195-00.ptx",
+        delete=(139, 587),
+        delete_check=("mov.u32 %r17, %laneid;", "st.shared.v2.u16 [%r570+1168], {%rs63, %rs64};"),
+        sust=True,
+        extern=""".extern .func d4r_dec0_head
+(
+	.param .b32 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
+)
+;
+""",
+        insert="""mov.u32 %r560, _ZZ22rrlite_dec0_4x4_foldedN6RRLite18Decoder0ParametersEE4smem;
+{
+	.param .b32 q0;
+	st.param.b32 [q0+0], %r560;
+	.param .b64 q1;
+	st.param.b64 [q1+0], %rd22;
+	.param .b64 q2;
+	st.param.b64 [q2+0], %rd24;
+	.param .b32 q3;
+	st.param.b32 [q3+0], %r475;
+	.param .b32 q4;
+	st.param.b32 [q4+0], %r476;
+	.param .b32 q5;
+	st.param.b32 [q5+0], %r3;
+	.param .b32 q6;
+	st.param.b32 [q6+0], %r4;
+	call.uni d4r_dec0_head, (q0, q1, q2, q3, q4, q5, q6);
+}""",
+    ),
+    # post and downsample: only the formatted surface stores become native (inline format conversion
+    # instead of ZLUDA's per-channel surface_formatted_store_bits calls)
+    "rrlite_post_3_2_mvhi_hdr_folded": dict(file="dlss-0176-00.ptx", sust=True, extern=""),
+    "rrlite_downsample_kernel_static_hdr": dict(file="dlss-0196-00.ptx", sust=True, extern=""),
+    # preset K (DLSS 4): the output kernel's 17 surface stores (13 formatted, 4 raw 32-bit)
+    "hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel": dict(file="dlss-0078-00.ptx", sust=True, rz_round=True, extern=""),
+    "hiluma_engine_input_depthinv_mvhi_hdr_v2_rel": dict(file="dlss-0062-00.ptx", sust=True, rz_round=True, extern=""),
+}
+
+
+SUST_EXTERN = """.extern .func d4r_sust_v2b16
+(
+	.param .b64 d4r_s0, .param .b32 d4r_s1, .param .b32 d4r_s2, .param .b16 d4r_s3, .param .b16 d4r_s4
+)
+;
+"""
+SUSTP_EXTERN = """.extern .func d4r_sust_p_v4b32
+(
+	.param .b64 d4r_t0, .param .b32 d4r_t1, .param .b32 d4r_t2, .param .b32 d4r_t3, .param .b32 d4r_t4,
+	.param .b32 d4r_t5, .param .b32 d4r_t6
+)
+;
+"""
+SUSTP_RE = re.compile(r"^sust\.p\.2d\.v4\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%f\d+),(%f\d+),(%f\d+),(%f\d+)\};$")
+SUSTB32_EXTERN = """.extern .func d4r_sust_b32
+(
+	.param .b64 d4r_w0, .param .b32 d4r_w1, .param .b32 d4r_w2, .param .b32 d4r_w3
+)
+;
+"""
+SUSTB32_RE = re.compile(r"^sust\.b\.2d\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%r\d+)\};$")
+SUST_RE = re.compile(r"^sust\.b\.2d\.v2\.b16\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%rs\d+),(%rs\d+)\};$")
+
+
+SUST_KINDS = set(os.environ.get("D4R_SUST_KINDS", "b16,p,b32").split(","))
+
+
+def replace_sust(lines):
+    out, n = [], {"b16": 0, "p": 0, "b32": 0}
+    for line in lines:
+        mb = SUSTB32_RE.match(line.strip()) if "b32" in SUST_KINDS else None
+        if mb:
+            s, x, y, v = mb.groups()
+            out.append(f"""{{
+	.param .b64 u0;
+	st.param.b64 [u0+0], {s};
+	.param .b32 u1;
+	st.param.b32 [u1+0], {x};
+	.param .b32 u2;
+	st.param.b32 [u2+0], {y};
+	.param .b32 u3;
+	st.param.b32 [u3+0], {v};
+	call d4r_sust_b32, (u0, u1, u2, u3);
+}}""")
+            n["b32"] += 1
+            continue
+        mp = SUSTP_RE.match(line.strip()) if "p" in SUST_KINDS else None
+        if mp:
+            s, x, y, *v = mp.groups()
+            st = "\n".join(f"\t.param .b32 v{i};\n\tst.param.b32 [v{i}+0], {r};" for i, r in enumerate(v))
+            out.append(f"""{{
+	.param .b64 u0;
+	st.param.b64 [u0+0], {s};
+	.param .b32 u1;
+	st.param.b32 [u1+0], {x};
+	.param .b32 u2;
+	st.param.b32 [u2+0], {y};
+{st}
+	call d4r_sust_p_v4b32, (u0, u1, u2, v0, v1, v2, v3);
+}}""")
+            n["p"] += 1
+            continue
+        m = SUST_RE.match(line.strip())
+        if not m:
+            out.append(line)
+            continue
+        s, x, y, a, b = m.groups()
+        out.append(f"""{{
+	.param .b64 u0;
+	st.param.b64 [u0+0], {s};
+	.param .b32 u1;
+	st.param.b32 [u1+0], {x};
+	.param .b32 u2;
+	st.param.b32 [u2+0], {y};
+	.param .b16 u3;
+	st.param.b16 [u3+0], {a};
+	.param .b16 u4;
+	st.param.b16 [u4+0], {b};
+	call d4r_sust_v2b16, (u0, u1, u2, u3, u4);
+}}""")
+        n["b16"] += 1
+    return out, n
+
+
+RZ_RE = re.compile(r"^add\.rz\.ftz\.f32 (%f\d+), (%f\d+), (%f\d+);$")
+CVT_RE = re.compile(r"^cvt\.rzi\.f32\.f32 (%f\d+), (%f\d+);$")
+
+
+def replace_rz_round(lines):
+    """roundf idiom add.rz(x, copysign(0.5, x)) + cvt.rzi -> exact round-half-away in default rounding:
+    t = trunc(x); r = |x - t| >= 0.5 ? t + 2 * copysign(0.5, x) : t (x - t is exact). Lets ZLUDA compile the
+    kernel without the strict-FP (constrained) mode that one explicit rounding mode forces."""
+    out, n, i = [], 0, 0
+    while i < len(lines):
+        m = RZ_RE.match(lines[i].strip())
+        if m:
+            d, x, sgn = m.groups()
+            j = i + 1
+            while j < len(lines) and not lines[j].strip():
+                j += 1
+            c = CVT_RE.match(lines[j].strip()) if j < len(lines) else None
+            if c and c.group(2) == d:
+                r = c.group(1)
+                out += [f"cvt.rzi.f32.f32 %rq{4*n}, {x};", f"sub.f32 %rq{4*n+1}, {x}, %rq{4*n};",
+                        f"abs.f32 %rq{4*n+1}, %rq{4*n+1};", f"setp.ge.f32 %pq{n}, %rq{4*n+1}, 0f3F000000;",
+                        f"add.f32 %rq{4*n+2}, {sgn}, {sgn};", f"add.f32 %rq{4*n+3}, %rq{4*n}, %rq{4*n+2};",
+                        f"selp.f32 {r}, %rq{4*n+3}, %rq{4*n}, %pq{n};"]
+                n += 1
+                i = j + 1
+                continue
+            raise SystemExit(f"add.rz at line {i} not followed by its cvt.rzi")
+        out.append(lines[i])
+        i += 1
+    return out, n
+
+
+def module_file(name, k):
+    """The extracted PTX module that defines NAME: the 310.7 file number, else any module that does
+    (other DLSS versions number their modules differently)."""
+    entry = re.compile(rf"\.entry\s+{re.escape(name)}\s*\(")
+    known = PTX_DIR / k["file"]
+    if known.exists() and entry.search(known.read_text()):
+        return known
+    for path in sorted(PTX_DIR.glob("*.ptx")):
+        if entry.search(path.read_text()):
+            return path
+    sys.exit(f"no PTX module in {PTX_DIR} defines {name}")
+
+
+def main():
+    name, out = sys.argv[1], Path(sys.argv[2])
+    k = KERNELS[name]
+    lines = module_file(name, k).read_text().replace("\r\n", "\n").split("\n")
+    if "cut_after" in k:
+        cut = k["cut_after"]
+        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 "delete" in k:
+        a, b = k["delete"]
+        if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]:
+            sys.exit(f"lines {a}/{b} are {lines[a - 1]!r}/{lines[b - 1]!r}, expected {k['delete_check']}")
+        head, tail = lines[:a - 1] + k["insert"].split("\n") + lines[b:], []
+    else:
+        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("rz_round"):
+        body, nrz = replace_rz_round(body)
+        # temporaries for the rewrite, declared after the kernel's own .reg lines
+        at = next(i for i, l in enumerate(body) if l.startswith(".reg"))
+        body = body[:at] + [f".reg .f32 %rq<{4 * nrz + 1}>;", f".reg .pred %pq<{nrz + 1}>;"] + body[at:]
+        n["rz"] = nrz
+    extern = k["extern"] + (SUST_EXTERN if n["b16"] else "") + (SUSTP_EXTERN if n["p"] else "") + (SUSTB32_EXTERN if n["b32"] else "")
+    text = "\n".join(head[:entry]) + "\n" + extern + "\n".join(body + tail)
+    out.write_text(text)
+    print(f"{out}: {len(text.splitlines())} lines, surface stores replaced: {n}")
+
+
+if __name__ == "__main__":
+    main()

kernels/tex/sust_only.hip

@@ -0,0 +1,2 @@
+// Native surface stores only (see make_ptx.py: post, downsample); the kernels are ZLUDA's compile.
+#include "tex_common.h"

kernels/tex/tex_common.h

@@ -0,0 +1,375 @@
+// Shared device code for native tails of the rrlite texture kernels (DLSS 310.7.0, preset M).
+// These are linked into ZLUDA's own compile of a hand-edited PTX kernel (D4R_ZLUDA_EXTRA_BC), so
+// only plain clang builtins are used (compiled with -nogpuinc -nogpulib).
+#pragma once
+#include <stdint.h>
+
+#pragma clang fp contract(off)
+
+#define DEV extern "C" __attribute__((device))
+// Kernels that ZLUDA compiles in strict-FP mode (constrained intrinsics) cannot take inlined FP code from
+// here (the AMDGPU backend fails on some constrained operations): build them with -DD4R_SUST_OUTLINE,
+// which keeps the FP bodies of the surface stores in noinline leaf functions (no calls, no stack).
+#ifdef D4R_SUST_OUTLINE
+#define FP_BODY __attribute__((device, noinline)) static
+#else
+#define FP_BODY __attribute__((device, always_inline)) static inline
+#endif
+#define KERNEL extern "C" __attribute__((global))
+
+typedef _Float16 half_t;
+typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
+typedef _Float16 h16 __attribute__((ext_vector_type(16)));
+typedef float f8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u4v __attribute__((ext_vector_type(4)));
+typedef uint32_t uint2_t __attribute__((ext_vector_type(2)));
+typedef uint16_t u16x2 __attribute__((ext_vector_type(2)));
+typedef int16_t i16x2 __attribute__((ext_vector_type(2)));
+typedef __attribute__((address_space(3))) uint8_t lds_u8;
+
+__attribute__((device)) static inline uint32_t lane_id()
+{
+    return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
+}
+
+// ZLUDA's e4m3x4_to_scaled_f16x2x2 (zluda_ptx_impl.cpp): e4m3 bytes -> f16 pairs holding value * 2^-8,
+// {bytes 0,1} and {bytes 2,3}.
+__attribute__((device)) static inline void e4m3x4_scaled(uint32_t packed, uint32_t& low, uint32_t& high)
+{
+    uint32_t magnitude = packed & 0x7f7f7f7fu;
+    magnitude |= (magnitude + 0x01010101u) & 0x80808080u;
+    uint32_t sign = packed & 0x80808080u;
+    low = (__builtin_amdgcn_perm(0u, magnitude, 0x0c010c00u) << 7) | __builtin_amdgcn_perm(0u, sign, 0x010c000cu);
+    high = (__builtin_amdgcn_perm(0u, magnitude, 0x0c030c02u) << 7) | __builtin_amdgcn_perm(0u, sign, 0x030c020cu);
+}
+
+// A k16 WMMA operand from 16 consecutive e4m3 bytes of NVIDIA's k order (4 dwords): register j holds
+// k 4j, 4j+1 and register 4 + j holds k 4j+2, 4j+3 -- ZLUDA's slot order (e4m3_a_k16_fragment).
+__attribute__((device)) static inline u8v operand16(u4v d)
+{
+    u8v v;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        uint32_t lo, hi;
+        e4m3x4_scaled(d[j], lo, hi);
+        v[j] = lo;
+        v[4 + j] = hi;
+    }
+    return v;
+}
+
+__attribute__((device)) static inline f8v wmma(u8v a, u8v b, f8v c)
+{
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
+}
+
+// ZLUDA's f16x2_to_e4m3x2_satfinite_bits: RNE satfinite e4m3 of both halves, low half -> low byte.
+__attribute__((device)) static inline uint32_t e4m3x2(uint32_t bits)
+{
+    u16x2 half = __builtin_bit_cast(u16x2, bits);
+    u16x2 magnitude = half & (uint16_t)0x7fffu;
+    u16x2 normal = (magnitude + ((magnitude >> (uint16_t)7) & (uint16_t)1) + (uint16_t)0xe03fu) >> (uint16_t)7;
+    u16x2 exponent = magnitude >> (uint16_t)10;
+    u16x2 clamped = __builtin_elementwise_max(__builtin_elementwise_min(exponent, (u16x2)8), (u16x2)1);
+    u16x2 shift = (uint16_t)16 - clamped;
+    u16x2 significand = (magnitude & (uint16_t)0x3ffu) |
+                        (__builtin_elementwise_min(magnitude, (u16x2)0x400) & (uint16_t)0x400u);
+    u16x2 subnormal = (significand + (((u16x2)1 << (shift - (uint16_t)1)) - (uint16_t)1) +
+                       ((significand >> shift) & (uint16_t)1)) >> shift;
+    u16x2 below = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(magnitude - (uint16_t)0x2400u)) >> (int16_t)15));
+    u16x2 code = (subnormal & below) | (normal & ~below);
+    code = __builtin_elementwise_min(code, (u16x2)0x7e);
+    u16x2 nan = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((uint16_t)0x7c00u - magnitude)) >> (int16_t)15));
+    code |= nan & (uint16_t)0x7fu;
+    code |= (half >> (uint16_t)8) & (uint16_t)0x80u;
+    return uint32_t(code.x) | (uint32_t(code.y) << 8);
+}
+
+__attribute__((device)) static inline uint32_t pack2(half_t a, half_t b)
+{
+    return __builtin_bit_cast(uint32_t, (hv2){a, b});
+}
+
+// value from the lane with the same index in the other half of the wave
+__attribute__((device)) static inline uint32_t other_half(uint32_t v)
+{
+    return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
+}
+
+// ---------------------------------------------------------------- surface stores
+typedef int32_t v2i __attribute__((ext_vector_type(2)));
+typedef int16_t v2s __attribute__((ext_vector_type(2)));
+typedef float v4f __attribute__((ext_vector_type(4)));
+typedef int32_t v4i_t __attribute__((ext_vector_type(4)));
+typedef __attribute__((address_space(4))) unsigned int tsharp_t;
+extern "C" __attribute__((device)) int __ockl_image_channel_data_type_2D(tsharp_t* image);
+extern "C" __attribute__((device)) int __ockl_image_channel_order_2D(tsharp_t* image);
+extern "C" __attribute__((device)) void __ockl_image_store_2D(tsharp_t* image, v2i coord, v4f value);
+extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t surface, v2i coord, v2s value);
+
+// ZLUDA's surface_half_bits_to_float (exact, independent of the f16 denormal mode)
+__attribute__((device)) static inline float half_bits_to_float(uint32_t bits)
+{
+    uint32_t sign = (bits & 0x8000u) << 16;
+    uint32_t exponent = (bits >> 10) & 0x1fu;
+    uint32_t mantissa = bits & 0x3ffu;
+    if (exponent == 0)
+    {
+        float magnitude = float(mantissa) * 0x1p-24f;
+        return __builtin_bit_cast(float, __builtin_bit_cast(uint32_t, magnitude) | sign);
+    }
+    if (exponent == 0x1f)
+        return __builtin_bit_cast(float, sign | 0x7f800000u | (mantissa << 13));
+    return __builtin_bit_cast(float, sign | ((exponent + 112u) << 23) | (mantissa << 13));
+}
+
+FP_BODY void sust_rg16f_body(tsharp_t* image, int32_t px, int32_t y, uint16_t a, uint16_t b)
+{
+    __ockl_image_store_2D(image, (v2i){px, y}, (v4f){half_bits_to_float(a), half_bits_to_float(b), 0.0f, 0.0f});
+}
+
+// sust.b.2d.v2.b16.zero: a whole RG16F pixel is one image store of the decoded halves (what ZLUDA's
+// surface_store_raw_2d does for that format); anything else goes through ZLUDA's generic helper.
+__attribute__((device, noinline)) static void sust_v2b16_generic(uint64_t surface, int32_t x, int32_t y, uint16_t a, uint16_t b)
+{
+    __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(surface, (v2i){x, y}, (v2s){(int16_t)a, (int16_t)b});
+}
+
+DEV __attribute__((always_inline)) void d4r_sust_v2b16(uint64_t surface, int32_t x, int32_t y, uint16_t a, uint16_t b)
+{
+    tsharp_t* image = (tsharp_t*)surface;
+    if (__ockl_image_channel_data_type_2D(image) == 14 && __ockl_image_channel_order_2D(image) == 3 && (x & 3) == 0)
+    {
+        if (x >= 0)
+            sust_rg16f_body(image, x >> 2, y, a, b);
+    }
+    else
+        sust_v2b16_generic(surface, x, y, a, b);
+}
+
+
+// f32 comparisons on the bits (no fcmp: kernels ZLUDA compiles in strict-FP mode cannot lower
+// constrained fcmp); same results as the float compares, -0.0 included
+__attribute__((device)) static inline bool f_isnan(uint32_t u) { return (u & 0x7fffffffu) > 0x7f800000u; }
+__attribute__((device)) static inline bool f_lt0(uint32_t u) { return (int32_t)u < 0 && (u & 0x7fffffffu) != 0 && !f_isnan(u); }
+// value > 1.0f (value not NaN)
+__attribute__((device)) static inline bool f_gt1(uint32_t u) { return (int32_t)u > 0x3f800000; }
+// value < -1.0f (value not NaN)
+__attribute__((device)) static inline bool f_ltm1(uint32_t u) { return (int32_t)u < 0 && (u & 0x7fffffffu) > 0x3f800000u; }
+
+// ZLUDA's surface_formatted_store_bits (zluda_ptx_impl.cpp), inlined: normalized CUDA formats convert
+// the f32 value, everything else stores the bits unchanged.
+__attribute__((device)) static inline uint32_t formatted_store_bits(uint32_t bits, uint32_t cuda_format, int channel)
+{
+    if (cuda_format == 80)
+    {
+        uint32_t u = bits;
+        if (f_isnan(u) || f_lt0(u))
+            u = 0u;
+        if (f_gt1(u))
+            u = 0x3f800000u;
+        const float value = __builtin_bit_cast(float, u);
+        const float scale = channel == 3 ? 3.0f : 1023.0f;
+        // float(int(floor(x))) == floor(x) here (0 <= x <= 1023.5): no int <-> float conversions, which
+        // strict-FP kernels cannot lower as constrained signed conversions
+        return __builtin_bit_cast(uint32_t, __builtin_floorf(value * scale + 0.5f) / scale);
+    }
+    float scale = 0.0f;
+    bool signed_normalized = false;
+    if (cuda_format >= 192 && cuda_format <= 194)
+        scale = 255.0f;
+    else if (cuda_format >= 195 && cuda_format <= 197)
+        scale = 65535.0f;
+    else if (cuda_format >= 198 && cuda_format <= 200)
+    {
+        scale = 127.0f;
+        signed_normalized = true;
+    }
+    else if (cuda_format >= 201 && cuda_format <= 203)
+    {
+        scale = 32767.0f;
+        signed_normalized = true;
+    }
+    if (!(cuda_format >= 192 && cuda_format <= 203))
+        return bits;
+    uint32_t u = bits;
+    if (f_isnan(u))
+        u = 0u;
+    if (signed_normalized ? f_ltm1(u) : f_lt0(u))
+        u = signed_normalized ? 0xbf800000u : 0u;
+    if (f_gt1(u))
+        u = 0x3f800000u;
+    const float scaled = __builtin_bit_cast(float, u) * scale;
+    const uint32_t su = __builtin_bit_cast(uint32_t, scaled);
+    int32_t integer = f_lt0(su) ? -int32_t(__builtin_floorf(-scaled + 0.5f)) : int32_t(__builtin_floorf(scaled + 0.5f));
+    return uint32_t(integer);
+}
+
+// sust.p.2d.v4.b32.zero: ZLUDA stores the raw bits unless the original CUDA format was normalized
+// (surface_formatted_store_bits; format word after the image and sampler descriptors).
+FP_BODY void sust_p_body(uint64_t surface, int32_t x, int32_t y, uint32_t a, uint32_t b, uint32_t c, uint32_t d)
+{
+    const uint32_t format = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
+    __ockl_image_store_2D((tsharp_t*)surface, (v2i){x, y},
+                          (v4f){__builtin_bit_cast(float, formatted_store_bits(a, format, 0)),
+                                __builtin_bit_cast(float, formatted_store_bits(b, format, 1)),
+                                __builtin_bit_cast(float, formatted_store_bits(c, format, 2)),
+                                __builtin_bit_cast(float, formatted_store_bits(d, format, 3))});
+}
+
+// d4r output redirect (wine_nvcuda_bridge.c d4rSetArrayRedirect): a tagged row pitch at byte 84 of the
+// ZLUDA surface object (zero padding otherwise) sends stores of that surface to the linear memory whose
+// address is at byte 88, rows `pitch` bytes apart. Out-of-range coordinates are dropped as by .zero stores.
+extern "C" __attribute__((device)) int __ockl_image_width_2D(tsharp_t* image);
+extern "C" __attribute__((device)) int __ockl_image_height_2D(tsharp_t* image);
+__attribute__((device)) static inline uint32_t redirect_pitch(uint64_t surface)
+{
+    // 'R2' tag in the high half (other handles may have data there), row pitch / 8 in the low half
+    const uint32_t word = ((const __attribute__((address_space(4))) uint32_t*)surface)[21];
+    return (word >> 16) == 0x5232u ? (word & 0xffffu) << 3 : 0u;
+}
+__attribute__((device)) static inline uint8_t* redirect_row(uint64_t surface, int32_t y)
+{
+    const uint64_t base = *(const __attribute__((address_space(4))) uint64_t*)(surface + 88);
+    return (uint8_t*)(base + (uint64_t)(uint32_t)y * redirect_pitch(surface));
+}
+__attribute__((device)) static inline bool redirect_inside(uint64_t surface, int32_t x, int32_t y)
+{
+    tsharp_t* image = (tsharp_t*)surface;
+    return x >= 0 && y >= 0 && x < __ockl_image_width_2D(image) && y < __ockl_image_height_2D(image);
+}
+
+DEV __attribute__((always_inline)) void d4r_sust_p_v4b32(uint64_t surface, int32_t x, int32_t y, uint32_t a, uint32_t b,
+                                                        uint32_t c, uint32_t d)
+{
+    const uint32_t format0 = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
+    if (redirect_pitch(surface) != 0 && format0 != 80 && !(format0 >= 192 && format0 <= 203))
+    {
+        // RGBA16F: the f32 values as halves, rounded toward zero like the image store's format conversion
+        if (redirect_inside(surface, x, y))
+        {
+            const uint32_t lo = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(__builtin_bit_cast(float, a), __builtin_bit_cast(float, b)));
+            const uint32_t hi = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(__builtin_bit_cast(float, c), __builtin_bit_cast(float, d)));
+            *(uint2_t*)(redirect_row(surface, y) + 8 * x) = (uint2_t){lo, hi};
+        }
+        return;
+    }
+    // formatted_store_bits keeps the bits unless the format is normalized: one uniform test per store
+    const uint32_t format = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
+    if (format != 80 && !(format >= 192 && format <= 203))
+        __ockl_image_store_2D((tsharp_t*)surface, (v2i){x, y},
+                              (v4f){__builtin_bit_cast(float, a), __builtin_bit_cast(float, b), __builtin_bit_cast(float, c),
+                                    __builtin_bit_cast(float, d)});
+    else
+        sust_p_body(surface, x, y, a, b, c, d);
+}
+
+// ---------------------------------------------------------------- raw 32-bit surface stores
+extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_b32_zero(uint64_t surface, v2i coord, uint32_t value);
+
+// exact int -> float for |v| < 2^24 without a signed conversion (see formatted_store_bits)
+__attribute__((device)) static inline float i2f(int32_t v)
+{
+    return v < 0 ? -(float)(uint32_t)(-v) : (float)(uint32_t)v;
+}
+
+// ZLUDA's surface_decode_channel (bits of one memory channel -> the float4 lane value the store converts back)
+__attribute__((device)) static inline float decode_channel(uint32_t bits, int data_type)
+{
+    switch (data_type)
+    {
+    case 0:
+        return i2f(__builtin_elementwise_max(__builtin_elementwise_min(int32_t(int8_t(bits)), 127), -127)) / 127.0f;
+    case 1:
+        return i2f(__builtin_elementwise_max(__builtin_elementwise_min(int32_t(int16_t(bits)), 32767), -32767)) / 32767.0f;
+    case 2:
+        return (float)(bits & 0xffu) / 255.0f;
+    case 3:
+        return (float)(bits & 0xffffu) / 65535.0f;
+    case 8:
+        return __builtin_bit_cast(float, int32_t(int8_t(bits)));
+    case 9:
+        return __builtin_bit_cast(float, int32_t(int16_t(bits)));
+    case 14:
+        return half_bits_to_float(bits);
+    default:
+        return __builtin_bit_cast(float, bits);
+    }
+}
+
+// out of line: inlining ZLUDA's generic raw store into large kernels crashes AMDGPU instruction
+// selection (see zluda_ptx_impl.cpp); it makes no calls, so the wrapper needs no stack
+__attribute__((device, noinline)) static void sust_b32_generic(uint64_t surface, int32_t x, int32_t y, uint32_t data)
+{
+    __zluda_ptx_impl_surfobj_b_2d_b32_zero(surface, (v2i){x, y}, data);
+}
+
+// sust.b.2d.b32.zero (x is a byte offset): whole 4-byte pixels and pairs of 2-byte pixels in R/RG/RGBA
+// channel order need no read-modify-write (ZLUDA's surface_store_raw_2d, same conversions); anything
+// else goes through ZLUDA's generic helper.
+FP_BODY void sust_b32_fast(tsharp_t* image, int32_t x, int32_t y, uint32_t data, int dt, int cb, int pb)
+{
+    if (pb == 4)
+    {
+        v4f v = {0.0f, 0.0f, 0.0f, 0.0f};
+        if (cb == 4)
+            v[0] = decode_channel(data, dt);
+        else if (cb == 2)
+        {
+            v[0] = decode_channel(data & 0xffffu, dt);
+            v[1] = decode_channel(data >> 16, dt);
+        }
+        else
+        {
+            v[0] = decode_channel(data & 0xffu, dt);
+            v[1] = decode_channel((data >> 8) & 0xffu, dt);
+            v[2] = decode_channel((data >> 16) & 0xffu, dt);
+            v[3] = decode_channel(data >> 24, dt);
+        }
+        __ockl_image_store_2D(image, (v2i){x >> 2, y}, v);
+    }
+    else
+    {
+        v4f v0 = {0.0f, 0.0f, 0.0f, 0.0f}, v1 = {0.0f, 0.0f, 0.0f, 0.0f};
+        if (cb == 2)
+        {
+            v0[0] = decode_channel(data & 0xffffu, dt);
+            v1[0] = decode_channel(data >> 16, dt);
+        }
+        else
+        {
+            v0[0] = decode_channel(data & 0xffu, dt);
+            v0[1] = decode_channel((data >> 8) & 0xffu, dt);
+            v1[0] = decode_channel((data >> 16) & 0xffu, dt);
+            v1[1] = decode_channel(data >> 24, dt);
+        }
+        __ockl_image_store_2D(image, (v2i){x >> 1, y}, v0);
+        __ockl_image_store_2D(image, (v2i){(x >> 1) + 1, y}, v1);
+    }
+}
+
+DEV __attribute__((always_inline)) void d4r_sust_b32(uint64_t surface, int32_t x, int32_t y, uint32_t data)
+{
+    tsharp_t* image = (tsharp_t*)surface;
+    if (redirect_pitch(surface) != 0)
+    {
+        // raw bits at byte offset x (RGBA16F rows are 8 bytes per pixel)
+        if (x >= 0 && (x & 3) == 0 && redirect_inside(surface, x >> 3, y))
+            *(uint32_t*)(redirect_row(surface, y) + x) = data;
+        return;
+    }
+    const int dt = __ockl_image_channel_data_type_2D(image), order = __ockl_image_channel_order_2D(image);
+    const int cb = (dt == 0 || dt == 2 || dt == 8 || dt == 11) ? 1 : (dt == 1 || dt == 3 || dt == 9 || dt == 12 || dt == 14) ? 2
+                   : (dt == 10 || dt == 13 || dt == 15) ? 4 : 0;
+    const int n = order == 1 ? 1 : order == 3 ? 2 : order == 8 ? 4 : 0; // R, RG, RGBA (identity lane order)
+    const int pb = cb * n;
+    if ((pb == 4 && (x & 3) == 0) || (pb == 2 && (x & 1) == 0))
+    {
+        if (x >= 0)
+            sust_b32_fast(image, x, y, data, dt, cb, pb);
+    }
+    else
+        sust_b32_generic(surface, x, y, data);
+}

kernels/tools/dump_runner.cpp

@@ -0,0 +1,167 @@
+// Launch a native kernel on a replay dump (manifest.txt, args.bin, alloc-N.bin), save every allocation
+// after the first launch, then time repeated launches.
+// usage: dump_runner <hsaco> <kernel> <dump dir> <out dir> [iterations]
+#include <hip/hip_runtime.h>
+#include <algorithm>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#define CHECK(x)                                                                                                        \
+    do                                                                                                                  \
+    {                                                                                                                   \
+        hipError_t e_ = (x);                                                                                            \
+        if (e_ != hipSuccess)                                                                                           \
+        {                                                                                                               \
+            fprintf(stderr, "%s failed: %s\n", #x, hipGetErrorString(e_));                                             \
+            return 1;                                                                                                   \
+        }                                                                                                               \
+    } while (0)
+
+static std::vector<char> read_file(const std::string& path)
+{
+    std::ifstream f(path, std::ios::binary);
+    return std::vector<char>((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
+}
+
+int main(int argc, char** argv)
+{
+    if (argc < 5)
+    {
+        fprintf(stderr, "usage: %s hsaco kernel dump out [iters]\n", argv[0]);
+        return 2;
+    }
+    const std::string dump = argv[3], out = argv[4];
+    const int iters = argc > 5 ? atoi(argv[5]) : 20;
+    hipModule_t mod;
+    hipFunction_t fn;
+    CHECK(hipModuleLoad(&mod, argv[1]));
+    CHECK(hipModuleGetFunction(&fn, mod, argv[2]));
+    hipFunction_t prep = nullptr;
+    if (hipModuleGetFunction(&prep, mod, (std::string(argv[2]) + "_prep").c_str()) != hipSuccess)
+        prep = nullptr;
+    const unsigned prepGrid = getenv("PREP_GRID") ? atoi(getenv("PREP_GRID")) : 88;
+
+    std::vector<char> args = read_file(dump + "/args.bin");
+    std::ifstream man(dump + "/manifest.txt");
+    std::string line;
+    unsigned grid[3] = {1, 1, 1}, block[3] = {1, 1, 1}, shared = 0;
+    std::vector<int> ids;
+    std::vector<void*> dev;
+    std::vector<size_t> sizes;
+    while (std::getline(man, line))
+    {
+        std::istringstream s(line);
+        std::string kind;
+        s >> kind;
+        if (kind == "launch")
+            s >> grid[0] >> grid[1] >> grid[2] >> block[0] >> block[1] >> block[2] >> shared;
+        else if (kind == "alloc")
+        {
+            int id;
+            std::string base;
+            size_t size;
+            s >> id >> base >> size;
+            std::vector<char> data = read_file(dump + "/alloc-" + std::to_string(id) + ".bin");
+            if (data.size() != size)
+            {
+                fprintf(stderr, "alloc %d size mismatch\n", id);
+                return 1;
+            }
+            void* d;
+            CHECK(hipMalloc(&d, size));
+            CHECK(hipMemcpy(d, data.data(), size, hipMemcpyHostToDevice));
+            ids.push_back(id);
+            dev.push_back(d);
+            sizes.push_back(size);
+        }
+        else if (kind == "pointer")
+        {
+            size_t off, delta;
+            int id;
+            s >> off >> id >> delta;
+            size_t i = std::find(ids.begin(), ids.end(), id) - ids.begin();
+            uint64_t v = (uint64_t)dev[i] + delta;
+            memcpy(args.data() + off, &v, 8);
+        }
+    }
+    {
+        hipDeviceptr_t bz;
+        size_t bzs = 0;
+        uint32_t v = 0;
+        if (hipModuleGetGlobal(&bz, &bzs, mod, "d4r_block_z") == hipSuccess && bzs == 4 &&
+            hipMemcpyDtoH(&v, bz, 4) == hipSuccess && v != 0)
+            block[2] = v;
+    }
+    {
+        hipDeviceptr_t gp;
+        size_t gs = 0;
+        uint32_t v = 0;
+        if (hipModuleGetGlobal(&gp, &gs, mod, "d4r_grid_x") == hipSuccess && gs == 4 &&
+            hipMemcpyDtoH(&v, gp, 4) == hipSuccess && v != 0)
+        {
+            grid[0] = v;
+            grid[1] = grid[2] = 1;
+        }
+    }
+    if (getenv("GRID_X"))
+        grid[0] = atoi(getenv("GRID_X"));
+    if (getenv("GRID_Y"))
+        grid[1] = atoi(getenv("GRID_Y"));
+    size_t argSize = args.size();
+    void* cfg[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, args.data(), HIP_LAUNCH_PARAM_BUFFER_SIZE, &argSize,
+                   HIP_LAUNCH_PARAM_END};
+    const bool prepOnce = getenv("PREP_ONCE") != nullptr;
+    int launches = 0;
+    auto launch = [&]() {
+        if (prep != nullptr && !(prepOnce && launches++ > 0))
+        {
+            hipError_t e = hipModuleLaunchKernel(prep, prepGrid, 1, 1, 128, 1, 1, 0, nullptr, nullptr, cfg);
+            if (e != hipSuccess)
+                return e;
+        }
+        return hipModuleLaunchKernel(fn, grid[0], grid[1], grid[2], block[0], block[1], block[2], shared, nullptr,
+                                     nullptr, cfg);
+    };
+    CHECK(launch());
+    CHECK(hipDeviceSynchronize());
+    for (size_t i = 0; i < ids.size(); ++i)
+    {
+        std::vector<char> h(sizes[i]);
+        CHECK(hipMemcpy(h.data(), dev[i], sizes[i], hipMemcpyDeviceToHost));
+        std::ofstream(out + "/alloc-" + std::to_string(ids[i]) + ".bin", std::ios::binary).write(h.data(), h.size());
+    }
+    if (const char* dbg = getenv("DBG_GLOBAL"))
+    {
+        hipDeviceptr_t ptr;
+        size_t size = 0;
+        if (hipModuleGetGlobal(&ptr, &size, mod, dbg) == hipSuccess)
+        {
+            std::vector<char> h(size);
+            CHECK(hipMemcpyDtoH(h.data(), ptr, size));
+            std::ofstream(out + "/dbg.bin", std::ios::binary).write(h.data(), h.size());
+        }
+    }
+    hipEvent_t a, b;
+    CHECK(hipEventCreate(&a));
+    CHECK(hipEventCreate(&b));
+    std::vector<float> ms;
+    for (int i = 0; i < iters; ++i)
+    {
+        CHECK(hipEventRecord(a));
+        CHECK(launch());
+        CHECK(hipEventRecord(b));
+        CHECK(hipEventSynchronize(b));
+        float t;
+        CHECK(hipEventElapsedTime(&t, a, b));
+        ms.push_back(t);
+    }
+    std::sort(ms.begin(), ms.end());
+    if (!ms.empty())
+        printf("%s: median %.4f ms, min %.4f ms over %d launches\n", argv[2], ms[ms.size() / 2], ms[0], iters);
+    return 0;
+}

kernels/tools/extract_dlss_ptx.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Extract the PTX modules embedded in NVIDIA's nvngx_dlss.dll.
+
+usage: extract_dlss_ptx.py NVNGX_DLSS_DLL OUT_DIR
+
+Writes OUT_DIR/dlss-NNNN-MM.ptx for PTX entry MM of the NNNN-th CUDA fatbin found in the DLL (in file
+order). The texture-kernel build (kernels/tex) edits some of these modules; the DLL is NVIDIA's and is
+never part of this repository, and neither are the extracted files.
+"""
+import os
+import struct
+import sys
+
+FATBIN_MAGIC = 0xBA55ED50
+
+
+def lz4_block(src: bytes, size: int) -> bytes:
+    dst = bytearray()
+    i = 0
+    while i < len(src):
+        token = src[i]
+        i += 1
+        literals = token >> 4
+        if literals == 15:
+            while True:
+                b = src[i]
+                i += 1
+                literals += b
+                if b != 255:
+                    break
+        dst += src[i:i + literals]
+        i += literals
+        if i >= len(src) or len(dst) >= size:
+            break
+        offset = src[i] | (src[i + 1] << 8)
+        i += 2
+        match = token & 15
+        if match == 15:
+            while True:
+                b = src[i]
+                i += 1
+                match += b
+                if b != 255:
+                    break
+        match += 4
+        for _ in range(match):
+            dst.append(dst[-offset])
+    return bytes(dst[:size])
+
+
+def fatbins(data: bytes):
+    """Yields (offset, bytes) of every fatbin container (magic, version 1) in file order."""
+    magic = struct.pack('<I', FATBIN_MAGIC)
+    start = data.find(magic)
+    while start >= 0:
+        _, version, header_size, size = struct.unpack_from('<IHHQ', data, start)
+        if version == 1 and header_size == 16 and start + header_size + size <= len(data):
+            yield start, data[start:start + header_size + size]
+            start = data.find(magic, start + header_size + size)
+        else:
+            start = data.find(magic, start + 4)
+
+
+def ptx_entries(fatbin: bytes):
+    """Yields the (decompressed) PTX text of each PTX entry of one fatbin."""
+    _, _, header_size, size = struct.unpack_from('<IHHQ', fatbin, 0)
+    offset = header_size
+    while offset < header_size + size:
+        kind, _, entry_header, entry_size = struct.unpack_from('<HHIQ', fatbin, offset)
+        flags = struct.unpack_from('<Q', fatbin, offset + 40)[0]
+        decompressed = struct.unpack_from('<Q', fatbin, offset + 56)[0] if entry_header >= 64 else 0
+        payload = fatbin[offset + entry_header:offset + entry_header + entry_size]
+        if kind == 1:  # PTX
+            yield lz4_block(payload, decompressed) if (flags & 0x2000 and decompressed) else payload
+        offset += entry_header + entry_size
+
+
+def main(argv):
+    if len(argv) != 3:
+        sys.exit(__doc__)
+    data = open(argv[1], 'rb').read()
+    os.makedirs(argv[2], exist_ok=True)
+    total = 0
+    for index, (_, fatbin) in enumerate(fatbins(data)):
+        for entry, ptx in enumerate(ptx_entries(fatbin)):
+            ptx = ptx.rstrip(b'\0')
+            with open(os.path.join(argv[2], f'dlss-{index:04d}-{entry:02d}.ptx'), 'wb') as out:
+                out.write(ptx)
+            total += 1
+    print(f'{total} PTX modules written to {argv[2]}')
+
+
+if __name__ == '__main__':
+    main(sys.argv)

kernels/tools/kernel_manifest.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python3
+"""Write the d4r-kernels.txt manifest of a native kernel directory.
+
+usage: kernel_manifest.py KERNEL_DIR NVNGX_DLSS_DLL [NVNGX_DLSS_DLL...]
+
+For every NAME.hsaco in KERNEL_DIR, finds the PTX module of each DLL that defines `.entry NAME` and
+lists the FNV-1a 64 hash of the module's text (trailing NUL bytes removed), one "NAME HASH" line per
+distinct hash. The Wine CUDA bridge serves a native kernel only while DLSS loads a module with a
+listed hash, so a DLSS version that changed the kernel runs ZLUDA's own compile of it instead.
+The manifest is written to KERNEL_DIR/d4r-kernels.txt; it holds hashes, not NVIDIA code.
+"""
+import hashlib
+import os
+import re
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from extract_dlss_ptx import fatbins, ptx_entries  # noqa: E402
+
+ENTRY = re.compile(rb'\.entry[ \t\n]+([A-Za-z0-9_$]+)[ \t\n\r]*\(')
+
+
+def fnv1a64(data: bytes) -> int:
+    value = 0xcbf29ce484222325
+    for byte in data:
+        value = ((value ^ byte) * 0x100000001b3) & 0xffffffffffffffff
+    return value
+
+
+def module_hashes(dll: str) -> dict:
+    """NAME -> set of module hashes for every PTX entry of one DLL."""
+    data = open(dll, 'rb').read()
+    found = {}
+    for _, fatbin in fatbins(data):
+        for ptx in ptx_entries(fatbin):
+            text = ptx.rstrip(b'\0')
+            value = fnv1a64(text)
+            for match in ENTRY.finditer(text):
+                found.setdefault(match.group(1).decode(), set()).add(value)
+    return found
+
+
+def main(argv):
+    if len(argv) < 3:
+        sys.exit(__doc__)
+    directory = argv[1]
+    names = sorted(name[:-len('.hsaco')] for name in os.listdir(directory) if name.endswith('.hsaco'))
+    lines = ['# d4r native kernels: NAME and the FNV-1a 64 hash of the DLSS PTX module each was written for',
+             '# (written by kernels/tools/kernel_manifest.py from:']
+    listed = {name: set() for name in names}
+    for dll in argv[2:]:
+        digest = hashlib.sha256(open(dll, 'rb').read()).hexdigest()[:16]
+        lines.append(f'#   {os.path.basename(dll)} sha256 {digest}...)')
+        hashes = module_hashes(dll)
+        for name in names:
+            listed[name] |= hashes.get(name, set())
+    missing = [name for name in names if not listed[name]]
+    if missing:
+        sys.exit(f'no DLL defines {", ".join(missing)}; not writing a manifest')
+    for name in names:
+        for value in sorted(listed[name]):
+            lines.append(f'{name} {value:016x}')
+    with open(os.path.join(directory, 'd4r-kernels.txt'), 'w') as out:
+        out.write('\n'.join(lines) + '\n')
+    print(f'{len(names)} kernels listed in {os.path.join(directory, "d4r-kernels.txt")}')
+
+
+if __name__ == '__main__':
+    main(sys.argv)

kernels/tools/kprof.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+# usage: kprof.py RUN.log FRAMES [WARM] : per-kernel median GPU ms and per-frame total from KERNEL_PROFILE lines
+import sys,re,statistics as st
+lines=[l for l in open(sys.argv[1],errors='replace') if 'KERNEL_PROFILE' in l]
+frames=int(sys.argv[2]); warm=int(sys.argv[3]) if len(sys.argv)>3 else frames//2
+rows=[(int(m.group(1)),float(m.group(2)),m.group(3)) for m in (re.search(r'launch=(\d+).*gpu_ms=([0-9.]+) name=(\S+)',l) for l in lines) if m]
+rows.sort(); n=len(rows); per=n/frames; cut=int(per*warm); rows=rows[cut:]; f=frames-warm
+agg={}
+for _,ms,name in rows: agg.setdefault(name,[]).append(ms)
+tot=sum(sum(v) for v in agg.values())/f
+print(f"launches/frame={per:.1f} total_gpu_ms/frame={tot:.3f} (frames {warm+1}-{frames})")
+for name,v in sorted(agg.items(),key=lambda kv:-sum(kv[1])):
+    print(f"  {sum(v)/f:8.3f} ms/frame  {len(v)/f:5.1f}x  med {st.median(v):7.3f}  {name}")

kernels/tools/psnr.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python3
+"""psnr.py DIR_A DIR_B: per-frame PSNR of rgba16f outputs (watermark box excluded, values clipped to [0, 1] after tonemap-free clamp)."""
+import sys, glob, os, numpy as np
+a, b = sys.argv[1:3]
+vals = []
+for f in sorted(glob.glob(a + '/frame-*.rgba16f')):
+    g = os.path.join(b, os.path.basename(f))
+    if not os.path.exists(g): continue
+    x = np.fromfile(f, np.float16).reshape(1440, 2560, 4)[..., :3].astype(np.float64)
+    y = np.fromfile(g, np.float16).reshape(1440, 2560, 4)[..., :3].astype(np.float64)
+    m = np.ones(x.shape[:2], bool); m[1370:1392, 1850:2460] = False
+    x, y = np.clip(x[m], 0, 1), np.clip(y[m], 0, 1)
+    mse = np.mean((x - y) ** 2)
+    vals.append(10 * np.log10(1.0 / mse) if mse > 0 else 99.0)
+print('frames', len(vals), 'PSNR dB min %.2f mean %.2f' % (min(vals), np.mean(vals)), ' '.join('%.1f' % v for v in vals))

kernels/tools/ptxload.c

@@ -0,0 +1,20 @@
+// usage: ptxload FILE.ptx... : cuModuleLoadData + cuModuleGetFunction for every .entry
+#include <dlfcn.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+typedef int (*I)(unsigned); typedef int (*DG)(int*,int); typedef int (*CC)(void**,unsigned,int);
+typedef int (*ML)(void**,const void*); typedef int (*GF)(void**,void*,const char*); typedef int (*MU)(void*);
+int main(int argc,char**argv){
+  void*l=dlopen("libcuda.so",RTLD_NOW); if(!l){fprintf(stderr,"%s\n",dlerror());return 1;}
+  I init=dlsym(l,"cuInit"); DG dg=dlsym(l,"cuDeviceGet"); CC cc=dlsym(l,"cuCtxCreate_v2");
+  ML ml=dlsym(l,"cuModuleLoadData"); GF gf=dlsym(l,"cuModuleGetFunction"); MU mu=dlsym(l,"cuModuleUnload");
+  int d; void*ctx; init(0); dg(&d,0); cc(&ctx,0,d); int bad=0;
+  for(int a=1;a<argc;a++){
+    FILE*f=fopen(argv[a],"rb"); fseek(f,0,2); long n=ftell(f); rewind(f); char*s=malloc(n+1); fread(s,1,n,f); s[n]=0; fclose(f);
+    void*m=0; int r=ml(&m,s); printf("%s load=%d",argv[a],r); if(r){bad++;puts("");free(s);continue;}
+    for(char*p=strstr(s,".entry ");p;p=strstr(p+1,".entry ")){
+      char name[256]; sscanf(p+7,"%255[A-Za-z0-9_$]",name); void*fn=0; int g=gf(&fn,m,name);
+      printf(" %s=%d",name,g); if(g) bad++; }
+    puts(""); mu(m); free(s);}
+  return bad?1:0;}

kernels/tools/ptxsim.py

@@ -0,0 +1,877 @@
+#!/usr/bin/env python3
+"""Minimal vectorised PTX interpreter for DLSS rrlite kernels.
+
+Runs one thread block of a fully unrolled kernel on the CPU (all threads of the
+block as numpy lanes), against memory restored from a bridge replay dump
+(D4R_CUDA_REPLAY_DUMP_DIR). It is a reference and an instrument: every
+statement can be observed (hooks), which is how the logical computation
+(GEMMs, weight layouts, normalisations) of a kernel is recovered.
+
+Only the instruction forms the rrlite kernels use are implemented; anything
+else raises NotImplementedError naming the statement.
+"""
+import re
+import struct
+import sys
+
+import numpy as np
+
+F16, F32, F64 = np.float16, np.float32, np.float64
+U64 = np.uint64
+PARAM_BASE = 1 << 60  # virtual address of the kernel parameter block
+MASK = {8: 0xff, 16: 0xffff, 32: 0xffffffff, 64: 0xffffffffffffffff}
+
+
+def _e4m3_decode_table():
+    t = np.zeros(256, F64)
+    for c in range(256):
+        s = -1.0 if c & 0x80 else 1.0
+        e, m = (c >> 3) & 0xf, c & 7
+        if (c & 0x7f) == 0x7f:
+            t[c] = np.nan
+        elif e == 0:
+            t[c] = s * m * 2.0 ** -9
+        else:
+            t[c] = s * (1 + m / 8) * 2.0 ** (e - 7)
+    return t
+
+
+E4M3 = _e4m3_decode_table()
+
+
+def _f16_to_e4m3_table():
+    """Exact RNE + satfinite f16 -> e4m3 for all 65536 f16 codes (ZLUDA's helper)."""
+    h = np.arange(65536, dtype=np.uint32)
+    mag = h & 0x7fff
+    normal = (mag - 0x2000 + 0x3f + ((mag >> 7) & 1)) >> 7
+    exp = mag >> 10
+    sig = (mag & 0x3ff) | np.where(exp != 0, 0x400, 0)
+    cl = np.clip(exp, 1, 8)
+    sh = 16 - cl
+    sub = (sig + (np.left_shift(1, sh - 1)) - 1 + ((sig >> sh) & 1)) >> sh
+    code = np.where(mag >= 0x2400, normal, sub)
+    code = np.where(mag > 0x5f00, 0x7e, code)
+    code = np.where(mag > 0x7c00, 0x7f, code)
+    return (((h >> 8) & 0x80) | code).astype(np.uint32)
+
+
+F16_TO_E4M3 = _f16_to_e4m3_table()
+
+
+def bits_to_f16(x):
+    return (x & 0xffff).astype(np.uint16).view(F16)
+
+
+def f16_to_bits(v):
+    return np.asarray(v, F16).view(np.uint16).astype(U64)
+
+
+def bits_to_f32(x):
+    return (x & 0xffffffff).astype(np.uint32).view(F32)
+
+
+def f32_to_bits(v):
+    return np.asarray(v, F32).view(np.uint32).astype(U64)
+
+
+class Stmt:
+    __slots__ = ('labels', 'pred', 'neg', 'op', 'args', 'text', 'index')
+
+
+def split_args(s):
+    out, depth, cur = [], 0, ''
+    for ch in s:
+        if ch in '{[':
+            depth += 1
+        elif ch in '}]':
+            depth -= 1
+        if ch == ',' and depth == 0:
+            out.append(cur.strip())
+            cur = ''
+        else:
+            cur += ch
+    if cur.strip():
+        out.append(cur.strip())
+    return out
+
+
+def parse_kernel(path, name):
+    src = open(path).read()
+    i = src.index('.entry ' + name + '(')
+    j = src.find('\n.visible', i + 10)
+    j = len(src) if j < 0 else j
+    body = src[i:j]
+    body = body[body.index(')') + 1:]
+    # module-scope .shared arrays (declared before the entry) are visible to it too
+    module_shared = {m.group(1): int(m.group(2)) for m in re.finditer(r'^\.shared[^\n;]*?(\w+)\[(\d+)\]', src[:i], re.M)}
+    body = re.sub(r'//[^\n]*', '', body)
+    body = re.sub(r'\.(maxntid|reqntid|minnctapersm|maxnreg)[^\n{]*', '', body)
+    stmts, types, shared = [], {}, {}
+    pending_labels = []
+    for raw in re.sub(r'\s*\n\s*', ' ', body).split(';'):
+        s = raw.strip()
+        while True:
+            s = s.strip()
+            if s.startswith('{') and not s.startswith('{%'):
+                s = s[1:]
+                continue
+            if s.startswith('}'):
+                s = s[1:]
+                continue
+            m = re.match(r'(\$[\w]+):\s*', s)
+            if m:
+                pending_labels.append(m.group(1))
+                s = s[m.end():]
+                continue
+            break
+        if not s:
+            continue
+        if s.startswith('.reg'):
+            m = re.match(r'\.reg\s*\.(\w+)\s+(.*)', s)
+            ty = m.group(1)
+            for decl in m.group(2).split(','):
+                decl = decl.strip()
+                m2 = re.match(r'(%[\w]+)<(\d+)>', decl)
+                if m2:
+                    types[('prefix', m2.group(1))] = ty
+                else:
+                    types[decl] = ty
+            continue
+        if s.startswith('.shared'):
+            m = re.search(r'(\w+)\[(\d+)\]', s)
+            shared[m.group(1)] = int(m.group(2))
+            continue
+        if s.startswith('.'):
+            continue
+        st = Stmt()
+        st.labels, pending_labels = pending_labels, []
+        st.pred, st.neg = None, False
+        m = re.match(r'@(!?)(%\w+)\s+', s)
+        if m:
+            st.neg, st.pred = m.group(1) == '!', m.group(2)
+            s = s[m.end():]
+        parts = s.split(None, 1)
+        st.op = parts[0]
+        st.args = split_args(parts[1]) if len(parts) > 1 else []
+        st.text = s
+        st.index = len(stmts)
+        stmts.append(st)
+    if pending_labels:
+        st = Stmt()
+        st.labels, st.pred, st.neg, st.op, st.args, st.text, st.index = pending_labels, None, False, 'nop', [], 'nop', len(stmts)
+        stmts.append(st)
+    for k, v in module_shared.items():
+        shared.setdefault(k, v)
+    return stmts, types, shared
+
+
+def reg_width(types, name):
+    ty = types.get(name)
+    if ty is None:
+        m = re.match(r'(%[a-z]+)', name)
+        ty = types.get(('prefix', m.group(1)))
+    return {'pred': 1, 'b16': 16, 'u16': 16, 's16': 16, 'f16': 16, 'b32': 32, 'u32': 32, 's32': 32, 'f32': 32,
+            'b64': 64, 'u64': 64, 's64': 64, 'f64': 64, 'b8': 8}[ty]
+
+
+class Memory:
+    """Global memory: the replay dump's allocations at their captured addresses."""
+
+    def __init__(self, dump_dir):
+        manifest = open(dump_dir + '/manifest.txt').read().split('\n')
+        self.args = open(dump_dir + '/args.bin', 'rb').read()
+        self.bases, self.data, self.ids = [], [], []
+        for line in manifest:
+            p = line.split()
+            if p and p[0] == 'alloc':
+                idx, base, size = int(p[1]), int(p[2], 16), int(p[3])
+                buf = np.fromfile(f'{dump_dir}/alloc-{idx}.bin', np.uint8)
+                assert buf.size == size
+                self.bases.append(base)
+                self.data.append(buf)
+                self.ids.append(idx)
+        order = np.argsort(self.bases)
+        self.bases = [self.bases[i] for i in order]
+        self.data = [self.data[i] for i in order]
+        self.ids = [self.ids[i] for i in order]
+        self.base_arr = np.array(self.bases, U64)
+        self.stores = []  # (addresses, bytes) of every global store
+
+    def locate(self, addr):
+        idx = np.searchsorted(self.base_arr, addr, side='right') - 1
+        if np.any(idx < 0):
+            raise ValueError('address below every allocation')
+        return idx
+
+    def load(self, addr, nbytes, mask):
+        out = np.zeros((addr.size, nbytes), np.uint8)
+        idx = self.locate(np.where(mask, addr, self.base_arr[0]))
+        for a in np.unique(idx[mask]):
+            sel = mask & (idx == a)
+            off = (addr[sel] - self.bases[a]).astype(np.int64)
+            if np.any(off + nbytes > self.data[a].size):
+                raise ValueError('load past the end of an allocation')
+            out[sel] = self.data[a][off[:, None] + np.arange(nbytes)]
+        return out
+
+    def store(self, addr, raw, mask):
+        idx = self.locate(np.where(mask, addr, self.base_arr[0]))
+        n = raw.shape[1]
+        for a in np.unique(idx[mask]):
+            sel = mask & (idx == a)
+            off = (addr[sel] - self.bases[a]).astype(np.int64)
+            self.data[a][off[:, None] + np.arange(n)] = raw[sel]
+            self.stores.append((self.ids[a], off.copy(), raw[sel].copy()))
+
+
+class Block:
+    def __init__(self, stmts, types, shared, mem, param_name, block_dim, grid_dim, ctaid):
+        self.stmts, self.types, self.mem = stmts, types, mem
+        self.param_name = param_name
+        bx, by, bz = block_dim
+        self.T = bx * by * bz
+        t = np.arange(self.T)
+        self.tid = {'x': t % bx, 'y': (t // bx) % by, 'z': t // (bx * by)}
+        self.lane = t % 32
+        self.warp = t // 32
+        self.ctaid = ctaid
+        self.grid = grid_dim
+        self.ntid = block_dim
+        self.regs = {}
+        self.shared = {k: np.zeros(v, np.uint8) for k, v in shared.items()}
+        self.shared_base = {}
+        base = 0
+        for k, v in shared.items():
+            self.shared_base[k] = base
+            base += (v + 15) & ~15
+        self.shared_mem = np.zeros(max(base, 16), np.uint8)
+        self.labels = {}
+        for st in stmts:
+            for l in st.labels:
+                self.labels[l] = st.index
+        self.hooks = []
+        self.retired = np.zeros(self.T, bool)
+        self.prov = {}
+
+    # --- operands -----------------------------------------------------------
+    def width(self, name):
+        return reg_width(self.types, name)
+
+    def special(self, name):
+        m = re.match(r'%(tid|ntid|ctaid|nctaid)\.([xyz])', name)
+        if m:
+            kind, ax = m.groups()
+            k = 'xyz'.index(ax)
+            if kind == 'tid':
+                return self.tid[ax].astype(U64)
+            if kind == 'ntid':
+                return np.full(self.T, self.ntid[k], U64)
+            if kind == 'ctaid':
+                return np.full(self.T, self.ctaid[k], U64)
+            return np.full(self.T, self.grid[k], U64)
+        if name == 'WARP_SZ':
+            return np.full(self.T, 32, U64)
+        if name == '%laneid':
+            return self.lane.astype(U64)
+        if name == '%warpid':
+            return self.warp.astype(U64)
+        return None
+
+    def imm(self, s, width=64):
+        if s.startswith('0f'):
+            return np.full(self.T, int(s[2:], 16), U64)
+        if s.startswith('0d'):
+            return np.full(self.T, int(s[2:], 16), U64)
+        v = int(s, 0)
+        return np.full(self.T, v & MASK[64], U64)
+
+    def get(self, s, width=None):
+        s = s.strip()
+        if s.startswith('%') or s in self.types or s == 'WARP_SZ':
+            sp = self.special(s) if (s.startswith('%') or s == 'WARP_SZ') else None
+            if sp is not None:
+                return sp
+            if s not in self.regs:
+                w = self.width(s)
+                self.regs[s] = np.zeros(self.T, bool if w == 1 else U64)
+            return self.regs[s]
+        if s in self.shared_base:
+            return np.full(self.T, self.shared_base[s], U64)
+        if s == self.param_name:
+            return np.full(self.T, PARAM_BASE, U64)
+        if s[0].isdigit() or s[0] == '-':
+            v = self.imm(s)
+            if width:
+                v = v & np.uint64(MASK[width])
+            return v
+        raise NotImplementedError('operand ' + s)
+
+    def put(self, name, value, mask, keep_prov=False):
+        name = name.strip()
+        if not keep_prov:
+            self.prov.pop(name, None)
+        w = self.width(name)
+        old = self.get(name)
+        if w == 1:
+            self.regs[name] = np.where(mask, value.astype(bool), old)
+        else:
+            self.regs[name] = np.where(mask, value.astype(U64) & np.uint64(MASK[w]), old)
+
+    def vec(self, s):
+        return [x.strip() for x in s.strip()[1:-1].split(',')]
+
+    def address(self, s):
+        """[reg+off] or [symbol+off] -> (space, address array)."""
+        inner = s.strip()[1:-1]
+        m = re.match(r'([%\w$.]+)\s*(\+\s*-?\d+)?', inner)
+        base, off = m.group(1), int(m.group(2).replace('+', '')) if m.group(2) else 0
+        if base == self.param_name:
+            return 'param', np.full(self.T, off, U64)
+        if base in self.shared_base:
+            return 'shared', np.full(self.T, self.shared_base[base] + off, U64)
+        return 'reg', (self.get(base) + np.uint64(off & MASK[64])) & np.uint64(MASK[64])
+
+    # --- execution ------------------------------------------------------------
+    def run(self, limit=None):
+        stmts = self.stmts
+        active = np.ones(self.T, bool)
+        pending = {}
+        pc = 0
+        steps = 0
+        while pc < len(stmts):
+            st = stmts[pc]
+            for l in st.labels:
+                if l in pending:
+                    active |= pending.pop(l)
+            if not active.any():
+                # jump to the nearest label someone waits at
+                targets = sorted(self.labels[l] for l in pending)
+                if not targets:
+                    break
+                pc = targets[0]
+                continue
+            active &= ~self.retired
+            mask = active
+            if st.pred is not None:
+                p = self.get(st.pred)
+                mask = active & (~p if st.neg else p)
+            if st.op == 'bra' or st.op.startswith('bra.'):
+                target = self.labels[st.args[0]]
+                taking = mask
+                if target <= pc:
+                    if (taking == active).all():
+                        pc = target
+                        continue
+                    if not taking.any():
+                        pc += 1
+                        continue
+                    raise NotImplementedError('divergent backward branch at ' + st.text)
+                active = active & ~taking
+                if taking.any():
+                    pending[st.args[0]] = pending.get(st.args[0], np.zeros(self.T, bool)) | taking
+                pc += 1
+                continue
+            if mask.any() or st.op.startswith(('mma', 'shfl', 'movmatrix', 'bar')):
+                self.execute(st, mask)
+                for hook in self.hooks:
+                    hook(self, st, mask)
+            pc += 1
+            steps += 1
+            if limit and steps >= limit:
+                break
+        return steps
+
+    def execute(self, st, mask):
+        op, a = st.op, st.args
+        parts = op.split('.')
+        base = parts[0]
+        fn = getattr(self, 'op_' + base, None)
+        if fn is None:
+            raise NotImplementedError(st.text)
+        try:
+            fn(st, parts, a, mask)
+        except Exception as e:
+            raise RuntimeError(f'{type(e).__name__} {e} in: {st.text}') from e
+
+    # --- helpers --------------------------------------------------------------
+    @staticmethod
+    def sext(v, w):
+        v = v.astype(np.int64) & ((1 << w) - 1) if w < 64 else v.view(np.int64)
+        if w < 64:
+            v = np.where(v >= (1 << (w - 1)), v - (1 << w), v)
+        return v
+
+    def typed(self, s, ty):
+        w = int(re.sub(r'\D', '', ty) or 32) if ty not in ('pred',) else 1
+        v = self.get(s, w)
+        if ty.startswith('s'):
+            return self.sext(v, w)
+        if ty.startswith('f'):
+            if ty == 'f16':
+                return bits_to_f16(v).astype(F64)
+            if ty == 'f32':
+                return bits_to_f32(v).astype(F64)
+            return v.view(F64)
+        return (v & np.uint64(MASK[w])).astype(np.int64) if w < 64 else v
+
+    def result(self, ty, value):
+        if ty == 'f16':
+            return f16_to_bits(np.asarray(value).astype(F16))
+        if ty == 'f32':
+            return f32_to_bits(np.asarray(value).astype(F32))
+        if ty == 'f64':
+            return np.asarray(value, F64).view(U64)
+        w = int(re.sub(r'\D', '', ty))
+        return (np.asarray(value).astype(np.int64).astype(U64)) & np.uint64(MASK[w])
+
+    # --- instructions -----------------------------------------------------------
+    def op_nop(self, st, parts, a, mask):
+        pass
+
+    def op_bar(self, st, parts, a, mask):
+        pass
+
+    def op_ret(self, st, parts, a, mask):
+        self.retired |= mask
+
+    def op_neg(self, st, parts, a, mask):
+        ty = parts[-1]
+        if ty == 'f16x2':
+            self.put(a[0], self.get(a[1]) ^ np.uint64(0x80008000), mask)
+        elif ty == 'f16':
+            self.put(a[0], self.get(a[1]) ^ np.uint64(0x8000), mask)
+        elif ty == 'f32':
+            self.put(a[0], self.get(a[1]) ^ np.uint64(0x80000000), mask)
+        else:
+            self.put(a[0], self.result('u' + ty[1:], -self.typed(a[1], ty)), mask)
+
+    def op_cvta(self, st, parts, a, mask):
+        self.put(a[0], self.get(a[1]), mask)
+
+    def op_mov(self, st, parts, a, mask):
+        d, s = a
+        if d.startswith('{'):
+            names = self.vec(d)
+            v = self.get(s)
+            w = self.width(names[0])
+            for i, n in enumerate(names):
+                self.put(n, (v >> np.uint64(i * w)) & np.uint64(MASK[w]), mask)
+            return
+        if s.startswith('{'):
+            names = self.vec(s)
+            w = self.width(names[0])
+            v = np.zeros(self.T, U64)
+            for i, n in enumerate(names):
+                v |= (self.get(n) & np.uint64(MASK[w])) << np.uint64(i * w)
+            self.put(d, v, mask)
+            return
+        self.put(d, self.get(s), mask)
+        if s in self.prov:
+            self.prov[d] = self.prov[s]
+
+    def op_ld(self, st, parts, a, mask):
+        space = 'global' if 'global' in parts else 'shared' if 'shared' in parts else 'param' if 'param' in parts else None
+        ty = parts[-1]
+        vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
+        w = int(re.sub(r'\D', '', ty))
+        nbytes = w // 8
+        kind, addr = self.address(a[1])
+        total = nbytes * vec
+        if kind == 'param':
+            raw = np.frombuffer(self.mem.args, np.uint8)
+            off = int(addr[0])
+            data = np.tile(raw[off:off + total], (self.T, 1))
+        elif space == 'param':
+            # ld.param through a computed address (mov.b64 %rd, <param>; add; ld.param [%rd])
+            raw = np.frombuffer(self.mem.args, np.uint8)
+            off = np.where(mask, addr - np.uint64(PARAM_BASE), 0).astype(np.int64)
+            data = raw[off[:, None] + np.arange(total)]
+        elif kind == 'shared' or space == 'shared':
+            off = addr.astype(np.int64)
+            data = self.shared_mem[off[:, None] + np.arange(total)]
+        else:
+            data = self.mem.load(addr, total, mask)
+        names = self.vec(a[0]) if a[0].startswith('{') else [a[0]]
+        for i, n in enumerate(names):
+            chunk = data[:, i * nbytes:(i + 1) * nbytes]
+            v = np.zeros(self.T, U64)
+            for b in range(nbytes):
+                v |= chunk[:, b].astype(U64) << np.uint64(8 * b)
+            if ty.startswith('s') and w < self.width(n):
+                v = (self.sext(v, w).astype(np.int64)).astype(U64)
+            self.put(n, v, mask)
+            self.prov[n] = (st.index, kind if kind != 'reg' else space, addr + np.uint64(i * nbytes), nbytes)
+
+    def op_st(self, st, parts, a, mask):
+        ty = parts[-1]
+        vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
+        w = int(re.sub(r'\D', '', ty))
+        nbytes = w // 8
+        kind, addr = self.address(a[0])
+        names = self.vec(a[1]) if a[1].startswith('{') else [a[1]]
+        raw = np.zeros((self.T, nbytes * len(names)), np.uint8)
+        for i, n in enumerate(names):
+            v = self.get(n, w)
+            for b in range(nbytes):
+                raw[:, i * nbytes + b] = ((v >> np.uint64(8 * b)) & np.uint64(0xff)).astype(np.uint8)
+        if kind == 'shared' or 'shared' in parts:
+            off = addr.astype(np.int64)[mask]
+            self.shared_mem[off[:, None] + np.arange(raw.shape[1])] = raw[mask]
+        else:
+            self.mem.store(addr, raw, mask)
+
+    def int_binop(self, parts, a, mask, fn):
+        ty = parts[-1]
+        x, y = self.typed(a[1], ty), self.typed(a[2], ty)
+        self.put(a[0], self.result(ty if not ty.startswith('b') else 'u' + ty[1:], fn(x, y)), mask)
+
+    def op_add(self, st, parts, a, mask):
+        ty = parts[-1]
+        if ty in ('f16', 'f32'):
+            self.put(a[0], self.result(ty, self.typed(a[1], ty) + self.typed(a[2], ty)), mask)
+        elif ty == 'f16x2':
+            self.put(a[0], self.h2(a[1], a[2], lambda x, y: x + y), mask)
+        else:
+            self.int_binop(parts, a, mask, lambda x, y: x + y)
+
+    def op_sub(self, st, parts, a, mask):
+        ty = parts[-1]
+        if ty in ('f16', 'f32'):
+            self.put(a[0], self.result(ty, self.typed(a[1], ty) - self.typed(a[2], ty)), mask)
+        elif ty == 'f16x2':
+            self.put(a[0], self.h2(a[1], a[2], lambda x, y: x - y), mask)
+        else:
+            self.int_binop(parts, a, mask, lambda x, y: x - y)
+
+    def op_mul(self, st, parts, a, mask):
+        ty = parts[-1]
+        if 'wide' in parts:
+            w = int(re.sub(r'\D', '', ty))
+            x, y = self.typed(a[1], ty), self.typed(a[2], ty)
+            self.put(a[0], (x * y).astype(np.int64).astype(U64), mask)
+        elif ty in ('f16', 'f32'):
+            self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty)), mask)
+        elif ty == 'f16x2':
+            self.put(a[0], self.h2(a[1], a[2], lambda x, y: x * y), mask)
+        else:
+            self.int_binop(parts, a, mask, lambda x, y: x * y)
+
+    def op_mad(self, st, parts, a, mask):
+        ty = parts[-1]
+        x, y, z = self.typed(a[1], ty), self.typed(a[2], ty), self.typed(a[3], ty)
+        self.put(a[0], self.result('u' + ty[1:], x * y + z), mask)
+
+    def h2(self, sa, sb, fn, sc=None):
+        va, vb = self.get(sa), self.get(sb)
+        out = np.zeros(self.T, U64)
+        for i in range(2):
+            x = bits_to_f16(va >> np.uint64(16 * i)).astype(F64)
+            y = bits_to_f16(vb >> np.uint64(16 * i)).astype(F64)
+            if sc is not None:
+                z = bits_to_f16(self.get(sc) >> np.uint64(16 * i)).astype(F64)
+                r = fn(x, y, z)
+            else:
+                r = fn(x, y)
+            out |= f16_to_bits(np.asarray(r).astype(F16)) << np.uint64(16 * i)
+        return out
+
+    def op_fma(self, st, parts, a, mask):
+        ty = parts[-1]
+        if ty == 'f16x2':
+            self.put(a[0], self.h2(a[1], a[2], lambda x, y, z: x * y + z, a[3]), mask)
+        elif ty in ('f16', 'f32'):
+            self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty) + self.typed(a[3], ty)), mask)
+        else:
+            raise NotImplementedError(st.text)
+
+    def minmax(self, parts, a, mask, is_max):
+        ty = parts[-1]
+        if ty == 'f16x2':
+            f = np.fmax if is_max else np.fmin
+            self.put(a[0], self.h2(a[1], a[2], f), mask)
+        elif ty in ('f16', 'f32'):
+            f = np.fmax if is_max else np.fmin
+            self.put(a[0], self.result(ty, f(self.typed(a[1], ty), self.typed(a[2], ty))), mask)
+        else:
+            f = np.maximum if is_max else np.minimum
+            self.int_binop(parts, a, mask, f)
+
+    def op_min(self, st, parts, a, mask):
+        self.minmax(parts, a, mask, False)
+
+    def op_max(self, st, parts, a, mask):
+        self.minmax(parts, a, mask, True)
+
+    def op_abs(self, st, parts, a, mask):
+        ty = parts[-1]
+        if ty in ('f16x2', 'f16', 'f32'):
+            m = {'f16x2': 0x7fff7fff, 'f16': 0x7fff, 'f32': 0x7fffffff}[ty]
+            self.put(a[0], self.get(a[1]) & np.uint64(m), mask)
+            return
+        self.put(a[0], self.result('u' + ty[1:], np.abs(self.typed(a[1], ty))), mask)
+
+    def logic(self, a, mask, fn):
+        if self.width(a[0]) == 1:
+            self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
+        else:
+            self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
+
+    def op_and(self, st, parts, a, mask):
+        self.logic(a, mask, lambda x, y: x & y)
+
+    def op_or(self, st, parts, a, mask):
+        self.logic(a, mask, lambda x, y: x | y)
+
+    def op_xor(self, st, parts, a, mask):
+        self.logic(a, mask, lambda x, y: x ^ y)
+
+    def op_not(self, st, parts, a, mask):
+        self.put(a[0], ~self.get(a[1]), mask)
+
+    def op_shl(self, st, parts, a, mask):
+        w = int(re.sub(r'\D', '', parts[-1]))
+        n = self.get(a[2]) & np.uint64(0xff)
+        v = np.where(n >= w, 0, (self.get(a[1]) << np.minimum(n, np.uint64(63))))
+        self.put(a[0], v.astype(U64), mask)
+
+    def op_shr(self, st, parts, a, mask):
+        ty = parts[-1]
+        w = int(re.sub(r'\D', '', ty))
+        n = (self.get(a[2]) & np.uint64(0xff)).astype(np.int64)
+        if ty.startswith('s'):
+            v = self.sext(self.get(a[1]), w)
+            r = v >> np.minimum(n, w - 1)
+            self.put(a[0], r.astype(np.int64).astype(U64), mask)
+        else:
+            v = self.get(a[1]) & np.uint64(MASK[w])
+            r = np.where(n >= w, 0, v >> np.minimum(n, 63).astype(U64))
+            self.put(a[0], r.astype(U64), mask)
+
+    def op_bfi(self, st, parts, a, mask):
+        f, b = self.get(a[1]), self.get(a[2])
+        pos = self.get(a[3]).astype(np.int64) & 0xff
+        ln = self.get(a[4]).astype(np.int64) & 0xff
+        w = int(re.sub(r'\D', '', parts[-1]))
+        out = np.array(b, U64)
+        for i in range(self.T):
+            p, l = int(pos[i]), int(ln[i])
+            if l == 0 or p >= w:
+                continue
+            l = min(l, w - p)
+            m = ((1 << l) - 1) << p
+            out[i] = (int(b[i]) & ~m & MASK[w]) | ((int(f[i]) << p) & m)
+        self.put(a[0], out, mask)
+
+    def op_setp(self, st, parts, a, mask):
+        cmp, ty = parts[1], parts[-1]
+        x, y = self.typed(a[1], ty), self.typed(a[2], ty)
+        r = {'lt': x < y, 'le': x <= y, 'gt': x > y, 'ge': x >= y, 'eq': x == y, 'ne': x != y,
+             'lo': x < y, 'ls': x <= y, 'hi': x > y, 'hs': x >= y}[cmp]
+        dests = a[0].split('|')
+        self.put(dests[0], r, mask)
+
+    def op_selp(self, st, parts, a, mask):
+        p = self.get(a[3])
+        self.put(a[0], np.where(p, self.get(a[1]), self.get(a[2])), mask)
+
+    def op_ex2(self, st, parts, a, mask):
+        ty = parts[-1]
+        v = self.get(a[1])
+        if ty == 'f16x2':
+            out = np.zeros(self.T, U64)
+            for i in range(2):
+                h = bits_to_f16((v >> np.uint64(16 * i)) & np.uint64(0xffff)).astype(F64)
+                out |= f16_to_bits(np.exp2(h).astype(F16)) << np.uint64(16 * i)
+            self.put(a[0], out, mask)
+        elif ty == 'f16':
+            self.put(a[0], f16_to_bits(np.exp2(bits_to_f16(v).astype(F64)).astype(F16)), mask)
+        elif ty == 'f32':
+            self.put(a[0], f32_to_bits(np.exp2(bits_to_f32(v).astype(F64)).astype(F32)), mask)
+        else:
+            raise NotImplementedError(st.text)
+
+    def op_rsqrt(self, st, parts, a, mask):
+        x = self.typed(a[1], 'f32').astype(F32)
+        with np.errstate(divide='ignore', invalid='ignore'):
+            self.put(a[0], f32_to_bits((1.0 / np.sqrt(x.astype(F64))).astype(F32)), mask)
+
+    def op_rcp(self, st, parts, a, mask):
+        x = self.typed(a[1], 'f32')
+        with np.errstate(divide='ignore'):
+            self.put(a[0], f32_to_bits((1.0 / x).astype(F32)), mask)
+
+    def op_div(self, st, parts, a, mask):
+        x, y = self.typed(a[1], 'f32'), self.typed(a[2], 'f32')
+        with np.errstate(divide='ignore', invalid='ignore'):
+            self.put(a[0], f32_to_bits((x / y).astype(F32)), mask)
+
+    def op_cvt(self, st, parts, a, mask):
+        types = [p for p in parts[1:] if p not in ('rn', 'rz', 'rm', 'rp', 'rni', 'rzi', 'sat', 'satfinite', 'relu', 'ftz')]
+        dt, stp = types[0], types[1]
+        relu = 'relu' in parts
+        v = self.get(a[1])
+        if dt == 'e4m3x2' and stp == 'f16x2':
+            out = np.zeros(self.T, U64)
+            for i in range(2):
+                h = (v >> np.uint64(16 * i)) & np.uint64(0xffff)
+                if relu:
+                    f = bits_to_f16(h)
+                    h = np.where((f < 0) & ~np.isnan(f.astype(F32)), 0, h).astype(U64)
+                out |= F16_TO_E4M3[h.astype(np.int64)].astype(U64) << np.uint64(8 * i)
+            self.put(a[0], out, mask)
+        elif dt == 'f16x2' and stp == 'e4m3x2':
+            out = np.zeros(self.T, U64)
+            for i in range(2):
+                c = ((v >> np.uint64(8 * i)) & np.uint64(0xff)).astype(np.int64)
+                out |= f16_to_bits(E4M3[c].astype(F16)) << np.uint64(16 * i)
+            self.put(a[0], out, mask)
+        elif dt == 'f16' and stp == 'f32':
+            self.put(a[0], f16_to_bits(bits_to_f32(v).astype(F16)), mask)
+        elif dt == 'f16' and stp == 'f64':
+            self.put(a[0], f16_to_bits(v.view(F64).astype(F16)), mask)
+        elif dt == 'f32' and stp == 'f16':
+            self.put(a[0], f32_to_bits(bits_to_f16(v).astype(F32)), mask)
+        elif dt in ('f16', 'f32') and stp[0] in 'us':
+            ws = int(re.sub(r'\D', '', stp))
+            iv = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
+            if dt == 'f16':
+                self.put(a[0], f16_to_bits(iv.astype(F64).astype(F16)), mask)
+            else:
+                self.put(a[0], f32_to_bits(iv.astype(F64).astype(F32)), mask)
+        elif dt[0] in 'us' and stp[0] in 'us':
+            ws = int(re.sub(r'\D', '', stp))
+            val = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
+            self.put(a[0], val.astype(np.int64).astype(U64), mask)
+        else:
+            raise NotImplementedError(st.text)
+
+    def op_shfl(self, st, parts, a, mask):
+        mode = parts[2]
+        v = self.get(a[1])
+        b = self.get(a[2]).astype(np.int64)
+        out = np.array(v)
+        for w in range(self.T // 32):
+            sl = slice(32 * w, 32 * w + 32)
+            lane = np.arange(32)
+            if mode == 'bfly':
+                src = lane ^ b[sl]
+            elif mode == 'down':
+                src = np.where(lane + b[sl] < 32, lane + b[sl], lane)
+            elif mode == 'up':
+                src = np.where(lane - b[sl] >= 0, lane - b[sl], lane)
+            elif mode == 'idx':
+                src = b[sl] & 31
+            out[sl] = v[sl][src]
+        self.put(a[0], out, mask)
+
+    def op_movmatrix(self, st, parts, a, mask):
+        v = self.get(a[1])
+        out = np.zeros(self.T, U64)
+        for w in range(self.T // 32):
+            m = np.zeros((8, 8), U64)
+            for lane in range(32):
+                x = int(v[32 * w + lane])
+                r, c = lane // 4, (lane % 4) * 2
+                m[r, c], m[r, c + 1] = x & 0xffff, (x >> 16) & 0xffff
+            mt = m.T
+            for lane in range(32):
+                r, c = lane // 4, (lane % 4) * 2
+                out[32 * w + lane] = int(mt[r, c]) | (int(mt[r, c + 1]) << 16)
+        self.put(a[0], out, mask)
+
+    def op_mma(self, st, parts, a, mask):
+        shape = parts[3]
+        d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
+        A = [self.get(n) for n in a_names]
+        Bv = [self.get(n) for n in b_names]
+        C = [self.get(n) for n in c_names]
+        dout = [np.array(self.get(n)) for n in d_names]
+        for w in range(self.T // 32):
+            base = 32 * w
+            if shape == 'm16n8k32':
+                Am = np.zeros((16, 32)); Bm = np.zeros((32, 8)); Cm = np.zeros((16, 8))
+                for lane in range(32):
+                    g, t = lane >> 2, lane & 3
+                    for i in range(16):
+                        byte = (int(A[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
+                        row = g if (i < 4 or 8 <= i < 12) else g + 8
+                        col = t * 4 + (i & 3) + (16 if i >= 8 else 0)
+                        Am[row, col] = E4M3[byte]
+                    for i in range(8):
+                        byte = (int(Bv[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
+                        Bm[t * 4 + (i & 3) + (16 if i >= 4 else 0), g] = E4M3[byte]
+            elif shape == 'm16n8k16':
+                Am = np.zeros((16, 16)); Bm = np.zeros((16, 8)); Cm = np.zeros((16, 8))
+                for lane in range(32):
+                    g, t = lane >> 2, lane & 3
+                    for i in range(8):
+                        h = (int(A[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
+                        row = g if (i < 2 or 4 <= i < 6) else g + 8
+                        col = t * 2 + (i & 1) + (8 if i >= 4 else 0)
+                        Am[row, col] = np.uint16(h).view(F16)
+                    for i in range(4):
+                        h = (int(Bv[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
+                        Bm[t * 2 + (i & 1) + (8 if i >= 2 else 0), g] = np.uint16(h).view(F16)
+            else:
+                raise NotImplementedError(st.text)
+            for lane in range(32):
+                g, t = lane >> 2, lane & 3
+                for i in range(4):
+                    h = (int(C[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
+                    Cm[g if i < 2 else g + 8, t * 2 + (i & 1)] = np.uint16(h).view(F16)
+            Dm = (Am @ Bm + Cm).astype(F16)
+            if self.mma_hook:
+                self.mma_hook(self, st, w, Am, Bm, Cm, Dm, [self.prov.get(n) for n in a_names], [self.prov.get(n) for n in b_names])
+            for lane in range(32):
+                g, t = lane >> 2, lane & 3
+                for r in range(2):
+                    lo = int(Dm[g + 8 * r, t * 2].view(np.uint16))
+                    hi = int(Dm[g + 8 * r, t * 2 + 1].view(np.uint16))
+                    dout[r][base + lane] = lo | (hi << 16)
+        for n, v in zip(d_names, dout):
+            self.put(n, v, mask)
+
+    mma_hook = None
+
+    # wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 with register fragments: two m16n8k16 MMAs
+    # sharing A registers 0..3, B registers 0,1 / 2,3 and C/D registers 0,1 / 2,3 (the upper A and B
+    # registers of the fragment ABI are duplicates; ZLUDA's LLVMZludaSplitWmmaM16n16k16 mapping)
+    def op_wmma(self, st, parts, a, mask):
+        if 'm16n16k16' not in parts or 'mma' not in parts:
+            raise NotImplementedError(st.text)
+        d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
+        for half in range(2):
+            sub = Stmt_like(st, 'mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16',
+                            ['{%s}' % ', '.join(d_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(a_names[0:4]),
+                             '{%s}' % ', '.join(b_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(c_names[2 * half:2 * half + 2])])
+            self.op_mma(sub, sub.op.split('.'), sub.args, mask)
+
+
+class Stmt_like:
+    def __init__(self, st, op, args):
+        self.op, self.args, self.text = op, args, st.text + f'  [as {op}]'
+        for k in Stmt.__slots__:
+            if k not in ('op', 'args', 'text'):
+                setattr(self, k, getattr(st, k, None))
+
+
+def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None):
+    stmts, types, shared = parse_kernel(ptx, kernel)
+    mem = Memory(dump)
+    launch = [l for l in open(dump + '/manifest.txt') if l.startswith('launch')][0].split()
+    grid = tuple(int(x) for x in launch[1:4])
+    block = tuple(int(x) for x in launch[4:7])
+    b = Block(stmts, types, shared, mem, kernel + '_param_0', block, grid, ctaid)
+    b.hooks = list(hooks)
+    b.mma_hook = mma_hook
+    steps = b.run(limit)
+    return b, mem, steps
+
+
+if __name__ == '__main__':
+    ptx, kernel, dump = sys.argv[1:4]
+    cx, cy = (int(v) for v in sys.argv[4].split(',')) if len(sys.argv) > 4 else (0, 0)
+    import time
+    t0 = time.time()
+    b, mem, steps = run_block(ptx, kernel, dump, (cx, cy, 0))
+    print(f'block ({cx},{cy}): {steps} statements in {time.time() - t0:.1f}s, {len(mem.stores)} store groups')

kernels/tools/pwin_model.py

@@ -0,0 +1,408 @@
+#!/usr/bin/env python3
+"""Numpy reference of DLSS 4 (preset K) DltssPaddedWinLayer blocks, built stage by stage against the
+PTX interpreter (kernels/tools/ptxsim.py). f16 arithmetic in the kernel's operation order."""
+import struct
+import numpy as np
+
+F16, F32, F64 = np.float16, np.float32, np.float64
+
+
+class Dump:
+    def __init__(self, d):
+        self.dir = d
+        man = open(d + '/manifest.txt').read().split('\n')
+        self.kernel = man[0].split()[1]
+        la = man[1].split()
+        self.grid, self.block = tuple(map(int, la[1:4])), tuple(map(int, la[4:7]))
+        self.args = open(d + '/args.bin', 'rb').read()
+        self.allocs = {}
+        for l in man:
+            p = l.split()
+            if p and p[0] == 'alloc':
+                self.allocs[int(p[1])] = (int(p[2], 16), np.fromfile(f'{d}/alloc-{p[1]}.bin', np.uint8))
+
+    def u64(self, off):
+        return struct.unpack_from('<Q', self.args, off)[0]
+
+    def i32x2(self, off):
+        return struct.unpack_from('<ii', self.args, off)
+
+    def ptr(self, off):
+        a = self.u64(off)
+        for i, (base, buf) in self.allocs.items():
+            if base <= a < base + buf.size:
+                return buf, a - base
+        raise ValueError(hex(a))
+
+    def f16(self, off_param, byte, n):
+        buf, o = self.ptr(off_param)
+        return buf[o + byte:o + byte + 2 * n].view(F16).copy()
+
+
+def f16(x):
+    return np.asarray(x, F64).astype(F16)
+
+
+def mirror(v, n):
+    v = abs(v)
+    return min(v, 2 * n - 2 - v)
+
+
+def load_window(P, bx, by, C):
+    """64 tokens (raster order of the 8x8 window) x C channels, f16, mirror padded."""
+    W, H = P.i32x2(0)
+    sx, sy = P.i32x2(56)
+    buf, o = P.ptr(8)
+    x = np.zeros((64, C), F16)
+    for r in range(64):
+        ty, tx = r >> 3, r & 7
+        Y = mirror(8 * by - sy + ty, H)
+        X = mirror(8 * bx - sx + tx, W)
+        a = o + 2 * ((Y * W + X) * C)
+        x[r] = buf[a:a + 2 * C].view(F16)
+    return x
+
+
+# --- reductions in the kernel's order ---------------------------------------------------------
+def sumsq_rows(x):
+    """f16 sum of squares per token, reproducing the kernel's tree for C = 64 (per lane: 16 products
+    over 8 words (t + 4k), pairwise f16x2 adds, butterfly over the 4 lanes of a row, final hi+lo)."""
+    T, C = x.shape
+    assert C == 64
+    sq = f16(x.astype(F64) * x.astype(F64))  # f16 products
+    out = np.zeros(T, F16)
+    for r in range(T):
+        lanes = []
+        for t in range(4):
+            # word k holds channels 2(t+4k), 2(t+4k)+1 as an f16 pair
+            w = [np.array([sq[r, 2 * (t + 4 * k)], sq[r, 2 * (t + 4 * k) + 1]], F16) for k in range(8)]
+            a = [f16(w[2 * i].astype(F64) + w[2 * i + 1].astype(F64)) for i in range(4)]  # (0+1),(2+3),(4+5),(6+7)
+            b = [f16(a[0].astype(F64) + a[2].astype(F64)), f16(a[1].astype(F64) + a[3].astype(F64))]
+            lanes.append(f16(b[0].astype(F64) + b[1].astype(F64)))
+        s1 = [f16(lanes[t].astype(F64) + lanes[t ^ 1].astype(F64)) for t in range(4)]
+        s2 = f16(s1[0].astype(F64) + s1[2].astype(F64))
+        out[r] = f16(s2[0].astype(F64) + s2[1].astype(F64))
+    return out
+
+
+def rsqrt16(s):
+    return f16(1.0 / np.sqrt(s.astype(F32).astype(F64)))
+
+
+def mma(A, B, C):
+    """one wmma k16 step: D = f16(A @ B + C) with exact sums (the interpreter's semantics)."""
+    return f16(A.astype(F64) @ B.astype(F64) + C.astype(F64))
+
+
+def frag_offset(k, n):
+    """byte offset of logical B[k][n] (k, n < 16) inside one 512-byte pre-swizzled f16 fragment"""
+    return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1)
+
+
+FRAG = np.array([[frag_offset(k, n) for n in range(16)] for k in range(16)])
+
+
+def weights(P, byte, K, N, nt_stride=None):
+    """logical W[K][N] from fragments at weights + byte + 512 * kt + nt_stride * nt"""
+    buf, o = P.ptr(64)
+    nt_stride = nt_stride or 512 * (K // 16)
+    W = np.zeros((K, N), F16)
+    for kt in range(K // 16):
+        for nt in range(N // 16):
+            base = o + byte + 512 * kt + nt_stride * nt
+            vals = buf[(base + FRAG)[..., None] + np.arange(2)].reshape(16, 16, 2).view(F16)[..., 0]
+            W[16 * kt:16 * kt + 16, 16 * nt:16 * nt + 16] = vals
+    return W
+
+
+def gemm(A, W, C0=None):
+    """chained wmma k16 steps (f16 rounding after each step), per 16x16 output tile"""
+    M, K = A.shape
+    N = W.shape[1]
+    D = np.zeros((M, N), F16) if C0 is None else C0.astype(F16).copy()
+    for kt in range(K // 16):
+        D = mma(A[:, 16 * kt:16 * kt + 16], W[16 * kt:16 * kt + 16], D)
+    return D
+
+
+def hv(bits):
+    return np.array(bits, np.uint16).view(F16)
+
+
+def fma16(a, b, c):
+    return f16(a.astype(F64) * b.astype(F64) + c.astype(F64))
+
+
+def add16(a, b):
+    return f16(a.astype(F64) + b.astype(F64))
+
+
+def mul16(a, b):
+    return f16(a.astype(F64) * b.astype(F64))
+
+
+def softmax_exp(S):
+    """e = ex2(t * (1 + a t^2) + c), t = clamp(S, +-21.64) (f16x2 ops of the kernel)"""
+    hi, lo = hv(0x4D69), hv(0xCD69)
+    a, one, c = hv(0x91D5), hv(0x3C00), hv(0xC50D)
+    t = np.maximum(np.minimum(S, hi), lo)
+    u = mul16(t, t)
+    v = fma16(u, a, one)
+    w = fma16(t, v, c)
+    return f16(np.exp2(w.astype(F64)))
+
+
+def gelu16(x):
+    """x * (0.5 + c (0.41216 - 0.081081 |c|)), c = clamp(x, +-2) (f16x2 ops of the kernel)"""
+    k1, k2 = f16(np.float32(np.uint32(0x3ED306EB).view(F32))), f16(np.float32(np.uint32(0x3DA60DD6).view(F32)))
+    half, two = F16(0.5), F16(2.0)
+    c = np.minimum(np.maximum(x, -two), two)
+    u = f16(k1.astype(F64) - mul16(np.full_like(c, k2), np.abs(c)).astype(F64))
+    return mul16(x, add16(np.full_like(c, half), mul16(c, u)))
+
+
+def bias_table(P, byte):
+    """64x64 f16 table stored in C-fragment order, 512 bytes per 16x16 tile (tile 4 mt + nt)"""
+    buf, o = P.ptr(64)
+    B = np.zeros((64, 64), F16)
+    for i in range(64):
+        for j in range(64):
+            mt, nt, r, c = i >> 4, j >> 4, i & 15, j & 15
+            off = 512 * (4 * nt + mt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3)
+            B[i, j] = buf[o + byte + off:o + byte + off + 2].view(F16)[0]
+    return B
+
+
+def vec(P, byte, n):
+    return P.f16(64, byte, n)
+
+
+def pwin_enc2(P, bx, by, pm_order=((0, 0), (0, 1), (1, 0), (1, 1)), b2_in_c0=True, dbg=None):
+    """logical enc2 layer (H 2, C 64 -> merged 96) for one window; returns (y [64 x 64], merged [16 x 96])"""
+    C, H, Dh = 64, 2, 32
+    x0 = load_window(P, bx, by, C)
+    s = f16(np.sum(x0.astype(F32).astype(F64) ** 2, axis=1))
+    r = rsqrt16(s)
+    h1 = mul16(x0, mul16(r[:, None], vec(P, 0, C)[None, :]))
+    Os = []
+    for h in range(H):
+        hb = 128 + 12288 * h
+        q, k, v = [gemm(h1[:, 32:64], weights(P, hb + 2048 * (j + 3), 32, 32, 1024),
+                        gemm(h1[:, 0:32], weights(P, hb + 2048 * j, 32, 32, 1024))) for j in range(3)]
+        S = gemm(q, k.T.copy(), bias_table(P, 24704 + 8192 * h))
+        E = softmax_exp(S)
+        rs = f16(np.sum(E.astype(F64), axis=1))
+        rc = f16(1.0 / rs.astype(F32).astype(F64))
+        Pm = mul16(E, rc[:, None])
+        Os.append(gemm(Pm, v))
+        if dbg is not None:
+            dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = q, k, v, S, Pm, Os[-1]
+    O = np.concatenate(Os, axis=1)
+    bo = vec(P, 41088, C)
+    Y = gemm(O[:, 32:64], weights(P, 45312, 32, 64, 1024), gemm(O[:, 0:32], weights(P, 41216, 32, 64, 1024), np.tile(bo, (64, 1))))
+    x1 = add16(Y, x0)
+    if dbg is not None:
+        dbg['Y'], dbg['x1'] = Y, x1
+    m = mul16(x1, vec(P, 49408, C)[None, :])
+    b2 = vec(P, 49536, C)
+    acc = add16(x1, np.tile(b2, (64, 1))) if b2_in_c0 else x1
+    for c in range(8):
+        Hc = gemm(m, weights(P, 49664 + 4096 * c, 64, 32, 2048), np.tile(vec(P, 82432 + 64 * c, 32), (64, 1)))
+        acc = gemm(gelu16(Hc), weights(P, 82944 + 4096 * c, 32, 64, 1024), acc)
+        if dbg is not None:
+            dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
+    y = acc
+    # patch merge: 16 merged tokens (4x4), K = 4 sub-tokens x 64
+    A = np.zeros((16, 256), F16)
+    for mt in range(16):
+        my, mx = mt >> 2, mt & 3
+        for q, (dy, dx) in enumerate(pm_order):
+            A[mt, 64 * q:64 * q + 64] = y[8 * (2 * my + dy) + 2 * mx + dx]
+    merged = np.concatenate([gemm(A, weights(P, 115712 + 24576 * w, 256, 48, 8192), np.tile(vec(P, 164864 + 96 * w, 48), (16, 1)))
+                             for w in range(2)], axis=1)
+    return y, merged
+
+
+class Layout:
+    """weight-buffer layout of a plain DltssPaddedWinLayer (encoder with patch merge), bytes"""
+    def __init__(self, H, C, COUT):
+        self.H, self.C, self.COUT = H, C, COUT
+        self.G1 = 0
+        self.QKV = 2 * C
+        self.HEAD = 192 * C
+        self.BIAS = self.QKV + self.HEAD * H
+        self.BO = self.BIAS + 8192 * H
+        self.WO = self.BO + 2 * C
+        self.G2 = self.WO + 64 * C * H
+        self.B2 = self.G2 + 2 * C
+        self.W1 = self.B2 + 2 * C
+        self.B1 = self.W1 + 8 * C * C
+        self.W2 = self.B1 + 8 * C
+        self.PM = self.W2 + 8 * C * C
+        self.NPW = COUT // H                  # merged channels per original warp
+        self.NPA = -(-self.NPW // 16) * 16    # allocated (padded to whole 16-column tiles)
+        self.PMB = self.PM + 8 * C * self.NPA * H
+
+
+def swin_core(P, L, x0, dbg=None, posattn=False):
+    """norm -> attention -> Wo + residual -> MLP (+ residual): the block output y [64 x C]"""
+    C, H = L.C, L.H
+    NCH = C // 32
+    s = f16(np.sum(x0.astype(F32).astype(F64) ** 2, axis=1))
+    r = rsqrt16(s)
+    h1 = mul16(x0, mul16(r[:, None], vec(P, L.G1, C)[None, :]))
+    Os = []
+    for h in range(H):
+        hb = L.QKV + L.HEAD * h
+        proj = []
+        for j in range(3):
+            acc = None
+            for ch in range(NCH):
+                acc = gemm(h1[:, 32 * ch:32 * ch + 32], weights(P, hb + 2048 * (3 * ch + j), 32, 32, 1024), acc)
+            proj.append(acc)
+        q, k, v = proj
+        if posattn:
+            # position-only attention (trait flag 2): the table is the probability matrix itself
+            Os.append(gemm(bias_table(P, L.BIAS + 8192 * h), v))
+            if dbg is not None:
+                dbg[f'v{h}'], dbg[f'O{h}'] = v, Os[-1]
+            continue
+        S = gemm(q, k.T.copy(), bias_table(P, L.BIAS + 8192 * h))
+        E = softmax_exp(S)
+        rs = f16(np.sum(E.astype(F64), axis=1))
+        rc = f16(1.0 / rs.astype(F32).astype(F64))
+        Os.append(gemm(mul16(E, rc[:, None]), v))
+        if dbg is not None:
+            dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'O{h}'], dbg[f'S{h}'], dbg[f'E{h}'], dbg[f'rs{h}'] = q, k, v, Os[-1], S, E, rs
+    Y = np.tile(vec(P, L.BO, C), (64, 1))
+    for h in range(H):
+        Y = gemm(Os[h], weights(P, L.WO + 64 * C * h, 32, C, 1024), Y)
+    x1 = add16(Y, x0)
+    m = mul16(x1, vec(P, L.G2, C)[None, :])
+    acc = add16(x1, np.tile(vec(P, L.B2, C), (64, 1)))
+    if dbg is not None:
+        dbg['Y'], dbg['x1'] = Y, x1
+    for c in range(C // 8):
+        Hc = gemm(m, weights(P, L.W1 + 64 * C * c, C, 32, 32 * C), np.tile(vec(P, L.B1 + 64 * c, 32), (64, 1)))
+        acc = gemm(gelu16(Hc), weights(P, L.W2 + 64 * C * c, 32, C, 1024), acc)
+        if dbg is not None:
+            dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
+    return acc
+
+
+def pwin_encoder(P, bx, by, H, C, COUT, pm_bias_padded=True, posattn=False):
+    L = Layout(H, C, COUT)
+    x0 = load_window(P, bx, by, C)
+    y = swin_core(P, L, x0, posattn=posattn)
+    A = np.zeros((16, 4 * C), F16)
+    for mt in range(16):
+        my, mx = mt >> 2, mt & 3
+        for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
+            A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
+    Wp = weights(P, L.PM, 4 * C, L.NPA * H, 128 * C)
+    bias = vec(P, L.PMB, L.NPA * H) if pm_bias_padded else None
+    cols = []
+    for w in range(H):
+        sl = slice(L.NPA * w, L.NPA * w + L.NPW)
+        b = bias[sl] if pm_bias_padded else vec(P, L.PMB + 2 * L.NPW * w, L.NPW)
+        cols.append(gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(np.concatenate([b, np.zeros(L.NPA - L.NPW, F16)]), (16, 1)))[:, :L.NPW])
+    merged = np.concatenate(cols, axis=1)
+    return y, merged
+
+
+class Shifted:
+    """view of the weight buffer starting at a byte offset (decoders: core weights follow the expand)"""
+    def __init__(self, P, base):
+        self.P, self.base = P, base
+
+    def ptr(self, off):
+        buf, o = self.P.ptr(off)
+        return buf, o + (self.base if off == 64 else 0)
+
+    def f16(self, off_param, byte, n):
+        return self.P.f16(off_param, byte + (self.base if off_param == 64 else 0), n)
+
+    def i32x2(self, off):
+        return self.P.i32x2(off)
+
+
+def load_lowres(P, bx, by, CL):
+    W, H = P.i32x2(0)
+    sx, sy = P.i32x2(56)
+    buf, o = P.ptr(8)
+    W2, H2 = W // 2, H // 2
+    lx0, ly0 = int((8 * bx - sx) / 2), int((8 * by - sy) / 2)
+    x = np.zeros((16, CL), F16)
+    for r in range(16):
+        Y = mirror(ly0 + (r >> 2), H2); X = mirror(lx0 + (r & 3), W2)
+        a = o + 2 * ((Y * W2 + X) * CL)
+        x[r] = buf[a:a + 2 * CL].view(F16)
+    return x
+
+
+def load_skip(P, bx, by, C):
+    W, H = P.i32x2(0)
+    sx, sy = P.i32x2(56)
+    buf, o = P.ptr(16)
+    x = np.zeros((64, C), F16)
+    for r in range(64):
+        Y = mirror(8 * by - sy + (r >> 3), H); X = mirror(8 * bx - sx + (r & 7), W)
+        a = o + 2 * ((Y * W + X) * C)
+        x[r] = buf[a:a + 2 * C].view(F16)
+    return x
+
+
+def pwin_decoder(P, bx, by, H, C, CL, order=((0, 0), (0, 1), (1, 0), (1, 1))):
+    xl = load_lowres(P, bx, by, CL)
+    EXP = 2 * CL * 4 * C
+    E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
+    skip = load_skip(P, bx, by, C)
+    x0 = np.zeros((64, C), F16)
+    for r in range(16):
+        ly, lx = r >> 2, r & 3
+        for q, (dy, dx) in enumerate(order):
+            t = 8 * (2 * ly + dy) + 2 * lx + dx
+            x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
+    L = Layout(H, C, C)
+    return swin_core(Shifted(P, EXP + 8 * C), L, x0)
+
+
+def pwin_enc0(P, bx, by, H=2, C=32, COUT=64, CIN=16):
+    """first layer: token embedding x0 = in (CIN) W_e + b_e, then the core (position-only attention) and merge"""
+    xin = load_window(P, bx, by, CIN)
+    x0 = gemm(xin, weights(P, 2 * C, CIN, C, 512 * (CIN // 16)), np.tile(vec(P, 0, C), (64, 1)))
+    x0 = np.maximum(x0, F16(0))   # ReLU after the embedding
+    CB = 2 * C + 2 * CIN * C
+    L = Layout(H, C, COUT)
+    Pw = Shifted(P, CB)
+    y = swin_core(Pw, L, x0, posattn=True)
+    A = np.zeros((16, 4 * C), F16)
+    for mt in range(16):
+        my, mx = mt >> 2, mt & 3
+        for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
+            A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
+    Wp = weights(Pw, L.PM, 4 * C, L.NPA * H, 128 * C)
+    bias = vec(Pw, L.PMB, L.NPA * H)
+    cols = [gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(bias[L.NPA * w:L.NPA * (w + 1)], (16, 1)))[:, :L.NPW] for w in range(H)]
+    return y, np.concatenate(cols, axis=1)
+
+
+def pwin_dec0(P, bx, by, H=2, C=32, CL=64, NOUT=40, NOUTA=48):
+    """last layer: patch expand + skip, core (position-only attention), output head y W_h + b_h (NOUT of NOUTA)"""
+    xl = load_lowres(P, bx, by, CL)
+    EXP = 2 * CL * 4 * C
+    E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
+    skip = load_skip(P, bx, by, C)
+    x0 = np.zeros((64, C), F16)
+    for r in range(16):
+        ly, lx = r >> 2, r & 3
+        for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
+            t = 8 * (2 * ly + dy) + 2 * lx + dx
+            x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
+    L = Layout(H, C, C)
+    Pw = Shifted(P, EXP + 8 * C)
+    y = swin_core(Pw, L, x0, posattn=True)
+    HB = L.PM              # head bias, then the head weights
+    Wh = weights(Pw, HB + 2 * NOUTA, C, NOUTA, 1024)
+    out = gemm(y, Wh, np.tile(vec(Pw, HB, NOUTA), (64, 1)))
+    return y, out[:, :NOUT]

kernels/tools/swin_model.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""Parametrised numpy model of the rrlite Swin blocks (enc3 spec generalised); checked stage by stage
+against ptxsim MMA captures."""
+import numpy as np
+from model_enc3 import (q8, f16, Params, E4M3, F16, exp_trick, row_sum16, C_SCALE, CLAMP, P1, P0,
+                        gelu_poly, mma_chain)
+
+def gp(c):
+    c = np.asarray(c)
+    return 32 * (c // 32) + 2 * ((c % 32) // 4) + (c % 2) + 16 * ((c % 4) // 2)
+
+def pair_to_nat(M):
+    return M[:, gp(np.arange(M.shape[1]))]
+
+def nat_to_pair(M):
+    out = np.empty_like(M); out[:, gp(np.arange(M.shape[1]))] = M; return out
+
+def woff_table(base, Ks, Ns, K, N):
+    k = np.arange(K)[:, None]; n = np.arange(N)[None, :]; j = n // 8
+    return base + Ks * (k // 32) + Ns * (j // 2) + 64 * (n % 8) + 16 * ((k % 16) // 4) + 8 * (j % 2) + 4 * ((k % 32) // 16) + k % 4
+
+def load_tokens(P, bx, by, planes):
+    X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, P.W), P.mirror(Y, P.H)
+    out = np.zeros((64, 32 * planes))
+    for p in range(planes):
+        off = P.inp - P.abase + ((p * P.H + Y) * P.W + X) * 32
+        out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
+    return out
+
+def sumsq_tokens(x):
+    """Kernel-exact f16 sum of squares of natural-order token rows (planes of 32; lane t holds words t, t+4)."""
+    sq = f16(x * x); planes = x.shape[1] // 32
+    def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
+    lanes = []
+    for t in range(4):
+        def pair(pp, word, k):
+            c = 32 * pp + 4 * word + 2 * k
+            return sq[:, c], sq[:, c + 1]
+        def plane_sum(word):
+            p = [add(pair(pp, word, 0), pair(pp, word, 1)) for pp in range(planes)]
+            if planes == 4: return add(add(p[0], p[1]), add(p[2], p[3]))
+            if planes == 3: return add(add(p[0], p[1]), p[2])
+            if planes == 2: return add(p[0], p[1])
+            raise ValueError(planes)
+        lanes.append(add(plane_sum(t), plane_sum(t + 4)))
+    a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
+    b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
+    return f16(b[0] + b[1])[:, None]
+
+def rms_from_sum(x, s, gamma):
+    r = f16(1.0 / np.sqrt(f16(s + f16(2.0 ** -13))))
+    return q8(f16(x * f16(r * gamma[None, :])))
+
+def rel_bias(P, off):
+    raw = P.f16vec(off, 256)
+    i, j = np.meshgrid(np.arange(16), np.arange(16), indexing='ij')
+    lane = 4 * (i % 8) + (j % 8) // 2; word = i // 8 + 2 * (j // 8)
+    return raw[8 * lane + 2 * word + (j % 2)]
+
+def attention(Qq, V, bias):
+    S = f16(Qq @ Qq.T + bias)
+    t = np.clip(f16(S * C_SCALE), -CLAMP, CLAMP)
+    e1 = f16(t * (-t) + P1); poly = f16(t * e1 + P0)
+    w = exp_trick(poly); s = row_sum16(w)
+    r = f16(np.float32(1.0) / s.astype(np.float32)); Pm = f16(w * r)
+    return f16(Pm @ V)
+
+def tree_norm2(x):
+    """Kernel-exact f16 sum of squares of pair-ordered C-fragment rows (NT8 = C/8 n8 tiles)."""
+    sq = f16(x * x); nt8 = x.shape[1] // 8
+    def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
+    lanes = []
+    for t in range(4):
+        T = [(sq[:, 8 * j + 2 * t], sq[:, 8 * j + 2 * t + 1]) for j in range(nt8)]
+        def quad(js): return add(add(T[js[0]], T[js[1]]), add(T[js[2]], T[js[3]]))
+        if nt8 == 16:
+            even = add(quad([0, 2, 4, 6]), quad([8, 10, 12, 14])); odd = add(quad([1, 3, 5, 7]), quad([9, 11, 13, 15]))
+        elif nt8 == 12:
+            even = add(quad([0, 2, 4, 6]), add(T[8], T[10])); odd = add(quad([1, 3, 5, 7]), add(T[9], T[11]))
+        elif nt8 == 8:
+            even = quad([0, 2, 4, 6]); odd = quad([1, 3, 5, 7])
+        else:
+            raise ValueError(nt8)
+        lanes.append(add(even, odd))
+    a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
+    b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
+    return f16(b[0] + b[1])[:, None]
+
+def layout(C, heads, pre=0):
+    """Weight offsets of a Swin block with C channels starting at byte `pre`."""
+    HS = 96 * C + 512
+    L = dict(g1=pre, head=[pre + 2 * C + HS * h for h in range(heads)])
+    E = pre + 2 * C + heads * HS
+    L.update(bo=E, g2=E + 2 * C)
+    M0 = E + 4 * C; NC = C // 8
+    L['b1'] = [M0 if c == 0 else M0 + 66 * C + 64 + (64 * C + 64) * (c - 1) for c in range(NC)]
+    L['b2'] = M0 + 64 * C + 64
+    L['end'] = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1)
+    return L
+
+def swin_block(P, x0, C, heads, pre=0):
+    """x0: 64 x C natural (f16 values). Returns (x2 pair order, out natural fp8 values)."""
+    fp8 = lambda T: E4M3[P.wbuf[T]]
+    L = layout(C, heads, pre)
+    h1 = rms_from_sum(x0, sumsq_tokens(x0), P.f16vec(L['g1'], C)[gp(np.arange(C))])
+    O = np.zeros((64, 32 * heads))
+    for h in range(heads):
+        hb = L['head'][h]
+        Q = mma_chain(h1, fp8(woff_table(hb, 1024, 512, C, 32)))
+        V = mma_chain(h1, fp8(woff_table(hb + 32 * C, 1024, 512, C, 32)))
+        B = rel_bias(P, hb + 64 * C)
+        for m in range(4):
+            O[16 * m:16 * m + 16, 32 * h:32 * h + 32] = attention(q8(Q[16 * m:16 * m + 16]), V[16 * m:16 * m + 16], B)
+    Wo = np.zeros((32 * heads, C), np.int64)
+    for s in range(heads):
+        Wo[32 * s:32 * s + 32] = woff_table(L['head'][s] + 64 * C + 512, 0, 512, 32, C)
+    x1 = mma_chain(pair_to_nat(q8(O)), fp8(Wo), C0=f16(nat_to_pair(x0) + P.f16vec(L['bo'], C)[None, :]))
+    h2 = pair_to_nat(rms_from_sum(x1, tree_norm2(x1), P.f16vec(L['g2'], C)))
+    x2 = f16(x1 + P.f16vec(L['b2'], C)[None, :])
+    idx = gp(np.arange(32))
+    for c, B1 in enumerate(L['b1']):
+        z = mma_chain(h2, fp8(woff_table(B1, 512, 16 * C, C, 32)), C0=np.broadcast_to(P.f16vec(B1 + 32 * C, 32), (64, 32)).copy())
+        g = q8(gelu_poly(z))[:, idx]
+        x2 = f16(g @ fp8(woff_table(B1 + 32 * C + 64, 0, 512, 32, C)) + x2)
+    return x2, q8(pair_to_nat(x2)), L
+
+TOK_X = (np.arange(64) & 3) + 4 * ((np.arange(64) >> 4) & 1)
+TOK_Y = ((np.arange(64) >> 2) & 3) + 4 * (np.arange(64) >> 5)
+
+def patch_merge_input(out, C):
+    """16 merged rows (r -> (r%4, r//4)) x 4C: column C*q + ch = out[token(2mx+q%2, 2my+q//2), ch]."""
+    A = np.zeros((16, 4 * C))
+    for r in range(16):
+        mx, my = r % 4, r // 4
+        for q in range(4):
+            t = np.nonzero((TOK_X == 2 * mx + (q & 1)) & (TOK_Y == 2 * my + (q >> 1)))[0][0]
+            A[r, C * q:C * q + C] = out[t]
+    return A
+
+def load_at(P, ptr, bx, by, planes, W=None, H=None, low=False):
+    W = W or P.W; H = H or P.H
+    X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, W), P.mirror(Y, H)
+    out = np.zeros((64, 32 * planes))
+    for p in range(planes):
+        off = ptr - P.abase + ((p * H + Y) * W + X) * 32
+        out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
+    return out
+
+def decoder_input(P, bx, by, C, CIN, NW, skip_ptr):
+    """x0 = f16(q8(patch_expand(low-res CIN)) + skip), expand output col G = (4C/NW) w + n -> q = G // C, ch = ginv(G % C)."""
+    fp8 = lambda T: E4M3[P.wbuf[T]]
+    W2, H2 = P.W // 2, P.H // 2
+    A = np.zeros((16, CIN))
+    for r in range(16):
+        MX = P.mirror(np.array((8 * bx - P.sx) // 2 + r % 4), W2); MY = P.mirror(np.array((8 * by - P.sy) // 2 + r // 4), H2)
+        for p in range(CIN // 32):
+            off = P.inp - P.abase + ((p * H2 + MY) * W2 + MX) * 32
+            A[r, 32 * p:32 * p + 32] = E4M3[P.arena[off:off + 32]]
+    cols = 4 * C // NW
+    E_ = np.concatenate([mma_chain(A, fp8(woff_table(w * CIN * cols, 512, 16 * CIN, CIN, cols)),
+                                   C0=np.broadcast_to(P.f16vec(4 * C * CIN + 2 * cols * w, cols), (16, cols)).copy()) for w in range(NW)], 1)
+    skip = load_at(P, skip_ptr, bx, by, C // 32)
+    x0 = np.zeros((64, C))
+    for T in range(64):
+        tx, ty = TOK_X[T], TOK_Y[T]; r = tx // 2 + 4 * (ty // 2); q = (tx & 1) + 2 * (ty & 1)
+        G = np.arange(C) + q * C          # global expand columns of sub-position q (pair order within C)
+        x0[T] = q8(E_[r, G])[gp(np.arange(C))]
+    return f16(x0 + skip)