update type conversion

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -124,9 +124,7 @@ def cast_int_int_div_trt_tensor(
     Returns:
         A list of lhs_val and rhs_val casted to the approriate datatype
     """
-    if (lhs_val.dtype == trt.int8 or lhs_val.dtype == trt.int32) and (
-        rhs_val.dtype == trt.int8 or rhs_val.dtype == trt.int32
-    ):
+    if lhs_val.dtype == trt.int32 and rhs_val.dtype == trt.int32:
         lhs_val = cast_trt_tensor(network, lhs_val, trt.float32, name)
         rhs_val = cast_trt_tensor(network, rhs_val, trt.float32, name)
     return [lhs_val, rhs_val]
@@ -190,14 +188,7 @@ def extend_attr_to_tuple(
     return val
 
 
-def trt_cast_int_to_float(network: TRTNetwork, name: str, tensor: TRTTensor):
-    if tensor.dtype == trt.int8 or tensor.dtype == trt.int32:
-        return cast_trt_tensor(network, tensor, trt.float32, name)
-
-    return tensor
-
-
-def trt_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):
     if tensor.dtype != trt.bool:
         return cast_trt_tensor(network, tensor, trt.bool, name)
 

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

@@ -5,8 +5,8 @@
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion.converter_utils import (
-    trt_cast_int_or_float_to_bool,
-    trt_cast_int_to_float,
+    cast_int_int_div_trt_tensor,
+    cast_int_or_float_to_bool,
 )
 from torch_tensorrt.dynamo.conversion.impl.elementwise.base import (
     convert_binary_elementwise,
@@ -324,11 +324,8 @@ def div(
     lhs_val: Union[TRTTensor, int, float],
     rhs_val: Union[TRTTensor, int, float],
 ) -> TRTTensor:
-    if isinstance(lhs_val, TRTTensor):
-        lhs_val = trt_cast_int_to_float(network, name, lhs_val)
-
-    if isinstance(rhs_val, TRTTensor):
-        rhs_val = trt_cast_int_to_float(network, name, rhs_val)
+    if isinstance(lhs_val, TRTTensor) and isinstance(rhs_val, TRTTensor):
+        lhs_val, rhs_val = cast_int_int_div_trt_tensor(network, lhs_val, rhs_val, name)
 
     return convert_binary_elementwise(
         network, target, source_ir, name, trt.ElementWiseOperation.DIV, lhs_val, rhs_val
@@ -343,11 +340,8 @@ def pow(
     lhs_val: Union[TRTTensor, int, float],
     rhs_val: Union[TRTTensor, int, float],
 ) -> TRTTensor:
-    if isinstance(lhs_val, TRTTensor):
-        lhs_val = trt_cast_int_to_float(network, name, lhs_val)
-
-    if isinstance(rhs_val, TRTTensor):
-        rhs_val = trt_cast_int_to_float(network, name, rhs_val)
+    if isinstance(lhs_val, TRTTensor) and isinstance(rhs_val, TRTTensor):
+        lhs_val, rhs_val = cast_int_int_div_trt_tensor(network, lhs_val, rhs_val, name)
 
     return convert_binary_elementwise(
         network, target, source_ir, name, trt.ElementWiseOperation.POW, lhs_val, rhs_val
@@ -382,10 +376,10 @@ def logical_and(
     rhs_val: Union[TRTTensor, int, float],
 ) -> TRTTensor:
     if isinstance(lhs_val, TRTTensor):
-        lhs_val = trt_cast_int_or_float_to_bool(network, name, lhs_val)
+        lhs_val = cast_int_or_float_to_bool(network, name, lhs_val)
 
     if isinstance(rhs_val, TRTTensor):
-        rhs_val = trt_cast_int_or_float_to_bool(network, name, rhs_val)
+        rhs_val = cast_int_or_float_to_bool(network, name, rhs_val)
 
     return convert_binary_elementwise(
         network, target, source_ir, name, trt.ElementWiseOperation.AND, lhs_val, rhs_val
@@ -401,10 +395,10 @@ def logical_or(
     rhs_val: Union[TRTTensor, int, float],
 ) -> TRTTensor:
     if isinstance(lhs_val, TRTTensor):
-        lhs_val = trt_cast_int_or_float_to_bool(network, name, lhs_val)
+        lhs_val = cast_int_or_float_to_bool(network, name, lhs_val)
 
     if isinstance(rhs_val, TRTTensor):
-        rhs_val = trt_cast_int_or_float_to_bool(network, name, rhs_val)
+        rhs_val = cast_int_or_float_to_bool(network, name, rhs_val)
 
     return convert_binary_elementwise(
         network, target, source_ir, name, trt.ElementWiseOperation.OR, lhs_val, rhs_val
@@ -420,10 +414,10 @@ def logical_xor(
     rhs_val: Union[TRTTensor, int, float],
 ) -> TRTTensor:
     if isinstance(lhs_val, TRTTensor):
-        lhs_val = trt_cast_int_or_float_to_bool(network, name, lhs_val)
+        lhs_val = cast_int_or_float_to_bool(network, name, lhs_val)
 
     if isinstance(rhs_val, TRTTensor):
-        rhs_val = trt_cast_int_or_float_to_bool(network, name, rhs_val)
+        rhs_val = cast_int_or_float_to_bool(network, name, rhs_val)
 
     return convert_binary_elementwise(
         network, target, source_ir, name, trt.ElementWiseOperation.XOR, lhs_val, rhs_val