aten::unsqueeze impl refactor

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

Author: borisfom

py/torch_tensorrt/fx/converters/acc_ops_converters.py

@@ -39,6 +39,7 @@
 from torch_tensorrt.fx.converters.impl.unary.base import convert_unary
 from torch_tensorrt.fx.converters.impl.shape import get_shape_with_dynamic_shape
 from torch_tensorrt.fx.converters.impl.squeeze import squeeze
+from torch_tensorrt.fx.converters.impl.unsqueeze import unsqueeze
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -2215,36 +2216,14 @@ def acc_ops_unsqueeze(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    input_t = kwargs["input"]
-    input_val = get_trt_tensor(network, input_t, f"{name}_input_t")
-    if not isinstance(input_val, TRTTensor):
-        raise RuntimeError(
-            f"unsqueeze received input {input_val} that is not part "
-            "of the TensorRT region!"
-        )
-
-    dim = cast(int, kwargs["dim"])
-    input_shape = input_val.shape
-    input_shape_size = (
-        len(input_val.shape) + 1
-        if network.has_implicit_batch_dimension
-        else len(input_val.shape)
-    )
-    dim = get_positive_dim(dim, input_shape_size + 1)
-
-    if network.has_implicit_batch_dimension:
-        assert dim != 0
-        dim -= 1
-
-    assert (
-        len(get_dynamic_dims(input_val.shape)) <= 1
-    ), "Currently we don't support unsqueeze with more than one dynamic dims."
-    layer = network.add_shuffle(input_val)
-    layer.reshape_dims = (
-        tuple(input_val.shape)[:dim] + (1,) + tuple(input_val.shape)[dim:]
+    return unsqueeze(
+        network,
+        target,
+        SourceIR.ACC,
+        name,
+        input_t=kwargs["input"],
+        dim=kwargs["dim"],
     )
-    set_layer_name(layer, target, name)
-    return layer.get_output(0)
 
 
 @tensorrt_converter(acc_ops.topk)

py/torch_tensorrt/fx/converters/aten_ops_converters.py

@@ -33,6 +33,7 @@
 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
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -485,6 +486,17 @@ def aten_ops_squeeze(
     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,

py/torch_tensorrt/fx/converters/impl/unsqueeze.py

@@ -0,0 +1,63 @@
+import operator
+import warnings
+from typing import Optional, cast, Any
+
+import numpy as np
+
+import tensorrt as trt
+import torch
+from torch.fx.node import Target
+
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor, Shape
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    get_positive_dim,
+    get_trt_tensor,
+)
+
+from torch_tensorrt.fx.converters.converter_utils import (
+    SourceIR,
+    get_positive_dim,
+    set_layer_name,
+)
+
+from torch_tensorrt.fx.utils import get_dynamic_dims
+
+
+def unsqueeze(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_t,
+    dim,
+) -> TRTTensor:
+    input_val = get_trt_tensor(network, input_t, f"{name}_input_t")
+    if not isinstance(input_val, TRTTensor):
+        raise RuntimeError(
+            f"unsqueeze received input {input_val} that is not part "
+            "of the TensorRT region!"
+        )
+
+    dim = cast(int, dim)
+    input_shape = input_val.shape
+    input_shape_size = (
+        len(input_val.shape) + 1
+        if network.has_implicit_batch_dimension
+        else len(input_val.shape)
+    )
+    dim = get_positive_dim(dim, input_shape_size + 1)
+
+    if network.has_implicit_batch_dimension:
+        assert dim != 0
+        dim -= 1
+
+    assert (
+        len(get_dynamic_dims(input_val.shape)) <= 1
+    ), "Currently we don't support unsqueeze with more than one dynamic dims."
+    layer = network.add_shuffle(input_val)
+    layer.reshape_dims = (
+        tuple(input_val.shape)[:dim] + (1,) + tuple(input_val.shape)[dim:]
+    )
+    set_layer_name(layer, target, name)
+    return layer.get_output(0)

py/torch_tensorrt/fx/test/converters/aten_op/test_unsqueeze_aten.py

@@ -0,0 +1,61 @@
+import torch
+import torch.fx
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
+
+
+class TestUnsqueeze(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("negative_dim", -2),
+            ("positive_dim", 2),
+        ]
+    )
+    def test_unsqueeze(self, _, dim):
+        class Unsqueeze(nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.dim = dim
+
+            def forward(self, x):
+                return torch.unsqueeze(x, self.dim)
+
+        inputs = [torch.randn(1, 2, 3)]
+        self.run_test(
+            Unsqueeze(dim), inputs, expected_ops={torch.ops.aten.unsqueeze.default}
+        )
+
+    # Testing with more than one dynamic dims results in following error:
+    # AssertionError: Currently we don't support unsqueeze with more than one dynamic dims.
+
+    @parameterized.expand(
+        [
+            ("negative_dim_dynamic", -4),
+            ("positive_dim_dynamic", 1),
+        ]
+    )
+    def test_unsqueeze_with_dynamic_shape(self, _, dim):
+        class Unsqueeze(nn.Module):
+            def __init__(self, dim):
+                super().__init__()
+                self.dim = dim
+
+            def forward(self, x):
+                return torch.unsqueeze(x, self.dim)
+
+        input_specs = [
+            InputTensorSpec(
+                shape=(-1, 2, 3),
+                dtype=torch.float32,
+                shape_ranges=[((1, 2, 3), (2, 2, 3), (3, 2, 3))],
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            Unsqueeze(dim), input_specs, expected_ops={torch.ops.aten.unsqueeze.default}
+        )
+
+
+if __name__ == "__main__":
+    run_tests()