[executorch][passes] Add config and pass to tag constants for external file

Pull Request resolved: #7193

- Add config 'external_constants' to ExecutorchBackendConfig.
- When set to True, run the 'external_constants_pass'
- This tags all constants as external, and moves them into a separate buffer to be serialized outside of the PTE file.

Note: users can write their own passes to tag weights to specific files / multiple files.

TODO: write example pass and test for the case where we have two constant files.
ghstack-source-id: 257353873

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

Author: lucylq

examples/portable/scripts/export.py

@@ -6,6 +6,8 @@
 
 # Example script for exporting simple models to flatbuffer
 
+# pyre-unsafe
+
 import argparse
 import logging
 
@@ -48,6 +50,15 @@ def main() -> None:
         required=False,
         help="specify segment alignment in hex. Default is 0x1000. Use 0x4000 for iOS",
     )
+
+    parser.add_argument(
+        "-e",
+        "--external_constants",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Save constants in external .ptd file. Default is False",
+    )
+
     parser.add_argument("-o", "--output_dir", default=".", help="output directory")
 
     args = parser.parse_args()
@@ -62,7 +73,7 @@ def main() -> None:
         *MODEL_NAME_TO_MODEL[args.model_name]
     )
 
-    backend_config = ExecutorchBackendConfig()
+    backend_config = ExecutorchBackendConfig(external_constants=args.external_constants)
     if args.segment_alignment is not None:
         backend_config.segment_alignment = int(args.segment_alignment, 16)
     if dynamic_shapes is not None:

exir/capture/_config.py

@@ -89,3 +89,7 @@ class ExecutorchBackendConfig:
     # If set to true, view_copy operations will be converted to lightweight
     # view operations in the ET runtime
     remove_view_copy: bool = True
+
+    # If set to true, all constant tensors will be stored in a separate file,
+    # external to the PTE file.
+    external_constants: bool = False

exir/emit/test/test_emit.py

@@ -1633,3 +1633,30 @@ def forward(self, x):
         input = torch.zeros(1)
         executorch_model(input)
         self.assertEqual(input, torch.ones(1))
+
+    def test_constant_tagged_tensors(self) -> None:
+        class LinearModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(5, 5)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        model = to_edge(export(LinearModule(), (torch.ones(5, 5),))).to_executorch(
+            config=ExecutorchBackendConfig(
+                external_constants=True,
+            )
+        )
+        emitter_output = model._emitter_output
+        # Check that constant_buffer is empty besides the non-constant placeholder 0.
+        self.assertEqual(len(emitter_output.program.constant_buffer), 1)
+        # Check that constant weights are in the external constant buffer.
+        self.assertEqual(len(emitter_output.external_constant_buffer), 2)
+        # Setting external_constants=True, saves all constants to the key
+        # '_default_external_constant'.
+        external_map = emitter_output.external_constant_map[
+            "_default_external_constant"
+        ]
+        self.assertEqual(external_map["linear.weight"], 0)
+        self.assertEqual(external_map["linear.bias"], 1)

exir/passes/TARGETS

@@ -10,6 +10,7 @@ python_library(
     deps = [
         ":const_prop_pass",
         ":debug_handle_generator_pass",
+        ":external_constants_pass",
         ":insert_write_back_for_buffers_pass",
         ":memory_format_ops_pass",
         ":memory_planning_pass",
@@ -54,6 +55,16 @@ python_library(
     ],
 )
 
+python_library(
+    name = "external_constants_pass",
+    srcs = [
+        "external_constants_pass.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+    ],
+)
+
 python_library(
     name = "insert_write_back_for_buffers_pass",
     srcs = [

exir/passes/external_constants_pass.py

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import torch
+from executorch.exir.tensor import TensorSpec
+from torch.export.exported_program import ExportedProgram
+
+
+def external_constants_pass(
+    ep: ExportedProgram,
+) -> ExportedProgram:
+    """
+    Move all constants to external file.
+    """
+    for module in ep.graph_module.modules():
+        if not isinstance(module, torch.fx.GraphModule):
+            continue
+
+        for node in module.graph.nodes:
+            if node.op == "placeholder":
+                spec = node.meta.get("spec")
+                if isinstance(spec, TensorSpec) and spec.const:
+                    node.meta["constant_tag"] = "_default_external_constant"
+    return ep

exir/program/_program.py

@@ -33,6 +33,7 @@
     MemoryFormatOpsPass,
     OpReplacePass,
 )
+from executorch.exir.passes.external_constants_pass import external_constants_pass
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
     insert_write_back_for_buffers_pass,
 )
@@ -1380,6 +1381,9 @@ def to_executorch(
                 )
             else:
                 new_gm_res = memory_planning_pass(new_gm)  # pyre-ignore[29]
+
+            if config.external_constants:
+                new_gm_res = external_constants_pass(new_gm_res)
             assert new_gm_res is not None
             new_gm = new_gm_res.graph_module