From e170b4ed20255d7ba368269b519c520eccdfb0a5 Mon Sep 17 00:00:00 2001 From: lucylq Date: Fri, 18 Apr 2025 13:52:27 -0700 Subject: [PATCH 1/2] Refactor export_delegated_program 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/) [ghstack-poisoned] --- runtime/executor/test/targets.bzl | 6 +-- test/models/export_delegated_program.py | 50 ++++++++++++++----------- test/models/targets.bzl | 30 +++++++-------- 3 files changed, 46 insertions(+), 40 deletions(-) diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 940601779cb..12b12052281 100644 --- a/runtime/executor/test/targets.bzl +++ b/runtime/executor/test/targets.bzl @@ -233,9 +233,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..139e6357b14 100644 --- a/test/models/export_delegated_program.py +++ b/test/models/export_delegated_program.py @@ -19,6 +19,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 +112,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 +136,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 +171,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 +200,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 +228,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 += f"-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 6d5b6753f3f..95ed0d4302f 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -150,7 +150,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", @@ -160,23 +160,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 = [ @@ -188,7 +188,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 = { From 3060e8dbaeded174e5dce8d18945340b4c76ba6d Mon Sep 17 00:00:00 2001 From: lucylq Date: Fri, 18 Apr 2025 13:59:06 -0700 Subject: [PATCH 2/2] Update on "Refactor export_delegated_program" 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/) [ghstack-poisoned]