d4r/docs/rr-engine.md

The swin-layer engine: design

Why this exists: a per-phase port of NVIDIA's cuda_dldn_engine_swin_* keeps the cost that makes those kernels slow. Disassembled, the translated swin_enc0 kernel is 83,373 instructions for 536 FP8 WMMAs. The top of the mix is operand assembly, not arithmetic:

instruction	count
v_cndmask_b32_e64	9,124
ds_bpermute_b32	7,976
v_lshrrev_b32_e32 / v_pk_lshrrev_b16	9,695
v_and_b32 / v_perm_b32 / v_pk_add_u16 / v_pk_min_u16	13,750
v_cvt_f16_f32 / v_cvt_f32_f16 (all forms)	3,568
v_wmma_f32_16x16x16_fp8_fp8	536

Those bit operations and cross-lane permutes exist only to reassemble NVIDIA's register fragments into WMMA operands. A kernel that keeps its activations in a staged e4m3 image and reads its operands directly from it does not pay them.

The primitives, already in the tree and already bit-identical to the translated path

kernels/tex/tex_common.h (with kernels/common/wmma_layout.h), used today by the working natives kernels/tex/enc0_tail.hip and kernels/tex/dec0_head.hip:

• kop_from16(u4v d) — this half's WMMA operand (8 e4m3 bytes) from 16 e4m3 bytes, in NVIDIA's k order. It is a half selection, not a shuffle.
• kop_image(const kslot* img, int idx) — the same from a prepared weight image, where a kslot is the 16 e4m3 bytes of one k16 operand row.
• kslot_from16(u4v d) — the weight image's own conversion, done once in a _prep kernel, so the main kernel never touches the weight layout.
• k32_e4m3(f8v c, kop a0, kop b0, kop a1, kop b1) — one m16n8k32 e4m3 step as two gfx12 FP8 WMMAs. With D4R_TEX_FP8 it is the native FP8 instruction with no scaling; the comment records that its result is bit-identical to ZLUDA's gfx12 lowering of NVIDIA's mma, which is what our reference runs through.

The one structural fact that makes this cheap: on gfx12 the accumulator layout is the next GEMM's operand layout (wmma_layout.h: "D VGPR i is row i + 8h, which is the gfx12 B-operand layout of the next GEMM (no exchange)"). So an epilogue requantising its accumulator to e4m3 and staging it as a row-major image in shared memory gives the next phase its operands with no cross-lane movement at all.

Shape of a layer

1. _prep kernel (once per launch, like enc0_tail's): read the layer's weight image at arg040 and write each phase's B operands as kslots into a global buffer.
2. Main kernel, per token patch:
   • load the layer's input (texture/surface args) and compute the first operand image;
   • for each phase in the descriptor: build A (kop_from16 from a staged image, or from the previous phase's requantised accumulator), run k32_e4m3 per (m, n) tile with the prepared B, apply the epilogue (e4m3 requantise into the next image, clamped cubic activation, f16 column-table add, RMS/LayerNorm, softmax, patch expand, plane scatter, lane merge, surface scatter);
   • write the layer's output.

Descriptor

Each layer is a kernels/rr/descriptor_<layer>.py with phases() (a list of dicts, one per phase: shapes, A source, B offset, C source, K tiling, epilogue kind and its operands) and layer() (launch geometry, argument offsets, token/tile geometry). Every offset must be a number with evidence behind it; anything unverified must say so.

Validation

By output surface against the translated reference: run the layer's native kernel alone (D4R_RRSWIN_ONLY=<layer> in its own native dir) and compare the layer's own output surfaces with kernels/tools/compare_launch_dump.py. Per-phase oracle comparison is a debugging aid for a phase that is already failing, not the acceptance gate. Note two traps that cost the agents hours: the debug payload is written into a window of a larger block (the containing argument and the window offset differ per layer, so read the argument the arena actually lives in, at arena_offset + base), and the harness reads its dump after a context synchronise, so a payload can be masked by later kernels in the frame.

Current gfx1201 evidence

The checked shared-memory primitive probe (rrswin_engine_k16_test.cpp) passes its GEMM, f16 accumulation rounding, requantization, weight preparation and phase-chain comparisons. These checks do not establish a complete RR implementation.

The phase-8 pilot now compares all 196,608 f16 results in 64 captured blocks without a mismatch. Its previous comparison was invalid: a six-token-tile phase overflowed a four-tile accumulator array, device output had no block stride, and the host block stride omitted the 32 lanes. Weight preparation now uses per-block shared memory rather than concurrent writes into one global buffer.

The host phase-8 model also needed all 32 lanes to assemble the activation image, an 80-register stride for the combined operand/result payload, and f16 output rounding. With those corrections, the documented fragment and weight maps match the captured oracle; the previous apparent operand-layout contradiction was not evidence of a wrong map.

An isolated native phase-8 splice preserves the other NVIDIA phases and produces a byte-identical final 3840×2160 RGBA16F frame with matching denormal settings. The initial splice is slower: enc0 measures 14.90 ms versus a 6.59 ms translated baseline. It is a correctness experiment, not a shipping performance improvement. The first implementation of patch 0010's native FP8 requantization measured 26.65 ms of GPU kernel time at 720p→4K after six warm-up frames. Its same-runtime software-codec control measured 33.90 ms. Compact integer saturation plus 0011's packed FP8 operand gathers measure 25.33 ms in that 20-frame run; the complete RGBA16F output remains byte-identical. The matching 80-frame control measures 24.764 ms after excluding the first 20 frames. Patch 0012's complete-pair shadow-store coalescing and eight-wave register budgeting together measure 19.146 ms, with all 80 saved frames byte-identical to that control. Patch 0013's signed native conversion reduces the same configuration to 16.234 ms with a byte-identical final frame; its plain and ReLU conversions match software for every half encoding in both packed positions. Removing the harness's 200 ms inter-frame idle gap measures 15.144 ms in steady state, also with the identical final image. The older installed-runtime FTZ output differs and is not a valid arithmetic control. Cyberpunk with 0010 + 0011 rendered a loaded 4K save with 720p RR input, coherent moving views and approximately 26.0 ms overlay upscaler time. With 0012 + 0013 the harness measures 16.234 ms (200 ms idle gap between frames) and 15.144 ms continuous, with byte-identical final frames.

Next direct-converter win: B operands from transposed accumulators

Validated mapping (read-only IR investigation, not yet implemented) for gfx12 native-FP8 MMAs whose A operand is a requantization of an earlier accumulator pair — the encoder's relu and plain chains in /tmp/d4r-packed-enc0-analysis/linked.ll:

• Producer: an MMA pair produces four <2 x i32> transposed accumulators (q=0/1, p=1/2); today each is converted with d4r_wmma_t_to_nv (4 bpermutes), stored to registers, requantized by d4r_cvt_rn_satfinite[_relu]_e4m3x2_f16x2, repacked into four words (%11374..%11377 for the relu chain), loaded again by every consumer and converted by d4r_wmma12f8_b_from_a (8 bpermutes per consumer). ≥60 consumer sites in enc0, ≥3 fragment groups.
• The composition t_to_nv → cvt → pack → b_from_a collapses algebraically: with h = lane/16, r = lane&7, t = 2h, the WMMA B word j of a lane is pack(cvt(f16x2{lo16(gather(A0, T_j)), lo16(gather(A1, T_j))}), ...), where A0/A1 = (r, r+8) and register .x for lanes < 16, (r+8, l) and .y otherwise. One helper per producer fragment — d4r_wmma12f8_b_from_t(uint2 d00, uint2 d01, uint2 d10, uint2 d11) — replaces 32 producer bpermutes + the alloca round trip + 8 bpermutes per consumer with 8 bpermutes total, all consumers CSE to one call.
• Slot order is chain-specific (out0/out1 ← q=0; out2/out3 ← q=1; low/high 16-bit half ← p=1/p=2); a matcher must read it off the actual insertelement/bitcast structure, and reject fragments with mixed cvt variants or partial updates (F32-shadow path needs its own derivation and is not covered).
• The same investigation found no sites where packed-f16 arithmetic could commute across toNv (the inserted values come from global loads, not shadow conversions), so that rewrite is not worth building.

Full derivation with anchors is in the session transcript (agent://Fp8ShadowOperandFusion); the probe extension to gate it (producer → requantize → consumer chain in rrswin_fp8_unpaired.ptx) is described there.

How a layer feeds the engine (design finding, not yet implemented)

The engine's phases read their activations from a staged image — rows of consecutive e4m3 k — not from NVIDIA's mma fragments. Two ways to get that image exist, and the cheap one is not the obvious one:

1. Rewrite the layer's prologue natively (load the input planes, the norm, the patch geometry) and have it stage the image directly. Correct, but it means re-deriving the layer's whole input path, which is where most of the per-layer reverse engineering sits.
2. Cut NVIDIA's PTX after its own loading and norming, exactly as the working bodies do today, and convert the fragments it hands over into the image once per phase boundary instead of per mma. Each lane's fragment bytes are documented (row and k per byte, in the body's header); writing them to shared memory at their (row, k) positions is a shared-memory scatter, which is cheap, whereas ZLUDA's per-mma reassembly is what the 7,976 ds_bpermute_b32 and 9,124 v_cndmask_b32_e64 in the translated kernel are paying for.

Route 2 reuses the cut machinery and the ABI the existing bodies already have, and it is the reason the engine does not need a prologue rewrite to be useful. What it does need, per layer, is: the fragment-to-image staging at the cut boundary, the descriptor (offset/layout per phase), the epilogue vocabulary (the engine has the clamped cubic activation; norm, softmax, column-table add, patch expand, plane scatter, lane merge and surface scatter are still to come), and the image-to-fragment conversion back at the exit, where the surviving PTX expects specific registers.

Rejected direct-converter route

The removed swin_gemm_zluda experiment passed widened NVIDIA A/B fragments straight into ZLUDA's native MMA helpers. Those types matched the helper ABI, but the data did not: ZLUDA first forms the WMMA A operand from NVIDIA B through t_a_from_b, and the WMMA B operand from NVIDIA A through t_b_from_a. Skipping those cross-lane gathers is not bit-exact by construction.

The staged-image route uses independently checked fragment maps and the existing kop_image / swin_operand / k32_e4m3 primitives instead. The isolated phase-8 recipe in swin_enc0_recipe.py preserves the surrounding NVIDIA code. It stages one token tile and the already-prefetched B fragments in shared memory, computes the two weight tiles, and hands the rounded result back through lane-major shared slots. Enc0's 3200-byte input_tensor is dead after prologue line 3938 and can be reused; the final surface epilogue still needs the separate input_noise array.

The generic helper model and phase-8 oracle comparison establish these mappings, not the unvalidated full-layer descriptors or the scalar full-body experiments.