example (#64)

Summary:
Pull Request resolved: #64

Adding export example for XNNPACK delegated models, also adding to executor runner to run

Differential Revision: D48371417

fbshipit-source-id: 00c7e069dc91771d70eadb007319694b3fe25504

Author: mcr229

examples/executor_runner/targets.bzl

@@ -16,6 +16,7 @@ def define_common_targets():
         name = "executor_runner",
         srcs = ["executor_runner.cpp"],
         deps = [
+            "//executorch/backends/xnnpack:xnnpack_backend",
             "//executorch/runtime/executor/test:test_backend_compiler_lib",
             "//executorch/runtime/executor:program",
             "//executorch/extension/data_loader:file_data_loader",

examples/xnnpack/export_xnnpack_example.py

@@ -0,0 +1,106 @@
+# 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.
+
+# Example script for exporting simple models to flatbuffer
+
+import argparse
+
+import executorch.exir as exir
+import torch
+from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
+    XnnpackFloatingPointPartitioner,
+)
+from executorch.exir.backend.backend_api import to_backend, validation_disabled
+
+from ..models import MODEL_NAME_TO_MODEL
+
+
+def export_add_module_with_lower_graph():
+    class AddModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, x):
+            return x + x
+
+    capture_config = exir.CaptureConfig(pt2_mode=True, enable_dynamic_shape=False)
+    edge_compile_config = exir.EdgeCompileConfig()
+    sample_inputs = (torch.ones(1, 2, 3, 4),)
+    print("Running the example to export a composite module with lowered graph...")
+    edge = exir.capture(AddModule(), sample_inputs, capture_config).to_edge(
+        edge_compile_config
+    )
+    print("Exported graph:\n", edge.exported_program.graph)
+
+    # Lower AddMulModule to the demo backend
+    print("Lowering to the demo backend...")
+    edge.exported_program = to_backend(
+        edge.exported_program, XnnpackFloatingPointPartitioner
+    )
+
+    # The graph module is still runnerable
+    edge.exported_program.graph_module(*sample_inputs)
+
+    print("Lowered graph:\n", edge.exported_program.graph)
+
+    exec_prog = edge.to_executorch()
+    buffer = exec_prog.buffer
+
+    model_name = "xnnpack_add"
+    filename = f"{model_name}.ff"
+    print(f"Saving exported program to {filename}")
+    with open(filename, "wb") as file:
+        file.write(buffer)
+
+
+def export_mv2_with_lower_graph():
+    mv2, example_inputs = MODEL_NAME_TO_MODEL["mv2"]()
+    mv2 = mv2.eval()
+    capture_config = exir.CaptureConfig(pt2_mode=True, enable_dynamic_shape=False)
+    edge_compile_config = exir.EdgeCompileConfig(_check_ir_validity=False)
+
+    edge = exir.capture(mv2, example_inputs, capture_config).to_edge(
+        edge_compile_config
+    )
+    with validation_disabled():
+        edge.exported_program = to_backend(
+            edge.exported_program, XnnpackFloatingPointPartitioner
+        )
+
+    edge.exported_program.graph_module(*example_inputs)
+    print("Lowered graph: \n", edge.exported_program.graph)
+
+    exec_prog = edge.to_executorch()
+    buffer = exec_prog.buffer
+
+    model_name = "xnnpack_mv2"
+    filename = f"{model_name}.ff"
+    print(f"Saving exported program to {filename}")
+    with open(filename, "wb") as file:
+        file.write(buffer)
+
+
+OPTIONS_TO_LOWER = {
+    "add": export_add_module_with_lower_graph,
+    "mv2": export_mv2_with_lower_graph,
+}
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--option",
+        required=True,
+        choices=list(OPTIONS_TO_LOWER.keys()),
+        help=f"Provide the flow name. Valid ones: {list(OPTIONS_TO_LOWER.keys())}",
+    )
+
+    args = parser.parse_args()
+
+    # Choose one option
+    option = OPTIONS_TO_LOWER[args.option]
+
+    # Run the example flow
+    option()