| """Build the second optimization report from recorded timings and real PNGs.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import statistics |
| from pathlib import Path |
|
|
| from h3_mlx.compare import image_metrics, make_comparison |
|
|
| ROOT = Path("reports/optimization-v2") |
| OUTPUT = Path("outputs/optimization-v2") |
| GALLERIES = Path("comparisons/optimization-v2") |
| NAMES = { |
| "quiet-baseline": "BF16 baseline (packing still active)", |
| "quiet-native": "Exact optimized BF16, final quiet session", |
| "quiet-packed": "Exact optimized packed BF16, final quiet session", |
| "wired-baseline": "BF16 baseline", |
| "wired-compiled": "Compiled BF16, retained refiners", |
| "packed12": "Lossless packed, retained refiners", |
| "prefetch": "Compiled BF16 + exact caches + prefetch", |
| "compact": "Lossless packed + exact caches + prefetch", |
| "fp16": "FP16 projection failure (non-finite)", |
| "five-pass": "Five original passes (skip head 1)", |
| } |
|
|
|
|
| def latest(case, profile): |
| files = sorted( |
| (OUTPUT / case).glob(f"{profile}-*.png"), |
| key=lambda p: int(p.stem.rsplit("-", 1)[1]), |
| ) |
| return files[-1] if files else None |
|
|
|
|
| def main(): |
| from PIL import Image |
|
|
| ROOT.mkdir(parents=True, exist_ok=True) |
| ledger = [ |
| json.loads(line) for line in (OUTPUT / "trials.jsonl").read_text().splitlines() |
| ] |
| summary = [] |
| for case in ("product", "portrait", "native-4k"): |
| reference = latest(case, "quiet-baseline") or latest(case, "wired-baseline") |
| if reference is None: |
| old = list((Path("outputs/suite") / case).glob("native-*.png")) |
| reference = sorted(old)[-1] if old else None |
| if reference is None: |
| continue |
| ref = Image.open(reference) |
| for profile in NAMES: |
| path = latest(case, profile) |
| if path is None: |
| continue |
| group = [r for r in ledger if r["case"] == case and r["profile"] == profile] |
| |
| group = list({r["output"]: r for r in group}.values()) |
| hot = [r for r in group if r["dit_was_resident"]] |
| import numpy as np |
| from safetensors import safe_open |
|
|
| with safe_open( |
| path.with_suffix(".trace.safetensors"), framework="numpy" |
| ) as trace: |
| finite = bool(np.isfinite(trace.get_tensor("latent")).all()) |
| row = { |
| "valid_denoising": finite, |
| "case": case, |
| "profile": profile, |
| "label": NAMES[profile], |
| "image": str(path), |
| "reference": str(reference), |
| "resident_samples": len(hot), |
| "load_included_s": [ |
| r["timings"]["total_s"] for r in group if not r["dit_was_resident"] |
| ], |
| "resident_total_s": [r["timings"]["total_s"] for r in hot], |
| "resident_dit_s": [ |
| sum(v for k, v in r["timings"].items() if k.startswith("step_")) |
| for r in hot |
| ], |
| "peak_mlx_gib": max(r["memory"]["peak_gb"] for r in group), |
| **image_metrics(ref, Image.open(path)), |
| } |
| row["resident_median_s"] = ( |
| statistics.median(row["resident_total_s"]) if hot else None |
| ) |
| row["resident_dit_median_s"] = ( |
| statistics.median(row["resident_dit_s"]) if hot else None |
| ) |
| summary.append(row) |
| if case in ("product", "portrait"): |
| candidates = [] |
| labels = ["Native six-pass BF16"] |
| for profile in ("prefetch", "compact", "fp16", "five-pass"): |
| selected = {"prefetch": "quiet-native", "compact": "quiet-packed"}.get( |
| profile, profile |
| ) |
| path = latest(case, selected) or latest(case, profile) |
| if path: |
| candidates.append(path) |
| labels.append(NAMES[profile]) |
| for profile, label in ( |
| ("q8", "Affine Q8 (earlier trial)"), |
| ("q4", "Affine Q4 (earlier trial)"), |
| ): |
| paths = sorted((Path("outputs/suite") / case).glob(f"{profile}-*.png")) |
| if paths: |
| candidates.append(paths[-1]) |
| labels.append(label) |
| make_comparison(reference, candidates, GALLERIES / case, labels=labels) |
| elif latest(case, "compact"): |
| make_comparison( |
| reference, |
| [latest(case, "compact")], |
| GALLERIES / case, |
| labels=["Native 4K six-pass BF16", "Optimized 4K lossless packed BF16"], |
| ) |
| (ROOT / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") |
| lines = [ |
| "| Case | Configuration | Warm samples | Warm median (s) | Peak MLX GiB | SSIM | Exact pixels |", |
| "|---|---|---:|---:|---:|---:|---|", |
| ] |
| for r in summary: |
| seconds = ( |
| "—" if r["resident_median_s"] is None else f"{r['resident_median_s']:.2f}" |
| ) |
| if not r["valid_denoising"]: |
| seconds = "invalid" |
| similarity = f"{r['ssim_rgb']:.6f}" if r["valid_denoising"] else "invalid state" |
| exact = ( |
| ("yes" if r["pixel_identical"] else "no") |
| if r["valid_denoising"] |
| else "non-finite" |
| ) |
| lines.append( |
| f"| {r['case']} | {r['label']} | {r['resident_samples']} | {seconds} | {r['peak_mlx_gib']:.2f} | {similarity} | {exact} |" |
| ) |
| (ROOT / "table.md").write_text("\n".join(lines) + "\n") |
| target = ROOT / "RESULTS.md" |
| if target.exists(): |
| text = target.read_text() |
| start = "<!-- RESULTS_TABLE -->" |
| end = "<!-- END_RESULTS_TABLE -->" |
| if start in text and end in text: |
| text = ( |
| text.split(start)[0] |
| + start |
| + "\n\n" |
| + "\n".join(lines) |
| + "\n\n" |
| + end |
| + text.split(end)[1] |
| ) |
| target.write_text(text) |
| try: |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| rows = [ |
| r |
| for r in summary |
| if r["case"] == "product" |
| and r["resident_median_s"] is not None |
| and r["profile"] in ("quiet-native", "quiet-packed") |
| ] |
| fig, ax = plt.subplots(figsize=(10, 6), layout="constrained") |
| for r in rows: |
| ax.scatter( |
| r["resident_median_s"], |
| r["peak_mlx_gib"], |
| s=90, |
| c=[r["ssim_rgb"]], |
| vmin=0.6, |
| vmax=1, |
| cmap="viridis", |
| edgecolor="black", |
| ) |
| ax.annotate( |
| r["profile"], |
| (r["resident_median_s"], r["peak_mlx_gib"]), |
| xytext=(5, 6), |
| textcoords="offset points", |
| fontsize=9, |
| ) |
| ax.set( |
| xlabel="Recorded warm end-to-end seconds (lower is faster)", |
| ylabel="Peak MLX allocator GiB (lower is smaller)", |
| title="M5 Pro 64 GB · 2048² product · exact and approximate options", |
| ) |
| ax.grid(alpha=0.2) |
| fig.savefig(ROOT / "tradeoffs.png", dpi=170) |
| plt.close(fig) |
| except ImportError: |
| pass |
| print(json.dumps(summary, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
