Changing torch_dtype_to_trt and torch_dtype_from_trt to unified_dtype_converter and adding convolution in acc_ops_converters.py and aten_ops_converters.py

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -26,7 +26,7 @@
     trt_transposed_matmul,
 )
 from torch_tensorrt.fx.tracer.acc_tracer.acc_ops import contiguous
-from torch_tensorrt.fx.converters.impl import activation
+from torch_tensorrt.fx.converters.impl import activation, convolution
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
 from torch_tensorrt.fx.converters.impl.unary import sign
 from torch_tensorrt.fx.converters.impl.elementwise.base import (

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -20,7 +20,7 @@
 
 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
+from torch_tensorrt.fx.converters.impl import activation, convolution
 from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
 from torch_tensorrt.fx.converters.impl.elementwise import rsqrt
 

py/torch_tensorrt/fx/converters/impl/elementwise/base.py

@@ -10,7 +10,7 @@
 from torch.fx.node import Target
 
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, TRTElementWiseOp
-from torch_tensorrt.fx.utils import torch_dtype_from_trt
+from torch_tensorrt.fx.utils import unified_dtype_converter, Frameworks
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     set_layer_name,
@@ -77,10 +77,10 @@ def convert_binary_elementwise(
     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:

py/torch_tensorrt/fx/converters/impl/elementwise/ops.py

@@ -10,7 +10,7 @@
 from torch.fx.node import Target
 
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, TRTElementWiseOp
-from torch_tensorrt.fx.utils import torch_dtype_from_trt
+from torch_tensorrt.fx.utils import unified_dtype_converter, Frameworks
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     get_trt_tensor,
@@ -70,7 +70,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 = convert_unary(

py/torch_tensorrt/fx/converters/impl/shape.py

@@ -10,7 +10,7 @@
 from torch.fx.node import Target
 
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, TRTElementWiseOp
-from torch_tensorrt.fx.utils import torch_dtype_from_trt
+from torch_tensorrt.fx.utils import unified_dtype_converter, Frameworks
 from torch_tensorrt.fx.converters.converter_utils import (
     SourceIR,
     set_layer_name,