ayo
f66cccaf28
- Build L texture variants and complete older release kernel bundles. - Add selective Swin encoder fallback to reduce L/M motion trails. - Add opt-in GPU motion-vector dilation for thin geometry and sky gaps. - Replace GPU input waits with cancellable readiness polling and use sleeping event queries for completion. - Retain recorded resources through command allocator completion, defer shutdown until cleanup, and preserve unfinished output synchronization. - Extend lifecycle harness checks and add event, model and motion tests. - Update configuration defaults, build scripts and release documentation.
329 lines
13 KiB
Python
Executable File
329 lines
13 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Hand-edited PTX for the texture kernels: keep NVIDIA's code up to a cut line, then call a native tail.
|
|
|
|
usage: make_ptx.py KERNEL OUT.ptx
|
|
"""
|
|
import os
|
|
import re
|
|
import sys
|
|
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"))
|
|
|
|
ENC0_TAIL = """{{
|
|
.param .b32 q0;
|
|
st.param.b32 [q0+0], {smem};
|
|
.param .b64 q1;
|
|
st.param.b64 [q1+0], {w};
|
|
.param .b64 q2;
|
|
st.param.b64 [q2+0], {out};
|
|
.param .b32 q3;
|
|
st.param.b32 [q3+0], {dimx};
|
|
.param .b32 q4;
|
|
st.param.b32 [q4+0], {dimy};
|
|
.param .b32 q5;
|
|
st.param.b32 [q5+0], {ox};
|
|
.param .b32 q6;
|
|
st.param.b32 [q6+0], {oy};
|
|
call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
|
|
}}
|
|
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(
|
|
file="dlss-0185-00.ptx",
|
|
cut_after=2548,
|
|
cut_check="bar.warp.sync -1;",
|
|
sust=True,
|
|
extern=""".extern .func d4r_enc0_tail
|
|
(
|
|
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3,
|
|
.param .b32 d4r_a4, .param .b32 d4r_a5, .param .b32 d4r_a6
|
|
)
|
|
;
|
|
""",
|
|
tail=ENC0_TAIL.format(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233"),
|
|
),
|
|
# dec0: the GEMM (from the B fragment loads to the e4m3 staging stores) becomes a native call; the
|
|
# A loads before it are left dead. %r560 (shared base) is the only register of the range used later.
|
|
"rrlite_dec0_4x4_folded": dict(
|
|
file="dlss-0195-00.ptx",
|
|
delete=(139, 587),
|
|
delete_check=("mov.u32 %r17, %laneid;", "st.shared.v2.u16 [%r570+1168], {%rs63, %rs64};"),
|
|
sust=True,
|
|
extern=""".extern .func d4r_dec0_head
|
|
(
|
|
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3,
|
|
.param .b32 d4r_a4, .param .b32 d4r_a5, .param .b32 d4r_a6
|
|
)
|
|
;
|
|
""",
|
|
insert="""mov.u32 %r560, _ZZ22rrlite_dec0_4x4_foldedN6RRLite18Decoder0ParametersEE4smem;
|
|
{
|
|
.param .b32 q0;
|
|
st.param.b32 [q0+0], %r560;
|
|
.param .b64 q1;
|
|
st.param.b64 [q1+0], %rd22;
|
|
.param .b64 q2;
|
|
st.param.b64 [q2+0], %rd24;
|
|
.param .b32 q3;
|
|
st.param.b32 [q3+0], %r475;
|
|
.param .b32 q4;
|
|
st.param.b32 [q4+0], %r476;
|
|
.param .b32 q5;
|
|
st.param.b32 [q5+0], %r3;
|
|
.param .b32 q6;
|
|
st.param.b32 [q6+0], %r4;
|
|
call.uni d4r_dec0_head, (q0, q1, q2, q3, q4, q5, q6);
|
|
}""",
|
|
),
|
|
# 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=""),
|
|
"rrlite_downsample_kernel_static_hdr": dict(file="dlss-0196-00.ptx", sust=True, extern=""),
|
|
# preset K (DLSS 4): the output kernel's 17 surface stores (13 formatted, 4 raw 32-bit)
|
|
"hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel": dict(file="dlss-0078-00.ptx", sust=True, rz_round=True, extern=""),
|
|
"hiluma_engine_input_depthinv_mvhi_hdr_v2_rel": dict(file="dlss-0062-00.ptx", sust=True, rz_round=True, extern=""),
|
|
}
|
|
|
|
|
|
SUST_EXTERN = """.extern .func d4r_sust_v2b16
|
|
(
|
|
.param .b64 d4r_s0, .param .b32 d4r_s1, .param .b32 d4r_s2, .param .b16 d4r_s3, .param .b16 d4r_s4
|
|
)
|
|
;
|
|
"""
|
|
SUSTP_EXTERN = """.extern .func d4r_sust_p_v4b32
|
|
(
|
|
.param .b64 d4r_t0, .param .b32 d4r_t1, .param .b32 d4r_t2, .param .b32 d4r_t3, .param .b32 d4r_t4,
|
|
.param .b32 d4r_t5, .param .b32 d4r_t6
|
|
)
|
|
;
|
|
"""
|
|
SUSTP_RE = re.compile(r"^sust\.p\.2d\.v4\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%f\d+),(%f\d+),(%f\d+),(%f\d+)\};$")
|
|
SUSTB32_EXTERN = """.extern .func d4r_sust_b32
|
|
(
|
|
.param .b64 d4r_w0, .param .b32 d4r_w1, .param .b32 d4r_w2, .param .b32 d4r_w3
|
|
)
|
|
;
|
|
"""
|
|
SUSTB32_RE = re.compile(r"^sust\.b\.2d\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%r\d+)\};$")
|
|
SUST_RE = re.compile(r"^sust\.b\.2d\.v2\.b16\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%rs\d+),(%rs\d+)\};$")
|
|
|
|
|
|
SUST_KINDS = set(os.environ.get("D4R_SUST_KINDS", "b16,p,b32").split(","))
|
|
|
|
|
|
def replace_sust(lines):
|
|
out, n = [], {"b16": 0, "p": 0, "b32": 0}
|
|
for line in lines:
|
|
mb = SUSTB32_RE.match(line.strip()) if "b32" in SUST_KINDS else None
|
|
if mb:
|
|
s, x, y, v = mb.groups()
|
|
out.append(f"""{{
|
|
.param .b64 u0;
|
|
st.param.b64 [u0+0], {s};
|
|
.param .b32 u1;
|
|
st.param.b32 [u1+0], {x};
|
|
.param .b32 u2;
|
|
st.param.b32 [u2+0], {y};
|
|
.param .b32 u3;
|
|
st.param.b32 [u3+0], {v};
|
|
call d4r_sust_b32, (u0, u1, u2, u3);
|
|
}}""")
|
|
n["b32"] += 1
|
|
continue
|
|
mp = SUSTP_RE.match(line.strip()) if "p" in SUST_KINDS else None
|
|
if mp:
|
|
s, x, y, *v = mp.groups()
|
|
st = "\n".join(f"\t.param .b32 v{i};\n\tst.param.b32 [v{i}+0], {r};" for i, r in enumerate(v))
|
|
out.append(f"""{{
|
|
.param .b64 u0;
|
|
st.param.b64 [u0+0], {s};
|
|
.param .b32 u1;
|
|
st.param.b32 [u1+0], {x};
|
|
.param .b32 u2;
|
|
st.param.b32 [u2+0], {y};
|
|
{st}
|
|
call d4r_sust_p_v4b32, (u0, u1, u2, v0, v1, v2, v3);
|
|
}}""")
|
|
n["p"] += 1
|
|
continue
|
|
m = SUST_RE.match(line.strip())
|
|
if not m:
|
|
out.append(line)
|
|
continue
|
|
s, x, y, a, b = m.groups()
|
|
out.append(f"""{{
|
|
.param .b64 u0;
|
|
st.param.b64 [u0+0], {s};
|
|
.param .b32 u1;
|
|
st.param.b32 [u1+0], {x};
|
|
.param .b32 u2;
|
|
st.param.b32 [u2+0], {y};
|
|
.param .b16 u3;
|
|
st.param.b16 [u3+0], {a};
|
|
.param .b16 u4;
|
|
st.param.b16 [u4+0], {b};
|
|
call d4r_sust_v2b16, (u0, u1, u2, u3, u4);
|
|
}}""")
|
|
n["b16"] += 1
|
|
return out, n
|
|
|
|
|
|
RZ_RE = re.compile(r"^add\.rz\.ftz\.f32 (%f\d+), (%f\d+), (%f\d+);$")
|
|
CVT_RE = re.compile(r"^cvt\.rzi\.f32\.f32 (%f\d+), (%f\d+);$")
|
|
|
|
|
|
def replace_rz_round(lines):
|
|
"""roundf idiom add.rz(x, copysign(0.5, x)) + cvt.rzi -> exact round-half-away in default rounding:
|
|
t = trunc(x); r = |x - t| >= 0.5 ? t + 2 * copysign(0.5, x) : t (x - t is exact). Lets ZLUDA compile the
|
|
kernel without the strict-FP (constrained) mode that one explicit rounding mode forces."""
|
|
out, n, i = [], 0, 0
|
|
while i < len(lines):
|
|
m = RZ_RE.match(lines[i].strip())
|
|
if m:
|
|
d, x, sgn = m.groups()
|
|
j = i + 1
|
|
while j < len(lines) and not lines[j].strip():
|
|
j += 1
|
|
c = CVT_RE.match(lines[j].strip()) if j < len(lines) else None
|
|
if c and c.group(2) == d:
|
|
r = c.group(1)
|
|
out += [f"cvt.rzi.f32.f32 %rq{4*n}, {x};", f"sub.f32 %rq{4*n+1}, {x}, %rq{4*n};",
|
|
f"abs.f32 %rq{4*n+1}, %rq{4*n+1};", f"setp.ge.f32 %pq{n}, %rq{4*n+1}, 0f3F000000;",
|
|
f"add.f32 %rq{4*n+2}, {sgn}, {sgn};", f"add.f32 %rq{4*n+3}, %rq{4*n}, %rq{4*n+2};",
|
|
f"selp.f32 {r}, %rq{4*n+3}, %rq{4*n}, %pq{n};"]
|
|
n += 1
|
|
i = j + 1
|
|
continue
|
|
raise SystemExit(f"add.rz at line {i} not followed by its cvt.rzi")
|
|
out.append(lines[i])
|
|
i += 1
|
|
return out, n
|
|
|
|
|
|
ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
|
|
# the other flag combinations of the kernels whose only native part is the surface-store rewrite
|
|
SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
|
|
r"rrlite_post_3_[12]_mv(hi|lo)_(hdr|ldr)(_folded)?|"
|
|
r"rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)|rrlite_dec0_4x4|"
|
|
r"rrlite_downsample_kernel_(static|dynamic)_(hdr|ldr))$")
|
|
ENC0_RE = re.compile(r"^rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)_folded$")
|
|
|
|
|
|
def read_lines(path):
|
|
return path.read_text().replace("\r\n", "\n").split("\n")
|
|
|
|
|
|
def enc0_registers(lines):
|
|
"""Cut line (1-based) and the tail call's registers of an enc0 variant. The code after the third
|
|
bar.warp.sync is the same GEMM in every flag combination, only register numbers differ: the shared
|
|
base and weights pointer are read off by aligning that block with the reference variant; the output
|
|
pointer and grid size are the kernel's loads of parameters +40 and +8; the block offsets are the two
|
|
values halved (shr 31 / add) by the epilogue."""
|
|
syncs = [i for i, l in enumerate(lines) if l.strip() == "bar.warp.sync -1;"]
|
|
if len(syncs) != 3:
|
|
sys.exit(f"enc0: expected 3 bar.warp.sync, found {len(syncs)}")
|
|
cut = syncs[2] + 1
|
|
before, after = lines[:cut], lines[cut:]
|
|
out = [m.group(1) for m in (re.match(r"^ld\.param\.u64 (%rd\d+), \[%rd\d+\+40\];$", l.strip()) for l in before) if m]
|
|
dims = [m.groups() for m in (re.match(r"^ld\.param\.v2\.u32 \{(%r\d+), (%r\d+)\}, \[%rd\d+\+8\];$", l.strip()) for l in before) if m]
|
|
halves = []
|
|
for a, b in zip(after, after[1:]):
|
|
m = re.match(r"^shr\.u32 (%r\d+), (%r\d+), 31;$", a.strip())
|
|
if m and re.match(rf"^add\.s32 %r\d+, {re.escape(m.group(2))}, {re.escape(m.group(1))};$", b.strip()):
|
|
halves.append(m.group(2))
|
|
if len(out) != 1 or len(dims) != 1 or len(halves) < 2:
|
|
sys.exit(f"enc0: parameter loads or epilogue not found ({out}, {dims}, {halves[:2]})")
|
|
return cut, dict(out=out[0], dimx=dims[0][0], dimy=dims[0][1], oy=halves[0], ox=halves[1])
|
|
|
|
|
|
REG_RE = re.compile(r"%(?:rd|rs|r|fd|f|p)\d+")
|
|
|
|
|
|
def enc0_variant(name):
|
|
ref = read_lines(module_file(ENC0_REF, KERNELS[ENC0_REF]))
|
|
path = module_file(name, {"file": ""})
|
|
lines = read_lines(path)
|
|
rcut, rregs = enc0_registers(ref)
|
|
k = KERNELS[ENC0_REF]
|
|
known = dict(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233")
|
|
if rcut != k["cut_after"] or any(rregs[key] != known[key] for key in rregs):
|
|
sys.exit(f"enc0: the register search does not reproduce {ENC0_REF} ({rcut}, {rregs})")
|
|
cut, regs = enc0_registers(lines)
|
|
# align the GEMM block (up to the reference's first closing brace) and map registers
|
|
n = next(i for i, l in enumerate(ref[rcut:]) if l.startswith("}"))
|
|
mapping = {}
|
|
for a, b in zip(ref[rcut:rcut + n], lines[cut:cut + n]):
|
|
if REG_RE.sub("R", a) != REG_RE.sub("R", b):
|
|
sys.exit(f"enc0: {name} differs from {ENC0_REF} after the cut: {a!r} / {b!r}")
|
|
for x, y in zip(REG_RE.findall(a), REG_RE.findall(b)):
|
|
if mapping.setdefault(x, y) != y:
|
|
sys.exit(f"enc0: inconsistent register mapping {x} -> {mapping[x]} / {y}")
|
|
regs.update(smem=mapping[known["smem"]], w=mapping[known["w"]])
|
|
return dict(file=path.name, cut_after=cut, cut_check="bar.warp.sync -1;", sust=True, extern=k["extern"],
|
|
tail=ENC0_TAIL.format(**regs))
|
|
|
|
|
|
def kernel_spec(name):
|
|
if name in KERNELS:
|
|
return KERNELS[name]
|
|
if ENC0_RE.match(name):
|
|
return enc0_variant(name)
|
|
if SUST_ONLY_RE.match(name):
|
|
return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_output"))
|
|
sys.exit(f"make_ptx: no recipe for {name}")
|
|
|
|
|
|
def module_file(name, k):
|
|
"""The extracted PTX module that defines NAME: the 310.7 file number, else any module that does
|
|
(other DLSS versions number their modules differently)."""
|
|
entry = re.compile(rf"\.entry\s+{re.escape(name)}\s*\(")
|
|
known = PTX_DIR / k["file"]
|
|
if k["file"] and known.is_file() and entry.search(known.read_text()):
|
|
return known
|
|
for path in sorted(PTX_DIR.glob("*.ptx")):
|
|
if entry.search(path.read_text()):
|
|
return path
|
|
sys.exit(f"no PTX module in {PTX_DIR} defines {name}")
|
|
|
|
|
|
def main():
|
|
name, out = sys.argv[1], Path(sys.argv[2])
|
|
k = kernel_spec(name)
|
|
lines = read_lines(module_file(name, k))
|
|
if "cut_after" in k:
|
|
cut = k["cut_after"]
|
|
if lines[cut - 1].strip() != k["cut_check"]:
|
|
sys.exit(f"line {cut} is {lines[cut - 1]!r}, expected {k['cut_check']!r}")
|
|
head, tail = lines[:cut], k["tail"].split("\n")
|
|
elif "delete" in k:
|
|
a, b = k["delete"]
|
|
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:], []
|
|
else:
|
|
head, tail = lines, []
|
|
entry = next(i for i, l in enumerate(head) if l.startswith(".visible .entry"))
|
|
body, n = replace_sust(head[entry:]) if k.get("sust") else (head[entry:], {"b16": 0, "p": 0, "b32": 0})
|
|
if k.get("rz_round"):
|
|
body, nrz = replace_rz_round(body)
|
|
# temporaries for the rewrite, declared after the kernel's own .reg lines
|
|
at = next(i for i, l in enumerate(body) if l.startswith(".reg"))
|
|
body = body[:at] + [f".reg .f32 %rq<{4 * nrz + 1}>;", f".reg .pred %pq<{nrz + 1}>;"] + body[at:]
|
|
n["rz"] = nrz
|
|
extern = k["extern"] + (SUST_EXTERN if n["b16"] 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}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|