fix: Refactor data type handling in FX

Author: gs-olive

py/torch_tensorrt/dynamo/backend/test/test_specialized_models.py

@@ -54,6 +54,7 @@ def forward(self, x):
             0,
             msg=f"MulInt TRT outputs don't match with the original model.",
         )
+        torch._dynamo.reset()
 
     def test_lowering_add_float(self):
         class AddFloat(torch.nn.Module):
@@ -106,6 +107,8 @@ def forward(self, x):
             msg=f"AddFloat TRT outputs don't match with the original model.",
         )
 
+        torch._dynamo.reset()
+
 
 if __name__ == "__main__":
     run_tests()

py/torch_tensorrt/dynamo/fx_ts_compat/fx2trt.py

@@ -16,7 +16,12 @@
 from torch_tensorrt.dynamo.fx_ts_compat import CONVERTERS
 from .input_tensor_spec import InputTensorSpec
 from torch_tensorrt.fx.observer import Observer
-from torch_tensorrt.fx.utils import get_dynamic_dims, LowerPrecision, torch_dtype_to_trt
+from torch_tensorrt.fx.utils import (
+    get_dynamic_dims,
+    LowerPrecision,
+    unified_dtype_converter,
+    Frameworks,
+)
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -305,7 +310,9 @@ def placeholder(self, target, args, kwargs):
                 self.optimization_profiles[i].set_shape(target, *shape_range)
 
         return self.network.add_input(
-            name=target, shape=tuple(shape), dtype=torch_dtype_to_trt(dtype)
+            name=target,
+            shape=tuple(shape),
+            dtype=unified_dtype_converter(dtype, Frameworks.TRT),
         )
 
     def call_module(self, target, args, kwargs):

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -18,7 +18,7 @@
 from torch.fx.immutable_collections import immutable_list
 from torch.fx.node import Argument, Target
 
-from ..utils import get_dynamic_dims, torch_dtype_from_trt, torch_dtype_to_trt
+from ..utils import get_dynamic_dims, unified_dtype_converter, Frameworks
 
 from .converter_utils import *  # noqa: F403
 from torch_tensorrt.fx.passes.lower_basic_pass import (
@@ -400,7 +400,7 @@ def acc_ops_pad_with_slice_layer(
         )
 
     # cast value to TRTensor
-    dt = torch_dtype_from_trt(input_val.dtype)
+    dt = unified_dtype_converter(input_val.dtype, Frameworks.TORCH)
     value = 0 if value == None else value
     value_const = get_trt_tensor(
         network, torch.tensor([value], dtype=dt), f"{name}_value"
@@ -1561,7 +1561,7 @@ def acc_ops_to_dtype(
     input_t = get_trt_tensor(network, input_val, f"{name}_input_t")
     if input_dtype:
         if isinstance(input_dtype, torch.dtype):
-            input_dtype = torch_dtype_to_trt(input_dtype)
+            input_dtype = unified_dtype_converter(input_dtype, Frameworks.TRT)
         input_t = type_cast(network, target, f"{name}_input", input_t, input_dtype)
     return input_t
 
@@ -1822,7 +1822,7 @@ def acc_ops_logical_xor(
 #             f"isinf received input {input_t} that is not part "
 #             "of the TensorRT region!"
 #         )
-#     tdtype = torch_dtype_from_trt(input_t.dtype)
+#     tdtype = unified_dtype_converter(input_t.dtype, Frameworks.TORCH)
 
 #     inf_t = torch.ones(tuple(input_t.shape))
 #     inf_t = inf_t * float("inf")
@@ -1860,7 +1860,7 @@ def acc_ops_any(
 
     if input_t.dtype in (trt.float32, trt.float16, trt.int32):
         comp_t = torch.zeros(tuple([*input_t.shape])).to(
-            torch_dtype_from_trt(input_t.dtype)
+            unified_dtype_converter(input_t.dtype, Frameworks.TORCH)
         )
         comp_t = get_trt_tensor(network, comp_t, f"{name}_comp_t")
         kwargs_new = {"input": input_t, "other": comp_t}
@@ -2749,7 +2749,7 @@ def acc_ops_masked_fill_tensor(
     if type(value_t) is torch.Tensor:
         value_t = value_t.cpu().numpy()
     # cast to input type
-    input_dtype = torch_dtype_from_trt(input_t.dtype)
+    input_dtype = unified_dtype_converter(input_t.dtype, Frameworks.TORCH)
     value_t = (torch.ones(shape) * value_t).to(input_dtype)
     input_val = get_trt_tensor(network, input_t, f"{name}_input")
     value_val = get_trt_tensor(network, value_t, f"{name}_input")
@@ -2883,7 +2883,11 @@ def add_clamp(network, input, val, op, name):
         # clamping scalar
         acc_ops_clamp_trt = get_trt_tensor(
             network,
-            squeeze_left(torch.tensor([val], dtype=torch_dtype_from_trt(input.dtype))),
+            squeeze_left(
+                torch.tensor(
+                    [val], dtype=unified_dtype_converter(input.dtype, Frameworks.TORCH)
+                )
+            ),
             f"{name}_clamp_{val}",
         )
     else:
@@ -2892,7 +2896,8 @@ def add_clamp(network, input, val, op, name):
             (
                 val
                 * torch.ones(
-                    acc_ops_clamp_shape, dtype=torch_dtype_from_trt(input.dtype)
+                    acc_ops_clamp_shape,
+                    dtype=unified_dtype_converter(input.dtype, Frameworks.TORCH),
                 )
             )
             .cpu()
@@ -3538,7 +3543,9 @@ def acc_ops_cumsum(
     iterator = loop.add_iterator(input_val, dim, False)
     data = iterator.get_output(0)
     new_dims = tuple(data.shape)
-    zero_tensor = torch.zeros(new_dims, dtype=trt_dtype_to_torch_dtype(input_val.dtype))
+    zero_tensor = torch.zeros(
+        new_dims, dtype=unified_dtype_converter(input_val.dtype, Frameworks.TORCH)
+    )
     zero_tensor = network.add_constant(
         zero_tensor.shape, to_numpy(zero_tensor)
     ).get_output(0)
@@ -3689,7 +3696,7 @@ def acc_ops_new_ones(
     dtype_val = kwargs.get("dtype")
     if dtype_val is None:
         dtype_val = input_val.dtype
-        dtype_val = torch_dtype_from_trt(dtype_val)
+        dtype_val = unified_dtype_converter(dtype_val, Frameworks.TORCH)
 
     device_val = kwargs.get("device")
     assert (
@@ -3713,7 +3720,7 @@ def acc_ops_new_empty(
     dtype_val = kwargs.get("dtype")
     if dtype_val is None:
         dtype_val = input_val.dtype
-        dtype_val = torch_dtype_from_trt(dtype_val)
+        dtype_val = unified_dtype_converter(dtype_val, Frameworks.TORCH)
 
     device_val = kwargs.get("device")
     assert (

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -18,8 +18,6 @@
 from torch.fx.immutable_collections import immutable_list
 from torch.fx.node import Argument, Target
 
-from ..utils import get_dynamic_dims, torch_dtype_from_trt, torch_dtype_to_trt
-
 from .converter_utils import *  # noqa: F403
 import torch_tensorrt.fx.tracer.acc_tracer.acc_utils as acc_utils
 from torch_tensorrt.fx.converters.impl import activation

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -20,7 +20,7 @@
     TRTPluginFieldCollection,
     TRTTensor,
 )
-from ..utils import torch_dtype_from_trt
+from ..utils import unified_dtype_converter, Frameworks
 
 
 class SourceIR(Enum):
@@ -151,38 +151,49 @@ def extend_mod_attr_to_tuple(mod: torch.nn.Module, name: str, size: int):
     return extend_attr_to_tuple(val, size)
 
 
-def to_numpy(value: Optional[Union[torch.Tensor, int, float]]) -> Optional[np.ndarray]:
+def to_numpy(
+    value: Optional[Union[torch.Tensor, np.ndarray, int, float]],
+    dtype: Optional[Union[torch.dtype, np.dtype, TRTDataType]] = None,
+) -> Optional[np.ndarray]:
     """
     Convert a PyTorch Tensor to a Numpy Array. If the tensor is
     quantized it will be dequantized first.
 
     Args:
-        value (Optional[Union[torch.Tensor, int, float]]): A PyTorch tensor, int, or float
+        value (Optional[Union[torch.Tensor, np.ndarray, int, float]]):
+            A PyTorch tensor, Numpy array, int, or float
 
     Returns:
         A Numpy array.
     """
+    output = None
 
-    if value is None:
-        return value
+    if value is None or isinstance(value, np.ndarray):
+        output = value
 
     elif isinstance(value, torch.Tensor):
         if value.is_quantized:
             value = value.dequantize()
 
-        return value.cpu().detach().contiguous().numpy()
+        output = value.cpu().detach().contiguous().numpy()
 
     elif isinstance(value, int):
-        return np.array([value], dtype=np.int32)
+        output = np.array([value], dtype=np.int32)
 
     elif isinstance(value, float):
-        return np.array([value], dtype=np.float32)
+        output = np.array([value], dtype=np.float32)
 
     else:
         raise AssertionError(
-            f"to_numpy can only be called on None, int, float, or torch.Tensor, got: {value}"
+            f"to_numpy can only be called on None, int, float, np.ndarray, or torch.Tensor, got: {value}"
         )
 
+    return (
+        output
+        if dtype is None
+        else output.astype(unified_dtype_converter(dtype, Frameworks.NUMPY))
+    )
+
 
 def has_dynamic_shape(shape: Shape) -> bool:
     """
@@ -234,35 +245,35 @@ def get_axes_for_reduce_op(
 
 def create_constant(
     network: TRTNetwork,
-    value: Union[int, float, torch.Tensor],
+    value: Union[int, float, np.ndarray, torch.Tensor],
     name: str,
-    dtype: Optional[torch.dtype],
+    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, torch.Tensor]): A literal value or a PyTorch tensor
-            that will be used as value of the added TensorRT 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[torch.dtype]): If a dtype is given, we will convert the type
-            of the given `value` to this dtype.
+        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.
     """
-
-    if dtype:
-        value = value.to(dtype)
-    constant = network.add_constant(value.shape, to_numpy(value))
+    constant = network.add_constant(value.shape, to_numpy(value, dtype))
     constant.name = name
     return constant.get_output(0)
 
 
 def get_trt_tensor(
-    network: TRTNetwork, input_val: Any, name: str, dtype: Optional[torch.dtype] = None
+    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.
@@ -274,33 +285,36 @@ def get_trt_tensor(
         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[torch.dtype]): If dtype is provided, the given value
-            will be converted to this dtype.
+        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, np.ndarray):
-        input_val = torch.from_numpy(input_val)
     if isinstance(input_val, bool):
         input_val = int(input_val)
-    if isinstance(input_val, torch.Tensor) and input_val.dtype == torch.bool:
-        input_val = input_val.to(torch.int32)
-    if isinstance(input_val, torch.Tensor) and input_val.dtype == torch.int64:
+
+    if 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.to(np.int32)
 
-    if isinstance(input_val, (torch.Tensor, int, float)):
+    if isinstance(input_val, (torch.Tensor, np.ndarray, int, float)):
         return create_constant(network, input_val, name, dtype)
-    elif not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"Received input {input_val} of name {name} that "
-            "is not part of the TensorRT region!"
-        )
-    else:
+    elif isinstance(input_val, TRTTensor):
         return input_val
 
+    raise RuntimeError(
+        f"Received input {input_val} of name {name} that "
+        "is not part of the TensorRT region!"
+    )
+
 
 def prepend_ones(
     network: TRTNetwork,
@@ -482,10 +496,10 @@ def add_binary_elementwise_layer(
     is_rhs_trt_tensor = False
 
     if isinstance(lhs_val, TRTTensor):
-        lhs_dtype = torch_dtype_from_trt(lhs_val.dtype)
+        lhs_dtype = unified_dtype_converter(lhs_val.dtype, Frameworks.TORCH)
         is_lhs_trt_tensor = True
     if isinstance(rhs_val, TRTTensor):
-        rhs_dtype = torch_dtype_from_trt(rhs_val.dtype)
+        rhs_dtype = unified_dtype_converter(rhs_val.dtype, Frameworks.TORCH)
         is_rhs_trt_tensor = True
 
     if not is_lhs_trt_tensor and not is_rhs_trt_tensor:
@@ -510,9 +524,13 @@ def add_binary_elementwise_layer(
     # dtype but we don't have a way to detect whether it makes sense for the
     # scalar to be float or half. Hence we go with the lhs dtype.
     if is_lhs_trt_tensor and isinstance(rhs_val, (float, int)):
-        rhs_val = torch.tensor([rhs_val], dtype=lhs_dtype)
+        rhs_val = np.array(
+            [rhs_val], dtype=unified_dtype_converter(lhs_val.dtype, Frameworks.NUMPY)
+        )
     if is_rhs_trt_tensor and isinstance(lhs_val, (float, int)):
-        lhs_val = torch.tensor([lhs_val], dtype=rhs_dtype)
+        lhs_val = np.array(
+            [lhs_val], dtype=unified_dtype_converter(rhs_val.dtype, Frameworks.NUMPY)
+        )
 
     # When lhs is scalar, and rhs has shape [1,], then currently the assert
     # will fail because lhs shape has fewer dimensions than rhs shape.  This
@@ -552,14 +570,19 @@ def add_binary_elementwise_layer(
     return output
 
 
-def squeeze_left(const: torch.Tensor):
+def squeeze_left(const: Union[torch.Tensor, np.ndarray]):
     """
     Squeeze the size-1 dimensions on the left side of the shape tuple.
     PyTorch's `squeeze()` doesn't support passing multiple `dim`s at once, so
     we do it iteratively.
     """
     while len(const.shape) > 0 and const.shape[0] == 1:
-        const = const.squeeze(dim=0)
+        if isinstance(const, torch.Tensor):
+            const = const.squeeze(dim=0)
+        elif isinstance(const, np.ndarray):
+            const = const.squeeze(axis=0)
+        else:
+            raise AssertionError(f"Expected torch Tensor or Numpy array, got: {const}")
     return const
 
 
@@ -786,7 +809,10 @@ def trunc_div(
         input = get_trt_tensor(network, input, f"{name}_input")
     if not isinstance(other, trt.tensorrt.ITensor):
         other = get_trt_tensor(
-            network, other, f"{name}_other", dtype=torch_dtype_from_trt(input.dtype)
+            network,
+            other,
+            f"{name}_other",
+            dtype=unified_dtype_converter(input.dtype, Frameworks.TORCH),
         )
 
     abs_input_output = add_unary_layer(
@@ -875,13 +901,3 @@ def type_cast(
     layer_i.set_output_type(0, cast_type)
     set_layer_name(layer_i, target, f"{name}_dtype_change")
     return layer_i.get_output(0)
-
-
-def trt_dtype_to_torch_dtype(trt_dtype):
-    table = {
-        trt.bool: torch.bool,
-        trt.int32: torch.int32,
-        trt.float16: torch.float16,
-        trt.float32: torch.float32,
-    }
-    return table[trt_dtype]