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

@@ -538,7 +538,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,
@@ -568,7 +568,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,
@@ -982,12 +982,13 @@ def aten_ops_isinf(
 
 
 def conv_param_validator(conv_node: Node) -> bool:
+    # Output padding and transposed convolutions not supported currently
     return (not conv_node.args[6]) and (conv_node.args[7] in ([0], [0, 0], [0, 0, 0]))
 
 
 @dynamo_tensorrt_converter(
     torch.ops.aten.convolution.default, capability_validator=conv_param_validator
-)
+)  # type: ignore[misc]
 def aten_ops_convolution(
     network: TRTNetwork,
     target: Target,

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -1,14 +1,16 @@
 import functools
 import logging
 import re
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional, Tuple, Union
 
+import numpy as np
 import tensorrt as trt
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.fx.converters.converter_utils import (
     Frameworks,
     get_axes_for_reduce_op,
+    to_numpy,
     unified_dtype_converter,
 )
 from torch_tensorrt.fx.types import TRTDataType, TRTNetwork, TRTTensor
@@ -187,4 +189,76 @@ def extend_attr_to_tuple(
 
     if isinstance(val, list):
         val = tuple(val)
-    return val
+
+    if isinstance(val, tuple):
+        return val
+    else:
+        raise AssertionError(f"Could not extend attribute {val}")
+
+
+def create_constant(
+    network: TRTNetwork,
+    value: Union[int, float, np.ndarray, torch.Tensor],
+    name: str,
+    dtype: Optional[Union[torch.dtype, np.dtype, TRTDataType]],
+) -> TRTTensor:
+    """
+    Add a TensorRT constant layer whose value is `value` to `network`.
+    Args:
+        network (TRTNetwork): A TensorRT network to which we want to add
+            a constant layer.
+        value (Union[int, float, np.ndarray, torch.Tensor]): A literal value, Numpy array,
+            or a PyTorch tensor that will be used as value of the added TensorRT Constant layer.
+        name (str): Name of the added TensorRT Constant layer.
+        dtype (Optional[Union[torch.dtype, np.dtype, TRTDataType]]):
+            If a dtype is given, we will convert the type of the given `value` to this dtype.
+    Returns:
+        A TensorRT ITensor that represents the given value.
+    """
+    constant = network.add_constant(
+        (1,) if isinstance(value, (int, float)) else value.shape,
+        to_numpy(value, dtype).copy(),
+    )
+    constant.name = name
+    return constant.get_output(0)
+
+
+def get_trt_tensor(
+    network: TRTNetwork,
+    input_val: Any,
+    name: str,
+    dtype: Optional[Union[torch.dtype, np.dtype, TRTDataType]] = None,
+) -> TRTTensor:
+    """
+    Given a value of random type, we try to convert it to a TensorRT ITensor.
+    An runtime error is raised if we're not able to do that.
+    Args:
+        network (TRTNetwork): A TensorRT network. If we want to
+            add a TensorRT Constant layer, we will add it to this network.
+        input_val (Any): An value that we want to convert to a TensorRT ITensor.
+        name (str): The name of the created TensorRT Constant layer if there's
+            one.
+        dtype (Optional[Union[torch.dtype, np.dtype, TRTDataType]]):
+            If dtype is provided, the given value will be converted to this dtype.
+    Returns:
+        A TensorRT ITensor that represents the given value.
+    """
+    # TRT can not add constant for bool type. We do a work around to 1) cast it to int and 2)cast to bool later
+    # This is useful for logical operations which require input to be bool type
+    if isinstance(input_val, bool):
+        input_val = int(input_val)
+    elif isinstance(input_val, torch.Tensor) and (
+        input_val.dtype == torch.bool or input_val.dtype == torch.int64
+    ):
+        input_val = input_val.to(torch.int32)
+    elif isinstance(input_val, np.ndarray) and (
+        input_val.dtype == np.bool_ or input_val.dtype == np.int64
+    ):
+        input_val = input_val.astype(np.int32)
+
+    if isinstance(input_val, (torch.Tensor, np.ndarray, int, float)):
+        return create_constant(network, input_val, name, dtype)
+    elif isinstance(input_val, TRTTensor):
+        return input_val
+    else:
+        raise AssertionError(f"Cannot convert {input_val} to TRT constant")

py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py

@@ -1,19 +1,17 @@
 from typing import Optional
 
+import tensorrt as trt
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
-from torch_tensorrt.dynamo.conversion.converter_utils import broadcastable
-from torch_tensorrt.dynamo.conversion.impl.slice import expand
-from torch_tensorrt.fx.converters.converter_utils import (
-    broadcast,
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    broadcastable,
     get_trt_tensor,
-    set_layer_name,
 )
+from torch_tensorrt.dynamo.conversion.impl.slice import expand
+from torch_tensorrt.fx.converters.converter_utils import broadcast, set_layer_name
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
-import tensorrt as trt
-
 
 def where(
     network: TRTNetwork,

py/torch_tensorrt/dynamo/conversion/impl/conv.py

@@ -7,11 +7,13 @@
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo.conversion import impl
-from torch_tensorrt.dynamo.conversion.converter_utils import extend_attr_to_tuple
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    extend_attr_to_tuple,
+    get_trt_tensor,
+)
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     get_dyn_range,
-    get_trt_tensor,
     has_dynamic_shape,
     mark_as_int8_layer,
     set_layer_name,
@@ -27,8 +29,8 @@ def convNd(
     name: str,
     is_conv1d: bool,
     input: TRTTensor,
-    weight: Union[TRTTensor, torch.Tensor],
-    bias: Optional[Union[TRTTensor, torch.Tensor]],
+    weight: Union[TRTTensor, torch.Tensor, np.ndarray],
+    bias: Optional[Union[TRTTensor, torch.Tensor, np.ndarray]],
     stride: Optional[Union[int, Sequence[int]]],
     padding: Optional[Union[int, Sequence[int]]],
     dilation: Optional[Union[int, Sequence[int]]],
@@ -97,19 +99,28 @@ def convNd(
     if isinstance(bias, TRTTensor):
         conv_layer.set_input(2, bias)
 
+    # Cast certain fields to tuples, in accordance with TRT requirements
+    padding = (padding,) if isinstance(padding, int) else padding
+    stride = (stride,) if isinstance(stride, int) else stride
+    dilation = (dilation,) if isinstance(dilation, int) else dilation
+
     # Expand parameters manually for Conv1D computations
     if is_conv1d:
-        padding = tuple(padding) + (0,)
-        stride = extend_attr_to_tuple(stride, 2)
-        dilation = extend_attr_to_tuple(dilation, 2)
+        padding = (tuple(padding) + (0,)) if padding is not None else padding
+        stride = extend_attr_to_tuple(stride, 2) if stride is not None else stride
+        dilation = (
+            extend_attr_to_tuple(dilation, 2) if dilation is not None else dilation
+        )
 
     set_layer_name(conv_layer, target, name, source_ir)
 
     # Set relevant attributes of convolution layer
-    conv_layer.padding_nd = padding
-    conv_layer.stride_nd = stride
-    conv_layer.dilation_nd = dilation
-
+    if padding is not None:
+        conv_layer.padding_nd = padding
+    if stride is not None:
+        conv_layer.stride_nd = stride
+    if dilation is not None:
+        conv_layer.dilation_nd = dilation
     if groups is not None:
         conv_layer.num_groups = groups
 

py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py

@@ -6,10 +6,12 @@
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
-from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor
+from torch_tensorrt.dynamo.conversion.converter_utils import (
+    cast_trt_tensor,
+    get_trt_tensor,
+)
 from torch_tensorrt.fx.converters.converter_utils import (
     broadcast,
-    get_trt_tensor,
     set_layer_name,
     squeeze_left,
 )

py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py

@@ -4,16 +4,13 @@
 import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
 from torch_tensorrt.dynamo.conversion.impl.elementwise.base import (
     convert_binary_elementwise,
 )
 from torch_tensorrt.dynamo.conversion.impl.unary import sign
 from torch_tensorrt.dynamo.conversion.impl.unary.base import convert_unary
-from torch_tensorrt.fx.converters.converter_utils import (
-    get_trt_tensor,
-    set_layer_name,
-    squeeze_left,
-)
+from torch_tensorrt.fx.converters.converter_utils import set_layer_name, squeeze_left
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 

py/torch_tensorrt/dynamo/conversion/impl/embedding.py

@@ -3,7 +3,8 @@
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
-from torch_tensorrt.fx.converters.converter_utils import get_trt_tensor, set_layer_name
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
+from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
 

py/torch_tensorrt/dynamo/conversion/impl/matmul.py

@@ -3,11 +3,8 @@
 import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
-from torch_tensorrt.fx.converters.converter_utils import (
-    broadcast,
-    get_trt_tensor,
-    set_layer_name,
-)
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
+from torch_tensorrt.fx.converters.converter_utils import broadcast, set_layer_name
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 

py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py

@@ -2,9 +2,9 @@
 
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.converter_utils import get_trt_tensor
 from torch_tensorrt.fx.converters.converter_utils import (
     get_positive_dim,
-    get_trt_tensor,
     set_layer_name,
 )
 from torch_tensorrt.fx.types import Shape, TRTNetwork, TRTTensor

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -2711,14 +2711,8 @@ def acc_ops_linear(
         "dim for linear and it can't be the last dim."
     )
 
-    if isinstance(kwargs["weight"], (torch.Tensor, np.ndarray)):
-        weight = get_trt_tensor(
-            network,
-            kwargs["weight"].t()
-            if isinstance(kwargs["weight"], torch.Tensor)
-            else kwargs["weight"].T,
-            f"{name}_weight",
-        )
+    if isinstance(kwargs["weight"], torch.Tensor):
+        weight = get_trt_tensor(network, kwargs["weight"].t(), f"{name}_weight")
         if target not in (acc_ops.linear, torch.ops.aten.linear):
             weight_op = trt.MatrixOperation.TRANSPOSE
         else:

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -271,7 +271,7 @@ def create_constant(
     """
     constant = network.add_constant(
         (1,) if isinstance(value, (int, float)) else value.shape,
-        to_numpy(value, dtype).copy(),
+        to_numpy(value, dtype),
     )
     constant.name = name
     return constant.get_output(0)
@@ -311,7 +311,7 @@ def get_trt_tensor(
     elif isinstance(input_val, np.ndarray) and (
         input_val.dtype == np.bool_ or input_val.dtype == np.int64
     ):
-        input_val = input_val.astype(np.int32)
+        input_val = input_val.to(np.int32)
 
     if isinstance(input_val, (torch.Tensor, np.ndarray, int, float)):
         return create_constant(network, input_val, name, dtype)

py/torch_tensorrt/fx/converters/impl/convolution.py

@@ -50,7 +50,7 @@ def convNd(
         )
 
     # Process bias terms
-    if isinstance(bias, (torch.Tensor, np.ndarray)):
+    if isinstance(bias, torch.Tensor):
         # Transform the bias constant into a Numpy array
         bias = to_numpy(bias)
 
@@ -75,7 +75,7 @@ def convNd(
                 network, target, tuple(), kwargs, name + "_unsqueeze_weight"
             )
 
-    elif isinstance(weight, (torch.Tensor, np.ndarray)):
+    elif isinstance(weight, torch.Tensor):
         # Transform the weight constant into a Numpy array
         weight = to_numpy(weight)