Refactor export_delegated_program

lucylq · web-flow · commit 6b49cb02b670 · 2025-04-21T14:03:59.000-07:00
Differential Revision: D73278562 Pull Request resolved: #10303
diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl
@@ -234,9 +234,9 @@ def define_common_targets(is_fbcode = False):
                 # Uses an fbcode target path because the authoring/export tools
                 # intentionally don't work in xplat (since they're host-only
                 # tools).
-                "ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul-nosegments-da1024.pte])",
-                "ET_MODULE_ADD_MUL_NOSEGMENTS_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul-nosegments.pte])",
-                "ET_MODULE_ADD_MUL_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul.pte])",
+                "ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul-nosegments-da1024.pte])",
+                "ET_MODULE_ADD_MUL_NOSEGMENTS_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul-nosegments.pte])",
+                "ET_MODULE_ADD_MUL_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul.pte])",
             },
         )
 
diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py
@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# pyre-unsafe
+
 import argparse
 import inspect
 import os
@@ -19,6 +21,7 @@
 from executorch.exir.backend.test.backend_with_compiler_demo import (
     BackendWithCompilerDemo,
 )
+from executorch.exir.program import ExecutorchProgramManager
 from torch import nn
 from torch.export import export
 
@@ -111,10 +114,10 @@ def export_module_to_program(
     *,
     backend_id: str,
     extract_delegate_segments: bool,
-    constant_tensor_alignemnt: Optional[int] = None,
+    constant_tensor_alignment: Optional[int] = None,
     delegate_alignment: Optional[int] = None,
     method: str = "forward",
-) -> bytes:
+) -> ExecutorchProgramManager:
     eager_module = module_class().eval()
     inputs = ()
     if hasattr(eager_module, "get_random_inputs"):
@@ -135,7 +138,7 @@ def forward(self, *args, **kwargs):
     edge_config = EdgeCompileConfig(_check_ir_validity=False)
     et_config = exir.ExecutorchBackendConfig(
         extract_delegate_segments=extract_delegate_segments,
-        constant_tensor_alignment=constant_tensor_alignemnt,
+        constant_tensor_alignment=constant_tensor_alignment,
         delegate_alignment=delegate_alignment,
     )
 
@@ -170,7 +173,7 @@ def forward(self, *args, **kwargs):
             export(composite_module, args=inputs, strict=True)
         ).to_executorch(config=et_config)
 
-    return executorch_program.buffer
+    return executorch_program
 
 
 def main() -> None:
@@ -199,6 +202,14 @@ def main() -> None:
         help="ID of the backend to use for delegation; "
         + f"one of {known_backend_ids}",
     )
+    parser.add_argument(
+        "--inline_delegate_segments",
+        action="store_true",
+        help="Store delegate data inside the flatbuffer.",
+    )
+    parser.add_argument(
+        "--delegate_alignment", type=int, default=None, help="Delegate alignment."
+    )
     parser.add_argument(
         "--outdir",
         type=str,
@@ -219,25 +230,22 @@ def main() -> None:
 
     # Export and write to the output files.
     os.makedirs(args.outdir, exist_ok=True)
+    suffix = ""
     for module_name, module_class in module_names_to_classes.items():
-        for extract_delegate_segments in (True, False):
-            suffix = "" if extract_delegate_segments else "-nosegments"
-            # Create files with the default alignment, and a large alignment.
-            # This alignment should be so large that it's extremely unlikely for
-            # the data to accidentally be aligned to it in the default case.
-            for delegate_alignment in (None, 1024):
-                suffix += f"-da{delegate_alignment}" if delegate_alignment else ""
-                outfile = os.path.join(args.outdir, f"{module_name}{suffix}.pte")
-                with open(outfile, "wb") as fp:
-                    fp.write(
-                        export_module_to_program(
-                            module_class,
-                            backend_id=args.backend_id,
-                            extract_delegate_segments=extract_delegate_segments,
-                            delegate_alignment=delegate_alignment,
-                        )
-                    )
-                print(f"Exported {module_name} and wrote program data to {outfile}")
+        if args.inline_delegate_segments:
+            suffix += "-nosegments"
+        if args.delegate_alignment is not None:
+            suffix += f"-da{args.delegate_alignment}"
+        outfile = os.path.join(args.outdir, f"{module_name}{suffix}.pte")
+        executorch_program = export_module_to_program(
+            module_class,
+            backend_id=args.backend_id,
+            extract_delegate_segments=not args.inline_delegate_segments,
+            delegate_alignment=args.delegate_alignment,
+        )
+        with open(outfile, "wb") as fp:
+            fp.write(executorch_program.buffer)
+        print(f"Exported {module_name} and wrote program data to {outfile}")
 
 
 if __name__ == "__main__":
diff --git a/test/models/targets.bzl b/test/models/targets.bzl
@@ -151,7 +151,7 @@ def define_common_targets():
         visibility = [],  # Private
     )
 
-    # Class names of nn.Modules for :exported_delegated_programs to export.
+    # Class names of nn.Modules available in export_delegated_program.py.
     DELEGATED_MODULES_TO_EXPORT = [
         "ModuleAddMul",
         "ModuleAddLarge",
@@ -161,23 +161,23 @@ def define_common_targets():
     # Name of the backend to use when exporting delegated programs.
     BACKEND_ID = "StubBackend"
 
-    # Generates Executorch .pte program files for various modules at build time.
+    # Generates Executorch .pte program files for the AddMul module at build time.
     # To use one, depend on a target like
-    # ":exported_delegated_programs[ModuleAdd.pte]" or
-    # ":exported_delegated_programs[ModuleAdd-nosegments.pte]" (which does not
+    # ":exported_delegated_add_mul[ModuleAdd.pte]" or
+    # ":exported_delegated_add_mul[ModuleAdd-nosegments.pte]" (which does not
     # extract the delegate data blobs into segments).
     runtime.genrule(
-        name = "exported_delegated_programs",
-        cmd = "$(exe :export_delegated_program)" +
-              " --modules " + ",".join(DELEGATED_MODULES_TO_EXPORT) +
-              " --backend_id " + BACKEND_ID +
-              " --outdir $OUT",
+        name = "exported_delegated_add_mul",
+        cmd = "$(exe :export_delegated_program) --modules ModuleAddMul --backend_id " + BACKEND_ID + " --outdir $OUT" +
+              " && $(exe :export_delegated_program) --modules ModuleAddMul --backend_id " + BACKEND_ID + " --inline_delegate_segments --outdir $OUT" +
+            # Create files with a large alignment as well as the default.
+            # This alignment should be so large that it's extremely unlikely for
+            # the data to accidentally be aligned to it in the default case.
+              " && $(exe :export_delegated_program) --modules ModuleAddMul --backend_id " + BACKEND_ID + " --inline_delegate_segments --delegate_alignment 1024 --outdir $OUT",
         outs = {
-            fname + seg_suffix + da_suffix + ".pte": [fname + seg_suffix + da_suffix + ".pte"]
-            for fname in DELEGATED_MODULES_TO_EXPORT
-            for seg_suffix in ["", "-nosegments"]
-            # "da" = delegate alignment
-            for da_suffix in ["", "-da1024"]
+            "ModuleAddMul.pte": ["ModuleAddMul.pte"],
+            "ModuleAddMul-nosegments.pte": ["ModuleAddMul-nosegments.pte"],
+            "ModuleAddMul-nosegments-da1024.pte": ["ModuleAddMul-nosegments-da1024.pte"],
         },
         default_outs = ["."],
         visibility = [
@@ -189,7 +189,7 @@ def define_common_targets():
     runtime.genrule(
         name = "exported_xnnp_delegated_programs",
         cmd = "$(exe :export_delegated_program)" +
-              " --modules " + ",".join(DELEGATED_MODULES_TO_EXPORT) +
+              " --modules ModuleAddLarge,ModuleSubLarge" +
               " --backend_id " + "XnnpackBackend" +
               " --outdir $OUT",
         outs = {
