Update AOTEagerandRecordGraphs backend (#138231)

Summary:
X-link: pytorch/pytorch#138231
Approved by: https://github.com/StrongerXi, https://github.com/mlazos, https://github.com/aakhundov

Reviewed By: clee2000

Differential Revision: D64581452

Pulled By: anijain2305

fbshipit-source-id: 3b9ff53abf2c4e1c525d7e62a52285279d2d4109

Author: anijain2305

userbenchmark/dynamo/dynamobench/_dynamo/testing.py

@@ -246,24 +246,34 @@ def __call__(
         return gm.forward
 
 
-# Equivalent to backend="aot_eager", but also records graphs that
-# we can assert on
-class AOTEagerAndRecordGraphs:
+class AotEagerAndRecordGraphs:
     def __init__(self) -> None:
         self.graphs: List[torch.fx.GraphModule] = []
+        self.fw_graphs: List[torch.fx.GraphModule] = []
+        self.bw_graphs: List[torch.fx.GraphModule] = []
 
     def __call__(
         self, gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]
     ) -> Callable[..., Any]:
-        def save_graph(gm: torch.fx.GraphModule, *args: Any, **kwargs: Any) -> Any:
-            self.graphs.append(gm)
+        self.graphs.append(gm)
+
+        def fw_compiler(
+            gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]
+        ) -> Callable[..., Any]:
+            self.fw_graphs.append(gm)
+            return gm.forward
+
+        def bw_compiler(
+            gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]
+        ) -> Callable[..., Any]:
+            self.bw_graphs.append(gm)
             return gm.forward
 
         return aot_eager(
             gm,
             example_inputs,
-            fw_compiler=save_graph,
-            bw_compiler=save_graph,
+            fw_compiler=fw_compiler,
+            bw_compiler=bw_compiler,
         )