woutermans
0fe3a28cb7
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).
221 lines
10 KiB
Python
Executable File
221 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""The splice of `cuda_dldn_engine_swin_enc2_kernel` (nvngx_dlssd.dll 310.7).
|
|
|
|
One warp per 8x8 token patch, grid (161, 93, 1), block (32, 1, 4); the 630 mma of the
|
|
38 phases are statements s4621-s13876, file lines 7669-33949 of
|
|
`~/.cache/d4r-rr-corpus/0015-PREPASS_ENTRYPOINT_NAME.ptx`. Everything outside that
|
|
region stays NVIDIA's: the input-texture prologue (s1-s4620) and the final epilogue
|
|
(s13877-s14199, the staged softmax window, the 3x3 texture gather and the three
|
|
surface writes). Inside it, the region becomes one `call.uni` into the native body
|
|
`kernels/rr/rrswin_enc2.hip`.
|
|
|
|
The call hands over phase 1's staged A fragments (six 16-row fragments read out of the
|
|
entry's own `.shared` tile), phase 1's B fragments, its zero C seed, the weight image
|
|
(arg040), the parameter block and the base of that same shared tile; the five
|
|
registers the epilogue reads but the region computes cross the cut through shared
|
|
memory, because ZLUDA lowers `call.uni` return parameters into the callee's LLVM
|
|
return type rather than a .param-space object (see kernels/rr/rrswin_enc0.hip).
|
|
|
|
Every byte offset, register list and insertion line below was read out of the PTX with
|
|
`kernels/rr/rr_layer_spec.py enc2` and the stage table is regenerated by hand from its
|
|
mma chain, so `D4R_RRSWIN_DEBUG=<stage>` splices a store of *this* layer's registers
|
|
into NVIDIA's own body and the native body's `-DD4R_RRSWIN_DEBUG_STAGE=<n>` writes the
|
|
same values in the same order.
|
|
|
|
Self-contained by contract: nothing is imported from kernels/tex/make_ptx.py.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# the cut
|
|
FILE = "0015-PREPASS_ENTRYPOINT_NAME.ptx"
|
|
# (first line of phase 1's first mma, first line after phase 38's last mma)
|
|
CUT_CALL = (7669, 33950)
|
|
|
|
# phase 1's A operands, in the first mma's operand order: six fragments of four,
|
|
# m-tile major. The entry stages eighteen of them (three K=32 slices x six m-tiles)
|
|
# at `smem + 512*(3*slice + m-tile) + 16*laneid`, so the body reads the other two
|
|
# slices from that same tile itself.
|
|
A_REGS = ["%r16120", "%r16121", "%r16122", "%r16123",
|
|
"%r16132", "%r16133", "%r16134", "%r16135",
|
|
"%r16144", "%r16145", "%r16146", "%r16147",
|
|
"%r16156", "%r16157", "%r16158", "%r16159",
|
|
"%r16168", "%r16169", "%r16170", "%r16171",
|
|
"%r16180", "%r16181", "%r16182", "%r16183"]
|
|
# phase 1's B operands: four fragments of two, the lane's 16 bytes of the two
|
|
# 512-byte weight tiles at arg040+192 and arg040+704
|
|
B_REGS = ["%r3251", "%r3252", "%r3253", "%r3254", "%r3255", "%r3256", "%r3257", "%r3258"]
|
|
# phase 1's C seed: %r8606 = 0, the whole region's zero accumulator start
|
|
C_REGS = ["%r8606"]
|
|
|
|
# what the epilogue reads but the region computes: param_0[0] (s13138), phase 38's
|
|
# two accumulator pairs (the staged softmax window) and the two values E38 itself
|
|
# needs from inside the region, %r17746 = 16*tid.z (s13139) and %r17758 = &smem (s13852)
|
|
RET_REGS = ["%r16472", "%r17496", "%r17506", "%r17746", "%r17758"]
|
|
SCRATCH = "_ZZ33cuda_dldn_engine_swin_enc2_kernel33DldnEngineSwinEncParamsStructBaseE4smem"
|
|
RET_STRIDE = 4 * len(RET_REGS)
|
|
|
|
CALLEE = "d4r_swin_enc2"
|
|
PARAM = "cuda_dldn_engine_swin_enc2_kernel_param_0"
|
|
|
|
|
|
def _extern():
|
|
"""The `.extern .func` declaration: the operand registers, the weight image, the
|
|
parameter block and the shared scratch, in that order."""
|
|
args = ", ".join([f".param .b32 d4r_a{i}" for i in range(len(A_REGS) + len(B_REGS) + len(C_REGS))]
|
|
+ [".param .b64 d4r_w", ".param .b64 d4r_v", ".param .b32 d4r_s"])
|
|
return f".extern .func {CALLEE}\n(\n\t{args}\n)\n;\n"
|
|
|
|
|
|
def _call():
|
|
"""The splice: hand the operands over, call, read the five accumulators back."""
|
|
regs = A_REGS + B_REGS + C_REGS
|
|
decl = "\n".join(f"\t.param .b32 p{i};" for i in range(len(regs)))
|
|
decl += "\n\t.param .b64 pw;\n\t.param .b64 pv;\n\t.param .b32 ps;"
|
|
store_args = "\n".join(f"\tst.param.b32 [p{i}+0], {r};" for i, r in enumerate(regs))
|
|
store_args += (f"\n\tst.param.b64 [pw+0], %rd2;\n\tst.param.b64 [pv+0], %rd8;"
|
|
f"\n\tmov.u32 %rswin0, {SCRATCH};\n\tst.param.b32 [ps+0], %rswin0;")
|
|
hel = ", ".join([f"p{i}" for i in range(len(regs))] + ["pw", "pv", "ps"])
|
|
gets = "\n".join(f"\tld.shared.b32 {r}, [%rswin1+{4 * i}];"
|
|
for i, r in enumerate(RET_REGS))
|
|
return f"""{{
|
|
.reg .b32 %rswin<2>;
|
|
{decl}
|
|
{store_args}
|
|
\tcall.uni {CALLEE}, ({hel});
|
|
\tmad.lo.s32 %rswin1, %laneid, {RET_STRIDE}, %rswin0;
|
|
{gets}
|
|
}}"""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Ground truth for a stage: NVIDIA's own body with a store of that stage's registers
|
|
# spliced into the plane arena, past everything the kernel writes, which the launch
|
|
# dump then holds for a block/lane.
|
|
#
|
|
# The arena is arg048, and the launch's arg048 sits at allocation offset 30,146,560 of
|
|
# a 65,945,600-byte allocation -- the bridge dumps that whole allocation under arg008
|
|
# (arg008 is the first argument pointing into it), so only 35,799,040 bytes of arena are
|
|
# dumped. The family default of 46,000,000 lands past the end of the allocation, which
|
|
# is why the first sweep of this layer found no payload at all. 4,000,000 is past
|
|
# everything E36 scatters -- its stores are 4 bytes at idx = 8*ex*y + 8*x, so
|
|
# idx < 8*320*184 = 471,040 -- and leaves room for 64 slots of up to 32 KiB.
|
|
DEBUG_BASE = int(os.environ.get("D4R_RRSWIN_DEBUG_BASE", 4000000))
|
|
DEBUG_ARENA_BYTES = 35799040
|
|
DEBUG_SLOT_FLOOR = 8192
|
|
GRID_X = 41 # grid (41, 24, 1), block (32, 1, 4): 320/8+1 by 184/8+1 patches
|
|
|
|
if DEBUG_BASE + 64 * 65536 >= DEBUG_ARENA_BYTES:
|
|
raise SystemExit("swin_enc2_recipe: the debug payload does not fit enc2's plane arena")
|
|
|
|
|
|
def _debug_block(regs):
|
|
"""PTX storing REGS (u32 each) to plane arena + DEBUG_BASE + block*slot + lane*len*4."""
|
|
want = max(DEBUG_SLOT_FLOOR, 32 * len(regs) * 4)
|
|
slot = 1 << (want - 1).bit_length()
|
|
stores = "\n".join(
|
|
f"@%pdbg1 st.global.v4.u32 [%rddbg3+{i * 16}], {{{', '.join(regs[i * 4:i * 4 + 4])}}};"
|
|
for i in range((len(regs) + 3) // 4))
|
|
return f"""// d4r debug stage (enc2)
|
|
{{
|
|
.reg .pred %pdbg<3>;
|
|
.reg .b32 %rdbg<9>;
|
|
.reg .b64 %rddbg<4>;
|
|
ld.param.u64 %rddbg1, [{PARAM}+48];
|
|
mov.u32 %rdbg1, %ctaid.x;
|
|
mov.u32 %rdbg2, %ctaid.y;
|
|
mad.lo.s32 %rdbg3, %rdbg2, {GRID_X}, %rdbg1;
|
|
// only wave 0 of the block may store: the four z-planes read different weight slabs
|
|
// (the 21504-byte stride), so their accumulators differ
|
|
mov.u32 %rdbg8, %tid.z;
|
|
setp.ne.u32 %pdbg2, %rdbg8, 0;
|
|
selp.b32 %rdbg3, 64, %rdbg3, %pdbg2;
|
|
setp.lt.u32 %pdbg1, %rdbg3, 64;
|
|
shl.b32 %rdbg4, %rdbg3, {slot.bit_length() - 1};
|
|
mov.u32 %rdbg5, %laneid;
|
|
mad.lo.s32 %rdbg6, %rdbg5, {len(regs) * 4}, %rdbg4;
|
|
add.s32 %rdbg7, %rdbg6, {DEBUG_BASE};
|
|
cvt.u64.u32 %rddbg2, %rdbg7;
|
|
add.s64 %rddbg3, %rddbg1, %rddbg2;
|
|
{stores}
|
|
}}
|
|
"""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Debug stages. "e2p<n>" stores phase n's accumulator registers in mma order (m-tile
|
|
# major, n-tile minor, each mma's (row g, row g+8) pair); the line is the last physical
|
|
# line of the phase's last mma, so the store lands between the mma block and its
|
|
# epilogue. Generated from kernels/rr/rr_layer_spec.py's own mma chain.
|
|
#
|
|
# The names carry the "e2" prefix because kernels/tex/make_ptx.py's swin_debug() falls
|
|
# back to enc0's *unkeyed* stage table when a name is absent from the keyed per-layer
|
|
# one, so a bare "p3" would splice enc0's line 7840 into enc2's PTX. With a prefix
|
|
# enc2's names miss that table, swin_debug() returns None and recipe() below handles
|
|
# the stage itself -- which is required anyway, because enc2's block is keyed by
|
|
# (ctaid.x, ctaid.y) with a plane-z guard (see _debug_block), not enc0's flat index.
|
|
STAGES = {
|
|
# the eighteen staged A fragments the entry reads before the cut, in mma operand
|
|
# order: phase 1's (slice 0), phase 3's (slice 1), phase 5's (slice 2), m-tile major
|
|
"e2a": (7668, [f"%r{16120 + 4 * q}" for q in range(72)]),
|
|
# E6(a): the 96 C seeds of phase 7, the 64x96 f16 score bias; the last load is line 8840
|
|
"e2e6c": (8840, sum(([f"%r{4745 + 4 * o + j}" for j in range(2)] for o in range(24)), [])),
|
|
# E6(b): phase 7's 16 A registers (four fragments, m-tile major) and its 24 B
|
|
# registers (twelve fragments, n-tile major); the last pack is line 9121
|
|
"e2e6a": (9121, sum(([f"%r{5033 + 120 * q + o}" for o in range(4)] for q in range(4)), [])),
|
|
"e2e6b": (9121, sum(([f"%r{5287 + 10 * q + o}" for o in range(2)] for q in range(12)), [])),
|
|
# E7(g): phase 8's 48 A registers (the softmax probabilities) and its 24 B
|
|
# registers (the movmatrix transpose of phase 6's D); the last pack is line 19600
|
|
"e2e7a": (19600, sum(([f"%r{8239 + 120 * q + o}" for o in range(4)] for q in range(12)), [])),
|
|
"e2e7b": (19600, sum(([f"%r{8533 + 10 * q + o}" for o in range(2)] for q in range(12)), [])),
|
|
}
|
|
|
|
|
|
def _phase_stages():
|
|
"""The accumulator registers of every phase, straight from the linearisation."""
|
|
import re
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
import rr_layer_spec as spec
|
|
|
|
path, entry = spec.resolve("enc2", None, None)
|
|
lines = open(path).read().split("\n")
|
|
first = last = None
|
|
for name, a, b in spec.find_entries(lines):
|
|
if name == entry:
|
|
first, last = a, b
|
|
break
|
|
stmts, _ = spec.linearise(lines, first, last)
|
|
phases = spec.build_phases(stmts, spec.collect_mma(stmts, spec.track_weights(stmts, entry)))
|
|
out = {}
|
|
for p in phases:
|
|
i = p.mma[-1].line
|
|
while not lines[i - 1].rstrip().endswith(";"):
|
|
i += 1
|
|
regs = []
|
|
for m in p.mma:
|
|
regs += [r.strip() for r in m.D.split(",")]
|
|
out[f"e2p{p.idx}"] = (i, regs)
|
|
return out
|
|
|
|
|
|
|
|
def stages():
|
|
"""Every stage this layer can capture: the operand stages above plus the phases."""
|
|
if "e2p1" not in STAGES:
|
|
STAGES.update(_phase_stages())
|
|
return STAGES
|
|
|
|
|
|
|
|
|
|
def recipe():
|
|
stage = os.environ.get("D4R_RRSWIN_DEBUG", "")
|
|
if stage:
|
|
if stage not in stages():
|
|
sys.exit(f"swin_enc2: no debug stage {stage!r} (have {', '.join(sorted(stages()))})")
|
|
line, regs = stages()[stage]
|
|
return dict(file=FILE, debug=(line, regs), debug_block=_debug_block)
|
|
return dict(file=FILE, cut_call=CUT_CALL, call=_call(), extern=_extern())
|