feat: Add maxpool lowering passes and experimental folder in Dynamo (#2358)

Author: gs-olive

.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
     rev: 'v1.4.1'
     hooks:
     -   id: mypy
-        exclude: "^py/torch_tensorrt/fx|^examples|^tests|^tools|^docs|noxfile.py|setup.py|versions.py"
+        exclude: "^py/torch_tensorrt/fx|^examples|^tests|^py/torch_tensorrt/dynamo/_experimental|^tools|^docs|noxfile.py|setup.py|versions.py"
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
     rev: v0.0.278

py/torch_tensorrt/dynamo/backend/backends.py

@@ -15,7 +15,6 @@
     get_decompositions,
     repair_input_aliasing,
 )
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import pre_aot_substitutions
 from torch_tensorrt.dynamo.utils import (
     parse_dynamo_kwargs,
     prepare_inputs,
@@ -68,9 +67,6 @@ def _pretraced_backend(
     try:
         logger.debug("Pre-AOT Autograd graph:\n" + str(gm.graph))
 
-        # Perform Pre-AOT Lowering for Module-Level Replacement
-        gm = pre_aot_substitutions(gm)
-
         fake_mode = detect_fake_mode(sample_inputs)
 
         # Place backend tracing within FakeTensor context allowing nonfake Tensors

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -1,7 +1,4 @@
 from ._decompositions import get_decompositions  # noqa: F401
 from ._fusers import *  # noqa: F401
-from ._pre_aot_lowering import SUBSTITUTION_REGISTRY  # noqa: F401
-from ._pre_aot_lowering import register_substitution  # noqa: F401
 from ._repair_input_aliasing import repair_input_aliasing
 from .passes import apply_lowering_passes
-from .substitutions import *  # noqa: F401

py/torch_tensorrt/dynamo/lowering/_pre_aot_lowering.py

@@ -9,6 +9,7 @@
 from .pass_manager import DynamoPassManager
 from .remove_input_alias_fixing_clones import remove_input_alias_fixing_clones
 from .repair_input_as_output import repair_input_as_output
+from .replace_max_pool_with_indices import replace_max_pool_with_indices
 
 ATEN_LOWERING_PASSES = DynamoPassManager.build_from_passlist(
     [
@@ -17,6 +18,7 @@
         repair_input_as_output,
         lower_efficient_attention,
         fuse_prims_broadcast,
+        replace_max_pool_with_indices,
     ]
 )
 

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

@@ -0,0 +1,60 @@
+import logging
+import operator
+from typing import Sequence
+
+import torch
+from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
+    clean_up_graph_after_modifications,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def replace_max_pool_with_indices(
+    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
+) -> torch.fx.GraphModule:
+    """Replace MaxPool nodes which return unused indices"""
+    replacement_dict = {
+        torch.ops.aten.max_pool1d_with_indices.default: torch.ops.aten.max_pool1d.default,
+        torch.ops.aten.max_pool2d_with_indices.default: torch.ops.aten.max_pool2d.default,
+        torch.ops.aten.max_pool3d_with_indices.default: torch.ops.aten.max_pool3d.default,
+    }
+
+    modified_graph = False
+
+    for node in gm.graph.nodes:
+        # If the node is a placeholder and its only user is a clone node
+        # it was modified by the input alias-fixing pass, and the change
+        # needs to be undone
+        if (
+            node.target in replacement_dict
+            and len(node.users) == 1
+            and list(node.users)[0].target == operator.getitem
+            and list(node.users)[0].args[1] == 0
+        ):
+            modified_graph = True
+
+            # Replace all uses of the clone with the placholder, delete the clone
+            getitem_node = list(node.users)[0]
+
+            with gm.graph.inserting_after(getitem_node):
+                maxpool_fused = gm.graph.call_function(
+                    replacement_dict[node.target],
+                    args=node.args,
+                    kwargs=node.kwargs,
+                )
+
+            logger.debug(
+                f"Replacing all uses of nodes {node}, {getitem_node} with fused maxpool node {maxpool_fused} "
+                f"is the only user of placeholder {node} and was inserted by the compiler."
+            )
+
+            getitem_node.replace_all_uses_with(maxpool_fused)
+            gm.graph.erase_node(getitem_node)
+            gm.graph.erase_node(node)
+
+    if modified_graph:
+        gm = clean_up_graph_after_modifications(gm)
+        logger.debug(f"Graph after fusing maxpool operators with indices:\n{gm.graph}")
+
+    return gm