kernels: build the denoiser's output kernel natively, authorized separately

Ray Reconstruction's eleven DLSS-D Swin layers still run ZLUDA's compile of
NVIDIA's PTX; at 1280x720 -> 3840x2160 on gfx1201 they are 30.2 of a 31.6 ms
frame. The denoiser's high-fidelity output kernel is the first one with a native
replacement: make_ptx.py replaces its whole-texel RGBA16F store
(sust.b.2d.v4.b16.zero) through d4r_sust_v4b16, kernels/build.sh builds the pair
as the `rr` family, and packaging builds it into every target folder including
the RDNA4 FP8 ones.

Both libraries define modules named cuda_dldn_engine_*, with different text, so
one manifest cannot take the union of their hashes: kernel_manifest.py now takes
the denoiser DLL as --rr-dll and lets only its hashes authorize a
cuda_dldn_engine_* binary, and package_release.sh passes it when the denoiser is
bundled.

Verified in the D3D12 harness with Ray Reconstruction at 1280x720 -> 3840x2160:
the output kernel drops from 1.44 to 1.00 ms/frame and the 66 MB raw output is
byte-identical to the run with every kernel translated.

docs/native-kernels.md gains the denoiser family, the per-kernel cost table and
the instruction counts that show why the Swin layers are the remaining work.

kernels/build.sh

@@ -2,7 +2,7 @@
 # Builds the native replacements for DLSS kernels into one directory that ZLUDA serves them from
 # (D4R_ZLUDA_NATIVE_DIR, or NativeKernelDirFast in d4r.ini). See docs/native-kernels.md.
 #
-# usage: kernels/build.sh [all|k|l|m|tex] [OUT_DIR]      (default: all, kernels/out/native)
+# usage: kernels/build.sh [all|k|l|m|tex|rr] [OUT_DIR] (default: all, kernels/out/native)
 #
 #   D4R_ROCM_DIR    ROCm installation with clang and the HIP device libraries (default /opt/rocm)
 #   D4R_GPU_ARCH    target GPU (default gfx1101): an RDNA3 (gfx110x) or RDNA4 (gfx120x) target; the kernels
@@ -17,6 +17,8 @@
 #   D4R_ZLUDA_EMIT  ZLUDA's d4r_emit (patches/zluda applied; `cargo build --release -p ptx --example d4r_emit`),
 #                   which compiles PTX offline for D4R_GPU_ARCH (no GPU of that kind needed)
 #   D4R_DLSS_PTX_DIR optional pre-extracted PTX directory for parallel per-target builds
+# Ray Reconstruction (rr) additionally needs D4R_DLSSD_DLL (CUDA-capable
+# nvngx_dlssd.dll); D4R_DLSSD_PTX_DIR may provide its pre-extracted modules.
 set -euo pipefail
 shopt -s extglob
 
@@ -38,7 +40,7 @@ esac
 CLANG="$ROCM/lib/llvm/bin/clang++"
 [[ -x "$CLANG" ]] || { echo "clang++ not found in $ROCM/lib/llvm/bin (set D4R_ROCM_DIR)" >&2; exit 2; }
 mkdir -p "$OUT"
-case "$WHAT" in all|k|l|m|tex) ;; *) echo "unknown kernel family: $WHAT" >&2; exit 2 ;; esac
+case "$WHAT" in all|k|l|m|tex|rr) ;; *) echo "unknown kernel family: $WHAT" >&2; exit 2 ;; esac
 # Never certify a directory containing older fast binaries as an accuracy set.
 if [[ "$ACCURACY" == 1 ]] && compgen -G "$OUT/*.hsaco" >/dev/null &&
     { [[ ! -f "$OUT/d4r-accuracy.txt" ]] || [[ "$(cat "$OUT/d4r-accuracy.txt")" != 1 ]]; }; then
@@ -185,5 +187,22 @@ if [[ "$WHAT" == all || "$WHAT" == tex || "$WHAT" == l ]]; then
         done
     fi
 fi
+if [[ "$WHAT" == rr || ( "$WHAT" == all && -n "${D4R_DLSSD_DLL:-}" ) ]]; then
+    echo "== Ray Reconstruction texture kernels"
+    : "${D4R_DLSSD_DLL:?set D4R_DLSSD_DLL to the CUDA-capable nvngx_dlssd.dll}"
+    : "${D4R_ZLUDA_EMIT:?set D4R_ZLUDA_EMIT to d4r_emit from a patched ZLUDA build}"
+    RR_PTX_DIR="${D4R_DLSSD_PTX_DIR:-$HERE/extracted/rr}"
+    if [[ -z "${D4R_DLSSD_PTX_DIR:-}" ]]; then
+        python3 "$HERE/tools/extract_dlss_ptx.py" "$D4R_DLSSD_DLL" "$RR_PTX_DIR"
+    elif [[ ! -d "$RR_PTX_DIR" ]]; then
+        echo "missing extracted RR PTX directory: $RR_PTX_DIR" >&2; exit 2
+    fi
+    for kernel in cuda_dldn_engine_hkpn_output_kernel_transformer \
+                  cuda_dldn_engine_hkpn_output_kernel_transformer_diamond_wallaby; do
+        D4R_PREFER_ACCURACY="$ACCURACY" D4R_ROCM_DIR="$ROCM" D4R_GPU_ARCH="$ARCH" \
+            D4R_ZLUDA_WAVE64="$((1 - ACCURACY))" D4R_TEX_FP8="$FP8" \
+            D4R_DLSS_PTX_DIR="$RR_PTX_DIR" "$HERE/tex/build_tex.sh" "$kernel" sust_only "$OUT"
+    done
+fi
 [[ "$ACCURACY" == 1 ]] && printf '1\n' > "$OUT/d4r-accuracy.txt"
 echo "native kernels in $OUT"

kernels/tex/make_ptx.py

@@ -99,6 +99,13 @@ SUST_EXTERN = """.extern .func d4r_sust_v2b16
 )
 ;
 """
+SUST4_EXTERN = """.extern .func d4r_sust_v4b16
+(
+	.param .b64 d4r_h0, .param .b32 d4r_h1, .param .b32 d4r_h2, .param .b16 d4r_h3,
+	.param .b16 d4r_h4, .param .b16 d4r_h5, .param .b16 d4r_h6
+)
+;
+"""
 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,
@@ -115,14 +122,31 @@ SUSTB32_EXTERN = """.extern .func d4r_sust_b32
 """
 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+)\};$")
+SUST4_RE = re.compile(r"^sust\.b\.2d\.v4\.b16\.zero \[(%rd\d+), \{(%r\d+),\s*(%r\d+)\}\], \{(%rs\d+),\s*(%rs\d+),\s*(%rs\d+),\s*(%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}
+    out, n = [], {"b16": 0, "v4b16": 0, "p": 0, "b32": 0}
     for line in lines:
+        mh = SUST4_RE.match(line.strip()) if "b16" in SUST_KINDS else None
+        if mh:
+            s, x, y, *v = mh.groups()
+            st = "\n".join(f"\t.param .b16 v{i};\n\tst.param.b16 [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_v4b16, (u0, u1, u2, v0, v1, v2, v3);
+}}""")
+            n["v4b16"] += 1
+            continue
         mb = SUSTB32_RE.match(line.strip()) if "b32" in SUST_KINDS else None
         if mb:
             s, x, y, v = mb.groups()
@@ -414,7 +438,8 @@ ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
 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))$")
+                          r"rrlite_downsample_kernel_(static|dynamic)_(hdr|ldr)|"
+                          r"cuda_dldn_engine_hkpn_output_kernel_transformer(_diamond_wallaby)?)$")
 ENC0_RE = re.compile(r"^rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)_folded$")
 
 
@@ -536,7 +561,7 @@ def main():
         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 "") + (SUSTP_EXTERN if n["p"] else "") + (SUSTB32_EXTERN if n["b32"] else "")
+    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 "")
     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}")

kernels/tex/tex_common.h

@@ -261,6 +261,39 @@ DEV __attribute__((always_inline)) void d4r_sust_v2b16(uint64_t surface, int32_t
         sust_v2b16_generic(surface, x, y, a, b);
 }
 
+typedef int16_t v4s __attribute__((ext_vector_type(4)));
+extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_v4_b16_zero(uint64_t surface, v2i coord, v4s value);
+
+__attribute__((device, noinline)) static void sust_v4b16_generic(uint64_t surface, int32_t x, int32_t y,
+                                                               uint16_t a, uint16_t b, uint16_t c, uint16_t d)
+{
+    __zluda_ptx_impl_surfobj_b_2d_v4_b16_zero(surface, (v2i){x, y},
+                                           (v4s){(int16_t)a, (int16_t)b, (int16_t)c, (int16_t)d});
+}
+
+FP_BODY void sust_rgba16f_body(tsharp_t* image, int32_t px, int32_t y,
+                             uint16_t a, uint16_t b, uint16_t c, uint16_t d)
+{
+    __ockl_image_store_2D(image, (v2i){px, y},
+                         (v4f){half_bits_to_float(a), half_bits_to_float(b),
+                               half_bits_to_float(c), half_bits_to_float(d)});
+}
+
+// RR stores complete RGBA16F texels. No channel reads or partial-pixel merges
+// are needed; other layouts retain the generic byte-addressed store.
+DEV __attribute__((always_inline)) void d4r_sust_v4b16(uint64_t surface, int32_t x, int32_t y,
+                                                     uint16_t a, uint16_t b, uint16_t c, uint16_t d)
+{
+    tsharp_t* image = (tsharp_t*)surface;
+    if (__ockl_image_channel_data_type_2D(image) == 14 && __ockl_image_channel_order_2D(image) == 8 && (x & 7) == 0)
+    {
+        if (x >= 0)
+            sust_rgba16f_body(image, x >> 3, y, a, b, c, d);
+    }
+    else
+        sust_v4b16_generic(surface, x, y, a, b, c, d);
+}
+
 
 // f32 comparisons on the bits (no fcmp: kernels ZLUDA compiles in strict-FP mode cannot lower
 // constrained fcmp); same results as the float compares, -0.0 included

kernels/tools/kernel_manifest.py

@@ -1,14 +1,19 @@
 #!/usr/bin/env python3
 """Write the d4r-kernels.txt manifest of a native kernel directory.
 
-usage: kernel_manifest.py KERNEL_DIR NVNGX_DLSS_DLL [NVNGX_DLSS_DLL...]
+usage: kernel_manifest.py KERNEL_DIR [NVNGX_DLSS_DLL...] [--rr-dll NVNGX_DLSSD_DLL]
 
 For every NAME.hsaco in KERNEL_DIR, finds the PTX module of each DLL that defines `.entry NAME` and
 lists the FNV-1a 64 hash of the module's text (trailing NUL bytes removed), one "NAME HASH" line per
 distinct hash. The Wine CUDA bridge serves a native kernel only while DLSS loads a module with a
 listed hash, so a DLSS version that changed the kernel runs ZLUDA's own compile of it instead.
 The manifest is written to KERNEL_DIR/d4r-kernels.txt; it holds hashes, not NVIDIA code.
+RR DLLs are explicit: only a DLL passed with --rr-dll may authorize a cuda_dldn_engine_* replacement,
+and the Super Resolution DLLs passed positionally authorize everything else. Both libraries define
+cuda_dldn_engine_* modules of their own (the shared capture and exposure helpers), so without the split
+a Super Resolution hash could stand in for the denoiser's different implementation of the same name.
 """
+import argparse
 import hashlib
 import os
 import re
@@ -41,19 +46,25 @@ def module_hashes(dll: str) -> dict:
 
 
 def main(argv):
-    if len(argv) < 3:
-        sys.exit(__doc__)
-    directory = argv[1]
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("directory")
+    parser.add_argument("dlls", nargs="*")
+    parser.add_argument("--rr-dll", action="append", default=[])
+    args = parser.parse_args(argv[1:])
+    if not args.dlls and not args.rr_dll:
+        parser.error("at least one SR DLL or --rr-dll is required")
+    directory = args.directory
     names = sorted(name[:-len('.hsaco')] for name in os.listdir(directory) if name.endswith('.hsaco'))
     lines = ['# d4r native kernels: NAME and the FNV-1a 64 hash of the DLSS PTX module each was written for',
              '# (written by kernels/tools/kernel_manifest.py from:']
     listed = {name: set() for name in names}
-    for dll in argv[2:]:
+    for dll, reconstruction in [(dll, False) for dll in args.dlls] + [(dll, True) for dll in args.rr_dll]:
         digest = hashlib.sha256(open(dll, 'rb').read()).hexdigest()[:16]
         lines.append(f'#   {os.path.basename(dll)} sha256 {digest}...)')
         hashes = module_hashes(dll)
         for name in names:
-            listed[name] |= hashes.get(name, set())
+            if name.startswith("cuda_dldn_engine_") == reconstruction:
+                listed[name] |= hashes.get(name, set())
     missing = [name for name in names if not listed[name]]
     if missing:
         sys.exit(f'no DLL defines {", ".join(missing)}; not writing a manifest')