#!/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=` splices a store of *this* layer's registers into NVIDIA's own body and the native body's `-DD4R_RRSWIN_DEBUG_STAGE=` 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" 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())