Add fx node name info to ODML Torch lowerings.

PiperOrigin-RevId: 725602099

Author: yijie-yang

ai_edge_torch/odml_torch/debuginfo/__init__.py

@@ -12,5 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ._build import build_mlir_debuginfo, build_mlir_file_debuginfo
+"""Debug info generation for ODML Torch."""
+
+from . import _build
 from ._op_polyfill import write_mlir_debuginfo_op
+
+build_nodename_debuginfo = _build.build_nodename_debuginfo
+build_mlir_file_debuginfo = _build.build_mlir_file_debuginfo
+build_mlir_debuginfo = _build.build_mlir_debuginfo

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -40,10 +40,16 @@ def _get_canonical_filename(filename):
 
   This should be factored out so that pattern is a global option that a user
   can override.
+
+  Args:
+    filename: The filename to canonicalize.
+
+  Returns:
+    The canonicalized filename.
   """
 
-  # TODO: We should add a config option to provide a regex to strip from the
-  # debug info. Currently absolute path is used.
+  # TODO(yijieyang): We should add a config option to provide a regex to strip
+  # from the debug info. Currently absolute path is used.
   return filename
 
 
@@ -55,9 +61,23 @@ def build_mlir_file_debuginfo(node: torch.fx.Node):
 
   # Note: This uses internal APIs and may break in the future.
   pt_trace = torch.fx.graph._parse_stack_trace(node.stack_trace)
+  if pt_trace is None:
+    return None, None
   return _get_canonical_filename(pt_trace.file), int(pt_trace.lineno)
 
 
+def build_nodename_debuginfo(node: torch.fx.Node):
+  """Build the fx node name for the given node's lowerings in MLIR."""
+  history = node.meta.get("from_node", [])
+  if not history:
+    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
+
+
 def build_mlir_debuginfo(node: torch.fx.Node):
   """Build the debuginfo string for the given node's lowerings in MLIR."""
 

ai_edge_torch/odml_torch/export.py

@@ -88,6 +88,21 @@ def __init__(self, module: torch.fx.GraphModule, lctx: LoweringContext):
     self.outputs = None
 
   def _build_loc(self, node: torch.fx.Node):
+    """Build MLIR location for the given node.
+
+    The location contains:
+    - layer info
+    - fx node name
+    - file and line info
+
+    Currently it's still under development and format is subject to change.
+
+    Args:
+      node: The torch.fx.Node to build the location for.
+
+    Returns:
+      The MLIR location for the given node.
+    """
 
     info = debuginfo.build_mlir_debuginfo(node)
     if info is None:
@@ -98,7 +113,12 @@ def _build_loc(self, node: torch.fx.Node):
     if file is not None:
       fileinfo = ir.Location.file(filename=file, line=line, col=0)
 
-    return ir.Location.name(name=info, childLoc=fileinfo)
+    node_name = debuginfo.build_nodename_debuginfo(node)
+    nodeinfo = None
+    if node_name is not None:
+      nodeinfo = ir.Location.name(name=node_name, childLoc=fileinfo)
+
+    return ir.Location.name(name=info, childLoc=nodeinfo)
 
   def run_node(self, node: torch.fx.Node):
     loc = self._build_loc(node)

ai_edge_torch/odml_torch/test/test_tf_integration.py

@@ -58,9 +58,11 @@ def test_resnet18(self):
 
     # Check value and debug info.
     self.assertTrue(np.allclose(lowering_output, torch_output, atol=1e-5))
-    self.assertIn(
-        "torchvision/models/resnet.py", lowered.get_text(enable_debug_info=True)
-    )
+    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)
 
 
 if __name__ == "__main__":