Files
d4r/kernels/rr/rrswin_enc0_epilogues.md
woutermans 0fe3a28cb7 kernels/rr: staged-image Swin layer engine prototypes
Per-layer staged-image engine for the eleven translated denoiser Swin
layers: descriptors (verified offsets per phase), cut recipes that splice an
isolated native phase into NVIDIA's PTX, wave templates, epilogue notes and
GPU probes. The isolated phase-8 splice produces a byte-identical final
4K frame but is slower than the translated body (14.90 vs 6.59 ms) - a
correctness experiment, not a shipping improvement; docs/rr-engine.md
records the design and the next direct-converter win (B operands built
from transposed accumulators).
2026-10-12 12:47:53 +02:00

57 KiB
Raw Permalink Blame History

cuda_dldn_engine_swin_enc0_kernel — the epilogues of phases 5…22

Companion to kernels/rr/rr_layer_spec.py enc0. Everything below was read out of ~/.cache/d4r-rr-corpus/0013-PREPASS_ENTRYPOINT_NAME.ptx, entry cuda_dldn_engine_swin_enc0_kernel.

0. Conventions

  • sNNNN is a statement index in the linearisation rr_layer_spec.py documents (statements of the entry body; a wrapped { … } block and a 4-line mma.sync count as one statement, numbered at their first line). LNNNNN is the physical line of the same instruction in the corpus file, so every claim can be checked with sed -n 'LNNNNN,LNNNNN p' …ptx.
  • An epilogue is the run of statements between one phase's last mma and the next phase's first mma. The phase table (rr_layer_spec.py enc0) gives the ranges; this file names what each one computes.
  • W denotes the prepared weight image: ld.param.u64 %rd367, [cuda_dldn_engine_swin_enc0_kernel_param_0+40] (s2221) and cvta.to.global.u64 %rd2, %rd1 (s12, %rd1 = the same +40 load at s11). Every ld.weak.global.ca.v4.u32 in an epilogue reads W + 16*laneid + imm. S denotes the plane arena: ld.param.u64 %rd276, [%rd368+48] (s16052).
  • The C/D register layout and the weight-addressing formula are fixed (see the settled notes): C/D reg0 = (row g, cols 2t,2t+1), reg1 = (row g+8, same cols), g = laneid>>2, t = laneid&3.
  • Anything I could not pin down from the PTX is in §6 under uncharacterised.

1. Shared vocabulary

1.1 Addresses and parameters an epilogue touches

symbol how it is obtained statements
W (weight image, 54560 bytes) ld.param.u64 %rd367, [param_0+40] / cvta.to.global.u64 %rd2, %rd1 s2221 / s11
S (plane arena) ld.param.u64 %rd276, [%rd368+48] then cvta.to.global.u64 %rd37, %rd276 s16052–s16053
%rd368 mov.b64 %rd368, cuda_dldn_engine_swin_enc0_kernel_param_0 s2220
extents (ex, ey) ld.param.v2.u32 {%r364, %r365}, [param_0+0] s13
%r1, %r6 (block origins) sub.s32 %r1, %r359, %r360 and sub.s32 %r6, %r367, %r366 s6, s16
per-lane weight address mul.wide.u32 %rd, %laneid, 16 + add.s64 %rd, W, %rd every load block
per-lane table address shl.b32 %r, %laneid, 2 → cvt.u64.u32 → and.b64 %rd, %r, 12 → add.s64 %rd, %rd2, %rd = W + ((laneid*4) & 12) s7881–s7885 and copies

The last line matters: the four ld.global.v2.u16 / ld.global.u32 loads in an epilogue read W + ((laneid<<2)&12) + off, i.e. the same 64-byte record is read by the four lanes with equal laneid&3, and the four loads pick four 16-byte words inside it. So those tables are indexed by column, not by row.

1.2 Constants that survive between epilogues

register first definition f32/f64 literal value used (after the cvt)
%f2085 s3984 mov.f32 %f2085, 0f3C8CCB50 0.017186790704727173 f16 0x2466 = 0.017181396484375
%fd767 s3988 mov.f64 %fd767, 0dBFE1CC0000000000 −0.55615234375 f16 0xB873 = −0.55615234375
%fd769 s3993 mov.f64 %fd769, 0d3FE1CC0000000000 +0.55615234375 f16 0x3873 = +0.55615234375
%f2087 s3999 mov.f32 %f2087, 0f3F6D6000 0.92724609375 f16 0x3B6B = 0.92724609375
%f2089 s4004 mov.f32 %f2089, 0f3FB00000 1.375 f16 0x3D80 = 1.375
%fd770 s2802 mov.f64 %fd770, 0d3F20000000000000 2⁻¹³ f16 0x0800 = 0.0001220703125
%f2598 s13192 mov.f32 %f2598, 0f3ED306EB 0.4121621549129486 f16 0x3698 = 0.412109375
%f2599 s13195 mov.f32 %f2599, 0f3DA60DD6 0.0810810774564743 f16 0x2D30 = 0.0810546875
%f2600 s13198 mov.f32 %f2600, 0f3F000000 0.5 f16 0x3800 = 0.5
%f2601 s13201 mov.f32 %f2601, 0f40000000 2.0 f16 0x4000 = 2.0

Also live across the region (defined in the prologue, read by every warp reduction): %r11301 = mov.u32 %r11301, 1 (s2286), %r11310 = mov.u32 %r11310, 2 (s2225), %r11312 = mov.u32 %r11312, -1 (s2621) — the three operands of every shfl.sync.bfly.b32.

%f2598…%f2601 are defined inside phase 13's epilogue and reused by the epilogues of phases 15, 17 and 19 (e.g. s14001 and s14693 read them).

1.3 Two quantisation helpers

  • cvt.rn.satfinite.e4m3x2.f16x2 — the requantisation used everywhere (f16x2_to_e4m3x2_satfinite_bits in ptx/lib/zluda_ptx_impl.cpp).
  • rcp.approx.ftz.f32, rsqrt.approx.ftz.f32, ex2.approx.ftz.f32, div.approx.ftz.f32 — always wrapped as cvt.f32.f16 → op → cvt.rn.f16.f32 (round each step through f16).

2. Summary

epilogue range stmts named operation
E5 s3983–s7825 3843 softmax of the 64×96 score block, V-transpose, pack
E6 s7874–s8033 160 residual + per-column table add, pack, next B
E7 s8050–s8056 7 next-B weight loads
E8 s8081–s8087 7 next-B weight loads
E9 s8112–s8282 171 score-bias C-seed loads, pack
E10 s8331–s12168 3838 softmax (2nd window, same code as E5)
E11 s12217–s12271 55 next-B weight loads, pack
E12 s12288–s13174 887 RMS norm (², sum, rsqrt, gain) + packs + next C
E13 s13191–s13899 709 clamped cubic activation + column bias + pack
E14 s13916–s13979 64 next B/C loads, pack
E15 s13996–s14595 600 clamped cubic activation + pack
E16 s14612–s14675 64 next B/C loads, pack
E17 s14692–s15291 600 clamped cubic activation + pack
E18 s15308–s15371 64 next B/C loads, pack
E19 s15388–s15987 600 clamped cubic activation + pack
E20 s16004–s16417 414 pack → plane-arena scatter → 8-lane merge → next B/C
E21 s16450–s16606 157 pack → surface scatter → %tid.x → next B/C
E22 s16615–s17222 608 cross-lane row max, softmax, shared-patch gather, store

3. The epilogues

E5 — s3983–s7825: softmax of the 64×96 score block

Phase 5 = 48 mma (s3935–s3982), D registers %r2566…%r3037 (96 f16x2 = 64 rows × 96 columns held 8 rows × 8 cols per lane). The epilogue has no memory access at all (0 ld.*, 0 st.*); it also reads phase 4's D (%r2149…%r2380, 48 registers) and writes 77 registers that survive: the 5 constants of §1.2, phase 6's 12 A fragments (%r5884 %r5885 %r5886 %r5887, %r5904…, … %r6284…%r6287) and its 12 B fragments (%r6178 %r6179 %r6188 %r6189 … %r6288 %r6289).

Sub-sections:

range stmts what
s3983–s6100 2118 96 per-element groups: scale, clamp, cubic, exponent extract
s6101–s6180 80 partial row-sum trees (in-lane)
s6181–s6312 132 16 shfl.sync.bfly → 8 row sums
s6313–s7081 769 192 rcp.approx.ftz.f32 → 96 reciprocals
s7082–s7561 480 192 mul.f16 → 96 probabilities
s7562–s7609 48 movmatrix transpose of phase 4's D
s7610–s7825 216 144 cvt.rn.satfinite.e4m3x2.f16x2 packs

(a) per-element exponent, s3983–s6100 (96 groups of 22 statements)

Evidence:

s3984  L8744: mov.f32 %f2085, 0f3C8CCB50
s3987  L8752: mul.f16x2 %r3048,%r2566,%r3047
s3992  L8766: max.f16x2 %r3051,%r3048,%r3053
s3997  L8780: min.f16x2 %r3054,%r3051,%r3056
s4003  L8798: fma.rn.f16x2 %r3059,%r3054,%r3057,%r3062
s4008  L8812: fma.rn.f16x2 %r3063,%r3054,%r3059,%r3066
s4009  L8815: shl.b32 %r16813, %r3063, 5
s4010  L8816: and.b32 %r4968, %r16813, 2145419232      // 0x7FE07FE0

For each of the 96 f16x2 D registers, per half (an f16 value x):

m  = clamp(f16(x * 0.017181396484375), -0.55615234375, +0.55615234375)   // s3987,s3992,s3997
t  = f16(m*(-m) + 0.92724609375)                                         // s4003 fma
u  = f16(m*t    + 1.375)                                                 // s4008 fma
    = f16(1.375 + m*(0.92724609375 - m^2))
E  = bits(u)[9:5]      // floor(32*frac(u))       -- the top five mantissa bits
F  = bits(u)[4:0]      // the low five mantissa bits
expval = f16 from bits  (E << 10) | (F << 5)      // = 2^(E-15) * (1 + F/32)

i.e. ((u_bits << 5) & 0x7FE07FE0) keeps u's mantissa and drops its sign and exponent, so the mantissa lands in the exponent field. Numerically expval ≈ 2^(32*frac(u) - 15) = 2^(32u - 47) and, because the clamp keeps u ∈ [1.03125, 1.71875], expval ∈ [2⁻¹⁴, 2⁸]. This is not exp and not exp2(x*log2e): the exponent is the cubic 32*(0.375 + 0.92724609375*m − m³), quantised to a 6-bit significand (1.5 % relative granularity). The 6-bit truncation is why the constant 1.375 is there: frac(u) for u∈[1,2) is the mantissa.

The clamp bounds are the guard: d/dm [0.92724609375m − m³] = 0.92724609375 − 3m² vanishes at m = √(0.92724609375/3) = 0.55615…, exactly the clamp, so the exponent saturates at 32·(0.92724609375·0.55615 − 0.55615³) + 12 = 23 (expval = 2⁸) at the top and at E = 1 (2⁻¹⁴) at the bottom. The two f64 literals are used only so the clamp is the exact f16 value.

(b) in-lane partial sums, s6101–s6180

96 expval registers are summed in three groups of twelve: each group does 12 add.f16x2 then (a+b)+c four times, giving 4 f16x2 (= 8 f16) per group, so each lane reduces its 96 values (8 rows × 12 registers) to 24 partial sums in 12 f16x2 registers.

(c) cross-lane row sums, s6181–s6312

s6182 L15601: add.f16x2 %r5208,%r5015,%r5021
s6187 L15612: shfl.sync.bfly.b32 %r5212,%r5208,%r11301,%r5215,%r11312
s6188 L15616: add.f16x2 %r5217,%r5208,%r5212
s6193 L15627: shfl.sync.bfly.b32 %r5221,%r5217,%r11310,%r5224,%r11312
s6194 L15631: add.f16x2 %r5226,%r5217,%r5221
s6196 L15636: add.f16 %rs2318,%rs2319,%rs2320
s6197 L15639: mov.b32 %r5423, {%rs2318, %rs2318}

Eight times (one per row g, g+8 of the four m-tiles): add the two f16x2 partials, butterfly over lane bits 0 and 1 (%r11301 = 1, %r11310 = 2), add the two halves of the result with add.f16, then broadcast into both halves. Result: 8 row sums %r5423 %r5427 %r5471 %r5475 %r5519 %r5523 %r5567 %r5571 (rows g, g+8 of m-tiles 0…3), covering all 96 columns of each row because the four lanes with equal laneid&7 hold different column pairs.

There is no row maximum anywhere in this epilogue: max.f16x2 (96×) and min.f16x2 (96×) are the clamp of (a); the only warp shuffles are the 16 shfl.sync.bfly above. See §6.

(d) reciprocal, s6313–s7081

192 rcp.approx.ftz.f32 = 96 f16x2 reciprocals, each computed as cvt.f32.f16 → rcp.approx.ftz.f32 → cvt.rn.f16.f32 on a row-sum register (duplicated, so each row sum is inverted four times — once per n-tile of that row):

s6314 L15935: mov.b32 {hl, hu}, %r5423
s6317 L15938: rcp.approx.ftz.f32 fl, fl
s6319 L15940: cvt.rn.f16.f32 hl, fl
s6321 L15942: mov.b32 %r5380, {hl, hu}

(e) probabilities, s7082–s7561

s7082 L17180: mov.b32 {%rs2343, %rs2346}, %r4968     // expval (row g)
s7083 L17181: mov.b32 {%rs2344, %rs2347}, %r5380     // 1/rowsum (row g)
s7084 L17183: mul.f16 %rs2345,%rs2346,%rs2347
s7085 L17187: mul.f16 %rs2342,%rs2343,%rs2344
s7086 L17190: mov.b32 %r5716, {%rs2342, %rs2345}

192 mul.f16 → 96 f16x2 results p = expval * (1/rowsum) — the row sums computed in (b)/(c) cancel the constant 2⁻⁴⁷ of (a), so p_ij ∝ 2^(32·u_ij), and Σ_j p_ij = 1 up to the rcp.approx error.

(f) V-transpose, s7562–s7609

s7562 L18237: movmatrix.sync.trans.aligned.m8n8.b16 %r5572, %r2149
s7563 L18240: movmatrix.sync.trans.aligned.m8n8.b16 %r5574, %r2150

48 movmatrix on 48 distinct phase-4 D registers (%r2149…%r2380, each exactly once). Each instruction transposes one 8×8 b16 tile built from the 32 threads' %r, i.e. 48 × 64 = 3072 values = the whole 96×32 D tile. (The epilogue therefore contains 48, not 96, movmatrix.)

(g) packs, s7610–s7825

s7611 L18384: cvt.rn.satfinite.e4m3x2.f16x2 %rs2918, %r5572
s7612 L18386: mov.b32 %r6178, {%rs2918, %rs2919}
s7683 L18552: cvt.rn.satfinite.e4m3x2.f16x2 %rs2966, %r5716
s7684 L18554: mov.b32 %r5884, {%rs2966, %rs2967}

144 packs in three groups:

  1. s7610–s7681 — from the movmatrix results → phase 6's 12 B fragments (%r6178…%r6289, 24 b32), i.e. the transposed V.
  2. s7682–s7799 — from the probabilities → phase 6's 12 A fragments (%r5884…%r6287, 48 b32). Each fragment consumes 8 consecutive products; the four halves are re-ordered (%r5884 ← (p0,p1), %r5886 ← (p2,p3), %r5885 ← (p4,p5), %r5887 ← (p6,p7)) to match the e4m3 fragment byte order.
  3. s7800–s7825 — more probability packs whose destinations are not read after s7825 (compiler spill; harmless).

Inputs: phase 5's 96 D f16x2, phase 4's 48 D f16x2, %r11301/%r11310/%r11312. Outputs: the 5 constants + phase 6's 48 A-register / 24 B-register operands. Memory read/written: none.


E6 — s7874–s8033: residual + per-column table, then pack

160 statements. Evidence:

s7877 L19225: add.s64 %rd161, %rd229, 20672
s7878 L19227: ld.weak.global.ca.v4.u32 { %r6293,%r6294,%r6295,%r6296},[%rd161]
s7880 L19231: ld.weak.global.ca.v4.u32 { %r6297,%r6298,%r6299,%r6300},[%rd162]   // +21184
s7886 L19240: ld.global.v2.u16 {%rs6191, %rs6192}, [%rd232+37072]
s7891 L19252: {add.f16 %rs3062,%rs6197,%rs3064;}
s7892 L19254: mov.b32 %r6342, {%rs3062, %rs3065}
s7987 L19536: cvt.rn.satfinite.e4m3x2.f16x2 %rs3254, %r5852

Four things happen, in this order:

  1. Next-B loads (s7874–s7880): two v4.u32 fragments from W + 16*laneid + {20672, 21184} → %r6293…%r6300, phase 7's B operand.

  2. Table loads (s7881–s7889): four ld.global.v2.u16 from W + ((laneid<<2)&12) + {37056, 37072, 37088, 37104} → 8 f16 (%rs6191…%rs6198). 16 bytes/lane × 32 lanes × 4 words = 512 bytes: a per-column vector covering 32 columns (4 words × 8 columns, each lane holding the two columns 2t,2t+1 of two of the four 8-column groups).

  3. Residual + column table add (s7890–s7985): 64 add.f16, i.e. 32 f16x2, computed in place on the f16 halves of phase 2's D (%r1174…%r1325, the halves unpacked at s2399–s2555 by mov.b32 {%rs3064, %rs3067}, %r1175 and its 38 siblings):

    new_low  = table[2t]   + D2_half_low        s7891
    new_high = table[2t+1] + D2_half_high       s7890
    → packed into %r6342, %r6343, %r6352, … %r6493   (32 registers)
    

    These 32 registers are phase 7's C operand, so phase 7 computes D7 = A7·B7 + (D2 + column_table) — an out-projection plus the residual.

  4. Pack (s7986–s8033): 32 packs of phase 6's D (%r5852, %r5862, %r5912, %r5922, %r5972, … — 32 of the 48 D registers, i.e. the 64×32 output) → 16 b32 = phase 7's A operand (%r6366 %r6367 %r6368 %r6369, %r6406…, %r6446…, %r6486…).

Outputs (50 live): %r6293 %r6297 (B), the 32 C seeds %r6342…%r6493, the 16 packed A registers. Memory: 2 weight loads + 4 table reads.


E7 — s8050–s8056: next-B loads

s8053 L19761: add.s64 %rd163, %rd234, 21696
s8054 L19763: ld.weak.global.ca.v4.u32 { %r6495,%r6496,%r6497,%r6498},[%rd163]
s8055 L19765: add.s64 %rd164, %rd234, 22208
s8056 L19767: ld.weak.global.ca.v4.u32 { %r6499,%r6500,%r6501,%r6502},[%rd164]

Two 16-byte/lane fragments of W + 16*laneid + {21696, 22208} = phase 8's B operand (weight block [21696, 22720) per the phase table). Nothing else.

E8 — s8081–s8087: next-B loads

Same shape, offsets +22720 and +23232:

s8084 L19942: add.s64 %rd165, %rd236, 22720
s8085 L19944: ld.weak.global.ca.v4.u32 { %r6744,%r6745,%r6746,%r6747},[%rd165]
s8086 L19946: add.s64 %rd166, %rd236, 23232
s8087 L19948: ld.weak.global.ca.v4.u32 { %r6748,%r6749,%r6750,%r6751},[%rd166]

→ phase 9's B operand, weight block [22720, 23744).

E9 — s8112–s8282: score-bias C-seed loads + pack of phase 8's D

171 statements, 24 loads + 80 packs:

s8116 L20125: ld.weak.global.ca.v4.u32 { %r6993,…%r6996},[%rd167]   // +23744
   …
s8162 L20217: ld.weak.global.ca.v4.u32 { %r7085,…%r7088},[%rd190]   // +23744+23*512 = +35520
  • 24 loads at W + 16*laneid + 23744 + 512*i, i = 0…23 — the phase-10 score-bias seeds, byte block [23744, 36032) = 12288 bytes = 6144 f16. This is the twin of phase 5's C-seed block ([8384, 20672)), loaded the same way in phase 4's epilogue (s3764–s3814).

  • 80 packs / 40 b32, whose sources are all registers written by phase 8's mma (%r6503…%r6744, the D pairs of s8057–s8080) — phase 9's D (%r6752…%r6983) is not read here; it is consumed by E10's movmatrix block instead (s11905: movmatrix.sync.trans.aligned.m8n8.b16 %r10175, %r6752). The 40 packed b32 split into

    • phase 10's four A-fragment sets: %r7281 %r7282 %r7283 %r7284, %r7401…, %r7521…, %r7641… (16 b32, e.g. s8237: mov.b32 %r7281, {%rs3334, %rs3335} from cvt of %r6503/%r6523 at s8235–s8236), and
    • phase 10's twelve B-fragment pairs %r7535 %r7536 … %r7645 %r7646 (24 b32, e.g. s8163–s8165: cvt …%r6523; cvt …%r6503; mov.b32 %r7535, {%rs3286, %rs3287}).

    The first A set and the first B set are bit-identical: both are {pack(%r6503), pack(%r6523)} (%r7281 from s8235–s8237, %r7535 from s8163–s8165). Whatever the score GEMM's two operands are, phase 10 gets one of them twice at fragment 0.

The C-seed layout (identical for phase 5 and phase 10): the four C registers of one mma are the four u32 of one 16-byte lane slot, and

C(m, n) = load_tile(4*(n>>1) + m) + 4*(n&1)      [register index within the block]
byte    = base + 512*(4*(n>>1) + m) + 16*laneid + 8*(n&1)

so the bias tile index runs 4*(n>>1) + m (stride 512 bytes) over 4 m-tiles and 6 (n>>1) groups = 24 tiles = 12288 bytes. Verified from phase 5's own mma operands: s3935 C=%r2390, s3936 C=%r2392, s3937 C=%r2406, s3947 C=%r2394, s3959 C=%r2398, s3971 C=%r2402, s3982 C=%r2484.


E10 — s8331–s12168: softmax, second window

Identical code to E5 with shifted registers; the op histogram is byte-for-byte the same (1048 mov.b32, 672 cvt.rn.f16.f32, 384 cvt.rn.f16.f64, 192 each fma.rn.f16x2 / cvt.f32.f16 / rcp.approx.ftz.f32 / mul.f16, 144 packs, 48 movmatrix). It re-uses E5's constants (%f2085 %fd767 %fd769 %f2087 %f2089, defined at s3984–s4004) and E5's shfl operands %r11301 %r11310 %r11312, so it contains no mov.f32/mov.f64 of its own.

Sub-sections:

range what
s8331–s10443 96 exponent groups (scale/clamp/cubic/and.b32 0x7FE07FE0)
s10444–s10523 80 in-lane partial-sum add.f16x2
s10524–s10655 16 shfl.sync.bfly → 8 row sums (s10530 L27701, s10651 L27990)
s10656–s11424 192 rcp.approx.ftz.f32 (s10660 L28027)
s11425–s11904 192 mul.f16 → 96 probabilities
s11905–s11952 48 movmatrix (s11905 L30326: movmatrix… %r10175, %r6752)
s11953–s12168 144 packs (s11953 L30470)

Inputs: phase 10's 96 D (%r7169…), phase 9's 48 D. Outputs (72 live): phase 11's 12 A fragments (%r10487…%r10892) + 12 B fragments (%r10781…)

  • the 4 constants. Memory: none.

E11 — s12217–s12271: next-B loads + pack

55 statements:

s12220 L31314: add.s64 %rd191, %rd240, 36032
s12221 L31316: ld.weak.global.ca.v4.u32 { %r10896,…%r10899},[%rd191]
s12222 L31318: add.s64 %rd192, %rd240, 36544
s12223 L31320: ld.weak.global.ca.v4.u32 { %r10900,…%r10903},[%rd192]
s12224 L31323: cvt.rn.satfinite.e4m3x2.f16x2 %rs4879, %r10515
s12226 L31328: mov.b32 %r10968, {%rs4878, %rs4879}
  • two weight fragments at W + 16*laneid + {36032, 36544} → phase 12's B;
  • 32 packs of phase 11's D (%r10455…%r10645) → 16 b32 = phase 12's A (%r10968 %r10969 %r10970 %r10971, %r11008…, %r11048…, %r11088…).

Unlike E6, no table: phase 11's output is requantised only, because the residual for the second attention branch is folded into phase 12's C operand (%r6334 %r6335 = phase 7's D, see the phase table).

Shape of the attention region (E5–E12)

P3  A = prologue tile, B = W+6336  C = 0        → D_K   (96×32)
P4  A = prologue tile, B = W+7360  C = 0        → D_QV  (96×32)   epilogue E4 packs both
     E4 packs D_K  → P5 B  (s3817: %r2932),     E4 loads P5 C = bias@8384 (s3764–s3814)
P5  A = pack(D_QV)   B = pack(D_K)  C = bias@8384    → scores 64×96
     E5: softmax over the 96 columns (no row max), 48 movmatrix(D_QV),
         pack → P6 A (48 b32) + P6 B (24 b32)
P6  A = pack(P)      B = pack(transpose(D_QV))  C = 0 → D_PV (64×32)
     E6: pack(PV) → P7 A (16 b32);  P7 C = D_QV2 + column table@37056;  loads P7 B
P7  C = residual       → D7                      epilogue E7 loads P8 B
P8  A = prologue tile, B = W+21696, C = 0  → D8    epilogue E9 packs D8 → P10 A and B
P9  A = prologue tile, B = W+22720, C = 0  → D9    (D9 only reaches E10's movmatrix)
P10 A = pack(D8)   B = pack(D8)   C = bias@23744 → scores 64×96
      E10: same softmax + transpose(D9) + packs → P11 A/B
P11 … → E11 packs → P12 A;  P12 C = D7 (residual)

Window 2 therefore mirrors window 1 except that phase 10's two operands both come out of phase 8 (see E9), while the transpose of phase 9's D supplies phase 11.


E12 — s12288–s13174: RMS norm before the MLP

887 statements, no stores, one ld.global.v2.u16 group, two weight groups. Structure:

range stmts op
s12288–s12415 128 32 × x·x squares of phase 12's D
s12416–s12431 16 in-lane pairwise add.f16x2
s12432–s12555 124 8 rows × (add + 2 shfl.sync.bfly + add + add.f16)
s12556–s12652 97 + eps, broadcast into 16 b32
s12653–s12909 257 32 × rsqrt.approx.ftz.f32
s12910–s12918 9 gain loads +37120/37136/37152/37168
s12919–s13014 96 inv · gain
s13015–s13110 96 x · (inv·gain)
s13111–s13117 7 next-B loads +37184/+37696
s13118–s13126 9 next-C loads +38208/38224/38240/38256
s13127–s13174 48 32 packs → phase 13's A

Evidence:

s12289 L31548: mul.f16 %rs4913,%rs5740,%rs5740          // square of a D half
s12438 L31945: shfl.sync.bfly.b32 %r11149,%r11145,%r11301,%r11152,%r11312
s12556 L32255: cvt.rn.f16.f64 %rs5126, %fd770           // eps = 2^-13
s12657 L32555: rsqrt.approx.ftz.f32 fl, fl
s12915 L32972: ld.global.v2.u16 {%rs6199, %rs6200}, [%rd243+37136]
s13015 L33265: mul.f16 %rs5514,%rs5740,%rs5322          // D_half * (inv*gain)
s13114 L33557: add.s64 %rd193, %rd245, 37184
s13123 L33572: ld.global.u32 %r11553, [%rd248+38208]
s13128 L33580: cvt.rn.satfinite.e4m3x2.f16x2 %rs5703, %r15389

What it computes. For each of the 8 rows a lane holds (8 columns each):

sumsq[row] = Σ_{c=0..31} x[row,c]^2                      (f16, 4-lane butterfly)
inv        = rsqrt.approx.ftz.f32(f32(f16(sumsq + 2^-13)))    // rounded to f16
out[row,c] = x[row,c] * inv * gain[c]
  • no mean subtraction and no 1/N: the quantity is the raw sum of the 32 squares in a row plus the constant 2^-13; the mean is not formed.
  • gain[c] is read from W + ((laneid<<2)&12) + {37120, 37136, 37152, 37168} (8 f16 per lane = 32 columns), so the gain is per column, constant across rows.
  • the inv register is duplicated into both halves of a b32 before use (mov.b32 %r11318, {%rs5127, %rs5130} at s12559), so one row-sum value serves all four n-tiles of that row (32 rsqrt for 8 row sums).
  • the norm's bias is not added here — it is folded into phase 13's C operand, loaded at the end of this epilogue (+38208…+38256).

Outputs (22): %r11383 %r11387 (phase 13 B, from +37184/+37696), %r11553 %r11563 %r11573 %r11583 (phase 13 C), 16 packed A registers (%r11456 %r11457 %r11458 %r11459, %r11496…, %r11536…%r11539), the 32 f16x2 normalised values %r15389…%r15420 (re-packed in E14), and the 32 half-registers %rs5737 %rs5743 … %rs5923 still read by E13's bias add.


E13 — s13191–s13899: clamped cubic activation, column bias, pack

709 statements. This is the only epilogue that materialises %f2598…%f2601 (the four activation constants of §1.2).

s13192 L33803: mov.f32 %f2598, 0f3ED306EB
s13198 L33817: mov.f32 %f2600, 0f3F000000          // 0.5
s13201 L33824: mov.f32 %f2601, 0f40000000          // 2.0
s13204 L33832: neg.f16x2 %r11589,%r11588
s13205 L33836: max.f16x2 %r11591,%r11424,%r11589
s13206 L33840: min.f16x2 %r11594,%r11591,%r11588
s13207 L33844: abs.f16x2 %r11597,%r11594
s13208 L33848: mul.f16x2 %r11599,%r11586,%r11597
s13209 L33852: sub.f16x2 %r11602,%r11585,%r11599
s13210 L33856: mul.f16x2 %r11605,%r11594,%r11602
s13211 L33860: add.f16x2 %r11608,%r11587,%r11605
s13212 L33864: mul.f16x2 %r11611,%r11424,%r11608

Per element, over each of the 32 f16x2 D registers of phase 13 (%r11424 %r11425 %r11504…, 64 f16 per lane):

y = clamp(x, -2, +2)
g = 0.5 + y * (0.412109375 - 0.0810546875 * |y|)      // f16 throughout
out = x * g

g is a cubic ramp, not a sigmoid: g(+2) = 1 ⇒ out = x for x ≥ 2; g(-2) = 0 ⇒ out = 0 for x ≤ -2; g(1) = 0.831055 where x·sigmoid(x) = 0.730957. Reproduce the formula, not SiLU/GELU.

Then:

range op
s13191–s13739 32 activation groups (each group 4 constants + 9 f16x2 ops)
s13740–s13746 next-B loads: +38272 → %r12514…%r12517, +38784 → %r12518…%r12521
s13747–s13755 lane-shift setup + ld.global.v2.u16 at W + ((laneid<<2)&12) + {39296, 39312, 39328, 39344}
s13756–s13851 64 add.f16 — table add, written over the activation inputs
s13852–s13899 32 packs of the activation outputs → phase 14's A

Evidence for the table add:

s13752 L35747: ld.global.v2.u16 {%rs6207, %rs6208}, [%rd253+39312]
s13756 L35752: {add.f16 %rs5738,%rs6214,%rs5740;}
s13757 L35756: {add.f16 %rs5735,%rs6213,%rs5737;}
s13758 L35759: mov.b32 %r12563, {%rs5735, %rs5738}
s13853 L36043: cvt.rn.satfinite.e4m3x2.f16x2 %rs5927, %r11611

%rs5737/%rs5740 are the halves of %r11424 (phase 13's D, split at s13191), so %r12563 = phase13_D + table[column]. The 32 such registers (%r12563 %r12564 %r12573 %r12574 … %r12713 %r12714) are phase 14's C operand; %r12514 %r12518 are phase 14's B fragments (weight image).

Outputs (54 live): %f2598…%f2601 (the constants — consumed again by E15, E17, E19), the 32 C seeds, the 16 packed A registers, %r12514 %r12518.


E14 — s13916–s13979: next B/C loads + pack

64 statements:

s13919 L36268: add.s64 %rd197, %rd255, 39360
s13920 L36270: ld.weak.global.ca.v4.u32 { %r12716,…%r12719},[%rd197]
s13921 L36272: add.s64 %rd198, %rd255, 39872
s13922 L36274: ld.weak.global.ca.v4.u32 { %r12720,…%r12723},[%rd198]
s13928 L36283: ld.global.u32 %r12886, [%rd258+40384]
s13929 L36284: ld.global.u32 %r12896, [%rd258+40400]
s13930 L36285: ld.global.u32 %r12906, [%rd258+40416]
s13931 L36286: ld.global.u32 %r12916, [%rd258+40432]
s13933 L36293: cvt.rn.satfinite.e4m3x2.f16x2 %rs5959, %r15389
s13934 L36295: mov.b32 %r12789, {%rs5959, %rs5960}
  • phase 15's B at W + 16*laneid + {39360, 39872};
  • phase 15's C/bias: 4 ld.global.u32 at W + ((laneid<<2)&12) + {40384, 40400, 40416, 40432} → %r12886 %r12896 %r12906 %r12916. Each of those four u32 is used in both halves of an mma C operand ({%r12886, %r12886} at s13980), i.e. a per-column bias replicated over the 16 rows of a tile and reused by all four m-tiles;
  • 32 packs of %r15389…%r15420 — the normalised values E12 produced, so phase 15's A operand is bit-identical to phase 13's A operand (phases 13 and 15 are two parallel GEMMs off the same normed tile).

E15 — s13996–s14595: clamped cubic activation, pack

600 statements, histogram: 96 mul.f16x2, 32 each neg/max/min/abs/sub/add f16x2, 128 cvt.rn.f16.f32, 144 mov.b32, 32 packs, 2 weight loads — and no constants of its own and no ld.global.

s13997 L36516: cvt.rn.f16.f32 low, %f2598            // constants from E13
s14005 L36539: neg.f16x2 %r12922,%r12921
s14006 L36543: max.f16x2 %r12924,%r12757,%r12922
s14009 L36555: mul.f16x2 %r12932,%r12919,%r12930
s14013 L36571: mul.f16x2 %r12944,%r12757,%r12941
s14545 L38441: ld.weak.global.ca.v4.u32 { %r13847,…%r13850},[%rd199]   // +40448
s14547 L38445: ld.weak.global.ca.v4.u32 { %r13851,…%r13854},[%rd200]   // +40960

Same function as E13: out = x·(0.5 + y·(0.412109375 − 0.0810546875|y|)), y = clamp(x,±2), over phase 15's 32 D registers (%r12757 %r12758 …), then 32 packs of the activation outputs (%r12944 at s14549, %r13060 at s14548, …) → phase 16's A (%r13919…%r14042), plus phase 17's B fragments from W + 16*laneid + {40448, 40960}. No column-bias add.

E16 — s14612–s14675: next B/C loads + pack

Identical shape to E14 with offsets +41472/+41984 (B) and +42496/42512/42528/42544 (C, 4 ld.global.u32 → %r14218 %r14228 %r14238 %r14248), followed by 32 packs of %r15389…%r15420 — again the normed tile E12 produced — giving phase 17's A (%r14121 %r14122 %r14123 %r14124, %r14161…, %r14201…).

Evidence: s14616 L38678: ld.weak.global…[%rd201], s14624 L38691: ld.global.u32 %r14218, [%rd265+42496], s14629 L38699: cvt.rn.satfinite.e4m3x2.f16x2 %rs6023, %r15389 (the pack source is E12's %r15389, defined at s13017).

E17 — s14692–s15291: clamped cubic activation, pack

600 statements; identical histogram to E15.

s14693 L38924: cvt.rn.f16.f32 low, %f2598
s14701 L38947: neg.f16x2 %r14254,%r14253
s14702 L38951: max.f16x2 %r14256,%r14089,%r14254      // phase 17's D
s14709 L38979: mul.f16x2 %r14276,%r14089,%r14273
s15241 L40849: ld.weak.global.ca.v4.u32 { %r15179,…},[%rd203]      // +42560
s15244: cvt.rn.satfinite.e4m3x2.f16x2 …

Activation on phase 17's D, 32 packs of those activation outputs (%r14276 at s15245, %r14392 at s15244, …) → phase 18's A, 2 weight loads at W + 16*laneid + {42560, 43072} → phase 18's B. No constants, no ld.global.

E18 — s15308–s15371: next B/C loads + pack

Same shape as E14/E16: B at +43584/+44096, C at +44608/44624/44640/44656 → %r15550 %r15560 %r15570 %r15580, then 32 packs of %r15389…%r15420 — the normed tile for the fourth time → phase 19's A (%r15453…%r15576). Evidence: s15312 L41086: ld.weak.global…, s15320 L41100: ld.global.u32 %r15550, [%rd272+44608], and s15325 L41107: cvt.rn.satfinite.e4m3x2.f16x2 %rs6087, %r15389 (source defined at s13017).

E19 — s15388–s15987: clamped cubic activation, pack

600 statements; identical histogram to E15/E17.

s15389 L41332: cvt.rn.f16.f32 low, %f2598
s15397 L41355: neg.f16x2 %r15586,%r15585
s15398 L41359: max.f16x2 %r15588,%r15421,%r15586      // phase 19's D
s15399 L41363: min.f16x2 %r15591,%r15588,%r15585
s15937 L43257: ld.weak.global.ca.v4.u32 { %r16511,…},[%rd207]      // +44672
s15940 L43270: cvt.rn.satfinite.e4m3x2.f16x2 …

Activation on phase 19's D, 32 packs of those activation outputs (%r15608 at s15941, %r15724 at s15940, …) → phase 20's A (%r16583…%r16706), 2 weight loads at W + 16*laneid + {44672, 45184} → phase 20's B.

Shape of the whole MLP region (E12–E20)

Putting the pack sources and the C operands together (all verified by tracing each mov.b32 %rX, {%rsA, %rsB} back through its two cvt statements):

E12  normed tile N (32 f16x2, %r15389…%r15420)
 ├─ pack → phase 13 A   (s13127)   phase 13: A=N, B=W13, C=bias@38208  → SiLU (E13)
 │                                              pack of SiLU → phase 14 A (s13852)
 │                                              phase 14: C = 0          → D14
 ├─ pack → phase 15 A   (s13932)   phase 15: A=N, B=W15, C=bias@40384  → SiLU (E15)
 │                                              pack of SiLU → phase 16 A (s14548)
 │                                              phase 16: C = D14        → D16
 ├─ pack → phase 17 A   (s14628)   phase 17: A=N, B=W17, C=bias@42496  → SiLU (E17)
 │                                              pack of SiLU → phase 18 A (s15244)
 │                                              phase 18: C = D16        → D18
 └─ pack → phase 19 A   (s15324)   phase 19: A=N, B=W19, C=bias@44608  → SiLU (E19)
                                                pack of SiLU → phase 20 A (s15940)
                                                phase 20: C = D18        → D20 (E20)

so the four branches share one input tile, each has its own weight+bias+SiLU, and their outputs are folded into a running residual (D16 = … + D14, D18 = … + D16, D20 = … + D18). The four SiLU sites are byte-identical blocks of the same f16 expression with the same four constants.


E20 — s16004–s16417: pack → plane-arena scatter → 8-lane merge → next B/C

414 statements. Four sub-sections:

(a) requantise, s16004–s16051

32 cvt.rn.satfinite.e4m3x2.f16x2 + 16 mov.b32 → %r262 … %r277 (phase 20's 64×32 output as 16 b32 = 32 e4m3 values per lane).

(b) plane-arena scatter, s16052–s16326

s16052 L43599: ld.param.u64 %rd276, [%rd368+48]
s16053 L43600: cvta.to.global.u64 %rd37, %rd276
s16055 L43604: shr.u32 %r278, %r16743, 2                 // g = laneid>>2
s16056 L43605: and.b32 %r279, %r16743, 3                 // t = laneid&3
s16059 L43608: add.s32 %r17110, %r6, %r17108             // y = %r6 + (laneid>>5)
s16065 L43614: add.s32 %r17111, %r1, %r17109             // x = %r1 + (g&7)
s16070 L43619: shl.b32 %r17112, %r17111, 3
s16071 L43620: mad.lo.s32 %r280, %r18179, %r17110, %r17112
s16072 L43621: add.s32 %r17113, %r279, %r280
s16073 L43622: selp.b32 %r281, -1, %r17113, %p35
s16074 L43623: setp.lt.s32 %p317, %r281, 0
s16075 L43624: @%p317 bra $L__BB1_258
s16078 L43628: st.global.u32 [%rd278], %r262
s16080 L43631: or.b32 %r282, %r279, 4
s16087 L43639: st.global.u32 [%rd280], %r263

The stride register is %r18179 = shl.b32 %r18179, %r364, 3 (s2219): 8·ex, and %r364/%r365 = ld.param.v2.u32 [param_0+0].

All sixteen stores (s = st.global.u32 [%rd37 + 4*index], value), where idx = stride*y + (x<<3) + d:

# statement value stride y x d
1 s16078 %r262 %r18179 %r6 + (laneid>>5) %r1 + ((laneid>>2)&7) t
2 s16087 %r263 %r18179 as above as above t|4
3 s16113 %r264 %r18180 = %r364<<3 %r6 + ((g+8)>>3) %r1 + ((g+8)&7) t
4 s16121 %r265 %r18180 as above as above t|4
5 s16147 %r266 %r18181 %r6 + ((g+16)>>3) %r1 + ((g+16)&7) t
6 s16155 %r267 %r18181 as above as above t|4
7 s16181 %r268 %r18182 %r6 + ((g+24)>>3) %r1 + ((g+24)&7) t
8 s16189 %r269 %r18182 as above as above t|4
9 s16215 %r270 %r18183 %r6 + ((g+32)>>3) %r1 + ((g+32)&7) t
10 s16223 %r271 %r18183 as above as above t|4
11 s16249 %r272 %r18184 %r6 + ((g+40)>>3) %r1 + ((g+40)&7) t
12 s16257 %r273 %r18184 as above as above t|4
13 s16283 %r274 %r18185 %r6 + ((g+48)>>3) %r1 + ((g+48)&7) t
14 s16291 %r275 %r18185 as above as above t|4
15 s16317 %r276 %r18186 %r6 + ((g+56)>>3) %r1 + ((g+56)&7) t
16 s16325 %r277 %r18186 as above as above t|4

with g = laneid>>2, t = laneid&3, every %r1818N equal to %r364<<3 (s16090, s16124, s16158, s16192, s16226, s16260, s16294), and each block guarded by p35…p43 = (y<0) | (y ≥ %r365) | (laneid>255) | (x<0) | (x ≥ %r364); a guard failure stores -1 as the index, which the following setp.lt.s32 %p, %idx, 0 + bra turns into "skip the store". Note (g+8j)&7 = g&7 and (g+8j)>>3 = (g>>3)+j, so blocks 1…7 only differ from block 0 in the row they land on (%r6 + g>>3 + j) and in the value written.

The 16 stored b32 are the requantised 64×32 output of phase 20, so this writes 8*32 = 256 bytes per row-stride unit of a %r364-wide arena.

(c) 8-lane merge, s16327–s16369

s16332 L43915: and.b32 %r17575, %r17171, 4              // laneid & 4
s16334 L43917: setp.eq.s32 %p389, %r17575, 0
s16338 L43921: shfl.sync.down.b32 %r17580|%p390, %r262, %r17576, %r17578, %r17579
s16339 L43922: shfl.sync.up.b32 %r17581|%p391, %r266, %r17577, %r17578, %r17579
s16340 L43923: selp.b32 %r17490, %r262, %r17581, %p389
s16341 L43924: selp.b32 %r17510, %r17580, %r266, %p389

Eight shfl.sync.down/up pairs, all with delta = 4 — the third operand (%r17576 = mov.u32 4, s16333) is the PTX b (delta/mask) operand and the fourth is the group/width operand c: 31 for the .down (%r17578, s16336) and 0 for the .up (%r17577, s16335) — with c = 0 the "same group" test ~(c-1) = 0 always succeeds, so the up-shuffle is unrestricted. Membermask is %r17579 = mov.u32 -1 (s16337). The selection predicate is %p389 = (laneid & 4) == 0:

down = a[lane+4]  (shfl.sync.down, delta 4)
up   = a[lane-4]  (shfl.sync.up,   delta 4)
dst1 = p389 ? own_low  : up        // lanes 0-3: own low ; lanes 4-7: lane-4's high
dst2 = p389 ? down     : own_high  // lanes 0-3: lane+4's low ; lanes 4-7: own high

so each of the eight destinations ends up holding, in its two halves, the values of both members of the pair {lane, lane±4}.

The eight pairs take (src_low, src_high) = (%r262,%r266) (%r264,%r268) (%r263,%r267) (%r265,%r269) (%r270,%r274) (%r272,%r276) (%r271,%r275) (%r273,%r277) and write, in order, %r17490 %r17510 %r17530 %r17550 %r17492 %r17512 %r17532 %r17552 %r17491 %r17511 %r17531 %r17551 %r17493 %r17513 %r17533 %r17553 (s16340–s16369). Those sixteen are phase 21's A operand: {%r17490 %r17491 %r17492 %r17493}, {%r17510 …}, {%r17530 …}, {%r17550 …} (four fragments of four, confirmed from the 32 mma at s16418–s16449); the same sixteen are re-read as %r17490 %r17491 %r17492 %r17493 %r17497 %r17510 %r17511 %r17512 %r17513 %r17530 %r17531 %r17532 %r17533 %r17550 %r17551 %r17552 %r17553 by the loads that follow. Net effect: the sixteen 32-bit payloads that eight lanes scattered are re-paired into four A fragments of four before the next GEMM.

(d) next B/C loads, s16370–s16417

s16373 L43958: add.s64 %rd309, %rd326, 45696
s16374 L43960: ld.weak.global.ca.v4.u32 { %r17173,…},[%rd309]   // … +45696+512*i, i=0..15
s16410 L44029: ld.global.u32 %r17247, [%rd329+53888]
s16411 L44030: ld.global.u32 %r17257, [%rd329+53904]
   … +53920, +53936, +53952, +53968, +53984, +54000
  • 16 weight fragments at W + 16*laneid + 45696 + 512*i → phase 21's B (%r17173 … %r17236);
  • 8 bias words at W + ((laneid<<2)&12) + {53888…54000} → %r17247 %r17257 %r17327 %r17337 %r17407 %r17417 %r17487 %r17497 = phase 21's C seed (128 bytes per lane = 64 f16 per lane).

Inputs: phase 20's D, %r1/%r6/%r364/%r365/%r18179, W, S. Outputs (58 live): %r262…%r277 (they are read again inside this same epilogue by the merge), the 16 merged A registers, phase 21's B and C, %rd369 %rd370 (cached weight pointer).


E21 — s16450–s16606: pack → surface scatter → %tid.x → next B/C

157 statements.

(a) requantise, s16450–s16472

s16450 L44261: add.s64 %rd38, %rd369, 54016
s16451 L44263: cvt.rn.satfinite.e4m3x2.f16x2 %rs6215, %r17298
s16453 L44269: cvt.rn.satfinite.e4m3x2.f16x2 %rs6218, %r17388
s16455 L44274: mov.b32 %r305, {%rs6217, %rs6218}

16 packs of phase 21's D → six b32 %r305…%r310 plus the two v2.u16 payloads {%rs6215,%rs6216} and {%rs6229,%rs6230}.

(b) coordinate / surface scatter, s16473–s16592

s16473 L44316: ld.param.v2.u32 {%r17597, %r17598}, [%rd370]      // param block +0
s16474 L44317: shr.u32 %r313, %r17598, 1                         // ey>>1
s16475 L44318: shr.u32 %r314, %r17597, 1                         // ex>>1
s16476 L44319: shl.b32 %r315, %r314, 3                           // 8*(ex>>1)
s16477 L44320: mul.lo.s32 %r316, %r313, %r315                    // row stride
s16481 L44326: ld.param.u64 %rd331, [%rd370+56]
s16482 L44327: cvta.to.global.u64 %rd39, %rd331                  // destination buffer
s16487 L44332: bfi.b32 %r17603, %r17602, %r17599, 1, 31
s16490 L44335: shr.s32 %r319, %r17605, 1                         // y0 = %r6 >> 1 (arith)
s16506 L44351: mad.lo.s32 %r321, %r315, %r17606, %r17610
s16507 L44352: or.b32 %r17611, %r318, %r321
s16513 L44359: st.global.v2.u16 [%rd333], {%rs6215, %rs6216}

The destination pointer is param block +56, not the plane arena; extents come from param+0 and are halved. Four index bases are formed, each guarded by p43 = (y<0)|(y≥(ey>>1))|(tile bits>3)|(x<0)|(x≥(ex>>1)):

base statement y x stores
%r321 s16506 (%r6>>1) + (((laneid>>5)&~2)|((g&1)<<1)) (%r1>>1) + ((laneid>>3)&3) s16513 v2.u16 {%rs6215,%rs6216} at +t; s16522 u32 %r305 at +(t|4)
%r325 s16542 recomputed from g+8 recomputed s16549 u32 %r306 at +t; s16557 u32 %r307 at +(t|4)
%r328 = %r321 + %r316 s16559 (as %r321) (as %r321) s16566 u32 %r308 at +t; s16574 u32 %r309 at +(t|4)
%r331 = %r325 + %r316 s16576 (as %r325) (as %r325) s16583 u32 %r310 at +t; s16591 v2.u16 {%rs6229,%rs6230} at +(t|4)

So: 8 guarded stores, 6 × st.global.u32 and 2 × st.global.v2.u16, index base + (t or t|4) in units of 4 bytes. (%r325 and %r328 are equal — see §6.)

(c) next B/C and %tid.x, s16593–s16606

s16595 L44456: mov.u32 %r18203, %tid.x
s16599 L44463: ld.weak.global.ca.v4.u32 { %r17629,…%r17632},[%rd348]   // +54016
s16605 L44472: ld.global.u32 %r17703, [%rd352+54528]
s16606 L44473: ld.global.u32 %r17713, [%rd352+54544]

phase 22's B fragment (W + 16*laneid + 54016) and two C seeds (W + ((laneid<<2)&12) + {54528, 54544}).

Outputs (5 live): %r17597 (the param+0 v2.u32), %r17629, %r17703 %r17713, %r18203 (%tid.x) — two of the eighteen values the cut returns.


E22 — s16615–s17222: the final epilogue

608 statements; it consumes the eighteen values returned by the cut (%r17597, %r18203 and the sixteen D registers %r17634 %r17635 %r17644 %r17645 %r17654 %r17655 %r17664 %r17665 %r17674 %r17675 %r17684 %r17685 %r17694 %r17695 %r17704 %r17705) and writes the surfaces.

(a) cross-lane row max, s16615–s16765

s16616 L44534: and.b32 %r17717, %r17714, 7                 // row key = laneid & 7
s16617 L44535: shr.u32 %r17718, %r17714, 3                 // column group = laneid >> 3
s16618 L44536: and.b32 %r17719, %r17714, 3
s16620 L44537: selp.b32 %r17720, %r17635, %r17634, %p430   // pick by (laneid&3)
s16636 L44554: bfi.b32 %r17733, %r17717, %r17732, 2, 3
s16639 L44556: shfl.sync.idx.b32 %r17736|%p434, %r17722, %r17733, %r17734, %r17735
s16657 L44574: selp.b32 %r334, %r17748, %r17750, %p440

Two identical blocks (rows g from the D pairs {r17634,r17635}, {r17654,r17655} and rows g+8 from {r17674,r17675}, {r17694,r17695}), each doing:

  1. a selp chain over (laneid&3) selects one of the four D pairs into %r17722/%r17725/%r17728/%r17731 (the four n-tiles of the row segment);
  2. four shfl.sync.idx.b32 with srcLane = (laneid&7) | ((((laneid>>3)+d)&3)<<2) for d = 0, 1, 2, -1 (%r17734 = mov.u32 31, %r17735 = mov.u32 -1);
  3. a selp chain over laneid>>3 picks between the four shuffled values.

Result: %r334…%r337 (row g) and %r339…%r342 (row g+8) hold the per-lane row maximum over the 8 columns; the following max.f16 / max.f16x2 chain (s16792–s16799, s16800…) reduces those four registers, then sub.f16x2 %r17849, %r334, %r17863 (s16800) subtracts the row max: this is a real row max, done with shfl.sync.idx, unlike E5's.

(b) bounds check, s16766–s16784

s16774 L44693: setp.ge.u32 %p462, %r346, %r17597           // x >= ex  (%r17597 = param+0 low)
s16775 L44694: setp.ge.u32 %p463, %r348, %r17598           // y >= ey
s16780 L44699: @%p466 bra $L__BB1_307
s16784 L44704: ld.param.u64 %rd353, [%rd372+112]            // surface A (param+112)
s16766 L44685: shr.s32 %r17836, %r18203, 31                // %r18203 = %tid.x (E21)

%r18203 is used to derive %r345 = tid.x - ((tid.x+round)>>3<<3) (row/tile offset), %r346 = %r345 + %r1 (x), %r348 = (…>>3) + %r6 (y); out-of-range lanes jump to $L__BB1_307 and skip (b)…(e).

(c) softmax, s16785–s16867

s16786 L44710: mul.ftz.f32 %f2630, %f2602, 0fBFB8AA3B      // -log2(e)
s16787 L44711: ex2.approx.ftz.f32 %f2631, %f2630
s16788 L44712: add.ftz.f32 %f2632, %f2631, 0f3F800000
s16790 L44713: div.approx.ftz.f32 %f2603, %f2633, %f2632   // 1/(1+2^(-max*log2e))
s16792 L44719: max.f16x2 %r17840,%r334,%r335
s16800 L44742: sub.f16x2 %r17849,%r334,%r17863             // x - rowmax
s16803 L44750: mul.ftz.f32 %f2634, %f2604, 0f3FB8AA3B      // *log2(e)
s16804 L44751: ex2.approx.ftz.f32 %f2606, %f2634           // 2^(…)
s16859 L44869: add.ftz.f32 %f2651, %f2649, %f2621          // row sum
s16860 L44870: div.approx.ftz.f32 %f2622, %f2633, %f2651   // 1/rowsum
s16863 L44875: mul.f16x2 %r17864,%r17865,%r17878

A textbook f32 softmax over each row's 8 values: max over the four lane registers, y = 2^((x-max)*log2e) in f32 (cvt.f32.f16 → mul by 0f3FB8AA3B → ex2.approx), accumulate the row sum in f32, invert with div.approx.ftz.f32 1.0/…, round to f16, and multiply the five f16x2 registers (%r17864 %r17867 %r17870 %r17873 %r17876). The %f2633 dividend is the mov.f32 %f2633, 0f3F800000 at s16789 (1.0). Note s16786 uses -log2(e) on the max to build 1/(1+2^(-max·log2e)) (a per-row scale kept in %rs6232), which is applied later to the second row group.

(d) shared-patch gather, s16868–s17143

s16871 L44899: mov.u32 %r17960, 8
s16868 L44896: mad.lo.s32 %r17969, %r347, 80, %r18206       // input_noise + 80*y
s16869 L44897: shl.b32 %r17970, %r345, 3
s16870 L44898: add.s32 %r17971, %r17969, %r17970
s16872 L44900: ld.shared.v4.u16 {%rs6327,…%rs6330}, [%r17971+8]
s16873 L44901: ld.shared.v4.u16 {%rs6331,…%rs6334}, [%r17971]
s16874 L44902: mov.b32 %r17880, {%rs6331, %rs6327}           // 2x2 transpose pair
s16875 L44904: fma.rn.f16x2 %r17879,%r17880,%r17864,%r17890  // acc += w * p
s16880 L44917: ld.shared.v4.u16 {%rs6341,…}, [%r17971+80]
s16881 L44918: ld.shared.v4.u16 {%rs6348,…}, [%r17971+16]
s16888 L44934: ld.shared.v4.u16 {%rs6355,…}, [%r17971+96]
s16889 L44935: ld.shared.v4.u16 {%rs6362,…}, [%r17971+88]
s16896 L44951: ld.shared.v4.u16 {…}, [%r17971+168]
s16897 L44952: ld.shared.v4.u16 {…}, [%r17971+160]
s16910 L44983: ld.shared.v4.u16 {…}, [%r17971+176]

18 ld.shared.v4.u16 at +0 +8 +16 | +80 +88 +96 | +160 +168 | +176, i.e. a 2×2 (plus one extra column) stencil of the shared array input_noise (_ZZ33cuda_dldn_engine_swin_enc0_kernel35DldnEngineFusedSwinEnc0ParamsStructE11input_noise), which is filled in the prologue (s1457–s1547) as st.shared.v4.u16 [input_noise + y*80 + x*8], {f16c0, f16c1, f16c2, 0} — a 10×10 grid of f16x2 triples (three channels + a zero pad), 800 bytes.

The 2×2 blocks are multiplied by the five softmax registers with fma.rn.f16x2, accumulated into %r17879 %r17883 %r17887 %r17891… (a 2×2 accumulator per row group), then the accumulator halves are summed with add.f16 (s17135–s17137) and one more shared value is folded in with fma.rn.f16 (s17140–s17142).

(e) sigmoid, difference, scale, cross-lane sum, s17143–s17200

s17145 L45494: mul.ftz.f32 %f2708, %f2679, 0fBFB8AA3B
s17146 L45495: ex2.approx.ftz.f32 %f2709, %f2708
s17148 L45497: div.approx.ftz.f32 %f2680, %f2689, %f2710    // sigmoid(acc)
s17156 L45515: mul.f16 %rs6445,%rs6394,%rs6429
s17160 L45530: mov.f32 %f2685, 0f3E800000                   // 0.25
s17162 L45536: sub.f16 %rs6454,%rs6429,%rs6445              // acc - sigmoid term
s17167 L45550: mul.f16x2 %r18080,%r18081,%r18082            // * 0.25
s17175 L45568: shfl.sync.down.b32 %r18084,%r18080,%r18105,%r18087,%r18115
s17176 L45574: add.f16x2 %r18089,%r18080,%r18084
s17191 L45613: add.f16 %rs6472,%rs6466,%rs6471
s17200 L45637: add.f16 %rs6478,%rs6472,%rs6477

sigmoid(v) = 1/(1 + 2^(-log2e·v)), then three acc − sigmoid(acc) differences scaled by 0.25, then a two-step cross-lane sum: shfl.sync.down with the delta operand %r18105 = 1 (s17175) then %r18113 = 8 (s17181), each followed by an add.f16x2; the same two steps repeat on the scalar path (s17188 delta 1, s17197 delta 8, add.f16 at s17191 and s17200). The group/width operand is 31 in all four, so this is a partial (not five-step) warp reduction.

(f) surface stores, s17201–s17222

s17201 L45640: ld.param.u64 %rd362, [%rd371+120]            // surface B
s17202 L45641: ld.param.u64 %rd365, [%rd371+128]            // surface half
s17204 L45644: sust.b.2d.v4.b16.zero [%rd360, {%r18119,%r354}], {%rs6454,%rs6457,%rs6460,%rs6549}
s17205 L45647: sust.b.2d.v4.b16.zero [%rd362, {%r18119,%r354}], {%rs6445,%rs6448,%rs6451,%rs6550}
s17210 L45653: @%p477 bra $L__BB1_310                       // gated on (%r355&1)==0
s17220 L45665: sust.b.2d.v4.b16.zero [%rd365, {%r18134,%r18135}], {%rs6551,%rs6552,%rs6478,%rs6554}
s17222 L45669: ret

Three 4-channel b16 surface writes (params +112, +120, +128); the third is taken only by even %r355 = %r315|%r354 lanes and stores 0 in its fourth channel.

Inputs: the eighteen returned values, W (+54016/54528), param+0/+112/+120/+128, the shared input_noise array, constants 0fBFB8AA3B, 0f3FB8AA3B, 0f3F800000, 0f3E800000. Outputs: none (the kernel returns).


4. The plane-arena writes (param block +48)

S = cvta.to.global.u64(%rd276) with %rd276 = ld.param.u64 [%rd368+48] (s16052–s16053) is the only occurrence of param_0+48 in the entry: I grep'ed every ld.param and this offset is read once.

There are exactly two store groups in the whole entry (24 st.global.* in the file, all after line 43628):

(i) E20, 16 × st.global.u32 [%rd37 + 4*idx], v — s16078…s16325 (see §3/E20 for the full table). In one line per store:

idx = (8*ex)*y + 8*x + d,   guard  0 <= y < ey  and  0 <= x < ex  and  laneid <= 255
store1  (s16078): v = %r262,  y = %r6 + (laneid>>5),                 x = %r1 + ((laneid>>2)&7),    d = t
store2  (s16087): v = %r263,  same                                    d = t|4
store3  (s16113): v = %r264,  y = %r6 + ((g+8)>>3),                   x = %r1 + ((g+8)&7),         d = t
store4  (s16121): v = %r265,  same                                    d = t|4
store5  (s16147): v = %r266,  y = %r6 + ((g+16)>>3),                  x = %r1 + ((g+16)&7),        d = t
store6  (s16155): v = %r267,  same                                    d = t|4
store7  (s16181): v = %r268,  y = %r6 + ((g+24)>>3),                  x = %r1 + ((g+24)&7),        d = t
store8  (s16189): v = %r269,  same                                    d = t|4
store9  (s16215): v = %r270,  y = %r6 + ((g+32)>>3),                  x = %r1 + ((g+32)&7),        d = t
store10 (s16223): v = %r271,  same                                    d = t|4
store11 (s16249): v = %r272,  y = %r6 + ((g+40)>>3),                  x = %r1 + ((g+40)&7),        d = t
store12 (s16257): v = %r273,  same                                    d = t|4
store13 (s16283): v = %r274,  y = %r6 + ((g+48)>>3),                  x = %r1 + ((g+48)&7),        d = t
store14 (s16291): v = %r275,  same                                    d = t|4
store15 (s16317): v = %r276,  y = %r6 + ((g+56)>>3),                  x = %r1 + ((g+56)&7),        d = t
store16 (s16325): v = %r277,  same                                    d = t|4
where g = laneid>>2, t = laneid&3, ex = %r364, ey = %r365,
      %r1 = x origin (s6), %r6 = y origin (s16)

(ii) E21, 8 stores to param+56, not to the arena — s16513 (v2.u16 of %rs6215/%rs6216), s16522, s16549, s16557, s16566, s16574, s16583 (u32 of %r305…%r310), s16591 (v2.u16 of %rs6229/%rs6230). See §3/E21 for the indices.

What is written: phase 20's requantised 64×32 output — 16 b32 per lane, one per store — i.e. 4 bytes × 32 lanes × 16 stores = 2048 bytes per block (4096 e4m3 elements' worth is not what lands in memory; only 16 of phase 20's 32 D registers per lane are scattered).

What reads it back: nothing inside this kernel. param_0+48 is never the base of an ld.*; P5 and P10 take their score bias from the weight image (E9's 24 loads at W+23744+512*i, and phase 4's epilogue's 24 loads at W+8384+512*i), and P5/P10 take their B operand from registers (phases 3/9 D → e4m3 packs). See §6.

5. The eight shfl reductions of the merge

The brief's "eight shfl reductions" are the eight shfl.sync.down/up pairs of E20(c), s16338–s16368 (statement range s16327–s16369 including the setup). Each pair uses the same delta operand %r17576 = 4 (the PTX b operand), with the group/width operand c = 31 (.down) or 0 (.up), and laneid & 4 as the selection predicate, so each of the eight 32-bit payloads (%r262…%r277) is re-paired between lanes 0-3 and 4-7 of every laneid&7 group before it becomes phase 21's A operand.

(There is a second, unrelated cross-lane reduction in E22: the four shfl.sync.idx per row group at s16639–s16651 / s16680–s16683 / s16714–s16726 / s16755–s16758 that implement the row maximum, plus the shfl.sync.down pair after the sigmoid at s17175/s17181 and its scalar copy at s17188/s17197.)

6. Uncharacterised

U1. "The plane-arena writes at stmts 15053–15198". In this file, with the rr_layer_spec.py numbering:

  • s15053–s15198 lie inside phase 17's epilogue (s14692–s15291). They are the per-output-channel activation coefficient materialisation (cvt.rn.f16.f32 %rs, %f2601 / neg / max / min / abs / mul / sub / mul / add / mul) and contain no store of any kind — the entry has zero st.* between line 1450 and line 43628.
  • The only writes to the plane arena are (i) and (ii) of §4, at s16078–s16325 and s16513–s16591.
  • No statement reads param_0+48, so nothing written there can be read back by P5 or P10 within this launch.

I could not reconcile 15053–15198 with this corpus file. The two other landmarks supplied with that range also disagree with it: %r17597 (the ld.param.v2.u32 of param+0) is s16473 / L44316, and %r18203 (%tid.x) is s16595 / L44456, not 15448 and 15570. Missing: the statement numbering that produced 15053/15198/15448/15570 — most likely a different linearisation (every ;-terminated physical line of the whole file gives 14946 for L44316 and 15375 for L43628; entry-only gives 14567 and 19913 for the non-; variant, neither matches). Give me that numbering (or the tool) and the range will map directly.

U2. The identity of the 12288-byte score bias. Both score phases add a 12288-byte f16 table as their GEMM C operand (P5: W+8384…W+20672, P10: W+23744…W+36032), indexed as §3/E9 describes. Whether it is a relative position bias, a learned score bias, or a per-lane constant cannot be decided from the PTX: it is an opaque blob in a permuted weight image. Missing: the weight image's layout description (the _prep permutation) or the parameter names from the model file.

U3. The identity of the per-column tables. W+37056/37072/37088/37104 (E6), W+37120/37136/37152/37168 (E12, used as the norm gain), and W+39296/39312/39328/39344 (E13) are all read with the same W + ((laneid<<2)&12) + off pattern and are 8 f16 per lane (32 columns). Their role is computable (E6: added to the residual before phase 7; E12: multiplies the inverse std; E13: added after the activation) but their names are not. The GEMM biases at +38208/40384/42496/44608/53888/54528 are the same case.

U4. E21's index bases. %r325 (s16542) and %r328 (s16559) are both %r321 + %r316, so the stores at s16549/s16557 and s16566/s16574 target the same two indices in program order — the later store wins. I could not tell from the PTX whether %r316's value or the g+8 recomputation in the second block is meant to separate them (it looks like the source of a dropped output, but it may equally be a degenerate case that the guard prevents). Missing: either a run of the kernel with a plane-arena/surface dump, or the corresponding high-level code.

U5. The geometry of P6 (M/N/K). E5 produces 12 A fragments and 12 B fragments for phase 6; the A side is the 96 probabilities and the B side the 48 transposed-V registers, so phase 6 is 64×32 = (64×96)·(96×32) with three k-steps of 32 (verified from the mma chain at s7826–s7873: D ← A·B + D_prev for exactly three steps). rr_layer_spec.py's phase table still prints M=192 N=96 K=96 for that phase; the epilogue contract above does not depend on resolving it.

U6. The logical meaning of E22's final expression. The tail computes out = (acc − sigmoid(acc)) * 0.25 on three channels gathered from the shared input patch. The PTX gives the arithmetic exactly (§3/E22(e)) but not what model term it corresponds to. Missing: the high-level source (a shrink / gate / residual term) or a reference implementation to compare against.

7. Quick acceptance index

requirement where
P5 epilogue: row max / exp / row sum / reciprocal / position bias §3-E5(a)–(e); no row max and no bias add are present — see §6 U1/U2
P5: the position bias added as the GEMM C operand (s3764–s3814, W+8384+512*i), not in the epilogue
P6: P*V, the transposes §3-E5(f) (48 movmatrix, s7562–s7609) + E6
second window P8–P11 §3-E8, E9, E10, E11
MLP P12–P20, the norm §3-E12
fc1/gelu/fc2 triples §3-E13…E19 (activation is the clamped cubic of E13, not GELU)
gains/biases at 38208/40384/42496/44608/53888 §1.1, §3-E12/E13/E14/E16/E18/E20
P21/P22 merge + eight shfl §5, §3-E20(c), E21
plane-arena writes, exact value and index §4
final epilogue after L44528 and the 18 values §3-E22