Execute gfx1201 K enc2 and generalize layer reference diagnostics

tools/windows/k_layer_validate.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python3
+"""Compare real gfx1201 K enc1/enc2 launches with the existing numpy reference."""
+import argparse
+import math
+import pathlib
+import sys
+
+import numpy as np
+
+SCRIPT = pathlib.Path(__file__).resolve()
+REFERENCE = SCRIPT.parent / "reference"
+if not (REFERENCE / "pwin_model.py").exists():
+    REFERENCE = SCRIPT.parents[2] / "kernels" / "tools"
+sys.path.insert(0, str(REFERENCE))
+from pwin_model import pwin_encoder  # noqa: E402
+
+
+class SyntheticParams:
+    def __init__(self, folder):
+        self.input = np.fromfile(folder / "input.bin", dtype=np.uint8)
+        self.weights = np.fromfile(folder / "weights.bin", dtype=np.uint8)
+
+    def i32x2(self, offset):
+        return (8, 8) if offset == 0 else (0, 0)
+
+    def ptr(self, offset):
+        if offset == 8:
+            return self.input, 0
+        if offset == 64:
+            return self.weights, 0
+        raise ValueError(f"unsupported synthetic pointer offset {offset}")
+
+    def f16(self, offset, byte, count):
+        if offset != 64:
+            raise ValueError(f"unsupported synthetic half pointer offset {offset}")
+        return self.weights[byte:byte + 2 * count].view(np.float16).copy()
+
+
+def compare(name, actual, expected):
+    if actual.shape != expected.shape:
+        raise RuntimeError(f"{name} shape mismatch: {actual.shape} != {expected.shape}")
+    if not np.isfinite(actual).all() or not np.isfinite(expected).all():
+        raise RuntimeError(f"{name} contains NaN or Inf")
+    diff = np.abs(actual.astype(np.float32) - expected.astype(np.float32))
+    maximum = float(diff.max())
+    relative = float((diff / np.maximum(np.abs(expected.astype(np.float32)), 1e-3)).max())
+    rms = float(np.sqrt(np.mean(np.square(diff.astype(np.float64)))))
+    peak = max(1.0, float(np.abs(expected.astype(np.float32)).max()))
+    psnr = math.inf if rms == 0 else 20.0 * math.log10(peak / rms)
+    mismatch = np.unravel_index(int(diff.argmax()), diff.shape)
+    print(f"K_REFERENCE {name} max_abs={maximum:.9g} max_rel={relative:.9g} "
+          f"rms={rms:.9g} psnr_db={psnr:.5g} worst={mismatch} "
+          f"expected={float(expected[mismatch]):.9g} actual={float(actual[mismatch]):.9g}")
+    return maximum, psnr
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fixture-dir", required=True, type=pathlib.Path)
+    parser.add_argument("--kernel-name", choices=("enc1", "enc2"), default="enc1")
+    args = parser.parse_args()
+    folder = args.fixture_dir
+    params = SyntheticParams(folder)
+    output_channels = 96 if args.kernel_name == "enc2" else 64
+    reference_full, reference_merged = pwin_encoder(
+        params, 0, 0, H=2, C=64, COUT=output_channels,
+        posattn=args.kernel_name == "enc1")
+    actual_full = np.fromfile(folder / "full.bin", np.float16).reshape(64, 64)
+    actual_merged = np.fromfile(folder / "merged.bin", np.float16).reshape(16, output_channels)
+    full_abs, full_psnr = compare("full", actual_full, reference_full)
+    merged_abs, merged_psnr = compare("merged", actual_merged, reference_merged)
+    # Initial acceptance for the synthetic nonzero FP16 fixture. Tighten this
+    # after each K stage has been compared on captured DLSS weights/activations.
+    if full_abs > 0.001 or full_psnr < 60.0 or merged_abs > 0.001 or merged_psnr < 60.0:
+        raise RuntimeError(f"K {args.kernel_name} numerical mismatch against pwin_model.py")
+    print(f"PASS K_REFERENCE architecture=gfx1201 kernel={args.kernel_name} "
+          "nonzero_weights=1 transformer_executed=1 nan_inf=0")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as error:
+        print(f"FAIL K_REFERENCE {error}", file=sys.stderr)
+        raise SystemExit(5)