From 3d512c99e287b99957288ef20c5bee6a3a9a644b Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Mon, 21 Apr 2025 11:52:45 -0700
Subject: [PATCH 1/5] Add pass to tag external constants for delegates

generate pte+ptd file for a delegated linear example

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

[ghstack-poisoned]
---
 backends/xnnpack/operators/node_visitor.py |  9 ++++++-
 backends/xnnpack/runtime/XNNCompiler.cpp   |  5 ++--
 exir/passes/external_constants_pass.py     | 29 ++++++++++++++++++++++
 test/models/export_delegated_program.py    | 26 +++++++++++++++++++
 test/models/targets.bzl                    | 19 ++++++++++++++
 5 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py
index 50781eade4d..d062c54eb72 100644
--- a/backends/xnnpack/operators/node_visitor.py
+++ b/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("delegate_constant_tag", None) is not None:
+            external_tag = tensor.meta["delegate_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
diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp
index c0204831c07..c72b0e806c8 100644
--- a/backends/xnnpack/runtime/XNNCompiler.cpp
+++ b/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 =
diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py
index e024fdcbcd2..d6941bcd38e 100644
--- a/exir/passes/external_constants_pass.py
+++ b/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["delegate_constant_tag"] = "_default_external_constant"
+                    mutated = True
+    return PassResult(gm, mutated)
diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py
index 25818a721a6..f23d12c2b1d 100644
--- a/test/models/export_delegated_program.py
+++ b/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__":
diff --git a/test/models/targets.bzl b/test/models/targets.bzl
index ab5fcc8a51d..d620a14f322 100644
--- a/test/models/targets.bzl
+++ b/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/...",
+        ],
+    )

From b12e5af73256ea963863c65e2e19658c2db34f91 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Mon, 21 Apr 2025 14:05:01 -0700
Subject: [PATCH 2/5] Update on "Add pass to tag external constants for
 delegates"

generate pte+ptd file for a delegated linear example

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

[ghstack-poisoned]
---
 backends/xnnpack/operators/node_visitor.py | 4 ++--
 exir/passes/external_constants_pass.py     | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py
index d062c54eb72..4c6cc327e12 100644
--- a/backends/xnnpack/operators/node_visitor.py
+++ b/backends/xnnpack/operators/node_visitor.py
@@ -594,8 +594,8 @@ def get_serialized_buffer_index(
         )
 
         external_tag = None
-        if tensor.meta.get("delegate_constant_tag", None) is not None:
-            external_tag = tensor.meta["delegate_constant_tag"]
+        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),
diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py
index d6941bcd38e..7da040e9066 100644
--- a/exir/passes/external_constants_pass.py
+++ b/exir/passes/external_constants_pass.py
@@ -100,6 +100,6 @@ def xnnpack_external_constants_pass(
             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["delegate_constant_tag"] = "_default_external_constant"
+                    node.meta["xnnpack_constant_tag"] = "_default_external_constant"
                     mutated = True
     return PassResult(gm, mutated)

From 5959eb91cff7a75595d25a540f267a7c2649628f Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Mon, 21 Apr 2025 14:53:47 -0700
Subject: [PATCH 3/5] Update on "Add pass to tag external constants for
 delegates"

generate pte+ptd file for a delegated linear example

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

[ghstack-poisoned]
---
 backends/xnnpack/operators/node_visitor.py |  4 +---
 exir/passes/external_constants_pass.py     | 14 ++++++--------
 test/models/export_delegated_program.py    | 13 ++++++-------
 test/models/targets.bzl                    |  2 +-
 4 files changed, 14 insertions(+), 19 deletions(-)

diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py
index 4c6cc327e12..01b2777bc61 100644
--- a/backends/xnnpack/operators/node_visitor.py
+++ b/backends/xnnpack/operators/node_visitor.py
@@ -593,9 +593,7 @@ def get_serialized_buffer_index(
             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"]
+        external_tag = tensor.meta.get("xnnpack_constant_tag", None)
         self._named_data_store.add_named_data(
             named_key,
             bytes(array),
diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py
index 7da040e9066..1eda6188a8f 100644
--- a/exir/passes/external_constants_pass.py
+++ b/exir/passes/external_constants_pass.py
@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from typing import List, Optional
+from typing import Callable, Optional
 
 import torch
 from executorch.exir.pass_base import PassResult
@@ -80,7 +80,7 @@ def external_mutable_weights_pass(
 
 def xnnpack_external_constants_pass(
     gm: GraphModule,
-    names: Optional[List[str]] = None,
+    filter_fn: Optional[Callable[[torch.fx.Node], str]] = None,
 ) -> PassResult:
     """
     Tag external constants before to_backend. Tagged constants will be saved
@@ -88,7 +88,7 @@ def xnnpack_external_constants_pass(
 
     Args:
         gm: GraphModule to tag.
-        names: List of constant names to tag. If None, tag all constants.
+        filter_fn: node -> str callable indicating the file (str) that a node should be saved to.
     Returns:
         PassResult: The resulting gm, and if it was mutated or not.
     """
@@ -97,9 +97,7 @@ def xnnpack_external_constants_pass(
         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
+            if node.op == "placeholder" and filter_fn is not None:
+                node.meta["xnnpack_constant_tag"] = filter_fn(node)
+                mutated = True
     return PassResult(gm, mutated)
diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py
index f23d12c2b1d..323f3aeaabf 100644
--- a/test/models/export_delegated_program.py
+++ b/test/models/export_delegated_program.py
@@ -166,7 +166,10 @@ def forward(self, *args, **kwargs):
 
         transform_passes = []
         if external_constants:
-            partial_function = partial(xnnpack_external_constants_pass, names=None)
+            partial_function = partial(
+                xnnpack_external_constants_pass,
+                filter_fn=lambda x: module_class.__name__,
+            )
             transform_passes.append(partial_function)
         executorch_program = to_edge_transform_and_lower(
             exported_program,
@@ -277,12 +280,8 @@ def main() -> None:
             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)
+            print(f"Saving external constants to {module_name}.ptd")
+            executorch_program.write_tensor_data_to_file(args.outdir)
 
 
 if __name__ == "__main__":
diff --git a/test/models/targets.bzl b/test/models/targets.bzl
index d620a14f322..aec3d1a865d 100644
--- a/test/models/targets.bzl
+++ b/test/models/targets.bzl
@@ -217,7 +217,7 @@ def define_common_targets():
         
         outs = {
             "ModuleLinear-e.pte": ["ModuleLinear-e.pte"],
-            "ModuleLinear-e.ptd": ["ModuleLinear-e.ptd"],
+            "ModuleLinear.ptd": ["ModuleLinear.ptd"],
         },
         default_outs = ["."],
         visibility = [

From 9018bd362ff4b942db017047a4adeb9af7f44ff8 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Mon, 21 Apr 2025 16:44:19 -0700
Subject: [PATCH 4/5] Update on "Add pass to tag external constants for
 delegates"

generate pte+ptd file for a delegated linear example

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

[ghstack-poisoned]
---
 backends/xnnpack/operators/node_visitor.py |  2 +-
 exir/passes/external_constants_pass.py     | 20 ++++++++++++++++----
 test/models/export_delegated_program.py    |  5 +++--
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py
index 01b2777bc61..95c5f315f5c 100644
--- a/backends/xnnpack/operators/node_visitor.py
+++ b/backends/xnnpack/operators/node_visitor.py
@@ -593,7 +593,7 @@ def get_serialized_buffer_index(
             ConstantDataOffset(offset=UINT64_MAX, size=size, named_key=named_key)
         )
 
-        external_tag = tensor.meta.get("xnnpack_constant_tag", None)
+        external_tag = tensor.meta.get("delegate_constant_tag", None)
         self._named_data_store.add_named_data(
             named_key,
             bytes(array),
diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py
index 1eda6188a8f..8aad572f41c 100644
--- a/exir/passes/external_constants_pass.py
+++ b/exir/passes/external_constants_pass.py
@@ -11,10 +11,20 @@
 import torch
 from executorch.exir.pass_base import PassResult
 from executorch.exir.tensor import TensorSpec
+
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
 from torch.export.exported_program import ExportedProgram, OutputKind
 from torch.fx import GraphModule
 
 
+def is_param_node(exp_prog: ExportedProgram, node: torch.fx.Node) -> bool:
+    return (
+        is_param(exp_prog, node)
+        or is_buffer(exp_prog, node)
+        or is_lifted_tensor_constant(exp_prog, node)
+    )
+
+
 def external_constants_pass(
     gm: GraphModule,
 ) -> PassResult:
@@ -78,8 +88,9 @@ def external_mutable_weights_pass(
     return PassResult(gm, mutated)
 
 
-def xnnpack_external_constants_pass(
+def delegate_external_constants_pass(
     gm: GraphModule,
+    ep: ExportedProgram,
     filter_fn: Optional[Callable[[torch.fx.Node], str]] = None,
 ) -> PassResult:
     """
@@ -97,7 +108,8 @@ def xnnpack_external_constants_pass(
         if not isinstance(module, torch.fx.GraphModule):
             continue
         for node in module.graph.nodes:
-            if node.op == "placeholder" and filter_fn is not None:
-                node.meta["xnnpack_constant_tag"] = filter_fn(node)
-                mutated = True
+            if node.op == "placeholder" and is_param_node(ep, node):
+                if filter_fn is not None:
+                    node.meta["delegate_constant_tag"] = filter_fn(node)
+                    mutated = True
     return PassResult(gm, mutated)
diff --git a/test/models/export_delegated_program.py b/test/models/export_delegated_program.py
index 323f3aeaabf..d5a780c9d62 100644
--- a/test/models/export_delegated_program.py
+++ b/test/models/export_delegated_program.py
@@ -24,7 +24,7 @@
     BackendWithCompilerDemo,
 )
 from executorch.exir.passes.external_constants_pass import (
-    xnnpack_external_constants_pass,
+    delegate_external_constants_pass,
 )
 from executorch.exir.program import ExecutorchProgramManager
 from torch import nn
@@ -167,7 +167,8 @@ def forward(self, *args, **kwargs):
         transform_passes = []
         if external_constants:
             partial_function = partial(
-                xnnpack_external_constants_pass,
+                delegate_external_constants_pass,
+                ep=exported_program,
                 filter_fn=lambda x: module_class.__name__,
             )
             transform_passes.append(partial_function)

From 233c6519f65d031ff2675fb7cb99fa2c6e6ae9f1 Mon Sep 17 00:00:00 2001
From: lucylq <lfq@meta.com>
Date: Mon, 21 Apr 2025 16:58:50 -0700
Subject: [PATCH 5/5] Update on "Add pass to tag external constants for
 delegates"

generate pte+ptd file for a delegated linear example

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

[ghstack-poisoned]
---
 exir/passes/external_constants_pass.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/exir/passes/external_constants_pass.py b/exir/passes/external_constants_pass.py
index 8aad572f41c..695de6f443d 100644
--- a/exir/passes/external_constants_pass.py
+++ b/exir/passes/external_constants_pass.py
@@ -97,8 +97,12 @@ def delegate_external_constants_pass(
     Tag external constants before to_backend. Tagged constants will be saved
     to an external file.
 
+    Note: this pass must be run after run_decompositions(), as tags on
+    constants are removed then.
+
     Args:
         gm: GraphModule to tag.
+        ep: ExportedProgram, to distinguish if a node is a constant.
         filter_fn: node -> str callable indicating the file (str) that a node should be saved to.
     Returns:
         PassResult: The resulting gm, and if it was mutated or not.