feat: Refactor pass manager and utilities

- Improve logging and pass manager utilities
- Add testing of new utilities

Author: gs-olive

examples/dynamo/dynamo_aten_lowering_passes.py

@@ -86,8 +86,24 @@ def repair_input_as_output(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
 # %%
 from torch_tensorrt.dynamo.lowering.passes import add_lowering_pass
 
+# Adds the lowering pass at the end of the pass list
 add_lowering_pass(repair_input_as_output)
 
+# Alternatively, specify an index to insert the lowering pass at a specific location
+add_lowering_pass(repair_input_as_output, 1)
+
+# To remove a lowering pass, specify the index of the pass to remove:
+from torch_tensorrt.dynamo.lowering.passes import remove_lowering_pass
+
+# Removes the lowering pass at index 1
+remove_lowering_pass(1)
+
+
+# To view all lowering passes, in the order they will be run, use the following
+from torch_tensorrt.dynamo.lowering.passes import dump_lowering_passes
+
+print(dump_lowering_passes())
+
 # %%
 # 3. Apply Available Lowering Passes
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

py/torch_tensorrt/dynamo/lowering/passes/__init__.py

@@ -1,27 +1,55 @@
-from typing import Callable
+import logging
+from typing import Callable, Optional
 
 import torch
-from torch.fx.passes.pass_manager import PassManager
 
+# Import and order lowering passes and pass manager
 from .constant_folding import constant_fold
+from .pass_manager import DynamoPassManager
 from .repair_input_as_output import repair_input_as_output
 
-ATEN_LOWERING_PASSES = PassManager.build_from_passlist(
+ATEN_LOWERING_PASSES = DynamoPassManager.build_from_passlist(
     [
         constant_fold,
         repair_input_as_output,
     ]
 )
 
+logger = logging.getLogger(__name__)
+
 
 def add_lowering_pass(
-    lowering_pass: Callable[[torch.fx.GraphModule], torch.fx.GraphModule]
+    lowering_pass: Callable[[torch.fx.GraphModule], torch.fx.GraphModule],
+    index: Optional[int] = None,
 ) -> None:
-    """Adds a lowering pass to the registry"""
-    ATEN_LOWERING_PASSES.add_pass(lowering_pass)
+    """Adds a lowering pass to the registry, at a specified index if desired
+
+    If no index is specified, the lowering pass is inserted at the end of the list
+    """
+    ATEN_LOWERING_PASSES.add_pass_with_index(lowering_pass, index)
+    logger.debug(
+        f"Added lowering pass {lowering_pass} to list at index {index}, current passlist: {ATEN_LOWERING_PASSES}"
+    )
+    return
+
+
+def remove_lowering_pass(index: int) -> None:
+    """Removes a lowering pass at a specific index from the registry"""
+    ATEN_LOWERING_PASSES.remove_pass_with_index(index)
+    logger.debug(
+        f"Removed lowering pass at index {index}, current passlist: {ATEN_LOWERING_PASSES}"
+    )
     return
 
 
 def apply_lowering_passes(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Applies the lowering passes to a graph module, returns the modified GraphModule"""
+    logging.debug(
+        f"Invoking DynamoPassManager and applying lowering passes: {ATEN_LOWERING_PASSES}"
+    )
     return ATEN_LOWERING_PASSES(gm)
+
+
+def dump_lowering_passes() -> str:
+    """Returns a string containing the lowering passes"""
+    return str(ATEN_LOWERING_PASSES)

py/torch_tensorrt/dynamo/lowering/passes/pass_manager.py

@@ -0,0 +1,42 @@
+from typing import Any, Callable, List, Optional
+
+import torch
+from torch.fx.passes.pass_manager import PassManager
+
+
+class DynamoPassManager(PassManager):  # type: ignore[misc]
+    def __init__(
+        self,
+        passes: Optional[
+            List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]]
+        ] = None,
+    ):
+        super().__init__(passes)
+
+    @classmethod
+    def build_from_passlist(
+        cls,
+        passes: Optional[List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]]],
+    ) -> Any:
+        pm = DynamoPassManager(passes)
+        return pm
+
+    def add_pass_with_index(
+        self,
+        lowering_pass: Callable[[torch.fx.GraphModule], torch.fx.GraphModule],
+        index: Optional[int] = None,
+    ) -> None:
+        if index is None:
+            self.passes.append(lowering_pass)
+            index = -1
+        else:
+            self.passes.insert(index, lowering_pass)
+
+    def remove_pass_with_index(self, index: int) -> None:
+        del self.passes[index]
+
+    def __call__(self, source: Any) -> Any:
+        return super().__call__(source)
+
+    def __str__(self) -> str:
+        return str(self.passes)

tests/py/dynamo/lowering/test_aten_lowering_passes.py

@@ -55,5 +55,43 @@ def forward(self, x, y):
         torch._dynamo.reset()
 
 
+class TestLoweringPassMembership(TestCase):
+    def insert_at_end(self):
+        from torch_tensorrt.dynamo.lowering.passes import (
+            ATEN_LOWERING_PASSES,
+            add_lowering_pass,
+            remove_lowering_pass,
+        )
+
+        def identity_pass(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+            return gm
+
+        add_lowering_pass(identity_pass)
+
+        self.assertEqual(identity_pass, ATEN_LOWERING_PASSES.passes[-1])
+
+        remove_lowering_pass(-1)
+
+        self.assertNotIn(identity_pass, ATEN_LOWERING_PASSES.passes)
+
+    def insert_at_index(self):
+        from torch_tensorrt.dynamo.lowering.passes import (
+            ATEN_LOWERING_PASSES,
+            add_lowering_pass,
+            remove_lowering_pass,
+        )
+
+        def identity_pass(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
+            return gm
+
+        add_lowering_pass(identity_pass, 0)
+
+        self.assertEqual(identity_pass, ATEN_LOWERING_PASSES.passes[0])
+
+        remove_lowering_pass(0)
+
+        self.assertNotIn(identity_pass, ATEN_LOWERING_PASSES.passes)
+
+
 if __name__ == "__main__":
     run_tests()