diff --git a/runtime/executor/test/targets.bzl b/runtime/executor/test/targets.bzl index 952a2433d8..75ea2674aa 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 4f4429aca8..44ae8df544 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 b8f3eeab0c..6538302c50 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 = {