Merge upstream RDNA4 layouts and validate native Windows K integration

kernels/build.sh

@@ -1,14 +1,19 @@
 #!/usr/bin/env bash
-# Builds the native RDNA3 replacements for DLSS kernels into one directory that ZLUDA serves them from
+# Builds the native 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)
+#   D4R_GPU_ARCH    target GPU (default gfx1101): an RDNA3 (gfx110x) or RDNA4 (gfx120x) target; the kernels
+#                   need wave32 WMMA, and kernels/common/wmma_layout.h picks the register layout per target
+#   D4R_NATIVE_FP8  1 (RDNA4 only): the variant for ZLUDA's native FP8 WMMA (d4r.ini [Kernels] NativeFp8):
+#                   M's layers on FP8 WMMA with NVIDIA's e4m3 activations, texture tails matching ZLUDA's
+#                   D4R_ZLUDA_WMMA_FP8_NATIVE lowering. K is the same in both variants.
 # 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)
+#   D4R_ZLUDA_EMIT  ZLUDA's d4r_emit (patches/zluda applied; `cargo build --release -p ptx --example d4r_emit`),
+#                   which compiles PTX offline for D4R_GPU_ARCH (no GPU of that kind needed)
 set -euo pipefail
 
 HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -16,17 +21,26 @@ WHAT="${1:-all}"
 OUT="$(realpath -m "${2:-$HERE/out/native}")"
 ROCM="${D4R_ROCM_DIR:-/opt/rocm}"
 ARCH="${D4R_GPU_ARCH:-gfx1101}"
+FP8="${D4R_NATIVE_FP8:-0}"
+case "$ARCH" in
+    gfx11[0-9][0-9]) [[ "$FP8" == 1 ]] && { echo "D4R_NATIVE_FP8 needs an RDNA4 (gfx12) target" >&2; exit 2; } ;;
+    gfx12[0-9][0-9]) ;;
+    *) echo "unsupported D4R_GPU_ARCH $ARCH (RDNA3 gfx110x or RDNA4 gfx120x)" >&2; exit 2 ;;
+esac
 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
+    local src="$1" extra="${2:-}" name
     name="$(basename "$src" .hip)"
-    local tmp
+    local tmp flags
     tmp="$(mktemp -d)"
-    (cd "$tmp" && "$CLANG" -x hip --offload-arch="$ARCH" --offload-device-only -O3 -I"$ROCM/include" \
+    # per-kernel compiler flags from a "// d4r-build-flags: ..." line in the source
+    flags="$(sed -n 's|^// d4r-build-flags: *||p' "$src")"
+    # shellcheck disable=SC2086
+    (cd "$tmp" && "$CLANG" -x hip --offload-arch="$ARCH" --offload-device-only -O3 $flags $extra -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' ')"
@@ -39,23 +53,39 @@ if [[ "$WHAT" == all || "$WHAT" == k ]]; then
 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
+    for src in "$HERE"/m/rrlite_*.hip; do build_hip "$src" "$([[ "$FP8" == 1 ]] && echo -DD4R_FP8_WMMA)"; 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}"
+    : "${D4R_ZLUDA_EMIT:?set D4R_ZLUDA_EMIT to d4r_emit from 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"
+    # kernel : source of its native part [: w64 = compiled as wave64 (two CUDA warps per wave; RDNA3 issues
+    # FP32 work for all 64 lanes at once). Only for kernels without MMAs; measured faster, bit-identical:
+    # post 0.89 -> 0.77 ms, hiluma output 0.905 -> 0.887 ms (downsample was slower as wave64).]
+    # Every flag combination DLSS selects (motion vectors hi/lo, HDR/LDR, depth inverted/regular, ...) gets
+    # the same treatment, so games other than the one measured run the native versions too.
+    specs=(rrlite_dec0_4x4_folded:dec0_head)
+    for mv in mvhi mvlo; do
+        for range in hdr ldr; do
+            specs+=("rrlite_enc0_4x4_${mv}_${range}_folded:enc0_tail")
+            for v in 3_1 3_2; do specs+=("rrlite_post_${v}_${mv}_${range}_folded:sust_only:w64"); done
+            for depth in depthinv depthreg; do
+                for kind in "" _max; do
+                    [[ -z "$kind" ]] && specs+=("hiluma_engine_output_${depth}_${mv}_${range}_v1_rel:sust_only:w64")
+                    specs+=("hiluma_engine_output_${depth}_${mv}_${range}${kind}_v2_rel:sust_only:w64")
+                done
+            done
+        done
+    done
+    for mode in static dynamic; do
+        for range in hdr ldr; do specs+=("rrlite_downsample_kernel_${mode}_${range}:sust_only"); done
+    done
+    for spec in "${specs[@]}"; do
+        IFS=: read -r kernel src mode <<< "$spec"
+        D4R_ZLUDA_WAVE64="$([[ "$mode" == w64 ]] && echo 1 || echo 0)" D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" \
+            D4R_TEX_FP8="$FP8" D4R_DLSS_PTX_DIR="$PTX_DIR" "$HERE/tex/build_tex.sh" "$kernel" "$src" "$OUT"
     done
 fi
 echo "native kernels in $OUT"

kernels/common/hip_device_minimal.h

@@ -5,12 +5,19 @@
 // needed. Definitions below mirror public HIP qualifiers and LLVM AMDGPU
 // work-item/grid builtins. The normal Linux HIP include path is untouched.
 #define __device__ __attribute__((device))
+#define __host__ __attribute__((host))
 #define __global__ __attribute__((global))
 #define __shared__ __attribute__((shared))
+#define __constant__ __attribute__((constant))
 #define __forceinline__ __attribute__((always_inline)) inline
 #define __launch_bounds__(n) __attribute__((amdgpu_flat_work_group_size(1, n)))
 
 struct d4r_device_dim3 { unsigned x, y, z; };
+// Public HIP vector storage ABI used by M's packed global/LDS accesses.
+struct alignas(8) uint2 { unsigned x, y; };
+struct alignas(16) uint4 { unsigned x, y, z, w; };
+static_assert(sizeof(uint2) == 8 && alignof(uint2) == 8, "HIP uint2 ABI");
+static_assert(sizeof(uint4) == 16 && alignof(uint4) == 16, "HIP uint4 ABI");
 __device__ __forceinline__ d4r_device_dim3 d4r_thread_idx() {
     return {__builtin_amdgcn_workitem_id_x(), __builtin_amdgcn_workitem_id_y(),
         __builtin_amdgcn_workitem_id_z()};

kernels/common/wmma_layout.h

@@ -0,0 +1,331 @@
+// WMMA register layouts of RDNA3 (gfx11) and RDNA4 (gfx12) for the native kernels (wave32, 16x16x16 f16 -> f32).
+// Every layout fact the kernels depend on is here; the kernel sources are written against these helpers.
+//
+// Notation: lane l, half h = l >> 4, r = l & 15. A operand lane l = row r of A, B operand lane l = column r
+// of B, D lane l = column r of D.
+//   gfx11: an operand holds all 16 K values of its row/column (8 dwords, the same data in both halves);
+//          D VGPR i is row 2i + h.
+//   gfx12: an operand holds 8 of the 16 K values (4 dwords), half h the K values 8h .. 8h+7 in order;
+//          D VGPR i is row i + 8h, which is the gfx12 B-operand layout of the next GEMM (no exchange).
+// A and B only have to agree on which K each slot holds; "half h = K 8h .. 8h+7" makes a gfx12 operand the
+// lower or upper 16 bytes of the gfx11 one, so row images and weight images serve both layouts.
+//
+// D4R_WMMA_LAYOUT selects the layout (default: the target's). D4R_WMMA_LAYOUT=12 on a gfx11 target is the
+// layout shim: the gfx12 register layout with every WMMA done by the gfx11 instruction on operands and
+// accumulators converted around it, in the gfx11 K order. Its results are bit-identical to the gfx11 build,
+// so the gfx12 indexing can be validated on RDNA3 hardware.
+//
+// Only clang builtins are used (the texture tails compile with -nogpuinc -nogpulib).
+#pragma once
+#include <stdint.h>
+
+#ifndef D4R_WMMA_LAYOUT
+#if defined(__GFX12__)
+#define D4R_WMMA_LAYOUT 12
+#else
+#define D4R_WMMA_LAYOUT 11
+#endif
+#endif
+#if D4R_WMMA_LAYOUT == 12 && !defined(__GFX12__)
+#define D4R_WMMA_SHIM 1
+#endif
+#if D4R_WMMA_LAYOUT == 11 && defined(__GFX12__)
+#error "the gfx11 WMMA layout does not exist on gfx12"
+#endif
+
+#define WM_FN __attribute__((device, always_inline)) static inline
+
+typedef _Float16 wm_f16;
+typedef _Float16 wm_hv2 __attribute__((ext_vector_type(2)));
+typedef _Float16 wm_h8 __attribute__((ext_vector_type(8)));
+typedef _Float16 wm_h16 __attribute__((ext_vector_type(16)));
+typedef float wm_f8v __attribute__((ext_vector_type(8)));
+typedef uint32_t wm_u8v __attribute__((ext_vector_type(8)));
+typedef uint32_t wm_u4v __attribute__((ext_vector_type(4)));
+typedef int wm_i2v __attribute__((ext_vector_type(2)));
+
+// ---------------------------------------------------------------- lanes
+WM_FN uint32_t wm_lane()
+{
+    return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
+}
+WM_FN uint32_t wm_half()
+{
+    return wm_lane() >> 4;
+}
+// value of the lane with the same index in the other half of the wave
+WM_FN uint32_t wm_other_half(uint32_t v)
+{
+    return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
+}
+WM_FN float wm_other_half_f(float v)
+{
+    return __builtin_bit_cast(float, wm_other_half(__builtin_bit_cast(uint32_t, v)));
+}
+WM_FN uint32_t wm_pack2(wm_f16 lo, wm_f16 hi)
+{
+    return __builtin_bit_cast(uint32_t, (wm_hv2){lo, hi});
+}
+WM_FN wm_f16 wm_lo16(uint32_t v)
+{
+    return __builtin_bit_cast(wm_hv2, v)[0];
+}
+WM_FN wm_f16 wm_hi16(uint32_t v)
+{
+    return __builtin_bit_cast(wm_hv2, v)[1];
+}
+
+// ---------------------------------------------------------------- accumulator rows
+// row (of the 16-row tile) that accumulator VGPR i holds in a lane of half h (prep code: logical lanes)
+WM_FN int wm_acc_row_h(int i, int h)
+{
+#if D4R_WMMA_LAYOUT == 12
+    return i + 8 * h;
+#else
+    return 2 * i + h;
+#endif
+}
+// ... in this lane
+WM_FN int wm_acc_row(int i)
+{
+    return wm_acc_row_h(i, (int)wm_half());
+}
+// rows are the K dimension of the next GEMM: VGPRs 2j, 2j + 1 are adjacent K values in gfx12 (and 16 bytes
+// of this lane's 8 rows are contiguous), rows 2i + h are interleaved between the halves in gfx11
+#if D4R_WMMA_LAYOUT == 12
+#define WM_ACC_CONTIGUOUS 1
+#else
+#define WM_ACC_CONTIGUOUS 0
+#endif
+
+// ---------------------------------------------------------------- f16 operands
+#if D4R_WMMA_LAYOUT == 12
+typedef wm_u4v wm_op; // this half's 8 K values
+#define WM_OP_DWORDS 4
+#else
+typedef wm_u8v wm_op; // all 16 K values
+#define WM_OP_DWORDS 8
+#endif
+
+// operand from 16 contiguous halves (K 0..15) in memory (16-byte aligned); gfx12 reads only this half's 8
+WM_FN wm_op wm_op_load(const wm_f16* p)
+{
+#if D4R_WMMA_LAYOUT == 12
+    return *(const wm_u4v*)(p + 8 * wm_half());
+#else
+    const wm_u4v a = *(const wm_u4v*)p, b = *(const wm_u4v*)(p + 8);
+    return (wm_u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+#endif
+}
+
+// operand slot `idx` of an image of 32-byte slots (16 halves each, K in order): gfx12 reads its 16-byte half
+WM_FN wm_op wm_op_image(const wm_u8v* __restrict__ img, int idx)
+{
+#if D4R_WMMA_LAYOUT == 12
+    return ((const wm_u4v*)img)[2 * idx + (int)wm_half()];
+#else
+    return img[idx];
+#endif
+}
+
+// element (K value) of an operand that holds accumulator row wm_acc_row(i) of the same lane
+WM_FN wm_f16 wm_op_acc_elem(const wm_op& v, int i)
+{
+#if D4R_WMMA_LAYOUT == 12
+    return (i & 1) ? wm_hi16(v[i >> 1]) : wm_lo16(v[i >> 1]);
+#else
+    const int c = 2 * i + (int)wm_half();
+    return (c & 1) ? wm_hi16(v[c >> 1]) : wm_lo16(v[c >> 1]);
+#endif
+}
+
+// the f16 accumulator values own[i] (row wm_acc_row(i) of this lane's column) as an operand of the next GEMM
+// (the rows become K; gfx11 fetches the other parity from lane l ^ 16)
+WM_FN wm_op wm_op_from_acc(const wm_f16 own[8])
+{
+#if D4R_WMMA_LAYOUT == 12
+    return (wm_u4v){wm_pack2(own[0], own[1]), wm_pack2(own[2], own[3]), wm_pack2(own[4], own[5]), wm_pack2(own[6], own[7])};
+#else
+    const uint32_t hf = wm_half();
+    uint32_t mine[4], theirs[4];
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        mine[j] = wm_pack2(own[2 * j], own[2 * j + 1]);
+        theirs[j] = wm_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;
+    wm_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
+    wm_u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+    {
+        const wm_f16 a = (j & 1) ? wm_hi16(mine[j >> 1]) : wm_lo16(mine[j >> 1]);
+        const wm_f16 b = (j & 1) ? wm_hi16(theirs[j >> 1]) : wm_lo16(theirs[j >> 1]);
+        r[j] = hf ? wm_pack2(b, a) : wm_pack2(a, b);
+    }
+    return r;
+#endif
+#endif
+}
+
+WM_FN wm_op wm_op_from_f8(wm_f8v d)
+{
+    wm_f16 h[8];
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        h[i] = (wm_f16)d[i];
+    return wm_op_from_acc(h);
+}
+
+// Stores an operand back as a row of 16 halves (K 0..15). In gfx11 both halves hold the whole row and the
+// lanes of half 0 store it; in gfx12 each half stores its 8 values. `count` (8 or 16) limits the row to its
+// first `count` halves; `ok` is the lane's own condition.
+WM_FN void wm_op_store(wm_f16* p, const wm_op& v, bool ok = true, int count = 16)
+{
+#if D4R_WMMA_LAYOUT == 12
+    if (ok && (wm_half() == 0 || count > 8))
+        *(wm_u4v*)(p + 8 * wm_half()) = v;
+#else
+    if (ok && wm_half() == 0)
+    {
+        *(wm_u4v*)p = (wm_u4v){v[0], v[1], v[2], v[3]};
+        if (count > 8)
+            *(wm_u4v*)(p + 8) = (wm_u4v){v[4], v[5], v[6], v[7]};
+    }
+#endif
+}
+
+// 8 halves of an f16 vector in accumulator row order: out[i] = v[wm_acc_row(i)], v at p (16 halves, 16-byte
+// aligned)
+WM_FN void wm_acc_vec_load(const wm_f16* p, wm_f16 out[8])
+{
+#if D4R_WMMA_LAYOUT == 12
+    const wm_u4v a = *(const wm_u4v*)(p + 8 * wm_half());
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        out[i] = (i & 1) ? wm_hi16(a[i >> 1]) : wm_lo16(a[i >> 1]);
+#else
+    const uint32_t hf = wm_half();
+    const wm_u4v a = *(const wm_u4v*)p, b = *(const wm_u4v*)(p + 8);
+    const wm_u8v v = (wm_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 ? wm_hi16(v[i]) : wm_lo16(v[i]);
+#endif
+}
+
+// ---------------------------------------------------------------- the WMMA itself
+#if D4R_WMMA_LAYOUT == 12 && defined(__GFX12__)
+WM_FN wm_f8v wm_mma(const wm_op& a, const wm_op& b, wm_f8v c)
+{
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(__builtin_bit_cast(wm_h8, a), __builtin_bit_cast(wm_h8, b), c);
+}
+#else
+WM_FN wm_f8v wm_mma11(const wm_u8v& a, const wm_u8v& b, wm_f8v c)
+{
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(wm_h16, a), __builtin_bit_cast(wm_h16, b), c);
+}
+#if D4R_WMMA_LAYOUT == 12
+// layout shim: gfx12 operand halves -> one gfx11 operand in K order (half 0's 8 K, then half 1's)
+WM_FN wm_u8v wm_shim_operand(const wm_u4v& v)
+{
+    const bool hf = wm_half() != 0;
+    wm_u8v r;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        const uint32_t o = wm_other_half(v[j]);
+        r[j] = hf ? o : v[j];
+        r[4 + j] = hf ? v[j] : o;
+    }
+    return r;
+}
+// gfx12 accumulator (VGPR i = row i + 8h) <-> gfx11 accumulator (VGPR i = row 2i + h)
+WM_FN wm_f8v wm_shim_acc_to11(const wm_f8v& c)
+{
+    const bool hf = wm_half() != 0;
+    wm_f8v o, r;
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        o[i] = wm_other_half_f(c[i]);
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+    {
+        // row 2i + h lives in half (2i + h) >> 3, VGPR (2i + h) & 7
+        const float h0 = i < 4 ? c[2 * i] : o[2 * i - 8];
+        const float h1 = i < 4 ? o[2 * i + 1] : c[2 * i - 7];
+        r[i] = hf ? h1 : h0;
+    }
+    return r;
+}
+WM_FN wm_f8v wm_shim_acc_to12(const wm_f8v& d)
+{
+    const bool hf = wm_half() != 0;
+    wm_f8v o, r;
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        o[i] = wm_other_half_f(d[i]);
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+    {
+        // row i + 8h lives in half i & 1, VGPR (i >> 1) + 4h
+        const float h0 = (i & 1) ? o[i >> 1] : d[i >> 1];
+        const float h1 = (i & 1) ? d[(i >> 1) + 4] : o[(i >> 1) + 4];
+        r[i] = hf ? h1 : h0;
+    }
+    return r;
+}
+WM_FN wm_f8v wm_mma(const wm_op& a, const wm_op& b, wm_f8v c)
+{
+    return wm_shim_acc_to12(wm_mma11(wm_shim_operand(a), wm_shim_operand(b), wm_shim_acc_to11(c)));
+}
+// ZLUDA's K order: its gfx11 helpers put K pair p (K 2p, 2p + 1) in operand dword p, its gfx12 helpers put
+// in half h, dword j the pair (j & 1) * 4 + 2h + (j >> 1) (see zluda_ptx_impl.cpp)
+WM_FN wm_u8v wm_shim_operand_zluda(const wm_u4v& v)
+{
+    const bool hf = wm_half() != 0;
+    wm_u4v o;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+        o[j] = wm_other_half(v[j]);
+    wm_u8v r;
+#pragma unroll
+    for (int p = 0; p < 8; ++p)
+    {
+        const int hp = (p & 3) >> 1, jp = 2 * (p & 1) + (p >> 2);
+        r[p] = (hp != 0) == hf ? v[jp] : o[jp];
+    }
+    return r;
+}
+WM_FN wm_f8v wm_mma_zluda(const wm_op& a, const wm_op& b, wm_f8v c)
+{
+    return wm_shim_acc_to12(wm_mma11(wm_shim_operand_zluda(a), wm_shim_operand_zluda(b), wm_shim_acc_to11(c)));
+}
+#else
+WM_FN wm_f8v wm_mma(const wm_op& a, const wm_op& b, wm_f8v c)
+{
+    return wm_mma11(a, b, c);
+}
+#endif
+#endif
+#if D4R_WMMA_LAYOUT == 12 && defined(__GFX12__)
+// operands whose K slots follow ZLUDA's gfx12 order: the same instruction (the order only matters to the shim)
+WM_FN wm_f8v wm_mma_zluda(const wm_op& a, const wm_op& b, wm_f8v c)
+{
+    return wm_mma(a, b, c);
+}
+#endif
+
+WM_FN wm_f8v wm_splat8(float v)
+{
+    return (wm_f8v){v, v, v, v, v, v, v, v};
+}

kernels/k/dltss_pwin_dec0_layer.hip

@@ -1,3 +1,5 @@
+// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
+// d4r-build-flags: -mcumode
 #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

@@ -1,3 +1,5 @@
+// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
+// d4r-build-flags: -mcumode
 #define PWIN_F32ACC
 #include "pwin_layer.h"
 PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)

kernels/k/dltss_pwin_dec2_layer.hip

@@ -1,3 +1,5 @@
+// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
+// d4r-build-flags: -mcumode
 #define PWIN_F32ACC
 #include "pwin_layer.h"
 PWIN_DECODER(dltss_pwin_dec2_layer, 2, 64, 96)

kernels/k/dltss_pwin_enc0_layer.hip

@@ -1,3 +1,5 @@
+// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
+// d4r-build-flags: -mcumode
 #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

@@ -1,3 +1,5 @@
+// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
+// d4r-build-flags: -mcumode
 #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

@@ -1,3 +1,5 @@
+// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
+// d4r-build-flags: -mcumode
 #define PWIN_F32ACC
 #include "pwin_layer.h"
 PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)

kernels/k/pwin_common.h

@@ -1,10 +1,11 @@
-// Native RDNA3 (gfx11 wave32 WMMA) kernels for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
+// Native kernels (gfx11 / gfx12 wave32 WMMA) 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").
+// Y^T = W^T X^T with the weights as the WMMA A operand (lane l: output channel l & 15) and the activations
+// as the B operand (lane l: token l & 15). The f32 result then has lane l = token l & 15 and VGPR i =
+// output channel wm_acc_row(i) of the 16-channel tile ("D^T"); ../common/wmma_layout.h has the per-target
+// register layouts (gfx11: channel 2i + (l >> 4); gfx12: channel i + 8 (l >> 4)).
 #pragma once
 #ifdef D4R_DEVICE_ONLY_MINIMAL
 #include "../common/hip_device_minimal.h"
@@ -12,7 +13,7 @@
 #include <hip/hip_runtime.h>
 #endif
 #include <stdint.h>
-
+#include "../common/wmma_layout.h"
 #pragma clang fp contract(off)
 
 typedef _Float16 half_t;
@@ -21,23 +22,19 @@ 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)));
-#include "../common/wmma_backend.h"
 
 __device__ __forceinline__ uint32_t lane_id()
 {
     return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
 }
 
-__device__ __forceinline__ f8v wmma(u8v a, u8v b, f8v c)
-{
-    return d4r_wmma_legacy_layout(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
-}
+typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
 
 // 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)
+__device__ __forceinline__ f8v mma16(const op_t& a, const op_t& b, f8v c)
 {
-    f8v d = wmma(a, b, c);
+    f8v d = wm_mma(a, b, c);
 #ifndef PWIN_F32ACC
 #pragma unroll
     for (int i = 0; i < 8; ++i)
@@ -70,50 +67,19 @@ __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])
+// D^T f16 values (own[i] = channel wm_acc_row(i) of token l & 15) -> operand of the same token (the
+// channels become K)
+__device__ __forceinline__ op_t 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
+    return wm_op_from_acc(own);
 }
 
-__device__ __forceinline__ u8v operand_from_f8(f8v d)
+__device__ __forceinline__ op_t 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);
+    return wm_op_from_f8(d);
 }
 
-// element c (0..15) of a B operand (16 halves)
+// element c (0..15) of a row of 16 halves (VALU code; not a WMMA operand)
 __device__ __forceinline__ half_t op_get(const u8v& v, int c)
 {
     return (c & 1) ? hi16(v[c >> 1]) : lo16(v[c >> 1]);
@@ -150,6 +116,58 @@ __device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
     gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
 }
 
+// WMMA operands from rows of 16 channels: LDS, global (PWIN_NT_LOAD as above), prep images (32-byte slots)
+__device__ __forceinline__ op_t op_lds(const half_t* p)
+{
+    return wm_op_load(p);
+}
+__device__ __forceinline__ op_t op_gload(const half_t* p)
+{
+#if D4R_WMMA_LAYOUT == 12
+#ifdef PWIN_NT_LOAD
+    return __builtin_nontemporal_load((const u4v*)(p + 8 * wm_half()));
+#else
+    return *(const u4v*)(p + 8 * wm_half());
+#endif
+#else
+    return gload_row16(p);
+#endif
+}
+__device__ __forceinline__ op_t op_img(const u8v* __restrict__ img, int idx)
+{
+    return wm_op_image(img, idx);
+}
+// operand -> slot idx of an operand image in LDS (read back with op_img)
+__device__ __forceinline__ void op_img_store(u8v* img, int idx, const op_t& v)
+{
+#if D4R_WMMA_LAYOUT == 12
+    ((u4v*)img)[2 * idx + wm_half()] = v;
+#else
+    if (!wm_half())
+        img[idx] = v;
+#endif
+}
+// operand (16 channels of this lane's token) -> row of `count` (8 or 16) channels in global memory, if ok
+__device__ __forceinline__ void op_gstore(half_t* dst, const op_t& v, bool ok, int count = 16)
+{
+#if D4R_WMMA_LAYOUT == 12
+    if (ok && (wm_half() == 0 || count >= 16))
+        gstore16(dst + 8 * wm_half(), v);
+#else
+    if (!wm_half() && ok)
+    {
+        gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
+        if (count >= 16)
+            gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
+    }
+#endif
+}
+// operand -> row of 16 channels in LDS
+__device__ __forceinline__ void op_store(half_t* dst, const op_t& v)
+{
+    wm_op_store(dst, v);
+}
+
 __device__ __forceinline__ void block_sync()
 {
     __syncthreads();
@@ -167,27 +185,3 @@ __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

@@ -72,14 +72,10 @@ __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])
+// channels 16 nt + wm_acc_row(i) (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, 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]);
+    wm_acc_vec_load((const half_t*)(w + byte), out);
 }
 
 __device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
@@ -131,7 +127,7 @@ template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u
 #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 key = 16 * kt + wm_acc_row_h(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);
     }
@@ -209,6 +205,45 @@ __device__ uint64_t g_stamp[16384][16];
 __device__ float g_dbg[12][64][160];
 #define DBG_ON (blockIdx.x == DBX && blockIdx.y == DBY)
 #endif
+// Softmax row sum over the 64 keys of this lane's query: this lane's 32 accumulator rows (summed into rs as
+// the scores are computed), plus the other half's. The layout shim (gfx12 layout on gfx11) instead redoes the
+// sum in gfx11's order (keys 16 kt + 2 i + h, then the other parity), which the f32 result depends on, so
+// that its output stays bit-identical to the gfx11 build.
+#ifdef D4R_WMMA_SHIM
+constexpr bool kShimRowSum = true;
+#else
+constexpr bool kShimRowSum = false;
+#endif
+__device__ __forceinline__ float row_sum_total(float rs, const half_t (&e)[4][8])
+{
+    if constexpr (!kShimRowSum)
+        return rs + __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+    // this lane: keys 16 kt + 8 h + j; the other half: keys 16 kt + 8 (1 - h) + j
+    const bool hf = wm_half() != 0;
+    float s[2] = {0.0f, 0.0f};
+#pragma unroll
+    for (int kt = 0; kt < 4; ++kt)
+    {
+        half_t lo[8], hi[8];
+#pragma unroll
+        for (int j = 0; j < 8; j += 2)
+        {
+            const uint32_t own = pack2(e[kt][j], e[kt][j + 1]), other = other_half(own);
+            const uint32_t l = hf ? other : own, u = hf ? own : other;
+            lo[j] = lo16(l), lo[j + 1] = hi16(l), hi[j] = lo16(u), hi[j + 1] = hi16(u);
+        }
+#pragma unroll
+        for (int par = 0; par < 2; ++par)
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const int key = 2 * i + par;
+                s[par] += (float)(key < 8 ? lo[key] : hi[key - 8]);
+            }
+    }
+    return s[0] + s[1];
+}
+
 // 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>
@@ -265,7 +300,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
 
     half_t vv_[8];
 
-    dvec8(w, L::BO + 32 * nt, hf, vv_);
+    dvec8(w, L::BO + 32 * nt, vv_);
 
 #pragma unroll
 
@@ -277,7 +312,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
     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];
+        op_t 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
@@ -287,28 +322,26 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
             {
                 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;
+                    d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
+                op_img_store(&vsh[h][nt][0][0], wv * 16 + m, operand_from_f8(d));
             }
             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];
+                    vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
         }
-        u8v pop[4];
+        op_t 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];
+                pop[kt] = op_img((const u8v*)bias, ((h * 4 + wv) * 4 + kt) * 16 + m);
         }
         else
         {
-        u8v qop[2];
+        op_t qop[2];
         // six independent chains (Q, K, V x 2 n-tiles), each over kt in order
         f8v dq[6];
 #pragma unroll
@@ -316,10 +349,10 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
             dq[u] = splat8(0.0f);
         for (int kt = 0; kt < KT; ++kt)
         {
-            const u8v row = lds_row16(hrow + 16 * kt);
+            const op_t row = op_lds(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]);
+                dq[u] = mma16(op_img(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)
@@ -331,15 +364,13 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
                     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);
+                    op_store(&kl[16 * wv + m][16 * nt], operand_from_f8(d));
                 }
                 else
                 {
 #pragma unroll
                     for (int i = 0; i < 8; ++i)
-                        vt[16 * nt + 2 * i + hf][16 * wv + m] = (half_t)d[i];
+                        vt[16 * nt + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
                 }
             }
         block_sync();
@@ -358,15 +389,16 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
             }
 #pragma unroll
             for (int ds = 0; ds < 2; ++ds)
-                d = mma16(lds_row16(&kl[16 * kt + m][16 * ds]), qop[ds], d);
+                d = mma16(op_lds(&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];
+                if constexpr (!kShimRowSum)
+                    rs += (float)e[kt][i];
             }
         }
-        rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+        rs = row_sum_total(rs, e);
         const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
 #pragma unroll
         for (int kt = 0; kt < 4; ++kt)
@@ -378,7 +410,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
             pop[kt] = operand_from_dt(pv);
         }
         }
-        u8v oop[2];
+        op_t oop[2];
 #pragma unroll
         for (int nt = 0; nt < 2; ++nt)
         {
@@ -389,20 +421,20 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
                 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);
+                    d = mma16(op_lds(&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];
+                    g_dbg[2 + h][16 * wv + m][16 * nt + wm_acc_row(i)] = (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]);
+                acc[nt] = mma16(op_img(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
     }
@@ -411,34 +443,32 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
     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];
+                g_dbg[1][16 * wv + m][16 * nt + wm_acc_row(i)] = 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);
+        const op_t xv = op_lds(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);
+        dvec8(w, L::G2 + 32 * nt, g2);
+        dvec8(w, L::B2 + 32 * nt, 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);
+            const half_t x1 = (half_t)acc[nt][i] + wm_op_acc_elem(xv, i);
             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);
+        op_store(hrow + 16 * nt, operand_from_dt(mv));
     }
     __builtin_amdgcn_wave_barrier();
 
     // ---- MLP: C/8 chunks of 32 hidden
     for (int c = 0; c < L::NMLP; ++c)
     {
-        u8v gop[2];
+        op_t gop[2];
 #pragma unroll
         for (int hn = 0; hn < 2; ++hn)
         {
@@ -447,7 +477,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
 
     half_t vv_[8];
 
-    dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv_);
+    dvec8(w, L::B1 + 64 * c + 32 * hn, vv_);
 
 #pragma unroll
 
@@ -457,7 +487,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
 
 }
             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);
+                d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
             half_t g[8];
             gelu8(d, g);
             gop[hn] = operand_from_dt(g);
@@ -466,16 +496,14 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
         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]);
+                acc[nt] = mma16(op_img(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);
+        op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
     }
     __builtin_amdgcn_wave_barrier();
 }
@@ -530,31 +558,29 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
         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];
+            op_t xe[L::KE];
 #pragma unroll
             for (int ks = 0; ks < L::KE; ++ks)
-                xe[ks] = gload_row16(xin + 16 * ks);
+                xe[ks] = op_gload(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_);
+                    dvec8(p0.w, 32 * nt, 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);
+                    d = mma16(op_img(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);
+                op_store(&act[tok][16 * nt], operand_from_dt(xo));
             }
         }
         else
@@ -592,7 +618,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
 
     half_t vv_[8];
 
-    dvec8(w, L::PMB + 32 * vt, hf, vv_);
+    dvec8(w, L::PMB + 32 * vt, vv_);
 
 #pragma unroll
 
@@ -606,18 +632,11 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
         {
             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);
+            d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&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]});
-        }
+        op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
     }
 }
 
@@ -749,7 +768,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
 
     half_t vv_[8];
 
-    dvec8(p.w, L::EXPB + 32 * gnt, hf, vv_);
+    dvec8(p.w, L::EXPB + 32 * gnt, vv_);
 
 #pragma unroll
 
@@ -759,15 +778,13 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
 
 }
             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);
+                d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
+            const op_t sv = op_gload(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);
+                xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
+            op_store(&act[t][16 * nt], operand_from_dt(xo));
         }
     }
     block_sync();
@@ -788,7 +805,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
 
     half_t vv_[8];
 
-    dvec8(q.w, L::HB + 32 * nt, hf, vv_);
+    dvec8(q.w, L::HB + 32 * nt, vv_);
 
 #pragma unroll
 
@@ -799,15 +816,8 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
 }
 #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]});
-            }
+                d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
+            op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
         }
     }
     else if (inb)

kernels/k/pwin_pos.h

@@ -53,7 +53,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
     for (int nt = 0; nt < KT; ++nt)
     {
         half_t vv[8];
-        dvec8(w, L::BO + 32 * nt, hf, vv);
+        dvec8(w, L::BO + 32 * nt, vv);
 #pragma unroll
         for (int mi = 0; mi < MT; ++mi)
 #pragma unroll
@@ -67,7 +67,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
 #endif
     for (int h = 0; h < H; ++h)
     {
-        u8v vtop[2][4];
+        op_t vtop[2][4];
 #ifdef POS_VSHARE
 #pragma unroll
         for (int nt = 0; nt < 2; ++nt)
@@ -78,17 +78,15 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
                 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;
+                    d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
+                op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_f8(d));
             }
         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];
+                vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
 #else
 #pragma unroll
         for (int nt = 0; nt < 2; ++nt)
@@ -98,19 +96,19 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
                 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);
+                    d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(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];
+        op_t oop[MT][2];
 #pragma unroll
         for (int mi = 0; mi < MT; ++mi)
         {
             const int qt = wv * MT + mi;
-            u8v pop[4];
+            op_t pop[4];
 #pragma unroll
             for (int kt = 0; kt < 4; ++kt)
-                pop[kt] = ((const u8v*)bias)[((h * 4 + qt) * 4 + kt) * 16 + m];
+                pop[kt] = op_img((const u8v*)bias, ((h * 4 + qt) * 4 + kt) * 16 + m);
 #pragma unroll
             for (int nt = 0; nt < 2; ++nt)
             {
@@ -126,7 +124,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
 #pragma unroll
             for (int ks = 0; ks < 2; ++ks)
             {
-                const u8v wo = img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m];
+                const op_t wo = op_img(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]);
@@ -139,24 +137,22 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
     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);
+        dvec8(w, L::G2 + 32 * nt, g2);
+        dvec8(w, L::B2 + 32 * nt, 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]);
+            const op_t xv = op_lds(&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);
+                const half_t x1 = (half_t)acc[mi][nt][i] + wm_op_acc_elem(xv, i);
                 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);
+            op_store(&hb[tok][16 * nt], operand_from_dt(mv));
         }
     }
     __builtin_amdgcn_wave_barrier();
@@ -165,12 +161,12 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
     // ---- 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];
+        op_t 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);
+            dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
             f8v d[MT];
 #pragma unroll
             for (int mi = 0; mi < MT; ++mi)
@@ -180,10 +176,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
 #pragma unroll
             for (int kt = 0; kt < KT; ++kt)
             {
-                const u8v w1 = img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m];
+                const op_t w1 = op_img(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]);
+                    d[mi] = mma16(w1, op_lds(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
             }
 #pragma unroll
             for (int mi = 0; mi < MT; ++mi)
@@ -198,7 +194,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
 #pragma unroll
             for (int ks = 0; ks < 2; ++ks)
             {
-                const u8v w2 = img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m];
+                const op_t w2 = op_img(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]);
@@ -211,9 +207,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
 #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);
+            op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
         }
     block_sync();
 }
@@ -239,29 +233,27 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
         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];
+            op_t xe[L::KE];
 #pragma unroll
             for (int ks = 0; ks < L::KE; ++ks)
-                xe[ks] = gload_row16(xin + 16 * ks);
+                xe[ks] = op_gload(xin + 16 * ks);
 #pragma unroll
             for (int nt = 0; nt < C / 16; ++nt)
             {
                 half_t vv[8];
-                dvec8(p0.w, 32 * nt, hf, vv);
+                dvec8(p0.w, 32 * nt, 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);
+                    d = mma16(op_img(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);
+                op_store(&act[tok][16 * nt], operand_from_dt(xo));
             }
         }
         else
@@ -296,7 +288,7 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
     for (int vt = wv; vt < L::PMT; vt += NW)
     {
         half_t vv[8];
-        dvec8(p.w, L::PMB + 32 * vt, hf, vv);
+        dvec8(p.w, L::PMB + 32 * vt, vv);
         f8v d;
 #pragma unroll
         for (int i = 0; i < 8; ++i)
@@ -305,17 +297,10 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
         {
             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);
+            d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&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]});
-        }
+        op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
     }
     STAMP(8);
 }
@@ -333,10 +318,10 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
         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];
+        op_t xls[L::KL];
 #pragma unroll
         for (int ks = 0; ks < L::KL; ++ks)
-            xls[ks] = gload_row16(xl + 16 * ks);
+            xls[ks] = op_gload(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;
@@ -347,22 +332,20 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
             {
                 const int gnt = q * (C / 16) + nt;
                 half_t vv[8];
-                dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
+                dvec8(p.w, L::EXPB + 32 * gnt, 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);
+                    d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), xls[ks], d);
+                const op_t sv = op_gload(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);
+                    xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
+                op_store(&act[t][16 * nt], operand_from_dt(xo));
             }
         }
     }
@@ -380,22 +363,15 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
         for (int nt = 0; nt < L::NOUTA / 16; ++nt)
         {
             half_t vv[8];
-            dvec8(q.w, L::HB + 32 * nt, hf, vv);
+            dvec8(q.w, L::HB + 32 * nt, 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]});
-            }
+                d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
+            op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
         }
     }
 }

kernels/k/pwin_wide.h

@@ -61,7 +61,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
                           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)
+    // residual input of the owned output tiles (D^T order: channel 16 nt + wm_acc_row(i))
     half_t x0[NTL][8];
 #pragma unroll
     for (int j = 0; j < NTL; ++j)
@@ -69,10 +69,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
         const int nt = g + NG * j;
         if (nt < KT)
         {
-            const u8v xv = lds_row16(arow + 16 * nt);
+            const op_t xv = op_lds(arow + 16 * nt);
 #pragma unroll
             for (int i = 0; i < 8; ++i)
-                x0[j][i] = op_get(xv, 2 * i + hf);
+                x0[j][i] = wm_op_acc_elem(xv, i);
         }
     }
     block_sync(); // hb complete, R free
@@ -86,7 +86,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
         if (nt < KT)
         {
             half_t vv[8];
-            dvec8(w, L::BO + 32 * nt, hf, vv);
+            dvec8(w, L::BO + 32 * nt, vv);
 #pragma unroll
             for (int i = 0; i < 8; ++i)
                 acc[j][i] = (float)vv[i];
@@ -99,7 +99,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
     for (int r = 0; r < D::ROUNDS; ++r)
     {
         const int h = g < NA ? r * NA + g : H;
-        u8v qop[2], oop[2];
+        op_t qop[2], oop[2];
         if (h < H)
         {
             // six independent chains (Q, K, V x 2 n-tiles), each over kt in order
@@ -109,10 +109,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
                 dq[u] = splat8(0.0f);
             for (int kt = 0; kt < KT; ++kt)
             {
-                const u8v row = lds_row16(hrow + 16 * kt);
+                const op_t row = op_lds(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]);
+                    dq[u] = mma16(op_img(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)
@@ -124,15 +124,13 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
                         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);
+                        op_store(&kl[g][tok][16 * nt], operand_from_f8(d));
                     }
                     else
                     {
 #pragma unroll
                         for (int i = 0; i < 8; ++i)
-                            vt[g][16 * nt + 2 * i + hf][tok] = (half_t)d[i];
+                            vt[g][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
                     }
                 }
         }
@@ -154,17 +152,18 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
                 }
 #pragma unroll
                 for (int ds = 0; ds < 2; ++ds)
-                    d = mma16(lds_row16(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
+                    d = mma16(op_lds(&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];
+                    if constexpr (!kShimRowSum)
+                        rs += (float)e[kt][i];
                 }
             }
-            rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+            rs = row_sum_total(rs, e);
             const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
-            u8v pop[4];
+            op_t pop[4];
 #pragma unroll
             for (int kt = 0; kt < 4; ++kt)
             {
@@ -180,15 +179,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
                 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);
+                    d = mma16(op_lds(&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)
+        if (h < H)
         {
-            store_row16(&ob[g][tok][0], oop[0]);
-            store_row16(&ob[g][tok][16], oop[1]);
+            op_store(&ob[g][tok][0], oop[0]);
+            op_store(&ob[g][tok][16], oop[1]);
         }
         block_sync();
         for (int jh = 0; jh < NA; ++jh)
@@ -196,15 +195,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
             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]);
+            const op_t o0 = op_lds(&ob[jh][tok][0]), o1 = op_lds(&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]);
+                    acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh) * KT + nt) * 16 + m), o0, acc[j]);
+                    acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m), o1, acc[j]);
                 }
             }
         }
@@ -220,8 +219,8 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
         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);
+            dvec8(w, L::G2 + 32 * nt, g2);
+            dvec8(w, L::B2 + 32 * nt, b2);
 #pragma unroll
             for (int i = 0; i < 8; ++i)
             {
@@ -229,9 +228,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
                 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);
+            op_store(hrow + 16 * nt, operand_from_dt(mv));
         }
     }
     block_sync();
@@ -249,18 +246,16 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
             f8v d;
             {
                 half_t vv[8];
-                dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv);
+                dvec8(w, L::B1 + 64 * c + 32 * hn, 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);
+                d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(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);
+            op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
         }
         block_sync();
         for (int jc = 0; jc < NG; ++jc)
@@ -268,15 +263,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
             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]);
+            const op_t h0 = op_lds(&hid[rr % D::NB][jc][tok][0]), h1 = op_lds(&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]);
+                    acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2) * KT + nt) * 16 + m), h0, acc[j]);
+                    acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m), h1, acc[j]);
                 }
             }
         }
@@ -290,9 +285,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
         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);
+            op_store(hrow + 16 * nt, operand_from_f8(acc[j]));
         }
     }
     block_sync();
@@ -349,7 +342,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
         f8v d;
         {
             half_t vv[8];
-            dvec8(w, L::PMB + 32 * vt, hf, vv);
+            dvec8(w, L::PMB + 32 * vt, vv);
 #pragma unroll
             for (int i = 0; i < 8; ++i)
                 d[i] = (float)vv[i];
@@ -358,17 +351,10 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
         {
             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);
+            d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&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]});
-        }
+        op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
     }
 }
 
@@ -409,21 +395,19 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
             f8v d;
             {
                 half_t vv[8];
-                dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
+                dvec8(p.w, L::EXPB + 32 * gnt, 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);
+                d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
+            const op_t sv = op_gload(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);
+                xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
+            op_store(&act[tt][16 * nt], operand_from_dt(xo));
         }
     }
     block_sync();

kernels/m/rrlite_enc1_4x4.hip

@@ -3,5 +3,6 @@
 #define SWIN_F32ACC
 #define SWIN_MAX_VGPR 192
 #define SWIN_MLP_R(NW) 1
+#define SWIN_SCALAR_BLOAD // scalar-base weight loads: faster here (no spills), slower in the 8-wave layers
 #include "swin_block.h"
 SWIN_MODULE_W(rrlite_enc1_4x4, 64, 2, 4, 96, false, CommonParams, 0)

kernels/m/rrlite_enc3_tube_4x4.hip

@@ -3,5 +3,6 @@
 #define SWIN_F32ACC
 #define SWIN_MAX_VGPR 192
 #define SWIN_MLP_R(NW) 2
+#define SWIN_PREP_SLOTS 8 // one weight image per tube block (six per frame)
 #include "swin_block.h"
 SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, 8, 0, true, TubeParams, 0)

kernels/m/swin_block.h

@@ -5,7 +5,10 @@
 //   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>
+
+// Immediate lane masks need only a compile-time integer, not the host C++
+// standard library (which is absent from device-only Windows HIP builds).
+template <int Mask> struct SwinLaneMask { static constexpr int value = Mask; };
 
 template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
 {
@@ -19,7 +22,7 @@ template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
     __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)
+    // weight image (slots of 16 K values, see expand_weights)
     __host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
     static constexpr int WO = 4 * C * NW;
     static constexpr int W1 = WO + 2 * NW * C;
@@ -33,6 +36,14 @@ template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
     static constexpr int NDESC = 2 * NW + 1 + 2 * NC + NPM / 32 + (CIN ? NW : 0);
 };
 
+// Start (in halves) of token T's row in the output staging buffer: stride C + 4 plus a shift by bits 4 and 5
+// of T. The patch merge's lanes read 8 bytes each from 16 tokens that differ in bits 1, 3, 4 and 5 (plus a
+// constant); this placement puts those 16 reads on 16 distinct 8-byte bank pairs for C = 64, 96 and 128.
+template <int C> __host__ __device__ constexpr int xo_row(int T)
+{
+    return (C + 4) * T + 4 * ((T >> 4) & 1) + 16 * (T >> 5);
+}
+
 template <int C, int NW, int NPM, int CIN> struct SwinDescs
 {
     GemmDesc d[SwinLayout<C, NW, NPM, CIN>::NDESC];
@@ -61,7 +72,7 @@ template <int C, int NW, int NPM, int CIN> constexpr SwinDescs<C, NW, NPM, CIN>
 
 // 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,
+__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const int bx,
                                            const int by, const int gx)
 {
     using L = SwinLayout<C, NH, NPM, CIN>;
@@ -77,8 +88,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
     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);
+    constexpr int U_ST4 = R * 64 * GST * 2;
+    constexpr int USIZE = SCR > U_ST4 ? SCR : U_ST4;
 
     __shared__ __attribute__((aligned(16))) half_t A[64 * AST];
     __shared__ __attribute__((aligned(16))) uint8_t U[USIZE];
@@ -118,11 +129,11 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 
     // 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");
+    constexpr int NSW = 64 * C / 4 / (32 * NW);
+    static_assert(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)
+    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));
@@ -130,26 +141,24 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
         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));
     }
+    // the input words as f16 into X0 (natural channel order); decoders add the patch expand to them later
+    auto stage_input_words = [&]() {
+#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);
+        }
+    };
 
     // ------------------------------------------------ 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)
         {
@@ -161,6 +170,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
             for (int b = 0; b < 4; ++b)
                 PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
         }
+        stage_input_words(); // the skip input
         __syncthreads();
         constexpr int NTEW = 4 * C / NW / 16; // expand tiles per wave (L::NTE per weight block)
         T16 e[NTEW];
@@ -170,7 +180,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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);
+            const gop a0 = lda(PA + l16 * PAST + 32 * kc), a1 = lda(PA + l16 * PAST + 32 * kc + 16);
 #pragma unroll
             for (int j = 0; j < NTEW; ++j)
             {
@@ -185,28 +195,16 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
             for (int i = 0; i < 8; ++i)
             {
-                const int r = 2 * i + hi;
+                const int r = mrow(i);
                 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]));
+                // e4m3_to_half is exact in f16, so adding the staged skip value is the same sum
+                X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + X0[T * AST + ch]);
             }
         }
     }
     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);
-        }
-    }
-    }
+        stage_input_words();
     __syncthreads();
 
     // residual accumulators of the output projection: f16(x0 + b_o) in the C-fragment layout
@@ -223,8 +221,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
             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});
+                x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + mrow(i)) * AST + ch],
+                                         X0[(16 * (mg * MT + mm) + mrow(i + 1)) * AST + ch]} + (hv2){bo, bo});
         }
     __syncthreads(); // stage 1 overwrites x0 rows with h1
 
@@ -308,7 +306,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
         for (int kc = 0; kc < C / 32; ++kc)
         {
-            h16 b[2][4];
+            gop b[2][4];
 #pragma unroll
             for (int s = 0; s < 2; ++s)
 #pragma unroll
@@ -317,8 +315,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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);
+                const gop a0 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
+                const gop a1 = lda(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]);
@@ -344,7 +342,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
         for (int i = 0; i < 8; ++i)
         {
-            const int ri = 2 * i + hi, cj = l16;
+            const int ri = mrow(i), 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))));
         }
@@ -359,14 +357,14 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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];
+                    Qs[mrow(i) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
+                    Qs[mrow(i + 1) * 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);
+            const gop a0 = lda(Qs + l16 * QST), a1 = lda(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
+            // rows mrow(2k), mrow(2k+1) 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;
@@ -392,21 +390,25 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
                 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 v1 = wg[k] + xr(wg[k], SwinLaneMask<8>{});
+                const hv2 v2 = v1 + xr(v1, SwinLaneMask<2>{});
+                const hv2 v3 = v2 + xr(v2, SwinLaneMask<4>{});
+                const hv2 sum = v3 + xr(v3, SwinLaneMask<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];
+                Qs[mrow(2 * k) * PST + l16] = pw[0];
+                Qs[mrow(2 * k + 1) * PST + l16] = pw[1];
             }
             wave_sync();
-            const h16 pa = lds16(Qs + l16 * PST);
+            const sop 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
+                // V as B operand: lane holds column l16 over the 16 tokens (rows, the K of P V)
+#if D4R_WMMA_LAYOUT == 12
+                // own pair k holds rows (2k + 8hi, 2k + 1 + 8hi): this half's K values in order
+                const sop vb = __builtin_bit_cast(sop, vp[m][nt]);
+#else
                 // 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;
@@ -419,14 +421,15 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
                     vw[2 * k + 1] = __builtin_amdgcn_perm(other, own, s_hi); // rows 4k+2, 4k+3
                 }
                 const h16 vb = __builtin_bit_cast(h16, vw);
+#endif
                 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];
+                    A[(16 * (m0 + m) + mrow(i)) * AST + ocol] = oq[0];
+                    A[(16 * (m0 + m) + mrow(i + 1)) * AST + ocol] = oq[1];
                 }
             }
             wave_sync();
@@ -438,7 +441,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
     for (int kc = 0; kc < NH; ++kc)
     {
-        h16 b[2][NTW];
+        gop b[2][NTW];
 #pragma unroll
         for (int s = 0; s < 2; ++s)
 #pragma unroll
@@ -447,8 +450,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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);
+            const gop a0 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
+            const gop a1 = lda(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]);
@@ -463,7 +466,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
         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);
+                A[(16 * (mg * MT + mm) + mrow(i)) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
     __syncthreads();
     if (tid < 64)
     {
@@ -514,12 +517,12 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
             for (int i = 0; i < 8; i += 2)
             {
-                const int row = 16 * (mg * MT + mm) + 2 * i + hi;
+                const int row = 16 * (mg * MT + mm) + mrow(i), row2 = 16 * (mg * MT + mm) + mrow(i + 1);
                 const hv2 v = x[mm][j].pair(i / 2);
-                const hv2 r2 = {R2[row], R2[row + 2]};
+                const hv2 r2 = {R2[row], R2[row2]};
                 const hv2 hq = q8x2(v * (hv2)(r2 * (hv2){g2, g2}));
                 A[row * AST + col] = hq[0];
-                A[(row + 2) * AST + col] = hq[1];
+                A[row2 * AST + col] = hq[1];
                 x[mm][j].set_pair(i / 2, v + (hv2){b2, b2});
             }
     }
@@ -546,7 +549,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
             for (int kc = 0; kc < C / 32; ++kc)
             {
-                h16 b[2][2];
+                gop b[2][2];
 #pragma unroll
                 for (int s = 0; s < 2; ++s)
 #pragma unroll
@@ -555,8 +558,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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);
+                    const gop a0 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
+                    const gop a1 = lda(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]);
                 }
@@ -579,8 +582,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
                                           (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];
+                        Gc[(16 * (mb + mm) + mrow(i)) * GST + col] = gq[0];
+                        Gc[(16 * (mb + mm) + mrow(i + 1)) * GST + col] = gq[1];
                     }
                 }
             __syncthreads();
@@ -588,7 +591,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
             for (int cc = 0; cc < R; ++cc)
             {
                 const int c2 = R * r + cc;
-                h16 b[2][NTW];
+                gop b[2][NTW];
 #pragma unroll
                 for (int s = 0; s < 2; ++s)
 #pragma unroll
@@ -597,8 +600,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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);
+                    const gop a0 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
+                    const gop a1 = lda(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]);
@@ -609,7 +612,10 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
     }
 
     // ------------------------------------------------ output: f16 staged in LDS, 16-value segments encoded by a loop
-    half_t* XO = A; // [64 tokens][C channels] natural order
+    // XO: [64 tokens][C channels] natural order; the row padding keeps the patch merge's 8-byte reads of 16
+    // tokens on distinct LDS banks (see xo_row)
+    half_t* XO = A;
+    static_assert(xo_row<C>(63) + C <= 64 * AST, "XO fits in A");
 #pragma unroll
     for (int j = 0; j < NTW; ++j)
     {
@@ -618,26 +624,28 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
         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);
+                XO[xo_row<C>(16 * (mg * MT + mm) + mrow(i)) + 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));
+        // rows are 8-byte aligned (xo_row), so the segment moves as four 8-byte words of two f16 pairs
+        uint2* xs = (uint2*)(XO + xo_row<C>(T) + c16);
         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 uint2 pr = xs[k];
+            hv2 qa, qb;
+            const uint32_t ca = enc8x2(__builtin_bit_cast(hv2, pr.x), qa), cb = enc8x2(__builtin_bit_cast(hv2, pr.y), qb);
+            w[k] = pack_codes(ca, cb);
+            // the patch merge consumes decode(code) = q8_exact(x), written back over x (this thread's own values)
+            if constexpr (NPM > 0)
+                xs[k] = (uint2){__builtin_bit_cast(uint32_t, qa), __builtin_bit_cast(uint32_t, qb)};
         }
         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;
@@ -663,23 +671,28 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
         for (int kc = 0; kc < C / 8; ++kc)
         {
-            h16 as[2];
+            gop as[2];
 #pragma unroll
             for (int s = 0; s < 2; ++s)
             {
-                h16& a = as[s];
+                sop a;
+                // the operand's K slots in runs of 4 (gfx12: this half's slots 8 hi .. 8 hi + 7)
+                constexpr int RUNS = sizeof(sop) / sizeof(half_t) / 4;
 #pragma unroll
-                for (int run = 0; run < 4; ++run)
+                for (int rr = 0; rr < RUNS; ++rr)
                 {
+                    const int run = RUNS == 4 ? rr : 2 * hi + rr;
                     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);
+                    typedef _Float16 h4 __attribute__((ext_vector_type(4)));
+                    const h4 v = *(const h4*)(XO + xo_row<C>(T) + ch);
 #pragma unroll
                     for (int b = 0; b < 4; ++b)
-                        a[4 * run + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
+                        a[4 * rr + b] = v[b];
                 }
+                as[s] = as_gop(a);
             }
 #pragma unroll
             for (int j = 0; j < NTP; ++j)
@@ -700,7 +713,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #pragma unroll
             for (int i = 0; i < 8; ++i)
             {
-                const int r = 2 * i + hi;
+                const int r = mrow(i);
                 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;
@@ -734,35 +747,79 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
 #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.
+#ifndef SWIN_PREP_SLOTS
+#define SWIN_PREP_SLOTS 0 // weight images kept per weights pointer (0: one image, re-prepared when the pointer changes)
+#endif
+
+// Weight-image slots for kernels launched with several weight sets (the six enc3 tube blocks). The ZLUDA
+// hook runs the prep only for a weights pointer it has not seen (d4r_prep_key_at / d4r_prep_key_slots);
+// the prep claims the first free slot for that pointer, and pointers beyond S share the overflow slot S,
+// which the hook re-prepares on every launch. Keys are written by one thread of the prep; any prep
+// thread that reads before that write finds the same (first free) slot.
+template <int S> __device__ __forceinline__ int prep_slot_find(const uint64_t* keys, uint64_t w)
+{
+    for (int i = 0; i < S; ++i)
+        if (__hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) == w)
+            return i;
+    return S;
+}
+template <int S> __device__ __forceinline__ int prep_slot_claim(uint64_t* keys, uint64_t w, bool writer)
+{
+    // one pass: the first slot holding w or nothing. Keys only ever go from 0 to a pointer, so every
+    // prep thread picks the same slot whether it reads that slot before or after the writer's store.
+    for (int i = 0; i < S; ++i)
+    {
+        const uint64_t k = __hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
+        if (k == w)
+            return i;
+        if (k == 0)
+        {
+            if (writer)
+                __hip_atomic_store(keys + i, w, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
+            return i;
+        }
+    }
+    return S;
+}
+
+// Module boilerplate: f16 weight image(s), 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];                                                                              \
+    constexpr int NAME##_S = SWIN_PREP_SLOTS;                                                                           \
+    __device__ wslot g_w16[(NAME##_S + 1) * NAME##_L::TOTAL];                                                           \
+    __device__ uint64_t g_prep_keys[NAME##_S + 1];                                                                      \
     __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;              \
+    /* the prep reads only p.w (offset 0): skip it while the image for that pointer is resident */                     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;                                      \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = NAME##_S;                                \
     extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PARAMS p)                                             \
     {                                                                                                                   \
         const int idx = blockIdx.x * 128 + threadIdx.x;                                                                 \
+        const uint8_t* w = ((const CommonParams*)&p)->w;                                                                \
+        const int slot = NAME##_S ? prep_slot_claim<NAME##_S>(g_prep_keys, (uint64_t)w, idx == 0) : 0;                 \
         if (idx < NAME##_L::TOTAL)                                                                                      \
-            expand_weights(((const CommonParams*)&p)->w, g_w16, g_descs.d, NAME##_L::NDESC, idx);                      \
+            expand_weights(w, g_w16 + slot * NAME##_L::TOTAL, 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;                                                              \
+        const int slot = NAME##_S ? prep_slot_find<NAME##_S>(g_prep_keys, (uint64_t)cp.w) : 0;                         \
+        const wslot* w16 = g_w16 + slot * NAME##_L::TOTAL;                                                              \
         /* 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);    \
+            swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, 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);            \
+            swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, b % gx, b / gx, gx);              \
             __syncthreads();                                                                                            \
         }                                                                                                               \
     }

kernels/m/swin_common.h

@@ -1,8 +1,33 @@
-// Shared device code for native RDNA3 (gfx11 wave32 WMMA) rrlite Swin-block kernels (DLSS 310.7.0).
+// Shared device code for native rrlite Swin-block kernels (DLSS 310.7.0; gfx11 / gfx12 wave32 WMMA).
 // Semantics: ../ENC3_SPEC.md generalised (numpy reference ../swin_model.py).
+//
+// Activations are the WMMA A operand (lane = token row), weights the B operand (lane = channel column), and
+// the f32 result has lane = channel column, VGPR i = token row wm_acc_row(i) (../common/wmma_layout.h).
+// D4R_FP8_WMMA (gfx12 layout only): the GEMMs whose operands are both e4m3 values run as native FP8 WMMAs on
+// e4m3 bytes (weights prepared as bytes, activations encoded when loaded); P V stays f16.
 #pragma once
+#ifdef D4R_DEVICE_ONLY_MINIMAL
+#include "../common/hip_device_minimal.h"
+#else
 #include <hip/hip_runtime.h>
+#endif
 #include <stdint.h>
+#include "../common/wmma_layout.h"
+#if defined(D4R_FP8_WMMA) && D4R_WMMA_LAYOUT != 12 && defined(__AMDGCN__)
+#error "D4R_FP8_WMMA needs the gfx12 WMMA layout"
+#endif
+// FP8 operands must hold e4m3 values, so the FP8 build re-quantises the intermediate activations as NVIDIA's
+// network does (the gfx11 "fast numerics" SWIN_NO_Q8 keeps them f16, which no gfx12 WMMA can multiply with
+// e4m3 weights)
+// (SWIN_FORCE_Q8 does the same on any target: the gfx11 reference of the FP8 build's numerics)
+#if (defined(D4R_FP8_WMMA) || defined(SWIN_FORCE_Q8) || defined(SWIN_EXACT)) && defined(SWIN_NO_Q8)
+#undef SWIN_NO_Q8
+#endif
+// SWIN_EXACT (tests): NVIDIA's numerics throughout, the semantics of tools/swin_model.py (e4m3 activations and
+// an f16 accumulator rounded after every k32 step)
+#if defined(SWIN_EXACT) && defined(SWIN_F32ACC)
+#undef SWIN_F32ACC
+#endif
 
 #pragma clang fp contract(off)
 
@@ -82,12 +107,9 @@ 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)
+// q8_exact on two halves with packed 16-bit integer / f16 ops (same per-half logic; no carries cross halves)
+__device__ __forceinline__ hv2 q8x2_exact(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;
@@ -102,6 +124,48 @@ __device__ __forceinline__ hv2 q8x2(hv2 h)
     return __builtin_bit_cast(hv2, (u16x2)(r | (u & (unsigned short)0x8000)));
 }
 
+__device__ __forceinline__ hv2 q8x2(hv2 h)
+{
+#ifdef SWIN_NO_Q8
+    return h;
+#endif
+    return q8x2_exact(h);
+}
+
+// enc8 of two halves with packed ops: returns the two e4m3 codes in the low bytes of the 16-bit
+// halves, and q = q8_exact(h) (the values the codes decode to)
+__device__ __forceinline__ uint32_t codes8x2(hv2 q);
+__device__ __forceinline__ uint32_t enc8x2(hv2 h, hv2& q)
+{
+    q = q8x2_exact(h);
+    return codes8x2(q);
+}
+
+// the e4m3 codes of two halves that already hold e4m3 values (q8_exact results), in the low bytes of the
+// 16-bit halves
+__device__ __forceinline__ uint32_t codes8x2(hv2 q)
+{
+    const u16x2 qb = __builtin_bit_cast(u16x2, q);
+    const u16x2 m = qb & (unsigned short)0x7fff;
+    // normals: the exponent moves from bias 15 to bias 7
+    const u16x2 norm = (u16x2)(m >> 7) - (unsigned short)64;
+    // subnormals are k * 2^-9 (k = 0..7): f16(k + 1024) has k in its low bits
+    const hv2 mag = __builtin_bit_cast(hv2, m);
+    const hv2 k = (hv2)(mag * (hv2){(half_t)512.0f, (half_t)512.0f}) + (hv2){(half_t)1024.0f, (half_t)1024.0f};
+    const u16x2 sub = __builtin_bit_cast(u16x2, k) - (unsigned short)0x6400;
+    const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
+    u16x2 code = (sub & lt) | (norm & ~lt);
+    const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
+    code = ((u16x2){0x7f, 0x7f} & gt) | (code & ~gt);
+    return __builtin_bit_cast(uint32_t, (u16x2)(((qb >> 8) & (unsigned short)0x80) | code));
+}
+
+// four e4m3 codes (a, b from enc8x2) as the bytes of one word in element order
+__device__ __forceinline__ uint32_t pack_codes(uint32_t a, uint32_t b)
+{
+    return __builtin_amdgcn_perm(b, a, 0x06040200u);
+}
+
 // 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)
 {
@@ -153,9 +217,105 @@ __device__ __forceinline__ int woff(int base, int Ks, int Ns, int k, int n)
            4 * ((k & 31) >> 4) + (k & 3);
 }
 
-__device__ __forceinline__ f8v wmma(h16 a, h16 b, f8v c)
+// ---------------------------------------------------------------- WMMA operands
+// half-wave of this lane: blocks are 32 x 1 x waves, so the lane is threadIdx.x (the same value as wm_half(),
+// but from the source the kernels already use for their lane index)
+__device__ __forceinline__ int m_half()
 {
+    return (int)threadIdx.x >> 4;
+}
+// token row that accumulator VGPR i holds in this lane
+__device__ __forceinline__ int mrow(int i)
+{
+    return wm_acc_row_h(i, m_half());
+}
+// sop: an f16 operand (gfx11: the row's 16 K values; gfx12: this half's 8). gop: an operand of the e4m3 GEMMs
+// (k32 steps): sop, or with D4R_FP8_WMMA the 8 e4m3 bytes of this half's K values.
+#if D4R_WMMA_LAYOUT == 12
+typedef h8 sop;
+#else
+typedef h16 sop;
+#endif
+typedef uint32_t u2v __attribute__((ext_vector_type(2)));
+typedef uint32_t u4v_t __attribute__((ext_vector_type(4)));
+#ifdef D4R_FP8_WMMA
+typedef u2v gop;
+typedef u4v_t wslot; // weight image slot: 16 e4m3 bytes (K slots 0..15)
+#else
+typedef sop gop;
+typedef h16 wslot; // weight image slot: 16 f16 values
+#endif
+
+__device__ __forceinline__ f8v wmma(sop a, sop b, f8v c)
+{
+#if D4R_WMMA_LAYOUT == 12
+    return wm_mma(__builtin_bit_cast(wm_op, a), __builtin_bit_cast(wm_op, b), c);
+#else
     return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
+#endif
+}
+
+// f16 operand from 16 contiguous halves in LDS (gfx12: this half's 8 of them)
+__device__ __forceinline__ sop lds16(const half_t* p)
+{
+#if D4R_WMMA_LAYOUT == 12
+    return *(const h8*)(p + 8 * m_half());
+#else
+    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);
+#endif
+}
+
+#ifdef D4R_FP8_WMMA
+// e4m3 bytes of an f16 operand whose values are e4m3 values (K order kept: byte j = element j)
+__device__ __forceinline__ gop to_fp8(sop v)
+{
+    uint32_t c[4];
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+        c[j] = codes8x2((hv2){v[2 * j], v[2 * j + 1]});
+    return (u2v){__builtin_amdgcn_perm(c[1], c[0], 0x06040200u), __builtin_amdgcn_perm(c[3], c[2], 0x06040200u)};
+}
+// f32 += A B over 16 K for e4m3 operands
+__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
+{
+#if defined(__GFX12__)
+    return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(wm_i2v, a), __builtin_bit_cast(wm_i2v, b), c);
+#else
+    // layout shim: the e4m3 values widened to f16 (exact), then the f16 WMMA
+    auto widen = [](gop q) {
+        sop r;
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+            r[j] = e4m3_to_half((q[j >> 2] >> (8 * (j & 3))) & 0xffu);
+        return r;
+    };
+    return wmma(widen(a), widen(b), c);
+#endif
+}
+__device__ __forceinline__ gop lda(const half_t* p)
+{
+    return to_fp8(lds16(p));
+}
+#else
+__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
+{
+    return wmma(a, b, c);
+}
+// GEMM A operand (activations) from LDS
+__device__ __forceinline__ gop lda(const half_t* p)
+{
+    return lds16(p);
+}
+#endif
+// an f16 operand built in registers -> GEMM operand
+__device__ __forceinline__ gop as_gop(sop v)
+{
+#ifdef D4R_FP8_WMMA
+    return to_fp8(v);
+#else
+    return v;
+#endif
 }
 
 __device__ __forceinline__ f8v round16(f8v v)
@@ -166,7 +326,7 @@ __device__ __forceinline__ f8v round16(f8v v)
     return v;
 }
 
-// A C-fragment tile: element i is row 2i + lane/16. Exact mode stores packed f16 (NVIDIA's f16
+// A C-fragment tile: element i is row wm_acc_row(i). 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
@@ -217,20 +377,20 @@ __device__ __forceinline__ void mixacc(T16& acc, f8v P)
 __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)
+__device__ __forceinline__ void k32(T16& acc, gop a0, gop b0, gop a1, gop 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));
+    acc.f = wmma8(a1, b1, wmma8(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);
+    f8v P = wmma8(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
+    P = wmma8(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));
+    c = wmma8(a1, b1, wmma8(a0, b0, c));
 #pragma unroll
     for (int i = 0; i < 8; ++i)
         acc.set(i, (half_t)c[i]);
@@ -273,12 +433,6 @@ __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");
@@ -296,21 +450,36 @@ __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.
+// ---------------------------------------------------------------- weight images
+// A GEMM's B operands in WMMA lane order: slot index dst + ((kc * 2 + s) * NT + nt) * 16 + lane16 holds the 16
+// K slots of logical k = 32 kc + kslot(s, i), n = 16 nt + lane16 (i = 0..15), as f16 values (wslot = h16) or,
+// with D4R_FP8_WMMA, as the e4m3 bytes. gfx12 lanes read the half of the slot with their 8 K slots.
 struct GemmDesc
 {
-    int dst; // in h16 units
+    int dst; // in slots
     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)
+// The tile index is wave-uniform (it depends only on the wave, the loop counters and constants), so it
+// is formed in scalar registers and the load uses a scalar base plus the lane's constant offset.
+__device__ __forceinline__ gop bload(const wslot* 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];
+#ifdef SWIN_SCALAR_BLOAD
+    const int tile = __builtin_amdgcn_readfirstlane(gemm_dst + ((kc * 2 + s) * NT + nt) * 16);
+    const wslot* t = w16 + tile;
+    const int at = l16;
+#else
+    const wslot* t = w16;
+    const int at = gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16;
+#endif
+#if D4R_WMMA_LAYOUT == 12
+    return ((const gop*)t)[2 * at + m_half()];
+#else
+    return t[at];
+#endif
 }
 
-__device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const GemmDesc* descs, int ndesc, int idx)
+__device__ __forceinline__ void expand_weights(const uint8_t* w, wslot* w16, const GemmDesc* descs, int ndesc, int idx)
 {
     int d = 0;
     while (d + 1 < ndesc && idx >= descs[d + 1].dst)
@@ -322,9 +491,22 @@ __device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const
     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;
+#ifdef D4R_FP8_WMMA
+    u4v_t v;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        uint32_t word = 0;
+#pragma unroll
+        for (int b = 0; b < 4; ++b)
+            word |= (uint32_t)w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, 4 * j + b), n)] << (8 * b);
+        v[j] = word;
+    }
+#else
     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)]);
+#endif
     w16[idx] = v;
 }

kernels/tex/build_tex.sh

@@ -2,25 +2,32 @@
 # 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.
+# Called by kernels/build.sh, which sets D4R_ROCM_DIR, D4R_GPU_ARCH, D4R_DLSS_PTX_DIR, D4R_ZLUDA_EMIT,
+# D4R_ZLUDA_WAVE64 (1 = wave64 code object) and D4R_TEX_FP8 (1 = the gfx12 native-FP8 variant).
+# ZLUDA compiles offline for D4R_GPU_ARCH with its d4r_emit tool, so any target builds on any machine.
 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}"
+FP8="${D4R_TEX_FP8:-0}"
 LLVM="$R/lib/llvm/bin"
 B="$(mktemp -d)"
 trap 'rm -rf "$B"' EXIT
+# D4R_TEX_CFLAGS: extra flags for SRC.hip (tests, e.g. -DD4R_WMMA_LAYOUT=12 for the gfx12-layout shim on gfx11)
+read -r -a FLAGS <<< "${D4R_TEX_CFLAGS:-}"
+[[ "$FP8" == 1 ]] && FLAGS+=(-DD4R_TEX_FP8)
 # 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"
+    --offload-arch="$ARCH" -fgpu-rdc -emit-llvm -c -Xclang -fdenormal-fp-math=dynamic -DD4R_KERNEL_NAME="$K" \
+    "${FLAGS[@]}" -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
+# the code-generation switches of the runtime (d4r.ini [Kernels]), which ZLUDA's module cache keys on
+env LD_LIBRARY_PATH="$R/lib" D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 \
+    D4R_ZLUDA_WMMA_FP8_NATIVE="$FP8" D4R_ZLUDA_IGNORE_DENORMAL=1 D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 \
+    D4R_ZLUDA_WAVE64="${D4R_ZLUDA_WAVE64:-0}" "$D4R_ZLUDA_EMIT" "$B/$K.ptx" "$B/dump" "$ARCH" >/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' ')"
+printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size|wavefront_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"

kernels/tex/dec0_head.hip

@@ -6,8 +6,14 @@
 
 // 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
+#if D4R_WMMA_LAYOUT == 12
+__attribute__((device)) kslot g_dec0_w16[8 * 2 * 2 * 16];
+#else
 __attribute__((device)) u8v g_dec0_w16[8 * 2 * 2 * 16];
+#endif
 DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
+// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
+DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
 
 struct Dec0Params
 {
@@ -27,7 +33,11 @@ KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4
 #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));
+#if D4R_WMMA_LAYOUT == 12
+    g_dec0_w16[idx] = kslot_from16(d);
+#else
     g_dec0_w16[idx] = operand16(d);
+#endif
 }
 
 __attribute__((device)) static inline uint32_t mirror_clamp(int32_t v, uint32_t full, uint32_t half)
@@ -47,6 +57,58 @@ DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx,
     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;
+#if D4R_WMMA_LAYOUT == 12
+    kop X[2][2];
+#pragma unroll
+    for (int c = 0; c < 2; ++c)
+#pragma unroll
+        for (int h = 0; h < 2; ++h)
+            X[c][h] = kop_from8(*(const uint2_t*)(in + plane * c + 4 * h + 2 * hf));
+
+    // E[p]: e4m3 codes of channels 16p + 8hf .. 16p + 8hf + 7
+    uint32_t E[8][2];
+#pragma unroll
+    for (int p = 0; p < 8; ++p)
+    {
+        f8v cf;
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+            cf[i] = (float)*(const half_t*)(w + 8192 + 2 * (16 * p + i + 8 * hf));
+        half_t acc[8];
+#pragma unroll
+        for (int c = 0; c < 2; ++c)
+        {
+            const int slot = ((p * 2 + c) * 2) * 16 + m;
+            cf = k32_e4m3(cf, kop_image(g_dec0_w16, slot), X[c][0], kop_image(g_dec0_w16, slot + 16), X[c][1]);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[i] = (half_t)cf[i];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                cf[i] = (float)acc[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 (layout as below): this lane holds tiles 4G + hf and 4G + 2 + hf (channels 16p + 8hf ..), so it
+    // writes H = hf itself
+    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* p0 = E[2 * G];
+        const uint32_t* p1 = E[2 * G + 1];
+        u4v words = {__builtin_amdgcn_perm(p1[0], p0[0], 0x05040100u), __builtin_amdgcn_perm(p1[0], p0[0], 0x07060302u),
+                     __builtin_amdgcn_perm(p1[1], p0[1], 0x05040100u), __builtin_amdgcn_perm(p1[1], p0[1], 0x07060302u)};
+        *(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
+    }
+#else
     u8v X[2][2];
 #pragma unroll
     for (int c = 0; c < 2; ++c)
@@ -109,4 +171,5 @@ DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx,
                      __builtin_amdgcn_perm(b23, a23, 0x05040100u), __builtin_amdgcn_perm(b23, a23, 0x07060302u)};
         *(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
     }
+#endif
 }

kernels/tex/enc0_tail.hip

@@ -1,4 +1,4 @@
-// Native tail of rrlite_enc0_4x4_mvhi_hdr_folded: the 2x2-pixel patch embedding GEMM (16 tokens x
+// Native tail of rrlite_enc0_4x4_*_folded (all flag combinations): 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
@@ -7,8 +7,14 @@
 
 // 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.
+#if D4R_WMMA_LAYOUT == 12
+__attribute__((device)) kslot g_enc0_w16[4 * 4 * 2 * 16];
+#else
 __attribute__((device)) u8v g_enc0_w16[4 * 4 * 2 * 16];
+#endif
 DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
+// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
+DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
 
 struct Enc0Params
 {
@@ -16,7 +22,13 @@ struct Enc0Params
     uint8_t rest[352];
 };
 
-KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4x4_mvhi_hdr_folded_prep(Enc0Params prm)
+// every flag combination (mvhi/mvlo, hdr/ldr) has this tail; build_tex.sh names the prep after the kernel
+#ifndef D4R_KERNEL_NAME
+#define D4R_KERNEL_NAME rrlite_enc0_4x4_mvhi_hdr_folded
+#endif
+#define D4R_CAT2(a, b) a##b
+#define D4R_CAT(a, b) D4R_CAT2(a, b)
+KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void D4R_CAT(D4R_KERNEL_NAME, _prep)(Enc0Params prm)
 {
     const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
     if (idx >= 512)
@@ -29,7 +41,11 @@ KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4
 #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));
+#if D4R_WMMA_LAYOUT == 12
+    g_enc0_w16[idx] = kslot_from16(d);
+#else
     g_enc0_w16[idx] = operand16(d);
+#endif
 }
 
 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)
@@ -38,6 +54,69 @@ DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx,
     const uint8_t* w = (const uint8_t*)wp;
     lds_u8* s = (lds_u8*)(uintptr_t)smem;
 
+#if D4R_WMMA_LAYOUT == 12
+    // activations: token m's 128 e4m3 features (addresses as below); this half's 8 bytes of each k16 part
+    const uint32_t rowoff = 32 * (m & 3) + 256 * (m >> 2);
+    const uint32_t choff[4] = {0, 16, 128, 144};
+    kop X[4][2];
+#pragma unroll
+    for (int c = 0; c < 4; ++c)
+#pragma unroll
+        for (int h = 0; h < 2; ++h)
+            X[c][h] = kop_from8(*(const __attribute__((address_space(3))) uint2_t*)(s + 1024 * h + rowoff + choff[c] + 8 * hf));
+
+    const hv2 lo = {(half_t)-6.2831855f, (half_t)-6.2831855f}, hi = {(half_t)6.2831855f, (half_t)6.2831855f};
+    // E[p]: e4m3 codes of channels 16p + 8hf .. 16p + 8hf + 7 (D^T element i = channel 16p + i + 8hf)
+    uint32_t E[4][2];
+#pragma unroll
+    for (int p = 0; p < 4; ++p)
+    {
+        f8v cf;
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+            cf[i] = (float)*(const half_t*)(w + 12288 + 2 * (16 * p + i + 8 * hf));
+        half_t acc[8];
+#pragma unroll
+        for (int c = 0; c < 4; ++c)
+        {
+            const int slot = ((p * 4 + c) * 2) * 16 + m;
+            cf = k32_e4m3(cf, kop_image(g_enc0_w16, slot), X[c][0], kop_image(g_enc0_w16, slot + 16), X[c][1]);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[i] = (half_t)cf[i];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                cf[i] = (float)acc[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));
+        }
+        E[p][0] = codes[0] | (codes[1] << 16);
+        E[p][1] = codes[2] | (codes[3] << 16);
+    }
+
+    // plane pl, word 4 j + q = columns 2q, 2q + 1 of tile (2 pl, j) then of tile (2 pl + 1, j) (tile (p, j) =
+    // channels 16p + 8j ..): this lane holds tiles j = hf, i.e. words 4 hf .. 4 hf + 3 of both planes
+    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)
+    {
+#pragma unroll
+        for (int pl = 0; pl < 2; ++pl)
+        {
+            const uint32_t* p0 = E[2 * pl];
+            const uint32_t* p1 = E[2 * pl + 1];
+            const u4v words = {__builtin_amdgcn_perm(p1[0], p0[0], 0x05040100u), __builtin_amdgcn_perm(p1[0], p0[0], 0x07060302u),
+                               __builtin_amdgcn_perm(p1[1], p0[1], 0x05040100u), __builtin_amdgcn_perm(p1[1], p0[1], 0x07060302u)};
+            *(u4v*)((uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * pl + 4 * hf) = words;
+        }
+    }
+#else
     // 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);
@@ -121,4 +200,5 @@ DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx,
         uint32_t* out = (uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * hf;
         *(u8v*)out = words;
     }
+#endif
 }

kernels/tex/make_ptx.py

@@ -11,6 +11,28 @@ 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"))
 
+ENC0_TAIL = """{{
+	.param .b32 q0;
+	st.param.b32 [q0+0], {smem};
+	.param .b64 q1;
+	st.param.b64 [q1+0], {w};
+	.param .b64 q2;
+	st.param.b64 [q2+0], {out};
+	.param .b32 q3;
+	st.param.b32 [q3+0], {dimx};
+	.param .b32 q4;
+	st.param.b32 [q4+0], {dimy};
+	.param .b32 q5;
+	st.param.b32 [q5+0], {ox};
+	.param .b32 q6;
+	st.param.b32 [q6+0], {oy};
+	call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
+}}
+ret;
+
+}}
+"""
+
 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(
@@ -25,27 +47,7 @@ KERNELS = {
 )
 ;
 """,
-        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;
-
-}
-""",
+        tail=ENC0_TAIL.format(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233"),
     ),
     # 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.
@@ -206,12 +208,83 @@ def replace_rz_round(lines):
     return out, n
 
 
+ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
+# the other flag combinations of the kernels whose only native part is the surface-store rewrite
+SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
+                          r"rrlite_post_3_[12]_mv(hi|lo)_(hdr|ldr)_folded|rrlite_downsample_kernel_(static|dynamic)_(hdr|ldr))$")
+ENC0_RE = re.compile(r"^rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)_folded$")
+
+
+def read_lines(path):
+    return path.read_text().replace("\r\n", "\n").split("\n")
+
+
+def enc0_registers(lines):
+    """Cut line (1-based) and the tail call's registers of an enc0 variant. The code after the third
+    bar.warp.sync is the same GEMM in every flag combination, only register numbers differ: the shared
+    base and weights pointer are read off by aligning that block with the reference variant; the output
+    pointer and grid size are the kernel's loads of parameters +40 and +8; the block offsets are the two
+    values halved (shr 31 / add) by the epilogue."""
+    syncs = [i for i, l in enumerate(lines) if l.strip() == "bar.warp.sync -1;"]
+    if len(syncs) != 3:
+        sys.exit(f"enc0: expected 3 bar.warp.sync, found {len(syncs)}")
+    cut = syncs[2] + 1
+    before, after = lines[:cut], lines[cut:]
+    out = [m.group(1) for m in (re.match(r"^ld\.param\.u64 (%rd\d+), \[%rd\d+\+40\];$", l.strip()) for l in before) if m]
+    dims = [m.groups() for m in (re.match(r"^ld\.param\.v2\.u32 \{(%r\d+), (%r\d+)\}, \[%rd\d+\+8\];$", l.strip()) for l in before) if m]
+    halves = []
+    for a, b in zip(after, after[1:]):
+        m = re.match(r"^shr\.u32 (%r\d+), (%r\d+), 31;$", a.strip())
+        if m and re.match(rf"^add\.s32 %r\d+, {re.escape(m.group(2))}, {re.escape(m.group(1))};$", b.strip()):
+            halves.append(m.group(2))
+    if len(out) != 1 or len(dims) != 1 or len(halves) < 2:
+        sys.exit(f"enc0: parameter loads or epilogue not found ({out}, {dims}, {halves[:2]})")
+    return cut, dict(out=out[0], dimx=dims[0][0], dimy=dims[0][1], oy=halves[0], ox=halves[1])
+
+
+REG_RE = re.compile(r"%(?:rd|rs|r|fd|f|p)\d+")
+
+
+def enc0_variant(name):
+    ref = read_lines(module_file(ENC0_REF, KERNELS[ENC0_REF]))
+    path = module_file(name, {"file": ""})
+    lines = read_lines(path)
+    rcut, rregs = enc0_registers(ref)
+    k = KERNELS[ENC0_REF]
+    known = dict(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233")
+    if rcut != k["cut_after"] or any(rregs[key] != known[key] for key in rregs):
+        sys.exit(f"enc0: the register search does not reproduce {ENC0_REF} ({rcut}, {rregs})")
+    cut, regs = enc0_registers(lines)
+    # align the GEMM block (up to the reference's first closing brace) and map registers
+    n = next(i for i, l in enumerate(ref[rcut:]) if l.startswith("}"))
+    mapping = {}
+    for a, b in zip(ref[rcut:rcut + n], lines[cut:cut + n]):
+        if REG_RE.sub("R", a) != REG_RE.sub("R", b):
+            sys.exit(f"enc0: {name} differs from {ENC0_REF} after the cut: {a!r} / {b!r}")
+        for x, y in zip(REG_RE.findall(a), REG_RE.findall(b)):
+            if mapping.setdefault(x, y) != y:
+                sys.exit(f"enc0: inconsistent register mapping {x} -> {mapping[x]} / {y}")
+    regs.update(smem=mapping[known["smem"]], w=mapping[known["w"]])
+    return dict(file=path.name, cut_after=cut, cut_check="bar.warp.sync -1;", sust=True, extern=k["extern"],
+                tail=ENC0_TAIL.format(**regs))
+
+
+def kernel_spec(name):
+    if name in KERNELS:
+        return KERNELS[name]
+    if ENC0_RE.match(name):
+        return enc0_variant(name)
+    if SUST_ONLY_RE.match(name):
+        return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_output"))
+    sys.exit(f"make_ptx: no recipe for {name}")
+
+
 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()):
+    if k["file"] and known.is_file() and entry.search(known.read_text()):
         return known
     for path in sorted(PTX_DIR.glob("*.ptx")):
         if entry.search(path.read_text()):
@@ -221,8 +294,8 @@ def module_file(name, k):
 
 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")
+    k = kernel_spec(name)
+    lines = read_lines(module_file(name, k))
     if "cut_after" in k:
         cut = k["cut_after"]
         if lines[cut - 1].strip() != k["cut_check"]:

kernels/tex/tex_common.h

@@ -3,6 +3,10 @@
 // only plain clang builtins are used (compiled with -nogpuinc -nogpulib).
 #pragma once
 #include <stdint.h>
+#include "../common/wmma_layout.h"
+#if defined(D4R_TEX_FP8) && D4R_WMMA_LAYOUT != 12 && defined(__HIP_DEVICE_COMPILE__)
+#error "D4R_TEX_FP8 needs the gfx12 WMMA layout"
+#endif
 
 #pragma clang fp contract(off)
 
@@ -65,6 +69,75 @@ __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);
 }
 
+// ---------------------------------------------------------------- gfx12 m16n8k32 e4m3 steps
+// On gfx12 the GEMM tails do a k32 step as ZLUDA's gfx12 lowering of NVIDIA's m16n8k32 e4m3 MMA does, so they
+// stay bit-identical to ZLUDA's compile of the kernel they replace: a lane's k16 operand holds this half's
+// e4m3 k 8h .. 8h + 7 (ZLUDA's gfx12 slot order), either widened as above with C * 2^-16 inside the f32
+// WMMAs (d4r_wmma12_*), or with D4R_TEX_FP8 (ZLUDA's D4R_ZLUDA_WMMA_FP8_NATIVE) as the bytes of a native FP8
+// WMMA with no scaling (d4r_wmma12f8_*). (The gfx11 tails use operand16 and wmma above.)
+#if D4R_WMMA_LAYOUT == 12
+// 8 e4m3 bytes (k in order) widened to value * 2^-8 f16 pairs, k in order
+__attribute__((device)) static inline u4v widen8(uint2_t d)
+{
+    uint32_t a0, a1, b0, b1;
+    e4m3x4_scaled(d.x, a0, a1);
+    e4m3x4_scaled(d.y, b0, b1);
+    return (u4v){a0, a1, b0, b1};
+}
+#ifdef D4R_TEX_FP8
+typedef uint2_t kop;
+typedef u4v kslot; // weight image slot: the 16 e4m3 bytes of one k16 operand row
+__attribute__((device)) static inline kop kop_from8(uint2_t d)
+{
+    return d;
+}
+__attribute__((device)) static inline kslot kslot_from16(u4v d)
+{
+    return d;
+}
+#else
+typedef u4v kop;
+typedef u8v kslot; // weight image slot: the 16 widened values of one k16 operand row, k in order
+__attribute__((device)) static inline kop kop_from8(uint2_t d)
+{
+    return widen8(d);
+}
+__attribute__((device)) static inline kslot kslot_from16(u4v d)
+{
+    const u4v lo = widen8((uint2_t){d[0], d[1]}), hi = widen8((uint2_t){d[2], d[3]});
+    return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
+}
+#endif
+// this half's k16 operand from 16 consecutive e4m3 bytes of NVIDIA's k order
+__attribute__((device)) static inline kop kop_from16(u4v d)
+{
+    return kop_from8(lane_id() >= 16 ? (uint2_t){d[2], d[3]} : (uint2_t){d[0], d[1]});
+}
+// this half's operand from slot idx of a weight image
+__attribute__((device)) static inline kop kop_image(const kslot* img, int idx)
+{
+    return ((const kop*)img)[2 * idx + (lane_id() >> 4)];
+}
+// f32 C (true scale) + one k32 step
+__attribute__((device)) static inline f8v k32_e4m3(f8v c, const kop& a0, const kop& b0, const kop& a1, const kop& b1)
+{
+#if defined(D4R_TEX_FP8) && defined(__GFX12__)
+    typedef int i2v __attribute__((ext_vector_type(2)));
+    c = __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(i2v, a0), __builtin_bit_cast(i2v, b0), c);
+    return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(i2v, a1), __builtin_bit_cast(i2v, b1), c);
+#elif defined(D4R_TEX_FP8)
+    // layout shim, like ZLUDA's: each FP8 WMMA as the scaled widening (exact powers of two)
+    c = wm_mma_zluda(widen8(a0), widen8(b0), c * 0x1p-16f) * 0x1p16f;
+    return wm_mma_zluda(widen8(a1), widen8(b1), c * 0x1p-16f) * 0x1p16f;
+#else
+    c = c * 0x1p-16f;
+    c = wm_mma_zluda(a0, b0, c);
+    c = wm_mma_zluda(a1, b1, c);
+    return c * 0x1p16f;
+#endif
+}
+#endif
+
 // 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)
 {

kernels/tools/dump_runner.cpp

@@ -44,7 +44,18 @@ int main(int argc, char** argv)
     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;
+    // the prep grid: PREP_GRID, else the module's d4r_prep_blocks (what the ZLUDA hook launches), else 88
+    unsigned prepGrid = 88;
+    {
+        hipDeviceptr_t pb;
+        size_t pbs = 0;
+        uint32_t v = 0;
+        if (hipModuleGetGlobal(&pb, &pbs, mod, "d4r_prep_blocks") == hipSuccess && pbs == 4 &&
+            hipMemcpyDtoH(&v, pb, 4) == hipSuccess && v != 0)
+            prepGrid = v;
+    }
+    if (getenv("PREP_GRID"))
+        prepGrid = atoi(getenv("PREP_GRID"));
 
     std::vector<char> args = read_file(dump + "/args.bin");
     std::ifstream man(dump + "/manifest.txt");
@@ -58,6 +69,13 @@ int main(int argc, char** argv)
         std::istringstream s(line);
         std::string kind;
         s >> kind;
+        if (kind == "texture" || kind == "surface")
+        {
+            // texture/surface object handles in the arguments would point at nothing here: the kernel faults
+            fprintf(stderr, "%s: the dump uses texture or surface objects, which dump_runner cannot recreate\n",
+                    dump.c_str());
+            return 3;
+        }
         if (kind == "launch")
             s >> grid[0] >> grid[1] >> grid[2] >> block[0] >> block[1] >> block[2] >> shared;
         else if (kind == "alloc")

kernels/tools/model_enc3.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""Numeric primitives and replay-dump access for swin_model.py (the rrlite Swin blocks of DLSS 4.5, preset M).
+
+Everything here is NVIDIA's arithmetic as the native kernels reproduce it (kernels/m/swin_common.h,
+swin_block.h): f16 rounding after every operation, e4m3 (OCP, satfinite, RNE) re-quantisation, one f16 rounding
+per k32 MMA step, and the softmax's bit-level exponent trick.
+"""
+import struct
+
+import numpy as np
+
+F16, F32, F64 = np.float16, np.float32, np.float64
+
+
+def f16(x):
+    """round to f16; the result stays float64 (on the f16 grid), so products and sums of such values are exact
+    until the next f16() (numpy would otherwise do f16 arithmetic, e.g. in matmul)"""
+    return np.asarray(x, F64).astype(F16).astype(F64)
+
+
+def _e4m3_table():
+    t = np.zeros(256, F64)
+    for c in range(256):
+        s, e, m = c >> 7, (c >> 3) & 15, c & 7
+        if e == 15 and m == 7:
+            v = np.nan
+        elif e == 0:
+            v = m * 2.0 ** -9
+        else:
+            v = (1 + m / 8) * 2.0 ** (e - 7)
+        t[c] = -v if s else v
+    return t
+
+
+E4M3 = _e4m3_table()
+_POS = E4M3[:127]  # codes 0..126: the finite non-negative values, ascending
+
+
+def q8(x):
+    """e4m3 RNE satfinite round trip of f16 values (NaN stays NaN), as f16"""
+    h = f16(x).astype(F64)
+    a = np.abs(h)
+    i = np.clip(np.searchsorted(_POS, a), 1, 126)
+    lo, hi = _POS[i - 1], _POS[i]
+    # nearest; ties to the even code
+    pick_hi = (hi - a < a - lo) | ((hi - a == a - lo) & ((i & 1) == 0))
+    v = np.where(pick_hi, hi, lo)
+    v = np.where(a >= _POS[126], _POS[126], v)
+    v = np.where(a <= 0, 0.0, v)
+    v = np.where(np.isnan(h), np.nan, np.copysign(v, h))
+    return f16(v)
+
+
+# softmax constants (swin_block.h stage 2)
+C_SCALE = f16(0.0972222164273262 / 5.656854152679443)
+CLAMP = F16(0.55615234375)
+P1 = f16(0.92730712890625)
+P0 = F16(1.375)
+
+
+def exp_trick(poly):
+    """per row, columns (2j, 2j+1): the 32-bit word even | odd << 16 of the f16 bits, << 5, + 0x7FF88000; the even
+    column takes its low half, the odd column its high half (as f16 bits)"""
+    b = poly.astype(F16).view(np.uint16).astype(np.uint64)
+    out = np.empty_like(b)
+    even, odd = b[:, 0::2], b[:, 1::2]
+    word = ((even | (odd << np.uint64(16))) << np.uint64(5)) + np.uint64(0x7FF88000)
+    word &= np.uint64(0xFFFFFFFF)
+    out[:, 0::2] = word & np.uint64(0xFFFF)
+    out[:, 1::2] = word >> np.uint64(16)
+    return out.astype(np.uint16).view(F16)
+
+
+def row_sum16(w):
+    """f16 butterfly over the 16 columns: partner c ^ 8, then ^ 2, ^ 4, ^ 1 (every column ends with the same sum)"""
+    s = w.astype(F16)
+    c = np.arange(16)
+    for m in (8, 2, 4, 1):
+        s = f16(s.astype(F64) + s[:, c ^ m].astype(F64))
+    return s
+
+
+def gelu_poly(z):
+    z = f16(z)
+    cz = np.clip(z, F16(-2.0), F16(2.0))
+    inner = f16(f16(0.41216981).astype(F64) - f16(f16(0.08108133).astype(F64) * np.abs(cz).astype(F64)).astype(F64))
+    return f16(z.astype(F64) * f16(0.5 + f16(cz.astype(F64) * inner.astype(F64)).astype(F64)).astype(F64))
+
+
+def mma_chain(A, W, C0=None):
+    """chained k32 MMA steps with f16 accumulation: D = f16(A[:, 32k:32k+32] W[32k:32k+32] + D), exact sums"""
+    M, K = A.shape
+    D = np.zeros((M, W.shape[1]), F64) if C0 is None else f16(C0)
+    A = np.asarray(A, F64)
+    W = np.asarray(W, F64)
+    for k in range(0, K, 32):
+        D = f16(A[:, k:k + 32] @ W[k:k + 32] + D.astype(F64))
+    return D
+
+
+class Params:
+    """one rrlite Swin-block launch from a replay dump (CommonParams: w 0, sx sy 8, tw th 16, in 24, p32 32,
+    out 40, p48 48)"""
+
+    def __init__(self, d):
+        man = open(d + '/manifest.txt').read().split('\n')
+        self.kernel = man[0].split()[1]
+        la = man[1].split()
+        self.grid = tuple(map(int, la[1:4]))
+        self.args = open(d + '/args.bin', 'rb').read()
+        self.allocs = {}
+        for line in man:
+            p = line.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))
+        self.w_ptr, = struct.unpack_from('<Q', self.args, 0)
+        self.sx, self.sy, self.W, self.H = struct.unpack_from('<iiii', self.args, 8)
+        self.inp, self.p32, self.out, self.p48 = struct.unpack_from('<QQQQ', self.args, 24)
+        wbase, wbuf = self.alloc_of(self.w_ptr)
+        self.wbuf = wbuf[self.w_ptr - wbase:]
+        self.abase, self.arena = self.alloc_of(self.inp)
+
+    def alloc_of(self, a):
+        for base, buf in self.allocs.values():
+            if base <= a < base + buf.size:
+                return base, buf
+        raise ValueError(hex(a))
+
+    def f16vec(self, off, n):
+        return self.wbuf[off:off + 2 * n].view(F16).copy()
+
+    def token_xy(self, bx, by):
+        T = np.arange(64)
+        tx = (T & 3) + 4 * ((T >> 4) & 1)
+        ty = ((T >> 2) & 3) + 4 * (T >> 5)
+        return 8 * bx - self.sx + tx, 8 * by - self.sy + ty
+
+    @staticmethod
+    def mirror(v, n):
+        v = np.abs(v)
+        return np.minimum(v, 2 * n - 2 - v)

kernels/tools/pwin_check.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""Compares native preset-K layer outputs with the numpy model (pwin_model.py) on sampled windows of a replay dump.
+
+usage: pwin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--windows N]
+  DUMP_DIR  replay dump of one dltss_pwin_* launch (inputs)
+  OUT_DIR   dump_runner output of a native kernel on that dump (alloc-N.bin after the launch)
+For every sampled window (the four corners, border windows and random interior ones) the model output and the
+kernel's stored values are compared per f16 element: exact matches, and the error in f16 ulps at the model
+value's magnitude (values below 2^-10 count with the ulp of 2^-10, so near-zero sign flips do not dominate).
+The kernels round as NVIDIA's f16 wmma does except in the PWIN_F32ACC layers, and gfx12 WMMA groups its f32
+sums differently from gfx11, so neither build is expected to match the model exactly.
+"""
+import os
+import random
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import pwin_model as pm  # noqa: E402
+
+F16 = np.float16
+
+# layer -> (kind, H, C, COUT/CL, extra)
+LAYERS = {
+    'enc0': ('enc0', 2, 32, 64, None), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False),
+    'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 16, None),
+    'dec4': ('dec', 4, 128, 160, None), 'dec3': ('dec', 4, 96, 128, None), 'dec2': ('dec', 2, 64, 96, None),
+    'dec1': ('dec', 2, 64, 64, None), 'dec0': ('dec0', 2, 32, 64, None),
+}
+
+
+class Outputs:
+    """the dump's pointer map over the allocations a native run saved"""
+    def __init__(self, dump, out_dir):
+        self.d = dump
+        self.allocs = {i: (base, np.fromfile(f'{out_dir}/alloc-{i}.bin', np.uint8)) for i, (base, _) in dump.allocs.items()}
+
+    def rows(self, off_param, index, width):
+        a = self.d.u64(off_param)
+        for base, buf in self.allocs.values():
+            if base <= a < base + buf.size:
+                o = a - base + 2 * index * width
+                return buf[o:o + 2 * width].view(F16)
+        raise ValueError(hex(a))
+
+
+def ulps(model, got):
+    m = model.astype(np.float64)
+    g = got.astype(np.float64)
+    scale = np.maximum(np.abs(m), 2.0 ** -10)
+    ulp = 2.0 ** (np.floor(np.log2(scale)) - 10)
+    return np.abs(m - g) / ulp
+
+
+def windows(grid, n, seed=1):
+    gx, gy = grid[0], grid[1]
+    w = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
+    r = random.Random(seed)
+    while len(w) < min(n, gx * gy):
+        w.add((r.randrange(gx), r.randrange(gy)))
+    return sorted(w)
+
+
+def model_window(P, layer, bx, by):
+    kind, H, C, X, extra = LAYERS[layer]
+    if kind == 'enc0':
+        return pm.pwin_enc0(P, bx, by)
+    if kind == 'enc':
+        return pm.pwin_encoder(P, bx, by, H, C, X, posattn=extra)
+    if kind == 'plain':
+        return pm.swin_core(P, pm.Layout(H, C, C), pm.load_window(P, bx, by, C)), None
+    if kind == 'dec':
+        return pm.pwin_decoder(P, bx, by, H, C, X), None
+    return pm.pwin_dec0(P, bx, by)
+
+
+def stored(P, O, layer, bx, by, y, extra_out):
+    """(model values, kernel values) of everything the kernel stores for window (bx, by)"""
+    kind, H, C, X, _ = LAYERS[layer]
+    W, Hh = P.i32x2(0)
+    sx, sy = P.i32x2(56)
+    ms, gs = [], []
+    full_off = 32 if kind in ('enc0', 'enc') else 24
+    for t in range(64):
+        Y0, X0 = 8 * by - sy + (t >> 3), 8 * bx - sx + (t & 7)
+        if not (0 <= Y0 < Hh and 0 <= X0 < W):
+            continue
+        if kind == 'dec0':
+            ms.append(extra_out[t]); gs.append(O.rows(24, Y0 * W + X0, 40))
+        else:
+            ms.append(y[t]); gs.append(O.rows(full_off, Y0 * W + X0, C))
+    if kind in ('enc0', 'enc'):
+        cout = extra_out.shape[1]
+        for r in range(16):
+            Ym, Xm = int((8 * by - sy) / 2) + (r >> 2), int((8 * bx - sx) / 2) + (r & 3)
+            if 0 <= Ym < Hh // 2 and 0 <= Xm < W // 2:
+                ms.append(extra_out[r]); gs.append(O.rows(24, Ym * (W // 2) + Xm, cout))
+    return np.concatenate(ms), np.concatenate(gs)
+
+
+def main():
+    args = [a for a in sys.argv[1:] if not a.startswith('--')]
+    n = 12
+    if '--windows' in sys.argv:
+        n = int(sys.argv[sys.argv.index('--windows') + 1])
+        args.remove(str(n))
+    dump_dir, outs = args[0], args[1:]
+    P = pm.Dump(dump_dir)
+    layer = P.kernel.replace('dltss_pwin_', '').replace('_layer', '')
+    wins = windows(P.grid, n)
+    models = {w: model_window(P, layer, *w) for w in wins}
+    for out in outs:
+        O = Outputs(P, out)
+        m, g = zip(*(stored(P, O, layer, bx, by, *models[(bx, by)]) for bx, by in wins))
+        m, g = np.concatenate(m), np.concatenate(g)
+        u = ulps(m, g)
+        exact = (m.view(np.uint16) == g.view(np.uint16)).mean()
+        md, gd = m.astype(np.float64), g.astype(np.float64)
+        psnr = 10 * np.log10(np.abs(md).max() ** 2 / max(((md - gd) ** 2).mean(), 1e-30))
+        print(f'{layer:5s} {os.path.basename(os.path.normpath(out)):10s} {len(wins)} windows {m.size} values: exact {100 * exact:5.1f}%  '
+              f'<=1ulp {100 * (u <= 1).mean():5.1f}%  <=4ulp {100 * (u <= 4).mean():6.2f}%  max {u.max():7.1f} ulp  '
+              f'mean {u.mean():.3f} ulp  psnr {psnr:.1f} dB  finite {np.isfinite(g.astype(np.float32)).mean() * 100:.0f}%')
+
+
+if __name__ == '__main__':
+    main()

kernels/tools/swin_check.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Compares native preset-M Swin-block outputs with the numpy model (swin_model.py) on sampled blocks of a dump.
+
+usage: swin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--blocks N]
+  DUMP_DIR  replay dump of one rrlite_{enc1,enc2,enc3_tube,dec1,dec2}_4x4 launch
+  OUT_DIR   dump_runner output of a native kernel on that dump
+The stored e4m3 outputs (block output, and the patch merge of encoders that have one) are compared code by code:
+the share of identical codes, of codes one e4m3 step apart, and the PSNR of the decoded values. The model is
+NVIDIA's arithmetic (e4m3 activations, f16 accumulation per k32 step); the production kernels keep activations
+f16 and accumulate in f32 (SWIN_NO_Q8, SWIN_F32ACC), so they are not expected to match it exactly.
+"""
+import os
+import random
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+import model_enc3 as me  # noqa: E402
+import swin_model as sm  # noqa: E402
+
+# kernel -> (C, heads, NPM, CIN)
+LAYERS = {'enc1': (64, 2, 96, 0), 'enc2': (96, 4, 128, 0), 'enc3_tube': (128, 4, 0, 0), 'dec2': (96, 4, 0, 128),
+          'dec1': (64, 2, 0, 96)}
+
+
+def blocks(grid, n, seed=1):
+    gx, gy = grid[0], grid[1]
+    b = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
+    r = random.Random(seed)
+    while len(b) < min(n, gx * gy):
+        b.add((r.randrange(gx), r.randrange(gy)))
+    return sorted(b)
+
+
+def model_block(P, layer, bx, by):
+    C, heads, NPM, CIN = LAYERS[layer]
+    pre = 4 * C * CIN + 8 * C if CIN else 0
+    x0 = sm.decoder_input(P, bx, by, C, CIN, heads, P.p32) if CIN else sm.load_tokens(P, bx, by, C // 32)
+    x2, out, L = sm.swin_block(P, x0, C, heads, pre)
+    merged = None
+    if NPM:
+        A = sm.patch_merge_input(out, C)
+        PM0 = L['end']
+        cols = []
+        for g in range(NPM // 32):
+            Wg = me.E4M3[P.wbuf[sm.woff_table(PM0 + g * 128 * C, 512, 64 * C, 4 * C, 32)]]
+            bias = P.f16vec(PM0 + 4 * C * NPM + 64 * g, 32)
+            cols.append(me.mma_chain(A, Wg, C0=np.broadcast_to(bias, (16, 32)).copy()))
+        merged = me.q8(sm.pair_to_nat(np.concatenate(cols, axis=1)))
+    return out, merged
+
+
+def codes_at(buf, base, ptr, W, H, X, Y, C):
+    """e4m3 codes of tokens (X, Y) (in bounds) from planes of 32 channels at ptr"""
+    c = np.arange(C)
+    off = ptr - base + (((c[None, :] >> 5) * H + Y[:, None]) * W + X[:, None]) * 32 + (c[None, :] & 31)
+    return buf[off]
+
+
+def main():
+    args = [a for a in sys.argv[1:] if not a.startswith('--')]
+    n = 8
+    if '--blocks' in sys.argv:
+        n = int(sys.argv[sys.argv.index('--blocks') + 1])
+        args.remove(str(n))
+    dump, outs = args[0], args[1:]
+    P = me.Params(dump)
+    layer = P.kernel.replace('rrlite_', '').replace('_4x4', '')
+    C, heads, NPM, CIN = LAYERS[layer]
+    bl = blocks(P.grid, n)
+    models = {b: model_block(P, layer, *b) for b in bl}
+    for out_dir in outs:
+        got, want = [], []
+        for (bx, by) in bl:
+            base, _ = P.alloc_of(P.out)
+            aid = next(i for i, (b0, buf) in P.allocs.items() if b0 == base)
+            buf = np.fromfile(f'{out_dir}/alloc-{aid}.bin', np.uint8)
+            X, Y = P.token_xy(bx, by)
+            ok = (X >= 0) & (X < P.W) & (Y >= 0) & (Y < P.H)
+            out, merged = models[(bx, by)]
+            got.append(codes_at(buf, base, P.out, P.W, P.H, X[ok], Y[ok], C).ravel())
+            want.append(out[ok].ravel())
+            if NPM:
+                W2, H2 = P.W // 2, P.H // 2
+                r = np.arange(16)
+                MX, MY = (8 * bx - P.sx) // 2 + (r & 3), (8 * by - P.sy) // 2 + (r >> 2)
+                okm = (MX >= 0) & (MX < W2) & (MY >= 0) & (MY < H2)
+                mbase, _ = P.alloc_of(P.p48)
+                mid = next(i for i, (b0, _) in P.allocs.items() if b0 == mbase)
+                mbuf = buf if mid == aid else np.fromfile(f'{out_dir}/alloc-{mid}.bin', np.uint8)
+                got.append(codes_at(mbuf, mbase, P.p48, W2, H2, MX[okm], MY[okm], NPM).ravel())
+                want.append(merged[okm].ravel())
+        g = me.E4M3[np.concatenate(got)]
+        w = np.concatenate(want).astype(np.float64)
+        rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v)  # noqa: E731
+        steps = np.abs(rank(g) - rank(w))
+        psnr = 10 * np.log10(np.nanmax(np.abs(w)) ** 2 / max(np.nanmean((g - w) ** 2), 1e-30))
+        print(f'{layer:9s} {os.path.basename(os.path.normpath(out_dir)):10s} {len(bl)} blocks {g.size} codes: '
+              f'identical {100 * (g == w).mean():5.1f}%  <=1 step {100 * (steps <= 1).mean():6.2f}%  '
+              f'max {int(steps.max())} steps  psnr {psnr:.1f} dB')
+
+
+if __name__ == '__main__':
+    main()

kernels/tools/tail_check.hip

@@ -0,0 +1,141 @@
+// Standalone check of the texture-kernel GEMM tails (kernels/tex/enc0_tail.hip, dec0_head.hip) outside the
+// ZLUDA-compiled kernels they are linked into: random e4m3 weights and activations, several tile positions
+// (including partial tiles at the image border), all outputs written to one file. Builds of different WMMA
+// layouts (gfx11, the gfx12-layout shim, gfx12, gfx12 native FP8) are then compared byte for byte (same
+// arithmetic) or by value (gfx12 rounding). The tails touch no images, so this also runs in rocjitsu.
+//
+// build: clang++ -x hip --offload-arch=ARCH -O3 [-DD4R_WMMA_LAYOUT=12] [-DD4R_TEX_FP8] -I kernels/tex
+//        kernels/tools/tail_check.hip -o tail_check           usage: tail_check OUT.bin
+#include <hip/hip_runtime.h>
+
+#include <cstdio>
+#include <cstdlib>
+#include <vector>
+
+// the surface-store helpers of tex_common.h reference ZLUDA's helper library, which the tails do not use
+extern "C" __device__ void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t, int __attribute__((ext_vector_type(2))),
+                                                                     short __attribute__((ext_vector_type(2))))
+{
+}
+
+// both tails define the same hook globals
+#define D4R_KERNEL_NAME tail_enc0
+#define d4r_prep_blocks d4r_prep_blocks_enc0
+#define d4r_prep_key_at d4r_prep_key_at_enc0
+#include "enc0_tail.hip"
+#undef d4r_prep_blocks
+#undef d4r_prep_key_at
+#define d4r_prep_blocks d4r_prep_blocks_dec0
+#define d4r_prep_key_at d4r_prep_key_at_dec0
+#include "dec0_head.hip"
+
+extern "C" __global__ void __launch_bounds__(32) run_enc0(const uint32_t* lds_in, const uint8_t* w, uint32_t* out,
+                                                          uint32_t dimx, uint32_t dimy, int ox, int oy)
+{
+    __shared__ uint32_t s[512 + 4];
+    for (int i = threadIdx.x; i < 512; i += 32)
+        s[i] = lds_in[i];
+    __syncthreads();
+    d4r_enc0_tail((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)out, dimx,
+                  dimy, ox, oy);
+}
+
+extern "C" __global__ void __launch_bounds__(32) run_dec0(const uint8_t* w, const uint32_t* in, uint32_t* staged,
+                                                          uint32_t dimx, uint32_t dimy, int ox, int oy)
+{
+    __shared__ uint32_t s[512 + 4];
+    for (int i = threadIdx.x; i < 512; i += 32)
+        s[i] = 0xdeadbeefu;
+    __syncthreads();
+    d4r_dec0_head((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)in, dimx,
+                  dimy, ox, oy);
+    __syncthreads();
+    for (int i = threadIdx.x; i < 512; i += 32)
+        staged[i] = s[i];
+}
+
+#define CHECK(x)                                                                                                   \
+    do                                                                                                             \
+    {                                                                                                              \
+        hipError_t e_ = (x);                                                                                       \
+        if (e_ != hipSuccess)                                                                                      \
+        {                                                                                                          \
+            fprintf(stderr, "%s: %s\n", #x, hipGetErrorString(e_));                                                \
+            return 1;                                                                                              \
+        }                                                                                                          \
+    } while (0)
+
+static uint8_t e4m3(uint32_t r)
+{
+    // Keep the matrix values in a useful range so saturation cannot hide a wrong WMMA result.
+    const uint8_t magnitude = (r & 31) == 0 ? 0 : (uint8_t)(0x18 + (r & 31));
+    return (uint8_t)(magnitude | ((r >> 8) & 0x80));
+}
+static uint16_t f16_bits(float v)
+{
+    return __builtin_bit_cast(uint16_t, (_Float16)v);
+}
+
+int main(int argc, char** argv)
+{
+    if (argc < 2)
+    {
+        fprintf(stderr, "usage: %s OUT.bin\n", argv[0]);
+        return 2;
+    }
+    srand(12345);
+    auto rnd = []() { return (uint32_t)rand() * 2654435761u + (uint32_t)rand(); };
+    // weights: e4m3 fragments, f16 biases (enc0 at 12288, dec0 at 8192)
+    std::vector<uint8_t> w(16384 + 512);
+    for (auto& b : w)
+        b = e4m3(rnd());
+    for (int i = 0; i < 64; ++i)
+        *(uint16_t*)&w[12288 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
+    for (int i = 0; i < 128; ++i)
+        *(uint16_t*)&w[8192 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
+    const uint32_t dimx = 20, dimy = 14, W2 = dimx / 2, H2 = dimy / 2, rowstride = W2 * 8, plane = H2 * rowstride;
+    std::vector<uint32_t> lds(512), in(2 * plane);
+    for (auto& v : lds)
+        v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
+    for (auto& v : in)
+        v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
+
+    uint8_t *dw, *dummy;
+    uint32_t *dlds, *din, *dout, *dstaged;
+    CHECK(hipMalloc(&dw, w.size()));
+    CHECK(hipMalloc(&dlds, lds.size() * 4));
+    CHECK(hipMalloc(&din, in.size() * 4));
+    CHECK(hipMalloc(&dout, 2 * plane * 4));
+    CHECK(hipMalloc(&dstaged, 512 * 4));
+    CHECK(hipMalloc(&dummy, 512));
+    CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
+    CHECK(hipMemcpy(dlds, lds.data(), lds.size() * 4, hipMemcpyHostToDevice));
+    CHECK(hipMemcpy(din, in.data(), in.size() * 4, hipMemcpyHostToDevice));
+    // the prep kernels read the weight pointer at parameter offset 0
+    Enc0Params pe{};
+    pe.w = (uint64_t)dw;
+    Dec0Params pd{};
+    pd.w = (uint64_t)dw;
+    hipLaunchKernelGGL(tail_enc0_prep, dim3(4), dim3(128), 0, 0, pe);
+    hipLaunchKernelGGL(rrlite_dec0_4x4_folded_prep, dim3(4), dim3(128), 0, 0, pd);
+    CHECK(hipDeviceSynchronize());
+
+    FILE* f = fopen(argv[1], "wb");
+    const int offsets[][2] = {{0, 0}, {8, 4}, {12, 6}, {16, 8}, {-4, -2}, {18, 12}};
+    std::vector<uint32_t> host(2 * plane);
+    for (const auto& o : offsets)
+    {
+        CHECK(hipMemset(dout, 0x5a, 2 * plane * 4));
+        hipLaunchKernelGGL(run_enc0, dim3(1), dim3(32), 0, 0, dlds, dw, dout, dimx, dimy, o[0], o[1]);
+        CHECK(hipDeviceSynchronize());
+        CHECK(hipMemcpy(host.data(), dout, 2 * plane * 4, hipMemcpyDeviceToHost));
+        fwrite(host.data(), 4, 2 * plane, f);
+        hipLaunchKernelGGL(run_dec0, dim3(1), dim3(32), 0, 0, dw, din, dstaged, dimx, dimy, o[0], o[1]);
+        CHECK(hipDeviceSynchronize());
+        CHECK(hipMemcpy(host.data(), dstaged, 512 * 4, hipMemcpyDeviceToHost));
+        fwrite(host.data(), 4, 512, f);
+    }
+    fclose(f);
+    printf("wrote %s\n", argv[1]);
+    return 0;
+}