woutermans
50fa65d87b
- engine/rr: d4r's own shaders for the denoiser frame (exposure, six encoders,
five decoders, hkpn output stage) at NGX layouts, plus PTX->GLSL/HIP packing
scripts, reference implementations and weight extraction
- runtime.h/runtime_rr.inc, game_rr.{h,cpp}: GameDenoiser wired into the engine
build (d4r-rr replay binary) and the D3D12 shim's feature-13 path
- shim: answer GetFeatureRequirements from d4r's own state only (the old
cross-binary CUDA handoff crashed OptiScaler's scan); drop the core's
CUDA GetFeatureRequirements import
- shim/harness: exact guide capture/replay of a frame's RR inputs
- kernels/rr: expanded enc0/engine prototypes; kernels/tools: ptxsim rewrite
plus rrsim and an engine test; vulkan: ptx2glsl FP8/cooperative-matrix
support, kvk inout buffers and KVK_DISABLE_OPTIMIZATION
- bridge: report the surface array in cuSurfObjectCreate traces
- packaging: RR settings and the ICU launch overrides in D4R_README
164 lines
12 KiB
Markdown
164 lines
12 KiB
Markdown
# `cuda_dldn_engine_swin_enc0_kernel` on the staged-image engine
|
||
|
||
This file is the implementation record for enc0's port from the mirror-cut scalar body to the
|
||
staged-image engine (`kernels/rr/rrswin_engine.h`). §1–§5 are the design contract the code
|
||
implements; §6 is the evidence table that accrues as each offset and layout claim is
|
||
validated (per-phase oracle captures, engine unit tests, output-surface compares).
|
||
|
||
Status: **design fixed; primitives and body port dispatched.** Gates and their results are
|
||
recorded at the bottom as they are measured.
|
||
|
||
## 1. What the port changes
|
||
|
||
The mirror-cut body `d4r_swin_enc0` (`kernels/rr/rrswin_enc0.hip`) keeps
|
||
|
||
* its cut ABI (called from NVIDIA's PTX between the prologue and E22; arguments
|
||
`a0..a23` = the staged block's 24 A-fragment words, `b0..b3` = phase 1's C seeds,
|
||
`w` = arg040 (the 54560-byte prepared weight image), `prm` = the 1032-byte parameter
|
||
block, `scratch` = the cut's `input_tensor`, 3200 bytes of shared memory),
|
||
* every epilogue (E5/E10 softmax, E6/E12 norms and table adds, E13/15/17/19 activations,
|
||
E20 arena scatter + lane merge, E21 surface scatter, E22 handover),
|
||
* the score/P·V phases 5, 6, 10, 11 (`score_gemm`, `pv_gemm`), phase 21 (`band21`) and
|
||
phase 22's GEMM as they are (scalar f32 dots over exact e4m3 decodes, 136 of 536 mma),
|
||
|
||
and replaces the sixteen plain GEMM phases **1, 2, 3, 4, 7, 8, 9, 12, 13, 14, 15, 16, 17,
|
||
18, 19, 20** (400 of 536 mma) with engine GEMMs: operands read from staged e4m3 images and
|
||
prepared slots, one `k32_e4m3` per (m, n) tile, accumulators rounded to f16 after each k32
|
||
step — the conservative chain that is bit-identical to the translated kernel's
|
||
`mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16` (see §2 of docs/rr-engine.md and
|
||
`rrswin_engine_k16_test.cpp` step 2).
|
||
|
||
Replaced with this port: the scalar `gemm()` calls for those sixteen phases, whose per-mma
|
||
cost is 64 byte loads from the weight image, 64 e4m3 decodes, 128 f32 FMAs and twelve DPP
|
||
permutes per lane.
|
||
|
||
## 2. Orientation and the two identities that make the port mechanical
|
||
|
||
Measured facts (all previously validated, not re-derived here):
|
||
|
||
* **D layout** (`kernels/tools/rrswin_operand_fit.cpp`): engine tile D — lane `l`, VGPR `i`
|
||
holds row `i + 8h`, column `l & 15`, `h = l >> 4`. With the engine's operand roles
|
||
(`swin_gemm_tile`: weights = A via `kop_image`, activations = B via `swin_operand`) the
|
||
tile's D is `[weight column wr][token tk]` — NVIDIA's D transposed:
|
||
`D[wr = i + 8h + 16*mt][tk = (l&15) + 16*nt]` for tile `(mt, nt)`.
|
||
* **Staging D into the next phase's A image** is exactly the existing
|
||
`swin_stage_operand(img, stride, 16*nt, mt, acc)`: it writes
|
||
`img[(16*nt + l&15)*stride + 16*mt + 8h + i] = e4m3(acc[i])`, i.e. image row = token, image
|
||
k = weight column — the A-operand form the next phase's `swin_operand` reads.
|
||
* **The residual identity**: the next phase's C operand is a cell-for-cell copy of the
|
||
producer's D. NVIDIA's mma computes `D = A·B + C` where `C[m][n]` sits at `(token, weight
|
||
column)`; the engine tile's D cell `(wr, tk)` is `C[tk][wr]`. So phase P+1's engine tile
|
||
seed is `c[i] = f16(P's own tile (mt, nt) at VGPR i)` — the producer's accumulator compacted
|
||
to f16 pairs, no exchange, no shuffle. This is why phases 12, 16, 18, 20 (residual =
|
||
previous phase's D) and 7 and 14 (residual + per-column table) need no C plumbing.
|
||
* **Per-column biases and tables are column-major f16**: every read of the form
|
||
`*(u32*)(w + base + 16*nt + ((laneid<<2)&12))` (E9's seeds, `wbias_regs`, E6/E13's tables)
|
||
places column `c` at byte `2*c` of the blob. Evidence: the word for `(nt, lane)` covers
|
||
columns `8nt+2t, 8nt+2t+1` at bytes `16nt+4t .. +3 = 2*(8nt+2t) .. +3`. The engine tile
|
||
seed for a per-column C is therefore `c[i] = f32(f16 at w + 2*(i + 8h + 16*mt))`.
|
||
|
||
## 3. The dcell exchange (the one cross-lane primitive the port adds)
|
||
|
||
Epilogues that consume the NVIDIA `d[mi][rh][nt]` view (`norm_as` for phases 2 and 12; the
|
||
scalar `score_gemm`/`pv_gemm` for phases 3, 4, 8, 9; E20's `r20` build for phase 20) get it
|
||
from the engine tiles through one shared-memory exchange.
|
||
|
||
* Writer (per tile `(mt, nt)`, lane `l`, its own four f16 pairs
|
||
`pairs[p] = f16(D[2p]) | f16(D[2p+1]) << 16`, `p = 0..3`):
|
||
`lds[(mt*n_nt + nt)*128 + lane*4 + p] = pairs[(mt*n_nt + nt)*4 + p]` — all of it in
|
||
**u32 units: 128 words (= 512 bytes) per tile**, 32 lanes × 4 words.
|
||
* Reader: `d[mi][rh][nt] = lds[((nt>>1)*n_nt + mi)*128 + (8rh + g + 16*(nt&1))*4 + t]`
|
||
with `g = laneid>>2`, `t = laneid&3` — same units.
|
||
|
||
Derivation (checkable against the two measured facts above): the wanted cell is
|
||
`D[wr = 8nt + 2t + cb][tk = 16mi + 8rh + g]`; its engine location is tile
|
||
`(mt_e = nt>>1, nt_e = mi)`, lane `(tk%16) + 16*((wr%16)>>3) = 8rh + g + 16*(nt&1)`,
|
||
VGPR `wr%8 = 2t + cb` — i.e. the reader's own pair index `t`, both halves. A full
|
||
twelve-tile exchange is 12 × 512 B = 6144 B and an eight-tile one 4096 B — both exceed
|
||
the scratch at exchange time, so **every exchange runs in 2048-byte chunks** (four tiles
|
||
at a time) overlaying the just-dead image region, with the chunk loop's barriers owned by
|
||
the phase code (§5).
|
||
|
||
## 4. New primitives (kernels/rr/rrswin_engine.h)
|
||
|
||
All are thin, checked by `kernels/tools/rrswin_enc0_engine_test.cpp` against independent
|
||
host models before the body relies on them.
|
||
|
||
| primitive | signature | semantics |
|
||
|---|---|---|
|
||
| `swin_slots_from_wimg` | `(uint8_t* slots, const uint8_t* w, int base, int n_rows)` | One phase's B fragments (8 u32 per lane at `w + base + 512*t + 16*lane + 4*word`, `t = 0..n_rows/16-1`, `word = 0..3`) scattered into slot layout: slot `(n_rows*(k>>4) + n)*16 + (k&15)`, with `(n, k)` from the verified map `n = 16t + (lane>>2) + 8*(word>>1)`, `k = 2t' + (j&1) + 16*(j>>1) + 8*(word&1)`. Matches `swin_gemm_tile`'s read with `slot_base = 16*mt`, `slot_step = n_rows`. |
|
||
| `swin_dcell_exchange` | writer `(lds_u32*, const uint32_t* pairs, int n_mt, int n_nt)` / reader `(const lds_u32*, uint32_t* d, int n_mi, int n_nt4, ...)` | §3. Writer and reader are separate calls; the wave barrier between them belongs to the caller. |
|
||
| `swin_stage_operand_p` | `(uint8_t* img, int stride, int row_base, int chunk, const uint32_t pairs[4], int mode)` | `swin_stage_operand` from a tile's four f16 pairs with a per-pair transform before `e4m3x2`: mode 0 plain, 1 `relu_pair` (ZLUDA's `relu_keep_nan_halves`), 2 `act_pair` (the clamped cubic). Byte `i` of the image row holds `e4m3(transform(pair[i/2]) half i%2)`. |
|
||
|
||
`swin_stage_fragment` needs no row-offset variant: the second half of a split 96-row GEMM
|
||
restages its fragment tile group at *local* rows 0–47 (`swin_stage_fragment(img, 32,
|
||
frag + 12, 3)`) and runs its GEMM tiles against local `row_base = 0`, emitting results to
|
||
the tiles' *global* indices (nt 3–5). The D output of every split phase is registers
|
||
(pairs), so no image write needs the global row position.
|
||
|
||
The relu and act transforms are per-element, so `e4m3x2(transform(pair))` equals the
|
||
translated kernel's `transform`-then-`cvt.rn.satfinite.e4m3x2.f16x2` element for element.
|
||
|
||
## 5. Shared-memory map (3200 bytes, the cut's `input_tensor`)
|
||
|
||
`input_noise` (800 B) is a separate `.shared` variable whose offset relative to
|
||
`input_tensor` the linked body cannot know, so the budget is the 3200-byte scratch only.
|
||
|
||
| region | size | live during |
|
||
|---|---|---|
|
||
| image buffer | 1536 (48 rows × 32 B) or 2048 (64 rows × 32 B) | one phase's A image, split into halves for the 96-row phases |
|
||
| slots | 1024 (64 slots × 16 B; 512 for phase 22, unused — phase 22 stays scalar) | one phase's GEMM |
|
||
| dcell exchange | 128 × (tile count) | between a phase's GEMM and its dcell-consuming epilogue |
|
||
|
||
The 96-row phases (1, 2, 3, 4, 8, 9) cannot hold a 3072-byte image and 1024 bytes of slots
|
||
at once, so their GEMM runs in two token-tile halves (nt 0–2, nt 3–5): stage rows 0–47,
|
||
run, restage rows 48–95, run. The halves' staging sources are registers (the cut's `as1`,
|
||
the normed tile `as3`), so restaging is cheap. Phases 7 and 12–20 have 64-row images
|
||
(2048 B) and fit without splitting. Every exchange overlays the just-dead image or slot
|
||
region of the same phase.
|
||
|
||
## 6. Per-phase implementation table
|
||
|
||
| phase | mma | A (engine B operand) | B (engine A operand) | C seed | D consumed as |
|
||
|---|---|---|---|---|---|
|
||
| 1 | 24 | `as1` fragments → image halves | slots from `w+4096` | per-column, `w+5120` | relu-staged into phase 2's image (mode 1), no dcell |
|
||
| 2 | 24 | relu-staged image of D1 | slots from `w+5184` | per-column, `w+6208` | dcell exchange → `norm_as` (gain `w+6272`, eps 2⁻¹³·2⁻⁴? see body) → `as3`; pairs kept for phase 7's C |
|
||
| 3 | 24 | `as3` halves | slots from `w+6336` | — | dcell exchange → `d3` |
|
||
| 4 | 24 | `as3` halves | slots from `w+7360` | — | dcell exchange → `d4` |
|
||
| 5–6 | 48+48 | scalar `score_gemm`/`pv_gemm` unchanged (`d3`, `d4`, bias `w+8384`) | | | |
|
||
| 7 | 16 | `as7` (from `d6`) → image | slots from `w+20672` | `d2` pairs + table `w+37056` | pairs kept (phase 12's C) |
|
||
| 8 | 24 | `as3` halves | slots from `w+21696` | — | dcell exchange → `d8` |
|
||
| 9 | 24 | `as3` halves | slots from `w+22720` | — | dcell exchange → `d9` |
|
||
| 10–11 | 48+48 | scalar `score_gemm`/`pv_gemm` unchanged (`d8`, `d9`, bias `w+23744`) | | | |
|
||
| 12 | 16 | `as12` (from `d11`) → image | slots from `w+36032` | `d7` pairs | pairs kept (phase 14's C); dcell exchange → `norm_as` (gain `w+37120`, eps 2⁻¹³) → `asN` |
|
||
| 13 | 16 | `asN` → image | slots from `w+37184` | per-column, `w+38208` | act-staged into phase 14's image (mode 2) |
|
||
| 14 | 16 | act-staged image of D13 | slots from `w+38272` | `d12` pairs + table `w+39296` | pairs kept (phase 16's C) |
|
||
| 15 | 16 | `asN` → image | slots from `w+39360` | per-column, `w+40384` | act-staged into phase 16's image |
|
||
| 16 | 16 | act-staged image of D15 | slots from `w+40448` | `d14` pairs | pairs kept (phase 18's C) |
|
||
| 17 | 16 | `asN` → image | slots from `w+41472` | per-column, `w+42496` | act-staged into phase 18's image |
|
||
| 18 | 16 | act-staged image of D17 | slots from `w+42560` | `d16` pairs | pairs kept (phase 20's C) |
|
||
| 19 | 16 | `asN` → image | slots from `w+43584` | per-column, `w+44608` | act-staged into phase 20's image |
|
||
| 20 | 16 | act-staged image of D19 | slots from `w+44672` | `d18` pairs | dcell exchange → E20's `r20` build (unchanged) |
|
||
| 21 | 32 | scalar `band21` unchanged (merged `m21`, weights `w+45696`, bias `w+53888`) | | | |
|
||
| 22 | 8 | scalar `gemm(..., nnt=2)` unchanged (`r20`, `w+54016`, bias `w+54528`) | | | |
|
||
|
||
Weight-image byte ranges are the body's own verified table (`rrswin_enc0.hip`'s
|
||
"Weight image layout of the 22 phases" comment); `waddr`'s formula is verified on all 1024
|
||
`(n, k)` pairs of a tile.
|
||
|
||
## 7. Validation gates and results
|
||
|
||
*(filled as measured)*
|
||
|
||
| gate | requirement | result |
|
||
|---|---|---|
|
||
| translated reference set | 22 oracle stages `p1..p22` captured (`/home/wouter/rr/gate-enc0/oracle-p<stage>/dumps`), 6 launches each | **captured** |
|
||
| harness determinism | `compare_launch_dump.py` oracle-p1 vs oracle-p3: 108 args, 6 differ — only arg048, first differing byte exactly at 46,000,000 (the debug payload); all surfaces/textures/weights byte-identical | **proven** |
|
||
| performance anchor (translated) | `D4R_HARNESS_FRAME_WAIT_MS=0` (continuous, matching the 16.2 ms config), 30 frames, frames 13–30: **enc0 = 3.879 ms/frame**, total 16.79 — the documented 3.79 ms baseline reproduced within ~2 % | **measured** |
|
||
| full-frame reference | `control-80`: 80-frame run of the installed runtime with `D4R_HARNESS_SAVE_FRAMES`, 80 saved frames (`/home/wouter/rr/gate-enc0/control-80/out.raw.frameN`) | **captured** |
|
||
| primitives | `rrswin_enc0_engine_test.cpp` steps vs independent host models | pending |
|
||
| per-phase oracle | engine `D4R_RRSWIN_DEBUG_STAGE=<n>` vs the translated stage, payload-window compare | pending |
|
||
| layer surfaces | `compare_launch_dump.py` translated vs engine enc0 output surfaces | pending |
|
||
| full frame | 80-frame engine run, frames 21–80 byte-identical to `control-80` | pending |
|
||
| performance | enc0 < 3.879 ms/frame in the anchored continuous configuration | pending |
|