Deleted sample_inputs

Author: cehongwang

py/torch_tensorrt/dynamo/_compiler.py

@@ -182,14 +182,14 @@ def compile(
         raise AssertionError(
             f"Input graph should be an ExportedProgram but got type {type(exported_program)}"
         )
-    exported_program = pre_export_lowering(exported_program, None)
+    exported_program = pre_export_lowering(exported_program)
     exported_program = exported_program.run_decompositions(
         get_decompositions(enable_experimental_decompositions)
     )
     gm = exported_program.module()
     logger.debug("Input graph: " + str(gm.graph))
     # Apply lowering on the graph module
-    gm = post_lowering(gm, None)
+    gm = post_lowering(gm)
     logger.debug("Lowered Input graph: " + str(gm.graph))
 
     compilation_options = {
@@ -602,7 +602,7 @@ def convert_module_to_trt_engine(
         "timing_cache_path": timing_cache_path,
     }
 
-    exported_program = pre_export_lowering(exported_program, torch_inputs)
+    exported_program = pre_export_lowering(exported_program)
     # Decompose the exported program
     exported_program = exported_program.run_decompositions(
         get_decompositions(enable_experimental_decompositions)
@@ -611,7 +611,7 @@ def convert_module_to_trt_engine(
     logger.debug("Input graph: " + str(gm.graph))
 
     # Apply lowering on the graph module
-    gm = post_lowering(gm, torch_inputs)
+    gm = post_lowering(gm)
     logger.debug("Lowered Input graph: " + str(gm.graph))
 
     settings = CompilationSettings(**compilation_options)

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

@@ -86,25 +86,21 @@ def _remove_lowering_pass(*, index: int) -> None:
     return
 
 
-def post_lowering(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def post_lowering(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Applies the lowering passes to a graph module after torch.export/ torch.compile and their decompositions, returns the modified GraphModule"""
     logging.debug(
         f"Invoking DynamoPassManager and applying lowering passes: {ATEN_POST_LOWERING_PASSES}"
     )
-    return ATEN_POST_LOWERING_PASSES(gm, sample_inputs)
+    return ATEN_POST_LOWERING_PASSES(gm)
 
 
-def pre_export_lowering(
-    ep: torch.export.ExportedProgram, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def pre_export_lowering(ep: torch.export.ExportedProgram) -> torch.fx.GraphModule:
     """Applies the lowering passes to a graph module after torch.export/ torch.compile and their decompositions, returns the modified GraphModule"""
     logging.debug(
         f"Invoking DynamoPassManager and applying lowering passes: {ATEN_PRE_LOWERING_PASSES}"
     )
     gm = ep.graph_module
-    gm = ATEN_PRE_LOWERING_PASSES(gm, sample_inputs)
+    gm = ATEN_PRE_LOWERING_PASSES(gm)
     return ep
 
 

py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Sequence
+from typing import Any
 
 import torch
 from torch_tensorrt._utils import sanitized_torch_version
@@ -19,9 +19,7 @@
 
 
 @torch.utils._python_dispatch._disable_current_modes()  # type: ignore
-def constant_fold(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def constant_fold(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Adapted from:
     https://github.com/pytorch/pytorch/blob/3a79621c9dce17f77fbddc06aab21f6bc477f313/torch/_inductor/freezing.py#L178-L197
 

py/torch_tensorrt/dynamo/lowering/passes/fuse_prims_broadcast.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Sequence
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
@@ -10,9 +9,7 @@
 
 
 # TODO: Add relevant prims to this fusion
-def fuse_prims_broadcast(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def fuse_prims_broadcast(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Fuses prim nodes which are effectively the ATen equivalents with keep_dim=True"""
     modified_graph = False
 

py/torch_tensorrt/dynamo/lowering/passes/lower_linear.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Callable, Sequence, Tuple
+from typing import Callable, Tuple
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
@@ -9,9 +9,7 @@
 logger = logging.getLogger(__name__)
 
 
-def lower_linear(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def lower_linear(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Replace aten.linear with an equivalent implementation which can be easily converted to TRT"""
     orig, replacement = linear_replacement()
 

py/torch_tensorrt/dynamo/lowering/passes/lower_scaled_dot_product_attention.py

@@ -16,7 +16,7 @@
 
 
 def lower_scaled_dot_product_attention(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
+    gm: torch.fx.GraphModule,
 ) -> torch.fx.GraphModule:
     """Replace specific versions of scaled_dot_product_attention with an equivalent
     implementation which can be easily converted to TRT

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

@@ -8,25 +8,15 @@ class DynamoPassManager(PassManager):  # type: ignore[misc]
     def __init__(
         self,
         passes: Optional[
-            List[
-                Callable[
-                    [torch.fx.GraphModule, Sequence[torch.Tensor]], torch.fx.GraphModule
-                ]
-            ]
+            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, Sequence[torch.Tensor]], torch.fx.GraphModule
-                ]
-            ]
-        ],
+        passes: Optional[List[Callable[[torch.fx.GraphModule], torch.fx.GraphModule]]],
     ) -> Any:
         pm = DynamoPassManager(passes)
         return pm
@@ -47,11 +37,11 @@ def add_pass_with_index(
     def remove_pass_with_index(self, index: int) -> None:
         del self.passes[index]
 
-    def __call__(self, gm: Any, sample_inputs: Any) -> Any:
+    def __call__(self, gm: Any) -> Any:
         self.validate()
-        out, example_inputs = gm, sample_inputs
+        out = gm
         for _pass in self.passes:
-            out = _pass(out, example_inputs)
+            out = _pass(out)
         return out
 
     def __str__(self) -> str:

py/torch_tensorrt/dynamo/lowering/passes/remove_detach.py

@@ -1,14 +1,11 @@
 import logging
-from typing import Sequence
 
 import torch
 
 logger = logging.getLogger(__name__)
 
 
-def remove_detach(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def remove_detach(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Remove detach ops in the graph"""
     count = 0
     for node in gm.graph.nodes:

py/torch_tensorrt/dynamo/lowering/passes/remove_input_alias_fixing_clones.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Sequence
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
@@ -10,9 +9,7 @@
 
 
 # TODO: Delete this lowering pass once aot_export_joint_simple is patched
-def remove_input_alias_fixing_clones(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def remove_input_alias_fixing_clones(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Remove the auxiliary clone nodes inserted to fix input aliasing
 
     See: https://github.com/pytorch/pytorch/issues/108079

py/torch_tensorrt/dynamo/lowering/passes/repair_input_as_output.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Sequence
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
@@ -10,9 +9,7 @@
 logger = logging.getLogger(__name__)
 
 
-def repair_input_as_output(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def repair_input_as_output(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Repair scenarios where inputs are also outputs of the graph
 
     TRT does not allow such cases, so we insert a clone (identity) layer

py/torch_tensorrt/dynamo/lowering/passes/replace_max_pool_with_indices.py

@@ -1,6 +1,5 @@
 import logging
 import operator
-from typing import Sequence
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
@@ -11,7 +10,7 @@
 
 
 def replace_max_pool_with_indices(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
+    gm: torch.fx.GraphModule,
 ) -> torch.fx.GraphModule:
     """Replace MaxPool nodes which return unused indices"""
     replacement_dict = {

py/torch_tensorrt/dynamo/lowering/passes/view_to_reshape.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List, Sequence
+from typing import List
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
@@ -11,9 +11,7 @@
 logger = logging.getLogger(__name__)
 
 
-def view_to_reshape(
-    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
-) -> torch.fx.GraphModule:
+def view_to_reshape(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """Replace aten.view with an equivalent implementation which avoids Tensor memory issues"""
     orig_op = torch.ops.aten.view.default
     replacement_op = torch.ops.aten.reshape.default

tests/py/dynamo/models/test_models_export_kwargs.py

@@ -66,7 +66,7 @@ def forward(self, x, b=5, c=None, d=None):
     cos_sim = cosine_similarity(model(*args, **kwargs), trt_mod(*args, **kwargs)[0])
     assertions.assertTrue(
         cos_sim > COSINE_THRESHOLD,
-        msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        msg=f"CustomKwargs Module TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
     )
 
     # Clean up model env