Add support for prelu dynamo converter (#2972)

Author: HolyWu

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -3395,3 +3395,29 @@ def aten_ops_native_dropout(
         args[1],
         args_bounds_check(args, 2, None),
     )
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten._prelu_kernel.default, supports_dynamic_shapes=True
+)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+        1: (TRTTensor,),
+    }
+)
+def aten_ops_prelu(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.prelu.prelu(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )

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

@@ -18,6 +18,7 @@
     pad,
     permutation,
     pool,
+    prelu,
     quantize,
     reduce,
     select,

py/torch_tensorrt/dynamo/conversion/impl/prelu.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import set_layer_name
+from torch_tensorrt.dynamo.types import TRTTensor
+
+
+def prelu(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    weight: TRTTensor,
+) -> TRTTensor:
+    layer = ctx.net.add_parametric_relu(input, weight)
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)

py/torch_tensorrt/dynamo/lowering/_decomposition_groups.py

@@ -108,7 +108,6 @@
     aten.norm,
     aten.ones,
     aten.ones_like,
-    aten._prelu_kernel,
     aten._prelu_kernel_backward,
     aten._reshape_alias,
     aten.rad2deg,

tests/py/dynamo/conversion/test_prelu_aten.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+from .harness import DispatchTestCase
+
+
+class TestPReLUConverter(DispatchTestCase):
+    def test_prelu(self):
+        class TestModule(nn.Module):
+            def forward(self, x, weight):
+                return torch.ops.aten._prelu_kernel.default(x, weight)
+
+        inputs = [torch.randn(1, 10), torch.randn(1, 1)]
+        self.run_test(TestModule(), inputs)
+
+    def test_prelu_with_dynamic_shape(self):
+        class TestModule(nn.Module):
+            def forward(self, x, weight):
+                return torch.ops.aten._prelu_kernel.default(x, weight)
+
+        input_specs = [
+            Input(
+                min_shape=(1, 1, 1),
+                opt_shape=(1, 2, 3),
+                max_shape=(3, 3, 3),
+                dtype=torch.float32,
+                name="x",
+            ),
+            Input(
+                min_shape=(1, 1, 1),
+                opt_shape=(1, 1, 1),
+                max_shape=(1, 1, 1),
+                dtype=torch.float32,
+                name="weight",
+            ),
+        ]
+        self.run_test_with_dynamic_shape(TestModule(), input_specs)
+
+    def test_prelu_with_dynamic_shape_four_dimensions(self):
+        class TestModule(nn.Module):
+            def forward(self, x, weight):
+                return torch.ops.aten._prelu_kernel.default(x, weight)
+
+        input_specs = [
+            Input(
+                min_shape=(1, 1, 1, 5),
+                opt_shape=(1, 2, 3, 5),
+                max_shape=(3, 3, 3, 5),
+                dtype=torch.float32,
+                name="x",
+            ),
+            Input(
+                min_shape=(1, 1, 1, 1),
+                opt_shape=(1, 2, 1, 1),
+                max_shape=(1, 3, 1, 1),
+                dtype=torch.float32,
+                name="weight",
+            ),
+        ]
+        self.run_test_with_dynamic_shape(TestModule(), input_specs)
+
+
+if __name__ == "__main__":
+    run_tests()