"""Export real H3 projections to ML Packages and audit all Core ML devices. Use the isolated .venv-coreml environment. This is a component feasibility benchmark, not a complete Core ML pipeline. Cache packages are reproducible. """ from __future__ import annotations import argparse import gc import json import platform import statistics import time from pathlib import Path import coremltools as ct import numpy as np import torch from coremltools.converters.mil import Builder as mb from coremltools.converters.mil.mil import types from safetensors import safe_open def weight_at(folder, key): for p in sorted(Path(folder).glob("*.safetensors")): with safe_open(p, framework="pt", device="cpu") as f: if key in f.keys(): return f.get_tensor(key).float().numpy() raise KeyError(key) def devices(plan): rows = [] def walk(block): for op in block.operations: usage = plan.get_compute_device_usage_for_mlprogram_operation(op) cost = plan.get_estimated_cost_for_mlprogram_operation(op) rows.append( { "operator": op.operator_name, "preferred": type(usage.preferred_compute_device).__name__ if usage else None, "supported": [ type(d).__name__ for d in usage.supported_compute_devices ] if usage else [], "estimated_cost": cost.weight if cost else None, } ) for child in op.blocks: walk(child) if plan.model_structure.program: for fn in plan.model_structure.program.functions.values(): walk(fn.block) return rows def projection_program(shape, dtype, weight, conv): @mb.program( input_specs=[mb.TensorSpec(shape=shape, dtype=dtype)], opset_version=ct.target.macOS15, ) def program(x): if conv: return mb.conv(x=x, weight=weight[:, :, None, None], name="projection") return mb.linear(x=x, weight=weight, name="projection") return program def main(a): dest = Path(a.output) dest.parent.mkdir(parents=True, exist_ok=True) root = Path(a.packages) root.mkdir(parents=True, exist_ok=True) report = { "platform": platform.platform(), "coremltools": ct.__version__, "repeats": a.repeats, "scope": "Real checkpoint projection weights; synthetic inputs. Synchronous predict includes NumPy I/O. Component tests, not full image latency.", "trials": [], } torch.set_num_threads(10) cases = [ ( "dit-fc2", Path(a.model) / "transformer", "blocks.0.mlp.fc2.weight", 4181, False, ), ( "dit-fc2-conv256", Path(a.model) / "transformer", "blocks.0.mlp.fc2.weight", 256, True, ), ] if a.vae: cases.extend( [ ( "vae-fc2", Path(a.vae), "decoder.transformer_blocks.0.ff.w2.weight", 2088, False, ), ( "vae-fc2-conv261", Path(a.vae), "decoder.transformer_blocks.0.ff.w2.weight", 261, True, ), ] ) for name, folder, key, tokens, conv in cases: if a.cases and name not in a.cases: continue w = weight_at(folder, key) rng = np.random.default_rng(731) # Round inputs to BF16 so the input precision matches the MLX study. x = rng.standard_normal((tokens, w.shape[1]), dtype=np.float32) u = x.view(np.uint32) x = ((u + np.uint32(0x7FFF) + ((u >> 16) & 1)) & np.uint32(0xFFFF0000)).view( np.float32 ) reference = (torch.from_numpy(x) @ torch.from_numpy(w).T).numpy() if conv: shape = (1, w.shape[1], 1, tokens) input_array = x.T[None, :, None, :] else: shape = x.shape input_array = x for precision in ("fp16", "fp32"): package = root / f"{name}-{precision}.mlpackage" dtype = types.fp16 if precision == "fp16" else types.fp32 if not package.exists(): program = projection_program( shape, dtype, w.astype(np.float16 if precision == "fp16" else np.float32), conv, ) model = ct.convert( program, convert_to="mlprogram", compute_precision=ct.precision.FLOAT16 if precision == "fp16" else ct.precision.FLOAT32, minimum_deployment_target=ct.target.macOS15, skip_model_load=True, ) model.save(str(package)) del model, program gc.collect() for unit in ( ct.ComputeUnit.CPU_ONLY, ct.ComputeUnit.CPU_AND_GPU, ct.ComputeUnit.CPU_AND_NE, ct.ComputeUnit.ALL, ): row = { "case": name, "precision": precision, "compute_units": unit.name, "shape": list(shape), "weight_shape": list(w.shape), "package_bytes": sum( p.stat().st_size for p in package.rglob("*") if p.is_file() ), } try: tick = time.perf_counter() model = ct.models.MLModel( str(package), compute_units=unit, optimization_hints={ "specializationStrategy": ct.SpecializationStrategy.FastPrediction }, ) row["load_compile_s"] = time.perf_counter() - tick plan = ct.models.compute_plan.MLComputePlan.load_from_path( model.get_compiled_model_path(), compute_units=unit ) row["placement"] = devices(plan) data = { "x": input_array.astype( np.float16 if precision == "fp16" else np.float32 ) } tick = time.perf_counter() result = model.predict(data) row["first_predict_ms"] = 1000 * (time.perf_counter() - tick) samples = [] for _ in range(a.repeats): tick = time.perf_counter() result = model.predict(data) samples.append(1000 * (time.perf_counter() - tick)) y = next(iter(result.values())).astype(np.float32) if conv: y = y[0, :, 0, :].T row.update( samples_ms=samples, median_ms=statistics.median(samples), relative_rmse=float( np.sqrt( np.mean((reference - y) ** 2) / np.mean(reference**2) ) ), max_abs=float(np.max(np.abs(reference - y))), ) del result, model, plan, data gc.collect() except Exception as exc: row["error"] = f"{type(exc).__name__}: {exc}" report["trials"].append(row) dest.write_text(json.dumps(report, indent=2) + "\n") print(json.dumps(row), flush=True) del w, x, reference gc.collect() if __name__ == "__main__": p = argparse.ArgumentParser(description=__doc__) p.add_argument("--model", required=True) p.add_argument("--vae") p.add_argument("--cases", nargs="+") p.add_argument("--repeats", type=int, default=5) p.add_argument("--packages", default="cache/coreml-v2") p.add_argument("--output", default="reports/optimization-v2/coreml.json") main(p.parse_args())