Move fixes into Dynamo directory

Author: gs-olive

py/torch_tensorrt/dynamo/backend/backends.py

@@ -9,7 +9,7 @@
 import torch.utils._pytree as pytree
 from torch._dynamo.utils import detect_fake_mode
 from torch._functorch.aot_autograd import _aot_export_function
-from torch._inductor.freezing import ConstantFolder, replace_node_with_constant
+from torch._inductor.constant_folding import ConstantFolder, replace_node_with_constant
 from torch._ops import OpOverload
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.compile import compile_module
@@ -100,7 +100,7 @@ def _pretraced_backend(
                 + "Returning GraphModule forward instead.",
                 exc_info=True,
             )
-            return gm.forward
+            return gm
         else:
             logger.critical(
                 "Halting compilation on build failure since "
@@ -114,6 +114,13 @@ def _pretraced_backend(
 
 @torch.utils._python_dispatch._disable_current_modes()  # type: ignore
 def constant_fold(gm: torch.fx.GraphModule) -> Any:
+    """Adapted from:
+    https://github.com/pytorch/pytorch/blob/3a79621c9dce17f77fbddc06aab21f6bc477f313/torch/_inductor/freezing.py#L178-L197
+
+    Folds constants in the graph module, not skipping constructors
+
+    Modifies the graph in-place and replaces node with constants
+    """
     cf = ConstantFolder(gm, skip_constructors=False)
     cf.run()
 
@@ -141,10 +148,13 @@ def aot_export_for_compile(
     decompositions: Optional[Dict[OpOverload, Callable[[Any], Any]]] = None,
 ) -> torch.fx.GraphModule:
     """Adapted from:
-    https://github.com/pytorch/pytorch/blob/054f3f1d8f9eb63ef8437991eba5b8f2aeee920f/torch/_functorch/aot_autograd.py#L4133-L4134
+    https://github.com/pytorch/pytorch/blob/1a5fdc2458b98697c75c32eb6f4b8b34d76429cf/torch/_functorch/aot_autograd.py#L4084-L4158
 
     Removed check for input aliasing in resultant subgraph - TRT is functional-only
+
+    Exports the function to ATen for torch compile
     """
+    # Trace function with input arguments and decompositions
     with torch.no_grad():
         fx_g, metadata, in_spec, out_spec = _aot_export_function(
             func,

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -361,7 +361,7 @@ def output(self, target: str, args: Any, kwargs: Any) -> List[Any]:
             outputs = (args[0],)
 
         for output_idx in range(len(outputs)):
-            from torch_tensorrt.fx.converters import get_trt_tensor
+            from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
 
             output = outputs[output_idx]
 

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -94,7 +94,7 @@ def aten_ops_fmod(
     return impl.elementwise.fmod(network, target, SourceIR.ATEN, name, args[0], args[1])
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.relu.default)
+@dynamo_tensorrt_converter(torch.ops.aten.relu.default)  # type: ignore[misc]
 def aten_ops_relu(
     network: TRTNetwork,
     target: Target,
@@ -111,7 +111,7 @@ def aten_ops_relu(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.sigmoid.default)
+@dynamo_tensorrt_converter(torch.ops.aten.sigmoid.default)  # type: ignore[misc]
 def aten_ops_sigmoid(
     network: TRTNetwork,
     target: Target,
@@ -128,7 +128,7 @@ def aten_ops_sigmoid(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.tanh.default)
+@dynamo_tensorrt_converter(torch.ops.aten.tanh.default)  # type: ignore[misc]
 def aten_ops_tanh(
     network: TRTNetwork,
     target: Target,
@@ -145,7 +145,7 @@ def aten_ops_tanh(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
+@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)  # type: ignore[misc]
 def aten_ops_leaky_relu(
     network: TRTNetwork,
     target: Target,
@@ -163,7 +163,7 @@ def aten_ops_leaky_relu(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.elu.default)
+@dynamo_tensorrt_converter(torch.ops.aten.elu.default)  # type: ignore[misc]
 def aten_ops_elu(
     network: TRTNetwork,
     target: Target,
@@ -182,7 +182,7 @@ def aten_ops_elu(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.softplus.default)
+@dynamo_tensorrt_converter(torch.ops.aten.softplus.default)  # type: ignore[misc]
 def aten_ops_softplus(
     network: TRTNetwork,
     target: Target,
@@ -200,7 +200,7 @@ def aten_ops_softplus(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.clip.default)
+@dynamo_tensorrt_converter(torch.ops.aten.clip.default)  # type: ignore[misc]
 def aten_ops_clip(
     network: TRTNetwork,
     target: Target,
@@ -219,7 +219,7 @@ def aten_ops_clip(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.hardsigmoid.default)
+@dynamo_tensorrt_converter(torch.ops.aten.hardsigmoid.default)  # type: ignore[misc]
 def aten_ops_hard_sigmoid(
     network: TRTNetwork,
     target: Target,
@@ -296,26 +296,20 @@ def aten_ops_rsqrt(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.neg.default)
+@dynamo_tensorrt_converter(torch.ops.aten.neg.default)  # type: ignore[misc]
 def aten_ops_neg(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = args[0]
-    if (isinstance(input_val, TRTTensor)) and (
-        input_val.dtype == trt.int8 or input_val.dtype == trt.int32
-    ):
-        input_val = cast_trt_tensor(network, input_val, trt.float32, name)
-
     return impl.unary.neg(
         network,
         target,
         SourceIR.ATEN,
         name,
-        input_val,
+        args[0],
     )
 
 
@@ -503,7 +497,7 @@ def aten_ops_clone(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.expand.default)
+@dynamo_tensorrt_converter(torch.ops.aten.expand.default)  # type: ignore[misc]
 def aten_ops_expand(
     network: TRTNetwork,
     target: Target,
@@ -533,7 +527,7 @@ def amax_param_validator(amax_node: Node) -> bool:
 
 @dynamo_tensorrt_converter(
     torch.ops.aten.amax.default, capability_validator=amax_param_validator
-)
+)  # type: ignore[misc]
 def aten_ops_amax(
     network: TRTNetwork,
     target: Target,
@@ -552,8 +546,8 @@ def aten_ops_amax(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.sum.default)
-@dynamo_tensorrt_converter(torch.ops.aten.sum.dim_IntList)
+@dynamo_tensorrt_converter(torch.ops.aten.sum.default)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.sum.dim_IntList)  # type: ignore[misc]
 def aten_ops_sum(
     network: TRTNetwork,
     target: Target,
@@ -946,8 +940,8 @@ def aten_ops_isinf(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.add.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.add.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.add.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.add.Scalar)  # type: ignore[misc]
 def aten_ops_add(
     network: TRTNetwork,
     target: Target,
@@ -978,8 +972,8 @@ def aten_ops_add(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.mul.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.mul.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.mul.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.mul.Scalar)  # type: ignore[misc]
 def aten_ops_mul(
     network: TRTNetwork,
     target: Target,
@@ -997,7 +991,7 @@ def aten_ops_mul(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.maximum.default)
+@dynamo_tensorrt_converter(torch.ops.aten.maximum.default)  # type: ignore[misc]
 def aten_ops_max(
     network: TRTNetwork,
     target: Target,
@@ -1015,7 +1009,7 @@ def aten_ops_max(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.minimum.default)
+@dynamo_tensorrt_converter(torch.ops.aten.minimum.default)  # type: ignore[misc]
 def aten_ops_min(
     network: TRTNetwork,
     target: Target,
@@ -1033,8 +1027,8 @@ def aten_ops_min(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.sub.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.sub.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.sub.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.sub.Scalar)  # type: ignore[misc]
 def aten_ops_sub(
     network: TRTNetwork,
     target: Target,
@@ -1065,10 +1059,10 @@ def aten_ops_sub(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.div.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.div.Tensor_mode)
-@dynamo_tensorrt_converter(torch.ops.aten.div.Scalar)
-@dynamo_tensorrt_converter(torch.ops.aten.div.Scalar_mode)
+@dynamo_tensorrt_converter(torch.ops.aten.div.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.div.Tensor_mode)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.div.Scalar)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.div.Scalar_mode)  # type: ignore[misc]
 def aten_ops_div(
     network: TRTNetwork,
     target: Target,
@@ -1111,9 +1105,9 @@ def aten_ops_div(
         )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.pow.Tensor_Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.pow.Scalar)
-@dynamo_tensorrt_converter(torch.ops.aten.pow.Tensor_Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.pow.Tensor_Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.pow.Scalar)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.pow.Tensor_Scalar)  # type: ignore[misc]
 def aten_ops_pow(
     network: TRTNetwork,
     target: Target,
@@ -1131,8 +1125,8 @@ def aten_ops_pow(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.floor_divide.default)
-@dynamo_tensorrt_converter(torch.ops.aten.floor_divide.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.floor_divide.default)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.floor_divide.Scalar)  # type: ignore[misc]
 def aten_ops_floor_div(
     network: TRTNetwork,
     target: Target,
@@ -1150,7 +1144,7 @@ def aten_ops_floor_div(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.logical_and.default)
+@dynamo_tensorrt_converter(torch.ops.aten.logical_and.default)  # type: ignore[misc]
 def aten_ops_logical_and(
     network: TRTNetwork,
     target: Target,
@@ -1168,7 +1162,7 @@ def aten_ops_logical_and(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.logical_or.default)
+@dynamo_tensorrt_converter(torch.ops.aten.logical_or.default)  # type: ignore[misc]
 def aten_ops_logical_or(
     network: TRTNetwork,
     target: Target,
@@ -1186,7 +1180,7 @@ def aten_ops_logical_or(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.logical_xor.default)
+@dynamo_tensorrt_converter(torch.ops.aten.logical_xor.default)  # type: ignore[misc]
 def aten_ops_logical_xor(
     network: TRTNetwork,
     target: Target,
@@ -1204,8 +1198,8 @@ def aten_ops_logical_xor(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.eq.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.eq.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.eq.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.eq.Scalar)  # type: ignore[misc]
 def aten_ops_equal(
     network: TRTNetwork,
     target: Target,
@@ -1223,8 +1217,8 @@ def aten_ops_equal(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.gt.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.gt.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.gt.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.gt.Scalar)  # type: ignore[misc]
 def aten_ops_greater(
     network: TRTNetwork,
     target: Target,
@@ -1242,8 +1236,8 @@ def aten_ops_greater(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.lt.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.lt.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.lt.Tensor)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.lt.Scalar)  # type: ignore[misc]
 def aten_ops_less(
     network: TRTNetwork,
     target: Target,
@@ -1267,7 +1261,7 @@ def conv_param_validator(conv_node: Node) -> bool:
 
 @dynamo_tensorrt_converter(
     torch.ops.aten.convolution.default, capability_validator=conv_param_validator
-)
+)  # type: ignore[misc]
 def aten_ops_convolution(
     network: TRTNetwork,
     target: Target,
@@ -1291,7 +1285,8 @@ def aten_ops_convolution(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.linear.default)
+@dynamo_tensorrt_converter(torch.ops.aten.linear.default)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.linear)  # type: ignore[misc]
 def aten_ops_linear(
     network: TRTNetwork,
     target: Target,