diff --git a/backends/cpu/BUCK b/backends/cpu/BUCK index 18e9006b51a..514d852768b 100644 --- a/backends/cpu/BUCK +++ b/backends/cpu/BUCK @@ -48,6 +48,23 @@ fbcode_target( "//executorch/backends/xnnpack/_passes:xnnpack_passes", "//executorch/backends/xnnpack/recipes:xnnpack_recipes", "//executorch/export:lib", + "//executorch/exir/native:native", + "//pytorch/ao:torchao", + "//pytorch/vision:torchvision", # @manual + ], +) + +fbcode_target( + _kind = runtime.python_test, + name = "test_workflow_export", + srcs = ["test/test_workflow_export.py", "workflows/export.py"], + deps = [ + ":lib", + "//caffe2:torch", + "//executorch/backends/xnnpack/_passes:xnnpack_passes", + "//executorch/backends/xnnpack/recipes:xnnpack_recipes", + "//executorch/export:lib", + "//executorch/exir/native:native", "//pytorch/ao:torchao", "//pytorch/vision:torchvision", # @manual ], diff --git a/backends/cpu/test/test_workflow_export.py b/backends/cpu/test/test_workflow_export.py new file mode 100644 index 00000000000..c189e088219 --- /dev/null +++ b/backends/cpu/test/test_workflow_export.py @@ -0,0 +1,75 @@ +# 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 __future__ import annotations + +import json +import tempfile +import unittest +import zipfile +from pathlib import Path +from unittest import mock + +import torch +from executorch.backends.cpu.workflows.export import main +from executorch.backends.native.serialization import deserialize_graph + + +class WorkflowExportTest(unittest.TestCase): + def test_native_export_uses_native_preprocessing(self) -> None: + model = torch.nn.Sequential( + torch.nn.Conv2d(3, 4, 1), + torch.nn.GELU(), + torch.nn.Flatten(1), + torch.nn.Linear(16, 2), + ).eval() + captured = torch.export.export(model, (torch.ones(1, 3, 2, 2),)) + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) + torch.export.save(captured, path / "capture.pt2") + argv = [ + "export", + "--runtime", + "native", + "--pt2", + str(path / "capture.pt2"), + "--output", + str(path), + "--source-revision", + "test", + ] + with mock.patch("sys.argv", argv): + main() + with zipfile.ZipFile(path / "model.ptn") as package: + graph = deserialize_graph(package.read("program.ptg")) + self.assertIn("program.safetensors", package.namelist()) + targets = {node.target for node in graph.nodes} + self.assertIn("torch.ops.aten.gelu_.default", targets) + self.assertIn("torch.ops.aten.view.default", targets) + manifest = json.loads((path / "export.json").read_text()) + self.assertEqual("NativeBackend", manifest["preprocess"]) + self.assertEqual("native_default", manifest["recipe"]["recipe"]) + self.assertEqual([1, 3, 2, 2], manifest["input"]["shape"]) + self.assertIsNone(manifest["cpu_delegate_version"]) + self.assertFalse((path / "model.pte").exists()) + + def test_native_export_rejects_cpu_recipe(self) -> None: + with tempfile.TemporaryDirectory() as directory: + argv = [ + "export", + "--runtime", + "native", + "--backend", + "xnnpack", + "--output", + directory, + "--source-revision", + "test", + ] + with mock.patch("sys.argv", argv): + with self.assertRaises(SystemExit) as error: + main() + self.assertEqual(2, error.exception.code) diff --git a/backends/cpu/workflows/README.md b/backends/cpu/workflows/README.md index e16b24679c1..de3f3d00f61 100644 --- a/backends/cpu/workflows/README.md +++ b/backends/cpu/workflows/README.md @@ -1,6 +1,6 @@ # CPU torchvision exports -List classifiers or export a PTE with the registered CPU or XNNPACK FP32 recipe: +List classifiers or export a model with the registered CPU or XNNPACK FP32 recipe: ```sh python -m executorch.backends.cpu.workflows.export --list-models @@ -19,9 +19,15 @@ for random weights; `--image-size` overrides it. For random-weight ViTs, this also configures the model's input size. Pretrained ViTs reject sizes that do not match their weights. Reused captures keep their recorded input shape, and the model name defaults to the capture's file stem. -Exports include `model.pte`, the captured `.pt2`, `input.bin`, `eager.bin`, -`golden.bin`, and an `export.json` manifest with artifact hashes and preprocessing -provenance. `input.bin` is the raw float32 input; `eager.bin` and `golden.bin` -hold the eager and edge-program logits. A runner can feed `input.bin` to -`model.pte` and compare its output against `eager.bin` to validate the PTE before -measuring latency. + +Exports default to `--runtime et` (`model.pte`), using the registered FP32 +`--backend cpu|xnnpack` recipe. `--runtime native --backend cpu` uses `to_native()` +with standard Native preprocessing to produce `model.ptn`. Reuse `--pt2` to compare +the same captured model through the two pipelines. + +Both formats include the captured `.pt2`, `input.bin`, complete eager logits in +`eager.bin`, `golden.bin`, and an `export.json` manifest with artifact hashes, +preprocessing provenance, and runtime/artifact identity. `golden.bin` holds the +edge-program logits for ET exports and matches `eager.bin` for Native exports. A +runner can feed `input.bin` to either artifact and compare its output against +`eager.bin` to validate it before measuring latency. diff --git a/backends/cpu/workflows/export.py b/backends/cpu/workflows/export.py index 705572b34fb..3e220141d66 100644 --- a/backends/cpu/workflows/export.py +++ b/backends/cpu/workflows/export.py @@ -15,15 +15,18 @@ from collections import Counter from pathlib import Path from typing import Any +from zipfile import ZipFile import torch import torchvision from executorch.backends.cpu.partitioner import CPU_DELEGATE_VERSION from executorch.backends.cpu.recipes import CPURecipeType +from executorch.backends.native.passes import get_default_passes from executorch.backends.native.serialization import deserialize_graph from executorch.backends.native.serialization.schema import OpKind from executorch.backends.xnnpack._passes import XNNPACKPassManager from executorch.backends.xnnpack.recipes import XNNPackRecipeType +from executorch.exir.native import to_native from executorch.export import export, ExportRecipe, ExportSession from torchvision.models import get_model, get_model_weights, list_models @@ -91,14 +94,20 @@ def write_manifest( args: argparse.Namespace, spec: dict[str, Any], exported: torch.export.ExportedProgram, - recipe: ExportRecipe, - session: ExportSession, + recipe: ExportRecipe | None, + session: ExportSession | None, image: torch.Tensor, ) -> None: + is_native = args.runtime == "native" is_cpu = spec["backend"] == "cpu" - passes = [] if is_cpu else XNNPACKPassManager(exported).passes + if is_native: + passes = [type(p) for p in get_default_passes()] + preprocess = "NativeBackend" + else: + passes = [] if is_cpu else XNNPACKPassManager(exported).passes + preprocess = "CpuBackend" if is_cpu else "XnnpackBackend" files = {} - artifact = "model.pte" + artifact = "model.ptn" if is_native else "model.pte" for name in (f"{args.model}.pt2", "input.bin", "eager.bin", "golden.bin", artifact): path = args.output / name with path.open("rb") as handle: @@ -106,24 +115,35 @@ def write_manifest( for chunk in iter(lambda: handle.read(8 << 20), b""): digest.update(chunk) files[name] = {"sha256": digest.hexdigest(), "bytes": path.stat().st_size} - program = session.get_executorch_program() - semantic_nodes = None - if is_cpu: + if session is None: + with ZipFile(args.output / artifact) as package: + graph = deserialize_graph(package.read("program.ptg")) semantic_nodes = [ - node.name - for delegate in program.execution_plan[0].delegates - if delegate.id == "CpuBackend" - for node in deserialize_graph( - bytes(program.backend_delegate_data[delegate.processed.index].data) - ).nodes - if node.op_kind == OpKind.CALL_FUNCTION + node.name for node in graph.nodes if node.op_kind == OpKind.CALL_FUNCTION ] + stages = ["torch.export", "to_native"] + delegates = {} + else: + program = session.get_executorch_program() + semantic_nodes = None + if is_cpu: + semantic_nodes = [ + node.name + for delegate in program.execution_plan[0].delegates + if delegate.id == "CpuBackend" + for node in deserialize_graph( + bytes(program.backend_delegate_data[delegate.processed.index].data) + ).nodes + if node.op_kind == OpKind.CALL_FUNCTION + ] + stages = [stage.name for stage in session.get_stage_artifacts()] + delegates = dict(Counter(d.id for d in program.execution_plan[0].delegates)) manifest = { "version": 1, "model": args.model, - "runtime": "et", + "runtime": args.runtime, "artifact": artifact, - "preprocess": "CpuBackend" if is_cpu else "XnnpackBackend", + "preprocess": preprocess, "source_revision": args.source_revision, "versions": { "torch": torch.__version__, @@ -131,18 +151,22 @@ def write_manifest( }, "recipe": spec, "quantization": None, - "cpu_delegate_version": CPU_DELEGATE_VERSION if is_cpu else None, + "cpu_delegate_version": ( + CPU_DELEGATE_VERSION if is_cpu and not is_native else None + ), "recipe_passes": {"aten": [], "edge": [], "edge_manager": []}, "backend_default_passes": [f"{p.__module__}.{p.__qualname__}" for p in passes], "preserved_ops": sorted( { str(op) - for partitioner in recipe.lowering_recipe.partitioners + for partitioner in ( + recipe.lowering_recipe.partitioners if recipe is not None else [] + ) for op in partitioner.ops_to_not_decompose(exported)[0] } ), - "stages": [stage.name for stage in session.get_stage_artifacts()], - "delegates": dict(Counter(d.id for d in program.execution_plan[0].delegates)), + "stages": stages, + "delegates": delegates, "semantic_nodes": semantic_nodes, "input": { "shape": list(image.shape), @@ -159,6 +183,7 @@ def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--output", type=Path) parser.add_argument("--backend", choices=tuple(RECIPE_TYPES), default="cpu") + parser.add_argument("--runtime", choices=("et", "native"), default="et") parser.add_argument("--pt2", type=Path) parser.add_argument("--model") parser.add_argument("--image-size", type=int, default=None) @@ -187,16 +212,30 @@ def main() -> None: parser.error(f"Unknown torchvision model {args.model!r}; see --list-models") if args.image_size is not None and args.image_size <= 0: parser.error("--image-size must be positive") - recipe_type = RECIPE_TYPES[args.backend] + if args.runtime == "native" and args.backend != "cpu": + parser.error( + "Native CPU export uses to_native defaults; CPU recipes are ET-only" + ) spec = { "version": 1, "backend": args.backend, - "recipe": recipe_type.value, + "recipe": ( + "native_default" + if args.runtime == "native" + else RECIPE_TYPES[args.backend].value + ), "kwargs": {}, } args.output.mkdir(parents=True, exist_ok=True) exported, image = capture(args) - recipe = ExportRecipe.get_recipe(recipe_type) + if args.runtime == "native": + artifact = args.output / "model.ptn" + to_native(exported).save(str(artifact)) + shutil.copyfile(args.output / "eager.bin", args.output / "golden.bin") + write_manifest(args, spec, exported, None, None, image) + logger.info("Exported Native program to %s", artifact) + return + recipe = ExportRecipe.get_recipe(RECIPE_TYPES[args.backend]) session = export(exported, export_recipe=recipe) edge_module = session.get_edge_program_manager().exported_program().module() with torch.no_grad():