From 2cf4eaaae6abe2181dc45c942d4c76c7a3d12e5a Mon Sep 17 00:00:00 2001 From: kee hyun an Date: Wed, 26 Jun 2024 11:20:38 +0900 Subject: [PATCH 1/3] dynamic shape for squeeze ops Not support for dds(data-dependent shape) input --- .../dynamo/conversion/impl/squeeze.py | 97 ++++++++++++++++--- .../py/dynamo/conversion/test_squeeze_aten.py | 39 ++++++-- 2 files changed, 117 insertions(+), 19 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py b/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py index 45bdefcd80..9e81a66d0a 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py @@ -1,12 +1,17 @@ from typing import Optional, Sequence, Union +import numpy as np +import tensorrt as trt from torch.fx.node import Target from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext -from torch_tensorrt.dynamo.conversion.converter_utils import get_positive_dim -from torch_tensorrt.fx.converters.converter_utils import set_layer_name +from torch_tensorrt.dynamo.conversion.converter_utils import ( + cast_trt_tensor, + get_positive_dim, + set_layer_name, +) +from torch_tensorrt.dynamo.conversion.impl.elementwise import ne from torch_tensorrt.fx.types import TRTTensor -from torch_tensorrt.fx.utils import get_dynamic_dims def squeeze( @@ -29,24 +34,90 @@ def squeeze( dims.append(dim) new_dims = [] + dim_has_dynamic_shape = False for dim in dims: dim = get_positive_dim( dim, len(input.shape), ) - assert input.shape[dim] != -1, "We don't support squeeze dynamic dim." - assert ( - len(get_dynamic_dims(input.shape)) <= 1 - ), "Currently more than one dynamic dim for input to squeeze is not supported." + if input.shape[dim] == -1: + dim_has_dynamic_shape = True new_dims.append(dim) - output_shape = [] - for i, s in enumerate(input.shape): - if (i in new_dims) and s == 1: - continue - output_shape.append(s) layer = ctx.net.add_shuffle(input) - layer.reshape_dims = tuple(output_shape) set_layer_name(layer, target, name, source_ir) + if dim_has_dynamic_shape: + num_shape = len(input.shape) + + tensor_shape_layer = ctx.net.add_shape(input) + tensor_shape = tensor_shape_layer.get_output(0) + tensor_shape = cast_trt_tensor( + ctx, tensor_shape, trt.int32, name + "shape_casted", "shape" + ) + + # change it to get_trt_tensor + one_layer = ctx.net.add_constant( + (num_shape,), + np.ascontiguousarray([1] * num_shape, np.int32), + ) + set_layer_name(one_layer, target, name + "_one", source_ir) + + zero_layer = ctx.net.add_constant( + (num_shape,), + np.zeros((num_shape,), dtype=np.int32), + ) + set_layer_name(zero_layer, target, name + "_zero", source_ir) + + # append last element value + num_append = num_shape - len(new_dims) + if num_append > 0: + new_dims += [new_dims[-1]] * num_append + + index_value = np.array(new_dims, dtype=np.int32) + index_layer = ctx.net.add_constant(index_value.shape, index_value) + set_layer_name(index_layer, target, name + "_index", source_ir) + + scatter_layer = ctx.net.add_scatter( + zero_layer.get_output(0), + index_layer.get_output(0), + one_layer.get_output(0), + trt.ScatterMode.ELEMENT, + ) + set_layer_name(scatter_layer, target, name + "_scatter", source_ir) + + # [1, 2, 1, 3, 1] + # [0, 0, 1, 1, 1] + # [t, t, f, t, f] + ne_tensor = ne( + ctx, + target, + source_ir, + name + "_ne", + tensor_shape, + scatter_layer.get_output(0), + ) + + # [t, t, f, t, f] -> [0, 1, 3] + non_zero_layer = ctx.net.add_non_zero(ne_tensor) + set_layer_name(non_zero_layer, target, name + "_non_zero", source_ir) + + non_zero_shuffle_layer = ctx.net.add_shuffle(non_zero_layer.get_output(0)) + set_layer_name(non_zero_shuffle_layer, target, name + "_shuffle", source_ir) + non_zero_shuffle_layer.second_transpose = (1, 0) + + # (1,2,1,3,1) + [0, 1, 3 ,4] -> [1, 2, 3, 1] + gather_layer = ctx.net.add_gather_v2( + tensor_shape, non_zero_shuffle_layer.get_output(0), mode=trt.GatherMode.ND + ) + set_layer_name(gather_layer, target, name + "_gather", source_ir) + + layer.set_input(1, gather_layer.get_output(0)) + else: + output_shape = [] + for i, s in enumerate(input.shape): + if (i in new_dims) and s == 1: + continue + output_shape.append(s) + layer.reshape_dims = tuple(output_shape) return layer.get_output(0) diff --git a/tests/py/dynamo/conversion/test_squeeze_aten.py b/tests/py/dynamo/conversion/test_squeeze_aten.py index 88483072ae..aa692f2710 100644 --- a/tests/py/dynamo/conversion/test_squeeze_aten.py +++ b/tests/py/dynamo/conversion/test_squeeze_aten.py @@ -46,25 +46,52 @@ def forward(self, x): class TestSqueezeConverter(DispatchTestCase): @parameterized.expand( [ - ("2d_dim", (1), (-1, 1), [((1, 1), (1, 1), (3, 1))]), - ("3d_one_dim", (1), (-1, 2, 1), [((1, 2, 1), (1, 2, 1), (3, 2, 1))]), + ( + "2d_dim", + (1), + (1, 1), + (1, 1), + (3, 1), + torch.half, + torch.half, + ), + ( + "3d_one_dim", + (1), + (1, 2, 1), + (1, 2, 1), + (3, 2, 1), + torch.float, + torch.float, + ), + ( + "3d_one_dim_dynamic", + (0), + (1, 2, 1), + (1, 2, 1), + (3, 2, 1), + torch.float, + torch.float, + ), ] ) - def test_squeeze(self, _, dim, init_size, shape_range): + def test_squeeze(self, _, dim, min_shape, opt_shape, max_shape, type, output_type): class Squeeze(nn.Module): def forward(self, x): return torch.ops.aten.squeeze.dim(x, dim) input_specs = [ Input( - shape=init_size, - dtype=torch.float32, - shape_ranges=shape_range, + min_shape=min_shape, + opt_shape=opt_shape, + max_shape=max_shape, + dtype=type, ), ] self.run_test_with_dynamic_shape( Squeeze(), input_specs, + output_dtypes=[output_type], ) From dafabaa9b3e12b76ca29d5800e1d564171921f85 Mon Sep 17 00:00:00 2001 From: kee hyun an Date: Wed, 10 Jul 2024 22:51:42 +0900 Subject: [PATCH 2/3] More than one dynamic dim for input to squeeze is supported --- .../dynamo/conversion/aten_ops_converters.py | 4 +- .../dynamo/conversion/impl/squeeze.py | 95 +++---------------- .../py/dynamo/conversion/test_squeeze_aten.py | 52 +++++----- 3 files changed, 43 insertions(+), 108 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index ec9d6b8c78..f6bafe24c1 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -622,8 +622,8 @@ def aten_ops_quantize_fp8( ) -@dynamo_tensorrt_converter(torch.ops.aten.squeeze.dim) -@dynamo_tensorrt_converter(torch.ops.aten.squeeze.dims) +@dynamo_tensorrt_converter(torch.ops.aten.squeeze.dim, supports_dynamic_shapes=True) +@dynamo_tensorrt_converter(torch.ops.aten.squeeze.dims, supports_dynamic_shapes=True) def aten_ops_squeeze( ctx: ConversionContext, target: Target, diff --git a/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py b/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py index 9e81a66d0a..7538ca9521 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py @@ -1,16 +1,12 @@ from typing import Optional, Sequence, Union -import numpy as np -import tensorrt as trt from torch.fx.node import Target from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext from torch_tensorrt.dynamo.conversion.converter_utils import ( - cast_trt_tensor, get_positive_dim, set_layer_name, ) -from torch_tensorrt.dynamo.conversion.impl.elementwise import ne from torch_tensorrt.fx.types import TRTTensor @@ -34,90 +30,29 @@ def squeeze( dims.append(dim) new_dims = [] - dim_has_dynamic_shape = False for dim in dims: dim = get_positive_dim( dim, len(input.shape), ) - if input.shape[dim] == -1: - dim_has_dynamic_shape = True + assert input.shape[dim] != -1, "We don't support squeeze dynamic dim." new_dims.append(dim) + dim_to_remove = [] + new_permutation = [] + for i, s in enumerate(input.shape): + if (i in new_dims) and s == 1: + dim_to_remove.append(i) + else: + new_permutation.append(i) + # If number of reshape dimensions is less than input, 0s are resolved by aligning + # the most significant dimensions of input + output_shape = tuple([0] * len(new_permutation)) + new_permutation += dim_to_remove + layer = ctx.net.add_shuffle(input) + layer.first_transpose = new_permutation + layer.reshape_dims = output_shape set_layer_name(layer, target, name, source_ir) - if dim_has_dynamic_shape: - num_shape = len(input.shape) - - tensor_shape_layer = ctx.net.add_shape(input) - tensor_shape = tensor_shape_layer.get_output(0) - tensor_shape = cast_trt_tensor( - ctx, tensor_shape, trt.int32, name + "shape_casted", "shape" - ) - - # change it to get_trt_tensor - one_layer = ctx.net.add_constant( - (num_shape,), - np.ascontiguousarray([1] * num_shape, np.int32), - ) - set_layer_name(one_layer, target, name + "_one", source_ir) - - zero_layer = ctx.net.add_constant( - (num_shape,), - np.zeros((num_shape,), dtype=np.int32), - ) - set_layer_name(zero_layer, target, name + "_zero", source_ir) - - # append last element value - num_append = num_shape - len(new_dims) - if num_append > 0: - new_dims += [new_dims[-1]] * num_append - - index_value = np.array(new_dims, dtype=np.int32) - index_layer = ctx.net.add_constant(index_value.shape, index_value) - set_layer_name(index_layer, target, name + "_index", source_ir) - - scatter_layer = ctx.net.add_scatter( - zero_layer.get_output(0), - index_layer.get_output(0), - one_layer.get_output(0), - trt.ScatterMode.ELEMENT, - ) - set_layer_name(scatter_layer, target, name + "_scatter", source_ir) - - # [1, 2, 1, 3, 1] - # [0, 0, 1, 1, 1] - # [t, t, f, t, f] - ne_tensor = ne( - ctx, - target, - source_ir, - name + "_ne", - tensor_shape, - scatter_layer.get_output(0), - ) - - # [t, t, f, t, f] -> [0, 1, 3] - non_zero_layer = ctx.net.add_non_zero(ne_tensor) - set_layer_name(non_zero_layer, target, name + "_non_zero", source_ir) - - non_zero_shuffle_layer = ctx.net.add_shuffle(non_zero_layer.get_output(0)) - set_layer_name(non_zero_shuffle_layer, target, name + "_shuffle", source_ir) - non_zero_shuffle_layer.second_transpose = (1, 0) - - # (1,2,1,3,1) + [0, 1, 3 ,4] -> [1, 2, 3, 1] - gather_layer = ctx.net.add_gather_v2( - tensor_shape, non_zero_shuffle_layer.get_output(0), mode=trt.GatherMode.ND - ) - set_layer_name(gather_layer, target, name + "_gather", source_ir) - - layer.set_input(1, gather_layer.get_output(0)) - else: - output_shape = [] - for i, s in enumerate(input.shape): - if (i in new_dims) and s == 1: - continue - output_shape.append(s) - layer.reshape_dims = tuple(output_shape) return layer.get_output(0) diff --git a/tests/py/dynamo/conversion/test_squeeze_aten.py b/tests/py/dynamo/conversion/test_squeeze_aten.py index aa692f2710..9dc786525b 100644 --- a/tests/py/dynamo/conversion/test_squeeze_aten.py +++ b/tests/py/dynamo/conversion/test_squeeze_aten.py @@ -43,55 +43,55 @@ def forward(self, x): ) -class TestSqueezeConverter(DispatchTestCase): +class TestSqueezeConverterDynamic(DispatchTestCase): @parameterized.expand( [ ( - "2d_dim", - (1), - (1, 1), - (1, 1), - (3, 1), - torch.half, - torch.half, + "5d_two_dynamic_shape_-1", + (0,), + (1, 1, 1, 1, 1), + (1, 2, 1, 2, 1), + (1, 4, 1, 3, 1), + ), + ( + "5d_two_dynamic_shape_-2", + (0, 2), + (1, 1, 1, 1, 1), + (1, 2, 1, 2, 1), + (1, 4, 1, 3, 1), ), ( - "3d_one_dim", - (1), - (1, 2, 1), - (1, 2, 1), - (3, 2, 1), - torch.float, - torch.float, + "5d_three_dynamic_shape_-2", + (0, 4), + (1, 1, 1, 1, 1), + (1, 2, 4, 2, 1), + (1, 4, 4, 3, 1), ), ( - "3d_one_dim_dynamic", - (0), - (1, 2, 1), - (1, 2, 1), - (3, 2, 1), - torch.float, - torch.float, + "4d_two_dynamic_shape_-2", + (0, 2), + (1, 1, 2, 1), + (1, 2, 2, 2), + (1, 4, 2, 3), ), ] ) - def test_squeeze(self, _, dim, min_shape, opt_shape, max_shape, type, output_type): + def test_squeeze(self, _, dim, min_shape, opt_shape, max_shape): class Squeeze(nn.Module): def forward(self, x): - return torch.ops.aten.squeeze.dim(x, dim) + return torch.ops.aten.squeeze.dims(x, dim) input_specs = [ Input( min_shape=min_shape, opt_shape=opt_shape, max_shape=max_shape, - dtype=type, + dtype=torch.float, ), ] self.run_test_with_dynamic_shape( Squeeze(), input_specs, - output_dtypes=[output_type], ) From 2cd7630fa02c23b73f6474dfd7550d667cede0fc Mon Sep 17 00:00:00 2001 From: kee hyun an Date: Sat, 13 Jul 2024 09:51:22 +0900 Subject: [PATCH 3/3] chore: rename dim variable in loop --- py/torch_tensorrt/dynamo/conversion/impl/squeeze.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py b/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py index 7538ca9521..dd6a2b9863 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/squeeze.py @@ -26,8 +26,8 @@ def squeeze( if isinstance(dim, int): dims.append(dim) else: - for dim in dim: - dims.append(dim) + for d in dim: + dims.append(d) new_dims = [] for dim in dims: