kernels/rr: recover a denoiser layer's phase table from its PTX
The eleven cuda_dldn_engine_swin_* layers have to be written natively before the denoiser can come close to 5 ms, and that needs their layout: which maximal runs of mma exist, what tile each covers, and which bytes of the 54 KB weight image it reads. rr_layer_spec.py derives that from the extracted PTX and nothing else - no GPU, no dependencies beyond the standard library. A phase is one maximal run of adjacent mma statements; M, N and K come from the number of distinct A and B fragments and the length of the accumulator chains partitioned by their D-to-C dataflow, rather than from a guess at the geometry; and each phase reports the weight and bias byte offsets it reads and the size of the epilogue that follows. Statement ranges carry their physical PTX line, so every claim can be checked with sed. enc0 comes out as 22 phases: two stem projections (M96 N32 K32), the fused Q/V projection (M96 N32 K32 x2), the score phase (M64 N96 K32, bias 8384..20172 - a 24-tile 16x16-per-head table), the banded P*V phase (M192 N96 K96, no weight operand at all), the output projection, then the second window's copy of the same pair, then the MLP phases and the M64 N256 K128 merge. The weight extents it reports agree with the harness dump to within a 12-byte unpadded tail (54548 against 54560 for enc0), and the widest phase gives each layer's patch-merge or patch-expand width.
This commit is contained in:
1 parent
4d45dda4a8
commit
f87df1cbc5
2 files changed
+906
No files matched your search
@@ -91,6 +91,8 @@ Measured on the RX 9070 XT (gfx1201, RDNA4), Ray Reconstruction preset E, 1280×
|
||||
|
||||
The denoiser's Swin block is the same block as presets L/M's: `cuda_dldn_engine_swin_enc0_kernel` computes position-only attention `S = bias + Q*Q^T` with no K projection, a 16x16 f16 bias table at 512 bytes per head, and one fused Q/V projection - the score `mma` at line 8406 takes both of its operands from the same accumulator registers, `%r1900`/`%r1920`, through their `cvt.rn.satfinite.e4m3x2.f16x2` at 8126/8129 and 8294/8297. That is `kernels/m/swin_block.h`'s block (`T16 S = bias; k32(S, a0, a0, a1, a1);`), not `kernels/k`'s learned-attention one. So the native denoiser layers can be built by instantiating that template - a prep that permutes the denoiser's own weight image into its slot layout, a parameter struct for the 144-byte block, and the surface stores through `tex_common.h` - rather than by writing a Swin block from scratch.
|
||||
|
||||
`kernels/rr/rr_layer_spec.py` recovers a layer's structure from the corpus PTX alone: `rr_layer_spec.py enc0` prints the phase table (each maximal run of adjacent `mma`, with M/N/K derived from the fragment counts and the accumulator dataflow, the weight and bias byte offsets it reads, and the size of the epilogue between phases) and the stage classification; `--geometry` does the same for all eleven layers and prints each one's weight image extent. It is how enc0's layout above was established, and it is the reference the native layers are written against.
|
||||
|
||||
The eleven Swin layers are 30.2 of those 31.6 ms. They are instruction-bound, not MMA-bound: the compiled `enc0` kernel is 83 373 instructions per thread for 536 FP8 WMMAs (0.6% of the stream), and the frame reaches about 6% of the card's FP8 WMMA peak. The stream is 13 k operand-gather shuffles (`ds_bpermute`/`v_perm`/`permlane`), 15 k packed 16-bit ops from the software `cvt.rn.satfinite.e4m3x2.f16x2` requantization, 9 k selects and 13 k ALU-dependency waits. `D4R_ZLUDA_IMPLICIT_MAX_BLOCK` (256 in the shipped config), `D4R_ZLUDA_SCHED_STRATEGY` and a 128-thread bound move the total by under 3%; a 128-thread bound additionally fails the launches of kernels that use larger blocks, and the whole-module wave64 switch produces no launches at all. Removing the gather and the codec — what a native layer does — is the remaining work.
|
||||
|
||||
## RDNA4
|
||||
|
||||
Reference in new issue
Block a user