feat: support activation dynamo converters

lint test file

fix bugs: circular import

delete gelu

change function calls from nn.Module to nn.functional

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -152,15 +152,15 @@ def aten_ops_fmod(
     return impl.elementwise.fmod(network, target, SourceIR.ATEN, name, args[0], args[1])
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.gelu.default)  # type: ignore[misc]
-def aten_ops_gelu(
+@dynamo_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 impl.activation.gelu(
+    return impl.activation.relu(
         network,
         target,
         SourceIR.ATEN,
@@ -169,55 +169,165 @@ def aten_ops_gelu(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.matmul)  # type: ignore[misc]
-@dynamo_tensorrt_converter(torch.ops.aten.mm.default)  # type: ignore[misc]
-def aten_ops_matmul(
+@dynamo_tensorrt_converter(torch.ops.aten.sigmoid.default)
+def aten_ops_sigmoid(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.matmul.matrix_multiply(
-        network, target, SourceIR.ATEN, name, args[0], args[1]
+    return impl.activation.sigmoid(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
-def aten_ops_layernorm(
+@dynamo_tensorrt_converter(torch.ops.aten.tanh.default)
+def aten_ops_tanh(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.normalization.layer_norm(
+    return impl.activation.tanh(
         network,
         target,
         SourceIR.ATEN,
         name,
         args[0],
-        args[1],
-        args[2],
-        args[3],
-        args[4],
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.relu.default)  # type: ignore[misc]
-def aten_ops_relu(
+@dynamo_tensorrt_converter(torch.ops.aten.leaky_relu.default)
+def aten_ops_leaky_relu(
     network: TRTNetwork,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.activation.relu(
+    return impl.activation.leaky_relu(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args_bounds_check(args, 1, 0.01),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.elu.default)
+def aten_ops_elu(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.activation.elu(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        alpha=args_bounds_check(args, 1, 1.0),
+        beta=args_bounds_check(args, 2, None),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.softplus.default)
+def aten_ops_softplus(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.activation.softplus(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        beta=args_bounds_check(args, 1, 1),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.clip.default)
+def aten_ops_clip(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.activation.clip(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        alpha=args_bounds_check(args, 1),
+        beta=args_bounds_check(args, 2),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.hardsigmoid.default)
+def aten_ops_hard_sigmoid(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.activation.hard_sigmoid(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        alpha=args_bounds_check(args, 1, 1 / 6),
+        beta=args_bounds_check(args, 2, 1 / 2),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.matmul)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.mm.default)  # type: ignore[misc]
+def aten_ops_matmul(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.matmul.matrix_multiply(
+        network, target, SourceIR.ATEN, name, args[0], args[1]
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
+def aten_ops_layernorm(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.normalization.layer_norm(
         network,
         target,
         SourceIR.ATEN,
         name,
         args[0],
+        args[1],
+        args[2],
+        args[3],
+        args[4],
     )
 
 

py/torch_tensorrt/dynamo/conversion/impl/activation.py

@@ -0,0 +1 @@
+from .ops import *

py/torch_tensorrt/dynamo/conversion/impl/activation/__init__.py

@@ -0,0 +1,42 @@
+from typing import Any, Callable, Optional
+
+import tensorrt as trt
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.fx.converters.converter_utils import (
+    mark_as_int8_layer,
+    set_layer_name,
+)
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+
+def convert_activation(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    operation_type: trt.ActivationType,
+    input_val: TRTTensor,
+    alpha: Optional[Any] = None,
+    beta: Optional[Any] = None,
+    dyn_range_fn: Optional[Callable[[float, float], Any]] = None,
+) -> TRTTensor:
+    """
+    Add a TensorRT Activation layer to `network`.
+    """
+    if not isinstance(input_val, TRTTensor):
+        raise RuntimeError(
+            f"{operation_type} received input {input_val} that is not part "
+            "of the TensorRT region!"
+        )
+    layer = network.add_activation(input_val, operation_type)
+    if alpha is not None:
+        layer.alpha = alpha
+    if beta is not None:
+        layer.beta = beta
+    set_layer_name(layer, target, name, source_ir)
+
+    if input_val.dynamic_range is not None:
+        dyn_range = dyn_range_fn(input_val.dynamic_range)
+        mark_as_int8_layer(layer, dyn_range)
+    return layer.get_output(0)