up

Author: metascroy

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

@@ -39,7 +39,12 @@ def pack_fat_blob(specializations: List[Tuple[str, bytes]]) -> bytes:
     entries = []
     offset = 0
     for backend_id, payload in specializations:
-        bid = backend_id.encode("utf-8")[:32].ljust(32, b"\x00")
+        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)
 

backends/native/partitioner.py

@@ -10,12 +10,15 @@
 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,
@@ -64,21 +67,33 @@ def is_node_supported(
 
 @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
 
-        self.delegation_spec = DelegationSpec(NativeBackend.__name__, [])
+        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 == "get_attr" and "lowered_module" in node.name:
+            if node.op == "call_function" and node.target is executorch_call_delegate:
                 return ([], None)
 
         present: List[torch._ops.OpOverload] = []
@@ -95,10 +110,6 @@ def ops_to_not_decompose(
         return (present, None)
 
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
-        from executorch.backends.native.preprocess import NativeBackend
-
-        NativeBackend._specialization_names = self.specializations
-
         partition_tags = {}
 
         capability_partitioner = CapabilityBasedPartitioner(

backends/native/preprocess.py

@@ -11,8 +11,9 @@
 memory planning → emit → serialize via _program_to_flatbuffer.
 """
 
+import json
 from functools import partial
-from typing import final, List, Optional, Sequence, Tuple
+from typing import final, List, Tuple
 
 import torch
 
@@ -65,6 +66,8 @@ def _build_backend_inplace_ops() -> frozenset:
 
 
 class _AnyOp(Verifier):
+    # "TRAINING" is the only dialect that allows non-functional (mutating) ops.
+    # After reinplace converts e.g. relu → relu_, the verifier must accept them.
     dialect = "TRAINING"
 
     def allowed_op_types(self):
@@ -93,16 +96,17 @@ def _apply_passes(program: ExportedProgram, passes) -> ExportedProgram:
 
 @final
 class NativeBackend(BackendDetails):
-    # Set by NativePartitioner before preprocess is called.
-    _specialization_names: Optional[Sequence[str]] = None
-
     @classmethod
     def preprocess(
         cls,
         program: ExportedProgram,
         module_compile_spec: List[CompileSpec],
     ) -> PreprocessResult:
-        names = cls._specialization_names
+        names = None
+        for spec in module_compile_spec:
+            if spec.key == "native_specializations":
+                names = json.loads(spec.value.decode("utf-8"))
+
         if not names:
             return cls._preprocess_native(program)
 
@@ -140,22 +144,6 @@ def _preprocess_native(
 
         program = _et_reinplace_pass(program, ops_to_inplace=BACKEND_INPLACE_OPS)
 
-        # reinplace may rename output nodes; fix the graph signature to match.
-        output_node = next(
-            n for n in program.graph_module.graph.nodes if n.op == "output"
-        )
-        actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")]
-        new_output_specs = []
-        for spec, name in zip(program.graph_signature.output_specs, actual_names):
-            new_output_specs.append(
-                type(spec)(
-                    kind=spec.kind,
-                    arg=type(spec.arg)(name=name),
-                    target=spec.target,
-                )
-            )
-        program._graph_signature.output_specs = new_output_specs
-
         # Re-run SpecPropPass: reinplace copies meta["val"] but not meta["spec"].
         program = _apply_passes(program, [SpecPropPass()])
 
@@ -212,6 +200,7 @@ def _preprocess_native(
             algo_list=[greedy_memory_planning]
         )
 
+        # Tells the emitter to use out-variant kernels for ops that support them.
         program.graph_module.encounter_to_out_var_failure = True
 
         program = _apply_passes(

backends/native/test/test_fat_pte.py

@@ -65,13 +65,10 @@ def test_empty_specializations(self):
         self.assertEqual(count, 0)
         self.assertEqual(len(blob), 12)
 
-    def test_backend_id_truncated_to_32_bytes(self):
-        long_name = "A" * 100
-        blob = pack_fat_blob([(long_name, b"x")])
-        header_size = 12
-        bid, _, _ = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size)
-        self.assertEqual(len(bid), 32)
-        self.assertEqual(bid, b"A" * 32)
+    def test_backend_id_over_32_bytes_raises(self):
+        long_name = "A" * 33
+        with self.assertRaises(ValueError):
+            pack_fat_blob([(long_name, b"x")])
 
 
 class TestBuildFatResult(unittest.TestCase):

backends/native/test/test_preprocess.py

@@ -12,7 +12,6 @@
 
 from executorch.backends.native.fat_pte import FAT_MAGIC
 from executorch.backends.native.partitioner import NativePartitioner
-from executorch.backends.native.preprocess import NativeBackend
 from executorch.backends.native.specializations import (
     _SPECIALIZATION_REGISTRY,
     register_specialization,
@@ -45,7 +44,6 @@ def setUp(self):
     def tearDown(self):
         _SPECIALIZATION_REGISTRY.clear()
         _SPECIALIZATION_REGISTRY.update(self._saved)
-        NativeBackend._specialization_names = None
 
     def test_register_and_lookup(self):
         register_specialization("TestBackend", lambda ep: b"test")

exir/passes/reinplace.py

@@ -507,4 +507,26 @@ def reinplace_pass(  # noqa: C901
             continue
 
         seen_nodes.update(node.all_input_nodes)
+
+    # Reinplace rewrites may rename output nodes (e.g. aten_relu_default →
+    # aten_relu__default) but replace_all_uses_with does not update the
+    # graph signature. Fix output_specs to match the actual output node args.
+    output_node = next((n for n in ep.graph.nodes if n.op == "output"), None)
+    if output_node is not None:
+        out_args = output_node.args[0]
+        if not isinstance(out_args, (tuple, list)):
+            out_args = (out_args,)
+        output_specs = ep.graph_signature.output_specs
+        assert len(output_specs) == len(out_args), (
+            f"reinplace: output spec count changed: "
+            f"{len(output_specs)} specs vs {len(out_args)} output args"
+        )
+        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
+
     return ep

exir/tests/test_reinplace_pass.py

@@ -326,6 +326,25 @@ def forward(
         # kernels in the other tests in this file.
         edge.to_executorch()
 
+    def test_output_specs_updated_after_reinplace(self) -> None:
+        """reinplace_pass must update graph_signature.output_specs to match
+        the renamed output nodes (e.g. relu -> relu_)."""
+
+        class M(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return torch.relu(x + 1.0)
+
+        ep = export(M(), (torch.randn(4),), strict=True)
+        edge_program = to_edge(ep).exported_program()
+
+        custom_set = {edge_ops.edge.aten.relu.default}
+        ep = reinplace_pass(edge_program, ops_to_inplace=custom_set)
+
+        output_node = next(n for n in ep.graph.nodes if n.op == "output")
+        actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")]
+        spec_names = [s.arg.name for s in ep.graph_signature.output_specs]
+        self.assertEqual(actual_names, spec_names)
+
     def test_broadcasting_self_not_reinplaced(self) -> None:
         """An op whose mutated arg (self) broadcasts up to a larger output
         must NOT be reinplaced: the in-place form cannot grow self