chore: dynamic shape support for rsqrt/erf ops (#2929)

Author: keehyuna

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -560,7 +560,7 @@ def aten_ops_matmul(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.rsqrt.default)
+@dynamo_tensorrt_converter(torch.ops.aten.rsqrt.default, supports_dynamic_shapes=True)
 def aten_ops_rsqrt(
     ctx: ConversionContext,
     target: Target,
@@ -634,7 +634,7 @@ def aten_ops_squeeze(
     return impl.squeeze.squeeze(ctx, target, SourceIR.ATEN, name, args[0], args[1])
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.erf.default)
+@dynamo_tensorrt_converter(torch.ops.aten.erf.default, supports_dynamic_shapes=True)
 def aten_ops_erf(
     ctx: ConversionContext,
     target: Target,

py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py

@@ -121,6 +121,10 @@ def rsqrt(
     name: str,
     input: TRTTensor,
 ) -> TRTTensor:
+    if (isinstance(input, TRTTensor)) and (
+        input.dtype == trt.int8 or input.dtype == trt.int32
+    ):
+        input = cast_trt_tensor(ctx, input, trt.float32, f"{name}_cast")
     sqrt_trt_output = convert_unary(
         ctx,
         target,

tests/py/dynamo/conversion/test_erf_aten.py

@@ -41,6 +41,54 @@ def forward(self, input):
             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,
+            ),
+            (
+                "3d_dim_dtype_int32",
+                (1, 1, 1),
+                (2, 2, 4),
+                (2, 3, 5),
+                torch.int32,
+                torch.float,
+            ),
+        ]
+    )
+    def test_dynamic_shape_erf(
+        self, _, min_shape, opt_shape, max_shape, type, output_type
+    ):
+        class erf(nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.erf.default(input)
+
+        input_specs = [
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            erf(), input_specs, output_dtypes=[output_type]
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/conversion/test_rsqrt_aten.py

@@ -10,21 +10,74 @@
 class TestRSqrtConverter(DispatchTestCase):
     @parameterized.expand(
         [
-            ("2d_dim_alpha", (2, 1), 2),
-            ("3d_dim_alpha", (2, 1, 2), 2),
+            ("2d_dim_float", (2, 1), torch.float),
+            ("3d_dim_float", (2, 1, 2), torch.float),
+            ("3d_dim_int32", (2, 1, 2), torch.int32),
         ]
     )
-    def test_rsqrt(self, _, x, alpha):
+    def test_rsqrt(self, _, x, type):
         class rsqrt(nn.Module):
             def forward(self, input):
                 return torch.ops.aten.rsqrt.default(input)
 
-        inputs = [torch.randn(x) + 1]
+        if type == torch.int32:
+            inputs = [torch.randint(255, x, dtype=torch.int32)]
+        else:
+            inputs = [torch.randn(x) + 1]
+
         self.run_test(
             rsqrt(),
             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,
+            ),
+            (
+                "3d_dim_dtype_int32",
+                (1, 1, 1),
+                (2, 2, 4),
+                (2, 3, 5),
+                torch.int32,
+                torch.float,
+            ),
+        ]
+    )
+    def test_dynamic_shape_rsqrt(
+        self, _, min_shape, opt_shape, max_shape, type, output_type
+    ):
+        class rsqrt(nn.Module):
+            def forward(self, input):
+                return torch.ops.aten.rsqrt.default(input)
+
+        input_specs = [
+            Input(
+                min_shape=min_shape,
+                opt_shape=opt_shape,
+                max_shape=max_shape,
+                dtype=type,
+            ),
+        ]
+
+        self.run_test_with_dynamic_shape(
+            rsqrt(), input_specs, output_dtypes=[output_type]
+        )
+
 
 if __name__ == "__main__":
     run_tests()