wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
417abc64b51cf26d8a3c14319ab7a26b779d8d0c
d4r/kernels/rr/descriptor_dec1.py
T
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

432 lines
22 KiB
Python
Executable File
Raw Blame History

#!/usr/bin/env python3
"""Per-layer descriptor for cuda_dldn_engine_swin_dec1_kernel (nvngx_dlssd 310.7).
The parameterised-engine deliverable: every number here is either (a) read out of
~/.cache/d4r-rr-corpus/0022-PREPASS_ENTRYPOINT_NAME.ptx, entry
`cuda_dldn_engine_swin_dec1_kernel`, and reproducible with
python3 kernels/rr/rr_layer_spec.py dec1
or (b) marked `"unverified"` with the reason. Phase boundaries, mma operand
register lists and the weight-image offsets were extracted mechanically from the
linearised statement list; the epilogue operations are named from
kernels/rr/rrswin_dec1_epilogues.md, which cites the statement ranges.
Geometry recovered by interpreting NVIDIA's own index arithmetic (not guessed)
is recorded in `e1_patch_grid`.
python3 -c "import descriptor_dec1 as d; print(d.layer())"
"""
# ---------------------------------------------------------------------------
# launch geometry, from the entry's prologue (s1-s325) and rr_layer_spec.py
# ---------------------------------------------------------------------------
def layer():
"""Launch geometry, parameter-block layout and tile shape.
`param_offsets` are byte offsets into the 144-byte parameter block; the entry
declares it as one `.param .align 8 .b8 ..._param_0[144]`, so every offset below
is a plain byte read (`rr_layer_spec.py dec1`: "param block 144 bytes (5
ld.param over 3 distinct offsets)", "param offsets [8, 40, 80]").
"""
return {
"kernel": "cuda_dldn_engine_swin_dec1_kernel",
"corpus": "0022-PREPASS_ENTRYPOINT_NAME.ptx",
"grid": (81, 47, 1),
"block": (32, 1, 2), # two waves per block; %tid.z is the window index
"waves": 2,
"laneid": 32,
"param_block_bytes": 144,
"param_offsets": {
0: {"kind": "u32x2", "name": "extents", "meaning": "ex, ey"},
8: {"kind": "ptr", "name": "in_plane",
"meaning": "input feature plane arena, 52 ld.global.u32 in the prologue"},
24: {"kind": "ptr", "name": "patch_gather",
"meaning": "the +24 skip input E1 gathers the patch from; reached as "
"[param_0+80 - 56], not as its own ld.param"},
40: {"kind": "ptr", "name": "weights",
"meaning": "NVIDIA's prepared weight image (arg040, 96,800 bytes)"},
48: {"kind": "ptr", "name": "out_plane",
"meaning": "output plane arena; E25's 16 guarded st.global.u32 write here"},
80: {"kind": "v2u16", "name": "block_origin",
"meaning": "the per-block origin table; %rd1 = param_0+80 and "
"param_0+24 is [rd1-56]"},
88: {"kind": "tex", "name": "colour_tex"},
96: {"kind": "tex", "name": "tex_int_0"},
104: {"kind": "tex", "name": "tex_int_1"},
112: {"kind": "surface", "name": "surface_0"},
120: {"kind": "surface", "name": "surface_1"},
136: {"kind": "f32x2", "name": "tex_coord_scale"},
},
"tid_z_strides": {
# two per-tid.z strides exist and are used by different groups
# (rrswin_dec1_epilogues.md §1.1, read off the load chains).
"phase1_B": {"bytes": 12288, "phase": 1},
"phase1_bias": {"bytes": 256, "phase": 1},
"phase2_5_B_and_phase6_bias": {"bytes": 18432, "phases": [2, 3, 4, 5, 6]},
},
"tid_z_range": [0, 1],
"shared_bytes": 6400,
"token_geometry": {
"patch": "8x8 tokens per warp; a 10x10 halo grid of 4-word cells",
"unverified": "the 10x10 grid's mapping to image coordinates was not "
"reduced to a closed form; see e1_patch_grid for what the "
"interpreter measured and rrswin_dec1_epilogues.md U4 for "
"what remains open.",
},
"constants_f16": {
# §1.2: value after the cvt to f16, i.e. what the arithmetic actually uses
"softsign_scale": 0x2466, # 0.017181396484375
"clamp_lo": 0xB873, # -0.55615234375
"clamp_hi": 0x3873, # +0.55615234375
"cubic_a": 0x3B6B, # 0.92724609375
"cubic_b": 0x3D80, # 1.375
"rms_eps": 0x0800, # 2^-13
"act_a": 0x3698, # 0.412109375
"act_b": 0x2D30, # 0.0810546875
"act_half": 0x3800, # 0.5
"act_two": 0x4000, # 2.0
},
}
# ---------------------------------------------------------------------------
# E1's patch grid, recovered by interpreting the PTX
# ---------------------------------------------------------------------------
def e1_patch_grid():
"""The 6400-byte shared tile E1 builds and reads, as measured.
`/tmp/d4r-dec1/interp.py` evaluates NVIDIA's own index arithmetic for all 32
lanes x tid.z in {0,1} over s689-s1367 (the 48 guarded `st.shared`) and
s1370-s2214 (the 24 `ld.shared.v2.u16`). The layout below is that output,
not an inference.
"""
return {
"sub_arrays": 4,
"sub_array_bytes": 1600,
"word_index": "sub*400 + row*40 + col*4 + t (row 0..9, col 0..9, t = laneid&3)",
"grid": "10 rows x 10 cols x 4 words",
"pack_index": "p = 16*mgrp + 4*(gp>>1) + 2*h + row; "
"pk[p] = e4m3x2(D[mgrp][2q][j=2][h][row]) "
"| e4m3x2(D[mgrp][2q+1][j=2][h][row])<<16, q = gp>>1",
"live_packs": [p for p in range(48) if p not in
(33, 35, 37, 39, 41, 43, 45, 47)],
"row_parity": {"tid_z_0": [1, 3, 5, 7, 9], "tid_z_1": [0, 2, 4, 6, 8]},
"coverage": "the 48 stores tile all 400 cells exactly once, zero duplicates",
"cells_per_g_tid_z_0": [28, 28, 28, 24, 20, 20, 24, 28],
"cells_per_g_tid_z_1": [24, 28, 28, 28, 28, 24, 20, 20],
"readback_positions": {
"count": 24,
"note": "6 grid positions x 4 sub-arrays",
"tid_z_0": "rows {1,2,3,4,0,9} at column g+1",
"tid_z_1": "rows {5,6,7,8} at column g+1, plus column halos "
"(g+1,0) and (g+1,9)",
},
"unverified": "the gather source arithmetic (param_0+24, index "
"clamp_row*(8*ex) + clamp_col*8 + clamp(t,0,7)), the "
"sum-of-squares trees and the restage were not decoded; the "
"epilogue cites them at s2215-s7095.",
}
# ---------------------------------------------------------------------------
# the 26 phases
# ---------------------------------------------------------------------------
def _gemm(mt, nt, k, mma, a, b, c, ksplit=None):
d = {
"m_tiles": mt, "n_tiles": nt, "k": k, "mma": mma,
"a": a, "b": b, "c": c,
"k_tiling": ksplit or "one 32-wide k step; the engine emits two 16-wide "
"gfx12 WMMA steps in NVIDIA's slot order",
}
return d
# Weight-image byte offsets are relative to W (param_0+40). `tid_z` says whether
# the offset carries a per-tid.z stride and which one.
_PATCH = "E1's restaged patch grid (shared), m-tile index selects the A fragment"
_NORMT = "the RMS-normalised tile N of E9 (32 f16x2), requantised to e4m3"
def phases():
"""One dict per phase, in execution order.
`m_tiles`, `n_tiles`, `k` and `mma` are exact: `mma` is the statement count and
the tiles come from counting the distinct A quads and B pairs the phase's mma
read (rr_layer_spec.py's phase table). `a`, `b`, `c` name the operand's source;
where the source is a weight-image read the byte offset and the per-tid.z
stride are given.
"""
out = []
P = out.append
P(dict(
index=1, stmts=(403, 546), epilogue="E1", epilogue_stmts=(547, 7095),
gemm=_gemm(9, 48, 96, 144,
{"source": "the staged 8x8 block patch (param_0+8)",
"operands": "9 .f16.e4m3 fragments of four b32, 36 registers, "
"passed across the cut",
"requantise": "already e4m3"},
{"source": "weight image",
"byte_offset": 0, "byte_span": 11788, "stride": "12288 * tid.z",
"note": "48 B fragments = 24 blocks x 16 columns; banded, the "
"three m-groups are a running sum along j, not a "
"k-accumulation"},
{"source": "per-chain bias seed",
"byte_offset": 24576, "byte_span": 240, "stride": "256 * tid.z",
"kind": "per-(chain, n-tile) seed pair broadcast into both C registers"},
)))
# ---- the four 96x32 projections -------------------------------------------------
for idx, a_mt, b_off, c_src in (
(2, [0, 2, 4, 6, 8, 10], 0, "zero"),
(3, [0, 2, 4, 6, 8, 10], 2048, "zero"),
(4, [1, 3, 5, 7, 9, 11], 1024, "phase 2"),
(5, [1, 3, 5, 7, 9, 11], 3072, "phase 3")):
st = {2: (7096, 7119), 3: (7127, 7150), 4: (7158, 7181), 5: (7189, 7212)}[idx]
P(dict(
index=idx, stmts=st, epilogue="E%d" % (idx + 1 if idx in (2, 3) else idx),
gemm=_gemm(6, 4, 32, 24,
{"source": _PATCH, "m_tiles_used": a_mt,
"requantise": "already e4m3 (read back from shared by E1)"},
{"source": "weight image",
"byte_offset": 25216 + b_off, "byte_span": 1024,
"stride": "18432 * tid.z"},
{"source": c_src,
"kind": "per-(m,n) seed pair" if c_src != "zero" else "zero register"},
)))
P(dict(
index=6, stmts=(7384, 7431), epilogue="E6", epilogue_stmts=(7432, 11274),
gemm=_gemm(4, 12, 32, 48,
{"source": "E5's pack of phase 2's and phase 3's accumulators",
"requantise": "cvt.rn.satfinite.e4m3x2.f16x2 then mov.b32"},
{"source": "E5's second pack stream of the same two accumulators",
"weight_bytes": None,
"note": "phase 6 reads no weight image at all (phase table: "
"weight column `-` for every mma)"},
{"source": "the score bias",
"byte_offset": 29312, "byte_span": 12288, "stride": "18432 * tid.z",
"kind": "per-(m,n) seed pair, 24 tiles x 512 bytes"},
)))
P(dict(
index=7, stmts=(11275, 11322), epilogue="E7", epilogue_stmts=(11323, 11505),
gemm=_gemm(12, 12, 96, 48,
{"source": "E6's pack of the softmax probabilities",
"requantise": "cvt.rn.satfinite.e4m3x2.f16x2"},
{"source": "E6's movmatrix.trans of phase 5's D",
"weight_bytes": None,
"note": "48 movmatrix.sync.trans.aligned.m8n8.b16 at s11011-s11058, "
"inputs are exactly phase 5's 48 D registers (%r6715..%r6946)"},
{"source": "zero"},
),
))
P(dict(
index=8, stmts=(11506, 11521), epilogue="E8", epilogue_stmts=(11522, 11537),
gemm=_gemm(2, 8, 32, 16,
{"source": "shared, E7's pack of phase 7's D at offsets 0/512/1024/1536"},
{"source": "weight image", "byte_offset": 41600, "byte_span": 1548},
{"source": "phase 7's D plus the per-column table",
"byte_offset": 62080, "byte_span": 112,
"kind": "per-column f16 pair, 8 words per lane"},
)))
P(dict(
index=9, stmts=(11538, 11553), epilogue="E9", epilogue_stmts=(11554, 12396),
gemm=_gemm(2, 8, 32, 16,
{"source": "shared, offsets +2048 and +2560"},
{"source": "weight image", "byte_offset": 60032, "byte_span": 1548},
{"source": "zero"},
)))
# ---- the MLP band, phases 10..25 ------------------------------------------------
# Odd phases are the 32x64 projections whose C is zero; even phases are the
# 64x64 ones whose C is a bias. Every odd phase's A is the same normalised tile N.
band = [
# (index, m_tiles, n_tiles, k, mma, weight_off, c_off, c_kind)
(10, 4, 8, 64, 16, 62336, 64384, "per-(m,n) seed pair"),
(11, 2, 8, 32, 16, 64448, None, "zero"),
(12, 4, 8, 64, 16, 66624, 68672, "per-(m,n) seed pair"),
(13, 2, 8, 32, 16, 68736, None, "zero"),
(14, 4, 8, 64, 16, 70784, 72832, "per-(m,n) seed pair"),
(15, 2, 8, 32, 16, 72896, None, "zero"),
(16, 4, 8, 64, 16, 74944, 76992, "per-(m,n) seed pair"),
(17, 2, 8, 32, 16, 77056, None, "zero"),
(18, 4, 8, 64, 16, 79104, 81152, "per-(m,n) seed pair"),
(19, 2, 8, 32, 16, 81216, None, "zero"),
(20, 4, 8, 64, 16, 83264, 85312, "per-(m,n) seed pair"),
(21, 2, 8, 32, 16, 85376, None, "zero"),
(22, 4, 8, 64, 16, 87424, 89472, "per-(m,n) seed pair"),
(23, 2, 8, 32, 16, 89536, None, "zero"),
(24, 4, 8, 64, 16, 91584, 93632, "per-(m,n) seed pair"),
(25, 2, 8, 32, 16, 93696, None, "zero"),
]
stmts = {10: (12397, 12412), 11: (12834, 12849), 12: (12918, 12933),
13: (13242, 13257), 14: (13326, 13341), 15: (13650, 13665),
16: (13734, 13749), 17: (14058, 14073), 18: (14142, 14157),
19: (14466, 14481), 20: (14550, 14565), 21: (14874, 14889),
22: (14958, 14973), 23: (15282, 15297), 24: (15366, 15381),
25: (15690, 15705)}
epi = {10: "E10", 11: "E11", 12: "E12", 13: "E13", 14: "E14", 15: "E15",
16: "E16", 17: "E17", 18: "E18", 19: "E19", 20: "E20", 21: "E21",
22: "E22", 23: "E23", 24: "E24", 25: "E25"}
for idx, mt, nt, k, mma, w, c, ckind in band:
even = idx % 2 == 0
P(dict(
index=idx, stmts=stmts[idx], epilogue=epi[idx],
gemm=_gemm(mt, nt, k, mma,
{"source": _NORMT,
"note": "every odd phase packs the same 32 f16x2 tile N; "
"E11/E13/E15/E17/E19/E21/E23 re-pack the identical "
"registers %r17567..%r17598",
"requantise": "cvt.rn.satfinite.e4m3x2.f16x2"},
{"source": "weight image", "byte_offset": w, "byte_span": 1548,
"stride": None},
{"source": "bias" if even else "zero",
"byte_offset": c, "kind": ckind} if even else
{"source": "zero"},
),
))
P(dict(
index=26, stmts=(15981, 15988), epilogue="E26", epilogue_stmts=(15989, 16304),
gemm=_gemm(4, 4, 64, 8,
{"source": "E25's requantised words %r489..%r504, scattered to "
"param_0+48 and read back as P26's A",
"requantise": "24 groups of 2 cvt + 1 mov.b32 at s15706-s15753"},
{"source": "weight image", "byte_offset": 95744, "byte_span": 524},
{"source": "bias", "byte_offset": 96768, "byte_span": 16,
"kind": "per-(m,n) seed pair"},
),
))
return out
# ---------------------------------------------------------------------------
# the epilogues, as an operation list the engine can dispatch on
# ---------------------------------------------------------------------------
def epilogues():
"""name -> the operations it performs, with the offsets that carry them.
Operation names are the engine's vocabulary: e4m3_requantise, pack_next_A,
clamped_cubic_activation, add_column_table, rms_norm, softmax, patch_expand,
plane_scatter, lane_merge, surface_scatter.
"""
return {
"E1": {"after": 1, "stmts": (547, 7095), "ops": [
{"op": "e4m3_requantise", "src": "phase 1's D", "count": 144,
"dst": "shared, 10x10x4 grid (see e1_patch_grid)"},
{"op": "patch_expand", "grid": [10, 10], "words_per_cell": 4},
{"op": "gather", "arg": 24, "index": "clamp_row*(8*ex) + clamp_col*8 + clamp(t,0,7)",
"stride_bytes": 4, "count": 24},
{"op": "e4m3_unpack", "count": 96},
{"op": "rms_norm", "eps": "2^-13", "reduction": "4-lane butterfly over laneid&3",
"gain_byte_offset": 25088, "gain_span": 112,
"note": "raw sum of squares + eps, no mean subtraction and no 1/N; "
"applied per column"},
{"op": "pack_next_A", "dst": "shared", "tiles": [0, 512, 1024, 1536, 4096, 4608],
"base": "tid.z*2048 + laneid*16", "guarded_by": "tid.z*16 < 32"},
{"op": "lane_merge", "note": "tid.z*2048 splits the two waves' A tiles; "
"each wave reads the other's"},
{"op": "load_B", "byte_offset": 25216, "stride": "18432 * tid.z"},
]},
"E2": {"after": 2, "ops": [{"op": "load_B", "byte_offset": 25216 + 2048,
"span": 1024, "stride": "18432 * tid.z"}]},
"E3": {"after": 3, "ops": [{"op": "load_B", "byte_offset": 25216 + 1024,
"span": 1024, "stride": "18432 * tid.z"}]},
"E4": {"after": 4, "ops": [{"op": "load_B", "byte_offset": 25216 + 3072,
"span": 1024, "stride": "18432 * tid.z"}]},
"E5": {"after": 5, "stmts": (7213, 7383), "ops": [
{"op": "load_bias", "byte_offset": 29312, "span": 12288,
"stride": "18432 * tid.z", "dst": "phase 6's C"},
{"op": "e4m3_requantise", "src": "phase 2's and phase 3's D", "count": 80,
"dst": "phase 6's A (80 regs) and B (80 regs)"},
]},
"E6": {"after": 6, "stmts": (7432, 11274), "ops": [
{"op": "softmax", "form": "cubic exponent, no row maximum, no bias add",
"scale": 0.017181396484375, "clamp": [-0.55615234375, 0.55615234375],
"cubic": [0.92724609375, 1.375],
"note": "expval = f16 from bits ((u<<5) & 0x7FE07FE0); the 8 row sums come "
"from two shfl.sync.bfly steps"},
{"op": "lane_merge", "note": "48 movmatrix.sync.trans.aligned.m8n8.b16 of "
"phase 5's D"},
{"op": "e4m3_requantise", "src": "the probabilities and the transposed V",
"count": 144, "dst": "phase 7's A and B"},
]},
"E7": {"after": 7, "ops": [
{"op": "e4m3_requantise", "count": 32, "dst": "shared"},
{"op": "lane_merge", "note": "pack D7 to shared, bar.sync, read back as phase 8's A"},
{"op": "add_column_table", "byte_offset": 62080, "span": 112,
"dst": "phase 8's C = phase 7's D + the table"},
{"op": "load_B", "byte_offset": 41600, "span": 1548},
]},
"E8": {"after": 8, "ops": [
{"op": "load_A", "from": "shared", "offsets": [2048, 2560]},
{"op": "load_B", "byte_offset": 60032, "span": 1548},
]},
"E9": {"after": 9, "stmts": (11554, 12396), "ops": [
{"op": "rms_norm", "eps": "2^-13", "reduction": "4-lane butterfly",
"gain_byte_offset": 62208, "gain_span": 112,
"note": "produces the normalised tile N, 32 f16x2, re-packed by E11, "
"E13, E15, E17, E19, E21 and E23"},
{"op": "e4m3_requantise", "count": 32, "dst": "phase 10's A"},
{"op": "load_B", "byte_offset": 62336, "span": 1548},
{"op": "load_bias", "byte_offset": 64384, "span": 48, "dst": "phase 10's C"},
]},
"E10": {"after": 10, "stmts": (12413, 12833), "ops": [
{"op": "clamped_cubic_activation",
"clamp": [-2.0, 2.0],
"form": "g = 0.5 + y*(0.412109375 - 0.0810546875*|y|); out = x*g",
"note": "a cubic ramp, not SiLU or GELU; g(+2)=1 and g(-2)=0"},
{"op": "add_column_table", "byte_offset": 66496, "span": 112,
"dst": "phase 11's C = phase 9's D + the table"},
{"op": "e4m3_requantise", "count": 16, "dst": "phase 12's A"},
{"op": "load_B", "byte_offset": 64448, "span": 1548},
]},
"E25": {"after": 25, "stmts": (15706, 15980), "ops": [
{"op": "e4m3_requantise", "count": 24, "kept": 16,
"dst": "phase 26's A"},
{"op": "plane_scatter", "arg": 48,
"index": "8*ex*y + 8*x + d, d = t or t|4",
"bases": ["%r488 + (laneid>>5)", "%r488 + ((g+8)>>3)",
"%r488 + ((g+16)>>3)", "%r488 + ((g+24)>>3)"],
"stores": 16, "guarded": "(y<0)|(y>=ey)|(slice>31)|(x<0)|(x>=ex)"},
{"op": "load_B", "byte_offset": 95744, "span": 524},
{"op": "load_bias", "byte_offset": 96768, "span": 16, "dst": "phase 26's C"},
]},
"E26": {"after": 26, "stmts": (15989, 16304), "ops": [
{"op": "lane_merge", "note": "12 guarded st.shared of the eight phase-26 D "
"registers, then a register-only row max over "
"5 u32 (10 f16) with no warp shuffle"},
{"op": "softmax", "form": "f32 exp2 with the row maximum subtracted",
"log2e": 0x3FB8AA3B},
{"op": "surface_scatter", "arg": 112,
"stencil": "3x3 colour tex at param_0+88 with the param_0+136 scales, "
"plus param_0+96 and param_0+104 at integer (x, y)",
"combine": "acc*sigmoid(v) + tex96*(1-sigmoid(v)) + tex104 (unless NaN)",
"channels": 3, "channel_pitch_bytes": 8, "fourth_channel": 0x0000},
]},
"_shared_band": {
"note": "E11/E13/E15/E17/E19/E21/E23 are 68-statement twins: 32 packs of "
"the same normalised tile, four v4 B loads and four per-(m,n) C "
"seed words. E12/E14/E16/E18/E20/E22/E24 are 308-statement twins "
"of E10's activation with no table add and no residual.",
"unverified": "the per-phase B and C byte offsets of phases 11-24 are "
"taken from rr_layer_spec.py's phase table (weight and bias "
"columns), which reads them off the load chains; they were "
"not re-derived from the ld.weak immediates here.",
},
}
if __name__ == "__main__":
import json
print(json.dumps({"layer": layer(), "patch_grid": e1_patch_grid(),
"phases": phases(), "epilogues": epilogues()}, indent=2))
Reference in new issue View Git Blame Copy Permalink