Core ML 경로와 뉴럴 엔진

#!/usr/bin/env python3
"""Build two Core ML models straight from MIL (no PyTorch needed), then time each
on every compute-unit setting and ask the compute plan which unit ran each op.
Model A: a decode-shaped matvec over a large fp16 matrix (bandwidth-bound).
Model B: a conv stack on an image (the Neural Engine's home turf)."""
import time, os, sys
import numpy as np
import coremltools as ct
from coremltools.converters.mil import Builder as mb
from coremltools.converters.mil.mil import types

OUT = os.path.expanduser("~/silicon-lab/coreml")
os.makedirs(OUT, exist_ok=True)
UNITS = {"CPU_ONLY": ct.ComputeUnit.CPU_ONLY, "CPU_AND_GPU": ct.ComputeUnit.CPU_AND_GPU,
         "CPU_AND_NE": ct.ComputeUnit.CPU_AND_NE, "ALL": ct.ComputeUnit.ALL}


def build_matvec(n: int):
    W = (np.random.randn(n, n) * 0.02).astype(np.float16)

    @mb.program(input_specs=[mb.TensorSpec(shape=(1, n), dtype=types.fp16)], opset_version=ct.target.iOS17)
    def prog(x):
        return mb.matmul(x=x, y=W, name="y")
    return ct.convert(prog, convert_to="mlprogram", minimum_deployment_target=ct.target.macOS14,
                      compute_precision=ct.precision.FLOAT16), {"x": np.random.randn(1, n).astype(np.float16)}, W.nbytes


def build_conv(c: int = 128, layers: int = 8, hw: int = 224):
    Ws = [(np.random.randn(c, c if i else 3, 3, 3) * 0.05).astype(np.float16) for i in range(layers)]

    @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, hw, hw), dtype=types.fp16)], opset_version=ct.target.iOS17)
    def prog(x):
        h = x
        for i, W in enumerate(Ws):
            h = mb.conv(x=h, weight=W, pad_type="same", name=f"conv{i}")
            h = mb.relu(x=h, name=f"relu{i}")
        return mb.reduce_mean(x=h, axes=[2, 3], name="y")
    flops = sum(2 * hw * hw * c * (c if i else 3) * 9 for i in range(layers))
    return ct.convert(prog, convert_to="mlprogram", minimum_deployment_target=ct.target.macOS14,
                      compute_precision=ct.precision.FLOAT16), {"x": np.random.randn(1, 3, hw, hw).astype(np.float16)}, flops


def plan_devices(path, unit):
    keep = ct.models.MLModel(path, compute_units=unit)          # the plan wants the .mlmodelc, which lives as long as this object
    plan = ct.models.compute_plan.MLComputePlan.load_from_path(path=keep.get_compiled_model_path(), compute_units=unit)
    used = {}
    for op in plan.model_structure.program.functions["main"].block.operations:
        du = plan.get_compute_device_usage_for_mlprogram_operation(op)
        if du is not None:
            k = type(du.preferred_compute_device).__name__.replace("ML", "").replace("ComputeDevice", "")
            used[k] = used.get(k, 0) + 1
    return used


def bench(name, model, inputs, reps=20):
    path = os.path.join(OUT, f"{name}.mlpackage")
    model.save(path)
    print(f"\n== {name}")
    for label, unit in UNITS.items():
        t = time.perf_counter(); m = ct.models.MLModel(path, compute_units=unit); load = time.perf_counter() - t
        for _ in range(3):
            m.predict(inputs)                                        # warm-up (the ANE compile happens here)
        t = time.perf_counter()
        for _ in range(reps):
            m.predict(inputs)
        print(f"{label:12} load {load:5.2f}s  predict {(time.perf_counter() - t) / reps * 1e3:8.2f} ms")
    print("planner, ALL:", plan_devices(path, ct.ComputeUnit.ALL), " CPU_AND_GPU:", plan_devices(path, ct.ComputeUnit.CPU_AND_GPU))


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "both"
    if which in ("matvec", "both"):
        for n in (4096, 8192, 16384):
            model, inputs, nbytes = build_matvec(n)
            print(f"\nmatvec {n}x{n} fp16 = {nbytes/1e6:.0f} MB of weights")
            bench(f"matvec{n}", model, inputs)
    if which in ("conv", "both"):
        model, inputs, flops = build_conv()
        print(f"\nconv stack: {flops/1e9:.1f} GFLOP per image")
        bench("conv8", model, inputs)

# office, M3 Ultra, coremltools 8.1, macOS 26.6.2, 2026-09-15 (predict ms; planner's preferred unit):
# matvec 16384^2 (537 MB): CPU_ONLY 8.04  CPU_AND_GPU 1.71  CPU_AND_NE 30.18  ALL 30.74   ALL -> NeuralEngine (load 24 s), CPU_AND_GPU -> GPU
# matvec  8192^2 (134 MB): CPU_ONLY 2.15  CPU_AND_GPU 2.18  CPU_AND_NE  2.56  ALL  2.10   ALL -> NeuralEngine,            CPU_AND_GPU -> CPU
# matvec  4096^2  (34 MB): CPU_ONLY 0.65  CPU_AND_GPU 0.66  CPU_AND_NE  0.62  ALL  0.63   ALL -> CPU,                     CPU_AND_GPU -> CPU
# conv8 (104 GFLOP):       CPU_ONLY 12.95 CPU_AND_GPU 4.65  CPU_AND_NE  6.15  ALL  6.12   ALL -> NeuralEngine x17,        CPU_AND_GPU -> GPU x17