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

feat: add d4r inference engines and Windows parity tooling

Browse Source 
Add native K/M/RR inference engines, HIP kernels, replay and validation
harnesses, and Windows DLSS setup and runtime support. Include GPU
profiling tools, build and packaging scripts, platform integration
patches, and implementation documentation.
This commit is contained in:
ayo committed 2026-10-16 08:54:47 -04:00
1 parent 25493b3198
commit 7c5d3c9157
430 files changed
+76191
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/build.sh
+125
View File
@@ -0,0 +1,125 @@
#!/usr/bin/env bash
# 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|l|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): 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.
# D4R_PREFER_ACCURACY 1: preserve FP8 quantization, per-MMA f16 rounding and denormal handling;
# build a separate set (default output: kernels/out/native/accuracy).
# 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_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)
# D4R_DLSS_PTX_DIR optional pre-extracted PTX directory for parallel per-target builds
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WHAT="${1:-all}"
ACCURACY="${D4R_PREFER_ACCURACY:-0}"
[[ "$ACCURACY" == 0 || "$ACCURACY" == 1 ]] || { echo "D4R_PREFER_ACCURACY must be 0 or 1" >&2; exit 2; }
DEFAULT_OUT="$HERE/out/native"
[[ "$ACCURACY" == 1 ]] && DEFAULT_OUT+="/accuracy"
OUT="$(realpath -m "${2:-$DEFAULT_OUT}")"
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"
case "$WHAT" in all|k|l|m|tex) ;; *) echo "unknown kernel family: $WHAT" >&2; exit 2 ;; esac
# Never certify a directory containing older fast binaries as an accuracy set.
if [[ "$ACCURACY" == 1 ]] && compgen -G "$OUT/*.hsaco" >/dev/null &&
{ [[ ! -f "$OUT/d4r-accuracy.txt" ]] || [[ "$(cat "$OUT/d4r-accuracy.txt")" != 1 ]]; }; then
echo "accuracy kernels need an empty output directory or an existing accuracy set: $OUT" >&2
exit 2
fi
rm -f "$OUT/d4r-accuracy.txt"
# HIP source -> raw code object (the kernel name inside matches the DLSS kernel it replaces)
build_hip() {
local src="$1" extra="${2:-}" name
name="$(basename "$src" .hip)"
local tmp flags
tmp="$(mktemp -d)"
# per-kernel compiler flags from a "// d4r-build-flags: ..." line in the source
flags="$(sed -n 's|^// d4r-build-flags: *||p' "$src")"
[[ "$ACCURACY" == 1 ]] && extra+=" -DD4R_ACCURACY"
# 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' ')"
rm -rf "$tmp"
}
if [[ "$WHAT" == all || "$WHAT" == k ]]; then
echo "== DLSS 4 (preset K) transformer layers"
for src in "$HERE"/k/dltss_pwin_*.hip; do build_hip "$src"; done
fi
if [[ "$WHAT" == all || "$WHAT" == m || "$WHAT" == l ]]; then
echo "== DLSS 4.5 (presets L/M) shared Swin layers"
for src in "$HERE"/m/rrlite_*.hip; do build_hip "$src" "$([[ "$FP8" == 1 ]] && echo -DD4R_FP8_WMMA)"; done
fi
if [[ "$WHAT" == all || "$WHAT" == tex || "$WHAT" == l ]]; then
echo "== texture kernels (NVIDIA PTX with native parts)"
: "${D4R_DLSS_DLL:?set D4R_DLSS_DLL to nvngx_dlss.dll (310.7)}"
: "${D4R_ZLUDA_EMIT:?set D4R_ZLUDA_EMIT to d4r_emit from a ZLUDA build with patches/zluda}"
PTX_DIR="${D4R_DLSS_PTX_DIR:-$HERE/extracted/ptx}"
if [[ -z "${D4R_DLSS_PTX_DIR:-}" ]]; then
python3 "$HERE/tools/extract_dlss_ptx.py" "$D4R_DLSS_DLL" "$PTX_DIR"
elif [[ ! -d "$PTX_DIR" ]]; then
echo "missing extracted PTX directory: $PTX_DIR" >&2; exit 2
fi
# 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=()
if [[ "$WHAT" != l ]]; then
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
fi
# L's unfolded variants keep their complete neural arithmetic. The folded M
# GEMM heads/tails are incompatible; substitute only the surface stores.
specs+=(rrlite_dec0_4x4:sust_only)
for mv in mvhi mvlo; do
for range in hdr ldr; do
specs+=("rrlite_enc0_4x4_${mv}_${range}:sust_only")
for v in 3_1 3_2; do specs+=("rrlite_post_${v}_${mv}_${range}:sust_only:w64"); 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_PREFER_ACCURACY="$ACCURACY" \
D4R_ZLUDA_WAVE64="$([[ "$mode" == w64 && "$ACCURACY" == 0 ]] && 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
[[ "$ACCURACY" == 1 ]] && printf '1\n' > "$OUT/d4r-accuracy.txt"
echo "native kernels in $OUT"
kernels/common/hip_device_minimal.h
+43
View File
@@ -0,0 +1,43 @@
#pragma once
// Device-only Windows HIP build shim. The diagnostic/kernel code objects are
// compiled with -nogpuinc/-nogpulib; no host HIP or C++ standard library is
// 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()};
}
__device__ __forceinline__ d4r_device_dim3 d4r_block_idx() {
return {__builtin_amdgcn_workgroup_id_x(), __builtin_amdgcn_workgroup_id_y(),
__builtin_amdgcn_workgroup_id_z()};
}
__device__ __forceinline__ d4r_device_dim3 d4r_grid_dim() {
return {__builtin_amdgcn_grid_size_x() / __builtin_amdgcn_workgroup_size_x(),
__builtin_amdgcn_grid_size_y() / __builtin_amdgcn_workgroup_size_y(),
__builtin_amdgcn_grid_size_z() / __builtin_amdgcn_workgroup_size_z()};
}
#define threadIdx d4r_thread_idx()
#define blockIdx d4r_block_idx()
#define gridDim d4r_grid_dim()
__device__ __attribute__((convergent, always_inline)) inline void __syncthreads() {
// Match HIP __work_group_barrier: execution rendezvous alone does not
// publish LDS stores or prevent loads moving across the barrier.
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
}
kernels/common/wave32_exchange.h
+31
View File
@@ -0,0 +1,31 @@
#pragma once
#include <stdint.h>
// Windows gfx11 portability workaround informed by realdody/d4r 1897b4a.
// The reported permlanex16 failure duplicates the even K value into the odd
// slot. Retain the native gfx11 layout and exchange through per-wave LDS.
// Other platforms and gfx12 keep their original instruction sequence.
__attribute__((device, always_inline)) static inline uint32_t d4r_wave32_other_half(uint32_t value)
{
#if defined(D4R_WINDOWS_GFX11_LANE_EXCHANGE) && defined(__GFX11__)
__attribute__((shared)) volatile uint32_t scratch[512];
const uint32_t tid = __builtin_amdgcn_workitem_id_x() + __builtin_amdgcn_workgroup_size_x() *
(__builtin_amdgcn_workitem_id_y() + __builtin_amdgcn_workgroup_size_y() * __builtin_amdgcn_workitem_id_z());
// All native K/M launch bounds are <=512. Do not silently corrupt LDS if
// this helper is later reused by a larger workgroup.
if (tid >= 512) __builtin_trap();
scratch[tid] = value;
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
__builtin_amdgcn_wave_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");
const uint32_t result = scratch[tid ^ 16u];
// Protect scratch reuse by the next exchange in this wave. No block-wide
// rendezvous: different waves can reach GEMMs in a different order.
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
__builtin_amdgcn_wave_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");
return result;
#else
return __builtin_amdgcn_permlanex16(value, value, 0x76543210u, 0xfedcba98u, false, false);
#endif
}
kernels/common/wmma_backend.h
+57
View File
@@ -0,0 +1,57 @@
#pragma once
#include <stdint.h>
#include "wave32_exchange.h"
// d4r's gfx11 interface: each lane supplies 16 FP16 K operands. The 8 FP32
// outputs have logical row 2*i + lane/16 and logical column lane%16.
// gfx12 consumes only 8 FP16 operands per lane and emits 8 consecutive rows
// per lane group. Keep the old interface for K/M/texture call sites until their
// surrounding fragments are ported and validated individually.
using d4r_wmma_h16 = _Float16 __attribute__((ext_vector_type(16)));
using d4r_wmma_h8 = _Float16 __attribute__((ext_vector_type(8)));
using d4r_wmma_f8 = float __attribute__((ext_vector_type(8)));
__attribute__((device, always_inline)) static inline float d4r_wmma_other_half(float value)
{
const uint32_t bits = __builtin_bit_cast(uint32_t, value);
const uint32_t swapped = d4r_wave32_other_half(bits);
return __builtin_bit_cast(float, swapped);
}
__attribute__((device, always_inline)) static inline d4r_wmma_f8
d4r_wmma_legacy_layout(d4r_wmma_h16 a, d4r_wmma_h16 b, d4r_wmma_f8 c)
{
#if defined(__gfx1200__) || defined(__gfx1201__)
const uint32_t group = __builtin_amdgcn_mbcnt_lo(~0u, 0u) >> 4;
d4r_wmma_h8 a12, b12;
d4r_wmma_f8 c12;
d4r_wmma_f8 remote_c;
#pragma unroll
for (int i = 0; i < 8; ++i)
remote_c[i] = d4r_wmma_other_half(c[i]);
#pragma unroll
for (int i = 0; i < 8; ++i) {
a12[i] = a[8u * group + i];
b12[i] = b[8u * group + i];
const uint32_t row = 8u * group + i;
const uint32_t old_index = row >> 1;
c12[i] = (row & 1u) == group ? c[old_index] : remote_c[old_index];
}
const d4r_wmma_f8 d12 = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a12, b12, c12);
d4r_wmma_f8 result, remote_d;
#pragma unroll
for (int i = 0; i < 8; ++i)
remote_d[i] = d4r_wmma_other_half(d12[i]);
#pragma unroll
for (int i = 0; i < 8; ++i) {
const uint32_t row = 2u * i + group;
const uint32_t new_index = row & 7u;
result[i] = (row >> 3) == group ? d12[new_index] : remote_d[new_index];
}
return result;
#elif defined(__GFX11__)
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
#else
#error "d4r WMMA backend requires a validated gfx11 or gfx12 target"
#endif
}
kernels/common/wmma_layout.h
+346
View File
@@ -0,0 +1,346 @@
// 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>
#include "wave32_exchange.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 d4r_wave32_other_half(v);
}
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);
}
#ifdef D4R_K_F16_WMMA
// Rejected K experiment: direct packed F16 changes the arithmetic on gfx1201.
// The public WMMA reproducer and enc1 replay are in tools/windows and docs.
// Experimental gfx12 packed FP16 C/D. K's strict baseline enters each K16
// step with an already FP16-rounded accumulator. Keep the existing logical
// eight-row contract; only this instruction/register representation changes.
WM_FN wm_f8v wm_mma_f16_step(const wm_op& a, const wm_op& b, wm_f8v c)
{
const wm_h8 hc = __builtin_convertvector(c, wm_h8);
const wm_h8 hd = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12(
__builtin_bit_cast(wm_h8, a), __builtin_bit_cast(wm_h8, b), hc);
return __builtin_convertvector(hd, wm_f8v);
}
#endif
#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
+10
View File
@@ -0,0 +1,10 @@
// 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
#ifndef D4R_K_POS_MT
#define D4R_K_POS_MT 2
#elif !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
#error "experimental position token tiling requires gfx12 strict K baseline"
#endif
#include "pwin_pos.h"
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, D4R_K_POS_MT)
kernels/k/dltss_pwin_dec1_layer.hip
+6
View File
@@ -0,0 +1,6 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// QKV chains two at a time (K-race port): -1.7% on top of op_imgt (paired gfx1200 replay, exact)
// d4r-build-flags: -mcumode -DPWIN_QKV_SPLIT
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)
kernels/k/dltss_pwin_dec2_layer.hip
+5
View File
@@ -0,0 +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_dec3_layer.hip
+2
View File
@@ -0,0 +1,2 @@
#include "pwin_wide.h"
WIDE_DECODER(dltss_pwin_dec3_layer, 4, 96, 128, 2)
kernels/k/dltss_pwin_dec4_layer.hip
+3
View File
@@ -0,0 +1,3 @@
// NG = 8 (32 waves per window): faster on gfx1200 in paired replay (local k9060xt-opt v3)
#include "pwin_wide.h"
WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 8)
kernels/k/dltss_pwin_dec5_layer.hip
+3
View File
@@ -0,0 +1,3 @@
// NG = 8 (32 waves per window): faster on gfx1200 in paired replay (local k9060xt-opt v3)
#include "pwin_wide.h"
WIDE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 8)
kernels/k/dltss_pwin_enc0_layer.hip
+20
View File
@@ -0,0 +1,20 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// op_imgt / uniform wave index: 0.6-6% slower here (paired gfx1200 replay), old address form kept
// d4r-build-flags: -mcumode -DPWIN_NO_IMGT -DPWIN_PLAIN_WV
#define PWIN_F32ACC
#ifndef D4R_K_POS_MT
#define D4R_K_POS_MT 2
#elif !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
#error "experimental position token tiling requires gfx12 strict K baseline"
#endif
// Reuse V across the window's waves. Paired gfx1201 replays retain every half
// code and reduce enc0 time by about 5%; gfx11 and arithmetic experiments keep
// their previous path. D4R_K_PRIVATE_V supplies the original A/B control.
#if defined(__GFX12__) && defined(D4R_K_FP16_BASELINE) && !defined(D4R_K_PRIVATE_V) && \
!defined(D4R_K_PACKED_ACC) && !defined(D4R_K_F16_WMMA)
#ifndef POS_VSHARE
#define POS_VSHARE
#endif
#endif
#include "pwin_pos.h"
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, D4R_K_POS_MT)
kernels/k/dltss_pwin_enc1_layer.hip
+16
View File
@@ -0,0 +1,16 @@
// 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
#ifdef D4R_K_ENC1_POS_MT
#if !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
#error "experimental enc1 token tiling requires gfx12 strict K baseline"
#endif
#if D4R_K_ENC1_POS_MT != 1 && D4R_K_ENC1_POS_MT != 2 && D4R_K_ENC1_POS_MT != 4
#error "enc1 token tiles per wave must be 1, 2 or 4"
#endif
#include "pwin_pos.h"
POS_ENCODER(dltss_pwin_enc1_layer, 2, 64, 64, 0, D4R_K_ENC1_POS_MT)
#else
#include "pwin_layer.h"
PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
#endif
kernels/k/dltss_pwin_enc2_layer.hip
+6
View File
@@ -0,0 +1,6 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// QKV chains two at a time (K-race port): -3% on top of op_imgt (paired gfx1200 replay, exact)
// d4r-build-flags: -mcumode -DPWIN_QKV_SPLIT
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)
kernels/k/dltss_pwin_enc3_layer.hip
+2
View File
@@ -0,0 +1,2 @@
#include "pwin_wide.h"
WIDE_ENCODER(dltss_pwin_enc3_layer, 4, 96, 128, 2)
kernels/k/dltss_pwin_enc4_layer.hip
+5
View File
@@ -0,0 +1,5 @@
// op_imgt / uniform wave index: 2-10% slower here (paired gfx1200 replay), old address form kept
// d4r-build-flags: -DPWIN_NO_IMGT -DPWIN_PLAIN_WV
// NG = 8 (32 waves per window): faster on gfx1200 in paired replay (local k9060xt-opt v3)
#include "pwin_wide.h"
WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 8)
kernels/k/pwin_common.h
+288
View File
@@ -0,0 +1,288 @@
// 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) 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"
#else
#include <hip/hip_runtime.h>
#endif
#include <stdint.h>
#include "../common/wmma_layout.h"
#if defined(D4R_ACCURACY) && defined(PWIN_F32ACC)
#undef PWIN_F32ACC
#endif
#pragma clang fp contract(off)
// LDS row padding in halves (port of the K-race research): activation rows of N halves are stored N + PWIN_PAD
// apart so the 16 token rows a WMMA operand load touches start on distinct banks. Layout only; values unchanged.
#ifndef PWIN_PAD
#define PWIN_PAD 8
#endif
#define LDSR(N) ((N) + PWIN_PAD)
typedef _Float16 half_t;
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
typedef _Float16 h16 __attribute__((ext_vector_type(16)));
typedef float f8v __attribute__((ext_vector_type(8)));
#ifdef D4R_K_PACKED_ACC
#if !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE) || defined(D4R_K_F16_WMMA)
#error "packed accumulator storage requires gfx12 strict F32 WMMA baseline"
#endif
// Keep already-rounded values packed between steps. Every matrix instruction
// still uses F32 C/D and rounds once to half afterwards, like the baseline.
typedef _Float16 acc8v __attribute__((ext_vector_type(8)));
#else
typedef f8v acc8v;
#endif
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
__device__ __forceinline__ uint32_t lane_id()
{
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
}
typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
// The CUDA fragment reduction rounds products and every tree node to f16.
// A float sum of the squared halves changes rsqrt and can amplify differences
// through attention. The strict baseline preserves its channel grouping.
template <int C>
__device__ __forceinline__ half_t l2_sum(const half_t* row)
{
#ifdef D4R_K_FP16_BASELINE
static_assert(C % 32 == 0 && C <= 160, "audited K channel counts");
hv2 partial[4];
#pragma unroll
for (int t = 0; t < 4; ++t)
{
hv2 halves[2];
#pragma unroll
for (int parity = 0; parity < 2; ++parity)
{
hv2 pairs[C / 32];
#pragma unroll
for (int group = 0; group < C / 32; ++group)
{
const hv2 a = *(const hv2*)(row + 2 * t + 8 * parity + 32 * group);
const hv2 b = *(const hv2*)(row + 2 * t + 8 * parity + 32 * group + 16);
pairs[group] = (a * a) + (b * b);
}
hv2 value = pairs[0];
if constexpr (C >= 64) value = value + pairs[1];
if constexpr (C == 96) value = value + pairs[2];
if constexpr (C >= 128) value = value + (pairs[2] + pairs[3]);
if constexpr (C == 160) value = value + pairs[4];
halves[parity] = value;
}
partial[t] = halves[0] + halves[1];
}
const hv2 sum = (partial[0] + partial[1]) + (partial[2] + partial[3]);
return sum[0] + sum[1];
#else
float sum = 0;
#pragma unroll
for (int c = 0; c < C; ++c) sum += (float)(half_t)(row[c] * row[c]);
return (half_t)sum;
#endif
}
// 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__ acc8v mma16(const op_t& a, const op_t& b, acc8v c)
{
#ifdef D4R_K_PACKED_ACC
const f8v fc = __builtin_convertvector(c, f8v);
const f8v d = wm_mma(a, b, fc);
return __builtin_convertvector(d, acc8v);
#elif defined(D4R_K_F16_WMMA)
#if !defined(__GFX12__) || D4R_WMMA_LAYOUT != 12 || !defined(D4R_K_FP16_BASELINE)
#error "experimental packed F16 WMMA requires gfx12 native layout and strict K baseline"
#endif
return wm_mma_f16_step(a, b, c);
#else
f8v d = wm_mma(a, b, c);
#if !defined(PWIN_F32ACC) || defined(D4R_K_FP16_BASELINE)
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)(half_t)d[i];
#endif
return d;
#endif
}
__device__ __forceinline__ acc8v splat8(float v)
{
#ifdef D4R_K_PACKED_ACC
const half_t h = (half_t)v;
return (acc8v){h, h, h, h, h, h, h, h};
#else
return (acc8v){v, v, v, v, v, v, v, v};
#endif
}
__device__ __forceinline__ uint32_t other_half(uint32_t v)
{
return wm_other_half(v);
}
__device__ __forceinline__ uint32_t pack2(half_t lo, half_t hi)
{
return __builtin_bit_cast(uint32_t, (hv2){lo, hi});
}
__device__ __forceinline__ half_t lo16(uint32_t v)
{
return __builtin_bit_cast(hv2, v)[0];
}
__device__ __forceinline__ half_t hi16(uint32_t v)
{
return __builtin_bit_cast(hv2, v)[1];
}
// D^T f16 values (own[i] = channel 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])
{
return wm_op_from_acc(own);
}
__device__ __forceinline__ op_t operand_from_f8(acc8v d)
{
#ifdef D4R_K_PACKED_ACC
const half_t h[8] = {d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7]};
return wm_op_from_acc(h);
#else
return wm_op_from_f8(d);
#endif
}
// 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]);
}
__device__ __forceinline__ u8v lds_row16(const half_t* p)
{
const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
}
// Global activation traffic (tokens in, rows out): PWIN_NT_LOAD / PWIN_NT_STORE mark it non-temporal so
// it does not evict the weight images from the caches. Same data either way.
__device__ __forceinline__ u8v gload_row16(const half_t* p)
{
#ifdef PWIN_NT_LOAD
const u4v a = __builtin_nontemporal_load((const u4v*)p), b = __builtin_nontemporal_load((const u4v*)(p + 8));
#else
const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
#endif
return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
}
__device__ __forceinline__ void gstore16(half_t* p, u4v v)
{
#ifdef PWIN_NT_STORE
__builtin_nontemporal_store(v, (u4v*)p);
#else
*(u4v*)p = v;
#endif
}
__device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
{
gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
}
// 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);
}
// wave index (threadIdx.z) as a wave-uniform scalar, so op_imgt's tile stays in SGPRs. PWIN_PLAIN_WV: the old
// per-lane read (with PWIN_NO_IMGT the layer compiles as before the port; enc0/enc4 are faster that way)
#ifdef PWIN_PLAIN_WV
#define PWIN_WAVE_INDEX threadIdx.z
#else
#define PWIN_WAVE_INDEX __builtin_amdgcn_readfirstlane(threadIdx.z)
#endif
// slot (tile, lane l & 15) of an operand image in global memory (port of K-race E14 step 1): the tile is
// wave-uniform, so the row address stays a scalar base and only the lane indexes it. Reads exactly the slot
// op_img(img, tile * 16 + m) reads; only the address arithmetic differs. PWIN_NO_IMGT: the old form.
__device__ __forceinline__ op_t op_imgt(const u8v* __restrict__ img, int tile)
{
#ifdef PWIN_NO_IMGT
return wm_op_image(img, tile * 16 + (lane_id() & 15)); // unsigned sum, exactly the old index expression
#else
return wm_op_image(img + 16 * tile, (int)(lane_id() & 15));
#endif
}
// 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();
}
__device__ __forceinline__ int mirror(int v, int n)
{
v = v < 0 ? -v : v;
return v < 2 * n - 2 - v ? v : 2 * n - 2 - v;
}
// ------------------------------------------------ pre-swizzled f16 weight fragments (NVIDIA)
// logical B[k][n] of one 16x16 fragment (512 bytes)
__device__ __forceinline__ int frag_offset(int k, int n)
{
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1);
}
kernels/k/pwin_layer.h
+869
View File
@@ -0,0 +1,869 @@
// Native DltssPaddedWinLayer (DLSS 4, preset K) template: 8x8 windows, H heads x 32, C channels.
// One block per window, 4 waves (one 16-token tile each). Semantics: kernels/tools/pwin_model.py (swin_core,
// pwin_encoder); weight layout: pwin_model.Layout.
#pragma once
#include "pwin_common.h"
#ifndef PWIN_VGPR
#define PWIN_VGPR
#endif
struct PwinParams
{
int W, H; // 0: token grid
const uint8_t* in; // 8
const uint8_t* skip; // 16 (decoders)
uint8_t* out24; // 24: encoders: merged output; decoders: output
uint8_t* out32; // 32: encoders: full-resolution output
uint64_t pad40, pad48; // 40
int sx, sy; // 56
const uint8_t* w; // 64
uint8_t rest[176 - 72];
};
static_assert(sizeof(PwinParams) == 176, "param block");
template <int H_, int C_, int COUT_, bool POS_ = false> struct PwinLayout
{
static constexpr bool POS = POS_; // position-only attention (trait flag 2): the table is P itself
static constexpr int H = H_, C = C_, COUT = COUT_, KT = C / 16, NCH = C / 32, NMLP = C / 8;
static constexpr int G1 = 0, QKV = 2 * C, HEAD = 192 * C, BIAS = QKV + HEAD * H, BO = BIAS + 8192 * H,
WO = BO + 2 * C, G2 = WO + 64 * C * H, B2 = G2 + 2 * C, W1 = B2 + 2 * C, B1 = W1 + 8 * C * C,
W2 = B1 + 8 * C, PM = W2 + 8 * C * C;
// patch merge: NVIDIA's warps own COUT / H channels each, allocated as whole 16-column tiles
static constexpr int NPW = COUT / H, NPA = (NPW + 15) / 16 * 16, PMT = H * NPA / 16, PMB = PM + 8 * C * NPA * H;
// prep image tiles (16 lanes of u8v each)
static constexpr int T_QKV = 0; // [h][j][kt KT][nt 2]
static constexpr int T_WO = T_QKV + H * 3 * KT * 2; // [ks 2H][nt KT]
static constexpr int T_W1 = T_WO + 2 * H * KT; // [c][kt KT][nt 2]
static constexpr int T_W2 = T_W1 + NMLP * KT * 2; // [c][ks 2][nt KT]
static constexpr int T_PM = T_W2 + NMLP * 2 * KT; // [ks C/4][nt COUT/16]
static constexpr int T_END = T_PM + (C / 4) * PMT;
static constexpr int NBIAS = H * 4 * 4 * 32; // u4v per lane
static constexpr int PREP_ITEMS = T_END * 16 + NBIAS;
__device__ static int src_of(int tile)
{
if (tile < T_WO)
{
const int nt = tile & 1, kt = (tile >> 1) % KT, j = (tile / (2 * KT)) % 3, h = tile / (6 * KT);
return QKV + HEAD * h + 2048 * (3 * (kt >> 1) + j) + 512 * (kt & 1) + 1024 * nt;
}
if (tile < T_W1)
{
const int t = tile - T_WO, nt = t % KT, ks = t / KT;
return WO + 64 * C * (ks >> 1) + 512 * (ks & 1) + 1024 * nt;
}
if (tile < T_W2)
{
const int t = tile - T_W1, nt = t & 1, kt = (t >> 1) % KT, c = t / (2 * KT);
return W1 + 64 * C * c + 512 * kt + 32 * C * nt;
}
if (tile < T_PM)
{
const int t = tile - T_W2, nt = t % KT, ks = (t / KT) & 1, c = t / (2 * KT);
return W2 + 64 * C * c + 512 * ks + 1024 * nt;
}
const int t = tile - T_PM, nt = t % PMT, ks = t / PMT;
return PM + 512 * ks + 128 * C * nt; // allocated tiles are contiguous
}
};
__device__ __forceinline__ half_t wload(const uint8_t* w, int byte)
{
return *(const half_t*)(w + byte);
}
// channels 16 nt + 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])
{
wm_acc_vec_load((const half_t*)(w + byte), out);
}
__device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
{
u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
r[j] = pack2(wload(w, src + frag_offset(2 * j, n)), wload(w, src + frag_offset(2 * j + 1, n)));
return r;
}
template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx)
{
const uint8_t* w = p.w;
if (idx < L::T_END * 16)
{
img[idx] = img_tile(w, L::src_of(idx >> 4), idx & 15);
return;
}
const int b = idx - L::T_END * 16;
if (b >= L::NBIAS)
return;
if constexpr (L::POS)
{
// B operand of P V: lane = query 16 qt + l16, 16 keys 16 kt .. 16 kt + 15 (u8v per lane; 16 lanes)
if (b >= L::H * 4 * 4 * 16)
return;
const int l16 = b & 15, kt = (b >> 4) & 3, qt = (b >> 6) & 3, h = b >> 8, r = l16;
u8v v;
#pragma unroll
for (int j = 0; j < 8; ++j)
{
half_t e[2];
#pragma unroll
for (int u = 0; u < 2; ++u)
{
const int c = 2 * j + u;
const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
e[u] = wload(w, L::BIAS + 8192 * h + off);
}
v[j] = pack2(e[0], e[1]);
}
((u8v*)bias)[b] = v;
return;
}
const int lane = b & 31, kt = (b >> 5) & 3, qt = (b >> 7) & 3, h = b >> 9;
const int q = 16 * qt + (lane & 15);
half_t v[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int key = 16 * kt + 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);
}
bias[b] = (u4v){pack2(v[0], v[1]), pack2(v[2], v[3]), pack2(v[4], v[5]), pack2(v[6], v[7])};
}
template <class L> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
{
pwin_prep_at<L>(p, img, bias, blockIdx.x * 128 + threadIdx.x);
}
__device__ __forceinline__ half_t hmin(half_t a, half_t b) { return a < b ? a : b; }
__device__ __forceinline__ half_t hmax(half_t a, half_t b) { return a > b ? a : b; }
__device__ __forceinline__ half_t softmax_e(half_t s)
{
const half_t hi = __builtin_bit_cast(half_t, (uint16_t)0x4D69), lo = __builtin_bit_cast(half_t, (uint16_t)0xCD69);
const half_t a = __builtin_bit_cast(half_t, (uint16_t)0x91D5), c = __builtin_bit_cast(half_t, (uint16_t)0xC50D);
const half_t t = hmax(hmin(s, hi), lo);
const half_t u = t * t;
const half_t v = __builtin_fmaf16(u, a, (half_t)1.0f);
const half_t x = __builtin_fmaf16(t, v, c);
return (half_t)__builtin_amdgcn_exp2f((float)x);
}
__device__ __forceinline__ half_t gelu(half_t x)
{
const half_t k1 = (half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
const half_t c = hmin(hmax(x, (half_t)-2.0f), (half_t)2.0f);
const half_t ac = c < (half_t)0.0f ? -c : c;
const half_t u = k1 - k2 * ac;
return x * ((half_t)0.5f + c * u);
}
// gelu of 8 accumulator values as 4 packed f16 pairs (the same f16 operations as gelu(), element-wise)
__device__ __forceinline__ void gelu8(const acc8v& d, half_t g[8])
{
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
#pragma unroll
for (int j = 0; j < 4; ++j)
{
const hv2 x = (hv2){(half_t)d[2 * j], (half_t)d[2 * j + 1]};
const hv2 c = __builtin_elementwise_min(__builtin_elementwise_max(x, (hv2)(half_t)-2.0f), (hv2)(half_t)2.0f);
const hv2 u = k1 - k2 * __builtin_elementwise_abs(c);
const hv2 r = x * ((hv2)(half_t)0.5f + c * u);
g[2 * j] = r[0];
g[2 * j + 1] = r[1];
}
}
// store the D^T tile of token m as 16 contiguous channels (lanes of one half-wave write)
__device__ __forceinline__ void store_row16(half_t* dst, const u8v& v)
{
*(u4v*)dst = (u4v){v[0], v[1], v[2], v[3]};
*(u4v*)(dst + 8) = (u4v){v[4], v[5], v[6], v[7]};
}
#ifdef PWIN_STAMP
// per-block stage timestamps of wave 0 (s_memtime), slot k
__device__ uint64_t g_stamp[16384][16];
#define STAMP(k) \
do \
{ \
if (threadIdx.z == 0 && lane_id() == 0) \
{ \
const int bid = blockIdx.y * gridDim.x + blockIdx.x; \
if (bid < 16384) \
g_stamp[bid][k] = __builtin_readcyclecounter(); \
} \
} while (0)
#else
#define STAMP(k) do {} while (0)
#endif
#ifdef PWIN_DBG
__device__ float g_dbg[12][64][160];
#define DBG_ON (blockIdx.x == DBX && blockIdx.y == DBY)
#endif
// 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>
__device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
half_t (*act)[LDSR(L::C)], half_t (*hb)[LDSR(L::C)], half_t (*kl)[LDSR(32)], half_t (*vt)[LDSR(64)])
{
constexpr int C = L::C, KT = L::KT, H = L::H;
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX;
const uint8_t* w = p.w;
half_t* arow = act[16 * wv + m];
half_t* hrow = hb[16 * wv + m];
STAMP(1);
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb
const half_t ss = l2_sum<C>(arow);
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
for (int kt = (int)hf; kt < KT; kt += 2)
{
const u8v xv = lds_row16(arow + 16 * kt);
u8v hv;
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt); // 16 contiguous gammas
#pragma unroll
for (int j = 0; j < 8; ++j)
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
store_row16(hrow + 16 * kt, hv);
}
__builtin_amdgcn_wave_barrier();
#ifdef PWIN_DBG
if (DBG_ON && !hf)
for (int c = 0; c < C; ++c)
g_dbg[0][16 * wv + m][c] = (float)hrow[c];
#endif
if constexpr (L::POS)
block_sync(); // V of all 64 tokens is computed by every wave from hb
STAMP(2);
// ---- attention, one head at a time; the output projection accumulates head by head
acc8v acc[KT];
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
half_t vv_[8];
dvec8(w, L::BO + 32 * nt, vv_);
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[nt][i] = (float)vv_[i];
}
for (int h = 0; h < H; ++h)
{
// position-only attention: V^T A operands straight from D = h1 Wv (lane = dim, 16 keys per tile)
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
__shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
acc8v d = splat8(0.0f);
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), 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] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
}
op_t pop[4];
if constexpr (L::POS)
{
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + wv) * 4 + kt);
}
else
{
op_t qop[2];
#ifdef PWIN_QKV_SPLIT
// (port of K-race dec1 QKV split) the six chains (V, K, Q x 2 n-tiles; each still over kt in order) two at
// a time: 2 live accumulators instead of 6. Same per-chain WMMA order, same stores.
#pragma unroll
for (int j = 2; j >= 0; --j)
{
acc8v dj[2] = {splat8(0.0f), splat8(0.0f)};
for (int kt = 0; kt < KT; ++kt)
{
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
dj[nt] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + j) * KT + kt) * 2 + nt), row, dj[nt]);
}
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const acc8v d = dj[nt];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
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 + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
}
}
}
#else
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
acc8v dq[6];
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[u]);
}
#pragma unroll
for (int j = 0; j < 3; ++j)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const acc8v d = dq[2 * j + nt];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
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 + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
}
}
#endif
block_sync();
half_t e[4][8];
float rs = 0.0f;
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
acc8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
{
d[2 * i] = (float)lo16(bi[i]);
d[2 * i + 1] = (float)hi16(bi[i]);
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(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]);
if constexpr (!kShimRowSum)
rs += (float)e[kt][i];
}
}
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)
{
half_t pv[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
pv[i] = e[kt][i] * rc;
pop[kt] = operand_from_dt(pv);
}
}
op_t oop[2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
acc8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
if constexpr (L::POS)
d = mma16(vtop[nt][kt], pop[kt], d);
else
d = mma16(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 + 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(op_imgt(img, L::T_WO + (2 * h + ks) * KT + nt), oop[ks], acc[nt]);
if constexpr (!L::POS)
block_sync(); // kl / vt are reused by the next head
}
#ifdef PWIN_DBG
if (DBG_ON)
for (int nt = 0; nt < KT; ++nt)
for (int i = 0; i < 8; ++i)
g_dbg[1][16 * wv + m][16 * nt + 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 op_t xv = op_lds(arow + 16 * nt);
half_t mv[8], g2[8], b2[8];
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] + wm_op_acc_elem(xv, i);
mv[i] = x1 * g2[i];
acc[nt][i] = (float)(half_t)(x1 + b2[i]);
}
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)
{
op_t gop[2];
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
acc8v d;
{
half_t vv_[8];
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(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_lds(hrow + 16 * kt), d);
half_t g[8];
gelu8(d, g);
gop[hn] = operand_from_dt(g);
}
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
acc[nt] = mma16(op_imgt(img, L::T_W2 + (c * 2 + ks) * KT + nt), gop[ks], acc[nt]);
}
STAMP(4);
// y -> act (rows of this wave)
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
}
__builtin_amdgcn_wave_barrier();
}
// first layer: token embedding x0 = relu(in (CIN) W_e + b_e) in front of the core (core weights at CB)
template <int H_, int C_, int COUT_, bool POS_, int CIN_> struct PwinEmbLayout : PwinLayout<H_, C_, COUT_, POS_>
{
using B = PwinLayout<H_, C_, COUT_, POS_>;
static constexpr int CIN = CIN_, KE = CIN / 16, CB = CIN ? 2 * C_ + 2 * CIN * C_ : 0;
static constexpr int T_EMB = B::T_END, T_END = T_EMB + KE * (C_ / 16);
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
};
template <class L> __device__ void pwin_emb_prep(const PwinParams& p, u8v* img, u4v* bias)
{
const int idx = blockIdx.x * 128 + threadIdx.x;
if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
{
const int t = (idx >> 4) - L::T_EMB, nt = t % (L::C / 16), ks = t / (L::C / 16);
img[idx] = img_tile(p.w, 2 * L::C + 512 * ks + 512 * L::KE * nt, idx & 15);
return;
}
PwinParams q = p;
q.w = p.w + L::CB;
// core tiles and the bias / P images (indices after T_END are shifted down to the base layout's)
const int base_idx = idx < L::T_EMB * 16 ? idx : idx - (L::T_END - L::T_EMB) * 16;
if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
return;
pwin_prep_at<typename L::B>(q, img, bias, base_idx);
}
// encoder: core + full-resolution output + 2x2 patch merge
template <class L>
__device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
PwinParams p = p0;
if constexpr (L::CB != 0)
p.w = p0.w + L::CB;
constexpr int C = L::C, COUT = L::COUT;
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t kl[64][LDSR(32)];
__shared__ __attribute__((aligned(16))) half_t vt[32][LDSR(64)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
const int bx = blockIdx.x, by = blockIdx.y;
STAMP(0);
// input rows (mirror padded) -> act (first layer: embedding + ReLU)
{
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
op_t xe[L::KE];
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
xe[ks] = op_gload(xin + 16 * ks);
#pragma unroll
for (int nt = 0; nt < C / 16; ++nt)
{
acc8v d;
{
half_t vv_[8];
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(op_imgt(img, L::T_EMB + ks * (C / 16) + nt), 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);
op_store(&act[tok][16 * nt], operand_from_dt(xo));
}
}
else
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
}
}
__builtin_amdgcn_wave_barrier();
pwin_core<L>(p, img, bias, act, hb, kl, vt);
STAMP(5);
// full-resolution output
const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
{
half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
}
block_sync();
STAMP(6);
// patch merge: 16 merged tokens, K = 4 sub-tokens x C, COUT/16 n-tiles over the 4 waves
const int my = m >> 2, mx = m & 3;
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
const uint8_t* w = p.w;
for (int vt = wv; vt < L::PMT; vt += 4)
{
acc8v d;
{
half_t vv_[8];
dvec8(w, L::PMB + 32 * vt, vv_);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv_[i];
STAMP(7);
}
for (int ks = 0; ks < C / 4; ++ks)
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&act[src][16 * (ks % L::KT)]), 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;
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
}
}
#define PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, CIN) \
using NAME##_L = PwinEmbLayout<H, C, COUT, POS, CIN>; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p) \
{ \
pwin_encoder<NAME##_L>(p, g_img, g_bias); \
}
#define PWIN_ENCODER(NAME, H, C, COUT) PWIN_ENCODER_EMB(NAME, H, C, COUT, false, 0)
#define PWIN_ENCODER_POS(NAME, H, C, COUT, POS) PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, 0)
// bottleneck (dec5): core only, output at the same resolution to +24
template <class L>
__device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C;
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t kl[64][LDSR(32)];
__shared__ __attribute__((aligned(16))) half_t vt[32][LDSR(64)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
const int bx = blockIdx.x, by = blockIdx.y;
{
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
}
__builtin_amdgcn_wave_barrier();
pwin_core<L>(p, img, bias, act, hb, kl, vt);
const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
{
half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
}
}
#define PWIN_KERNEL(NAME, H, C, COUT, BODY) \
using NAME##_L = PwinLayout<H, C, COUT>; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
pwin_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) \
{ \
BODY<NAME##_L>(p, g_img, g_bias); \
}
// ---------------------------------------------------------------- decoders (patch expand + skip)
template <int H_, int C_, int CL_, bool POS_ = false, int NOUT_ = 0, int NOUTA_ = 0>
struct PwinDecLayout : PwinLayout<H_, C_, 16, POS_>
{
using B = PwinLayout<H_, C_, 16, POS_>;
// last layer: output head out = y W_h + b_h (NOUTA computed, NOUT stored); b_h at B::PM, W_h after it
static constexpr int NOUT = NOUT_, NOUTA = NOUTA_, HB = B::PM, HW = B::PM + 2 * NOUTA_;
static constexpr int CL = CL_, KL = CL / 16, EXPN = 4 * C_ / 16; // expand: K CL, N 4C
static constexpr int EXPB = 2 * CL * 4 * C_, CB = EXPB + 8 * C_; // expand bias, core weights base
static constexpr int T_EXP = B::T_PM; // expand tiles replace the merge tiles
static constexpr int T_HEAD = T_EXP + KL * EXPN;
static constexpr int T_END = T_HEAD + B::KT * (NOUTA_ / 16);
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
};
template <class L> __device__ void pwin_dec_prep(const PwinParams& p, u8v* img, u4v* bias)
{
const int idx = blockIdx.x * 128 + threadIdx.x;
if (idx >= L::T_EXP * 16 && idx < L::T_HEAD * 16)
{
const int t = (idx >> 4) - L::T_EXP, nt = t % L::EXPN, ks = t / L::EXPN;
img[idx] = img_tile(p.w, 512 * ks + 32 * L::CL * nt, idx & 15);
return;
}
if (idx >= L::T_HEAD * 16 && idx < L::T_END * 16)
{
const int t = (idx >> 4) - L::T_HEAD, nt = t % (L::NOUTA / 16), ks = t / (L::NOUTA / 16);
img[idx] = img_tile(p.w + L::CB, L::HW + 512 * ks + 512 * L::KT * nt, idx & 15);
return;
}
PwinParams q = p;
q.w = p.w + L::CB;
// core tiles lie below T_EXP; bias / P images follow the base layout's (unused) merge tiles
pwin_prep_at<typename L::B>(q, img, bias, idx < L::T_EXP * 16 ? idx : L::B::T_END * 16 + (idx - L::T_END * 16));
}
template <class L>
__device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C;
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : 32;
__shared__ __attribute__((aligned(16))) half_t kl[KVR][LDSR(32)];
__shared__ __attribute__((aligned(16))) half_t vt[KVC][LDSR(64)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, bx = blockIdx.x, by = blockIdx.y;
// ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens
{
const int W2 = p.W / 2, H2 = p.H / 2;
const int ly = m >> 2, lx = m & 3;
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
const int dy = wv >> 1, dx = wv & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
for (int nt = 0; nt < C / 16; ++nt)
{
const int gnt = wv * (C / 16) + nt; // expand output column tile
acc8v d;
{
half_t vv_[8];
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(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), 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] + wm_op_acc_elem(sv, i);
op_store(&act[t][16 * nt], operand_from_dt(xo));
}
}
block_sync();
PwinParams q = p;
q.w = p.w + L::CB;
pwin_core<L>(q, img, bias, act, hb, kl, vt);
const int tok = 16 * wv + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
if constexpr (L::NOUTA != 0)
{
half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
#pragma unroll
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
acc8v d;
{
half_t vv_[8];
dvec8(q.w, L::HB + 32 * nt, vv_);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv_[i];
}
#pragma unroll
for (int kt = 0; kt < L::KT; ++kt)
d = mma16(op_imgt(img, L::T_HEAD + kt * (L::NOUTA / 16) + nt), op_lds(&act[tok][16 * kt]), d);
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
}
}
else if (inb)
{
half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
}
}
#define PWIN_DECODER_EX(NAME, H, C, CL, POS, NOUT, NOUTA) \
using NAME##_L = PwinDecLayout<H, C, CL, POS, NOUT, NOUTA>; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
pwin_dec_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p) \
{ \
pwin_decoder<NAME##_L>(p, g_img, g_bias); \
}
#define PWIN_DECODER(NAME, H, C, CL) PWIN_DECODER_EX(NAME, H, C, CL, false, 0, 0)
kernels/k/pwin_pos.h
+399
View File
@@ -0,0 +1,399 @@
// Position-only-attention layers (trait flag 2: enc0, enc1, dec0) with MT token tiles per wave, so every
// weight operand loaded from L2 is used MT times (4 / MT waves per window). Same math as pwin_core.
#pragma once
#include "pwin_layer.h"
template <class L, int MT>
__device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
half_t (*act)[LDSR(L::C)], half_t (*hb)[LDSR(L::C)])
{
constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
static_assert(MT == 1 || MT == 2 || MT == 4, "position token tiles per wave");
static_assert(L::POS, "position-only attention layers");
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX;
const uint8_t* w = p.w;
// ---- L2 norm of every tile of this wave -> hb
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const half_t* arow = act[16 * (wv * MT + mi) + m];
half_t* hrow = hb[16 * (wv * MT + mi) + m];
const half_t ss = l2_sum<C>(arow);
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
for (int kt = (int)hf; kt < KT; kt += 2)
{
const u8v xv = lds_row16(arow + 16 * kt);
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
u8v hv;
#pragma unroll
for (int j = 0; j < 8; ++j)
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
store_row16(hrow + 16 * kt, hv);
}
}
block_sync(); // V of all 64 tokens is computed from hb
STAMP(2);
// ---- attention: O = P V per head (P from the table), output projection accumulated head by head
acc8v acc[MT][KT];
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, vv);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[mi][nt][i] = (float)vv[i];
}
#ifdef POS_VSHARE
// V^T operands of every head: each wave computes those of its own key tiles, shared through LDS
// (one slot per head, so a head's slot is never rewritten while it is read)
__shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
#endif
for (int h = 0; h < H; ++h)
{
op_t vtop[2][4];
#ifdef POS_VSHARE
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int kt = wv * MT + mi;
acc8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), 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] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
#else
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
acc8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt), d);
vtop[nt][kt] = operand_from_f8(d);
}
#endif
op_t oop[MT][2];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int qt = wv * MT + mi;
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + qt) * 4 + kt);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
acc8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(vtop[nt][kt], pop[kt], d);
oop[mi][nt] = operand_from_f8(d);
}
}
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const op_t wo = op_imgt(img, L::T_WO + (2 * h + ks) * KT + nt);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
}
}
STAMP(3);
// ---- residual, MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
half_t g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, 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 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] + wm_op_acc_elem(xv, i);
mv[i] = x1 * g2[i];
acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
}
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
}
}
__builtin_amdgcn_wave_barrier();
STAMP(4);
// ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
for (int c = 0; c < L::NMLP; ++c)
{
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, bb);
acc8v d[MT];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
#pragma unroll
for (int i = 0; i < 8; ++i)
d[mi][i] = (float)bb[i];
#pragma unroll
for (int kt = 0; kt < KT; ++kt)
{
const op_t w1 = op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn);
#pragma unroll
for (int mi = 0; mi < MT; ++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)
{
half_t g[8];
gelu8(d[mi], g);
gop[mi][hn] = operand_from_dt(g);
}
}
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const op_t w2 = op_imgt(img, L::T_W2 + (c * 2 + ks) * KT + nt);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
}
}
STAMP(5);
// y -> act
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
}
block_sync();
}
// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
template <class L, int MT>
__device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C, NW = 4 / MT;
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
PwinParams p = p0;
if constexpr (L::CB != 0)
p.w = p0.w + L::CB;
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, bx = blockIdx.x, by = blockIdx.y;
STAMP(0);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int tok = 16 * (wv * MT + mi) + m, ty = tok >> 3, tx = tok & 7;
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
op_t xe[L::KE];
#pragma unroll
for (int ks = 0; ks < L::KE; ++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, vv);
acc8v 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(op_imgt(img, L::T_EMB + ks * (C / 16) + nt), 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);
op_store(&act[tok][16 * nt], operand_from_dt(xo));
}
}
else
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
}
}
__builtin_amdgcn_wave_barrier();
STAMP(1);
pos_core<L, MT>(p, img, bias, act, hb);
STAMP(6);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int tok = 16 * (wv * MT + mi) + m;
const int Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
{
half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
for (int kt = (int)hf; kt < L::KT; kt += 2)
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
}
}
STAMP(7);
const int my = m >> 2, mx = m & 3;
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
for (int vt = wv; vt < L::PMT; vt += NW)
{
half_t vv[8];
dvec8(p.w, L::PMB + 32 * vt, vv);
acc8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
for (int ks = 0; ks < C / 4; ++ks)
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&act[src][16 * (ks % L::KT)]), d);
}
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
}
STAMP(8);
}
// dec0: patch expand + skip, core, output head
template <class L, int MT>
__device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C, NW = 4 / MT;
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, bx = blockIdx.x, by = blockIdx.y;
{
const int W2 = p.W / 2, H2 = p.H / 2, ly = m >> 2, lx = m & 3;
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
op_t xls[L::KL];
#pragma unroll
for (int ks = 0; ks < L::KL; ++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;
const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
#pragma unroll
for (int nt = 0; nt < C / 16; ++nt)
{
const int gnt = q * (C / 16) + nt;
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, vv);
acc8v 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(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), 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] + wm_op_acc_elem(sv, i);
op_store(&act[t][16 * nt], operand_from_dt(xo));
}
}
}
block_sync();
PwinParams q = p;
q.w = p.w + L::CB;
pos_core<L, MT>(q, img, bias, act, hb);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int tok = 16 * (wv * MT + mi) + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
#pragma unroll
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
half_t vv[8];
dvec8(q.w, L::HB + 32 * nt, vv);
acc8v 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(op_imgt(img, L::T_HEAD + kt * (L::NOUTA / 16) + nt), op_lds(&act[tok][16 * kt]), d);
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
}
}
}
#define POS_ENCODER(NAME, H, C, COUT, CIN, MT) \
using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
{ \
pos_encoder<NAME##_L, MT>(p, g_img, g_bias); \
}
#define POS_DECODER(NAME, H, C, CL, NOUT, NOUTA, MT) \
using NAME##_L = PwinDecLayout<H, C, CL, true, NOUT, NOUTA>; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
pwin_dec_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
{ \
pos_decoder<NAME##_L, MT>(p, g_img, g_bias); \
}
kernels/k/pwin_wide.h
+430
View File
@@ -0,0 +1,430 @@
// Wide variant of the DltssPaddedWinLayer template for the deep, latency-bound levels (few windows):
// 4 * NG waves per window. Wave wv = threadIdx.z works on token tile t = wv & 3 and channel group
// g = wv >> 2: heads g, g + NG, ... (one head per group and round), output tiles nt = g + NG j of the
// output projection and MLP, and MLP hidden chunks c = NG rr + g. Every output element goes through
// the same f16 operations in the same order as pwin_layer.h (bit-identical results).
//
// Shared memory: hb (64 x C) = normalised input, then MLP input, then the block output y; R = union of
// x0 (64 x C), the K / V images of the NG heads of an attention round, their outputs O, and the double
// buffered MLP hidden chunks of a round.
#pragma once
#include "pwin_layer.h"
template <class L, int NG> struct WideDims
{
// NA: heads per attention round (their K / V images share R)
static constexpr int NA = NG < 4 ? NG : 4;
static constexpr int C = L::C, KT = L::KT, H = L::H, NTL = (KT + NG - 1) / NG, ROUNDS = (H + NA - 1) / NA;
static constexpr int KL_BYTES = 64 * LDSR(32) * 2, VT_BYTES = 32 * LDSR(64) * 2;
static constexpr int X0_BYTES = 64 * LDSR(C) * 2, KV_BYTES = NA * (KL_BYTES + VT_BYTES);
// MLP hidden chunks: one per group and round, double-buffered while that fits
static constexpr int NB = NG <= 4 ? 2 : 1, MROUNDS = (L::NMLP + NG - 1) / NG, HID_BYTES = NB * NG * 64 * LDSR(32) * 2;
static constexpr int R_BYTES = X0_BYTES > KV_BYTES ? (X0_BYTES > HID_BYTES ? X0_BYTES : HID_BYTES)
: (KV_BYTES > HID_BYTES ? KV_BYTES : HID_BYTES);
static_assert(!L::POS, "position-only attention layers use pwin_pos.h");
};
// Core for one window. On entry R holds x0 (64 x C); on return hb holds the block output y.
template <class L, int NG>
__device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
half_t (*hb)[LDSR(L::C)], uint8_t* R)
{
using D = WideDims<L, NG>;
constexpr int C = L::C, KT = L::KT, H = L::H, NTL = D::NTL;
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const uint8_t* w = p.w;
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
const half_t* arow = act[tok];
half_t* hrow = hb[tok];
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb (all groups of a token tile compute the sum)
const half_t ss = l2_sum<C>(arow);
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
for (int kt = 2 * g + (int)hf; kt < KT; kt += 2 * NG)
{
const u8v xv = lds_row16(arow + 16 * kt);
u8v hv;
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
#pragma unroll
for (int j = 0; j < 8; ++j)
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
store_row16(hrow + 16 * kt, hv);
}
// residual input of the owned output tiles (D^T order: channel 16 nt + wm_acc_row(i))
half_t x0[NTL][8];
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
const int nt = g + NG * j;
if (nt < KT)
{
const op_t xv = op_lds(arow + 16 * nt);
#pragma unroll
for (int i = 0; i < 8; ++i)
x0[j][i] = wm_op_acc_elem(xv, i);
}
}
block_sync(); // hb complete, R free
// ---- attention: NG heads per round; the output projection accumulates head by head
acc8v acc[NTL];
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
const int nt = g + NG * j;
if (nt < KT)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[j][i] = (float)vv[i];
}
}
half_t(*kl)[64][LDSR(32)] = (half_t(*)[64][LDSR(32)])R; // [group][token][dim]
constexpr int NA = D::NA;
half_t(*vt)[32][LDSR(64)] = (half_t(*)[32][LDSR(64)])(R + NA * D::KL_BYTES); // [group][dim][token]
half_t(*ob)[64][LDSR(32)] = (half_t(*)[64][LDSR(32)])R; // [group][token][dim] (after K / V)
for (int r = 0; r < D::ROUNDS; ++r)
{
const int h = g < NA ? r * NA + g : H;
op_t qop[2], oop[2];
if (h < H)
{
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
acc8v dq[6];
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[u]);
}
#pragma unroll
for (int j = 0; j < 3; ++j)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const acc8v d = dq[2 * j + nt];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
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 + wm_acc_row(i)][tok] = (half_t)d[i];
}
}
}
block_sync();
if (h < H)
{
half_t e[4][8];
float rs = 0.0f;
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
acc8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
{
d[2 * i] = (float)lo16(bi[i]);
d[2 * i + 1] = (float)hi16(bi[i]);
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(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]);
if constexpr (!kShimRowSum)
rs += (float)e[kt][i];
}
}
rs = row_sum_total(rs, e);
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
half_t pv[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
pv[i] = e[kt][i] * rc;
pop[kt] = operand_from_dt(pv);
}
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
acc8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
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)
{
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)
{
const int hh = r * NA + jh;
if (hh >= H)
break;
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(op_imgt(img, L::T_WO + (2 * hh) * KT + nt), o0, acc[j]);
acc[j] = mma16(op_imgt(img, L::T_WO + (2 * hh + 1) * KT + nt), o1, acc[j]);
}
}
}
if (r + 1 < D::ROUNDS)
block_sync(); // ob is overwritten by the next round's K / V
}
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
const int nt = g + NG * j;
if (nt < KT)
{
half_t mv[8], g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const half_t x1 = (half_t)acc[j][i] + x0[j][i];
mv[i] = x1 * g2[i];
acc[j][i] = (float)(half_t)(x1 + b2[i]);
}
op_store(hrow + 16 * nt, operand_from_dt(mv));
}
}
block_sync();
// ---- MLP: rounds of NG chunks of 32 hidden (double-buffered in R)
half_t(*hid)[NG][64][LDSR(32)] = (half_t(*)[NG][64][LDSR(32)])R;
for (int rr = 0; rr < D::MROUNDS; ++rr)
{
const int c = rr * NG + g;
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
if (c >= L::NMLP)
break;
acc8v d;
{
half_t vv[8];
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(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_lds(hrow + 16 * kt), d);
half_t gg[8];
gelu8(d, gg);
op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
}
block_sync();
for (int jc = 0; jc < NG; ++jc)
{
const int cc = rr * NG + jc;
if (cc >= L::NMLP)
break;
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(op_imgt(img, L::T_W2 + (cc * 2) * KT + nt), h0, acc[j]);
acc[j] = mma16(op_imgt(img, L::T_W2 + (cc * 2 + 1) * KT + nt), h1, acc[j]);
}
}
}
if (D::NB == 1 && rr + 1 < D::MROUNDS)
block_sync();
}
// y -> hb (the last reads of hb were the W1 operands before the final barrier)
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
const int nt = g + NG * j;
if (nt < KT)
{
op_store(hrow + 16 * nt, operand_from_f8(acc[j]));
}
}
block_sync();
}
// x0 rows of a window (mirror padded) -> R
template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uint8_t* R)
{
constexpr int C = L::C;
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W), Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
}
// y rows of the window's own tokens -> dst (C per token)
template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[LDSR(L::C)], uint8_t* dst)
{
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int Y0 = 8 * (int)blockIdx.y - p.sy + (tok >> 3), X0 = 8 * (int)blockIdx.x - p.sx + (tok & 7);
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
{
half_t* yout = (half_t*)dst + (size_t)(Y0 * p.W + X0) * L::C;
for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
gstore_row16(yout + 16 * kt, lds_row16(&hb[tok][16 * kt]));
}
}
template <class L, int NG>
__device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C, COUT = L::COUT;
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
wide_load_x0<L, NG>(p, R);
block_sync();
wide_core<L, NG>(p, img, bias, hb, R);
wide_store_y<L, NG>(p, hb, p.out32);
// patch merge: 16 merged tokens, K = 4 sub-tokens x C, the allocated n-tiles over the waves
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX;
const int my = m >> 2, mx = m & 3;
const int Ym = (8 * (int)blockIdx.y - p.sy) / 2 + my, Xm = (8 * (int)blockIdx.x - p.sx) / 2 + mx;
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
const uint8_t* w = p.w;
for (int vt = wv; vt < L::PMT; vt += 4 * NG)
{
acc8v d;
{
half_t vv[8];
dvec8(w, L::PMB + 32 * vt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
}
for (int ks = 0; ks < C / 4; ++ks)
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(op_imgt(img, L::T_PM + ks * L::PMT + vt), op_lds(&hb[src][16 * (ks % L::KT)]), d);
}
const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
}
}
// bottleneck (dec5): core only, output at the same resolution to +24
template <class L, int NG>
__device__ void wide_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(L::C)];
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
wide_load_x0<L, NG>(p, R);
block_sync();
wide_core<L, NG>(p, img, bias, hb, R);
wide_store_y<L, NG>(p, hb, p.out24);
}
template <class L, int NG>
__device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C;
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = PWIN_WAVE_INDEX, q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
// ---- patch expand: waves of sub-token q = (dy, dx) produce its tiles nt = g + NG j for the 16 low-res tokens
{
const int W2 = p.W / 2, H2 = p.H / 2;
const int ly = m >> 2, lx = m & 3;
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
const int dy = q >> 1, dx = q & 1, tt = 8 * (2 * ly + dy) + 2 * lx + dx;
const int SX = mirror(8 * bx - p.sx + (tt & 7), p.W), SY = mirror(8 * by - p.sy + (tt >> 3), p.H);
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
for (int nt = g; nt < C / 16; nt += NG)
{
const int gnt = q * (C / 16) + nt;
acc8v d;
{
half_t vv[8];
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(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), 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] + wm_op_acc_elem(sv, i);
op_store(&act[tt][16 * nt], operand_from_dt(xo));
}
}
block_sync();
PwinParams qp = p;
qp.w = p.w + L::CB;
wide_core<L, NG>(qp, img, bias, hb, R);
static_assert(L::NOUTA == 0, "output head: use pwin_layer.h");
wide_store_y<L, NG>(p, hb, p.out24);
}
#define WIDE_KERNEL(NAME, PREP, BODY, NG, ...) \
using NAME##_L = __VA_ARGS__; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 * (NG); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
PREP<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(128 * (NG)) PWIN_VGPR NAME(PwinParams p) \
{ \
BODY<NAME##_L, NG>(p, g_img, g_bias); \
}
#define WIDE_ENCODER(NAME, H, C, COUT, NG) WIDE_KERNEL(NAME, pwin_prep, wide_encoder, NG, PwinLayout<H, C, COUT>)
#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep, wide_plain, NG, PwinLayout<H, C, 16>)
#define WIDE_DECODER(NAME, H, C, CL, NG) WIDE_KERNEL(NAME, pwin_dec_prep, wide_decoder, NG, PwinDecLayout<H, C, CL>)
kernels/m/rrlite_dec1_4x4.hip
+7
View File
@@ -0,0 +1,7 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 128
#define SWIN_MLP_R(NW) 2
#include "swin_block.h"
SWIN_MODULE_W(rrlite_dec1_4x4, 64, 2, 8, 0, false, CommonParams, 96)
kernels/m/rrlite_dec2_4x4.hip
+7
View File
@@ -0,0 +1,7 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 2
#include "swin_block.h"
SWIN_MODULE_W(rrlite_dec2_4x4, 96, 4, 8, 0, false, CommonParams, 128)
kernels/m/rrlite_enc1_4x4.hip
+8
View File
@@ -0,0 +1,8 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 1
#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_enc2_4x4.hip
+7
View File
@@ -0,0 +1,7 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 2
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc2_4x4, 96, 4, 8, 128, false, CommonParams, 0)
kernels/m/rrlite_enc3_tube_4x4.hip
+8
View File
@@ -0,0 +1,8 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 2
#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
+897
View File
@@ -0,0 +1,897 @@
// Templated rrlite Swin block (64 tokens = 4 windows of 4x4 per block, NW waves = heads of 32 dims).
// C channels (64 / 96 / 128), NW waves (2 or 4)
// NPM patch-merge output channels (0 = none): 2x2 merge of the block output, 4C -> NPM
// TUBE 88-byte parameter block with the cross-launch flag protocol
// CIN patch-expand input channels (0 = none): decoder input = f16(q8(expand(low-res CIN)) + skip)
#pragma once
#include "swin_common.h"
#ifdef D4R_SWIN_DIAGNOSTICS
extern "C" {
__device__ int d4r_swin_debug_block[2] = {0, 4};
__device__ half_t d4r_swin_debug_values[4 * 64 * 128];
}
template <int C, int AST> __device__ __forceinline__ void swin_snapshot(
const half_t* activation, int stage, int bx, int by, int tid, bool pair_order) {
if (bx == d4r_swin_debug_block[0] && by == d4r_swin_debug_block[1] && tid < 64)
for (int c = 0; c < C; ++c)
d4r_swin_debug_values[(stage * 64 + tid) * 128 + c] =
activation[tid * AST + (pair_order ? gperm(c) : apos(c))];
__syncthreads();
}
#endif
// 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
{
static constexpr int PREB = CIN ? 4 * C * CIN + 8 * C : 0; // patch-expand weights + bias first
static constexpr int NTE = CIN ? 4 * C / NW / 16 : 0;
static constexpr int NPL = C / 32, NT = C / 16, NT8 = C / 8, INNER = 32 * NW, NC = C / 8;
static constexpr int HS = 96 * C + 512;
static constexpr int G1 = PREB, E = PREB + 2 * C + NW * HS, BO = E, G2 = E + 2 * C, M0 = E + 4 * C;
static constexpr int B2 = M0 + 64 * C + 64;
__host__ __device__ static constexpr int head(int h) { return PREB + 2 * C + HS * h; }
__host__ __device__ static constexpr int b1(int c) { return c == 0 ? M0 : M0 + 66 * C + 64 + (64 * C + 64) * (c - 1); }
static constexpr int PM0 = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1);
static constexpr int PMB = PM0 + 4 * C * NPM;
// weight image (slots of 16 K values, see expand_weights)
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
static constexpr int WO = 4 * C * NW;
static constexpr int W1 = WO + 2 * NW * C;
__host__ __device__ static constexpr int w1(int c) { return W1 + 4 * C * c; }
__host__ __device__ static constexpr int w2(int c) { return W1 + 4 * C * c + 2 * C; }
static constexpr int PM = W1 + 4 * C * NC;
__host__ __device__ static constexpr int pm(int g) { return PM + 8 * C * g; }
static constexpr int PE = PM + (NPM / 32) * 8 * C;
__host__ __device__ static constexpr int pe(int w) { return PE + w * (CIN / 32) * 2 * NTE * 16; }
static constexpr int TOTAL = PE + (CIN ? NW * (CIN / 32) * 2 * NTE * 16 : 0);
static constexpr int NDESC = 2 * NW + 1 + 2 * NC + NPM / 32 + (CIN ? NW : 0);
};
// Start (in halves) of token T's row in the output staging buffer: stride C + 4 plus a shift by bits 4 and 5
// of T. The patch merge's lanes read 8 bytes each from 16 tokens that differ in bits 1, 3, 4 and 5 (plus a
// constant); this placement puts those 16 reads on 16 distinct 8-byte bank pairs for C = 64, 96 and 128.
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];
};
template <int C, int NW, int NPM, int CIN> constexpr SwinDescs<C, NW, NPM, CIN> make_descs()
{
using L = SwinLayout<C, NW, NPM, CIN>;
SwinDescs<C, NW, NPM, CIN> r{};
int n = 0;
for (int h = 0; h < NW; ++h)
for (int q = 0; q < 2; ++q)
r.d[n++] = GemmDesc{L::qv(h, q), L::head(h) + (q ? 32 * C : 0), 1024, 512, C / 32, 2};
r.d[n++] = GemmDesc{L::WO, L::head(0) + 64 * C + 512, L::HS, 512, NW, C / 16};
for (int c = 0; c < L::NC; ++c)
{
r.d[n++] = GemmDesc{L::w1(c), L::b1(c), 512, 16 * C, C / 32, 2};
r.d[n++] = GemmDesc{L::w2(c), L::b1(c) + 32 * C + 64, 0, 512, 1, C / 16};
}
for (int g = 0; g < NPM / 32; ++g)
r.d[n++] = GemmDesc{L::pm(g), L::PM0 + g * 128 * C, 512, 64 * C, C / 8, 2};
for (int w = 0; w < (CIN ? NW : 0); ++w)
r.d[n++] = GemmDesc{L::pe(w), w * CIN * 4 * C / NW, 512, 16 * CIN, CIN / 32, L::NTE};
return r;
}
// NH heads (NVIDIA's warps per block) run on NW waves (NW = NH, or 2 NH to halve LDS per wave)
template <int C, int NH, int NW, int NPM, bool TUBE, int CIN>
__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const int bx,
const int by, const int gx)
{
using L = SwinLayout<C, NH, NPM, CIN>;
static_assert(NW % NH == 0 && NW / NH <= 4, "waves per head");
constexpr int NPL = L::NPL, NT = L::NT, INNER = L::INNER;
constexpr int AST = (C > INNER ? C : INNER) + 8;
constexpr int GST = 40, PST = 24;
#ifdef SWIN_EXACT_PV
constexpr int QST = 56;
#else
constexpr int QST = 40;
#endif
#ifndef SWIN_MLP_R
#define SWIN_MLP_R(NW) ((NW) / 2)
#endif
constexpr int R = SWIN_MLP_R(NW); // MLP chunks per round
constexpr int MG = (NT % NW == 0) ? 1 : ((2 * NT) % NW == 0 ? 2 : 4), MT = 4 / MG, NTW = NT * MG / NW;
static_assert(NTW * NW == NT * MG, "column split");
constexpr int PAST = CIN + 8;
constexpr int SCR = NW * 16 * QST * 2 > (CIN ? 16 * PAST * 2 : 0) ? NW * 16 * QST * 2 : 16 * PAST * 2;
constexpr int U_ST4 = R * 64 * GST * 2;
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];
__shared__ half_t R2[64];
half_t* X0 = A; // stages 0-1: block input f16 [64][AST], natural order
half_t* QP = (half_t*)U; // stage 2: per-wave Q / P scratch [NW][16][QST]
half_t* PA = QP; // stage 0: low-res tokens f16 [16][PAST] (apos order)
half_t* G = (half_t*)U; // stage 4: fc1 outputs [R][64][GST]
const int lane = threadIdx.x, wv = threadIdx.z, tid = lane + 32 * wv;
const int l16 = lane & 15, hi = lane >> 4;
const uint8_t* W = p.w;
if constexpr (TUBE)
{
if (tp->in_flags != nullptr)
{
if (lane <= 3 && wv == 0)
{
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
{
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
// bounded spin: a missing producer must never hang the GPU
for (int it = 0; it < (1 << 22); ++it)
{
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
break;
__builtin_amdgcn_s_sleep(1);
}
}
}
__syncthreads();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
}
}
// block input words (encoders: fp8 input, decoders: fp8 skip), issued first so their latency
// overlaps the patch expand; word k of thread tid is idx = tid + 32 NW k
constexpr int NSW = 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 < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
const int X = mirror(8 * bx - p.sx + tok_x(T), p.tw), Y = mirror(8 * by - p.sy + tok_y(T), p.th);
const uint8_t* src = CIN ? p.p32 : p.in;
inw[k] = *(const uint32_t*)(src + ((size_t)((c4 >> 5) * p.th + Y) * p.tw + X) * 32 + (c4 & 31));
}
// 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)
{
const int W2 = p.tw / 2, H2 = p.th / 2;
for (int idx = tid; idx < 16 * CIN / 4; idx += 32 * NW)
{
const int r = idx / (CIN / 4), c4 = 4 * (idx % (CIN / 4));
const int MX = mirror((8 * bx - p.sx) / 2 + (r & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (r >> 2), H2);
const uint32_t bytes =
*(const uint32_t*)(p.in + ((size_t)((c4 >> 5) * H2 + MY) * W2 + MX) * 32 + (c4 & 31));
#pragma unroll
for (int b = 0; b < 4; ++b)
PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
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];
#pragma unroll
for (int j = 0; j < NTEW; ++j)
e[j] = t16_splat(hload(W, 4 * C * CIN + 2 * (16 * (NTEW * wv + j) + l16)));
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
{
const 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)
{
const int t = NTEW * wv + j, blk = t / L::NTE, lt = t % L::NTE;
k32(e[j], a0, bload(w16, L::pe(blk), L::NTE, kc, 0, lt, l16), a1, bload(w16, L::pe(blk), L::NTE, kc, 1, lt, l16));
}
}
#pragma unroll
for (int j = 0; j < NTEW; ++j)
{
const int Gc = 16 * (NTEW * wv + j) + l16, q = Gc / C, ch = ginv(Gc % C);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 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);
// 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
stage_input_words();
__syncthreads();
// residual accumulators of the output projection: f16(x0 + b_o) in the C-fragment layout
// (wave owns pair columns of n tiles ng*NTW.. for m tiles mg*MT..)
const int ng = wv % (NW / MG), mg = wv / (NW / MG);
T16 x[MT][NTW]; // packed f16 accumulators: residual f16(x0 + b_o), live through stage 2
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int j = 0; j < NTW; ++j)
{
const int n = 16 * (ng * NTW + j) + l16;
const half_t bo = hload(W, L::BO + 2 * n);
const int ch = ginv(n);
#pragma unroll
for (int i = 0; i < 8; i += 2)
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + 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
// ------------------------------------------------ stage 1: norm1 (kernel-exact tree)
{
const int g = lane >> 2, t = lane & 3;
#pragma unroll
for (int r = 0; r < 8 / NW; ++r)
{
const int T = (64 / NW) * wv + g + 8 * r;
half_t v[NPL][2][4];
#pragma unroll
for (int pl = 0; pl < NPL; ++pl)
#pragma unroll
for (int ww = 0; ww < 2; ++ww)
#pragma unroll
for (int b = 0; b < 4; ++b)
v[pl][ww][b] = X0[T * AST + 32 * pl + 4 * (t + 4 * ww) + b];
// the 4 lanes of a token read the whole row before any of them writes h1 over it
wave_sync();
half_t ws[2][2];
#pragma unroll
for (int ww = 0; ww < 2; ++ww)
{
half_t pp[NPL][2];
#pragma unroll
for (int pl = 0; pl < NPL; ++pl)
#pragma unroll
for (int e = 0; e < 2; ++e)
{
half_t s0 = v[pl][ww][e] * v[pl][ww][e];
half_t s1 = v[pl][ww][2 + e] * v[pl][ww][2 + e];
pp[pl][e] = s0 + s1;
}
#pragma unroll
for (int e = 0; e < 2; ++e)
{
if constexpr (NPL == 4)
ws[ww][e] = (pp[0][e] + pp[1][e]) + (pp[2][e] + pp[3][e]);
else if constexpr (NPL == 3)
ws[ww][e] = (pp[0][e] + pp[1][e]) + pp[2][e];
else
ws[ww][e] = pp[0][e] + pp[1][e];
}
}
half_t a0 = ws[0][0] + ws[1][0], a1 = ws[0][1] + ws[1][1];
a0 = a0 + hxor<1>(a0);
a1 = a1 + hxor<1>(a1);
a0 = a0 + hxor<2>(a0);
a1 = a1 + hxor<2>(a1);
const half_t rn = rsqrt16(a0 + a1);
const hv2 rn2 = {rn, rn};
#pragma unroll
for (int pl = 0; pl < NPL; ++pl)
#pragma unroll
for (int ww = 0; ww < 2; ++ww)
#pragma unroll
for (int b = 0; b < 4; b += 2)
{
const int c = 32 * pl + 4 * (t + 4 * ww) + b;
// gperm(c + 1) = gperm(c) + 1 for even c: one 32-bit gamma load per pair
const hv2 gm = __builtin_bit_cast(hv2, *(const uint32_t*)(W + L::G1 + 2 * gperm(c)));
const hv2 hq = q8x2((hv2){v[pl][ww][b], v[pl][ww][b + 1]} * (hv2)(rn2 * gm));
// apos(c + 1) = apos(c) + 1 for even c
*(hv2*)(A + T * AST + apos(c)) = hq;
}
}
}
__syncthreads();
// ------------------------------------------------ stage 2: Q/V projection + attention (wave = head)
#ifdef D4R_SWIN_DIAGNOSTICS
swin_snapshot<C, AST>(A, 0, bx, by, tid, false);
#endif
{
constexpr int MW = 4 / (NW / NH); // windows per wave
const int h = NW == NH ? wv : wv % NH, m0 = NW == NH ? 0 : MW * (wv / NH);
T16 acc[MW][4]; // [window][Q nt0, Q nt1, V nt0, V nt1]
#pragma unroll
for (int m = 0; m < MW; ++m)
#pragma unroll
for (int q = 0; q < 4; ++q)
acc[m][q] = t16_splat((half_t)0.0f);
#pragma unroll
for (int kc = 0; kc < C / 32; ++kc)
{
gop b[2][4];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
for (int q = 0; q < 4; ++q)
b[s][q] = bload(w16, L::qv(h, q >> 1), 2, kc, s, q & 1, l16);
#pragma unroll
for (int m = 0; m < MW; ++m)
{
const 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]);
}
}
// pack: Q as q8 pairs, V as f16 pairs (all that attention needs), halving the live registers
hv2 qp[MW][2][4], vp[MW][2][4];
#pragma unroll
for (int m = 0; m < MW; ++m)
#pragma unroll
for (int q = 0; q < 2; ++q)
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
qp[m][q][i / 2] = q8x2(acc[m][q].pair(i / 2));
vp[m][q][i / 2] = acc[m][2 + q].pair(i / 2);
}
__syncthreads(); // every wave is done reading h1; O may overwrite A
half_t* Qs = QP + wv * 16 * QST;
const int bbase = L::head(h) + 64 * C;
T16 bias;
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int ri = 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))));
}
const half_t cs = f16(0.0972222164273262f / 5.656854152679443f);
const half_t cl = (half_t)0.55615234375f;
const half_t p1 = f16(0.92730712890625f), p0 = (half_t)1.375f;
#pragma unroll
for (int m = 0; m < MW; ++m)
{
#pragma unroll
for (int q = 0; q < 2; ++q)
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
Qs[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 gop a0 = lda(Qs + l16 * QST), a1 = lda(Qs + l16 * QST + 16);
T16 S = bias;
k32(S, a0, a0, a1, a1);
// 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;
const uint32_t sel_a = oddc ? 0x01000504u : 0x05040100u, sel_b = oddc ? 0x03020706u : 0x07060302u;
const uint32_t sel_r = oddc ? 0x07060302u : 0x05040100u;
#pragma unroll
for (int k = 0; k < 4; ++k)
{
hv2 tv = S.pair(k) * (hv2){cs, cs};
tv = __builtin_elementwise_max(__builtin_elementwise_min(tv, (hv2){cl, cl}), (hv2){-cl, -cl});
const hv2 e1 = __builtin_elementwise_fma(tv, -tv, (hv2){p1, p1});
const hv2 poly = __builtin_elementwise_fma(tv, e1, (hv2){p0, p0});
const uint32_t mine = __builtin_bit_cast(uint32_t, poly);
const uint32_t other = xor_lane<1>(mine);
const uint32_t ua = (__builtin_amdgcn_perm(other, mine, sel_a) << 5) + 0x7FF88000u;
const uint32_t ub = (__builtin_amdgcn_perm(other, mine, sel_b) << 5) + 0x7FF88000u;
wg[k] = __builtin_bit_cast(hv2, __builtin_amdgcn_perm(ub, ua, sel_r));
}
wave_sync(); // Q scratch is dead: P overwrites it
#pragma unroll
for (int k = 0; k < 4; ++k)
{
auto xr = [](hv2 v, auto mask) {
return __builtin_bit_cast(hv2, xor_lane<decltype(mask)::value>(__builtin_bit_cast(uint32_t, v)));
};
const hv2 v1 = wg[k] + xr(wg[k], 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[mrow(2 * k) * PST + l16] = pw[0];
Qs[mrow(2 * k + 1) * PST + l16] = pw[1];
}
wave_sync();
#ifdef SWIN_EXACT_PV
half_t* Vs = Qs + 16 * PST;
for (int nt = 0; nt < 2; ++nt)
for (int i = 0; i < 8; ++i)
Vs[mrow(i) * 32 + 16 * nt + l16] = vp[m][nt][i / 2][i & 1];
wave_sync();
#endif
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, 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;
#pragma unroll
for (int k = 0; k < 4; ++k)
{
const uint32_t own = __builtin_bit_cast(uint32_t, vp[m][nt][k]);
const uint32_t other = xor_lane<16>(own);
vw[2 * k] = __builtin_amdgcn_perm(other, own, s_lo); // rows 4k, 4k+1
vw[2 * k + 1] = __builtin_amdgcn_perm(other, own, s_hi); // rows 4k+2, 4k+3
}
const h16 vb = __builtin_bit_cast(h16, vw);
#endif
f8v o = wmma(pa, vb, splat(0.0f));
#ifdef SWIN_EXACT_PV
// FP16 P/V products need more than 24 significant bits in a
// cancellation-heavy dot. An f32 WMMA can round to the wrong
// FP16 neighbour at a midpoint, then flip the FP8 output code.
// This baseline sums the 16 FP16 products before one FP16 round.
for (int i = 0; i < 8; ++i) {
double dot = 0.0;
for (int k = 0; k < 16; ++k)
dot += (double)Qs[mrow(i) * PST + k] * (double)Vs[k * 32 + 16 * nt + l16];
o[i] = (float)(half_t)dot;
}
#endif
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) + mrow(i)) * AST + ocol] = oq[0];
A[(16 * (m0 + m) + mrow(i + 1)) * AST + ocol] = oq[1];
}
}
wave_sync();
}
}
__syncthreads();
// ------------------------------------------------ stage 3: output projection + residual
#ifdef D4R_SWIN_DIAGNOSTICS
swin_snapshot<INNER, AST>(A, 3, bx, by, tid, false);
#endif
#pragma unroll
for (int kc = 0; kc < NH; ++kc)
{
gop b[2][NTW];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
for (int j = 0; j < NTW; ++j)
b[s][j] = bload(w16, L::WO, NT, kc, s, ng * NTW + j, l16);
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
{
const 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]);
}
}
__syncthreads(); // X0b and O are dead
// ------------------------------------------------ norm2 (kernel-exact tree), x1 staged in A
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int j = 0; j < NTW; ++j)
#pragma unroll
for (int i = 0; i < 8; ++i)
A[(16 * (mg * MT + mm) + mrow(i)) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
__syncthreads();
if (tid < 64)
{
// the row as 128-bit LDS loads: element 8 j + c of n8 tile j
#ifdef D4R_SWIN_DIAGNOSTICS
if (bx == d4r_swin_debug_block[0] && by == d4r_swin_debug_block[1])
for (int c = 0; c < C; ++c)
d4r_swin_debug_values[(64 + tid) * 128 + c] = A[tid * AST + gperm(c)];
#endif
h8 row[L::NT8];
#pragma unroll
for (int j = 0; j < L::NT8; ++j)
row[j] = *(const h8*)(A + tid * AST + 8 * j);
half_t L8[8];
#pragma unroll
for (int c = 0; c < 8; ++c)
{
half_t sq[L::NT8];
#pragma unroll
for (int j = 0; j < L::NT8; ++j)
sq[j] = row[j][c] * row[j][c];
half_t even, odd;
if constexpr (L::NT8 == 16)
{
even = ((sq[0] + sq[2]) + (sq[4] + sq[6])) + ((sq[8] + sq[10]) + (sq[12] + sq[14]));
odd = ((sq[1] + sq[3]) + (sq[5] + sq[7])) + ((sq[9] + sq[11]) + (sq[13] + sq[15]));
}
else if constexpr (L::NT8 == 12)
{
even = ((sq[0] + sq[2]) + (sq[4] + sq[6])) + (sq[8] + sq[10]);
odd = ((sq[1] + sq[3]) + (sq[5] + sq[7])) + (sq[9] + sq[11]);
}
else
{
even = (sq[0] + sq[2]) + (sq[4] + sq[6]);
odd = (sq[1] + sq[3]) + (sq[5] + sq[7]);
}
L8[c] = even + odd;
}
const half_t b0 = (L8[0] + L8[2]) + (L8[4] + L8[6]);
const half_t b1 = (L8[1] + L8[3]) + (L8[5] + L8[7]);
R2[tid] = rsqrt16(b0 + b1);
}
__syncthreads();
#pragma unroll
for (int j = 0; j < NTW; ++j)
{
const int n = 16 * (ng * NTW + j) + l16;
const half_t g2 = hload(W, L::G2 + 2 * n), b2 = hload(W, L::B2 + 2 * n);
const int col = apos(ginv(n));
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
const int row = 16 * (mg * MT + mm) + 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[row2]};
const hv2 hq = q8x2(v * (hv2)(r2 * (hv2){g2, g2}));
A[row * AST + col] = hq[0];
A[row2 * AST + col] = hq[1];
x[mm][j].set_pair(i / 2, v + (hv2){b2, b2});
}
}
__syncthreads();
// ------------------------------------------------ stage 4: MLP, R chunks per round
#ifdef D4R_SWIN_DIAGNOSTICS
swin_snapshot<C, AST>(A, 2, bx, by, tid, false);
#endif
{
constexpr int WPC = NW / R, MPW = 4 / WPC; // waves per chunk, fc1 m tiles per wave
const int cw = wv / WPC, mb = MPW * (wv % WPC);
#pragma unroll 1
for (int r = 0; r < L::NC / R; ++r)
{
const int c = R * r + cw;
const int B1 = L::b1(c);
T16 z[MPW][2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const half_t b1 = hload(W, B1 + 32 * C + 2 * (16 * nt + l16));
#pragma unroll
for (int mm = 0; mm < MPW; ++mm)
z[mm][nt] = t16_splat(b1);
}
#pragma unroll
for (int kc = 0; kc < C / 32; ++kc)
{
gop b[2][2];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
b[s][nt] = bload(w16, L::w1(c), 2, kc, s, nt, l16);
#pragma unroll
for (int mm = 0; mm < MPW; ++mm)
{
const 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]);
}
}
half_t* Gc = G + cw * 64 * GST;
#pragma unroll
for (int mm = 0; mm < MPW; ++mm)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const int col = apos(ginv(16 * nt + l16));
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
const hv2 zv = z[mm][nt].pair(i / 2);
const hv2 cz = __builtin_elementwise_max(__builtin_elementwise_min(zv, (hv2){2.0f16, 2.0f16}),
(hv2){-2.0f16, -2.0f16});
const hv2 az = __builtin_bit_cast(hv2, (u16x2)(__builtin_bit_cast(u16x2, cz) & (unsigned short)0x7fff));
const hv2 inner = (hv2){f16(0.41216981f), f16(0.41216981f)} -
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
const hv2 gl = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
const hv2 gq = q8x2(gl);
Gc[(16 * (mb + mm) + mrow(i)) * GST + col] = gq[0];
Gc[(16 * (mb + mm) + mrow(i + 1)) * GST + col] = gq[1];
}
}
__syncthreads();
#pragma unroll
for (int cc = 0; cc < R; ++cc)
{
const int c2 = R * r + cc;
gop b[2][NTW];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
for (int j = 0; j < NTW; ++j)
b[s][j] = bload(w16, L::w2(c2), NT, 0, s, ng * NTW + j, l16);
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
{
const 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]);
}
}
__syncthreads();
}
}
// ------------------------------------------------ output: f16 staged in LDS, 16-value segments encoded by a loop
// 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)
{
const int ch = ginv(16 * (ng * NTW + j) + l16);
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int i = 0; i < 8; ++i)
XO[xo_row<C>(16 * (mg * MT + mm) + mrow(i)) + ch] = x[mm][j].get(i);
}
__syncthreads();
#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)
{
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]};
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
*(uint4*)(p.out + ((size_t)((c16 >> 5) * p.th + Y) * p.tw + X) * 32 + (c16 & 31)) = v;
}
if constexpr (NPM > 0)
__syncthreads();
// ------------------------------------------------ patch merge: 2x2 tokens (4C) -> NPM, half resolution
if constexpr (NPM > 0)
{
// merge tiles split evenly over the largest wave count dividing them (remaining waves idle)
constexpr int PT = NPM / 16 ? NPM / 16 : 1;
constexpr int PMW = PT % NW == 0 ? NW : (PT % 4 == 0 && NW >= 4 ? 4 : (PT % 3 == 0 && NW >= 3 ? 3 : (PT % 2 == 0 && NW >= 2 ? 2 : 1)));
constexpr int NTP = PT / PMW;
constexpr bool PM_EVEN = true;
if (wv < PMW)
{
T16 y[NTP];
#pragma unroll
for (int j = 0; j < NTP; ++j)
y[j] = t16_splat(hload(W, L::PMB + 2 * (16 * (PM_EVEN ? wv * NTP + j : (wv * NTP + j) % (NPM / 16)) + l16)));
const int mx = l16 & 3, my = l16 >> 2;
#pragma unroll
for (int kc = 0; kc < C / 8; ++kc)
{
gop as[2];
#pragma unroll
for (int s = 0; s < 2; ++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 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);
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 * rr + b] = v[b];
}
as[s] = as_gop(a);
}
#pragma unroll
for (int j = 0; j < NTP; ++j)
{
// out-of-range tiles (uneven split) recompute a valid one and are not stored
const int t16 = PM_EVEN ? wv * NTP + j : (wv * NTP + j) % (NPM / 16);
k32(y[j], as[0], bload(w16, L::pm(t16 >> 1), 2, kc, 0, t16 & 1, l16), as[1],
bload(w16, L::pm(t16 >> 1), 2, kc, 1, t16 & 1, l16));
}
}
const int W2 = p.tw / 2, H2 = p.th / 2;
#pragma unroll
for (int j = 0; j < NTP; ++j)
{
if (!PM_EVEN && wv * NTP + j >= NPM / 16)
continue;
const int ch = ginv(16 * (wv * NTP + j) + l16);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 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;
p.p48[((size_t)((ch >> 5) * H2 + MY) * W2 + MX) * 32 + (ch & 31)] =
(uint8_t)enc8(
#ifdef SWIN_EXACT
// The audited 310.9 patch-merge path bounds the f16
// result to +/- 2*pi before e4m3 quantisation.
__builtin_elementwise_max(__builtin_elementwise_min(y[j].get(i), (half_t)6.283185307179586f),
(half_t)-6.283185307179586f)
#else
y[j].get(i)
#endif
);
}
}
}
}
if constexpr (TUBE)
{
if (tp->out_flags != nullptr)
{
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "agent");
__syncthreads();
if (tid == 0)
__hip_atomic_store(tp->out_flags + (size_t)(by * gx + bx) * 32, (uint8_t)1, __ATOMIC_RELEASE,
__HIP_MEMORY_SCOPE_AGENT);
}
}
}
#if defined(SWIN_WAVES_PER_EU)
#define SWIN_VGPR_ATTR __attribute__((amdgpu_waves_per_eu(SWIN_WAVES_PER_EU)))
#elif defined(SWIN_MAX_VGPR)
#define SWIN_VGPR_ATTR __attribute__((amdgpu_num_vgpr(SWIN_MAX_VGPR)))
#else
#define SWIN_VGPR_ATTR
#endif
#ifndef SWIN_PERSIST
#define SWIN_PERSIST 0 // persistent workgroups in x (0: one workgroup per original block)
#endif
#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
#ifndef SWIN_PREP_KEY_AT
#define SWIN_PREP_KEY_AT 1
#endif
#ifndef SWIN_PREP_IDENTITY_AT
#define SWIN_PREP_IDENTITY_AT 0
#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>; \
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(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, 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, w16, b % gx, b / gx, gx); \
__syncthreads(); \
} \
}
#if defined(SWIN_A8) && defined(SWIN_TOKEN_LANES)
#include "swin_tok.h"
#undef SWIN_MODULE_W
#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN) SWIN_TOK_MODULE(NAME, C, NH, NPM, TUBE, PARAMS, CIN)
#endif
kernels/m/swin_common.h
+573
View File
@@ -0,0 +1,573 @@
// 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"
#ifdef D4R_ACCURACY
#define SWIN_EXACT
#endif
#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
// Experimental RDNA4 byte-activation path. Keep disabled unless the caller
// explicitly builds native FP8 WMMA for gfx12.
#if defined(D4R_FP8_WMMA) && defined(__GFX12__) && !defined(SWIN_NO_A8)
#define SWIN_A8 1
#endif
#if defined(SWIN_A8) && defined(SWIN_EXACT) && !defined(K32_NO_MIX) && !defined(K32_MIX)
#define K32_MIX
#endif
#pragma clang fp contract(off)
typedef _Float16 half_t;
typedef _Float16 h16 __attribute__((ext_vector_type(16)));
typedef _Float16 h8 __attribute__((ext_vector_type(8)));
typedef float f8v __attribute__((ext_vector_type(8)));
struct CommonParams
{
const uint8_t* w; // 0
int sx, sy; // 8
int tw, th; // 16
const uint8_t* in; // 24
const uint8_t* p32; // 32 (decoders: skip input)
uint8_t* out; // 40
uint8_t* p48; // 48 (encoders with patch merge: merged output)
};
static_assert(sizeof(CommonParams) == 56, "param block");
struct TubeParams
{
CommonParams c;
int r69, r70;
int r71, r72;
const uint8_t* in_flags;
uint8_t* out_flags;
};
static_assert(sizeof(TubeParams) == 88, "param block");
// ---------------------------------------------------------------- numerics
__device__ __forceinline__ half_t e4m3_to_half(uint32_t code)
{
uint16_t bits = (uint16_t)(((code & 0x7fu) << 7) | ((code & 0x80u) << 8));
return __builtin_bit_cast(half_t, bits) * (half_t)256.0f;
}
__device__ __forceinline__ uint32_t half_to_e4m3(half_t h)
{
uint32_t half = __builtin_bit_cast(uint16_t, h);
uint32_t magnitude = half & 0x7fffu;
uint32_t normal = (magnitude - 0x2000u + 0x3fu + ((magnitude >> 7) & 1u)) >> 7;
uint32_t exponent = magnitude >> 10;
uint32_t significand = (magnitude & 0x3ffu) | (exponent != 0 ? 0x400u : 0u);
uint32_t clamped = exponent < 1u ? 1u : (exponent > 8u ? 8u : exponent);
uint32_t shift = 16u - clamped;
uint32_t subnormal = (significand + (1u << (shift - 1u)) - 1u + ((significand >> shift) & 1u)) >> shift;
uint32_t code = magnitude >= 0x2400u ? normal : subnormal;
code = magnitude > 0x5f00u ? 0x7eu : code;
code = magnitude > 0x7c00u ? 0x7fu : code;
return ((half >> 8) & 0x80u) | code;
}
// e4m3 RNE satfinite round trip evaluated in the f16 domain (exhaustively equal to decode(encode(h)))
__device__ __forceinline__ half_t q8_exact(half_t h);
__device__ __forceinline__ half_t q8(half_t h)
{
#ifdef SWIN_NO_Q8
return h; // fast mode: intermediate activations stay f16 (stored outputs are still FP8)
#endif
return q8_exact(h);
}
__device__ __forceinline__ half_t q8_exact(half_t h)
{
const uint32_t u = __builtin_bit_cast(uint16_t, h);
const uint32_t m = u & 0x7fffu;
uint32_t r = (m + 0x3fu + ((m >> 7) & 1u)) & 0x7f80u;
r = r < 0x5f00u ? r : 0x5f00u;
const half_t mag = __builtin_bit_cast(half_t, (uint16_t)m);
const uint32_t sub = __builtin_bit_cast(uint16_t, (half_t)((half_t)(mag + (half_t)2.0f) - (half_t)2.0f));
r = m < 0x2400u ? sub : r;
r = m > 0x7c00u ? m : r;
return __builtin_bit_cast(half_t, (uint16_t)((u & 0x8000u) | r));
}
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
typedef unsigned short u16x2 __attribute__((ext_vector_type(2)));
typedef short i16x2 __attribute__((ext_vector_type(2)));
// q8_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)
{
const u16x2 u = __builtin_bit_cast(u16x2, h);
const u16x2 m = u & (unsigned short)0x7fff;
u16x2 r = (m + (unsigned short)0x3f + ((m >> 7) & (unsigned short)1)) & (unsigned short)0x7f80;
r = __builtin_elementwise_min(r, (u16x2){0x5f00, 0x5f00});
const hv2 mag = __builtin_bit_cast(hv2, m);
const hv2 two = {(half_t)2.0f, (half_t)2.0f};
const u16x2 sub = __builtin_bit_cast(u16x2, (hv2)((hv2)(mag + two) - two));
const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
r = (sub & lt) | (r & ~lt);
const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
r = (m & gt) | (r & ~gt);
return __builtin_bit_cast(hv2, (u16x2)(r | (u & (unsigned short)0x8000)));
}
__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);
}
// e4m3 byte codes of four f16 values; caller provides exact e4m3 values.
#ifdef SWIN_A8
__device__ __forceinline__ uint32_t codes4(hv2 a, hv2 b)
{
const hv2 lim = {(half_t)448.0f, (half_t)448.0f};
a = __builtin_elementwise_maximum(__builtin_elementwise_minimum(a, lim), -lim);
b = __builtin_elementwise_maximum(__builtin_elementwise_minimum(b, lim), -lim);
const uint32_t w = (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32((float)a[0], (float)a[1], 0, false);
return (uint32_t)__builtin_amdgcn_cvt_pk_fp8_f32((float)b[0], (float)b[1], (int)w, true);
}
#endif
// 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)
{
const uint32_t qb = __builtin_bit_cast(uint16_t, q8_exact(h));
const uint32_t m = qb & 0x7fffu;
const half_t mag = __builtin_bit_cast(half_t, (uint16_t)m);
uint32_t code = m >= 0x2400u ? (m >> 7) - 64u : (uint32_t)(uint16_t)(mag * (half_t)512.0f);
code = m > 0x7c00u ? 0x7fu : code;
return ((qb >> 8) & 0x80u) | code;
}
__device__ __forceinline__ half_t f16(float v)
{
return (half_t)v;
}
__device__ __forceinline__ half_t rsqrt16(half_t s)
{
return f16(1.0f / __builtin_sqrtf((float)(s + (half_t)(1.0f / 8192.0f))));
}
__device__ __forceinline__ half_t recip16(half_t s)
{
return f16(1.0f / (float)s);
}
__device__ __forceinline__ hv2 merge_clamp(hv2 v)
{
const hv2 lim = {(half_t)6.2831855f, (half_t)6.2831855f};
return __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, -lim), lim);
}
__device__ __forceinline__ half_t merge_clamp(half_t v)
{
return merge_clamp((hv2){v, v})[0];
}
// natural channel c -> pair-order column, and inverse (per 32-channel group)
__device__ __forceinline__ int gperm(int c)
{
return 32 * (c >> 5) + 2 * ((c & 31) >> 2) + (c & 1) + 16 * ((c & 3) >> 1);
}
__device__ __forceinline__ int ginv(int n)
{
int m = n & 31, e = m >> 4, a = (m & 15) >> 1, b = m & 1;
return (n & ~31) + 4 * a + 2 * e + b;
}
// WMMA k16 step s of a k32 chunk: operand slot i holds logical k = kslot(s, i) (matches the weight
// memory runs). Activations live in LDS at apos(k), so an A operand is 16 contiguous halves.
__device__ __forceinline__ int kslot(int s, int i)
{
int a = 2 * s + (i >> 3), r = i & 7;
return (r < 4) ? (4 * a + r) : (16 + 4 * a + (r - 4));
}
__device__ __forceinline__ int apos(int k)
{
return (k & ~31) + 8 * ((k & 15) >> 2) + 4 * ((k >> 4) & 1) + (k & 3);
}
__device__ __forceinline__ int woff(int base, int Ks, int Ns, int k, int n)
{
int j = n >> 3;
return base + Ks * (k >> 5) + Ns * (j >> 1) + 64 * (n & 7) + 16 * ((k & 15) >> 2) + 8 * (j & 1) +
4 * ((k & 31) >> 4) + (k & 3);
}
// ---------------------------------------------------------------- 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)
{
#if (defined(D4R_FP8_HW_PACK) || defined(SWIN_FP8_HWCVT)) && defined(__GFX12__)
// Inputs already contain exact e4m3 values. gfx12's native packed convert
// avoids the software subnormal/sign/exponent reconstruction below. Keep
// opt-in until all finite encodings and layer/network outputs pass replay.
gop result;
#pragma unroll
for (int word = 0; word < 2; ++word) {
int packed = __builtin_amdgcn_cvt_pk_fp8_f32((float)v[4 * word], (float)v[4 * word + 1], 0, false);
packed = __builtin_amdgcn_cvt_pk_fp8_f32((float)v[4 * word + 2], (float)v[4 * word + 3], packed, true);
result[word] = (uint32_t)packed;
}
return result;
#else
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)};
#endif
}
__device__ __forceinline__ uint32_t input_word(const uint32_t* p)
{
#if defined(SWIN_NT_INPUT)
return __builtin_nontemporal_load(p);
#else
return *p;
#endif
}
// 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)
{
#pragma unroll
for (int i = 0; i < 8; ++i)
v[i] = (float)(half_t)v[i];
return v;
}
// A C-fragment tile: element i is row 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
{
#ifdef SWIN_F32ACC
f8v f;
__device__ hv2 pair(int k) const { return (hv2){(half_t)f[2 * k], (half_t)f[2 * k + 1]}; }
__device__ void set_pair(int k, hv2 p) { f[2 * k] = (float)p[0]; f[2 * k + 1] = (float)p[1]; }
__device__ half_t get(int i) const { return (half_t)f[i]; }
__device__ void set(int i, half_t h) { f[i] = (float)h; }
#else
hv2 v[4];
__device__ hv2 pair(int k) const { return v[k]; }
__device__ void set_pair(int k, hv2 p) { v[k] = p; }
__device__ half_t get(int i) const { return v[i / 2][i & 1]; }
__device__ void set(int i, half_t h) { v[i / 2][i & 1] = h; }
#endif
};
__device__ __forceinline__ T16 t16_splat(half_t h)
{
T16 t;
#pragma unroll
for (int k = 0; k < 4; ++k)
t.set_pair(k, (hv2){h, h});
return t;
}
#ifndef MIX_PAD
#define MIX_PAD ""
#endif
// acc = f16(P + acc), rounded once (v_fma_mix: verified == f16(exact sum) on 8.4M cases)
__device__ __forceinline__ hv2 mix2(float p0, float p1, hv2 c)
{
hv2 r = c;
asm(MIX_PAD "v_fma_mixlo_f16 %0, %1, 1.0, %2 op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p0), "v"(c));
asm("v_fma_mixhi_f16 %0, %1, 1.0, %2 op_sel:[0,0,1] op_sel_hi:[0,0,1]" : "+v"(r) : "v"(p1), "v"(c));
return r;
}
__device__ __forceinline__ void mixacc(T16& acc, f8v P)
{
#pragma unroll
for (int k = 0; k < 4; ++k)
acc.set_pair(k, mix2(P[2 * k], P[2 * k + 1], acc.pair(k)));
}
#ifdef K32_CHECK
__device__ unsigned d4r_k32_bad[4 + 64 * 4];
#endif
// f16 accumulate of one k32 step (NVIDIA m16n8k32 semantics: products summed, one rounding with C)
__device__ __forceinline__ void k32(T16& acc, gop a0, gop b0, gop a1, gop b1)
{
#if defined(SWIN_F32ACC)
acc.f = wmma8(a1, b1, wmma8(a0, b0, acc.f));
#elif defined(K32_MIX)
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 = wmma8(a1, b1, wmma8(a0, b0, c));
#pragma unroll
for (int i = 0; i < 8; ++i)
acc.set(i, (half_t)c[i]);
#endif
}
__device__ __forceinline__ f8v splat(float v)
{
return (f8v){v, v, v, v, v, v, v, v};
}
// lane xor shuffles without LDS: DPP within rows of 16, permlanex16 across the two rows
template <int MASK> __device__ __forceinline__ uint32_t xor_lane(uint32_t v)
{
if constexpr (MASK == 1)
return __builtin_amdgcn_update_dpp(0, (int)v, 0xB1, 0xf, 0xf, false);
else if constexpr (MASK == 2)
return __builtin_amdgcn_update_dpp(0, (int)v, 0x4E, 0xf, 0xf, false);
else if constexpr (MASK == 4)
return __builtin_amdgcn_update_dpp(0, (int)v, 0x164, 0xf, 0xf, false);
else if constexpr (MASK == 8)
return __builtin_amdgcn_update_dpp(0, (int)v, 0x168, 0xf, 0xf, false);
else
return __builtin_amdgcn_permlanex16((int)v, (int)v, 0x76543210, 0xfedcba98, false, false);
}
template <int MASK> __device__ __forceinline__ half_t hxor(half_t v)
{
return __builtin_bit_cast(half_t, (uint16_t)xor_lane<MASK>(__builtin_bit_cast(uint16_t, v)));
}
__device__ __forceinline__ int mirror(int v, int n)
{
v = v < 0 ? -v : v;
int o = 2 * n - 2 - v;
return v < o ? v : o;
}
__device__ __forceinline__ half_t hload(const uint8_t* w, int off)
{
return __builtin_bit_cast(half_t, *(const uint16_t*)(w + off));
}
__device__ __forceinline__ void wave_sync()
{
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
__builtin_amdgcn_wave_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "wavefront");
}
// Token T (0..63) of a block: 8x8 tokens, windows of 4x4 (T >> 4)
__device__ __forceinline__ int tok_x(int T)
{
return (T & 3) + 4 * ((T >> 4) & 1);
}
__device__ __forceinline__ int tok_y(int T)
{
return ((T >> 2) & 3) + 4 * (T >> 5);
}
// ---------------------------------------------------------------- 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 slots
int base, Ks, Ns, KC, NT;
};
// 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)
{
#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, wslot* w16, const GemmDesc* descs, int ndesc, int idx)
{
int d = 0;
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
++d;
const GemmDesc g = descs[d];
const int local = idx - g.dst;
if (local >= g.KC * 2 * g.NT * 16)
return;
const int l16 = local & 15, rest = local >> 4;
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
const int n = 16 * nt + l16;
#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/m/swin_tok.h
+801
View File
@@ -0,0 +1,801 @@
// Token-lane rrlite Swin block for RDNA4 native FP8 (SWIN_A8), selected with SWIN_TOKEN_LANES.
//
// swin_block.h computes every GEMM as activations x weights: the result has lane = channel column and
// VGPR = token row, so each result is transposed through LDS (byte stores, barriers) before it can be the
// next GEMM's operand. Here the weights are the first WMMA operand and the activations the second, which
// gives the transposed result: lane = token, VGPR i = row 8 half + i of the output. A gfx12 operand is this
// half's 8 K values, so an encoded result IS the next GEMM's operand and stays in registers. The products
// and the K slot of every product are the same as in swin_block.h, so the values are the same (operand
// roles swapped: bit-identical on an RX 9070 XT for the FP8 and the f16 WMMA).
//
// One wave owns one 4x4 window (16 tokens) for the whole block: norms, attention and the MLP need no LDS
// and no barriers. Only the patch expand (decoders) and the patch merge (encoders) cross windows and go
// through an e4m3 byte image in LDS.
//
// A half-wave owns the channels whose bit 2 equals the half (dword 4 t + 2 g + half of tile t, g = 0, 1):
// both halves of a k32 operand's K slots are then two of the lane's own code dwords.
//
// SWIN_TOK_NATURAL: a half-wave owns the channels whose bit 3 equals the half (dword 4 t + 2 half + g), and
// the two WMMAs of a k32 step hold K 0..15 and K 16..31 in order: a tile's two code dwords are one operand.
// The products are the same; what changes is which of them share a WMMA and in which K slot. That is the
// order of ZLUDA's lowering of NVIDIA's m16n8k32, so a kernel built this way (with K32_NO_MIX, which also adds
// C inside the first WMMA as that lowering does) accumulates in f32 exactly as ZLUDA's compile of the PTX
// does, including the rare sums that are not exact in f32 (../tex/enc0l_tail.hip).
#pragma once
enum : int { TOK_ID, TOK_PI, TOK_X, TOK_NU, TOK_PE };
// A GEMM's weight tiles as in GemmDesc (slot dst + ((kc * 2 + s) * NT + nt) * 16 + r holds the 16 K slots of
// row r of tile nt), with the raw weight column of (nt, r) chosen by perm:
// TOK_ID 16 nt + r (Q: the attention's K order)
// TOK_PI gperm(kslot(nt, r)) (V, fc1: row r is K slot r of step nt of the GEMM that consumes it)
// TOK_X gperm(tok_xnat(nt, r)) (output projection, fc2: the x layout)
// TOK_NU gperm(16 nt + r) (patch merge: natural channel order, 8 contiguous output bytes)
// TOK_PE col0 + gperm(16 nt + r) over weight blocks of blkcols columns / blkbytes bytes (patch expand)
struct TokDesc
{
int dst, base, Ks, Ns, KC, NT, perm, col0, blkcols, blkbytes;
};
enum : int { TC_X, TC_PI, TC_NU, TC_ATT };
// Per-channel f16 vectors as 8-value slots: slot dst + 2 j + half holds rows 8 half + i of tile j in the
// layout of kind (TC_ATT: j = query token, row = key token, the transposed attention bias)
struct TokConst
{
int dst, kind, base;
};
// x layout: natural channel of row r (= 8 half + i) of tile t
__device__ __forceinline__ int tok_xnat(int t, int r)
{
#ifdef SWIN_TOK_NATURAL
return 16 * t + r;
#else
return 16 * t + 8 * ((r >> 2) & 1) + 4 * (r >> 3) + (r & 3);
#endif
}
// logical K (of a 32-value chunk) in operand slot i of WMMA s of a k32 step
__device__ __forceinline__ int tok_kslot(int s, int i)
{
#ifdef SWIN_TOK_NATURAL
return 16 * s + i;
#else
return kslot(s, i);
#endif
}
// The 4 codes (tile t, g) of half-wave hf within a token's C channel bytes are at tok_half(hf) + tok_word(t, g)
__device__ __forceinline__ int tok_half(int hf)
{
#ifdef SWIN_TOK_NATURAL
return 8 * hf;
#else
return 4 * hf;
#endif
}
__device__ __forceinline__ int tok_word(int t, int g)
{
#ifdef SWIN_TOK_NATURAL
return 16 * t + 4 * g;
#else
return 16 * t + 8 * g;
#endif
}
// first of the four values k (of a merged token's 4C inputs) in K slots 8 half + 4 rr .. + 3 of WMMA s of chunk kc
// (without the half-wave's own offset tok_half(hf))
__device__ __forceinline__ int tok_merge_k(int kc, int s, int rr)
{
#ifdef SWIN_TOK_NATURAL
return 32 * kc + 16 * s + 4 * rr;
#else
return 32 * kc + 16 * rr + 8 * s;
#endif
}
template <int C, int NH, int NPM, int CIN> struct TokLayout
{
using L = SwinLayout<C, NH, NPM, CIN>;
static constexpr int NT = C / 16, NC = C / 8, NTP = NPM / 16;
// weight image (slots)
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
static constexpr int WO = 4 * C * NH;
static constexpr int W1 = WO + 2 * NH * C;
__host__ __device__ static constexpr int w1(int c) { return W1 + 4 * C * c; }
__host__ __device__ static constexpr int w2(int c) { return W1 + 4 * C * c + 2 * C; }
static constexpr int PM = W1 + 4 * C * NC;
static constexpr int PE = PM + NPM * C / 4;
__host__ __device__ static constexpr int pe(int q) { return PE + q * (CIN * C / 16); }
static constexpr int TOTAL = PE + CIN * C / 4;
static constexpr int NDESC = 2 * NH + 1 + 2 * NC + (NPM ? 1 : 0) + (CIN ? 4 : 0);
// constant image (slots of 8 f16)
static constexpr int CBO = 0, CG1 = 2 * NT, CG2 = 4 * NT, CB2 = 6 * NT, CB1 = 8 * NT;
static constexpr int CATT = CB1 + 4 * NC, CPMB = CATT + 32 * NH, CPEB = CPMB + 2 * NTP;
static constexpr int CTOTAL = CPEB + (CIN ? 8 * NT : 0);
static constexpr int NCONST = 4 + NC + NH + (NPM ? 1 : 0) + (CIN ? 4 : 0);
};
template <int C, int NH, int NPM, int CIN> struct TokTables
{
TokDesc d[TokLayout<C, NH, NPM, CIN>::NDESC];
TokConst c[TokLayout<C, NH, NPM, CIN>::NCONST];
};
template <int C, int NH, int NPM, int CIN> constexpr TokTables<C, NH, NPM, CIN> make_tok_tables()
{
using L = SwinLayout<C, NH, NPM, CIN>;
using K = TokLayout<C, NH, NPM, CIN>;
TokTables<C, NH, NPM, CIN> r{};
int n = 0;
for (int h = 0; h < NH; ++h)
{
r.d[n++] = TokDesc{K::qv(h, 0), L::head(h), 1024, 512, C / 32, 2, TOK_ID, 0, 0, 0};
r.d[n++] = TokDesc{K::qv(h, 1), L::head(h) + 32 * C, 1024, 512, C / 32, 2, TOK_PI, 0, 0, 0};
}
r.d[n++] = TokDesc{K::WO, L::head(0) + 64 * C + 512, L::HS, 512, NH, C / 16, TOK_X, 0, 0, 0};
for (int c = 0; c < K::NC; ++c)
{
r.d[n++] = TokDesc{K::w1(c), L::b1(c), 512, 16 * C, C / 32, 2, TOK_PI, 0, 0, 0};
r.d[n++] = TokDesc{K::w2(c), L::b1(c) + 32 * C + 64, 0, 512, 1, C / 16, TOK_X, 0, 0, 0};
}
if (NPM)
r.d[n++] = TokDesc{K::PM, L::PM0, 512, 64 * C, C / 8, NPM / 16, TOK_NU, 0, 0, 0};
for (int q = 0; q < (CIN ? 4 : 0); ++q)
r.d[n++] = TokDesc{K::pe(q), 0, 512, 16 * CIN, CIN / 32, C / 16, TOK_PE, q * C, 4 * C / NH, CIN * 4 * C / NH};
int m = 0;
r.c[m++] = TokConst{K::CBO, TC_X, L::BO};
r.c[m++] = TokConst{K::CG1, TC_X, L::G1};
r.c[m++] = TokConst{K::CG2, TC_X, L::G2};
r.c[m++] = TokConst{K::CB2, TC_X, L::B2};
for (int c = 0; c < K::NC; ++c)
r.c[m++] = TokConst{K::CB1 + 4 * c, TC_PI, L::b1(c) + 32 * C};
for (int h = 0; h < NH; ++h)
r.c[m++] = TokConst{K::CATT + 32 * h, TC_ATT, L::head(h) + 64 * C};
if (NPM)
r.c[m++] = TokConst{K::CPMB, TC_NU, L::PMB};
for (int q = 0; q < (CIN ? 4 : 0); ++q)
r.c[m++] = TokConst{K::CPEB + 2 * K::NT * q, TC_NU, 4 * C * CIN + 2 * q * C};
return r;
}
__device__ __forceinline__ void expand_tok_weights(const uint8_t* w, wslot* w16, const TokDesc* descs, int ndesc, int idx)
{
int d = 0;
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
++d;
const TokDesc g = descs[d];
const int local = idx - g.dst;
if (local >= g.KC * 2 * g.NT * 16)
return;
const int r = local & 15, rest = local >> 4;
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
int n = 16 * nt + r, base = g.base;
if (g.perm == TOK_PI)
n = gperm(tok_kslot(nt, r));
else if (g.perm == TOK_X)
n = gperm(tok_xnat(nt, r));
else if (g.perm == TOK_NU || g.perm == TOK_PE)
n = gperm(16 * nt + r);
if (g.perm == TOK_PE)
{
const int gc = g.col0 + n;
base += (gc / g.blkcols) * g.blkbytes;
n = gc % g.blkcols;
}
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(base, g.Ks, g.Ns, 32 * kc + tok_kslot(s, 4 * j + b), n)] << (8 * b);
v[j] = word;
}
w16[idx] = v;
}
__device__ __forceinline__ void expand_tok_consts(const uint8_t* w, h8* c16, const TokConst* cs, int ncs, int idx)
{
int d = 0;
while (d + 1 < ncs && idx >= cs[d + 1].dst)
++d;
const TokConst g = cs[d];
const int local = idx - g.dst, hf = local & 1, j = local >> 1;
h8 v;
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 8 * hf + i;
int off;
if (g.kind == TC_X)
off = 2 * gperm(tok_xnat(j, r));
else if (g.kind == TC_PI)
off = 2 * gperm(tok_kslot(j, r));
else if (g.kind == TC_NU)
off = 2 * gperm(16 * j + r);
else
{
// the attention bias of (query j, key r)
const int bl = 4 * (j & 7) + ((r & 7) >> 1), word = (j >> 3) + 2 * (r >> 3);
off = 2 * (8 * bl + 2 * word + (r & 1));
}
v[i] = hload(w, g.base + off);
}
c16[idx] = v;
}
// ---------------------------------------------------------------- register helpers
typedef float tok_f2v __attribute__((ext_vector_type(2)));
// two f32 -> f16 (round to nearest even) packed into one register
__device__ __forceinline__ hv2 tok_pk16(float a, float b)
{
return (hv2){(half_t)a, (half_t)b};
}
// two e4m3 codes of a word (bytes 0, 1 or 2, 3) as f16. Equal to e4m3_to_half for every code except the NaN
// codes 0x7f / 0xff, which the conversion decodes as NaN (e4m3_to_half: +-480). The FP8 WMMAs read the
// same bytes with the hardware's decoding, and no encoder here stores those codes for a non-NaN value.
template <bool UPPER> __device__ __forceinline__ hv2 tok_dec2(uint32_t word)
{
const tok_f2v f = __builtin_amdgcn_cvt_pk_f32_fp8((int)word, UPPER);
return tok_pk16(f[0], f[1]);
}
// the same lane of the other half-wave
__device__ __forceinline__ hv2 tok_other(hv2 v)
{
return __builtin_bit_cast(hv2, xor_lane<16>(__builtin_bit_cast(uint32_t, v)));
}
__device__ __forceinline__ T16 tok_t16(h8 c)
{
T16 t;
#pragma unroll
for (int k = 0; k < 4; ++k)
t.set_pair(k, (hv2){c[2 * k], c[2 * k + 1]});
return t;
}
// the two operands (WMMA s = 0, 1) of a k32 step from the code words [tile][g] of the chunk's two tiles
__device__ __forceinline__ gop tok_k32_operand(const uint32_t (*c)[2], int s)
{
#ifdef SWIN_TOK_NATURAL
return (u2v){c[s][0], c[s][1]};
#else
return (u2v){c[0][s], c[1][s]};
#endif
}
// e4m3 codes of a tile's 8 rows: this half's 8 K slots of an operand
__device__ __forceinline__ gop tok_codes(const T16& t)
{
return (u2v){codes4(t.pair(0), t.pair(1)), codes4(t.pair(2), t.pair(3))};
}
// 1 / sqrt(sum of squares + eps) of a token's C values, with the summation tree of norm1 and norm2 of
// swin_block.h (both are this tree in the x layout): per (g, element) the squares of a dword's two pairs,
// the 32-channel groups, the two 16-channel halves of a group, then the half-waves, g and the elements
template <int NT> __device__ __forceinline__ half_t tok_rsq(const hv2 (&v)[NT][4])
{
constexpr int NPL = NT / 2;
hv2 u[2];
#pragma unroll
for (int g = 0; g < 2; ++g)
{
hv2 ws[2];
#pragma unroll
for (int ww = 0; ww < 2; ++ww)
{
hv2 pp[NPL];
#pragma unroll
for (int pl = 0; pl < NPL; ++pl)
{
const hv2 a = v[2 * pl + ww][2 * g], b = v[2 * pl + ww][2 * g + 1];
pp[pl] = (hv2)(a * a) + (hv2)(b * b);
}
if constexpr (NPL == 4)
ws[ww] = (pp[0] + pp[1]) + (pp[2] + pp[3]);
else if constexpr (NPL == 3)
ws[ww] = (pp[0] + pp[1]) + pp[2];
else if constexpr (NPL == 2)
ws[ww] = pp[0] + pp[1];
else
ws[ww] = pp[0];
}
#ifdef SWIN_TOK_NATURAL
u[g] = ws[0] + ws[1];
}
// the tree's next partner (dword ^ 1) is this lane's other g, the one after it (dword ^ 2) the other half
const hv2 a = u[0] + u[1];
const hv2 s = a + tok_other(a);
#else
const hv2 a = ws[0] + ws[1];
u[g] = a + tok_other(a);
}
const hv2 s = u[0] + u[1];
#endif
return rsqrt16(s[0] + s[1]);
}
// The block from its input x0 to the e4m3 codes of its output, for the 16 tokens of one 4x4 window (lane l16 =
// token, both half-waves). v: x0 as f16 pairs, [tile][2 g + pair of the dword]; it is overwritten (x1).
// co: the output codes, [tile][g] (see tok_word). ch / wl: this lane's bases into the constant image (c16 + half)
// and the weight image (row l16 of a tile, this half's 8 K slots). x2 (optional): the output before its e4m3
// encoding, as f16 pairs like v.
template <int C, int NH, int NPM, int CIN>
__device__ __forceinline__ void swin_tok_core(const h8* const ch, const gop* const wl, const int l16, hv2 (&v)[C / 16][4],
uint32_t (&co)[C / 16][2], hv2 (*x2)[4] = nullptr)
{
using K = TokLayout<C, NH, NPM, CIN>;
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8;
auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
// residual of the output projection: f16(x0 + b_o)
hv2 res[NT][4];
#pragma unroll
for (int t = 0; t < NT; ++t)
{
const h8 bo = ch[K::CBO + 2 * t];
#pragma unroll
for (int k = 0; k < 4; ++k)
res[t][k] = v[t][k] + (hv2){bo[2 * k], bo[2 * k + 1]};
}
// ------------------------------------------------ norm1: h1 codes, the K operands of Q and V
gop hop[NKC][2];
{
const half_t rn = tok_rsq<NT>(v);
const hv2 rn2 = {rn, rn};
uint32_t c1[NT][2];
#pragma unroll
for (int t = 0; t < NT; ++t)
{
const h8 g1 = ch[K::CG1 + 2 * t];
#pragma unroll
for (int g = 0; g < 2; ++g)
c1[t][g] = codes4(v[t][2 * g] * (hv2)(rn2 * (hv2){g1[4 * g], g1[4 * g + 1]}),
v[t][2 * g + 1] * (hv2)(rn2 * (hv2){g1[4 * g + 2], g1[4 * g + 3]}));
}
#pragma unroll
for (int kc = 0; kc < NKC; ++kc)
#pragma unroll
for (int s = 0; s < 2; ++s)
hop[kc][s] = tok_k32_operand(&c1[2 * kc], s);
}
// ------------------------------------------------ attention per head: O codes, the K operands of the projection
gop oc[NH][2];
#pragma unroll
for (int hd = 0; hd < NH; ++hd)
{
// Q transposed (lane = token, rows = Q's 32 columns), V as in swin_block.h (lane = column, rows = tokens)
T16 qt[2], vt[2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
qt[nt] = t16_splat((half_t)0.0f);
vt[nt] = t16_splat((half_t)0.0f);
}
#pragma unroll
for (int kc = 0; kc < NKC; ++kc)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
k32(qt[nt], wt(K::qv(hd, 0), 2, kc, 0, nt), hop[kc][0], wt(K::qv(hd, 0), 2, kc, 1, nt), hop[kc][1]);
k32(vt[nt], hop[kc][0], wt(K::qv(hd, 1), 2, kc, 0, nt), hop[kc][1], wt(K::qv(hd, 1), 2, kc, 1, nt));
}
const gop q0 = tok_codes(qt[0]), q1 = tok_codes(qt[1]);
// S transposed: lane = query token, rows = key tokens
T16 S = tok_t16(ch[K::CATT + 32 * hd + 2 * l16]);
k32(S, q0, q0, q1, q1);
const half_t cs = f16(0.0972222164273262f / 5.656854152679443f);
const half_t cl = (half_t)0.55615234375f;
const half_t p1 = f16(0.92730712890625f), p0 = (half_t)1.375f;
hv2 wg[4], v1[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
{
hv2 tv = S.pair(k) * (hv2){cs, cs};
tv = __builtin_elementwise_max(__builtin_elementwise_min(tv, (hv2){cl, cl}), (hv2){-cl, -cl});
const hv2 e1 = __builtin_elementwise_fma(tv, -tv, (hv2){p1, p1});
const hv2 poly = __builtin_elementwise_fma(tv, e1, (hv2){p0, p0});
// the pair is the (even, odd) key word of the exponent trick
wg[k] = __builtin_bit_cast(hv2, (__builtin_bit_cast(uint32_t, poly) << 5) + 0x7FF88000u);
v1[k] = wg[k] + tok_other(wg[k]); // keys c and c ^ 8
}
const hv2 v3 = (v1[0] + v1[1]) + (v1[2] + v1[3]); // c ^ 2, then c ^ 4
const half_t sum = v3[0] + v3[1]; // c ^ 1
const half_t rinv = recip16(sum);
hv2 pw[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
pw[k] = wg[k] * (hv2){rinv, rinv};
#ifdef SWIN_PV_ZLUDA_ORDER
// P V is the block's only product of two f16 operands: its f32 sum over the 16 keys is not exact, so the
// rounded result depends on which K slot holds which key. Here the slots follow ZLUDA's lowering of
// NVIDIA's m16n8k16 (dword j of half h = key pair (j & 1) * 4 + 2 h + (j >> 1)) instead of this
// layout's own (key pair 4 h + j), which costs two half-wave exchanges per operand. A kernel that has to
// reproduce ZLUDA's compile byte for byte needs it; the two orders differ in about one token in 30,000.
auto zk = [&](hv2 (&w)[4]) {
const bool upper = m_half() != 0;
const hv2 r0 = tok_other(upper ? w[0] : w[2]), r1 = tok_other(upper ? w[1] : w[3]);
const hv2 a = upper ? r0 : w[0], b = upper ? w[2] : r0, c = upper ? r1 : w[1], d = upper ? w[3] : r1;
w[0] = a, w[1] = b, w[2] = c, w[3] = d;
};
zk(pw);
#endif
const sop pb = __builtin_bit_cast(sop, pw);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
hv2 vp[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
vp[k] = vt[nt].pair(k);
#ifdef SWIN_PV_ZLUDA_ORDER
zk(vp);
#endif
const f8v o = wmma(__builtin_bit_cast(sop, vp), pb, splat(0.0f));
oc[hd][nt] = (u2v){codes4(tok_pk16(o[0], o[1]), tok_pk16(o[2], o[3])), codes4(tok_pk16(o[4], o[5]), tok_pk16(o[6], o[7]))};
}
}
// ------------------------------------------------ output projection + residual
T16 x[NT];
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int k = 0; k < 4; ++k)
x[t].set_pair(k, res[t][k]);
#pragma unroll
for (int kc = 0; kc < NH; ++kc)
#pragma unroll
for (int t = 0; t < NT; ++t)
k32(x[t], wt(K::WO, NT, kc, 0, t), oc[kc][0], wt(K::WO, NT, kc, 1, t), oc[kc][1]);
// ------------------------------------------------ norm2: h2 codes, the K operands of fc1
gop fop[NKC][2];
{
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int k = 0; k < 4; ++k)
v[t][k] = x[t].pair(k);
const half_t rn = tok_rsq<NT>(v);
const hv2 rn2 = {rn, rn};
uint32_t c2[NT][2];
#pragma unroll
for (int t = 0; t < NT; ++t)
{
const h8 g2 = ch[K::CG2 + 2 * t], b2 = ch[K::CB2 + 2 * t];
#pragma unroll
for (int g = 0; g < 2; ++g)
c2[t][g] = codes4(v[t][2 * g] * (hv2)(rn2 * (hv2){g2[4 * g], g2[4 * g + 1]}),
v[t][2 * g + 1] * (hv2)(rn2 * (hv2){g2[4 * g + 2], g2[4 * g + 3]}));
#pragma unroll
for (int k = 0; k < 4; ++k)
x[t].set_pair(k, v[t][k] + (hv2){b2[2 * k], b2[2 * k + 1]});
}
#pragma unroll
for (int kc = 0; kc < NKC; ++kc)
#pragma unroll
for (int s = 0; s < 2; ++s)
fop[kc][s] = tok_k32_operand(&c2[2 * kc], s);
}
// ------------------------------------------------ MLP: fc1 transposed, its codes are fc2's operands
#pragma unroll 1
for (int c = 0; c < NC; ++c)
{
const int w1 = K::W1 + 4 * C * c;
// the round's operands first: 2 bias slots, fc1's and fc2's weight tiles
h8 zb[2];
gop wa[NKC][2][2], wb[2][NT];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
zb[nt] = ch[K::CB1 + 4 * c + 2 * nt];
#pragma unroll
for (int kc = 0; kc < NKC; ++kc)
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
wa[kc][s][nt] = wt(w1, 2, kc, s, nt);
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
for (int t = 0; t < NT; ++t)
wb[s][t] = wt(w1 + 2 * C, NT, 0, s, t);
#ifndef SWIN_TOK_NO_LOAD_SYNC
// one wait for the whole batch here instead of one before each WMMA
asm("" : "+v"(wb[1][NT - 1]));
#endif
T16 z[2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
z[nt] = tok_t16(zb[nt]);
#pragma unroll
for (int kc = 0; kc < NKC; ++kc)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
k32(z[nt], wa[kc][0][nt], fop[kc][0], wa[kc][1][nt], fop[kc][1]);
gop gz[2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
hv2 gl[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
{
const hv2 zv = z[nt].pair(k);
// minimum / maximum instead of swin_block.h's minnum / maxnum: they differ only for a NaN z, and
// then the product with z below is NaN either way (no canonicalising instruction per pair)
const hv2 cz = __builtin_elementwise_maximum(__builtin_elementwise_minimum(zv, (hv2){2.0f16, 2.0f16}),
(hv2){-2.0f16, -2.0f16});
const hv2 az = __builtin_bit_cast(hv2, (u16x2)(__builtin_bit_cast(u16x2, cz) & (unsigned short)0x7fff));
const hv2 inner = (hv2){f16(0.41216981f), f16(0.41216981f)} -
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
gl[k] = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
}
gz[nt] = (u2v){codes4(gl[0], gl[1]), codes4(gl[2], gl[3])};
}
#pragma unroll
for (int t = 0; t < NT; ++t)
k32(x[t], wb[0][t], gz[0], wb[1][t], gz[1]);
}
// ------------------------------------------------ output codes
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
co[t][g] = codes4(x[t].pair(2 * g), x[t].pair(2 * g + 1));
if (x2 != nullptr)
{
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int k = 0; k < 4; ++k)
x2[t][k] = x[t].pair(k);
}
}
template <int C, int NH, int NPM, bool TUBE, int CIN>
__device__ __forceinline__ void swin_tok_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const h8* c16,
const int bx, const int by, const int gx)
{
using K = TokLayout<C, NH, NPM, CIN>;
constexpr int NT = C / 16, NKC = C / 32, NC = C / 8;
static_assert(C % 32 == 0, "32-channel groups");
const int lane = threadIdx.x, wv = threadIdx.z;
const int l16 = lane & 15, hf = lane >> 4;
const int T = 16 * wv + l16; // this lane's token: window wv, position l16
// The lane's bases into the two images, hidden from constant folding: every tile and constant is then
// one load at a fixed offset from a register (folded, each load rebuilt the image's address first).
uintptr_t chv = (uintptr_t)(c16 + hf), wlv = (uintptr_t)((const gop*)w16 + 2 * l16 + hf);
asm("" : "+v"(chv));
asm("" : "+v"(wlv));
const h8* const ch = (const h8*)chv;
// row l16 of a weight tile: this half's 8 K slots
const gop* const wl = (const gop*)wlv;
auto wt = [&](int dst, int nt_count, int kc, int s, int nt) { return wl[2 * (dst + ((kc * 2 + s) * nt_count + nt) * 16)]; };
if constexpr (TUBE)
{
if (tp->in_flags != nullptr)
{
if (lane <= 3 && wv == 0)
{
int nx = bx + (lane & 1) + tp->r71, ny = by + (lane >> 1) + tp->r72;
if (nx >= 0 && nx < tp->r69 && ny >= 0 && ny < tp->r70)
{
uint8_t* f = (uint8_t*)tp->in_flags + (size_t)(ny * tp->r69 + nx) * 32;
// bounded spin: a missing producer must never hang the GPU
for (int it = 0; it < (1 << 22); ++it)
{
if (__hip_atomic_load(f, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) != 0)
break;
__builtin_amdgcn_s_sleep(1);
}
}
}
__syncthreads();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "agent");
}
}
// ------------------------------------------------ block input x0 (encoders: the fp8 input; decoders: skip)
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
const int Xm = mirror(X, p.tw), Ym = mirror(Y, p.th);
uint32_t inw[NT][2];
{
const uint8_t* const src = CIN ? p.p32 : p.in;
#pragma unroll
for (int pl = 0; pl < NKC; ++pl)
{
const uint8_t* b = src + ((size_t)(pl * p.th + Ym) * p.tw + Xm) * 32 + tok_half(hf);
#pragma unroll
for (int j = 0; j < 4; ++j)
inw[2 * pl + (j >> 1)][j & 1] = input_word((const uint32_t*)(b + tok_word(j >> 1, j & 1)));
}
}
hv2 v[NT][4]; // x0, then x1, as f16 pairs: [tile][2 g + pair of the dword]
if constexpr (CIN > 0)
{
// patch expand: lanes = the block's 16 low-res tokens, wave = quadrant; the codes of q8(expand) go to
// their child tokens' rows of a byte image, where the windows' waves pick them up
__shared__ __attribute__((aligned(16))) uint8_t E8[64 * C];
const int W2 = p.tw / 2, H2 = p.th / 2;
const int MX = mirror((8 * bx - p.sx) / 2 + (l16 & 3), W2), MY = mirror((8 * by - p.sy) / 2 + (l16 >> 2), H2);
gop pa[CIN / 32][2];
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
{
const uint8_t* b = p.in + ((size_t)(kc * H2 + MY) * W2 + MX) * 32 + tok_half(hf);
#pragma unroll
for (int s = 0; s < 2; ++s)
#ifdef SWIN_TOK_NATURAL
pa[kc][s] = (u2v){input_word((const uint32_t*)(b + tok_word(s, 0))), input_word((const uint32_t*)(b + tok_word(s, 1)))};
#else
pa[kc][s] = (u2v){input_word((const uint32_t*)(b + tok_word(0, s))), input_word((const uint32_t*)(b + tok_word(1, s)))};
#endif
}
const int tx = 2 * (l16 & 3) + (wv & 1), ty = 2 * (l16 >> 2) + (wv >> 1);
const int Tc = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
const int pew = K::pe(0) + wv * (CIN * C / 16);
#pragma unroll
for (int ct = 0; ct < NT; ++ct)
{
T16 e = tok_t16(ch[K::CPEB + 2 * (NT * wv + ct)]);
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
k32(e, wt(pew, NT, kc, 0, ct), pa[kc][0], wt(pew, NT, kc, 1, ct), pa[kc][1]);
*(u2v*)(E8 + Tc * C + 16 * ct + 8 * hf) = tok_codes(e);
}
__syncthreads();
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
{
const uint32_t ew = *(const uint32_t*)(E8 + T * C + tok_word(t, g) + tok_half(hf));
v[t][2 * g] = tok_dec2<false>(ew) + tok_dec2<false>(inw[t][g]);
v[t][2 * g + 1] = tok_dec2<true>(ew) + tok_dec2<true>(inw[t][g]);
}
}
else
{
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
{
v[t][2 * g] = tok_dec2<false>(inw[t][g]);
v[t][2 * g + 1] = tok_dec2<true>(inw[t][g]);
}
}
uint32_t co[NT][2];
swin_tok_core<C, NH, NPM, CIN>(ch, wl, l16, v, co);
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
{
#pragma unroll
for (int pl = 0; pl < NKC; ++pl)
{
uint8_t* b = p.out + ((size_t)(pl * p.th + Y) * p.tw + X) * 32 + tok_half(hf);
#pragma unroll
for (int j = 0; j < 4; ++j)
*(uint32_t*)(b + tok_word(j >> 1, j & 1)) = co[2 * pl + (j >> 1)][j & 1];
}
}
// ------------------------------------------------ patch merge: 2x2 tokens (4C) -> NPM, half resolution
if constexpr (NPM > 0)
{
// lanes = the block's 16 merged tokens; the merge reads the output codes from a byte image
constexpr int XB = C + 4, NTP = NPM / 16;
__shared__ __attribute__((aligned(16))) uint8_t U[64 * XB];
#pragma unroll
for (int t = 0; t < NT; ++t)
#pragma unroll
for (int g = 0; g < 2; ++g)
*(uint32_t*)(U + T * XB + tok_word(t, g) + tok_half(hf)) = co[t][g];
__syncthreads();
if (wv < NTP)
{
const int mx = l16 & 3, my = l16 >> 2;
// K slot 8 half + 4 rr + b of step (kc, s) is value k = 32 kc + 16 rr + 8 s + 4 half + b of the merged
// token's 4C inputs: channel k % C of its source token k / C. The four source tokens are consecutive
// in x and 4 apart in y, and 4 half never crosses a token, so every run of four codes is at a fixed
// offset from the lane's first source token.
const uint8_t* const ub = U + ((2 * mx & 3) + 4 * (2 * my & 3) + 16 * (mx >> 1) + 32 * (my >> 1)) * XB + tok_half(hf);
gop as[C / 8][2];
#pragma unroll
for (int kc = 0; kc < C / 8; ++kc)
#pragma unroll
for (int s = 0; s < 2; ++s)
{
uint32_t run4[2];
#pragma unroll
for (int rr = 0; rr < 2; ++rr)
{
const int k0 = tok_merge_k(kc, s, rr), q = k0 / C;
run4[rr] = *(const uint32_t*)(ub + ((q & 1) + 4 * (q >> 1)) * XB + k0 % C);
}
as[kc][s] = (u2v){run4[0], run4[1]};
}
const int W2 = p.tw / 2, H2 = p.th / 2;
const int MX = (8 * bx - p.sx) / 2 + mx, MY = (8 * by - p.sy) / 2 + my;
const bool inside = MX >= 0 && MX < W2 && MY >= 0 && MY < H2;
// output tiles u = wv, wv + 4, ...: rows = 16 natural channels, 8 contiguous bytes per half-wave
#pragma unroll
for (int j = 0; j < (NTP + 3) / 4; ++j)
{
const int u = wv + 4 * j;
if (u < NTP)
{
T16 y = tok_t16(ch[K::CPMB + 2 * u]);
#pragma unroll
for (int kc = 0; kc < C / 8; ++kc)
k32(y, wt(K::PM, NTP, kc, 0, u), as[kc][0], wt(K::PM, NTP, kc, 1, u), as[kc][1]);
const u2v code = {codes4(merge_clamp(y.pair(0)), merge_clamp(y.pair(1))),
codes4(merge_clamp(y.pair(2)), merge_clamp(y.pair(3)))};
if (inside)
*(u2v*)(p.p48 + ((size_t)((u >> 1) * H2 + MY) * W2 + MX) * 32 + 16 * (u & 1) + 8 * hf) = code;
}
}
}
}
if constexpr (TUBE)
{
if (tp->out_flags != nullptr)
{
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "agent");
__syncthreads();
if (lane == 0 && wv == 0)
__hip_atomic_store(tp->out_flags + (size_t)(by * gx + bx) * 32, (uint8_t)1, __ATOMIC_RELEASE,
__HIP_MEMORY_SCOPE_AGENT);
}
}
}
#if defined(SWIN_TOK_MAX_VGPR)
#define SWIN_TOK_VGPR_ATTR __attribute__((amdgpu_num_vgpr(SWIN_TOK_MAX_VGPR)))
#else
#define SWIN_TOK_VGPR_ATTR
#endif
// Module boilerplate: the weight and constant images, their prep kernel, and the 4-wave main kernel
#define SWIN_TOK_MODULE(NAME, C, NH, NPM, TUBE, PARAMS, CIN) \
using NAME##_K = TokLayout<C, NH, NPM, CIN>; \
constexpr int NAME##_S = SWIN_PREP_SLOTS; \
__device__ wslot g_w16[(NAME##_S + 1) * NAME##_K::TOTAL]; \
__device__ h8 g_c16[(NAME##_S + 1) * NAME##_K::CTOTAL]; \
extern "C" __device__ __attribute__((used)) uint64_t g_prep_keys[NAME##_S + 1] = {}; \
__constant__ TokTables<C, NH, NPM, CIN> g_tok = make_tok_tables<C, NH, NPM, CIN>(); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = \
((NAME##_K::TOTAL > NAME##_K::CTOTAL ? NAME##_K::TOTAL : NAME##_K::CTOTAL) + 127) / 128; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = SWIN_PREP_KEY_AT; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = NAME##_S; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_identity_at = SWIN_PREP_IDENTITY_AT; \
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##_K::TOTAL) \
expand_tok_weights(w, g_w16 + slot * NAME##_K::TOTAL, g_tok.d, NAME##_K::NDESC, idx); \
if (idx < NAME##_K::CTOTAL) \
expand_tok_consts(w, g_c16 + slot * NAME##_K::CTOTAL, g_tok.c, NAME##_K::NCONST, idx); \
} \
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = 0; \
extern "C" __global__ void __launch_bounds__(128) SWIN_TOK_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 int gx = (cp.tw + cp.sx + 7) / 8; \
swin_tok_block<C, NH, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16 + slot * NAME##_K::TOTAL, \
g_c16 + slot * NAME##_K::CTOTAL, blockIdx.x, blockIdx.y, gx); \
}
kernels/support/motion_dilate.hip
+64
View File
@@ -0,0 +1,64 @@
#include "../../tools/d4r_motion_dilation.h"
#pragma clang fp contract(off)
__attribute__((device)) static int clamp_index(int x, uint32_t limit)
{
return x < 0 ? 0 : x >= (int)limit ? (int)limit - 1 : x;
}
extern "C" __attribute__((global, amdgpu_flat_work_group_size(128, 128)))
void d4r_motion_dilate(D4rMotionDilationParams p)
{
const uint32_t id = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
if ((uint64_t)id >= (uint64_t)p.motion_width * p.motion_height)
return;
const uint32_t x = id % p.motion_width, y = id / p.motion_width;
const int cx = clamp_index((int)__builtin_floorf((x + 0.5f) * p.render_width / p.motion_width + p.jitter_x), p.render_width);
const int cy = clamp_index((int)__builtin_floorf((y + 0.5f) * p.render_height / p.motion_height + p.jitter_y), p.render_height);
const uint8_t* depth = (const uint8_t*)p.depth;
const uint8_t* motion = (const uint8_t*)p.motion;
uint8_t* output = (uint8_t*)p.output;
const float current = *(const float*)(depth + (uint64_t)cy * p.depth_pitch + 4 * cx);
float best = current;
int bx = cx, by = cy;
for (int dy = -(int)p.radius; dy <= (int)p.radius; ++dy)
for (int dx = -(int)p.radius; dx <= (int)p.radius; ++dx)
{
const int nx = clamp_index(cx + dx, p.render_width), ny = clamp_index(cy + dy, p.render_height);
const float candidate = *(const float*)(depth + (uint64_t)ny * p.depth_pitch + 4 * nx);
if (p.reversed_depth ? candidate > best : candidate < best)
best = candidate, bx = nx, by = ny;
}
// Sky gaps need wider coverage than opaque edges. Select the closest
// foreground pixel within eight render pixels, rather than the closest
// depth (which can borrow a distant roof/chimney's velocity).
const bool sky = p.reversed_depth ? current <= 1.e-7f : current >= 1.0f - 1.e-7f;
bool sky_filled = false;
if (sky)
{
int best_distance = 65;
for (int dy = -8; dy <= 8; ++dy)
for (int dx = -8; dx <= 8; ++dx)
{
const int nx = cx + dx, ny = cy + dy;
const int distance = dx * dx + dy * dy;
if (nx < 0 || ny < 0 || nx >= (int)p.render_width || ny >= (int)p.render_height ||
distance > 64 || distance > best_distance)
continue;
const float candidate = *(const float*)(depth + (uint64_t)ny * p.depth_pitch + 4 * nx);
const bool foreground = p.reversed_depth ? candidate > 1.e-7f : candidate < 1.0f - 1.e-7f;
// Stable coordinate tie-break matches nearest-foreground fill.
if (foreground && (distance < best_distance || nx < bx || (nx == bx && ny < by)))
best_distance = distance, best = candidate, bx = nx, by = ny, sky_filled = true;
}
}
uint32_t value = *(const uint32_t*)(motion + (uint64_t)y * p.motion_pitch + 4 * x);
const bool nearer = p.reversed_depth ? best > current * 1.02f + 1.e-7f : best < current * 0.98f - 1.e-7f;
if (sky_filled || nearer)
{
const int mx = clamp_index((int)__builtin_floorf((bx + 0.5f - p.jitter_x) * p.motion_width / p.render_width), p.motion_width);
const int my = clamp_index((int)__builtin_floorf((by + 0.5f - p.jitter_y) * p.motion_height / p.render_height), p.motion_height);
value = *(const uint32_t*)(motion + (uint64_t)my * p.motion_pitch + 4 * mx);
}
*(uint32_t*)(output + (uint64_t)y * p.motion_pitch + 4 * x) = value;
}
kernels/tex/build_tex.sh
+37
View File
@@ -0,0 +1,37 @@
#!/usr/bin/env bash
# build_tex.sh KERNEL SRC OUT_DIR: one texture kernel. NVIDIA's PTX for KERNEL, edited by make_ptx.py so
# part of it calls functions from SRC.hip, is compiled by ZLUDA with SRC's bitcode linked in
# (D4R_ZLUDA_EXTRA_BC) and the resulting code object is saved as OUT_DIR/KERNEL.hsaco.
# Called by kernels/build.sh, which sets D4R_ROCM_DIR, D4R_GPU_ARCH, D4R_DLSS_PTX_DIR, D4R_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}"
ACCURACY="${D4R_PREFER_ACCURACY:-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)
[[ "$ACCURACY" == 1 ]] && FLAGS+=(-DD4R_ACCURACY)
# 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 -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" -
D4R_PREFER_ACCURACY="$ACCURACY" python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
# the code-generation switches of the runtime (d4r.ini [Kernels]), which ZLUDA's module cache keys on
env LD_LIBRARY_PATH="$R/lib" D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 \
D4R_ZLUDA_WMMA_FP8_NATIVE="$FP8" D4R_ZLUDA_IGNORE_DENORMAL="$((1 - ACCURACY))" \
D4R_ZLUDA_WMMA_F32ACC=0 D4R_ZLUDA_FAST_MATH=0 D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 \
D4R_ZLUDA_WAVE64="$([[ "$ACCURACY" == 1 ]] && echo 0 || echo "${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|wavefront_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
kernels/tex/dec0_head.hip
+175
View File
@@ -0,0 +1,175 @@
// Native head of rrlite_dec0_4x4_folded: the token GEMM (16 half-resolution tokens x K 64 (two
// e4m3 planes) x N 128, f16 bias, f16 accumulation per k32 step) and the e4m3 staging of its result
// in shared memory, which the rest of the kernel (ZLUDA's compile of NVIDIA's PTX, dec0_head.ptx)
// reads. Arithmetic as ZLUDA's m16n8k32 lowering (see enc0_tail.hip); tokens are mirror padded.
#include "tex_common.h"
// weights as WMMA A images: index ((p * 2 + c) * 2 + h) * 16 + n16 for channel pair p (16p..16p+15),
// k32 chunk c (input plane), k16 half h
#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
{
uint64_t w;
uint8_t rest[160]; // 168-byte parameter block
};
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4x4_folded_prep(Dec0Params prm)
{
const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
if (idx >= 512)
return;
const uint8_t* w = (const uint8_t*)prm.w;
const int n16 = idx & 15, h = (idx >> 4) & 1, c = (idx >> 5) & 1, p = idx >> 6;
const int g = n16 & 7, second = n16 >> 3;
u4v d;
#pragma unroll
for (int j = 0; j < 4; ++j)
d[j] = *(const uint32_t*)(w + 512 * (2 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
#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)
{
const int32_t a = v < 0 ? -v : v;
const int32_t m = __builtin_elementwise_min(a, (int32_t)(full & ~1u) - a - 2);
return __builtin_elementwise_min(half - 1u, (uint32_t)m);
}
DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx, uint32_t dimy, int32_t ox, int32_t oy)
{
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const uint8_t* w = (const uint8_t*)wp;
lds_u8* s = (lds_u8*)(uintptr_t)smem;
const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
const uint32_t Y = mirror_clamp(oy / 2 + (int32_t)(m >> 2), dimy, H2);
const uint32_t Xc = mirror_clamp(ox / 2 + (int32_t)(m & 3), dimx, W2);
const uint32_t* in = (const uint32_t*)inp + rowstride * Y + 8 * Xc;
#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)
#pragma unroll
for (int h = 0; h < 2; ++h)
X[c][h] = operand16(*(const u4v*)(in + plane * c + 4 * h));
uint32_t E[8][2];
#pragma unroll
for (int p = 0; p < 8; ++p)
{
half_t acc[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = *(const half_t*)(w + 8192 + 2 * (16 * p + 2 * i + hf));
#pragma unroll
for (int c = 0; c < 2; ++c)
{
f8v cf;
#pragma unroll
for (int i = 0; i < 8; ++i)
cf[i] = (float)acc[i];
cf *= 0x1p-16f;
const u8v* wv = &g_dec0_w16[((p * 2 + c) * 2) * 16 + m];
cf = wmma(wv[0], X[c][0], cf);
cf = wmma(wv[16], X[c][1], cf);
cf = cf * 0x1p16f;
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = (half_t)cf[i];
}
uint32_t codes[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
codes[k] = e4m3x2(pack2(acc[2 * k], acc[2 * k + 1]));
E[p][0] = codes[0] | (codes[1] << 16);
E[p][1] = codes[2] | (codes[3] << 16);
}
// staging: word t of token m at 1024 * H + 32 * (m & 3) + 256 * (m >> 2) + {0, 16, 128, 144}[G] + 4t
// holds columns 2t, 2t + 1 of tiles 4G + H and 4G + 2 + H; lanes 0..15 write H = 0, 16..31 H = 1
uint32_t EV[8], OD[8];
#pragma unroll
for (int p = 0; p < 8; ++p)
{
const uint32_t theirs = other_half(hf ? E[p][0] : E[p][1]);
EV[p] = hf ? theirs : E[p][0];
OD[p] = hf ? E[p][1] : theirs;
}
const uint32_t base = 1024 * hf + 32 * (m & 3) + 256 * (m >> 2);
const uint32_t goff[4] = {0, 16, 128, 144};
#pragma unroll
for (int G = 0; G < 4; ++G)
{
const uint32_t a01 = __builtin_amdgcn_perm(OD[2 * G], EV[2 * G], 0x05010400u);
const uint32_t a23 = __builtin_amdgcn_perm(OD[2 * G], EV[2 * G], 0x07030602u);
const uint32_t b01 = __builtin_amdgcn_perm(OD[2 * G + 1], EV[2 * G + 1], 0x05010400u);
const uint32_t b23 = __builtin_amdgcn_perm(OD[2 * G + 1], EV[2 * G + 1], 0x07030602u);
u4v words = {__builtin_amdgcn_perm(b01, a01, 0x05040100u), __builtin_amdgcn_perm(b01, a01, 0x07060302u),
__builtin_amdgcn_perm(b23, a23, 0x05040100u), __builtin_amdgcn_perm(b23, a23, 0x07060302u)};
*(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
}
#endif
}
kernels/tex/enc0_tail.hip
+204
View File
@@ -0,0 +1,204 @@
// 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
// (scaled e4m3 widening, same k slot order, C * 2^-16 inside the f32 WMMA, f16 rounding per k32).
#include "tex_common.h"
// B operands (weights) as WMMA A images: index ((p * 4 + c) * 2 + h) * 16 + n16 for channel pair p
// (channels 16p..16p+15), k32 chunk c, k16 half h.
#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
{
uint64_t w;
uint8_t rest[352];
};
// 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)
return;
const uint8_t* w = (const uint8_t*)prm.w;
const int n16 = idx & 15, h = (idx >> 4) & 1, c = (idx >> 5) & 3, p = idx >> 7;
// NVIDIA B fragments: v4 load 4p + c of lane 4 * column + j holds {b0, b1} of tile 2p then 2p + 1
const int g = n16 & 7, second = n16 >> 3;
u4v d;
#pragma unroll
for (int j = 0; j < 4; ++j)
d[j] = *(const uint32_t*)(w + 4096 + 512 * (4 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
#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)
{
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const uint8_t* w = (const uint8_t*)wp;
lds_u8* s = (lds_u8*)(uintptr_t)smem;
#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);
const uint32_t choff[4] = {0, 16, 128, 144};
u8v X[4][2];
#pragma unroll
for (int c = 0; c < 4; ++c)
#pragma unroll
for (int h = 0; h < 2; ++h)
X[c][h] = operand16(*(const __attribute__((address_space(3))) u4v*)(s + 1024 * h + rowoff + choff[c]));
const hv2 lo = {(half_t)-6.2831855f, (half_t)-6.2831855f}, hi = {(half_t)6.2831855f, (half_t)6.2831855f};
uint32_t E[4][2];
#pragma unroll
for (int p = 0; p < 4; ++p)
{
// D^T layout: element i is channel 16p + 2i + hf of token m
half_t acc[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = *(const half_t*)(w + 12288 + 2 * (16 * p + 2 * i + hf));
#pragma unroll
for (int c = 0; c < 4; ++c)
{
f8v cf;
#pragma unroll
for (int i = 0; i < 8; ++i)
cf[i] = (float)acc[i];
cf *= 0x1p-16f;
const u8v* wv = &g_enc0_w16[((p * 4 + c) * 2) * 16 + m];
cf = wmma(wv[0], X[c][0], cf);
cf = wmma(wv[16], X[c][1], cf);
cf = cf * 0x1p16f;
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = (half_t)cf[i];
}
uint32_t codes[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
{
hv2 v = {acc[2 * k], acc[2 * k + 1]};
v = __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, lo), hi);
codes[k] = e4m3x2(__builtin_bit_cast(uint32_t, v));
}
// byte i of E[p][0] (E[p][1]) is column 2i + hf (2(i - 4) + hf) of tile 2p (2p + 1)
E[p][0] = codes[0] | (codes[1] << 16);
E[p][1] = codes[2] | (codes[3] << 16);
}
// lanes 0..15 write plane 0 (tiles 0..3), lanes 16..31 plane 1 (tiles 4..7): swap the other
// plane's codes with the partner lane holding the other column parity
uint32_t mine[2][2], theirs[2][2];
#pragma unroll
for (int a = 0; a < 2; ++a)
#pragma unroll
for (int j = 0; j < 2; ++j)
{
const uint32_t send = hf ? E[a][j] : E[2 + a][j];
theirs[a][j] = other_half(send);
mine[a][j] = hf ? E[2 + a][j] : E[a][j];
}
u8v words;
#pragma unroll
for (int j = 0; j < 2; ++j)
{
const uint32_t ev0 = hf ? theirs[0][j] : mine[0][j], od0 = hf ? mine[0][j] : theirs[0][j];
const uint32_t ev1 = hf ? theirs[1][j] : mine[1][j], od1 = hf ? mine[1][j] : theirs[1][j];
const uint32_t a01 = __builtin_amdgcn_perm(od0, ev0, 0x05010400u), a23 = __builtin_amdgcn_perm(od0, ev0, 0x07030602u);
const uint32_t b01 = __builtin_amdgcn_perm(od1, ev1, 0x05010400u), b23 = __builtin_amdgcn_perm(od1, ev1, 0x07030602u);
words[4 * j + 0] = __builtin_amdgcn_perm(b01, a01, 0x05040100u);
words[4 * j + 1] = __builtin_amdgcn_perm(b01, a01, 0x07060302u);
words[4 * j + 2] = __builtin_amdgcn_perm(b23, a23, 0x05040100u);
words[4 * j + 3] = __builtin_amdgcn_perm(b23, a23, 0x07060302u);
}
const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
const int32_t Y = oy / 2 + (int32_t)(m >> 2), Xc = ox / 2 + (int32_t)(m & 3);
if ((uint32_t)Y < H2 && (uint32_t)Xc < W2)
{
uint32_t* out = (uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * hf;
*(u8v*)out = words;
}
#endif
}
kernels/tex/make_ptx.py
+329
View File
@@ -0,0 +1,329 @@
#!/usr/bin/env python3
"""Hand-edited PTX for the texture kernels: keep NVIDIA's code up to a cut line, then call a native tail.
usage: make_ptx.py KERNEL OUT.ptx
"""
import os
import re
import sys
from pathlib import Path
# PTX extracted from nvngx_dlss.dll by kernels/tools/extract_dlss_ptx.py (kernels/build.sh does this)
PTX_DIR = Path(os.environ.get("D4R_DLSS_PTX_DIR", Path(__file__).resolve().parent.parent / "extracted" / "ptx"))
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(
file="dlss-0185-00.ptx",
cut_after=2548,
cut_check="bar.warp.sync -1;",
sust=True,
extern=""".extern .func d4r_enc0_tail
(
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3,
.param .b32 d4r_a4, .param .b32 d4r_a5, .param .b32 d4r_a6
)
;
""",
tail=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.
"rrlite_dec0_4x4_folded": dict(
file="dlss-0195-00.ptx",
delete=(139, 587),
delete_check=("mov.u32 %r17, %laneid;", "st.shared.v2.u16 [%r570+1168], {%rs63, %rs64};"),
sust=True,
extern=""".extern .func d4r_dec0_head
(
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3,
.param .b32 d4r_a4, .param .b32 d4r_a5, .param .b32 d4r_a6
)
;
""",
insert="""mov.u32 %r560, _ZZ22rrlite_dec0_4x4_foldedN6RRLite18Decoder0ParametersEE4smem;
{
.param .b32 q0;
st.param.b32 [q0+0], %r560;
.param .b64 q1;
st.param.b64 [q1+0], %rd22;
.param .b64 q2;
st.param.b64 [q2+0], %rd24;
.param .b32 q3;
st.param.b32 [q3+0], %r475;
.param .b32 q4;
st.param.b32 [q4+0], %r476;
.param .b32 q5;
st.param.b32 [q5+0], %r3;
.param .b32 q6;
st.param.b32 [q6+0], %r4;
call.uni d4r_dec0_head, (q0, q1, q2, q3, q4, q5, q6);
}""",
),
# post and downsample: only the formatted surface stores become native (inline format conversion
# instead of ZLUDA's per-channel surface_formatted_store_bits calls)
"rrlite_post_3_2_mvhi_hdr_folded": dict(file="dlss-0176-00.ptx", sust=True, extern=""),
"rrlite_downsample_kernel_static_hdr": dict(file="dlss-0196-00.ptx", sust=True, extern=""),
# preset K (DLSS 4): the output kernel's 17 surface stores (13 formatted, 4 raw 32-bit)
"hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel": dict(file="dlss-0078-00.ptx", sust=True, rz_round=True, extern=""),
"hiluma_engine_input_depthinv_mvhi_hdr_v2_rel": dict(file="dlss-0062-00.ptx", sust=True, rz_round=True, extern=""),
}
SUST_EXTERN = """.extern .func d4r_sust_v2b16
(
.param .b64 d4r_s0, .param .b32 d4r_s1, .param .b32 d4r_s2, .param .b16 d4r_s3, .param .b16 d4r_s4
)
;
"""
SUSTP_EXTERN = """.extern .func d4r_sust_p_v4b32
(
.param .b64 d4r_t0, .param .b32 d4r_t1, .param .b32 d4r_t2, .param .b32 d4r_t3, .param .b32 d4r_t4,
.param .b32 d4r_t5, .param .b32 d4r_t6
)
;
"""
SUSTP_RE = re.compile(r"^sust\.p\.2d\.v4\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%f\d+),(%f\d+),(%f\d+),(%f\d+)\};$")
SUSTB32_EXTERN = """.extern .func d4r_sust_b32
(
.param .b64 d4r_w0, .param .b32 d4r_w1, .param .b32 d4r_w2, .param .b32 d4r_w3
)
;
"""
SUSTB32_RE = re.compile(r"^sust\.b\.2d\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%r\d+)\};$")
SUST_RE = re.compile(r"^sust\.b\.2d\.v2\.b16\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%rs\d+),(%rs\d+)\};$")
SUST_KINDS = set(os.environ.get("D4R_SUST_KINDS", "b16,p,b32").split(","))
def replace_sust(lines):
out, n = [], {"b16": 0, "p": 0, "b32": 0}
for line in lines:
mb = SUSTB32_RE.match(line.strip()) if "b32" in SUST_KINDS else None
if mb:
s, x, y, v = mb.groups()
out.append(f"""{{
.param .b64 u0;
st.param.b64 [u0+0], {s};
.param .b32 u1;
st.param.b32 [u1+0], {x};
.param .b32 u2;
st.param.b32 [u2+0], {y};
.param .b32 u3;
st.param.b32 [u3+0], {v};
call d4r_sust_b32, (u0, u1, u2, u3);
}}""")
n["b32"] += 1
continue
mp = SUSTP_RE.match(line.strip()) if "p" in SUST_KINDS else None
if mp:
s, x, y, *v = mp.groups()
st = "\n".join(f"\t.param .b32 v{i};\n\tst.param.b32 [v{i}+0], {r};" for i, r in enumerate(v))
out.append(f"""{{
.param .b64 u0;
st.param.b64 [u0+0], {s};
.param .b32 u1;
st.param.b32 [u1+0], {x};
.param .b32 u2;
st.param.b32 [u2+0], {y};
{st}
call d4r_sust_p_v4b32, (u0, u1, u2, v0, v1, v2, v3);
}}""")
n["p"] += 1
continue
m = SUST_RE.match(line.strip())
if not m:
out.append(line)
continue
s, x, y, a, b = m.groups()
out.append(f"""{{
.param .b64 u0;
st.param.b64 [u0+0], {s};
.param .b32 u1;
st.param.b32 [u1+0], {x};
.param .b32 u2;
st.param.b32 [u2+0], {y};
.param .b16 u3;
st.param.b16 [u3+0], {a};
.param .b16 u4;
st.param.b16 [u4+0], {b};
call d4r_sust_v2b16, (u0, u1, u2, u3, u4);
}}""")
n["b16"] += 1
return out, n
RZ_RE = re.compile(r"^add\.rz\.ftz\.f32 (%f\d+), (%f\d+), (%f\d+);$")
CVT_RE = re.compile(r"^cvt\.rzi\.f32\.f32 (%f\d+), (%f\d+);$")
def replace_rz_round(lines):
"""roundf idiom add.rz(x, copysign(0.5, x)) + cvt.rzi -> exact round-half-away in default rounding:
t = trunc(x); r = |x - t| >= 0.5 ? t + 2 * copysign(0.5, x) : t (x - t is exact). Lets ZLUDA compile the
kernel without the strict-FP (constrained) mode that one explicit rounding mode forces."""
out, n, i = [], 0, 0
while i < len(lines):
m = RZ_RE.match(lines[i].strip())
if m:
d, x, sgn = m.groups()
j = i + 1
while j < len(lines) and not lines[j].strip():
j += 1
c = CVT_RE.match(lines[j].strip()) if j < len(lines) else None
if c and c.group(2) == d:
r = c.group(1)
out += [f"cvt.rzi.f32.f32 %rq{4*n}, {x};", f"sub.f32 %rq{4*n+1}, {x}, %rq{4*n};",
f"abs.f32 %rq{4*n+1}, %rq{4*n+1};", f"setp.ge.f32 %pq{n}, %rq{4*n+1}, 0f3F000000;",
f"add.f32 %rq{4*n+2}, {sgn}, {sgn};", f"add.f32 %rq{4*n+3}, %rq{4*n}, %rq{4*n+2};",
f"selp.f32 {r}, %rq{4*n+3}, %rq{4*n}, %pq{n};"]
n += 1
i = j + 1
continue
raise SystemExit(f"add.rz at line {i} not followed by its cvt.rzi")
out.append(lines[i])
i += 1
return out, n
ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
# Surface-store and round-idiom recipes for the other flag combinations.
SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
r"hiluma_engine_input_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)_v2_rel|"
r"rrlite_post_3_[12]_mv(hi|lo)_(hdr|ldr)(_folded)?|"
r"rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)|rrlite_dec0_4x4|"
r"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_"))
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 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()):
return path
sys.exit(f"no PTX module in {PTX_DIR} defines {name}")
def main():
name, out = sys.argv[1], Path(sys.argv[2])
k = 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"]:
sys.exit(f"line {cut} is {lines[cut - 1]!r}, expected {k['cut_check']!r}")
head, tail = lines[:cut], k["tail"].split("\n")
elif "delete" in k:
a, b = k["delete"]
if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]:
sys.exit(f"lines {a}/{b} are {lines[a - 1]!r}/{lines[b - 1]!r}, expected {k['delete_check']}")
head, tail = lines[:a - 1] + k["insert"].split("\n") + lines[b:], []
else:
head, tail = lines, []
entry = next(i for i, l in enumerate(head) if l.startswith(".visible .entry"))
body, n = replace_sust(head[entry:]) if k.get("sust") else (head[entry:], {"b16": 0, "p": 0, "b32": 0})
if k.get("rz_round"):
body, nrz = replace_rz_round(body)
# temporaries for the rewrite, declared after the kernel's own .reg lines
at = next(i for i, l in enumerate(body) if l.startswith(".reg"))
body = body[:at] + [f".reg .f32 %rq<{4 * nrz + 1}>;", f".reg .pred %pq<{nrz + 1}>;"] + body[at:]
n["rz"] = nrz
extern = k["extern"] + (SUST_EXTERN if n["b16"] else "") + (SUSTP_EXTERN if n["p"] else "") + (SUSTB32_EXTERN if n["b32"] else "")
text = "\n".join(head[:entry]) + "\n" + extern + "\n".join(body + tail)
out.write_text(text)
print(f"{out}: {len(text.splitlines())} lines, surface stores replaced: {n}")
if __name__ == "__main__":
main()
kernels/tex/sust_only.hip
+2
View File
@@ -0,0 +1,2 @@
// Native surface stores only (see make_ptx.py: post, downsample); the kernels are ZLUDA's compile.
#include "tex_common.h"
kernels/tex/tex_common.h
+448
View File
@@ -0,0 +1,448 @@
// Shared device code for native tails of the rrlite texture kernels (DLSS 310.7.0, preset M).
// These are linked into ZLUDA's own compile of a hand-edited PTX kernel (D4R_ZLUDA_EXTRA_BC), so
// only plain clang builtins are used (compiled with -nogpuinc -nogpulib).
#pragma once
#include <stdint.h>
#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)
#define DEV extern "C" __attribute__((device))
// Kernels that ZLUDA compiles in strict-FP mode (constrained intrinsics) cannot take inlined FP code from
// here (the AMDGPU backend fails on some constrained operations): build them with -DD4R_SUST_OUTLINE,
// which keeps the FP bodies of the surface stores in noinline leaf functions (no calls, no stack).
#if defined(D4R_SUST_OUTLINE) || defined(D4R_ACCURACY)
#define FP_BODY __attribute__((device, noinline)) static
#else
#define FP_BODY __attribute__((device, always_inline)) static inline
#endif
#define KERNEL extern "C" __attribute__((global))
typedef _Float16 half_t;
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
typedef _Float16 h16 __attribute__((ext_vector_type(16)));
typedef float f8v __attribute__((ext_vector_type(8)));
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
typedef uint32_t uint2_t __attribute__((ext_vector_type(2)));
typedef uint16_t u16x2 __attribute__((ext_vector_type(2)));
typedef int16_t i16x2 __attribute__((ext_vector_type(2)));
typedef __attribute__((address_space(3))) uint8_t lds_u8;
__attribute__((device)) static inline uint32_t lane_id()
{
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
}
// ZLUDA's e4m3x4_to_scaled_f16x2x2 (zluda_ptx_impl.cpp): e4m3 bytes -> f16 pairs holding value * 2^-8,
// {bytes 0,1} and {bytes 2,3}.
__attribute__((device)) static inline void e4m3x4_scaled(uint32_t packed, uint32_t& low, uint32_t& high)
{
uint32_t magnitude = packed & 0x7f7f7f7fu;
magnitude |= (magnitude + 0x01010101u) & 0x80808080u;
uint32_t sign = packed & 0x80808080u;
low = (__builtin_amdgcn_perm(0u, magnitude, 0x0c010c00u) << 7) | __builtin_amdgcn_perm(0u, sign, 0x010c000cu);
high = (__builtin_amdgcn_perm(0u, magnitude, 0x0c030c02u) << 7) | __builtin_amdgcn_perm(0u, sign, 0x030c020cu);
}
// A k16 WMMA operand from 16 consecutive e4m3 bytes of NVIDIA's k order (4 dwords): register j holds
// k 4j, 4j+1 and register 4 + j holds k 4j+2, 4j+3 -- ZLUDA's slot order (e4m3_a_k16_fragment).
__attribute__((device)) static inline u8v operand16(u4v d)
{
u8v v;
#pragma unroll
for (int j = 0; j < 4; ++j)
{
uint32_t lo, hi;
e4m3x4_scaled(d[j], lo, hi);
v[j] = lo;
v[4 + j] = hi;
}
return v;
}
__attribute__((device)) static inline f8v wmma(u8v a, u8v b, f8v c)
{
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
}
// ---------------------------------------------------------------- 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)
{
u16x2 half = __builtin_bit_cast(u16x2, bits);
u16x2 magnitude = half & (uint16_t)0x7fffu;
u16x2 normal = (magnitude + ((magnitude >> (uint16_t)7) & (uint16_t)1) + (uint16_t)0xe03fu) >> (uint16_t)7;
u16x2 exponent = magnitude >> (uint16_t)10;
u16x2 clamped = __builtin_elementwise_max(__builtin_elementwise_min(exponent, (u16x2)8), (u16x2)1);
u16x2 shift = (uint16_t)16 - clamped;
u16x2 significand = (magnitude & (uint16_t)0x3ffu) |
(__builtin_elementwise_min(magnitude, (u16x2)0x400) & (uint16_t)0x400u);
u16x2 subnormal = (significand + (((u16x2)1 << (shift - (uint16_t)1)) - (uint16_t)1) +
((significand >> shift) & (uint16_t)1)) >> shift;
u16x2 below = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(magnitude - (uint16_t)0x2400u)) >> (int16_t)15));
u16x2 code = (subnormal & below) | (normal & ~below);
code = __builtin_elementwise_min(code, (u16x2)0x7e);
u16x2 nan = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((uint16_t)0x7c00u - magnitude)) >> (int16_t)15));
code |= nan & (uint16_t)0x7fu;
code |= (half >> (uint16_t)8) & (uint16_t)0x80u;
return uint32_t(code.x) | (uint32_t(code.y) << 8);
}
__attribute__((device)) static inline uint32_t pack2(half_t a, half_t b)
{
return __builtin_bit_cast(uint32_t, (hv2){a, b});
}
// value from the lane with the same index in the other half of the wave
__attribute__((device)) static inline uint32_t other_half(uint32_t v)
{
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
}
// ---------------------------------------------------------------- surface stores
typedef int32_t v2i __attribute__((ext_vector_type(2)));
typedef int16_t v2s __attribute__((ext_vector_type(2)));
typedef float v4f __attribute__((ext_vector_type(4)));
typedef int32_t v4i_t __attribute__((ext_vector_type(4)));
typedef __attribute__((address_space(4))) unsigned int tsharp_t;
extern "C" __attribute__((device)) int __ockl_image_channel_data_type_2D(tsharp_t* image);
extern "C" __attribute__((device)) int __ockl_image_channel_order_2D(tsharp_t* image);
extern "C" __attribute__((device)) void __ockl_image_store_2D(tsharp_t* image, v2i coord, v4f value);
extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t surface, v2i coord, v2s value);
// ZLUDA's surface_half_bits_to_float (exact, independent of the f16 denormal mode)
__attribute__((device)) static inline float half_bits_to_float(uint32_t bits)
{
uint32_t sign = (bits & 0x8000u) << 16;
uint32_t exponent = (bits >> 10) & 0x1fu;
uint32_t mantissa = bits & 0x3ffu;
if (exponent == 0)
{
float magnitude = float(mantissa) * 0x1p-24f;
return __builtin_bit_cast(float, __builtin_bit_cast(uint32_t, magnitude) | sign);
}
if (exponent == 0x1f)
return __builtin_bit_cast(float, sign | 0x7f800000u | (mantissa << 13));
return __builtin_bit_cast(float, sign | ((exponent + 112u) << 23) | (mantissa << 13));
}
FP_BODY void sust_rg16f_body(tsharp_t* image, int32_t px, int32_t y, uint16_t a, uint16_t b)
{
__ockl_image_store_2D(image, (v2i){px, y}, (v4f){half_bits_to_float(a), half_bits_to_float(b), 0.0f, 0.0f});
}
// sust.b.2d.v2.b16.zero: a whole RG16F pixel is one image store of the decoded halves (what ZLUDA's
// surface_store_raw_2d does for that format); anything else goes through ZLUDA's generic helper.
__attribute__((device, noinline)) static void sust_v2b16_generic(uint64_t surface, int32_t x, int32_t y, uint16_t a, uint16_t b)
{
__zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(surface, (v2i){x, y}, (v2s){(int16_t)a, (int16_t)b});
}
DEV __attribute__((always_inline)) void d4r_sust_v2b16(uint64_t surface, int32_t x, int32_t y, uint16_t a, uint16_t b)
{
tsharp_t* image = (tsharp_t*)surface;
if (__ockl_image_channel_data_type_2D(image) == 14 && __ockl_image_channel_order_2D(image) == 3 && (x & 3) == 0)
{
if (x >= 0)
sust_rg16f_body(image, x >> 2, y, a, b);
}
else
sust_v2b16_generic(surface, x, y, a, b);
}
// f32 comparisons on the bits (no fcmp: kernels ZLUDA compiles in strict-FP mode cannot lower
// constrained fcmp); same results as the float compares, -0.0 included
__attribute__((device)) static inline bool f_isnan(uint32_t u) { return (u & 0x7fffffffu) > 0x7f800000u; }
__attribute__((device)) static inline bool f_lt0(uint32_t u) { return (int32_t)u < 0 && (u & 0x7fffffffu) != 0 && !f_isnan(u); }
// value > 1.0f (value not NaN)
__attribute__((device)) static inline bool f_gt1(uint32_t u) { return (int32_t)u > 0x3f800000; }
// value < -1.0f (value not NaN)
__attribute__((device)) static inline bool f_ltm1(uint32_t u) { return (int32_t)u < 0 && (u & 0x7fffffffu) > 0x3f800000u; }
// ZLUDA's surface_formatted_store_bits (zluda_ptx_impl.cpp), inlined: normalized CUDA formats convert
// the f32 value, everything else stores the bits unchanged.
__attribute__((device)) static inline uint32_t formatted_store_bits(uint32_t bits, uint32_t cuda_format, int channel)
{
if (cuda_format == 80)
{
uint32_t u = bits;
if (f_isnan(u) || f_lt0(u))
u = 0u;
if (f_gt1(u))
u = 0x3f800000u;
const float value = __builtin_bit_cast(float, u);
const float scale = channel == 3 ? 3.0f : 1023.0f;
// float(int(floor(x))) == floor(x) here (0 <= x <= 1023.5): no int <-> float conversions, which
// strict-FP kernels cannot lower as constrained signed conversions
return __builtin_bit_cast(uint32_t, __builtin_floorf(value * scale + 0.5f) / scale);
}
float scale = 0.0f;
bool signed_normalized = false;
if (cuda_format >= 192 && cuda_format <= 194)
scale = 255.0f;
else if (cuda_format >= 195 && cuda_format <= 197)
scale = 65535.0f;
else if (cuda_format >= 198 && cuda_format <= 200)
{
scale = 127.0f;
signed_normalized = true;
}
else if (cuda_format >= 201 && cuda_format <= 203)
{
scale = 32767.0f;
signed_normalized = true;
}
if (!(cuda_format >= 192 && cuda_format <= 203))
return bits;
uint32_t u = bits;
if (f_isnan(u))
u = 0u;
if (signed_normalized ? f_ltm1(u) : f_lt0(u))
u = signed_normalized ? 0xbf800000u : 0u;
if (f_gt1(u))
u = 0x3f800000u;
const float scaled = __builtin_bit_cast(float, u) * scale;
const uint32_t su = __builtin_bit_cast(uint32_t, scaled);
int32_t integer = f_lt0(su) ? -int32_t(__builtin_floorf(-scaled + 0.5f)) : int32_t(__builtin_floorf(scaled + 0.5f));
return uint32_t(integer);
}
// sust.p.2d.v4.b32.zero: ZLUDA stores the raw bits unless the original CUDA format was normalized
// (surface_formatted_store_bits; format word after the image and sampler descriptors).
FP_BODY void sust_p_body(uint64_t surface, int32_t x, int32_t y, uint32_t a, uint32_t b, uint32_t c, uint32_t d)
{
const uint32_t format = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
__ockl_image_store_2D((tsharp_t*)surface, (v2i){x, y},
(v4f){__builtin_bit_cast(float, formatted_store_bits(a, format, 0)),
__builtin_bit_cast(float, formatted_store_bits(b, format, 1)),
__builtin_bit_cast(float, formatted_store_bits(c, format, 2)),
__builtin_bit_cast(float, formatted_store_bits(d, format, 3))});
}
// d4r output redirect (wine_nvcuda_bridge.c d4rSetArrayRedirect): a tagged row pitch at byte 84 of the
// ZLUDA surface object (zero padding otherwise) sends stores of that surface to the linear memory whose
// address is at byte 88, rows `pitch` bytes apart. Out-of-range coordinates are dropped as by .zero stores.
extern "C" __attribute__((device)) int __ockl_image_width_2D(tsharp_t* image);
extern "C" __attribute__((device)) int __ockl_image_height_2D(tsharp_t* image);
__attribute__((device)) static inline uint32_t redirect_pitch(uint64_t surface)
{
// 'R2' tag in the high half (other handles may have data there), row pitch / 8 in the low half
const uint32_t word = ((const __attribute__((address_space(4))) uint32_t*)surface)[21];
return (word >> 16) == 0x5232u ? (word & 0xffffu) << 3 : 0u;
}
__attribute__((device)) static inline uint8_t* redirect_row(uint64_t surface, int32_t y)
{
const uint64_t base = *(const __attribute__((address_space(4))) uint64_t*)(surface + 88);
return (uint8_t*)(base + (uint64_t)(uint32_t)y * redirect_pitch(surface));
}
__attribute__((device)) static inline bool redirect_inside(uint64_t surface, int32_t x, int32_t y)
{
tsharp_t* image = (tsharp_t*)surface;
return x >= 0 && y >= 0 && x < __ockl_image_width_2D(image) && y < __ockl_image_height_2D(image);
}
DEV __attribute__((always_inline)) void d4r_sust_p_v4b32(uint64_t surface, int32_t x, int32_t y, uint32_t a, uint32_t b,
uint32_t c, uint32_t d)
{
const uint32_t format0 = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
if (redirect_pitch(surface) != 0 && format0 != 80 && !(format0 >= 192 && format0 <= 203))
{
// RGBA16F: the f32 values as halves, rounded toward zero like the image store's format conversion
if (redirect_inside(surface, x, y))
{
const uint32_t lo = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(__builtin_bit_cast(float, a), __builtin_bit_cast(float, b)));
const uint32_t hi = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(__builtin_bit_cast(float, c), __builtin_bit_cast(float, d)));
*(uint2_t*)(redirect_row(surface, y) + 8 * x) = (uint2_t){lo, hi};
}
return;
}
// formatted_store_bits keeps the bits unless the format is normalized: one uniform test per store
const uint32_t format = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
if (format != 80 && !(format >= 192 && format <= 203))
__ockl_image_store_2D((tsharp_t*)surface, (v2i){x, y},
(v4f){__builtin_bit_cast(float, a), __builtin_bit_cast(float, b), __builtin_bit_cast(float, c),
__builtin_bit_cast(float, d)});
else
sust_p_body(surface, x, y, a, b, c, d);
}
// ---------------------------------------------------------------- raw 32-bit surface stores
extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_b32_zero(uint64_t surface, v2i coord, uint32_t value);
// exact int -> float for |v| < 2^24 without a signed conversion (see formatted_store_bits)
__attribute__((device)) static inline float i2f(int32_t v)
{
return v < 0 ? -(float)(uint32_t)(-v) : (float)(uint32_t)v;
}
// ZLUDA's surface_decode_channel (bits of one memory channel -> the float4 lane value the store converts back)
__attribute__((device)) static inline float decode_channel(uint32_t bits, int data_type)
{
switch (data_type)
{
case 0:
return i2f(__builtin_elementwise_max(__builtin_elementwise_min(int32_t(int8_t(bits)), 127), -127)) / 127.0f;
case 1:
return i2f(__builtin_elementwise_max(__builtin_elementwise_min(int32_t(int16_t(bits)), 32767), -32767)) / 32767.0f;
case 2:
return (float)(bits & 0xffu) / 255.0f;
case 3:
return (float)(bits & 0xffffu) / 65535.0f;
case 8:
return __builtin_bit_cast(float, int32_t(int8_t(bits)));
case 9:
return __builtin_bit_cast(float, int32_t(int16_t(bits)));
case 14:
return half_bits_to_float(bits);
default:
return __builtin_bit_cast(float, bits);
}
}
// out of line: inlining ZLUDA's generic raw store into large kernels crashes AMDGPU instruction
// selection (see zluda_ptx_impl.cpp); it makes no calls, so the wrapper needs no stack
__attribute__((device, noinline)) static void sust_b32_generic(uint64_t surface, int32_t x, int32_t y, uint32_t data)
{
__zluda_ptx_impl_surfobj_b_2d_b32_zero(surface, (v2i){x, y}, data);
}
// sust.b.2d.b32.zero (x is a byte offset): whole 4-byte pixels and pairs of 2-byte pixels in R/RG/RGBA
// channel order need no read-modify-write (ZLUDA's surface_store_raw_2d, same conversions); anything
// else goes through ZLUDA's generic helper.
FP_BODY void sust_b32_fast(tsharp_t* image, int32_t x, int32_t y, uint32_t data, int dt, int cb, int pb)
{
if (pb == 4)
{
v4f v = {0.0f, 0.0f, 0.0f, 0.0f};
if (cb == 4)
v[0] = decode_channel(data, dt);
else if (cb == 2)
{
v[0] = decode_channel(data & 0xffffu, dt);
v[1] = decode_channel(data >> 16, dt);
}
else
{
v[0] = decode_channel(data & 0xffu, dt);
v[1] = decode_channel((data >> 8) & 0xffu, dt);
v[2] = decode_channel((data >> 16) & 0xffu, dt);
v[3] = decode_channel(data >> 24, dt);
}
__ockl_image_store_2D(image, (v2i){x >> 2, y}, v);
}
else
{
v4f v0 = {0.0f, 0.0f, 0.0f, 0.0f}, v1 = {0.0f, 0.0f, 0.0f, 0.0f};
if (cb == 2)
{
v0[0] = decode_channel(data & 0xffffu, dt);
v1[0] = decode_channel(data >> 16, dt);
}
else
{
v0[0] = decode_channel(data & 0xffu, dt);
v0[1] = decode_channel((data >> 8) & 0xffu, dt);
v1[0] = decode_channel((data >> 16) & 0xffu, dt);
v1[1] = decode_channel(data >> 24, dt);
}
__ockl_image_store_2D(image, (v2i){x >> 1, y}, v0);
__ockl_image_store_2D(image, (v2i){(x >> 1) + 1, y}, v1);
}
}
DEV __attribute__((always_inline)) void d4r_sust_b32(uint64_t surface, int32_t x, int32_t y, uint32_t data)
{
tsharp_t* image = (tsharp_t*)surface;
if (redirect_pitch(surface) != 0)
{
// raw bits at byte offset x (RGBA16F rows are 8 bytes per pixel)
if (x >= 0 && (x & 3) == 0 && redirect_inside(surface, x >> 3, y))
*(uint32_t*)(redirect_row(surface, y) + x) = data;
return;
}
const int dt = __ockl_image_channel_data_type_2D(image), order = __ockl_image_channel_order_2D(image);
const int cb = (dt == 0 || dt == 2 || dt == 8 || dt == 11) ? 1 : (dt == 1 || dt == 3 || dt == 9 || dt == 12 || dt == 14) ? 2
: (dt == 10 || dt == 13 || dt == 15) ? 4 : 0;
const int n = order == 1 ? 1 : order == 3 ? 2 : order == 8 ? 4 : 0; // R, RG, RGBA (identity lane order)
const int pb = cb * n;
if ((pb == 4 && (x & 3) == 0) || (pb == 2 && (x & 1) == 0))
{
if (x >= 0)
sust_b32_fast(image, x, y, data, dt, cb, pb);
}
else
sust_b32_generic(surface, x, y, data);
}
kernels/tools/dump_runner.cpp
+185
View File
@@ -0,0 +1,185 @@
// Launch a native kernel on a replay dump (manifest.txt, args.bin, alloc-N.bin), save every allocation
// after the first launch, then time repeated launches.
// usage: dump_runner <hsaco> <kernel> <dump dir> <out dir> [iterations]
#include <hip/hip_runtime.h>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#define CHECK(x) \
do \
{ \
hipError_t e_ = (x); \
if (e_ != hipSuccess) \
{ \
fprintf(stderr, "%s failed: %s\n", #x, hipGetErrorString(e_)); \
return 1; \
} \
} while (0)
static std::vector<char> read_file(const std::string& path)
{
std::ifstream f(path, std::ios::binary);
return std::vector<char>((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
}
int main(int argc, char** argv)
{
if (argc < 5)
{
fprintf(stderr, "usage: %s hsaco kernel dump out [iters]\n", argv[0]);
return 2;
}
const std::string dump = argv[3], out = argv[4];
const int iters = argc > 5 ? atoi(argv[5]) : 20;
hipModule_t mod;
hipFunction_t fn;
CHECK(hipModuleLoad(&mod, argv[1]));
CHECK(hipModuleGetFunction(&fn, mod, argv[2]));
hipFunction_t prep = nullptr;
if (hipModuleGetFunction(&prep, mod, (std::string(argv[2]) + "_prep").c_str()) != hipSuccess)
prep = nullptr;
// 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");
std::string line;
unsigned grid[3] = {1, 1, 1}, block[3] = {1, 1, 1}, shared = 0;
std::vector<int> ids;
std::vector<void*> dev;
std::vector<size_t> sizes;
while (std::getline(man, line))
{
std::istringstream s(line);
std::string kind;
s >> kind;
if (kind == "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")
{
int id;
std::string base;
size_t size;
s >> id >> base >> size;
std::vector<char> data = read_file(dump + "/alloc-" + std::to_string(id) + ".bin");
if (data.size() != size)
{
fprintf(stderr, "alloc %d size mismatch\n", id);
return 1;
}
void* d;
CHECK(hipMalloc(&d, size));
CHECK(hipMemcpy(d, data.data(), size, hipMemcpyHostToDevice));
ids.push_back(id);
dev.push_back(d);
sizes.push_back(size);
}
else if (kind == "pointer")
{
size_t off, delta;
int id;
s >> off >> id >> delta;
size_t i = std::find(ids.begin(), ids.end(), id) - ids.begin();
uint64_t v = (uint64_t)dev[i] + delta;
memcpy(args.data() + off, &v, 8);
}
}
{
hipDeviceptr_t bz;
size_t bzs = 0;
uint32_t v = 0;
if (hipModuleGetGlobal(&bz, &bzs, mod, "d4r_block_z") == hipSuccess && bzs == 4 &&
hipMemcpyDtoH(&v, bz, 4) == hipSuccess && v != 0)
block[2] = v;
}
{
hipDeviceptr_t gp;
size_t gs = 0;
uint32_t v = 0;
if (hipModuleGetGlobal(&gp, &gs, mod, "d4r_grid_x") == hipSuccess && gs == 4 &&
hipMemcpyDtoH(&v, gp, 4) == hipSuccess && v != 0)
{
grid[0] = v;
grid[1] = grid[2] = 1;
}
}
if (getenv("GRID_X"))
grid[0] = atoi(getenv("GRID_X"));
if (getenv("GRID_Y"))
grid[1] = atoi(getenv("GRID_Y"));
size_t argSize = args.size();
void* cfg[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, args.data(), HIP_LAUNCH_PARAM_BUFFER_SIZE, &argSize,
HIP_LAUNCH_PARAM_END};
const bool prepOnce = getenv("PREP_ONCE") != nullptr;
int launches = 0;
auto launch = [&]() {
if (prep != nullptr && !(prepOnce && launches++ > 0))
{
hipError_t e = hipModuleLaunchKernel(prep, prepGrid, 1, 1, 128, 1, 1, 0, nullptr, nullptr, cfg);
if (e != hipSuccess)
return e;
}
return hipModuleLaunchKernel(fn, grid[0], grid[1], grid[2], block[0], block[1], block[2], shared, nullptr,
nullptr, cfg);
};
CHECK(launch());
CHECK(hipDeviceSynchronize());
for (size_t i = 0; i < ids.size(); ++i)
{
std::vector<char> h(sizes[i]);
CHECK(hipMemcpy(h.data(), dev[i], sizes[i], hipMemcpyDeviceToHost));
std::ofstream(out + "/alloc-" + std::to_string(ids[i]) + ".bin", std::ios::binary).write(h.data(), h.size());
}
if (const char* dbg = getenv("DBG_GLOBAL"))
{
hipDeviceptr_t ptr;
size_t size = 0;
if (hipModuleGetGlobal(&ptr, &size, mod, dbg) == hipSuccess)
{
std::vector<char> h(size);
CHECK(hipMemcpyDtoH(h.data(), ptr, size));
std::ofstream(out + "/dbg.bin", std::ios::binary).write(h.data(), h.size());
}
}
hipEvent_t a, b;
CHECK(hipEventCreate(&a));
CHECK(hipEventCreate(&b));
std::vector<float> ms;
for (int i = 0; i < iters; ++i)
{
CHECK(hipEventRecord(a));
CHECK(launch());
CHECK(hipEventRecord(b));
CHECK(hipEventSynchronize(b));
float t;
CHECK(hipEventElapsedTime(&t, a, b));
ms.push_back(t);
}
std::sort(ms.begin(), ms.end());
if (!ms.empty())
printf("%s: median %.4f ms, min %.4f ms over %d launches\n", argv[2], ms[ms.size() / 2], ms[0], iters);
return 0;
}
kernels/tools/extract_dlss_ptx.py
+94
View File
@@ -0,0 +1,94 @@
#!/usr/bin/env python3
"""Extract the PTX modules embedded in NVIDIA's nvngx_dlss.dll.
usage: extract_dlss_ptx.py NVNGX_DLSS_DLL OUT_DIR
Writes OUT_DIR/dlss-NNNN-MM.ptx for PTX entry MM of the NNNN-th CUDA fatbin found in the DLL (in file
order). The texture-kernel build (kernels/tex) edits some of these modules; the DLL is NVIDIA's and is
never part of this repository, and neither are the extracted files.
"""
import os
import struct
import sys
FATBIN_MAGIC = 0xBA55ED50
def lz4_block(src: bytes, size: int) -> bytes:
dst = bytearray()
i = 0
while i < len(src):
token = src[i]
i += 1
literals = token >> 4
if literals == 15:
while True:
b = src[i]
i += 1
literals += b
if b != 255:
break
dst += src[i:i + literals]
i += literals
if i >= len(src) or len(dst) >= size:
break
offset = src[i] | (src[i + 1] << 8)
i += 2
match = token & 15
if match == 15:
while True:
b = src[i]
i += 1
match += b
if b != 255:
break
match += 4
for _ in range(match):
dst.append(dst[-offset])
return bytes(dst[:size])
def fatbins(data: bytes):
"""Yields (offset, bytes) of every fatbin container (magic, version 1) in file order."""
magic = struct.pack('<I', FATBIN_MAGIC)
start = data.find(magic)
while start >= 0:
_, version, header_size, size = struct.unpack_from('<IHHQ', data, start)
if version == 1 and header_size == 16 and start + header_size + size <= len(data):
yield start, data[start:start + header_size + size]
start = data.find(magic, start + header_size + size)
else:
start = data.find(magic, start + 4)
def ptx_entries(fatbin: bytes):
"""Yields the (decompressed) PTX text of each PTX entry of one fatbin."""
_, _, header_size, size = struct.unpack_from('<IHHQ', fatbin, 0)
offset = header_size
while offset < header_size + size:
kind, _, entry_header, entry_size = struct.unpack_from('<HHIQ', fatbin, offset)
flags = struct.unpack_from('<Q', fatbin, offset + 40)[0]
decompressed = struct.unpack_from('<Q', fatbin, offset + 56)[0] if entry_header >= 64 else 0
payload = fatbin[offset + entry_header:offset + entry_header + entry_size]
if kind == 1: # PTX
yield lz4_block(payload, decompressed) if (flags & 0x2000 and decompressed) else payload
offset += entry_header + entry_size
def main(argv):
if len(argv) != 3:
sys.exit(__doc__)
data = open(argv[1], 'rb').read()
os.makedirs(argv[2], exist_ok=True)
total = 0
for index, (_, fatbin) in enumerate(fatbins(data)):
for entry, ptx in enumerate(ptx_entries(fatbin)):
ptx = ptx.rstrip(b'\0')
with open(os.path.join(argv[2], f'dlss-{index:04d}-{entry:02d}.ptx'), 'wb') as out:
out.write(ptx)
total += 1
print(f'{total} PTX modules written to {argv[2]}')
if __name__ == '__main__':
main(sys.argv)
kernels/tools/kernel_manifest.py
+69
View File
@@ -0,0 +1,69 @@
#!/usr/bin/env python3
"""Write the d4r-kernels.txt manifest of a native kernel directory.
usage: kernel_manifest.py KERNEL_DIR NVNGX_DLSS_DLL [NVNGX_DLSS_DLL...]
For every NAME.hsaco in KERNEL_DIR, finds the PTX module of each DLL that defines `.entry NAME` and
lists the FNV-1a 64 hash of the module's text (trailing NUL bytes removed), one "NAME HASH" line per
distinct hash. The Wine CUDA bridge serves a native kernel only while DLSS loads a module with a
listed hash, so a DLSS version that changed the kernel runs ZLUDA's own compile of it instead.
The manifest is written to KERNEL_DIR/d4r-kernels.txt; it holds hashes, not NVIDIA code.
"""
import hashlib
import os
import re
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from extract_dlss_ptx import fatbins, ptx_entries # noqa: E402
ENTRY = re.compile(rb'\.entry[ \t\n]+([A-Za-z0-9_$]+)[ \t\n\r]*\(')
def fnv1a64(data: bytes) -> int:
value = 0xcbf29ce484222325
for byte in data:
value = ((value ^ byte) * 0x100000001b3) & 0xffffffffffffffff
return value
def module_hashes(dll: str) -> dict:
"""NAME -> set of module hashes for every PTX entry of one DLL."""
data = open(dll, 'rb').read()
found = {}
for _, fatbin in fatbins(data):
for ptx in ptx_entries(fatbin):
text = ptx.rstrip(b'\0')
value = fnv1a64(text)
for match in ENTRY.finditer(text):
found.setdefault(match.group(1).decode(), set()).add(value)
return found
def main(argv):
if len(argv) < 3:
sys.exit(__doc__)
directory = argv[1]
names = sorted(name[:-len('.hsaco')] for name in os.listdir(directory) if name.endswith('.hsaco'))
lines = ['# d4r native kernels: NAME and the FNV-1a 64 hash of the DLSS PTX module each was written for',
'# (written by kernels/tools/kernel_manifest.py from:']
listed = {name: set() for name in names}
for dll in argv[2:]:
digest = hashlib.sha256(open(dll, 'rb').read()).hexdigest()[:16]
lines.append(f'# {os.path.basename(dll)} sha256 {digest}...)')
hashes = module_hashes(dll)
for name in names:
listed[name] |= hashes.get(name, set())
missing = [name for name in names if not listed[name]]
if missing:
sys.exit(f'no DLL defines {", ".join(missing)}; not writing a manifest')
for name in names:
for value in sorted(listed[name]):
lines.append(f'{name} {value:016x}')
with open(os.path.join(directory, 'd4r-kernels.txt'), 'w') as out:
out.write('\n'.join(lines) + '\n')
print(f'{len(names)} kernels listed in {os.path.join(directory, "d4r-kernels.txt")}')
if __name__ == '__main__':
main(sys.argv)
kernels/tools/kprof.py
+13
View File
@@ -0,0 +1,13 @@
#!/usr/bin/env python3
# usage: kprof.py RUN.log FRAMES [WARM] : per-kernel median GPU ms and per-frame total from KERNEL_PROFILE lines
import sys,re,statistics as st
lines=[l for l in open(sys.argv[1],errors='replace') if 'KERNEL_PROFILE' in l]
frames=int(sys.argv[2]); warm=int(sys.argv[3]) if len(sys.argv)>3 else frames//2
rows=[(int(m.group(1)),float(m.group(2)),m.group(3)) for m in (re.search(r'launch=(\d+).*gpu_ms=([0-9.]+) name=(\S+)',l) for l in lines) if m]
rows.sort(); n=len(rows); per=n/frames; cut=int(per*warm); rows=rows[cut:]; f=frames-warm
agg={}
for _,ms,name in rows: agg.setdefault(name,[]).append(ms)
tot=sum(sum(v) for v in agg.values())/f
print(f"launches/frame={per:.1f} total_gpu_ms/frame={tot:.3f} (frames {warm+1}-{frames})")
for name,v in sorted(agg.items(),key=lambda kv:-sum(kv[1])):
print(f" {sum(v)/f:8.3f} ms/frame {len(v)/f:5.1f}x med {st.median(v):7.3f} {name}")
kernels/tools/model_enc3.py
+141
View File
@@ -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/psnr.py
+15
View File
@@ -0,0 +1,15 @@
#!/usr/bin/env python3
"""psnr.py DIR_A DIR_B: per-frame PSNR of rgba16f outputs (watermark box excluded, values clipped to [0, 1] after tonemap-free clamp)."""
import sys, glob, os, numpy as np
a, b = sys.argv[1:3]
vals = []
for f in sorted(glob.glob(a + '/frame-*.rgba16f')):
g = os.path.join(b, os.path.basename(f))
if not os.path.exists(g): continue
x = np.fromfile(f, np.float16).reshape(1440, 2560, 4)[..., :3].astype(np.float64)
y = np.fromfile(g, np.float16).reshape(1440, 2560, 4)[..., :3].astype(np.float64)
m = np.ones(x.shape[:2], bool); m[1370:1392, 1850:2460] = False
x, y = np.clip(x[m], 0, 1), np.clip(y[m], 0, 1)
mse = np.mean((x - y) ** 2)
vals.append(10 * np.log10(1.0 / mse) if mse > 0 else 99.0)
print('frames', len(vals), 'PSNR dB min %.2f mean %.2f' % (min(vals), np.mean(vals)), ' '.join('%.1f' % v for v in vals))
kernels/tools/ptxload.c
+20
View File
@@ -0,0 +1,20 @@
// usage: ptxload FILE.ptx... : cuModuleLoadData + cuModuleGetFunction for every .entry
#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
typedef int (*I)(unsigned); typedef int (*DG)(int*,int); typedef int (*CC)(void**,unsigned,int);
typedef int (*ML)(void**,const void*); typedef int (*GF)(void**,void*,const char*); typedef int (*MU)(void*);
int main(int argc,char**argv){
void*l=dlopen("libcuda.so",RTLD_NOW); if(!l){fprintf(stderr,"%s\n",dlerror());return 1;}
I init=dlsym(l,"cuInit"); DG dg=dlsym(l,"cuDeviceGet"); CC cc=dlsym(l,"cuCtxCreate_v2");
ML ml=dlsym(l,"cuModuleLoadData"); GF gf=dlsym(l,"cuModuleGetFunction"); MU mu=dlsym(l,"cuModuleUnload");
int d; void*ctx; init(0); dg(&d,0); cc(&ctx,0,d); int bad=0;
for(int a=1;a<argc;a++){
FILE*f=fopen(argv[a],"rb"); fseek(f,0,2); long n=ftell(f); rewind(f); char*s=malloc(n+1); fread(s,1,n,f); s[n]=0; fclose(f);
void*m=0; int r=ml(&m,s); printf("%s load=%d",argv[a],r); if(r){bad++;puts("");free(s);continue;}
for(char*p=strstr(s,".entry ");p;p=strstr(p+1,".entry ")){
char name[256]; sscanf(p+7,"%255[A-Za-z0-9_$]",name); void*fn=0; int g=gf(&fn,m,name);
printf(" %s=%d",name,g); if(g) bad++; }
puts(""); mu(m); free(s);}
return bad?1:0;}
kernels/tools/ptxsim.py
+883
View File
@@ -0,0 +1,883 @@
#!/usr/bin/env python3
"""Minimal vectorised PTX interpreter for DLSS rrlite kernels.
Runs one thread block of a fully unrolled kernel on the CPU (all threads of the
block as numpy lanes), against memory restored from a bridge replay dump
(D4R_CUDA_REPLAY_DUMP_DIR). It is a reference and an instrument: every
statement can be observed (hooks), which is how the logical computation
(GEMMs, weight layouts, normalisations) of a kernel is recovered.
Only the instruction forms the rrlite kernels use are implemented; anything
else raises NotImplementedError naming the statement.
"""
import re
import struct
import sys
import numpy as np
F16, F32, F64 = np.float16, np.float32, np.float64
U64 = np.uint64
PARAM_BASE = 1 << 60 # virtual address of the kernel parameter block
MASK = {8: 0xff, 16: 0xffff, 32: 0xffffffff, 64: 0xffffffffffffffff}
def _e4m3_decode_table():
t = np.zeros(256, F64)
for c in range(256):
s = -1.0 if c & 0x80 else 1.0
e, m = (c >> 3) & 0xf, c & 7
if (c & 0x7f) == 0x7f:
t[c] = np.nan
elif e == 0:
t[c] = s * m * 2.0 ** -9
else:
t[c] = s * (1 + m / 8) * 2.0 ** (e - 7)
return t
E4M3 = _e4m3_decode_table()
def _f16_to_e4m3_table():
"""Exact RNE + satfinite f16 -> e4m3 for all 65536 f16 codes (ZLUDA's helper)."""
h = np.arange(65536, dtype=np.uint32)
mag = h & 0x7fff
normal = (mag - 0x2000 + 0x3f + ((mag >> 7) & 1)) >> 7
exp = mag >> 10
sig = (mag & 0x3ff) | np.where(exp != 0, 0x400, 0)
cl = np.clip(exp, 1, 8)
sh = 16 - cl
sub = (sig + (np.left_shift(1, sh - 1)) - 1 + ((sig >> sh) & 1)) >> sh
code = np.where(mag >= 0x2400, normal, sub)
code = np.where(mag > 0x5f00, 0x7e, code)
code = np.where(mag > 0x7c00, 0x7f, code)
return (((h >> 8) & 0x80) | code).astype(np.uint32)
F16_TO_E4M3 = _f16_to_e4m3_table()
def bits_to_f16(x):
return (x & 0xffff).astype(np.uint16).view(F16)
def f16_to_bits(v):
return np.asarray(v, F16).view(np.uint16).astype(U64)
def bits_to_f32(x):
return (x & 0xffffffff).astype(np.uint32).view(F32)
def f32_to_bits(v):
return np.asarray(v, F32).view(np.uint32).astype(U64)
class Stmt:
__slots__ = ('labels', 'pred', 'neg', 'op', 'args', 'text', 'index')
def split_args(s):
out, depth, cur = [], 0, ''
for ch in s:
if ch in '{[':
depth += 1
elif ch in '}]':
depth -= 1
if ch == ',' and depth == 0:
out.append(cur.strip())
cur = ''
else:
cur += ch
if cur.strip():
out.append(cur.strip())
return out
def parse_kernel(path, name):
with open(path) as file:
src = file.read()
i = src.index('.entry ' + name + '(')
j = src.find('\n.visible', i + 10)
j = len(src) if j < 0 else j
body = src[i:j]
body = body[body.index(')') + 1:]
# module-scope .shared arrays (declared before the entry) are visible to it too
module_shared = {m.group(1): int(m.group(2)) for m in re.finditer(r'^\.shared[^\n;]*?(\w+)\[(\d+)\]', src[:i], re.M)}
body = re.sub(r'//[^\n]*', '', body)
body = re.sub(r'\.(maxntid|reqntid|minnctapersm|maxnreg)[^\n{]*', '', body)
stmts, types, shared = [], {}, {}
pending_labels = []
for raw in re.sub(r'\s*\n\s*', ' ', body).split(';'):
s = raw.strip()
while True:
s = s.strip()
if s.startswith('{') and not s.startswith('{%'):
s = s[1:]
continue
if s.startswith('}'):
s = s[1:]
continue
m = re.match(r'(\$[\w]+):\s*', s)
if m:
pending_labels.append(m.group(1))
s = s[m.end():]
continue
break
if not s:
continue
if s.startswith('.reg'):
m = re.match(r'\.reg\s*\.(\w+)\s+(.*)', s)
ty = m.group(1)
for decl in m.group(2).split(','):
decl = decl.strip()
m2 = re.match(r'(%[\w]+)<(\d+)>', decl)
if m2:
types[('prefix', m2.group(1))] = ty
else:
types[decl] = ty
continue
if s.startswith('.shared'):
m = re.search(r'(\w+)\[(\d+)\]', s)
shared[m.group(1)] = int(m.group(2))
continue
if s.startswith('.'):
continue
st = Stmt()
st.labels, pending_labels = pending_labels, []
st.pred, st.neg = None, False
m = re.match(r'@(!?)(%\w+)\s+', s)
if m:
st.neg, st.pred = m.group(1) == '!', m.group(2)
s = s[m.end():]
parts = s.split(None, 1)
st.op = parts[0]
st.args = split_args(parts[1]) if len(parts) > 1 else []
st.text = s
st.index = len(stmts)
stmts.append(st)
if pending_labels:
st = Stmt()
st.labels, st.pred, st.neg, st.op, st.args, st.text, st.index = pending_labels, None, False, 'nop', [], 'nop', len(stmts)
stmts.append(st)
for k, v in module_shared.items():
shared.setdefault(k, v)
return stmts, types, shared
def reg_width(types, name):
ty = types.get(name)
if ty is None:
m = re.match(r'(%[a-z]+)', name)
ty = types.get(('prefix', m.group(1)))
return {'pred': 1, 'b16': 16, 'u16': 16, 's16': 16, 'f16': 16, 'b32': 32, 'u32': 32, 's32': 32, 'f32': 32,
'b64': 64, 'u64': 64, 's64': 64, 'f64': 64, 'b8': 8}[ty]
class Memory:
"""Global memory: the replay dump's allocations at their captured addresses."""
def __init__(self, dump_dir):
manifest = open(dump_dir + '/manifest.txt').read().split('\n')
self.args = open(dump_dir + '/args.bin', 'rb').read()
self.bases, self.data, self.ids = [], [], []
for line in manifest:
p = line.split()
if p and p[0] == 'alloc':
idx, base, size = int(p[1]), int(p[2], 16), int(p[3])
buf = np.fromfile(f'{dump_dir}/alloc-{idx}.bin', np.uint8)
assert buf.size == size
self.bases.append(base)
self.data.append(buf)
self.ids.append(idx)
order = np.argsort(self.bases)
self.bases = [self.bases[i] for i in order]
self.data = [self.data[i] for i in order]
self.ids = [self.ids[i] for i in order]
self.base_arr = np.array(self.bases, U64)
self.stores = [] # (addresses, bytes) of every global store
def locate(self, addr):
idx = np.searchsorted(self.base_arr, addr, side='right') - 1
if np.any(idx < 0):
raise ValueError('address below every allocation')
return idx
def load(self, addr, nbytes, mask):
out = np.zeros((addr.size, nbytes), np.uint8)
idx = self.locate(np.where(mask, addr, self.base_arr[0]))
for a in np.unique(idx[mask]):
sel = mask & (idx == a)
off = (addr[sel] - self.bases[a]).astype(np.int64)
if np.any(off + nbytes > self.data[a].size):
raise ValueError('load past the end of an allocation')
out[sel] = self.data[a][off[:, None] + np.arange(nbytes)]
return out
def store(self, addr, raw, mask):
idx = self.locate(np.where(mask, addr, self.base_arr[0]))
n = raw.shape[1]
for a in np.unique(idx[mask]):
sel = mask & (idx == a)
off = (addr[sel] - self.bases[a]).astype(np.int64)
self.data[a][off[:, None] + np.arange(n)] = raw[sel]
self.stores.append((self.ids[a], off.copy(), raw[sel].copy()))
class Block:
def __init__(self, stmts, types, shared, mem, param_name, block_dim, grid_dim, ctaid):
self.stmts, self.types, self.mem = stmts, types, mem
self.param_name = param_name
bx, by, bz = block_dim
self.T = bx * by * bz
t = np.arange(self.T)
self.tid = {'x': t % bx, 'y': (t // bx) % by, 'z': t // (bx * by)}
self.lane = t % 32
self.warp = t // 32
self.ctaid = ctaid
self.grid = grid_dim
self.ntid = block_dim
self.regs = {}
self.shared = {k: np.zeros(v, np.uint8) for k, v in shared.items()}
self.shared_base = {}
base = 0
for k, v in shared.items():
self.shared_base[k] = base
base += (v + 15) & ~15
self.shared_mem = np.zeros(max(base, 16), np.uint8)
self.labels = {}
for st in stmts:
for l in st.labels:
self.labels[l] = st.index
self.hooks = []
self.retired = np.zeros(self.T, bool)
self.prov = {}
# --- operands -----------------------------------------------------------
def width(self, name):
return reg_width(self.types, name)
def special(self, name):
m = re.match(r'%(tid|ntid|ctaid|nctaid)\.([xyz])', name)
if m:
kind, ax = m.groups()
k = 'xyz'.index(ax)
if kind == 'tid':
return self.tid[ax].astype(U64)
if kind == 'ntid':
return np.full(self.T, self.ntid[k], U64)
if kind == 'ctaid':
return np.full(self.T, self.ctaid[k], U64)
return np.full(self.T, self.grid[k], U64)
if name == 'WARP_SZ':
return np.full(self.T, 32, U64)
if name == '%laneid':
return self.lane.astype(U64)
if name == '%warpid':
return self.warp.astype(U64)
return None
def imm(self, s, width=64):
if s.startswith('0f'):
return np.full(self.T, int(s[2:], 16), U64)
if s.startswith('0d'):
return np.full(self.T, int(s[2:], 16), U64)
v = int(s, 0)
return np.full(self.T, v & MASK[64], U64)
def get(self, s, width=None):
s = s.strip()
if s.startswith('%') or s in self.types or s == 'WARP_SZ':
sp = self.special(s) if (s.startswith('%') or s == 'WARP_SZ') else None
if sp is not None:
return sp
if s not in self.regs:
w = self.width(s)
self.regs[s] = np.zeros(self.T, bool if w == 1 else U64)
return self.regs[s]
if s in self.shared_base:
return np.full(self.T, self.shared_base[s], U64)
if s == self.param_name:
return np.full(self.T, PARAM_BASE, U64)
if s[0].isdigit() or s[0] == '-':
v = self.imm(s)
if width:
v = v & np.uint64(MASK[width])
return v
raise NotImplementedError('operand ' + s)
def put(self, name, value, mask, keep_prov=False):
name = name.strip()
if not keep_prov:
self.prov.pop(name, None)
w = self.width(name)
old = self.get(name)
if w == 1:
self.regs[name] = np.where(mask, value.astype(bool), old)
else:
self.regs[name] = np.where(mask, value.astype(U64) & np.uint64(MASK[w]), old)
def vec(self, s):
return [x.strip() for x in s.strip()[1:-1].split(',')]
def address(self, s):
"""[reg+off] or [symbol+off] -> (space, address array)."""
inner = s.strip()[1:-1]
m = re.match(r'([%\w$.]+)\s*(\+\s*-?\d+)?', inner)
base, off = m.group(1), int(m.group(2).replace('+', '')) if m.group(2) else 0
if base == self.param_name:
return 'param', np.full(self.T, off, U64)
if base in self.shared_base:
return 'shared', np.full(self.T, self.shared_base[base] + off, U64)
return 'reg', (self.get(base) + np.uint64(off & MASK[64])) & np.uint64(MASK[64])
# --- execution ------------------------------------------------------------
def run(self, limit=None):
stmts = self.stmts
active = np.ones(self.T, bool)
pending = {}
pc = 0
steps = 0
while pc < len(stmts):
st = stmts[pc]
if pc in pending:
active |= pending.pop(pc)
if not active.any():
# Reconverge at either a forward target or a loop fallthrough.
targets = sorted(pending)
if not targets:
break
pc = targets[0]
continue
active &= ~self.retired
mask = active
if st.pred is not None:
p = self.get(st.pred)
mask = active & (~p if st.neg else p)
if st.op == 'bra' or st.op.startswith('bra.'):
target = self.labels[st.args[0]]
taking = mask
if target <= pc:
if (taking == active).all():
pc = target
continue
if not taking.any():
pc += 1
continue
# Lanes finishing a loop wait at its fallthrough while the
# other lanes execute their remaining iterations. Their
# registers and memory writes stay masked until reconvergence.
pending[pc + 1] = pending.get(pc + 1, np.zeros(self.T, bool)) | (active & ~taking)
active = taking.copy()
pc = target
continue
active = active & ~taking
if taking.any():
pending[target] = pending.get(target, np.zeros(self.T, bool)) | taking
pc += 1
continue
if mask.any() or st.op.startswith(('mma', 'shfl', 'movmatrix', 'bar')):
self.execute(st, mask)
for hook in self.hooks:
hook(self, st, mask)
pc += 1
steps += 1
if limit and steps >= limit:
break
return steps
def execute(self, st, mask):
op, a = st.op, st.args
parts = op.split('.')
base = parts[0]
fn = getattr(self, 'op_' + base, None)
if fn is None:
raise NotImplementedError(st.text)
try:
fn(st, parts, a, mask)
except Exception as e:
raise RuntimeError(f'{type(e).__name__} {e} in: {st.text}') from e
# --- helpers --------------------------------------------------------------
@staticmethod
def sext(v, w):
v = v.astype(np.int64) & ((1 << w) - 1) if w < 64 else v.view(np.int64)
if w < 64:
v = np.where(v >= (1 << (w - 1)), v - (1 << w), v)
return v
def typed(self, s, ty):
w = int(re.sub(r'\D', '', ty) or 32) if ty not in ('pred',) else 1
v = self.get(s, w)
if ty.startswith('s'):
return self.sext(v, w)
if ty.startswith('f'):
if ty == 'f16':
return bits_to_f16(v).astype(F64)
if ty == 'f32':
return bits_to_f32(v).astype(F64)
return v.view(F64)
return (v & np.uint64(MASK[w])).astype(np.int64) if w < 64 else v
def result(self, ty, value):
if ty == 'f16':
return f16_to_bits(np.asarray(value).astype(F16))
if ty == 'f32':
return f32_to_bits(np.asarray(value).astype(F32))
if ty == 'f64':
return np.asarray(value, F64).view(U64)
w = int(re.sub(r'\D', '', ty))
return (np.asarray(value).astype(np.int64).astype(U64)) & np.uint64(MASK[w])
# --- instructions -----------------------------------------------------------
def op_nop(self, st, parts, a, mask):
pass
def op_bar(self, st, parts, a, mask):
pass
def op_ret(self, st, parts, a, mask):
self.retired |= mask
def op_neg(self, st, parts, a, mask):
ty = parts[-1]
if ty == 'f16x2':
self.put(a[0], self.get(a[1]) ^ np.uint64(0x80008000), mask)
elif ty == 'f16':
self.put(a[0], self.get(a[1]) ^ np.uint64(0x8000), mask)
elif ty == 'f32':
self.put(a[0], self.get(a[1]) ^ np.uint64(0x80000000), mask)
else:
self.put(a[0], self.result('u' + ty[1:], -self.typed(a[1], ty)), mask)
def op_cvta(self, st, parts, a, mask):
self.put(a[0], self.get(a[1]), mask)
def op_mov(self, st, parts, a, mask):
d, s = a
if d.startswith('{'):
names = self.vec(d)
v = self.get(s)
w = self.width(names[0])
for i, n in enumerate(names):
self.put(n, (v >> np.uint64(i * w)) & np.uint64(MASK[w]), mask)
return
if s.startswith('{'):
names = self.vec(s)
w = self.width(names[0])
v = np.zeros(self.T, U64)
for i, n in enumerate(names):
v |= (self.get(n) & np.uint64(MASK[w])) << np.uint64(i * w)
self.put(d, v, mask)
return
self.put(d, self.get(s), mask)
if s in self.prov:
self.prov[d] = self.prov[s]
def op_ld(self, st, parts, a, mask):
space = 'global' if 'global' in parts else 'shared' if 'shared' in parts else 'param' if 'param' in parts else None
ty = parts[-1]
vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
w = int(re.sub(r'\D', '', ty))
nbytes = w // 8
kind, addr = self.address(a[1])
total = nbytes * vec
if kind == 'param':
raw = np.frombuffer(self.mem.args, np.uint8)
off = int(addr[0])
data = np.tile(raw[off:off + total], (self.T, 1))
elif space == 'param':
# ld.param through a computed address (mov.b64 %rd, <param>; add; ld.param [%rd])
raw = np.frombuffer(self.mem.args, np.uint8)
off = np.where(mask, addr - np.uint64(PARAM_BASE), 0).astype(np.int64)
data = raw[off[:, None] + np.arange(total)]
elif kind == 'shared' or space == 'shared':
off = addr.astype(np.int64)
data = self.shared_mem[off[:, None] + np.arange(total)]
else:
data = self.mem.load(addr, total, mask)
names = self.vec(a[0]) if a[0].startswith('{') else [a[0]]
for i, n in enumerate(names):
chunk = data[:, i * nbytes:(i + 1) * nbytes]
v = np.zeros(self.T, U64)
for b in range(nbytes):
v |= chunk[:, b].astype(U64) << np.uint64(8 * b)
if ty.startswith('s') and w < self.width(n):
v = (self.sext(v, w).astype(np.int64)).astype(U64)
self.put(n, v, mask)
self.prov[n] = (st.index, kind if kind != 'reg' else space, addr + np.uint64(i * nbytes), nbytes)
def op_st(self, st, parts, a, mask):
ty = parts[-1]
vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
w = int(re.sub(r'\D', '', ty))
nbytes = w // 8
kind, addr = self.address(a[0])
names = self.vec(a[1]) if a[1].startswith('{') else [a[1]]
raw = np.zeros((self.T, nbytes * len(names)), np.uint8)
for i, n in enumerate(names):
v = self.get(n, w)
for b in range(nbytes):
raw[:, i * nbytes + b] = ((v >> np.uint64(8 * b)) & np.uint64(0xff)).astype(np.uint8)
if kind == 'shared' or 'shared' in parts:
off = addr.astype(np.int64)[mask]
self.shared_mem[off[:, None] + np.arange(raw.shape[1])] = raw[mask]
else:
self.mem.store(addr, raw, mask)
def int_binop(self, parts, a, mask, fn):
ty = parts[-1]
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
self.put(a[0], self.result(ty if not ty.startswith('b') else 'u' + ty[1:], fn(x, y)), mask)
def op_add(self, st, parts, a, mask):
ty = parts[-1]
if ty in ('f16', 'f32'):
self.put(a[0], self.result(ty, self.typed(a[1], ty) + self.typed(a[2], ty)), mask)
elif ty == 'f16x2':
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x + y), mask)
else:
self.int_binop(parts, a, mask, lambda x, y: x + y)
def op_sub(self, st, parts, a, mask):
ty = parts[-1]
if ty in ('f16', 'f32'):
self.put(a[0], self.result(ty, self.typed(a[1], ty) - self.typed(a[2], ty)), mask)
elif ty == 'f16x2':
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x - y), mask)
else:
self.int_binop(parts, a, mask, lambda x, y: x - y)
def op_mul(self, st, parts, a, mask):
ty = parts[-1]
if 'wide' in parts:
w = int(re.sub(r'\D', '', ty))
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
self.put(a[0], (x * y).astype(np.int64).astype(U64), mask)
elif ty in ('f16', 'f32'):
self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty)), mask)
elif ty == 'f16x2':
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x * y), mask)
else:
self.int_binop(parts, a, mask, lambda x, y: x * y)
def op_mad(self, st, parts, a, mask):
ty = parts[-1]
x, y, z = self.typed(a[1], ty), self.typed(a[2], ty), self.typed(a[3], ty)
self.put(a[0], self.result('u' + ty[1:], x * y + z), mask)
def h2(self, sa, sb, fn, sc=None):
va, vb = self.get(sa), self.get(sb)
out = np.zeros(self.T, U64)
for i in range(2):
x = bits_to_f16(va >> np.uint64(16 * i)).astype(F64)
y = bits_to_f16(vb >> np.uint64(16 * i)).astype(F64)
if sc is not None:
z = bits_to_f16(self.get(sc) >> np.uint64(16 * i)).astype(F64)
r = fn(x, y, z)
else:
r = fn(x, y)
out |= f16_to_bits(np.asarray(r).astype(F16)) << np.uint64(16 * i)
return out
def op_fma(self, st, parts, a, mask):
ty = parts[-1]
if ty == 'f16x2':
self.put(a[0], self.h2(a[1], a[2], lambda x, y, z: x * y + z, a[3]), mask)
elif ty in ('f16', 'f32'):
self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty) + self.typed(a[3], ty)), mask)
else:
raise NotImplementedError(st.text)
def minmax(self, parts, a, mask, is_max):
ty = parts[-1]
if ty == 'f16x2':
f = np.fmax if is_max else np.fmin
self.put(a[0], self.h2(a[1], a[2], f), mask)
elif ty in ('f16', 'f32'):
f = np.fmax if is_max else np.fmin
self.put(a[0], self.result(ty, f(self.typed(a[1], ty), self.typed(a[2], ty))), mask)
else:
f = np.maximum if is_max else np.minimum
self.int_binop(parts, a, mask, f)
def op_min(self, st, parts, a, mask):
self.minmax(parts, a, mask, False)
def op_max(self, st, parts, a, mask):
self.minmax(parts, a, mask, True)
def op_abs(self, st, parts, a, mask):
ty = parts[-1]
if ty in ('f16x2', 'f16', 'f32'):
m = {'f16x2': 0x7fff7fff, 'f16': 0x7fff, 'f32': 0x7fffffff}[ty]
self.put(a[0], self.get(a[1]) & np.uint64(m), mask)
return
self.put(a[0], self.result('u' + ty[1:], np.abs(self.typed(a[1], ty))), mask)
def logic(self, a, mask, fn):
if self.width(a[0]) == 1:
self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
else:
self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
def op_and(self, st, parts, a, mask):
self.logic(a, mask, lambda x, y: x & y)
def op_or(self, st, parts, a, mask):
self.logic(a, mask, lambda x, y: x | y)
def op_xor(self, st, parts, a, mask):
self.logic(a, mask, lambda x, y: x ^ y)
def op_not(self, st, parts, a, mask):
self.put(a[0], ~self.get(a[1]), mask)
def op_shl(self, st, parts, a, mask):
w = int(re.sub(r'\D', '', parts[-1]))
n = self.get(a[2]) & np.uint64(0xff)
v = np.where(n >= w, 0, (self.get(a[1]) << np.minimum(n, np.uint64(63))))
self.put(a[0], v.astype(U64), mask)
def op_shr(self, st, parts, a, mask):
ty = parts[-1]
w = int(re.sub(r'\D', '', ty))
n = (self.get(a[2]) & np.uint64(0xff)).astype(np.int64)
if ty.startswith('s'):
v = self.sext(self.get(a[1]), w)
r = v >> np.minimum(n, w - 1)
self.put(a[0], r.astype(np.int64).astype(U64), mask)
else:
v = self.get(a[1]) & np.uint64(MASK[w])
r = np.where(n >= w, 0, v >> np.minimum(n, 63).astype(U64))
self.put(a[0], r.astype(U64), mask)
def op_bfi(self, st, parts, a, mask):
f, b = self.get(a[1]), self.get(a[2])
pos = self.get(a[3]).astype(np.int64) & 0xff
ln = self.get(a[4]).astype(np.int64) & 0xff
w = int(re.sub(r'\D', '', parts[-1]))
out = np.array(b, U64)
for i in range(self.T):
p, l = int(pos[i]), int(ln[i])
if l == 0 or p >= w:
continue
l = min(l, w - p)
m = ((1 << l) - 1) << p
out[i] = (int(b[i]) & ~m & MASK[w]) | ((int(f[i]) << p) & m)
self.put(a[0], out, mask)
def op_setp(self, st, parts, a, mask):
cmp, ty = parts[1], parts[-1]
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
r = {'lt': x < y, 'le': x <= y, 'gt': x > y, 'ge': x >= y, 'eq': x == y, 'ne': x != y,
'lo': x < y, 'ls': x <= y, 'hi': x > y, 'hs': x >= y}[cmp]
dests = a[0].split('|')
self.put(dests[0], r, mask)
def op_selp(self, st, parts, a, mask):
p = self.get(a[3])
self.put(a[0], np.where(p, self.get(a[1]), self.get(a[2])), mask)
def op_ex2(self, st, parts, a, mask):
ty = parts[-1]
v = self.get(a[1])
if ty == 'f16x2':
out = np.zeros(self.T, U64)
for i in range(2):
h = bits_to_f16((v >> np.uint64(16 * i)) & np.uint64(0xffff)).astype(F64)
out |= f16_to_bits(np.exp2(h).astype(F16)) << np.uint64(16 * i)
self.put(a[0], out, mask)
elif ty == 'f16':
self.put(a[0], f16_to_bits(np.exp2(bits_to_f16(v).astype(F64)).astype(F16)), mask)
elif ty == 'f32':
self.put(a[0], f32_to_bits(np.exp2(bits_to_f32(v).astype(F64)).astype(F32)), mask)
else:
raise NotImplementedError(st.text)
def op_rsqrt(self, st, parts, a, mask):
x = self.typed(a[1], 'f32').astype(F32)
with np.errstate(divide='ignore', invalid='ignore'):
self.put(a[0], f32_to_bits((1.0 / np.sqrt(x.astype(F64))).astype(F32)), mask)
def op_rcp(self, st, parts, a, mask):
x = self.typed(a[1], 'f32')
with np.errstate(divide='ignore'):
self.put(a[0], f32_to_bits((1.0 / x).astype(F32)), mask)
def op_div(self, st, parts, a, mask):
x, y = self.typed(a[1], 'f32'), self.typed(a[2], 'f32')
with np.errstate(divide='ignore', invalid='ignore'):
self.put(a[0], f32_to_bits((x / y).astype(F32)), mask)
def op_cvt(self, st, parts, a, mask):
types = [p for p in parts[1:] if p not in ('rn', 'rz', 'rm', 'rp', 'rni', 'rzi', 'sat', 'satfinite', 'relu', 'ftz')]
dt, stp = types[0], types[1]
relu = 'relu' in parts
v = self.get(a[1])
if dt == 'e4m3x2' and stp == 'f16x2':
out = np.zeros(self.T, U64)
for i in range(2):
h = (v >> np.uint64(16 * i)) & np.uint64(0xffff)
if relu:
f = bits_to_f16(h)
h = np.where((f < 0) & ~np.isnan(f.astype(F32)), 0, h).astype(U64)
out |= F16_TO_E4M3[h.astype(np.int64)].astype(U64) << np.uint64(8 * i)
self.put(a[0], out, mask)
elif dt == 'f16x2' and stp == 'e4m3x2':
out = np.zeros(self.T, U64)
for i in range(2):
c = ((v >> np.uint64(8 * i)) & np.uint64(0xff)).astype(np.int64)
out |= f16_to_bits(E4M3[c].astype(F16)) << np.uint64(16 * i)
self.put(a[0], out, mask)
elif dt == 'f16' and stp == 'f32':
self.put(a[0], f16_to_bits(bits_to_f32(v).astype(F16)), mask)
elif dt == 'f16' and stp == 'f64':
self.put(a[0], f16_to_bits(v.view(F64).astype(F16)), mask)
elif dt == 'f32' and stp == 'f16':
self.put(a[0], f32_to_bits(bits_to_f16(v).astype(F32)), mask)
elif dt in ('f16', 'f32') and stp[0] in 'us':
ws = int(re.sub(r'\D', '', stp))
iv = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
if dt == 'f16':
self.put(a[0], f16_to_bits(iv.astype(F64).astype(F16)), mask)
else:
self.put(a[0], f32_to_bits(iv.astype(F64).astype(F32)), mask)
elif dt[0] in 'us' and stp[0] in 'us':
ws = int(re.sub(r'\D', '', stp))
val = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
self.put(a[0], val.astype(np.int64).astype(U64), mask)
else:
raise NotImplementedError(st.text)
def op_shfl(self, st, parts, a, mask):
mode = parts[2]
v = self.get(a[1])
b = self.get(a[2]).astype(np.int64)
out = np.array(v)
for w in range(self.T // 32):
sl = slice(32 * w, 32 * w + 32)
lane = np.arange(32)
if mode == 'bfly':
src = lane ^ b[sl]
elif mode == 'down':
src = np.where(lane + b[sl] < 32, lane + b[sl], lane)
elif mode == 'up':
src = np.where(lane - b[sl] >= 0, lane - b[sl], lane)
elif mode == 'idx':
src = b[sl] & 31
out[sl] = v[sl][src]
self.put(a[0], out, mask)
def op_movmatrix(self, st, parts, a, mask):
v = self.get(a[1])
out = np.zeros(self.T, U64)
for w in range(self.T // 32):
m = np.zeros((8, 8), U64)
for lane in range(32):
x = int(v[32 * w + lane])
r, c = lane // 4, (lane % 4) * 2
m[r, c], m[r, c + 1] = x & 0xffff, (x >> 16) & 0xffff
mt = m.T
for lane in range(32):
r, c = lane // 4, (lane % 4) * 2
out[32 * w + lane] = int(mt[r, c]) | (int(mt[r, c + 1]) << 16)
self.put(a[0], out, mask)
def op_mma(self, st, parts, a, mask):
shape = parts[3]
d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
A = [self.get(n) for n in a_names]
Bv = [self.get(n) for n in b_names]
C = [self.get(n) for n in c_names]
dout = [np.array(self.get(n)) for n in d_names]
for w in range(self.T // 32):
base = 32 * w
if shape == 'm16n8k32':
Am = np.zeros((16, 32)); Bm = np.zeros((32, 8)); Cm = np.zeros((16, 8))
for lane in range(32):
g, t = lane >> 2, lane & 3
for i in range(16):
byte = (int(A[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
row = g if (i < 4 or 8 <= i < 12) else g + 8
col = t * 4 + (i & 3) + (16 if i >= 8 else 0)
Am[row, col] = E4M3[byte]
for i in range(8):
byte = (int(Bv[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
Bm[t * 4 + (i & 3) + (16 if i >= 4 else 0), g] = E4M3[byte]
elif shape == 'm16n8k16':
Am = np.zeros((16, 16)); Bm = np.zeros((16, 8)); Cm = np.zeros((16, 8))
for lane in range(32):
g, t = lane >> 2, lane & 3
for i in range(8):
h = (int(A[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
row = g if (i < 2 or 4 <= i < 6) else g + 8
col = t * 2 + (i & 1) + (8 if i >= 4 else 0)
Am[row, col] = np.uint16(h).view(F16)
for i in range(4):
h = (int(Bv[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
Bm[t * 2 + (i & 1) + (8 if i >= 2 else 0), g] = np.uint16(h).view(F16)
else:
raise NotImplementedError(st.text)
for lane in range(32):
g, t = lane >> 2, lane & 3
for i in range(4):
h = (int(C[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
Cm[g if i < 2 else g + 8, t * 2 + (i & 1)] = np.uint16(h).view(F16)
Dm = (Am @ Bm + Cm).astype(F16)
if self.mma_hook:
self.mma_hook(self, st, w, Am, Bm, Cm, Dm, [self.prov.get(n) for n in a_names], [self.prov.get(n) for n in b_names])
for lane in range(32):
g, t = lane >> 2, lane & 3
for r in range(2):
lo = int(Dm[g + 8 * r, t * 2].view(np.uint16))
hi = int(Dm[g + 8 * r, t * 2 + 1].view(np.uint16))
dout[r][base + lane] = lo | (hi << 16)
for n, v in zip(d_names, dout):
self.put(n, v, mask)
mma_hook = None
# wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 with register fragments: two m16n8k16 MMAs
# sharing A registers 0..3, B registers 0,1 / 2,3 and C/D registers 0,1 / 2,3 (the upper A and B
# registers of the fragment ABI are duplicates; ZLUDA's LLVMZludaSplitWmmaM16n16k16 mapping)
def op_wmma(self, st, parts, a, mask):
if 'm16n16k16' not in parts or 'mma' not in parts:
raise NotImplementedError(st.text)
d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
for half in range(2):
sub = Stmt_like(st, 'mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16',
['{%s}' % ', '.join(d_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(a_names[0:4]),
'{%s}' % ', '.join(b_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(c_names[2 * half:2 * half + 2])])
self.op_mma(sub, sub.op.split('.'), sub.args, mask)
class Stmt_like:
def __init__(self, st, op, args):
self.op, self.args, self.text = op, args, st.text + f' [as {op}]'
for k in Stmt.__slots__:
if k not in ('op', 'args', 'text'):
setattr(self, k, getattr(st, k, None))
def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None):
stmts, types, shared = parse_kernel(ptx, kernel)
mem = Memory(dump)
launch = [l for l in open(dump + '/manifest.txt') if l.startswith('launch')][0].split()
grid = tuple(int(x) for x in launch[1:4])
block = tuple(int(x) for x in launch[4:7])
b = Block(stmts, types, shared, mem, kernel + '_param_0', block, grid, ctaid)
b.hooks = list(hooks)
b.mma_hook = mma_hook
steps = b.run(limit)
return b, mem, steps
if __name__ == '__main__':
ptx, kernel, dump = sys.argv[1:4]
cx, cy = (int(v) for v in sys.argv[4].split(',')) if len(sys.argv) > 4 else (0, 0)
import time
t0 = time.time()
b, mem, steps = run_block(ptx, kernel, dump, (cx, cy, 0))
print(f'block ({cx},{cy}): {steps} statements in {time.time() - t0:.1f}s, {len(mem.stores)} store groups')
kernels/tools/pwin_check.py
+127
View File
@@ -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/pwin_model.py
+440
View File
@@ -0,0 +1,440 @@
#!/usr/bin/env python3
"""Numpy reference of DLSS 4 (preset K) DltssPaddedWinLayer blocks, built stage by stage against the
PTX interpreter (kernels/tools/ptxsim.py). f16 arithmetic in the kernel's operation order."""
import struct
import numpy as np
F16, F32, F64 = np.float16, np.float32, np.float64
# Legacy replay reports retain their original approximate norm by default.
# Windows strict K validation selects the PTX fragment's rounded reduction.
NORM_FRAGMENT_ORDER = False
def norm_rows(x):
if not NORM_FRAGMENT_ORDER:
return f16(np.sum(x.astype(F32).astype(F64) ** 2, axis=1))
tokens, channels = x.shape
assert channels in (32, 64, 96, 128, 160)
square = f16(x.astype(F64) ** 2)
partial = []
for lane in range(4):
halves = []
for parity in range(2):
pairs = [add16(square[:, 2 * lane + 8 * parity + 32 * g:2 * lane + 8 * parity + 32 * g + 2],
square[:, 2 * lane + 8 * parity + 32 * g + 16:2 * lane + 8 * parity + 32 * g + 18])
for g in range(channels // 32)]
value = pairs[0]
if channels >= 64:
value = add16(value, pairs[1])
if channels == 96:
value = add16(value, pairs[2])
if channels >= 128:
value = add16(value, add16(pairs[2], pairs[3]))
if channels == 160:
value = add16(value, pairs[4])
halves.append(value)
partial.append(add16(halves[0], halves[1]))
total = add16(add16(partial[0], partial[1]), add16(partial[2], partial[3]))
return add16(total[:, 0], total[:, 1])
class Dump:
def __init__(self, d):
self.dir = d
man = open(d + '/manifest.txt').read().split('\n')
self.kernel = man[0].split()[1]
la = man[1].split()
self.grid, self.block = tuple(map(int, la[1:4])), tuple(map(int, la[4:7]))
self.args = open(d + '/args.bin', 'rb').read()
self.allocs = {}
for l in man:
p = l.split()
if p and p[0] == 'alloc':
self.allocs[int(p[1])] = (int(p[2], 16), np.fromfile(f'{d}/alloc-{p[1]}.bin', np.uint8))
def u64(self, off):
return struct.unpack_from('<Q', self.args, off)[0]
def i32x2(self, off):
return struct.unpack_from('<ii', self.args, off)
def ptr(self, off):
a = self.u64(off)
for i, (base, buf) in self.allocs.items():
if base <= a < base + buf.size:
return buf, a - base
raise ValueError(hex(a))
def f16(self, off_param, byte, n):
buf, o = self.ptr(off_param)
return buf[o + byte:o + byte + 2 * n].view(F16).copy()
def f16(x):
return np.asarray(x, F64).astype(F16)
def mirror(v, n):
v = abs(v)
return min(v, 2 * n - 2 - v)
def load_window(P, bx, by, C):
"""64 tokens (raster order of the 8x8 window) x C channels, f16, mirror padded."""
W, H = P.i32x2(0)
sx, sy = P.i32x2(56)
buf, o = P.ptr(8)
x = np.zeros((64, C), F16)
for r in range(64):
ty, tx = r >> 3, r & 7
Y = mirror(8 * by - sy + ty, H)
X = mirror(8 * bx - sx + tx, W)
a = o + 2 * ((Y * W + X) * C)
x[r] = buf[a:a + 2 * C].view(F16)
return x
# --- reductions in the kernel's order ---------------------------------------------------------
def sumsq_rows(x):
"""f16 sum of squares per token, reproducing the kernel's tree for C = 64 (per lane: 16 products
over 8 words (t + 4k), pairwise f16x2 adds, butterfly over the 4 lanes of a row, final hi+lo)."""
T, C = x.shape
assert C == 64
sq = f16(x.astype(F64) * x.astype(F64)) # f16 products
out = np.zeros(T, F16)
for r in range(T):
lanes = []
for t in range(4):
# word k holds channels 2(t+4k), 2(t+4k)+1 as an f16 pair
w = [np.array([sq[r, 2 * (t + 4 * k)], sq[r, 2 * (t + 4 * k) + 1]], F16) for k in range(8)]
a = [f16(w[2 * i].astype(F64) + w[2 * i + 1].astype(F64)) for i in range(4)] # (0+1),(2+3),(4+5),(6+7)
b = [f16(a[0].astype(F64) + a[2].astype(F64)), f16(a[1].astype(F64) + a[3].astype(F64))]
lanes.append(f16(b[0].astype(F64) + b[1].astype(F64)))
s1 = [f16(lanes[t].astype(F64) + lanes[t ^ 1].astype(F64)) for t in range(4)]
s2 = f16(s1[0].astype(F64) + s1[2].astype(F64))
out[r] = f16(s2[0].astype(F64) + s2[1].astype(F64))
return out
def rsqrt16(s):
return f16(1.0 / np.sqrt(s.astype(F32).astype(F64)))
def mma(A, B, C):
"""one wmma k16 step: D = f16(A @ B + C) with exact sums (the interpreter's semantics)."""
return f16(A.astype(F64) @ B.astype(F64) + C.astype(F64))
def frag_offset(k, n):
"""byte offset of logical B[k][n] (k, n < 16) inside one 512-byte pre-swizzled f16 fragment"""
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1)
FRAG = np.array([[frag_offset(k, n) for n in range(16)] for k in range(16)])
def weights(P, byte, K, N, nt_stride=None):
"""logical W[K][N] from fragments at weights + byte + 512 * kt + nt_stride * nt"""
buf, o = P.ptr(64)
nt_stride = nt_stride or 512 * (K // 16)
W = np.zeros((K, N), F16)
for kt in range(K // 16):
for nt in range(N // 16):
base = o + byte + 512 * kt + nt_stride * nt
vals = buf[(base + FRAG)[..., None] + np.arange(2)].reshape(16, 16, 2).view(F16)[..., 0]
W[16 * kt:16 * kt + 16, 16 * nt:16 * nt + 16] = vals
return W
def gemm(A, W, C0=None):
"""chained wmma k16 steps (f16 rounding after each step), per 16x16 output tile"""
M, K = A.shape
N = W.shape[1]
D = np.zeros((M, N), F16) if C0 is None else C0.astype(F16).copy()
for kt in range(K // 16):
D = mma(A[:, 16 * kt:16 * kt + 16], W[16 * kt:16 * kt + 16], D)
return D
def hv(bits):
return np.array(bits, np.uint16).view(F16)
def fma16(a, b, c):
return f16(a.astype(F64) * b.astype(F64) + c.astype(F64))
def add16(a, b):
return f16(a.astype(F64) + b.astype(F64))
def mul16(a, b):
return f16(a.astype(F64) * b.astype(F64))
def softmax_exp(S):
"""e = ex2(t * (1 + a t^2) + c), t = clamp(S, +-21.64) (f16x2 ops of the kernel)"""
hi, lo = hv(0x4D69), hv(0xCD69)
a, one, c = hv(0x91D5), hv(0x3C00), hv(0xC50D)
t = np.maximum(np.minimum(S, hi), lo)
u = mul16(t, t)
v = fma16(u, a, one)
w = fma16(t, v, c)
return f16(np.exp2(w.astype(F64)))
def gelu16(x):
"""x * (0.5 + c (0.41216 - 0.081081 |c|)), c = clamp(x, +-2) (f16x2 ops of the kernel)"""
k1, k2 = f16(np.float32(np.uint32(0x3ED306EB).view(F32))), f16(np.float32(np.uint32(0x3DA60DD6).view(F32)))
half, two = F16(0.5), F16(2.0)
c = np.minimum(np.maximum(x, -two), two)
u = f16(k1.astype(F64) - mul16(np.full_like(c, k2), np.abs(c)).astype(F64))
return mul16(x, add16(np.full_like(c, half), mul16(c, u)))
def bias_table(P, byte):
"""64x64 f16 table stored in C-fragment order, 512 bytes per 16x16 tile (tile 4 mt + nt)"""
buf, o = P.ptr(64)
B = np.zeros((64, 64), F16)
for i in range(64):
for j in range(64):
mt, nt, r, c = i >> 4, j >> 4, i & 15, j & 15
off = 512 * (4 * nt + mt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3)
B[i, j] = buf[o + byte + off:o + byte + off + 2].view(F16)[0]
return B
def vec(P, byte, n):
return P.f16(64, byte, n)
def pwin_enc2(P, bx, by, pm_order=((0, 0), (0, 1), (1, 0), (1, 1)), b2_in_c0=True, dbg=None):
"""logical enc2 layer (H 2, C 64 -> merged 96) for one window; returns (y [64 x 64], merged [16 x 96])"""
C, H, Dh = 64, 2, 32
x0 = load_window(P, bx, by, C)
s = norm_rows(x0)
r = rsqrt16(s)
h1 = mul16(x0, mul16(r[:, None], vec(P, 0, C)[None, :]))
Os = []
for h in range(H):
hb = 128 + 12288 * h
q, k, v = [gemm(h1[:, 32:64], weights(P, hb + 2048 * (j + 3), 32, 32, 1024),
gemm(h1[:, 0:32], weights(P, hb + 2048 * j, 32, 32, 1024))) for j in range(3)]
S = gemm(q, k.T.copy(), bias_table(P, 24704 + 8192 * h))
E = softmax_exp(S)
rs = f16(np.sum(E.astype(F64), axis=1))
rc = f16(1.0 / rs.astype(F32).astype(F64))
Pm = mul16(E, rc[:, None])
Os.append(gemm(Pm, v))
if dbg is not None:
dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = q, k, v, S, Pm, Os[-1]
O = np.concatenate(Os, axis=1)
bo = vec(P, 41088, C)
Y = gemm(O[:, 32:64], weights(P, 45312, 32, 64, 1024), gemm(O[:, 0:32], weights(P, 41216, 32, 64, 1024), np.tile(bo, (64, 1))))
x1 = add16(Y, x0)
if dbg is not None:
dbg['Y'], dbg['x1'] = Y, x1
m = mul16(x1, vec(P, 49408, C)[None, :])
b2 = vec(P, 49536, C)
acc = add16(x1, np.tile(b2, (64, 1))) if b2_in_c0 else x1
for c in range(8):
Hc = gemm(m, weights(P, 49664 + 4096 * c, 64, 32, 2048), np.tile(vec(P, 82432 + 64 * c, 32), (64, 1)))
acc = gemm(gelu16(Hc), weights(P, 82944 + 4096 * c, 32, 64, 1024), acc)
if dbg is not None:
dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
y = acc
# patch merge: 16 merged tokens (4x4), K = 4 sub-tokens x 64
A = np.zeros((16, 256), F16)
for mt in range(16):
my, mx = mt >> 2, mt & 3
for q, (dy, dx) in enumerate(pm_order):
A[mt, 64 * q:64 * q + 64] = y[8 * (2 * my + dy) + 2 * mx + dx]
merged = np.concatenate([gemm(A, weights(P, 115712 + 24576 * w, 256, 48, 8192), np.tile(vec(P, 164864 + 96 * w, 48), (16, 1)))
for w in range(2)], axis=1)
return y, merged
class Layout:
"""weight-buffer layout of a plain DltssPaddedWinLayer (encoder with patch merge), bytes"""
def __init__(self, H, C, COUT):
self.H, self.C, self.COUT = H, C, COUT
self.G1 = 0
self.QKV = 2 * C
self.HEAD = 192 * C
self.BIAS = self.QKV + self.HEAD * H
self.BO = self.BIAS + 8192 * H
self.WO = self.BO + 2 * C
self.G2 = self.WO + 64 * C * H
self.B2 = self.G2 + 2 * C
self.W1 = self.B2 + 2 * C
self.B1 = self.W1 + 8 * C * C
self.W2 = self.B1 + 8 * C
self.PM = self.W2 + 8 * C * C
self.NPW = COUT // H # merged channels per original warp
self.NPA = -(-self.NPW // 16) * 16 # allocated (padded to whole 16-column tiles)
self.PMB = self.PM + 8 * C * self.NPA * H
def swin_core(P, L, x0, dbg=None, posattn=False):
"""norm -> attention -> Wo + residual -> MLP (+ residual): the block output y [64 x C]"""
C, H = L.C, L.H
NCH = C // 32
s = norm_rows(x0)
r = rsqrt16(s)
h1 = mul16(x0, mul16(r[:, None], vec(P, L.G1, C)[None, :]))
Os = []
for h in range(H):
hb = L.QKV + L.HEAD * h
proj = []
for j in range(3):
acc = None
for ch in range(NCH):
acc = gemm(h1[:, 32 * ch:32 * ch + 32], weights(P, hb + 2048 * (3 * ch + j), 32, 32, 1024), acc)
proj.append(acc)
q, k, v = proj
if posattn:
# position-only attention (trait flag 2): the table is the probability matrix itself
Os.append(gemm(bias_table(P, L.BIAS + 8192 * h), v))
if dbg is not None:
dbg[f'v{h}'], dbg[f'O{h}'] = v, Os[-1]
continue
S = gemm(q, k.T.copy(), bias_table(P, L.BIAS + 8192 * h))
E = softmax_exp(S)
rs = f16(np.sum(E.astype(F64), axis=1))
rc = f16(1.0 / rs.astype(F32).astype(F64))
Os.append(gemm(mul16(E, rc[:, None]), v))
if dbg is not None:
dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'O{h}'], dbg[f'S{h}'], dbg[f'E{h}'], dbg[f'rs{h}'] = q, k, v, Os[-1], S, E, rs
Y = np.tile(vec(P, L.BO, C), (64, 1))
for h in range(H):
Y = gemm(Os[h], weights(P, L.WO + 64 * C * h, 32, C, 1024), Y)
x1 = add16(Y, x0)
m = mul16(x1, vec(P, L.G2, C)[None, :])
acc = add16(x1, np.tile(vec(P, L.B2, C), (64, 1)))
if dbg is not None:
dbg['Y'], dbg['x1'] = Y, x1
for c in range(C // 8):
Hc = gemm(m, weights(P, L.W1 + 64 * C * c, C, 32, 32 * C), np.tile(vec(P, L.B1 + 64 * c, 32), (64, 1)))
acc = gemm(gelu16(Hc), weights(P, L.W2 + 64 * C * c, 32, C, 1024), acc)
if dbg is not None:
dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
return acc
def pwin_encoder(P, bx, by, H, C, COUT, pm_bias_padded=True, posattn=False):
L = Layout(H, C, COUT)
x0 = load_window(P, bx, by, C)
y = swin_core(P, L, x0, posattn=posattn)
A = np.zeros((16, 4 * C), F16)
for mt in range(16):
my, mx = mt >> 2, mt & 3
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
Wp = weights(P, L.PM, 4 * C, L.NPA * H, 128 * C)
bias = vec(P, L.PMB, L.NPA * H) if pm_bias_padded else None
cols = []
for w in range(H):
sl = slice(L.NPA * w, L.NPA * w + L.NPW)
b = bias[sl] if pm_bias_padded else vec(P, L.PMB + 2 * L.NPW * w, L.NPW)
cols.append(gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(np.concatenate([b, np.zeros(L.NPA - L.NPW, F16)]), (16, 1)))[:, :L.NPW])
merged = np.concatenate(cols, axis=1)
return y, merged
class Shifted:
"""view of the weight buffer starting at a byte offset (decoders: core weights follow the expand)"""
def __init__(self, P, base):
self.P, self.base = P, base
def ptr(self, off):
buf, o = self.P.ptr(off)
return buf, o + (self.base if off == 64 else 0)
def f16(self, off_param, byte, n):
return self.P.f16(off_param, byte + (self.base if off_param == 64 else 0), n)
def i32x2(self, off):
return self.P.i32x2(off)
def load_lowres(P, bx, by, CL):
W, H = P.i32x2(0)
sx, sy = P.i32x2(56)
buf, o = P.ptr(8)
W2, H2 = W // 2, H // 2
lx0, ly0 = int((8 * bx - sx) / 2), int((8 * by - sy) / 2)
x = np.zeros((16, CL), F16)
for r in range(16):
Y = mirror(ly0 + (r >> 2), H2); X = mirror(lx0 + (r & 3), W2)
a = o + 2 * ((Y * W2 + X) * CL)
x[r] = buf[a:a + 2 * CL].view(F16)
return x
def load_skip(P, bx, by, C):
W, H = P.i32x2(0)
sx, sy = P.i32x2(56)
buf, o = P.ptr(16)
x = np.zeros((64, C), F16)
for r in range(64):
Y = mirror(8 * by - sy + (r >> 3), H); X = mirror(8 * bx - sx + (r & 7), W)
a = o + 2 * ((Y * W + X) * C)
x[r] = buf[a:a + 2 * C].view(F16)
return x
def pwin_decoder(P, bx, by, H, C, CL, order=((0, 0), (0, 1), (1, 0), (1, 1))):
xl = load_lowres(P, bx, by, CL)
EXP = 2 * CL * 4 * C
E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
skip = load_skip(P, bx, by, C)
x0 = np.zeros((64, C), F16)
for r in range(16):
ly, lx = r >> 2, r & 3
for q, (dy, dx) in enumerate(order):
t = 8 * (2 * ly + dy) + 2 * lx + dx
x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
L = Layout(H, C, C)
return swin_core(Shifted(P, EXP + 8 * C), L, x0)
def pwin_enc0(P, bx, by, H=2, C=32, COUT=64, CIN=16):
"""first layer: token embedding x0 = in (CIN) W_e + b_e, then the core (position-only attention) and merge"""
xin = load_window(P, bx, by, CIN)
x0 = gemm(xin, weights(P, 2 * C, CIN, C, 512 * (CIN // 16)), np.tile(vec(P, 0, C), (64, 1)))
x0 = np.maximum(x0, F16(0)) # ReLU after the embedding
CB = 2 * C + 2 * CIN * C
L = Layout(H, C, COUT)
Pw = Shifted(P, CB)
y = swin_core(Pw, L, x0, posattn=True)
A = np.zeros((16, 4 * C), F16)
for mt in range(16):
my, mx = mt >> 2, mt & 3
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
Wp = weights(Pw, L.PM, 4 * C, L.NPA * H, 128 * C)
bias = vec(Pw, L.PMB, L.NPA * H)
cols = [gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(bias[L.NPA * w:L.NPA * (w + 1)], (16, 1)))[:, :L.NPW] for w in range(H)]
return y, np.concatenate(cols, axis=1)
def pwin_dec0(P, bx, by, H=2, C=32, CL=64, NOUT=40, NOUTA=48):
"""last layer: patch expand + skip, core (position-only attention), output head y W_h + b_h (NOUT of NOUTA)"""
xl = load_lowres(P, bx, by, CL)
EXP = 2 * CL * 4 * C
E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
skip = load_skip(P, bx, by, C)
x0 = np.zeros((64, C), F16)
for r in range(16):
ly, lx = r >> 2, r & 3
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
t = 8 * (2 * ly + dy) + 2 * lx + dx
x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
L = Layout(H, C, C)
Pw = Shifted(P, EXP + 8 * C)
y = swin_core(Pw, L, x0, posattn=True)
HB = L.PM # head bias, then the head weights
Wh = weights(Pw, HB + 2 * NOUTA, C, NOUTA, 1024)
out = gemm(y, Wh, np.tile(vec(Pw, HB, NOUTA), (64, 1)))
return y, out[:, :NOUT]
kernels/tools/swin_check.py
+113
View File
@@ -0,0 +1,113 @@
#!/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)}
# Preserve the existing Linux oracle by default; Windows strict replay enables
# the audited 310.9 f16 patch-merge bound before final e4m3 quantisation.
CLAMP_PATCH_MERGE = False
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 = sm.pair_to_nat(np.concatenate(cols, axis=1))
if CLAMP_PATCH_MERGE:
limit = me.f16(2 * np.pi)
merged = np.clip(merged, -limit, limit)
merged = me.q8(merged)
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/swin_model.py
+170
View File
@@ -0,0 +1,170 @@
#!/usr/bin/env python3
"""Parametrised numpy model of the rrlite Swin blocks (enc3 spec generalised); checked stage by stage
against ptxsim MMA captures."""
import numpy as np
from model_enc3 import (q8, f16, Params, E4M3, F16, exp_trick, row_sum16, C_SCALE, CLAMP, P1, P0,
gelu_poly, mma_chain)
def gp(c):
c = np.asarray(c)
return 32 * (c // 32) + 2 * ((c % 32) // 4) + (c % 2) + 16 * ((c % 4) // 2)
def pair_to_nat(M):
return M[:, gp(np.arange(M.shape[1]))]
def nat_to_pair(M):
out = np.empty_like(M); out[:, gp(np.arange(M.shape[1]))] = M; return out
def woff_table(base, Ks, Ns, K, N):
k = np.arange(K)[:, None]; n = np.arange(N)[None, :]; j = n // 8
return base + Ks * (k // 32) + Ns * (j // 2) + 64 * (n % 8) + 16 * ((k % 16) // 4) + 8 * (j % 2) + 4 * ((k % 32) // 16) + k % 4
def load_tokens(P, bx, by, planes):
X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, P.W), P.mirror(Y, P.H)
out = np.zeros((64, 32 * planes))
for p in range(planes):
off = P.inp - P.abase + ((p * P.H + Y) * P.W + X) * 32
out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
return out
def sumsq_tokens(x):
"""Kernel-exact f16 sum of squares of natural-order token rows (planes of 32; lane t holds words t, t+4)."""
sq = f16(x * x); planes = x.shape[1] // 32
def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
lanes = []
for t in range(4):
def pair(pp, word, k):
c = 32 * pp + 4 * word + 2 * k
return sq[:, c], sq[:, c + 1]
def plane_sum(word):
p = [add(pair(pp, word, 0), pair(pp, word, 1)) for pp in range(planes)]
if planes == 4: return add(add(p[0], p[1]), add(p[2], p[3]))
if planes == 3: return add(add(p[0], p[1]), p[2])
if planes == 2: return add(p[0], p[1])
raise ValueError(planes)
lanes.append(add(plane_sum(t), plane_sum(t + 4)))
a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
return f16(b[0] + b[1])[:, None]
def rms_from_sum(x, s, gamma):
r = f16(1.0 / np.sqrt(f16(s + f16(2.0 ** -13))))
return q8(f16(x * f16(r * gamma[None, :])))
def rel_bias(P, off):
raw = P.f16vec(off, 256)
i, j = np.meshgrid(np.arange(16), np.arange(16), indexing='ij')
lane = 4 * (i % 8) + (j % 8) // 2; word = i // 8 + 2 * (j // 8)
return raw[8 * lane + 2 * word + (j % 2)]
def attention(Qq, V, bias):
S = f16(Qq @ Qq.T + bias)
t = np.clip(f16(S * C_SCALE), -CLAMP, CLAMP)
e1 = f16(t * (-t) + P1); poly = f16(t * e1 + P0)
w = exp_trick(poly); s = row_sum16(w)
r = f16(np.float32(1.0) / s.astype(np.float32)); Pm = f16(w * r)
return f16(Pm @ V)
def tree_norm2(x):
"""Kernel-exact f16 sum of squares of pair-ordered C-fragment rows (NT8 = C/8 n8 tiles)."""
sq = f16(x * x); nt8 = x.shape[1] // 8
def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
lanes = []
for t in range(4):
T = [(sq[:, 8 * j + 2 * t], sq[:, 8 * j + 2 * t + 1]) for j in range(nt8)]
def quad(js): return add(add(T[js[0]], T[js[1]]), add(T[js[2]], T[js[3]]))
if nt8 == 16:
even = add(quad([0, 2, 4, 6]), quad([8, 10, 12, 14])); odd = add(quad([1, 3, 5, 7]), quad([9, 11, 13, 15]))
elif nt8 == 12:
even = add(quad([0, 2, 4, 6]), add(T[8], T[10])); odd = add(quad([1, 3, 5, 7]), add(T[9], T[11]))
elif nt8 == 8:
even = quad([0, 2, 4, 6]); odd = quad([1, 3, 5, 7])
else:
raise ValueError(nt8)
lanes.append(add(even, odd))
a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
return f16(b[0] + b[1])[:, None]
def layout(C, heads, pre=0):
"""Weight offsets of a Swin block with C channels starting at byte `pre`."""
HS = 96 * C + 512
L = dict(g1=pre, head=[pre + 2 * C + HS * h for h in range(heads)])
E = pre + 2 * C + heads * HS
L.update(bo=E, g2=E + 2 * C)
M0 = E + 4 * C; NC = C // 8
L['b1'] = [M0 if c == 0 else M0 + 66 * C + 64 + (64 * C + 64) * (c - 1) for c in range(NC)]
L['b2'] = M0 + 64 * C + 64
L['end'] = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1)
return L
def swin_block(P, x0, C, heads, pre=0, stages=None):
"""x0: 64 x C natural (f16 values). Returns (x2 pair order, out natural fp8 values)."""
fp8 = lambda T: E4M3[P.wbuf[T]]
L = layout(C, heads, pre)
h1 = rms_from_sum(x0, sumsq_tokens(x0), P.f16vec(L['g1'], C)[gp(np.arange(C))])
O = np.zeros((64, 32 * heads))
for h in range(heads):
hb = L['head'][h]
Q = mma_chain(h1, fp8(woff_table(hb, 1024, 512, C, 32)))
V = mma_chain(h1, fp8(woff_table(hb + 32 * C, 1024, 512, C, 32)))
B = rel_bias(P, hb + 64 * C)
for m in range(4):
O[16 * m:16 * m + 16, 32 * h:32 * h + 32] = attention(q8(Q[16 * m:16 * m + 16]), V[16 * m:16 * m + 16], B)
Wo = np.zeros((32 * heads, C), np.int64)
for s in range(heads):
Wo[32 * s:32 * s + 32] = woff_table(L['head'][s] + 64 * C + 512, 0, 512, 32, C)
x1 = mma_chain(pair_to_nat(q8(O)), fp8(Wo), C0=f16(nat_to_pair(x0) + P.f16vec(L['bo'], C)[None, :]))
h2 = pair_to_nat(rms_from_sum(x1, tree_norm2(x1), P.f16vec(L['g2'], C)))
if stages is not None:
stages.update(h1=h1, x1=pair_to_nat(x1), h2=h2, O=pair_to_nat(q8(O)))
x2 = f16(x1 + P.f16vec(L['b2'], C)[None, :])
idx = gp(np.arange(32))
for c, B1 in enumerate(L['b1']):
z = mma_chain(h2, fp8(woff_table(B1, 512, 16 * C, C, 32)), C0=np.broadcast_to(P.f16vec(B1 + 32 * C, 32), (64, 32)).copy())
g = q8(gelu_poly(z))[:, idx]
x2 = f16(g @ fp8(woff_table(B1 + 32 * C + 64, 0, 512, 32, C)) + x2)
return x2, q8(pair_to_nat(x2)), L
TOK_X = (np.arange(64) & 3) + 4 * ((np.arange(64) >> 4) & 1)
TOK_Y = ((np.arange(64) >> 2) & 3) + 4 * (np.arange(64) >> 5)
def patch_merge_input(out, C):
"""16 merged rows (r -> (r%4, r//4)) x 4C: column C*q + ch = out[token(2mx+q%2, 2my+q//2), ch]."""
A = np.zeros((16, 4 * C))
for r in range(16):
mx, my = r % 4, r // 4
for q in range(4):
t = np.nonzero((TOK_X == 2 * mx + (q & 1)) & (TOK_Y == 2 * my + (q >> 1)))[0][0]
A[r, C * q:C * q + C] = out[t]
return A
def load_at(P, ptr, bx, by, planes, W=None, H=None, low=False):
W = W or P.W; H = H or P.H
X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, W), P.mirror(Y, H)
out = np.zeros((64, 32 * planes))
for p in range(planes):
off = ptr - P.abase + ((p * H + Y) * W + X) * 32
out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
return out
def decoder_input(P, bx, by, C, CIN, NW, skip_ptr):
"""x0 = f16(q8(patch_expand(low-res CIN)) + skip), expand output col G = (4C/NW) w + n -> q = G // C, ch = ginv(G % C)."""
fp8 = lambda T: E4M3[P.wbuf[T]]
W2, H2 = P.W // 2, P.H // 2
A = np.zeros((16, CIN))
for r in range(16):
MX = P.mirror(np.array((8 * bx - P.sx) // 2 + r % 4), W2); MY = P.mirror(np.array((8 * by - P.sy) // 2 + r // 4), H2)
for p in range(CIN // 32):
off = P.inp - P.abase + ((p * H2 + MY) * W2 + MX) * 32
A[r, 32 * p:32 * p + 32] = E4M3[P.arena[off:off + 32]]
cols = 4 * C // NW
E_ = np.concatenate([mma_chain(A, fp8(woff_table(w * CIN * cols, 512, 16 * CIN, CIN, cols)),
C0=np.broadcast_to(P.f16vec(4 * C * CIN + 2 * cols * w, cols), (16, cols)).copy()) for w in range(NW)], 1)
skip = load_at(P, skip_ptr, bx, by, C // 32)
x0 = np.zeros((64, C))
for T in range(64):
tx, ty = TOK_X[T], TOK_Y[T]; r = tx // 2 + 4 * (ty // 2); q = (tx & 1) + 2 * (ty & 1)
G = np.arange(C) + q * C # global expand columns of sub-position q (pair order within C)
x0[T] = q8(E_[r, G])[gp(np.arange(C))]
return f16(x0 + skip)
kernels/tools/tail_check.hip
+141
View File
@@ -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;
}