diff --git a/test/end2end/exported_module.py b/test/end2end/exported_module.py index 67a03b59a8..d3dcc22910 100644 --- a/test/end2end/exported_module.py +++ b/test/end2end/exported_module.py @@ -69,6 +69,7 @@ def export( capture_config=None, skip_type_promotion: bool = False, export_joint_graph: bool = False, + external_constants: bool = False, ) -> "ExportedModule": """ Creates a new ExportedModule for the specified module class. @@ -206,6 +207,7 @@ def __init__(self, method): dynamic_memory_planning_mode=dynamic_memory_planning_mode, memory_planning_pass=memory_planning_pass, to_out_var_pass=ToOutVarPass(ignore_to_out_var_failure), + external_constants=external_constants, ) ) diff --git a/test/models/export_program.py b/test/models/export_program.py index caea394f33..7d5d560683 100644 --- a/test/models/export_program.py +++ b/test/models/export_program.py @@ -15,6 +15,7 @@ import torch from executorch.exir import CaptureConfig from executorch.exir.passes import MemoryPlanningPass +from executorch.exir.program._program import ExecutorchProgramManager from torch import nn from torch.export import Dim @@ -190,7 +191,8 @@ def export_joint(): def export_module_to_program( module_class: Type[nn.Module], skip_type_promotion: bool, -): + external_constants: bool = False, +) -> ExecutorchProgramManager: """Exports the module and returns the serialized program data.""" torch.manual_seed(0) # Look for an optional @staticmethod that defines custom trace params. @@ -211,9 +213,10 @@ def export_module_to_program( methods, skip_type_promotion=skip_type_promotion, export_joint_graph=export_joint, + external_constants=external_constants, **export_kwargs, ) - return module.executorch_program.buffer + return module.executorch_program def main() -> None: @@ -235,7 +238,12 @@ def main() -> None: "--outdir", type=str, required=True, - help="Path to the directory to write .pte files to", + help="Path to the directory to write .pte files and .ptd files to", + ) + parser.add_argument( + "--external-constants", + action="store_true", + help="Export the model with external constants", ) args = parser.parse_args() @@ -257,14 +265,16 @@ def main() -> None: # Type promotion will convert to fp32. skip_type_promotion = True outfile = os.path.join(args.outdir, f"{module_name}.pte") + prog = export_module_to_program( + module_class, + skip_type_promotion=skip_type_promotion, + external_constants=args.external_constants, + ) with open(outfile, "wb") as fp: - fp.write( - export_module_to_program( - module_class, - skip_type_promotion=skip_type_promotion, - ) - ) - print(f"Exported {module_name} and wrote program data to {outfile}") + prog.write_to_file(fp) + print(f"Exported {module_name} and wrote program data to {outfile}") + + prog.write_tensor_data_to_file(args.outdir) if __name__ == "__main__": diff --git a/test/models/targets.bzl b/test/models/targets.bzl index f291a17c62..e2cd0f264b 100644 --- a/test/models/targets.bzl +++ b/test/models/targets.bzl @@ -75,9 +75,28 @@ def define_common_targets(): name = "exported_programs", cmd = "$(exe :export_program) --modules " + ",".join(MODULES_TO_EXPORT) + " --outdir $OUT", outs = { - fname + seg_suffix + ".pte": [fname + seg_suffix + ".pte"] + fname + ".pte": [fname + ".pte"] for fname in MODULES_TO_EXPORT - for seg_suffix in ["", "-no-constant-segment"] + }, + default_outs = ["."], + visibility = [ + "//executorch/...", + # This genrule can't run in xplat since it uses EXIR, so make its + # output visible to xplat tests. This is an exceptional case, and + # typically shouldn't be done. + "fbsource//xplat/executorch/...", + ], + # Allow the xplat entry in the visibility list. This is an exceptional + # case, and typically shouldn't be done. + _is_external_target = True, + ) + + runtime.genrule( + name = "exported_programs_with_data_separated", + cmd = "$(exe :export_program) --modules ModuleLinear --external-constants --outdir $OUT", + outs = { + "ModuleLinear.pte": ["ModuleLinear.pte"], + "ModuleLinear.ptd": ["_default_external_constant.ptd"], }, default_outs = ["."], visibility = [