feat: Add permute operation implementation

Author: gs-olive

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -358,3 +358,21 @@ def aten_ops_slice(
         args[3],
         args_bounds_check(args, 4, replacement=1),
     )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.permute.default)
+def aten_ops_slice(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.permutation.permute(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )

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

@@ -11,3 +11,4 @@
 from . import shape
 from . import squeeze
 from . import unsqueeze
+from . import permutation

py/torch_tensorrt/dynamo/conversion/impl/permutation.py

@@ -0,0 +1,34 @@
+from typing import Optional, Sequence, cast
+
+
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+from torch_tensorrt.dynamo.conversion import SourceIR
+from torch_tensorrt.fx.converters.converter_utils import (
+    set_layer_name,
+    get_positive_dim,
+)
+
+
+def permute(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    permutation: Sequence[int],
+) -> TRTTensor:
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"permute received input {input} that is not a TensorRT ITensor"
+        )
+
+    permutation = [
+        get_positive_dim(i, len(input.shape)) for i in cast(Sequence[int], permutation)
+    ]
+
+    layer = network.add_shuffle(input)
+    layer.second_transpose = tuple(permutation)
+    set_layer_name(layer, target, name, source_ir)
+    return layer.get_output(0)

tests/py/dynamo/converters/test_permutation_aten.py

@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.dynamo.test_utils import DispatchTestCase
+from torch_tensorrt import Input
+
+
+class TestPermuteConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("positive", [0, 2, 1]),
+            ("negative", [0, -1, -2]),
+        ]
+    )
+    def test_permute_list(self, _, permutation):
+        class Permute(nn.Module):
+            def forward(self, x):
+                return x.permute(permutation)
+
+        inputs = [torch.randn(1, 3, 2)]
+        self.run_test(Permute(), inputs, expected_ops={torch.ops.aten.permute.default})
+
+    @parameterized.expand(
+        [
+            ("positive", [0, 2, 1]),
+            ("negative", [0, -1, -2]),
+        ]
+    )
+    def test_permute(self, _, permutation):
+        class Permute(nn.Module):
+            def forward(self, x):
+                return x.permute(*permutation)
+
+        inputs = [torch.randn(1, 3, 2)]
+        self.run_test(Permute(), inputs, expected_ops={torch.ops.aten.permute.default})
+
+    def test_permute_with_dynamic_shape(self):
+        class Permute(nn.Module):
+            def forward(self, x):
+                return x.permute(1, 2, 0)
+
+        input_specs = [
+            Input(
+                shape=(-1, -1, -1),
+                dtype=torch.float32,
+                shape_ranges=[((1, 1, 1), (1, 2, 3), (3, 3, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            Permute(), input_specs, expected_ops={torch.ops.aten.permute.default}
+        )
+
+    def test_permute_with_dynamic_shape_four_dimensions(self):
+        class Permute(nn.Module):
+            def forward(self, x):
+                return x.permute(1, 2, 3, 0)
+
+        input_specs = [
+            Input(
+                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(
+            Permute(), input_specs, expected_ops={torch.ops.aten.permute.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()