[ExecuTorch][Core AI][2/x] Core backend: partitioner, preprocess, IO compat, copy-op pass

backends/apple/coreai/README.md

@@ -0,0 +1,7 @@
+# Core AI backend (ExecuTorch)
+
+> ⚠️ **Under construction — not for use.**
+> This backend is a work in progress. It is **not** ready for production or
+> general use, and provides **no backward- or forward-compatibility guarantees**
+> (no BC/FC). APIs, compile specs, serialized manifests, and file layouts may
+> change or break at any time without notice. Do not depend on it.

backends/apple/coreai/__init__.py

@@ -0,0 +1,48 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import List
+
+from executorch.backends.apple.coreai.compiler.preprocess import (
+    AOTCompileConfig,
+    coreai_sidecar_dir,
+    CoreAIBackend,
+)
+from executorch.backends.apple.coreai.partition.partitioner import CoreAIPartitioner
+from executorch.exir import EdgeCompileConfig
+from executorch.exir.pass_base import PassType
+
+
+def get_default_passes() -> List[PassType]:
+    """Default edge transform passes for Core AI lowering.
+
+    ``NarrowToCoreAIDtypesPass`` casts int64/float64 graph inputs to 32-bit at
+    the boundary (preserving the model's external I/O dtype) so index-style
+    inputs can be delegated. Core AI supports only up-to-32-bit dtypes.
+    """
+    from executorch.backends.apple.coreai.passes import NarrowToCoreAIDtypesPass
+
+    return [NarrowToCoreAIDtypesPass()]
+
+
+def get_default_compile_config() -> EdgeCompileConfig:
+    """Default ``EdgeCompileConfig`` for Core AI lowering.
+
+    ``_skip_dim_order=True`` keeps ExecuTorch on ``aten._to_copy`` instead of
+    emitting ``dim_order_ops._to_dim_order_copy``, which coreai-torch cannot
+    lower (its validator requires dim-order ops be decomposed).
+    """
+    return EdgeCompileConfig(_skip_dim_order=True)
+
+
+__all__ = [
+    "AOTCompileConfig",
+    "CoreAIBackend",
+    "CoreAIPartitioner",
+    "coreai_sidecar_dir",
+    "get_default_compile_config",
+    "get_default_passes",
+]

backends/apple/coreai/compiler/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.

backends/apple/coreai/compiler/constants.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Shared constants for the Core AI backend compiler."""
+
+# The single entrypoint / graph name of a per-delegate coreai AIProgram.  Each
+# ExecuTorch delegate is converted to its own coreai program whose only graph is
+# named "main" (``save_asset`` emits ``main.hash`` / ``main.mlirb`` accordingly).
+MAIN_ENTRYPOINT = "main"

backends/apple/coreai/compiler/io_compat.py

@@ -0,0 +1,168 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Delegate-boundary compatibility checks for the Core AI backend.
+
+ExecuTorch calls a delegate with plain tensors whose dtypes/shapes come from the
+edge subgraph boundary.  The converted ``.aimodel`` must declare I/O that matches
+what ExecuTorch will feed/read, or the runtime will mismatch.  :func:`
+assert_io_compatible` compares the two *before* compilation and fails loudly.
+"""
+
+import logging
+from typing import Any, List, Sequence, Tuple
+
+import torch
+from executorch.backends.apple.coreai.compiler.constants import MAIN_ENTRYPOINT
+from torch.export.exported_program import ExportedProgram
+
+logger = logging.getLogger(__name__)
+
+# coreai's element-type spelling for each torch dtype (from coreai_torch's
+# _type_mapping: signed ints -> siN, unsigned -> uiN, bool -> i1, floats -> fN).
+# These are the un-narrowed spellings, asserted exactly. coreai narrows
+# int64->si32 and float64->f32 in get_tensor_type, so those two dtypes surface
+# here as mismatches by design (that is the signal).
+_TORCH_TO_COREAI_ELT = {
+    torch.float16: "f16",
+    torch.float32: "f32",
+    torch.float64: "f64",
+    torch.bfloat16: "bf16",
+    torch.bool: "i1",
+    torch.int8: "si8",
+    torch.uint8: "ui8",
+    torch.int16: "si16",
+    torch.uint16: "ui16",
+    torch.int32: "si32",
+    torch.uint32: "ui32",
+    torch.int64: "si64",
+}
+
+
+def _coreai_io(program, entrypoint: str = MAIN_ENTRYPOINT):
+    """Core AI graph I/O as lists of (element_type_str, shape_or_None).
+
+    Ranked tensors -> (element_type_str, shape_tuple).  Any non-tensor coreai
+    type is represented as (str(type), None) so the comparison can flag it,
+    rather than assuming a tensor and crashing on ``.element_type`` / ``.shape``.
+    """
+    from coreai._compiler.ir import RankedTensorType
+
+    func_type = program.get_graph(entrypoint).function_type.value
+
+    def _describe(t):
+        if RankedTensorType.isinstance(t):
+            return (str(t.element_type), tuple(t.shape))
+        return (str(t), None)
+
+    inputs = [_describe(t) for t in func_type.inputs]
+    outputs = [_describe(t) for t in func_type.results]
+    return inputs, outputs
+
+
+def _edge_io(edge_program: ExportedProgram):
+    """Edge subgraph I/O as lists of (dtype_or_typename, shape_or_None).
+
+    Tensors -> (torch_dtype, shape_tuple).  Non-tensor I/O (symint / const
+    int/float) -> (type_name_str, None), kept as entries so both sides stay
+    positionally aligned with :func:`_coreai_io` (which also emits non-tensor
+    entries).
+    """
+    placeholders = {
+        n.name: n for n in edge_program.graph.nodes if n.op == "placeholder"
+    }
+    inputs = []
+    for name in edge_program.graph_signature.user_inputs:
+        # A non-tensor input appears in user_inputs as its literal value, not a
+        # placeholder name.
+        node = placeholders.get(name)
+        val = node.meta.get("val") if node is not None else name
+        if hasattr(val, "dtype"):
+            inputs.append((val.dtype, tuple(val.shape)))
+        else:
+            inputs.append((type(val).__name__, None))
+    outputs = []
+    for arg in edge_program.graph.output_node().args[0]:
+        val = arg.meta.get("val") if hasattr(arg, "meta") else arg
+        if hasattr(val, "dtype"):
+            outputs.append((val.dtype, tuple(val.shape)))
+        else:
+            outputs.append((type(val).__name__, None))
+    return inputs, outputs
+
+
+def io_mismatches(
+    coreai_io: Sequence[Tuple[str, Tuple[Any, ...]]],
+    edge_io: Sequence[Tuple[Any, Tuple[Any, ...]]],
+    kind: str,
+) -> List[str]:
+    """Return human-readable mismatch messages between coreai and edge I/O.
+
+    Compares count, per-tensor dtype (all mapped types: floats and ints), rank,
+    and static (concrete) dims. Symbolic (dynamic) dims are skipped: coreai is
+    specialized to a concrete size while the edge dim stays symbolic.
+    """
+    if len(coreai_io) != len(edge_io):
+        return [
+            f"{kind} count mismatch: .aimodel has {len(coreai_io)}, "
+            f"ExecuTorch feeds {len(edge_io)}"
+        ]
+    errors: List[str] = []
+    for i, ((elt, cshape), (dtype, eshape)) in enumerate(zip(coreai_io, edge_io)):
+        coreai_nontensor = cshape is None
+        edge_nontensor = eshape is None
+        if coreai_nontensor or edge_nontensor:
+            # At least one side is non-tensor; only a class mismatch matters.
+            if coreai_nontensor != edge_nontensor:
+                errors.append(
+                    f"{kind} {i}: tensor/non-tensor mismatch "
+                    f"(.aimodel={elt}, ExecuTorch feeds {dtype})"
+                )
+            continue
+        expected = _TORCH_TO_COREAI_ELT.get(dtype)
+        if expected is not None and elt != expected:
+            errors.append(
+                f"{kind} {i}: dtype mismatch (.aimodel={elt}, ExecuTorch feeds "
+                f"{dtype}, expected {expected})"
+            )
+        if len(cshape) != len(eshape):
+            errors.append(
+                f"{kind} {i}: rank mismatch (.aimodel={cshape}, "
+                f"ExecuTorch={eshape})"
+            )
+            continue
+        for d, (cdim, edim) in enumerate(zip(cshape, eshape)):
+            if isinstance(edim, int) and cdim != edim:
+                errors.append(
+                    f"{kind} {i} dim {d}: static shape mismatch "
+                    f"(.aimodel={cdim}, ExecuTorch={edim})"
+                )
+    return errors
+
+
+def assert_io_compatible(program, edge_program: ExportedProgram) -> None:
+    """Raise if the ``.aimodel`` boundary I/O is incompatible with ExecuTorch."""
+    coreai_in, coreai_out = _coreai_io(program)
+    edge_in, edge_out = _edge_io(edge_program)
+    logger.info(
+        "Core AI delegate boundary I/O:\n"
+        "  inputs:  .aimodel=%s\n"
+        "           ExecuTorch=%s\n"
+        "  outputs: .aimodel=%s\n"
+        "           ExecuTorch=%s",
+        coreai_in,
+        edge_in,
+        coreai_out,
+        edge_out,
+    )
+    errors = io_mismatches(coreai_in, edge_in, "input") + io_mismatches(
+        coreai_out, edge_out, "output"
+    )
+    if errors:
+        raise ValueError(
+            "Core AI delegate boundary is incompatible with ExecuTorch:\n  "
+            + "\n  ".join(errors)
+        )