Update on "[et] generate debug handle before opeartor decomposition"

This diff update the debug handle generation, from each node in the edge program having a individual debug handle, to all nodes having a same ancestor in export graph sharing a same debug handle, which update the start point of tracing our node transformation from edge graph to exported graph.

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

[ghstack-poisoned]

Author: Gasoonjia

exir/tests/test_passes.py

@@ -67,7 +67,7 @@
 from executorch.exir.tensor import TensorSpec
 from executorch.exir.tests.common import register_additional_test_aten_ops
 from executorch.exir.tests.control_flow_models import FTCondDeadCode, FTMapBasic
-from executorch.exir.tests.models import MLP, Mul
+from executorch.exir.tests.models import FeedForwardBlock, MLP, Mul
 from functorch.experimental import control_flow
 
 from torch import nn
@@ -870,6 +870,69 @@ def test_debug_handle_generator_pass(self) -> None:
             if node.op != "placeholder" and node.op != "output":
                 self.assertIn("debug_handle", node.meta)
 
+    def test_debug_handle_generator_pass_generate_same_debug_handle_on_ops_sharing_same_source(
+        self,
+    ) -> None:
+        eager_model = FeedForwardBlock(256, 512)
+        inputs = (torch.randn(12, 256),)
+
+        graph_module = (
+            to_edge(export(eager_model, inputs, strict=True))
+            .exported_program()
+            .graph_module
+        )
+
+        same_source_nodes = {
+            "aten_native_layer_norm_default": (
+                "aten_native_layer_norm_default",
+                "getitem",
+            ),
+            "getitem": ("aten_native_layer_norm_default", "getitem"),
+            "aten_permute_copy_default": (
+                "aten_permute_copy_default",
+                "aten_addmm_default",
+            ),
+            "aten_addmm_default": ("aten_permute_copy_default", "aten_addmm_default"),
+            "aten_native_dropout_default": ("aten_native_dropout_default", "getitem_1"),
+            "getitem_1": ("aten_native_dropout_default", "getitem_1"),
+            "aten_relu_default": ("aten_relu_default",),
+            "aten_permute_copy_default_1": (
+                "aten_permute_copy_default_1",
+                "aten_addmm_default_1",
+            ),
+            "aten_addmm_default_1": (
+                "aten_permute_copy_default_1",
+                "aten_addmm_default_1",
+            ),
+            "aten_native_dropout_default_1": (
+                "aten_native_dropout_default_1",
+                "getitem_2",
+            ),
+            "getitem_2": ("aten_native_dropout_default_1", "getitem_2"),
+        }
+
+        node_name_to_debug_handle = {}
+
+        # Node having same source should have same debug handle
+        for node in graph_module.graph.nodes:
+            if node.op != "placeholder" and node.op != "output":
+                self.assertIn("debug_handle", node.meta)
+                if node.name in node_name_to_debug_handle:
+                    for node_name_with_same_debug_handle in same_source_nodes[
+                        node.name
+                    ]:
+                        self.assertEqual(
+                            node_name_to_debug_handle[node_name_with_same_debug_handle],
+                            node.meta["debug_handle"],
+                        )
+                else:
+                    for node_name_with_same_debug_handle in same_source_nodes[
+                        node.name
+                    ]:
+                        node_name_to_debug_handle[node_name_with_same_debug_handle] = (
+                            node.meta["debug_handle"]
+                        )
+
     def test_generate_missing_debug_handles(self) -> None:
         eager_model = MLP(2, output_size=4)
         inputs = eager_model.get_random_inputs()