Validate native Windows M against real PTX and correct patch-merge bounds

tools/windows/m_ptx_validate.py

@@ -0,0 +1,62 @@
+"""Independent PTX interpreter validation of real M captures; keep outputs private."""
+import argparse
+import pathlib
+import re
+import sys
+
+import numpy as np
+
+from m_replay_validate import me, LAYERS, blocks, model_block, stored, metrics
+import ptxsim
+from extract_dlss_ptx import fatbins, ptx_entries
+
+
+def validate(args):
+    name = f'rrlite_{args.layer}_4x4'
+    parameters = me.Params(str(args.capture_dir))
+    if parameters.kernel != name:
+        raise RuntimeError('Capture kernel does not match requested M layer')
+    native = {i: np.fromfile(args.native_dir / f'alloc-{i}.bin', np.uint8) for i in parameters.allocs}
+    candidates = []
+    for bx, by in blocks(parameters.grid, 12):
+        model = model_block(parameters, args.layer, bx, by)
+        expected, actual = stored(parameters, args.layer, bx, by, native, model)
+        difference = float(np.abs(expected - actual).max())
+        candidates.append((difference, bx, by, model, expected, actual))
+    selected = sorted(candidates, key=lambda c: (-c[0], c[1], c[2]))[:args.blocks]
+    pattern = re.compile(rb'\.entry\s+' + name.encode() + rb'\s*\(')
+    sources = [source.rstrip(b'\0') for _, fatbin in fatbins(args.dlss_dll.read_bytes())
+               for source in ptx_entries(fatbin) if pattern.search(source)]
+    if not sources:
+        raise RuntimeError('Local DLSS DLL does not contain this M kernel')
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    path = args.output_dir / (name + '.ptx')
+    path.write_bytes(sources[-1])
+    for _, bx, by, model, numpy_expected, actual in selected:
+        _, memory, steps = ptxsim.run_block(str(path), name, str(args.capture_dir), (bx, by, 0))
+        allocation_data = dict(zip(memory.ids, memory.data))
+        _, expected = stored(parameters, args.layer, bx, by, allocation_data, model)
+        print(f'PTX_REFERENCE layer={args.layer} block={bx},{by} statements={steps}', flush=True)
+        _, _, psnr, fp8_steps = metrics(expected, actual, 'PTX_VS_NATIVE')
+        metrics(expected, numpy_expected, 'PTX_VS_NUMPY')
+        if fp8_steps.max() > 1 or psnr < 50:
+            raise RuntimeError('Native M differs from independent PTX reference')
+    print(f'PASS M_PTX_REFERENCE layer={args.layer} blocks={len(selected)} nan_inf=0', flush=True)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--layer', choices=list(LAYERS), required=True)
+    parser.add_argument('--capture-dir', type=pathlib.Path, required=True)
+    parser.add_argument('--native-dir', type=pathlib.Path, required=True)
+    parser.add_argument('--dlss-dll', type=pathlib.Path, required=True)
+    parser.add_argument('--output-dir', type=pathlib.Path, required=True)
+    parser.add_argument('--blocks', type=int, default=2)
+    args = parser.parse_args()
+    if args.blocks < 1:
+        parser.error('--blocks must be positive')
+    try:
+        validate(args)
+    except Exception as error:
+        print(f'FAIL M_PTX_REFERENCE {error}', file=sys.stderr)
+        raise SystemExit(5)