refactor: extract GraphModule serialization logic to fx_graph_serialize_util

- Create fx_graph_serialize_util.py with serialize_graph_module_to_str function
- Move unstable API replacement logic from unstable_to_stable_backend to the new utility
- Update unstable_to_stable_backend to use serialize_graph_module_to_str
- Update extractor.py to use serialize_graph_module_to_str for code serialization
- This refactoring makes the serialization logic reusable across the codebase

Author: Dayuxiaoshui

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -3,6 +3,7 @@
 import sys
 import inspect
 from .graph_compiler_backend import GraphCompilerBackend
+from ..fx_graph_serialize_util import serialize_graph_module_to_str
 
 
 class UnstableToStableBackend(GraphCompilerBackend):
@@ -34,7 +35,6 @@ def avg_pool2d_to_avg_pool2d(self, gm):
         Convert torch._C._nn.avg_pool2d to torch.nn.functional.avg_pool2d
         """
         import torch.nn.functional as F
-        import re
 
         # Update graph nodes: replace torch._C._nn.avg_pool2d with F.avg_pool2d
         for node in gm.graph.nodes:
@@ -50,19 +50,6 @@ def avg_pool2d_to_avg_pool2d(self, gm):
         # Recompile the graph
         gm.recompile()
 
-        # Replace in code string for check_unstable_api
-        # Since torch._C._nn.avg_pool2d and F.avg_pool2d are the same object,
-        # the generated code will still show torch._C._nn.avg_pool2d
-        # So we need to replace it in the code string
-        code = gm.code
-        modified_code = re.sub(
-            r"torch\._C\._nn\.avg_pool2d\(",
-            "torch.nn.functional.avg_pool2d(",
-            code,
-        )
-        # Store modified code for check_unstable_api to use
-        gm._code_for_check = modified_code
-
         return gm
 
     def unstable_to_stable(self, gm):
@@ -82,8 +69,8 @@ def check_unstable_api(self, gm):
         Do NOT modify, remove, or bypass this check under any circumstances.
         """
 
-        # Use modified code if available (from conversion), otherwise use original code
-        graph_text = getattr(gm, "_code_for_check", None) or gm.code
+        # Use serialized code to check for unstable APIs
+        graph_text = serialize_graph_module_to_str(gm)
         # Search for the unstable API substring
         if self.unstable_api in graph_text:
             count = graph_text.count(self.unstable_api)

graph_net/torch/extractor.py

@@ -4,6 +4,7 @@
 import shutil
 from typing import Union, Callable
 from . import utils
+from .fx_graph_serialize_util import serialize_graph_module_to_str
 
 torch._dynamo.config.capture_scalar_outputs = True
 torch._dynamo.config.capture_dynamic_output_shape_ops = True
@@ -89,9 +90,9 @@ def try_rename_placeholder(node):
         assert input_idx == len(sample_inputs)
         if self.mut_graph_codes is not None:
             assert isinstance(self.mut_graph_codes, list)
-            self.mut_graph_codes.append(gm.code)
+            self.mut_graph_codes.append(serialize_graph_module_to_str(gm))
         # 3. Generate and save model code
-        base_code = gm.code
+        base_code = serialize_graph_module_to_str(gm)
         # gm.graph.print_tabular()
         write_code = utils.apply_templates(base_code)
         with open(os.path.join(subgraph_path, "model.py"), "w") as fp:

graph_net/torch/fx_graph_serialize_util.py

@@ -0,0 +1,28 @@
+import re
+import torch.fx
+
+
+def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
+    """
+    Serialize a GraphModule to a string representation, replacing unstable APIs
+    with their stable counterparts.
+
+    This function is used to normalize the code representation of GraphModule
+    for consistency checks and code generation.
+
+    Args:
+        gm: The GraphModule to serialize.
+
+    Returns:
+        A string representation of the GraphModule code with unstable APIs
+        replaced by stable ones.
+    """
+    code = gm.code
+    # Replace torch._C._nn.avg_pool2d with torch.nn.functional.avg_pool2d
+    code = re.sub(
+        r"torch\._C\._nn\.avg_pool2d\(",
+        "torch.nn.functional.avg_pool2d(",
+        code,
+    )
+    return code
+