chore: dynamic shape support for clamp/min/max/floor_div/logical_and

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -733,10 +733,10 @@ def aten_ops_where(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.clamp.default)
-@dynamo_tensorrt_converter(torch.ops.aten.clamp.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.clip.default)
-@dynamo_tensorrt_converter(torch.ops.aten.clip.Tensor)
+@dynamo_tensorrt_converter(torch.ops.aten.clamp.default, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(torch.ops.aten.clamp.Tensor, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(torch.ops.aten.clip.default, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(torch.ops.aten.clip.Tensor, supports_dynamic_shapes=True)
 def aten_ops_clamp(
     ctx: ConversionContext,
     target: Target,
@@ -1878,7 +1878,7 @@ def aten_ops_mul(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.maximum.default)
+@dynamo_tensorrt_converter(torch.ops.aten.maximum.default, supports_dynamic_shapes=True)
 def aten_ops_maximum(
     ctx: ConversionContext,
     target: Target,
@@ -1896,7 +1896,7 @@ def aten_ops_maximum(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.minimum.default)
+@dynamo_tensorrt_converter(torch.ops.aten.minimum.default, supports_dynamic_shapes=True)
 def aten_ops_minimum(
     ctx: ConversionContext,
     target: Target,
@@ -2017,8 +2017,13 @@ def aten_ops_pow(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.floor_divide.default)
-@dynamo_tensorrt_converter(torch.ops.aten.floor_divide.Scalar)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.floor_divide.default, supports_dynamic_shapes=True
+)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.floor_divide.Scalar, supports_dynamic_shapes=True
+)
+@dynamo_tensorrt_converter(operator.floordiv, supports_dynamic_shapes=True)
 def aten_ops_floor_div(
     ctx: ConversionContext,
     target: Target,
@@ -2036,7 +2041,9 @@ def aten_ops_floor_div(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.logical_and.default)
+@dynamo_tensorrt_converter(
+    torch.ops.aten.logical_and.default, supports_dynamic_shapes=True
+)
 def aten_ops_logical_and(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/conversion/test_clamp_aten.py

@@ -54,12 +54,12 @@ def forward(self, x):
 
         input_specs = [
             Input(
-                shape=(-1, -1, 3, 3),
-                dtype=torch.float32,
-                shape_ranges=[((1, 1, 3, 3), (3, 3, 3, 3), (5, 5, 3, 3))],
+                min_shape=(1, 1, 3, 3),
+                opt_shape=(3, 3, 3, 3),
+                max_shape=(5, 5, 3, 3),
+                dtype=torch.float,
             ),
         ]
-
         self.run_test_with_dynamic_shape(TestModule(), input_specs)
         self.run_test_with_dynamic_shape(TestScalarModule(), input_specs)
 

tests/py/dynamo/conversion/test_floor_div_aten.py

@@ -61,6 +61,58 @@ def forward(self, lhs_val):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            (
+                "2d_dim_dtype_half",
+                (1, 1),
+                (2, 2),
+                (4, 4),
+                torch.half,
+                torch.half,
+            ),
+            (
+                "3d_dim_dtype_float",
+                (1, 1, 1),
+                (1, 2, 3),
+                (3, 3, 3),
+                torch.float,
+                torch.float,
+            ),
+        ]
+    )
+    def test_floor_div_dynamic_shape(
+        self, _, min_shape, opt_shape, max_shape, type, output_type
+    ):
+        class floor_div(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.floor_divide.default(lhs_val, rhs_val)
+
+        class floor_div_operator(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return lhs_val // rhs_val
+
+        input_specs = [
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            floor_div(), input_specs, output_dtypes=[output_type]
+        )
+        self.run_test_with_dynamic_shape(
+            floor_div_operator(), input_specs, output_dtypes=[output_type]
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/conversion/test_logical_and_aten.py

@@ -25,6 +25,45 @@ def forward(self, lhs_val, rhs_val):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            (
+                "2d_dim_dtype_float",
+                (1, 1),
+                (2, 2),
+                (4, 4),
+                torch.float,
+            ),
+            (
+                "3d_dim_dtype_bool",
+                (1, 1, 1),
+                (1, 2, 3),
+                (3, 3, 3),
+                torch.bool,
+            ),
+        ]
+    )
+    def test_logical_and_dynamic_shape(self, _, min_shape, opt_shape, max_shape, type):
+        class logical_and(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.logical_and.default(lhs_val, rhs_val)
+
+        input_specs = [
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+        ]
+        self.run_test_with_dynamic_shape(logical_and(), input_specs)
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/conversion/test_maximum_aten.py

@@ -17,7 +17,7 @@ class TestMaximumConverter(DispatchTestCase):
     def test_maximum(self, _, shape):
         class Maximum(nn.Module):
             def forward(self, lhs_val, rhs_val):
-                return torch.maximum(lhs_val, rhs_val)
+                return torch.ops.aten.maximum.default(lhs_val, rhs_val)
 
         inputs = [torch.randn(shape), torch.randn(shape)]
         self.run_test(
@@ -26,6 +26,51 @@ def forward(self, lhs_val, rhs_val):
             use_dynamo_tracer=True,
         )
 
+    @parameterized.expand(
+        [
+            (
+                "2d_dim_dtype_half",
+                (1, 1),
+                (2, 2),
+                (4, 4),
+                torch.half,
+                torch.half,
+            ),
+            (
+                "3d_dim_dtype_float",
+                (1, 1, 1),
+                (1, 2, 3),
+                (3, 3, 3),
+                torch.float,
+                torch.float,
+            ),
+        ]
+    )
+    def test_maximum_dynamic_shape(
+        self, _, min_shape, opt_shape, max_shape, type, output_type
+    ):
+        class Maximum(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.maximum.default(lhs_val, rhs_val)
+
+        input_specs = [
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            Maximum(), input_specs, output_dtypes=[output_type]
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/conversion/test_minimum_aten.py

@@ -17,7 +17,7 @@ class TestMinimumConverter(DispatchTestCase):
     def test_minimum(self, _, shape):
         class Minimum(nn.Module):
             def forward(self, lhs_val, rhs_val):
-                return torch.minimum(lhs_val, rhs_val)
+                return torch.ops.aten.minimum.default(lhs_val, rhs_val)
 
         inputs = [torch.randn(shape), torch.randn(shape)]
         self.run_test(
@@ -26,6 +26,51 @@ def forward(self, lhs_val, rhs_val):
             use_dynamo_tracer=True,
         )
 
+    @parameterized.expand(
+        [
+            (
+                "2d_dim_dtype_half",
+                (1, 1),
+                (2, 2),
+                (4, 4),
+                torch.half,
+                torch.half,
+            ),
+            (
+                "3d_dim_dtype_float",
+                (1, 1, 1),
+                (1, 2, 3),
+                (3, 3, 3),
+                torch.float,
+                torch.float,
+            ),
+        ]
+    )
+    def test_minimum_dynamic_shape(
+        self, _, min_shape, opt_shape, max_shape, type, output_type
+    ):
+        class Minimum(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.minimum.default(lhs_val, rhs_val)
+
+        input_specs = [
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+        ]
+        self.run_test_with_dynamic_shape(
+            Minimum(), input_specs, output_dtypes=[output_type]
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/partitioning/test_dynamic_partitioning.py

@@ -8,11 +8,8 @@
 
 from ..testing_utilities import lower_graph_testing
 
-# This testcase assumes that torch.ops.aten.clamp.default converter doesn't support
-# dynamic shapes. One should remove this testcase when the support is added.
-# This testcase tests if the graph is partitioned correctly into a TRT segment
-# and a Pytorch segment when the torch.ops.aten.clamp.default converter gets disabled
-# due to lack of dynamic shape support.
+# This testcase assumes that torch.ops.aten.clamp.default supports dynamic shapes.
+# This testcase tests if the graph is partitioned correctly into a TRT segment only.
 
 
 class TestDynamicPartitioning(TestCase):
@@ -53,51 +50,6 @@ def forward(self, x):
             1,
             f"Number of TRT segments should be 1 but got {trt_segments}",
         )
-        self.assertEqual(
-            pyt_segments,
-            1,
-            f"Number of PyTorch segments should be 1 but got {pyt_segments}",
-        )
-
-    def test_assume_dynamic_shape_support_converters(self):
-        class Clamp(torch.nn.Module):
-            def __init__(self, *args, **kwargs) -> None:
-                super().__init__(*args, **kwargs)
-                self.relu = torch.nn.ReLU()
-
-            def forward(self, x):
-                x = self.relu(x)
-                return torch.ops.aten.clamp.default(x, min=2.5, max=6.5)
-
-        model = Clamp().eval().cuda()
-        trt_model = torch_tensorrt.compile(
-            model,
-            inputs=[
-                torch_tensorrt.Input(
-                    min_shape=(1, 3, 8, 8),
-                    opt_shape=(4, 3, 8, 8),
-                    max_shape=(8, 3, 8, 8),
-                    dtype=torch.float32,
-                    name="x",
-                )
-            ],
-            dryrun=True,
-            assume_dynamic_shape_support=True,
-            min_block_size=1,
-        )
-
-        trt_segments, pyt_segments = 0, 0
-        for submod in list(trt_model.named_children()):
-            if "_run_on_acc" in submod[0]:
-                trt_segments += 1
-            elif "_run_on_gpu" in submod[0]:
-                pyt_segments += 1
-
-        self.assertEqual(
-            trt_segments,
-            1,
-            f"Number of TRT segments should be 2 but got {trt_segments}",
-        )
         self.assertEqual(
             pyt_segments,
             0,