Add pass to tag external constants for delegates

Pull Request resolved: #10328

generate pte+ptd file for a delegated linear example
ghstack-source-id: 279382617

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

Author: lucylq

backends/xnnpack/operators/node_visitor.py

@@ -592,8 +592,15 @@ def get_serialized_buffer_index(
         xnn_graph.constant_data.append(
             ConstantDataOffset(offset=UINT64_MAX, size=size, named_key=named_key)
         )
+
+        external_tag = None
+        if tensor.meta.get("xnnpack_constant_tag", None) is not None:
+            external_tag = tensor.meta["xnnpack_constant_tag"]
         self._named_data_store.add_named_data(
-            named_key, bytes(array), alignment=CONSTANT_TENSOR_ALIGNMENT
+            named_key,
+            bytes(array),
+            alignment=CONSTANT_TENSOR_ALIGNMENT,
+            external_tag=external_tag,
         )
 
         return buffer_idx

backends/xnnpack/runtime/XNNCompiler.cpp

@@ -204,8 +204,9 @@ const uint8_t* getConstantDataPtr(
         if (!buffer.ok()) {
           ET_LOG(
               Error,
-              "Failed to get constant data for key %s",
-              data_name.c_str());
+              "Failed to get constant data for key %s from named_data_map. Error code: %u",
+              data_name.c_str(),
+              static_cast<uint32_t>(buffer.error()));
           return nullptr;
         }
         const uint8_t* data_ptr =

exir/passes/external_constants_pass.py

@@ -6,6 +6,8 @@
 
 # pyre-strict
 
+from typing import List, Optional
+
 import torch
 from executorch.exir.pass_base import PassResult
 from executorch.exir.tensor import TensorSpec
@@ -74,3 +76,30 @@ def external_mutable_weights_pass(
                     node.meta["constant_tag"] = "_default_external_constant"
                     mutated = True
     return PassResult(gm, mutated)
+
+
+def xnnpack_external_constants_pass(
+    gm: GraphModule,
+    names: Optional[List[str]] = None,
+) -> PassResult:
+    """
+    Tag external constants before to_backend. Tagged constants will be saved
+    to an external file.
+
+    Args:
+        gm: GraphModule to tag.
+        names: List of constant names to tag. If None, tag all constants.
+    Returns:
+        PassResult: The resulting gm, and if it was mutated or not.
+    """
+    mutated = False
+    for module in gm.modules():
+        if not isinstance(module, torch.fx.GraphModule):
+            continue
+        for node in module.graph.nodes:
+            if node.op == "placeholder":
+                # Move specified constants to external file. If none, move all constants.
+                if names is None or node.name in names:
+                    node.meta["xnnpack_constant_tag"] = "_default_external_constant"
+                    mutated = True
+    return PassResult(gm, mutated)

test/models/export_delegated_program.py

@@ -10,6 +10,8 @@
 import inspect
 import os
 import sys
+
+from functools import partial
 from typing import Dict, final, Optional, Sequence, Type
 
 import executorch.exir as exir
@@ -21,6 +23,9 @@
 from executorch.exir.backend.test.backend_with_compiler_demo import (
     BackendWithCompilerDemo,
 )
+from executorch.exir.passes.external_constants_pass import (
+    xnnpack_external_constants_pass,
+)
 from executorch.exir.program import ExecutorchProgramManager
 from torch import nn
 from torch.export import export
@@ -129,6 +134,7 @@ def export_module_to_program(
     constant_tensor_alignment: Optional[int] = None,
     delegate_alignment: Optional[int] = None,
     method_name: str = "forward",
+    external_constants: bool = False,
 ) -> ExecutorchProgramManager:
     eager_module = module_class().eval()
     inputs = ()
@@ -158,8 +164,13 @@ def forward(self, *args, **kwargs):
             XnnpackPartitioner,
         )
 
+        transform_passes = []
+        if external_constants:
+            partial_function = partial(xnnpack_external_constants_pass, names=None)
+            transform_passes.append(partial_function)
         executorch_program = to_edge_transform_and_lower(
             exported_program,
+            transform_passes=transform_passes,
             compile_config=edge_config,
             partitioner=[XnnpackPartitioner()],
         ).to_executorch(config=et_config)
@@ -221,6 +232,11 @@ def main() -> None:
     parser.add_argument(
         "--delegate_alignment", type=int, default=None, help="Delegate alignment."
     )
+    parser.add_argument(
+        "--external_constants",
+        action="store_true",
+        help="Export the model with all constants saved to an external file.",
+    )
     parser.add_argument(
         "--outdir",
         type=str,
@@ -247,16 +263,26 @@ def main() -> None:
             suffix += "-nosegments"
         if args.delegate_alignment is not None:
             suffix += f"-da{args.delegate_alignment}"
+        if args.external_constants:
+            suffix += f"-e"
         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,
+            external_constants=args.external_constants,
         )
         with open(outfile, "wb") as fp:
             fp.write(executorch_program.buffer)
         print(f"Exported {module_name} and wrote program data to {outfile}")
+        if args.external_constants:
+            # current infra doesnt easily allow renaming this file, so just hackily do it here.
+            executorch_program._tensor_data[f"{module_name}{suffix}"] = (
+                executorch_program._tensor_data.pop("_default_external_constant")
+            )
+            print(f"Saving external constants to {module_name}{suffix}.ptd")
+        executorch_program.write_tensor_data_to_file(args.outdir)
 
 
 if __name__ == "__main__":

test/models/targets.bzl

@@ -206,3 +206,22 @@ def define_common_targets():
         ],
         env = {"PYTORCH_DISABLE_JUSTKNOBS": "1",},
     )
+
+    runtime.genrule(
+        name = "exported_program_data",
+        cmd = "$(exe :export_delegated_program)" +
+            " --modules ModuleLinear" + 
+            " --backend_id XnnpackBackend" +
+            " --external_constants" +
+            " --outdir $OUT",
+        
+        outs = {
+            "ModuleLinear-e.pte": ["ModuleLinear-e.pte"],
+            "ModuleLinear-e.ptd": ["ModuleLinear-e.ptd"],
+        },
+        default_outs = ["."],
+        visibility = [
+            "//executorch/runtime/executor/test/...",
+            "//executorch/test/...",
+        ],
+    )