Converter reorg and gelu

Linting error

Author: apbose

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -3401,33 +3401,14 @@ def acc_ops_gelu(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_val = kwargs["input"]
-    approximate = kwargs["approximate"]
-    if approximate != "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
+    return activation.gelu(
+        network,
+        target,
+        SourceIR.ACC,
+        name,
+        kwargs["input"],
+        kwargs["approximate"],
     )
-    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)
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)
 
 
 @tensorrt_converter(acc_ops.chunk)

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -242,6 +242,23 @@ def aten_ops_hardtanh(
     )
 
 
+@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.fmod.Tensor)
 def aten_ops_fmod(
     network: TRTNetwork,

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

@@ -2,6 +2,7 @@
 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
@@ -11,11 +12,15 @@
 
 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
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    get_trt_plugin,
+)
 
 from torch_tensorrt.fx.types import (
     TRTNetwork,
     TRTTensor,
+    TRTPluginFieldCollection,
 )
 
 
@@ -250,3 +255,47 @@ 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/nn_ops_converters.py

@@ -68,6 +68,21 @@ 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/test/converters/aten_op/test_gelu_aten.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestGeLUConverter(DispatchTestCase):
+    def test_gelu(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.gelu(x)
+
+        inputs = [torch.randn(1, 10)]
+        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.gelu.default})
+
+    def test_gelu_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.gelu(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.gelu.default}
+        )
+
+    def test_gelu_with_dynamic_shape_four_dimensions(self):
+        class TestModule(nn.Module):
+            def forward(self, x):
+                return nn.functional.gelu(x)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, -1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1, 5), (1, 2, 3, 5), (3, 3, 3, 5))],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            TestModule(), input_specs, expected_ops={torch.ops.aten.gelu.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()