Clean up history stack before fx passes in ODML torch lowering

PiperOrigin-RevId: 731019293

Author: yijie-yang

ai_edge_torch/_convert/conversion.py

@@ -125,6 +125,10 @@ def export(**kwargs):
     else:
       exported_program = torch.export.export(**kwargs, strict=True)
 
+    exported_program = fx_infra.graph_utils.reset_from_node_meta(
+        exported_program
+    )
+
     exported_program = fx_infra.safe_run_decompositions(
         exported_program,
         fx_infra.decomp.pre_convert_decomp(),

ai_edge_torch/fx_infra/graph_utils.py

@@ -13,7 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 """FX graph utilities."""
+
+from packaging import version
 import torch
+from torch.fx import traceback
 
 
 def remove_dangling_args(graph_module: torch.fx.GraphModule):
@@ -40,3 +43,32 @@ def remove_assert_tensor_metadata_nodes(graph_module: torch.fx.GraphModule):
   graph_module.graph.lint()
   graph_module.recompile()
   return graph_module
+
+
+def is_torch_version_under(torch_version: str) -> bool:
+  """Checks if the current torch version is under the given version."""
+  if not torch_version:
+    raise ValueError("torch_version cannot be empty.")
+  current_version = version.parse(torch.__version__)
+  compared_version = version.parse(torch_version)
+  return current_version < compared_version
+
+
+def reset_from_node_meta(ep: torch.export.ExportedProgram):
+  """Resets the "from_node" meta field to fx node name only for the exported program."""
+
+  for node in ep.graph.nodes:
+    if not hasattr(node, "meta") or "from_node" not in node.meta:
+      continue
+    if is_torch_version_under("2.6.0.dev0"):
+      # For torch version under 2.6.0, the history stack is a list of tuple. We
+      # will only keep the current node's name in the history stack.
+      history = [(node.name,)]
+    else:
+      # Clean up the history stack by keeping only the current node info (fx
+      # node name and graph id) in a list of size 1. Clear the "from_node" field
+      # to prevent redundant additions to the history stack.
+      history = [traceback.NodeSource(node)]
+      history[0].from_node = []
+    node.meta["from_node"] = history
+  return ep

ai_edge_torch/odml_torch/debuginfo/_build.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-import re
+from ai_edge_torch.fx_infra import graph_utils
 import torch
 
 
@@ -57,9 +57,13 @@ 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 not history:
+    return None
 
-  if len(history) > 1:  # Compatible with torch version under 2.6.0
-    return history[1][0]
+  # Compatible with torch version under 2.6.0. The history stack is a list of
+  # tuple. The first element of the first tuple is the node name.
+  if graph_utils.is_torch_version_under("2.6.0.dev0"):
+    return history[0][0]
 
   if not hasattr(history[0], "name"):
     return None
@@ -68,12 +72,10 @@ def _get_canoical_nodename(node: torch.fx.Node) -> str:
     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
+  # The history stack is generated by tracing the node's transformation history
+  # during lowering. The last name in the history stack is used to map to the
+  # original torch fx node name.
+  return names[-1]
 
 
 def build_mlir_file_debuginfo(node: torch.fx.Node):

ai_edge_torch/odml_torch/test/test_tf_integration.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from ai_edge_torch import fx_infra
 from ai_edge_torch import odml_torch
 import numpy as np
 import torch
@@ -20,6 +21,23 @@
 from absl.testing import absltest as googletest
 
 
+def _reset_from_node_meta_and_lower(ep: torch.export.ExportedProgram):
+  """Lower the exported program with canonical history stack."""
+  ep = fx_infra.graph_utils.reset_from_node_meta(ep)
+  return odml_torch.export.exported_program_to_mlir(ep)
+
+
+def _is_aten_op(node: torch.fx.Node) -> bool:
+  return node.op == "call_function" and not node.name.startswith("getitem")
+
+
+class AddModel(torch.nn.Module):
+  """A simple model that does addition."""
+
+  def forward(self, x, y):
+    return x + y + x + y
+
+
 class TensorflowIntegrationTest(googletest.TestCase):
 
   def setUp(self):
@@ -29,11 +47,6 @@ def setUp(self):
   def test_mlir_lowered_call(self):
     """Test a simple model with MLIR lowered call."""
 
-    class AddModel(torch.nn.Module):
-
-      def forward(self, x, y):
-        return x + y + x + y
-
     model = AddModel().eval()
     forward_args = lambda: (torch.rand((10, 10)), torch.rand((10, 10)))
     ep = torch.export.export(model, forward_args())
@@ -53,8 +66,7 @@ def test_resnet18(self):
     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 = _reset_from_node_meta_and_lower(ep)
 
     args = forward_args()
     torch_output = model(*args).detach().numpy()
@@ -70,7 +82,7 @@ def test_debuginfo_from_export_lower(self):
     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 = _reset_from_node_meta_and_lower(ep)
 
     lowered_text = lowered.get_text(enable_debug_info=True)
     # Check the file info.
@@ -79,10 +91,35 @@ def test_debuginfo_from_export_lower(self):
     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"):
+      if _is_aten_op(n):
         # Ensure strings like `loc("relu__1"` are present in the lowered text.
         self.assertIn(f'loc("{n.name}"', lowered_text)
 
+  def test_debuginfo_from_loaded_reexport_lower(self):
+    """Test the debuginfo with loaded reexport lower."""
+
+    model = AddModel().eval()
+    forward_args = lambda: (torch.rand((10, 10)), torch.rand((10, 10)))
+
+    # Ensure the debuginfo is preserved after saving, loading and reexporting.
+    ep = torch.export.export(model, forward_args())
+    torch.export.save(ep, "/tmp/add_model.pt2")
+    loaded_ep = torch.export.load("/tmp/add_model.pt2")
+    reexported_ep = torch.export.export(loaded_ep.module(), forward_args())
+    lowered = _reset_from_node_meta_and_lower(reexported_ep)
+
+    lowered_text = lowered.get_text(enable_debug_info=True)
+    # Check the file info.
+    self.assertIn(
+        "ai_edge_torch/odml_torch/test/test_tf_integration.py", lowered_text
+    )
+    # Check the fx node names.
+    for n in reexported_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 _is_aten_op(n):
+        self.assertIn(f'loc("{n.name}"', lowered_text)
+
 
 if __name__ == "__main__":
   googletest.main()