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

Add RDNA4 (gfx12) support and a 0.1.2 optimization round

Browse Source 
RDNA4 (experimental; untested on RDNA4 hardware)
- kernels/common/wmma_layout.h picks the WMMA operand/accumulator layout
  per target; K (pwin_*) and M (swin_*) index through it, so the same
  sources build for gfx11 and gfx12. D4R_WMMA_LAYOUT=12 on gfx11 is a test
  shim that runs the gfx12 indexing on gfx11 WMMA.

- M can use native e4m3 FP8 WMMA on gfx12 (D4R_NATIVE_FP8=1). Texture
  tails (enc0_tail, dec0_head) get gfx12 f16 and FP8 paths that match
  ZLUDA's gfx12 lowering.

- ZLUDA patch 0007: gfx12 WMMA helpers, optional native FP8 MMA
  (D4R_ZLUDA_WMMA_FP8_NATIVE), a layout override for testing, and an
  architecture argument for d4r_emit. Texture kernels are now compiled
  offline per target through d4r_emit (D4R_ZLUDA_EMIT) instead of on the
  local GPU.

- New d4r.ini option [Kernels] NativeFp8 (default on), wired through the
  shim and d4r_config.py. On gfx12 the bridge serves kernels/<arch>-fp8
  when it is on.

- The bridge picks the KFD GPU with the most SIMDs (so an iGPU is not
  chosen over a discrete GPU); D4R_GPU_ARCH overrides it.

- kernels/build.sh and package_release.sh build gfx1100-gfx1103 and
  gfx1200/gfx1201 (plus -fp8 variants), and D4R_BUNDLE_TEX takes
  per-target folders.

Performance (RDNA3; harness byte-identical unless noted)
- ZLUDA patch 0006: m16n8k8 f16 MMAs on WMMA (E; not bit-exact, 52-55 dB
  per layer, D4R_ZLUDA_WMMA_K8=0 disables it), per-key weight-image slots
  for native prep kernels (d4r_prep_key_at, d4r_prep_key_slots), and
  wave64 builds of MMA-free kernels (D4R_ZLUDA_WAVE64).

- M: packed FP8 encoding, conflict-free patch-merge reads, scalar weight
  loads in enc1, and weight prep once per weight set for the six enc3
  tube blocks.

- K: CU mode for the six non-wide pwin layers.

- Texture kernels: every DLSS flag variant is built natively (24 K output,
  8 M post, 4 M downsample, 4 M enc0), not only Townfall's; post and K
  output run as wave64.

- M direct output: rrlite_downsample_kernel writes the shim's output
  buffer; the shim's default direct-output presets are now 11,13.

Tools
- dump_runner takes the prep grid from d4r_prep_blocks and refuses
  texture/surface dumps.

- New pwin_check.py, swin_check.py, model_enc3.py and tail_check.hip.

Docs: README GPU table and RDNA4 notes, building (patches 0002-0007,
d4r_emit), native-kernels (RDNA4 section), performance (new numbers),
packaging README. Version 0.1.2.
This commit is contained in:
ayo committed 2026-09-29 20:36:10 -04:00
1 parent f310d52b69
commit 4334a0c3c4
42 files changed
+3068 -529
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

README.md
+52 -22
View File
@@ -1,10 +1,10 @@
# d4r (dlss 4 radeon)
d4r runs NVIDIA's official DLSS Super Resolution library (`nvngx_dlss.dll`) in Windows games on AMD Radeon GPUs under Linux and Proton. The game asks for DLSS as usual; the DLSS network runs on the AMD GPU through [ZLUDA](https://github.com/vosen/ZLUDA) (CUDA on ROCm/HIP), with the heaviest DLSS kernels replaced by hand-written RDNA3 code.
d4r runs NVIDIA's official DLSS Super Resolution library (`nvngx_dlss.dll`) in Windows games on AMD Radeon GPUs under Linux and Proton. The game asks for DLSS as usual; the DLSS network runs on the AMD GPU through [ZLUDA](https://github.com/vosen/ZLUDA) (CUDA on ROCm/HIP), with the heaviest DLSS kernels replaced by hand-written RDNA3 and RDNA4 code.
**Supported DLSS models:** DLSS 3 CNN (E), DLSS 4 transformer (K, default), and DLSS 4.5 transformer (M). DLSS 5 is not supported.
It is an experiment. It has been tested on one GPU in a handful of games, it depends on unreleased patches to ZLUDA and vkd3d-proton.
**Proof of concept:** d4r shows that DLSS can run on an AMD GPU, but it is not really that practical for everyday use yet. It has been tested on one GPU in a handful of games and depends on unreleased patches to ZLUDA and vkd3d-proton.
See [supported games](SUPPORTED_GAMES.md) for the tested games and DLSS models.
@@ -16,12 +16,14 @@ See [supported games](SUPPORTED_GAMES.md) for the tested games and DLSS models.
Radeon RX 7700 XT (RDNA3, gfx1101), SILENT HILL Townfall at 2560×1440, every DLSS result is presented in the frame it belongs to (no added latency).
| Mode | DLSS 3 CNN (preset E) | DLSS 4 (preset K) | DLSS 4.5 (preset M) | FSR 4 |
| Mode | DLSS 3 CNN (preset E) | DLSS 4 (preset K) | DLSS 4.5 (preset M) | FSR 4\* |
|---|---|---|---|---|
| Quality (1705×960) | **72.4** fps | 69.4 | 51.5 | 76.0 |
| Balanced (1488×837) | **80.5** | 76.6 | 59.6 | 84.5 |
| Performance (1280×720) | **87.8** | 84.0 | 69.9 | 94.0 |
| Ultra Performance (853×480) | 88.9 | **94.5** | 90.3 | 107.3 |
| Quality (1705×960) | **71.7** fps | 67.9 | 49.3 | 76.0 |
| Balanced (1488×837) | **80.8** | 75.9 | 58.0 | 84.5 |
| Performance (1280×720) | **89.7** | 84.1 | 69.0 | 94.0 |
| Ultra Performance (853×480) | 89.1 | **95.5** | 92.2 | 107.3 |
\* FSR 4 was measured in an earlier session (2026-09-27); the DLSS columns on 2026-09-29, when the same machine ran about 2–4% slower overall. Run back to back on the same day, this release is faster than 0.1.1 at Quality with every model (E 71.2 → 71.9, K 67.9 → 68.1, M 48.4 → 49.3 fps), and more in games whose DLSS settings differ from Townfall's, whose kernel variants now run natively too.
For reference, native 2560×1440 without upscaling (the game's TSR at 100%) runs at 49.1 fps on the same walk.
@@ -33,16 +35,19 @@ On the same GPU DLSS starts at a disadvantage: its networks were designed for NV
## GPU support
d4r needs RDNA3 (Radeon RX 7000 series), but full speed depends on the chip: the release ships its native kernels only for **gfx1101 (Navi 32)**, and on other GPUs DLSS 4 and 4.5 fall back to the much slower translated NVIDIA code.
d4r builds for RDNA3 and RDNA4. A newly built release compiles native DLSS 4 and 4.5 network kernels for the targets below and selects the matching set at runtime. Only the RX 7700 XT has been tested on a real GPU; RDNA4 runtime and performance remain unverified on hardware.
| GPU | Chip | DLSS 3 CNN (E) | DLSS 4 (K) | DLSS 4.5 (M) |
|---|---|---|---|---|
| RX 7700 XT (tested), RX 7800 XT, RX 7700, Radeon PRO W7700 | gfx1101 | full speed | full speed | full speed |
| RX 7900 GRE / XT / XTX, Radeon PRO W7800 / W7900 | gfx1100 | full speed | runs without its transformer (worse image, slower) | about half speed |
| RX 7600 / 7600 XT, RX 7650 GRE | gfx1102 | full speed | runs without its transformer (worse image, slower) | about half speed |
| RDNA4, RDNA2 and older | | not supported | not supported | not supported |
| GPU | Chip | FP8 math | Runtime testing |
|---|---|---|---|
| RX 7700 XT, RX 7800 XT, RX 7700, Radeon PRO W7700 | gfx1101 | widened to f16 | RX 7700 XT only |
| RX 7900 GRE / XT / XTX, Radeon PRO W7800 / W7900 | gfx1100 | widened to f16 | untested |
| RX 7600 / 7600 XT, RX 7650 GRE | gfx1102 | widened to f16 | untested |
| RDNA3 integrated GPUs | gfx1103 | widened to f16 | untested |
| RX 9070 XT / 9070 / 9070 GRE, Radeon AI PRO R9700 | gfx1201 | native FP8 (`NativeFp8`, default on) | emulator only |
| RX 9060 XT / 9060 | gfx1200 | native FP8 (`NativeFp8`, default on) | compile only |
| RDNA2 and older | | | unsupported |
Only the RX 7700 XT has been tested; the other gfx1101 cards use the same chip. Without the native kernels, NVIDIA's translated DLSS 4 layers produce invalid values and DLSS discards the network's output, so K loses the transformer's image quality. Preset E is unaffected on any RDNA3 GPU. The kernels can be built for gfx1100 and gfx1102 (`D4R_GPU_ARCHS` in `scripts/package_release.sh`), but those builds are untested.
Native network kernels are built for all listed gfx11/gfx12 targets. Texture kernels are compiled for each target by `d4r_emit` when supplied to the package script; RDNA4 also has a `-fp8` variant. The bridge selects the KFD GPU with the most SIMDs, avoiding an integrated GPU when a discrete GPU is present; `D4R_GPU_ARCH` overrides that choice. Missing native kernels fall back to ZLUDA. Without the native network kernels, NVIDIA's translated DLSS 4 layers produce invalid values and DLSS discards the network's output, so K loses the transformer's image quality. Preset E does not use the native network kernels.
## How it works
@@ -57,13 +62,13 @@ game (D3D12) ──► OptiScaler (DLSS inputs) ──► d4r_nvngx.dll (NGX D3
▼
ZLUDA (patches/zluda: PTX → AMDGPU, WMMA, native kernel overrides)
▼
native RDNA3 kernels (kernels/: DLSS 4 and 4.5 network layers)
native RDNA3 and RDNA4 kernels (kernels/: DLSS 4 and 4.5 network layers)
```
- **The shim** (`d4r_nvngx.dll`) implements the D3D12 NGX entry points OptiScaler calls, copies the game's colour, depth and motion vectors into buffers shared with HIP, evaluates DLSS through the CUDA version of NGX, and writes the result back into the game's output texture.
- **The bridge** is a Wine builtin `nvcuda.dll` that forwards the CUDA driver API to ZLUDA on the Linux side, plus a few helpers the shim needs (Vulkan memory import, asynchronous array copies, GPU-side waits).
- **ZLUDA** compiles NVIDIA's PTX for the AMD GPU. The patches add what DLSS needs (textures, surfaces, FP8 and tensor-core MMA on RDNA3 WMMA) and a hook that serves hand-written kernels in place of selected PTX kernels.
- **Native kernels** reimplement the DLSS 4 and 4.5 network layers for RDNA3 and replace parts of a few texture-heavy kernels. They are 2–3× faster than the translated versions and are checked against numpy reference models of the layers.
- **ZLUDA** compiles NVIDIA's PTX for the AMD GPU. The patches add what DLSS needs (textures, surfaces, FP8 and tensor-core MMA on RDNA3/RDNA4 WMMA) and a hook that serves hand-written kernels in place of selected PTX kernels.
- **Native kernels** reimplement the DLSS 4 and 4.5 network layers for RDNA3 and RDNA4 and replace parts of a few texture-heavy kernels. The 2–3× speedup was measured on the RX 7700 XT; RDNA4 speed is unverified.
Details: [docs/architecture.md](docs/architecture.md) and [docs/native-kernels.md](docs/native-kernels.md).
@@ -81,7 +86,7 @@ ROCm does not need to be installed: the zip includes its runtime (from AMD's Ubu
## Requirements (building from source)
- Linux with an AMD RDNA3 GPU. The native kernels use gfx11 WMMA; they are built for gfx1101 by default (`D4R_GPU_ARCH` for other RDNA3 chips, untested). Without them DLSS runs through ZLUDA alone: much more slowly, and DLSS 4 (K) without its transformer (see [GPU support](#gpu-support)).
- Linux with an AMD RDNA3 or RDNA4 GPU. The native kernels use gfx11 or gfx12 WMMA. `scripts/package_release.sh` builds network kernels for gfx1100–gfx1103 and gfx1200–gfx1201 by default; `D4R_GPU_ARCH` selects one target for a standalone `kernels/build.sh` invocation. Runtime support outside the RX 7700 XT remains unverified on hardware (see [GPU support](#gpu-support)).
- ROCm with HIP and its clang (tested with ROCm 7.2).
- GE-Proton with OptiScaler integration (tested with GE-Proton11-3).
- Build tools: a Rust toolchain and git-lfs (ZLUDA), meson and ninja (vkd3d-proton), `winegcc`/`winebuild` (bridge), `x86_64-w64-mingw32-g++` and `clang-cl` (shim), Python 3.
@@ -91,27 +96,52 @@ 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`–`0005` applied.
1. Build ZLUDA `ee2f25a` with `patches/zluda/0002`–`0007` 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`.
5. Build the native kernels: `D4R_ROCM_DIR=… D4R_DLSS_DLL=… D4R_ZLUDA_BUILD=… kernels/build.sh`.
5. Build the native kernels: `D4R_ROCM_DIR=… D4R_DLSS_DLL=… D4R_ZLUDA_EMIT=… kernels/build.sh`.
6. Fill in `~/.config/d4r/d4r.ini` (created from `config/d4r.ini.default` on first launch) and start the game with `scripts/d4r_play.sh`.
The launcher backs up and restores everything it touches in the Proton prefix: OptiScaler.ini, the prefix's `d3d12.dll`/`d3d12core.dll`, and the game's Engine.ini when cvars are set.
## Preset videos
Gameplay recordings of each DLSS model at the Quality, Performance and Ultra Performance presets.
**Ready or Not**
| Model | Quality | Performance |
|---|---|---|
| DLSS 4.5 | [video](https://cdn.ayois.gay/sk-ant-api03-OPG1vqqsmDkNduG5fXJBelDI-1APyL9_ytJrnOSekZxoh-2NL8NqNuqL3pCI24kJQsudaB896GYpu5TFv93tabAgK1H2sAA) | [video](https://cdn.ayois.gay/sk-ant-api03-d-4XS1it5mtMkm-QRz0Qnkfu9FAjz_dFJURvsh_3gPZTnC1sljpWc28F6fyjUWIC6xlZm-FckWGjaP-KuIipYmroa3dQ8AA) |
| DLSS 4 | [video](https://cdn.ayois.gay/sk-ant-api03-uXTWm4hN5jZksiSm4fvgIOr0dw43hGJ0TpZLe2ZghIbGb7aR9CGGlTV-1paBYO7G6lzIf83tOrCdaKdJerJ4re6epRWDEAA) | [video](https://cdn.ayois.gay/sk-ant-api03-fcZCkctdFOv513z1u6rPBrRseY47Qnw43Gnv5mZf_GekHw-yUoWYcBixjp2qqqbBIf7sVmht0Y676U96n2cwmFIsRgWkHAA) |
| DLSS 3 | [video](https://cdn.ayois.gay/sk-ant-api03-_vuywLcsMzZiJxcNgucBeeXFzWi44mes5A1777bg_4XQ2X0MH22J9M7JFpTS0_q6yaGWMXf3FvL8MDiv_dc7UANxrCmlGAA) | [video](https://cdn.ayois.gay/sk-ant-api03-WzaIy34WIesQeOOTVQNXZwpPPE8khmbUynvMw9B-Luj8GFcZ_sJvGXfg6aW00TD0Rr3nWu43ELJGx2K6qs_HL6FkLdmb1AA) |
**Townfall**
| Model | Quality | Ultra Performance |
|---|---|---|
| DLSS 4.5 | [video](https://cdn.ayois.gay/sk-ant-api03-p8uL0hOUklvA-BSeEafS8LQU-gayzbqAE3NacWAmckLoXOSf2sXGf7OeG1LmvOqeGFRh_ceCWi4pLBl9yxJlvYwQWH3Y1AA) | [video](https://cdn.ayois.gay/sk-ant-api03-Mz0BzYTrLjaYSuQXtOEHTw4ydYVbPvHmMjflm8UhFp976ftXqMExD8BCMYlo93JLEGwVL3rMWoZR0CH_deJJU7jX8_s3GAA) |
| DLSS 4 | [video](https://cdn.ayois.gay/sk-ant-api03-WgNoUPTwH0vt55Ag7RaWprXXINU5OOOEtuDXOvIjHgfAScSiIPIJGVaERoF-EA3JO7LS8xkohUlb6AFsW8AVZ2vMB14JwAA) | [video](https://cdn.ayois.gay/sk-ant-api03-skXZVD6G4DUdAfQw0RLz_UJW8GOfBlpZx1N4ezaF_UIXm7OIY1ciQysfRu4zRsxSXA8plZJi4NSzkXkV7Oxmg9ytg9zaaAA) |
## Repository layout
| Path | Contents |
|---|---|
| `tools/` | the NGX shim, the Wine CUDA bridge, a D3D12 DLSS harness, probes and kernel replay tools |
| `kernels/` | native RDNA3 kernels (`k/` DLSS 4, `m/` DLSS 4.5, `tex/` texture-kernel parts), their build script, validation tools and numpy reference models |
| `kernels/` | native RDNA3/RDNA4 kernels (`k/` DLSS 4, `m/` DLSS 4.5, `tex/` texture-kernel parts, `common/` WMMA layouts), their build script, validation tools and numpy reference models |
| `patches/` | ZLUDA and vkd3d-proton patches |
| `scripts/` | build, install, launch and probe scripts |
| `config/` | the default `d4r.ini` for the developer launcher |
| `packaging/` | the release's `d4r.ini`, OptiScaler settings, user guide and install check |
| `docs/` | architecture, build, native kernel and performance notes |
## Stars
[![Star History Chart](https://api.star-history.com/svg?repos=countervolts/d4r&type=Date)](https://star-history.com/#countervolts/d4r&Date)
## License
Apache License 2.0 (see [LICENSE](LICENSE)). The patches in `patches/` are offered under the licenses of the projects they modify: ZLUDA (Apache-2.0 or MIT) and vkd3d-proton (LGPL-2.1). See [NOTICE](NOTICE).
To contact me my discord is `._ayo`.
config/d4r.ini.default
+4 -1
View File
@@ -36,7 +36,7 @@ Model = K
FrameAge = 0
[Kernels]
; Native RDNA3 kernels for the DLSS 4 and 4.5 network layers (replace ZLUDA's translated code).
; Native RDNA3 / RDNA4 kernels for the DLSS 4 and 4.5 network layers (replace ZLUDA's translated code).
; fast = native kernels with relaxed numerics (visually identical to exact, ~15% faster)
; exact = native kernels that reproduce NVIDIA's FP8/f16 rounding bit for bit
; off = ZLUDA-translated kernels only
@@ -44,6 +44,9 @@ FrameAge = 0
NativeKernels = off
; FP8 tensor-core emulation on RDNA3 WMMA (much faster than the scalar path).
Fp8Wmma = true
; RDNA4 only: FP8 MMAs on the native FP8 WMMA (and the -fp8 native kernel variants); off widens FP8 to f16 as
; on RDNA3. Ignored on other GPUs. (D4R_ZLUDA_WMMA_FP8_NATIVE)
NativeFp8 = true
; Drop per-instruction flush-to-zero mode switches in ZLUDA-compiled kernels
; (only values below ~1e-38 differ; a little faster).
IgnoreDenormals = true
docs/building.md
+13 -6
View File
@@ -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`, `0003`, `0004`, `0005` in that order |
| ZLUDA | `ee2f25a` (upstream), plus `patches/zluda/0002` through `0007` in order |
| vkd3d-proton | `3dfc6f07` (the base GE-Proton11-3 ships), plus `patches/vkd3d-proton/0001` |
| ROCm | 7.2 (HIP runtime, clang, device libraries) |
| Proton | GE-Proton11-3 (with its OptiScaler integration) |
@@ -17,10 +17,11 @@ 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; do git apply /path/to/d4r/patches/zluda/$p-*.patch; done
for p in 0002 0003 0004 0005 0006 0007; 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
LIBRARY_PATH=/opt/rocm/lib cargo build --release -p ptx --example d4r_emit
mkdir -p ~/.cache/d4r-zluda-current
cp target/release/libnvcuda.so ~/.cache/d4r-zluda-current/
ln -sf libnvcuda.so ~/.cache/d4r-zluda-current/libcuda.so
@@ -34,6 +35,10 @@ What the patches add:
- **0003**: an implicit 256-thread launch bound for kernels without PTX bounds (removes massive register spilling, `D4R_ZLUDA_IMPLICIT_MAX_BLOCK`) and f16 tensor-core MMA on RDNA3 WMMA (`D4R_ZLUDA_WMMA`).
- **0004**: the native kernel override hook (`D4R_ZLUDA_NATIVE_DIR`, see [native-kernels.md](native-kernels.md)), FP8 MMA on WMMA (`D4R_ZLUDA_WMMA_FP8`), optional elision of per-instruction denormal mode switches (`D4R_ZLUDA_IGNORE_DENORMAL`), inlined image helpers, and linking extra bitcode into a PTX module (`D4R_ZLUDA_EXTRA_BC`, used by the texture-kernel build).
- **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.
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.
## 2. vkd3d-proton
@@ -65,11 +70,11 @@ This stages the shim, the bridge, the NGX core and the DLSS feature library in `
```sh
D4R_ROCM_DIR=/opt/rocm \
D4R_DLSS_DLL=/path/to/nvngx_dlss.dll \
D4R_ZLUDA_BUILD=~/.cache/d4r-zluda-current \
D4R_ZLUDA_EMIT=/path/to/zluda/target/release/examples/d4r_emit \
kernels/build.sh all kernels/out/native
```
`kernels/out/native` then holds one code object per replaced DLSS kernel; set `NativeKernelDirFast` in d4r.ini (or `D4R_ZLUDA_NATIVE_DIR`) to it. `kernels/build.sh k` or `m` builds only the network layers and needs neither the DLL nor ZLUDA. The texture kernels (`tex`) extract PTX from your DLL into `kernels/extracted/`; that directory and the build output are git-ignored and must not be redistributed.
`kernels/out/native` then holds one code object per replaced DLSS kernel; set `NativeKernelDirFast` in d4r.ini (or `D4R_ZLUDA_NATIVE_DIR`) to it. `kernels/build.sh k` or `m` builds only the network layers and needs neither the DLL nor ZLUDA. The texture kernels (`tex`) extract PTX from your DLL into `kernels/extracted/` and compile it offline for `D4R_GPU_ARCH` through `D4R_ZLUDA_EMIT`; that directory and the build output are git-ignored and must not be redistributed. On gfx12, build a second variant with `D4R_NATIVE_FP8=1` and place it in an `<arch>-fp8` folder.
## 6. Configure and play
@@ -93,14 +98,16 @@ scripts/package_release.sh # -> dist/d4r-<version>.zip
```
The script:
- builds the shim, the bridge and the network-layer kernels;
- builds the shim, the bridge and the network-layer kernels for gfx1100–gfx1103 and gfx1200–gfx1201 by default (`D4R_GPU_ARCHS` can select fewer targets); each gfx12 target also gets an `<arch>-fp8` variant;
- writes the kernel manifest from the DLLs you list (it records hashes of their PTX, nothing else);
- stages OptiScaler as `dxgi.dll` with the settings in `packaging/optiscaler.settings`;
- adds NVIDIA's two DLLs and the texture kernels (built from NVIDIA's PTX by `kernels/build.sh tex`);
- adds NVIDIA's two DLLs and any supplied texture kernels (built from NVIDIA's PTX by `kernels/build.sh tex`);
- zips the result together with the ZLUDA and vkd3d-proton builds, the ROCm runtime (as `d4r/rocm`), `packaging/d4r.ini`, the licenses and the patches.
The NVIDIA files are not covered by d4r's license; redistributing them is up to whoever publishes the zip. `D4R_BUNDLE_NVIDIA=0` builds `d4r-<version>-nonvidia.zip` without them and without the texture kernels. [architecture.md](architecture.md#portable-installs-the-release-zip) describes how the installed files work together.
`D4R_BUNDLE_TEX` accepts the older flat directory for gfx1101, or a directory with per-target subdirectories (`gfx1100/`–`gfx1103/`, `gfx1200/`, `gfx1201/`, plus `gfx1200-fp8/` and `gfx1201-fp8/`). Build each texture set with `D4R_GPU_ARCH` and `D4R_ZLUDA_EMIT`; the `-fp8` sets also need `D4R_NATIVE_FP8=1`. `d4r_emit` targets those GPUs offline. Missing texture kernels fall back to ZLUDA on that target; the network-layer kernels are still included. Only the RX 7700 XT has been tested on real hardware.
## Checks
- `scripts/check_environment.sh` lists the tools, GPU and Proton builds it finds.
docs/native-kernels.md
+39 -21
View File
@@ -1,6 +1,6 @@
# Native kernels
The native kernels are RDNA3 (gfx11) code objects that ZLUDA runs in place of specific DLSS kernels. Each one has the same name, launch shape and parameter block as the PTX kernel it replaces, and reads the same buffers, including the network weights inside NVIDIA's DLL. No weights or NVIDIA code are stored here.
The native kernels are RDNA3 (gfx11) and RDNA4 (gfx12) code objects that ZLUDA runs in place of specific DLSS kernels. Each one has the same name, launch shape and parameter block as the PTX kernel it replaces, and reads the same buffers, including the network weights inside NVIDIA's DLL. No weights or NVIDIA code are stored here.
## The override hook (patches/zluda/0004)
@@ -10,7 +10,9 @@ With `D4R_ZLUDA_NATIVE_DIR=DIR`, whenever NGX loads a PTX module and asks for ke
|---|---|
| `NAME_prep` kernel | launched first on the same stream, with the same parameters (typically re-lays out weights) |
| `u32 d4r_prep_blocks` | grid size of the prep kernel (128 threads per block) |
| `u32 d4r_prep_key_offset` | byte offset of a u64 in the parameter block; the prep is skipped while that value (the weights pointer) repeats |
| `u32 d4r_prep_key_offset` | byte offset of a u64 in the parameter block; the prep is skipped while that value (the weights pointer) repeats (0 = no key) |
| `u32 d4r_prep_key_at` | the same as offset + 1, so a key at offset 0 can be named (M's layers keep their weights pointer there) |
| `u32 d4r_prep_key_slots` | N > 0: the kernel keeps one prepared image per key for up to N keys, so the prep runs once per weight set (the six enc3 tube blocks); keys beyond N share an overflow image and are prepared on every launch |
| `u32 d4r_block_z` | replaces the launch's block z dimension (more waves per window) |
| `u32 d4r_grid_x` | replaces the grid (persistent kernels) |
@@ -19,7 +21,7 @@ Kernels without a file in `DIR` are compiled from PTX as usual, so a partial set
**Release directories.** A directory that holds `d4r-kernels.txt` (or `<gfx target>/d4r-kernels.txt`) is checked before use:
- `kernels/tools/kernel_manifest.py DIR nvngx_dlss.dll...` writes the manifest. For each `NAME.hsaco` it records the FNV-1a 64 hash of the PTX module that defines `.entry NAME`, taken from each DLL given.
- The CUDA bridge serves a kernel only after NGX has loaded a module with one of the listed hashes.
- It picks the target folder from the GPU's KFD topology entry.
- It picks the KFD GPU with the most SIMDs, then selects that GPU's target folder. `D4R_GPU_ARCH` overrides the choice. On gfx12, `D4R_ZLUDA_WMMA_FP8_NATIVE=1` selects the `<arch>-fp8` folder.
The release is built this way, with hashes from DLSS 310.7.0 and 310.9.1. The PTX of every replaced kernel is identical in those two versions.
@@ -34,8 +36,20 @@ The release is built this way, with hashes from DLSS 310.7.0 and 310.9.1. The PT
- **Weights:** prep kernels expand the weights into WMMA operand images once.
**DLSS 4.5, preset 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.
- **Network:** the Swin blocks of DLSS 4.5, whose weights are FP8. Their prep kernels expand the weights to f16 WMMA operands, once per weight set.
- **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.
## RDNA4
`kernels/common/wmma_layout.h` selects the WMMA operand and accumulator layout. gfx12 uses eight f16 values per lane in each K half; the accumulator rows are `i + 8·half`. The gfx11 path retains its original layout. `D4R_WMMA_LAYOUT=12` on gfx11 is a test shim: it exercises gfx12 indexing but executes gfx11 WMMA, so its replay output can be compared byte for byte with the existing gfx11 build.
The gfx12 M variant can use native e4m3 FP8 WMMA with `D4R_NATIVE_FP8=1`. Its activations are requantized to e4m3 at operand load, and its prepared weights remain bytes. The matching ZLUDA and texture-tail variant uses `D4R_ZLUDA_WMMA_FP8_NATIVE=1` / `D4R_TEX_FP8=1`; the release bridge chooses the `-fp8` kernel folder. `NativeFp8` in d4r.ini controls that choice and defaults to on. The f16 widening path remains available.
The gfx11 refactor and gfx12-layout shim matched the recorded gfx11 replays and 88-frame K/M harness images byte for byte on the RX 7700 XT. rocjitsu gfx1201 replay comparisons put the native K layers within 96–109 dB PSNR of gfx11, and M differed in 0.001–0.3% of e4m3 bytes. These are emulator results, not RDNA4 hardware results. rocjitsu's FP8 arithmetic is not hardware-checked; bit-exactness, speed and driver stability on an RDNA4 GPU remain unverified.
For ZLUDA-compiled code, compare rocjitsu runs with other rocjitsu runs: the emulator does not reproduce the denormal-flush mode of those code objects, so a GPU byte comparison is not a valid arithmetic check. Captures with texture or surface objects must not be sent to the replay tools; the standalone `tail_check.hip` and the full D3D12 harness cover those tails.
**Texture kernels (`kernels/tex`).**
- **The problem:** some DLSS kernels are mostly texture sampling and scalar maths that ZLUDA already compiles well, apart from a few slow parts.
@@ -43,8 +57,10 @@ The release is built this way, with hashes from DLSS 310.7.0 and 310.9.1. The PT
- surface stores become calls into `tex_common.h`;
- the `roundf` idiom that forces ZLUDA into strict-FP mode is rewritten exactly;
- for M's enc0 and dec0, a range of the kernel is replaced by a native tail or head.
- **Build:** ZLUDA compiles the edited PTX with the HIP bitcode linked in (`D4R_ZLUDA_EXTRA_BC`), and the resulting code object is saved.
- **Output redirect:** the K output kernel (`hiluma_engine_output_*`) can also write its result straight into the shim's output buffer instead of a CUDA array. See `D4R_SHIM_OUTPUT_DIRECT` in [architecture.md](architecture.md).
- **Build:** ZLUDA's `d4r_emit` compiles the edited PTX for `D4R_GPU_ARCH` with the HIP bitcode linked in (`D4R_ZLUDA_EXTRA_BC`), and the resulting code object is saved. This is an offline per-target build.
- **Every flag combination:** DLSS picks kernel variants by the game's settings (motion vectors at render or display resolution, HDR or LDR input, inverted or regular depth, ...). `kernels/build.sh` builds all of them: 24 K output kernels, 8 M post, 4 M downsample and 4 M enc0 kernels. enc0's native tail is spliced into each variant at its own line, with the registers read off the variant's PTX (`make_ptx.py` checks that the code after the cut matches the reference variant line for line).
- **Wave64:** kernels without MMAs can be compiled as wave64 (`D4R_ZLUDA_WAVE64=1`, ZLUDA patch 0006): each wave runs two CUDA warps, and RDNA3 issues FP32 work for all 64 lanes at once. The post and K output kernels are built this way (post 0.89 → 0.77 ms), with bit-identical results.
- **Output redirect:** the kernel that writes a preset's result can write it straight into the shim's output buffer instead of a CUDA array: `hiluma_engine_output_*` for K, `rrlite_downsample_kernel_*` for M (M's post kernel works at 3840×2160; the downsample produces the output). See `D4R_SHIM_OUTPUT_DIRECT` in [architecture.md](architecture.md).
## Numerics
@@ -61,7 +77,7 @@ kernels/build.sh tex # texture kernels (needs D4R_DLSS_DLL and D4R_ZLUDA_
kernels/build.sh all DIR # everything into DIR (default kernels/out/native)
```
`D4R_GPU_ARCH` selects the target (default gfx1101). All kernels need gfx11 WMMA.
`D4R_GPU_ARCH` selects one target for this standalone build (default gfx1101). The release script builds the network layers for gfx1100–gfx1103 and gfx1200–gfx1201 by default and puts each set in its own directory. The network kernels need gfx11 or gfx12 WMMA. Texture-kernel builds use `D4R_ZLUDA_EMIT` to generate code objects for the selected target without that GPU. Set `D4R_NATIVE_FP8=1` on gfx12 for the matching FP8 M and texture variants.
## Validating a kernel
@@ -72,23 +88,25 @@ kernels/build.sh all DIR # everything into DIR (default kernels/out/native)
PREP_GRID=<d4r_prep_blocks> dump_runner NAME.hsaco NAME <capture> <out> 30
```
It runs the prep and main kernels, saves every buffer after the first launch, and times 30 launches.
It runs the prep and main kernels, saves every buffer after the first launch, and times 30 launches. The replay tool refuses captures with texture or surface objects; use the full harness for those kernels.
3. **Compare against a reference.** `kernels/tools/pwin_model.py` (K) and `swin_model.py` (M) are numpy models of the layers, written stage by stage against `ptxsim.py`, a small vectorised PTX interpreter that runs NVIDIA's kernels on the CPU. A bit-exact variant should match the model within f16 rounding; after any change, compare the new build's buffers with the previous build's.
4. **Check the whole pipeline.** Replay captured frames through the harness with both kernel sets and compare the outputs with `kernels/tools/psnr.py`. `D4R_CUDA_KERNEL_PROFILE=1` makes the bridge time every kernel; `kernels/tools/kprof.py` summarises the log per frame.
## Per-kernel cost
GPU time per frame at Quality (1705×960 → 2560×1440, RX 7700 XT), from `D4R_CUDA_KERNEL_PROFILE`:
Median GPU time per frame at Quality (1706×960 → 2560×1440, RX 7700 XT), from `D4R_CUDA_KERNEL_PROFILE` in the D3D12 harness (flags as in Townfall; frames 21–40). The profiler's per-launch events add a little to every kernel, so totals are slightly above the in-game cost. "Before" is release 0.1.1.
| DLSS 4 (K) | ms | DLSS 4.5 (M) | ms |
|---|---|---|---|
| hiluma output (texture kernel) | 0.82 | enc3 tube (6 launches) | 1.42 |
| dec0 | 0.30 | enc1 | 1.27 |
| enc0 | 0.28 | dec1 | 1.18 |
| dec1 | 0.24 | enc2 | 0.79 |
| hiluma input (ZLUDA) | 0.24 | post (texture kernel) | 0.77 |
| enc1 | 0.21 | dec2 | 0.69 |
| enc2, dec2 | 0.10 each | enc0 (texture kernel) | 0.60 |
| enc3, dec3, enc4, dec4, dec5 | 0.05–0.08 each | dec0 (texture kernel) | 0.46 |
| NGX exposure and misc | 0.1 | downsample, NGX misc | 0.36 |
| **total** | **2.8** | **total** | **7.5** |
| DLSS 4 (K) | before | now | DLSS 4.5 (M) | before | now |
|---|---|---|---|---|---|
| hiluma output (texture kernel) | 0.92 | 0.89 | enc3 tube (6 launches) | 1.62 | 1.57 |
| dec0 | 0.33 | 0.32 | enc1 | 1.41 | 1.27 |
| enc0 | 0.31 | 0.31 | dec1 | 1.44 | 1.43 |
| dec1 | 0.26 | 0.26 | enc2 | 0.87 | 0.77 |
| hiluma input (ZLUDA) | 0.25 | 0.24 | post (texture kernel) | 0.88 | 0.77 |
| enc1 | 0.23 | 0.22 | dec2 | 0.84 | 0.82 |
| enc2, dec2 | 0.12, 0.10 | 0.12, 0.10 | enc0 (texture kernel) | 0.64 | 0.63 |
| enc3, dec3, enc4, dec4, dec5 | 0.07–0.09 each | 0.07–0.09 each | dec0 (texture kernel) | 0.64 | 0.63 |
| NGX exposure and misc | 0.1 | 0.1 | downsample, NGX misc | 0.32 | 0.33 |
| **total** | **3.10** | **2.97** | **total** | **9.04** | **8.23** |
Runs of the same build vary by about ±3% per kernel. M additionally skips its output copy now (direct output through the downsample kernel), which is not a kernel and not in the table. Games whose flags differ from Townfall's (motion vectors at render resolution, LDR input, regular depth) gain more: those kernel variants used to run as plain ZLUDA compiles (K 3.6–4.6 ms → 2.8–3.0 ms, M −0.5 to −0.7 ms per frame).
docs/performance.md
+10 -6
View File
@@ -1,5 +1,7 @@
# Performance
All frame rates and kernel timings on this page were measured on an RX 7700 XT (gfx1101). RDNA4 (gfx1200/gfx1201) has compile and emulator validation only; no RDNA4 frame rate, latency, or kernel timing has been measured. Native FP8 may change both arithmetic and cost, so the gfx11 results cannot be projected to it.
## Method
- **Hardware and settings:** Radeon RX 7700 XT (RDNA3, 54 CUs), Ryzen 9 5900XT, Linux 6.18. SILENT HILL Townfall at 2560×1440 output.
@@ -8,18 +10,20 @@
- **Pairing:** runs are compared back-to-back in one session, because repeated runs of an identical setup vary by about 1%.
- **Screen recording:** it costs a few fps, so recorded videos show slightly lower numbers than the table.
- **Latency:** all DLSS numbers are same-frame (frame age 0). Every frame shows its own DLSS result, as with a native upscaler.
- **Other upscaler modes:** Balanced, Performance and Ultra Performance use OptiScaler's render-ratio override (1.72, 2.0, 3.0) on top of the game's Quality mode.
- **Upscaler modes:** every mode uses OptiScaler's render-ratio override (1.5, 1.72, 2.0, 3.0), so the result does not depend on the DLSS mode chosen in the game.
## Results
| Mode (render resolution) | DLSS 3 CNN (E) | DLSS 4 (K) | DLSS 4.5 (M) | FSR 4 |
| Mode (render resolution) | DLSS 3 CNN (E) | DLSS 4 (K) | DLSS 4.5 (M) | FSR 4\* |
|---|---|---|---|---|
| Quality (1705×960) | **72.4** | 69.4 | 51.5 | 76.0 |
| Balanced (1488×837) | **80.5** | 76.6 | 59.6 | 84.5 |
| Performance (1280×720) | **87.8** | 84.0 | 69.9 | 94.0 |
| Ultra Performance (853×480) | 88.9 | **94.5** | 90.3 | 107.3 |
| Quality (1705×960) | **71.7** | 67.9 | 49.3 | 76.0 |
| Balanced (1488×837) | **80.8** | 75.9 | 58.0 | 84.5 |
| Performance (1280×720) | **89.7** | 84.1 | 69.0 | 94.0 |
| Ultra Performance (853×480) | 89.1 | **95.5** | 92.2 | 107.3 |
| Native 2560×1440, no upscaling (TSR at 100%) | | 49.1 | | |
\* FSR 4 and the native row are from 2026-09-27; the DLSS columns from 2026-09-29, when the machine ran about 2–4% slower overall (release 0.1.1 at Quality that day: E 71.2, K 67.9, M 48.4 fps, against 72.4, 69.4 and 51.5 two days earlier). The render resolution of every mode is set with OptiScaler's ratio override for whichever DLSS mode the game is set to (ratios 1.5, 1.72, 2.0 and 3.0).
DLSS 4's cost is almost constant across modes (about 2.8 ms of GPU time per frame). Its network runs on a grid set by the output resolution, not the render resolution, so the gap to FSR 4 widens as the render resolution drops. The CNN gets more expensive at the 3× ratio, so K is the better choice at Ultra Performance.
## What each step contributed
kernels/build.sh
+47 -17
View File
@@ -1,14 +1,19 @@
#!/usr/bin/env bash
# Builds the native RDNA3 replacements for DLSS kernels into one directory that ZLUDA serves them from
# Builds the native replacements for DLSS kernels into one directory that ZLUDA serves them from
# (D4R_ZLUDA_NATIVE_DIR, or NativeKernelDirFast in d4r.ini). See docs/native-kernels.md.
#
# usage: kernels/build.sh [all|k|m|tex] [OUT_DIR] (default: all, kernels/out/native)
#
# D4R_ROCM_DIR ROCm installation with clang and the HIP device libraries (default /opt/rocm)
# D4R_GPU_ARCH target GPU (default gfx1101; the kernels need RDNA3 / gfx11 WMMA)
# D4R_GPU_ARCH target GPU (default gfx1101): an RDNA3 (gfx110x) or RDNA4 (gfx120x) target; the kernels
# need wave32 WMMA, and kernels/common/wmma_layout.h picks the register layout per target
# D4R_NATIVE_FP8 1 (RDNA4 only): the variant for ZLUDA's native FP8 WMMA (d4r.ini [Kernels] NativeFp8):
# M's layers on FP8 WMMA with NVIDIA's e4m3 activations, texture tails matching ZLUDA's
# D4R_ZLUDA_WMMA_FP8_NATIVE lowering. K is the same in both variants.
# Texture kernels (tex) are NVIDIA's PTX with parts replaced, so they additionally need:
# D4R_DLSS_DLL nvngx_dlss.dll 310.7 (its PTX is extracted into kernels/extracted/, never committed)
# D4R_ZLUDA_BUILD a ZLUDA build with patches/zluda applied (directory holding libcuda.so)
# D4R_ZLUDA_EMIT ZLUDA's d4r_emit (patches/zluda applied; `cargo build --release -p ptx --example d4r_emit`),
# which compiles PTX offline for D4R_GPU_ARCH (no GPU of that kind needed)
set -euo pipefail
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -16,17 +21,26 @@ WHAT="${1:-all}"
OUT="$(realpath -m "${2:-$HERE/out/native}")"
ROCM="${D4R_ROCM_DIR:-/opt/rocm}"
ARCH="${D4R_GPU_ARCH:-gfx1101}"
FP8="${D4R_NATIVE_FP8:-0}"
case "$ARCH" in
gfx11[0-9][0-9]) [[ "$FP8" == 1 ]] && { echo "D4R_NATIVE_FP8 needs an RDNA4 (gfx12) target" >&2; exit 2; } ;;
gfx12[0-9][0-9]) ;;
*) echo "unsupported D4R_GPU_ARCH $ARCH (RDNA3 gfx110x or RDNA4 gfx120x)" >&2; exit 2 ;;
esac
CLANG="$ROCM/lib/llvm/bin/clang++"
[[ -x "$CLANG" ]] || { echo "clang++ not found in $ROCM/lib/llvm/bin (set D4R_ROCM_DIR)" >&2; exit 2; }
mkdir -p "$OUT"
# HIP source -> raw code object (the kernel name inside matches the DLSS kernel it replaces)
build_hip() {
local src="$1" name
local src="$1" extra="${2:-}" name
name="$(basename "$src" .hip)"
local tmp
local tmp flags
tmp="$(mktemp -d)"
(cd "$tmp" && "$CLANG" -x hip --offload-arch="$ARCH" --offload-device-only -O3 -I"$ROCM/include" \
# per-kernel compiler flags from a "// d4r-build-flags: ..." line in the source
flags="$(sed -n 's|^// d4r-build-flags: *||p' "$src")"
# shellcheck disable=SC2086
(cd "$tmp" && "$CLANG" -x hip --offload-arch="$ARCH" --offload-device-only -O3 $flags $extra -I"$ROCM/include" \
--rocm-path="$ROCM" --rocm-device-lib-path="$ROCM/amdgcn/bitcode" -I"$(dirname "$src")" \
-o "$OUT/$name.hsaco" "$src" -save-temps=cwd --no-gpu-bundle-output)
printf '%-52s %s\n' "$name" "$(grep -hE '^; (NumVgprs|ScratchSize|Occupancy)' "$tmp"/*.s | tail -3 | paste -sd' ')"
@@ -39,23 +53,39 @@ if [[ "$WHAT" == all || "$WHAT" == k ]]; then
fi
if [[ "$WHAT" == all || "$WHAT" == m ]]; then
echo "== DLSS 4.5 (preset M) Swin layers"
for src in "$HERE"/m/rrlite_*.hip; do build_hip "$src"; done
for src in "$HERE"/m/rrlite_*.hip; do build_hip "$src" "$([[ "$FP8" == 1 ]] && echo -DD4R_FP8_WMMA)"; done
fi
if [[ "$WHAT" == all || "$WHAT" == tex ]]; then
echo "== texture kernels (NVIDIA PTX with native parts)"
: "${D4R_DLSS_DLL:?set D4R_DLSS_DLL to nvngx_dlss.dll (310.7)}"
: "${D4R_ZLUDA_BUILD:?set D4R_ZLUDA_BUILD to a ZLUDA build with patches/zluda}"
: "${D4R_ZLUDA_EMIT:?set D4R_ZLUDA_EMIT to d4r_emit from a ZLUDA build with patches/zluda}"
PTX_DIR="$HERE/extracted/ptx"
python3 "$HERE/tools/extract_dlss_ptx.py" "$D4R_DLSS_DLL" "$PTX_DIR"
PTXLOAD="$HERE/out/tools/ptxload"
mkdir -p "$(dirname "$PTXLOAD")"
cc -O2 -o "$PTXLOAD" "$HERE/tools/ptxload.c" -ldl
# kernel : source of its native part
for pair in hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel:sust_only rrlite_post_3_2_mvhi_hdr_folded:sust_only \
rrlite_downsample_kernel_static_hdr:sust_only rrlite_enc0_4x4_mvhi_hdr_folded:enc0_tail \
rrlite_dec0_4x4_folded:dec0_head; do
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" D4R_DLSS_PTX_DIR="$PTX_DIR" D4R_PTXLOAD="$PTXLOAD" \
"$HERE/tex/build_tex.sh" "${pair%%:*}" "${pair#*:}" "$OUT"
# kernel : source of its native part [: w64 = compiled as wave64 (two CUDA warps per wave; RDNA3 issues
# FP32 work for all 64 lanes at once). Only for kernels without MMAs; measured faster, bit-identical:
# post 0.89 -> 0.77 ms, hiluma output 0.905 -> 0.887 ms (downsample was slower as wave64).]
# Every flag combination DLSS selects (motion vectors hi/lo, HDR/LDR, depth inverted/regular, ...) gets
# the same treatment, so games other than the one measured run the native versions too.
specs=(rrlite_dec0_4x4_folded:dec0_head)
for mv in mvhi mvlo; do
for range in hdr ldr; do
specs+=("rrlite_enc0_4x4_${mv}_${range}_folded:enc0_tail")
for v in 3_1 3_2; do specs+=("rrlite_post_${v}_${mv}_${range}_folded:sust_only:w64"); done
for depth in depthinv depthreg; do
for kind in "" _max; do
[[ -z "$kind" ]] && specs+=("hiluma_engine_output_${depth}_${mv}_${range}_v1_rel:sust_only:w64")
specs+=("hiluma_engine_output_${depth}_${mv}_${range}${kind}_v2_rel:sust_only:w64")
done
done
done
done
for mode in static dynamic; do
for range in hdr ldr; do specs+=("rrlite_downsample_kernel_${mode}_${range}:sust_only"); done
done
for spec in "${specs[@]}"; do
IFS=: read -r kernel src mode <<< "$spec"
D4R_ZLUDA_WAVE64="$([[ "$mode" == w64 ]] && echo 1 || echo 0)" D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" \
D4R_TEX_FP8="$FP8" D4R_DLSS_PTX_DIR="$PTX_DIR" "$HERE/tex/build_tex.sh" "$kernel" "$src" "$OUT"
done
fi
echo "native kernels in $OUT"
kernels/common/wmma_layout.h
+331
View File
@@ -0,0 +1,331 @@
// WMMA register layouts of RDNA3 (gfx11) and RDNA4 (gfx12) for the native kernels (wave32, 16x16x16 f16 -> f32).
// Every layout fact the kernels depend on is here; the kernel sources are written against these helpers.
//
// Notation: lane l, half h = l >> 4, r = l & 15. A operand lane l = row r of A, B operand lane l = column r
// of B, D lane l = column r of D.
// gfx11: an operand holds all 16 K values of its row/column (8 dwords, the same data in both halves);
// D VGPR i is row 2i + h.
// gfx12: an operand holds 8 of the 16 K values (4 dwords), half h the K values 8h .. 8h+7 in order;
// D VGPR i is row i + 8h, which is the gfx12 B-operand layout of the next GEMM (no exchange).
// A and B only have to agree on which K each slot holds; "half h = K 8h .. 8h+7" makes a gfx12 operand the
// lower or upper 16 bytes of the gfx11 one, so row images and weight images serve both layouts.
//
// D4R_WMMA_LAYOUT selects the layout (default: the target's). D4R_WMMA_LAYOUT=12 on a gfx11 target is the
// layout shim: the gfx12 register layout with every WMMA done by the gfx11 instruction on operands and
// accumulators converted around it, in the gfx11 K order. Its results are bit-identical to the gfx11 build,
// so the gfx12 indexing can be validated on RDNA3 hardware.
//
// Only clang builtins are used (the texture tails compile with -nogpuinc -nogpulib).
#pragma once
#include <stdint.h>
#ifndef D4R_WMMA_LAYOUT
#if defined(__GFX12__)
#define D4R_WMMA_LAYOUT 12
#else
#define D4R_WMMA_LAYOUT 11
#endif
#endif
#if D4R_WMMA_LAYOUT == 12 && !defined(__GFX12__)
#define D4R_WMMA_SHIM 1
#endif
#if D4R_WMMA_LAYOUT == 11 && defined(__GFX12__)
#error "the gfx11 WMMA layout does not exist on gfx12"
#endif
#define WM_FN __attribute__((device, always_inline)) static inline
typedef _Float16 wm_f16;
typedef _Float16 wm_hv2 __attribute__((ext_vector_type(2)));
typedef _Float16 wm_h8 __attribute__((ext_vector_type(8)));
typedef _Float16 wm_h16 __attribute__((ext_vector_type(16)));
typedef float wm_f8v __attribute__((ext_vector_type(8)));
typedef uint32_t wm_u8v __attribute__((ext_vector_type(8)));
typedef uint32_t wm_u4v __attribute__((ext_vector_type(4)));
typedef int wm_i2v __attribute__((ext_vector_type(2)));
// ---------------------------------------------------------------- lanes
WM_FN uint32_t wm_lane()
{
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
}
WM_FN uint32_t wm_half()
{
return wm_lane() >> 4;
}
// value of the lane with the same index in the other half of the wave
WM_FN uint32_t wm_other_half(uint32_t v)
{
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
}
WM_FN float wm_other_half_f(float v)
{
return __builtin_bit_cast(float, wm_other_half(__builtin_bit_cast(uint32_t, v)));
}
WM_FN uint32_t wm_pack2(wm_f16 lo, wm_f16 hi)
{
return __builtin_bit_cast(uint32_t, (wm_hv2){lo, hi});
}
WM_FN wm_f16 wm_lo16(uint32_t v)
{
return __builtin_bit_cast(wm_hv2, v)[0];
}
WM_FN wm_f16 wm_hi16(uint32_t v)
{
return __builtin_bit_cast(wm_hv2, v)[1];
}
// ---------------------------------------------------------------- accumulator rows
// row (of the 16-row tile) that accumulator VGPR i holds in a lane of half h (prep code: logical lanes)
WM_FN int wm_acc_row_h(int i, int h)
{
#if D4R_WMMA_LAYOUT == 12
return i + 8 * h;
#else
return 2 * i + h;
#endif
}
// ... in this lane
WM_FN int wm_acc_row(int i)
{
return wm_acc_row_h(i, (int)wm_half());
}
// rows are the K dimension of the next GEMM: VGPRs 2j, 2j + 1 are adjacent K values in gfx12 (and 16 bytes
// of this lane's 8 rows are contiguous), rows 2i + h are interleaved between the halves in gfx11
#if D4R_WMMA_LAYOUT == 12
#define WM_ACC_CONTIGUOUS 1
#else
#define WM_ACC_CONTIGUOUS 0
#endif
// ---------------------------------------------------------------- f16 operands
#if D4R_WMMA_LAYOUT == 12
typedef wm_u4v wm_op; // this half's 8 K values
#define WM_OP_DWORDS 4
#else
typedef wm_u8v wm_op; // all 16 K values
#define WM_OP_DWORDS 8
#endif
// operand from 16 contiguous halves (K 0..15) in memory (16-byte aligned); gfx12 reads only this half's 8
WM_FN wm_op wm_op_load(const wm_f16* p)
{
#if D4R_WMMA_LAYOUT == 12
return *(const wm_u4v*)(p + 8 * wm_half());
#else
const wm_u4v a = *(const wm_u4v*)p, b = *(const wm_u4v*)(p + 8);
return (wm_u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
#endif
}
// operand slot `idx` of an image of 32-byte slots (16 halves each, K in order): gfx12 reads its 16-byte half
WM_FN wm_op wm_op_image(const wm_u8v* __restrict__ img, int idx)
{
#if D4R_WMMA_LAYOUT == 12
return ((const wm_u4v*)img)[2 * idx + (int)wm_half()];
#else
return img[idx];
#endif
}
// element (K value) of an operand that holds accumulator row wm_acc_row(i) of the same lane
WM_FN wm_f16 wm_op_acc_elem(const wm_op& v, int i)
{
#if D4R_WMMA_LAYOUT == 12
return (i & 1) ? wm_hi16(v[i >> 1]) : wm_lo16(v[i >> 1]);
#else
const int c = 2 * i + (int)wm_half();
return (c & 1) ? wm_hi16(v[c >> 1]) : wm_lo16(v[c >> 1]);
#endif
}
// the f16 accumulator values own[i] (row wm_acc_row(i) of this lane's column) as an operand of the next GEMM
// (the rows become K; gfx11 fetches the other parity from lane l ^ 16)
WM_FN wm_op wm_op_from_acc(const wm_f16 own[8])
{
#if D4R_WMMA_LAYOUT == 12
return (wm_u4v){wm_pack2(own[0], own[1]), wm_pack2(own[2], own[3]), wm_pack2(own[4], own[5]), wm_pack2(own[6], own[7])};
#else
const uint32_t hf = wm_half();
uint32_t mine[4], theirs[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
mine[j] = wm_pack2(own[2 * j], own[2 * j + 1]);
theirs[j] = wm_other_half(mine[j]);
}
#ifndef PWIN_NO_PERM
// v_perm_b32(theirs, mine): bytes 0-3 = mine, 4-7 = theirs. Even j: the low halves, odd j: the high
// halves, in the order (mine, theirs) for hf = 0 and (theirs, mine) for hf = 1.
const uint32_t sel_lo = hf ? 0x01000504u : 0x05040100u, sel_hi = hf ? 0x03020706u : 0x07060302u;
wm_u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
r[j] = __builtin_amdgcn_perm(theirs[j >> 1], mine[j >> 1], (j & 1) ? sel_hi : sel_lo);
return r;
#else
wm_u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
{
const wm_f16 a = (j & 1) ? wm_hi16(mine[j >> 1]) : wm_lo16(mine[j >> 1]);
const wm_f16 b = (j & 1) ? wm_hi16(theirs[j >> 1]) : wm_lo16(theirs[j >> 1]);
r[j] = hf ? wm_pack2(b, a) : wm_pack2(a, b);
}
return r;
#endif
#endif
}
WM_FN wm_op wm_op_from_f8(wm_f8v d)
{
wm_f16 h[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
h[i] = (wm_f16)d[i];
return wm_op_from_acc(h);
}
// Stores an operand back as a row of 16 halves (K 0..15). In gfx11 both halves hold the whole row and the
// lanes of half 0 store it; in gfx12 each half stores its 8 values. `count` (8 or 16) limits the row to its
// first `count` halves; `ok` is the lane's own condition.
WM_FN void wm_op_store(wm_f16* p, const wm_op& v, bool ok = true, int count = 16)
{
#if D4R_WMMA_LAYOUT == 12
if (ok && (wm_half() == 0 || count > 8))
*(wm_u4v*)(p + 8 * wm_half()) = v;
#else
if (ok && wm_half() == 0)
{
*(wm_u4v*)p = (wm_u4v){v[0], v[1], v[2], v[3]};
if (count > 8)
*(wm_u4v*)(p + 8) = (wm_u4v){v[4], v[5], v[6], v[7]};
}
#endif
}
// 8 halves of an f16 vector in accumulator row order: out[i] = v[wm_acc_row(i)], v at p (16 halves, 16-byte
// aligned)
WM_FN void wm_acc_vec_load(const wm_f16* p, wm_f16 out[8])
{
#if D4R_WMMA_LAYOUT == 12
const wm_u4v a = *(const wm_u4v*)(p + 8 * wm_half());
#pragma unroll
for (int i = 0; i < 8; ++i)
out[i] = (i & 1) ? wm_hi16(a[i >> 1]) : wm_lo16(a[i >> 1]);
#else
const uint32_t hf = wm_half();
const wm_u4v a = *(const wm_u4v*)p, b = *(const wm_u4v*)(p + 8);
const wm_u8v v = (wm_u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
#pragma unroll
for (int i = 0; i < 8; ++i)
out[i] = hf ? wm_hi16(v[i]) : wm_lo16(v[i]);
#endif
}
// ---------------------------------------------------------------- the WMMA itself
#if D4R_WMMA_LAYOUT == 12 && defined(__GFX12__)
WM_FN wm_f8v wm_mma(const wm_op& a, const wm_op& b, wm_f8v c)
{
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(__builtin_bit_cast(wm_h8, a), __builtin_bit_cast(wm_h8, b), c);
}
#else
WM_FN wm_f8v wm_mma11(const wm_u8v& a, const wm_u8v& b, wm_f8v c)
{
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(wm_h16, a), __builtin_bit_cast(wm_h16, b), c);
}
#if D4R_WMMA_LAYOUT == 12
// layout shim: gfx12 operand halves -> one gfx11 operand in K order (half 0's 8 K, then half 1's)
WM_FN wm_u8v wm_shim_operand(const wm_u4v& v)
{
const bool hf = wm_half() != 0;
wm_u8v r;
#pragma unroll
for (int j = 0; j < 4; ++j)
{
const uint32_t o = wm_other_half(v[j]);
r[j] = hf ? o : v[j];
r[4 + j] = hf ? v[j] : o;
}
return r;
}
// gfx12 accumulator (VGPR i = row i + 8h) <-> gfx11 accumulator (VGPR i = row 2i + h)
WM_FN wm_f8v wm_shim_acc_to11(const wm_f8v& c)
{
const bool hf = wm_half() != 0;
wm_f8v o, r;
#pragma unroll
for (int i = 0; i < 8; ++i)
o[i] = wm_other_half_f(c[i]);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
// row 2i + h lives in half (2i + h) >> 3, VGPR (2i + h) & 7
const float h0 = i < 4 ? c[2 * i] : o[2 * i - 8];
const float h1 = i < 4 ? o[2 * i + 1] : c[2 * i - 7];
r[i] = hf ? h1 : h0;
}
return r;
}
WM_FN wm_f8v wm_shim_acc_to12(const wm_f8v& d)
{
const bool hf = wm_half() != 0;
wm_f8v o, r;
#pragma unroll
for (int i = 0; i < 8; ++i)
o[i] = wm_other_half_f(d[i]);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
// row i + 8h lives in half i & 1, VGPR (i >> 1) + 4h
const float h0 = (i & 1) ? o[i >> 1] : d[i >> 1];
const float h1 = (i & 1) ? d[(i >> 1) + 4] : o[(i >> 1) + 4];
r[i] = hf ? h1 : h0;
}
return r;
}
WM_FN wm_f8v wm_mma(const wm_op& a, const wm_op& b, wm_f8v c)
{
return wm_shim_acc_to12(wm_mma11(wm_shim_operand(a), wm_shim_operand(b), wm_shim_acc_to11(c)));
}
// ZLUDA's K order: its gfx11 helpers put K pair p (K 2p, 2p + 1) in operand dword p, its gfx12 helpers put
// in half h, dword j the pair (j & 1) * 4 + 2h + (j >> 1) (see zluda_ptx_impl.cpp)
WM_FN wm_u8v wm_shim_operand_zluda(const wm_u4v& v)
{
const bool hf = wm_half() != 0;
wm_u4v o;
#pragma unroll
for (int j = 0; j < 4; ++j)
o[j] = wm_other_half(v[j]);
wm_u8v r;
#pragma unroll
for (int p = 0; p < 8; ++p)
{
const int hp = (p & 3) >> 1, jp = 2 * (p & 1) + (p >> 2);
r[p] = (hp != 0) == hf ? v[jp] : o[jp];
}
return r;
}
WM_FN wm_f8v wm_mma_zluda(const wm_op& a, const wm_op& b, wm_f8v c)
{
return wm_shim_acc_to12(wm_mma11(wm_shim_operand_zluda(a), wm_shim_operand_zluda(b), wm_shim_acc_to11(c)));
}
#else
WM_FN wm_f8v wm_mma(const wm_op& a, const wm_op& b, wm_f8v c)
{
return wm_mma11(a, b, c);
}
#endif
#endif
#if D4R_WMMA_LAYOUT == 12 && defined(__GFX12__)
// operands whose K slots follow ZLUDA's gfx12 order: the same instruction (the order only matters to the shim)
WM_FN wm_f8v wm_mma_zluda(const wm_op& a, const wm_op& b, wm_f8v c)
{
return wm_mma(a, b, c);
}
#endif
WM_FN wm_f8v wm_splat8(float v)
{
return (wm_f8v){v, v, v, v, v, v, v, v};
}
kernels/k/dltss_pwin_dec0_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#include "pwin_pos.h"
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)
kernels/k/dltss_pwin_dec1_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)
kernels/k/dltss_pwin_dec2_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_DECODER(dltss_pwin_dec2_layer, 2, 64, 96)
kernels/k/dltss_pwin_enc0_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#include "pwin_pos.h"
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)
kernels/k/dltss_pwin_enc1_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
kernels/k/dltss_pwin_enc2_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
// d4r-build-flags: -mcumode
#define PWIN_F32ACC
#include "pwin_layer.h"
PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)
kernels/k/pwin_common.h
+68 -73
View File
@@ -1,14 +1,15 @@
// Native RDNA3 (gfx11 wave32 WMMA) kernels for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
// Native kernels (gfx11 / gfx12 wave32 WMMA) for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
// Semantics: kernels/tools/pwin_model.py (verified against the PTX interpreter, kernels/tools/ptxsim.py).
//
// Layout conventions (transposed WMMA, as in the preset M kernels): a GEMM Y = X W is computed as
// Y^T = W^T X^T with the weights as the WMMA A operand (lane l: output channel l & 15, 16 K values)
// and the activations as the B operand (lane l: token l & 15, 16 K values). The f32 result then has
// lane l = token l & 15 and VGPR i = output channel 2i + (l >> 4) of the 16-channel tile ("D^T").
// Y^T = W^T X^T with the weights as the WMMA A operand (lane l: output channel l & 15) and the activations
// as the B operand (lane l: token l & 15). The f32 result then has lane l = token l & 15 and VGPR i =
// output channel wm_acc_row(i) of the 16-channel tile ("D^T"); ../common/wmma_layout.h has the per-target
// register layouts (gfx11: channel 2i + (l >> 4); gfx12: channel i + 8 (l >> 4)).
#pragma once
#include <hip/hip_runtime.h>
#include <stdint.h>
#include "../common/wmma_layout.h"
#pragma clang fp contract(off)
typedef _Float16 half_t;
@@ -23,16 +24,13 @@ __device__ __forceinline__ uint32_t lane_id()
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
}
__device__ __forceinline__ f8v wmma(u8v a, u8v b, f8v c)
{
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
}
typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
// one k16 step with the f16 accumulator of NVIDIA's f16 wmma (rounded after the step)
// PWIN_F32ACC: keep the accumulator in f32 through the chain (rounded to f16 where the values are used)
__device__ __forceinline__ f8v mma16(u8v a, u8v b, f8v c)
__device__ __forceinline__ f8v mma16(const op_t& a, const op_t& b, f8v c)
{
f8v d = wmma(a, b, c);
f8v d = wm_mma(a, b, c);
#ifndef PWIN_F32ACC
#pragma unroll
for (int i = 0; i < 8; ++i)
@@ -65,50 +63,19 @@ __device__ __forceinline__ half_t hi16(uint32_t v)
return __builtin_bit_cast(hv2, v)[1];
}
// D^T f16 values (own[i] = channel 2i + hf of token l & 15) -> B operand of the same token (all 16
// channels in order): the other parity comes from lane l ^ 16
__device__ __forceinline__ u8v operand_from_dt(const half_t own[8])
// D^T f16 values (own[i] = channel wm_acc_row(i) of token l & 15) -> operand of the same token (the
// channels become K)
__device__ __forceinline__ op_t operand_from_dt(const half_t own[8])
{
const uint32_t hf = lane_id() >> 4;
uint32_t mine[4], theirs[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
mine[j] = pack2(own[2 * j], own[2 * j + 1]);
theirs[j] = other_half(mine[j]);
}
#ifndef PWIN_NO_PERM
// v_perm_b32(theirs, mine): bytes 0-3 = mine, 4-7 = theirs. Even j: the low halves, odd j: the high
// halves, in the order (mine, theirs) for hf = 0 and (theirs, mine) for hf = 1.
const uint32_t sel_lo = hf ? 0x01000504u : 0x05040100u, sel_hi = hf ? 0x03020706u : 0x07060302u;
u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
r[j] = __builtin_amdgcn_perm(theirs[j >> 1], mine[j >> 1], (j & 1) ? sel_hi : sel_lo);
return r;
#else
u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
{
const half_t a = (j & 1) ? hi16(mine[j >> 1]) : lo16(mine[j >> 1]);
const half_t b = (j & 1) ? hi16(theirs[j >> 1]) : lo16(theirs[j >> 1]);
r[j] = hf ? pack2(b, a) : pack2(a, b);
}
return r;
#endif
return wm_op_from_acc(own);
}
__device__ __forceinline__ u8v operand_from_f8(f8v d)
__device__ __forceinline__ op_t operand_from_f8(f8v d)
{
half_t h[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
h[i] = (half_t)d[i];
return operand_from_dt(h);
return wm_op_from_f8(d);
}
// element c (0..15) of a B operand (16 halves)
// element c (0..15) of a row of 16 halves (VALU code; not a WMMA operand)
__device__ __forceinline__ half_t op_get(const u8v& v, int c)
{
return (c & 1) ? hi16(v[c >> 1]) : lo16(v[c >> 1]);
@@ -145,6 +112,58 @@ __device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
}
// WMMA operands from rows of 16 channels: LDS, global (PWIN_NT_LOAD as above), prep images (32-byte slots)
__device__ __forceinline__ op_t op_lds(const half_t* p)
{
return wm_op_load(p);
}
__device__ __forceinline__ op_t op_gload(const half_t* p)
{
#if D4R_WMMA_LAYOUT == 12
#ifdef PWIN_NT_LOAD
return __builtin_nontemporal_load((const u4v*)(p + 8 * wm_half()));
#else
return *(const u4v*)(p + 8 * wm_half());
#endif
#else
return gload_row16(p);
#endif
}
__device__ __forceinline__ op_t op_img(const u8v* __restrict__ img, int idx)
{
return wm_op_image(img, idx);
}
// operand -> slot idx of an operand image in LDS (read back with op_img)
__device__ __forceinline__ void op_img_store(u8v* img, int idx, const op_t& v)
{
#if D4R_WMMA_LAYOUT == 12
((u4v*)img)[2 * idx + wm_half()] = v;
#else
if (!wm_half())
img[idx] = v;
#endif
}
// operand (16 channels of this lane's token) -> row of `count` (8 or 16) channels in global memory, if ok
__device__ __forceinline__ void op_gstore(half_t* dst, const op_t& v, bool ok, int count = 16)
{
#if D4R_WMMA_LAYOUT == 12
if (ok && (wm_half() == 0 || count >= 16))
gstore16(dst + 8 * wm_half(), v);
#else
if (!wm_half() && ok)
{
gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
if (count >= 16)
gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
}
#endif
}
// operand -> row of 16 channels in LDS
__device__ __forceinline__ void op_store(half_t* dst, const op_t& v)
{
wm_op_store(dst, v);
}
__device__ __forceinline__ void block_sync()
{
__syncthreads();
@@ -162,27 +181,3 @@ __device__ __forceinline__ int frag_offset(int k, int n)
{
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1);
}
// A-operand image of a K x N weight: tile (kt, nt) = 16 lanes x 8 dwords (lane = output channel)
struct WDesc
{
int src; // byte offset of the fragments in the weight buffer
int kt_stride; // bytes between k tiles
int nt_stride; // bytes between n tiles
int KT, NT; // tiles
int dst; // u8v index of the image in the prep buffer
};
// weight operand tile `tile` of an A-operand image (16 lanes x 32 bytes): lanes l and l + 16 need the same
// row, so each half-wave loads one 16-byte half and the halves are swapped across (halves L2 traffic)
__device__ __forceinline__ u8v wtile(const u8v* __restrict__ img, int tile)
{
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const u4v mine = ((const u4v*)&img[tile * 16 + m])[hf];
u4v other;
#pragma unroll
for (int j = 0; j < 4; ++j)
other[j] = other_half(mine[j]);
const u4v lo = hf ? other : mine, hi = hf ? mine : other;
return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
}
kernels/k/pwin_layer.h
+90 -80
View File
@@ -72,14 +72,10 @@ __device__ __forceinline__ half_t wload(const uint8_t* w, int byte)
return *(const half_t*)(w + byte);
}
// channels 16 nt + 2i + hf (i = 0..7) of an f16 vector at byte offset `byte` (16 nt already added)
__device__ __forceinline__ void dvec8(const uint8_t* w, int byte, uint32_t hf, half_t out[8])
// channels 16 nt + wm_acc_row(i) (i = 0..7) of an f16 vector at byte offset `byte` (16 nt already added)
__device__ __forceinline__ void dvec8(const uint8_t* w, int byte, half_t out[8])
{
const u4v a = *(const u4v*)(w + byte), b = *(const u4v*)(w + byte + 16);
const u8v v = (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
#pragma unroll
for (int i = 0; i < 8; ++i)
out[i] = hf ? hi16(v[i]) : lo16(v[i]);
wm_acc_vec_load((const half_t*)(w + byte), out);
}
__device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
@@ -131,7 +127,7 @@ template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int key = 16 * kt + 2 * i + (lane >> 4), r = q & 15, c = key & 15;
const int key = 16 * kt + wm_acc_row_h(i, lane >> 4), r = q & 15, c = key & 15;
const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
v[i] = wload(w, L::BIAS + 8192 * h + off);
}
@@ -209,6 +205,45 @@ __device__ uint64_t g_stamp[16384][16];
__device__ float g_dbg[12][64][160];
#define DBG_ON (blockIdx.x == DBX && blockIdx.y == DBY)
#endif
// Softmax row sum over the 64 keys of this lane's query: this lane's 32 accumulator rows (summed into rs as
// the scores are computed), plus the other half's. The layout shim (gfx12 layout on gfx11) instead redoes the
// sum in gfx11's order (keys 16 kt + 2 i + h, then the other parity), which the f32 result depends on, so
// that its output stays bit-identical to the gfx11 build.
#ifdef D4R_WMMA_SHIM
constexpr bool kShimRowSum = true;
#else
constexpr bool kShimRowSum = false;
#endif
__device__ __forceinline__ float row_sum_total(float rs, const half_t (&e)[4][8])
{
if constexpr (!kShimRowSum)
return rs + __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
// this lane: keys 16 kt + 8 h + j; the other half: keys 16 kt + 8 (1 - h) + j
const bool hf = wm_half() != 0;
float s[2] = {0.0f, 0.0f};
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
half_t lo[8], hi[8];
#pragma unroll
for (int j = 0; j < 8; j += 2)
{
const uint32_t own = pack2(e[kt][j], e[kt][j + 1]), other = other_half(own);
const uint32_t l = hf ? other : own, u = hf ? own : other;
lo[j] = lo16(l), lo[j + 1] = hi16(l), hi[j] = lo16(u), hi[j + 1] = hi16(u);
}
#pragma unroll
for (int par = 0; par < 2; ++par)
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int key = 2 * i + par;
s[par] += (float)(key < 8 ? lo[key] : hi[key - 8]);
}
}
return s[0] + s[1];
}
// The Swin core for one window: x0 rows of this wave's 16 tokens are in act (f16, C per token).
// On return act holds the block output y. hb: normalised input rows; K/V of one head at a time in kl/vt.
template <class L>
@@ -265,7 +300,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
half_t vv_[8];
dvec8(w, L::BO + 32 * nt, hf, vv_);
dvec8(w, L::BO + 32 * nt, vv_);
#pragma unroll
@@ -277,7 +312,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int h = 0; h < H; ++h)
{
// position-only attention: V^T A operands straight from D = h1 Wv (lane = dim, 16 keys per tile)
u8v vtop[2][4];
op_t vtop[2][4];
if constexpr (L::POS)
{
// each wave computes the V^T operands of its own key tile; shared through LDS, one slot per head
@@ -287,28 +322,26 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
f8v d = splat8(0.0f);
for (int ks = 0; ks < KT; ++ks)
d = mma16(lds_row16(&hb[16 * wv + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
const u8v vo = operand_from_f8(d);
if (!hf)
vsh[h][nt][wv][m] = vo;
d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
op_img_store(&vsh[h][nt][0][0], wv * 16 + m, operand_from_f8(d));
}
block_sync();
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
vtop[nt][kt] = vsh[h][nt][kt][m];
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
}
u8v pop[4];
op_t pop[4];
if constexpr (L::POS)
{
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = ((const u8v*)bias)[((h * 4 + wv) * 4 + kt) * 16 + m];
pop[kt] = op_img((const u8v*)bias, ((h * 4 + wv) * 4 + kt) * 16 + m);
}
else
{
u8v qop[2];
op_t qop[2];
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
f8v dq[6];
#pragma unroll
@@ -316,10 +349,10 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const u8v row = lds_row16(hrow + 16 * kt);
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[u]);
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
}
#pragma unroll
for (int j = 0; j < 3; ++j)
@@ -331,15 +364,13 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
const u8v kop = operand_from_f8(d);
if (!hf)
store_row16(&kl[16 * wv + m][16 * nt], kop);
op_store(&kl[16 * wv + m][16 * nt], operand_from_f8(d));
}
else
{
#pragma unroll
for (int i = 0; i < 8; ++i)
vt[16 * nt + 2 * i + hf][16 * wv + m] = (half_t)d[i];
vt[16 * nt + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
}
}
block_sync();
@@ -358,15 +389,16 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(lds_row16(&kl[16 * kt + m][16 * ds]), qop[ds], d);
d = mma16(op_lds(&kl[16 * kt + m][16 * ds]), qop[ds], d);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
e[kt][i] = softmax_e((half_t)d[i]);
rs += (float)e[kt][i];
if constexpr (!kShimRowSum)
rs += (float)e[kt][i];
}
}
rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
rs = row_sum_total(rs, e);
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
@@ -378,7 +410,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
pop[kt] = operand_from_dt(pv);
}
}
u8v oop[2];
op_t oop[2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
@@ -389,20 +421,20 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
if constexpr (L::POS)
d = mma16(vtop[nt][kt], pop[kt], d);
else
d = mma16(lds_row16(&vt[16 * nt + m][16 * kt]), pop[kt], d);
d = mma16(op_lds(&vt[16 * nt + m][16 * kt]), pop[kt], d);
}
oop[nt] = operand_from_f8(d);
#ifdef PWIN_DBG
if (DBG_ON)
for (int i = 0; i < 8; ++i)
g_dbg[2 + h][16 * wv + m][16 * nt + 2 * i + hf] = (float)(half_t)d[i];
g_dbg[2 + h][16 * wv + m][16 * nt + wm_acc_row(i)] = (float)(half_t)d[i];
#endif
}
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
acc[nt] = mma16(img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m], oop[ks], acc[nt]);
acc[nt] = mma16(op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m), oop[ks], acc[nt]);
if constexpr (!L::POS)
block_sync(); // kl / vt are reused by the next head
}
@@ -411,34 +443,32 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (DBG_ON)
for (int nt = 0; nt < KT; ++nt)
for (int i = 0; i < 8; ++i)
g_dbg[1][16 * wv + m][16 * nt + 2 * i + hf] = acc[nt][i];
g_dbg[1][16 * wv + m][16 * nt + wm_acc_row(i)] = acc[nt][i];
#endif
STAMP(3);
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
const u8v xv = lds_row16(arow + 16 * nt);
const op_t xv = op_lds(arow + 16 * nt);
half_t mv[8], g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, hf, g2);
dvec8(w, L::B2 + 32 * nt, hf, b2);
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const half_t x1 = (half_t)acc[nt][i] + op_get(xv, 2 * i + hf);
const half_t x1 = (half_t)acc[nt][i] + wm_op_acc_elem(xv, i);
mv[i] = x1 * g2[i];
acc[nt][i] = (float)(half_t)(x1 + b2[i]);
}
const u8v mo = operand_from_dt(mv);
if (!hf)
store_row16(hrow + 16 * nt, mo);
op_store(hrow + 16 * nt, operand_from_dt(mv));
}
__builtin_amdgcn_wave_barrier();
// ---- MLP: C/8 chunks of 32 hidden
for (int c = 0; c < L::NMLP; ++c)
{
u8v gop[2];
op_t gop[2];
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
@@ -447,7 +477,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
half_t vv_[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv_);
dvec8(w, L::B1 + 64 * c + 32 * hn, vv_);
#pragma unroll
@@ -457,7 +487,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
for (int kt = 0; kt < KT; ++kt)
d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
half_t g[8];
gelu8(d, g);
gop[hn] = operand_from_dt(g);
@@ -466,16 +496,14 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
acc[nt] = mma16(img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m], gop[ks], acc[nt]);
acc[nt] = mma16(op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m), gop[ks], acc[nt]);
}
STAMP(4);
// y -> act (rows of this wave)
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
const u8v yo = operand_from_f8(acc[nt]);
if (!hf)
store_row16(arow + 16 * nt, yo);
op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
}
__builtin_amdgcn_wave_barrier();
}
@@ -530,31 +558,29 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
u8v xe[L::KE];
op_t xe[L::KE];
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
xe[ks] = gload_row16(xin + 16 * ks);
xe[ks] = op_gload(xin + 16 * ks);
#pragma unroll
for (int nt = 0; nt < C / 16; ++nt)
{
f8v d;
{
half_t vv_[8];
dvec8(p0.w, 32 * nt, hf, vv_);
dvec8(p0.w, 32 * nt, vv_);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv_[i];
}
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
d = mma16(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
half_t xo[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
xo[i] = hmax((half_t)d[i], (half_t)0.0f);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[tok][16 * nt], xv);
op_store(&act[tok][16 * nt], operand_from_dt(xo));
}
}
else
@@ -592,7 +618,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
half_t vv_[8];
dvec8(w, L::PMB + 32 * vt, hf, vv_);
dvec8(w, L::PMB + 32 * vt, vv_);
#pragma unroll
@@ -606,18 +632,11 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
}
const u8v mo = operand_from_f8(d);
// allocated column 16 vt + c belongs to original warp g, local column 16 u + c
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
half_t* dst = mout + L::NPW * g + 16 * u;
if (!hf && minb)
{
gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
if (valid >= 16)
gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
}
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
}
}
@@ -749,7 +768,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
half_t vv_[8];
dvec8(p.w, L::EXPB + 32 * gnt, hf, vv_);
dvec8(p.w, L::EXPB + 32 * gnt, vv_);
#pragma unroll
@@ -759,15 +778,13 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
}
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
const u8v sv = gload_row16(sk + 16 * nt);
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
const op_t sv = op_gload(sk + 16 * nt);
half_t xo[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[t][16 * nt], xv);
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
op_store(&act[t][16 * nt], operand_from_dt(xo));
}
}
block_sync();
@@ -788,7 +805,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
half_t vv_[8];
dvec8(q.w, L::HB + 32 * nt, hf, vv_);
dvec8(q.w, L::HB + 32 * nt, vv_);
#pragma unroll
@@ -799,15 +816,8 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
}
#pragma unroll
for (int kt = 0; kt < L::KT; ++kt)
d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
const u8v ho = operand_from_f8(d);
const int valid = L::NOUT - 16 * nt;
if (!hf && inb)
{
gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
if (valid >= 16)
gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
}
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
}
}
else if (inb)
kernels/k/pwin_pos.h
+39 -63
View File
@@ -53,7 +53,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
for (int nt = 0; nt < KT; ++nt)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, hf, vv);
dvec8(w, L::BO + 32 * nt, vv);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
#pragma unroll
@@ -67,7 +67,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#endif
for (int h = 0; h < H; ++h)
{
u8v vtop[2][4];
op_t vtop[2][4];
#ifdef POS_VSHARE
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
@@ -78,17 +78,15 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
f8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
const u8v vo = operand_from_f8(d);
if (!hf)
vsh[h][nt][kt][m] = vo;
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_f8(d));
}
block_sync();
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
vtop[nt][kt] = vsh[h][nt][kt][m];
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
#else
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
@@ -98,19 +96,19 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
f8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
vtop[nt][kt] = operand_from_f8(d);
}
#endif
u8v oop[MT][2];
op_t oop[MT][2];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int qt = wv * MT + mi;
u8v pop[4];
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = ((const u8v*)bias)[((h * 4 + qt) * 4 + kt) * 16 + m];
pop[kt] = op_img((const u8v*)bias, ((h * 4 + qt) * 4 + kt) * 16 + m);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
@@ -126,7 +124,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const u8v wo = img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m];
const op_t wo = op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
@@ -139,24 +137,22 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
for (int nt = 0; nt < KT; ++nt)
{
half_t g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, hf, g2);
dvec8(w, L::B2 + 32 * nt, hf, b2);
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int tok = 16 * (wv * MT + mi) + m;
const u8v xv = lds_row16(&act[tok][16 * nt]);
const op_t xv = op_lds(&act[tok][16 * nt]);
half_t mv[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const half_t x1 = (half_t)acc[mi][nt][i] + op_get(xv, 2 * i + hf);
const half_t x1 = (half_t)acc[mi][nt][i] + wm_op_acc_elem(xv, i);
mv[i] = x1 * g2[i];
acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
}
const u8v mo = operand_from_dt(mv);
if (!hf)
store_row16(&hb[tok][16 * nt], mo);
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
}
}
__builtin_amdgcn_wave_barrier();
@@ -165,12 +161,12 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
// ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
for (int c = 0; c < L::NMLP; ++c)
{
u8v gop[MT][2];
op_t gop[MT][2];
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
half_t bb[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, hf, bb);
dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
f8v d[MT];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
@@ -180,10 +176,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int kt = 0; kt < KT; ++kt)
{
const u8v w1 = img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m];
const op_t w1 = op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
d[mi] = mma16(w1, lds_row16(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
d[mi] = mma16(w1, op_lds(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
}
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
@@ -198,7 +194,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const u8v w2 = img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m];
const op_t w2 = op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
@@ -211,9 +207,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
const u8v yo = operand_from_f8(acc[mi][nt]);
if (!hf)
store_row16(&act[16 * (wv * MT + mi) + m][16 * nt], yo);
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
}
block_sync();
}
@@ -239,29 +233,27 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
u8v xe[L::KE];
op_t xe[L::KE];
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
xe[ks] = gload_row16(xin + 16 * ks);
xe[ks] = op_gload(xin + 16 * ks);
#pragma unroll
for (int nt = 0; nt < C / 16; ++nt)
{
half_t vv[8];
dvec8(p0.w, 32 * nt, hf, vv);
dvec8(p0.w, 32 * nt, vv);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
d = mma16(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
half_t xo[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
xo[i] = hmax((half_t)d[i], (half_t)0.0f);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[tok][16 * nt], xv);
op_store(&act[tok][16 * nt], operand_from_dt(xo));
}
}
else
@@ -296,7 +288,7 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
for (int vt = wv; vt < L::PMT; vt += NW)
{
half_t vv[8];
dvec8(p.w, L::PMB + 32 * vt, hf, vv);
dvec8(p.w, L::PMB + 32 * vt, vv);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
@@ -305,17 +297,10 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
}
const u8v mo = operand_from_f8(d);
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
half_t* dst = mout + L::NPW * g + 16 * u;
if (!hf && minb)
{
gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
if (valid >= 16)
gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
}
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
}
STAMP(8);
}
@@ -333,10 +318,10 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
const int W2 = p.W / 2, H2 = p.H / 2, ly = m >> 2, lx = m & 3;
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
u8v xls[L::KL];
op_t xls[L::KL];
#pragma unroll
for (int ks = 0; ks < L::KL; ++ks)
xls[ks] = gload_row16(xl + 16 * ks);
xls[ks] = op_gload(xl + 16 * ks);
for (int q = wv; q < 4; q += NW)
{
const int dy = q >> 1, dx = q & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
@@ -347,22 +332,20 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
{
const int gnt = q * (C / 16) + nt;
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
dvec8(p.w, L::EXPB + 32 * gnt, vv);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
#pragma unroll
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], xls[ks], d);
const u8v sv = gload_row16(sk + 16 * nt);
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), xls[ks], d);
const op_t sv = op_gload(sk + 16 * nt);
half_t xo[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[t][16 * nt], xv);
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
op_store(&act[t][16 * nt], operand_from_dt(xo));
}
}
}
@@ -380,22 +363,15 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
half_t vv[8];
dvec8(q.w, L::HB + 32 * nt, hf, vv);
dvec8(q.w, L::HB + 32 * nt, vv);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
#pragma unroll
for (int kt = 0; kt < L::KT; ++kt)
d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
const u8v ho = operand_from_f8(d);
const int valid = L::NOUT - 16 * nt;
if (!hf && inb)
{
gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
if (valid >= 16)
gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
}
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
}
}
}
kernels/k/pwin_wide.h
+39 -55
View File
@@ -61,7 +61,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
store_row16(hrow + 16 * kt, hv);
}
// residual input of the owned output tiles (D^T order: channel 16 nt + 2 i + hf)
// residual input of the owned output tiles (D^T order: channel 16 nt + wm_acc_row(i))
half_t x0[NTL][8];
#pragma unroll
for (int j = 0; j < NTL; ++j)
@@ -69,10 +69,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int nt = g + NG * j;
if (nt < KT)
{
const u8v xv = lds_row16(arow + 16 * nt);
const op_t xv = op_lds(arow + 16 * nt);
#pragma unroll
for (int i = 0; i < 8; ++i)
x0[j][i] = op_get(xv, 2 * i + hf);
x0[j][i] = wm_op_acc_elem(xv, i);
}
}
block_sync(); // hb complete, R free
@@ -86,7 +86,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (nt < KT)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, hf, vv);
dvec8(w, L::BO + 32 * nt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[j][i] = (float)vv[i];
@@ -99,7 +99,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int r = 0; r < D::ROUNDS; ++r)
{
const int h = g < NA ? r * NA + g : H;
u8v qop[2], oop[2];
op_t qop[2], oop[2];
if (h < H)
{
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
@@ -109,10 +109,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const u8v row = lds_row16(hrow + 16 * kt);
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[u]);
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
}
#pragma unroll
for (int j = 0; j < 3; ++j)
@@ -124,15 +124,13 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
const u8v kop = operand_from_f8(d);
if (!hf)
store_row16(&kl[g][tok][16 * nt], kop);
op_store(&kl[g][tok][16 * nt], operand_from_f8(d));
}
else
{
#pragma unroll
for (int i = 0; i < 8; ++i)
vt[g][16 * nt + 2 * i + hf][tok] = (half_t)d[i];
vt[g][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
}
}
}
@@ -154,17 +152,18 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(lds_row16(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
d = mma16(op_lds(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
e[kt][i] = softmax_e((half_t)d[i]);
rs += (float)e[kt][i];
if constexpr (!kShimRowSum)
rs += (float)e[kt][i];
}
}
rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
rs = row_sum_total(rs, e);
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
u8v pop[4];
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
@@ -180,15 +179,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
f8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(lds_row16(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
d = mma16(op_lds(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
oop[nt] = operand_from_f8(d);
}
}
block_sync(); // K / V of the round consumed
if (h < H && !hf)
if (h < H)
{
store_row16(&ob[g][tok][0], oop[0]);
store_row16(&ob[g][tok][16], oop[1]);
op_store(&ob[g][tok][0], oop[0]);
op_store(&ob[g][tok][16], oop[1]);
}
block_sync();
for (int jh = 0; jh < NA; ++jh)
@@ -196,15 +195,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int hh = r * NA + jh;
if (hh >= H)
break;
const u8v o0 = lds_row16(&ob[jh][tok][0]), o1 = lds_row16(&ob[jh][tok][16]);
const op_t o0 = op_lds(&ob[jh][tok][0]), o1 = op_lds(&ob[jh][tok][16]);
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
const int nt = g + NG * j;
if (nt < KT)
{
acc[j] = mma16(img[(L::T_WO + (2 * hh) * KT + nt) * 16 + m], o0, acc[j]);
acc[j] = mma16(img[(L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m], o1, acc[j]);
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh) * KT + nt) * 16 + m), o0, acc[j]);
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m), o1, acc[j]);
}
}
}
@@ -220,8 +219,8 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (nt < KT)
{
half_t mv[8], g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, hf, g2);
dvec8(w, L::B2 + 32 * nt, hf, b2);
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
@@ -229,9 +228,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
mv[i] = x1 * g2[i];
acc[j][i] = (float)(half_t)(x1 + b2[i]);
}
const u8v mo = operand_from_dt(mv);
if (!hf)
store_row16(hrow + 16 * nt, mo);
op_store(hrow + 16 * nt, operand_from_dt(mv));
}
}
block_sync();
@@ -249,18 +246,16 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
f8v d;
{
half_t vv[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv);
dvec8(w, L::B1 + 64 * c + 32 * hn, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
}
for (int kt = 0; kt < KT; ++kt)
d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
half_t gg[8];
gelu8(d, gg);
const u8v gop = operand_from_dt(gg);
if (!hf)
store_row16(&hid[rr % D::NB][g][tok][16 * hn], gop);
op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
}
block_sync();
for (int jc = 0; jc < NG; ++jc)
@@ -268,15 +263,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int cc = rr * NG + jc;
if (cc >= L::NMLP)
break;
const u8v h0 = lds_row16(&hid[rr % D::NB][jc][tok][0]), h1 = lds_row16(&hid[rr % D::NB][jc][tok][16]);
const op_t h0 = op_lds(&hid[rr % D::NB][jc][tok][0]), h1 = op_lds(&hid[rr % D::NB][jc][tok][16]);
#pragma unroll
for (int j = 0; j < NTL; ++j)
{
const int nt = g + NG * j;
if (nt < KT)
{
acc[j] = mma16(img[(L::T_W2 + (cc * 2) * KT + nt) * 16 + m], h0, acc[j]);
acc[j] = mma16(img[(L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m], h1, acc[j]);
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2) * KT + nt) * 16 + m), h0, acc[j]);
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m), h1, acc[j]);
}
}
}
@@ -290,9 +285,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int nt = g + NG * j;
if (nt < KT)
{
const u8v yo = operand_from_f8(acc[j]);
if (!hf)
store_row16(hrow + 16 * nt, yo);
op_store(hrow + 16 * nt, operand_from_f8(acc[j]));
}
}
block_sync();
@@ -349,7 +342,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
f8v d;
{
half_t vv[8];
dvec8(w, L::PMB + 32 * vt, hf, vv);
dvec8(w, L::PMB + 32 * vt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
@@ -358,17 +351,10 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&hb[src][16 * (ks % L::KT)]), d);
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&hb[src][16 * (ks % L::KT)]), d);
}
const u8v mo = operand_from_f8(d);
const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
half_t* dst = mout + L::NPW * gw + 16 * u;
if (!hf && minb)
{
gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
if (valid >= 16)
gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
}
op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
}
}
@@ -409,21 +395,19 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
f8v d;
{
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
dvec8(p.w, L::EXPB + 32 * gnt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
}
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
const u8v sv = gload_row16(sk + 16 * nt);
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
const op_t sv = op_gload(sk + 16 * nt);
half_t xo[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[tt][16 * nt], xv);
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
op_store(&act[tt][16 * nt], operand_from_dt(xo));
}
}
block_sync();
kernels/m/rrlite_enc1_4x4.hip
+1
View File
@@ -3,5 +3,6 @@
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 1
#define SWIN_SCALAR_BLOAD // scalar-base weight loads: faster here (no spills), slower in the 8-wave layers
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc1_4x4, 64, 2, 4, 96, false, CommonParams, 0)
kernels/m/rrlite_enc3_tube_4x4.hip
+1
View File
@@ -3,5 +3,6 @@
#define SWIN_F32ACC
#define SWIN_MAX_VGPR 192
#define SWIN_MLP_R(NW) 2
#define SWIN_PREP_SLOTS 8 // one weight image per tube block (six per frame)
#include "swin_block.h"
SWIN_MODULE_W(rrlite_enc3_tube_4x4, 128, 4, 8, 0, true, TubeParams, 0)
kernels/m/swin_block.h
+144 -90
View File
@@ -19,7 +19,7 @@ template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
__host__ __device__ static constexpr int b1(int c) { return c == 0 ? M0 : M0 + 66 * C + 64 + (64 * C + 64) * (c - 1); }
static constexpr int PM0 = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1);
static constexpr int PMB = PM0 + 4 * C * NPM;
// f16 weight image (h16 units)
// weight image (slots of 16 K values, see expand_weights)
__host__ __device__ static constexpr int qv(int h, int q) { return (2 * h + q) * 2 * C; }
static constexpr int WO = 4 * C * NW;
static constexpr int W1 = WO + 2 * NW * C;
@@ -33,6 +33,14 @@ template <int C, int NW, int NPM, int CIN = 0> struct SwinLayout
static constexpr int NDESC = 2 * NW + 1 + 2 * NC + NPM / 32 + (CIN ? NW : 0);
};
// Start (in halves) of token T's row in the output staging buffer: stride C + 4 plus a shift by bits 4 and 5
// of T. The patch merge's lanes read 8 bytes each from 16 tokens that differ in bits 1, 3, 4 and 5 (plus a
// constant); this placement puts those 16 reads on 16 distinct 8-byte bank pairs for C = 64, 96 and 128.
template <int C> __host__ __device__ constexpr int xo_row(int T)
{
return (C + 4) * T + 4 * ((T >> 4) & 1) + 16 * (T >> 5);
}
template <int C, int NW, int NPM, int CIN> struct SwinDescs
{
GemmDesc d[SwinLayout<C, NW, NPM, CIN>::NDESC];
@@ -61,7 +69,7 @@ template <int C, int NW, int NPM, int CIN> constexpr SwinDescs<C, NW, NPM, CIN>
// NH heads (NVIDIA's warps per block) run on NW waves (NW = NH, or 2 NH to halve LDS per wave)
template <int C, int NH, int NW, int NPM, bool TUBE, int CIN>
__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const h16* w16, const int bx,
__device__ __forceinline__ void swin_block(const CommonParams& p, const TubeParams* tp, const wslot* w16, const int bx,
const int by, const int gx)
{
using L = SwinLayout<C, NH, NPM, CIN>;
@@ -77,8 +85,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
static_assert(NTW * NW == NT * MG, "column split");
constexpr int PAST = CIN + 8;
constexpr int SCR = NW * 16 * QST * 2 > (CIN ? 16 * PAST * 2 : 0) ? NW * 16 * QST * 2 : 16 * PAST * 2;
constexpr int U_ST4 = R * 64 * GST * 2, U_OB = NPM ? 64 * C : 0;
constexpr int USIZE = SCR > U_ST4 ? (SCR > U_OB ? SCR : U_OB) : (U_ST4 > U_OB ? U_ST4 : U_OB);
constexpr int U_ST4 = R * 64 * GST * 2;
constexpr int USIZE = SCR > U_ST4 ? SCR : U_ST4;
__shared__ __attribute__((aligned(16))) half_t A[64 * AST];
__shared__ __attribute__((aligned(16))) uint8_t U[USIZE];
@@ -118,11 +126,11 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
// block input words (encoders: fp8 input, decoders: fp8 skip), issued first so their latency
// overlaps the patch expand; word k of thread tid is idx = tid + 32 NW k
constexpr int NSW = CIN ? 1 : 64 * C / 4 / (32 * NW);
static_assert(CIN || NSW * 32 * NW == 64 * C / 4, "input words per thread");
constexpr int NSW = 64 * C / 4 / (32 * NW);
static_assert(NSW * 32 * NW == 64 * C / 4, "input words per thread");
uint32_t inw[NSW];
#pragma unroll
for (int k = 0; k < (CIN ? 0 : NSW); ++k)
for (int k = 0; k < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
@@ -130,26 +138,24 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const uint8_t* src = CIN ? p.p32 : p.in;
inw[k] = *(const uint32_t*)(src + ((size_t)((c4 >> 5) * p.th + Y) * p.tw + X) * 32 + (c4 & 31));
}
// the input words as f16 into X0 (natural channel order); decoders add the patch expand to them later
auto stage_input_words = [&]() {
#pragma unroll
for (int k = 0; k < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
const uint32_t bytes = inw[k];
#pragma unroll
for (int b = 0; b < 4; ++b)
X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
};
// ------------------------------------------------ stage 0: block input x0 (f16) into A
// encoders: the fp8 input; decoders: q8(patch expand(low-res)) + skip
if constexpr (CIN > 0)
{
constexpr int NTEW_ = 4 * C / NW / 16;
uint32_t sk[NTEW_][8];
#pragma unroll
for (int j = 0; j < NTEW_; ++j)
{
const int Gc = 16 * (NTEW_ * wv + j) + l16, q = Gc / C, ch = ginv(Gc % C);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 2 * i + hi;
const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
const int X = mirror(8 * bx - p.sx + tx, p.tw), Y = mirror(8 * by - p.sy + ty, p.th);
sk[j][i] = p.p32[((size_t)((ch >> 5) * p.th + Y) * p.tw + X) * 32 + (ch & 31)];
}
}
const int W2 = p.tw / 2, H2 = p.th / 2;
for (int idx = tid; idx < 16 * CIN / 4; idx += 32 * NW)
{
@@ -161,6 +167,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int b = 0; b < 4; ++b)
PA[r * PAST + apos(c4 + b)] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
stage_input_words(); // the skip input
__syncthreads();
constexpr int NTEW = 4 * C / NW / 16; // expand tiles per wave (L::NTE per weight block)
T16 e[NTEW];
@@ -170,7 +177,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < CIN / 32; ++kc)
{
const h16 a0 = lds16(PA + l16 * PAST + 32 * kc), a1 = lds16(PA + l16 * PAST + 32 * kc + 16);
const gop a0 = lda(PA + l16 * PAST + 32 * kc), a1 = lda(PA + l16 * PAST + 32 * kc + 16);
#pragma unroll
for (int j = 0; j < NTEW; ++j)
{
@@ -185,28 +192,16 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 2 * i + hi;
const int r = mrow(i);
const int tx = 2 * (r & 3) + (q & 1), ty = 2 * (r >> 2) + (q >> 1);
const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + e4m3_to_half(sk[j][i]));
// e4m3_to_half is exact in f16, so adding the staged skip value is the same sum
X0[T * AST + ch] = (half_t)(q8(e[j].get(i)) + X0[T * AST + ch]);
}
}
}
else
{
#pragma unroll
for (int k = 0; k < NSW; ++k)
{
const int idx = tid + 32 * NW * k;
const int T = idx / (C / 4), c4 = 4 * (idx % (C / 4));
const uint32_t bytes = inw[k];
#pragma unroll
for (int b = 0; b < 4; ++b)
{
X0[T * AST + c4 + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
}
}
}
stage_input_words();
__syncthreads();
// residual accumulators of the output projection: f16(x0 + b_o) in the C-fragment layout
@@ -223,8 +218,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const int ch = ginv(n);
#pragma unroll
for (int i = 0; i < 8; i += 2)
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + 2 * i + hi) * AST + ch],
X0[(16 * (mg * MT + mm) + 2 * i + 2 + hi) * AST + ch]} + (hv2){bo, bo});
x[mm][j].set_pair(i / 2, (hv2){X0[(16 * (mg * MT + mm) + mrow(i)) * AST + ch],
X0[(16 * (mg * MT + mm) + mrow(i + 1)) * AST + ch]} + (hv2){bo, bo});
}
__syncthreads(); // stage 1 overwrites x0 rows with h1
@@ -308,7 +303,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < C / 32; ++kc)
{
h16 b[2][4];
gop b[2][4];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -317,8 +312,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int m = 0; m < MW; ++m)
{
const h16 a0 = lds16(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
const h16 a1 = lds16(A + (16 * (m0 + m) + l16) * AST + 32 * kc + 16);
const gop a0 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc);
const gop a1 = lda(A + (16 * (m0 + m) + l16) * AST + 32 * kc + 16);
#pragma unroll
for (int q = 0; q < 4; ++q)
k32(acc[m][q], a0, b[0][q], a1, b[1][q]);
@@ -344,7 +339,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int ri = 2 * i + hi, cj = l16;
const int ri = mrow(i), cj = l16;
const int bl = 4 * (ri & 7) + ((cj & 7) >> 1), word = (ri >> 3) + 2 * (cj >> 3);
bias.set(i, hload(W, bbase + 2 * (8 * bl + 2 * word + (cj & 1))));
}
@@ -359,14 +354,14 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
Qs[(2 * i + hi) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
Qs[(2 * i + 2 + hi) * QST + 16 * q + l16] = qp[m][q][i / 2][1];
Qs[mrow(i) * QST + 16 * q + l16] = qp[m][q][i / 2][0];
Qs[mrow(i + 1) * QST + 16 * q + l16] = qp[m][q][i / 2][1];
}
wave_sync();
const h16 a0 = lds16(Qs + l16 * QST), a1 = lds16(Qs + l16 * QST + 16);
const gop a0 = lda(Qs + l16 * QST), a1 = lda(Qs + l16 * QST + 16);
T16 S = bias;
k32(S, a0, a0, a1, a1);
// rows 2k+hi and 2k+2+hi of this lane's column, packed: scale, clamp and the cubic
// rows mrow(2k), mrow(2k+1) of this lane's column, packed: scale, clamp and the cubic
// polynomial as f16x2, the exponent trick per row on the (even, odd) column word
hv2 wg[4];
const bool oddc = l16 & 1;
@@ -398,15 +393,19 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
const hv2 sum = v3 + xr(v3, std::integral_constant<int, 1>{});
const hv2 rinv = {f16(1.0f / (float)sum[0]), f16(1.0f / (float)sum[1])};
const hv2 pw = wg[k] * rinv;
Qs[(4 * k + hi) * PST + l16] = pw[0];
Qs[(4 * k + 2 + hi) * PST + l16] = pw[1];
Qs[mrow(2 * k) * PST + l16] = pw[0];
Qs[mrow(2 * k + 1) * PST + l16] = pw[1];
}
wave_sync();
const h16 pa = lds16(Qs + l16 * PST);
const sop pa = lds16(Qs + l16 * PST);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
// V as B operand: lane holds column l16 over the 16 tokens (rows); own rows 2i+hi
// V as B operand: lane holds column l16 over the 16 tokens (rows, the K of P V)
#if D4R_WMMA_LAYOUT == 12
// own pair k holds rows (2k + 8hi, 2k + 1 + 8hi): this half's K values in order
const sop vb = __builtin_bit_cast(sop, vp[m][nt]);
#else
// own pair k holds rows (4k+hi, 4k+2+hi); the other half-wave has (4k+1-hi, 4k+3-hi)
uint32_t vw[8];
const uint32_t s_lo = hi ? 0x01000504u : 0x05040100u, s_hi = hi ? 0x03020706u : 0x07060302u;
@@ -419,14 +418,15 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
vw[2 * k + 1] = __builtin_amdgcn_perm(other, own, s_hi); // rows 4k+2, 4k+3
}
const h16 vb = __builtin_bit_cast(h16, vw);
#endif
const f8v o = wmma(pa, vb, splat(0.0f));
const int ocol = apos(32 * h + ginv(16 * nt + l16));
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
const hv2 oq = q8x2((hv2){(half_t)o[i], (half_t)o[i + 1]});
A[(16 * (m0 + m) + 2 * i + hi) * AST + ocol] = oq[0];
A[(16 * (m0 + m) + 2 * i + 2 + hi) * AST + ocol] = oq[1];
A[(16 * (m0 + m) + mrow(i)) * AST + ocol] = oq[0];
A[(16 * (m0 + m) + mrow(i + 1)) * AST + ocol] = oq[1];
}
}
wave_sync();
@@ -438,7 +438,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < NH; ++kc)
{
h16 b[2][NTW];
gop b[2][NTW];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -447,8 +447,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
{
const h16 a0 = lds16(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
const h16 a1 = lds16(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc + 16);
const gop a0 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc);
const gop a1 = lda(A + (16 * (mg * MT + mm) + l16) * AST + 32 * kc + 16);
#pragma unroll
for (int j = 0; j < NTW; ++j)
k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
@@ -463,7 +463,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int j = 0; j < NTW; ++j)
#pragma unroll
for (int i = 0; i < 8; ++i)
A[(16 * (mg * MT + mm) + 2 * i + hi) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
A[(16 * (mg * MT + mm) + mrow(i)) * AST + 16 * (ng * NTW + j) + l16] = x[mm][j].get(i);
__syncthreads();
if (tid < 64)
{
@@ -514,12 +514,12 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; i += 2)
{
const int row = 16 * (mg * MT + mm) + 2 * i + hi;
const int row = 16 * (mg * MT + mm) + mrow(i), row2 = 16 * (mg * MT + mm) + mrow(i + 1);
const hv2 v = x[mm][j].pair(i / 2);
const hv2 r2 = {R2[row], R2[row + 2]};
const hv2 r2 = {R2[row], R2[row2]};
const hv2 hq = q8x2(v * (hv2)(r2 * (hv2){g2, g2}));
A[row * AST + col] = hq[0];
A[(row + 2) * AST + col] = hq[1];
A[row2 * AST + col] = hq[1];
x[mm][j].set_pair(i / 2, v + (hv2){b2, b2});
}
}
@@ -546,7 +546,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < C / 32; ++kc)
{
h16 b[2][2];
gop b[2][2];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -555,8 +555,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int mm = 0; mm < MPW; ++mm)
{
const h16 a0 = lds16(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
const h16 a1 = lds16(A + (16 * (mb + mm) + l16) * AST + 32 * kc + 16);
const gop a0 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc);
const gop a1 = lda(A + (16 * (mb + mm) + l16) * AST + 32 * kc + 16);
k32(z[mm][0], a0, b[0][0], a1, b[1][0]);
k32(z[mm][1], a0, b[0][1], a1, b[1][1]);
}
@@ -579,8 +579,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
(hv2)((hv2){f16(0.08108133f), f16(0.08108133f)} * az);
const hv2 gl = zv * (hv2)((hv2){0.5f16, 0.5f16} + (hv2)(cz * inner));
const hv2 gq = q8x2(gl);
Gc[(16 * (mb + mm) + 2 * i + hi) * GST + col] = gq[0];
Gc[(16 * (mb + mm) + 2 * i + 2 + hi) * GST + col] = gq[1];
Gc[(16 * (mb + mm) + mrow(i)) * GST + col] = gq[0];
Gc[(16 * (mb + mm) + mrow(i + 1)) * GST + col] = gq[1];
}
}
__syncthreads();
@@ -588,7 +588,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int cc = 0; cc < R; ++cc)
{
const int c2 = R * r + cc;
h16 b[2][NTW];
gop b[2][NTW];
#pragma unroll
for (int s = 0; s < 2; ++s)
#pragma unroll
@@ -597,8 +597,8 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int mm = 0; mm < MT; ++mm)
{
const h16 a0 = lds16(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
const h16 a1 = lds16(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST + 16);
const gop a0 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST);
const gop a1 = lda(G + (cc * 64 + 16 * (mg * MT + mm) + l16) * GST + 16);
#pragma unroll
for (int j = 0; j < NTW; ++j)
k32(x[mm][j], a0, b[0][j], a1, b[1][j]);
@@ -609,7 +609,10 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
}
// ------------------------------------------------ output: f16 staged in LDS, 16-value segments encoded by a loop
half_t* XO = A; // [64 tokens][C channels] natural order
// XO: [64 tokens][C channels] natural order; the row padding keeps the patch merge's 8-byte reads of 16
// tokens on distinct LDS banks (see xo_row)
half_t* XO = A;
static_assert(xo_row<C>(63) + C <= 64 * AST, "XO fits in A");
#pragma unroll
for (int j = 0; j < NTW; ++j)
{
@@ -618,26 +621,28 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
for (int mm = 0; mm < MT; ++mm)
#pragma unroll
for (int i = 0; i < 8; ++i)
XO[(16 * (mg * MT + mm) + 2 * i + hi) * C + ch] = x[mm][j].get(i);
XO[xo_row<C>(16 * (mg * MT + mm) + mrow(i)) + ch] = x[mm][j].get(i);
}
__syncthreads();
uint8_t* OB = U; // e4m3 codes [64][C] for the patch merge
#pragma unroll 1
for (int seg = tid; seg < 64 * C / 16; seg += 32 * NW)
{
const int T = seg / (C / 16), c16 = 16 * (seg % (C / 16));
// rows are 8-byte aligned (xo_row), so the segment moves as four 8-byte words of two f16 pairs
uint2* xs = (uint2*)(XO + xo_row<C>(T) + c16);
uint32_t w[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
{
w[k] = 0;
#pragma unroll
for (int b = 0; b < 4; ++b)
w[k] |= enc8(XO[T * C + c16 + 4 * k + b]) << (8 * b);
const uint2 pr = xs[k];
hv2 qa, qb;
const uint32_t ca = enc8x2(__builtin_bit_cast(hv2, pr.x), qa), cb = enc8x2(__builtin_bit_cast(hv2, pr.y), qb);
w[k] = pack_codes(ca, cb);
// the patch merge consumes decode(code) = q8_exact(x), written back over x (this thread's own values)
if constexpr (NPM > 0)
xs[k] = (uint2){__builtin_bit_cast(uint32_t, qa), __builtin_bit_cast(uint32_t, qb)};
}
const uint4 v = {w[0], w[1], w[2], w[3]};
if constexpr (NPM > 0)
*(uint4*)(OB + T * C + c16) = v;
const int X = 8 * bx - p.sx + tok_x(T), Y = 8 * by - p.sy + tok_y(T);
if (X >= 0 && X < p.tw && Y >= 0 && Y < p.th)
*(uint4*)(p.out + ((size_t)((c16 >> 5) * p.th + Y) * p.tw + X) * 32 + (c16 & 31)) = v;
@@ -663,23 +668,28 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int kc = 0; kc < C / 8; ++kc)
{
h16 as[2];
gop as[2];
#pragma unroll
for (int s = 0; s < 2; ++s)
{
h16& a = as[s];
sop a;
// the operand's K slots in runs of 4 (gfx12: this half's slots 8 hi .. 8 hi + 7)
constexpr int RUNS = sizeof(sop) / sizeof(half_t) / 4;
#pragma unroll
for (int run = 0; run < 4; ++run)
for (int rr = 0; rr < RUNS; ++rr)
{
const int run = RUNS == 4 ? rr : 2 * hi + rr;
const int k = 32 * kc + kslot(s, 4 * run);
const int q = k / C, ch = k % C;
const int tx = 2 * mx + (q & 1), ty = 2 * my + (q >> 1);
const int T = (tx & 3) + 4 * (ty & 3) + 16 * (tx >> 2) + 32 * (ty >> 2);
const uint32_t bytes = *(const uint32_t*)(OB + T * C + ch);
typedef _Float16 h4 __attribute__((ext_vector_type(4)));
const h4 v = *(const h4*)(XO + xo_row<C>(T) + ch);
#pragma unroll
for (int b = 0; b < 4; ++b)
a[4 * run + b] = e4m3_to_half((bytes >> (8 * b)) & 0xffu);
a[4 * rr + b] = v[b];
}
as[s] = as_gop(a);
}
#pragma unroll
for (int j = 0; j < NTP; ++j)
@@ -700,7 +710,7 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int r = 2 * i + hi;
const int r = mrow(i);
const int MX = (8 * bx - p.sx) / 2 + (r & 3), MY = (8 * by - p.sy) / 2 + (r >> 2);
if (MX < 0 || MX >= W2 || MY < 0 || MY >= H2)
continue;
@@ -734,35 +744,79 @@ __device__ __forceinline__ void swin_block(const CommonParams& p, const TubePara
#ifndef SWIN_PERSIST
#define SWIN_PERSIST 0 // persistent workgroups in x (0: one workgroup per original block)
#endif
// Module boilerplate: f16 weight image, prep kernel and its grid for the ZLUDA override hook.
#ifndef SWIN_PREP_SLOTS
#define SWIN_PREP_SLOTS 0 // weight images kept per weights pointer (0: one image, re-prepared when the pointer changes)
#endif
// Weight-image slots for kernels launched with several weight sets (the six enc3 tube blocks). The ZLUDA
// hook runs the prep only for a weights pointer it has not seen (d4r_prep_key_at / d4r_prep_key_slots);
// the prep claims the first free slot for that pointer, and pointers beyond S share the overflow slot S,
// which the hook re-prepares on every launch. Keys are written by one thread of the prep; any prep
// thread that reads before that write finds the same (first free) slot.
template <int S> __device__ __forceinline__ int prep_slot_find(const uint64_t* keys, uint64_t w)
{
for (int i = 0; i < S; ++i)
if (__hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT) == w)
return i;
return S;
}
template <int S> __device__ __forceinline__ int prep_slot_claim(uint64_t* keys, uint64_t w, bool writer)
{
// one pass: the first slot holding w or nothing. Keys only ever go from 0 to a pointer, so every
// prep thread picks the same slot whether it reads that slot before or after the writer's store.
for (int i = 0; i < S; ++i)
{
const uint64_t k = __hip_atomic_load(keys + i, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
if (k == w)
return i;
if (k == 0)
{
if (writer)
__hip_atomic_store(keys + i, w, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);
return i;
}
}
return S;
}
// Module boilerplate: f16 weight image(s), prep kernel and its grid for the ZLUDA override hook.
#define SWIN_MODULE(NAME, C, NW, NPM, TUBE, PARAMS, CIN) SWIN_MODULE_W(NAME, C, NW, NW, NPM, TUBE, PARAMS, CIN)
// NH heads on NWAVES waves; d4r_block_z tells the ZLUDA hook the block z size to launch with
#define SWIN_MODULE_W(NAME, C, NH, NWAVES, NPM, TUBE, PARAMS, CIN) \
using NAME##_L = SwinLayout<C, NH, NPM, CIN>; \
__device__ h16 g_w16[NAME##_L::TOTAL]; \
constexpr int NAME##_S = SWIN_PREP_SLOTS; \
__device__ wslot g_w16[(NAME##_S + 1) * NAME##_L::TOTAL]; \
__device__ uint64_t g_prep_keys[NAME##_S + 1]; \
__constant__ SwinDescs<C, NH, NPM, CIN> g_descs = make_descs<C, NH, NPM, CIN>(); \
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = NWAVES; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::TOTAL + 127) / 128; \
/* the prep reads only p.w (offset 0): skip it while the image for that pointer is resident */ \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1; \
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_slots = NAME##_S; \
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PARAMS p) \
{ \
const int idx = blockIdx.x * 128 + threadIdx.x; \
const uint8_t* w = ((const CommonParams*)&p)->w; \
const int slot = NAME##_S ? prep_slot_claim<NAME##_S>(g_prep_keys, (uint64_t)w, idx == 0) : 0; \
if (idx < NAME##_L::TOTAL) \
expand_weights(((const CommonParams*)&p)->w, g_w16, g_descs.d, NAME##_L::NDESC, idx); \
expand_weights(w, g_w16 + slot * NAME##_L::TOTAL, g_descs.d, NAME##_L::NDESC, idx); \
} \
extern "C" __device__ __attribute__((used)) uint32_t d4r_grid_x = SWIN_PERSIST; \
extern "C" __global__ void __launch_bounds__(32 * NWAVES) SWIN_VGPR_ATTR NAME(PARAMS p) \
{ \
const CommonParams& cp = *(const CommonParams*)&p; \
const int slot = NAME##_S ? prep_slot_find<NAME##_S>(g_prep_keys, (uint64_t)cp.w) : 0; \
const wslot* w16 = g_w16 + slot * NAME##_L::TOTAL; \
/* the original grid: 8x8-token blocks covering the shifted token grid */ \
const int gx = (cp.tw + cp.sx + 7) / 8, gy = (cp.th + cp.sy + 7) / 8; \
if (SWIN_PERSIST == 0) \
{ \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, blockIdx.x, blockIdx.y, gx); \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, blockIdx.x, blockIdx.y, gx); \
return; \
} \
_Pragma("unroll 1") for (int b = blockIdx.x; b < gx * gy; b += gridDim.x) \
{ \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, g_w16, b % gx, b / gx, gx); \
swin_block<C, NH, NWAVES, NPM, TUBE, CIN>(cp, (const TubeParams*)&p, w16, b % gx, b / gx, gx); \
__syncthreads(); \
} \
}
kernels/m/swin_common.h
+204 -26
View File
@@ -1,8 +1,29 @@
// Shared device code for native RDNA3 (gfx11 wave32 WMMA) rrlite Swin-block kernels (DLSS 310.7.0).
// Shared device code for native rrlite Swin-block kernels (DLSS 310.7.0; gfx11 / gfx12 wave32 WMMA).
// Semantics: ../ENC3_SPEC.md generalised (numpy reference ../swin_model.py).
//
// Activations are the WMMA A operand (lane = token row), weights the B operand (lane = channel column), and
// the f32 result has lane = channel column, VGPR i = token row wm_acc_row(i) (../common/wmma_layout.h).
// D4R_FP8_WMMA (gfx12 layout only): the GEMMs whose operands are both e4m3 values run as native FP8 WMMAs on
// e4m3 bytes (weights prepared as bytes, activations encoded when loaded); P V stays f16.
#pragma once
#include <hip/hip_runtime.h>
#include <stdint.h>
#include "../common/wmma_layout.h"
#if defined(D4R_FP8_WMMA) && D4R_WMMA_LAYOUT != 12 && defined(__AMDGCN__)
#error "D4R_FP8_WMMA needs the gfx12 WMMA layout"
#endif
// FP8 operands must hold e4m3 values, so the FP8 build re-quantises the intermediate activations as NVIDIA's
// network does (the gfx11 "fast numerics" SWIN_NO_Q8 keeps them f16, which no gfx12 WMMA can multiply with
// e4m3 weights)
// (SWIN_FORCE_Q8 does the same on any target: the gfx11 reference of the FP8 build's numerics)
#if (defined(D4R_FP8_WMMA) || defined(SWIN_FORCE_Q8) || defined(SWIN_EXACT)) && defined(SWIN_NO_Q8)
#undef SWIN_NO_Q8
#endif
// SWIN_EXACT (tests): NVIDIA's numerics throughout, the semantics of tools/swin_model.py (e4m3 activations and
// an f16 accumulator rounded after every k32 step)
#if defined(SWIN_EXACT) && defined(SWIN_F32ACC)
#undef SWIN_F32ACC
#endif
#pragma clang fp contract(off)
@@ -82,12 +103,9 @@ typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
typedef unsigned short u16x2 __attribute__((ext_vector_type(2)));
typedef short i16x2 __attribute__((ext_vector_type(2)));
// q8 on two halves with packed 16-bit integer / f16 ops (same per-half logic as q8; no carries cross halves)
__device__ __forceinline__ hv2 q8x2(hv2 h)
// q8_exact on two halves with packed 16-bit integer / f16 ops (same per-half logic; no carries cross halves)
__device__ __forceinline__ hv2 q8x2_exact(hv2 h)
{
#ifdef SWIN_NO_Q8
return h;
#endif
const u16x2 u = __builtin_bit_cast(u16x2, h);
const u16x2 m = u & (unsigned short)0x7fff;
u16x2 r = (m + (unsigned short)0x3f + ((m >> 7) & (unsigned short)1)) & (unsigned short)0x7f80;
@@ -102,6 +120,48 @@ __device__ __forceinline__ hv2 q8x2(hv2 h)
return __builtin_bit_cast(hv2, (u16x2)(r | (u & (unsigned short)0x8000)));
}
__device__ __forceinline__ hv2 q8x2(hv2 h)
{
#ifdef SWIN_NO_Q8
return h;
#endif
return q8x2_exact(h);
}
// enc8 of two halves with packed ops: returns the two e4m3 codes in the low bytes of the 16-bit
// halves, and q = q8_exact(h) (the values the codes decode to)
__device__ __forceinline__ uint32_t codes8x2(hv2 q);
__device__ __forceinline__ uint32_t enc8x2(hv2 h, hv2& q)
{
q = q8x2_exact(h);
return codes8x2(q);
}
// the e4m3 codes of two halves that already hold e4m3 values (q8_exact results), in the low bytes of the
// 16-bit halves
__device__ __forceinline__ uint32_t codes8x2(hv2 q)
{
const u16x2 qb = __builtin_bit_cast(u16x2, q);
const u16x2 m = qb & (unsigned short)0x7fff;
// normals: the exponent moves from bias 15 to bias 7
const u16x2 norm = (u16x2)(m >> 7) - (unsigned short)64;
// subnormals are k * 2^-9 (k = 0..7): f16(k + 1024) has k in its low bits
const hv2 mag = __builtin_bit_cast(hv2, m);
const hv2 k = (hv2)(mag * (hv2){(half_t)512.0f, (half_t)512.0f}) + (hv2){(half_t)1024.0f, (half_t)1024.0f};
const u16x2 sub = __builtin_bit_cast(u16x2, k) - (unsigned short)0x6400;
const u16x2 lt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(m - (unsigned short)0x2400)) >> 15));
u16x2 code = (sub & lt) | (norm & ~lt);
const u16x2 gt = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((unsigned short)0x7c00 - m)) >> 15));
code = ((u16x2){0x7f, 0x7f} & gt) | (code & ~gt);
return __builtin_bit_cast(uint32_t, (u16x2)(((qb >> 8) & (unsigned short)0x80) | code));
}
// four e4m3 codes (a, b from enc8x2) as the bytes of one word in element order
__device__ __forceinline__ uint32_t pack_codes(uint32_t a, uint32_t b)
{
return __builtin_amdgcn_perm(b, a, 0x06040200u);
}
// e4m3 code of h (RNE, satfinite) derived from q8(h): normals shift the exponent, subnormals are k * 2^-9
__device__ __forceinline__ uint32_t enc8(half_t h)
{
@@ -153,9 +213,105 @@ __device__ __forceinline__ int woff(int base, int Ks, int Ns, int k, int n)
4 * ((k & 31) >> 4) + (k & 3);
}
__device__ __forceinline__ f8v wmma(h16 a, h16 b, f8v c)
// ---------------------------------------------------------------- WMMA operands
// half-wave of this lane: blocks are 32 x 1 x waves, so the lane is threadIdx.x (the same value as wm_half(),
// but from the source the kernels already use for their lane index)
__device__ __forceinline__ int m_half()
{
return (int)threadIdx.x >> 4;
}
// token row that accumulator VGPR i holds in this lane
__device__ __forceinline__ int mrow(int i)
{
return wm_acc_row_h(i, m_half());
}
// sop: an f16 operand (gfx11: the row's 16 K values; gfx12: this half's 8). gop: an operand of the e4m3 GEMMs
// (k32 steps): sop, or with D4R_FP8_WMMA the 8 e4m3 bytes of this half's K values.
#if D4R_WMMA_LAYOUT == 12
typedef h8 sop;
#else
typedef h16 sop;
#endif
typedef uint32_t u2v __attribute__((ext_vector_type(2)));
typedef uint32_t u4v_t __attribute__((ext_vector_type(4)));
#ifdef D4R_FP8_WMMA
typedef u2v gop;
typedef u4v_t wslot; // weight image slot: 16 e4m3 bytes (K slots 0..15)
#else
typedef sop gop;
typedef h16 wslot; // weight image slot: 16 f16 values
#endif
__device__ __forceinline__ f8v wmma(sop a, sop b, f8v c)
{
#if D4R_WMMA_LAYOUT == 12
return wm_mma(__builtin_bit_cast(wm_op, a), __builtin_bit_cast(wm_op, b), c);
#else
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a, b, c);
#endif
}
// f16 operand from 16 contiguous halves in LDS (gfx12: this half's 8 of them)
__device__ __forceinline__ sop lds16(const half_t* p)
{
#if D4R_WMMA_LAYOUT == 12
return *(const h8*)(p + 8 * m_half());
#else
h8 lo = *(const h8*)p, hi = *(const h8*)(p + 8);
return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
#endif
}
#ifdef D4R_FP8_WMMA
// e4m3 bytes of an f16 operand whose values are e4m3 values (K order kept: byte j = element j)
__device__ __forceinline__ gop to_fp8(sop v)
{
uint32_t c[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
c[j] = codes8x2((hv2){v[2 * j], v[2 * j + 1]});
return (u2v){__builtin_amdgcn_perm(c[1], c[0], 0x06040200u), __builtin_amdgcn_perm(c[3], c[2], 0x06040200u)};
}
// f32 += A B over 16 K for e4m3 operands
__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
{
#if defined(__GFX12__)
return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(wm_i2v, a), __builtin_bit_cast(wm_i2v, b), c);
#else
// layout shim: the e4m3 values widened to f16 (exact), then the f16 WMMA
auto widen = [](gop q) {
sop r;
#pragma unroll
for (int j = 0; j < 8; ++j)
r[j] = e4m3_to_half((q[j >> 2] >> (8 * (j & 3))) & 0xffu);
return r;
};
return wmma(widen(a), widen(b), c);
#endif
}
__device__ __forceinline__ gop lda(const half_t* p)
{
return to_fp8(lds16(p));
}
#else
__device__ __forceinline__ f8v wmma8(gop a, gop b, f8v c)
{
return wmma(a, b, c);
}
// GEMM A operand (activations) from LDS
__device__ __forceinline__ gop lda(const half_t* p)
{
return lds16(p);
}
#endif
// an f16 operand built in registers -> GEMM operand
__device__ __forceinline__ gop as_gop(sop v)
{
#ifdef D4R_FP8_WMMA
return to_fp8(v);
#else
return v;
#endif
}
__device__ __forceinline__ f8v round16(f8v v)
@@ -166,7 +322,7 @@ __device__ __forceinline__ f8v round16(f8v v)
return v;
}
// A C-fragment tile: element i is row 2i + lane/16. Exact mode stores packed f16 (NVIDIA's f16
// A C-fragment tile: element i is row wm_acc_row(i). Exact mode stores packed f16 (NVIDIA's f16
// accumulator, rounded after every k32 step); SWIN_F32ACC keeps f32 through whole GEMM chains
// (more precise than NVIDIA; values are rounded to f16 where they are consumed).
struct T16
@@ -217,20 +373,20 @@ __device__ __forceinline__ void mixacc(T16& acc, f8v P)
__device__ unsigned d4r_k32_bad[4 + 64 * 4];
#endif
// f16 accumulate of one k32 step (NVIDIA m16n8k32 semantics: products summed, one rounding with C)
__device__ __forceinline__ void k32(T16& acc, h16 a0, h16 b0, h16 a1, h16 b1)
__device__ __forceinline__ void k32(T16& acc, gop a0, gop b0, gop a1, gop b1)
{
#if defined(SWIN_F32ACC)
acc.f = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, acc.f));
acc.f = wmma8(a1, b1, wmma8(a0, b0, acc.f));
#elif defined(K32_MIX)
f8v P = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
P = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, P);
f8v P = wmma8(a0, b0, (f8v){0, 0, 0, 0, 0, 0, 0, 0});
P = wmma8(a1, b1, P);
mixacc(acc, P);
#else // exact: C added inside the f32 WMMA, f16 rounding after the k32 step
f8v c;
#pragma unroll
for (int i = 0; i < 8; ++i)
c[i] = (float)acc.get(i);
c = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a1, b1, __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a0, b0, c));
c = wmma8(a1, b1, wmma8(a0, b0, c));
#pragma unroll
for (int i = 0; i < 8; ++i)
acc.set(i, (half_t)c[i]);
@@ -273,12 +429,6 @@ __device__ __forceinline__ half_t hload(const uint8_t* w, int off)
return __builtin_bit_cast(half_t, *(const uint16_t*)(w + off));
}
__device__ __forceinline__ h16 lds16(const half_t* p)
{
h8 lo = *(const h8*)p, hi = *(const h8*)(p + 8);
return __builtin_shufflevector(lo, hi, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15);
}
__device__ __forceinline__ void wave_sync()
{
__builtin_amdgcn_fence(__ATOMIC_RELEASE, "wavefront");
@@ -296,21 +446,36 @@ __device__ __forceinline__ int tok_y(int T)
return ((T >> 2) & 3) + 4 * (T >> 5);
}
// ---------------------------------------------------------------- f16 weight images
// A GEMM's B operands expanded to f16 in WMMA lane order: h16 index
// dst + ((kc * 2 + s) * NT + nt) * 16 + lane16, holding logical k = 32 kc + kslot(s, i), n = 16 nt + lane16.
// ---------------------------------------------------------------- weight images
// A GEMM's B operands in WMMA lane order: slot index dst + ((kc * 2 + s) * NT + nt) * 16 + lane16 holds the 16
// K slots of logical k = 32 kc + kslot(s, i), n = 16 nt + lane16 (i = 0..15), as f16 values (wslot = h16) or,
// with D4R_FP8_WMMA, as the e4m3 bytes. gfx12 lanes read the half of the slot with their 8 K slots.
struct GemmDesc
{
int dst; // in h16 units
int dst; // in slots
int base, Ks, Ns, KC, NT;
};
__device__ __forceinline__ h16 bload(const h16* w16, int gemm_dst, int NT, int kc, int s, int nt, int l16)
// The tile index is wave-uniform (it depends only on the wave, the loop counters and constants), so it
// is formed in scalar registers and the load uses a scalar base plus the lane's constant offset.
__device__ __forceinline__ gop bload(const wslot* w16, int gemm_dst, int NT, int kc, int s, int nt, int l16)
{
return w16[gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16];
#ifdef SWIN_SCALAR_BLOAD
const int tile = __builtin_amdgcn_readfirstlane(gemm_dst + ((kc * 2 + s) * NT + nt) * 16);
const wslot* t = w16 + tile;
const int at = l16;
#else
const wslot* t = w16;
const int at = gemm_dst + ((kc * 2 + s) * NT + nt) * 16 + l16;
#endif
#if D4R_WMMA_LAYOUT == 12
return ((const gop*)t)[2 * at + m_half()];
#else
return t[at];
#endif
}
__device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const GemmDesc* descs, int ndesc, int idx)
__device__ __forceinline__ void expand_weights(const uint8_t* w, wslot* w16, const GemmDesc* descs, int ndesc, int idx)
{
int d = 0;
while (d + 1 < ndesc && idx >= descs[d + 1].dst)
@@ -322,9 +487,22 @@ __device__ __forceinline__ void expand_weights(const uint8_t* w, h16* w16, const
const int l16 = local & 15, rest = local >> 4;
const int nt = rest % g.NT, s = (rest / g.NT) & 1, kc = rest / g.NT / 2;
const int n = 16 * nt + l16;
#ifdef D4R_FP8_WMMA
u4v_t v;
#pragma unroll
for (int j = 0; j < 4; ++j)
{
uint32_t word = 0;
#pragma unroll
for (int b = 0; b < 4; ++b)
word |= (uint32_t)w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, 4 * j + b), n)] << (8 * b);
v[j] = word;
}
#else
h16 v;
#pragma unroll
for (int i = 0; i < 16; ++i)
v[i] = e4m3_to_half(w[woff(g.base, g.Ks, g.Ns, 32 * kc + kslot(s, i), n)]);
#endif
w16[idx] = v;
}
kernels/tex/build_tex.sh
+14 -7
View File
@@ -2,25 +2,32 @@
# build_tex.sh KERNEL SRC OUT_DIR: one texture kernel. NVIDIA's PTX for KERNEL, edited by make_ptx.py so
# part of it calls functions from SRC.hip, is compiled by ZLUDA with SRC's bitcode linked in
# (D4R_ZLUDA_EXTRA_BC) and the resulting code object is saved as OUT_DIR/KERNEL.hsaco.
# Called by kernels/build.sh, which sets D4R_ROCM_DIR, D4R_GPU_ARCH, D4R_DLSS_PTX_DIR, D4R_PTXLOAD and
# needs D4R_ZLUDA_BUILD.
# Called by kernels/build.sh, which sets D4R_ROCM_DIR, D4R_GPU_ARCH, D4R_DLSS_PTX_DIR, D4R_ZLUDA_EMIT,
# D4R_ZLUDA_WAVE64 (1 = wave64 code object) and D4R_TEX_FP8 (1 = the gfx12 native-FP8 variant).
# ZLUDA compiles offline for D4R_GPU_ARCH with its d4r_emit tool, so any target builds on any machine.
set -euo pipefail
K=$1 SRC=$2 OUT=$3
D="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
R="${D4R_ROCM_DIR:-/opt/rocm}"
ARCH="${D4R_GPU_ARCH:-gfx1101}"
FP8="${D4R_TEX_FP8:-0}"
LLVM="$R/lib/llvm/bin"
B="$(mktemp -d)"
trap 'rm -rf "$B"' EXIT
# D4R_TEX_CFLAGS: extra flags for SRC.hip (tests, e.g. -DD4R_WMMA_LAYOUT=12 for the gfx12-layout shim on gfx11)
read -r -a FLAGS <<< "${D4R_TEX_CFLAGS:-}"
[[ "$FP8" == 1 ]] && FLAGS+=(-DD4R_TEX_FP8)
# SRC.hip -> bitcode without target attributes, so ZLUDA's own target settings apply after linking
"$LLVM/clang++" -x hip -std=c++20 -nogpuinc -nogpulib -O3 -mno-wavefrontsize64 --offload-device-only \
--offload-arch="$ARCH" -fgpu-rdc -emit-llvm -c -Xclang -fdenormal-fp-math=dynamic -o "$B/raw.bc" "$D/$SRC.hip"
--offload-arch="$ARCH" -fgpu-rdc -emit-llvm -c -Xclang -fdenormal-fp-math=dynamic -DD4R_KERNEL_NAME="$K" \
"${FLAGS[@]}" -o "$B/raw.bc" "$D/$SRC.hip"
"$LLVM/llvm-dis" "$B/raw.bc" -o "$B/raw.ll"
sed -E -e '/@llvm.used/d' -e '/wchar_size/d' -e '/llvm.module.flags/d' -e '/__hip_cuid/d' -e 's/optnone//g' \
-e "s/\"target-cpu\"=\"$ARCH\"//g" -e 's/"target-features"="[^"]+"//g' "$B/raw.ll" | "$LLVM/llvm-as" -o "$B/extra.bc" -
python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
env LD_LIBRARY_PATH="$D4R_ZLUDA_BUILD:$R/lib" XDG_CACHE_HOME="$B/jit" D4R_ZLUDA_DUMP_DIR="$B/dump" \
D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 D4R_ZLUDA_IGNORE_DENORMAL=1 \
D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 "$D4R_PTXLOAD" "$B/$K.ptx" >/dev/null
# the code-generation switches of the runtime (d4r.ini [Kernels]), which ZLUDA's module cache keys on
env LD_LIBRARY_PATH="$R/lib" D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 \
D4R_ZLUDA_WMMA_FP8_NATIVE="$FP8" D4R_ZLUDA_IGNORE_DENORMAL=1 D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 \
D4R_ZLUDA_WAVE64="${D4R_ZLUDA_WAVE64:-0}" "$D4R_ZLUDA_EMIT" "$B/$K.ptx" "$B/dump" "$ARCH" >/dev/null
cp "$B/dump/module.hsaco" "$OUT/$K.hsaco"
printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size|wavefront_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
kernels/tex/dec0_head.hip
+63
View File
@@ -6,8 +6,14 @@
// weights as WMMA A images: index ((p * 2 + c) * 2 + h) * 16 + n16 for channel pair p (16p..16p+15),
// k32 chunk c (input plane), k16 half h
#if D4R_WMMA_LAYOUT == 12
__attribute__((device)) kslot g_dec0_w16[8 * 2 * 2 * 16];
#else
__attribute__((device)) u8v g_dec0_w16[8 * 2 * 2 * 16];
#endif
DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
struct Dec0Params
{
@@ -27,7 +33,11 @@ KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4
#pragma unroll
for (int j = 0; j < 4; ++j)
d[j] = *(const uint32_t*)(w + 512 * (2 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
#if D4R_WMMA_LAYOUT == 12
g_dec0_w16[idx] = kslot_from16(d);
#else
g_dec0_w16[idx] = operand16(d);
#endif
}
__attribute__((device)) static inline uint32_t mirror_clamp(int32_t v, uint32_t full, uint32_t half)
@@ -47,6 +57,58 @@ DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx,
const uint32_t Y = mirror_clamp(oy / 2 + (int32_t)(m >> 2), dimy, H2);
const uint32_t Xc = mirror_clamp(ox / 2 + (int32_t)(m & 3), dimx, W2);
const uint32_t* in = (const uint32_t*)inp + rowstride * Y + 8 * Xc;
#if D4R_WMMA_LAYOUT == 12
kop X[2][2];
#pragma unroll
for (int c = 0; c < 2; ++c)
#pragma unroll
for (int h = 0; h < 2; ++h)
X[c][h] = kop_from8(*(const uint2_t*)(in + plane * c + 4 * h + 2 * hf));
// E[p]: e4m3 codes of channels 16p + 8hf .. 16p + 8hf + 7
uint32_t E[8][2];
#pragma unroll
for (int p = 0; p < 8; ++p)
{
f8v cf;
#pragma unroll
for (int i = 0; i < 8; ++i)
cf[i] = (float)*(const half_t*)(w + 8192 + 2 * (16 * p + i + 8 * hf));
half_t acc[8];
#pragma unroll
for (int c = 0; c < 2; ++c)
{
const int slot = ((p * 2 + c) * 2) * 16 + m;
cf = k32_e4m3(cf, kop_image(g_dec0_w16, slot), X[c][0], kop_image(g_dec0_w16, slot + 16), X[c][1]);
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = (half_t)cf[i];
#pragma unroll
for (int i = 0; i < 8; ++i)
cf[i] = (float)acc[i];
}
uint32_t codes[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
codes[k] = e4m3x2(pack2(acc[2 * k], acc[2 * k + 1]));
E[p][0] = codes[0] | (codes[1] << 16);
E[p][1] = codes[2] | (codes[3] << 16);
}
// staging (layout as below): this lane holds tiles 4G + hf and 4G + 2 + hf (channels 16p + 8hf ..), so it
// writes H = hf itself
const uint32_t base = 1024 * hf + 32 * (m & 3) + 256 * (m >> 2);
const uint32_t goff[4] = {0, 16, 128, 144};
#pragma unroll
for (int G = 0; G < 4; ++G)
{
const uint32_t* p0 = E[2 * G];
const uint32_t* p1 = E[2 * G + 1];
u4v words = {__builtin_amdgcn_perm(p1[0], p0[0], 0x05040100u), __builtin_amdgcn_perm(p1[0], p0[0], 0x07060302u),
__builtin_amdgcn_perm(p1[1], p0[1], 0x05040100u), __builtin_amdgcn_perm(p1[1], p0[1], 0x07060302u)};
*(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
}
#else
u8v X[2][2];
#pragma unroll
for (int c = 0; c < 2; ++c)
@@ -109,4 +171,5 @@ DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx,
__builtin_amdgcn_perm(b23, a23, 0x05040100u), __builtin_amdgcn_perm(b23, a23, 0x07060302u)};
*(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
}
#endif
}
kernels/tex/enc0_tail.hip
+82 -2
View File
@@ -1,4 +1,4 @@
// Native tail of rrlite_enc0_4x4_mvhi_hdr_folded: the 2x2-pixel patch embedding GEMM (16 tokens x
// Native tail of rrlite_enc0_4x4_*_folded (all flag combinations): the 2x2-pixel patch embedding GEMM (16 tokens x
// K 128 x N 64, e4m3 inputs, f16 accumulation per k32 step, f16 bias), clamp to +-2pi and e4m3 store.
// The feature code before it is ZLUDA's compile of the original PTX (enc0_tail.ptx calls this after
// the kernel's own shared-memory staging); the arithmetic reproduces ZLUDA's m16n8k32 lowering
@@ -7,8 +7,14 @@
// B operands (weights) as WMMA A images: index ((p * 4 + c) * 2 + h) * 16 + n16 for channel pair p
// (channels 16p..16p+15), k32 chunk c, k16 half h.
#if D4R_WMMA_LAYOUT == 12
__attribute__((device)) kslot g_enc0_w16[4 * 4 * 2 * 16];
#else
__attribute__((device)) u8v g_enc0_w16[4 * 4 * 2 * 16];
#endif
DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
struct Enc0Params
{
@@ -16,7 +22,13 @@ struct Enc0Params
uint8_t rest[352];
};
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4x4_mvhi_hdr_folded_prep(Enc0Params prm)
// every flag combination (mvhi/mvlo, hdr/ldr) has this tail; build_tex.sh names the prep after the kernel
#ifndef D4R_KERNEL_NAME
#define D4R_KERNEL_NAME rrlite_enc0_4x4_mvhi_hdr_folded
#endif
#define D4R_CAT2(a, b) a##b
#define D4R_CAT(a, b) D4R_CAT2(a, b)
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void D4R_CAT(D4R_KERNEL_NAME, _prep)(Enc0Params prm)
{
const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
if (idx >= 512)
@@ -29,7 +41,11 @@ KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4
#pragma unroll
for (int j = 0; j < 4; ++j)
d[j] = *(const uint32_t*)(w + 4096 + 512 * (4 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
#if D4R_WMMA_LAYOUT == 12
g_enc0_w16[idx] = kslot_from16(d);
#else
g_enc0_w16[idx] = operand16(d);
#endif
}
DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx, uint32_t dimy, int32_t ox, int32_t oy)
@@ -38,6 +54,69 @@ DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx,
const uint8_t* w = (const uint8_t*)wp;
lds_u8* s = (lds_u8*)(uintptr_t)smem;
#if D4R_WMMA_LAYOUT == 12
// activations: token m's 128 e4m3 features (addresses as below); this half's 8 bytes of each k16 part
const uint32_t rowoff = 32 * (m & 3) + 256 * (m >> 2);
const uint32_t choff[4] = {0, 16, 128, 144};
kop X[4][2];
#pragma unroll
for (int c = 0; c < 4; ++c)
#pragma unroll
for (int h = 0; h < 2; ++h)
X[c][h] = kop_from8(*(const __attribute__((address_space(3))) uint2_t*)(s + 1024 * h + rowoff + choff[c] + 8 * hf));
const hv2 lo = {(half_t)-6.2831855f, (half_t)-6.2831855f}, hi = {(half_t)6.2831855f, (half_t)6.2831855f};
// E[p]: e4m3 codes of channels 16p + 8hf .. 16p + 8hf + 7 (D^T element i = channel 16p + i + 8hf)
uint32_t E[4][2];
#pragma unroll
for (int p = 0; p < 4; ++p)
{
f8v cf;
#pragma unroll
for (int i = 0; i < 8; ++i)
cf[i] = (float)*(const half_t*)(w + 12288 + 2 * (16 * p + i + 8 * hf));
half_t acc[8];
#pragma unroll
for (int c = 0; c < 4; ++c)
{
const int slot = ((p * 4 + c) * 2) * 16 + m;
cf = k32_e4m3(cf, kop_image(g_enc0_w16, slot), X[c][0], kop_image(g_enc0_w16, slot + 16), X[c][1]);
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[i] = (half_t)cf[i];
#pragma unroll
for (int i = 0; i < 8; ++i)
cf[i] = (float)acc[i];
}
uint32_t codes[4];
#pragma unroll
for (int k = 0; k < 4; ++k)
{
hv2 v = {acc[2 * k], acc[2 * k + 1]};
v = __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, lo), hi);
codes[k] = e4m3x2(__builtin_bit_cast(uint32_t, v));
}
E[p][0] = codes[0] | (codes[1] << 16);
E[p][1] = codes[2] | (codes[3] << 16);
}
// plane pl, word 4 j + q = columns 2q, 2q + 1 of tile (2 pl, j) then of tile (2 pl + 1, j) (tile (p, j) =
// channels 16p + 8j ..): this lane holds tiles j = hf, i.e. words 4 hf .. 4 hf + 3 of both planes
const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
const int32_t Y = oy / 2 + (int32_t)(m >> 2), Xc = ox / 2 + (int32_t)(m & 3);
if ((uint32_t)Y < H2 && (uint32_t)Xc < W2)
{
#pragma unroll
for (int pl = 0; pl < 2; ++pl)
{
const uint32_t* p0 = E[2 * pl];
const uint32_t* p1 = E[2 * pl + 1];
const u4v words = {__builtin_amdgcn_perm(p1[0], p0[0], 0x05040100u), __builtin_amdgcn_perm(p1[0], p0[0], 0x07060302u),
__builtin_amdgcn_perm(p1[1], p0[1], 0x05040100u), __builtin_amdgcn_perm(p1[1], p0[1], 0x07060302u)};
*(u4v*)((uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * pl + 4 * hf) = words;
}
}
#else
// activations: token m's 128 e4m3 features; A[m][k] of chunk c at
// 1024 * (k >> 4) + 32 * (m & 3) + 256 * (m >> 2) + (k & 15) + {0, 16, 128, 144}[c]
const uint32_t rowoff = 32 * (m & 3) + 256 * (m >> 2);
@@ -121,4 +200,5 @@ DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx,
uint32_t* out = (uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * hf;
*(u8v*)out = words;
}
#endif
}
kernels/tex/make_ptx.py
+97 -24
View File
@@ -11,6 +11,28 @@ from pathlib import Path
# PTX extracted from nvngx_dlss.dll by kernels/tools/extract_dlss_ptx.py (kernels/build.sh does this)
PTX_DIR = Path(os.environ.get("D4R_DLSS_PTX_DIR", Path(__file__).resolve().parent.parent / "extracted" / "ptx"))
ENC0_TAIL = """{{
.param .b32 q0;
st.param.b32 [q0+0], {smem};
.param .b64 q1;
st.param.b64 [q1+0], {w};
.param .b64 q2;
st.param.b64 [q2+0], {out};
.param .b32 q3;
st.param.b32 [q3+0], {dimx};
.param .b32 q4;
st.param.b32 [q4+0], {dimy};
.param .b32 q5;
st.param.b32 [q5+0], {ox};
.param .b32 q6;
st.param.b32 [q6+0], {oy};
call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
}}
ret;
}}
"""
KERNELS = {
# enc0: cut after the bar.warp.sync that ends the smem staging of the GEMM A operand
"rrlite_enc0_4x4_mvhi_hdr_folded": dict(
@@ -25,27 +47,7 @@ KERNELS = {
)
;
""",
tail="""{
.param .b32 q0;
st.param.b32 [q0+0], %r259;
.param .b64 q1;
st.param.b64 [q1+0], %rd106;
.param .b64 q2;
st.param.b64 [q2+0], %rd78;
.param .b32 q3;
st.param.b32 [q3+0], %r1066;
.param .b32 q4;
st.param.b32 [q4+0], %r1067;
.param .b32 q5;
st.param.b32 [q5+0], %r1230;
.param .b32 q6;
st.param.b32 [q6+0], %r1233;
call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
}
ret;
}
""",
tail=ENC0_TAIL.format(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233"),
),
# dec0: the GEMM (from the B fragment loads to the e4m3 staging stores) becomes a native call; the
# A loads before it are left dead. %r560 (shared base) is the only register of the range used later.
@@ -206,12 +208,83 @@ def replace_rz_round(lines):
return out, n
ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
# the other flag combinations of the kernels whose only native part is the surface-store rewrite
SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
r"rrlite_post_3_[12]_mv(hi|lo)_(hdr|ldr)_folded|rrlite_downsample_kernel_(static|dynamic)_(hdr|ldr))$")
ENC0_RE = re.compile(r"^rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)_folded$")
def read_lines(path):
return path.read_text().replace("\r\n", "\n").split("\n")
def enc0_registers(lines):
"""Cut line (1-based) and the tail call's registers of an enc0 variant. The code after the third
bar.warp.sync is the same GEMM in every flag combination, only register numbers differ: the shared
base and weights pointer are read off by aligning that block with the reference variant; the output
pointer and grid size are the kernel's loads of parameters +40 and +8; the block offsets are the two
values halved (shr 31 / add) by the epilogue."""
syncs = [i for i, l in enumerate(lines) if l.strip() == "bar.warp.sync -1;"]
if len(syncs) != 3:
sys.exit(f"enc0: expected 3 bar.warp.sync, found {len(syncs)}")
cut = syncs[2] + 1
before, after = lines[:cut], lines[cut:]
out = [m.group(1) for m in (re.match(r"^ld\.param\.u64 (%rd\d+), \[%rd\d+\+40\];$", l.strip()) for l in before) if m]
dims = [m.groups() for m in (re.match(r"^ld\.param\.v2\.u32 \{(%r\d+), (%r\d+)\}, \[%rd\d+\+8\];$", l.strip()) for l in before) if m]
halves = []
for a, b in zip(after, after[1:]):
m = re.match(r"^shr\.u32 (%r\d+), (%r\d+), 31;$", a.strip())
if m and re.match(rf"^add\.s32 %r\d+, {re.escape(m.group(2))}, {re.escape(m.group(1))};$", b.strip()):
halves.append(m.group(2))
if len(out) != 1 or len(dims) != 1 or len(halves) < 2:
sys.exit(f"enc0: parameter loads or epilogue not found ({out}, {dims}, {halves[:2]})")
return cut, dict(out=out[0], dimx=dims[0][0], dimy=dims[0][1], oy=halves[0], ox=halves[1])
REG_RE = re.compile(r"%(?:rd|rs|r|fd|f|p)\d+")
def enc0_variant(name):
ref = read_lines(module_file(ENC0_REF, KERNELS[ENC0_REF]))
path = module_file(name, {"file": ""})
lines = read_lines(path)
rcut, rregs = enc0_registers(ref)
k = KERNELS[ENC0_REF]
known = dict(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233")
if rcut != k["cut_after"] or any(rregs[key] != known[key] for key in rregs):
sys.exit(f"enc0: the register search does not reproduce {ENC0_REF} ({rcut}, {rregs})")
cut, regs = enc0_registers(lines)
# align the GEMM block (up to the reference's first closing brace) and map registers
n = next(i for i, l in enumerate(ref[rcut:]) if l.startswith("}"))
mapping = {}
for a, b in zip(ref[rcut:rcut + n], lines[cut:cut + n]):
if REG_RE.sub("R", a) != REG_RE.sub("R", b):
sys.exit(f"enc0: {name} differs from {ENC0_REF} after the cut: {a!r} / {b!r}")
for x, y in zip(REG_RE.findall(a), REG_RE.findall(b)):
if mapping.setdefault(x, y) != y:
sys.exit(f"enc0: inconsistent register mapping {x} -> {mapping[x]} / {y}")
regs.update(smem=mapping[known["smem"]], w=mapping[known["w"]])
return dict(file=path.name, cut_after=cut, cut_check="bar.warp.sync -1;", sust=True, extern=k["extern"],
tail=ENC0_TAIL.format(**regs))
def kernel_spec(name):
if name in KERNELS:
return KERNELS[name]
if ENC0_RE.match(name):
return enc0_variant(name)
if SUST_ONLY_RE.match(name):
return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_output"))
sys.exit(f"make_ptx: no recipe for {name}")
def module_file(name, k):
"""The extracted PTX module that defines NAME: the 310.7 file number, else any module that does
(other DLSS versions number their modules differently)."""
entry = re.compile(rf"\.entry\s+{re.escape(name)}\s*\(")
known = PTX_DIR / k["file"]
if known.exists() and entry.search(known.read_text()):
if k["file"] and known.is_file() and entry.search(known.read_text()):
return known
for path in sorted(PTX_DIR.glob("*.ptx")):
if entry.search(path.read_text()):
@@ -221,8 +294,8 @@ def module_file(name, k):
def main():
name, out = sys.argv[1], Path(sys.argv[2])
k = KERNELS[name]
lines = module_file(name, k).read_text().replace("\r\n", "\n").split("\n")
k = kernel_spec(name)
lines = read_lines(module_file(name, k))
if "cut_after" in k:
cut = k["cut_after"]
if lines[cut - 1].strip() != k["cut_check"]:
kernels/tex/tex_common.h
+73
View File
@@ -3,6 +3,10 @@
// only plain clang builtins are used (compiled with -nogpuinc -nogpulib).
#pragma once
#include <stdint.h>
#include "../common/wmma_layout.h"
#if defined(D4R_TEX_FP8) && D4R_WMMA_LAYOUT != 12 && defined(__HIP_DEVICE_COMPILE__)
#error "D4R_TEX_FP8 needs the gfx12 WMMA layout"
#endif
#pragma clang fp contract(off)
@@ -65,6 +69,75 @@ __attribute__((device)) static inline f8v wmma(u8v a, u8v b, f8v c)
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
}
// ---------------------------------------------------------------- gfx12 m16n8k32 e4m3 steps
// On gfx12 the GEMM tails do a k32 step as ZLUDA's gfx12 lowering of NVIDIA's m16n8k32 e4m3 MMA does, so they
// stay bit-identical to ZLUDA's compile of the kernel they replace: a lane's k16 operand holds this half's
// e4m3 k 8h .. 8h + 7 (ZLUDA's gfx12 slot order), either widened as above with C * 2^-16 inside the f32
// WMMAs (d4r_wmma12_*), or with D4R_TEX_FP8 (ZLUDA's D4R_ZLUDA_WMMA_FP8_NATIVE) as the bytes of a native FP8
// WMMA with no scaling (d4r_wmma12f8_*). (The gfx11 tails use operand16 and wmma above.)
#if D4R_WMMA_LAYOUT == 12
// 8 e4m3 bytes (k in order) widened to value * 2^-8 f16 pairs, k in order
__attribute__((device)) static inline u4v widen8(uint2_t d)
{
uint32_t a0, a1, b0, b1;
e4m3x4_scaled(d.x, a0, a1);
e4m3x4_scaled(d.y, b0, b1);
return (u4v){a0, a1, b0, b1};
}
#ifdef D4R_TEX_FP8
typedef uint2_t kop;
typedef u4v kslot; // weight image slot: the 16 e4m3 bytes of one k16 operand row
__attribute__((device)) static inline kop kop_from8(uint2_t d)
{
return d;
}
__attribute__((device)) static inline kslot kslot_from16(u4v d)
{
return d;
}
#else
typedef u4v kop;
typedef u8v kslot; // weight image slot: the 16 widened values of one k16 operand row, k in order
__attribute__((device)) static inline kop kop_from8(uint2_t d)
{
return widen8(d);
}
__attribute__((device)) static inline kslot kslot_from16(u4v d)
{
const u4v lo = widen8((uint2_t){d[0], d[1]}), hi = widen8((uint2_t){d[2], d[3]});
return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
}
#endif
// this half's k16 operand from 16 consecutive e4m3 bytes of NVIDIA's k order
__attribute__((device)) static inline kop kop_from16(u4v d)
{
return kop_from8(lane_id() >= 16 ? (uint2_t){d[2], d[3]} : (uint2_t){d[0], d[1]});
}
// this half's operand from slot idx of a weight image
__attribute__((device)) static inline kop kop_image(const kslot* img, int idx)
{
return ((const kop*)img)[2 * idx + (lane_id() >> 4)];
}
// f32 C (true scale) + one k32 step
__attribute__((device)) static inline f8v k32_e4m3(f8v c, const kop& a0, const kop& b0, const kop& a1, const kop& b1)
{
#if defined(D4R_TEX_FP8) && defined(__GFX12__)
typedef int i2v __attribute__((ext_vector_type(2)));
c = __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(i2v, a0), __builtin_bit_cast(i2v, b0), c);
return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(i2v, a1), __builtin_bit_cast(i2v, b1), c);
#elif defined(D4R_TEX_FP8)
// layout shim, like ZLUDA's: each FP8 WMMA as the scaled widening (exact powers of two)
c = wm_mma_zluda(widen8(a0), widen8(b0), c * 0x1p-16f) * 0x1p16f;
return wm_mma_zluda(widen8(a1), widen8(b1), c * 0x1p-16f) * 0x1p16f;
#else
c = c * 0x1p-16f;
c = wm_mma_zluda(a0, b0, c);
c = wm_mma_zluda(a1, b1, c);
return c * 0x1p16f;
#endif
}
#endif
// ZLUDA's f16x2_to_e4m3x2_satfinite_bits: RNE satfinite e4m3 of both halves, low half -> low byte.
__attribute__((device)) static inline uint32_t e4m3x2(uint32_t bits)
{
kernels/tools/dump_runner.cpp
+19 -1
View File
@@ -44,7 +44,18 @@ int main(int argc, char** argv)
hipFunction_t prep = nullptr;
if (hipModuleGetFunction(&prep, mod, (std::string(argv[2]) + "_prep").c_str()) != hipSuccess)
prep = nullptr;
const unsigned prepGrid = getenv("PREP_GRID") ? atoi(getenv("PREP_GRID")) : 88;
// the prep grid: PREP_GRID, else the module's d4r_prep_blocks (what the ZLUDA hook launches), else 88
unsigned prepGrid = 88;
{
hipDeviceptr_t pb;
size_t pbs = 0;
uint32_t v = 0;
if (hipModuleGetGlobal(&pb, &pbs, mod, "d4r_prep_blocks") == hipSuccess && pbs == 4 &&
hipMemcpyDtoH(&v, pb, 4) == hipSuccess && v != 0)
prepGrid = v;
}
if (getenv("PREP_GRID"))
prepGrid = atoi(getenv("PREP_GRID"));
std::vector<char> args = read_file(dump + "/args.bin");
std::ifstream man(dump + "/manifest.txt");
@@ -58,6 +69,13 @@ int main(int argc, char** argv)
std::istringstream s(line);
std::string kind;
s >> kind;
if (kind == "texture" || kind == "surface")
{
// texture/surface object handles in the arguments would point at nothing here: the kernel faults
fprintf(stderr, "%s: the dump uses texture or surface objects, which dump_runner cannot recreate\n",
dump.c_str());
return 3;
}
if (kind == "launch")
s >> grid[0] >> grid[1] >> grid[2] >> block[0] >> block[1] >> block[2] >> shared;
else if (kind == "alloc")
kernels/tools/model_enc3.py
+141
View File
@@ -0,0 +1,141 @@
#!/usr/bin/env python3
"""Numeric primitives and replay-dump access for swin_model.py (the rrlite Swin blocks of DLSS 4.5, preset M).
Everything here is NVIDIA's arithmetic as the native kernels reproduce it (kernels/m/swin_common.h,
swin_block.h): f16 rounding after every operation, e4m3 (OCP, satfinite, RNE) re-quantisation, one f16 rounding
per k32 MMA step, and the softmax's bit-level exponent trick.
"""
import struct
import numpy as np
F16, F32, F64 = np.float16, np.float32, np.float64
def f16(x):
"""round to f16; the result stays float64 (on the f16 grid), so products and sums of such values are exact
until the next f16() (numpy would otherwise do f16 arithmetic, e.g. in matmul)"""
return np.asarray(x, F64).astype(F16).astype(F64)
def _e4m3_table():
t = np.zeros(256, F64)
for c in range(256):
s, e, m = c >> 7, (c >> 3) & 15, c & 7
if e == 15 and m == 7:
v = np.nan
elif e == 0:
v = m * 2.0 ** -9
else:
v = (1 + m / 8) * 2.0 ** (e - 7)
t[c] = -v if s else v
return t
E4M3 = _e4m3_table()
_POS = E4M3[:127] # codes 0..126: the finite non-negative values, ascending
def q8(x):
"""e4m3 RNE satfinite round trip of f16 values (NaN stays NaN), as f16"""
h = f16(x).astype(F64)
a = np.abs(h)
i = np.clip(np.searchsorted(_POS, a), 1, 126)
lo, hi = _POS[i - 1], _POS[i]
# nearest; ties to the even code
pick_hi = (hi - a < a - lo) | ((hi - a == a - lo) & ((i & 1) == 0))
v = np.where(pick_hi, hi, lo)
v = np.where(a >= _POS[126], _POS[126], v)
v = np.where(a <= 0, 0.0, v)
v = np.where(np.isnan(h), np.nan, np.copysign(v, h))
return f16(v)
# softmax constants (swin_block.h stage 2)
C_SCALE = f16(0.0972222164273262 / 5.656854152679443)
CLAMP = F16(0.55615234375)
P1 = f16(0.92730712890625)
P0 = F16(1.375)
def exp_trick(poly):
"""per row, columns (2j, 2j+1): the 32-bit word even | odd << 16 of the f16 bits, << 5, + 0x7FF88000; the even
column takes its low half, the odd column its high half (as f16 bits)"""
b = poly.astype(F16).view(np.uint16).astype(np.uint64)
out = np.empty_like(b)
even, odd = b[:, 0::2], b[:, 1::2]
word = ((even | (odd << np.uint64(16))) << np.uint64(5)) + np.uint64(0x7FF88000)
word &= np.uint64(0xFFFFFFFF)
out[:, 0::2] = word & np.uint64(0xFFFF)
out[:, 1::2] = word >> np.uint64(16)
return out.astype(np.uint16).view(F16)
def row_sum16(w):
"""f16 butterfly over the 16 columns: partner c ^ 8, then ^ 2, ^ 4, ^ 1 (every column ends with the same sum)"""
s = w.astype(F16)
c = np.arange(16)
for m in (8, 2, 4, 1):
s = f16(s.astype(F64) + s[:, c ^ m].astype(F64))
return s
def gelu_poly(z):
z = f16(z)
cz = np.clip(z, F16(-2.0), F16(2.0))
inner = f16(f16(0.41216981).astype(F64) - f16(f16(0.08108133).astype(F64) * np.abs(cz).astype(F64)).astype(F64))
return f16(z.astype(F64) * f16(0.5 + f16(cz.astype(F64) * inner.astype(F64)).astype(F64)).astype(F64))
def mma_chain(A, W, C0=None):
"""chained k32 MMA steps with f16 accumulation: D = f16(A[:, 32k:32k+32] W[32k:32k+32] + D), exact sums"""
M, K = A.shape
D = np.zeros((M, W.shape[1]), F64) if C0 is None else f16(C0)
A = np.asarray(A, F64)
W = np.asarray(W, F64)
for k in range(0, K, 32):
D = f16(A[:, k:k + 32] @ W[k:k + 32] + D.astype(F64))
return D
class Params:
"""one rrlite Swin-block launch from a replay dump (CommonParams: w 0, sx sy 8, tw th 16, in 24, p32 32,
out 40, p48 48)"""
def __init__(self, d):
man = open(d + '/manifest.txt').read().split('\n')
self.kernel = man[0].split()[1]
la = man[1].split()
self.grid = tuple(map(int, la[1:4]))
self.args = open(d + '/args.bin', 'rb').read()
self.allocs = {}
for line in man:
p = line.split()
if p and p[0] == 'alloc':
self.allocs[int(p[1])] = (int(p[2], 16), np.fromfile(f'{d}/alloc-{p[1]}.bin', np.uint8))
self.w_ptr, = struct.unpack_from('<Q', self.args, 0)
self.sx, self.sy, self.W, self.H = struct.unpack_from('<iiii', self.args, 8)
self.inp, self.p32, self.out, self.p48 = struct.unpack_from('<QQQQ', self.args, 24)
wbase, wbuf = self.alloc_of(self.w_ptr)
self.wbuf = wbuf[self.w_ptr - wbase:]
self.abase, self.arena = self.alloc_of(self.inp)
def alloc_of(self, a):
for base, buf in self.allocs.values():
if base <= a < base + buf.size:
return base, buf
raise ValueError(hex(a))
def f16vec(self, off, n):
return self.wbuf[off:off + 2 * n].view(F16).copy()
def token_xy(self, bx, by):
T = np.arange(64)
tx = (T & 3) + 4 * ((T >> 4) & 1)
ty = ((T >> 2) & 3) + 4 * (T >> 5)
return 8 * bx - self.sx + tx, 8 * by - self.sy + ty
@staticmethod
def mirror(v, n):
v = np.abs(v)
return np.minimum(v, 2 * n - 2 - v)
kernels/tools/pwin_check.py
+127
View File
@@ -0,0 +1,127 @@
#!/usr/bin/env python3
"""Compares native preset-K layer outputs with the numpy model (pwin_model.py) on sampled windows of a replay dump.
usage: pwin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--windows N]
DUMP_DIR replay dump of one dltss_pwin_* launch (inputs)
OUT_DIR dump_runner output of a native kernel on that dump (alloc-N.bin after the launch)
For every sampled window (the four corners, border windows and random interior ones) the model output and the
kernel's stored values are compared per f16 element: exact matches, and the error in f16 ulps at the model
value's magnitude (values below 2^-10 count with the ulp of 2^-10, so near-zero sign flips do not dominate).
The kernels round as NVIDIA's f16 wmma does except in the PWIN_F32ACC layers, and gfx12 WMMA groups its f32
sums differently from gfx11, so neither build is expected to match the model exactly.
"""
import os
import random
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pwin_model as pm # noqa: E402
F16 = np.float16
# layer -> (kind, H, C, COUT/CL, extra)
LAYERS = {
'enc0': ('enc0', 2, 32, 64, None), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False),
'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 16, None),
'dec4': ('dec', 4, 128, 160, None), 'dec3': ('dec', 4, 96, 128, None), 'dec2': ('dec', 2, 64, 96, None),
'dec1': ('dec', 2, 64, 64, None), 'dec0': ('dec0', 2, 32, 64, None),
}
class Outputs:
"""the dump's pointer map over the allocations a native run saved"""
def __init__(self, dump, out_dir):
self.d = dump
self.allocs = {i: (base, np.fromfile(f'{out_dir}/alloc-{i}.bin', np.uint8)) for i, (base, _) in dump.allocs.items()}
def rows(self, off_param, index, width):
a = self.d.u64(off_param)
for base, buf in self.allocs.values():
if base <= a < base + buf.size:
o = a - base + 2 * index * width
return buf[o:o + 2 * width].view(F16)
raise ValueError(hex(a))
def ulps(model, got):
m = model.astype(np.float64)
g = got.astype(np.float64)
scale = np.maximum(np.abs(m), 2.0 ** -10)
ulp = 2.0 ** (np.floor(np.log2(scale)) - 10)
return np.abs(m - g) / ulp
def windows(grid, n, seed=1):
gx, gy = grid[0], grid[1]
w = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
r = random.Random(seed)
while len(w) < min(n, gx * gy):
w.add((r.randrange(gx), r.randrange(gy)))
return sorted(w)
def model_window(P, layer, bx, by):
kind, H, C, X, extra = LAYERS[layer]
if kind == 'enc0':
return pm.pwin_enc0(P, bx, by)
if kind == 'enc':
return pm.pwin_encoder(P, bx, by, H, C, X, posattn=extra)
if kind == 'plain':
return pm.swin_core(P, pm.Layout(H, C, C), pm.load_window(P, bx, by, C)), None
if kind == 'dec':
return pm.pwin_decoder(P, bx, by, H, C, X), None
return pm.pwin_dec0(P, bx, by)
def stored(P, O, layer, bx, by, y, extra_out):
"""(model values, kernel values) of everything the kernel stores for window (bx, by)"""
kind, H, C, X, _ = LAYERS[layer]
W, Hh = P.i32x2(0)
sx, sy = P.i32x2(56)
ms, gs = [], []
full_off = 32 if kind in ('enc0', 'enc') else 24
for t in range(64):
Y0, X0 = 8 * by - sy + (t >> 3), 8 * bx - sx + (t & 7)
if not (0 <= Y0 < Hh and 0 <= X0 < W):
continue
if kind == 'dec0':
ms.append(extra_out[t]); gs.append(O.rows(24, Y0 * W + X0, 40))
else:
ms.append(y[t]); gs.append(O.rows(full_off, Y0 * W + X0, C))
if kind in ('enc0', 'enc'):
cout = extra_out.shape[1]
for r in range(16):
Ym, Xm = int((8 * by - sy) / 2) + (r >> 2), int((8 * bx - sx) / 2) + (r & 3)
if 0 <= Ym < Hh // 2 and 0 <= Xm < W // 2:
ms.append(extra_out[r]); gs.append(O.rows(24, Ym * (W // 2) + Xm, cout))
return np.concatenate(ms), np.concatenate(gs)
def main():
args = [a for a in sys.argv[1:] if not a.startswith('--')]
n = 12
if '--windows' in sys.argv:
n = int(sys.argv[sys.argv.index('--windows') + 1])
args.remove(str(n))
dump_dir, outs = args[0], args[1:]
P = pm.Dump(dump_dir)
layer = P.kernel.replace('dltss_pwin_', '').replace('_layer', '')
wins = windows(P.grid, n)
models = {w: model_window(P, layer, *w) for w in wins}
for out in outs:
O = Outputs(P, out)
m, g = zip(*(stored(P, O, layer, bx, by, *models[(bx, by)]) for bx, by in wins))
m, g = np.concatenate(m), np.concatenate(g)
u = ulps(m, g)
exact = (m.view(np.uint16) == g.view(np.uint16)).mean()
md, gd = m.astype(np.float64), g.astype(np.float64)
psnr = 10 * np.log10(np.abs(md).max() ** 2 / max(((md - gd) ** 2).mean(), 1e-30))
print(f'{layer:5s} {os.path.basename(os.path.normpath(out)):10s} {len(wins)} windows {m.size} values: exact {100 * exact:5.1f}% '
f'<=1ulp {100 * (u <= 1).mean():5.1f}% <=4ulp {100 * (u <= 4).mean():6.2f}% max {u.max():7.1f} ulp '
f'mean {u.mean():.3f} ulp psnr {psnr:.1f} dB finite {np.isfinite(g.astype(np.float32)).mean() * 100:.0f}%')
if __name__ == '__main__':
main()
kernels/tools/swin_check.py
+105
View File
@@ -0,0 +1,105 @@
#!/usr/bin/env python3
"""Compares native preset-M Swin-block outputs with the numpy model (swin_model.py) on sampled blocks of a dump.
usage: swin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--blocks N]
DUMP_DIR replay dump of one rrlite_{enc1,enc2,enc3_tube,dec1,dec2}_4x4 launch
OUT_DIR dump_runner output of a native kernel on that dump
The stored e4m3 outputs (block output, and the patch merge of encoders that have one) are compared code by code:
the share of identical codes, of codes one e4m3 step apart, and the PSNR of the decoded values. The model is
NVIDIA's arithmetic (e4m3 activations, f16 accumulation per k32 step); the production kernels keep activations
f16 and accumulate in f32 (SWIN_NO_Q8, SWIN_F32ACC), so they are not expected to match it exactly.
"""
import os
import random
import sys
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import model_enc3 as me # noqa: E402
import swin_model as sm # noqa: E402
# kernel -> (C, heads, NPM, CIN)
LAYERS = {'enc1': (64, 2, 96, 0), 'enc2': (96, 4, 128, 0), 'enc3_tube': (128, 4, 0, 0), 'dec2': (96, 4, 0, 128),
'dec1': (64, 2, 0, 96)}
def blocks(grid, n, seed=1):
gx, gy = grid[0], grid[1]
b = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
r = random.Random(seed)
while len(b) < min(n, gx * gy):
b.add((r.randrange(gx), r.randrange(gy)))
return sorted(b)
def model_block(P, layer, bx, by):
C, heads, NPM, CIN = LAYERS[layer]
pre = 4 * C * CIN + 8 * C if CIN else 0
x0 = sm.decoder_input(P, bx, by, C, CIN, heads, P.p32) if CIN else sm.load_tokens(P, bx, by, C // 32)
x2, out, L = sm.swin_block(P, x0, C, heads, pre)
merged = None
if NPM:
A = sm.patch_merge_input(out, C)
PM0 = L['end']
cols = []
for g in range(NPM // 32):
Wg = me.E4M3[P.wbuf[sm.woff_table(PM0 + g * 128 * C, 512, 64 * C, 4 * C, 32)]]
bias = P.f16vec(PM0 + 4 * C * NPM + 64 * g, 32)
cols.append(me.mma_chain(A, Wg, C0=np.broadcast_to(bias, (16, 32)).copy()))
merged = me.q8(sm.pair_to_nat(np.concatenate(cols, axis=1)))
return out, merged
def codes_at(buf, base, ptr, W, H, X, Y, C):
"""e4m3 codes of tokens (X, Y) (in bounds) from planes of 32 channels at ptr"""
c = np.arange(C)
off = ptr - base + (((c[None, :] >> 5) * H + Y[:, None]) * W + X[:, None]) * 32 + (c[None, :] & 31)
return buf[off]
def main():
args = [a for a in sys.argv[1:] if not a.startswith('--')]
n = 8
if '--blocks' in sys.argv:
n = int(sys.argv[sys.argv.index('--blocks') + 1])
args.remove(str(n))
dump, outs = args[0], args[1:]
P = me.Params(dump)
layer = P.kernel.replace('rrlite_', '').replace('_4x4', '')
C, heads, NPM, CIN = LAYERS[layer]
bl = blocks(P.grid, n)
models = {b: model_block(P, layer, *b) for b in bl}
for out_dir in outs:
got, want = [], []
for (bx, by) in bl:
base, _ = P.alloc_of(P.out)
aid = next(i for i, (b0, buf) in P.allocs.items() if b0 == base)
buf = np.fromfile(f'{out_dir}/alloc-{aid}.bin', np.uint8)
X, Y = P.token_xy(bx, by)
ok = (X >= 0) & (X < P.W) & (Y >= 0) & (Y < P.H)
out, merged = models[(bx, by)]
got.append(codes_at(buf, base, P.out, P.W, P.H, X[ok], Y[ok], C).ravel())
want.append(out[ok].ravel())
if NPM:
W2, H2 = P.W // 2, P.H // 2
r = np.arange(16)
MX, MY = (8 * bx - P.sx) // 2 + (r & 3), (8 * by - P.sy) // 2 + (r >> 2)
okm = (MX >= 0) & (MX < W2) & (MY >= 0) & (MY < H2)
mbase, _ = P.alloc_of(P.p48)
mid = next(i for i, (b0, _) in P.allocs.items() if b0 == mbase)
mbuf = buf if mid == aid else np.fromfile(f'{out_dir}/alloc-{mid}.bin', np.uint8)
got.append(codes_at(mbuf, mbase, P.p48, W2, H2, MX[okm], MY[okm], NPM).ravel())
want.append(merged[okm].ravel())
g = me.E4M3[np.concatenate(got)]
w = np.concatenate(want).astype(np.float64)
rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v) # noqa: E731
steps = np.abs(rank(g) - rank(w))
psnr = 10 * np.log10(np.nanmax(np.abs(w)) ** 2 / max(np.nanmean((g - w) ** 2), 1e-30))
print(f'{layer:9s} {os.path.basename(os.path.normpath(out_dir)):10s} {len(bl)} blocks {g.size} codes: '
f'identical {100 * (g == w).mean():5.1f}% <=1 step {100 * (steps <= 1).mean():6.2f}% '
f'max {int(steps.max())} steps psnr {psnr:.1f} dB')
if __name__ == '__main__':
main()
kernels/tools/tail_check.hip
+141
View File
@@ -0,0 +1,141 @@
// Standalone check of the texture-kernel GEMM tails (kernels/tex/enc0_tail.hip, dec0_head.hip) outside the
// ZLUDA-compiled kernels they are linked into: random e4m3 weights and activations, several tile positions
// (including partial tiles at the image border), all outputs written to one file. Builds of different WMMA
// layouts (gfx11, the gfx12-layout shim, gfx12, gfx12 native FP8) are then compared byte for byte (same
// arithmetic) or by value (gfx12 rounding). The tails touch no images, so this also runs in rocjitsu.
//
// build: clang++ -x hip --offload-arch=ARCH -O3 [-DD4R_WMMA_LAYOUT=12] [-DD4R_TEX_FP8] -I kernels/tex
// kernels/tools/tail_check.hip -o tail_check usage: tail_check OUT.bin
#include <hip/hip_runtime.h>
#include <cstdio>
#include <cstdlib>
#include <vector>
// the surface-store helpers of tex_common.h reference ZLUDA's helper library, which the tails do not use
extern "C" __device__ void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t, int __attribute__((ext_vector_type(2))),
short __attribute__((ext_vector_type(2))))
{
}
// both tails define the same hook globals
#define D4R_KERNEL_NAME tail_enc0
#define d4r_prep_blocks d4r_prep_blocks_enc0
#define d4r_prep_key_at d4r_prep_key_at_enc0
#include "enc0_tail.hip"
#undef d4r_prep_blocks
#undef d4r_prep_key_at
#define d4r_prep_blocks d4r_prep_blocks_dec0
#define d4r_prep_key_at d4r_prep_key_at_dec0
#include "dec0_head.hip"
extern "C" __global__ void __launch_bounds__(32) run_enc0(const uint32_t* lds_in, const uint8_t* w, uint32_t* out,
uint32_t dimx, uint32_t dimy, int ox, int oy)
{
__shared__ uint32_t s[512 + 4];
for (int i = threadIdx.x; i < 512; i += 32)
s[i] = lds_in[i];
__syncthreads();
d4r_enc0_tail((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)out, dimx,
dimy, ox, oy);
}
extern "C" __global__ void __launch_bounds__(32) run_dec0(const uint8_t* w, const uint32_t* in, uint32_t* staged,
uint32_t dimx, uint32_t dimy, int ox, int oy)
{
__shared__ uint32_t s[512 + 4];
for (int i = threadIdx.x; i < 512; i += 32)
s[i] = 0xdeadbeefu;
__syncthreads();
d4r_dec0_head((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)in, dimx,
dimy, ox, oy);
__syncthreads();
for (int i = threadIdx.x; i < 512; i += 32)
staged[i] = s[i];
}
#define CHECK(x) \
do \
{ \
hipError_t e_ = (x); \
if (e_ != hipSuccess) \
{ \
fprintf(stderr, "%s: %s\n", #x, hipGetErrorString(e_)); \
return 1; \
} \
} while (0)
static uint8_t e4m3(uint32_t r)
{
// Keep the matrix values in a useful range so saturation cannot hide a wrong WMMA result.
const uint8_t magnitude = (r & 31) == 0 ? 0 : (uint8_t)(0x18 + (r & 31));
return (uint8_t)(magnitude | ((r >> 8) & 0x80));
}
static uint16_t f16_bits(float v)
{
return __builtin_bit_cast(uint16_t, (_Float16)v);
}
int main(int argc, char** argv)
{
if (argc < 2)
{
fprintf(stderr, "usage: %s OUT.bin\n", argv[0]);
return 2;
}
srand(12345);
auto rnd = []() { return (uint32_t)rand() * 2654435761u + (uint32_t)rand(); };
// weights: e4m3 fragments, f16 biases (enc0 at 12288, dec0 at 8192)
std::vector<uint8_t> w(16384 + 512);
for (auto& b : w)
b = e4m3(rnd());
for (int i = 0; i < 64; ++i)
*(uint16_t*)&w[12288 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
for (int i = 0; i < 128; ++i)
*(uint16_t*)&w[8192 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
const uint32_t dimx = 20, dimy = 14, W2 = dimx / 2, H2 = dimy / 2, rowstride = W2 * 8, plane = H2 * rowstride;
std::vector<uint32_t> lds(512), in(2 * plane);
for (auto& v : lds)
v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
for (auto& v : in)
v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
uint8_t *dw, *dummy;
uint32_t *dlds, *din, *dout, *dstaged;
CHECK(hipMalloc(&dw, w.size()));
CHECK(hipMalloc(&dlds, lds.size() * 4));
CHECK(hipMalloc(&din, in.size() * 4));
CHECK(hipMalloc(&dout, 2 * plane * 4));
CHECK(hipMalloc(&dstaged, 512 * 4));
CHECK(hipMalloc(&dummy, 512));
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(dlds, lds.data(), lds.size() * 4, hipMemcpyHostToDevice));
CHECK(hipMemcpy(din, in.data(), in.size() * 4, hipMemcpyHostToDevice));
// the prep kernels read the weight pointer at parameter offset 0
Enc0Params pe{};
pe.w = (uint64_t)dw;
Dec0Params pd{};
pd.w = (uint64_t)dw;
hipLaunchKernelGGL(tail_enc0_prep, dim3(4), dim3(128), 0, 0, pe);
hipLaunchKernelGGL(rrlite_dec0_4x4_folded_prep, dim3(4), dim3(128), 0, 0, pd);
CHECK(hipDeviceSynchronize());
FILE* f = fopen(argv[1], "wb");
const int offsets[][2] = {{0, 0}, {8, 4}, {12, 6}, {16, 8}, {-4, -2}, {18, 12}};
std::vector<uint32_t> host(2 * plane);
for (const auto& o : offsets)
{
CHECK(hipMemset(dout, 0x5a, 2 * plane * 4));
hipLaunchKernelGGL(run_enc0, dim3(1), dim3(32), 0, 0, dlds, dw, dout, dimx, dimy, o[0], o[1]);
CHECK(hipDeviceSynchronize());
CHECK(hipMemcpy(host.data(), dout, 2 * plane * 4, hipMemcpyDeviceToHost));
fwrite(host.data(), 4, 2 * plane, f);
hipLaunchKernelGGL(run_dec0, dim3(1), dim3(32), 0, 0, dw, din, dstaged, dimx, dimy, o[0], o[1]);
CHECK(hipDeviceSynchronize());
CHECK(hipMemcpy(host.data(), dstaged, 512 * 4, hipMemcpyDeviceToHost));
fwrite(host.data(), 4, 512, f);
}
fclose(f);
printf("wrote %s\n", argv[1]);
return 0;
}
packaging/D4R_README.txt
+11 -6
View File
@@ -1,10 +1,10 @@
d4r @VERSION@: NVIDIA DLSS on AMD Radeon RDNA3 under Linux and Proton
d4r @VERSION@: NVIDIA DLSS on AMD Radeon RDNA3 and RDNA4 under Linux and Proton
=====================================================================
d4r runs NVIDIA's own DLSS Super Resolution (DLSS 4 by default) on AMD Radeon RX 7000 GPUs in
d4r runs NVIDIA's own DLSS Super Resolution (DLSS 4 by default) on AMD Radeon RDNA3 and RDNA4 GPUs in
DirectX 12 games under Proton. OptiScaler (included) catches the game's DLSS calls and hands them to
d4r. d4r then runs NVIDIA's DLSS on the Radeon through ZLUDA, with the heaviest parts replaced by
kernels written for RDNA3.
kernels written for RDNA3 and RDNA4.
This is an early release, tested on one GPU (Radeon RX 7700 XT) in about a dozen DirectX 12 games.
SUPPORTED_GAMES.md in the source repository lists them, with per-game setup notes. Expect problems
@@ -13,9 +13,11 @@ in other games.
What you need
-------------
- An AMD Radeon RX 7000 series GPU (RDNA3). The fast kernels are built for gfx1101 (RX 7700 XT,
RX 7800 XT, RX 7700). Other RDNA3 cards run DLSS without them: DLSS 3 (E) at full speed,
DLSS 4.5 (M) at about half speed, and DLSS 4 (K) without its transformer (worse image).
- An AMD RDNA3 or RDNA4 GPU. This build includes native DLSS 4 and 4.5 network kernels for
gfx1100–gfx1103 and gfx1200–gfx1201 unless the builder selected fewer targets. Only the RX 7700 XT
(gfx1101) has been tested on a real GPU. gfx1201 has emulator testing; gfx1200 has compile testing
only. Neither establishes RDNA4 runtime support, performance or driver stability. Texture-kernel
optimisations are included only for targets supplied by the builder.
- Linux with the amdgpu kernel driver (/dev/kfd). ROCm itself is not needed: the zip includes the
ROCm 7.2.4 runtime in d4r/rocm. RocmDir in d4r/d4r.ini selects another ROCm installation instead.
- GE-Proton 11 (tested: GE-Proton11-3), selected for the game in Steam.
@@ -58,6 +60,9 @@ Settings
--------
- d4r/d4r.ini: d4r's settings for this game. The main one is the DLSS model:
K (DLSS 4, default), E (DLSS 3 CNN, a little faster) or M (DLSS 4.5, slowest).
- On RDNA4, [Kernels] NativeFp8 is on by default. It selects native FP8 WMMA and the matching
gfx12-fp8 kernel folder; set it to false to use f16 widening. The FP8 path has not been tested
on a real RDNA4 GPU.
- OptiScaler.ini: OptiScaler's settings, such as its fps overlay and the render resolution of each
quality mode. Press Insert in game for OptiScaler's menu.
packaging/VERSION
+1 -1
View File
@@ -1 +1 @@
0.1.1
0.1.2
packaging/d4r.ini
+7 -3
View File
@@ -21,10 +21,14 @@ ShowWatermark = false
FrameAge = 0
[Kernels]
; Hand-written RDNA3 kernels for the DLSS 4 and 4.5 networks (d4r\kernels). Each one is used only
; when your DLSS library's code for that kernel matches the code it was written for; otherwise that
; kernel is compiled from NVIDIA's code instead, which is slower but still correct.
; Hand-written RDNA3 and RDNA4 kernels for the DLSS 4 and 4.5 networks (d4r\kernels). Each one is used
; only when your DLSS library's code for that kernel matches the code it was written for; otherwise
; that kernel is compiled from NVIDIA's code instead, which is slower but still correct.
NativeKernels = on
; RDNA4 (RX 9000 series) only: run DLSS's FP8 math on the GPU's native FP8 matrix units, with NVIDIA's
; own FP8 rounding in the DLSS 4.5 (M) layers. off runs it as 16-bit math, as on RDNA3. Ignored on
; other GPUs.
NativeFp8 = on
[Interop]
; These keep DLSS's inputs and output in video memory and move the waits between the game and DLSS
patches/zluda/0006-k8-wmma-prep-slots-wave64.patch
+444
View File
@@ -0,0 +1,444 @@
diff --git a/llvm_zluda/src/compile.rs b/llvm_zluda/src/compile.rs
index 5d8c0a5..b2870e5 100644
--- a/llvm_zluda/src/compile.rs
+++ b/llvm_zluda/src/compile.rs
@@ -70,8 +70,10 @@ fn create_oclc_constants(ctx: &Context, gcn_arch: &str) -> Result<Module, String
let module = Module::new(ctx, c"oclc_constants");
// used by ockl
- add_constant(ctx, &module, c"__oclc_wavefrontsize64", 0);
- add_constant(ctx, &module, c"__oclc_wavefrontsize_log2", 5);
+ // d4r: D4R_ZLUDA_WAVE64=1 (see ptx::pass::wave64_requested)
+ let wave64 = std::env::var("D4R_ZLUDA_WAVE64").map_or(false, |v| v.trim() == "1");
+ add_constant(ctx, &module, c"__oclc_wavefrontsize64", if wave64 { 1 } else { 0 });
+ add_constant(ctx, &module, c"__oclc_wavefrontsize_log2", if wave64 { 6 } else { 5 });
add_constant(ctx, &module, c"__oclc_ABI_version", 500);
add_constant(
ctx,
@@ -174,7 +176,13 @@ pub fn compile(
let triple = c"amdgcn-amd-amdhsa";
let cpu = std::ffi::CString::new(gcn_arch).map_err(|_| ("invalid gcn_arch").to_string())?;
- let features = c"-wavefrontsize64,+cumode";
+ // functions without their own target features (the linked helper libraries) take these, so under
+ // D4R_ZLUDA_WAVE64 they must be optimised as wave64 too (lane-id and ballot folding)
+ let features = if std::env::var("D4R_ZLUDA_WAVE64").map_or(false, |v| v.trim() == "1") {
+ c"-wavefrontsize32,+wavefrontsize64,+cumode"
+ } else {
+ c"-wavefrontsize64,+cumode"
+ };
let mut target = unsafe { std::mem::zeroed() };
let mut err = std::ptr::null_mut();
diff --git a/llvm_zluda/src/lib.cpp b/llvm_zluda/src/lib.cpp
index 88848b1..0d1f9cf 100644
--- a/llvm_zluda/src/lib.cpp
+++ b/llvm_zluda/src/lib.cpp
@@ -869,10 +869,111 @@ unsigned LLVMZludaSplitWmmaM16n16k16(LLVMModuleRef ModuleRef)
return static_cast<unsigned>(Calls.size());
}
+// d4r: an m16n8k8 f16 MMA is an m16n8k16 one whose second K half is zero (its A fragment is the k16
+// A's registers 0-1, its B the k16 B's register 0), and adding zero products is exact. Each k8 call
+// becomes a k16 call so the WMMA rewrite below lowers it too. Zero halves are loads of one zeroed
+// register, so the operands stay "loads of plain registers" and two k8 MMAs that share A still fuse
+// into one WMMA. Returns the number of calls converted.
+static unsigned widenMmaK8(Module *M)
+{
+ Function *K8 = M->getFunction("__zluda_ptx_impl_mma_sync_aligned_m16n8k8_row_col_f16_f16_f16_f16");
+ if (K8 == nullptr || K8->use_empty())
+ return 0;
+ LLVMContext &Context = M->getContext();
+ Type *I32 = Type::getInt32Ty(Context);
+ Type *V2 = FixedVectorType::get(I32, 2);
+ Type *V4 = FixedVectorType::get(I32, 4);
+ AttributeList Attributes = AttributeList::get(Context, AttributeList::FunctionIndex,
+ AttrBuilder(Context, K8->getAttributes().getFnAttrs()));
+ FunctionCallee K16 = M->getOrInsertFunction(kMmaF16, Attributes, V2, V4, V2, V2);
+ SmallVector<CallInst *, 64> Calls;
+ for (User *U : K8->users())
+ if (auto *Call = dyn_cast<CallInst>(U); Call != nullptr && Call->getCalledFunction() == K8)
+ Calls.push_back(Call);
+ DenseMap<Function *, AllocaInst *> Zeros;
+ for (CallInst *Call : Calls)
+ {
+ Function *F = Call->getFunction();
+ AllocaInst *&Zero = Zeros[F];
+ if (Zero == nullptr)
+ {
+ // an alloca like ZLUDA's register allocas (same address space), stored once at entry
+ AllocaInst *Like = nullptr;
+ for (Instruction &I : F->getEntryBlock())
+ if ((Like = dyn_cast<AllocaInst>(&I)) != nullptr)
+ break;
+ IRBuilder<> Entry(&*F->getEntryBlock().getFirstInsertionPt());
+ Zero = Entry.CreateAlloca(I32, Like != nullptr ? Like->getAddressSpace() : 5, nullptr, "d4r.k8.zero");
+ Entry.CreateStore(Entry.getInt32(0), Zero);
+ }
+ IRBuilder<> Builder(Call);
+ SmallVector<AllocaInst *, 2> ARegisters, BRegisters;
+ auto build = [&](ArrayRef<Value *> Elements) {
+ Value *Vector = UndefValue::get(FixedVectorType::get(I32, Elements.size()));
+ for (unsigned Index = 0; Index < Elements.size(); ++Index)
+ Vector = Builder.CreateInsertElement(Vector, Elements[Index], Builder.getInt8(Index));
+ return Vector;
+ };
+ auto zero = [&]() { return Builder.CreateLoad(I32, Zero); };
+ // the registers may be reloaded at the call only if no store reaches them after their original loads
+ auto unchangedSinceLoads = [&](Value *Vector) {
+ for (Value *V = Vector; auto *Insert = dyn_cast<InsertElementInst>(V); V = Insert->getOperand(0))
+ {
+ auto *Load = cast<LoadInst>(Insert->getOperand(1));
+ if (Load->getParent() != Call->getParent() ||
+ storesBetween(Load, Call, {cast<AllocaInst>(Load->getPointerOperand())}))
+ return false;
+ }
+ return true;
+ };
+ // B is one register: a <1 x i32> vector or a plain i32
+ Value *BArg = Call->getArgOperand(1);
+ auto bRegister = [&]() {
+ if (BArg->getType()->isVectorTy())
+ return matchLoadedVector(BArg, 1, BRegisters) && unchangedSinceLoads(BArg);
+ auto *Load = dyn_cast<LoadInst>(BArg);
+ auto *Register = Load != nullptr ? dyn_cast<AllocaInst>(Load->getPointerOperand()) : nullptr;
+ if (Register == nullptr || !Register->getAllocatedType()->isIntegerTy(32) || !isPlainRegister(Register) ||
+ Load->getParent() != Call->getParent() || storesBetween(Load, Call, {Register}))
+ return false;
+ BRegisters.assign(1, Register);
+ return true;
+ };
+ auto element = [&](Value *V, unsigned Index) -> Value * {
+ return V->getType()->isVectorTy() ? Builder.CreateExtractElement(V, Builder.getInt8(Index)) : V;
+ };
+ Value *A, *B;
+ if (matchLoadedVector(Call->getArgOperand(0), 2, ARegisters) && bRegister() &&
+ unchangedSinceLoads(Call->getArgOperand(0)))
+ {
+ // reload the same registers right at the call (nothing lies between these loads and it)
+ Value *A0 = Builder.CreateLoad(I32, ARegisters[0]), *A1 = Builder.CreateLoad(I32, ARegisters[1]);
+ Value *B0 = Builder.CreateLoad(I32, BRegisters[0]);
+ A = build({A0, A1, zero(), zero()});
+ B = build({B0, zero()});
+ }
+ else
+ {
+ Value *A0 = Builder.CreateExtractElement(Call->getArgOperand(0), Builder.getInt8(0));
+ Value *A1 = Builder.CreateExtractElement(Call->getArgOperand(0), Builder.getInt8(1));
+ Value *B0 = element(BArg, 0);
+ A = build({A0, A1, Builder.getInt32(0), Builder.getInt32(0)});
+ B = build({B0, Builder.getInt32(0)});
+ }
+ CallInst *Wide = Builder.CreateCall(K16, {A, B, Call->getArgOperand(2)});
+ Call->replaceAllUsesWith(Wide);
+ Call->eraseFromParent();
+ }
+ return static_cast<unsigned>(Calls.size());
+}
+
// Flags bit 0: also lower m16n8k32 e4m3 MMAs (D4R_ZLUDA_WMMA_FP8=1).
unsigned LLVMZludaRewriteMmaF16Wmma(LLVMModuleRef ModuleRef, unsigned Flags)
{
Module *M = unwrap(ModuleRef);
+ // D4R_ZLUDA_WMMA_K8=0 keeps m16n8k8 on the dot-product fallback
+ if (const char *K8Env = getenv("D4R_ZLUDA_WMMA_K8"); K8Env == nullptr || K8Env[0] != '0')
+ widenMmaK8(M);
Function *Mma = M->getFunction(kMmaF16);
Function *MmaE4m3 = (Flags & 1) != 0 ? M->getFunction(kMmaE4m3) : nullptr;
if (Mma != nullptr && Mma->use_empty())
diff --git a/ptx/lib/zluda_ptx_impl.cpp b/ptx/lib/zluda_ptx_impl.cpp
index cdefdf1..74cf36f 100644
--- a/ptx/lib/zluda_ptx_impl.cpp
+++ b/ptx/lib/zluda_ptx_impl.cpp
@@ -127,9 +127,32 @@ extern "C"
{
extern "C" __attribute__((constant)) uint32_t __oclc_ISA_version __device__;
+ // d4r: kernels may be compiled as wave64 (D4R_ZLUDA_WAVE64, offline texture-kernel builds, linked with
+ // the zluda_ptx_impl*_w64.bc helpers built with -DD4R_WAVE64), where each wave holds two CUDA warps:
+ // lanes 0-31 and 32-63. Warp-level helpers then act on the lane's own 32-lane half.
+#ifdef D4R_WAVE64
+ __device__ static inline uint32_t d4r_wave_lane()
+ {
+ return __builtin_amdgcn_mbcnt_hi(~0u, __builtin_amdgcn_mbcnt_lo(~0u, 0u));
+ }
+ __device__ static inline uint32_t d4r_warp_base()
+ {
+ return d4r_wave_lane() & 32u;
+ }
+#else
+ __device__ static inline uint32_t d4r_warp_base()
+ {
+ return 0u;
+ }
+#endif
+
uint32_t FUNC(activemask)()
{
+#ifdef D4R_WAVE64
+ return (uint32_t)(__builtin_amdgcn_read_exec() >> d4r_warp_base());
+#else
return __builtin_amdgcn_read_exec_lo();
+#endif
}
size_t __ockl_get_local_id(uint32_t) __device__;
@@ -158,7 +181,11 @@ extern "C"
uint32_t FUNC(sreg_laneid)()
{
+#ifdef D4R_WAVE64
+ return d4r_wave_lane() & 31u;
+#else
return __lane_id();
+#endif
}
uint32_t FUNC(sreg_lanemask_le)()
@@ -322,7 +349,7 @@ extern "C"
int32_t section_mask = (opts >> 8) & 0b11111; \
int32_t warp_end = opts & 0b11111; \
delta &= 0b11111; \
- int32_t self = (int32_t)__lane_id(); \
+ int32_t self = (int32_t)__zluda_ptx_impl_sreg_laneid(); \
int32_t subsection = section_mask & self; \
int32_t subsection_end = subsection | (~section_mask & warp_end); \
int32_t idx = calculate_index; \
@@ -331,7 +358,7 @@ extern "C"
{ \
idx = self; \
} \
- int32_t output = __builtin_amdgcn_ds_bpermute(idx << 2, (int32_t)input); \
+ int32_t output = __builtin_amdgcn_ds_bpermute((int32_t)(d4r_warp_base() | (uint32_t)idx) << 2, (int32_t)input); \
return {(uint32_t)output, uint32_t(!out_of_bounds)}; \
} \
\
@@ -850,7 +877,11 @@ extern "C"
__device__ static inline uint32_t ballot(bool value, bool negate)
{
__builtin_amdgcn_wave_barrier();
+#ifdef D4R_WAVE64
+ return (uint32_t)(__builtin_amdgcn_ballot_w64(negate ? !value : value) >> d4r_warp_base());
+#else
return __builtin_amdgcn_ballot_w32(negate ? !value : value);
+#endif
}
bool FUNC(vote_sync_any_pred)(bool value, uint32_t membermask __attribute__((unused)))
@@ -866,7 +897,7 @@ extern "C"
// IMPORTANT: exec mask must be a subset of membermask, the behavior is undefined otherwise
bool FUNC(vote_sync_all_pred)(bool value, uint32_t membermask __attribute__((unused)))
{
- return ballot(value, false) == __builtin_amdgcn_read_exec_lo();
+ return ballot(value, false) == __zluda_ptx_impl_activemask();
}
// also known as "none"
diff --git a/ptx/src/pass/llvm/emit.rs b/ptx/src/pass/llvm/emit.rs
index 8294bb4..6af2495 100644
--- a/ptx/src/pass/llvm/emit.rs
+++ b/ptx/src/pass/llvm/emit.rs
@@ -397,7 +397,12 @@ impl<'a, 'input> ModuleEmitContext<'a, 'input> {
fn emit_target_features(&mut self, fn_: LLVMValueRef) {
let value = format!(
- "+wavefrontsize32,-wavefrontsize64,+cumode{}",
+ "{},+cumode{}",
+ if crate::pass::wave64_requested() {
+ "-wavefrontsize32,+wavefrontsize64"
+ } else {
+ "+wavefrontsize32,-wavefrontsize64"
+ },
if cfg!(debug_assertions) {
",+precise-memory"
} else {
diff --git a/ptx/src/pass/mod.rs b/ptx/src/pass/mod.rs
index 63aa98c..e3baa72 100644
--- a/ptx/src/pass/mod.rs
+++ b/ptx/src/pass/mod.rs
@@ -37,6 +37,18 @@ mod test;
static ZLUDA_PTX_IMPL: &'static [u8] = include_bytes!("../../lib/zluda_ptx_impl.bc");
static ZLUDA_PTX_IMPL_CONSTRAINED: &'static [u8] =
include_bytes!("../../lib/zluda_ptx_impl_constrained.bc");
+static ZLUDA_PTX_IMPL_W64: &'static [u8] = include_bytes!("../../lib/zluda_ptx_impl_w64.bc");
+static ZLUDA_PTX_IMPL_CONSTRAINED_W64: &'static [u8] =
+ include_bytes!("../../lib/zluda_ptx_impl_constrained_w64.bc");
+
+/// d4r: D4R_ZLUDA_WAVE64=1 compiles kernels as wave64: each wave runs two CUDA warps (lanes 0-31 and
+/// 32-63; the warp-level helpers act on the lane's own half) and RDNA3 issues most FP32 VALU work for all
+/// 64 lanes at once. Meant for offline builds of texture kernels whose warp operations are limited to
+/// lane ids, shuffles and votes; not for MMA/matrix kernels (their helpers assume one warp per wave).
+pub fn wave64_requested() -> bool {
+ static VALUE: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
+ *VALUE.get_or_init(|| std::env::var("D4R_ZLUDA_WAVE64").map_or(false, |v| v.trim() == "1"))
+}
const ZLUDA_PTX_PREFIX: &'static str = "__zluda_ptx_impl_";
quick_error! {
@@ -194,10 +206,11 @@ pub struct Module {
impl Module {
pub fn linked_bitcode(&self) -> &'static [u8] {
- if self.constrained_fp {
- ZLUDA_PTX_IMPL_CONSTRAINED
- } else {
- ZLUDA_PTX_IMPL
+ match (self.constrained_fp, wave64_requested()) {
+ (true, false) => ZLUDA_PTX_IMPL_CONSTRAINED,
+ (false, false) => ZLUDA_PTX_IMPL,
+ (true, true) => ZLUDA_PTX_IMPL_CONSTRAINED_W64,
+ (false, true) => ZLUDA_PTX_IMPL_W64,
}
}
}
diff --git a/zluda/src/impl/function.rs b/zluda/src/impl/function.rs
index bacb93d..033d9a0 100644
--- a/zluda/src/impl/function.rs
+++ b/zluda/src/impl/function.rs
@@ -18,9 +18,14 @@ pub(crate) struct Function {
/// d4r native override: persistent kernels launch this many workgroups in x (0 = as launched)
pub(crate) grid_x: u32,
/// d4r native override: byte offset of the u64 parameter that keys the prep kernel's output
- /// (0 = run the prep before every launch) and the key of the last prep launched
- pub(crate) prep_key: u32,
+ /// (None = run the prep before every launch) and the key of the last prep launched
+ pub(crate) prep_key: Option<u32>,
pub(crate) prep_last: std::sync::atomic::AtomicU64,
+ /// d4r native override: with N > 0 prep slots the kernel keeps one prepared image per key for
+ /// up to N keys, so the prep runs only for keys not seen yet (and for every key beyond N, which
+ /// 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<Vec<u64>>,
}
impl<'a, E: zluda_common::CudaErrorType> zluda_common::FromCuda<'a, CUfunction, E>
@@ -140,7 +145,7 @@ pub(crate) fn launch_kernel(
translated_extra_storage.as_mut_ptr()
};
// the parameter block: kernel_params[0] (a single struct parameter) or the extra buffer
- let prep_key = if f.prep.is_some() && f.prep_key != 0 {
+ let prep_key = if let (Some(_), Some(key_offset)) = (f.prep, f.prep_key) {
let mut block: *const u8 = std::ptr::null();
if !extra.is_null() {
let mut i = 0usize;
@@ -160,12 +165,27 @@ pub(crate) fn launch_kernel(
if block.is_null() {
0
} else {
- unsafe { std::ptr::read_unaligned(block.add(f.prep_key as usize) as *const u64) }
+ unsafe { std::ptr::read_unaligned(block.add(key_offset as usize) as *const u64) }
}
} else {
0
};
- let run_prep = prep_key == 0 || f.prep_last.load(std::sync::atomic::Ordering::Relaxed) != prep_key;
+ let run_prep = if prep_key == 0 {
+ true
+ } else if f.prep_slots == 0 {
+ f.prep_last.load(std::sync::atomic::Ordering::Relaxed) != prep_key
+ } else {
+ // prepared once per key while slots last; keys beyond them are prepared on every launch
+ let mut seen = f.prep_seen.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
+ if seen.contains(&prep_key) {
+ false
+ } else {
+ if seen.len() < f.prep_slots as usize {
+ seen.push(prep_key);
+ }
+ true
+ }
+ };
if let (Some((prep, blocks)), true) = (f.prep, run_prep) {
unsafe {
hipModuleLaunchKernel(
diff --git a/zluda/src/impl/module.rs b/zluda/src/impl/module.rs
index 5db5ace..c5a39c9 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) = match native_override(entry.key()) {
+ let (func_handle, prep, block_z, grid_x, prep_key, prep_slots) = 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, 0)
+ (func_handle, None, 0, 0, None, 0)
}
};
let func = Box::new(Function {
@@ -54,6 +54,8 @@ impl ModuleMutable {
grid_x,
prep_key,
prep_last: std::sync::atomic::AtomicU64::new(0),
+ prep_slots,
+ prep_seen: std::sync::Mutex::new(Vec::new()),
});
&*entry.insert(func)
}
@@ -67,7 +69,9 @@ 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.
-fn native_override(name: &CStr) -> Option<(hipFunction_t, Option<(hipFunction_t, u32)>, u32, u32, u32)> {
+fn native_override(
+ name: &CStr,
+) -> Option<(hipFunction_t, Option<(hipFunction_t, u32)>, u32, u32, Option<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()?;
@@ -103,34 +107,38 @@ fn native_override(name: &CStr) -> Option<(hipFunction_t, Option<(hipFunction_t,
} else {
None
};
- let read_u32 = |symbol: &CStr| -> u32 {
+ let read_u32_opt = |symbol: &CStr| -> Option<u32> {
let mut ptr = unsafe { mem::zeroed() };
let mut size = 0usize;
let mut value = 0u32;
- if unsafe { hipModuleGetGlobal(&mut ptr, &mut size, native_module, symbol.as_ptr()) }.is_ok()
+ (unsafe { hipModuleGetGlobal(&mut ptr, &mut size, native_module, symbol.as_ptr()) }.is_ok()
&& size == 4
- && unsafe { hipMemcpyDtoH((&mut value as *mut u32).cast(), ptr, 4) }.is_ok()
- {
- value
- } else {
- 0
- }
+ && unsafe { hipMemcpyDtoH((&mut value as *mut u32).cast(), ptr, 4) }.is_ok())
+ .then_some(value)
};
+ let read_u32 = |symbol: &CStr| -> u32 { read_u32_opt(symbol).unwrap_or(0) };
let block_z = read_u32(c"d4r_block_z");
let grid_x = read_u32(c"d4r_grid_x");
// byte offset of a u64 in the parameter block that fully determines the prep kernel's output
- // (the weights pointer): the prep is skipped while it repeats (0 = prep before every launch)
- let prep_key = read_u32(c"d4r_prep_key_offset");
+ // (the weights pointer): the prep is skipped while it repeats. d4r_prep_key_at holds offset + 1
+ // (so a key at offset 0 is possible); the older d4r_prep_key_offset uses 0 for "no key".
+ let prep_key = match read_u32_opt(c"d4r_prep_key_at") {
+ Some(at) if at > 0 => Some(at - 1),
+ _ => Some(read_u32(c"d4r_prep_key_offset")).filter(|&offset| offset != 0),
+ };
+ // 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");
eprintln!(
- "[d4r-native] {:?} -> {} (prep: {:?}, block_z: {}, grid_x: {}, prep_key: {})",
+ "[d4r-native] {:?} -> {} (prep: {:?}, block_z: {}, grid_x: {}, prep_key: {:?}, prep_slots: {})",
name,
path.display(),
prep.map(|p| p.1),
block_z,
grid_x,
- prep_key
+ prep_key,
+ prep_slots
);
- Some((func, prep, block_z, grid_x, prep_key))
+ Some((func, prep, block_z, grid_x, prep_key, prep_slots))
}
impl ZludaObject for Module {
patches/zluda/0007-gfx12-wmma-native-fp8.patch
+551
View File
@@ -0,0 +1,551 @@
--- a/llvm_zluda/src/compile.rs
+++ b/llvm_zluda/src/compile.rs
@@ -113,9 +113,57 @@
*VALUE.get_or_init(|| std::env::var("D4R_ZLUDA_PARAM_PROVENANCE").map_or(true, |value| value.trim() != "0"))
}
-fn wmma_rewrite_enabled(gcn_arch: &str) -> bool {
- // The helpers assume RDNA3's wave32 WMMA register layout.
- wmma_rewrite_requested() && gcn_arch.starts_with("gfx11")
+/// D4R_ZLUDA_WMMA_FP8_NATIVE=1: on gfx12 (RDNA4), e4m3 MMAs use the native FP8 WMMA instead of widening
+/// to f16. Ignored on every other GPU (gfx11 has no FP8 WMMA).
+pub fn wmma_fp8_native_requested() -> bool {
+ static VALUE: OnceLock<bool> = OnceLock::new();
+ *VALUE.get_or_init(|| std::env::var("D4R_ZLUDA_WMMA_FP8_NATIVE").map_or(false, |value| value.trim() == "1"))
+}
+
+/// The WMMA register layout family of a GPU: 11 (RDNA3) or 12 (RDNA4); the MMA helpers exist for these.
+/// D4R_ZLUDA_WMMA_LAYOUT=12 on gfx11 is for tests only: it takes the gfx12 helper names, which then must
+/// come from a helper library built with D4R_WMMA12_SHIM (gfx12 data movement, gfx11 WMMAs).
+fn wmma_family(gcn_arch: &str) -> Option<u32> {
+ let family = if gcn_arch.starts_with("gfx11") {
+ 11
+ } else if gcn_arch.starts_with("gfx12") {
+ 12
+ } else {
+ return None;
+ };
+ if family == 11 && std::env::var("D4R_ZLUDA_WMMA_LAYOUT").map_or(false, |value| value.trim() == "12") {
+ return Some(12);
+ }
+ Some(family)
+}
+
+/// Whether e4m3 MMAs compiled for this GPU use native FP8 WMMA (the setting only acts on gfx12).
+pub fn wmma_fp8_native_effective(gcn_arch: &str) -> bool {
+ wmma_rewrite_flags(gcn_arch).map_or(false, |flags| flags & 4 != 0)
+}
+
+/// Flags for LLVMZludaRewriteMmaF16Wmma, or None to keep the dot-product fallback (logged once when the
+/// rewrite was requested but this GPU has no WMMA lowering).
+fn wmma_rewrite_flags(gcn_arch: &str) -> Option<u32> {
+ if !wmma_rewrite_requested() {
+ return None;
+ }
+ let Some(family) = wmma_family(gcn_arch) else {
+ static LOGGED: OnceLock<()> = OnceLock::new();
+ LOGGED.get_or_init(|| {
+ eprintln!("[ZLUDA] D4R_ZLUDA_WMMA: no WMMA lowering for {gcn_arch}; MMAs use the dot-product fallback")
+ });
+ return None;
+ };
+ let mut flags = if wmma_fp8_rewrite_requested() { 1 } else { 0 };
+ if family == 12 {
+ flags |= 2;
+ // native FP8 needs gfx12 hardware; the layout shim on gfx11 takes it only for tests (D4R_WMMA12_SHIM)
+ if flags & 1 != 0 && wmma_fp8_native_requested() {
+ flags |= 4;
+ }
+ }
+ Some(flags)
}
pub fn compile(
@@ -139,8 +187,7 @@
// wmma.mma m16n16k16 always becomes m16n8k16 MMAs (there is no helper for it).
unsafe { crate::ffi::LLVMZludaSplitWmmaM16n16k16(main.get()) };
- if wmma_rewrite_enabled(gcn_arch) {
- let flags = if wmma_fp8_rewrite_requested() { 1 } else { 0 };
+ if let Some(flags) = wmma_rewrite_flags(gcn_arch) {
unsafe { crate::ffi::LLVMZludaRewriteMmaF16Wmma(main.get(), flags) };
}
linked.link(main)?;
--- a/llvm_zluda/src/lib.cpp
+++ b/llvm_zluda/src/lib.cpp
@@ -968,6 +968,9 @@
}
// Flags bit 0: also lower m16n8k32 e4m3 MMAs (D4R_ZLUDA_WMMA_FP8=1).
+// Flags bit 1: gfx12 (RDNA4) WMMA layout: the d4r_wmma12_* helpers.
+// Flags bit 2 (with bit 1): e4m3 MMAs on native FP8 WMMA: the d4r_wmma12f8_* e4m3 helpers.
+// The helper sets share one contract (see zluda_ptx_impl.cpp), so only their names change here.
unsigned LLVMZludaRewriteMmaF16Wmma(LLVMModuleRef ModuleRef, unsigned Flags)
{
Module *M = unwrap(ModuleRef);
@@ -990,33 +993,39 @@
Type *V8 = FixedVectorType::get(I32, 8);
AttributeList Attributes = AttributeList::get(Context, AttributeList::FunctionIndex,
AttrBuilder(Context, Reference->getAttributes().getFnAttrs()));
+ // arch-specific helpers (operand gathers, WMMAs), e4m3 MMAs, e4m3 fragment widening; the conversions
+ // between NVIDIA fragments, the transposed layout and f32 shadows are shared
+ const bool Gfx12 = (Flags & 2) != 0, NativeFp8 = Gfx12 && (Flags & 4) != 0;
+ const std::string Arch = Gfx12 ? "__zluda_ptx_impl_d4r_wmma12_" : "__zluda_ptx_impl_d4r_wmma_";
+ const std::string E4m3 = NativeFp8 ? "__zluda_ptx_impl_d4r_wmma12f8_" : Arch;
+ const std::string Widen = NativeFp8 ? E4m3 : "__zluda_ptx_impl_d4r_wmma_";
Helpers H;
- H.AFromB = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_a_from_b", Attributes, V8, V2);
- H.BFromA = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_b_from_a", Attributes, V8, V4);
- H.Mma = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_mma_f16", Attributes, V2, V8, V8, V2);
+ H.AFromB = M->getOrInsertFunction((Arch + "t_a_from_b").c_str(), Attributes, V8, V2);
+ H.BFromA = M->getOrInsertFunction((Arch + "t_b_from_a").c_str(), Attributes, V8, V4);
+ H.Mma = M->getOrInsertFunction((Arch + "t_mma_f16").c_str(), Attributes, V2, V8, V8, V2);
H.FromNv = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_from_nv", Attributes, V2, V2);
H.ToNv = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_to_nv", Attributes, V2, V2);
- H.AFromBPair = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_a_from_b_pair", Attributes, V8, V2, V2);
- H.MmaPair = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_mma_f16_pair", Attributes, V4, V8, V8, V2, V2);
+ H.AFromBPair = M->getOrInsertFunction((Arch + "t_a_from_b_pair").c_str(), Attributes, V8, V2, V2);
+ H.MmaPair = M->getOrInsertFunction((Arch + "t_mma_f16_pair").c_str(), Attributes, V4, V8, V8, V2, V2);
{
Type *F = Type::getFloatTy(Context);
Type *V4F = FixedVectorType::get(F, 4);
Type *V8F = FixedVectorType::get(F, 8);
H.SFromT = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_s_from_t", Attributes, V4F, V2, F);
H.TFromS = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_from_s", Attributes, V2, V4F, F);
- H.MmaS = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_s_mma_f16", Attributes, V4F, V8, V8, V4F, F, F);
- H.MmaPairS = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_s_mma_f16_pair", Attributes, V8F, V8, V8, V4F, V4F, F, F);
- H.MmaE4m3S = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_s_mma_e4m3", Attributes, V4F, V8, V8, V8, V8, V4F);
+ H.MmaS = M->getOrInsertFunction((Arch + "s_mma_f16").c_str(), Attributes, V4F, V8, V8, V4F, F, F);
+ H.MmaPairS = M->getOrInsertFunction((Arch + "s_mma_f16_pair").c_str(), Attributes, V8F, V8, V8, V4F, V4F, F, F);
+ H.MmaE4m3S = M->getOrInsertFunction((E4m3 + "s_mma_e4m3").c_str(), Attributes, V4F, V8, V8, V8, V8, V4F);
H.MmaE4m3PairS =
- M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_s_mma_e4m3_pair", Attributes, V8F, V8, V8, V8, V8, V4F, V4F);
+ M->getOrInsertFunction((E4m3 + "s_mma_e4m3_pair").c_str(), Attributes, V8F, V8, V8, V8, V8, V4F, V4F);
}
if (MmaE4m3 != nullptr)
{
- H.E4m3AToF16 = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_e4m3_a_to_f16", Attributes, V8, V4);
- H.E4m3BToF16 = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_e4m3_b_to_f16", Attributes, V4, V2);
- H.MmaE4m3 = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_mma_e4m3", Attributes, V2, V8, V8, V8, V8, V2);
+ H.E4m3AToF16 = M->getOrInsertFunction((Widen + "e4m3_a_to_f16").c_str(), Attributes, V8, V4);
+ H.E4m3BToF16 = M->getOrInsertFunction((Widen + "e4m3_b_to_f16").c_str(), Attributes, V4, V2);
+ H.MmaE4m3 = M->getOrInsertFunction((E4m3 + "t_mma_e4m3").c_str(), Attributes, V2, V8, V8, V8, V8, V2);
H.MmaE4m3Pair =
- M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_mma_e4m3_pair", Attributes, V4, V8, V8, V8, V8, V2, V2);
+ M->getOrInsertFunction((E4m3 + "t_mma_e4m3_pair").c_str(), Attributes, V4, V8, V8, V8, V8, V2, V2);
}
unsigned Rewritten = 0;
for (Function &F : *M)
--- a/ptx/examples/d4r_emit.rs
+++ b/ptx/examples/d4r_emit.rs
@@ -1,9 +1,13 @@
use std::{env, fs, path::PathBuf};
use ptx_parser::parse_module_checked;
+// usage: d4r_emit PTX OUT_DIR [GPU_ARCH]: translates and compiles a PTX module offline (no GPU needed) for
+// GPU_ARCH (default gfx1101), with the same D4R_ZLUDA_* code-generation switches as the runtime, and writes
+// the IR stages, asm.s and module.hsaco to OUT_DIR
fn main() {
let input = env::args().nth(1).expect("PTX path");
let outdir = PathBuf::from(env::args().nth(2).unwrap_or_else(|| "/tmp/d4r-autoexposure-emitted".into()));
+ let arch = env::args().nth(3).unwrap_or_else(|| "gfx1101".into());
fs::create_dir_all(&outdir).expect("create output directory");
let bytes = fs::read(&input).expect("read PTX");
let text = std::str::from_utf8(&bytes).expect("PTX UTF-8").trim_end_matches('\0');
@@ -18,13 +22,14 @@
"opt.ll" => "opt.ll",
"asm" => "asm.s",
"o" => "object.o",
+ "elf" => "linked.elf", // the code object compile() returns (written again as module.hsaco)
other => panic!("unexpected stage: {other}"),
};
fs::write(outdir.join(filename), data).expect("write compiler stage");
};
let elf = llvm_zluda::compile(
&module.context,
- "gfx1101",
+ &arch,
module.llvm_ir,
linked_helpers,
module.attributes_ir,
--- a/ptx/lib/zluda_ptx_impl.cpp
+++ b/ptx/lib/zluda_ptx_impl.cpp
@@ -1696,6 +1696,349 @@
}
}
+// gfx12 (RDNA4) lowering of the same MMAs (LLVMZludaRewriteMmaF16Wmma with the gfx12 flag: helpers
+// d4r_wmma12_*, and d4r_wmma12f8_* for the e4m3 ones with native FP8 WMMA). The contract is the gfx11 one
+// above: NVIDIA fragments in, the transposed accumulator layout (lane l holds D[l % 16][2v + l / 16] as two
+// f16x2 words, or four floats for the f32 shadows), so the rewrite only switches helper names.
+//
+// A gfx12 WMMA operand holds 8 of the 16 K values per lane: half h = lane / 16 its own 8 (4 dwords, returned
+// in the low half of a wmma_uint8). With A' = B^T and B' = A^T as above, dword j of half h holds the k-pair of
+// NVIDIA thread t = 2h + (j >> 1): the low pair (k 2t, 2t + 1) for even j, the high pair (k 2t + 8, 2t + 9)
+// for odd j. Then every lane of a gather reads the same source register (4 bpermutes per operand; B' also
+// needs a DPP swap per dword for rows 8..15). For e4m3 fragments the same slots hold, after the widening
+// above, e4m3 k 8h .. 8h + 7 in order.
+// The result D' has VGPR i = row n = i + 8h (columns m = lane % 16): unpaired, rows 0..7 are the h = 0 lanes
+// and the h = 1 lanes are padding; paired (A' rows 8..15 = B1^T), D1 is in the h = 1 lanes. Converting from
+// and to the transposed layout takes 2 permlanex16 each way (4 for f32 shadows).
+//
+// D4R_WMMA12_SHIM (tests only): these helpers compiled for gfx11, every gfx12 WMMA done by the gfx11 one on
+// operands rebuilt in the gfx11 helpers' K order (and for e4m3 with the same scaled widening), so the results
+// are bit-identical to the gfx11 helpers and the gfx12 data movement can be checked on RDNA3.
+typedef _Float16 wmma_half8 __attribute__((ext_vector_type(8)));
+typedef uint32_t wmma_uint4 __attribute__((ext_vector_type(4)));
+typedef int wmma_int2 __attribute__((ext_vector_type(2)));
+#ifndef D4R_WMMA12_SHIM
+__device__ wmma_float8 __zluda_wmma12_f32_16x16x16_f16(wmma_half8 a, wmma_half8 b, wmma_float8 c) __asm("llvm.amdgcn.wmma.f32.16x16x16.f16.v8f32.v8f16");
+__device__ wmma_float8 __zluda_wmma12_f32_16x16x16_fp8_fp8(wmma_int2 a, wmma_int2 b, wmma_float8 c) __asm("llvm.amdgcn.wmma.f32.16x16x16.fp8.fp8.v8f32.v2i32");
+#endif
+
+__device__ static uint32_t wmma12_half()
+{
+ return wmma_lane() / 16;
+}
+
+__device__ static uint32_t wmma12_other_half(uint32_t value)
+{
+ return uint32_t(__builtin_amdgcn_permlanex16(int(value), int(value), 0x76543210, 0xfedcba98, false, false));
+}
+
+__device__ static float wmma12_other_half_f(float value)
+{
+ return std::bit_cast<float>(wmma12_other_half(std::bit_cast<uint32_t>(value)));
+}
+
+// dword j of an operand gathered from two candidate registers: rows r < 8 take `low` from lane 4r + t, rows
+// r >= 8 take `high` from lane 4(r - 8) + t (t = 2h + (j >> 1)). The source lane is in the first half-wave,
+// while t differs between the two destination half-waves, so the two candidates must be gathered independently.
+__device__ static uint32_t wmma12_gather2(uint32_t low, uint32_t high, uint32_t j)
+{
+ const uint32_t row = wmma_lane() % 16;
+ const uint32_t t = (j >> 1) + 2 * wmma12_half();
+ const uint32_t lo = wmma_from_lane(row * 4 + t, low);
+ const uint32_t hi = wmma_from_lane((row % 8) * 4 + t, high);
+ return row < 8 ? lo : hi;
+}
+
+__device__ static wmma_uint4 wmma12_low4(wmma_uint8 v)
+{
+ return {v[0], v[1], v[2], v[3]};
+}
+
+// f32 D' = A' B' + C' (f16 operands in the slot order above)
+__device__ static wmma_float8 wmma12_f16(wmma_uint8 a, wmma_uint8 b, wmma_float8 c)
+{
+#ifndef D4R_WMMA12_SHIM
+ return __zluda_wmma12_f32_16x16x16_f16(std::bit_cast<wmma_half8>(wmma12_low4(a)), std::bit_cast<wmma_half8>(wmma12_low4(b)), c);
+#else
+ // gfx11 operand dword p (K 2p, 2p + 1 in the gfx11 helpers' order) is gfx12 dword jp of half hp
+ const bool high = wmma12_half() != 0;
+ wmma_uint8 a11, b11;
+ for (uint32_t p = 0; p < 8; p++)
+ {
+ const uint32_t hp = (p & 3) >> 1, jp = 2 * (p & 1) + (p >> 2);
+ const uint32_t ao = wmma12_other_half(a[jp]), bo = wmma12_other_half(b[jp]);
+ a11[p] = (hp != 0) == high ? a[jp] : ao;
+ b11[p] = (hp != 0) == high ? b[jp] : bo;
+ }
+ // accumulators: gfx12 VGPR i = row i + 8h, gfx11 VGPR i = row 2i + h
+ wmma_float8 co, c11;
+ for (uint32_t i = 0; i < 8; i++)
+ co[i] = wmma12_other_half_f(c[i]);
+ for (uint32_t i = 0; i < 8; i++)
+ c11[i] = high ? (i < 4 ? co[2 * i + 1] : c[2 * i - 7]) : (i < 4 ? c[2 * i] : co[2 * i - 8]);
+ wmma_float8 d11 = __zluda_wmma_f32_16x16x16_f16(std::bit_cast<wmma_half16>(a11), std::bit_cast<wmma_half16>(b11), c11);
+ wmma_float8 dout, d;
+ for (uint32_t i = 0; i < 8; i++)
+ dout[i] = wmma12_other_half_f(d11[i]);
+ for (uint32_t i = 0; i < 8; i++)
+ d[i] = high ? ((i & 1) ? d11[(i >> 1) + 4] : dout[(i >> 1) + 4]) : ((i & 1) ? dout[i >> 1] : d11[i >> 1]);
+ return d;
+#endif
+}
+
+// f32 D' = A' B' + C' for operands whose dwords hold two e4m3 bytes in their low 16 bits (the d4r_wmma12f8
+// e4m3 "widening" below keeps the bytes): true e4m3 values, no scaling
+__device__ static wmma_float8 wmma12_e4m3(wmma_uint8 a, wmma_uint8 b, wmma_float8 c)
+{
+ const wmma_int2 a8 = {int(__builtin_amdgcn_perm(a[1], a[0], 0x05040100u)), int(__builtin_amdgcn_perm(a[3], a[2], 0x05040100u))};
+ const wmma_int2 b8 = {int(__builtin_amdgcn_perm(b[1], b[0], 0x05040100u)), int(__builtin_amdgcn_perm(b[3], b[2], 0x05040100u))};
+#ifndef D4R_WMMA12_SHIM
+ return __zluda_wmma12_f32_16x16x16_fp8_fp8(a8, b8, c);
+#else
+ // the gfx11 helpers' scaled widening (value * 2^-8 per operand, C * 2^-16): exact powers of two
+ auto widen = [](wmma_int2 q) {
+ auto lo = e4m3x4_to_scaled_f16x2x2(uint32_t(q.x)), hi = e4m3x4_to_scaled_f16x2x2(uint32_t(q.y));
+ return wmma_uint8{lo.x, lo.y, hi.x, hi.y, 0u, 0u, 0u, 0u};
+ };
+ return wmma12_f16(widen(a8), widen(b8), c * 0x1p-16f) * 0x1p16f;
+#endif
+}
+
+// transposed f16 accumulator(s) <-> D' / C'
+__device__ static wmma_float8 wmma12_c_from_t(uint2::Native_vec_ c_t)
+{
+ // lane h = 0 needs rows 0..7: its own rows 0, 2, 4, 6 and the other half's 1, 3, 5, 7 (h = 1 lanes: padding)
+ const uint32_t o0 = wmma12_other_half(c_t.x), o1 = wmma12_other_half(c_t.y);
+ f16x2 e01 = std::bit_cast<f16x2>(c_t.x), e23 = std::bit_cast<f16x2>(c_t.y);
+ f16x2 d01 = std::bit_cast<f16x2>(o0), d23 = std::bit_cast<f16x2>(o1);
+ return {float(e01.x), float(d01.x), float(e01.y), float(d01.y), float(e23.x), float(d23.x), float(e23.y), float(d23.y)};
+}
+
+__device__ static wmma_float8 wmma12_c_from_t_pair(uint2::Native_vec_ c0_t, uint2::Native_vec_ c1_t)
+{
+ // h = 0: C0 rows 0..7 (own even, other odd); h = 1: C1 rows 0..7 (other even, own odd)
+ const bool high = wmma12_half() != 0;
+ const uint2::Native_vec_ own = high ? c1_t : c0_t, send = high ? c0_t : c1_t;
+ const uint32_t r0 = wmma12_other_half(send.x), r1 = wmma12_other_half(send.y);
+ f16x2 o01 = std::bit_cast<f16x2>(own.x), o23 = std::bit_cast<f16x2>(own.y);
+ f16x2 p01 = std::bit_cast<f16x2>(r0), p23 = std::bit_cast<f16x2>(r1);
+ f16x2 e01 = high ? p01 : o01, e23 = high ? p23 : o23, d01 = high ? o01 : p01, d23 = high ? o23 : p23;
+ return {float(e01.x), float(d01.x), float(e01.y), float(d01.y), float(e23.x), float(d23.x), float(e23.y), float(d23.y)};
+}
+
+__device__ static void wmma12_rows_f16(wmma_float8 d, uint2::Native_vec_ &even, uint2::Native_vec_ &odd)
+{
+ f16x2 e01{static_cast<f16>(d[0]), static_cast<f16>(d[2])}, e23{static_cast<f16>(d[4]), static_cast<f16>(d[6])};
+ f16x2 o01{static_cast<f16>(d[1]), static_cast<f16>(d[3])}, o23{static_cast<f16>(d[5]), static_cast<f16>(d[7])};
+ even = {std::bit_cast<uint32_t>(e01), std::bit_cast<uint32_t>(e23)};
+ odd = {std::bit_cast<uint32_t>(o01), std::bit_cast<uint32_t>(o23)};
+}
+
+__device__ static uint2::Native_vec_ wmma12_t_from_d(wmma_float8 d)
+{
+ // h = 0 keeps rows 0, 2, 4, 6; h = 1 takes rows 1, 3, 5, 7 from the other half
+ uint2::Native_vec_ even, odd;
+ wmma12_rows_f16(d, even, odd);
+ const uint2::Native_vec_ other = {wmma12_other_half(odd.x), wmma12_other_half(odd.y)};
+ return wmma12_half() != 0 ? other : even;
+}
+
+__device__ static uint4::Native_vec_ wmma12_t_from_d_pair(wmma_float8 d)
+{
+ // h = 0 holds D0 rows 0..7, h = 1 D1 rows 0..7; lane h wants rows 2v + h of both
+ const bool high = wmma12_half() != 0;
+ uint2::Native_vec_ even, odd;
+ wmma12_rows_f16(d, even, odd);
+ const uint2::Native_vec_ send = high ? even : odd;
+ const uint2::Native_vec_ recv = {wmma12_other_half(send.x), wmma12_other_half(send.y)};
+ return high ? uint4::Native_vec_{recv.x, recv.y, odd.x, odd.y} : uint4::Native_vec_{even.x, even.y, recv.x, recv.y};
+}
+
+// the same for f32 shadows (element v = row 2v + h)
+__device__ static wmma_float8 wmma12_c_from_s(wmma_float4 c)
+{
+ wmma_float4 o;
+ for (uint32_t v = 0; v < 4; v++)
+ o[v] = wmma12_other_half_f(c[v]);
+ return {c[0], o[0], c[1], o[1], c[2], o[2], c[3], o[3]};
+}
+
+__device__ static wmma_float8 wmma12_c_from_s_pair(wmma_float4 c0, wmma_float4 c1)
+{
+ const bool high = wmma12_half() != 0;
+ wmma_float8 r;
+ for (uint32_t v = 0; v < 4; v++)
+ {
+ const float own = high ? c1[v] : c0[v], recv = wmma12_other_half_f(high ? c0[v] : c1[v]);
+ r[2 * v] = high ? recv : own;
+ r[2 * v + 1] = high ? own : recv;
+ }
+ return r;
+}
+
+__device__ static wmma_float4 wmma12_s_from_d(wmma_float8 d)
+{
+ const bool high = wmma12_half() != 0;
+ wmma_float4 r;
+ for (uint32_t v = 0; v < 4; v++)
+ {
+ const float odd = wmma12_other_half_f(d[2 * v + 1]);
+ r[v] = high ? odd : d[2 * v];
+ }
+ return r;
+}
+
+__device__ static wmma_float8 wmma12_s_from_d_pair(wmma_float8 d)
+{
+ const bool high = wmma12_half() != 0;
+ wmma_float8 r;
+ for (uint32_t v = 0; v < 4; v++)
+ {
+ const float recv = wmma12_other_half_f(high ? d[2 * v] : d[2 * v + 1]);
+ r[v] = high ? recv : d[2 * v];
+ r[4 + v] = high ? d[2 * v + 1] : recv;
+ }
+ return r;
+}
+
+extern "C"
+{
+ wmma_uint8 FUNC(d4r_wmma12_t_a_from_b)(uint2::Native_vec_ b)
+ {
+ __builtin_amdgcn_sched_barrier(0);
+ uint32_t n = wmma_lane() % 16;
+ wmma_uint8 result = {};
+ for (uint32_t j = 0; j < 4; j++)
+ {
+ uint32_t value = wmma_from_lane((n % 8) * 4 + (j >> 1) + 2 * wmma12_half(), (j & 1) ? b.y : b.x);
+ result[j] = n < 8 ? value : 0u;
+ }
+ return result;
+ }
+
+ wmma_uint8 FUNC(d4r_wmma12_t_b_from_a)(uint4::Native_vec_ a)
+ {
+ __builtin_amdgcn_sched_barrier(0);
+ wmma_uint8 result = {};
+ for (uint32_t j = 0; j < 4; j++)
+ result[j] = wmma12_gather2((j & 1) ? a.z : a.x, (j & 1) ? a.w : a.y, j);
+ return result;
+ }
+
+ wmma_uint8 FUNC(d4r_wmma12_t_a_from_b_pair)(uint2::Native_vec_ b0, uint2::Native_vec_ b1)
+ {
+ __builtin_amdgcn_sched_barrier(0);
+ wmma_uint8 result = {};
+ for (uint32_t j = 0; j < 4; j++)
+ result[j] = wmma12_gather2((j & 1) ? b0.y : b0.x, (j & 1) ? b1.y : b1.x, j);
+ return result;
+ }
+
+ uint2::Native_vec_ FUNC(d4r_wmma12_t_mma_f16)(wmma_uint8 a_t, wmma_uint8 b_t, uint2::Native_vec_ c_t)
+ {
+ return wmma12_t_from_d(wmma12_f16(a_t, b_t, wmma12_c_from_t(c_t)));
+ }
+
+ uint4::Native_vec_ FUNC(d4r_wmma12_t_mma_f16_pair)(wmma_uint8 a_t, wmma_uint8 b_t, uint2::Native_vec_ c0_t, uint2::Native_vec_ c1_t)
+ {
+ return wmma12_t_from_d_pair(wmma12_f16(a_t, b_t, wmma12_c_from_t_pair(c0_t, c1_t)));
+ }
+
+ wmma_float4 FUNC(d4r_wmma12_s_mma_f16)(wmma_uint8 a_t, wmma_uint8 b_t, wmma_float4 c, float up, float down)
+ {
+ return wmma12_s_from_d(wmma12_f16(a_t, b_t, wmma12_c_from_s(c * up))) * down;
+ }
+
+ wmma_float8 FUNC(d4r_wmma12_s_mma_f16_pair)(wmma_uint8 a_t, wmma_uint8 b_t, wmma_float4 c0, wmma_float4 c1, float up,
+ float down)
+ {
+ return wmma12_s_from_d_pair(wmma12_f16(a_t, b_t, wmma12_c_from_s_pair(c0 * up, c1 * up))) * down;
+ }
+
+ // e4m3 with the f16 widening (d4r_wmma_e4m3_a_to_f16 / b_to_f16 above): two chained f16 WMMAs in the
+ // 2^-16 scaled domain, like the gfx11 helpers
+ uint2::Native_vec_ FUNC(d4r_wmma12_t_mma_e4m3)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, uint2::Native_vec_ c_t)
+ {
+ wmma_float8 c = wmma12_c_from_t(c_t) * 0x1p-16f;
+ c = wmma12_f16(a_t_high, b_t_high, wmma12_f16(a_t_low, b_t_low, c));
+ return wmma12_t_from_d(c * 0x1p16f);
+ }
+
+ uint4::Native_vec_ FUNC(d4r_wmma12_t_mma_e4m3_pair)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, uint2::Native_vec_ c0_t,
+ uint2::Native_vec_ c1_t)
+ {
+ wmma_float8 c = wmma12_c_from_t_pair(c0_t, c1_t) * 0x1p-16f;
+ c = wmma12_f16(a_t_high, b_t_high, wmma12_f16(a_t_low, b_t_low, c));
+ return wmma12_t_from_d_pair(c * 0x1p16f);
+ }
+
+ wmma_float4 FUNC(d4r_wmma12_s_mma_e4m3)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, wmma_float4 c)
+ {
+ return wmma12_s_from_d(wmma12_f16(a_t_high, b_t_high, wmma12_f16(a_t_low, b_t_low, wmma12_c_from_s(c))));
+ }
+
+ wmma_float8 FUNC(d4r_wmma12_s_mma_e4m3_pair)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, wmma_float4 c0, wmma_float4 c1)
+ {
+ return wmma12_s_from_d_pair(
+ wmma12_f16(a_t_high, b_t_high, wmma12_f16(a_t_low, b_t_low, wmma12_c_from_s_pair(c0, c1))));
+ }
+
+ // Native FP8 (d4r_wmma12f8_*, D4R_ZLUDA_WMMA_FP8_NATIVE): the "widening" keeps the e4m3 bytes, two per
+ // dword in the low 16 bits and in the same slots the f16 widening fills, so the f16 gathers above move them;
+ // the MMA packs this half's 8 bytes (e4m3 k 8h .. 8h + 7 in order) into the FP8 operand. OCP e4m3 (gfx12)
+ // is NVIDIA's e4m3, so the bytes need no conversion. Shadows keep the 2^-16 scaled domain of the e4m3
+ // functions (exact power-of-two rescaling around the WMMA).
+ __device__ static uint2::Native_vec_ e4m3x4_split(uint32_t packed)
+ {
+ return {__builtin_amdgcn_perm(0u, packed, 0x0c0c0100u), __builtin_amdgcn_perm(0u, packed, 0x0c0c0302u)};
+ }
+
+ wmma_uint8 FUNC(d4r_wmma12f8_e4m3_a_to_f16)(uint4::Native_vec_ a)
+ {
+ auto x = e4m3x4_split(a.x), y = e4m3x4_split(a.y), z = e4m3x4_split(a.z), w = e4m3x4_split(a.w);
+ return {x.x, y.x, x.y, y.y, z.x, w.x, z.y, w.y};
+ }
+
+ uint4::Native_vec_ FUNC(d4r_wmma12f8_e4m3_b_to_f16)(uint2::Native_vec_ b)
+ {
+ auto x = e4m3x4_split(b.x), y = e4m3x4_split(b.y);
+ return {x.x, x.y, y.x, y.y};
+ }
+
+ uint2::Native_vec_ FUNC(d4r_wmma12f8_t_mma_e4m3)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, uint2::Native_vec_ c_t)
+ {
+ return wmma12_t_from_d(wmma12_e4m3(a_t_high, b_t_high, wmma12_e4m3(a_t_low, b_t_low, wmma12_c_from_t(c_t))));
+ }
+
+ uint4::Native_vec_ FUNC(d4r_wmma12f8_t_mma_e4m3_pair)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, uint2::Native_vec_ c0_t,
+ uint2::Native_vec_ c1_t)
+ {
+ return wmma12_t_from_d_pair(
+ wmma12_e4m3(a_t_high, b_t_high, wmma12_e4m3(a_t_low, b_t_low, wmma12_c_from_t_pair(c0_t, c1_t))));
+ }
+
+ wmma_float4 FUNC(d4r_wmma12f8_s_mma_e4m3)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, wmma_float4 c)
+ {
+ wmma_float8 d = wmma12_e4m3(a_t_high, b_t_high, wmma12_e4m3(a_t_low, b_t_low, wmma12_c_from_s(c * 0x1p16f)));
+ return wmma12_s_from_d(d) * 0x1p-16f;
+ }
+
+ wmma_float8 FUNC(d4r_wmma12f8_s_mma_e4m3_pair)(wmma_uint8 a_t_low, wmma_uint8 b_t_low, wmma_uint8 a_t_high,
+ wmma_uint8 b_t_high, wmma_float4 c0, wmma_float4 c1)
+ {
+ wmma_float8 d = wmma12_e4m3(a_t_high, b_t_high,
+ wmma12_e4m3(a_t_low, b_t_low, wmma12_c_from_s_pair(c0 * 0x1p16f, c1 * 0x1p16f)));
+ return wmma12_s_from_d_pair(d) * 0x1p-16f;
+ }
+}
+
template <typename T, typename A, typename B, int BOffset, typename OcklFunc>
__device__ T dp2a_impl(T a, T b, T c, OcklFunc ockl_func)
{
--- a/zluda/src/impl/module.rs
+++ b/zluda/src/impl/module.rs
@@ -356,6 +356,7 @@
wmma: llvm_zluda::compile::wmma_rewrite_requested(),
wmma_fp8: llvm_zluda::compile::wmma_fp8_rewrite_requested(),
param_provenance: llvm_zluda::compile::param_provenance_enabled(),
+ wmma_fp8_native: llvm_zluda::compile::wmma_fp8_native_effective(gcn_arch),
};
let mut cache_with_key = match (text, global_state.cache_path.as_ref()) {
(Some(text), Some(p)) => (|| {
@@ -397,6 +398,10 @@
wmma: bool,
wmma_fp8: bool,
param_provenance: bool,
+ // d4r: native FP8 WMMA where it takes effect (gfx12); absent from the key everywhere else, so the setting
+ // (d4r.ini NativeFp8) never forces a recompile on GPUs that ignore it
+ #[serde(skip_serializing_if = "std::ops::Not::not")]
+ wmma_fp8_native: bool,
}
fn get_hip_properties<'a>() -> Result<hipDeviceProp_tR0600, CUerror> {
@@ -430,6 +435,7 @@
&& !matches!(
key.as_str(),
"D4R_ZLUDA_NATIVE_DIR" | "D4R_ZLUDA_DIR" | "D4R_ZLUDA_DUMP_DIR" | "D4R_ZLUDA_DUMP_PTX_DIR" | "D4R_ZLUDA_LIBCUDA"
+ | "D4R_ZLUDA_WMMA_FP8_NATIVE" // keyed by effect (ExtraCacheAttributes)
)
})
.map(|(key, value)| format!("{key}={}", value.trim()))
scripts/build_zluda_ptx_helpers.sh
+7 -1
View File
@@ -9,6 +9,9 @@ LLVM_DIS="${LLVM_DIS:-$ROCM_ROOT/lib/llvm/bin/llvm-dis}"
LLVM_AS="${LLVM_AS:-$ROCM_ROOT/lib/llvm/bin/llvm-as}"
OCML="${OCML:-$ROCM_ROOT/amdgcn/bitcode/ocml.bc}"
WORK_ROOT="${D4R_ZLUDA_PTX_HELPER_WORK:-${TMPDIR:-/tmp}/d4r-zluda-ptx-helpers}"
# D4R_ZLUDA_HELPER_FLAGS: extra compiler flags, e.g. -DD4R_WMMA12_SHIM for the gfx12-layout test build (never
# ship it: its gfx12 helpers use gfx11 WMMAs)
read -r -a EXTRA_FLAGS <<< "${D4R_ZLUDA_HELPER_FLAGS:-}"
if [[ -z "$ZLUDA_SOURCE_ROOT" || ! -f "$ZLUDA_SOURCE_ROOT/ptx/lib/zluda_ptx_impl.cpp" ]]; then
printf 'Set ZLUDA_SOURCE_ROOT to a ZLUDA source checkout\n' >&2
@@ -27,7 +30,7 @@ PTX_LIB="$ZLUDA_SOURCE_ROOT/ptx/lib"
COMMON_FLAGS=(-DHIP_ENABLE_WARP_SYNC_BUILTINS -std=c++20 -Xclang -fdenormal-fp-math=dynamic
-Wall -Wextra -Wsign-compare -Wconversion -x hip "$SOURCE" -nogpulib -O3
-mno-wavefrontsize64 --offload-device-only --offload-arch=gfx1030 -emit-llvm -c
-Xclang -mlink-bitcode-file -Xclang "$OCML")
-Xclang -mlink-bitcode-file -Xclang "$OCML" "${EXTRA_FLAGS[@]}")
build_one() {
local suffix="$1"
@@ -55,3 +58,6 @@ build_one() {
build_one ""
build_one "_constrained" -ffp-model=strict -ffp-exception-behavior=ignore
# wave64 helpers (D4R_ZLUDA_WAVE64: two CUDA warps per wave; see d4r_warp_base in zluda_ptx_impl.cpp)
build_one "_w64" -mwavefrontsize64 -DD4R_WAVE64
build_one "_constrained_w64" -ffp-model=strict -ffp-exception-behavior=ignore -mwavefrontsize64 -DD4R_WAVE64
scripts/d4r_config.py
+3
View File
@@ -129,6 +129,9 @@ def main(argv):
value = flag("Kernels", key)
if value is not None:
env[var] = "1" if value else "0"
# RDNA4's native FP8 WMMA (on unless set off; ZLUDA and the bridge ignore it on other GPUs)
value = flag("Kernels", "NativeFp8")
env["D4R_ZLUDA_WMMA_FP8_NATIVE"] = "0" if value is False else "1"
value = get("Kernels", "ImplicitMaxBlock")
if value:
env["D4R_ZLUDA_IMPLICIT_MAX_BLOCK"] = value
scripts/d4r_proton_env.sh
+2 -2
View File
@@ -22,8 +22,8 @@ export D4R_ZLUDA_LIBCUDA="$D4R_ZLUDA_DIR/libcuda.so"
# most 256-thread blocks, as CUDA itself allows 255 registers for them; the
# default 1024 bound caps gfx11 waves at 96 VGPRs and makes them spill.
export D4R_ZLUDA_IMPLICIT_MAX_BLOCK="${D4R_ZLUDA_IMPLICIT_MAX_BLOCK:-256}"
# f16 tensor-core MMAs (DLSS's convolutions) run on gfx11 WMMA instead of a
# dot-product fallback (patches/zluda/0003); 0 selects the fallback.
# f16 tensor-core MMAs (DLSS's convolutions) run on gfx11/gfx12 WMMA instead of a
# dot-product fallback (patches/zluda/0003 and 0007); 0 selects the fallback.
export D4R_ZLUDA_WMMA="${D4R_ZLUDA_WMMA:-1}"
# DLSS frames in flight. The presented upscaled image is this many game frames
# old at most: 3 = highest frame rate (Townfall ~86 fps, age 2-3), 2 = age 2
scripts/package_release.sh
+33 -11
View File
@@ -13,7 +13,8 @@
# D4R_ROCM_DIR ROCm with clang, for the kernels default /opt/rocm
# D4R_ROCM_RUNTIME ROCm runtime bundled as d4r/rocm (scripts/fetch_rocm_runtime.sh)
# default ~/.cache/d4r-rocm-runtime
# D4R_GPU_ARCHS GPU targets to build kernels for default gfx1101
# D4R_GPU_ARCHS GPU targets to build kernels for default gfx1100 gfx1101 gfx1102 gfx1103 gfx1200 gfx1201
# (RDNA4 targets also get a <target>-fp8 folder: the kernels for d4r.ini NativeFp8 = on)
# D4R_OPTISCALER_LICENSE OptiScaler's LICENSE (GPL-3.0) text; default: the system's SPDX copy
# D4R_VKD3D_SRC vkd3d-proton source checkout, for its license files default ~/.cache/d4r-vkd3d-proton
# D4R_ZLUDA_SRC ZLUDA source checkout, for its license files default: D4R_ZLUDA_DIR's ../d4r-zluda-upstream
@@ -21,7 +22,9 @@
# The zip also contains NVIDIA's files and the texture kernels built from NVIDIA's PTX; redistributing
# those is up to whoever publishes it (they are not covered by d4r's license):
# D4R_BUNDLE_DLSS nvngx_dlss.dll to include D4R_BUNDLE_NGX _nvngx.dll to include
# D4R_BUNDLE_TEX directory with the texture-kernel code objects (kernels/build.sh tex), for gfx1101
# D4R_BUNDLE_TEX directory with texture-kernel code objects (kernels/build.sh tex), one subdirectory
# per target folder (gfx1101, gfx1201, gfx1201-fp8, ...), or a flat gfx1101 directory for
# older builds
# D4R_BUNDLE_NVIDIA=0 leaves them out (d4r-VERSION-nonvidia.zip; users then add the two DLLs themselves).
set -euo pipefail
@@ -37,7 +40,7 @@ ZLUDA="${D4R_ZLUDA_DIR:-$HOME/.cache/d4r-zluda-current}"
VKD3D="${D4R_VKD3D_DIR:-$HOME/.cache/d4r-vkd3d-d4r}"
ROCM="${D4R_ROCM_DIR:-/opt/rocm}"
ROCM_RUNTIME="${D4R_ROCM_RUNTIME:-$HOME/.cache/d4r-rocm-runtime}"
ARCHS="${D4R_GPU_ARCHS:-gfx1101}"
ARCHS="${D4R_GPU_ARCHS:-gfx1100 gfx1101 gfx1102 gfx1103 gfx1200 gfx1201}"
: "${D4R_OPTISCALER:?set D4R_OPTISCALER to the OptiScaler release archive or folder}"
: "${D4R_DLSS_DLLS:?set D4R_DLSS_DLLS to the nvngx_dlss.dll files the kernel manifest accepts}"
for f in "$ZLUDA/libnvcuda.so" "$VKD3D/d3d12.dll" "$VKD3D/d3d12core.dll" "$ROCM_RUNTIME/lib/libamdhip64.so.7"; do
@@ -108,16 +111,35 @@ else
printf 'Put NVIDIA'"'"'s NGX runtime, _nvngx.dll, in this folder (see D4R_README.txt).\r\n' > "$STAGE/d4r/ngx/README.txt"
fi
IFS=: read -r -a DLLS <<< "$D4R_DLSS_DLLS"
# one folder per target; RDNA4 targets also get <target>-fp8 (native FP8 WMMA, d4r.ini NativeFp8), which the
# bridge serves instead when that setting is on
folders=()
for arch in $ARCHS; do
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$arch" "$ROOT/kernels/build.sh" k "$STAGE/d4r/kernels/$arch" >/dev/null
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$arch" "$ROOT/kernels/build.sh" m "$STAGE/d4r/kernels/$arch" >/dev/null
rm -f "$STAGE/d4r/kernels/$arch"/*.resolution.txt # empty LTO notes from clang's -save-temps
if [[ "$VARIANT" == full && "$arch" == gfx1101 ]]; then
for f in "$D4R_BUNDLE_TEX"/*.hsaco; do # texture kernels only; the layers above are built from source
[[ -e "$STAGE/d4r/kernels/$arch/$(basename "$f")" ]] || cp "$f" "$STAGE/d4r/kernels/$arch/"
done
folders+=("$arch")
[[ "$arch" == gfx12* ]] && folders+=("$arch-fp8")
done
for folder in "${folders[@]}"; do
arch="${folder%-fp8}"
fp8=0
[[ "$folder" == *-fp8 ]] && fp8=1
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$arch" D4R_NATIVE_FP8="$fp8" "$ROOT/kernels/build.sh" k "$STAGE/d4r/kernels/$folder" >/dev/null
D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$arch" D4R_NATIVE_FP8="$fp8" "$ROOT/kernels/build.sh" m "$STAGE/d4r/kernels/$folder" >/dev/null
rm -f "$STAGE/d4r/kernels/$folder"/*.resolution.txt # empty LTO notes from clang's -save-temps
if [[ "$VARIANT" == full ]]; then
tex_dir=""
if [[ -d "$D4R_BUNDLE_TEX/$folder" ]]; then
tex_dir="$D4R_BUNDLE_TEX/$folder"
elif [[ "$folder" == gfx1101 ]]; then
tex_dir="$D4R_BUNDLE_TEX" # older, flat gfx1101 texture-kernel builds
fi
if [[ -n "$tex_dir" ]]; then
for f in "$tex_dir"/*.hsaco; do # texture kernels only; the layers above are built from source
[[ -f "$f" ]] || continue
[[ -e "$STAGE/d4r/kernels/$folder/$(basename "$f")" ]] || cp "$f" "$STAGE/d4r/kernels/$folder/"
done
fi
fi
python3 "$ROOT/kernels/tools/kernel_manifest.py" "$STAGE/d4r/kernels/$arch" "${DLLS[@]}"
python3 "$ROOT/kernels/tools/kernel_manifest.py" "$STAGE/d4r/kernels/$folder" "${DLLS[@]}"
done
# licenses and sources
tools/d4r_nvngx_shim.cpp
+12 -3
View File
@@ -535,6 +535,8 @@ static void load_portable_config()
portable_set_unix("D4R_ZLUDA_NATIVE_DIR", g_portable.unixDir + "/kernels");
portable_set_unix("D4R_ZLUDA_WMMA", ini_flag(ini, "kernels", "Wmma", 1) ? "1" : "0");
portable_set_unix("D4R_ZLUDA_WMMA_FP8", ini_flag(ini, "kernels", "Fp8Wmma", 1) ? "1" : "0");
// RDNA4's native FP8 WMMA; ZLUDA and the bridge ignore it on other GPUs
portable_set_unix("D4R_ZLUDA_WMMA_FP8_NATIVE", ini_flag(ini, "kernels", "NativeFp8", 1) ? "1" : "0");
portable_set_unix("D4R_ZLUDA_IGNORE_DENORMAL", ini_flag(ini, "kernels", "IgnoreDenormals", 1) ? "1" : "0");
const std::string maxBlock = ini_value(ini, "kernels", "ImplicitMaxBlock");
portable_set_unix("D4R_ZLUDA_IMPLICIT_MAX_BLOCK", maxBlock.empty() ? "256" : maxBlock);
@@ -3774,6 +3776,12 @@ static void run_evaluation(Feature* feature, int slotIndex, uint32_t frame, Fram
g.cu.setArrayRedirect(feature->output.array, target, static_cast<uint32_t>(rowBytes)) == 0;
if (!redirect && feature->outputRedirected)
g.cu.setArrayRedirect(feature->output.array, 0, 0);
static bool logged = false;
if (redirect && !logged)
{
logf("frame %u: direct output (the native output kernel writes the shared buffer; no array copy)", frame);
logged = true;
}
feature->outputRedirected = redirect;
}
const int eventStartResult = timing.enabled && feature->profileEventsReady
@@ -4250,11 +4258,12 @@ D4R_EXPORT NgxResult NVSDK_NGX_D3D12_CreateFeature(ID3D12GraphicsCommandList*, u
preset = value;
}
feature->preset = presets[1];
// Direct output needs the native output kernel of the preset in use (verified: K = 11; M's result is
// not written through the redirected stores). Only a forced preset qualifies.
// Direct output needs the native output kernel of the preset in use (verified: K = 11 through
// hiluma_engine_output, M = 13 through rrlite_downsample_kernel, its last kernel). Only a forced
// preset qualifies.
if (!forced.empty())
{
char allowed[128] = "11";
char allowed[128] = "11,13";
if (const std::string list = env_string("D4R_SHIM_OUTPUT_DIRECT_PRESETS"); !list.empty())
snprintf(allowed, sizeof(allowed), "%s", list.c_str());
for (char* token = strtok(allowed, ","); token != nullptr; token = strtok(nullptr, ","))
tools/wine_nvcuda_bridge.c
+41 -8
View File
@@ -389,11 +389,21 @@ static void native_cleanup(void)
remove_directory(native_served);
}
/* The first GPU's gfx target from the KFD topology ("gfx1101"), "" when unknown. */
/* The gfx target ("gfx1101") of the GPU DLSS runs on, "" when unknown: D4R_GPU_ARCH if set, else the KFD
topology's GPU with the most SIMDs. With an integrated and a discrete GPU (Ryzen 7000/9000 desktops) the
first KFD GPU node can be the integrated one, whose kernels would not match the discrete GPU. */
static void gpu_architecture(char* out, size_t size)
{
out[0] = '\0';
for (int node = 0; node < 16 && out[0] == '\0'; ++node)
const char* forced = getenv("D4R_GPU_ARCH");
if (forced != NULL && strncmp(forced, "gfx", 3) == 0)
{
snprintf(out, size, "%s", forced);
return;
}
unsigned long best_simds = 0;
int gpus = 0;
for (int node = 0; node < 16; ++node)
{
char path[128];
snprintf(path, sizeof(path), "/sys/class/kfd/kfd/topology/nodes/%d/properties", node);
@@ -408,9 +418,17 @@ static void gpu_architecture(char* out, size_t size)
sscanf(line, "gfx_target_version %lu", &target);
}
fclose(properties);
if (simds != 0 && target != 0)
if (simds == 0 || target == 0)
continue;
++gpus;
if (simds > best_simds)
{
best_simds = simds;
snprintf(out, size, "gfx%lu%lu%lx", target / 10000, (target / 100) % 100, target % 100);
}
}
if (gpus > 1)
tracef("native kernels: %d GPUs; using the one with the most SIMDs (%s); D4R_GPU_ARCH overrides", gpus, out);
}
static void prepare_native_kernels(const char* cache_home)
@@ -421,9 +439,22 @@ static void prepare_native_kernels(const char* cache_home)
/* a release directory holds one manifest and set of code objects per GPU target */
char architecture[32], source[1024], path[1200];
gpu_architecture(architecture, sizeof(architecture));
snprintf(source, sizeof(source), "%s/%s", configured, architecture);
snprintf(path, sizeof(path), "%s/d4r-kernels.txt", source);
FILE* manifest = architecture[0] != '\0' ? fopen(path, "r") : NULL;
FILE* manifest = NULL;
/* RDNA4 with native FP8 WMMA (D4R_ZLUDA_WMMA_FP8_NATIVE=1, d4r.ini NativeFp8): the <target>-fp8 variant,
whose kernels match ZLUDA's FP8 lowering; no other GPU has such a directory */
const char* fp8 = getenv("D4R_ZLUDA_WMMA_FP8_NATIVE");
if (fp8 != NULL && strcmp(fp8, "1") == 0 && strncmp(architecture, "gfx12", 5) == 0)
{
snprintf(source, sizeof(source), "%s/%s-fp8", configured, architecture);
snprintf(path, sizeof(path), "%s/d4r-kernels.txt", source);
manifest = fopen(path, "r");
}
if (manifest == NULL)
{
snprintf(source, sizeof(source), "%s/%s", configured, architecture);
snprintf(path, sizeof(path), "%s/d4r-kernels.txt", source);
manifest = architecture[0] != '\0' ? fopen(path, "r") : NULL;
}
if (manifest == NULL)
{
snprintf(source, sizeof(source), "%s", configured);
@@ -679,7 +710,7 @@ static void verify_native_kernels(const void* image)
pthread_mutex_unlock(&native_lock);
}
/* DLSS's output kernels (hiluma_engine_output_*): NGX looks up every variant, then launches the one
/* DLSS's output kernels (hiluma_engine_output_*, rrlite_downsample_kernel_*): NGX looks up every variant, then launches the one
matching the feature's flags. The shim's direct output relies on stores that only d4r's native
output kernel makes, so it asks whether the last output kernel launched was native (a file in
D4R_ZLUDA_NATIVE_DIR, which is the verified per-process directory in a portable install). */
@@ -694,7 +725,9 @@ static int last_output_native = -1;
static void note_function_lookup(CUfunction function, const char* name)
{
if (name == NULL || strncmp(name, "hiluma_engine_output", 20) != 0)
/* DLSS 4 writes its result with hiluma_engine_output_*, DLSS 4.5 (rrlite) with its final
rrlite_downsample_kernel_* (the post kernel works at a larger internal resolution) */
if (name == NULL || (strncmp(name, "hiluma_engine_output", 20) != 0 && strncmp(name, "rrlite_downsample_kernel", 24) != 0))
return;
const char* directory = getenv("D4R_ZLUDA_NATIVE_DIR");
int native = 0;