From 4e5334ade114b2e2884e1a118d998e4fd6537352 Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Tue, 4 Jun 2024 16:07:05 -0700 Subject: [PATCH 1/3] add dynamic shape support for sin/cos --- .../dynamo/conversion/aten_ops_converters.py | 4 +- tests/py/dynamo/conversion/test_cos_aten.py | 49 +++++++++++++++++++ tests/py/dynamo/conversion/test_sin_aten.py | 49 +++++++++++++++++++ 3 files changed, 100 insertions(+), 2 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 09cc78f7db..90a11c272a 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1327,7 +1327,7 @@ def aten_ops_abs( ) -@dynamo_tensorrt_converter(torch.ops.aten.sin.default) +@dynamo_tensorrt_converter(torch.ops.aten.sin.default, supports_dynamic_shapes=True) def aten_ops_sin( ctx: ConversionContext, target: Target, @@ -1344,7 +1344,7 @@ def aten_ops_sin( ) -@dynamo_tensorrt_converter(torch.ops.aten.cos.default) +@dynamo_tensorrt_converter(torch.ops.aten.cos.default, supports_dynamic_shapes=True) def aten_ops_cos( ctx: ConversionContext, target: Target, diff --git a/tests/py/dynamo/conversion/test_cos_aten.py b/tests/py/dynamo/conversion/test_cos_aten.py index 505f303219..152254324e 100644 --- a/tests/py/dynamo/conversion/test_cos_aten.py +++ b/tests/py/dynamo/conversion/test_cos_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 @@ -44,6 +45,54 @@ def forward(self, input): inputs, ) + @parameterized.expand( + [ + ( + "3d_dim_dtype_int32", + (3, 2, 1), + (3, 2, 3), + (3, 3, 4), + torch.int32, + torch.float32, + ), + ( + "2d_dim_dtype_float16", + (1, 1), + (2, 2), + (4, 4), + torch.float16, + torch.float16, + ), + ( + "3d_dim_dtype_float", + (1, 1, 1), + (1, 2, 3), + (3, 3, 3), + torch.float, + torch.float, + ), + ] + ) + def test_dynamic_shape_cos( + self, _, min_shape, opt_shape, max_shape, type, output_type + ): + class cos(nn.Module): + def forward(self, input): + return torch.ops.aten.cos.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( + cos(), input_specs, output_dtypes=[output_type] + ) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_sin_aten.py b/tests/py/dynamo/conversion/test_sin_aten.py index 70a3dc0315..7d5ca2204d 100644 --- a/tests/py/dynamo/conversion/test_sin_aten.py +++ b/tests/py/dynamo/conversion/test_sin_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 @@ -44,6 +45,54 @@ def forward(self, input): inputs, ) + @parameterized.expand( + [ + ( + "3d_dim_dtype_int32", + (3, 2, 1), + (3, 2, 3), + (3, 3, 4), + torch.int32, + torch.float32, + ), + ( + "2d_dim_dtype_float16", + (1, 1), + (2, 2), + (4, 4), + torch.float16, + torch.float16, + ), + ( + "3d_dim_dtype_float", + (1, 1, 1), + (1, 2, 3), + (3, 3, 3), + torch.float, + torch.float, + ), + ] + ) + def test_dynamic_shape_sin( + self, _, min_shape, opt_shape, max_shape, type, output_type + ): + class sin(nn.Module): + def forward(self, input): + return torch.ops.aten.sin.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( + sin(), input_specs, output_dtypes=[output_type] + ) + if __name__ == "__main__": run_tests() From 9b0527221cdba96101cbf4067aecd793a0445706 Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Tue, 4 Jun 2024 16:07:05 -0700 Subject: [PATCH 2/3] add dynamic shape support for cat/sin/cos --- .../dynamo/conversion/aten_ops_converters.py | 6 +-- tests/py/dynamo/conversion/test_cos_aten.py | 49 +++++++++++++++++++ tests/py/dynamo/conversion/test_sin_aten.py | 49 +++++++++++++++++++ 3 files changed, 101 insertions(+), 3 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 09cc78f7db..2ad12a3b2d 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -219,7 +219,7 @@ def aten_ops_group_norm( ) -@dynamo_tensorrt_converter(torch.ops.aten.cat.default) +@dynamo_tensorrt_converter(torch.ops.aten.cat.default, supports_dynamic_shapes=True) def aten_ops_cat( ctx: ConversionContext, target: Target, @@ -1327,7 +1327,7 @@ def aten_ops_abs( ) -@dynamo_tensorrt_converter(torch.ops.aten.sin.default) +@dynamo_tensorrt_converter(torch.ops.aten.sin.default, supports_dynamic_shapes=True) def aten_ops_sin( ctx: ConversionContext, target: Target, @@ -1344,7 +1344,7 @@ def aten_ops_sin( ) -@dynamo_tensorrt_converter(torch.ops.aten.cos.default) +@dynamo_tensorrt_converter(torch.ops.aten.cos.default, supports_dynamic_shapes=True) def aten_ops_cos( ctx: ConversionContext, target: Target, diff --git a/tests/py/dynamo/conversion/test_cos_aten.py b/tests/py/dynamo/conversion/test_cos_aten.py index 505f303219..152254324e 100644 --- a/tests/py/dynamo/conversion/test_cos_aten.py +++ b/tests/py/dynamo/conversion/test_cos_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 @@ -44,6 +45,54 @@ def forward(self, input): inputs, ) + @parameterized.expand( + [ + ( + "3d_dim_dtype_int32", + (3, 2, 1), + (3, 2, 3), + (3, 3, 4), + torch.int32, + torch.float32, + ), + ( + "2d_dim_dtype_float16", + (1, 1), + (2, 2), + (4, 4), + torch.float16, + torch.float16, + ), + ( + "3d_dim_dtype_float", + (1, 1, 1), + (1, 2, 3), + (3, 3, 3), + torch.float, + torch.float, + ), + ] + ) + def test_dynamic_shape_cos( + self, _, min_shape, opt_shape, max_shape, type, output_type + ): + class cos(nn.Module): + def forward(self, input): + return torch.ops.aten.cos.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( + cos(), input_specs, output_dtypes=[output_type] + ) + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_sin_aten.py b/tests/py/dynamo/conversion/test_sin_aten.py index 70a3dc0315..7d5ca2204d 100644 --- a/tests/py/dynamo/conversion/test_sin_aten.py +++ b/tests/py/dynamo/conversion/test_sin_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 @@ -44,6 +45,54 @@ def forward(self, input): inputs, ) + @parameterized.expand( + [ + ( + "3d_dim_dtype_int32", + (3, 2, 1), + (3, 2, 3), + (3, 3, 4), + torch.int32, + torch.float32, + ), + ( + "2d_dim_dtype_float16", + (1, 1), + (2, 2), + (4, 4), + torch.float16, + torch.float16, + ), + ( + "3d_dim_dtype_float", + (1, 1, 1), + (1, 2, 3), + (3, 3, 3), + torch.float, + torch.float, + ), + ] + ) + def test_dynamic_shape_sin( + self, _, min_shape, opt_shape, max_shape, type, output_type + ): + class sin(nn.Module): + def forward(self, input): + return torch.ops.aten.sin.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( + sin(), input_specs, output_dtypes=[output_type] + ) + if __name__ == "__main__": run_tests() From ed18394d7c5dff21325c3355a720bfe76cf92dd3 Mon Sep 17 00:00:00 2001 From: Lan Luo Date: Wed, 5 Jun 2024 14:22:00 -0700 Subject: [PATCH 3/3] Replace shape_ranges to min_shape/opt_shape/max_shape --- tests/py/dynamo/conversion/test_cat_aten.py | 50 +++++++++++++++++---- 1 file changed, 42 insertions(+), 8 deletions(-) diff --git a/tests/py/dynamo/conversion/test_cat_aten.py b/tests/py/dynamo/conversion/test_cat_aten.py index a8d8bae42f..a9e4a45c81 100644 --- a/tests/py/dynamo/conversion/test_cat_aten.py +++ b/tests/py/dynamo/conversion/test_cat_aten.py @@ -38,14 +38,16 @@ def forward(self, x, y): input_specs = [ Input( - shape=(16, -1, 3), dtype=torch.float32, - shape_ranges=[((16, 2, 3), (16, 3, 3), (16, 32, 3))], + min_shape=(16, 2, 3), + opt_shape=(16, 3, 3), + max_shape=(16, 32, 3), ), Input( - shape=(16, -1, 3), dtype=torch.float32, - shape_ranges=[((16, 2, 3), (16, 16, 3), (16, 32, 3))], + min_shape=(16, 2, 3), + opt_shape=(16, 16, 3), + max_shape=(16, 32, 3), ), ] self.run_test_with_dynamic_shape( @@ -71,14 +73,46 @@ def forward(self, x, y): input_specs = [ Input( - shape=(-1, 16, 3), dtype=torch.float32, - shape_ranges=[((2, 16, 3), (3, 16, 3), (32, 16, 3))], + min_shape=(2, 16, 3), + opt_shape=(3, 16, 3), + max_shape=(32, 16, 3), ), Input( - shape=(-1, 16, 3), dtype=torch.float32, - shape_ranges=[((2, 16, 3), (3, 16, 3), (32, 16, 3))], + min_shape=(2, 16, 3), + opt_shape=(3, 16, 3), + max_shape=(32, 16, 3), + ), + ] + self.run_test_with_dynamic_shape( + Cat(), + input_specs, + ) + + @parameterized.expand( + [ + ("pos", 1), + ("neg", -2), + ] + ) + def test_cat_dynamic_shape_dim(self, _, dim): + class Cat(nn.Module): + def forward(self, x, y): + return torch.ops.aten.cat.default((x, y), dim) + + input_specs = [ + Input( + dtype=torch.float32, + min_shape=(2, 1, 1), + opt_shape=(3, 1, 2), + max_shape=(4, 1, 3), + ), + Input( + dtype=torch.float32, + min_shape=(2, 2, 1), + opt_shape=(3, 3, 2), + max_shape=(4, 4, 3), ), ] self.run_test_with_dynamic_shape(