Correct fx node name in ODML torch lowering

PiperOrigin-RevId: 729584237

Author: yijie-yang

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import torch
 import re
+import torch
 
 
 def _class_fullname(cls):
@@ -35,7 +35,7 @@ def _get_hierarchy(node: torch.fx.Node):
   return hierachy_str
 
 
-def _get_canonical_filename(filename):
+def _get_canonical_filename(filename: str):
   """Remove unnecessary path prefix to make the filename more readable.
 
   This should be factored out so that pattern is a global option that a user
@@ -53,6 +53,29 @@ def _get_canonical_filename(filename):
   return filename
 
 
+def _get_canoical_nodename(node: torch.fx.Node) -> str:
+  """Get the canonical node name from the node's history."""
+
+  history = node.meta.get("from_node", [])
+
+  if len(history) > 1:  # Compatible with torch version under 2.6.0
+    return history[1][0]
+
+  if not hasattr(history[0], "name"):
+    return None
+  names = []
+  while history:
+    names.append(history[0].name)
+    history = history[0].from_node
+
+  # Based on the experiment, the third to last name in the history stack
+  # can be mapped to the original torch node name. The history stack is
+  # generated by tracing the node's transformation history during lowering.
+  if len(names) > 2:
+    return names[-3]
+  return None
+
+
 def build_mlir_file_debuginfo(node: torch.fx.Node):
   """Build the file and line info for the given node's lowerings in MLIR."""
 
@@ -66,16 +89,13 @@ def build_mlir_file_debuginfo(node: torch.fx.Node):
   return _get_canonical_filename(pt_trace.file), int(pt_trace.lineno)
 
 
-def build_nodename_debuginfo(node: torch.fx.Node):
+def build_nodename_debuginfo(node: torch.fx.Node) -> str:
   """Build the fx node name for the given node's lowerings in MLIR."""
-  history = node.meta.get("from_node", [])
-  if not history:
+
+  if not hasattr(node, "meta") or "from_node" not in node.meta:
     return None
-  if len(history) > 1:
-    return history[1][0]
-  if hasattr(history[0], "name"):  # torch 2.6.0+
-    return history[0].name
-  return None
+
+  return _get_canoical_nodename(node)
 
 
 def build_mlir_debuginfo(node: torch.fx.Node):

ai_edge_torch/odml_torch/test/test_tf_integration.py

@@ -27,6 +27,8 @@ def setUp(self):
     torch.manual_seed(0)
 
   def test_mlir_lowered_call(self):
+    """Test a simple model with MLIR lowered call."""
+
     class AddModel(torch.nn.Module):
 
       def forward(self, x, y):
@@ -45,6 +47,8 @@ def forward(self, x, y):
     self.assertTrue(np.allclose(lowering_output, torch_output))
 
   def test_resnet18(self):
+    """Test Resnet18 model with MLIR lowered call."""
+
     model = torchvision.models.resnet18().eval()
     forward_args = lambda: (torch.rand((1, 3, 224, 224)),)
 
@@ -58,11 +62,26 @@ def test_resnet18(self):
 
     # Check value and debug info.
     self.assertTrue(np.allclose(lowering_output, torch_output, atol=1e-5))
+
+  def test_debuginfo_from_export_lower(self):
+    """Test the debuginfo with export lower."""
+
+    model = torchvision.models.resnet18().eval()
+    forward_args = lambda: (torch.rand((1, 3, 224, 224)),)
+
+    ep = torch.export.export(model, forward_args())
+    lowered = odml_torch.export.exported_program_to_mlir(ep)
+
     lowered_text = lowered.get_text(enable_debug_info=True)
     # Check the file info.
     self.assertIn("torchvision/models/resnet.py", lowered_text)
-    # Check the fx node name.
-    self.assertIn("relu_1", lowered_text)
+    # Check the fx node names.
+    for n in ep.graph.nodes:
+      # Record all aten op nodes from the original graph and check if they
+      # are lowered to the same name in the lowered graph.
+      if n.op == "call_function" and not n.name.startswith("getitem"):
+        # Ensure strings like `loc("relu__1"` are present in the lowered text.
+        self.assertIn(f'loc("{n.name}"', lowered_text)
 
 
 if __name__ == "__main__":