[executorch] Add program-data separation example to test files (pytorch#7887)

Pull Request resolved: pytorch#7765

Add to test env to use for testing in next diff.
ghstack-source-id: 262611426
@exported-using-ghexport

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

Author: lucylq

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,
             )
         )
 

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 <classname>.pte files to",
+        help="Path to the directory to write <classname>.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__":

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