add dynamic shape support for aten.ops.gt and aten.ops.ge (#2883)

Author: lanluo-nvidia

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -2161,8 +2161,8 @@ def aten_ops_ne(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.gt.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.gt.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.gt.Tensor, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(torch.ops.aten.gt.Scalar, supports_dynamic_shapes=True)
 @enforce_tensor_types(
     {
         0: (TRTTensor,),
@@ -2185,8 +2185,8 @@ def aten_ops_gt(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.ge.Tensor)
-@dynamo_tensorrt_converter(torch.ops.aten.ge.Scalar)
+@dynamo_tensorrt_converter(torch.ops.aten.ge.Tensor, supports_dynamic_shapes=True)
+@dynamo_tensorrt_converter(torch.ops.aten.ge.Scalar, supports_dynamic_shapes=True)
 @enforce_tensor_types(
     {
         0: (TRTTensor,),

tests/py/dynamo/conversion/test_ge_aten.py

@@ -2,6 +2,7 @@
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase
 
@@ -61,6 +62,56 @@ def forward(self, lhs_val):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            ("2d_2d", (5, 3), (5, 1)),
+            ("3d_2d", (5, 3, 2), (3, 1)),
+            ("4d_3d", (5, 3, 4, 1), (3, 1, 1)),
+        ]
+    )
+    def test_ge_tensor_broadcast(self, _, lshape, rshape):
+        class ge(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.ge.Tensor(lhs_val, rhs_val)
+
+        inputs = [
+            torch.randint(0, 3, lshape, dtype=torch.int32),
+            torch.randint(0, 3, rshape, dtype=torch.int32),
+        ]
+        self.run_test(
+            ge(),
+            inputs,
+        )
+
+    @parameterized.expand(
+        [
+            ("2d_2d", (2, 3), (4, 3), (5, 3), (2, 3), (4, 3), (5, 3)),
+            ("3d_2d", (2, 2, 2), (2, 3, 2), (2, 4, 2), (2, 1), (3, 1), (4, 1)),
+        ]
+    )
+    def test_ge_dynamic_tensor(self, *args):
+        class ge(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.ge.Tensor(lhs_val, rhs_val)
+
+        input_specs = [
+            Input(
+                min_shape=args[1],
+                opt_shape=args[2],
+                max_shape=args[3],
+            ),
+            Input(
+                min_shape=args[4],
+                opt_shape=args[5],
+                max_shape=args[6],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            ge(),
+            input_specs,
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/conversion/test_gt_aten.py

@@ -2,6 +2,7 @@
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
 
 from .harness import DispatchTestCase
 
@@ -58,6 +59,56 @@ def forward(self, lhs_val):
             inputs,
         )
 
+    @parameterized.expand(
+        [
+            ("2d_2d", (5, 3), (5, 1)),
+            ("3d_2d", (5, 3, 2), (3, 1)),
+            ("4d_3d", (5, 3, 4, 1), (3, 1, 1)),
+        ]
+    )
+    def test_gt_tensor_broadcast(self, _, lshape, rshape):
+        class gt(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.gt.Tensor(lhs_val, rhs_val)
+
+        inputs = [
+            torch.randint(0, 3, lshape, dtype=torch.int32),
+            torch.randint(0, 3, rshape, dtype=torch.int32),
+        ]
+        self.run_test(
+            gt(),
+            inputs,
+        )
+
+    @parameterized.expand(
+        [
+            ("2d_2d", (2, 3), (4, 3), (5, 3), (2, 3), (4, 3), (5, 3)),
+            ("3d_2d", (2, 2, 2), (2, 3, 2), (2, 4, 2), (2, 1), (3, 1), (4, 1)),
+        ]
+    )
+    def test_gt_dynamic_tensor(self, *args):
+        class gt(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.gt.Tensor(lhs_val, rhs_val)
+
+        input_specs = [
+            Input(
+                min_shape=args[1],
+                opt_shape=args[2],
+                max_shape=args[3],
+            ),
+            Input(
+                min_shape=args[4],
+                opt_shape=args[5],
+                max_shape=args[6],
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            gt(),
+            input_specs,
+        )
+
 
 if __name__ == "__main__":
     run_tests()