File size: 7,654 Bytes
3da88e3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """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]
# If a suite was rerun with the same filenames, keep its latest row.
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()
|