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).
44 lines
1.9 KiB
Python
44 lines
1.9 KiB
Python
"""Enc0 cut recipes; isolated phase 8 preserves NVIDIA's other phases."""
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def recipe():
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "tex"))
|
|
import make_ptx
|
|
|
|
if not os.environ.get("D4R_RRSWIN_PHASE8"):
|
|
if os.environ.get("D4R_RRSWIN_CUT"):
|
|
return dict(file="0013-PREPASS_ENTRYPOINT_NAME.ptx", cut_call=(3964, 44528),
|
|
extern=make_ptx.swin_cut_extern())
|
|
return make_ptx.KERNELS["cuda_dldn_engine_swin_enc0_kernel"]
|
|
|
|
calls = []
|
|
for mi, base in enumerate((6784, 6824, 6864, 6904, 6944, 6984)):
|
|
operands = [f"%r{base + i}" for i in range(4)]
|
|
operands += [f"%r{6495 + i}" for i in range(8)]
|
|
outputs = [f"%r{6503 + 40 * mi + 10 * nt + rh}" for nt in range(4) for rh in range(2)]
|
|
declarations = "\n".join(f"\t.param .b32 p8a{i};" for i in range(12))
|
|
stores = "\n".join(f"\tst.param.b32 [p8a{i}], {reg};" for i, reg in enumerate(operands))
|
|
loads = "\n".join(f"\tld.shared.b32 {reg}, [%p8out+{4 * i}];" for i, reg in enumerate(outputs))
|
|
args = ", ".join([f"p8a{i}" for i in range(12)] + ["p8s"])
|
|
calls.append(f"""{{
|
|
\t.reg .b32 %p8out;
|
|
{declarations}
|
|
\t.param .b32 p8s;
|
|
{stores}
|
|
\tst.param.b32 [p8s], %p8base;
|
|
\tcall.uni d4r_swin_enc0_p8, ({args});
|
|
\tmad.lo.u32 %p8out, %laneid, 32, %p8base;
|
|
{loads}
|
|
}}""")
|
|
# input_tensor's last read is in the prologue (line 3938); only the
|
|
# disjoint input_noise array is read by the final surface epilogue.
|
|
call = "{\n\t.reg .b32 %p8base;\n"
|
|
call += f"\tmov.u32 %p8base, {make_ptx.SWIN_RET_SCRATCH};\n" + "\n".join(calls) + "\n}"
|
|
params = ",\n\t".join([f".param .b32 a{i}" for i in range(12)]
|
|
+ [".param .b32 scratch"])
|
|
return dict(file="0013-PREPASS_ENTRYPOINT_NAME.ptx", cut_call=(19770, 19935),
|
|
call=call, extern=f".extern .func d4r_swin_enc0_p8\n(\n\t{params}\n)\n;\n")
|