| """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) |
| |
| 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()) |
|
|
