Fix delegate debug handle generation (#6134)

Fix delegate debug handle generation (#5953) Summary: Pull Request resolved: #5953 Improve the logic used to generate a debug handle for delegate nodes. The existing logic is causing conflicts with other nodes debug handles. Reviewed By: Olivia-liu Differential Revision: D63779522 fbshipit-source-id: ff1583ee21bf05748c91a1c8b8e489996f80d7d9 (cherry picked from commit aad548c) Co-authored-by: Tarun Karuturi <[email protected]>
pytorch · Oct 11, 2024 · eca44f0 · eca44f0
1 parent edb43d8
commit eca44f0
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py
@@ -261,6 +261,15 @@ def _partition_and_lower_one_graph_module(
                     call_delegate_args.append(inp_node)
                     break
 
+        def generate_debug_handle(ep: ExportedProgram) -> int:
+            """
+            Generate a debug handle for the given ExportedProgram.
+            """
+            debug_handle = 0
+            for node in ep.graph_module.graph.nodes:
+                debug_handle = max(debug_handle, node.meta.get("debug_handle", 0))
+            return debug_handle + 1
+
         # Replace the partitioned submodule with a lowered submodule
         # Add call_method node with function "forward"
         with tagged_graph_module.graph.inserting_before(call_module_node):
@@ -273,8 +282,8 @@ def _partition_and_lower_one_graph_module(
                 (lowered_node,) + tuple(call_delegate_args),
                 call_module_node.kwargs,
             )
-            call_delegate_node.meta["debug_handle"] = len(
-                tagged_graph_module.graph.nodes
+            call_delegate_node.meta["debug_handle"] = generate_debug_handle(
+                owning_program
             )
             call_delegate_node.meta["val"] = submodule_output_node.meta["val"]
             call_module_node.replace_all_uses_with(call_delegate_node)