FX graph visualization

[cehongwang]

Description

Debugging FX graphs can be challenging due to the complexity of analyzing node connections directly from the FX table. Therefore, providing a clear visualization of the FX graph is essential to facilitate effective debugging.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[narendasan]

Why doesnt this just edit the list internally instead of calling out to a function that calls back in?

[cehongwang]

I was thinking to keep the debug pass in a separate file, as all other lowering passes are in separated files. Also each lowering pass are different because the name depends on the position and pre/post lowering, which need to be dynamically generated using the function get_draw_fx_graph_pass_pre/post_lowering. Do you think I should put everything such as get_draw_fx_graph_pass_post_lowering function in pass_manager?

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/pass_manager.py	2025-05-23 04:32:05.196604+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/passes/pass_manager.py	2025-05-23 04:32:28.295008+00:00
@@ -31,11 +31,11 @@
                Callable[
                    [torch.fx.GraphModule, CompilationSettings], torch.fx.GraphModule
                ]
            ]
        ] = None,
-        constraints: Optional[List[Callable]] = None
+        constraints: Optional[List[Callable]] = None,
    ):
        super().__init__(passes, constraints)

    @classmethod
    def build_from_passlist(
@@ -66,11 +66,11 @@

    def remove_pass_with_index(self, index: int) -> None:
        del self.passes[index]

    def insert_debug_pass_before(
-        self, passes: List[str], output_path_prefix: str=tempfile.gettempdir()
+        self, passes: List[str], output_path_prefix: str = tempfile.gettempdir()
    ) -> None:
        """Insert debug passes in the PassManager pass sequence prior to the execution of a particular pass.

        Args:
            passes: List of pass names to insert debug passes before
@@ -80,18 +80,22 @@
        in the pass sequence.
        """
        new_pass_list = []
        for ps in self.passes:
            if ps.__name__ in passes:
-                new_pass_list.append(_generate_draw_fx_graph_pass(output_path_prefix, f"before_{ps.__name__}"))
+                new_pass_list.append(
+                    _generate_draw_fx_graph_pass(
+                        output_path_prefix, f"before_{ps.__name__}"
+                    )
+                )
            new_pass_list.append(ps)

        self.passes = new_pass_list
        self._validated = False

    def insert_debug_pass_after(
-        self, passes: List[str], output_path_prefix: str=tempfile.gettempdir()
+        self, passes: List[str], output_path_prefix: str = tempfile.gettempdir()
    ) -> None:
        """Insert debug passes in the PassManager pass sequence after the execution of a particular pass.

        Args:
            passes: List of pass names to insert debug passes after
@@ -102,12 +106,15 @@
        """
        new_pass_list = []
        for ps in self.passes:
            new_pass_list.append(ps)
            if ps.__name__ in passes:
-                new_pass_list.append(_generate_draw_fx_graph_pass(output_path_prefix, f"after_{ps.__name__}"))
-
+                new_pass_list.append(
+                    _generate_draw_fx_graph_pass(
+                        output_path_prefix, f"after_{ps.__name__}"
+                    )
+                )

        self.passes = new_pass_list
        self._validated = False

    def __call__(self, gm: Any, settings: CompilationSettings) -> Any: