Merge nucleic/sleek-ember-seal-uady into dev

inspect_coreml.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""Inspect Core ML operation placement and estimated accelerator cost share."""
+
+from __future__ import annotations
+
+import argparse
+import platform
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any, Sequence
+
+from purpose_data import DataError, write_json
+
+
+def _device_category(device: Any) -> str:
+    name = type(device).__name__.lower()
+    description = str(device).lower()
+    combined = f"{name} {description}"
+    if "neural" in combined:
+        return "neuralEngine"
+    if "gpu" in combined:
+        return "gpu"
+    if "cpu" in combined:
+        return "cpu"
+    return "unknown"
+
+
+def _compute_unit(coremltools: Any, requested: str) -> Any:
+    values = {
+        "all": coremltools.ComputeUnit.ALL,
+        "cpu-only": coremltools.ComputeUnit.CPU_ONLY,
+        "cpu-and-gpu": coremltools.ComputeUnit.CPU_AND_GPU,
+        "cpu-and-ne": coremltools.ComputeUnit.CPU_AND_NE,
+    }
+    return values[requested]
+
+
+def inspect(args: argparse.Namespace) -> dict[str, Any]:
+    if not args.model.exists():
+        raise DataError(f"{args.model}: Core ML model is missing")
+    try:
+        import coremltools as ct
+    except ImportError as exc:
+        raise DataError(
+            "Core ML inspection requires requirements-coreml.txt on macOS"
+        ) from exc
+
+    compiled = ct.models.utils.compile_model(str(args.model))
+    compute_plan = ct.models.compute_plan.MLComputePlan.load_from_path(
+        path=str(compiled),
+        compute_units=_compute_unit(ct, args.compute_units),
+    )
+    program = compute_plan.model_structure.program
+    if program is None or "main" not in program.functions:
+        raise DataError("Core ML package is not an ML Program with a main function")
+    operations = list(program.functions["main"].block.operations)
+    if not operations:
+        raise DataError("Core ML compute plan contains no operations")
+
+    preferred_counts: Counter[str] = Counter()
+    preferred_costs: dict[str, float] = defaultdict(float)
+    supported_counts: Counter[str] = Counter()
+    operation_reports = []
+    operations_with_usage = 0
+    operations_with_cost = 0
+    total_cost = 0.0
+    for operation in operations:
+        usage = compute_plan.get_compute_device_usage_for_mlprogram_operation(
+            operation
+        )
+        cost = compute_plan.get_estimated_cost_for_mlprogram_operation(operation)
+        preferred = "unknown"
+        supported: list[str] = []
+        if usage is not None:
+            operations_with_usage += 1
+            preferred = _device_category(usage.preferred_compute_device)
+            preferred_counts[preferred] += 1
+            supported = sorted(
+                {_device_category(device) for device in usage.supported_compute_devices}
+            )
+            supported_counts.update(supported)
+        weight = None
+        if cost is not None:
+            operations_with_cost += 1
+            weight = float(cost.weight)
+            total_cost += weight
+            preferred_costs[preferred] += weight
+        operation_reports.append(
+            {
+                "operatorName": str(operation.operator_name),
+                "preferredDevice": preferred,
+                "supportedDevices": supported,
+                "estimatedCostWeight": weight,
+            }
+        )
+
+    ane_operations = preferred_counts["neuralEngine"]
+    ane_cost = preferred_costs["neuralEngine"]
+    report = {
+        "schemaVersion": 1,
+        "model": str(args.model),
+        "coremltoolsVersion": ct.__version__,
+        "machine": platform.machine(),
+        "macOS": platform.mac_ver()[0],
+        "computeUnits": args.compute_units,
+        "operations": len(operations),
+        "operationsWithDeviceUsage": operations_with_usage,
+        "operationsWithEstimatedCost": operations_with_cost,
+        "preferredOperationCounts": dict(sorted(preferred_counts.items())),
+        "supportedOperationCounts": dict(sorted(supported_counts.items())),
+        "preferredEstimatedCosts": dict(sorted(preferred_costs.items())),
+        "neuralEngineOperationShare": (
+            ane_operations / operations_with_usage if operations_with_usage else 0.0
+        ),
+        "neuralEngineEstimatedCostShare": (
+            ane_cost / total_cost if total_cost else 0.0
+        ),
+        "operationDetails": operation_reports,
+    }
+    write_json(args.report, report)
+    print(
+        "Core ML placement: "
+        f"ANE operations={report['neuralEngineOperationShare']:.2%} "
+        f"ANE estimated cost={report['neuralEngineEstimatedCostShare']:.2%}"
+    )
+    return report
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", type=Path, required=True)
+    parser.add_argument("--report", type=Path, required=True)
+    parser.add_argument(
+        "--compute-units",
+        choices=("all", "cpu-only", "cpu-and-gpu", "cpu-and-ne"),
+        default="cpu-and-ne",
+    )
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    try:
+        inspect(args)
+    except (DataError, OSError, RuntimeError, ValueError) as exc:
+        print(f"error: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())