【Hackathon 9th No.113】torch._C._fft.fft_irfft API转换 torch.fft.irfft (#323)

* 英文注释，撤销原始日志文件

* 修正api

Author: G2uge

graph_net/torch/backend/unstable_to_stable_backend.py

@@ -30,6 +30,36 @@ def my_backend(gm, sample_inputs):
     **Stable API reference link:**
     """
 
+    def fft_irfft_to_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 avg_pool2d_to_avg_pool2d(self, gm):
         """
         Convert torch._C._nn.avg_pool2d to torch.nn.functional.avg_pool2d
@@ -56,6 +86,8 @@ 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_irfft":
+            gm = self.fft_irfft_to_irfft(gm)
         return gm
 
     def check_unstable_api(self, gm):

graph_net/torch/fx_graph_serialize_util.py

@@ -24,4 +24,9 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
         "torch.nn.functional.avg_pool2d(",
         code,
     )
+    code = re.sub(
+        r"torch\._C\._fft\.fft_irfft\(",
+        "torch.fft.irfft(",
+        code,
+    )
     return code