Revert all changes to py/torch_tensorrt/fx

Revert "fix: Add automatic type promotion for FX ops"

This reverts commit f1f3716.

Revert "Moved clamp to impl"

This reverts commit df401dd.

Revert "aten::unsqueeze impl refactor"

This reverts commit b424735.

Revert "Converter reorg and where operator"

This reverts commit b4da15e.

Revert "converter reorg and matmul"

This reverts commit 7551eee.

Revert "converter reorg and slice"

This reverts commit 9bbdc9e.

Revert "Converter reorg and select operation"

This reverts commit fb70253.

Revert "Converter reorg and squeeze operator"

This reverts commit 294545c.

Revert "Converter reorg and gelu"

This reverts commit 37d1168.

Revert "Converter reorg and softmax operation"

This reverts commit 1ba6d13.

Revert "layer_norm converter"

This reverts commit e0b34b1.

Revert "Converter reorg batch norm"

This reverts commit 59354e5.

Revert "Converter reorg and rsub"

This reverts commit db15d27.

Revert "Converter reorg fmod"

This reverts commit ce3fa67.

Revert "Moving elementwise core to impl - rsqrt (FX Converter Refactor [9/N]) <Target: converter_reorg_elementwise> (#1905)"

This reverts commit 7158ca5.

Revert "refactor: Moving elementwise and unary core to impl"

This reverts commit 45e43ca.

Author: gs-olive

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -21,28 +21,9 @@
 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, convolution
-from torch_tensorrt.fx.converters.impl.elementwise import trunc_div
-from torch_tensorrt.fx.converters.impl.elementwise import rsqrt
-from torch_tensorrt.fx.converters.impl.elementwise import fmod
-from torch_tensorrt.fx.converters.impl.elementwise import rsub
-from torch_tensorrt.fx.converters.impl.normalization import batch_norm
-from torch_tensorrt.fx.converters.impl.normalization import layer_norm
-from torch_tensorrt.fx.converters.impl.normalization import softmax
-from torch_tensorrt.fx.converters.impl.squeeze import squeeze
-from torch_tensorrt.fx.converters.impl.select import select
-from torch_tensorrt.fx.converters.impl.slice import slice_op
-from torch_tensorrt.fx.converters.impl.matmul import matrix_multiply
-from torch_tensorrt.fx.converters.impl.condition import where
-from torch_tensorrt.fx.converters.impl.unsqueeze import unsqueeze
-from torch_tensorrt.fx.converters.impl.elementwise import clamp
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
-
-def or_none(args, i):
-    return args[i] if len(args) > i else None
-
-
 ## converter list in alphabetic order
 @tensorrt_converter(torch.ops.aten.add.Tensor)
 def aten_ops_add(
@@ -108,19 +89,18 @@ def aten_ops_batch_norm(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return batch_norm(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        args[1],
-        args[2],
-        args[3],
-        args[4],
-        args[5],
-        args[6],
-        args[7],
+    kwargs_new = {
+        "input": args[0],
+        "weight": args[1],
+        "bias": args[2],
+        "running_mean": args[3],
+        "running_var": args[4],
+        "training": args[5],
+        "momentum": args[6],
+        "eps": args[7],
+    }
+    return acc_ops_converters.acc_ops_batch_norm(
+        network, target, None, kwargs_new, name
     )
 
 
@@ -202,7 +182,9 @@ def aten_ops_div(
             network, target, None, kwargs_new, name
         )
     elif rounding_mode == "trunc":
-        return trunc_div(network, target, SourceIR.ATEN, name, args[0], args[1])
+        return acc_ops_converters.acc_ops_trunc_div(
+            network, target, None, kwargs_new, name
+        )
     else:
         raise RuntimeError(
             f"Target {target} does not support rounding mode {rounding_mode}"
@@ -260,7 +242,11 @@ def aten_ops_fmod(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return fmod(network, target, SourceIR.ATEN, name, args[0], args[1])
+    kwargs_new = {
+        "input": args[0],
+        "other": args[1],
+    }
+    return acc_ops_converters.acc_ops_fmod(network, target, None, kwargs_new, name)
 
 
 @tensorrt_converter(torch.ops.aten.hardtanh.default)
@@ -271,40 +257,12 @@ def aten_ops_hardtanh(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
+
     return activation.hardtanh(
         network, target, SourceIR.ATEN, name, args[0], args[1], args[2]
     )
 
 
-@tensorrt_converter(torch.ops.aten.gelu.default)
-def aten_ops_gelu(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return activation.gelu(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-    )
-
-
-@tensorrt_converter(torch.ops.aten.matmul)
-@tensorrt_converter(torch.ops.aten.mm.default)
-def aten_ops_matmul(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return matrix_multiply(network, target, SourceIR.ATEN, name, args[0], args[1])
-
-
 @tensorrt_converter(torch.ops.aten.fmod.Tensor)
 def aten_ops_fmod(
     network: TRTNetwork,
@@ -328,28 +286,8 @@ def aten_ops_leaky_relu(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
-
 
-@tensorrt_converter(torch.ops.aten.layer_norm.default)
-def aten_ops_layernorm(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return layer_norm(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        args[1],
-        args[2],
-        args[3],
-        args[4],
-    )
+    return activation.leaky_relu(network, target, SourceIR.ATEN, name, args[0], args[1])
 
 
 @tensorrt_converter(torch.ops.aten.linear)
@@ -452,42 +390,6 @@ def aten_ops_relu(
     )
 
 
-@tensorrt_converter(torch.ops.aten.relu.default)
-def aten_ops_relu(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-
-    return activation.relu(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-    )
-
-
-@tensorrt_converter(torch.ops.aten.rsqrt.default)
-def aten_ops_rsqrt(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-
-    return rsqrt(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-    )
-
-
 @tensorrt_converter(torch.ops.aten.sub.Tensor)
 def aten_ops_sub(
     network: TRTNetwork,
@@ -503,29 +405,6 @@ def aten_ops_sub(
     return acc_ops_converters.acc_ops_sub(network, target, None, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.squeeze.dim)
-@tensorrt_converter(torch.ops.aten.squeeze.dims)
-def aten_ops_squeeze(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return squeeze(network, target, SourceIR.ATEN, name, args[0], args[1])
-
-
-@tensorrt_converter(torch.ops.aten.unsqueeze.default)
-def aten_ops_unsqueeze(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return unsqueeze(network, target, SourceIR.ATEN, name, input_t=args[0], dim=args[1])
-
-
 @tensorrt_converter(torch.ops.aten.view.default)
 def aten_ops_reshape(
     network: TRTNetwork,
@@ -563,31 +442,6 @@ def aten_ops_reshape(
     return layer.get_output(0)
 
 
-@tensorrt_converter(torch.ops.aten.rsub.Tensor)
-def aten_ops_rsub(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    alpha = None
-    if "alpha" in kwargs:
-        alpha = kwargs["alpha"]
-    return rsub(network, target, SourceIR.ATEN, name, args[0], args[1], alpha)
-
-
-@tensorrt_converter(torch.ops.aten._softmax.default)
-def aten_ops_softmax(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return softmax(network, target, SourceIR.ATEN, name, args[0], args[1])
-
-
 @tensorrt_converter(torch.ops.aten.tanh.default)
 def aten_ops_tanh(
     network: TRTNetwork,
@@ -596,30 +450,12 @@ def aten_ops_tanh(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return activation.tanh(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-    )
 
-
-@tensorrt_converter(torch.ops.aten.where.self)
-def aten_ops_where(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return where(
+    return activation.tanh(
         network,
         target,
         SourceIR.ATEN,
         name,
-        args[1],
-        args[2],
         args[0],
     )
 
@@ -639,25 +475,6 @@ def aten_ops_cat(
     return acc_ops_converters.acc_ops_cat(network, target, None, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.clamp.default)
-def aten_ops_clamp(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return clamp.clamp(
-        network,
-        target,
-        SourceIR.ACC,
-        name,
-        input_val=args[0],
-        min_val=or_none(args, 1),
-        max_val=or_none(args, 2),
-    )
-
-
 @tensorrt_converter(torch.ops.aten.expand.default)
 def aten_ops_expand(
     network: TRTNetwork,
@@ -720,17 +537,6 @@ def aten_ops_operator_add(
     return acc_ops_converters.acc_ops_add(network, target, None, kwargs_new, name)
 
 
-@tensorrt_converter(torch.ops.aten.select.int)
-def aten_ops_select(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return select(network, target, SourceIR.ATEN, name, args[0], args[1], args[2])
-
-
 @tensorrt_converter(operator.sub)
 def aten_ops_operator_sub(
     network: TRTNetwork,
@@ -766,27 +572,6 @@ def aten_ops_sym_numel(
     return reduce_layer.get_output(0)
 
 
-@tensorrt_converter(torch.ops.aten.slice.Tensor)
-def aten_ops_slice(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return slice_op(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        args[1],
-        args[2],
-        args[3],
-        args[4],
-    )
-
-
 @tensorrt_converter(torch.ops.aten.sym_size)
 def aten_ops_sym_size(
     network: TRTNetwork,

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -2,7 +2,6 @@
 import operator
 import warnings
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
-import math
 
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
@@ -12,15 +11,11 @@
 
 from torch_tensorrt.fx.converters.converter_utils import mark_as_int8_layer
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
-from torch_tensorrt.fx.converters.converter_utils import (
-    SourceIR,
-    get_trt_plugin,
-)
+from torch_tensorrt.fx.converters.converter_utils import SourceIR
 
 from torch_tensorrt.fx.types import (
     TRTNetwork,
     TRTTensor,
-    TRTPluginFieldCollection,
 )
 
 
@@ -255,47 +250,3 @@ def elu_dyn_range_fn(dyn_range):
         input_val,
         dyn_range_fn=elu_dyn_range_fn,
     )
-
-
-def gelu(
-    network: TRTNetwork,
-    target: Target,
-    source_ir: Optional[SourceIR],
-    name: str,
-    input_val: TRTTensor,
-    alpha: Optional[Any] = None,
-):
-    approximate = alpha
-    if approximate is not None:
-        raise RuntimeError("GeLU converter currently doesn't support fast gelu compute")
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"GELU received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-    if network.has_implicit_batch_dimension:
-        raise RuntimeError(
-            "GeLU converter currently doesn't support implicit batch dimension"
-        )
-    plugin_name = "CustomGeluPluginDynamic"
-    # type_id 0 for float32, 1 for  float16
-    type_id = trt.PluginField(
-        "type_id", np.array(0, dtype=np.int32), trt.PluginFieldType.INT32
-    )
-    field_collection = TRTPluginFieldCollection([type_id])
-    plugin_version = "1"
-
-    plugin = get_trt_plugin(plugin_name, field_collection, plugin_version)
-
-    layer = network.add_plugin_v2([input_val], plugin)
-
-    def gelu_dyn_range_fn(dyn_range):
-        return (
-            dyn_range[0] * 0.5 * (1.0 + torch.erf(dyn_range[0] / math.sqrt(2.0)))
-        ), (dyn_range[1] * 0.5 * (1.0 + torch.erf(dyn_range[0] / math.sqrt(2.0))))
-
-    if input_val.dynamic_range is not None:
-        dyn_range = gelu_dyn_range_fn(input_val.dynamic_range)
-        mark_as_int8_layer(layer, dyn_range)
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)

py/torch_tensorrt/fx/converters/converter_utils.py

@@ -69,21 +69,6 @@ def tanh(network, submod, args, kwargs, layer_name):
     )
 
 
-@tensorrt_converter(torch.nn.functional.gelu)
-@tensorrt_converter(torch.nn.modules.activation.GELU)
-def gelu(network, submod, args, kwargs, layer_name):
-    # args/kwargs should have already been normalized to kwargs
-    assert len(args) == 0
-
-    return activation.gelu(
-        network=network,
-        target="torch.nn.functional.gelu",
-        source_ir=SourceIR.NN,
-        name=layer_name,
-        input_val=kwargs["input"],
-    )
-
-
 @tensorrt_converter(torch.nn.functional.leaky_relu)
 @tensorrt_converter(torch.nn.modules.activation.LeakyReLU)
 def leaky_relu(network, submod, args, kwargs, layer_name):

py/torch_tensorrt/fx/converters/impl/activation.py

@@ -57,23 +57,6 @@ def forward(self, x):
         inputs = [torch.rand(1, 1) + 1]
         self.run_test(m, inputs, expected_ops={expected_op})
 
-    @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in elementwise_ops])
-    def test_elementwise_ops_mismatched_dtypes(
-        self, name, orig_op: Callable, expected_op
-    ):
-        class TestModule(nn.Module):
-            def __init__(self, orig_op):
-                super().__init__()
-                self.orig_op = orig_op
-
-            def forward(self, x):
-                return self.orig_op(x.int(), x)
-
-        m = TestModule(orig_op)
-        # Avoid dividing by 0.
-        inputs = [2 * torch.rand(1, 1, dtype=torch.float) + 1]
-        self.run_test(m, inputs, expected_ops={expected_op})
-
     @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in elementwise_ops])
     def test_elementwise_ops_with_one_constant(
         self, name, orig_op: Callable, expected_op

py/torch_tensorrt/fx/converters/impl/condition/__init__.py

@@ -12,12 +12,7 @@ def forward(self, x):
                 return x.squeeze(2)
 
         inputs = [torch.randn(1, 2, 1)]
-        self.run_test(
-            Squeeze(),
-            inputs,
-            expected_ops={acc_ops.squeeze},
-            test_implicit_batch_dim=False,
-        )
+        self.run_test(Squeeze(), inputs, expected_ops={acc_ops.squeeze})
 
     # Testing with shape=(-1, -1, -1, -1) results in error:
     # AssertionError: We don't support squeeze dynamic dim.
@@ -38,9 +33,7 @@ def forward(self, x):
             ),
         ]
         self.run_test_with_dynamic_shape(
-            Squeeze(),
-            input_specs,
-            expected_ops={acc_ops.squeeze},
+            Squeeze(), input_specs, expected_ops={acc_ops.squeeze}
         )