解决冲突问题

Author: G2uge

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -30,7 +30,37 @@ def my_backend(gm, sample_inputs):
     **Stable API reference link:**
     """
 
-    def avg_pool2d_to_avg_pool2d(self, gm):
+    def _impl_unstable_to_stable_irfft(self, gm):
+        def replace_in_graph(graph_mod):
+            # Register stable implementation on GraphModule, codegen can use self.irfft
+            try:
+                setattr(graph_mod, "irfft", torch.fft.irfft)
+            except Exception:
+                pass
+
+            for node in graph_mod.graph.nodes:
+                if node.op == "call_function":
+                    # Match for all forms of target names
+                    if "fft_irfft" in str(node.target):
+                        # Directly point target to Python layer function
+                        node.target = torch.fft.irfft
+            # Validate and recompile the graph
+            graph_mod.graph.lint()
+            graph_mod.recompile()
+
+        # Process main gm and all nested GraphModules
+        modules = [gm]
+        modules += [
+            m
+            for _, m in gm.named_modules()
+            if isinstance(m, torch.fx.GraphModule) and m is not gm
+        ]
+        for m in modules:
+            replace_in_graph(m)
+
+        return gm
+
+    def _impl_unstable_to_stable_avg_pool2d(self, gm):
         """
         Convert torch._C._nn.avg_pool2d to torch.nn.functional.avg_pool2d
         """
@@ -52,7 +82,7 @@ def avg_pool2d_to_avg_pool2d(self, gm):
 
         return gm
 
-    def fft_rfft_to_rfft(self, gm):
+    def _impl_unstable_to_stable_rfft(self, gm):
         """
         Convert torch._C._fft.fft_rfft to torch.fft.rfft
         """
@@ -75,11 +105,13 @@ def fft_rfft_to_rfft(self, gm):
         return gm
 
     def unstable_to_stable(self, gm):
-        # Convert based on unstable_api environment variable
-        if self.unstable_api == "torch._C._nn.avg_pool2d":
-            gm = self.avg_pool2d_to_avg_pool2d(gm)
-        elif self.unstable_api == "torch._C._fft.fft_rfft":
-            gm = self.fft_rfft_to_rfft(gm)
+        methods = (
+            name
+            for name in vars(type(self)).keys()
+            if name.startswith("_impl_unstable_to_stable")
+        )
+        for method in methods:
+            gm = getattr(self, method)(gm)
         return gm
 
     def check_unstable_api(self, gm):

graph_net/torch/fx_graph_serialize_util.py

@@ -19,14 +19,12 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
     """
     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,
-    )
-    code = re.sub(
-        r"torch\._C\._fft\.fft_rfft\(",
-        "torch.fft.rfft(",
-        code,
-    )
+    replacements = [
+        (r"torch\._C\._nn\.avg_pool2d\(", "torch.nn.functional.avg_pool2d("),
+        (r"torch\._C\._fft\.fft_irfft\(", "torch.fft.irfft("),
+        (r"torch\._C\._fft\.fft_rfft\(", "torch.fft.rfft("),
+        # Add new rules to this list as needed
+    ]
+    for pattern, repl in replacements:
+        code = re.sub(pattern, repl, code)
     return code