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

Port K kernel tuning and M gfx1101 flags

Browse Source 
Preset K, byte-identical to the previous output in all captured buffers
and the harness image, measured on an RX 7700 XT:

- LDS row padding so WMMA operand loads hit distinct banks

- WGP mode for enc1/enc2/dec1/dec2, ds_bpermute half-wave exchange for
  enc0/dec0/dec3, more waves per window in enc4/dec4/dec5

- cheaper operand plumbing in common/wmma_layout.h (packed f16
  conversions, no dynamic vector indexing), raw buffer weight loads,
  accumulator-init vectors prepared once

- enc1 as a single-buffer recompute layer, dec1 Q/K/V chains split

- dec5 as four dependent phase kernels
This commit is contained in:
ayo committed 2026-10-05 14:48:37 -04:00
1 parent 8ce8618181
commit 22026ede12
45 files changed
+4514 -354
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/build.sh
+34 -1
View File
@@ -52,8 +52,15 @@ build_hip() {
name="$(basename "$src" .hip)"
local tmp flags
tmp="$(mktemp -d)"
# per-kernel compiler flags from a "// d4r-build-flags: ..." line in the source
# per-kernel compiler flags from a "// d4r-build-flags: ..." line in the source. A line for the target
# ("// d4r-build-flags-gfx1101:", or "// d4r-build-flags-gfx12:" for any gfx12 target) replaces it.
flags="$(sed -n 's|^// d4r-build-flags: *||p' "$src")"
for tag in "$ARCH" "${ARCH%??}"; do
if grep -q "^// d4r-build-flags-$tag:" "$src"; then
flags="$(sed -n "s|^// d4r-build-flags-$tag: *||p" "$src")"
break
fi
done
[[ "$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" \
@@ -114,12 +121,38 @@ if [[ "$WHAT" == all || "$WHAT" == tex || "$WHAT" == l ]]; then
for mode in static dynamic; do
for range in hdr ldr; do specs+=("rrlite_downsample_kernel_${mode}_${range}:sust_only"); done
done
# The fast set serves two K kernels from native code translated from the PTX (kernels/native): these builds
# replace the sust_only build of the same name. The accuracy set keeps ZLUDA's compile (denormals preserved).
native_out=hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel
native_in=hiluma_engine_input_depthinv_mvhi_hdr_v2_rel
for spec in "${specs[@]}"; do
IFS=: read -r kernel src mode <<< "$spec"
[[ "$ACCURACY" == 0 && "$kernel" == "$native_out" ]] && continue
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
if [[ "$WHAT" != l && "$ACCURACY" == 0 ]]; then
# NVIDIA's PTX translated to HIP (native/ptx2hip.py) and rebuilt natively; the generated sources stay in a
# temporary directory. Both kernels run through the same texture/surface helpers as the other texture
# kernels and are byte-identical to ZLUDA's compile of the PTX.
ptx_of() { grep -l -E "\.entry[[:space:]]+$1[[:space:]]*\(" "$PTX_DIR"/*.ptx | head -1; }
NAT="$HERE/native"
NT="$(mktemp -d)"
trap 'rm -rf "$NT"' EXIT
ptx="$(ptx_of "$native_out")"; [[ -n "$ptx" ]] || { echo "no PTX defines $native_out" >&2; exit 2; }
# hiluma output (path B for Quality..UltraPerf), LDS tile as a function, global-memory redirect stores
python3 "$NAT/ptx2hip.py" --ctile --as1 "$ptx" "$native_out" > "$NT/output_translated.hip"
python3 "$NAT/rewrite_output.py" "$NT/output_translated.hip" > "$NT/$native_out.hip"
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" "$NAT/build_native.sh" "$NT/$native_out.hip" "$OUT/$native_out.hsaco" \
-DD4R_NATIVE_F2I -DD4R_K_NO_DEPTH_TILE -DD4R_AS1_STORES
# hiluma input: scratch arrays in registers, wave32
ptx="$(ptx_of "$native_in")"; [[ -n "$ptx" ]] || { echo "no PTX defines $native_in" >&2; exit 2; }
python3 "$NAT/ptx2hip.py" --ifconv --sink "$ptx" "$native_in" > "$NT/in_ifs.hip"
python3 "$NAT/in_fuse.py" "$NT/in_ifs.hip" > "$NT/in_unit.h"
cp "$NAT/in_regs.hip" "$NT/in_regs.hip"
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" WAVE=32 "$NAT/build_native.sh" "$NT/in_regs.hip" "$OUT/$native_in.hsaco"
fi
fi
[[ "$ACCURACY" == 1 ]] && printf '1\n' > "$OUT/d4r-accuracy.txt"
echo "native kernels in $OUT"
kernels/common/wmma_layout.h
+97 -3
View File
@@ -56,7 +56,21 @@ WM_FN uint32_t wm_half()
// 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)
{
#ifdef WM_HALF_SWIZZLE
// Bit-mode AND=31, OR=0, XOR=16: the same half-wave partner, no address VGPR.
return (uint32_t)__builtin_amdgcn_ds_swizzle((int)v, 0x401f);
#elif defined(WM_HALF_BPERMUTE)
// through the LDS crossbar (no VALU issue; the same value)
return (uint32_t)__builtin_amdgcn_ds_bpermute((int)((__builtin_amdgcn_mbcnt_lo(~0u, 0u) ^ 16u) << 2), (int)v);
#elif defined(WM_PERMLANE_ASM)
// the builtin ties the destination to an "old" input (unused here: all lanes active, every select valid), which
// costs a v_mov per exchange; a write-only asm output avoids it. Same instruction, same result.
uint32_t r;
__asm__("v_permlanex16_b32 %0, %1, %2, %3" : "=&v"(r) : "v"(v), "s"(0x76543210u), "s"(0xfedcba98u));
return r;
#else
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
#endif
}
WM_FN float wm_other_half_f(float v)
{
@@ -134,8 +148,9 @@ 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]);
// element 2 i + h is half h of dword i: one shift by 16 h (written as v[(2 i + h) >> 1] the compiler could not
// see that the dword index is the constant i and indexed the operand dynamically, ~9 VALU per element)
return __builtin_bit_cast(wm_f16, (uint16_t)(v[i] >> (wm_half() << 4)));
#endif
}
@@ -146,6 +161,24 @@ 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
#ifdef WM_OP_VIA_LDS
// Through a per-wave LDS row buffer (WM_OP_VIA_LDS = waves per block, wave index = work-item z): every lane
// stores its 8 rows of its token's row, then loads the whole 16-value row. 8 VALU conversions + 10 LDS
// instructions instead of ~20 VALU (pack, exchange, perm). The same values. Rows 48 bytes apart (16 mod 32).
// A wave's LDS instructions complete in order, so the loads see all lanes' stores and the next call's stores
// cannot overtake these loads.
__shared__ __attribute__((aligned(16))) wm_f16 wm_xbuf[WM_OP_VIA_LDS][16][24];
{
const uint32_t ln = __builtin_amdgcn_mbcnt_lo(~0u, 0u);
wm_f16* row = wm_xbuf[__builtin_amdgcn_workitem_id_z()][ln & 15];
const uint32_t h = ln >> 4;
#pragma unroll
for (int i = 0; i < 8; ++i)
row[2 * i + h] = own[i];
const wm_u4v a = *(const wm_u4v*)row, b = *(const wm_u4v*)(row + 8);
return (wm_u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
}
#endif
const uint32_t hf = wm_half();
uint32_t mine[4], theirs[4];
#pragma unroll
@@ -177,6 +210,66 @@ WM_FN wm_op wm_op_from_acc(const wm_f16 own[8])
#endif
}
// two f32 -> f16 conversions (round to nearest even) packed into one register (low = a, high = b). WM_CVT_ASM: the
// same v_cvt_f16_f32 instructions as inline asm writing %0.l then %0.h; the compiler otherwise re-packs the halves
// with v_pack_b32_f16 copies
WM_FN uint32_t wm_cvt2(float a, float b)
{
#ifdef WM_CVT_ASM
uint32_t r;
asm("v_cvt_f16_f32 %0.l, %1" : "=v"(r) : "v"(a)); // r may reuse a's register (a is read first)
asm("v_cvt_f16_f32 %0.h, %1" : "+v"(r) : "v"(b)); // keeps the low half
return r;
#else
return wm_pack2((wm_f16)a, (wm_f16)b);
#endif
}
#if D4R_WMMA_LAYOUT == 12
// gfx12: an operand is this half's 8 K values in 4 dwords, accumulator VGPR i = row (K) 8 h + i: the packed pairs
// ARE the operand (no exchange, no byte permutation)
WM_FN wm_op wm_op_from_mine(const uint32_t mine[4])
{
return (wm_u4v){mine[0], mine[1], mine[2], mine[3]};
}
// the pair (element 2 j, 2 j + 1) of wm_op_acc_elem(v, .): operand dword j
WM_FN uint32_t wm_op_acc_pair(const wm_op& v, int j)
{
return v[j];
}
#else
// operand from this lane's 8 accumulator-row halves already packed in pairs (mine[j] = own[2 j], own[2 j + 1]);
// the same exchange + byte permutation as wm_op_from_acc
WM_FN wm_op wm_op_from_mine(const uint32_t mine[4])
{
const uint32_t hf = wm_half();
uint32_t theirs[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
theirs[j] = wm_other_half(mine[j]);
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;
}
// the pair (element 2 j, 2 j + 1) of wm_op_acc_elem(v, .) as one packed register (one byte permutation)
WM_FN uint32_t wm_op_acc_pair(const wm_op& v, int j)
{
return __builtin_amdgcn_perm(v[2 * j + 1], v[2 * j], wm_half() ? 0x07060302u : 0x05040100u);
}
#endif
#if defined(WM_CVT_ASM) && (D4R_WMMA_LAYOUT == 12 || (!defined(WM_OP_VIA_LDS) && !defined(PWIN_NO_PERM)))
#define WM_HAVE_FROM_MINE 1
WM_FN wm_op wm_op_from_f8(wm_f8v d)
{
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
mine[j] = wm_cvt2(d[2 * j], d[2 * j + 1]);
return wm_op_from_mine(mine);
}
#else
WM_FN wm_op wm_op_from_f8(wm_f8v d)
{
wm_f16 h[8];
@@ -185,6 +278,7 @@ WM_FN wm_op wm_op_from_f8(wm_f8v d)
h[i] = (wm_f16)d[i];
return wm_op_from_acc(h);
}
#endif
// 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
@@ -219,7 +313,7 @@ WM_FN void wm_acc_vec_load(const wm_f16* p, wm_f16 out[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]);
out[i] = __builtin_bit_cast(wm_f16, (uint16_t)(v[i] >> (hf << 4)));
#endif
}
kernels/k/dltss_pwin_dec0_layer.hip
+8 -2
View File
@@ -1,5 +1,11 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers). The WMMA half
// exchange goes through ds_bpermute (VALU-bound layer: -1..-3%)
// f32 -> f16 pairs converted straight into one register (WM_CVT_ASM: no v_pack copies; -1..-3 %)
// norm / residual / embedding ReLU on packed f16 pairs (PWIN_PK_ELEM with the packed g2/b2 table; same ops, -1 %)
// d4r-build-flags: -mcumode -DWM_HALF_BPERMUTE -DWM_CVT_ASM -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
// RDNA4: no WM_CVT_ASM (its inline asm needs the compiler's real-true16 mode on gfx12; the plain conversion
// costs one pack per pair and gfx12 operands need no lane exchange anyway)
// d4r-build-flags-gfx12: -mcumode -DWM_HALF_BPERMUTE -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
#define PWIN_F32ACC
#include "pwin_pos.h"
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)
kernels/k/dltss_pwin_dec1_layer.hip
+5 -2
View File
@@ -1,5 +1,8 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
// WGP mode: with the padded LDS rows (PWIN_PAD) this layer is 2-5% faster in WGP mode than in CU mode
// (CU mode was 1-6% faster before the padding)
// accumulator init vectors from the weights (the prepared f32 table is 2-4% slower here)
// Q, K, V chains two at a time (fewer live accumulators, spills 164 -> 68 B; -2 %)
// d4r-build-flags: -DPWIN_VEC_FROM_WEIGHTS -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
+4 -2
View File
@@ -1,5 +1,7 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
// WGP mode: with the padded LDS rows (PWIN_PAD) this layer is 2-5% faster in WGP mode than in CU mode
// (CU mode was 1-6% faster before the padding)
// weight images through plain global loads here (buffer loads were 1-3% slower in this layer)
// d4r-build-flags: -DPWIN_NO_BUFFER_IMG -DPWIN_VEC_FROM_WEIGHTS
#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
@@ -1,2 +1,4 @@
// the WMMA half exchange through ds_bpermute instead of v_permlanex16: -1.5% here
// d4r-build-flags: -DWM_HALF_BPERMUTE
#include "pwin_wide.h"
WIDE_DECODER(dltss_pwin_dec3_layer, 4, 96, 128, 2)
kernels/k/dltss_pwin_dec4_layer.hip
+2 -1
View File
@@ -1,2 +1,3 @@
// NG = 8 (32 waves per window): with the padded LDS rows faster than NG = 4 here
#include "pwin_wide.h"
WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 4)
WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 8)
kernels/k/dltss_pwin_dec5_layer.hip
+4 -2
View File
@@ -1,2 +1,4 @@
#include "pwin_wide.h"
WIDE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 4)
// dec5 (bottleneck, H 8, C 160) as four phase kernels (pwin_phase.h). Needs the patched
// ZLUDA hook (post kernels). MAXW 64 windows (1440p: 9).
#include "pwin_phase.h"
PHASE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 64, 128, 96, 384, 96)
kernels/k/dltss_pwin_enc0_layer.hip
+10 -2
View File
@@ -1,5 +1,13 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers). The WMMA half
// exchange goes through ds_bpermute (VALU-bound layer: -1..-3%)
// f32 -> f16 pairs converted straight into one register (WM_CVT_ASM: no v_pack copies; -1..-3 %)
// accumulator-init vectors: f32 table read with scalar loads + one select per value (PWIN_VEC_SCALAR, -2 %; replaces
// PWIN_VEC_FROM_WEIGHTS)
// norm / residual / embedding ReLU on packed f16 pairs (PWIN_PK_ELEM with the packed g2/b2 table; same ops, -1 %)
// d4r-build-flags: -mcumode -DWM_HALF_BPERMUTE -DPWIN_VEC_SCALAR -DWM_CVT_ASM -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
// RDNA4: no WM_CVT_ASM (its inline asm needs the compiler's real-true16 mode on gfx12; the plain conversion
// costs one pack per pair and gfx12 operands need no lane exchange anyway)
// d4r-build-flags-gfx12: -mcumode -DWM_HALF_BPERMUTE -DPWIN_VEC_SCALAR -DPWIN_PK_ELEM -DPWIN_PACKED_RESIDUAL
#define PWIN_F32ACC
#include "pwin_pos.h"
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)
kernels/k/dltss_pwin_enc1_layer.hip
+4 -4
View File
@@ -1,5 +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
// Exact single-buffer enc1: reload the original input for the residual, preserve WGP mode.
// d4r-build-flags: -DPOS_VSHARE1 -DPWIN_VEC_SCALAR -DPWIN_PK_ELEM
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
#include "pwin_pos_recompute.h"
POS_RECOMPUTE_ENCODER(dltss_pwin_enc1_layer, 2, 64, 64, 0, 1)
kernels/k/dltss_pwin_enc2_layer.hip
+4 -2
View File
@@ -1,5 +1,7 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
// WGP mode: with the padded LDS rows (PWIN_PAD) this layer is 2-5% faster in WGP mode than in CU mode
// (CU mode was 1-6% faster before the padding)
// weight images through plain global loads here (buffer loads were 1-3% slower in this layer)
// d4r-build-flags: -DPWIN_NO_BUFFER_IMG -DPWIN_VEC_FROM_WEIGHTS
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)
kernels/k/dltss_pwin_enc4_layer.hip
+2 -1
View File
@@ -1,2 +1,3 @@
// NG = 8 (32 waves per window): with the padded LDS rows faster than NG = 4 here
#include "pwin_wide.h"
WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 4)
WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 8)
kernels/k/pwin_common.h
+104
View File
@@ -13,7 +13,10 @@
#if defined(D4R_ACCURACY) && defined(PWIN_F32ACC)
#undef PWIN_F32ACC
#endif
// D4R_FAST_MATH: the opt-in approximate variant (FMA contraction allowed; not bit-identical to NVIDIA's rounding)
#ifndef D4R_FAST_MATH
#pragma clang fp contract(off)
#endif
typedef _Float16 half_t;
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
@@ -22,6 +25,32 @@ typedef float f8v __attribute__((ext_vector_type(8)));
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
__device__ __forceinline__ half_t half_reciprocal(half_t h)
{
#ifdef PWIN_NATIVE_HRECIP
return __builtin_amdgcn_rcph(h);
#else
return (half_t)__builtin_amdgcn_rcpf((float)h);
#endif
}
__device__ __forceinline__ half_t half_rsqrt(half_t h)
{
#ifdef PWIN_NATIVE_HRECIP
return __builtin_amdgcn_rsqh(h);
#else
return (half_t)__builtin_amdgcn_rsqf((float)h);
#endif
}
// LDS row padding in halves: activation rows of N halves are stored N + PWIN_PAD apart, so the 16 token rows a
// WMMA operand load (op_lds / lds_row16) touches start on distinct banks (row stride = 16 mod 32 bytes). With
// unpadded 128- or 256-byte rows every lane hits the same banks (measured: 6x slower WMMA feed). 0 = old layout.
#ifndef PWIN_PAD
#define PWIN_PAD 8
#endif
#define LDSR(N) ((N) + PWIN_PAD)
__device__ __forceinline__ uint32_t lane_id()
{
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
@@ -33,11 +62,24 @@ typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
// PWIN_F32ACC: keep the accumulator in f32 through the chain (rounded to f16 where the values are used)
__device__ __forceinline__ f8v mma16(const op_t& a, const op_t& b, f8v c)
{
#if defined(PWIN_NATIVE_F16ACC) && !defined(PWIN_F32ACC) && D4R_WMMA_LAYOUT == 11
h16 hc = {};
#pragma unroll
for (int i = 0; i < 8; ++i)
hc[2 * i] = (half_t)c[i];
const h16 hd = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(
__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), hc, false);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)hd[2 * i];
#else
f8v d = wm_mma(a, b, c);
#ifndef PWIN_F32ACC
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)(half_t)d[i];
#endif
#endif
return d;
}
@@ -136,6 +178,46 @@ __device__ __forceinline__ op_t op_img(const u8v* __restrict__ img, int idx)
{
return wm_op_image(img, idx);
}
// 8 floats of image slot `slot` + this lane's half (f32 accumulator-init vectors the prep kernel stores after the
// operand tiles: slot pair (2 v, 2 v + 1) = rows wm_acc_row_h(., h) of vector v), slot uniform
__device__ __forceinline__ f8v img_vec(const u8v* __restrict__ img, int slot)
{
const uint32_t h = lane_id() >> 4;
#if defined(PWIN_VEC_SCALAR) // both layouts: slot pair (2 v, 2 v + 1) = the two halves' rows
// both halves' slots with wave-uniform (scalar) loads, then one select per value (instead of a per-lane
// vector load, or 8 shifts + 8 f16 -> f32 conversions from the weights); the same f32 values
const u8v x0 = img[slot], x1 = img[slot + 1];
u8v x;
#pragma unroll
for (int i = 0; i < 8; ++i)
x[i] = h ? x1[i] : x0[i];
return __builtin_bit_cast(f8v, x);
#elif D4R_WMMA_LAYOUT == 12 || defined(PWIN_NO_BUFFER_IMG)
const u8v x = img[slot + (int)h];
return __builtin_bit_cast(f8v, x);
#else
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc((void*)img, (short)0, 0x7fffffff, 0x31004000);
const uint32_t vo = h * 32u, so = 32u * (uint32_t)slot;
const u4v a = __builtin_amdgcn_raw_buffer_load_b128(r, vo, so, 0), b = __builtin_amdgcn_raw_buffer_load_b128(r, vo + 16, so, 0);
return __builtin_bit_cast(f8v, (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]});
#endif
}
// slot (tile, lane l & 15) of an operand image in global memory: tile is wave-uniform, so the row address stays a
// scalar base and only the lane indexes it (global loads with an SGPR base and a VGPR offset; written as one
// index tile * 16 + m every load paid 64-bit VALU address arithmetic)
__device__ __forceinline__ op_t op_imgt(const u8v* __restrict__ img, int tile)
{
#if D4R_WMMA_LAYOUT == 12 || defined(PWIN_NO_BUFFER_IMG)
return wm_op_image(img + 16 * tile, (int)(lane_id() & 15));
#else
// gfx11: a raw buffer load (resource from img: base in SGPRs, 0x31004000 = gfx11 raw dword3, validated), the
// tile offset as the scalar offset and the lane's 32-byte slot as the vector offset: no VALU address math
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc((void*)img, (short)0, 0x7fffffff, 0x31004000);
const uint32_t vo = (lane_id() & 15) * 32u, so = 512u * (uint32_t)tile;
const u4v a = __builtin_amdgcn_raw_buffer_load_b128(r, vo, so, 0), b = __builtin_amdgcn_raw_buffer_load_b128(r, vo + 16, so, 0);
return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
#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)
{
@@ -184,3 +266,25 @@ __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);
}
// Diagnostic phase stamps (-DPWIN_PROF): per-wave SHADER_CYCLES deltas summed into d4r_dbg[k], wave counts into
// d4r_dbg[32 + k] (read back by bench/native/pwbench). Region k ends at PROF(k). Off: no code.
#ifdef PWIN_PROF
__device__ unsigned long long d4r_dbg[64];
#define PROF_INIT uint32_t prof_last_ = (uint32_t)__builtin_readcyclecounter()
#define PROF(k) \
do \
{ \
const uint32_t n_ = (uint32_t)__builtin_readcyclecounter(); \
if (lane_id() == 0) \
{ \
__hip_atomic_fetch_add(&d4r_dbg[k], (unsigned long long)((n_ - prof_last_) & 0xfffffu), __ATOMIC_RELAXED, \
__HIP_MEMORY_SCOPE_AGENT); \
__hip_atomic_fetch_add(&d4r_dbg[32 + (k)], 1ull, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); \
} \
prof_last_ = n_; \
} while (0)
#else
#define PROF_INIT do {} while (0)
#define PROF(k) do {} while (0)
#endif
kernels/k/pwin_gelu_lut.h
+37
View File
@@ -0,0 +1,37 @@
// PWIN_GELU_LUT: exact f16 GELU by exhaustive lookup, populated on the prep stream.
// Include after gelu8_math. The POS prep wrappers reserve at least 65536 work-items.
// Every index comes from a uint16_t f16 bit pattern, so it is inside this allocation.
#pragma once
__device__ __attribute__((aligned(256))) uint16_t g_gelu_lut[65536];
constexpr int pwin_gelu_prep_items(int n) { return n < 65536 ? 65536 : n; }
__device__ __forceinline__ void pwin_gelu_prep()
{
const unsigned i = blockIdx.x * 128u + threadIdx.x;
if (i < 65536u)
{
const half_t x = __builtin_bit_cast(half_t, (uint16_t)i);
half_t g[8];
gelu8_math(splat8((float)x), g);
g_gelu_lut[i] = __builtin_bit_cast(uint16_t, g[0]);
}
}
__device__ __forceinline__ void gelu8_lut(const f8v& d, half_t g[8])
{
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const uint16_t bits = __builtin_bit_cast(uint16_t, (half_t)d[i]);
#ifdef PWIN_GELU_LUT_BUFFER
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc(
(void*)g_gelu_lut, (short)0, sizeof(g_gelu_lut), 0x31004000);
const uint16_t value = __builtin_amdgcn_raw_buffer_load_b16(r, (unsigned)bits * 2u, 0u, 0);
#else
const uint16_t value = g_gelu_lut[bits];
#endif
g[i] = __builtin_bit_cast(half_t, value);
}
}
kernels/k/pwin_layer.h
+501 -124
View File
@@ -38,7 +38,12 @@ template <int H_, int C_, int COUT_, bool POS_ = false> struct PwinLayout
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;
// f32 accumulator-init vectors in image slots after the tiles (see img_vec): Wo bias, MLP fc1 biases, merge bias
static constexpr int V_BO = 0, V_B1 = V_BO + KT, V_PM = V_B1 + 2 * NMLP, NV = V_PM + PMT;
static constexpr int VEC0 = T_END * 16, IMG_SLOTS = VEC0 + 2 * NV;
static constexpr int PREP_ITEMS = T_END * 16 + NBIAS + 2 * NV;
__device__ static const uint8_t* extra_vec(const uint8_t*, const uint8_t*, int) { return nullptr; }
__device__ static int src_of(int tile)
{
@@ -87,9 +92,101 @@ __device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
return r;
}
// one (vector, half) slot of the f32 vector table: the 8 rows wm_acc_row_h(i, h) of the 16 halves at src, as floats
__device__ __forceinline__ void vec_slot(u8v* img, int slot, const uint8_t* src, int h)
{
f8v f;
#pragma unroll
for (int i = 0; i < 8; ++i)
f[i] = (float)wload(src, 2 * wm_acc_row_h(i, h));
img[slot] = __builtin_bit_cast(u8v, f);
}
// Optional 640-byte module table for G2/B2 in each half-wave's accumulator order.
#if defined(PWIN_PACKED_RESIDUAL) && D4R_WMMA_LAYOUT == 11
__device__ __attribute__((aligned(16))) u4v g_residual_vectors[40];
#endif
template <class L> __device__ __forceinline__ void residual_vector_prep(const uint8_t* w, int idx)
{
#if defined(PWIN_PACKED_RESIDUAL) && D4R_WMMA_LAYOUT == 11
static_assert(L::KT <= 10, "residual table capacity");
if (idx >= 0 && idx < 4 * L::KT)
{
const int kind = idx / (2 * L::KT), nt = (idx >> 1) % L::KT, h = idx & 1;
const int base = (kind ? L::B2 : L::G2) + 32 * nt;
u4v v;
#pragma unroll
for (int j = 0; j < 4; ++j)
v[j] = pack2(wload(w, base + 2 * (4 * j + h)), wload(w, base + 2 * (4 * j + 2 + h)));
g_residual_vectors[idx] = v;
}
#endif
}
template <class L> __device__ __forceinline__ void residual_vector(const uint8_t* w, int kind, int nt, half_t out[8])
{
#if defined(PWIN_PACKED_RESIDUAL) && D4R_WMMA_LAYOUT == 11
static_assert(L::KT <= 10, "residual table capacity");
const u4v v = g_residual_vectors[2 * L::KT * kind + 2 * nt + (lane_id() >> 4)];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
out[2 * j] = lo16(v[j]);
out[2 * j + 1] = hi16(v[j]);
}
#else
dvec8(w, (kind ? L::B2 : L::G2) + 32 * nt, out);
#endif
}
// prep items of the f32 vector table (j = 0 .. 2 NV - 1): core vectors from the core weights wc, the derived
// layouts' extra vectors from w0 (the layer's weights pointer). MERGE: the layer has a patch merge (its merge bias
// exists; for decoders and the bottleneck that offset lies past the layer's weights, so it is never read).
template <class L, bool MERGE> __device__ void pwin_vec_prep(const uint8_t* w0, const uint8_t* wc, u8v* img, int j)
{
const int v = j >> 1, h = j & 1;
const uint8_t* src = nullptr;
if (v < L::V_B1)
src = wc + L::BO + 32 * (v - L::V_BO);
else if (v < L::V_PM)
src = wc + L::B1 + 32 * (v - L::V_B1);
else if (v < L::V_PM + L::PMT)
{
if (!MERGE)
return;
src = wc + L::PMB + 32 * (v - L::V_PM);
}
else
src = L::extra_vec(w0, wc, v);
if (src)
vec_slot(img, L::VEC0 + j, src, h);
}
// accumulator init vector v (16 halves at w + byte as floats in accumulator row order): from the prepared f32 table
// (one 32-byte load per lane), or with PWIN_VEC_FROM_WEIGHTS from the weights (load + 8 shifts + 8 conversions;
// faster in some layers). The same values.
template <class L> __device__ __forceinline__ f8v acc_vec(const u8v* img, const uint8_t* w, int byte, int v)
{
#ifdef PWIN_VEC_FROM_WEIGHTS
half_t t[8];
dvec8(w, byte, t);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)t[i];
(void)img; (void)v;
return d;
#else
(void)w; (void)byte;
return img_vec(img, L::VEC0 + 2 * v);
#endif
}
template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx)
{
const uint8_t* w = p.w;
residual_vector_prep<L>(w, idx);
if (idx < L::T_END * 16)
{
img[idx] = img_tile(w, L::src_of(idx >> 4), idx & 15);
@@ -134,9 +231,20 @@ template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u
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)
template <class L, bool MERGE = true> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
{
pwin_prep_at<L>(p, img, bias, blockIdx.x * 128 + threadIdx.x);
const int idx = blockIdx.x * 128 + threadIdx.x, j = idx - (L::T_END * 16 + L::NBIAS);
if (j >= 0)
{
if (j < 2 * L::NV)
pwin_vec_prep<L, MERGE>(p.w, p.w, img, j);
return;
}
pwin_prep_at<L>(p, img, bias, idx);
}
template <class L> __device__ void pwin_prep_nomerge(const PwinParams& p, u8v* img, u4v* bias)
{
pwin_prep<L, false>(p, img, bias);
}
__device__ __forceinline__ half_t hmin(half_t a, half_t b) { return a < b ? a : b; }
@@ -150,7 +258,11 @@ __device__ __forceinline__ half_t softmax_e(half_t s)
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);
#ifdef PWIN_NATIVE_HEXP
return __builtin_elementwise_exp2(x);
#else
return (half_t)__builtin_amdgcn_exp2f((float)x);
#endif
}
__device__ __forceinline__ half_t gelu(half_t x)
@@ -163,7 +275,7 @@ __device__ __forceinline__ half_t gelu(half_t x)
}
// gelu of 8 accumulator values as 4 packed f16 pairs (the same f16 operations as gelu(), element-wise)
__device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
__device__ __forceinline__ void gelu8_math(const f8v& d, half_t g[8])
{
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
#pragma unroll
@@ -178,6 +290,121 @@ __device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
}
}
// Exact GELU table experiment; enabled only by the POS wrappers that initialize it in prep.
#ifdef PWIN_GELU_LUT
#include "pwin_gelu_lut.h"
#else
constexpr int pwin_gelu_prep_items(int n) { return n; }
__device__ __forceinline__ void pwin_gelu_prep() {}
#endif
__device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
{
#ifdef PWIN_GELU_LUT
gelu8_lut(d, g);
#else
gelu8_math(d, g);
#endif
}
// gelu of the 8 accumulator values as the next GEMM's operand. With WM_CVT_ASM the pairs stay packed from the
// conversion through the packed f16 math to the operand exchange (the same f16 operations as gelu8_math).
__device__ __forceinline__ op_t gelu_op(const f8v& d)
{
#if defined(PWIN_GELU_PACKED) && defined(WM_HAVE_FROM_MINE) && !defined(PWIN_GELU_LUT)
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
#ifdef PWIN_GELU_PACKED_C
const hv2 x = (hv2){(half_t)d[2 * j], (half_t)d[2 * j + 1]}; // C conversions (no asm in the MLP)
#else
const hv2 x = __builtin_bit_cast(hv2, wm_cvt2(d[2 * j], d[2 * j + 1]));
#endif
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);
mine[j] = __builtin_bit_cast(uint32_t, x * ((hv2)(half_t)0.5f + c * u));
}
return wm_op_from_mine(mine);
#else
half_t g[8];
gelu8(d, g);
return operand_from_dt(g);
#endif
}
// ---- accumulator abstraction for the POS layers (pwin_pos.h).
// Exact builds: pacc = f8v and every helper is the operation the code used before (bit-identical).
// PWIN_F16ACC_NATIVE (opt-in approximate variant): accumulators stay in f16 registers through
// the native f16-output WMMA (value i in the low half of dword i); this rounds every k16 step to f16 instead of
// keeping f32 through the chain (harness PSNR ~62 dB vs the exact output), and removes most f32 -> f16 conversions.
#ifdef PWIN_F16ACC_NATIVE
static_assert(D4R_WMMA_LAYOUT == 11, "f16 accumulator variant: gfx11 layout");
typedef h16 pacc;
__device__ __forceinline__ pacc acc_zero() { return (pacc){}; }
template <class L> __device__ __forceinline__ pacc acc_vec_p(const u8v* img, const uint8_t* w, int byte, int v)
{
const f8v f = acc_vec<L>(img, w, byte, v);
pacc r = {};
#pragma unroll
for (int i = 0; i < 8; ++i)
r[2 * i] = (half_t)f[i];
return r;
}
__device__ __forceinline__ pacc mma16p(const op_t& a, const op_t& b, pacc c)
{
return __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c, false);
}
// elements 2 j, 2 j + 1 as one packed register (low halves of dwords 2 j and 2 j + 1: one byte permutation)
__device__ __forceinline__ uint32_t acc_pair(const pacc& d, int j)
{
const u8v u = __builtin_bit_cast(u8v, d);
return __builtin_amdgcn_perm(u[2 * j + 1], u[2 * j], 0x05040100u);
}
__device__ __forceinline__ void acc_set_pair(pacc& d, int j, hv2 y)
{
d[4 * j] = y[0];
d[4 * j + 2] = y[1];
}
__device__ __forceinline__ op_t operand_from_acc_p(const pacc& d)
{
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
mine[j] = acc_pair(d, j);
return wm_op_from_mine(mine);
}
__device__ __forceinline__ op_t gelu_op_p(const pacc& d)
{
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
const hv2 x = __builtin_bit_cast(hv2, acc_pair(d, j));
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);
mine[j] = __builtin_bit_cast(uint32_t, x * ((hv2)(half_t)0.5f + c * u));
}
return wm_op_from_mine(mine);
}
#else
typedef f8v pacc;
__device__ __forceinline__ pacc acc_zero() { return splat8(0.0f); }
template <class L> __device__ __forceinline__ pacc acc_vec_p(const u8v* img, const uint8_t* w, int byte, int v) { return acc_vec<L>(img, w, byte, v); }
__device__ __forceinline__ pacc mma16p(const op_t& a, const op_t& b, pacc c) { return mma16(a, b, c); }
__device__ __forceinline__ uint32_t acc_pair(const pacc& d, int j) { return wm_cvt2(d[2 * j], d[2 * j + 1]); }
__device__ __forceinline__ void acc_set_pair(pacc& d, int j, hv2 y)
{
d[2 * j] = (float)y[0];
d[2 * j + 1] = (float)y[1];
}
__device__ __forceinline__ op_t operand_from_acc_p(const pacc& d) { return operand_from_f8(d); }
__device__ __forceinline__ op_t gelu_op_p(const pacc& d) { return gelu_op(d); }
#endif
// 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)
{
@@ -198,6 +425,8 @@ __device__ uint64_t g_stamp[16384][16];
g_stamp[bid][k] = __builtin_readcyclecounter(); \
} \
} while (0)
#elif defined(PWIN_PROF)
#define STAMP(k) PROF(k)
#else
#define STAMP(k) do {} while (0)
#endif
@@ -246,17 +475,59 @@ __device__ __forceinline__ float row_sum_total(float rs, const half_t (&e)[4][8]
// 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>
#ifdef PWIN_VT_DIRECT
// V^T operands in the vt buffer: slot r (0..127) at byte 48 r (32 B operand + 16 B pad: 16 lanes reading 16
// consecutive slots with b128 loads hit disjoint LDS banks; an unpadded 32 B pitch conflicts 2-way)
#define PWIN_VT_ROWS 43 /* 43 x LDSR(64) halves = 6192 B >= 128 x 48 B */
static_assert(PWIN_VT_ROWS * LDSR(64) * 2 >= 128 * 48, "vt operand slots");
static_assert(sizeof(op_t) == 32, "gfx11 operand layout");
__device__ __forceinline__ void vt_op_store(half_t* vt, int r, const op_t& v)
{
if (!wm_half())
{
u4v* q = (u4v*)((uint8_t*)vt + 48 * r);
q[0] = (u4v){v[0], v[1], v[2], v[3]};
q[1] = (u4v){v[4], v[5], v[6], v[7]};
}
}
__device__ __forceinline__ op_t vt_op_load(const half_t* vt, int r)
{
const u4v* q = (const u4v*)((const uint8_t*)vt + 48 * r);
const u4v a = q[0], b = q[1];
return (op_t){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
}
#else
#define PWIN_VT_ROWS 32
#endif
struct PwinNoX0
{
static constexpr bool active = false;
__device__ void operator()() const {}
};
template <class F, class B> struct PwinX0
{
static constexpr bool active = true;
F& f;
B b;
__device__ void operator()() const { f(b); }
};
// X0F (PWIN_DEC_RECOMPUTE): act and hb are one buffer, so the residual's x0 is recomputed after attention by x0f()
// into x0b (rows of this layer's tokens, aliasing kl / vt, which are dead by then)
template <class L, class X0F = PwinNoX0>
__device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
half_t (*act)[L::C], half_t (*hb)[L::C], half_t (*kl)[32], half_t (*vt)[64])
half_t (*act)[LDSR(L::C)], half_t (*hb)[LDSR(L::C)], half_t (*kl)[LDSR(32)], half_t (*vt)[LDSR(64)],
const X0F& x0f = X0F(), half_t (*x0b)[LDSR(L::C)] = nullptr)
{
constexpr int C = L::C, KT = L::KT, H = L::H;
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
const uint8_t* w = p.w;
half_t* arow = act[16 * wv + m];
half_t* hrow = hb[16 * wv + m];
PROF_INIT;
STAMP(1);
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb
float ss = 0.0f;
@@ -264,13 +535,23 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
const u8v xv = lds_row16(arow + 16 * kt);
#pragma unroll
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_SS)
for (int j = 0; j < 8; ++j)
{
const uint32_t xj = xv[j]; // (no __builtin_bit_cast of a vector subscript: clang bug, REWRITE.md)
const hv2 v = __builtin_bit_cast(hv2, xj), q = v * v;
ss += (float)q[0];
ss += (float)q[1];
}
#else
for (int c = 0; c < 16; ++c)
{
const half_t v = op_get(xv, c);
ss += (float)(half_t)(v * v);
}
#endif
}
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
const half_t r16 = half_rsqrt((half_t)ss);
for (int kt = (int)hf; kt < KT; kt += 2)
{
const u8v xv = lds_row16(arow + 16 * kt);
@@ -278,8 +559,15 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
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)
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_APPLY)
{
const uint32_t xj = xv[j], gj = gv[j];
hv[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, xj) * ((hv2)r16 * __builtin_bit_cast(hv2, gj)));
}
#else
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)));
#endif
store_row16(hrow + 16 * kt, hv);
}
__builtin_amdgcn_wave_barrier();
@@ -296,19 +584,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
f8v 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];
}
{ acc[nt] = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt); }
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)
@@ -322,7 +598,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
f8v d = splat8(0.0f);
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
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();
@@ -337,11 +613,59 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = op_img((const u8v*)bias, ((h * 4 + wv) * 4 + kt) * 16 + m);
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + wv) * 4 + kt);
}
else
{
op_t qop[2];
#ifdef PWIN_BIAS_PREFETCH
// this head's score bias tiles, loaded before the QKV chains (used right after the next barrier)
u4v bpre[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
bpre[kt] = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
#endif
#ifdef PWIN_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
#pragma unroll
for (int j = 2; j >= 0; --j)
{
f8v 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)
{
const op_t wt = op_imgt(img, L::T_QKV + ((h * 3 + j) * KT + kt) * 2 + nt);
#ifdef PWIN_VT_DIRECT
dj[nt] = j == 2 ? mma16(row, wt, dj[nt]) : mma16(wt, row, dj[nt]);
#else
dj[nt] = mma16(wt, row, dj[nt]);
#endif
}
}
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const f8v 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
{
#ifdef PWIN_VT_DIRECT
vt_op_store(&vt[0][0], 64 * nt + wv * 16 + m, 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
}
}
}
#else
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
f8v dq[6];
#pragma unroll
@@ -350,9 +674,19 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int kt = 0; kt < KT; ++kt)
{
const op_t row = op_lds(hrow + 16 * kt);
#ifdef PWIN_VT_DIRECT
// V chains with the operands swapped (as in the POS layers): D = V^T of this wave's 16 keys, lane = dim
#pragma unroll
for (int u = 0; u < 4; ++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 u = 4; u < 6; ++u)
dq[u] = mma16(row, op_imgt(img, L::T_QKV + ((h * 3 + 2) * KT + kt) * 2 + (u & 1)), dq[u]);
#else
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
dq[u] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[u]);
#endif
}
#pragma unroll
for (int j = 0; j < 3; ++j)
@@ -368,18 +702,28 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
else
{
#ifdef PWIN_VT_DIRECT
// V^T operand of key tile wv -> slot nt * 64 + wv * 16 + m of the vt buffer (128 x 48 B)
vt_op_store(&vt[0][0], 64 * nt + wv * 16 + m, 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
}
}
#endif
block_sync();
half_t e[4][8];
float rs = 0.0f;
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
#ifdef PWIN_BIAS_PREFETCH
const u4v bi = bpre[kt];
#else
const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
#endif
f8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
@@ -399,7 +743,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
}
rs = row_sum_total(rs, e);
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
const half_t rc = half_reciprocal((half_t)rs);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
@@ -421,7 +765,13 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
if constexpr (L::POS)
d = mma16(vtop[nt][kt], pop[kt], d);
else
{
#ifdef PWIN_VT_DIRECT
d = mma16(vt_op_load(&vt[0][0], 64 * nt + kt * 16 + m), pop[kt], d);
#else
d = mma16(op_lds(&vt[16 * nt + m][16 * kt]), pop[kt], d);
#endif
}
}
oop[nt] = operand_from_f8(d);
#ifdef PWIN_DBG
@@ -434,7 +784,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
acc[nt] = mma16(op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m), oop[ks], acc[nt]);
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
}
@@ -447,13 +797,33 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
#endif
STAMP(3);
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
if constexpr (X0F::active)
{
x0f(); // every wave's attention is done (block barrier at the end of the last head)
block_sync(); // x0 rows of all waves' tokens written
}
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
const op_t xv = op_lds(arow + 16 * nt);
const op_t xv = X0F::active ? op_lds(&x0b[16 * wv + m][16 * nt]) : 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);
residual_vector<L>(w, 0, nt, g2);
residual_vector<L>(w, 1, nt, b2);
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_RESID)
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
const hv2 x1 = __builtin_bit_cast(hv2, wm_cvt2(acc[nt][2 * j], acc[nt][2 * j + 1])) +
__builtin_bit_cast(hv2, wm_op_acc_pair(xv, j));
mine[j] = __builtin_bit_cast(uint32_t, x1 * (hv2){g2[2 * j], g2[2 * j + 1]});
const hv2 y = x1 + (hv2){b2[2 * j], b2[2 * j + 1]};
acc[nt][2 * j] = (float)y[0];
acc[nt][2 * j + 1] = (float)y[1];
}
(void)mv;
op_store(hrow + 16 * nt, wm_op_from_mine(mine));
#else
#pragma unroll
for (int i = 0; i < 8; ++i)
{
@@ -462,6 +832,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
acc[nt][i] = (float)(half_t)(x1 + b2[i]);
}
op_store(hrow + 16 * nt, operand_from_dt(mv));
#endif
}
__builtin_amdgcn_wave_barrier();
@@ -473,30 +844,16 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int hn = 0; hn < 2; ++hn)
{
f8v 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];
}
{ d = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn); }
for (int kt = 0; kt < KT; ++kt)
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
half_t g[8];
gelu8(d, g);
gop[hn] = operand_from_dt(g);
d = mma16(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_lds(hrow + 16 * kt), d);
gop[hn] = gelu_op(d);
}
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
acc[nt] = mma16(op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m), gop[ks], acc[nt]);
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)
@@ -506,6 +863,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
}
__builtin_amdgcn_wave_barrier();
PROF(5);
}
// first layer: token embedding x0 = relu(in (CIN) W_e + b_e) in front of the core (core weights at CB)
@@ -514,12 +872,22 @@ template <int H_, int C_, int COUT_, bool POS_, int CIN_> struct PwinEmbLayout :
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;
// + the embedding bias (C / 16 vectors at the layer's weights pointer)
static constexpr int V_EMB = B::NV, NV = V_EMB + (CIN ? C_ / 16 : 0);
static constexpr int VEC0 = T_END * 16, IMG_SLOTS = VEC0 + 2 * NV;
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS + 2 * NV;
__device__ static const uint8_t* extra_vec(const uint8_t* w0, const uint8_t*, int v) { return w0 + 32 * (v - V_EMB); }
};
template <class L> __device__ void pwin_emb_prep(const PwinParams& p, u8v* img, u4v* bias)
{
const int idx = blockIdx.x * 128 + threadIdx.x;
const int idx = blockIdx.x * 128 + threadIdx.x, j = idx - (L::T_END * 16 + L::B::NBIAS);
if (j >= 0)
{
if (j < 2 * L::NV)
pwin_vec_prep<L, true>(p.w, p.w + L::CB, img, j);
return;
}
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);
@@ -543,14 +911,15 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
if constexpr (L::CB != 0)
p.w = p0.w + L::CB;
constexpr int C = L::C, COUT = L::COUT;
__shared__ __attribute__((aligned(16))) half_t act[64][C];
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
__shared__ __attribute__((aligned(16))) half_t kl[64][32];
__shared__ __attribute__((aligned(16))) half_t vt[32][64];
__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[PWIN_VT_ROWS][LDSR(64)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
const int bx = blockIdx.x, by = blockIdx.y;
PROF_INIT;
STAMP(0);
// input rows (mirror padded) -> act (first layer: embedding + ReLU)
{
@@ -566,16 +935,10 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
for (int nt = 0; nt < C / 16; ++nt)
{
f8v 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];
}
{ d = acc_vec<L>(img, p0.w, 32 * nt, L::V_EMB + nt); }
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
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)
@@ -614,25 +977,14 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
for (int vt = wv; vt < L::PMT; vt += 4)
{
f8v d;
{
{ d = acc_vec<L>(img, w, L::PMB + 32 * vt, L::V_PM + vt);
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);
}
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_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
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;
@@ -642,7 +994,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
#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__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__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; \
@@ -663,12 +1015,12 @@ template <class L>
__device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C;
__shared__ __attribute__((aligned(16))) half_t act[64][C];
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
__shared__ __attribute__((aligned(16))) half_t kl[64][32];
__shared__ __attribute__((aligned(16))) half_t vt[32][64];
__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[PWIN_VT_ROWS][LDSR(64)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
const int bx = blockIdx.x, by = blockIdx.y;
{
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
@@ -689,7 +1041,7 @@ __device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, con
#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__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__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; \
@@ -715,12 +1067,25 @@ struct PwinDecLayout : PwinLayout<H_, C_, 16, POS_>
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;
// + the expand biases (EXPN vectors at w0 + EXPB) and the output head bias (NOUTA / 16 at core + HB)
static constexpr int V_EXP = B::NV, V_HB = V_EXP + EXPN, NV = V_HB + NOUTA_ / 16;
static constexpr int VEC0 = T_END * 16, IMG_SLOTS = VEC0 + 2 * NV;
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS + 2 * NV;
__device__ static const uint8_t* extra_vec(const uint8_t* w0, const uint8_t* wc, int v)
{
return v < V_HB ? w0 + EXPB + 32 * (v - V_EXP) : wc + HB + 32 * (v - V_HB);
}
};
template <class L> __device__ void pwin_dec_prep(const PwinParams& p, u8v* img, u4v* bias)
{
const int idx = blockIdx.x * 128 + threadIdx.x;
const int idx = blockIdx.x * 128 + threadIdx.x, j = idx - (L::T_END * 16 + L::B::NBIAS);
if (j >= 0)
{
if (j < 2 * L::NV)
pwin_vec_prep<L, false>(p.w, p.w + L::CB, img, j);
return;
}
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;
@@ -743,16 +1108,35 @@ template <class L>
__device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C;
__shared__ __attribute__((aligned(16))) half_t act[64][C];
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : 32;
__shared__ __attribute__((aligned(16))) half_t kl[KVR][32];
__shared__ __attribute__((aligned(16))) half_t vt[KVC][64];
__shared__ __attribute__((aligned(16))) half_t act[64][LDSR(C)];
constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : PWIN_VT_ROWS;
#ifdef PWIN_DEC_RECOMPUTE
// one buffer for act / hb; kl, vt and (after attention) the recomputed x0 rows share one region
static_assert(!L::POS, "PWIN_DEC_RECOMPUTE: full-attention decoders");
constexpr int KLB = KVR * LDSR(32) * 2, VTB = KVC * LDSR(64) * 2, X0B = 64 * LDSR(C) * 2;
constexpr int KVXB = (KLB + VTB) > X0B ? (KLB + VTB) : X0B;
__shared__ __attribute__((aligned(16))) uint8_t kvx[KVXB];
half_t (*hb)[LDSR(C)] = act;
half_t (*kl)[LDSR(32)] = (half_t (*)[LDSR(32)])kvx;
half_t (*vt)[LDSR(64)] = (half_t (*)[LDSR(64)])(kvx + KLB);
half_t (*x0b)[LDSR(C)] = (half_t (*)[LDSR(C)])kvx;
#else
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t kl[KVR][LDSR(32)];
__shared__ __attribute__((aligned(16))) half_t vt[KVC][LDSR(64)];
#endif
#ifdef PWIN_LDS_PAD
// occupancy-sensitivity probe only: PWIN_LDS_PAD bytes of unused LDS per block (one dword written)
__shared__ uint32_t lds_pad[PWIN_LDS_PAD / 4];
if (threadIdx.x == 0 && threadIdx.z == 0)
((volatile uint32_t*)lds_pad)[0] = 0u;
#endif
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
PROF_INIT;
// ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens
{
// ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens -> dst rows
auto expand = [&](half_t (*dst)[LDSR(C)]) {
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);
@@ -764,33 +1148,37 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
{
const int gnt = wv * (C / 16) + nt; // expand output column tile
f8v 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];
}
{ d = acc_vec<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt); }
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
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);
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_EXPAND)
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
mine[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, wm_cvt2(d[2 * j], d[2 * j + 1])) +
__builtin_bit_cast(hv2, wm_op_acc_pair(sv, j)));
op_store(&dst[t][16 * nt], wm_op_from_mine(mine));
#else
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));
op_store(&dst[t][16 * nt], operand_from_dt(xo));
#endif
}
}
};
expand(act);
block_sync();
PROF(8);
PwinParams q = p;
q.w = p.w + L::CB;
#ifdef PWIN_DEC_RECOMPUTE
pwin_core<L>(q, img, bias, act, hb, kl, vt, PwinX0<decltype(expand), half_t (*)[LDSR(C)]>{expand, x0b}, x0b);
#else
pwin_core<L>(q, img, bias, act, hb, kl, vt);
#endif
PROF(9);
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;
@@ -801,22 +1189,10 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
f8v 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];
}
{ d = acc_vec<L>(img, q.w, L::HB + 32 * nt, L::V_HB + nt); }
#pragma unroll
for (int kt = 0; kt < L::KT; ++kt)
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
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);
}
}
@@ -826,11 +1202,12 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
for (int kt = (int)hf; kt < L::KT; kt += 2)
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
}
PROF(10);
}
#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__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__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; \
kernels/k/pwin_phase.h
+402
View File
@@ -0,0 +1,402 @@
// Phase-split DltssPaddedWinLayer: the deep layers have few windows (dec5: 9), so one block
// per window leaves most of the GPU idle. Here a layer runs as four kernels (the main kernel plus NAME_post1..3,
// launched back to back by the patched ZLUDA hook, patches/zluda/0008-native-post-phase-kernels.patch), each a fixed-size grid
// looping over independent units:
// P0 (main) per (window, head): L2 norm + QKV + attention -> O (f16, global)
// P1 (post1) per (window, out tile, tok tile): Wo chain over heads + residual -> m, a (f16, global)
// P2 (post2) per (window, hidden tile, tok): fc1 chain + gelu -> hid (f16, global)
// P3 (post3) per (window, out tile, tok tile): fc2 chain over hidden from a -> y (layer output)
// Every value goes through the same f16 / f32 operations in the same order as pwin_wide.h (and pwin_layer.h), and
// everything passed between phases is an f16 value already, so the output is bit-identical. Units are distributed
// with compile-time grid sizes (no gridDim), the window count comes from the parameters (W, H, shift).
// Scratch buffers are module globals sized for MAXW windows; a launch with more windows does nothing (guard
// against out-of-bounds writes; the layer output is then not written).
#pragma once
#include "pwin_wide.h"
// layer kinds: core weights at p.w (+ CB for decoders); x0 rows from the layer input (mirror padded) or, for
// decoders, from the expand phase's scratch rows gX[window][token][C]
template <class L> struct PhaseCB
{
template <class T> static constexpr int cb(decltype(T::CB)*) { return T::CB; }
template <class T> static constexpr int cb(...) { return 0; }
static constexpr int value = cb<L>(nullptr);
};
template <class L> __device__ __forceinline__ const half_t* phase_x0_row(const PwinParams& p, const half_t* gX, int wi, int bx,
int by, int tok)
{
if (gX)
return gX + ((size_t)wi * 64 + tok) * L::C;
const int X = mirror(8 * bx - p.sx + (tok & 7), p.W), Y = mirror(8 * by - p.sy + (tok >> 3), p.H);
return (const half_t*)p.in + (size_t)(Y * p.W + X) * L::C;
}
template <class L> struct PhaseWin
{
int GX, GY, n;
__device__ PhaseWin(const PwinParams& p) : GX((p.W + p.sx + 7) / 8), GY((p.H + p.sy + 7) / 8), n(GX * GY) {}
};
// P0: unit u = (window w, head h); 4 waves = the 4 token tiles of the window
template <class L, int MAXW, int GRID>
__device__ void phase_attn(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias, half_t* gO,
const half_t* gX = nullptr, int ub = -1)
{
constexpr int C = L::C, KT = L::KT, H = L::H;
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) half_t kl[1][64][LDSR(32)];
__shared__ __attribute__((aligned(16))) half_t vt[1][32][LDSR(64)];
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int t = __builtin_amdgcn_readfirstlane(threadIdx.z), tok = 16 * t + m;
const PhaseWin<L> win(p);
if (win.n > MAXW)
return;
const uint8_t* w = p.w + PhaseCB<L>::value;
half_t* hrow = hb[tok];
// ub >= 0 (single-unit launches): exactly unit ub; otherwise grid-strided over all units
for (int u = ub >= 0 ? ub : (int)blockIdx.x; u < win.n * H; u += ub >= 0 ? 0x3fffffff : GRID)
{
const int wi = u / H, h = u % H, bx = wi % win.GX, by = wi / win.GX;
const half_t* xin = phase_x0_row<L>(p, gX, wi, bx, by, tok);
// L2 norm (wide_core): h1 = x0 * (rsqrt(sum x0^2) * gamma1)
float ss = 0.0f;
for (int kt = 0; kt < KT; ++kt)
{
const u8v xv = gload_row16(xin + 16 * kt);
#pragma unroll
for (int c = 0; c < 16; ++c)
{
const half_t v = op_get(xv, c);
ss += (float)(half_t)(v * v);
}
}
const half_t r16 = half_rsqrt((half_t)ss);
for (int kt = (int)hf; kt < KT; kt += 2)
{
const u8v xv = gload_row16(xin + 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);
}
__builtin_amdgcn_wave_barrier();
op_t qop[2], oop[2];
wide_qkv<L, 0, 6>(img, hrow, h, 0, tok, qop, kl, vt);
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 + t) * 4 + kt) * 32 + l];
f8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
{
d[2 * i] = (float)lo16(bi[i]);
d[2 * i + 1] = (float)hi16(bi[i]);
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(op_lds(&kl[0][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_reciprocal((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)
{
f8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(op_lds(&vt[0][16 * nt + m][16 * kt]), pop[kt], d);
oop[nt] = operand_from_f8(d);
}
}
half_t* orow = gO + ((size_t)(wi * H + h) * 64 + tok) * 32;
op_gstore(orow, oop[0], true);
op_gstore(orow + 16, oop[1], true);
block_sync(); // hb / kl / vt are rewritten by the next unit
}
}
// global wave index helpers for the per-wave-unit phases (blocks of 4 waves, GRID blocks)
// (ub >= 0: only the 4 wave units of chunk ub, see phase_attn)
#define PHASE_WAVE_LOOP(UNITS, GRID) \
for (int uw = (ub >= 0 ? ub : (int)blockIdx.x) * 4 + (int)__builtin_amdgcn_readfirstlane(threadIdx.z); uw < (UNITS); \
uw += ub >= 0 ? 0x3fffffff : (GRID) * 4)
// P1: unit = (window, output tile nt, token tile): Wo chain over heads, residual; writes the MLP input m and the
// fc2 accumulator init a = f16(x1 + b2) (rows of C halves per token)
template <class L, int MAXW, int GRID>
__device__ void phase_wo(const PwinParams& p, const u8v* __restrict__ img, const half_t* gO, half_t* gM, half_t* gA,
const half_t* gX = nullptr, int ub = -1)
{
constexpr int C = L::C, KT = L::KT, H = L::H;
const uint32_t m = lane_id() & 15;
const PhaseWin<L> win(p);
if (win.n > MAXW)
return;
const uint8_t* w = p.w + PhaseCB<L>::value;
PHASE_WAVE_LOOP(win.n * KT * 4, GRID)
{
const int wi = uw / (KT * 4), r = uw % (KT * 4), nt = r >> 2, tt = r & 3, tok = 16 * tt + m;
const int bx = wi % win.GX, by = wi / win.GX;
f8v acc = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
for (int hh = 0; hh < H; ++hh)
{
const half_t* orow = gO + ((size_t)(wi * H + hh) * 64 + tok) * 32;
acc = mma16(op_imgt(img, L::T_WO + (2 * hh) * KT + nt), op_gload(orow), acc);
acc = mma16(op_imgt(img, L::T_WO + (2 * hh + 1) * KT + nt), op_gload(orow + 16), acc);
}
const op_t xv = op_gload(phase_x0_row<L>(p, gX, wi, bx, by, tok) + 16 * nt);
half_t mv[8], av[8], g2[8], b2[8];
residual_vector<L>(w, 0, nt, g2);
residual_vector<L>(w, 1, nt, b2);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const half_t x1 = (half_t)acc[i] + wm_op_acc_elem(xv, i);
mv[i] = x1 * g2[i];
av[i] = (half_t)(x1 + b2[i]);
}
const size_t row = ((size_t)wi * 64 + tok) * C + 16 * nt;
op_gstore(gM + row, operand_from_dt(mv), true);
op_gstore(gA + row, operand_from_dt(av), true);
}
}
// P2: unit = (window, hidden tile j = 2 c + hn, token tile): fc1 chain over kt + gelu -> hidden rows (4C halves)
template <class L, int MAXW, int GRID>
__device__ void phase_fc1(const PwinParams& p, const u8v* __restrict__ img, const half_t* gM, half_t* gH, int ub = -1)
{
constexpr int C = L::C, KT = L::KT, NJ = 2 * L::NMLP;
const uint32_t m = lane_id() & 15;
const PhaseWin<L> win(p);
if (win.n > MAXW)
return;
const uint8_t* w = p.w + PhaseCB<L>::value;
PHASE_WAVE_LOOP(win.n * NJ * 4, GRID)
{
const int wi = uw / (NJ * 4), r = uw % (NJ * 4), j = r >> 2, tt = r & 3, tok = 16 * tt + m;
const int c = j >> 1, hn = j & 1;
f8v d = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
const half_t* mrow = gM + ((size_t)wi * 64 + tok) * C;
for (int kt = 0; kt < KT; ++kt)
d = mma16(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_gload(mrow + 16 * kt), d);
op_gstore(gH + ((size_t)wi * 64 + tok) * (4 * C) + 16 * j, gelu_op(d), true);
}
}
// P3: unit = (window, output tile nt, token tile): fc2 chain over the hidden chunks from a; y -> layer output
template <class L, int MAXW, int GRID>
__device__ void phase_fc2(const PwinParams& p, const u8v* __restrict__ img, const half_t* gA, const half_t* gH, uint8_t* dst,
half_t* gY = nullptr, int ub = -1)
{
constexpr int C = L::C, KT = L::KT;
const uint32_t m = lane_id() & 15;
const PhaseWin<L> win(p);
if (win.n > MAXW)
return;
PHASE_WAVE_LOOP(win.n * KT * 4, GRID)
{
const int wi = uw / (KT * 4), r = uw % (KT * 4), nt = r >> 2, tt = r & 3, tok = 16 * tt + m;
const int bx = wi % win.GX, by = wi / win.GX;
const op_t av = op_gload(gA + ((size_t)wi * 64 + tok) * C + 16 * nt);
f8v acc;
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = (float)wm_op_acc_elem(av, i);
const half_t* hrow = gH + ((size_t)wi * 64 + tok) * (4 * C);
for (int cc = 0; cc < L::NMLP; ++cc)
{
acc = mma16(op_imgt(img, L::T_W2 + (cc * 2) * KT + nt), op_gload(hrow + 32 * cc), acc);
acc = mma16(op_imgt(img, L::T_W2 + (cc * 2 + 1) * KT + nt), op_gload(hrow + 32 * cc + 16), acc);
}
const int 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* yout = (half_t*)dst + (size_t)(Y0 * p.W + X0) * C;
const op_t y = operand_from_f8(acc);
op_gstore(yout + 16 * nt, y, inb);
if (gY)
op_gstore(gY + ((size_t)wi * 64 + tok) * C + 16 * nt, y, true);
}
}
// decoder expand (wide_decoder): unit = (window, sub-token q, tile nt) per wave; lanes = the 16 low-resolution
// tokens; x0 = patch expand of the previous decoder output + encoder skip -> gX rows (all 64 tokens of the window)
template <class L, int MAXW, int GRID>
__device__ void phase_expand(const PwinParams& p, const u8v* __restrict__ img, half_t* gX, int ub = -1)
{
constexpr int C = L::C, KT = L::KT;
const uint32_t m = lane_id() & 15;
const PhaseWin<L> win(p);
if (win.n > MAXW)
return;
PHASE_WAVE_LOOP(win.n * 4 * KT, GRID)
{
const int wi = uw / (4 * KT), r = uw % (4 * KT), q = r / KT, nt = r % KT;
const int bx = wi % win.GX, by = wi / win.GX;
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;
const int gnt = q * (C / 16) + nt;
f8v d = acc_vec<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt);
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_gstore(gX + ((size_t)wi * 64 + tt) * C + 16 * nt, operand_from_dt(xo), true);
}
}
// encoder patch merge (wide_encoder): unit = (window, merge tile vt) per wave; lanes = the 16 merged tokens
template <class L, int MAXW, int GRID>
__device__ void phase_merge(const PwinParams& p, const u8v* __restrict__ img, const half_t* gY, int ub = -1)
{
constexpr int C = L::C, KT = L::KT;
const uint32_t m = lane_id() & 15;
const PhaseWin<L> win(p);
if (win.n > MAXW)
return;
const uint8_t* w = p.w;
PHASE_WAVE_LOOP(win.n * L::PMT, GRID)
{
const int wi = uw / L::PMT, vt = uw % L::PMT;
const int bx = wi % win.GX, by = wi / win.GX;
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;
f8v d = acc_vec<L>(img, w, L::PMB + 32 * vt, L::V_PM + vt);
for (int ks = 0; ks < C / 4; ++ks)
{
const int sub = ks / 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_gload(gY + ((size_t)wi * 64 + src) * C + 16 * (ks % 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 layer (dec5 shape: core only, output to +24) as four phase kernels
#define PHASE_PLAIN(NAME, H, C, MAXW, G0, G1, G2, G3) \
using NAME##_L = PwinLayout<H, C, 16>; \
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
__device__ half_t g_O[(MAXW) * (H) * 64 * 32]; \
__device__ half_t g_M[(MAXW) * 64 * (C)]; \
__device__ half_t g_A[(MAXW) * 64 * (C)]; \
__device__ half_t g_H[(MAXW) * 64 * 4 * (C)]; \
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_grid_x = (G0); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post1_grid_x = (G1); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post1_block_z = 4; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post2_grid_x = (G2); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post2_block_z = 4; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post3_grid_x = (G3); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post3_block_z = 4; \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
{ \
pwin_prep_nomerge<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) \
{ \
phase_attn<NAME##_L, MAXW, G0>(p, g_img, g_bias, g_O); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) \
{ \
phase_wo<NAME##_L, MAXW, G1>(p, g_img, g_O, g_M, g_A); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) \
{ \
phase_fc1<NAME##_L, MAXW, G2>(p, g_img, g_M, g_H); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) \
{ \
phase_fc2<NAME##_L, MAXW, G3>(p, g_img, g_A, g_H, p.out24); \
}
#define PHASE_GLOBALS(NAME, H, C, MAXW) \
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
__device__ half_t g_O[(MAXW) * (H) * 64 * 32]; \
__device__ half_t g_M[(MAXW) * 64 * (C)]; \
__device__ half_t g_A[(MAXW) * 64 * (C)]; \
__device__ half_t g_H[(MAXW) * 64 * 4 * (C)]; \
__device__ half_t g_X[(MAXW) * 64 * (C)]; \
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;
#define PHASE_POST(K, G) \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post##K##_grid_x = (G); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_post##K##_block_z = 4;
// decoder (wide shape, patch expand + skip, no output head): expand, attention, Wo, fc1, fc2
#define PHASE_DECODER(NAME, H, C, CL, MAXW, GE, G0, G1, G2, G3) \
using NAME##_L = PwinDecLayout<H, C, CL>; \
PHASE_GLOBALS(NAME, H, C, MAXW) \
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = (GE); \
PHASE_POST(1, G0) PHASE_POST(2, G1) PHASE_POST(3, G2) PHASE_POST(4, G3) \
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) NAME(PwinParams p) { phase_expand<NAME##_L, MAXW, GE>(p, g_img, g_X); } \
extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) \
{ \
phase_attn<NAME##_L, MAXW, G0>(p, g_img, g_bias, g_O, g_X); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) \
{ \
phase_wo<NAME##_L, MAXW, G1>(p, g_img, g_O, g_M, g_A, g_X); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) { phase_fc1<NAME##_L, MAXW, G2>(p, g_img, g_M, g_H); } \
extern "C" __global__ void __launch_bounds__(128) NAME##_post4(PwinParams p) \
{ \
phase_fc2<NAME##_L, MAXW, G3>(p, g_img, g_A, g_H, p.out24); \
}
// encoder (wide shape): attention, Wo, fc1, fc2 (+ full-resolution output), patch merge
#define PHASE_ENCODER(NAME, H, C, COUT, MAXW, G0, G1, G2, G3, GM) \
using NAME##_L = PwinLayout<H, C, COUT>; \
PHASE_GLOBALS(NAME, H, C, MAXW) \
__device__ half_t g_Y[(MAXW) * 64 * (C)]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = (G0); \
PHASE_POST(1, G1) PHASE_POST(2, G2) PHASE_POST(3, G3) PHASE_POST(4, GM) \
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) { phase_attn<NAME##_L, MAXW, G0>(p, g_img, g_bias, g_O); } \
extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) { phase_wo<NAME##_L, MAXW, G1>(p, g_img, g_O, g_M, g_A); } \
extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) { phase_fc1<NAME##_L, MAXW, G2>(p, g_img, g_M, g_H); } \
extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) \
{ \
phase_fc2<NAME##_L, MAXW, G3>(p, g_img, g_A, g_H, p.out32, g_Y); \
} \
extern "C" __global__ void __launch_bounds__(128) NAME##_post4(PwinParams p) { phase_merge<NAME##_L, MAXW, GM>(p, g_img, g_Y); }
kernels/k/pwin_pos.h
+171 -85
View File
@@ -1,17 +1,21 @@
// 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
#if defined(PWIN_F16ACC_NATIVE) && !defined(PWIN_PK_ELEM)
#error "PWIN_F16ACC_NATIVE needs PWIN_PK_ELEM (the element-wise paths use acc_pair / acc_set_pair)"
#endif
#include "pwin_layer.h"
template <class L, int MT>
__device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
half_t (*act)[L::C], half_t (*hb)[L::C])
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(L::POS, "position-only attention layers");
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
const uint8_t* w = p.w;
PROF_INIT;
// ---- L2 norm of every tile of this wave -> hb
#pragma unroll
@@ -25,13 +29,23 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
{
const u8v xv = lds_row16(arow + 16 * kt);
#pragma unroll
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_SS)
for (int j = 0; j < 8; ++j)
{
const uint32_t xj = xv[j]; // (__builtin_bit_cast of a vector subscript reads element 0: clang bug)
const hv2 v = __builtin_bit_cast(hv2, xj), q = v * v;
ss += (float)q[0];
ss += (float)q[1];
}
#else
for (int c = 0; c < 16; ++c)
{
const half_t v = op_get(xv, c);
ss += (float)(half_t)(v * v);
}
#endif
}
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
const half_t r16 = half_rsqrt((half_t)ss);
for (int kt = (int)hf; kt < KT; kt += 2)
{
const u8v xv = lds_row16(arow + 16 * kt);
@@ -39,8 +53,15 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
u8v hv;
#pragma unroll
for (int j = 0; j < 8; ++j)
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_APPLY)
{
const uint32_t xj = xv[j], gj = gv[j]; // (see above: no __builtin_bit_cast of a vector subscript)
hv[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, xj) * ((hv2)r16 * __builtin_bit_cast(hv2, gj)));
}
#else
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)));
#endif
store_row16(hrow + 16 * kt, hv);
}
}
@@ -48,38 +69,58 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
STAMP(2);
// ---- attention: O = P V per head (P from the table), output projection accumulated head by head
f8v acc[MT][KT];
pacc acc[MT][KT];
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, vv);
const pacc bo = acc_vec_p<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
#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];
acc[mi][nt] = bo;
}
#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
#ifdef POS_VSHARE1
// V^T operands shared through ONE 4 KB slot reused per head (a barrier before it is rewritten); each wave
// computes its own key tiles, the PV loop reads the operands from LDS when it needs them (no 8 operands live
// in registers). Same values as computing all of V in every wave.
__shared__ __attribute__((aligned(16))) u8v vsh1[2][4][16];
#endif
for (int h = 0; h < H; ++h)
{
op_t vtop[2][4];
#ifdef POS_VSHARE
#ifdef POS_VSHARE1
if (h > 0)
block_sync(); // every wave has finished reading the previous head's V^T
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int kt = wv * MT + mi;
f8v d = splat8(0.0f);
pacc d = acc_zero();
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_f8(d));
d = mma16p(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(&vsh1[nt][0][0], kt * 16 + m, operand_from_acc_p(d));
}
block_sync();
(void)vtop;
#elif defined(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;
pacc d = acc_zero();
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16p(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_acc_p(d));
}
block_sync();
#pragma unroll
@@ -93,11 +134,11 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
f8v d = splat8(0.0f);
pacc d = acc_zero();
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
vtop[nt][kt] = operand_from_f8(d);
d = mma16p(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_acc_p(d);
}
#endif
op_t oop[MT][2];
@@ -108,15 +149,21 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = op_img((const u8v*)bias, ((h * 4 + qt) * 4 + kt) * 16 + m);
pop[kt] = op_imgt((const u8v*)bias, (h * 4 + qt) * 4 + kt);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
f8v d = splat8(0.0f);
pacc d = acc_zero();
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(vtop[nt][kt], pop[kt], d);
oop[mi][nt] = operand_from_f8(d);
{
#ifdef POS_VSHARE1
d = mma16p(op_img(&vsh1[nt][0][0], kt * 16 + m), pop[kt], d);
#else
d = mma16p(vtop[nt][kt], pop[kt], d);
#endif
}
oop[mi][nt] = operand_from_acc_p(d);
}
}
#pragma unroll
@@ -124,10 +171,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const op_t wo = op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m);
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]);
acc[mi][nt] = mma16p(wo, oop[mi][ks], acc[mi][nt]);
}
}
@@ -137,13 +184,27 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
for (int nt = 0; nt < KT; ++nt)
{
half_t g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
residual_vector<L>(w, 0, nt, g2);
residual_vector<L>(w, 1, 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]);
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_RESID)
// the same f16 operations on element pairs (packed f16 math; pairs stay packed into the operand)
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
const hv2 x1 = __builtin_bit_cast(hv2, acc_pair(acc[mi][nt], j)) +
__builtin_bit_cast(hv2, wm_op_acc_pair(xv, j));
mine[j] = __builtin_bit_cast(uint32_t, x1 * (hv2){g2[2 * j], g2[2 * j + 1]});
const hv2 y = x1 + (hv2){b2[2 * j], b2[2 * j + 1]};
acc_set_pair(acc[mi][nt], j, y);
}
op_store(&hb[tok][16 * nt], wm_op_from_mine(mine));
#else
half_t mv[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
@@ -153,6 +214,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
}
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
#endif
}
}
__builtin_amdgcn_wave_barrier();
@@ -165,28 +227,23 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
half_t bb[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
f8v d[MT];
const pacc bb = acc_vec_p<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
pacc 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];
d[mi] = bb;
#pragma unroll
for (int kt = 0; kt < KT; ++kt)
{
const op_t w1 = op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m);
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]);
d[mi] = mma16p(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);
gop[mi][hn] = gelu_op_p(d[mi]);
}
}
#pragma unroll
@@ -194,10 +251,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const op_t w2 = op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m);
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]);
acc[mi][nt] = mma16p(w2, gop[mi][ks], acc[mi][nt]);
}
}
STAMP(5);
@@ -207,9 +264,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#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]));
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_acc_p(acc[mi][nt]));
}
block_sync();
PROF(6);
}
// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
@@ -217,14 +275,15 @@ template <class L, int MT>
__device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C, NW = 4 / MT;
__shared__ __attribute__((aligned(16))) half_t act[64][C];
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
__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 = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
STAMP(0);
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
PROF_INIT;
STAMP(20);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
@@ -232,28 +291,34 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
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];
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
#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);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
pacc d = acc_vec_p<L>(img, p0.w, 32 * nt, L::V_EMB + nt);
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
d = mma16p(op_imgt(img, L::T_EMB + ks * (C / 16) + nt), xe[ks], d);
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_RELU)
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
const hv2 x = __builtin_bit_cast(hv2, acc_pair(d, j));
mine[j] = __builtin_bit_cast(uint32_t, (hv2){hmax(x[0], (half_t)0.0f), hmax(x[1], (half_t)0.0f)});
}
op_store(&act[tok][16 * nt], wm_op_from_mine(mine));
#else
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));
#endif
}
}
else
@@ -264,9 +329,9 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
}
}
__builtin_amdgcn_wave_barrier();
STAMP(1);
STAMP(21);
pos_core<L, MT>(p, img, bias, act, hb);
STAMP(6);
STAMP(26);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
@@ -280,29 +345,24 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
}
}
STAMP(7);
STAMP(27);
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);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
pacc d = acc_vec_p<L>(img, p.w, L::PMB + 32 * vt, L::V_PM + vt);
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_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
d = mma16p(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);
op_gstore(mout + L::NPW * g + 16 * u, operand_from_acc_p(d), minb, valid);
}
STAMP(8);
STAMP(28);
}
// dec0: patch expand + skip, core, output head
@@ -310,10 +370,11 @@ template <class L, int MT>
__device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C, NW = 4 / MT;
__shared__ __attribute__((aligned(16))) half_t act[64][C];
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
__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 = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
PROF_INIT;
{
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);
@@ -322,8 +383,24 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
#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)
#ifndef POS_NO_SKIP_PREFETCH
// the skip rows of all this wave's sub-tokens, loaded up front with xls (their latency overlaps)
op_t svs[MT][C / 16];
#pragma unroll
for (int qi = 0; qi < MT; ++qi)
{
const int q = wv + NW * qi, 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)
svs[qi][nt] = op_gload(sk + 16 * nt);
}
#endif
#pragma unroll
for (int qi = 0; qi < MT; ++qi)
{
const int q = wv + NW * qi;
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;
@@ -331,28 +408,39 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
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);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
pacc d = acc_vec_p<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt);
#pragma unroll
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), xls[ks], d);
d = mma16p(op_imgt(img, L::T_EXP + ks * L::EXPN + gnt), xls[ks], d);
#ifndef POS_NO_SKIP_PREFETCH
const op_t sv = svs[qi][nt];
(void)sk;
#else
const op_t sv = op_gload(sk + 16 * nt);
#endif
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_EXPAND)
uint32_t mine[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
mine[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, acc_pair(d, j)) +
__builtin_bit_cast(hv2, wm_op_acc_pair(sv, j)));
op_store(&act[t][16 * nt], wm_op_from_mine(mine));
#else
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));
#endif
}
}
}
block_sync();
PROF(8);
PwinParams q = p;
q.w = p.w + L::CB;
pos_core<L, MT>(q, img, bias, act, hb);
PROF(9);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
@@ -362,29 +450,26 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
#pragma unroll
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
half_t vv[8];
dvec8(q.w, L::HB + 32 * nt, vv);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
pacc d = acc_vec_p<L>(img, q.w, L::HB + 32 * nt, L::V_HB + nt);
#pragma unroll
for (int kt = 0; kt < L::KT; ++kt)
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
d = mma16p(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_acc_p(d), inb, L::NOUT - 16 * nt);
}
}
PROF(10);
}
#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__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__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_prep_blocks = (pwin_gelu_prep_items(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_gelu_prep(); \
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
@@ -394,13 +479,14 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
#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__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__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_prep_blocks = (pwin_gelu_prep_items(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_gelu_prep(); \
pwin_dec_prep<NAME##_L>(p, g_img, g_bias); \
} \
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
kernels/k/pwin_pos_recompute.h
+379
View File
@@ -0,0 +1,379 @@
// Single-buffer position-only-attention encoder experiment (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_recompute_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(L::POS, "position-only attention layers");
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
const uint8_t* w = p.w;
PROF_INIT;
// ---- L2 norm of every tile of this wave -> hb
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const half_t* arow = act[16 * (wv * MT + mi) + m];
half_t* hrow = hb[16 * (wv * MT + mi) + m];
float ss = 0.0f;
#pragma unroll
for (int kt = 0; kt < KT; ++kt)
{
const u8v xv = lds_row16(arow + 16 * kt);
#pragma unroll
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_SS)
for (int j = 0; j < 8; ++j)
{
const uint32_t xj = xv[j];
const hv2 v = __builtin_bit_cast(hv2, xj), q = v * v;
ss += (float)q[0];
ss += (float)q[1];
}
#else
for (int c = 0; c < 16; ++c)
{
const half_t v = op_get(xv, c);
ss += (float)(half_t)(v * v);
}
#endif
}
const half_t r16 = half_rsqrt((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)
#if defined(PWIN_PK_ELEM) && !defined(PWIN_PK_NO_APPLY)
{
const uint32_t xj = xv[j], gj = gv[j];
hv[j] = __builtin_bit_cast(uint32_t, __builtin_bit_cast(hv2, xj) * ((hv2)r16 * __builtin_bit_cast(hv2, gj)));
}
#else
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)));
#endif
store_row16(hrow + 16 * kt, hv);
}
}
block_sync(); // V of all 64 tokens is computed from hb
STAMP(2);
// ---- attention: O = P V per head (P from the table), output projection accumulated head by head
f8v acc[MT][KT];
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
const f8v bo = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
acc[mi][nt] = bo;
}
#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
#ifdef POS_VSHARE1
// V^T operands shared through ONE 4 KB slot reused per head (a barrier before it is rewritten); each wave
// computes its own key tiles, the PV loop reads the operands from LDS when it needs them (no 8 operands live
// in registers). Same values as computing all of V in every wave.
__shared__ __attribute__((aligned(16))) u8v vsh1[2][4][16];
#endif
for (int h = 0; h < H; ++h)
{
op_t vtop[2][4];
#ifdef POS_VSHARE1
if (h > 0)
block_sync(); // every wave has finished reading the previous head's V^T
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int kt = wv * MT + mi;
f8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(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(&vsh1[nt][0][0], kt * 16 + m, operand_from_f8(d));
}
block_sync();
(void)vtop;
#elif defined(POS_VSHARE)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int kt = wv * MT + mi;
f8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(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)
{
f8v 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)
{
f8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
#ifdef POS_VSHARE1
d = mma16(op_img(&vsh1[nt][0][0], kt * 16 + m), pop[kt], d);
#else
d = mma16(vtop[nt][kt], pop[kt], d);
#endif
}
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);
// All attention readers finish before the shared input rows are reused.
block_sync();
// Recover x0 exactly from the unchanged input (repeat the original embedding,
// or reload a no-embedding layer), then write the MLP input into the same LDS rows.
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
half_t g2[8], b2[8];
residual_vector<L>(w, 0, nt, g2);
residual_vector<L>(w, 1, nt, b2);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int tok = 16 * (wv * MT + mi) + m;
const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W);
const int Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
half_t x0[8];
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
f8v d = acc_vec<L>(img, w - L::CB, 32 * nt, L::V_EMB + nt);
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
d = mma16(op_imgt(img, L::T_EMB + ks * KT + nt), op_gload(xin + 16 * ks), d);
#pragma unroll
for (int i = 0; i < 8; ++i)
x0[i] = hmax((half_t)d[i], (half_t)0.0f);
}
else
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
const op_t xv = op_gload(xin + 16 * nt);
#pragma unroll
for (int i = 0; i < 8; ++i)
x0[i] = wm_op_acc_elem(xv, i);
}
half_t mv[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const half_t x1 = (half_t)acc[mi][nt][i] + x0[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)
{
const f8v bb = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
f8v d[MT];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
d[mi] = bb;
#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)
{
gop[mi][hn] = gelu_op(d[mi]);
}
}
#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();
PROF(6);
}
// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
template <class L, int MT>
__device__ void pos_recompute_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)];
half_t (*hb)[LDSR(C)] = act; // one allocation: x0 -> norm -> MLP input -> final output
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 = __builtin_amdgcn_readfirstlane(threadIdx.z), bx = blockIdx.x, by = blockIdx.y;
PROF_INIT;
STAMP(20);
#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)
{
f8v d = acc_vec<L>(img, p0.w, 32 * nt, L::V_EMB + nt);
#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(21);
pos_recompute_core<L, MT>(p, img, bias, act, hb);
STAMP(26);
#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(27);
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)
{
f8v d = acc_vec<L>(img, p.w, L::PMB + 32 * vt, L::V_PM + vt);
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(28);
}
#define POS_RECOMPUTE_ENCODER(NAME, H, C, COUT, CIN, MT) \
using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>; \
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__device__ u4v g_bias[NAME##_L::NBIAS]; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (pwin_gelu_prep_items(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_gelu_prep(); \
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_recompute_encoder<NAME##_L, MT>(p, g_img, g_bias); \
}
kernels/k/pwin_wide.h
+124 -102
View File
@@ -15,27 +15,71 @@ template <class L, int NG> struct WideDims
// NA: heads per attention round (their K / V images share R)
static constexpr int NA = NG < 4 ? NG : 4;
static constexpr int C = L::C, KT = L::KT, H = L::H, NTL = (KT + NG - 1) / NG, ROUNDS = (H + NA - 1) / NA;
static constexpr int X0_BYTES = 64 * C * 2, KV_BYTES = NA * 2 * 64 * 32 * 2;
// K rows (64 x 32), V^T rows (32 x 64), O rows (64 x 32), hidden rows (64 x 32): all padded (LDSR)
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 * 32 * 2;
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");
};
// QKV chains u = U0 .. U0 + NU - 1 of head h for this wave's token tile (u = 2 j + nt; j: Q, K, V): Q stays in
// registers (qop), K and V^T go to LDS slot s. Each chain runs over kt in order (exact).
template <class L, int U0, int NU>
__device__ __forceinline__ void wide_qkv(const u8v* __restrict__ img, const half_t* hrow, int h, int s, int tok,
op_t (&qop)[2], half_t (*kl)[64][LDSR(32)], half_t (*vt)[32][LDSR(64)])
{
constexpr int KT = L::KT;
const uint32_t m = lane_id() & 15;
f8v dq[NU];
#pragma unroll
for (int uu = 0; uu < NU; ++uu)
dq[uu] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int uu = 0; uu < NU; ++uu)
{
const int u = U0 + uu;
dq[uu] = mma16(op_imgt(img, L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)), row, dq[uu]);
}
}
#pragma unroll
for (int uu = 0; uu < NU; ++uu)
{
const int u = U0 + uu, j = u >> 1, nt = u & 1;
const f8v d = dq[uu];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
op_store(&kl[s][tok][16 * nt], operand_from_f8(d));
else
{
#pragma unroll
for (int i = 0; i < 8; ++i)
vt[s][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
}
}
}
// Core for one window. On entry R holds x0 (64 x C); on return hb holds the block output y.
template <class L, int NG>
__device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
half_t (*hb)[L::C], uint8_t* R)
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 = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const uint8_t* w = p.w;
half_t(*act)[C] = (half_t(*)[C])R;
half_t(*act)[LDSR(C)] = (half_t(*)[LDSR(C)])R;
const half_t* arow = act[tok];
half_t* hrow = hb[tok];
PROF_INIT;
PROF(1); // entry (x0 load + barrier)
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb (all groups of a token tile compute the sum)
float ss = 0.0f;
@@ -49,7 +93,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
ss += (float)(half_t)(v * v);
}
}
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
const half_t r16 = half_rsqrt((half_t)ss);
for (int kt = 2 * g + (int)hf; kt < KT; kt += 2 * NG)
{
const u8v xv = lds_row16(arow + 16 * kt);
@@ -76,6 +120,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
}
block_sync(); // hb complete, R free
PROF(2); // norm
// ---- attention: NG heads per round; the output projection accumulates head by head
f8v acc[NTL];
@@ -84,65 +129,57 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
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];
}
acc[j] = acc_vec<L>(img, w, L::BO + 32 * nt, L::V_BO + nt);
}
half_t(*kl)[64][32] = (half_t(*)[64][32])R; // [group][token][dim]
half_t(*kl)[64][LDSR(32)] = (half_t(*)[64][LDSR(32)])R; // [group][token][dim]
constexpr int NA = D::NA;
half_t(*vt)[32][64] = (half_t(*)[32][64])(R + NA * 64 * 32 * 2); // [group][dim][token]
half_t(*ob)[64][32] = (half_t(*)[64][32])R; // [group][token][dim] (after K / V)
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)
// NG >= 2 NA: the groups beyond the NA heads of a round would idle during QKV and attention; instead group
// g + NA computes half of the six QKV chains of group g's head (K tile 1 and V), group g the other half
// (Q, kept in registers for its attention, and K tile 0)
constexpr bool SPLIT = NG >= 2 * NA;
const int hs = SPLIT ? g % NA : g, part = SPLIT ? g / NA : 0;
for (int r = 0; r < D::ROUNDS; ++r)
{
const int h = g < NA ? r * NA + g : H;
const int h = (SPLIT ? g < 2 * NA : g < NA) ? r * NA + hs : H;
const bool att = h < H && part == 0;
op_t qop[2], oop[2];
#ifdef PWIN_BIAS_PREFETCH
u4v bpre[4];
if (att)
{
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
bpre[kt] = bias[((h * 4 + t) * 4 + kt) * 32 + l];
}
#endif
if (h < H)
{
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
f8v dq[6];
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
if constexpr (SPLIT)
{
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
if (part == 0)
wide_qkv<L, 0, 3>(img, hrow, h, hs, tok, qop, kl, vt);
else
wide_qkv<L, 3, 3>(img, hrow, h, hs, tok, qop, kl, vt);
}
#pragma unroll
for (int j = 0; j < 3; ++j)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
const f8v d = dq[2 * j + nt];
if (j == 0)
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
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];
}
}
else
wide_qkv<L, 0, 6>(img, hrow, h, hs, tok, qop, kl, vt);
}
block_sync();
if (h < H)
PROF(3); // QKV
if (att)
{
half_t e[4][8];
float rs = 0.0f;
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
#ifdef PWIN_BIAS_PREFETCH
const u4v bi = bpre[kt];
#else
const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
#endif
f8v d;
#pragma unroll
for (int i = 0; i < 4; ++i)
@@ -152,7 +189,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(op_lds(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
d = mma16(op_lds(&kl[hs][16 * kt + m][16 * ds]), qop[ds], d);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
@@ -162,7 +199,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
}
rs = row_sum_total(rs, e);
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
const half_t rc = half_reciprocal((half_t)rs);
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
@@ -179,15 +216,17 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
f8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(op_lds(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
d = mma16(op_lds(&vt[hs][16 * nt + m][16 * kt]), pop[kt], d);
oop[nt] = operand_from_f8(d);
}
}
PROF(4); // attention
block_sync(); // K / V of the round consumed
if (h < H)
PROF(5); // barrier after attention
if (att)
{
op_store(&ob[g][tok][0], oop[0]);
op_store(&ob[g][tok][16], oop[1]);
op_store(&ob[hs][tok][0], oop[0]);
op_store(&ob[hs][tok][16], oop[1]);
}
block_sync();
for (int jh = 0; jh < NA; ++jh)
@@ -202,13 +241,14 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int nt = g + NG * j;
if (nt < KT)
{
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh) * KT + nt) * 16 + m), o0, acc[j]);
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m), o1, acc[j]);
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
PROF(6); // O exchange + Wo
}
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
@@ -219,8 +259,8 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (nt < KT)
{
half_t mv[8], g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
residual_vector<L>(w, 0, nt, g2);
residual_vector<L>(w, 1, nt, b2);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
@@ -232,9 +272,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
}
block_sync();
PROF(7); // residual
// ---- MLP: rounds of NG chunks of 32 hidden (double-buffered in R)
half_t(*hid)[NG][64][32] = (half_t(*)[NG][64][32])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;
@@ -243,21 +284,14 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
if (c >= L::NMLP)
break;
f8v 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];
}
f8v d = acc_vec<L>(img, w, L::B1 + 64 * c + 32 * hn, L::V_B1 + 2 * c + hn);
for (int kt = 0; kt < KT; ++kt)
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
half_t gg[8];
gelu8(d, gg);
op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
d = mma16(op_imgt(img, L::T_W1 + (c * KT + kt) * 2 + hn), op_lds(hrow + 16 * kt), d);
op_store(&hid[rr % D::NB][g][tok][16 * hn], gelu_op(d));
}
PROF(8); // MLP W1 + gelu
block_sync();
PROF(9); // MLP barrier
for (int jc = 0; jc < NG; ++jc)
{
const int cc = rr * NG + jc;
@@ -270,13 +304,14 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int nt = g + NG * j;
if (nt < KT)
{
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2) * KT + nt) * 16 + m), h0, acc[j]);
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m), h1, acc[j]);
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();
PROF(10); // MLP W2
}
// y -> hb (the last reads of hb were the W1 operands before the final barrier)
#pragma unroll
@@ -289,15 +324,16 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
}
block_sync();
PROF(11); // y store + barrier
}
// x0 rows of a window (mirror padded) -> R
template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uint8_t* R)
{
constexpr int C = L::C;
half_t(*act)[C] = (half_t(*)[C])R;
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 = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W), Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
@@ -305,10 +341,10 @@ template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uin
}
// y rows of the window's own tokens -> dst (C per token)
template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[L::C], uint8_t* dst)
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 = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), t = wv & 3, g = wv >> 2, tok = 16 * t + m;
const int Y0 = 8 * (int)blockIdx.y - p.sy + (tok >> 3), X0 = 8 * (int)blockIdx.x - p.sx + (tok & 7);
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
{
@@ -322,7 +358,7 @@ template <class L, int NG>
__device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C, COUT = L::COUT;
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
__shared__ __attribute__((aligned(16))) 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();
@@ -331,7 +367,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
// patch merge: 16 merged tokens, K = 4 sub-tokens x C, the allocated n-tiles over the waves
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const int wv = threadIdx.z;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z);
const int my = m >> 2, mx = m & 3;
const int Ym = (8 * (int)blockIdx.y - p.sy) / 2 + my, Xm = (8 * (int)blockIdx.x - p.sx) / 2 + mx;
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
@@ -339,19 +375,12 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
const uint8_t* w = p.w;
for (int vt = wv; vt < L::PMT; vt += 4 * NG)
{
f8v d;
{
half_t vv[8];
dvec8(w, L::PMB + 32 * vt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
}
f8v d = acc_vec<L>(img, w, L::PMB + 32 * vt, L::V_PM + vt);
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_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&hb[src][16 * (ks % L::KT)]), d);
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);
@@ -362,7 +391,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
template <class L, int NG>
__device__ void wide_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
__shared__ __attribute__((aligned(16))) half_t hb[64][L::C];
__shared__ __attribute__((aligned(16))) 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();
@@ -374,11 +403,11 @@ template <class L, int NG>
__device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
{
constexpr int C = L::C;
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
__shared__ __attribute__((aligned(16))) half_t hb[64][LDSR(C)];
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
half_t(*act)[C] = (half_t(*)[C])R;
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 = threadIdx.z, q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
const int wv = __builtin_amdgcn_readfirstlane(threadIdx.z), q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
// ---- patch expand: waves of sub-token q = (dy, dx) produce its tiles nt = g + NG j for the 16 low-res tokens
{
@@ -392,16 +421,9 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
for (int nt = g; nt < C / 16; nt += NG)
{
const int gnt = q * (C / 16) + nt;
f8v d;
{
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
}
f8v d = acc_vec<L>(img, p.w, L::EXPB + 32 * gnt, L::V_EXP + gnt);
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
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
@@ -420,7 +442,7 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
#define WIDE_KERNEL(NAME, PREP, BODY, NG, ...) \
using NAME##_L = __VA_ARGS__; \
__device__ u8v g_img[NAME##_L::T_END * 16]; \
__device__ u8v g_img[NAME##_L::IMG_SLOTS]; \
__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); \
@@ -435,5 +457,5 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
}
#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_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep_nomerge, 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
+10 -1
View File
@@ -1,7 +1,16 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
// RX 7700 XT: VGPR cap 160 removes the scratch use; byte-identical to the default build, gfx1101 only
// d4r-build-flags-gfx1101: -DSWIN_MAX_VGPR=160
#define SWIN_NO_Q8
#define SWIN_F32ACC
#ifndef SWIN_MAX_VGPR
#define SWIN_MAX_VGPR 128
#endif
#ifndef SWIN_MLP_R
#define SWIN_MLP_R(NW) 2
#endif
#include "swin_block.h"
SWIN_MODULE_W(rrlite_dec1_4x4, 64, 2, 8, 0, false, CommonParams, 96)
#ifndef SWIN_NW
#define SWIN_NW 8
#endif
SWIN_MODULE_W(rrlite_dec1_4x4, 64, 2, SWIN_NW, 0, false, CommonParams, 96)
kernels/m/rrlite_dec2_4x4.hip
+10 -1
View File
@@ -1,7 +1,16 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
// RX 7700 XT: memory-clause scheduling at the 192 VGPR cap; byte-identical to the default build, gfx1101 only
// d4r-build-flags-gfx1101: -mllvm -misched=gcn-max-memory-clause
#define SWIN_NO_Q8
#define SWIN_F32ACC
#ifndef SWIN_MAX_VGPR
#define SWIN_MAX_VGPR 192
#endif
#ifndef SWIN_MLP_R
#define SWIN_MLP_R(NW) 2
#endif
#include "swin_block.h"
SWIN_MODULE_W(rrlite_dec2_4x4, 96, 4, 8, 0, false, CommonParams, 128)
#ifndef SWIN_NW
#define SWIN_NW 8
#endif
SWIN_MODULE_W(rrlite_dec2_4x4, 96, 4, SWIN_NW, 0, false, CommonParams, 128)
kernels/m/rrlite_enc1_4x4.hip
+8 -1
View File
@@ -1,8 +1,15 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#ifndef SWIN_MAX_VGPR
#define SWIN_MAX_VGPR 192
#endif
#ifndef SWIN_MLP_R
#define SWIN_MLP_R(NW) 1
#endif
#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)
#ifndef SWIN_NW
#define SWIN_NW 4
#endif
SWIN_MODULE_W(rrlite_enc1_4x4, 64, 2, SWIN_NW, 96, false, CommonParams, 0)
kernels/m/rrlite_enc2_4x4.hip
+11 -1
View File
@@ -1,7 +1,17 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
// RX 7700 XT: packed FP8 input decode, packed norm1 sums, half-wave exchange via ds_bpermute, MLP loop unrolled 4x;
// byte-identical to the default build, gfx1101 only (not measured on other targets)
// d4r-build-flags-gfx1101: -DSWIN_PACKED_INPUT -DSWIN_PACKED_NORM1 -DSWIN_HALF_BPERMUTE -DSWIN_UNROLL_MLP=4
#define SWIN_NO_Q8
#define SWIN_F32ACC
#ifndef SWIN_MAX_VGPR
#define SWIN_MAX_VGPR 192
#endif
#ifndef SWIN_MLP_R
#define SWIN_MLP_R(NW) 2
#endif
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc2_4x4, 96, 4, 8, 128, false, CommonParams, 0)
#ifndef SWIN_NW
#define SWIN_NW 8
#endif
SWIN_MODULE_W(rrlite_enc2_4x4, 96, 4, SWIN_NW, 128, false, CommonParams, 0)
kernels/m/rrlite_enc3_tube_4x4.hip
+8 -1
View File
@@ -1,11 +1,18 @@
// fast numerics: intermediate activations f16 (no FP8 re-quantisation), f32 accumulation through GEMM chains
#define SWIN_NO_Q8
#define SWIN_F32ACC
#ifndef SWIN_MAX_VGPR
#define SWIN_MAX_VGPR 192
#endif
#ifndef SWIN_MLP_R
#define SWIN_MLP_R(NW) 2
#endif
// NGX can recycle a weights address for another tube block after feature
// recreation. Address-only caching then serves stale packed weights. Prepare
// each launch; the default single image is sufficient on the launch stream.
#define SWIN_PREP_KEY_AT 0
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, 8, 0, true, TubeParams, 0)
#ifndef SWIN_NW
#define SWIN_NW 8
#endif
SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, SWIN_NW, 0, true, TubeParams, 0)
kernels/m/swin_block.h
+157 -7
View File
@@ -6,6 +6,8 @@
#pragma once
#include "swin_common.h"
#include <type_traits>
#define SWIN_PRAGMA_IMPL(x) _Pragma(#x)
#define SWIN_PRAGMA(x) SWIN_PRAGMA_IMPL(x)
template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
{
@@ -136,7 +138,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
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));
inw[k] = input_word((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 = [&]() {
@@ -146,9 +148,13 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
const uint32_t bytes = inw[k];
#ifdef SWIN_PACKED_INPUT
*(h4*)(X0 + T * AST + c4) = decode4(bytes);
#else
#pragma unroll
for (int b = 0; b < 4; ++b)
X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
#endif
}
};
@@ -162,10 +168,14 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
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));
input_word((const uint32_t*)(p.in + ((size_t)((c4 >> 5) * H2 + MY) * W2 + MX) * 32 + (c4 & 31)));
#ifdef SWIN_PACKED_INPUT
*(h4*)(PA + r * PAST + apos(c4)) = decode4(bytes);
#else
#pragma unroll
for (int b = 0; b < 4; ++b)
PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
#endif
}
stage_input_words(); // the skip input
__syncthreads();
@@ -207,7 +217,11 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
// 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
#ifdef SWIN_PACKED_RESIDUAL
hv2 residual[MT][NTW][4]; // Only f16 is needed until the output projection.
#else
T16 x[MT][NTW];
#endif
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
@@ -218,8 +232,15 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const int ch = ginv(n);
#pragma unroll
for (int i = 0; i < 8; i += 2)
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + mrow(i)) * AST + ch],
X0[(16 * (mg * MT + mm) + mrow(i + 1)) * AST + ch]} + (hv2){bo, bo});
{
const hv2 v = (hv2){X0[(16 * (mg * MT + mm) + mrow(i)) * AST + ch],
X0[(16 * (mg * MT + mm) + mrow(i + 1)) * AST + ch]} + (hv2){bo, bo};
#ifdef SWIN_PACKED_RESIDUAL
residual[mm][j][i / 2] = v;
#else
x[mm][j].set_pair(i / 2, v);
#endif
}
}
__syncthreads(); // stage 1 overwrites x0 rows with h1
@@ -240,6 +261,30 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
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();
#ifdef SWIN_PACKED_NORM1
hv2 ws[2];
#pragma unroll
for (int ww = 0; ww < 2; ++ww)
{
hv2 pp[NPL];
#pragma unroll
for (int pl = 0; pl < NPL; ++pl)
{
hv2 a = {v[pl][ww][0], v[pl][ww][1]}, b = {v[pl][ww][2], v[pl][ww][3]};
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
ws[ww] = pp[0] + pp[1];
}
hv2 a = ws[0] + ws[1];
a += __builtin_bit_cast(hv2, xor_lane<1>(__builtin_bit_cast(uint32_t, a)));
a += __builtin_bit_cast(hv2, xor_lane<2>(__builtin_bit_cast(uint32_t, a)));
const half_t rn = rsqrt16(a[0] + a[1]);
#else
half_t ws[2][2];
#pragma unroll
for (int ww = 0; ww < 2; ++ww)
@@ -271,6 +316,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
a0 = a0 + hxor<2>(a0);
a1 = a1 + hxor<2>(a1);
const half_t rn = rsqrt16(a0 + a1);
#endif
const hv2 rn2 = {rn, rn};
#pragma unroll
for (int pl = 0; pl < NPL; ++pl)
@@ -294,6 +340,48 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
{
constexpr int MW = 4 / (NW / NH); // windows per wave
const int h = NW == NH ? wv : wv % NH, m0 = NW == NH ? 0 : MW * (wv / NH);
hv2 qp[MW][2][4], vp[MW][2][4];
#ifdef SWIN_QV_SPLIT
// Q and V have independent k32 chains. Retain Q as packed halves while
// calculating V, instead of keeping both sets of f32 chains live.
auto project = [&](auto value_tag, hv2 (&dest)[MW][2][4]) {
constexpr bool VALUE = decltype(value_tag)::value;
T16 part[MW][2];
#pragma unroll
for (int m = 0; m < MW; ++m)
#pragma unroll
for (int q = 0; q < 2; ++q)
part[m][q] = t16_splat((half_t)0.0f);
#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 q = 0; q < 2; ++q)
b[s][q] = bload(w16, L::qv(h, VALUE ? 1 : 0), 2, kc, s, q, 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 < 2; ++q)
k32(part[m][q], a0, b[0][q], a1, b[1][q]);
}
}
#pragma unroll
for (int m = 0; m < MW; ++m)
#pragma unroll
for (int q = 0; q < 2; ++q)
#pragma unroll
for (int k = 0; k < 4; ++k) {
if constexpr (VALUE) dest[m][q][k] = part[m][q].pair(k);
else dest[m][q][k] = q8x2(part[m][q].pair(k));
}
};
project(std::false_type{}, qp);
project(std::true_type{}, vp);
#else
T16 acc[MW][4]; // [window][Q nt0, Q nt1, V nt0, V nt1]
#pragma unroll
for (int m = 0; m < MW; ++m)
@@ -320,7 +408,6 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
}
}
// 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
@@ -331,6 +418,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
qp[m][q][i / 2] = q8x2(acc[m][q].pair(i / 2));
vp[m][q][i / 2] = acc[m][2 + q].pair(i / 2);
}
#endif
__syncthreads(); // every wave is done reading h1; O may overwrite A
half_t* Qs = QP + wv * 16 * QST;
@@ -435,6 +523,16 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
__syncthreads();
// ------------------------------------------------ stage 3: output projection + residual
#ifdef SWIN_PACKED_RESIDUAL
T16 x[MT][NTW];
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int j = 0; j < NTW; ++j)
#pragma unroll
for (int k = 0; k < 4; ++k)
x[mm][j].set_pair(k, residual[mm][j][k]);
#endif
#pragma unroll
for (int kc = 0; kc < NH; ++kc)
{
@@ -465,6 +563,43 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
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();
#ifdef SWIN_PARALLEL_NORM2
// Four lanes per token. Preserve the original per-column and even/odd trees.
#pragma unroll
for (int r = 0; r < 8 / NW; ++r)
{
const int T = (64 / NW) * wv + (lane >> 2) + 8 * r;
const int c = 2 * (lane & 3);
hv2 sq[L::NT8];
#pragma unroll
for (int j = 0; j < L::NT8; ++j)
{
const hv2 v = *(const hv2*)(A + T * AST + 8 * j + c);
sq[j] = v * v;
}
hv2 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]);
}
hv2 sum = even + odd;
sum += __builtin_bit_cast(hv2, xor_lane<1>(__builtin_bit_cast(uint32_t, sum)));
sum += __builtin_bit_cast(hv2, xor_lane<2>(__builtin_bit_cast(uint32_t, sum)));
if ((lane & 3) == 0)
R2[T] = rsqrt16(sum[0] + sum[1]);
}
#else
if (tid < 64)
{
// the row as 128-bit LDS loads: element 8 j + c of n8 tile j
@@ -502,6 +637,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const half_t b1 = (L8[1] + L8[3]) + (L8[5] + L8[7]);
R2[tid] = rsqrt16(b0 + b1);
}
#endif
__syncthreads();
#pragma unroll
for (int j = 0; j < NTW; ++j)
@@ -529,7 +665,11 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
{
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);
#ifdef SWIN_UNROLL_MLP
SWIN_PRAGMA(unroll SWIN_UNROLL_MLP)
#else
#pragma unroll 1
#endif
for (int r = 0; r < L::NC / R; ++r)
{
const int c = R * r + cw;
@@ -636,7 +776,17 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
{
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);
uint32_t ca, cb;
#ifdef SWIN_DIRECT_CODES
if constexpr (NPM == 0) {
ca = direct_codes8x2(__builtin_bit_cast(hv2, pr.x));
cb = direct_codes8x2(__builtin_bit_cast(hv2, pr.y));
} else
#endif
{
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)
kernels/m/swin_common.h
+71 -1
View File
@@ -105,6 +105,18 @@ __device__ __forceinline__ half_t q8_exact(half_t h)
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)));
typedef _Float16 h4 __attribute__((ext_vector_type(4)));
__device__ __forceinline__ h4 decode4(uint32_t bytes)
{
uint32_t m = bytes & 0x7f7f7f7fu, sign = bytes & 0x80808080u;
uint32_t lo = (__builtin_amdgcn_perm(0u, m, 0x0c010c00u) << 7) |
__builtin_amdgcn_perm(0u, sign, 0x010c000cu);
uint32_t hi = (__builtin_amdgcn_perm(0u, m, 0x0c030c02u) << 7) |
__builtin_amdgcn_perm(0u, sign, 0x030c020cu);
const hv2 scale = {(half_t)256.0f, (half_t)256.0f};
hv2 a = __builtin_bit_cast(hv2, lo) * scale, b = __builtin_bit_cast(hv2, hi) * scale;
return (h4){a[0], a[1], b[0], b[1]};
}
// 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)
@@ -140,6 +152,34 @@ __device__ __forceinline__ uint32_t enc8x2(hv2 h, hv2& q)
return codes8x2(q);
}
// Direct encoder for stores that do not also consume the rounded f16 values.
// The codes remain in the low byte of each 16-bit half, as codes8x2 expects.
__device__ __forceinline__ uint32_t direct_codes8x2(hv2 h)
{
const u16x2 bits = __builtin_bit_cast(u16x2, h), m = bits & (unsigned short)0x7fff;
const u16x2 normal = (m + ((m >> 7) & (unsigned short)1) + (unsigned short)0xe03f) >> 7;
#ifdef SWIN_CODES_SHIFT
const u16x2 exp = m >> 10;
const u16x2 bounded = __builtin_elementwise_max(__builtin_elementwise_min(exp, (u16x2)8), (u16x2)1);
const u16x2 shift = (unsigned short)16 - bounded;
const u16x2 sig = (m & (unsigned short)0x3ff) |
(__builtin_elementwise_min(m, (u16x2)0x400) & (unsigned short)0x400);
const u16x2 sub = (sig + (((u16x2)1 << (shift - (unsigned short)1)) - (unsigned short)1) +
((sig >> shift) & (unsigned short)1)) >> shift;
#else
// Below 2^-6, f16(2 + |h|) has the rounded e4m3 subnormal code
// directly in its low bits (ULP at 2 is 2^-9). No decode round trip.
const hv2 biased = __builtin_bit_cast(hv2, m) + (hv2){(half_t)2.0f, (half_t)2.0f};
const u16x2 sub = __builtin_bit_cast(u16x2, biased) - (unsigned short)0x4000;
#endif
const u16x2 tiny = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
u16x2 code = __builtin_elementwise_min((sub & tiny) | (normal & ~tiny), (u16x2)0x7e);
const u16x2 nan = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
code |= nan & (unsigned short)0x7f;
code |= (bits >> 8) & (unsigned short)0x80;
return __builtin_bit_cast(uint32_t, code);
}
// 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)
@@ -332,7 +372,13 @@ 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__ hv2 pair(int k) const {
#ifdef SWIN_CVT_ASM
return __builtin_bit_cast(hv2, wm_cvt2(f[2 * k], f[2 * k + 1]));
#else
return (hv2){(half_t)f[2 * k], (half_t)f[2 * k + 1]};
#endif
}
__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; }
@@ -413,7 +459,11 @@ template <int MASK> __device__ __forceinline__ uint32_t xor_lane(uint32_t v)
else if constexpr (MASK == 8)
return __builtin_amdgcn_update_dpp(0, (int)v, 0x168, 0xf, 0xf, false);
else
#ifdef SWIN_HALF_BPERMUTE
return __builtin_amdgcn_ds_bpermute(((int)threadIdx.x ^ 16) * 4, v);
#else
return __builtin_amdgcn_permlanex16((int)v, (int)v, 0x76543210, 0xfedcba98, false, false);
#endif
}
template <int MASK> __device__ __forceinline__ half_t hxor(half_t v)
{
@@ -463,6 +513,16 @@ struct GemmDesc
// 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)
{
#if defined(SWIN_BUFFER_BLOAD) && D4R_WMMA_LAYOUT == 11
// Same 32-byte slot as the original load; only address formation changes.
const int tile = __builtin_amdgcn_readfirstlane(gemm_dst + ((kc * 2 + s) * NT + nt) * 16);
const __amdgpu_buffer_rsrc_t r = __builtin_amdgcn_make_buffer_rsrc((void*)w16, (short)0, 0x7fffffff, 0x31004000);
const uint32_t vo = (uint32_t)l16 * 32u, so = (uint32_t)tile * 32u;
const u4v_t a = __builtin_amdgcn_raw_buffer_load_b128(r, vo, so, 0);
const u4v_t b = __builtin_amdgcn_raw_buffer_load_b128(r, vo + 16, so, 0);
typedef uint32_t u8 __attribute__((ext_vector_type(8)));
return __builtin_bit_cast(gop, (u8){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]});
#else
#ifdef SWIN_SCALAR_BLOAD
const int tile = __builtin_amdgcn_readfirstlane(gemm_dst + ((kc * 2 + s) * NT + nt) * 16);
const wslot* t = w16 + tile;
@@ -476,6 +536,16 @@ __device__ __forceinline__ gop bload(const wslot* w16, int gemm_dst, int NT, int
#else
return t[at];
#endif
#endif
}
__device__ __forceinline__ uint32_t input_word(const uint32_t* p)
{
#ifdef SWIN_NT_INPUT
return __builtin_nontemporal_load(p);
#else
return *p;
#endif
}
__device__ __forceinline__ void expand_weights(const uint8_t* w, wslot* w16, const GemmDesc* descs, int ndesc, int idx)
kernels/native/build_native.sh
Executable
+18
View File
@@ -0,0 +1,18 @@
#!/usr/bin/env bash
# build_native.sh SRC.hip OUT.hsaco [extra clang flags]: natively compiled DLSS texture kernel (gfx1101, wave64 default)
set -euo pipefail
ND=$(cd "$(dirname "$0")" && pwd)
SRC=$(realpath "$1") OUT=$(realpath -m "$2"); shift 2
[[ -s "$SRC" ]] || { echo "empty native source: $SRC" >&2; exit 2; }
R=${D4R_ROCM_DIR:-/opt/rocm}
ARCH=${D4R_GPU_ARCH:-gfx1101}
W=${WAVE:-64}
WF=$([[ $W == 64 ]] && echo -mwavefrontsize64 || echo -mno-wavefrontsize64)
D=$(cd "$(dirname "$SRC")" && pwd)
T=$(mktemp -d); trap 'rm -rf "$T"' EXIT
(cd "$T" && "$R/lib/llvm/bin/clang++" -x hip -std=c++20 --offload-arch=$ARCH --offload-device-only -O3 $WF -mcumode \
-ffp-contract=off -fgpu-flush-denormals-to-zero -fno-fast-math -I"$D" -I"$ND" -I"$R/include" --rocm-path="$R" \
--rocm-device-lib-path="$R/amdgcn/bitcode" "$@" -o "$OUT" "$SRC" -save-temps=cwd --no-gpu-bundle-output)
S=$(ls "$T"/*.s | head -1)
cp "$S" "${OUT%.hsaco}.s"
printf '%s: %s\n' "$(basename "$OUT")" "$(grep -hE '^\s+\.(vgpr_count|sgpr_count|private_segment_fixed_size|vgpr_spill_count|group_segment_fixed_size):' "$S" | tr -s ' ' | paste -sd' ')"
kernels/native/ctile.py
Executable
+411
View File
@@ -0,0 +1,411 @@
"""Color-tap LDS tile pass for the hiluma output kernels (used by ptx2hip.py --ctile).
Path A of hiluma_engine_output_* (taken when a block's render footprint is wider than path B's 24-texel LDS tile,
e.g. at Quality) fetches 16 color texels per thread and converts each one (YCoCg, exposure, tonemap weight) right
after the fetch. The conversion of a texel depends only on the texel and the exposure %f1, and neighbouring
threads convert the same texels (~9x at ratio 1.5). This pass:
1. finds every tex.base on the color texture inside path A,
2. slices the instructions that compute the tap's packed result (cvt.rn.f16x2 x2 or cvt.rn.f16.f32) from the
fetched texel (forward taint + backward need, including the tonemap branch diamond), and checks that no other
value escapes the slice,
3. replaces the slice with `d4r.ctap.G` (tile lookup, falling back to fetch + the same conversion), and
4. emits one conversion function per slice shape (generated from the slice's own instructions, so the tile holds
bit-identical values) and a cooperative tile fill at the start of path A.
The tile is keyed by the final (clamped) texel coordinates, so any tap outside it just takes the fallback.
"""
import re
REG = re.compile(r'%[A-Za-z_$][\w$]*')
CAP = 1024 # tile texels (TW*TH); a larger footprint falls back to direct fetches
class CtileErr(Exception):
pass
def parse(body, tokenize):
"""-> list of dicts: kind ('stmt'|'label'|'open'|'close'), text, plus op/dst/src/pred for stmts"""
items = []
for kind, text in tokenize(body):
it = {'kind': kind, 'text': text}
if kind == 'stmt':
s = text
m = re.match(r'@(!?)(%p\d+)\s+(.*)$', s)
it['pred'] = m.group(2) if m else None
if m:
s = m.group(3)
parts = s.split(None, 1)
it['op'] = parts[0]
rest = parts[1] if len(parts) > 1 else ''
it['rest'] = rest
dst, src = [], []
op = it['op']
if op.startswith('.') or op.startswith('bra') or op.startswith('bar') or op.startswith('ret'):
src = REG.findall(rest)
elif op.startswith('st.') or op.startswith('sust'):
src = REG.findall(rest)
else:
# first operand = destination(s)
depth, i = 0, 0
while i < len(rest):
c = rest[i]
if c in '{[':
depth += 1
elif c in '}]':
depth -= 1
elif c == ',' and depth == 0:
break
i += 1
dst = REG.findall(rest[:i])
src = REG.findall(rest[i:])
if it['pred']:
src = src + [it['pred']]
it['dst'], it['src'] = dst, src
items.append(it)
return items
def branch_target(it):
return it['rest'].strip().rstrip(';').strip().lstrip('$') if it['kind'] == 'stmt' and it['op'].startswith('bra') else None
def find_taps(items, color_aliases):
return [i for i, it in enumerate(items) if it['kind'] == 'stmt' and it['op'] == 'tex.base.2d.v4.f32.s32'
and re.search(r'\[\s*(%rd\d+)', it['rest']).group(1) in color_aliases]
def slice_tap(items, t):
"""-> (removed index set, results list, kind 'in'|'out', externals set)"""
tex = items[t]
taint = set(tex['dst'])
tainted_idx = [t]
results, kind = [], None
i = t + 1
while i < len(items):
it = items[i]
if it['kind'] == 'stmt' and any(s in taint for s in it['src']):
op = it['op']
if op.startswith('cvt.rn.f16x2.f32') and kind in (None, 'in'):
kind = 'in'
results.append((i, it['dst'][0]))
tainted_idx.append(i)
if len(results) == 2:
break
elif op.startswith('cvt.rn.f16.f32') and kind is None:
kind = 'out'
results.append((i, it['dst'][0]))
tainted_idx.append(i)
break
else:
taint.update(it['dst'])
tainted_idx.append(i)
i += 1
else:
raise CtileErr(f'tap at {t}: no packed result found')
last = results[-1][0]
# backward need from the results (within [t, last])
need_regs = set()
for ri, _ in results:
need_regs.update(items[ri]['src'])
needed = {ri for ri, _ in results}
# conditional branches on tainted predicates (the tonemap branch) are part of the slice
for j in tainted_idx:
it = items[j]
if it['op'].startswith('bra') and it['pred'] is not None:
needed.add(j)
need_regs.add(it['pred'])
for j in range(last, t - 1, -1):
it = items[j]
if j in needed:
continue
if it['kind'] == 'stmt' and j in tainted_idx and any(d in need_regs for d in it['dst']):
needed.add(j)
need_regs.update(it['src'])
extra = [j for j in tainted_idx if j not in needed]
if extra:
raise CtileErr(f'tap at {t}: tainted but not needed: {[items[j]["text"] for j in extra]}')
removed = set(needed)
# branch diamonds: region from each conditional branch to its merge label
for j in sorted(needed):
it = items[j]
if not it['op'].startswith('bra'):
continue
targets, seen = {branch_target(it)}, set()
k = j + 1
while True:
x = items[k]
if x['kind'] == 'label':
lab = x['text'].lstrip('$')
if lab in targets and targets - seen == {lab} and lab not in {branch_target(it)}:
removed.add(k)
break
seen.add(lab)
removed.add(k)
elif x['kind'] == 'stmt':
if x['op'].startswith('bra'):
if x['pred'] is not None and k not in needed:
raise CtileErr(f'tap at {t}: foreign conditional branch in diamond')
targets.add(branch_target(x))
removed.add(k)
elif k not in needed:
# constant moves used only inside the slice are part of it
if x['op'].startswith('mov') and not any(REG.match(s) for s in x['src']):
removed.add(k)
else:
raise CtileErr(f'tap at {t}: foreign instruction in diamond: {x["text"]}')
k += 1
if k > last:
raise CtileErr(f'tap at {t}: diamond does not close before the result')
# escape check: registers defined in the slice (except results) used outside it
defs = set()
for j in removed:
if items[j]['kind'] == 'stmt':
defs.update(items[j]['dst'])
res_regs = {r for _, r in results}
inner_defs = defs - res_regs
for j, it in enumerate(items):
if j in removed or it['kind'] != 'stmt':
continue
bad = inner_defs.intersection(it['src'])
if bad:
raise CtileErr(f'tap at {t}: {bad} escapes to {it["text"]}')
if defs.intersection(it['dst']) - res_regs or res_regs.intersection(it['dst']):
raise CtileErr(f'tap at {t}: slice register redefined by {it["text"]}')
# labels removed must not be targeted from outside
labels = {items[j]['text'].lstrip('$') for j in removed if items[j]['kind'] == 'label'}
for j, it in enumerate(items):
if j not in removed and branch_target(it) in labels:
raise CtileErr(f'tap at {t}: label targeted from outside the slice')
used = set()
for j in removed:
if items[j]['kind'] == 'stmt' and j != t:
used.update(items[j]['src'])
externals = used - defs - set(tex['dst'])
externals.discard(None)
return removed, results, kind, externals
def canon(items, removed, t, results):
"""canonical text of the slice (registers by first appearance, texel by position)"""
names, labels = {}, {}
for k, r in enumerate(items[t]['dst']):
names[r] = f'T{k}'
for k, (_, r) in enumerate(results):
names[r] = f'R{k}'
out = []
for j in sorted(removed):
if j == t:
continue
it = items[j]
if it['kind'] == 'label':
lab = it['text'].lstrip('$')
labels.setdefault(lab, f'L{len(labels)}')
out.append(labels[lab] + ':')
continue
if it['kind'] != 'stmt':
continue
s = it['text']
def rn(m):
r = m.group(0)
if r not in names:
names[r] = f'V{len(names)}'
return names[r]
s = REG.sub(rn, s)
s = re.sub(r'\$(L__\w+)', lambda m: labels.setdefault(m.group(1), f'L{len(labels)}'), s)
out.append(s)
return '\n'.join(out), names
def apply(tr, tokenize):
"""Rewrite tr.body in place for the color tile. Returns the HIP prelude (conversion functions) and info."""
body = tr.body
items = parse(body, tokenize)
# color texture handle: ld.param.u64 [%rd1+-56] (param+200, %rd10 / %rd69); taps use it or its mov aliases
color = {it['dst'][0] for it in items if it['kind'] == 'stmt' and it['op'] == 'ld.param.u64'
and re.search(r'\[%rd1\+-56\]', it['rest'])}
aliases = set(color)
for it in items:
if it['kind'] == 'stmt' and it['op'] == 'mov.u64' and len(it['src']) == 1 and it['src'][0] in color:
aliases.add(it['dst'][0])
# path A's color taps lie between $L__BB0_1 and $L__BB0_77 (the per-tap blocks); path B's are elsewhere
labels = {it['text'].lstrip('$'): i for i, it in enumerate(items) if it['kind'] == 'label'}
a0, a1 = labels['L__BB0_1'], labels['L__BB0_77']
taps = [t for t in find_taps(items, aliases) if a0 < t < a1]
if len(taps) != 16:
raise CtileErr(f'expected 16 path A color taps, found {len(taps)}')
groups, per_tap = {}, []
all_removed = set()
for t in taps:
removed, results, kind, ext = slice_tap(items, t)
if removed & all_removed:
raise CtileErr('overlapping slices')
all_removed |= removed
text, names = canon(items, removed, t, results)
ext_names = sorted(ext, key=lambda r: names.get(r, r))
key = (kind, text, tuple(names.get(r) for r in ext_names))
g = groups.setdefault(key, len(groups))
per_tap.append((t, removed, results, kind, ext_names, g))
if len(groups) > 3:
raise CtileErr(f'{len(groups)} slice shapes')
# externals: %f1 (exposure) or constant registers (single mov-immediate definition)
const_def = {}
for it in items:
if it['kind'] == 'stmt' and it['op'].startswith('mov') and len(it['dst']) == 1 and not it['src']:
const_def.setdefault(it['dst'][0], []).append(it)
group_info = {}
for (t, removed, results, kind, ext_names, g) in per_tap:
if g in group_info:
continue
params, consts = [], []
for r in ext_names:
if r == '%f1':
params.append(r)
elif r in const_def and len(const_def[r]) == 1:
consts.append(const_def[r][0]['text'])
else:
raise CtileErr(f'external {r} of tap {t}')
group_info[g] = dict(t=t, removed=removed, results=results, kind=kind, params=params, consts=consts)
return items, per_tap, group_info, color
def emit(tr, tokenize, Translator):
"""Apply the pass to translator tr (before translate()). Installs the pseudo-op handler; returns prelude text."""
items, per_tap, group_info, color = apply(tr, tokenize)
# ---- conversion functions, one per group, translated from the group's first slice
prelude = []
for g, info in sorted(group_info.items()):
t, removed, results = info['t'], info['removed'], info['results']
sub = Translator.__new__(Translator)
sub.kernel = tr.kernel
sub.param_name = tr.param_name
sub.regs, sub.shared, sub.locals, sub.lines = {}, [], [], []
sub.scope_names, sub.scope_id = [{}], 0
tex = items[t]
# texel inputs
for k, r in enumerate(tex['dst']):
sub.reg(r)
sub.lines.append(f' {sub.reg(r)} = t_[{k}];')
for c in info['consts']:
sub.stmt(c)
for j in sorted(removed):
if j == t:
continue
it = items[j]
if it['kind'] == 'label':
sub.lines.append(f'{it["text"].replace("$", "")}:;')
elif it['kind'] == 'stmt':
sub.stmt(it['text'])
outs = [sub.reg(r) for _, r in results]
decl = []
for name, ty in sorted(sub.regs.items()):
if name in [sub.reg(p) for p in info['params']]:
continue
decl.append(f' {ty} {name};')
args = ', '.join(['f4 t_'] + [f'float {sub.reg(p)}' for p in info['params']] + [f'uint32_t* o{k}' for k in range(len(outs))])
prelude.append(f'RT void d4r_ctap_conv{g}({args})\n{{')
prelude.extend(decl)
prelude.extend(sub.lines)
for k, o in enumerate(outs):
prelude.append(f' *o{k} = (uint32_t){o};')
prelude.append('}')
# ---- tile storage (aliases the kernel's LDS through d4r_lds)
offs, off = {}, 0
for g, info in sorted(group_info.items()):
for k in range(len(info['results'])):
offs[(g, k)] = off
off += CAP * 4
# ---- rewrite the body: removed statements dropped, pseudo-op at each tap, fill at the start of path A
tap_at = {t: (results, ext, g) for (t, removed, results, kind, ext, g) in per_tap}
removed_all = set()
for (t, removed, *_r) in per_tap:
removed_all |= removed
out = []
for i, it in enumerate(items):
if it['kind'] == 'open':
out.append('{'); continue
if it['kind'] == 'close':
out.append('}'); continue
if it['kind'] == 'label':
if i in removed_all:
continue
out.append(it['text'] + ':')
if it['text'].lstrip('$') == 'L__BB0_1':
out.append('d4r.ctile_fill;')
continue
if i in tap_at:
results, ext, g = tap_at[i]
m = re.search(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', it['rest'])
regs = ', '.join(r for _, r in results)
out.append(f'd4r.ctap.{g} {regs}, {m.group(1)}, {m.group(2)};')
continue
if i in removed_all:
continue
out.append(it['text'] + ';')
tr.body = '\n'.join(out)
summary = dict(groups=group_info, offs=offs, lds_bytes=off, taps=len(per_tap))
orig_op = tr.op
def op(opname, a):
if opname == 'd4r.ctile_fill':
tr.emit('d4r_ctile_fill(KARGP, (int32_t)r3, (int32_t)r4, (int32_t)r5, (int32_t)r6, f1, &ct_x0, &ct_y0, &ct_w, &ct_h);')
return
if opname.startswith('d4r.ctap.'):
g = int(opname.split('.')[2])
nres = len(group_info[g]['results'])
res = [tr.reg(x) for x in a[:nres]]
h = tr.val(a[nres], 'u64')
x, y = tr.val(a[nres + 1], 'u32'), tr.val(a[nres + 2], 'u32')
prm = ''.join(f', {tr.val(p, "f32")}' for p in group_info[g]['params'])
tr.emit(f'{{ int32_t dx_ = (int32_t){x} - ct_x0, dy_ = (int32_t){y} - ct_y0;')
tr.emit(f' if ((uint32_t)dx_ < (uint32_t)ct_w && (uint32_t)dy_ < (uint32_t)ct_h) {{ uint32_t i_ = (uint32_t)(dy_ * ct_w + dx_);')
for k, r in enumerate(res):
tr.emit(f' {r} = (uint32_t)((const AS3 uint32_t*)(d4r_lds + {offs[(g, k)]}))[i_];')
outs = ', '.join(f'o_ + {k}' for k in range(nres))
tr.emit(f' }} else {{ uint32_t o_[2]; d4r_ctap_conv{g}(tex_fetch({h}, (int32_t){x}, (int32_t){y}){prm}, {outs});')
for k, r in enumerate(res):
tr.emit(f' {r} = o_[{k}];')
tr.emit(' } }')
return
orig_op(opname, a)
tr.op = op
# result register types: f16x2 packs are u32 regs, f16 packs are u16 regs: both assigned via (uint32_t) casts
# fill function (needs the conversion functions and offsets)
fill = ['RT void d4r_ctile_fill(const AS4 uint8_t* KARGP, int32_t x0l, int32_t y0l, int32_t x1l, int32_t y1l, float f1,',
' int32_t* ox0, int32_t* oy0, int32_t* ow, int32_t* oh)',
'{',
' // final texel coordinate of a tile-relative logical coordinate u: u > 0 ? min(lo + u, hi) : lo',
' // (the clamp every color tap applies; params +208/+212 = lo, +216/+220 = hi)',
' int32_t lox = *(const AS4 int32_t*)(KARGP + 208), loy = *(const AS4 int32_t*)(KARGP + 212);',
' int32_t hix = *(const AS4 int32_t*)(KARGP + 216), hiy = *(const AS4 int32_t*)(KARGP + 220);',
' int32_t fx0 = x0l > 0 ? __builtin_elementwise_min(lox + x0l, hix) : lox;',
' int32_t fy0 = y0l > 0 ? __builtin_elementwise_min(loy + y0l, hiy) : loy;',
' int32_t fx1 = x1l > 0 ? __builtin_elementwise_min(lox + x1l, hix) : lox;',
' int32_t fy1 = y1l > 0 ? __builtin_elementwise_min(loy + y1l, hiy) : loy;',
' int32_t w = fx1 - fx0 + 1, h = fy1 - fy0 + 1;',
f' if (w <= 0 || h <= 0 || w * h > {CAP}) {{ w = 0; h = 0; }}',
' uint64_t tex = *(const AS4 uint64_t*)(KARGP + 200);',
' int32_t n = w * h;',
' for (int32_t i = (int32_t)(__builtin_amdgcn_workitem_id_y() * 16u + __builtin_amdgcn_workitem_id_x()); i < n; i += 256)',
' {',
' int32_t ty = i / w, tx = i - ty * w;',
' f4 t_ = tex_fetch(tex, fx0 + tx, fy0 + ty);']
for g, info in sorted(group_info.items()):
nres = len(info['results'])
prm = ''.join(', f1' for _ in info['params'])
outs = ', '.join(f'o_ + {k}' for k in range(nres))
fill.append(f' {{ uint32_t o_[2]; d4r_ctap_conv{g}(t_{prm}, {outs});')
for k in range(nres):
fill.append(f' ((AS3 uint32_t*)(d4r_lds + {offs[(g, k)]}))[i] = o_[{k}];')
fill.append(' }')
fill += [' }',
' bar_sync();',
' *ox0 = fx0; *oy0 = fy0; *ow = w; *oh = h;',
'}']
for p in sum((info['params'] for info in group_info.values()), []):
if p != '%f1':
raise CtileErr('only %f1 is supported as a conversion parameter')
return prelude, fill, summary
kernels/native/ifconv.py
Executable
+234
View File
@@ -0,0 +1,234 @@
"""If-conversion of pure branch diamonds in PTX (used by ptx2hip.py --ifconv).
NVIDIA's DLSS kernels branch around small pure arithmetic alternatives (e.g. the inverse tonemap
`v < 0.998 ? 1 / (1 - v) : (exp2(...) - 1) / (9 v)`), laid out as
@%p bra $A; (or @!%p)
bra.uni $B;
$A: <pure ops> bra.uni $M;
$B: <pure ops> (falls through)
$M:
Branches like these split the code into basic blocks that no compiler schedules memory operations across, so
independent fetch chains on either side are serialized. This pass rewrites such a diamond into straight-line code:
both arms compute into fresh registers and every register an arm assigns is then selected by the predicate. Each lane
keeps exactly the values of the arm it would have taken (GPU arithmetic has no traps), so results are bit-identical.
Only arms made of register arithmetic are converted (no memory, texture, barrier, shuffle, branch or return; no
integer division, which is undefined behaviour in C for some inputs).
"""
import re
REG = re.compile(r'%([a-zA-Z]+)(\d+)\b')
PURE_BASES = {'add', 'sub', 'mul', 'fma', 'mad', 'div', 'rcp', 'sqrt', 'rsqrt', 'ex2', 'lg2', 'min', 'max', 'neg',
'abs', 'cvt', 'mov', 'selp', 'setp', 'and', 'or', 'xor', 'not', 'shl', 'shr', 'set', 'bfi', 'copysign',
'sin', 'cos', 'cvta'}
SEL = {'f': 'selp.f32', 'r': 'selp.b32', 'rd': 'selp.b64', 'rs': 'selp.b16'}
def _pure(it):
if it['kind'] != 'stmt' or it['pred'] is not None:
return False
op = it['op']
base = op.split('.')[0]
if base not in PURE_BASES:
return False
if base == 'div' and not op.endswith('.f32') and 'approx' not in op:
return False # integer division (and IEEE f32 division, which the translator rejects anyway)
if base == 'cvta':
return False
return True
def _target(it):
if it['kind'] == 'stmt' and it['op'].startswith('bra'):
return it['rest'].strip().rstrip(';').strip()
return None
def apply(body, tokenize, parse, split_operands):
items = parse(body, tokenize)
# label -> number of branches targeting it
refs = {}
for it in items:
t = _target(it)
if t:
refs[t] = refs.get(t, 0) + 1
# fresh register numbers per kind
top = {}
for m in REG.finditer(body):
k, n = m.group(1), int(m.group(2))
top[k] = max(top.get(k, 0), n)
fresh = {k: v + 100000 for k, v in top.items()}
def new_reg(kind):
fresh[kind] = fresh.get(kind, 100000) + 1
return f'%{kind}{fresh[kind]}'
use_sites = {}
for x in items:
if x['kind'] == 'stmt':
for r in x['src']:
use_sites.setdefault(r, []).append(id(x))
def uses_outside(r, arm_ids):
return any(u not in arm_ids for u in use_sites.get(r, []))
defined_before = set()
out, i, n_conv = [], 0, 0
while i < len(items):
it = items[i]
conv = None
if it['kind'] == 'stmt' and it['op'].startswith('bra') and it['pred'] is not None:
conv = _match(items, i, refs)
if conv is None:
if it['kind'] == 'stmt':
defined_before.update(it['dst'])
out.append(_text(it))
i += 1
continue
neg, pred, arm_a, arm_b, merge_idx = conv
lines, last = [], [{}, {}]
for side, arm in enumerate((arm_a, arm_b)):
ren = {}
for st in arm:
op, rest = st['text'].split(None, 1)
ops = split_operands(rest)
if len(st['dst']) != 1 or ops[0] != st['dst'][0]:
raise ValueError('unexpected destination form: ' + st['text'])
srcs = [REG.sub(lambda mm: ren.get(mm.group(0), mm.group(0)), o) for o in ops[1:]]
d = st['dst'][0]
kind = REG.match(d).group(1)
if kind not in SEL:
raise ValueError('predicate assigned in a diamond arm: ' + st['text'])
nr = new_reg(kind)
ren[d] = nr
lines.append(f"{op} {', '.join([nr] + srcs)}")
last[side] = ren
# only registers read outside the two arms need a select (arm temporaries are dropped)
arm_ids = {id(x) for x in arm_a} | {id(x) for x in arm_b}
assigned = sorted(r for r in set(last[0]) | set(last[1]) if uses_outside(r, arm_ids))
cond = pred
for r in assigned:
kind = REG.match(r).group(1)
a = last[0].get(r, r if r in defined_before else last[1].get(r))
b = last[1].get(r, r if r in defined_before else last[0].get(r))
# selp d, a, b, p: d = p ? a : b ; arm A is taken when (p xor neg)
if neg:
a, b = b, a
lines.append(f'{SEL[kind]} {r}, {a}, {b}, {cond}')
defined_before.update(assigned)
out.extend(l + ';' for l in lines)
n_conv += 1
i = merge_idx # the merge label itself is emitted by the main loop
return '\n'.join(out), n_conv
def _text(it):
if it['kind'] == 'open':
return '{'
if it['kind'] == 'close':
return '}'
if it['kind'] == 'label':
return it['text'] + ':'
return it['text'] + ';'
def _match(items, i, refs):
"""@p bra A; bra.uni B; A: ... bra.uni M; B: ... M: -> (neg, pred, armA, armB, index of M)"""
br = items[i]
m = re.match(r'@(!?)(%p\d+)\s', br['text'])
if not m:
return None
neg, pred = m.group(1) == '!', m.group(2)
a_lab = _target(br)
if i + 2 >= len(items):
return None
j = items[i + 1]
if not (j['kind'] == 'stmt' and j['op'].startswith('bra') and j['pred'] is None):
return None
b_lab = _target(j)
if refs.get(a_lab) != 1 or refs.get(b_lab) != 1:
return None
def arm(start, label):
if not (items[start]['kind'] == 'label' and items[start]['text'] == label):
return None
body, k = [], start + 1
while k < len(items):
x = items[k]
if x['kind'] == 'label':
return body, k, None
if x['kind'] in ('open', 'close'):
return None
if x['kind'] == 'stmt' and x['op'].startswith('bra'):
if x['pred'] is not None:
return None
return body, k + 1, _target(x)
if not _pure(x):
return None
body.append(x)
k += 1
return None
first = arm(i + 2, a_lab)
order = 'AB'
if first is None:
first = arm(i + 2, b_lab)
order = 'BA'
if first is None:
return None
body1, k, jump1 = first
if jump1 is None:
return None
second = arm(k, b_lab if order == 'AB' else a_lab)
if second is None:
return None
body2, k2, jump2 = second
if jump2 is not None and jump2 != jump1:
return None # the second arm may also jump to the merge label explicitly
if not (k2 < len(items) and items[k2]['kind'] == 'label' and items[k2]['text'] == jump1):
return None
arm_a, arm_b = (body1, body2) if order == 'AB' else (body2, body1)
return neg, pred, arm_a, arm_b, k2
def sink_stores(body, tokenize, parse):
"""Drop labels nothing branches to, then move each st.shared / st.local down to just before the next statement
that could observe or reorder it: a shared/local load or any other store, a barrier, shuffle, atomic, label,
branch or return, or a statement that writes one of the store's registers. Stores keep their relative order, so
every memory location sees the same sequence of writes; texture loads (read-only data) and arithmetic move ahead
of them, which lets independent fetch chains issue together. Returns (body, number of stores moved)."""
items = parse(body, tokenize)
refs = set()
for it in items:
t = _target(it)
if t:
refs.add(t)
items = [it for it in items if not (it['kind'] == 'label' and it['text'] not in refs)]
def is_fence(it):
if it['kind'] != 'stmt':
return True # labels and scope braces
op = it['op']
base = op.split('.')[0]
if base in ('bar', 'shfl', 'atom', 'red', 'bra', 'ret', 'membar', 'fence', 'st', 'sust', 'd4r'):
return True
if base == 'ld' and ('.shared' in op or '.local' in op):
return True
return False
moved = 0
i = len(items) - 1
while i >= 0:
it = items[i]
if it['kind'] == 'stmt' and it['pred'] is None and it['op'].startswith('st.') and ('.shared' in it['op'] or '.local' in it['op']):
regs = set(it['src'])
j = i + 1
while j < len(items) and not is_fence(items[j]) and not (regs & set(items[j]['dst'])):
j += 1
if j > i + 1:
# after the pop the fence that stopped the scan is at j - 1: the store goes right before it
items.insert(j - 1, items.pop(i))
moved += 1
i -= 1
return '\n'.join(_text(it) for it in items), moved
kernels/native/in_fuse.py
Executable
+123
View File
@@ -0,0 +1,123 @@
#!/usr/bin/env python3
"""in_fuse.py IN.hip > in_unit.h: the hiluma input kernel (translated by ptx2hip.py) as a device function
in_unit(KARGP, gx, gy, out[4])
computing what thread (gx, gy) of the original launch (blocks of 16x16) writes: the 4 halves at network pixel
(gx >> 2, gy), channels 4 (gx & 3) .. 4 (gx & 3) + 3 (used by in_regs.hip).
Changes, all value-preserving:
- thread ids: workgroup/workitem ids -> gx >> 4, gx & 15, gy >> 4, gy & 15 (the kernel only uses gx = 16 wg_x + wi_x,
gy = 16 wg_y + wi_y and gx & 1, gx & 2);
- the per-thread LDS scratch (4 planes of 3 floats at slot 12 * (16 wi_y + wi_x), written in straight-line code
before any branch, read at fixed offsets and once at plane r14 in 0..3) -> 12 registers + a 4-way select;
- the 32-byte local depot (4 float pairs written in straight-line code, one pair read at plane r14) -> 8 registers
+ a 4-way select;
- the 4 final f16 stores -> out[0..3].
Refuses any input that does not match these patterns exactly.
"""
import re
import sys
src = open(sys.argv[1]).read()
lines = src.split('\n')
k0 = next(i for i, l in enumerate(lines) if l.startswith('extern "C" __attribute__((global))'))
assert lines[k0 + 1] == '{'
end = len(lines) - 1
while lines[end].strip() != '}':
end -= 1
body = lines[k0 + 2:end]
first_label = next(i for i, l in enumerate(body) if re.match(r'^L__BB0_\d+:;$', l))
out = []
SUBS = {
'r39 = (uint32_t)((uint32_t)__builtin_amdgcn_workgroup_id_x());': 'r39 = gx >> 4;',
'r41 = (uint32_t)((uint32_t)__builtin_amdgcn_workitem_id_x());': 'r41 = gx & 15u;',
'r46 = (uint32_t)((uint32_t)__builtin_amdgcn_workgroup_id_y());': 'r46 = gy >> 4;',
'r48 = (uint32_t)((uint32_t)__builtin_amdgcn_workitem_id_y());': 'r48 = gy & 15u;',
'r191 = (uint32_t)((uint32_t)__builtin_amdgcn_workitem_id_y());': 'r191 = 0u;',
'r190 = (uint32_t)((uint32_t)__builtin_amdgcn_workitem_id_x());': 'r190 = 0u;',
'r90 = (uint32_t)((uint32_t)(uintptr_t)(AS3 uint8_t*)sharedWarpedSpaceHistoryYcocg);': 'r90 = 0u; // (LDS scratch replaced by registers)',
}
seen = set()
ST = re.compile(r'^ \*\(AS3 float\*\)\(uintptr_t\)\(r20 \+ (\d+)u\) = \(float\)\((f\d+)\);$')
LD = re.compile(r'^ (f\d+) = \*\(AS3 float\*\)\(uintptr_t\)\(r20 \+ (\d+)u\);$')
LD153 = re.compile(r'^ (f\d+) = \*\(AS3 float\*\)\(uintptr_t\)\(r153 \+ 0u\);$')
STORE = re.compile(r'^ \*\(uint16_t\*\)\(uintptr_t\)\(rd48 \+ (\d)\) = \(uint16_t\)\((rs\d+)\);$')
stored = set()
nst = 0
# the 32-byte local depot: 4 float pairs (offsets 0..28) written through rd6 / rd59 / rd61 = depot + 0 in
# straight-line code, read once as the pair at rd61 + 8 r14 (r14 in 0..3) -> 8 registers + a 4-way select
DST = re.compile(r'^ \*\(AS5 float\*\)\(uintptr_t\)\(uint32_t\)\((rd6|rd59|rd61) \+ (\d+)\) = \(float\)\((f\d+)\);$')
DLD = re.compile(r'^ (f\d+) = \*\(AS5 float\*\)\(uintptr_t\)\(uint32_t\)\(rd44 \+ ([04])\);$')
dstored = set()
for i, l in enumerate(body):
t = l.strip()
if t.startswith('const AS4 uint8_t* KARGP =') or t.startswith('const uint64_t KARGU ='):
continue
if t in SUBS:
out.append(' ' + SUBS[t])
seen.add(t)
continue
m = ST.match(l)
if m:
off = int(m.group(1))
assert i < first_label and off % 4 == 0 and off // 3072 < 4 and (off % 3072) // 4 < 3, l
out.append(f' shv{off // 3072}_{(off % 3072) // 4} = {m.group(2)};')
stored.add(off)
continue
m = LD.match(l)
if m:
off = int(m.group(2))
assert off in stored, f'load before store: {l}'
out.append(f' {m.group(1)} = shv{off // 3072}_{(off % 3072) // 4};')
continue
m = LD153.match(l)
if m:
assert all(o in stored for o in (0, 3072, 6144, 9216)), 'plane loads need all planes stored'
out.append(f' {m.group(1)} = r14 == 0u ? shv0_0 : r14 == 1u ? shv1_0 : r14 == 2u ? shv2_0 : shv3_0; // r14 in 0..3')
continue
m = DST.match(l)
if m:
off = int(m.group(2))
assert i < first_label and off % 4 == 0 and off < 32, l
out.append(f' dpv{off // 8}_{(off % 8) // 4} = {m.group(3)};')
dstored.add(off)
continue
m = DLD.match(l)
if m:
assert len(dstored) == 8, 'depot read before all pairs are stored'
j = int(m.group(2)) // 4
out.append(f' {m.group(1)} = r14 == 0u ? dpv0_{j} : r14 == 1u ? dpv1_{j} : r14 == 2u ? dpv2_{j} : dpv3_{j}; // r14 in 0..3')
continue
m = STORE.match(l)
if m:
out.append(f' out[{int(m.group(1)) // 2}] = {m.group(2)};')
nst += 1
continue
assert 'AS3' not in l and 'workitem_id' not in l and 'workgroup_id' not in l, f'unhandled: {l}'
assert not re.search(r'\(rd48 \+', l), f'unhandled store: {l}'
out.append(l)
assert len(seen) == len(SUBS), set(SUBS) - seen
assert nst == 4, nst
assert len(stored) == 12, sorted(stored)
assert len(dstored) == 8, sorted(dstored)
for k in (' rd6 = (uint64_t)(l2_SPL + (uint64_t)0ull);', ' rd59 = (uint64_t)(l2_SPL + (uint64_t)0ull);',
' rd61 = (uint64_t)(l2_SPL + (uint64_t)0ull);', ' rd43 = (uint64_t)((int64_t)(int32_t)r14 * (int64_t)(int32_t)8);',
' rd44 = (uint64_t)(rd61 + rd43);'):
assert k in body, k
# r14 must be 2 r58 + (1 - r65) with r58, r65 parity bits (checked textually)
assert ' r14 = (uint32_t)(r68 + r67);' in body and ' r68 = (uint32_t)(r58 << (1u & 31));' in body
assert ' r67 = (uint32_t)(r26 - r65);' in body and ' r26 = (uint32_t)(1u);' in body
hdr = ['// Generated by kernels/native/in_fuse.py from ' + sys.argv[1].split('/')[-1] + ' (do not edit).',
'#pragma once',
'#include "native_rt.h"',
'// Thread (gx, gy) of hiluma_engine_input_depthinv_mvhi_hdr_v2_rel (16x16 blocks): the 4 halves it stores at',
'// network pixel (gx >> 2, gy), channels 4 (gx & 3) ... KARGP: the input kernel\'s 248-byte parameter block.',
'__attribute__((device)) inline __attribute__((always_inline)) void in_unit(const AS4 uint8_t* KARGP, uint32_t gx, uint32_t gy, uint16_t out[4])',
'{',
' const uint64_t KARGU = (uint64_t)(uintptr_t)KARGP;',
' float ' + ', '.join(f'shv{p}_{j}' for p in range(4) for j in range(3)) + ';',
' float ' + ', '.join(f'dpv{p}_{j}' for p in range(4) for j in range(2)) + ';']
print('\n'.join(hdr + out + ['}']))
kernels/native/in_regs.hip
+22
View File
@@ -0,0 +1,22 @@
// hiluma input kernel rebuilt from in_unit.h (kernels/native/in_fuse.py on in_ifs.hip): the translated kernel's
// per-thread LDS scratch and 32-byte local array live in registers (the local array compiled to scratch-memory
// flat stores/loads; -21 % in the replay bench). Same launch (16x16 blocks), same output.
// Build: kernels/build.sh (generates in_unit.h from the extracted PTX); by hand: WAVE=32 build_native.sh in_regs.hip OUT.hsaco
#include "in_unit.h"
struct __attribute__((aligned(8))) hiluma_engine_input_depthinv_mvhi_hdr_v2_rel_params { uint8_t b[248]; };
extern "C" __attribute__((global)) __attribute__((amdgpu_flat_work_group_size(1, 256))) void hiluma_engine_input_depthinv_mvhi_hdr_v2_rel(hiluma_engine_input_depthinv_mvhi_hdr_v2_rel_params)
{
const AS4 uint8_t* KARGP = (const AS4 uint8_t*)__builtin_amdgcn_kernarg_segment_ptr();
const uint32_t gx = __builtin_amdgcn_workgroup_id_x() * 16u + __builtin_amdgcn_workitem_id_x();
const uint32_t gy = __builtin_amdgcn_workgroup_id_y() * 16u + __builtin_amdgcn_workitem_id_y();
uint16_t o[4];
in_unit(KARGP, gx, gy, o);
// the original's store: network pixel (gx >> 2, gy), channels 4 (gx & 3) .. + 3; row stride W (u32 at +64),
// buffer at +240 (16 halves per pixel)
const uint32_t W = *(const AS4 uint32_t*)(KARGP + 64);
uint16_t* d = (uint16_t*)*(const AS4 uint64_t*)(KARGP + 240) + ((size_t)(W * gy + (gx >> 2)) * 16 + 4 * (gx & 3));
d[0] = o[0];
d[1] = o[1];
d[2] = o[2];
d[3] = o[3];
}
kernels/native/m_ptx_rt.h
+28
View File
@@ -0,0 +1,28 @@
// M texture-translation helpers. Kept separate from K runtime code.
#pragma once
#include "ptx_rt.h"
extern "C" __attribute__((device)) h1 __ocml_tanh_f16(h1);
extern "C" __attribute__((device)) h2 m_min2(h2, h2) __asm("llvm.minimumnum.v2f16");
extern "C" __attribute__((device)) h2 m_max2(h2, h2) __asm("llvm.maximumnum.v2f16");
RT uint32_t prmt32(uint32_t a, uint32_t b, uint32_t select)
{
uint32_t out = 0;
#pragma unroll
for (int i = 0; i < 4; ++i) {
uint32_t s = (select >> (4 * i)) & 15u;
uint32_t byte = ((s & 4u) ? b : a) >> (8 * (s & 3u));
byte = (s & 8u) ? ((byte & 128u) ? 255u : 0u) : (byte & 255u);
out |= byte << (8 * i);
}
return out;
}
RT uint32_t shfl_idx(uint32_t input, uint32_t index, uint32_t opts, bool* in_bounds)
{
uint32_t mask = (opts >> 8) & 31u;
uint32_t subsection = laneid() & mask;
uint32_t end = subsection | (opts & ~mask & 31u);
uint32_t src = subsection | (index & ~mask & 31u);
bool valid = src <= end;
if (in_bounds) *in_bounds = valid;
return __builtin_amdgcn_ds_bpermute(4u * (warp_base() | (valid ? src : laneid())), input);
}
kernels/native/native_rt.h
+185
View File
@@ -0,0 +1,185 @@
// Runtime for natively compiled DLSS texture kernels: PTX semantics (ptx_rt.h) plus surface stores
// (tex/tex_common.h, with ZLUDA's generic raw-store fallback ported here, since no ZLUDA bitcode is linked).
#pragma once
#include "ptx_rt.h"
#include "../tex/tex_common.h"
extern "C" __attribute__((device)) v4f __ockl_image_load_2D(tsharp_t* image, v2i coord);
namespace zsurf {
// zluda_ptx_impl.cpp surface_* helpers (raw sust.b stores), verbatim semantics
RT int channel_bytes(int dt)
{
switch (dt)
{
case 0: case 2: case 8: case 11: return 1;
case 1: case 3: case 9: case 12: case 14: return 2;
case 10: case 13: case 15: return 4;
default: return 0;
}
}
RT int channel_count(int order)
{
switch (order)
{
case 3: case 4: case 5: return 2;
case 6: case 12: return 3;
case 7: case 8: case 9: case 10: case 11: case 13: case 14: case 15: return 4;
default: return 1;
}
}
RT int memory_channel_lane(int order, int channel)
{
switch (order)
{
case 0: return 3;
case 5: return channel == 0 ? 0 : 3;
case 9: case 15: return channel == 0 ? 2 : channel == 2 ? 0 : channel;
case 10: return channel == 0 ? 3 : channel - 1;
case 11: return 3 - channel;
default: return channel;
}
}
RT uint32_t float_to_half_bits(float value)
{
uint32_t bits = __builtin_bit_cast(uint32_t, value);
uint32_t sign = (bits >> 16) & 0x8000u, exponent = (bits >> 23) & 0xffu, mantissa = bits & 0x7fffffu;
if (exponent == 0xff)
return sign | 0x7c00u | (mantissa != 0 ? 0x200u | (mantissa >> 13) : 0u);
int half_exponent = int(exponent) - 112;
if (half_exponent >= 31)
return sign | 0x7c00u;
uint32_t shift = 13;
if (half_exponent <= 0)
{
if (half_exponent < -10)
return sign;
mantissa |= 0x800000u;
shift = uint32_t(14 - half_exponent);
half_exponent = 0;
}
uint32_t half_mantissa = mantissa >> shift, remainder = mantissa & ((1u << shift) - 1u), halfway = 1u << (shift - 1u);
uint32_t result = sign | (uint32_t(half_exponent) << 10);
result += half_mantissa & (half_exponent == 0 ? 0x7ffu : 0x3ffu);
if (remainder > halfway || (remainder == halfway && (half_mantissa & 1u)))
result++;
return result;
}
RT int32_t clampi(int32_t v, int32_t lo, int32_t hi) { return v < lo ? lo : (v > hi ? hi : v); }
RT float decode(uint32_t bits, int dt)
{
switch (dt)
{
case 0: return float(clampi(int32_t(int8_t(bits)), -127, 127)) / 127.0f;
case 1: return float(clampi(int32_t(int16_t(bits)), -32767, 32767)) / 32767.0f;
case 2: return float(int32_t(bits & 0xffu)) / 255.0f;
case 3: return float(int32_t(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);
}
}
RT uint32_t encode(float value, int dt)
{
switch (dt)
{
case 0: return uint32_t(clampi(int32_t(__builtin_floorf(value * 127.0f + 0.5f)), -127, 127)) & 0xffu;
case 1: return uint32_t(clampi(int32_t(__builtin_floorf(value * 32767.0f + 0.5f)), -32767, 32767)) & 0xffffu;
case 2: return uint32_t(clampi(int32_t(__builtin_floorf(value * 255.0f + 0.5f)), 0, 255));
case 3: return uint32_t(clampi(int32_t(__builtin_floorf(value * 65535.0f + 0.5f)), 0, 65535));
case 14: return float_to_half_bits(value);
default: return __builtin_bit_cast(uint32_t, value);
}
}
RT int lane_channel(int order, int lane, int channels)
{
switch (order)
{
case 0: return lane == 3 ? 0 : -1;
case 5: return lane == 0 ? 0 : (lane == 3 ? 1 : -1);
case 9: case 15: return lane == 0 ? 2 : (lane == 2 ? 0 : lane);
case 10: return lane == 3 ? 0 : lane + 1;
case 11: return 3 - lane;
default: return lane < channels ? lane : -1;
}
}
RT uint32_t sel_channel(uint32_t c0, uint32_t c1, uint32_t c2, uint32_t c3, int ch)
{
return ch == 0 ? c0 : (ch == 1 ? c1 : (ch == 2 ? c2 : c3));
}
RT float sel_lane(v4f v, int lane) { return lane == 0 ? v.x : (lane == 1 ? v.y : (lane == 2 ? v.z : v.w)); }
RT uint32_t merge_channel(uint32_t bits, int channel, int cb, uint32_t data, int byte_offset, int byte_count)
{
for (int b = 0; b < 4; b++)
{
int relative = channel * cb + b - byte_offset;
if (b < cb && relative >= 0 && relative < byte_count)
{
uint32_t byte = (data >> (8 * relative)) & 0xffu;
bits = (bits & ~(0xffu << (8 * b))) | (byte << (8 * b));
}
}
return bits;
}
RT float lane_value(uint32_t c0, uint32_t c1, uint32_t c2, uint32_t c3, int order, int lane, int channels, int dt)
{
int ch = lane_channel(order, lane, channels);
return ch < 0 ? 0.0f : decode(sel_channel(c0, c1, c2, c3, ch), dt);
}
__attribute__((device, noinline)) static void store_raw_2d(tsharp_t* image, int32_t cx, int32_t cy, uint32_t data, int byte_count)
{
int dt = __ockl_image_channel_data_type_2D(image), order = __ockl_image_channel_order_2D(image);
int cb = channel_bytes(dt);
if (cb == 0 || cx < 0)
return;
int channels = channel_count(order), pixel_bytes = cb * channels;
for (int consumed = 0; consumed < byte_count;)
{
int byte_x = cx + consumed, pixel_x = byte_x / pixel_bytes, byte_offset = byte_x - pixel_x * pixel_bytes;
int chunk = pixel_bytes - byte_offset;
if (chunk > byte_count - consumed)
chunk = byte_count - consumed;
uint32_t fragment = data >> (8 * consumed);
uint32_t c0 = 0, c1 = 0, c2 = 0, c3 = 0;
if (byte_offset != 0 || chunk < pixel_bytes)
{
v4f loaded = __ockl_image_load_2D(image, (v2i){pixel_x, cy});
c0 = encode(sel_lane(loaded, memory_channel_lane(order, 0)), dt);
c1 = encode(sel_lane(loaded, memory_channel_lane(order, 1)), dt);
c2 = encode(sel_lane(loaded, memory_channel_lane(order, 2)), dt);
c3 = encode(sel_lane(loaded, memory_channel_lane(order, 3)), dt);
}
c0 = merge_channel(c0, 0, cb, fragment, byte_offset, chunk);
c1 = merge_channel(c1, 1, cb, fragment, byte_offset, chunk);
c2 = merge_channel(c2, 2, cb, fragment, byte_offset, chunk);
c3 = merge_channel(c3, 3, cb, fragment, byte_offset, chunk);
__ockl_image_store_2D(image, (v2i){pixel_x, cy},
(v4f){lane_value(c0, c1, c2, c3, order, 0, channels, dt), lane_value(c0, c1, c2, c3, order, 1, channels, dt),
lane_value(c0, c1, c2, c3, order, 2, channels, dt), lane_value(c0, c1, c2, c3, order, 3, channels, dt)});
consumed += chunk;
}
}
} // namespace zsurf
extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_b32_zero(uint64_t surface, v2i coord, uint32_t value)
{
zsurf::store_raw_2d((tsharp_t*)(uintptr_t)surface, coord.x, coord.y, value, 4);
}
extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t surface, v2i coord, v2s value)
{
zsurf::store_raw_2d((tsharp_t*)(uintptr_t)surface, coord.x, coord.y, uint32_t(uint16_t(value.x)) | (uint32_t(uint16_t(value.y)) << 16), 4);
}
// A sample that every lane takes with the same handle, coordinates and lod (e.g. a 1x1 exposure texture at (0.5,
// 0.5)): only the first active lane samples and the result is broadcast with readfirstlane. Same values, but the
// texture unit processes one lane instead of the whole wave.
RT f4 tex_level_uniform(uint64_t t, float x, float y, float lod)
{
f4 r = (f4){0.0f, 0.0f, 0.0f, 0.0f};
if (wave_lane() == (uint32_t)__builtin_ctzll(__builtin_amdgcn_read_exec()))
r = tex_level(t, x, y, lod);
return (f4){asf(__builtin_amdgcn_readfirstlane(asu(r[0]))), asf(__builtin_amdgcn_readfirstlane(asu(r[1]))),
asf(__builtin_amdgcn_readfirstlane(asu(r[2]))), asf(__builtin_amdgcn_readfirstlane(asu(r[3])))};
}
kernels/native/output_tile.h
+71
View File
@@ -0,0 +1,71 @@
// Structured path-B tile fill. Same color/depth values and 24-column LDS layout
// as the original kernel; work-items traverse x/y directly without integer division.
#pragma once
#ifndef D4R_OUTPUT_TILE_PITCH
#define D4R_OUTPUT_TILE_PITCH 24
#endif
static_assert(D4R_OUTPUT_TILE_PITCH >= 24 && D4R_OUTPUT_TILE_PITCH <= 32, "LDS tile pitch capacity");
RT void output_color_pack(f4 color, float exposure, uint32_t* a, uint32_t* b)
{
const float r = fmaxn(color[0], 0.0f), g = fmaxn(color[1], 0.0f), blue = fmaxn(color[2], 0.0f);
const float qr = r * 0.25f, hg = g * 0.5f, qb = blue * 0.25f;
const float y = exposure * ((qr + hg) + qb);
const float co = __builtin_fmaf(r, 0.5f, g * 0.0f) - blue * 0.5f;
const float cg = (hg - qr) - qb;
float weight;
if (y < 1000.0f)
weight = rcp_a(y + 1.0f);
else
{
const float z = __builtin_fmaf(y, 9.0f, 1.0f);
const float ln = lg2_a(z) * asf(0x3F317218u);
const float q = div_a(ln, 9.0f) + asf(0xBC4FAE29u);
weight = rcp_a(y) * q;
}
*a = pack16(ashu((h1)(y * weight)), ashu((h1)((exposure * co) * weight)));
*b = pack16(ashu((h1)((exposure * cg) * weight)), ashu((h1)color[3]));
}
RT int output_tile_coord(int relative, int lo, int hi)
{
return relative > 0 ? __builtin_elementwise_min(lo + relative, hi) : lo;
}
RT void output_tile_fill(const AS4 uint8_t* args, int x0, int y0, int width, int height,
float exposure, AS3 uint8_t* color_tile, AS3 uint8_t* depth_tile)
{
const uint64_t color = *(const AS4 uint64_t*)(args + 200);
const uint64_t depth = *(const AS4 uint64_t*)(args + 248);
const int cx0 = *(const AS4 int32_t*)(args + 208), cy0 = *(const AS4 int32_t*)(args + 212);
const int cx1 = *(const AS4 int32_t*)(args + 216), cy1 = *(const AS4 int32_t*)(args + 220);
const int dx0 = *(const AS4 int32_t*)(args + 96), dy0 = *(const AS4 int32_t*)(args + 100);
const int dx1 = *(const AS4 int32_t*)(args + 104), dy1 = *(const AS4 int32_t*)(args + 108);
bool fill_depth = true;
#ifdef D4R_K_NO_DEPTH_TILE
// Pinned K LUT has |slope|<=1. Its 4x3 color stencil stays in the four-texel
// halo for positive upsampling scale<=1, so no color read aliases depth LDS.
// Keep the original fill outside the proven geometry. Input depth is unchanged.
const float scale_x = *(const AS4 float*)(args + 40), scale_y = *(const AS4 float*)(args + 44);
fill_depth = !(scale_x > 0.0f && scale_x <= 1.0f && scale_y > 0.0f && scale_y <= 1.0f);
#endif
// Path B selects a footprint at most 24 x 24. At each iteration 0<=x<width,
// 0<=y<height: color accesses stay within 4608 bytes, depth within 2304 bytes.
for (int y = __builtin_amdgcn_workitem_id_y(); y < height; y += 16)
for (int x = __builtin_amdgcn_workitem_id_x(); x < width; x += 16)
{
const int sx = x0 + x, sy = y0 + y;
const f4 c = tex_fetch(color, output_tile_coord(sx, cx0, cx1), output_tile_coord(sy, cy0, cy1));
uint32_t a, b;
output_color_pack(c, exposure, &a, &b);
AS3 uint32_t* dst = (AS3 uint32_t*)(color_tile + y * (8 * D4R_OUTPUT_TILE_PITCH) + x * 8);
dst[0] = a;
dst[1] = b;
if (fill_depth)
{
const f4 d = tex_fetch(depth, output_tile_coord(sx, dx0, dx1), output_tile_coord(sy, dy0, dy1));
*(AS3 float*)(depth_tile + y * (4 * D4R_OUTPUT_TILE_PITCH) + x * 4) = d[0];
}
}
}
kernels/native/ptx2hip.py
Executable
+857
View File
@@ -0,0 +1,857 @@
#!/usr/bin/env python3
"""Translate one NVIDIA PTX kernel (DLSS texture kernels) into native HIP C++.
usage: ptx2hip.py MODULE.ptx KERNEL > KERNEL.hip
The output is a plain HIP kernel with the same name and parameter block. Every PTX instruction becomes
the C expression of ZLUDA's lowering of it (see ptx_rt.h), so the translation computes what ZLUDA's own
compile of the PTX computes, but it is compiled natively by clang and can then be optimised as C++.
Registers become locals, basic blocks become labels, branches become gotos.
"""
import re
import sys
TYPES = {'f': 'float', 'r': 'uint32_t', 'rd': 'uint64_t', 'rs': 'uint16_t', 'p': 'bool', 'fd': 'double',
'SP': 'uint64_t', 'SPL': 'uint64_t'}
SREG = {'%tid.x': '__builtin_amdgcn_workitem_id_x()', '%tid.y': '__builtin_amdgcn_workitem_id_y()',
'%tid.z': '__builtin_amdgcn_workitem_id_z()', '%ctaid.x': '__builtin_amdgcn_workgroup_id_x()',
'%ctaid.y': '__builtin_amdgcn_workgroup_id_y()', '%ctaid.z': '__builtin_amdgcn_workgroup_id_z()',
'%laneid': 'laneid()',
'%ntid.x': '__builtin_amdgcn_workgroup_size_x()', '%ntid.y': '__builtin_amdgcn_workgroup_size_y()',
'%nctaid.x': '(__builtin_amdgcn_grid_size_x() / __builtin_amdgcn_workgroup_size_x())',
'%nctaid.y': '(__builtin_amdgcn_grid_size_y() / __builtin_amdgcn_workgroup_size_y())'}
REG_RE = re.compile(r'%([a-zA-Z]+)(\d+)$')
class Err(Exception):
pass
def strip_comments(text):
return re.sub(r'//[^\n]*', '', text)
def tokenize(body):
"""-> list of ('open'|'close'|'label'|'stmt', text)"""
out, i, n = [], 0, len(body)
while i < n:
c = body[i]
if c.isspace():
i += 1
continue
if c == '{':
out.append(('open', None)); i += 1; continue
if c == '}':
out.append(('close', None)); i += 1; continue
m = re.match(r'(\$?[A-Za-z_][\w$]*):', body[i:])
if m and not body[i:].startswith('%'):
out.append(('label', m.group(1))); i += m.end(); continue
j = body.index(';', i)
out.append(('stmt', ' '.join(body[i:j].split())))
i = j + 1
return out
def split_operands(s):
ops, depth, cur = [], 0, ''
for ch in s:
if ch in '{[':
depth += 1
elif ch in '}]':
depth -= 1
if ch == ',' and depth == 0:
ops.append(cur.strip()); cur = ''
else:
cur += ch
if cur.strip():
ops.append(cur.strip())
return ops
class Translator:
def __init__(self, ptx, kernel):
self.kernel = kernel
text = strip_comments(ptx)
m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S)
if not m:
raise Err('kernel not found')
params = m.group(1)
pm = re.search(r'\.param\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', params)
self.param_name, self.param_size, self.param_align = pm.group(2), int(pm.group(3)), int(pm.group(1))
rest = text[m.end():]
self.attrs = rest[:rest.index('{')]
start = rest.index('{')
depth, k = 0, start
while True:
if rest[k] == '{':
depth += 1
elif rest[k] == '}':
depth -= 1
if depth == 0:
break
k += 1
self.body = rest[start + 1:k]
self.globals_text = text[:m.start()]
self.regs = {} # c name -> type
self.shared = [] # (name, size, align)
self.locals = []
self.scope_names = [{}]
self.scope_id = 0
self.lines = []
# ------------------------------------------------------------ operands
def reg(self, tok):
tok = tok.strip()
for scope in reversed(self.scope_names):
if tok in scope:
return scope[tok]
m = REG_RE.match(tok)
if not m:
raise Err('not a register: ' + tok)
kind, num = m.groups()
name = f'{kind}{num}'
if kind not in TYPES:
raise Err('register kind ' + kind)
self.regs[name] = TYPES[kind]
return name
def rtype(self, name):
return self.regs[name]
def val(self, tok, as_type):
"""operand as a C expression of type as_type ('f32','u32','s32','u64','s64','u16','s16','pred','f16x2')"""
tok = tok.strip()
if tok in SREG:
return f'(uint32_t){SREG[tok]}'
if tok.startswith('%') or any(tok in s for s in self.scope_names):
name = self.reg(tok)
t = self.rtype(name)
if as_type == 'f32':
return name if t == 'float' else f'asf({name})'
if as_type in ('u32', 'b32'):
return f'asu({name})' if t == 'float' else name
if as_type == 's32':
return f'(int32_t)asu({name})' if t == 'float' else f'(int32_t){name}'
if as_type in ('u64', 'b64'):
return name
if as_type == 's64':
return f'(int64_t){name}'
if as_type in ('u16', 'b16'):
return name
if as_type == 's16':
return f'(int16_t){name}'
if as_type == 'f16':
return f'ash({name})'
if as_type == 'f16x2':
return f'ash2({name})'
if as_type == 'pred':
return name
raise Err('type ' + as_type)
if tok in self.shared_names():
return f'(uint32_t)(uintptr_t)(AS3 uint8_t*){tok}'
if tok == '__local_depot0':
# private (address space 5) address of the depot, zero-extended: ld/st.local use it as an AS5 pointer
return '(uint64_t)(uint32_t)(uintptr_t)(AS5 uint8_t*)__local_depot0'
# immediates
if re.match(r'0[fF][0-9A-Fa-f]{8}$', tok):
bits = int(tok[2:], 16)
return f'asf(0x{bits:08X}u)' if as_type == 'f32' else f'0x{bits:08X}u'
if re.match(r'-?\d+$', tok) or re.match(r'0x[0-9a-fA-F]+$', tok):
v = int(tok, 0)
if as_type == 'f32':
raise Err('int immediate as float')
if as_type in ('u64', 's64', 'b64'):
return f'(uint64_t){v}ull' if v >= 0 else f'(uint64_t)({v}ll)'
if as_type in ('u16', 'b16', 's16'):
return f'(uint16_t){v & 0xffff}'
if as_type == 's32':
v = v & 0xffffffff
v = v - (1 << 32) if v >= (1 << 31) else v
return f'(int32_t){v}' if v != -(1 << 31) else '(int32_t)0x80000000u'
return f'{v & 0xffffffff}u'
raise Err('operand ' + tok)
def shared_names(self):
return [s[0] for s in self.shared]
def assign(self, dst, expr, src_type):
"""dst = expr, expr of C type src_type ('float','u32','u64','u16','bool','f16x2','f16')"""
name = self.reg(dst)
t = self.rtype(name)
if src_type == 'float':
e = expr if t == 'float' else f'asu({expr})'
elif src_type == 'u32':
e = f'asf({expr})' if t == 'float' else f'(uint32_t)({expr})'
elif src_type == 'u64':
e = f'(uint64_t)({expr})'
elif src_type == 'u16':
e = f'(uint16_t)({expr})'
elif src_type == 'f16':
e = f'ashu({expr})'
elif src_type == 'f16x2':
e = f'ash2u({expr})'
elif src_type == 'bool':
e = f'({expr})'
else:
raise Err('src type')
self.emit(f'{name} = {e};')
def emit(self, s):
self.lines.append(' ' + s)
def addr(self, tok):
"""[reg+off] / [name+off] -> (base_expr, offset_int, space_hint)"""
m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok)
if not m:
raise Err('address ' + tok)
base, off = m.group(1), int(m.group(2) or 0)
if base == self.param_name:
return 'KARG', off
if base in self.shared_names():
return f'(uint32_t)(uintptr_t)(AS3 uint8_t*){base}', off
return self.val(base, 'u64' if not base.startswith('%r') or base.startswith('%rd') else 'u32'), off
# ------------------------------------------------------------ instructions
def translate(self):
for kind, text in tokenize(self.body):
if kind == 'open':
self.scope_names.append({})
continue
if kind == 'close':
self.scope_names.pop()
continue
if kind == 'label':
self.lines.append(f'{text.replace("$", "")}:;')
lab = text.replace('$', '')
if lab in getattr(self, 'stamps', []):
self.lines.append(f' D4R_STAMP({self.stamps.index(lab) + 1});')
if getattr(self, 'bbcount', False):
self.bblabels.append(lab)
self.lines.append(f' D4R_BBC({len(self.bblabels) - 1});')
continue
self.stmt(text)
def stmt(self, s):
if s.startswith('.reg'):
m = re.match(r'\.reg\s+\.(\w+)\s+(.*)$', s)
ty, names = m.groups()
if '<' in names:
return # numbered register banks: declared on use
self.scope_id += 1
for nm in names.split(','):
nm = nm.strip()
cname = f'l{self.scope_id}_{nm.lstrip("%")}'
ctype = {'f16': 'uint16_t', 'b16': 'uint16_t', 'f32': 'float', 'b32': 'uint32_t', 'pred': 'bool',
'f16x2': 'uint32_t', 'u32': 'uint32_t', 's32': 'uint32_t', 'b64': 'uint64_t', 'u64': 'uint64_t'}[ty]
self.scope_names[-1][nm] = cname
self.regs[cname] = ctype
return
if s.startswith('.shared'):
m = re.match(r'\.shared\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
self.shared.append((m.group(2), int(m.group(3)), int(m.group(1))))
return
if s.startswith('.local'):
m = re.match(r'\.local\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
self.locals.append((m.group(2), int(m.group(3)), int(m.group(1))))
return
if s.startswith('.pragma'):
return
pred = None
m = re.match(r'@(!?)(%p\d+)\s+(.*)$', s)
if m:
pred = ('!' if m.group(1) else '') + self.reg(m.group(2))
s = m.group(3)
parts = s.split(None, 1)
op = parts[0]
args = split_operands(parts[1]) if len(parts) > 1 else []
if pred is not None:
if op.startswith('bra'):
self.emit(f'if ({pred}) goto {args[0].replace("$", "")};')
return
self.emit(f'if ({pred}) {{')
self.op(op, args)
self.emit('}')
return
self.op(op, args)
def op(self, op, a):
o = op.split('.')
base = o[0]
mods = set(o[1:])
E, A = self.emit, self.assign
V = self.val
def ftype():
for t in ('f32', 'f16x2', 'f16', 'f64'):
if t in mods:
return t
return None
if base == 'bra':
E(f'goto {a[0].replace("$", "")};'); return
if base == 'ret':
if getattr(self, 'stamps', None):
E(f'D4R_STAMP({len(self.stamps) + 1});')
E('return;'); return
if base == 'bar':
E('bar_sync();'); return
if base == 'mov':
dst, src = a
if dst.startswith('{'):
names = [x.strip() for x in dst[1:-1].split(',')]
srcv = V(src, 'u32')
A(names[0], f'{srcv} & 0xffffu', 'u16')
A(names[1], f'{srcv} >> 16', 'u16')
return
if src.startswith('{'):
names = [x.strip() for x in src[1:-1].split(',')]
if 'b64' in mods:
A(dst, f'(uint64_t){V(names[0], "u32")} | ((uint64_t){V(names[1], "u32")} << 32)', 'u64'); return
A(dst, f'pack16({V(names[0], "u16")}, {V(names[1], "u16")})', 'u32'); return
if src == self.param_name:
A(dst, 'KARGU', 'u64'); return
if 'f32' in mods:
A(dst, V(src, 'f32'), 'float'); return
if mods & {'u64', 'b64', 's64'}:
A(dst, V(src, 'u64'), 'u64'); return
if mods & {'u16', 'b16', 's16'}:
A(dst, V(src, 'u16'), 'u16'); return
A(dst, V(src, 'u32'), 'u32'); return
ft = ftype()
# ---------------------------------------------------- float arithmetic
if ft == 'f32' and base in ('add', 'sub', 'mul'):
if 'rz' in mods:
raise Err('add.rz outside the roundf idiom')
sym = {'add': '+', 'sub': '-', 'mul': '*'}[base]
A(a[0], f'{V(a[1], "f32")} {sym} {V(a[2], "f32")}', 'float'); return
if ft == 'f32' and base == 'fma':
A(a[0], f'__builtin_fmaf({V(a[1], "f32")}, {V(a[2], "f32")}, {V(a[3], "f32")})', 'float'); return
if ft == 'f32' and base in ('max', 'min'):
if 'NaN' in mods:
raise Err('NaN min/max')
A(a[0], f'{"fmaxn" if base == "max" else "fminn"}({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return
if ft == 'f32' and base == 'neg':
A(a[0], f'neg_ftz({V(a[1], "f32")})' if 'ftz' in mods else f'-{V(a[1], "f32")}', 'float'); return
if ft == 'f32' and base == 'abs':
A(a[0], f'abs_ftz({V(a[1], "f32")})' if 'ftz' in mods else f'__builtin_fabsf({V(a[1], "f32")})', 'float'); return
if ft == 'f32' and base in ('rcp', 'ex2', 'lg2', 'rsqrt', 'sqrt'):
if 'approx' not in mods:
raise Err(op)
fn = {'rcp': 'rcp_a', 'ex2': 'ex2_a', 'lg2': 'lg2_a', 'rsqrt': 'rsq_a', 'sqrt': '__builtin_amdgcn_sqrtf'}[base]
A(a[0], f'{fn}({V(a[1], "f32")})', 'float'); return
if ft == 'f32' and base == 'div':
if 'approx' not in mods:
raise Err(op)
A(a[0], f'div_a({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return
# ---------------------------------------------------- f16 / f16x2
if ft == 'f16x2' and base in ('add', 'sub', 'mul', 'fma'):
x = [V(t, 'f16x2') for t in a[1:]]
e = f'h2fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})'
if 'sat' in mods:
e = f'h2sat({e})'
A(a[0], e, 'f16x2'); return
if ft == 'f16x2' and base == 'abs':
A(a[0], f'h2abs({V(a[1], "f16x2")})', 'f16x2'); return
if ft == 'f16x2' and base == 'ex2':
A(a[0], f'h2ex2({V(a[1], "f16x2")})', 'f16x2'); return
if ft in ('f16', 'f16x2') and base in ('min', 'max'):
if 'NaN' in mods:
raise Err('NaN-propagating half min/max is not supported')
fn = 'hminn' if base == 'min' else 'hmaxn'
if ft == 'f16':
A(a[0], f'{fn}({V(a[1], "f16")}, {V(a[2], "f16")})', 'f16'); return
x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2')
A(a[0], f'(h2){{{fn}({x}.x, {y}.x), {fn}({x}.y, {y}.y)}}', 'f16x2'); return
if ft == 'f16' and base in ('add', 'sub', 'mul', 'fma'):
x = [V(t, 'f16') for t in a[1:]]
e = f'hfma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'(h1)({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})'
if 'sat' in mods:
e = f'hsat({e})'
A(a[0], e, 'f16'); return
if op == 'tanh.approx.f16':
A(a[0], f'__ocml_tanh_f16({V(a[1], "f16")})', 'f16'); return
if base == 'set' and 'f16x2' in mods and 'u32' in mods:
cmp = self.cmp_op(mods)
x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2')
A(a[0], f'pack16(({cmp.format(f"{x}.x", f"{y}.x")}) ? 0xffffu : 0u, ({cmp.format(f"{x}.y", f"{y}.y")}) ? 0xffffu : 0u)', 'u32'); return
# ---------------------------------------------------- compares / selects
if base == 'setp':
cmp = self.cmp_op(mods)
if mods & {'f32'}:
x, y = V(a[1], 'f32'), V(a[2], 'f32')
elif 'f16' in mods:
x, y = V(a[1], 'f16'), V(a[2], 'f16')
elif mods & {'s32'}:
x, y = V(a[1], 's32'), V(a[2], 's32')
elif mods & {'u32', 'b32'}:
x, y = V(a[1], 'u32'), V(a[2], 'u32')
elif mods & {'s16'}:
x, y = V(a[1], 's16'), V(a[2], 's16')
elif mods & {'u16', 'b16'}:
x, y = V(a[1], 'u16'), V(a[2], 'u16')
elif mods & {'s64'}:
x, y = V(a[1], 's64'), V(a[2], 's64')
elif mods & {'u64'}:
x, y = V(a[1], 'u64'), V(a[2], 'u64')
else:
raise Err(op)
if len(a) != 3 or '|' in a[0]:
raise Err('setp form ' + op)
A(a[0], cmp.format(x, y), 'bool'); return
if base == 'selp':
t = 'f32' if 'f32' in mods else 'u16' if mods & {'u16', 'b16', 's16'} else 'u64' if mods & {'u64', 'b64', 's64'} else 'u32'
ct = {'f32': 'float', 'u16': 'u16', 'u64': 'u64', 'u32': 'u32'}[t]
A(a[0], f'{V(a[3], "pred")} ? {V(a[1], t)} : {V(a[2], t)}', ct); return
# ---------------------------------------------------- predicates / integer
if 'pred' in mods and base in ('and', 'or', 'xor', 'not'):
if base == 'not':
A(a[0], f'!{V(a[1], "pred")}', 'bool'); return
sym = {'and': '&&', 'or': '||', 'xor': '!='}[base]
A(a[0], f'{V(a[1], "pred")} {sym} {V(a[2], "pred")}', 'bool'); return
if base in ('and', 'or', 'xor', 'not', 'shl', 'shr', 'add', 'sub', 'mul', 'mad', 'min', 'max', 'div', 'rem', 'neg', 'abs'):
if ft is not None:
raise Err('unhandled floating operation: ' + op)
width = 64 if mods & {'u64', 's64', 'b64'} else 16 if mods & {'u16', 's16', 'b16'} else 32
signed = bool(mods & {'s32', 's64', 's16'})
ut = {64: 'u64', 32: 'u32', 16: 'u16'}[width]
st = {64: 's64', 32: 's32', 16: 's16'}[width]
ct = {64: 'u64', 32: 'u32', 16: 'u16'}[width]
if base == 'mul' and 'hi' in mods and width == 32:
wide = 'int64_t' if signed else 'uint64_t'
typ = 's32' if signed else 'u32'
A(a[0], f'(({wide}){V(a[1], typ)} * ({wide}){V(a[2], typ)}) >> 32', 'u32'); return
if base == 'mul' and 'wide' in mods:
if mods & {'s32'}:
A(a[0], f'(int64_t){V(a[1], "s32")} * (int64_t){V(a[2], "s32")}', 'u64'); return
if mods & {'u32'}:
A(a[0], f'(uint64_t){V(a[1], "u32")} * (uint64_t){V(a[2], "u32")}', 'u64'); return
if mods & {'u16'}:
A(a[0], f'(uint32_t){V(a[1], "u16")} * (uint32_t){V(a[2], "u16")}', 'u32'); return
raise Err(op)
if base == 'not':
A(a[0], f'~{V(a[1], ut)}', ct); return
if base in ('neg', 'abs'):
if base == 'abs' and signed:
A(a[0], f'{V(a[1], st)} < 0 ? (uint{width}_t)(0 - {V(a[1], ut)}) : {V(a[1], ut)}', ct); return
if base == 'neg':
A(a[0], f'(uint{width}_t)(0 - {V(a[1], ut)})', ct); return
raise Err(op)
x = V(a[1], st if signed else ut)
if base == 'shl':
A(a[0], f'{V(a[1], ut)} << ({V(a[2], "u32")} & {width - 1})', ct); return
if base == 'shr':
if signed:
A(a[0], f'{V(a[1], st)} >> ({V(a[2], "u32")} & {width - 1})', ct); return
A(a[0], f'{V(a[1], ut)} >> ({V(a[2], "u32")} & {width - 1})', ct); return
y = V(a[2], st if signed else ut)
if base in ('and', 'or', 'xor'):
sym = {'and': '&', 'or': '|', 'xor': '^'}[base]
A(a[0], f'{V(a[1], ut)} {sym} {V(a[2], ut)}', ct); return
if base in ('add', 'sub'):
A(a[0], f'{V(a[1], ut)} {"+" if base == "add" else "-"} {V(a[2], ut)}', ct); return
if base == 'mul':
if 'lo' not in mods:
raise Err(op)
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)}', ct); return
if base == 'mad':
if 'lo' not in mods:
raise Err(op)
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)} + {V(a[3], ut)}', ct); return
if base in ('min', 'max'):
fn = 'min' if base == 'min' else 'max'
A(a[0], f'__builtin_elementwise_{fn}({x}, {y})', ct); return
if base == 'div':
A(a[0], f'{x} / {y}', ct); return
if base == 'rem':
A(a[0], f'{x} % {y}', ct); return
if base == 'bfi':
A(a[0], f'bfi32({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")}, {V(a[4], "u32")})', 'u32'); return
if op == 'prmt.b32':
A(a[0], f'prmt32({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")})', 'u32'); return
# ---------------------------------------------------- conversions
if base == 'cvt':
return self.cvt(op, o[1:], a)
if base == 'cvta':
A(a[0], V(a[1], 'u64'), 'u64'); return
# ---------------------------------------------------- memory
if base == 'ld':
return self.ld(o[1:], a)
if base == 'st':
return self.st(o[1:], a)
if base == 'tex':
dst = [x.strip() for x in a[0][1:-1].split(',')]
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', a[1])
h = V(m.group(1), 'u64')
c = [x.strip() for x in m.group(2).split(',')]
if 'base' in mods and mods >= {'2d', 'v4', 's32'} and mods & {'f32', 'f16'}:
E(f'{{ ::f4 t_ = tex_fetch({h}, (int32_t){V(c[0], "u32")}, (int32_t){V(c[1], "u32")});')
elif mods >= {'level', '2d', 'v4', 'f32'} and op.endswith('f32.f32'):
E(f'{{ ::f4 t_ = tex_level({h}, {V(c[0], "f32")}, {V(c[1], "f32")}, {V(a[2], "f32")});')
else:
raise Err(op)
for i, d in enumerate(dst):
A(d, f'(h1)t_[{i}]', 'f16') if 'f16' in mods else A(d, f't_[{i}]', 'float')
E('}')
return
if base == 'tld4':
if not (mods >= {'r', '2d', 'v4', 'f32'}):
raise Err(op)
dst = [x.strip() for x in a[0][1:-1].split(',')]
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', a[1])
c = [x.strip() for x in m.group(2).split(',')]
E(f'{{ ::f4 t_ = tex_gather_r({V(m.group(1), "u64")}, {V(c[0], "f32")}, {V(c[1], "f32")});')
for i, d in enumerate(dst):
A(d, f't_[{i}]', 'float')
E('}')
return
if base == 'sust':
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', a[0])
h = V(m.group(1), 'u64')
c = [V(x.strip(), 'u32') for x in m.group(2).split(',')]
v = [V(x.strip(), 'u32') for x in a[1][1:-1].split(',')]
if op == 'sust.p.2d.v4.b32.zero':
E(f'd4r_sust_p_v4b32({h}, (int32_t){c[0]}, (int32_t){c[1]}, {v[0]}, {v[1]}, {v[2]}, {v[3]});'); return
if op == 'sust.b.2d.b32.zero':
E(f'd4r_sust_b32({h}, (int32_t){c[0]}, (int32_t){c[1]}, {v[0]});'); return
if op == 'sust.b.2d.v2.b16.zero':
v = [V(x.strip(), 'u16') for x in a[1][1:-1].split(',')]
E(f'd4r_sust_v2b16({h}, (int32_t){c[0]}, (int32_t){c[1]}, {v[0]}, {v[1]});'); return
raise Err(op)
if base == 'shfl':
if op == 'shfl.sync.idx.b32':
d = a[0].split('|')
pred = f'&{self.reg(d[1])}' if len(d) > 1 else 'nullptr'
A(d[0], f'shfl_idx({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")}, {pred})', 'u32'); return
if op != 'shfl.sync.bfly.b32':
raise Err(op)
d = a[0].split('|')
p = f'&{self.reg(d[1])}' if len(d) > 1 else 'nullptr'
A(d[0], f'shfl_bfly({V(a[1], "u32")}, (int32_t){V(a[2], "u32")}, {V(a[3], "u32")}, {p})', 'u32'); return
if op == 'vote.sync.ballot.b32':
A(a[0], f'(uint32_t)(__builtin_amdgcn_ballot_w64({V(a[1], "pred")}) >> (wave_lane() & 32u))', 'u32'); return
raise Err('unhandled ' + op)
def cmp_op(self, mods):
table = {'lt': '{} < {}', 'gt': '{} > {}', 'le': '{} <= {}', 'ge': '{} >= {}', 'eq': '{} == {}',
'ne': '{} != {}',
# unordered: true if either is NaN
'ltu': '!({} >= {})', 'gtu': '!({} <= {})', 'leu': '!({} > {})', 'geu': '!({} < {})',
'equ': '!({} != {})', 'neu': '!({} == {})'}
for k in table:
if k in mods:
if k in ('ne', 'equ') and mods & {'f32', 'f16', 'f16x2', 'f64'}:
# PTX setp.ne.f32 is ordered (false on NaN) and equ is unordered (true on NaN), while C's
# != is unordered: spell both out. Integer ne is plain.
return '({0} < {1} || {0} > {1})' if k == 'ne' else '!({0} < {1} || {0} > {1})'
return table[k]
raise Err('compare ' + str(mods))
def cvt(self, op, o, a):
A, V = self.assign, self.val
mods = set(o)
if op == 'cvt.rn.satfinite.e4m3x2.f16x2':
A(a[0], f'e4m3x2({V(a[1], "u32")})', 'u16'); return
if op == 'cvt.rn.f16x2.e4m3x2':
A(a[0], f'd4r_decode_e4m3x2({V(a[1], "u16")})', 'u32'); return
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
dt, stp = types[0], types[1]
if dt == 'f32' and stp == 'f32':
x = V(a[1], 'f32')
if 'sat' in mods and 'ftz' in mods and not (mods & {'rzi', 'rmi', 'rni', 'rpi'}):
A(a[0], f'sat_ftz({x})', 'float'); return
if mods & {'sat'}:
raise Err(op)
if 'rzi' in mods:
A(a[0], f'__builtin_truncf({x})', 'float'); return
if 'rmi' in mods:
A(a[0], f'__builtin_floorf({x})', 'float'); return
if 'rni' in mods:
A(a[0], f'__builtin_rintf({x})', 'float'); return
if 'rpi' in mods:
A(a[0], f'__builtin_ceilf({x})', 'float'); return
if 'ftz' in mods:
A(a[0], f'canon({x})', 'float'); return
A(a[0], x, 'float'); return
if dt == 's32' and stp == 'f32':
x = V(a[1], 'f32')
if 'rzi' in mods:
A(a[0], f'(uint32_t)f2i_rz({x})', 'u32'); return
if 'rmi' in mods:
A(a[0], f'(uint32_t)f2i_rz(__builtin_floorf({x}))', 'u32'); return
if 'rni' in mods:
A(a[0], f'(uint32_t)f2i_rz(__builtin_rintf({x}))', 'u32'); return
raise Err(op)
if dt == 'u32' and stp == 'f32' and 'rzi' in mods:
A(a[0], f'f2u_rz({V(a[1], "f32")})', 'u32'); return
if dt == 'f32' and stp == 's32':
A(a[0], f'(float)(int32_t){V(a[1], "u32")}', 'float'); return
if dt == 'f32' and stp == 'u32':
A(a[0], f'(float){V(a[1], "u32")}', 'float'); return
if dt == 'f32' and stp == 'f16':
A(a[0], f'(float){V(a[1], "f16")}', 'float'); return
if dt == 'f16' and stp == 'f32':
if 'rn' not in mods:
raise Err(op)
A(a[0], f'(h1){V(a[1], "f32")}', 'f16'); return
if dt == 'f16x2' and stp == 'f32':
# d = {cvt(a) in the high half, cvt(b) in the low half}
A(a[0], f'pack16(ashu((h1){V(a[2], "f32")}), ashu((h1){V(a[1], "f32")}))', 'u32'); return
if dt == 'f16' and stp == 'f16':
x = V(a[1], 'f16')
if 'rni' in mods:
A(a[0], f'h_roundeven({x})', 'f16'); return
if 'rmi' in mods:
A(a[0], f'h_floor({x})', 'f16'); return
raise Err(op)
if dt == 's32' and stp == 'f16':
x = V(a[1], 'f16')
if 'rni' in mods:
A(a[0], f'(uint32_t)h2i_sat(h_roundeven({x}))', 'u32'); return
if 'rmi' in mods:
A(a[0], f'(uint32_t)h2i_sat(h_floor({x}))', 'u32'); return
raise Err(op)
if dt == 'f16' and stp == 's32':
A(a[0], f'(h1)(int32_t){V(a[1], "u32")}', 'f16'); return
if dt == 's64' and stp == 's32':
A(a[0], f'(uint64_t)(int64_t)(int32_t){V(a[1], "u32")}', 'u64'); return
if dt == 'u64' and stp == 'u32':
A(a[0], f'(uint64_t){V(a[1], "u32")}', 'u64'); return
if dt in ('u16', 's16') and stp in ('u32', 's32'):
A(a[0], f'(uint16_t){V(a[1], "u32")}', 'u16'); return
if dt in ('u32', 's32') and stp in ('u16',):
A(a[0], f'(uint32_t){V(a[1], "u16")}', 'u32'); return
if dt in ('u32', 's32') and stp in ('s16',):
A(a[0], f'(uint32_t)(int32_t)(int16_t){V(a[1], "u16")}', 'u32'); return
if dt == 'u32' and stp == 'u64':
A(a[0], f'(uint32_t){V(a[1], "u64")}', 'u32'); return
raise Err('cvt ' + op)
def mem_type(self, o):
for t in ('u8', 's8', 'u16', 's16', 'b16', 'u32', 's32', 'b32', 'f32', 'u64', 's64', 'b64', 'f16'):
if t in o:
return t
raise Err('mem type ' + str(o))
def ctype_of(self, t):
return {'u8': 'uint8_t', 's8': 'int8_t', 'u16': 'uint16_t', 's16': 'int16_t', 'b16': 'uint16_t',
'u32': 'uint32_t', 's32': 'int32_t', 'b32': 'uint32_t', 'f32': 'float', 'u64': 'uint64_t',
's64': 'int64_t', 'b64': 'uint64_t'}[t]
def ptr(self, space, base, off, ct):
if base == 'KARG':
return f'*(const AS4 {ct}*)(KARGP + {off})'
if space == 'param':
return f'*(const AS4 {ct}*)(uintptr_t)({base} + {off})'
if space == 'shared':
return f'*(AS3 {ct}*)(uintptr_t)({base} + {off}u)'
if space == 'local':
return f'*(AS5 {ct}*)(uintptr_t)(uint32_t)({base} + {off})'
if space == 'global' and getattr(self, 'as1', False):
# --as1: ld/st.global through address-space-1 pointers (global_load / global_store instead of flat)
return f'*(AS1 {ct}*)(uintptr_t)({base} + {off})'
if space in ('global', 'local'):
return f'*({ct}*)(uintptr_t)({base} + {off})'
raise Err('space ' + space)
def ld(self, o, a):
space = next(s for s in ('param', 'global', 'shared', 'local') if s in o)
t = self.mem_type(o)
ct = self.ctype_of(t)
base, off = self.addr(a[1])
dsts = [x.strip() for x in a[0][1:-1].split(',')] if a[0].startswith('{') else [a[0]]
width = {'uint8_t': 1, 'int8_t': 1, 'uint16_t': 2, 'int16_t': 2, 'uint32_t': 4, 'int32_t': 4, 'float': 4,
'uint64_t': 8, 'int64_t': 8}[ct]
if len(dsts) in (2, 4) and width == 4 and space in ('global', 'shared') and base != 'KARG':
# ld.v2/v4 (naturally aligned in PTX) as one vector load instead of per-element loads
vt = 'u4' if len(dsts) == 4 else 'u2v'
asq = 'AS3 ' if space == 'shared' else ('AS1 ' if getattr(self, 'as1', False) else '')
ptr = f'(uintptr_t)({base} + {off}u)' if space == 'shared' else f'(uintptr_t)({base} + {off})'
self.emit(f'{{ const {vt} v_ = *({asq}const {vt}*){ptr};')
for i, d in enumerate(dsts):
self.assign(d, f'v_[{i}]', 'u32')
self.emit('}')
return
for i, d in enumerate(dsts):
e = self.ptr(space, base, off + i * width, ct)
if ct == 'float':
self.assign(d, e, 'float')
elif width == 8:
self.assign(d, e, 'u64')
else:
name = self.reg(d)
if self.rtype(name) == 'uint16_t':
self.assign(d, e, 'u16')
else:
self.assign(d, f'(uint32_t){e}' if ct != 'int8_t' and ct != 'int16_t' else f'(uint32_t)(int32_t){e}', 'u32')
def st(self, o, a):
space = next(s for s in ('param', 'global', 'shared', 'local') if s in o)
t = self.mem_type(o)
ct = self.ctype_of(t)
base, off = self.addr(a[0])
srcs = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]]
width = {'uint8_t': 1, 'uint16_t': 2, 'uint32_t': 4, 'int32_t': 4, 'float': 4, 'uint64_t': 8}[ct]
for i, s in enumerate(srcs):
v = self.val(s, 'f32' if ct == 'float' else 'u64' if width == 8 else 'u16' if width == 2 else 'u32')
self.emit(f'{self.ptr(space, base, off + i * width, ct)} = ({ct})({v});')
# ------------------------------------------------------------ roundf idiom
def rewrite_rz(self):
"""add.rz.ftz.f32 t, x, h; cvt.rzi.f32.f32 y, t (h = copysign(0.5, x)) -> y = roundf(x)"""
text = self.body
pat = re.compile(r'add\.rz\.ftz\.f32\s+(%f\d+),\s*(%f\d+),\s*(%f\d+);\s*(?://[^\n]*\s*)*cvt\.rzi\.f32\.f32\s+(%f\d+),\s*(%f\d+);')
count = 0
def repl(m):
nonlocal count
t, x, h, y, t2 = m.groups()
if t != t2:
raise Err('rz idiom mismatch')
uses = len(re.findall(re.escape(t) + r'\b', text))
if uses != 2:
raise Err(f'rz temp {t} used {uses} times')
# check h = copysign(0.5, x): mov.b32 rA, x; and.b32 rB, rA, -2147483648; or.b32 rC, rB, 1056964608; mov.b32 h, rC
count += 1
return f'd4r.roundf {y}, {x};'
self.body = pat.sub(repl, text)
return count
def main():
args = [a for a in sys.argv[1:] if not a.startswith('--')]
flags = {a for a in sys.argv[1:] if a.startswith('--')}
if len(args) != 2:
sys.exit(__doc__)
ptx, kernel = open(args[0]).read(), args[1]
mpart = None
if '--m-tail' in flags:
lines = ptx.replace('\r\n', '\n').split('\n')
if kernel == 'rrlite_enc0_4x4_mvhi_hdr_folded':
if lines[2547].strip() != 'bar.warp.sync -1;':
raise Err('enc0 cut changed')
lines = lines[:2548] + ['d4r.enc0 %r259, %rd106, %rd78, %r1066, %r1067, %r1230, %r1233;', 'ret;', '}']
mpart = 'enc0_tail'
elif kernel == 'rrlite_dec0_4x4_folded':
if lines[138].strip() != 'mov.u32 %r17, %laneid;' or lines[586].strip() != 'st.shared.v2.u16 [%r570+1168], {%rs63, %rs64};':
raise Err('dec0 cut changed')
lines[138:587] = ['mov.u32 %r560, _ZZ22rrlite_dec0_4x4_foldedN6RRLite18Decoder0ParametersEE4smem;',
'd4r.dec0 %r560, %rd22, %rd24, %r475, %r476, %r3, %r4;']
mpart = 'dec0_head'
else:
raise Err('--m-tail requires folded M enc0 or dec0')
ptx = '\n'.join(lines)
tr = Translator(ptx, kernel)
n_rz = tr.rewrite_rz()
if '--ifconv' in flags:
import ifconv, ctile as _ct
tr.body, n_if = ifconv.apply(tr.body, tokenize, _ct.parse, split_operands)
sys.stderr.write(f'ifconv: {n_if} diamonds converted\n')
stamps = [f.split('=', 1)[1].split(',') for f in flags if f.startswith('--stamps=')]
tr.stamps = stamps[0] if stamps else []
tr.bbcount = '--bbcount' in flags
tr.as1 = '--as1' in flags
tr.bblabels = []
ct = None
if '--ctile' in flags:
import ctile
ct = ctile.emit(tr, tokenize, Translator)
if '--sink' in flags:
import ifconv, ctile as _ct
tr.body, n_sink = ifconv.sink_stores(tr.body, tokenize, _ct.parse)
sys.stderr.write(f'sink: {n_sink} stores moved\n')
orig_op = tr.op
def op_with_round(op, a):
if op in ('d4r.hmin2', 'd4r.hmax2'):
fn = 'm_min2' if op == 'd4r.hmin2' else 'm_max2'
left = ', '.join(tr.val(v, 'f16') for v in a[2:4])
right = ', '.join(tr.val(v, 'f16') for v in a[4:6])
tr.emit(f'{{ uint32_t packed_ = ash2u({fn}((h2){{{left}}}, (h2){{{right}}}));')
tr.assign(a[0], 'packed_ & 0xffffu', 'u16')
tr.assign(a[1], 'packed_ >> 16', 'u16')
tr.emit('}')
return
if op in ('d4r.enc0', 'd4r.dec0'):
vals = [tr.val(v, t) for v, t in zip(a, ('u32', 'u64', 'u64', 'u32', 'u32', 's32', 's32'))]
tr.emit(f'd4r_{"enc0_tail" if op == "d4r.enc0" else "dec0_head"}({", ".join(vals)});')
return
if op == 'd4r.roundf':
tr.assign(a[0], f'__builtin_roundf({tr.val(a[1], "f32")})', 'float')
return
orig_op(op, a)
tr.op = op_with_round
tr.translate()
out = []
out.append(f'// Generated by kernels/native/ptx2hip.py from NVIDIA PTX ({kernel}); {n_rz} roundf idioms.')
out.append('#include "native_rt.h"')
if kernel.startswith('rrlite_'):
out.append('#include "m_ptx_rt.h"')
if 'cvt.rn.f16x2.e4m3x2' in ptx:
out.append('RT uint32_t d4r_decode_e4m3x2(uint16_t v) { h2 h; for (int i = 0; i < 2; ++i) { '
'uint32_t b = (v >> (8*i)) & 255u; '
'h[i] = (b & 127u) == 127u ? ash((uint16_t)(0x7fffu | ((b & 128u) << 8))) : '
'(h1)(ash((uint16_t)(((b & 127u) << 7) | ((b & 128u) << 8))) * (h1)256.0f); } return ash2u(h); }')
if mpart:
out.append(f'#include "../tex/{mpart}.hip"')
if ct is None:
for name, size, align in tr.shared:
out.append(f'__attribute__((shared)) __attribute__((aligned({max(align, 16)}))) static uint8_t {name}[{size}];')
else:
# one LDS block: the kernel's own arrays (path B) and the color tile (path A) are never live together
off = 0
for name, size, align in tr.shared:
off = (off + 15) & ~15
out.append(f'#define {name} (d4r_lds + {off})')
off += size
prelude, fill, info = ct
out.insert(2, f'__attribute__((shared)) __attribute__((aligned(16))) static uint8_t d4r_lds[{max(off, info["lds_bytes"])}];')
out.insert(2, f'// --ctile: {info["taps"]} color taps through an LDS tile, {len(info["groups"])} conversion shapes')
out.extend(prelude)
out.extend(fill)
out.append(f'struct __attribute__((aligned({tr.param_align}))) {kernel}_params {{ uint8_t b[{tr.param_size}]; }};')
maxntid = re.search(r'\.maxntid\s+(\d+)', tr.attrs)
lb = f'__attribute__((amdgpu_flat_work_group_size(1, {maxntid.group(1)})))' if maxntid else ''
bounds = [f.split('=', 1)[1] for f in flags if f.startswith('--max-block=')]
if bounds:
n = int(bounds[0])
if not 1 <= n <= 1024:
raise Err('invalid --max-block')
lb = f'__attribute__((amdgpu_flat_work_group_size(1, {n})))'
out.append(f'extern "C" __attribute__((global)) {lb} void {kernel}({kernel}_params)')
out.append('{')
out.append(' const AS4 uint8_t* KARGP = (const AS4 uint8_t*)__builtin_amdgcn_kernarg_segment_ptr();')
out.append(' const uint64_t KARGU = (uint64_t)(uintptr_t)KARGP;')
for name, size, align in tr.locals:
out.append(f' __attribute__((aligned({align}))) uint8_t {name}[{size}];')
if ct is not None:
out.append(' int32_t ct_x0 = 0, ct_y0 = 0, ct_w = 0, ct_h = 0;')
if tr.bbcount:
# diagnostic build: d4r_bbc[k] = number of waves that entered basic block k (label list in a comment)
out.insert(2, '__attribute__((device)) unsigned long long d4r_bbc[512];\n'
'#define D4R_BBC(k) do { if (wave_lane() == (uint32_t)__builtin_ctzll(__builtin_amdgcn_read_exec())) '
'__hip_atomic_fetch_add(&d4r_bbc[k], 1ull, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); } while (0)\n'
'// d4r_bbc labels: ' + ' '.join(tr.bblabels))
if tr.stamps:
# diagnostic build: per-region wave cycles (20-bit SHADER_CYCLES deltas) summed into d4r_dbg[k],
# wave counts into d4r_dbg[32 + k]; region k ends at stamp label k (k = len + 1: kernel return)
out.insert(2, '__attribute__((device)) unsigned long long d4r_dbg[64];\n'
'#define D4R_STAMP(k) do { uint32_t n_ = (uint32_t)__builtin_readcyclecounter(); '
'if (wave_lane() == (uint32_t)__builtin_ctzll(__builtin_amdgcn_read_exec())) { '
'__hip_atomic_fetch_add(&d4r_dbg[k], (unsigned long long)((n_ - st_last) & 0xfffffu), __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); '
'__hip_atomic_fetch_add(&d4r_dbg[32 + (k)], 1ull, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); } '
'st_last = n_; } while (0)')
out.append(' uint32_t st_last = (uint32_t)__builtin_readcyclecounter();')
by_type = {}
for name, t in sorted(tr.regs.items()):
by_type.setdefault(t, []).append(name)
for t, names in by_type.items():
for i in range(0, len(names), 16):
out.append(f' {t} {", ".join(names[i:i + 16])};')
out.extend(tr.lines)
out.append('}')
print('\n'.join(out))
if __name__ == '__main__':
main()
kernels/native/ptx_rt.h
+166
View File
@@ -0,0 +1,166 @@
// Native runtime for PTX kernels translated to HIP by ptx2hip.py (no ZLUDA compile involved).
// Every helper reproduces ZLUDA's lowering of the PTX instruction (ptx/src/pass/llvm/emit.rs and
// ptx/lib/zluda_ptx_impl.cpp, patches/zluda), so a translated kernel computes the same values as
// ZLUDA's own compile of NVIDIA's PTX. Texture and surface objects are ZLUDA handles: the address of
// a HIP image descriptor (8 dwords), its sampler (4 dwords at +48 bytes) and the CUDA format (+80).
// Compile flags (see build_native.sh): -ffp-contract=off, f32 denormals flushed (ZLUDA's
// "denormal-fp-math-f32"="preserve-sign"), f16 denormals kept.
#pragma once
#include <stdint.h>
#pragma clang fp contract(off)
#define RT static __attribute__((device, always_inline)) inline
#define AS4 __attribute__((address_space(4)))
#define AS3 __attribute__((address_space(3)))
#define AS1 __attribute__((address_space(1)))
#define AS5 __attribute__((address_space(5)))
typedef _Float16 h1;
typedef _Float16 h2 __attribute__((ext_vector_type(2)));
typedef float f4 __attribute__((ext_vector_type(4)));
typedef float f2 __attribute__((ext_vector_type(2)));
typedef int i2 __attribute__((ext_vector_type(2)));
typedef int i8v __attribute__((ext_vector_type(8)));
typedef int i4v __attribute__((ext_vector_type(4)));
typedef uint32_t u4 __attribute__((ext_vector_type(4)));
typedef uint32_t u2v __attribute__((ext_vector_type(2)));
// ---------------------------------------------------------------- bit casts
RT float asf(uint32_t u) { return __builtin_bit_cast(float, u); }
RT uint32_t asu(float f) { return __builtin_bit_cast(uint32_t, f); }
RT h1 ash(uint16_t u) { return __builtin_bit_cast(h1, u); }
RT uint16_t ashu(h1 h) { return __builtin_bit_cast(uint16_t, h); }
RT h2 ash2(uint32_t u) { return __builtin_bit_cast(h2, u); }
RT uint32_t ash2u(h2 h) { return __builtin_bit_cast(uint32_t, h); }
RT uint32_t pack16(uint16_t lo, uint16_t hi) { return (uint32_t)lo | ((uint32_t)hi << 16); }
// ---------------------------------------------------------------- f32
RT float canon(float x) { return __builtin_canonicalizef(x); }
RT float fmaxn(float a, float b) { return __builtin_fmaximum_numf(a, b); }
RT float fminn(float a, float b) { return __builtin_fminimum_numf(a, b); }
RT float sat(float x) { return fminn(fmaxn(x, 0.0f), 1.0f); }
RT float sat_ftz(float x) { return sat(canon(x)); } // cvt.ftz.sat.f32.f32
RT float abs_ftz(float x) { return __builtin_fabsf(canon(x)); }
RT float neg_ftz(float x) { return -canon(x); }
RT float rcp_a(float x) { return __builtin_amdgcn_rcpf(x); } // rcp.approx(.ftz).f32
RT float div_a(float a, float b) { return a * __builtin_amdgcn_rcpf(b); } // div.approx.ftz.f32 (fdiv arcp afn)
RT float ex2_a(float x) { return __builtin_amdgcn_exp2f(x); }
RT float lg2_a(float x) { return __builtin_amdgcn_logf(x); }
RT float rsq_a(float x) { return __builtin_amdgcn_rsqf(x); }
// cvt.rzi.s32.f32: saturating, NaN -> 0 (llvm.fptosi.sat, v_cvt_i32_f32); spelled out, a plain C cast is UB out of range
// Branch-free: the operand is clamped into int range first (fmaxf(NaN, lo) = lo), so the cast is always defined and
// the compiler can use selects (the guarded form compiled to a branch per conversion, splitting basic blocks).
RT int32_t f2i_rz(float x)
{
#ifdef D4R_NATIVE_F2I
// AMD V_CVT_I32_F32 truncates, saturates out-of-range inputs and maps NaN to zero.
// Explicit instruction avoids redundant clamps/selects emitted for a C saturating cast.
int32_t r;
__asm__("v_cvt_i32_f32 %0, %1" : "=v"(r) : "v"(x));
return r;
#else
const float c = __builtin_fminf(__builtin_fmaxf(x, -2147483648.0f), 2147483520.0f);
int32_t r = (int32_t)c;
r = x >= 2147483648.0f ? 2147483647 : r;
return x != x ? 0 : r;
#endif
}
// cvt.rzi.u32.f32: ZLUDA emits a plain fptoui (frozen) and relies on V_CVT_U32_F32 (truncates; negative and NaN -> 0,
// too large -> 0xffffffff), so the native helper is that instruction
RT uint32_t f2u_rz(float x)
{
uint32_t r;
__asm__("v_cvt_u32_f32 %0, %1" : "=v"(r) : "v"(x));
return r;
}
// ---------------------------------------------------------------- f16 (denormals kept)
RT h1 hfma(h1 a, h1 b, h1 c) { return __builtin_fmaf16(a, b, c); }
RT h2 h2fma(h2 a, h2 b, h2 c) { return __builtin_elementwise_fma(a, b, c); }
RT h1 hmaxn(h1 a, h1 b) { return __builtin_fmaximum_numf16(a, b); }
RT h1 hminn(h1 a, h1 b) { return __builtin_fminimum_numf16(a, b); }
RT h1 hsat(h1 x) { return hminn(hmaxn(x, (h1)0.0f), (h1)1.0f); }
RT h2 h2sat(h2 x) { return (h2){hsat(x.x), hsat(x.y)}; }
RT h2 h2abs(h2 x) { return __builtin_elementwise_abs(x); }
// ex2.approx.f16x2: v_exp_f32 per half (zluda_ptx_impl.cpp)
RT h2 h2ex2(h2 a) { return (h2){(h1)__builtin_amdgcn_exp2f((float)a.x), (h1)__builtin_amdgcn_exp2f((float)a.y)}; }
// cvt.rni/rmi.s32.f16: llvm.fptosi.sat.i32.f16 after the rounding (|f16| < 2^31, NaN -> 0)
RT int32_t h2i_sat(h1 h)
{
#ifdef D4R_NATIVE_H2I
return f2i_rz((float)h);
#else
return h != h ? 0 : (int32_t)(float)h;
#endif
}
RT h1 h_roundeven(h1 h) { return __builtin_elementwise_roundeven(h); }
RT h1 h_floor(h1 h) { return __builtin_elementwise_floor(h); }
// ---------------------------------------------------------------- threads (wave64: two CUDA warps per wave)
RT uint32_t wave_lane() { return __builtin_amdgcn_mbcnt_hi(~0u, __builtin_amdgcn_mbcnt_lo(~0u, 0u)); }
RT uint32_t warp_base()
{
#if __AMDGCN_WAVEFRONT_SIZE == 64
return wave_lane() & 32u;
#else
return 0u;
#endif
}
RT uint32_t laneid() { return wave_lane() & 31u; }
// shfl.sync.bfly.b32 d|p, a, delta, opts, mask (ZLUDA's SHFL_SYNC_IMPL)
RT uint32_t shfl_bfly(uint32_t input, int32_t delta, uint32_t opts, bool* in_bounds)
{
int32_t section_mask = (opts >> 8) & 31, warp_end = opts & 31;
delta &= 31;
int32_t self = (int32_t)laneid();
int32_t subsection = section_mask & self;
int32_t subsection_end = subsection | (~section_mask & warp_end);
int32_t idx = self ^ delta;
bool oob = idx > subsection_end;
if (oob)
idx = self;
if (in_bounds)
*in_bounds = !oob;
return (uint32_t)__builtin_amdgcn_ds_bpermute((int32_t)(warp_base() | (uint32_t)idx) << 2, (int32_t)input);
}
RT uint32_t bfi32(uint32_t insert, uint32_t base, uint32_t pos_32, uint32_t len_32)
{
uint32_t pos = pos_32 & 0xffu, len = len_32 & 0xffu;
if (pos >= 32)
return base;
uint32_t mask = len >= 32 ? (0xffffffffu << pos) : (((len >= 32 ? 0u : (1u << len)) - 1u) << pos);
return (~mask & base) | (mask & (insert << pos));
}
// ---------------------------------------------------------------- textures (ZLUDA handles)
__attribute__((device)) static uint32_t d4r_null_texture[32];
RT uint64_t tex_or_null(uint64_t t) { return t != 0 ? t : (uint64_t)(uintptr_t)d4r_null_texture; }
RT i8v tdesc(uint64_t t) { return *(const AS4 i8v*)(uintptr_t)tex_or_null(t); }
RT i4v sdesc(uint64_t t) { return *(const AS4 i4v*)(uintptr_t)(tex_or_null(t) + 48); }
__attribute__((device)) f4 llvm_image_load_mip_2d(uint32_t, int32_t, int32_t, int32_t, i8v, int, int)
__asm("llvm.amdgcn.image.load.mip.2d.v4f32.i32");
__attribute__((device)) f4 llvm_image_sample_l_2d(uint32_t, float, float, float, i8v, i4v, bool, int, int)
__asm("llvm.amdgcn.image.sample.l.2d.v4f32.f32");
extern "C" __attribute__((device)) f4 __ockl_image_gather4r_2D(const AS4 uint32_t* image, const AS4 uint32_t* sampler, f2 coord);
// tex.base.2d.v4.f32.s32: texel load, mip 0
RT f4 tex_fetch(uint64_t t, int32_t x, int32_t y) { return llvm_image_load_mip_2d(0xf, x, y, 0, tdesc(t), 0, 0); }
// tex.level.2d.v4.f32.f32
RT f4 tex_level(uint64_t t, float x, float y, float lod) { return llvm_image_sample_l_2d(0xf, x, y, lod, tdesc(t), sdesc(t), false, 0, 0); }
// tld4.r.2d.v4.f32.f32
RT f4 tex_gather_r(uint64_t t, float x, float y)
{
const AS4 uint32_t* image = (const AS4 uint32_t*)(uintptr_t)tex_or_null(t);
return __ockl_image_gather4r_2D(image, image + 12, (f2){x, y});
}
// bar.sync 0 (HIP's __syncthreads)
RT void bar_sync()
{
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
}
kernels/native/rewrite_output.py
Executable
+66
View File
@@ -0,0 +1,66 @@
#!/usr/bin/env python3
"""Build the staged output rewrite from the exact translated kernel.
Replaces path B's cooperative tile fill with output_tile.h. All other phases and
the DLAA path remain intact. Refuses to generate if the region's register contract
changes, so new translator output cannot silently reuse an incompatible mapping.
"""
import argparse
import re
from pathlib import Path
REGISTER = re.compile(r"\b(?:rd|rs|r|f|fd|p)\d+\b")
ASSIGN = re.compile(r"^\s*((?:rd|rs|r|f|fd|p)\d+)\s*=")
LABEL = re.compile(r"^(L__BB0_\d+):;$")
def rewrite(source):
head, main = source.split('extern "C" __attribute__((global))', 1)
lines = main.splitlines(keepends=True)
labels = [(i, LABEL.match(line.strip()).group(1)) for i, line in enumerate(lines) if LABEL.match(line.strip())]
blocks = {name: (start, labels[k + 1][0] if k + 1 < len(labels) else len(lines))
for k, (start, name) in enumerate(labels)}
pending, region = ['L__BB0_124'], set()
while pending:
name = pending.pop()
if name in region or name == 'L__BB0_138':
continue
region.add(name)
start, end = blocks[name]
body = ''.join(lines[start + 1:end])
pending += re.findall(r'goto (L__BB0_\d+);', body)
if not re.search(r'^\s*goto L__BB0_\d+;\s*$', lines[end - 1]):
# This region's fallthroughs are only its internal branch joins.
pending += [n for i, n in labels if i == end]
expected = {'L__BB0_124'} | {f'L__BB0_{k}' for k in range(127, 138)}
if region != expected:
raise ValueError(f'tile CFG changed: {sorted(region)}')
removed = {i for name in region for i in range(*blocks[name])}
definitions = {m.group(1) for i in removed if (m := ASSIGN.match(lines[i]))}
# Definitions declared at function scope count as uses. Ignore declaration-only lines.
uses_outside = set(REGISTER.findall(''.join(line for i, line in enumerate(lines)
if i not in removed and not re.match(r'^\s*(?:uint\d+_t|float|double|bool)\s+(?:rd|rs|r|f|fd|p)\d+(?:,|;)', line))))
live_out = definitions & uses_outside
if live_out != {'r215', 'r217'}:
raise ValueError(f'tile register contract changed: {sorted(live_out)}')
replacement = '''L__BB0_124:;
output_tile_fill(KARGP, (int32_t)r3, (int32_t)r4, (int32_t)r7, (int32_t)(r8 + 1u),
f1, (AS3 uint8_t*)sharedColor, (AS3 uint8_t*)sharedDepth);
r215 = (uint32_t)__builtin_amdgcn_workitem_id_y();
r217 = (uint32_t)__builtin_amdgcn_workitem_id_x();
goto L__BB0_138;
'''
result = []
for i, line in enumerate(lines):
if i == blocks['L__BB0_124'][0]:
result.append(replacement)
if i not in removed:
result.append(line)
return head + '#include "output_tile.h"\n' + 'extern "C" __attribute__((global))' + ''.join(result)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('source', type=Path)
args = parser.parse_args()
print(rewrite(args.source.read_text()), end='')
kernels/tex/tex_common.h
+9 -2
View File
@@ -303,6 +303,13 @@ __attribute__((device)) static inline uint32_t redirect_pitch(uint64_t surface)
const uint32_t word = ((const __attribute__((address_space(4))) uint32_t*)surface)[21];
return (word >> 16) == 0x5232u ? (word & 0xffffu) << 3 : 0u;
}
// D4R_AS1_STORES: redirect stores through address-space-1 pointers (global_store instead of flat_store; the
// redirect target is linear device memory)
#ifdef D4R_AS1_STORES
#define D4R_REDIRECT_PTR(T, address) ((__attribute__((address_space(1))) T*)(uintptr_t)(address))
#else
#define D4R_REDIRECT_PTR(T, address) ((T*)(address))
#endif
__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);
@@ -325,7 +332,7 @@ DEV __attribute__((always_inline)) void d4r_sust_p_v4b32(uint64_t surface, int32
{
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};
*D4R_REDIRECT_PTR(uint2_t, redirect_row(surface, y) + 8 * x) = (uint2_t){lo, hi};
}
return;
}
@@ -430,7 +437,7 @@ DEV __attribute__((always_inline)) void d4r_sust_b32(uint64_t surface, int32_t x
{
// 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;
*D4R_REDIRECT_PTR(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);