Refactor export_delegated_program (pytorch#10334)

Currently, we generate every combination of inputs for each module with the export_delegate_program script:
- extract_segments=True, False
- delegate_alignment=None,1024

Planning to add another flag, 'external_constants', which will move constants into a separate file to test program-data separation for delegated programs.

This test only requires pte, ptd, with default settings. So refactoring the export script to only generate based on the args, and update genrule to generate what the test requires.

Differential Revision: [D73278562](https://our.internmc.facebook.com/intern/diff/D73278562/)

Author: pytorchbot

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])",
             },
         )
 

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__":

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 = {