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).
57 KiB
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
sNNNNis a statement index in the linearisationrr_layer_spec.pydocuments (statements of the entry body; a wrapped{ … }block and a 4-linemma.synccount as one statement, numbered at their first line).LNNNNNis the physical line of the same instruction in the corpus file, so every claim can be checked withsed -n 'LNNNNN,LNNNNN p' …ptx.- An epilogue is the run of statements between one phase's last
mmaand the next phase's firstmma. The phase table (rr_layer_spec.py enc0) gives the ranges; this file names what each one computes. Wdenotes the prepared weight image:ld.param.u64 %rd367, [cuda_dldn_engine_swin_enc0_kernel_param_0+40](s2221) andcvta.to.global.u64 %rd2, %rd1(s12,%rd1= the same+40load at s11). Everyld.weak.global.ca.v4.u32in an epilogue readsW + 16*laneid + imm.Sdenotes 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_bitsinptx/lib/zluda_ptx_impl.cpp).rcp.approx.ftz.f32,rsqrt.approx.ftz.f32,ex2.approx.ftz.f32,div.approx.ftz.f32— always wrapped ascvt.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:
- s7610–s7681 — from the movmatrix results → phase 6's 12 B fragments
(
%r6178…%r6289, 24 b32), i.e. the transposed V. - 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. - 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:
-
Next-B loads (s7874–s7880): two
v4.u32fragments fromW + 16*laneid + {20672, 21184}→%r6293…%r6300, phase 7's B operand. -
Table loads (s7881–s7889): four
ld.global.v2.u16fromW + ((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 columns2t,2t+1of two of the four 8-column groups). -
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 bymov.b32 {%rs3064, %rs3067}, %r1175and 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. -
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'smovmatrixblock 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}fromcvtof%r6503/%r6523at 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)}(%r7281from s8235–s8237,%r7535from s8163–s8165). Whatever the score GEMM's two operands are, phase 10 gets one of them twice at fragment 0. - phase 10's four A-fragment sets:
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 constant2^-13; the mean is not formed. gain[c]is read fromW + ((laneid<<2)&12) + {37120, 37136, 37152, 37168}(8 f16 per lane = 32 columns), so the gain is per column, constant across rows.- the
invregister 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 (32rsqrtfor 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.u32atW + ((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:
- a
selpchain over(laneid&3)selects one of the four D pairs into%r17722/%r17725/%r17728/%r17731(the four n-tiles of the row segment); - four
shfl.sync.idx.b32withsrcLane = (laneid&7) | ((((laneid>>3)+d)&3)<<2)ford = 0, 1, 2, -1(%r17734 = mov.u32 31,%r17735 = mov.u32 -1); - a
selpchain overlaneid>>3picks 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–s15198lie 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 zerost.*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 |