fix unsqueeze cannot work on more than 1 dynamic_shape dimensions (#2933)

lanluo-nvidia · web-flow · commit 363d0bde3058 · 2024-06-18T14:10:57.000-07:00
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py b/py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py
@@ -9,7 +9,6 @@
 )
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import Shape, TRTTensor
-from torch_tensorrt.fx.utils import get_dynamic_dims
 
 
 def unsqueeze(
@@ -32,13 +31,49 @@ def unsqueeze(
     input_shape_size = len(input_val.shape)
     dim = get_positive_dim(dim, input_shape_size + 1)
 
-    assert (
-        len(get_dynamic_dims(input_val.shape)) <= 1
-    ), "Currently we don't support unsqueeze with more than one dynamic dims."
-    layer = ctx.net.add_shuffle(input_val)
-    layer.reshape_dims = (
-        tuple(input_val.shape)[:dim] + (1,) + tuple(input_val.shape)[dim:]
+    intermediate_dim = 0
+    dynamic_shape_cnt = 0
+    # if unsqueeze the last dimensions, we can directly append to the shape
+    if dim == input_shape_size:
+        intermediate_dim = dim
+    else:
+        # since maximum of one dimension is permitted to be specified as -1
+        # find the intermediate_dim which has only 1 dynamic_shape_cnt
+        # and then we can add a transpose after reshape if it is not the final shape we want
+        for i, s in reversed(list(enumerate(input_val.shape))):
+            if i >= dim:
+                if s == -1:
+                    dynamic_shape_cnt += 1
+                if dynamic_shape_cnt > 1:
+                    intermediate_dim = i + 1
+                    break
+                if i == dim:
+                    intermediate_dim = i
+                    break
+    # calculate the new_shape for the shuffle layer's reshape_dims
+    new_shape = list(
+        tuple(input_val.shape)[:intermediate_dim]
+        + (1,)
+        + tuple(input_val.shape)[intermediate_dim:]
     )
+    for i, s in enumerate(new_shape):
+        if i < intermediate_dim and s == -1:
+            new_shape[i] = 0
+    layer = ctx.net.add_shuffle(input_val)
+    layer.reshape_dims = tuple(new_shape)
+    # if the intermediate_dim is not the final dim we want to unsqueeze, add a second_transpose after reshape
+    if intermediate_dim != dim:
+        # calculate the second_transpose for the shuffle layer
+        permutation = [*range(0, len(new_shape))]
+        # for example: if the reshape_dims is (3, 3, 5, 1, 5) and the final shape we want is (3, 1, 3, 5, 5)
+        # here intermediate_dim=3, dim=1, we need to move intermediate_dim before [dim: intermediate_dim)
+        new_permutation = (
+            tuple(permutation[:dim])
+            + (intermediate_dim,)
+            + tuple(permutation[dim:intermediate_dim])
+            + tuple(permutation[intermediate_dim + 1 :])
+        )
+        layer.second_transpose = new_permutation
     set_layer_name(layer, target, name, source_ir)
     return layer.get_output(0)
 
diff --git a/tests/py/dynamo/conversion/test_unsqueeze_aten.py b/tests/py/dynamo/conversion/test_unsqueeze_aten.py
@@ -3,9 +3,8 @@
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
-from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException
-
 from torch_tensorrt import Input
+from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException
 
 from .harness import DispatchTestCase
 
@@ -29,16 +28,32 @@ def forward(self, x):
         inputs = [torch.randn(1, 2, 3)]
         self.run_test(Unsqueeze(dim), inputs)
 
-    # 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),
+            ("1_dynamic_shape_2d_-3", -3, (2, 5), (3, 5), (4, 5)),
+            ("1_dynamic_shape_2d_-2", -2, (2, 3), (2, 4), (2, 5)),
+            ("1_dynamic_shape_2d_-1", -1, (2, 3), (2, 4), (2, 5)),
+            ("1_dynamic_shape_2d_0", 0, (2, 3), (2, 4), (2, 5)),
+            ("1_dynamic_shape_2d_1", 1, (2, 3), (2, 4), (2, 5)),
+            ("1_dynamic_shape_2d_2", 2, (2, 3), (2, 4), (2, 5)),
+            ("2_dynamic_shape_3d_-1", -1, (2, 2, 3), (4, 3, 3), (5, 5, 3)),
+            ("2_dynamic_shape_3d_0", 2, (2, 2, 3), (4, 3, 3), (5, 5, 3)),
+            ("2_dynamic_shape_3d_1", 1, (2, 2, 3), (4, 3, 3), (5, 6, 3)),
+            ("2_dynamic_shape_3d_2", 2, (2, 2, 3), (4, 3, 3), (6, 5, 3)),
+            ("4_dynamic_shape_4d_-4", -4, (1, 2, 3, 4), (2, 2, 3, 5), (3, 3, 5, 5)),
+            ("4_dynamic_shape_4d_-3", -3, (1, 2, 3, 4), (2, 2, 3, 5), (3, 3, 5, 5)),
+            ("4_dynamic_shape_4d_-2", -2, (1, 2, 3, 4), (2, 2, 3, 5), (4, 3, 5, 6)),
+            ("4_dynamic_shape_4d_-1", -1, (1, 2, 3, 4), (2, 2, 3, 5), (4, 3, 5, 6)),
+            ("4_dynamic_shape_4d_0", 0, (1, 2, 3, 4), (2, 2, 5, 7), (2, 3, 6, 8)),
+            ("4_dynamic_shape_4d_1", 1, (1, 2, 3, 4), (2, 2, 3, 5), (3, 3, 5, 5)),
+            ("4_dynamic_shape_4d_2", 2, (1, 2, 3, 4), (2, 2, 3, 5), (3, 3, 5, 5)),
+            ("4_dynamic_shape_4d_3", 3, (1, 2, 3, 4), (2, 2, 3, 5), (3, 3, 5, 5)),
+            ("4_dynamic_shape_4d_4", 4, (1, 2, 3, 4), (2, 2, 3, 5), (3, 3, 5, 5)),
         ]
     )
-    def test_unsqueeze_with_dynamic_shape(self, _, dim):
+    def test_unsqueeze_with_dynamic_shape(
+        self, _, dim, min_shape, opt_shape, max_shape
+    ):
         class Unsqueeze(nn.Module):
             def __init__(self, dim):
                 super().__init__()
@@ -49,9 +64,10 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, 2, 3),
                 dtype=torch.float32,
-                shape_ranges=[((1, 2, 3), (2, 2, 3), (3, 2, 3))],
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
             ),
         ]
         self.run_test_with_dynamic_shape(Unsqueeze(dim), input_specs)
