fix delegate cache duplicate bug (pytorch#7281)

Summary:

Reported by pytorch#7175 that the delegate is not deduplicate when they're exactly the same

Differential Revision: D67067997

Author: cccclai

exir/_serialize/_program.py

@@ -224,6 +224,7 @@ def _extract_delegate_segments(
     """
     remaining_inline: List[BackendDelegateInlineData] = []
     inline_indices_seen: set[int] = set()
+    segment_index_map: dict[bytes, int] = {}
     for plan in program.execution_plan:
         for delegate in plan.delegates:
             if delegate.processed.location != DataLocation.INLINE:
@@ -249,8 +250,11 @@ def _extract_delegate_segments(
             inline_indices_seen.add(delegate.processed.index)
             if inline.data:
                 # Move the delegate data out of the program.
-                segment_index = len(segments)
-                segments.append(Cord(inline.data))
+                segment_index = segment_index_map.get(inline.data)
+                if segment_index is None:
+                    segment_index = len(segments)
+                    segments.append(Cord(inline.data))
+                    segment_index_map[inline.data] = segment_index
                 delegate.processed = BackendDelegateDataReference(
                     location=DataLocation.SEGMENT,
                     index=segment_index,

exir/backend/test/demos/rpc/TARGETS

@@ -28,6 +28,7 @@ runtime.python_library(
     ],
     visibility = [
         "//executorch/exir/backend/test/...",
+        "//executorch/exir/emit/test/...",
     ],
     deps = [
         ":executor_backend_preprocess",

exir/emit/_emitter.py

@@ -122,6 +122,8 @@ class _ProgramState:
     # Delegate data stored directly in the flatbuffer. Pointed to by BackendDelegateDataReference,
     # and should be copied to Program.backend_delegate_data.
     backend_delegate_data: List[BackendDelegateInlineData] = field(default_factory=list)
+    # Delegate cache that is used across all entry points. Key is the hash of the delegated payload.
+    backend_delegate_data_cache: Dict[str, int] = field(default_factory=dict)
 
     # Constants are optionally stored in external files.
     # Aggregate unique external constants into one buffer.
@@ -144,7 +146,8 @@ class _EmitterState:
     operators: List[Operator]
     delegates: List[BackendDelegate]
     operator_cache: Dict[Tuple[str, str], int]
-    delegate_cache: Dict[bytes, int]
+    # delegate_cache: (key: hash(delegated_payload), value: index in delegates)
+    delegate_cache: Dict[str, int]
     emit_stacktrace: bool
 
     spec2id_dict: Dict[TensorSpec, int] = field(default_factory=dict)
@@ -1073,8 +1076,8 @@ def _emit_delegate(
         """Emit the delegates inputs and outputs as specified by the schema, then emit the
         delegate's blob."""
         processed_bytes = lowered_module.processed_bytes
-
-        delegate_index = self.emitter_state.delegate_cache.get(processed_bytes)
+        hashed = hashlib.sha256(processed_bytes).hexdigest()
+        delegate_index = self.emitter_state.delegate_cache.get(hashed)
         delegate_ret = None
 
         if isinstance(self.node.meta["spec"], list):
@@ -1112,10 +1115,16 @@ def _emit_delegate(
         if delegate_index is None:
             # Allocate an entry for the data. TODO(T150113674): Reuse any duplicate entries if
             # present.
-            data_index: int = len(self.program_state.backend_delegate_data)
-            self.program_state.backend_delegate_data.append(
-                BackendDelegateInlineData(data=processed_bytes)
+            hashed = hashlib.sha256(processed_bytes).hexdigest()
+            data_index: Optional[int] = (
+                self.program_state.backend_delegate_data_cache.get(hashed)
             )
+            if data_index is None:
+                data_index = len(self.program_state.backend_delegate_data)
+                self.program_state.backend_delegate_data_cache[hashed] = data_index
+                self.program_state.backend_delegate_data.append(
+                    BackendDelegateInlineData(data=processed_bytes)
+                )
 
             backend_delegate = BackendDelegate(
                 id=lowered_module.backend_id,
@@ -1126,7 +1135,7 @@ def _emit_delegate(
             )
             delegate_index = len(self.emitter_state.delegate_cache)
             self.emitter_state.delegates.append(backend_delegate)
-            self.emitter_state.delegate_cache[processed_bytes] = delegate_index
+            self.emitter_state.delegate_cache[hashed] = delegate_index
 
         # TODO(angelayi) Will need to emit the kwargs too, in the correct order according to the
         # function's spec and with default arguments. This requires us to store the function's spec

exir/emit/test/TARGETS

@@ -16,6 +16,7 @@ python_unittest(
         "//executorch/exir:lib",
         "//executorch/exir:print_program",
         "//executorch/exir:schema",
+        "//executorch/exir/backend/test/demos/rpc:executor_backend_partitioner",
         "//executorch/exir/backend:backend_api",
         "//executorch/exir/emit:lib",
         "//executorch/exir/passes:const_prop_pass",

exir/emit/test/test_emit.py

@@ -27,6 +27,9 @@
 from executorch.exir._serialize._program import deserialize_pte_binary
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult
+from executorch.exir.backend.test.demos.rpc.executor_backend_partitioner import (
+    ExecutorBackendPartitioner,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.emit import emit_program  # noqa
 from executorch.exir.error import InternalError
@@ -63,7 +66,7 @@
 from functorch.experimental import control_flow
 from torch import nn
 
-from torch.export import Dim, export
+from torch.export import Dim, export, export_for_training
 
 
 class WrapperModule(torch.nn.Module):
@@ -1679,3 +1682,54 @@ def forward(self, x):
         ]
         self.assertEqual(external_map["linear.weight"], 0)
         self.assertEqual(external_map["linear.bias"], 1)
+
+    def test_delegate_deduplicate(self) -> None:
+        class SharedModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(2, 2)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class Module1(torch.nn.Module):
+            def __init__(self, shared_module):
+                super().__init__()
+                self.shared_module = shared_module
+
+            def forward(self, x):
+                return self.shared_module(x)
+
+        class Module2(torch.nn.Module):
+            def __init__(self, shared_module):
+                super().__init__()
+                self.shared_module = shared_module
+
+            def forward(self, x):
+                return self.shared_module(x)
+
+        shared_module = SharedModule()
+        module_1 = Module1(shared_module)
+        module_2 = Module2(shared_module)
+        example_inputs = (torch.randn(2, 2),)
+        module_1(*example_inputs)
+        module_2(*example_inputs)
+
+        ep1 = export_for_training(module_1, example_inputs)
+        ep2 = export_for_training(module_2, example_inputs)
+
+        edge_program_manager = exir.to_edge(
+            {"forward1": ep1, "forward2": ep2},
+            compile_config=exir.EdgeCompileConfig(
+                _check_ir_validity=False, _use_edge_ops=True
+            ),
+        )
+
+        edge_program_manager = edge_program_manager.to_backend(
+            ExecutorBackendPartitioner()
+        ).to_executorch()
+
+        # Check that there is only one delegate because two methods are exactly the same
+        self.assertEqual(
+            len(edge_program_manager.executorch_program.backend_delegate_data), 1
+        )