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).
231 lines
9.3 KiB
Python
231 lines
9.3 KiB
Python
"""Splice recipe for `cuda_dldn_engine_swin_dec0_kernel`.
|
|
|
|
Self-contained: `kernels/tex/make_ptx.py` imports this module by path and returns
|
|
`recipe()` for the kernel name, so nothing here may import make_ptx. Everything
|
|
below is derived from the corpus PTX with `kernels/rr/rr_layer_spec.py`'s own
|
|
parser (the sibling module, imported for its lineariser and mma collector): the
|
|
mma region is `s326-s18231` (file lines 1393-46263), 544 mma in 19 phases, the
|
|
weight image is 47376 bytes at arg040, the parameter block is 152 bytes, and the
|
|
launch is grid 161x93 with block 32x1x1 (one wave).
|
|
|
|
**Shape of the splice.** NVIDIA's plane/texture prologue (lines 1-1392) and its
|
|
surface epilogue (46264-end) are kept verbatim, and so is every epilogue between
|
|
the phases: E1's shared-window staging, the RMS norms, both attention softmaxes
|
|
and the clamped-cubic activations all stay NVIDIA's PTX, so the body only has to
|
|
reproduce the arithmetic the mma statements perform. Each phase's contiguous mma
|
|
run is replaced by one `call.uni` into `d4r_swin_dec0` with that phase's operands.
|
|
|
|
**Why the calls need spill slots.** `call.uni` has no ABI: every register the
|
|
surrounding PTX still needs across the call is destroyed. So each block
|
|
|
|
1. stores the registers that are live across this call (defined before the
|
|
phase's first mma, read after its last) into `d4r_spill`,
|
|
2. copies the phase's A/B/C operand registers into the parameter bank,
|
|
3. calls,
|
|
4. reads the phase's D registers back out of `d4r_spill` into the exact
|
|
register names the mma would have written, and
|
|
5. reloads the spilled registers.
|
|
|
|
The largest live-across set is 508 bytes per lane (phase 3), so the spill area
|
|
is `SPILL_SLOT` bytes per lane, and the largest D return is phase 1's 192 b32,
|
|
so the return area is `RET_SLOT` bytes per lane. Both are indexed off
|
|
`laneid * LANE_STRIDE` inside a single module-scope `.shared` array; the array is
|
|
declared in the `extern` string, which `make_ptx.py` emits at module scope (a
|
|
`.shared` declaration is not legal inside a block).
|
|
|
|
**Operand encoding.** The callee receives, in this order, the phase number, the
|
|
A fragments (four b32 each, in first-use order), the B fragments (two b32 each,
|
|
in first-use order), the C words (two b32 per distinct C operand, in first-use
|
|
order), the weight image, the parameter block and its shared row base. B
|
|
fragments and C operands that `rr_layer_spec` resolves into the weight image are
|
|
*not* passed: the callee reads them from arg040 at the offsets the mma resolved.
|
|
Chain C operands (a previous k step's D) and the constant `%r13609` are not
|
|
passed either: the callee computes them.
|
|
"""
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
_HERE = Path(__file__).resolve().parent
|
|
if str(_HERE) not in sys.path:
|
|
sys.path.insert(0, str(_HERE))
|
|
import rr_layer_spec as SPEC # noqa: E402 (path is fixed up above)
|
|
|
|
MODULE = "0023-PREPASS_ENTRYPOINT_NAME.ptx"
|
|
ENTRY = "cuda_dldn_engine_swin_dec0_kernel"
|
|
PARAM = ENTRY + "_param_0"
|
|
|
|
# the cut: phase 1's first mma (line 1393) through phase 19's last (line 46266);
|
|
# make_ptx.py replaces lines[before-1:after-1]
|
|
CUT = (1393, 46267)
|
|
|
|
# shared-memory geometry, bytes per lane
|
|
LANE_STRIDE = 1280
|
|
RET_SLOT = 768 # 192 b32, phase 1's D
|
|
SPILL_SLOT = 512 # the widest live-across set is 508
|
|
SPILL = RET_SLOT # spill area starts after the return area
|
|
SPILL_BYTES = LANE_STRIDE * 32
|
|
|
|
# the callee's parameter bank: MAXP b32 then b64, b64, b32. Phase 4 passes the
|
|
# most (137); rrswin_dec0.hip's d4r_swin_dec0 must declare the same count.
|
|
MAXP = 144
|
|
|
|
SPILL_SYMBOL = "d4r_spill"
|
|
|
|
|
|
def _ptx_dir():
|
|
return Path(os.environ.get("D4R_DLSS_PTX_DIR",
|
|
_HERE.parent / "extracted" / "ptx"))
|
|
|
|
|
|
def _key(reg):
|
|
return (reg[:3], int(reg[3:]) if reg[3:] else 0)
|
|
|
|
|
|
def _analyse():
|
|
"""Every number the splice needs, read out of the corpus PTX."""
|
|
lines = (_ptx_dir() / MODULE).read_text().split("\n")
|
|
body = next(((a, b) for n, a, b in SPEC.find_entries(lines) if n == ENTRY))
|
|
if body is None:
|
|
raise SystemExit(f"{MODULE}: no entry {ENTRY}")
|
|
stmts, _ = SPEC.linearise(lines, *body)
|
|
defs = SPEC.track_weights(stmts, ENTRY)
|
|
mmas = SPEC.collect_mma(stmts, defs)
|
|
phases = SPEC.build_phases(stmts, mmas)
|
|
|
|
last_w, last_r = {}, {}
|
|
for s in stmts:
|
|
for r in SPEC.dest_regs(s.text):
|
|
last_w[r] = s.idx
|
|
for r in SPEC.read_regs(s.text):
|
|
last_r[r] = max(last_r.get(r, 0), s.idx)
|
|
for group in SPEC.RE_OPERAND.findall(s.text):
|
|
for r in SPEC.REG_SRC.findall(group):
|
|
last_r[r] = max(last_r.get(r, 0), s.idx)
|
|
|
|
out = []
|
|
for p in phases:
|
|
lo, hi = p.mma[0].idx, p.mma[-1].idx
|
|
a_frag, b_frag, c_frag = [], [], []
|
|
c_index, d_regs = [], []
|
|
for m in p.mma:
|
|
a = tuple(m.regs("A"))
|
|
b = tuple(m.regs("B"))
|
|
if a not in a_frag:
|
|
a_frag.append(a)
|
|
if b not in b_frag:
|
|
b_frag.append(b)
|
|
d_regs.extend(m.regs("D"))
|
|
cregs = tuple(m.regs("C"))
|
|
if all(r in d_regs or r == "%r13609" for r in cregs):
|
|
c_index.append(-1) # chain or the zero constant
|
|
elif all(o is None for o in m.co):
|
|
c_index.append(-1) # resolved into the weight image
|
|
else:
|
|
if cregs not in c_frag:
|
|
c_frag.append(cregs)
|
|
c_index.append(c_frag.index(cregs))
|
|
out.append(dict(
|
|
idx=p.idx, lo=lo, hi=hi,
|
|
line0=p.mma[0].line, line1=p.mma[-1].line + 3,
|
|
a=[r for f in a_frag for r in f],
|
|
b=[r for f in b_frag for r in f if defs.from_weight(r, lo) is None],
|
|
c=[r for f in c_frag for r in f],
|
|
c_index=c_index, d=d_regs,
|
|
spill=sorted((r for r in last_w if last_w[r] < lo and last_r.get(r, 0) > hi),
|
|
key=_key),
|
|
))
|
|
return lines, out
|
|
|
|
|
|
def _extern():
|
|
"""The callee's declaration and the shared spill array, at module scope."""
|
|
params = ", ".join(f".param .b32 d4r_p{i}" for i in range(MAXP))
|
|
return (f".shared .align 4 .b8 {SPILL_SYMBOL}[{SPILL_BYTES}];\n"
|
|
f".extern .func d4r_swin_dec0\n(\n\t{params},\n"
|
|
"\t.param .b64 d4r_w, .param .b64 d4r_v, .param .b32 d4r_s\n)\n;\n")
|
|
|
|
|
|
def _spill_ops(ph):
|
|
"""The st.shared / ld.shared pair that carries one live-across register."""
|
|
lines = []
|
|
for j, reg in enumerate(ph["spill"]):
|
|
off = SPILL + 4 * j
|
|
if reg.startswith("%rd"):
|
|
lines.append(f"\tmov.b64 {{%rtlo,%rthi}}, {reg};")
|
|
lines.append(f"\tst.shared.v2.b32 [%rk+{off}], {{%rtlo,%rthi}};")
|
|
elif reg.startswith("%p"):
|
|
lines.append(f"\tselp.b32 %rtlo, 1, 0, {reg};")
|
|
lines.append(f"\tst.shared.u32 [%rk+{off}], %rtlo;")
|
|
else:
|
|
lines.append(f"\tst.shared.u32 [%rk+{off}], {reg};")
|
|
return lines
|
|
|
|
|
|
def _reload_ops(ph):
|
|
lines = []
|
|
for j, reg in enumerate(ph["spill"]):
|
|
off = SPILL + 4 * j
|
|
if reg.startswith("%rd"):
|
|
lines.append(f"\tld.shared.v2.b32 {{%rtlo,%rthi}}, [%rk+{off}];")
|
|
lines.append(f"\tmov.b64 {reg}, {{%rtlo,%rthi}};")
|
|
elif reg.startswith("%p"):
|
|
lines.append(f"\tld.shared.u32 %rtlo, [%rk+{off}];")
|
|
lines.append(f"\tsetp.ne.u32 {reg}, %rtlo, 0;")
|
|
else:
|
|
lines.append(f"\tld.shared.u32 {reg}, [%rk+{off}];")
|
|
return lines
|
|
|
|
|
|
def _block(ph):
|
|
"""The `{ ... }` that replaces one phase's mma run."""
|
|
k = ph["idx"]
|
|
n = 1 + len(ph["a"]) + len(ph["b"]) + len(ph["c"])
|
|
if n > MAXP:
|
|
raise SystemExit(f"phase {k} needs {n} parameters, MAXP is {MAXP}")
|
|
decl = "\n".join(f"\t.param .b32 q{k}_{i};" for i in range(n))
|
|
decl += f"\n\t.param .b64 qw{k};\n\t.param .b64 qv{k};\n\t.param .b32 qs{k};"
|
|
zero = "\n".join(f"\tst.param.b32 [q{k}_{i}+0], 0;" for i in range(n))
|
|
operands = [f"\tst.param.b32 [q{k}_{i + 1}+0], {r};"
|
|
for i, r in enumerate(ph["a"] + ph["b"] + ph["c"])]
|
|
args = ", ".join([f"q{k}_{i}" for i in range(n)] + [f"qw{k}", f"qv{k}", f"qs{k}"])
|
|
body = [
|
|
f"mov.u32 %rk, {SPILL_SYMBOL};",
|
|
f"mad.lo.s32 %rk, %laneid, {LANE_STRIDE}, %rk;",
|
|
]
|
|
body += _spill_ops(ph)
|
|
body += [decl, zero, "\n".join(operands)]
|
|
body += [
|
|
f"ld.param.u64 %rv, [{PARAM}+40];",
|
|
f"st.param.b64 [qw{k}+0], %rv;",
|
|
f"mov.b64 %rv, {PARAM};",
|
|
f"st.param.b64 [qv{k}+0], %rv;",
|
|
f"st.param.b32 [qs{k}+0], %rk;",
|
|
f"call.uni d4r_swin_dec0, ({args});",
|
|
]
|
|
body += [f"\tld.shared.u32 {ph['d'][i]}, [%rk+{4 * i}];" for i in range(len(ph["d"]))]
|
|
body += _reload_ops(ph)
|
|
return "\n".join(["{",
|
|
".reg .b32 %rk, %rv, %rtlo, %rthi;",
|
|
"\n".join(body),
|
|
"}"])
|
|
|
|
|
|
def _call():
|
|
lines, phases = _analyse()
|
|
out = []
|
|
cursor = CUT[0]
|
|
for ph in phases:
|
|
if not ph["spill"] and not ph["c"] and len(ph["a"]) + len(ph["b"]) > 0:
|
|
pass
|
|
out += lines[cursor - 1:ph["line0"] - 1]
|
|
out.append(_block(ph))
|
|
cursor = ph["line1"] + 1
|
|
out += lines[cursor - 1:CUT[1] - 1]
|
|
return "\n".join(out)
|
|
|
|
|
|
def recipe():
|
|
return dict(file=MODULE, cut_call=CUT, extern=_extern(), call=_call()) |