From 22026ede12b2a8c912d1ca35750fb7c8dc94b52a Mon Sep 17 00:00:00 2001 From: ayo <93222141+countervolts@users.noreply.github.com> Date: Mon, 5 Oct 2026 14:48:37 -0400 Subject: [PATCH] Port K kernel tuning and M gfx1101 flags 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 --- README.md | 2 +- docs/building.md | 7 +- docs/native-kernels.md | 12 + docs/performance.md | 14 + kernels/build.sh | 35 +- kernels/common/wmma_layout.h | 100 +- kernels/k/dltss_pwin_dec0_layer.hip | 10 +- kernels/k/dltss_pwin_dec1_layer.hip | 7 +- kernels/k/dltss_pwin_dec2_layer.hip | 6 +- kernels/k/dltss_pwin_dec3_layer.hip | 2 + kernels/k/dltss_pwin_dec4_layer.hip | 3 +- kernels/k/dltss_pwin_dec5_layer.hip | 6 +- kernels/k/dltss_pwin_enc0_layer.hip | 12 +- kernels/k/dltss_pwin_enc1_layer.hip | 8 +- kernels/k/dltss_pwin_enc2_layer.hip | 6 +- kernels/k/dltss_pwin_enc4_layer.hip | 3 +- kernels/k/pwin_common.h | 104 +++ kernels/k/pwin_gelu_lut.h | 37 + kernels/k/pwin_layer.h | 625 ++++++++++--- kernels/k/pwin_phase.h | 402 ++++++++ kernels/k/pwin_pos.h | 256 ++++-- kernels/k/pwin_pos_recompute.h | 379 ++++++++ kernels/k/pwin_wide.h | 226 ++--- kernels/m/rrlite_dec1_4x4.hip | 11 +- kernels/m/rrlite_dec2_4x4.hip | 11 +- kernels/m/rrlite_enc1_4x4.hip | 9 +- kernels/m/rrlite_enc2_4x4.hip | 12 +- kernels/m/rrlite_enc3_tube_4x4.hip | 9 +- kernels/m/swin_block.h | 164 +++- kernels/m/swin_common.h | 72 +- kernels/native/build_native.sh | 18 + kernels/native/ctile.py | 411 +++++++++ kernels/native/ifconv.py | 234 +++++ kernels/native/in_fuse.py | 123 +++ kernels/native/in_regs.hip | 22 + kernels/native/m_ptx_rt.h | 28 + kernels/native/native_rt.h | 185 ++++ kernels/native/output_tile.h | 71 ++ kernels/native/ptx2hip.py | 857 ++++++++++++++++++ kernels/native/ptx_rt.h | 166 ++++ kernels/native/rewrite_output.py | 66 ++ kernels/tex/tex_common.h | 11 +- packaging/SOURCES.txt | 2 +- .../0008-native-post-phase-kernels.patch | 122 +++ scripts/build_release_glibc241.sh | 2 +- 45 files changed, 4514 insertions(+), 354 deletions(-) create mode 100644 kernels/k/pwin_gelu_lut.h create mode 100644 kernels/k/pwin_phase.h create mode 100644 kernels/k/pwin_pos_recompute.h create mode 100755 kernels/native/build_native.sh create mode 100755 kernels/native/ctile.py create mode 100755 kernels/native/ifconv.py create mode 100755 kernels/native/in_fuse.py create mode 100644 kernels/native/in_regs.hip create mode 100644 kernels/native/m_ptx_rt.h create mode 100644 kernels/native/native_rt.h create mode 100644 kernels/native/output_tile.h create mode 100755 kernels/native/ptx2hip.py create mode 100644 kernels/native/ptx_rt.h create mode 100755 kernels/native/rewrite_output.py create mode 100644 patches/zluda/0008-native-post-phase-kernels.patch diff --git a/README.md b/README.md index fa2223f..73fabad 100644 --- a/README.md +++ b/README.md @@ -123,7 +123,7 @@ ROCm does not need to be installed: the zip includes its runtime (from AMD's Ubu The full sequence is in [docs/building.md](docs/building.md). In short: -1. Build ZLUDA `ee2f25a` with `patches/zluda/0002`–`0007` applied, including its `d4r_emit` example for offline texture builds. +1. Build ZLUDA `ee2f25a` with `patches/zluda/0002`–`0008` applied, including its `d4r_emit` example for offline texture builds. 2. Build the patched vkd3d-proton: `scripts/build_vkd3d_proton_d4r.sh OUT_DIR`. 3. Build the shim and bridge: `scripts/build_d4r_nvngx_shim.sh`, `scripts/build_wine_nvcuda_bridge.sh`. 4. Stage the runtime with your NVIDIA files: `scripts/install_d4r_runtime.sh _nvngx.dll nvngx_dlss.dll`. diff --git a/docs/building.md b/docs/building.md index 60cd488..3747463 100644 --- a/docs/building.md +++ b/docs/building.md @@ -4,7 +4,7 @@ These steps produce the pieces the launcher needs: a patched ZLUDA, a patched vk | Component | Tested version | |---|---| -| ZLUDA | `ee2f25a` (upstream), plus `patches/zluda/0002` through `0007` in order | +| ZLUDA | `ee2f25a` (upstream), plus `patches/zluda/0002` through `0008` in order | | vkd3d-proton | `3dfc6f07` (the base GE-Proton11-3 ships), plus `patches/vkd3d-proton/0001` and `0002` | | ROCm | 7.2 (HIP runtime, clang, device libraries) | | Proton | GE-Proton11-3 (with its OptiScaler integration) | @@ -17,7 +17,7 @@ git clone https://github.com/vosen/ZLUDA zluda && cd zluda git checkout ee2f25a git submodule update --init --recursive git lfs pull -for p in 0002 0003 0004 0005 0006 0007; do git apply /path/to/d4r/patches/zluda/$p-*.patch; done +for p in 0002 0003 0004 0005 0006 0007 0008; do git apply /path/to/d4r/patches/zluda/$p-*.patch; done # rebuild the device helpers the patches changed (ptx/lib/zluda_ptx_impl*.bc) ZLUDA_SOURCE_ROOT=$PWD ROCM_ROOT=/opt/rocm /path/to/d4r/scripts/build_zluda_ptx_helpers.sh LIBRARY_PATH=/opt/rocm/lib cargo build --release -p zluda @@ -44,6 +44,7 @@ What the patches add: - **0005**: a null texture object (CUDA handle 0) reads as zeros, as on NVIDIA GPUs, instead of faulting the GPU. DLSS samples absent optional inputs that way in some configurations (low-resolution motion vectors without HDR, as in Ghost of Tsushima). - **0006**: `m16n8k8` f16 MMAs (DLSS 3 CNN, presets E/F) on RDNA3 WMMA like the k16 ones (`D4R_ZLUDA_WMMA_K8=0` disables it); weight-image slots for native prep kernels (`d4r_prep_key_at`, `d4r_prep_key_slots`); and wave64 compilation for offline texture-kernel builds (`D4R_ZLUDA_WAVE64=1`, with wave64 builds of the helper bitcode). `scripts/build_zluda_ptx_helpers.sh` now writes the `_w64` helper variants too. - **0007**: gfx12 WMMA layout lowering, optional native e4m3 FP8 WMMA (`D4R_ZLUDA_WMMA_FP8_NATIVE=1`), and an architecture argument for `d4r_emit`. `D4R_ZLUDA_WMMA_LAYOUT=12` on gfx11 is a validation shim, not a release setting. +- **0008**: post-phase kernels for native overrides (`NAME_post1` … `NAME_post8`, see [native-kernels.md](native-kernels.md)). Preset K's dec5 layer runs as four phase kernels, so the K set needs this patch: an older `libnvcuda.so` would skip the phases. When linking on a system with ROCm libraries outside the default search path, include their library directory in `LIBRARY_PATH`. The build also needs the appropriate ROCm link libraries. `CARGO_BUILD_JOBS=8` caps parallel Rust compilation if memory is limited. @@ -131,7 +132,7 @@ Mount the repository at `/work`, ROCm's compiler/headers/device libraries at `/o | Variable | Input | |---|---| | `D4R_RELEASE_BUILD_ROOT` | writable directory under `/work/build/` | -| `D4R_ZLUDA_SRC` | isolated checkout inside that directory, with both pinned LLVM and HiGHS submodules, the real OCKL LFS payload, and patches `0002`–`0007` applied | +| `D4R_ZLUDA_SRC` | isolated checkout inside that directory, with both pinned LLVM and HiGHS submodules, the real OCKL LFS payload, and patches `0002`–`0008` applied | | `D4R_VKD3D_SRC` | isolated vkd3d-proton checkout inside that directory, with its submodules and patches `0001` and `0002` applied | | `CARGO_HOME` | writable build-local Cargo cache; prefetch the locked dependencies for an offline build | | `D4R_ROCM_DIR` | `/opt/rocm` | diff --git a/docs/native-kernels.md b/docs/native-kernels.md index 2711395..e09bd98 100644 --- a/docs/native-kernels.md +++ b/docs/native-kernels.md @@ -15,6 +15,7 @@ With `D4R_ZLUDA_NATIVE_DIR=DIR`, whenever NGX loads a PTX module and asks for ke | `u32 d4r_prep_key_slots` | N > 0: the kernel keeps one prepared image per key for up to N keys; keys beyond N share an overflow image. This requires stable weight-address identity. Enc3 disables key caching and prepares every launch instead. | | `u32 d4r_block_z` | replaces the launch's block z dimension (more waves per window) | | `u32 d4r_grid_x` | replaces the grid (persistent kernels) | +| `NAME_post1` … `NAME_post8` kernels | launched after `NAME` on the same stream with the same parameters (patches/zluda/0008); `u32 d4r_postK_grid_x` / `_grid_y` replace the grid (0 = the launch's), `d4r_postK_block_x` / `_y` / `_z` the block (0 = 32 / 1 / 1). One layer can run as several dependent phases. | Kernels without a file in `DIR` are compiled from PTX as usual, so a partial set works. @@ -34,12 +35,23 @@ The release is built this way, with hashes from DLSS 310.7.0 and 310.9.1. The PT - `pwin_pos.h`: position-only layers with two token tiles per wave, which halves weight traffic. - `pwin_wide.h`: the deep layers, which have only 6–77 windows per frame. It uses 4·NG waves per window, split by token tile and channel group, so the GPU stays busy. - **Weights:** prep kernels expand the weights into WMMA operand images once. +- **Exact tuning round (2026-10-04, RX 7700 XT):** every change keeps all captured buffers and the harness image byte-identical to the previous build. + - LDS rows are padded (`PWIN_PAD`) so the 16 lanes of a WMMA operand load hit distinct banks. + - enc1/enc2/dec1/dec2 build in WGP mode; enc0, dec0 and dec3 move the half-wave exchange to `ds_bpermute`. + - `common/wmma_layout.h` removes the dynamic vector indexing of accumulator elements and adds packed f16 conversions (`wm_cvt2`, `WM_CVT_ASM`, `PWIN_PK_ELEM`). + - Weight tiles load through raw buffer loads with a uniform offset (`op_imgt`); the f32 accumulator-init vectors are prepared once into image slots. + - The deep layers use more waves per window (`NG = 8` for enc4, dec4, dec5). + - dec5 runs as four dependent phase kernels (`pwin_phase.h`, post kernels, needs patch 0008). + - enc1 keeps one LDS buffer and reloads its input for the residual (`pwin_pos_recompute.h`); dec1 computes its Q, K and V chains two at a time. + - Tried and not kept (slower or not exact): interleaved WMMA chains, a GELU lookup table, streaming attention, f16-accumulating WMMA, persistent or megakernel layers, phase kernels for the other deep layers, fast math. +- **Texture kernels as native code:** `kernels/native/ptx2hip.py` translates a DLSS PTX kernel into HIP, each instruction as ZLUDA's lowering of it (`ptx_rt.h`, `native_rt.h`). `kernels/build.sh` uses it, at build time and from your own DLL, for two kernels of the fast set: `hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel` (`--ctile`, the path-B LDS tile as a function in `output_tile.h`, no unused depth prefetch, global-memory redirect stores) and `hiluma_engine_input_depthinv_mvhi_hdr_v2_rel` (`--ifconv --sink`, `in_fuse.py` moves its scratch arrays into registers, `in_regs.hip`, wave32). Both match ZLUDA's compile byte for byte in the harness. The generated sources are never stored. Other flag combinations of these kernels still use the PTX route above, and the accuracy set keeps ZLUDA's compile because the native code flushes f32 denormals. **DLSS 4.5, presets L/M (`kernels/m`, `rrlite_*`).** - **Network:** the Swin blocks of DLSS 4.5, whose weights are FP8. Their prep kernels expand the weights to f16 WMMA operands. The enc3 tube layer prepares on every launch: NGX can recycle a weights address for a different tube block after feature recreation, so persistent address-only reuse can return stale weights and corrupt the image after quality changes. Other layers retain their existing preparation policy. - **Template:** `swin_block.h` covers encoders, the tube-shaped enc3 and decoders. - **Output encoding:** the FP8 output is encoded two values at a time with packed 16-bit operations (`enc8x2`, exhaustively equal to the scalar encoder), and the 2×2 patch merge reads the rounded f16 values the codes decode to, from a row layout without LDS bank conflicts, instead of decoding the bytes again in every wave. - **Weight loads:** enc1 forms each weight tile's address in scalar registers (`SWIN_SCALAR_BLOAD`); the 8-wave layers are faster without it. +- **gfx1101 tuning (`// d4r-build-flags-gfx1101:` lines):** enc2 uses a packed FP8 decode, packed norm sums, a `ds_bpermute` half-wave exchange and a 4x unrolled MLP loop; dec1 caps VGPRs at 160 (no scratch); dec2 uses the memory-clause scheduler. All are byte-identical to the default build; the total gain is small (about 0.07 ms of 7.9 ms). **Preset L texture variants.** L uses unfolded `rrlite_enc0_4x4_*`, `rrlite_dec0_4x4`, and `rrlite_post_3_*` kernels. The L build keeps their full neural arithmetic and replaces surface stores with the existing native format-conversion functions. M's folded enc0 tail and dec0 head cannot be used for these variants. All four input flag combinations and all eight post variants are built, along with the shared downsample kernels. `all` and `tex` include these variants; `l` builds the shared Swin layers plus only the L texture set. L remains experimental; these replacements do not establish RTX parity or game performance. diff --git a/docs/performance.md b/docs/performance.md index cbeb9a5..8b861c1 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -44,6 +44,20 @@ DLSS 4 (K) at Quality, frames per second on the walk: DLSS 4.5 (M) at Quality went from 28 fps (ZLUDA only) to 42 with the first native Swin layers, 46 with more waves per window, 50 with fast numerics and the texture-kernel tails, and 51.5 with the hand-off changes. +## DLSS 4 (K) tuning round, 2026-10-04 + +A separate work copy of this tree tuned the K kernels with byte-identical output (see [native-kernels.md](native-kernels.md)). Numbers from that copy, RX 7700 XT, D3D12 harness, motion scene, FrameAge 0, warmed medians of back-to-back runs; the shipped 0.1.4 kernels against the tuned set: + +| | shipped | tuned | +|---|---|---| +| 2560×1440 Quality: DLSS GPU time (`gpu_eval`) | 2.89 ms | 2.41 ms | +| 1920×1080 Quality: DLSS GPU time | 1.86 ms | 1.50 ms | +| 2560×1440 Quality: sum of kernels | 2.90 ms | 2.40 ms | + +In SILENT HILL Townfall (static scene, K Quality 1706×960 → 2560×1440, MangoHud, two alternating rounds): 63.1 / 63.3 fps → 65.4 / 65.1 fps with default settings (frame time −0.5 ms), DLSS GPU time 2.92 → 2.41 ms. At 1920×1080: 86.1 → 89.3 fps. DLSS 4.5 (M) gained little from its own round (DLSS GPU time 7.96 → 7.91 ms at 2560×1440 Quality, Townfall 1080p +0.7%). + +These were measured in the work copy. They have not been re-measured from this repository since the port. + ## What did not help - **Accumulating in f16 on the WMMA units.** The error grows about 20× and the image degrades. diff --git a/kernels/build.sh b/kernels/build.sh index 3e279e4..efbc1be 100755 --- a/kernels/build.sh +++ b/kernels/build.sh @@ -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" diff --git a/kernels/common/wmma_layout.h b/kernels/common/wmma_layout.h index e417b0c..23fe9fa 100644 --- a/kernels/common/wmma_layout.h +++ b/kernels/common/wmma_layout.h @@ -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 } diff --git a/kernels/k/dltss_pwin_dec0_layer.hip b/kernels/k/dltss_pwin_dec0_layer.hip index 061cf87..c5b373a 100644 --- a/kernels/k/dltss_pwin_dec0_layer.hip +++ b/kernels/k/dltss_pwin_dec0_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_dec1_layer.hip b/kernels/k/dltss_pwin_dec1_layer.hip index 75e5015..77a9587 100644 --- a/kernels/k/dltss_pwin_dec1_layer.hip +++ b/kernels/k/dltss_pwin_dec1_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_dec2_layer.hip b/kernels/k/dltss_pwin_dec2_layer.hip index 37d5d42..81fb6ea 100644 --- a/kernels/k/dltss_pwin_dec2_layer.hip +++ b/kernels/k/dltss_pwin_dec2_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_dec3_layer.hip b/kernels/k/dltss_pwin_dec3_layer.hip index aa31730..9096cb5 100644 --- a/kernels/k/dltss_pwin_dec3_layer.hip +++ b/kernels/k/dltss_pwin_dec3_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_dec4_layer.hip b/kernels/k/dltss_pwin_dec4_layer.hip index a05d36a..0dbcf96 100644 --- a/kernels/k/dltss_pwin_dec4_layer.hip +++ b/kernels/k/dltss_pwin_dec4_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_dec5_layer.hip b/kernels/k/dltss_pwin_dec5_layer.hip index 29762d8..946de7d 100644 --- a/kernels/k/dltss_pwin_dec5_layer.hip +++ b/kernels/k/dltss_pwin_dec5_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_enc0_layer.hip b/kernels/k/dltss_pwin_enc0_layer.hip index 9c9086f..406b1ac 100644 --- a/kernels/k/dltss_pwin_enc0_layer.hip +++ b/kernels/k/dltss_pwin_enc0_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_enc1_layer.hip b/kernels/k/dltss_pwin_enc1_layer.hip index 8cd91bb..a1946c5 100644 --- a/kernels/k/dltss_pwin_enc1_layer.hip +++ b/kernels/k/dltss_pwin_enc1_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_enc2_layer.hip b/kernels/k/dltss_pwin_enc2_layer.hip index d7a7aac..ae28b47 100644 --- a/kernels/k/dltss_pwin_enc2_layer.hip +++ b/kernels/k/dltss_pwin_enc2_layer.hip @@ -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) diff --git a/kernels/k/dltss_pwin_enc4_layer.hip b/kernels/k/dltss_pwin_enc4_layer.hip index e99b551..e33ba50 100644 --- a/kernels/k/dltss_pwin_enc4_layer.hip +++ b/kernels/k/dltss_pwin_enc4_layer.hip @@ -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) diff --git a/kernels/k/pwin_common.h b/kernels/k/pwin_common.h index 4948d71..3702d6f 100644 --- a/kernels/k/pwin_common.h +++ b/kernels/k/pwin_common.h @@ -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 diff --git a/kernels/k/pwin_gelu_lut.h b/kernels/k/pwin_gelu_lut.h new file mode 100644 index 0000000..ec5f8cd --- /dev/null +++ b/kernels/k/pwin_gelu_lut.h @@ -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); + } +} diff --git a/kernels/k/pwin_layer.h b/kernels/k/pwin_layer.h index f0ad301..4b35fff 100644 --- a/kernels/k/pwin_layer.h +++ b/kernels/k/pwin_layer.h @@ -38,7 +38,12 @@ template 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 __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 __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 __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 __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 __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx) { const uint8_t* w = p.w; + residual_vector_prep(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 __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 __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias) +template __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias) { - pwin_prep_at(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(p.w, p.w, img, j); + return; + } + pwin_prep_at(p, img, bias, idx); +} +template __device__ void pwin_prep_nomerge(const PwinParams& p, u8v* img, u4v* bias) +{ + pwin_prep(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 __device__ __forceinline__ pacc acc_vec_p(const u8v* img, const uint8_t* w, int byte, int v) +{ + const f8v f = acc_vec(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 __device__ __forceinline__ pacc acc_vec_p(const u8v* img, const uint8_t* w, int byte, int v) { return acc_vec(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 +#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 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 __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(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(w, 0, nt, g2); + residual_vector(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(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 struct PwinEmbLayout : using B = PwinLayout; 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 __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(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(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(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; \ - __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 __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; \ - __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 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 __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(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 __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(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(q, img, bias, act, hb, kl, vt, PwinX0{expand, x0b}, x0b); +#else pwin_core(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(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; \ - __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; \ diff --git a/kernels/k/pwin_phase.h b/kernels/k/pwin_phase.h new file mode 100644 index 0000000..3c50bc7 --- /dev/null +++ b/kernels/k/pwin_phase.h @@ -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 struct PhaseCB +{ + template static constexpr int cb(decltype(T::CB)*) { return T::CB; } + template static constexpr int cb(...) { return 0; } + static constexpr int value = cb(nullptr); +}; +template __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 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 +__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 win(p); + if (win.n > MAXW) + return; + const uint8_t* w = p.w + PhaseCB::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(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(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 +__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 win(p); + if (win.n > MAXW) + return; + const uint8_t* w = p.w + PhaseCB::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(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(p, gX, wi, bx, by, tok) + 16 * nt); + half_t mv[8], av[8], g2[8], b2[8]; + residual_vector(w, 0, nt, g2); + residual_vector(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 +__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 win(p); + if (win.n > MAXW) + return; + const uint8_t* w = p.w + PhaseCB::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(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 +__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 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 +__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 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(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 +__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 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(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; \ + __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(p, g_img, g_bias); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) \ + { \ + phase_attn(p, g_img, g_bias, g_O); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) \ + { \ + phase_wo(p, g_img, g_O, g_M, g_A); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) \ + { \ + phase_fc1(p, g_img, g_M, g_H); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) \ + { \ + phase_fc2(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; \ + 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(p, g_img, g_bias); } \ + extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) { phase_expand(p, g_img, g_X); } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) \ + { \ + phase_attn(p, g_img, g_bias, g_O, g_X); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) \ + { \ + phase_wo(p, g_img, g_O, g_M, g_A, g_X); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) { phase_fc1(p, g_img, g_M, g_H); } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post4(PwinParams p) \ + { \ + phase_fc2(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; \ + 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(p, g_img, g_bias); } \ + extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) { phase_attn(p, g_img, g_bias, g_O); } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post1(PwinParams p) { phase_wo(p, g_img, g_O, g_M, g_A); } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post2(PwinParams p) { phase_fc1(p, g_img, g_M, g_H); } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post3(PwinParams p) \ + { \ + phase_fc2(p, g_img, g_A, g_H, p.out32, g_Y); \ + } \ + extern "C" __global__ void __launch_bounds__(128) NAME##_post4(PwinParams p) { phase_merge(p, g_img, g_Y); } diff --git a/kernels/k/pwin_pos.h b/kernels/k/pwin_pos.h index 9159f21..4c44e89 100644 --- a/kernels/k/pwin_pos.h +++ b/kernels/k/pwin_pos.h @@ -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 __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(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(w, 0, nt, g2); + residual_vector(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(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 __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(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(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(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 __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(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(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(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; \ - __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(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; \ - __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(p, g_img, g_bias); \ } \ extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \ diff --git a/kernels/k/pwin_pos_recompute.h b/kernels/k/pwin_pos_recompute.h new file mode 100644 index 0000000..52c6309 --- /dev/null +++ b/kernels/k/pwin_pos_recompute.h @@ -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 +__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(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(w, 0, nt, g2); + residual_vector(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(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(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 +__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(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(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(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; \ + __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(p, g_img, g_bias); \ + } \ + extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \ + { \ + pos_recompute_encoder(p, g_img, g_bias); \ + } + diff --git a/kernels/k/pwin_wide.h b/kernels/k/pwin_wide.h index 298fd48..a3aa2c8 100644 --- a/kernels/k/pwin_wide.h +++ b/kernels/k/pwin_wide.h @@ -15,27 +15,71 @@ template 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 +__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 __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; 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(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(img, hrow, h, hs, tok, qop, kl, vt); + else + wide_qkv(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(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(w, 0, nt, g2); + residual_vector(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(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 __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 __device__ void wide_load_x0(const PwinParams& p, uin } // y rows of the window's own tokens -> dst (C per token) -template __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[L::C], uint8_t* dst) +template __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 __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::R_BYTES]; wide_load_x0(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(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 __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::R_BYTES]; wide_load_x0(p, R); block_sync(); @@ -374,11 +403,11 @@ template __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::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(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) -#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep, wide_plain, NG, PwinLayout) +#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep_nomerge, wide_plain, NG, PwinLayout) #define WIDE_DECODER(NAME, H, C, CL, NG) WIDE_KERNEL(NAME, pwin_dec_prep, wide_decoder, NG, PwinDecLayout) diff --git a/kernels/m/rrlite_dec1_4x4.hip b/kernels/m/rrlite_dec1_4x4.hip index 79ba4db..4dccc6f 100644 --- a/kernels/m/rrlite_dec1_4x4.hip +++ b/kernels/m/rrlite_dec1_4x4.hip @@ -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) diff --git a/kernels/m/rrlite_dec2_4x4.hip b/kernels/m/rrlite_dec2_4x4.hip index 860681d..22458f8 100644 --- a/kernels/m/rrlite_dec2_4x4.hip +++ b/kernels/m/rrlite_dec2_4x4.hip @@ -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) diff --git a/kernels/m/rrlite_enc1_4x4.hip b/kernels/m/rrlite_enc1_4x4.hip index 10ef58a..60f4dea 100644 --- a/kernels/m/rrlite_enc1_4x4.hip +++ b/kernels/m/rrlite_enc1_4x4.hip @@ -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) diff --git a/kernels/m/rrlite_enc2_4x4.hip b/kernels/m/rrlite_enc2_4x4.hip index dd57e54..72cd14f 100644 --- a/kernels/m/rrlite_enc2_4x4.hip +++ b/kernels/m/rrlite_enc2_4x4.hip @@ -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) diff --git a/kernels/m/rrlite_enc3_tube_4x4.hip b/kernels/m/rrlite_enc3_tube_4x4.hip index 0c001c6..6a3c222 100644 --- a/kernels/m/rrlite_enc3_tube_4x4.hip +++ b/kernels/m/rrlite_enc3_tube_4x4.hip @@ -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) diff --git a/kernels/m/swin_block.h b/kernels/m/swin_block.h index ab6ec12..29cf769 100644 --- a/kernels/m/swin_block.h +++ b/kernels/m/swin_block.h @@ -6,6 +6,8 @@ #pragma once #include "swin_common.h" #include +#define SWIN_PRAGMA_IMPL(x) _Pragma(#x) +#define SWIN_PRAGMA(x) SWIN_PRAGMA_IMPL(x) template 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) diff --git a/kernels/m/swin_common.h b/kernels/m/swin_common.h index d544800..f07c28f 100644 --- a/kernels/m/swin_common.h +++ b/kernels/m/swin_common.h @@ -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 __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 __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) diff --git a/kernels/native/build_native.sh b/kernels/native/build_native.sh new file mode 100755 index 0000000..73d2ff3 --- /dev/null +++ b/kernels/native/build_native.sh @@ -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' ')" diff --git a/kernels/native/ctile.py b/kernels/native/ctile.py new file mode 100755 index 0000000..e72f935 --- /dev/null +++ b/kernels/native/ctile.py @@ -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 diff --git a/kernels/native/ifconv.py b/kernels/native/ifconv.py new file mode 100755 index 0000000..af1d81f --- /dev/null +++ b/kernels/native/ifconv.py @@ -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: bra.uni $M; + $B: (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 diff --git a/kernels/native/in_fuse.py b/kernels/native/in_fuse.py new file mode 100755 index 0000000..34c232c --- /dev/null +++ b/kernels/native/in_fuse.py @@ -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 + ['}'])) diff --git a/kernels/native/in_regs.hip b/kernels/native/in_regs.hip new file mode 100644 index 0000000..8953d50 --- /dev/null +++ b/kernels/native/in_regs.hip @@ -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]; +} diff --git a/kernels/native/m_ptx_rt.h b/kernels/native/m_ptx_rt.h new file mode 100644 index 0000000..229b052 --- /dev/null +++ b/kernels/native/m_ptx_rt.h @@ -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); +} diff --git a/kernels/native/native_rt.h b/kernels/native/native_rt.h new file mode 100644 index 0000000..861ce06 --- /dev/null +++ b/kernels/native/native_rt.h @@ -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])))}; +} + diff --git a/kernels/native/output_tile.h b/kernels/native/output_tile.h new file mode 100644 index 0000000..dd9c786 --- /dev/null +++ b/kernels/native/output_tile.h @@ -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 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() diff --git a/kernels/native/ptx_rt.h b/kernels/native/ptx_rt.h new file mode 100644 index 0000000..3761833 --- /dev/null +++ b/kernels/native/ptx_rt.h @@ -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 + +#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"); +} diff --git a/kernels/native/rewrite_output.py b/kernels/native/rewrite_output.py new file mode 100755 index 0000000..ae31309 --- /dev/null +++ b/kernels/native/rewrite_output.py @@ -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='') diff --git a/kernels/tex/tex_common.h b/kernels/tex/tex_common.h index 67b4531..3910696 100644 --- a/kernels/tex/tex_common.h +++ b/kernels/tex/tex_common.h @@ -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); diff --git a/packaging/SOURCES.txt b/packaging/SOURCES.txt index 8caff54..aea1ba1 100644 --- a/packaging/SOURCES.txt +++ b/packaging/SOURCES.txt @@ -10,7 +10,7 @@ Built from source by d4r (Apache License 2.0, d4r/licenses/d4r-LICENSE.txt): Third-party software, unmodified or patched by d4r: d4r/zluda/libcuda.so ZLUDA (Apache License 2.0 or MIT), https://github.com/vosen/ZLUDA - commit @ZLUDA_COMMIT@ with d4r/source/patches/zluda/0002 to 0007 applied. + commit @ZLUDA_COMMIT@ with d4r/source/patches/zluda/0002 to 0008 applied. It contains LLVM (Apache License 2.0 with LLVM exceptions). Licenses: d4r/licenses/ZLUDA-LICENSE-*.txt, d4r/licenses/LLVM-LICENSE.txt d3d12.dll, d3d12core.dll diff --git a/patches/zluda/0008-native-post-phase-kernels.patch b/patches/zluda/0008-native-post-phase-kernels.patch new file mode 100644 index 0000000..021bde2 --- /dev/null +++ b/patches/zluda/0008-native-post-phase-kernels.patch @@ -0,0 +1,122 @@ +ZLUDA: post-phase kernels for native overrides. A native code object may define NAME_post1 .. NAME_post8; the hook +launches them after NAME on the same stream with the same parameters (u32 globals d4r_postK_grid_x/_grid_y, 0 = the +launch's grid; d4r_postK_block_x/_y/_z, 0 = 32/1/1). Preset K's bottleneck layer (dec5) runs as four phase kernels. + +diff --git a/zluda/src/impl/function.rs b/zluda/src/impl/function.rs +index 033d9a0..69c70c9 100644 +--- a/zluda/src/impl/function.rs ++++ b/zluda/src/impl/function.rs +@@ -26,6 +26,9 @@ pub(crate) struct Function { + /// the kernel prepares into a shared overflow slot). 0 = only the last key is kept. + pub(crate) prep_slots: u32, + pub(crate) prep_seen: std::sync::Mutex>, ++ /// d4r phase kernels NAME_post1.. launched after the main kernel, on the same stream with the same ++ /// parameters: (function, grid x, grid y, block x, block y, block z); a grid entry of 0 keeps the launch's value ++ pub(crate) post: Vec<(hipFunction_t, u32, u32, u32, u32, u32)>, + } + + impl<'a, E: zluda_common::CudaErrorType> zluda_common::FromCuda<'a, CUfunction, E> +@@ -220,7 +223,25 @@ pub(crate) fn launch_kernel( + kernel_params, + extra_for_hip, + ) ++ }?; ++ for &(post, gx, gy, bx, by, bz) in f.post.iter() { ++ unsafe { ++ hipModuleLaunchKernel( ++ post, ++ if gx != 0 { gx } else { grid_dim_x }, ++ if gy != 0 { gy } else { grid_dim_y }, ++ if gx != 0 || gy != 0 { 1 } else { grid_dim_z }, ++ bx, ++ by, ++ bz, ++ 0, ++ stream, ++ kernel_params, ++ extra_for_hip, ++ ) ++ }?; + } ++ Ok(()) + } + + pub(crate) unsafe fn set_attribute( +diff --git a/zluda/src/impl/module.rs b/zluda/src/impl/module.rs +index 9f72f2f..3173500 100644 +--- a/zluda/src/impl/module.rs ++++ b/zluda/src/impl/module.rs +@@ -37,12 +37,12 @@ impl ModuleMutable { + Ok(match self.functions.entry(name) { + hash_map::Entry::Occupied(entry) => &*entry.into_mut(), + hash_map::Entry::Vacant(entry) => { +- let (func_handle, prep, block_z, grid_x, prep_key, prep_slots) = match native_override(entry.key()) { ++ let (func_handle, prep, block_z, grid_x, prep_key, prep_slots, post) = match native_override(entry.key()) { + Some(native) => native, + None => { + let mut func_handle = unsafe { std::mem::zeroed() }; + unsafe { hipModuleGetFunction(&mut func_handle, module, entry.key().as_ptr()) }?; +- (func_handle, None, 0, 0, None, 0) ++ (func_handle, None, 0, 0, None, 0, Vec::new()) + } + }; + let func = Box::new(Function { +@@ -56,6 +56,7 @@ impl ModuleMutable { + prep_last: std::sync::atomic::AtomicU64::new(0), + prep_slots, + prep_seen: std::sync::Mutex::new(Vec::new()), ++ post, + }); + &*entry.insert(func) + } +@@ -69,9 +70,10 @@ impl ModuleMutable { + /// same stream with `d4r_prep_blocks` (a u32 global in the code object) blocks of 128 threads. + /// An optional u32 global `d4r_block_z` replaces the launch's block z dimension, and `d4r_grid_x` + /// (persistent kernels, which derive the original grid from their parameters) the grid. ++#[allow(clippy::type_complexity)] + fn native_override( + name: &CStr, +-) -> Option<(hipFunction_t, Option<(hipFunction_t, u32)>, u32, u32, Option, u32)> { ++) -> Option<(hipFunction_t, Option<(hipFunction_t, u32)>, u32, u32, Option, u32, Vec<(hipFunction_t, u32, u32, u32, u32, u32)>)> { + let dir = std::env::var_os("D4R_ZLUDA_NATIVE_DIR")?; + let path = std::path::Path::new(&dir).join(format!("{}.hsaco", name.to_str().ok()?)); + let image = fs::read(&path).ok()?; +@@ -128,17 +130,35 @@ fn native_override( + }; + // d4r_prep_key_slots: the kernel keeps one prepared image per key (see Function::prep_slots) + let prep_slots = read_u32(c"d4r_prep_key_slots"); ++ // d4r: phase kernels NAME_post1 .. NAME_post8, each with u32 globals d4r_postK_grid_x / _grid_y (0 = the ++ // launch's grid) and d4r_postK_block_x / _y / _z (0 = 32 / 1 / 1) ++ let mut post = Vec::new(); ++ for k in 1..=8u32 { ++ let Ok(post_name) = CString::new(format!("{}_post{}", name.to_str().ok()?, k)) else { break }; ++ let mut post_func = unsafe { mem::zeroed() }; ++ if unsafe { hipModuleGetFunction(&mut post_func, native_module, post_name.as_ptr()) }.is_err() { ++ break; ++ } ++ let sym = |field: &str| CString::new(format!("d4r_post{}_{}", k, field)).unwrap(); ++ let gx = read_u32(&sym("grid_x")); ++ let gy = read_u32(&sym("grid_y")); ++ let bx = read_u32(&sym("block_x")); ++ let by = read_u32(&sym("block_y")); ++ let bz = read_u32(&sym("block_z")); ++ post.push((post_func, gx, gy, if bx == 0 { 32 } else { bx }, if by == 0 { 1 } else { by }, if bz == 0 { 1 } else { bz })); ++ } + eprintln!( +- "[d4r-native] {:?} -> {} (prep: {:?}, block_z: {}, grid_x: {}, prep_key: {:?}, prep_slots: {})", ++ "[d4r-native] {:?} -> {} (prep: {:?}, block_z: {}, grid_x: {}, prep_key: {:?}, prep_slots: {}, post: {})", + name, + path.display(), + prep.map(|p| p.1), + block_z, + grid_x, + prep_key, +- prep_slots ++ prep_slots, ++ post.len() + ); +- Some((func, prep, block_z, grid_x, prep_key, prep_slots)) ++ Some((func, prep, block_z, grid_x, prep_key, prep_slots, post)) + } + + impl ZludaObject for Module { diff --git a/scripts/build_release_glibc241.sh b/scripts/build_release_glibc241.sh index 69be734..28e5b31 100755 --- a/scripts/build_release_glibc241.sh +++ b/scripts/build_release_glibc241.sh @@ -4,7 +4,7 @@ set -euo pipefail ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" : "${D4R_RELEASE_BUILD_ROOT:?set the build output directory}" -: "${D4R_ZLUDA_SRC:?set a build-local ZLUDA checkout with patches 0002-0007 applied}" +: "${D4R_ZLUDA_SRC:?set a build-local ZLUDA checkout with patches 0002-0008 applied}" : "${D4R_VKD3D_SRC:?set a build-local vkd3d-proton checkout with patches 0001 and 0002 applied}" : "${D4R_ROCM_DIR:?set ROCm with clang and device libraries}" BUILD="$(realpath -m "$D4R_RELEASE_BUILD_ROOT")"