From b79d88d91b6df42fc8b0cca2d3775f4a79888568 Mon Sep 17 00:00:00 2001 From: lucylq Date: Mon, 21 Apr 2025 11:29:33 -0700 Subject: [PATCH] Refactor export_delegated_program Pull Request resolved: https://github.com/pytorch/executorch/pull/10303 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. ghstack-source-id: 279342112 Differential Revision: [D73278562](https://our.internmc.facebook.com/intern/diff/D73278562/) --- runtime/executor/test/targets.bzl | 6 +-- test/models/export_delegated_program.py | 52 ++++++++++++++----------- test/models/targets.bzl | 30 +++++++------- 3 files changed, 48 insertions(+), 40 deletions(-) diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 952a2433d83..75ea2674aa7 100644 --- 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 index 4f4429aca88..44ae8df544f 100644 --- 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 index b8f3eeab0c1..6538302c507 100644 --- 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 = {