merged the file to pass_manager

Author: cehongwang

py/torch_tensorrt/dynamo/lowering/passes/draw_fx_graph.py

@@ -1,12 +1,40 @@
-from typing import Any, Callable, List, Optional, Sequence
+from typing import Any, Callable, List, Optional
 
 import torch
+from torch.fx import passes
 from torch.fx.passes.pass_manager import PassManager
 from torch_tensorrt.dynamo._settings import CompilationSettings
-from torch_tensorrt.dynamo.lowering.passes.draw_fx_graph import (
-    get_draw_fx_graph_pass_post_lowering,
-    get_draw_fx_graph_pass_pre_lowering,
-)
+
+
+def get_draw_fx_graph_pass_lowering(
+    idx: int, path_prefix: str, post: bool
+) -> Callable[[torch.fx.GraphModule, CompilationSettings], torch.fx.GraphModule]:
+    from torch_tensorrt.dynamo.lowering.passes import (
+        post_lowering_pass_list,
+        pre_lowering_pass_list,
+    )
+
+    PRE_DEBUG_NAME = {
+        i + 1: f"after_{p.__name__}" for i, p in enumerate(pre_lowering_pass_list)
+    }
+    PRE_DEBUG_NAME[0] = "exported_program"
+
+    POST_DEBUG_NAME = {
+        i + 1: f"after_{p.__name__}" for i, p in enumerate(post_lowering_pass_list)
+    }
+    POST_DEBUG_NAME[0] = "after_decomposition"
+
+    def draw_fx_graph_pass(
+        gm: torch.fx.GraphModule, settings: CompilationSettings
+    ) -> torch.fx.GraphModule:
+        DEBUG_NAME = POST_DEBUG_NAME[idx] if post else PRE_DEBUG_NAME[idx]
+        path = f"{path_prefix}_{DEBUG_NAME}.svg"
+        g = passes.graph_drawer.FxGraphDrawer(gm, DEBUG_NAME)
+        with open(path, "wb") as f:
+            f.write(g.get_dot_graph().create_svg())
+        return gm
+
+    return draw_fx_graph_pass
 
 
 class DynamoPassManager(PassManager):  # type: ignore[misc]
@@ -39,8 +67,7 @@ def build_from_passlist(
     def add_pass_with_index(
         self,
         lowering_pass: Callable[
-            [torch.fx.GraphModule, CompilationSettings, Sequence[torch.Tensor]],
-            torch.fx.GraphModule,
+            [torch.fx.GraphModule, CompilationSettings], torch.fx.GraphModule
         ],
         index: Optional[int] = None,
     ) -> None:
@@ -58,14 +85,10 @@ def insert_debug_pass(
     ) -> None:
 
         for i in range(len(index)):
-            if post:
-                debug_pass = get_draw_fx_graph_pass_post_lowering(
-                    index[i], filename_prefix
-                )
-            else:
-                debug_pass = get_draw_fx_graph_pass_pre_lowering(
-                    index[i], filename_prefix
-                )
+
+            debug_pass = get_draw_fx_graph_pass_lowering(
+                index[i], filename_prefix, post
+            )
             self.add_pass_with_index(debug_pass, index[i] + i)
 
     def __call__(self, gm: Any, settings: CompilationSettings) -> Any: