Add file-line info to ODML Torch lowerings.

PiperOrigin-RevId: 724070541

Author: GleasonK

ai_edge_torch/odml_torch/debuginfo/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ._build import build_mlir_debuginfo
+from ._build import build_mlir_debuginfo, build_mlir_file_debuginfo
 from ._op_polyfill import write_mlir_debuginfo_op

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import torch
+import re
 
 
 def _class_fullname(cls):
@@ -34,6 +35,29 @@ def _get_hierarchy(node: torch.fx.Node):
   return hierachy_str
 
 
+def _get_canonical_filename(filename):
+  """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
+  can override.
+  """
+
+  # TODO: We should add a config option to provide a regex to strip from the
+  # debug info. Currently absolute path is used.
+  return filename
+
+
+def build_mlir_file_debuginfo(node: torch.fx.Node):
+  """Build the file and line info for the given node's lowerings in MLIR."""
+
+  if not node.stack_trace:
+    return None, None
+
+  # Note: This uses internal APIs and may break in the future.
+  pt_trace = torch.fx.graph._parse_stack_trace(node.stack_trace)
+  return _get_canonical_filename(pt_trace.file), int(pt_trace.lineno)
+
+
 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

@@ -93,7 +93,12 @@ def _build_loc(self, node: torch.fx.Node):
     if info is None:
       return ir.Location.unknown()
 
-    return ir.Location.name(name=info)
+    (file, line) = debuginfo.build_mlir_file_debuginfo(node)
+    fileinfo = None
+    if file is not None:
+      fileinfo = ir.Location.file(filename=file, line=line, col=0)
+
+    return ir.Location.name(name=info, childLoc=fileinfo)
 
   def run_node(self, node: torch.fx.Node):
     loc = self._build_loc(node)

ai_edge_torch/odml_torch/test/test_tf_integration.py

@@ -56,7 +56,11 @@ def test_resnet18(self):
     torch_output = model(*args).detach().numpy()
     lowering_output = np.array(lowered(*args))
 
+    # 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)
+    )
 
 
 if __name__ == "__main__":