diff --git a/backends/mlx/passes.py b/backends/mlx/passes.py index e88cd83f0bc..c22b8c0b15e 100644 --- a/backends/mlx/passes.py +++ b/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): diff --git a/backends/native/fat_pte.py b/backends/native/fat_pte.py new file mode 100644 index 00000000000..d2700fcd019 --- /dev/null +++ b/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(), + ) diff --git a/backends/native/partitioner.py b/backends/native/partitioner.py new file mode 100644 index 00000000000..65577dee3b1 --- /dev/null +++ b/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, + ) diff --git a/backends/native/preprocess.py b/backends/native/preprocess.py new file mode 100644 index 00000000000..627af8f0e2d --- /dev/null +++ b/backends/native/preprocess.py @@ -0,0 +1,226 @@ +# 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. + +""" +NativeBackend — AOT preprocess for the native portable runtime. + +Pipeline: SpecProp → reinplace → writeback → external constants → +memory planning → emit → serialize via _program_to_flatbuffer. +""" + +import json +from functools import partial +from typing import final, List, Tuple + +import torch + +from executorch.backends.native.fat_pte import build_fat_result +from executorch.backends.native.specializations import _SPECIALIZATION_REGISTRY +from executorch.exir._serialize._flatbuffer_program import _program_to_flatbuffer + +from executorch.exir.backend.backend_details import ( + BackendDetails, + CompileSpec, + ExportedProgram, + PreprocessResult, +) + +# A specialization recipe: callable that takes an ExportedProgram and returns +# serialized bytes. Each recipe owns the full +# to_edge_transform_and_lower → to_executorch → serialize flow. + + +from executorch.exir.dialects._ops import ops as _edge_ops +from executorch.exir.emit import emit_program +from executorch.exir.memory_planning import greedy, MemoryPlanningAlgorithmSuite +from executorch.exir.passes import MemoryPlanningPass, SpecPropPass +from executorch.exir.passes.insert_write_back_for_buffers_pass import ( + insert_write_back_for_buffers_pass, +) +from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass +from executorch.exir.program._program import _transform + +from torch._export.verifier import Verifier + +_edge = _edge_ops.edge.aten +BACKEND_INPLACE_OPS: frozenset = frozenset( + [ + _edge.relu.default, + _edge.gelu.default, + _edge.sigmoid.default, + _edge.index_put.default, + _edge.index_copy.default, + ] +) + + +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): + from typing import Callable + + return (Callable,) + + +def _apply_passes(program: ExportedProgram, passes) -> ExportedProgram: + from executorch.exir.pass_base import ExportPass, PassBase + + for p in passes: + if isinstance(p, MemoryPlanningPass) and hasattr(p, "run"): + p.run(program.graph_module) + elif issubclass(type(p), (ExportPass, PassBase)): + if hasattr(p, "_exported_program"): + p._exported_program = program + program = _transform(program, p, override_verifiers=[_AnyOp]) + if isinstance(p, SpecPropPass): + p.update_placeholder_tensor_specs(program, program.graph_module) + else: + program = p(program) + + return program + + +@final +class NativeBackend(BackendDetails): + @classmethod + def preprocess( + cls, + program: ExportedProgram, + module_compile_spec: List[CompileSpec], + ) -> PreprocessResult: + 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) + + for name in names: + if name not in _SPECIALIZATION_REGISTRY: + raise ValueError( + f"Specialization '{name}' is not registered. " + f"Registered: {sorted(_SPECIALIZATION_REGISTRY)}" + ) + + import copy + + native_result = cls._preprocess_native(copy.deepcopy(program)) + + results: List[Tuple[str, PreprocessResult]] = [ + ("NativeBackend", native_result), + ] + for name in names: + spec_program = copy.deepcopy(program) + result = _SPECIALIZATION_REGISTRY[name](spec_program) + results.append((name, result)) + + return build_fat_result(results) + + @classmethod + def _preprocess_native( + cls, + program: ExportedProgram, + ) -> PreprocessResult: + program = _apply_passes(program, [SpecPropPass()]) + + from executorch.exir.passes.reinplace import ( + reinplace_pass as _et_reinplace_pass, + ) + + program = _et_reinplace_pass(program, ops_to_inplace=BACKEND_INPLACE_OPS) + + # Re-run SpecPropPass: reinplace copies meta["val"] but not meta["spec"]. + program = _apply_passes(program, [SpecPropPass()]) + + from torch.export.graph_signature import InputKind, OutputKind + + has_buffer_mutation = any( + ospec.kind == OutputKind.BUFFER_MUTATION + for ospec in program.graph_signature.output_specs + ) + if has_buffer_mutation: + gm, new_sig = insert_write_back_for_buffers_pass(program) + program._graph_module = gm + program._graph_signature = new_sig + program = _apply_passes(program, [SpecPropPass()]) + + # Spec-share buffer placeholders with their mutation sources so the + # memory planner assigns them the same storage slot. + sig = program.graph_signature + nodes_by_name = {n.name: n for n in program.graph_module.graph.nodes} + buf_target_to_node = { + ispec.target: nodes_by_name.get(ispec.arg.name) + for ispec in sig.input_specs + if ispec.kind == InputKind.BUFFER and ispec.target + } + for ospec in sig.output_specs: + if ospec.kind != OutputKind.BUFFER_MUTATION: + continue + buf_node = buf_target_to_node.get(ospec.target) + wb_node = nodes_by_name.get(ospec.arg.name) + if ( + buf_node is None + or wb_node is None + or wb_node.op != "call_function" + or wb_node.target != torch.ops.aten.copy_.default + or len(wb_node.args) < 2 + ): + continue + src_node = wb_node.args[1] + if not hasattr(src_node, "meta"): + continue + buf_spec = buf_node.meta.get("spec") + if buf_spec is None or "spec" not in src_node.meta: + continue + src_node.meta["spec"] = buf_spec + + from executorch.exir.passes.external_constants_pass import ( + external_constants_pass, + ) + + external_constants_pass(program.graph_module) + + greedy_memory_planning = partial(greedy, allow_overlapping_allocations=True) + mem_planning_suite = MemoryPlanningAlgorithmSuite( + 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( + program, + [ + ConstraintBasedSymShapeEvalPass(), + MemoryPlanningPass(memory_planning_algo=mem_planning_suite), + ], + ) + + emitter_output = emit_program(program) + + from executorch.exir._serialize._named_data_store import NamedDataStore + + named_data_store = NamedDataStore() + if emitter_output.external_constant_buffer: + for _tag, fqn_to_idx in emitter_output.external_constant_map.items(): + for fqn, idx in fqn_to_idx.items(): + data = emitter_output.external_constant_buffer[idx] + named_data_store.add_named_data(fqn, data) + + # Use _program_to_flatbuffer (not serialize_pte_binary) to keep + # constants inline — the delegate can't see outer PTE segments. + fb_result = _program_to_flatbuffer(emitter_output.program) + serialized_bytes = bytes(fb_result.data) + + return PreprocessResult( + processed_bytes=serialized_bytes, + debug_handle_map=emitter_output.debug_handle_map, + data_store_output=named_data_store.get_named_data_store_output(), + ) diff --git a/backends/native/specializations.py b/backends/native/specializations.py new file mode 100644 index 00000000000..fa4499e5147 --- /dev/null +++ b/backends/native/specializations.py @@ -0,0 +1,60 @@ +# 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. + +""" +Specialization registry and recipes for the native backend fat PTE. + +Each recipe is a callable that takes an ExportedProgram and returns +a PreprocessResult. The processed_bytes contain the delegate payload +and data_store_output (if any) carries externalized constants that +build_fat_result merges into the shared named-data store, enabling +cross-specialization dedup of identical weight data. + +Automatically imported by NativePartitioner to populate the registry. +""" + +from typing import Callable, Dict + +from executorch.exir.backend.backend_details import PreprocessResult +from torch.export import ExportedProgram + +SpecializationRecipe = Callable[[ExportedProgram], PreprocessResult] + +_SPECIALIZATION_REGISTRY: Dict[str, SpecializationRecipe] = {} + + +def register_specialization(name: str, recipe: SpecializationRecipe) -> None: + _SPECIALIZATION_REGISTRY[name] = recipe + + +# --------------------------------------------------------------------------- +# Built-in recipes +# --------------------------------------------------------------------------- + + +def _mlx_recipe(ep: ExportedProgram) -> PreprocessResult: + from executorch.backends.mlx.partitioner import MLXPartitioner + from executorch.backends.mlx.passes import get_default_passes + from executorch.exir import to_edge_transform_and_lower + from executorch.exir.lowered_backend_module import get_lowered_backend_modules + + lowered = to_edge_transform_and_lower( + ep, + transform_passes=get_default_passes(), + partitioner=[MLXPartitioner()], + ) + lbms = get_lowered_backend_modules( + lowered.exported_program().graph_module, + ) + assert len(lbms) == 1, f"Expected 1 MLX delegate, got {len(lbms)}" + lbm = lbms[0] + return PreprocessResult( + processed_bytes=lbm.processed_bytes, + data_store_output=lbm.named_data_store_output, + ) + + +register_specialization("MLXBackend", _mlx_recipe) diff --git a/backends/native/test/__init__.py b/backends/native/test/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/native/test/test_fat_pte.py b/backends/native/test/test_fat_pte.py new file mode 100644 index 00000000000..283556115ba --- /dev/null +++ b/backends/native/test/test_fat_pte.py @@ -0,0 +1,133 @@ +# 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. + +import struct +import unittest + +from executorch.backends.native.fat_pte import ( + _ENTRY_FMT, + _ENTRY_SIZE, + build_fat_result, + FAT_MAGIC, + FAT_VERSION, + pack_fat_blob, +) +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir.backend.backend_details import PreprocessResult + + +class TestPackFatBlob(unittest.TestCase): + def test_single_specialization(self): + payload = b"hello" + blob = pack_fat_blob([("TestBackend", payload)]) + + magic, version, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(magic, FAT_MAGIC) + self.assertEqual(version, FAT_VERSION) + self.assertEqual(count, 1) + + header_size = 12 + bid, offset, size = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size) + self.assertEqual(bid.rstrip(b"\x00"), b"TestBackend") + self.assertEqual(offset, 0) + self.assertEqual(size, len(payload)) + + data_start = header_size + _ENTRY_SIZE + self.assertEqual(blob[data_start:], payload) + + def test_multiple_specializations(self): + entries = [("A", b"aaa"), ("B", b"bbbb"), ("C", b"cc")] + blob = pack_fat_blob(entries) + + _, _, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(count, 3) + + header_size = 12 + offsets_sizes = [] + for i in range(3): + bid, off, sz = struct.unpack_from( + "<" + _ENTRY_FMT, blob, header_size + i * _ENTRY_SIZE + ) + offsets_sizes.append((bid.rstrip(b"\x00"), off, sz)) + + self.assertEqual(offsets_sizes[0], (b"A", 0, 3)) + self.assertEqual(offsets_sizes[1], (b"B", 3, 4)) + self.assertEqual(offsets_sizes[2], (b"C", 7, 2)) + + data_start = header_size + 3 * _ENTRY_SIZE + self.assertEqual(blob[data_start:], b"aaabbbbcc") + + def test_empty_specializations(self): + blob = pack_fat_blob([]) + magic, version, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(count, 0) + self.assertEqual(len(blob), 12) + + 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): + def test_merges_results(self): + r1 = PreprocessResult(processed_bytes=b"native", debug_handle_map={1: [1]}) + r2 = PreprocessResult(processed_bytes=b"accel", debug_handle_map={2: [2]}) + + result = build_fat_result([("NativeBackend", r1), ("AccelBackend", r2)]) + + magic, version, count = struct.unpack_from("<4sII", result.processed_bytes, 0) + self.assertEqual(magic, FAT_MAGIC) + self.assertEqual(count, 2) + + # debug_handle_map defaults to first result's + self.assertEqual(result.debug_handle_map, {1: [1]}) + + def test_explicit_debug_handle_map(self): + r1 = PreprocessResult(processed_bytes=b"a", debug_handle_map={1: [1]}) + custom_map = {99: [99]} + result = build_fat_result([("X", r1)], debug_handle_map=custom_map) + self.assertEqual(result.debug_handle_map, custom_map) + + def test_single_result(self): + r = PreprocessResult(processed_bytes=b"only") + result = build_fat_result([("Solo", r)]) + _, _, count = struct.unpack_from("<4sII", result.processed_bytes, 0) + self.assertEqual(count, 1) + + def test_shared_constants_deduped(self): + """Two results with identical constant data share one buffer.""" + weight_data = b"\x01\x02\x03\x04" * 64 + + store_a = NamedDataStore() + store_a.add_named_data("model.weight", weight_data, alignment=16) + r_native = PreprocessResult( + processed_bytes=b"native", + data_store_output=store_a.get_named_data_store_output(), + ) + + store_b = NamedDataStore() + store_b.add_named_data("ab12/model.weight", weight_data, alignment=16) + r_accel = PreprocessResult( + processed_bytes=b"accel", + data_store_output=store_b.get_named_data_store_output(), + ) + + result = build_fat_result([("Native", r_native), ("Accel", r_accel)]) + merged = result.data_store_output + + self.assertIn("model.weight", merged.pte_data) + self.assertIn("ab12/model.weight", merged.pte_data) + # Different keys, but same buffer index (deduped by content). + self.assertEqual( + merged.pte_data["model.weight"].buffer_index, + merged.pte_data["ab12/model.weight"].buffer_index, + ) + self.assertEqual(len(merged.buffers), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/native/test/test_partitioner.py b/backends/native/test/test_partitioner.py new file mode 100644 index 00000000000..b045455d09f --- /dev/null +++ b/backends/native/test/test_partitioner.py @@ -0,0 +1,125 @@ +# 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. + +import unittest +from unittest.mock import MagicMock + +import torch +import torch.nn as nn + +from executorch.backends.native.partitioner import ( + _SUPPORTED_NON_CORE_OPS, + NativePartitioner, + NativeSupportedOperators, +) + + +class TestNativeSupportedOperators(unittest.TestCase): + def _make_node(self, op, target): + node = MagicMock() + node.op = op + node.target = target + return node + + def test_rejects_placeholder(self): + sup = NativeSupportedOperators() + node = self._make_node("placeholder", None) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_output(self): + sup = NativeSupportedOperators() + node = self._make_node("output", None) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_get_attr(self): + sup = NativeSupportedOperators() + node = self._make_node("get_attr", None) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_accepts_core_aten_op(self): + sup = NativeSupportedOperators() + node = self._make_node("call_function", torch.ops.aten.add.Tensor) + self.assertTrue(sup.is_node_supported({}, node)) + + def test_accepts_non_core_supported_op(self): + sup = NativeSupportedOperators() + for op in _SUPPORTED_NON_CORE_OPS: + node = self._make_node("call_function", op) + self.assertTrue( + sup.is_node_supported({}, node), + f"{op} should be supported as a non-core op", + ) + + def test_rejects_non_core_unsupported_op(self): + sup = NativeSupportedOperators() + op = torch.ops.aten.linalg_solve_triangular.default + if torch.Tag.core not in op.tags: + node = self._make_node("call_function", op) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_higher_order_operator(self): + sup = NativeSupportedOperators() + hop = torch.ops.higher_order.cond + node = self._make_node("call_function", hop) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_rejects_non_opoverload_callable(self): + sup = NativeSupportedOperators() + node = self._make_node("call_function", lambda x: x) + self.assertFalse(sup.is_node_supported({}, node)) + + def test_accepts_edge_op_overlay(self): + """EdgeOpOverload wraps OpOverload; partitioner should unwrap and accept.""" + from executorch.exir.dialects._ops import ops as _edge_ops + + sup = NativeSupportedOperators() + edge_add = _edge_ops.edge.aten.add.Tensor + node = self._make_node("call_function", edge_add) + self.assertTrue(sup.is_node_supported({}, node)) + + def test_rejects_non_core_edge_op(self): + from executorch.exir.dialects._ops import ops as _edge_ops + + sup = NativeSupportedOperators() + edge_op = _edge_ops.edge.aten.linalg_solve_triangular.default + if torch.Tag.core not in edge_op._op.tags: + node = self._make_node("call_function", edge_op) + self.assertFalse(sup.is_node_supported({}, node)) + + +class TestNativePartitionerE2E(unittest.TestCase): + def test_linear_delegates_all_ops(self): + from executorch.exir import to_edge_transform_and_lower + + model = nn.Linear(4, 4) + ep = torch.export.export(model, (torch.randn(1, 4),)) + lowered = to_edge_transform_and_lower(ep, partitioner=[NativePartitioner()]) + graph = lowered._edge_programs["forward"].graph + delegate_calls = [ + n + for n in graph.nodes + if n.op == "call_function" and "executorch_call_delegate" in str(n.target) + ] + non_delegate_ops = [ + n + for n in graph.nodes + if n.op == "call_function" + and "executorch_call_delegate" not in str(n.target) + and "getitem" not in str(n.target) + ] + self.assertGreater( + len(delegate_calls), 0, "Expected at least one delegate call" + ) + self.assertEqual( + len(non_delegate_ops), + 0, + f"All ops should be delegated, but found: " + f"{[str(n.target) for n in non_delegate_ops]}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/native/test/test_preprocess.py b/backends/native/test/test_preprocess.py new file mode 100644 index 00000000000..5995a02c792 --- /dev/null +++ b/backends/native/test/test_preprocess.py @@ -0,0 +1,170 @@ +# 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. + +import struct +import unittest + +import torch +import torch.nn as nn + +from executorch.backends.native.fat_pte import FAT_MAGIC +from executorch.backends.native.partitioner import NativePartitioner +from executorch.backends.native.specializations import ( + _SPECIALIZATION_REGISTRY, + register_specialization, +) +from executorch.exir import to_edge_transform_and_lower +from executorch.exir._serialize._flatbuffer_program import _flatbuffer_to_program +from executorch.exir.backend.backend_details import PreprocessResult + + +def _lower(model, example_inputs, specializations=None): + ep = torch.export.export(model, example_inputs) + edge = to_edge_transform_and_lower( + ep, + partitioner=[NativePartitioner(specializations=specializations)], + ) + return edge + + +def _get_delegate_blob(edge): + """Extract the single delegate's processed bytes from the lowered program.""" + et = edge.to_executorch() + delegates = et.executorch_program.backend_delegate_data + assert len(delegates) == 1, f"Expected 1 delegate blob, got {len(delegates)}" + return bytes(delegates[0].data) + + +class TestSpecializationRegistry(unittest.TestCase): + def setUp(self): + self._saved = dict(_SPECIALIZATION_REGISTRY) + + def tearDown(self): + _SPECIALIZATION_REGISTRY.clear() + _SPECIALIZATION_REGISTRY.update(self._saved) + + def test_register_and_lookup(self): + register_specialization( + "TestBackend", lambda ep: PreprocessResult(processed_bytes=b"test") + ) + self.assertIn("TestBackend", _SPECIALIZATION_REGISTRY) + + def test_unregistered_backend_rejected(self): + with self.assertRaises(ValueError) as ctx: + _lower( + nn.Linear(2, 2), + (torch.randn(1, 2),), + specializations=["UnregisteredBackend"], + ) + self.assertIn("UnregisteredBackend", str(ctx.exception)) + self.assertIn("not registered", str(ctx.exception)) + + def test_no_specializations_produces_plain_flatbuffer(self): + edge = _lower(nn.Linear(2, 2), (torch.randn(1, 2),)) + blob = _get_delegate_blob(edge) + self.assertNotEqual(blob[:4], FAT_MAGIC) + program = _flatbuffer_to_program(blob) + self.assertGreater(len(program.execution_plan[0].operators), 0) + + def test_registered_backend_produces_fat_pte(self): + register_specialization( + "TestAccel", lambda ep: PreprocessResult(processed_bytes=b"accel_payload") + ) + edge = _lower( + nn.Linear(2, 2), (torch.randn(1, 2),), specializations=["TestAccel"] + ) + blob = _get_delegate_blob(edge) + self.assertEqual(blob[:4], FAT_MAGIC) + + +class TestPreprocessSerialization(unittest.TestCase): + def _lower_and_deserialize(self, model, example_inputs): + edge = _lower(model, example_inputs) + blob = _get_delegate_blob(edge) + return _flatbuffer_to_program(blob) + + def test_linear_op_names(self): + program = self._lower_and_deserialize(nn.Linear(4, 4), (torch.randn(1, 4),)) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::linear", op_names) + + def test_add_op_name(self): + class AddModel(nn.Module): + def forward(self, x, y): + return x + y + + program = self._lower_and_deserialize( + AddModel(), (torch.randn(2, 3), torch.randn(2, 3)) + ) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::add", op_names) + + def test_reinplace_converts_relu(self): + class ReluModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(8, 8) + + def forward(self, x): + return torch.relu(self.linear(x)) + + program = self._lower_and_deserialize(ReluModel(), (torch.randn(1, 8),)) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::linear", op_names) + self.assertIn("aten::relu_", op_names) + + def test_fat_pte_contains_both_payloads(self): + saved = dict(_SPECIALIZATION_REGISTRY) + register_specialization( + "FakeAccel", lambda ep: PreprocessResult(processed_bytes=b"FAKE_ACCEL_DATA") + ) + try: + edge = _lower( + nn.Linear(2, 2), + (torch.randn(1, 2),), + specializations=["FakeAccel"], + ) + blob = _get_delegate_blob(edge) + magic, version, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(magic, FAT_MAGIC) + self.assertEqual(count, 2) + self.assertIn(b"FAKE_ACCEL_DATA", blob) + finally: + _SPECIALIZATION_REGISTRY.clear() + _SPECIALIZATION_REGISTRY.update(saved) + + def test_fat_pte_native_payload_is_valid(self): + """The native slice inside a fat PTE is a valid flatbuffer program.""" + saved = dict(_SPECIALIZATION_REGISTRY) + register_specialization( + "FakeAccel", lambda ep: PreprocessResult(processed_bytes=b"FAKE") + ) + try: + edge = _lower( + nn.Linear(2, 2), + (torch.randn(1, 2),), + specializations=["FakeAccel"], + ) + blob = _get_delegate_blob(edge) + _, _, count = struct.unpack_from("<4sII", blob, 0) + self.assertEqual(count, 2) + + from executorch.backends.native.fat_pte import _ENTRY_FMT, _ENTRY_SIZE + + header_size = 12 + bid, offset, size = struct.unpack_from("<" + _ENTRY_FMT, blob, header_size) + data_start = header_size + count * _ENTRY_SIZE + native_blob = blob[data_start + offset : data_start + offset + size] + program = _flatbuffer_to_program(native_blob) + op_names = [op.name for op in program.execution_plan[0].operators] + self.assertIn("aten::linear", op_names) + finally: + _SPECIALIZATION_REGISTRY.clear() + _SPECIALIZATION_REGISTRY.update(saved) + + +if __name__ == "__main__": + unittest.main() diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index 332878b5190..1d14a29b081 100644 --- a/exir/passes/reinplace.py +++ b/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 diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 6f1e07b7f32..d4244d1e4e8 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/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 diff --git a/pytest.ini b/pytest.ini index 05aea9d4da6..949d918a963 100644 --- a/pytest.ini +++ b/pytest.ini @@ -76,6 +76,7 @@ testpaths = # backends backends/apple/coreml/test + backends/native/test backends/test/harness/tests backends/test/suite/tests backends/transforms