wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

kernels/rr: staged-image Swin layer engine prototypes

Browse Source 
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).
This commit is contained in:
woutermans committed 2026-10-12 12:47:53 +02:00
1 parent 798e25611c
commit 0fe3a28cb7
58 files changed
+28330 -5
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/tex/make_ptx.py
+327 -2
View File
@@ -3,6 +3,7 @@
usage: make_ptx.py KERNEL OUT.ptx
"""
import importlib.util
import os
import re
import sys
@@ -11,6 +12,76 @@ from pathlib import Path
# PTX extracted from nvngx_dlss.dll by kernels/tools/extract_dlss_ptx.py (kernels/build.sh does this)
PTX_DIR = Path(os.environ.get("D4R_DLSS_PTX_DIR", Path(__file__).resolve().parent.parent / "extracted" / "ptx"))
# Ray Reconstruction enc0 mirror cut: PTX keeps the prologue (textures, input staging, the A/B/bias
# register setup up to line 3963) and the surface epilogue (from line 44528); the 536 mma of the 22
# phases (lines 3964-44527) become one call to the native body, which is handed the six staged A
# fragments, the four bias seeds, the weight image and the parameter block. Phase 22's eight mma each
# write a *pair* of b32 (four f16 per lane), and the epilogue's `selp` chains (stmts 15595-15609,
# 15643-15676, 15718-15720) select between the halves on `laneid & 3`, so both halves of every pair
# must cross the cut -- returning only the even halves (the previous list) left the odd ones undefined
# in the spliced PTX and made byte-exact output impossible.
#
# The eighteen accumulators come back through shared memory, not `call.uni`'s return parameters.
# ZLUDA lowers those to the *function's* LLVM return type (llvm_zluda/ptx/src/pass/llvm/emit.rs,
# emit_call: one value, or an `get_or_create_struct_type` struct for several, distributed with
# LLVMBuildExtractValue). There is no .param-space object behind them and no pointer is ever formed,
# so the callee sees 0x0 in all eighteen slots and stores through them fault the launch; an 18 x i32
# C++ struct return lowers to an sret pointer, which is a different LLVM signature and will not link.
#
# So the callee is handed the address of `input_tensor` (3200 bytes of .shared, dead after the cut --
# its last ld.shared is line 3938 and no register derived from its base is read past it) and writes
# the eighteen values at `base + 72*laneid + 4*i`; the spliced PTX reads them back with eighteen
# ld.shared.b32. 72 * 32 lanes = 2304 bytes, inside the 3200. Per-block LDS is what makes this safe:
# a global scratch slot would be shared by all 14973 blocks, and sizing it per block and lane
# (14973 * 32 * 72 = 34.5 MB) would not fit past the 45,219,840 arena bytes the kernel writes either.
# enc0_tail already uses this shape (d4r_enc0_tail takes the shared base as a plain b32 argument).
SWIN_A_REGS = ["%r120", "%r136", "%r128", "%r144", "%r152", "%r168", "%r160", "%r176",
"%r184", "%r200", "%r192", "%r208", "%r216", "%r232", "%r224", "%r240",
"%r246", "%r257", "%r253", "%r1042", "%r1079", "%r1080", "%r1081", "%r1082"]
SWIN_BIAS_REGS = ["%r1086", "%r1096", "%r1106", "%r1116"]
SWIN_RET_REGS = ["%r17597", "%r18203"] + sum(
([f"%r{d}", f"%r{d + 1}"] for d in range(17634, 17706, 10)), [])
# the .shared scratch the accumulators cross the cut through: 18 b32 per lane, 72 bytes apart
SWIN_RET_SCRATCH = "_ZZ33cuda_dldn_engine_swin_enc0_kernel35DldnEngineFusedSwinEnc0ParamsStructE12input_tensor"
SWIN_RET_STRIDE = 4 * len(SWIN_RET_REGS)
SWIN_P2_A_REGS = ["%r1207", "%r1208", "%r1209", "%r1210", "%r1247", "%r1248", "%r1249", "%r1250",
"%r1287", "%r1288", "%r1289", "%r1290", "%r1327", "%r1328", "%r1329", "%r1330",
"%r1367", "%r1368", "%r1369", "%r1370", "%r1407", "%r1408", "%r1409", "%r1410"]
def swin_cut_call():
"""The native body is handed the staged A fragments, the four bias seeds, the weight image, the
parameter block and the address of the shared scratch; it writes the eighteen accumulators there
and this splice reads them back into the epilogue's registers."""
args = SWIN_A_REGS + SWIN_BIAS_REGS
decl = "\n".join(f"\t.param .b32 p{i};" for i in range(len(args)))
store_args = "\n".join(f"\tst.param.b32 [p{i}+0], {r};" for i, r in enumerate(args))
decl += "\n\t.param .b64 pw;\n\t.param .b64 pv;\n\t.param .b32 ps;"
store_args += ("\n\tst.param.b64 [pw+0], %rd367;\n\tst.param.b64 [pv+0], %rd368;"
f"\n\tmov.u32 %rswin0, {SWIN_RET_SCRATCH};"
"\n\tst.param.b32 [ps+0], %rswin0;")
hel = ", ".join([f"p{i}" for i in range(len(args))] + ["pw", "pv", "ps"])
gets = "\n".join(f"\tld.shared.b32 {r}, [%rswin1+{4 * i}];"
for i, r in enumerate(SWIN_RET_REGS))
return f"""{{
.reg .b32 %rswin<2>;
{decl}
{store_args}
\tcall.uni d4r_swin_enc0, ({hel});
\tmad.lo.s32 %rswin1, %laneid, {SWIN_RET_STRIDE}, %rswin0;
{gets}
}}"""
def swin_cut_extern():
args = ", ".join([f".param .b32 d4r_a{i}" for i in range(len(SWIN_A_REGS) + len(SWIN_BIAS_REGS))]
+ [".param .b64 d4r_w", ".param .b64 d4r_v", ".param .b32 d4r_s"])
return f".extern .func d4r_swin_enc0\n(\n\t{args}\n)\n;\n"
ENC0_TAIL = """{{
.param .b32 q0;
st.param.b32 [q0+0], {smem};
@@ -33,6 +104,30 @@ ret;
}}
"""
SWIN_EXTERN = """.extern .func d4r_swin_enc0
(
.param .b64 d4r_e0, .param .b32 d4r_e1, .param .b32 d4r_e2
)
;
"""
# NVIDIA's own entry takes the address of its parameter block (`mov.b64 %rd368, <param_0>`), so the
# native body is handed that pointer plus the two shared arrays' addresses. The registers are ones the
# entry already declares; the body between them is deleted.
SWIN_STUB = """mov.b64 %rd376, cuda_dldn_engine_swin_enc0_kernel_param_0;
mov.u32 %r18228, _ZZ33cuda_dldn_engine_swin_enc0_kernel35DldnEngineFusedSwinEnc0ParamsStructE11input_noise;
mov.u32 %r18227, _ZZ33cuda_dldn_engine_swin_enc0_kernel35DldnEngineFusedSwinEnc0ParamsStructE12input_tensor;
{
.param .b64 q0;
st.param.b64 [q0+0], %rd376;
.param .b32 q1;
st.param.b32 [q1+0], %r18228;
.param .b32 q2;
st.param.b32 [q2+0], %r18227;
call.uni d4r_swin_enc0, (q0, q1, q2);
}
ret;"""
KERNELS = {
# enc0: cut after the bar.warp.sync that ends the smem staging of the GEMM A operand
"rrlite_enc0_4x4_mvhi_hdr_folded": dict(
@@ -83,6 +178,15 @@ KERNELS = {
call.uni d4r_dec0_head, (q0, q1, q2, q3, q4, q5, q6);
}""",
),
# Ray Reconstruction denoiser Swin layers (kernels/rr/rrswin_*.hip): the whole entry body of the
# module's Shwin entry is replaced by a call into the native implementation, linked in as bitcode.
# PREPASS_ENTRYPOINT_NAME, the module's other entry, is left untouched.
"cuda_dldn_engine_swin_enc0_kernel": dict(
file="0013-PREPASS_ENTRYPOINT_NAME.ptx",
replace_body=True,
extern=SWIN_EXTERN,
stub=SWIN_STUB,
),
# post and downsample: only the formatted surface stores become native (inline format conversion
# instead of ZLUDA's per-channel surface_formatted_store_bits calls)
"rrlite_post_3_2_mvhi_hdr_folded": dict(file="dlss-0176-00.ptx", sust=True, extern=""),
@@ -497,7 +601,193 @@ def enc0_variant(name):
tail=ENC0_TAIL.format(**regs))
SWIN_DEBUG_BASE = 46000000 # bytes into the plane arena (arg048) the kernel never writes
SWIN_DEBUG_SLOT = 8192 # the per-block slot floor: 32 lanes x 64 registers x 4 bytes
def swin_debug_slot(nregs):
"""Per-block slot for a stage of NREGS registers: 32 lanes x NREGS x 4 bytes, at least
SWIN_DEBUG_SLOT, rounded up to a power of two so the splice's `shl` can form block * slot."""
want = max(SWIN_DEBUG_SLOT, 32 * nregs * 4)
return 1 << (want - 1).bit_length()
def swin_p5_regs():
"""Phase 5's 48 mma D pairs in the mma statement order: m-tile major, n-tile minor, the row
pair (row g, row g+8) each mma writes. 2566 + 120*mi + 10*tt is the first mma of group mi."""
return sum(([f"%r{2566 + 120 * mi + 10 * tt + rh}" for rh in (0, 1)]
for mi in range(4) for tt in range(12)), [])
def swin_p5_operand_regs():
"""Phase 5's 48 mma C pairs (the position bias: 24 tiles of four words, tile T at 2390 + 4*T) and
the twelve B fragment pairs (the e4m3 keys, 2932 + 10*tt) -- the operands the 48 mma read."""
return (sum(([f"%r{2390 + 4 * t + w}" for w in range(4)] for t in range(24)), [])
+ sum(([f"%r{2932 + 10 * t}", f"%r{2933 + 10 * t}"] for t in range(12)), []))
def swin_e5_regs():
"""Phase 6's twelve A fragments (four b32 each, mma operand order) and twelve B fragment pairs:
the softmax probabilities and the transposed V that E5 packs."""
return (sum(([f"%r{5884 + 120 * (a // 3) + 20 * (a % 3) + o}" for o in range(4)]
for a in range(12)), [])
+ sum(([f"%r{6178 + 10 * b}", f"%r{6179 + 10 * b}"] for b in range(12)), []))
def swin_e5a_regs():
"""Phase 6's twelve A fragments alone: the softmax probabilities E5 packs."""
return sum(([f"%r{5884 + 120 * (a // 3) + 20 * (a % 3) + o}" for o in range(4)]
for a in range(12)), [])
def swin_p6_regs():
"""Phase 6's sixteen mma D pairs, each the third (last) k-step of its chain: m-tile major,
n-tile minor, the row pair (row g, row g+8)."""
return sum(([f"%r{5852 + 120 * mi + 10 * nt + 40 * (nt >> 1) + rh}" for rh in (0, 1)]
for mi in range(4) for nt in range(4)), [])
def swin_e6_regs():
"""Phase 7's sixteen A registers (four fragments of four) and its thirty-two C registers (the
residual plus the per-column table), m-tile major and n-tile minor."""
return (sum(([f"%r{6366 + 40 * mi + o}" for o in range(4)] for mi in range(4)), [])
+ sum(([f"%r{6342 + 40 * mi + 10 * nt + rh}" for rh in (0, 1)]
for mi in range(4) for nt in range(4)), []))
def swin_p7_regs():
"""Phase 7's sixteen mma D pairs, m-tile major and n-tile minor."""
return sum(([f"%r{6334 + 10 * (4 * mi + nt) + rh}" for rh in (0, 1)]
for mi in range(4) for nt in range(4)), [])
# a debug stage: (line after which to insert, registers to store); the per-block slot follows from
# the register count, so a stage is reproducible from its tuple alone
SWIN_DEBUG_STAGES = {
# the staged block input: the 24 A registers in the first mma's operand order plus the four bias seeds
"a": (3962, ["%r120", "%r136", "%r128", "%r144", "%r152", "%r168", "%r160", "%r176",
"%r184", "%r200", "%r192", "%r208", "%r216", "%r232", "%r224", "%r240",
"%r246", "%r257", "%r253", "%r1042", "%r1079", "%r1080", "%r1081", "%r1082",
"%r1086", "%r1096", "%r1106", "%r1116"]),
# phase 1's 24 accumulator register pairs (m-tile major, n-tile minor)
"p1": (4128, ["%r877", "%r878", "%r887", "%r888", "%r897", "%r898", "%r907", "%r908",
"%r917", "%r918", "%r927", "%r928", "%r937", "%r938", "%r947", "%r948",
"%r957", "%r958", "%r967", "%r968", "%r977", "%r978", "%r987", "%r988",
"%r997", "%r998", "%r1007", "%r1008", "%r1017", "%r1018", "%r1027", "%r1028",
"%r1037", "%r1038", "%r1047", "%r1048", "%r1057", "%r1058", "%r1067", "%r1068",
"%r1077", "%r1078", "%r1087", "%r1088", "%r1097", "%r1098", "%r1107", "%r1108"]),
# phase 2's accumulator pairs: its last mma's `;` is line 4489
"p2": (4489, sum(([f"%r{d + 10 * i + j}" for i in range(4) for j in range(2)]
for d in range(1175, 1406, 40)), [])),
# phase 3's (the fused Q projection); its last mma ends on line 7840
"p3": (7840, sum(([f"%r{d + 10 * i + j}" for i in range(4) for j in range(2)]
for d in range(1900, 2131, 40)), [])),
# phase 4's (V); its last mma ends on line 8021
"p4": (8021, sum(([f"%r{d + 10 * i + j}" for i in range(4) for j in range(2)]
for d in range(2149, 2380, 40)), [])),
# phase 2's A operands as the epilogue packed them (the last bias load is line 4323)
"a2": (4323, SWIN_P2_A_REGS),
# phase 5 (the score GEMM, Q·Kᵀ with the position bias); its last mma ends on line 8741
"p5": (8741, swin_p5_regs()),
# phase 5's C (bias) and B (key) operand registers, as the mma reads them
"p5o": (8741, swin_p5_operand_regs()),
# E5's output: phase 6's A and B fragments; the last pack is line 18883
"e5": (18883, swin_e5_regs()),
# E5's probabilities alone: the fragment's B side is a movmatrix of phase 4's D, which the
# scalar body only reaches through phase 6's accumulators, so it is captured by p6 instead
"e5a": (18883, swin_e5a_regs()),
# phase 6 (P·V); its last mma ends on line 19220
"p6": (19220, swin_p6_regs()),
# E6's output: phase 7's A fragments and C seeds; the last pack is line 19643
"e6": (19643, swin_e6_regs()),
# phase 7 (the output projection); its last mma ends on line 19756
"p7": (19756, swin_p7_regs()),
}
# Per-layer extra debug stages, so a layer can add its own without editing this file:
# kernels/rr/swin_<layer>_debug_stages.py exposing stages() -> {name: (line, [regs])}. They are keyed
# by (layer, stage) because two layers can use the same stage name; the entries above stay unkeyed
# and belong to enc0.
for _extra in sorted((Path(__file__).resolve().parent.parent / "rr").glob("swin_*_debug_stages.py")):
_layer = _extra.stem[len("swin_"):-len("_debug_stages")]
_spec = importlib.util.spec_from_file_location(_extra.stem, _extra)
_mod = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_mod)
for _stage, _where in _mod.stages().items():
SWIN_DEBUG_STAGES[(_layer, _stage)] = _where
def swin_debug_block(regs, entry="cuda_dldn_engine_swin_enc0_kernel", base=None):
"""PTX that stores REGS (u32 each) to the output buffer + base + block*slot + lane*len(regs)*4.
The region must be past the bytes the kernel itself writes; enc0's plane arena is 45,219,840 bytes
and a layer whose output buffer is smaller sets D4R_RRSWIN_DEBUG_BASE lower (both sides of the
comparison read the same value). Blocks above 63 are skipped so the first 64 blocks (one warp
each) land in distinct slots. Works on either side of the splice: it is inserted into NVIDIA's own
PTX for the reference and reproduced by the native body's debug flag.
"""
if base is None:
base = int(os.environ.get("D4R_RRSWIN_DEBUG_BASE", SWIN_DEBUG_BASE))
slot = swin_debug_slot(len(regs))
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
{{
.reg .pred %pdbg<2>;
.reg .b32 %rdbg<8>;
.reg .b64 %rddbg<4>;
ld.param.u64 %rddbg1, [{entry}_param_0+48];
mov.u32 %rdbg1, %ctaid.x;
mov.u32 %rdbg2, %ctaid.y;
mad.lo.s32 %rdbg3, %rdbg2, 161, %rdbg1;
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, {base};
cvt.u64.u32 %rddbg2, %rdbg7;
add.s64 %rddbg3, %rddbg1, %rddbg2;
{stores}
}}
"""
def swin_debug(name):
"""Debug recipe: keep NVIDIA's whole body and store one stage's accumulator registers."""
stage = os.environ.get("D4R_RRSWIN_DEBUG", "")
match = re.fullmatch(r"cuda_dldn_engine_swin_(\w+)_kernel", name)
if not stage or match is None:
return None
where = SWIN_DEBUG_STAGES.get((match.group(1), stage))
if where is None:
where = SWIN_DEBUG_STAGES.get(stage)
if where is None:
return None
return dict(file=module_file(name, {"file": ""}).name, debug=where,
debug_block=lambda regs: swin_debug_block(regs, name))
def kernel_spec(name):
d = swin_debug(name)
if d:
return d
# Per-layer recipes: kernels/rr/swin_<layer>_recipe.py owns one denoiser layer's splice, so the
# ten layers can be written in parallel without sharing this file. Falls through to the built-in
# enc0 entry below when the module is absent.
swin = re.fullmatch(r"cuda_dldn_engine_swin_(\w+)_kernel", name)
if swin is not None:
module_path = Path(__file__).resolve().parent.parent / "rr" / f"swin_{swin.group(1)}_recipe.py"
if module_path.is_file():
spec = importlib.util.spec_from_file_location(f"swin_{swin.group(1)}_recipe", module_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module.recipe()
if name == "cuda_dldn_engine_swin_enc0_kernel" and os.environ.get("D4R_RRSWIN_CUT"):
return dict(file="0013-PREPASS_ENTRYPOINT_NAME.ptx", cut_call=(3964, 44528), extern=swin_cut_extern())
if name in KERNELS:
return KERNELS[name]
if ENC0_RE.match(name):
@@ -542,6 +832,41 @@ def main():
if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]:
sys.exit(f"lines {a}/{b} are {lines[a - 1]!r}/{lines[b - 1]!r}, expected {k['delete_check']}")
head, tail = lines[:a - 1] + k["insert"].split("\n") + lines[b:], []
elif "cut_call" in k:
before, after = k["cut_call"]
if not lines[before - 1].startswith("mma.sync.aligned.m16n8k32"):
sys.exit(f"{name}: cut line {before} is not an mma: {lines[before - 1]!r}")
# the line the cut ends on must be the last line of an mma statement; which mma that is
# depends on the layer, so this only checks the shape, not enc0's registers
if not lines[after - 2].rstrip().endswith(";"):
sys.exit(f"{name}: cut end {after} is not the end of an mma statement: {lines[after - 2]!r}")
head = lines[:before - 1] + k.get("call", swin_cut_call()).split("\n") + lines[after - 1:]
tail = []
elif "debug" in k:
after, regs = k["debug"]
window = "\n".join(lines[max(0, after - 4000):after])
prev = lines[after - 1].rstrip()
if regs[0] not in window or not (prev.endswith(";") or prev.endswith("//") or not prev):
sys.exit(f"{name}: debug line {after} is not phase 1's last mma: {lines[after - 1]!r}")
head, tail = lines[:after] + k.get("debug_block", swin_debug_block)(regs).split("\n") + lines[after:], []
elif "replace_body" in k:
e = next(i for i, l in enumerate(lines) if l.startswith(f".visible .entry {name}("))
b0 = next(i for i in range(e, len(lines)) if lines[i].strip() == "{")
j = b0 + 1
while j < len(lines) and (not lines[j].strip()
or lines[j].lstrip().startswith((".reg", ".shared", ".param", ".pragma", "//"))):
j += 1
if not lines[j - 1].lstrip().startswith("//") and lines[j - 1].strip() != "":
sys.exit(f"{name}: could not find the end of the entry declarations (line {j} is {lines[j - 1]!r})")
depth, b = 0, None
for i in range(b0, len(lines)):
depth += lines[i].count("{") - lines[i].count("}")
if i > b0 and depth == 0:
b = i
break
if b is None or lines[b].strip() != "}":
sys.exit(f"{name}: entry closing brace not found")
head, tail = lines[:j] + k["stub"].split("\n") + [""] + lines[b:], []
else:
head, tail = lines, []
entry = next(i for i, l in enumerate(head) if l.startswith(".visible .entry"))
@@ -560,8 +885,8 @@ def main():
body, n["dec8"] = replace_dec8(body)
if n["dec8"] != k["dec8"]:
sys.exit(f"{name}: {n['dec8']} e4m3 decodes, expected {k['dec8']}")
k = dict(k, extern=k["extern"] + DEC8_EXTERN)
extern = k["extern"] + (SUST_EXTERN if n["b16"] else "") + (SUST4_EXTERN if n.get("v4b16") else "") + (SUSTP_EXTERN if n["p"] else "") + (SUSTB32_EXTERN if n["b32"] else "")
k = dict(k, extern=k.get("extern", "") + DEC8_EXTERN)
extern = k.get("extern", "") + (SUST_EXTERN if n["b16"] else "") + (SUST4_EXTERN if n.get("v4b16") else "") + (SUSTP_EXTERN if n["p"] else "") + (SUSTB32_EXTERN if n["b32"] else "")
text = "\n".join(head[:entry]) + "\n" + extern + "\n".join(body + tail)
out.write_text(text)
print(f"{out}: {len(text.splitlines())} lines, surface stores replaced: {n}")