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/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
@@ -185,11 +187,85 @@ 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 cast_int_or_float_to_bool(network: TRTNetwork, name: str, tensor: TRTTensor):
+def cast_int_or_float_to_bool(
+    network: TRTNetwork, name: str, tensor: TRTTensor
+) -> TRTTensor:
     if tensor.dtype != trt.bool:
         return cast_trt_tensor(network, tensor, trt.bool, name)
 
     return tensor
+
+
+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

@@ -7,10 +7,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

@@ -7,17 +7,14 @@
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     cast_int_int_div_trt_tensor,
     cast_int_or_float_to_bool,
+    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/linear.py

@@ -5,7 +5,7 @@
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo.conversion import impl
-from torch_tensorrt.fx.converters.converter_utils import SourceIR, get_trt_tensor
+from torch_tensorrt.dynamo.conversion.converter_utils import SourceIR, get_trt_tensor
 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/dynamo/lowering/_pre_aot_lowering.py

@@ -124,7 +124,6 @@ def pre_aot_substitutions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
 
                 # Replace all original node uses and clean up graph
                 n.replace_all_uses_with(new_node)
-                gm.graph.eliminate_dead_code()
                 gm.graph.lint()
                 gm.recompile()
 
@@ -138,7 +137,6 @@ def pre_aot_substitutions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
                 continue
 
     # Perform cleanup and recompilation before returning module
-    gm.graph.eliminate_dead_code()
     gm.graph.lint()
     gm.recompile()