Native serialization

backends/mlx/passes.py

@@ -99,49 +99,8 @@ def call(self, exported_program) -> ExportedProgramPassResult:
         }
         if ops_to_inplace:
             reinplace_pass(exported_program, ops_to_inplace=ops_to_inplace)
-            self._resync_output_specs(exported_program)
         return ExportedProgramPassResult(exported_program, True)
 
-    @staticmethod
-    def _resync_output_specs(exported_program) -> None:
-        """Re-sync graph-signature output names after reinplace.
-
-        ``reinplace_pass`` rewrites an output-producing node (e.g. the final
-        ``exp`` -> ``exp_``) via ``replace_all_uses_with`` + erase, but does not
-        update ``graph_signature.output_specs``. Output order is preserved, so we
-        positionally re-sync each spec's argument name to the current output node
-        arg; otherwise ``ExportedProgram.validate()`` (run by the pass manager)
-        raises a SpecViolationError.
-
-        This positional pairing is only valid because reinplace does a 1:1
-        ``replace_all_uses_with`` + erase and never drops, adds, or reorders
-        outputs. We assert ``len(output_specs) == len(out_args)`` so that a
-        future change violating that invariant fails loudly here instead of
-        silently mis-pairing names (``zip`` would otherwise truncate). Specs
-        whose ``arg`` is not a named tensor (e.g. ``ConstantArgument``) carry no
-        ``name`` and are skipped by the ``getattr`` guard below.
-        """
-        out_node = next(
-            n for n in reversed(exported_program.graph.nodes) if n.op == "output"
-        )
-        out_args = out_node.args[0]
-        if not isinstance(out_args, (tuple, list)):
-            out_args = (out_args,)
-        output_specs = exported_program.graph_signature.output_specs
-        assert len(output_specs) == len(out_args), (
-            "reinplace changed graph output count: "
-            f"{len(output_specs)} output_specs vs {len(out_args)} output args. "
-            "Positional output-spec re-sync assumes a 1:1, order-preserving "
-            "rewrite."
-        )
-        for spec, arg in zip(output_specs, out_args):
-            if (
-                isinstance(arg, torch.fx.Node)
-                and getattr(spec.arg, "name", None) is not None
-                and spec.arg.name != arg.name
-            ):
-                spec.arg.name = arg.name
-
 
 @dataclass
 class RMSNormMatch(PatternMatch):

backends/native/fat_pte.py

@@ -0,0 +1,77 @@
+# 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.
+
+"""
+Fat PTE: packs multiple backend specializations into one delegate payload.
+Runtime selects the best specialization at load time. All specializations
+share a single PTD (named data store).
+
+Binary layout:
+  [4B] magic "NFAT"
+  [4B] version (1)
+  [4B] num_specializations
+  Per specialization:
+    [32B] backend_id (utf-8, null-padded)
+    [8B]  offset into data section
+    [8B]  length
+  [payload bytes ...]
+"""
+
+import struct
+from typing import Dict, List, Optional, Tuple
+
+from executorch.exir._serialize._named_data_store import NamedDataStore
+from executorch.exir.backend.backend_details import PreprocessResult
+
+FAT_MAGIC = b"NFAT"
+FAT_VERSION = 1
+_ENTRY_FMT = "32sQQ"
+_ENTRY_SIZE = struct.calcsize(_ENTRY_FMT)
+
+
+def pack_fat_blob(specializations: List[Tuple[str, bytes]]) -> bytes:
+    """Pack (backend_id, payload) pairs into a fat blob."""
+    header = struct.pack("<4sII", FAT_MAGIC, FAT_VERSION, len(specializations))
+
+    entries = []
+    offset = 0
+    for backend_id, payload in specializations:
+        encoded = backend_id.encode("ascii")
+        if len(encoded) > 32:
+            raise ValueError(
+                f"Backend ID '{backend_id}' is {len(encoded)} bytes; max is 32"
+            )
+        bid = encoded.ljust(32, b"\x00")
+        entries.append(struct.pack("<" + _ENTRY_FMT, bid, offset, len(payload)))
+        offset += len(payload)
+
+    return header + b"".join(entries) + b"".join(p for _, p in specializations)
+
+
+def build_fat_result(
+    results: List[Tuple[str, PreprocessResult]],
+    debug_handle_map: Optional[Dict] = None,
+) -> PreprocessResult:
+    """Merge (backend_id, PreprocessResult) pairs into a single fat result."""
+    fat_entries: List[Tuple[str, bytes]] = []
+    merged_data_store = NamedDataStore()
+
+    for backend_id, result in results:
+        fat_entries.append((backend_id, result.processed_bytes))
+
+        if result.data_store_output:
+            for key, entry in result.data_store_output.pte_data.items():
+                buf = result.data_store_output.buffers[entry.buffer_index]
+                merged_data_store.add_named_data(key, buf, alignment=entry.alignment)
+
+    if debug_handle_map is None and results:
+        debug_handle_map = results[0][1].debug_handle_map
+
+    return PreprocessResult(
+        processed_bytes=pack_fat_blob(fat_entries),
+        debug_handle_map=debug_handle_map,
+        data_store_output=merged_data_store.get_named_data_store_output(),
+    )

backends/native/partitioner.py

@@ -0,0 +1,135 @@
+# 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.
+
+"""
+Native backend partitioner.
+
+Claims core ATen ops (torch.Tag.core) plus an explicit opt-in set.
+"""
+
+import json
+
+from typing import Callable, final, List, Mapping, Optional, Tuple
+
+import executorch.backends.native.specializations  # noqa: F401 — register recipes
+
+import torch
+
+from executorch.exir.backend.backend_details import CompileSpec
+from executorch.exir.backend.partitioner import (
+    DelegationSpec,
+    Partitioner,
+    PartitionResult,
+)
+from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
+
+from torch.export.exported_program import ExportedProgram
+from torch.fx import Node
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
+from torch.fx.passes.operator_support import OperatorSupportBase
+
+# Non-core ops the native backend supports. Also preserved (not decomposed).
+_SUPPORTED_NON_CORE_OPS = [
+    torch.ops.aten.matmul.default,
+    torch.ops.aten.linear.default,
+    torch.ops.aten.addmm.default,
+    torch.ops.aten.scaled_dot_product_attention.default,
+]
+
+
+class NativeSupportedOperators(OperatorSupportBase):
+    _NON_CORE = set(_SUPPORTED_NON_CORE_OPS)
+
+    def is_node_supported(
+        self, submodules: Mapping[str, torch.nn.Module], node: Node
+    ) -> bool:
+        if node.op in ("placeholder", "output", "get_attr"):
+            return False
+        if node.op != "call_function":
+            return False
+        if isinstance(node.target, torch._ops.HigherOrderOperator):
+            return False
+
+        from executorch.exir.dialects.edge._ops import EdgeOpOverload
+
+        target = node.target
+        if isinstance(target, EdgeOpOverload):
+            target = target._op
+        if isinstance(target, torch._ops.OpOverload):
+            if target in self._NON_CORE:
+                return True
+            return torch.Tag.core in target.tags or torch.Tag.view_copy in target.tags
+        return False
+
+
+@final
+class NativePartitioner(Partitioner):
+    _SPECIALIZATIONS_KEY = "native_specializations"
+
+    def __init__(
+        self,
+        specializations: Optional[List[str]] = None,
+    ) -> None:
+        self.specializations = specializations
+        from executorch.backends.native.preprocess import NativeBackend
+
+        compile_specs = []
+        if specializations:
+            compile_specs.append(
+                CompileSpec(
+                    self._SPECIALIZATIONS_KEY,
+                    json.dumps(specializations).encode("utf-8"),
+                )
+            )
+        self.delegation_spec = DelegationSpec(NativeBackend.__name__, compile_specs)
+
+    def ops_to_not_decompose(
+        self, ep: ExportedProgram
+    ) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[Node], bool]]]:
+        # Already-partitioned graph -> nothing to preserve.
+        from executorch.exir.lowered_backend_module import executorch_call_delegate
+
+        for node in ep.graph.nodes:
+            if node.op == "call_function" and node.target is executorch_call_delegate:
+                return ([], None)
+
+        present: List[torch._ops.OpOverload] = []
+        seen = set()
+        for node in ep.graph.nodes:
+            if node.op != "call_function":
+                continue
+            if not isinstance(node.target, torch._ops.OpOverload):
+                continue
+            if node.target in _SUPPORTED_NON_CORE_OPS and node.target not in seen:
+                present.append(node.target)
+                seen.add(node.target)
+
+        return (present, None)
+
+    def partition(self, exported_program: ExportedProgram) -> PartitionResult:
+        partition_tags = {}
+
+        capability_partitioner = CapabilityBasedPartitioner(
+            exported_program.graph_module,
+            NativeSupportedOperators(),
+            allows_single_node_partition=True,
+        )
+
+        partition_list = capability_partitioner.propose_partitions()
+
+        for partition in partition_list:
+            for node in partition.nodes:
+                tag = f"tag{partition.id}"
+                node.meta["delegation_tag"] = tag
+                partition_tags[tag] = self.delegation_spec
+
+        tag_constant_data(exported_program)
+        tag_mutated_buffer(exported_program)
+
+        return PartitionResult(
+            tagged_exported_program=exported_program,
+            partition_tags=partition_tags,
+        )