From 436406b2f79d4034226960fae36924803d937099 Mon Sep 17 00:00:00 2001 From: Evan Li Date: Fri, 29 Sep 2023 19:30:14 -0700 Subject: [PATCH 1/9] feat: support prod, max, min, and mean via reduce layer --- .../dynamo/conversion/aten_ops_converters.py | 84 ++++++++++++- .../dynamo/conversion/impl/reduce.py | 109 ++++++++++++++++ tests/py/dynamo/conversion/test_prod_aten.py | 117 ++++++++++++++++++ 3 files changed, 308 insertions(+), 2 deletions(-) create mode 100644 tests/py/dynamo/conversion/test_prod_aten.py diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 67ce83469f..5e55f37a53 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -692,6 +692,86 @@ def aten_ops_sum( ) +@dynamo_tensorrt_converter(torch.ops.aten.prod.default) # type: ignore[misc] +@dynamo_tensorrt_converter(torch.ops.aten.prod.dim_int) # type: ignore[misc] +def aten_ops_prod( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.prod( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args_bounds_check(args, 1, replacement=None), + args_bounds_check(args, 2, replacement=False), + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.max.default) # type: ignore[misc] +@dynamo_tensorrt_converter(torch.ops.aten.max.dim) # type: ignore[misc] +def aten_ops_max( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.max( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args_bounds_check(args, 1, replacement=None), + args_bounds_check(args, 2, replacement=False), + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.min.default) # type: ignore[misc] +@dynamo_tensorrt_converter(torch.ops.aten.min.dim) # type: ignore[misc] +def aten_ops_min( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.min( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args_bounds_check(args, 1, replacement=None), + args_bounds_check(args, 2, replacement=False), + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.mean.default) # type: ignore[misc] +@dynamo_tensorrt_converter(torch.ops.aten.mean.dim) # type: ignore[misc] +def aten_ops_mean( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.mean( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args_bounds_check(args, 1, replacement=None), + args_bounds_check(args, 2, replacement=False), + ) + + @dynamo_tensorrt_converter(torch.ops.aten.exp.default) # type: ignore[misc] def aten_ops_exp( ctx: ConversionContext, @@ -1118,7 +1198,7 @@ def aten_ops_mul( @dynamo_tensorrt_converter(torch.ops.aten.maximum.default) # type: ignore[misc] -def aten_ops_max( +def aten_ops_maximum( ctx: ConversionContext, target: Target, args: Tuple[Argument, ...], @@ -1136,7 +1216,7 @@ def aten_ops_max( @dynamo_tensorrt_converter(torch.ops.aten.minimum.default) # type: ignore[misc] -def aten_ops_min( +def aten_ops_minimum( ctx: ConversionContext, target: Target, args: Tuple[Argument, ...], diff --git a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py index 0357962be5..0fbf94f822 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py @@ -53,6 +53,7 @@ def sum( if dim is None: dim = tuple(range(len(input_val.shape))) + layer = ctx.net.add_reduce( input_val, trt.ReduceOperation.SUM, @@ -61,3 +62,111 @@ def sum( ) set_layer_name(layer, target, name, source_ir) return layer.get_output(0) + + +def prod( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_val: TRTTensor, + dim: Optional[Union[int, Sequence[int]]] = None, + keepdim: bool = False, +) -> TRTTensor: + if (isinstance(input_val, TRTTensor)) and ( + input_val.dtype == trt.int8 or input_val.dtype == trt.int32 + ): + input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + + if dim is None: + dim = tuple(range(len(input_val.shape))) + + layer = ctx.net.add_reduce( + input_val, + trt.ReduceOperation.PROD, + axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), + keep_dims=keepdim, + ) + set_layer_name(layer, target, name, source_ir) + return layer.get_output(0) + + +def max( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_val: TRTTensor, + dim: Optional[Union[int, Sequence[int]]] = None, + keepdim: bool = False, +) -> TRTTensor: + if (isinstance(input_val, TRTTensor)) and ( + input_val.dtype == trt.int8 or input_val.dtype == trt.int32 + ): + input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + + if dim is None: + dim = tuple(range(len(input_val.shape))) + + layer = ctx.net.add_reduce( + input_val, + trt.ReduceOperation.MAX, + axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), + keep_dims=keepdim, + ) + set_layer_name(layer, target, name, source_ir) + return layer.get_output(0) + + +def min( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_val: TRTTensor, + dim: Optional[Union[int, Sequence[int]]] = None, + keepdim: bool = False, +) -> TRTTensor: + if (isinstance(input_val, TRTTensor)) and ( + input_val.dtype == trt.int8 or input_val.dtype == trt.int32 + ): + input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + + if dim is None: + dim = tuple(range(len(input_val.shape))) + + layer = ctx.net.add_reduce( + input_val, + trt.ReduceOperation.MIN, + axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), + keep_dims=keepdim, + ) + set_layer_name(layer, target, name, source_ir) + return layer.get_output(0) + + +def mean( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_val: TRTTensor, + dim: Optional[Union[int, Sequence[int]]] = None, + keepdim: bool = False, +) -> TRTTensor: + if (isinstance(input_val, TRTTensor)) and ( + input_val.dtype == trt.int8 or input_val.dtype == trt.int32 + ): + input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + + if dim is None: + dim = tuple(range(len(input_val.shape))) + + layer = ctx.net.add_reduce( + input_val, + trt.ReduceOperation.AVG, + axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), + keep_dims=keepdim, + ) + set_layer_name(layer, target, name, source_ir) + return layer.get_output(0) diff --git a/tests/py/dynamo/conversion/test_prod_aten.py b/tests/py/dynamo/conversion/test_prod_aten.py new file mode 100644 index 0000000000..ada30d6bea --- /dev/null +++ b/tests/py/dynamo/conversion/test_prod_aten.py @@ -0,0 +1,117 @@ +import torch +import torch.nn as nn +from parameterized import parameterized +from torch.testing._internal.common_utils import run_tests + +from .harness import DispatchTestCase + + +class TestProdConverter(DispatchTestCase): + @parameterized.expand( + [ + ((3, 2, 4),), + ((2, 3, 4, 5),), + ((2, 3, 4, 5),), + ((6, 7, 5, 4, 5),), + ] + ) + def test_prod_dim_int_default(self, input_shape): + class Prod(nn.Module): + def forward(self, x): + return torch.prod(x) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Prod(), + inputs, + expected_ops={torch.ops.aten.prod.default}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), 1, True), + ((2, 3, 4, 5), 3, True), + ((2, 3, 4, 5), None, False), + ((6, 7, 5, 4, 5), 4, False), + ((1, 5, 2, 1), -3, False), + ((1, 5, 2, 3), -2, True), + ] + ) + def test_prod_dim_int(self, input_shape, dim, keep_dims): + class Prod(nn.Module): + def forward(self, x): + return torch.prod(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Prod(), + inputs, + expected_ops={torch.ops.aten.prod.dim_int}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), [1], True), + ((2, 1, 4, 5), None, True), + ((2, 3, 4, 5), [0, 1, 2, 3], False), + ((6, 7, 5, 4, 5), [1, 3, 4], False), + ((6, 7, 5, 4, 5), [-5, -4, -2], False), + ] + ) + def test_prod_dim_tuple(self, input_shape, dim, keep_dims): + class Prod(nn.Module): + def forward(self, x): + return torch.prod(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Prod(), + inputs, + expected_ops={torch.ops.aten.prod.dim_int}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), 1, True, torch.int, 0, 5), + ((2, 3, 4, 5), None, True, torch.int, -10, 10), + ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), + ] + ) + def test_prod_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Prod(nn.Module): + def forward(self, x): + return torch.prod(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Prod(), + inputs, + expected_ops={torch.ops.aten.prod.dim_int}, + check_dtype=False, + ) + + @parameterized.expand( + [ + ((3, 2, 4), [1], True, torch.int, 0, 5), + ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), + ((2, 3, 4, 5), None, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), + ] + ) + def test_prod_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Prod(nn.Module): + def forward(self, x): + return torch.prod(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Prod(), + inputs, + expected_ops={torch.ops.aten.prod.dim_int}, + check_dtype=False, + ) + + +if __name__ == "__main__": + run_tests() From 306a9cd1cf73d8f217d6ae4055986aaec11bf97e Mon Sep 17 00:00:00 2001 From: Evan Li Date: Mon, 2 Oct 2023 14:50:22 -0700 Subject: [PATCH 2/9] add tests --- tests/py/dynamo/conversion/test_max_aten.py | 106 ++++++++++++++++-- .../py/dynamo/conversion/test_maximum_aten.py | 31 +++++ tests/py/dynamo/conversion/test_min_aten.py | 105 +++++++++++++++-- .../py/dynamo/conversion/test_minimum_aten.py | 31 +++++ 4 files changed, 254 insertions(+), 19 deletions(-) create mode 100644 tests/py/dynamo/conversion/test_maximum_aten.py create mode 100644 tests/py/dynamo/conversion/test_minimum_aten.py diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index d2247f61dd..3fbec86208 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -2,7 +2,6 @@ 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 @@ -10,20 +9,107 @@ class TestMaxConverter(DispatchTestCase): @parameterized.expand( [ - ("2d", (2, 1)), - ("3d", (2, 1, 2)), + ((3, 2, 4),), + ((2, 3, 4, 5),), + ((2, 3, 4, 5),), + ((6, 7, 5, 4, 5),), ] ) - def test_max(self, _, shape): - class max(nn.Module): - def forward(self, lhs_val, rhs_val): - return torch.ops.aten.maximum.default(lhs_val, rhs_val) + def test_max_dim_int_default(self, input_shape): + class Max(nn.Module): + def forward(self, x): + return torch.max(x) - inputs = [torch.randn(shape), torch.randn(shape)] + inputs = [torch.randn(*input_shape)] self.run_test( - max(), + Max(), inputs, - # expected_ops={torch.ops.aten.maximum.default}, + expected_ops={torch.ops.aten.max.default}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), 1, True), + ((2, 3, 4, 5), 3, True), + ((2, 3, 4, 5), None, False), + ((6, 7, 5, 4, 5), 4, False), + ((1, 5, 2, 1), -3, False), + ((1, 5, 2, 3), -2, True), + ] + ) + def test_max_dim_int(self, input_shape, dim, keep_dims): + class Max(nn.Module): + def forward(self, x): + return torch.max(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.dim}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), [1], True), + ((2, 1, 4, 5), None, True), + ((2, 3, 4, 5), [0, 1, 2, 3], False), + ((6, 7, 5, 4, 5), [1, 3, 4], False), + ((6, 7, 5, 4, 5), [-5, -4, -2], False), + ] + ) + def test_max_dim_tuple(self, input_shape, dim, keep_dims): + class Max(nn.Module): + def forward(self, x): + return torch.max(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.dim}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), 1, True, torch.int, 0, 5), + ((2, 3, 4, 5), None, True, torch.int, -10, 10), + ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), + ] + ) + def test_max_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Max(nn.Module): + def forward(self, x): + return torch.max(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.dim}, + check_dtype=False, + ) + + @parameterized.expand( + [ + ((3, 2, 4), [1], True, torch.int, 0, 5), + ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), + ((2, 3, 4, 5), None, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), + ] + ) + def test_max_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Max(nn.Module): + def forward(self, x): + return torch.max(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.dim}, + check_dtype=False, ) diff --git a/tests/py/dynamo/conversion/test_maximum_aten.py b/tests/py/dynamo/conversion/test_maximum_aten.py new file mode 100644 index 0000000000..2be1d9c74b --- /dev/null +++ b/tests/py/dynamo/conversion/test_maximum_aten.py @@ -0,0 +1,31 @@ +import torch +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 + + +class TestMaxConverter(DispatchTestCase): + @parameterized.expand( + [ + ("2d", (2, 1)), + ("3d", (2, 1, 2)), + ] + ) + def test_max(self, _, shape): + class max(nn.Module): + def forward(self, lhs_val, rhs_val): + return torch.max(lhs_val, rhs_val) + + inputs = [torch.randn(shape), torch.randn(shape)] + self.run_test( + max(), + inputs, + expected_ops={torch.ops.aten.maximum.default}, + ) + + +if __name__ == "__main__": + run_tests() diff --git a/tests/py/dynamo/conversion/test_min_aten.py b/tests/py/dynamo/conversion/test_min_aten.py index 49853cb111..7e125a7d7d 100644 --- a/tests/py/dynamo/conversion/test_min_aten.py +++ b/tests/py/dynamo/conversion/test_min_aten.py @@ -2,7 +2,6 @@ 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 @@ -10,19 +9,107 @@ class TestMinConverter(DispatchTestCase): @parameterized.expand( [ - ("2d", (2, 1)), - ("3d", (2, 1, 2)), + ((3, 2, 4),), + ((2, 3, 4, 5),), + ((2, 3, 4, 5),), + ((6, 7, 5, 4, 5),), ] ) - def test_min(self, _, shape): - class min(nn.Module): - def forward(self, lhs_val, rhs_val): - return torch.ops.aten.minimum.default(lhs_val, rhs_val) + def test_min_dim_int_default(self, input_shape): + class Min(nn.Module): + def forward(self, x): + return torch.min(x) - inputs = [torch.randn(shape), torch.randn(shape)] + inputs = [torch.randn(*input_shape)] self.run_test( - min(), + Min(), inputs, + expected_ops={torch.ops.aten.min.default}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), 1, True), + ((2, 3, 4, 5), 3, True), + ((2, 3, 4, 5), None, False), + ((6, 7, 5, 4, 5), 4, False), + ((1, 5, 2, 1), -3, False), + ((1, 5, 2, 3), -2, True), + ] + ) + def test_min_dim_int(self, input_shape, dim, keep_dims): + class Min(nn.Module): + def forward(self, x): + return torch.min(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.dim}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), [1], True), + ((2, 1, 4, 5), None, True), + ((2, 3, 4, 5), [0, 1, 2, 3], False), + ((6, 7, 5, 4, 5), [1, 3, 4], False), + ((6, 7, 5, 4, 5), [-5, -4, -2], False), + ] + ) + def test_min_dim_tuple(self, input_shape, dim, keep_dims): + class Min(nn.Module): + def forward(self, x): + return torch.min(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randn(*input_shape)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.dim}, + ) + + @parameterized.expand( + [ + ((3, 2, 4), 1, True, torch.int, 0, 5), + ((2, 3, 4, 5), None, True, torch.int, -10, 10), + ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), + ] + ) + def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Min(nn.Module): + def forward(self, x): + return torch.min(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.dim}, + check_dtype=False, + ) + + @parameterized.expand( + [ + ((3, 2, 4), [1], True, torch.int, 0, 5), + ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), + ((2, 3, 4, 5), None, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), + ] + ) + def test_min_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Min(nn.Module): + def forward(self, x): + return torch.min(x, dim=dim, keepdim=keep_dims) + + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.dim}, + check_dtype=False, ) diff --git a/tests/py/dynamo/conversion/test_minimum_aten.py b/tests/py/dynamo/conversion/test_minimum_aten.py new file mode 100644 index 0000000000..35d0d7163f --- /dev/null +++ b/tests/py/dynamo/conversion/test_minimum_aten.py @@ -0,0 +1,31 @@ +import torch +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 + + +class TestMinConverter(DispatchTestCase): + @parameterized.expand( + [ + ("2d", (2, 1)), + ("3d", (2, 1, 2)), + ] + ) + def test_min(self, _, shape): + class min(nn.Module): + def forward(self, lhs_val, rhs_val): + return torch.min(lhs_val, rhs_val) + + inputs = [torch.randn(shape), torch.randn(shape)] + self.run_test( + min(), + inputs, + expected_ops={torch.ops.aten.minimum.default}, + ) + + +if __name__ == "__main__": + run_tests() From 5723d3abbdcd0c3371f4c634ad54a995ad0119ee Mon Sep 17 00:00:00 2001 From: Evan Li Date: Tue, 3 Oct 2023 19:15:32 -0700 Subject: [PATCH 3/9] comment out max and min converters, will revisit later --- .../dynamo/conversion/aten_ops_converters.py | 83 +++---- .../dynamo/conversion/impl/reduce.py | 104 ++++----- tests/py/dynamo/conversion/test_max_aten.py | 206 ++++++++--------- tests/py/dynamo/conversion/test_min_aten.py | 208 +++++++++--------- 4 files changed, 305 insertions(+), 296 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 5e55f37a53..0609ddcde3 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -712,44 +712,51 @@ def aten_ops_prod( ) -@dynamo_tensorrt_converter(torch.ops.aten.max.default) # type: ignore[misc] -@dynamo_tensorrt_converter(torch.ops.aten.max.dim) # type: ignore[misc] -def aten_ops_max( - ctx: ConversionContext, - target: Target, - args: Tuple[Argument, ...], - kwargs: Dict[str, Argument], - name: str, -) -> Union[TRTTensor, Sequence[TRTTensor]]: - return impl.reduce.max( - ctx, - target, - SourceIR.ATEN, - name, - args[0], - args_bounds_check(args, 1, replacement=None), - args_bounds_check(args, 2, replacement=False), - ) - - -@dynamo_tensorrt_converter(torch.ops.aten.min.default) # type: ignore[misc] -@dynamo_tensorrt_converter(torch.ops.aten.min.dim) # type: ignore[misc] -def aten_ops_min( - ctx: ConversionContext, - target: Target, - args: Tuple[Argument, ...], - kwargs: Dict[str, Argument], - name: str, -) -> Union[TRTTensor, Sequence[TRTTensor]]: - return impl.reduce.min( - ctx, - target, - SourceIR.ATEN, - name, - args[0], - args_bounds_check(args, 1, replacement=None), - args_bounds_check(args, 2, replacement=False), - ) +# def one_user_validator(min_node: Node) -> bool: +# # Validate only one user, which is a getitem node that accesses the first element in the list +# return (len(min_node.users) == 1 and +# list(min_node.users)[0].target == operator.getitem and +# list(min_node.users)[0].args[1] == 0) + + +# @dynamo_tensorrt_converter(torch.ops.aten.max.default) # type: ignore[misc] +# @dynamo_tensorrt_converter(torch.ops.aten.max.dim) # type: ignore[misc] +# def aten_ops_max( +# ctx: ConversionContext, +# target: Target, +# args: Tuple[Argument, ...], +# kwargs: Dict[str, Argument], +# name: str, +# ) -> Union[TRTTensor, Sequence[TRTTensor]]: +# return impl.reduce.max( +# ctx, +# target, +# SourceIR.ATEN, +# name, +# args[0], +# args_bounds_check(args, 1, replacement=None), +# args_bounds_check(args, 2, replacement=False), +# ) + + +# @dynamo_tensorrt_converter(torch.ops.aten.min.default) # type: ignore[misc] +# @dynamo_tensorrt_converter(torch.ops.aten.min.dim) # type: ignore[misc] +# def aten_ops_min( +# ctx: ConversionContext, +# target: Target, +# args: Tuple[Argument, ...], +# kwargs: Dict[str, Argument], +# name: str, +# ) -> Union[TRTTensor, Sequence[TRTTensor]]: +# return impl.reduce.min( +# ctx, +# target, +# SourceIR.ATEN, +# name, +# args[0], +# args_bounds_check(args, 1, replacement=None), +# args_bounds_check(args, 2, replacement=False), +# ) @dynamo_tensorrt_converter(torch.ops.aten.mean.default) # type: ignore[misc] diff --git a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py index 0fbf94f822..17dc2c988f 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py @@ -91,58 +91,58 @@ def prod( return layer.get_output(0) -def max( - ctx: ConversionContext, - target: Target, - source_ir: Optional[SourceIR], - name: str, - input_val: TRTTensor, - dim: Optional[Union[int, Sequence[int]]] = None, - keepdim: bool = False, -) -> TRTTensor: - if (isinstance(input_val, TRTTensor)) and ( - input_val.dtype == trt.int8 or input_val.dtype == trt.int32 - ): - input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) - - if dim is None: - dim = tuple(range(len(input_val.shape))) - - layer = ctx.net.add_reduce( - input_val, - trt.ReduceOperation.MAX, - axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), - keep_dims=keepdim, - ) - set_layer_name(layer, target, name, source_ir) - return layer.get_output(0) - - -def min( - ctx: ConversionContext, - target: Target, - source_ir: Optional[SourceIR], - name: str, - input_val: TRTTensor, - dim: Optional[Union[int, Sequence[int]]] = None, - keepdim: bool = False, -) -> TRTTensor: - if (isinstance(input_val, TRTTensor)) and ( - input_val.dtype == trt.int8 or input_val.dtype == trt.int32 - ): - input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) - - if dim is None: - dim = tuple(range(len(input_val.shape))) - - layer = ctx.net.add_reduce( - input_val, - trt.ReduceOperation.MIN, - axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), - keep_dims=keepdim, - ) - set_layer_name(layer, target, name, source_ir) - return layer.get_output(0) +# def max( +# ctx: ConversionContext, +# target: Target, +# source_ir: Optional[SourceIR], +# name: str, +# input_val: TRTTensor, +# dim: Optional[Union[int, Sequence[int]]] = None, +# keepdim: bool = False, +# ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: +# if (isinstance(input_val, TRTTensor)) and ( +# input_val.dtype == trt.int8 or input_val.dtype == trt.int32 +# ): +# input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + +# if dim is None: +# dim = tuple(range(len(input_val.shape))) + +# layer = ctx.net.add_reduce( +# input_val, +# trt.ReduceOperation.MAX, +# axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), +# keep_dims=keepdim, +# ) +# set_layer_name(layer, target, name, source_ir) +# return layer.get_output(0), torch.tensor(0) + + +# def min( +# ctx: ConversionContext, +# target: Target, +# source_ir: Optional[SourceIR], +# name: str, +# input_val: TRTTensor, +# dim: Optional[Union[int, Sequence[int]]] = None, +# keepdim: bool = False, +# ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: +# if (isinstance(input_val, TRTTensor)) and ( +# input_val.dtype == trt.int8 or input_val.dtype == trt.int32 +# ): +# input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + +# if dim is None: +# dim = tuple(range(len(input_val.shape))) + +# layer = ctx.net.add_reduce( +# input_val, +# trt.ReduceOperation.MIN, +# axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), +# keep_dims=keepdim, +# ) +# set_layer_name(layer, target, name, source_ir) +# return layer.get_output(0), torch.randint(0, 3, list(layer.get_output(0).shape)) def mean( diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index 3fbec86208..e9b66987dd 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -1,117 +1,117 @@ -import torch -import torch.nn as nn -from parameterized import parameterized -from torch.testing._internal.common_utils import run_tests +# import torch +# import torch.nn as nn +# from parameterized import parameterized +# from torch.testing._internal.common_utils import run_tests -from .harness import DispatchTestCase +# from .harness import DispatchTestCase -class TestMaxConverter(DispatchTestCase): - @parameterized.expand( - [ - ((3, 2, 4),), - ((2, 3, 4, 5),), - ((2, 3, 4, 5),), - ((6, 7, 5, 4, 5),), - ] - ) - def test_max_dim_int_default(self, input_shape): - class Max(nn.Module): - def forward(self, x): - return torch.max(x) +# class TestMaxConverter(DispatchTestCase): +# @parameterized.expand( +# [ +# ((3, 2, 4),), +# ((2, 3, 4, 5),), +# ((2, 3, 4, 5),), +# ((6, 7, 5, 4, 5),), +# ] +# ) +# def test_max_dim_int_default(self, input_shape): +# class Max(nn.Module): +# def forward(self, x): +# return torch.max(x) - inputs = [torch.randn(*input_shape)] - self.run_test( - Max(), - inputs, - expected_ops={torch.ops.aten.max.default}, - ) +# inputs = [torch.randn(*input_shape)] +# self.run_test( +# Max(), +# inputs, +# expected_ops={torch.ops.aten.max.default}, +# ) - @parameterized.expand( - [ - ((3, 2, 4), 1, True), - ((2, 3, 4, 5), 3, True), - ((2, 3, 4, 5), None, False), - ((6, 7, 5, 4, 5), 4, False), - ((1, 5, 2, 1), -3, False), - ((1, 5, 2, 3), -2, True), - ] - ) - def test_max_dim_int(self, input_shape, dim, keep_dims): - class Max(nn.Module): - def forward(self, x): - return torch.max(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), 1, True), +# ((2, 3, 4, 5), 3, True), +# ((2, 3, 4, 5), None, False), +# ((6, 7, 5, 4, 5), 4, False), +# ((1, 5, 2, 1), -3, False), +# ((1, 5, 2, 3), -2, True), +# ] +# ) +# def test_max_dim_int(self, input_shape, dim, keep_dims): +# class Max(nn.Module): +# def forward(self, x): +# return torch.max(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randn(*input_shape)] - self.run_test( - Max(), - inputs, - expected_ops={torch.ops.aten.max.dim}, - ) +# inputs = [torch.randn(*input_shape)] +# self.run_test( +# Max(), +# inputs, +# expected_ops={torch.ops.aten.max.dim}, +# ) - @parameterized.expand( - [ - ((3, 2, 4), [1], True), - ((2, 1, 4, 5), None, True), - ((2, 3, 4, 5), [0, 1, 2, 3], False), - ((6, 7, 5, 4, 5), [1, 3, 4], False), - ((6, 7, 5, 4, 5), [-5, -4, -2], False), - ] - ) - def test_max_dim_tuple(self, input_shape, dim, keep_dims): - class Max(nn.Module): - def forward(self, x): - return torch.max(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), [1], True), +# ((2, 1, 4, 5), None, True), +# ((2, 3, 4, 5), [0, 1, 2, 3], False), +# ((6, 7, 5, 4, 5), [1, 3, 4], False), +# ((6, 7, 5, 4, 5), [-5, -4, -2], False), +# ] +# ) +# def test_max_dim_tuple(self, input_shape, dim, keep_dims): +# class Max(nn.Module): +# def forward(self, x): +# return torch.max(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randn(*input_shape)] - self.run_test( - Max(), - inputs, - expected_ops={torch.ops.aten.max.dim}, - ) +# inputs = [torch.randn(*input_shape)] +# self.run_test( +# Max(), +# inputs, +# expected_ops={torch.ops.aten.max.dim}, +# ) - @parameterized.expand( - [ - ((3, 2, 4), 1, True, torch.int, 0, 5), - ((2, 3, 4, 5), None, True, torch.int, -10, 10), - ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), - ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), - ] - ) - def test_max_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): - class Max(nn.Module): - def forward(self, x): - return torch.max(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), 1, True, torch.int, 0, 5), +# ((2, 3, 4, 5), None, True, torch.int, -10, 10), +# ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), +# ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), +# ] +# ) +# def test_max_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): +# class Max(nn.Module): +# def forward(self, x): +# return torch.max(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randint(low, high, input_shape, dtype=dtype)] - self.run_test( - Max(), - inputs, - expected_ops={torch.ops.aten.max.dim}, - check_dtype=False, - ) +# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] +# self.run_test( +# Max(), +# inputs, +# expected_ops={torch.ops.aten.max.dim}, +# check_dtype=False, +# ) - @parameterized.expand( - [ - ((3, 2, 4), [1], True, torch.int, 0, 5), - ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), - ((2, 3, 4, 5), None, False, torch.int32, -5, 0), - ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), - ] - ) - def test_max_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): - class Max(nn.Module): - def forward(self, x): - return torch.max(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), [1], True, torch.int, 0, 5), +# ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), +# ((2, 3, 4, 5), None, False, torch.int32, -5, 0), +# ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), +# ] +# ) +# def test_max_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): +# class Max(nn.Module): +# def forward(self, x): +# return torch.max(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randint(low, high, input_shape, dtype=dtype)] - self.run_test( - Max(), - inputs, - expected_ops={torch.ops.aten.max.dim}, - check_dtype=False, - ) +# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] +# self.run_test( +# Max(), +# inputs, +# expected_ops={torch.ops.aten.max.dim}, +# check_dtype=False, +# ) -if __name__ == "__main__": - run_tests() +# if __name__ == "__main__": +# run_tests() diff --git a/tests/py/dynamo/conversion/test_min_aten.py b/tests/py/dynamo/conversion/test_min_aten.py index 7e125a7d7d..7507407263 100644 --- a/tests/py/dynamo/conversion/test_min_aten.py +++ b/tests/py/dynamo/conversion/test_min_aten.py @@ -1,117 +1,119 @@ -import torch -import torch.nn as nn -from parameterized import parameterized -from torch.testing._internal.common_utils import run_tests +# import torch +# import torch.nn as nn +# from parameterized import parameterized +# from torch.testing._internal.common_utils import run_tests -from .harness import DispatchTestCase +# from .harness import DispatchTestCase -class TestMinConverter(DispatchTestCase): - @parameterized.expand( - [ - ((3, 2, 4),), - ((2, 3, 4, 5),), - ((2, 3, 4, 5),), - ((6, 7, 5, 4, 5),), - ] - ) - def test_min_dim_int_default(self, input_shape): - class Min(nn.Module): - def forward(self, x): - return torch.min(x) +# class TestMinConverter(DispatchTestCase): +# @parameterized.expand( +# [ +# ((3, 2, 4),), +# ((2, 3, 4, 5),), +# ((2, 3, 4, 5),), +# ((6, 7, 5, 4, 5),), +# ] +# ) +# def test_min_dim_int_default(self, input_shape): +# class Min(nn.Module): +# def forward(self, x): +# return torch.min(x) - inputs = [torch.randn(*input_shape)] - self.run_test( - Min(), - inputs, - expected_ops={torch.ops.aten.min.default}, - ) +# inputs = [torch.randn(*input_shape)] +# self.run_test( +# Min(), +# inputs, +# expected_ops={torch.ops.aten.min.default}, +# ) - @parameterized.expand( - [ - ((3, 2, 4), 1, True), - ((2, 3, 4, 5), 3, True), - ((2, 3, 4, 5), None, False), - ((6, 7, 5, 4, 5), 4, False), - ((1, 5, 2, 1), -3, False), - ((1, 5, 2, 3), -2, True), - ] - ) - def test_min_dim_int(self, input_shape, dim, keep_dims): - class Min(nn.Module): - def forward(self, x): - return torch.min(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), 1, True), +# ((2, 3, 4, 5), 3, True), +# ((2, 3, 4, 5), None, False), +# ((6, 7, 5, 4, 5), 4, False), +# ((1, 5, 2, 1), -3, False), +# ((1, 5, 2, 3), -2, True), +# ] +# ) +# def test_min_dim_int(self, input_shape, dim, keep_dims): +# class Min(nn.Module): +# def forward(self, x): +# return torch.min(x, dim=dim, keepdim=keep_dims) +# # return torch.ops.aten.min.dim(x, dim, keep_dims) - inputs = [torch.randn(*input_shape)] - self.run_test( - Min(), - inputs, - expected_ops={torch.ops.aten.min.dim}, - ) +# inputs = [torch.randn(*input_shape)] +# self.run_test( +# Min(), +# inputs, +# expected_ops={torch.ops.aten.min.dim}, +# disable_passes=True, +# ) - @parameterized.expand( - [ - ((3, 2, 4), [1], True), - ((2, 1, 4, 5), None, True), - ((2, 3, 4, 5), [0, 1, 2, 3], False), - ((6, 7, 5, 4, 5), [1, 3, 4], False), - ((6, 7, 5, 4, 5), [-5, -4, -2], False), - ] - ) - def test_min_dim_tuple(self, input_shape, dim, keep_dims): - class Min(nn.Module): - def forward(self, x): - return torch.min(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), [1], True), +# ((2, 1, 4, 5), None, True), +# ((2, 3, 4, 5), [0, 1, 2, 3], False), +# ((6, 7, 5, 4, 5), [1, 3, 4], False), +# ((6, 7, 5, 4, 5), [-5, -4, -2], False), +# ] +# ) +# def test_min_dim_tuple(self, input_shape, dim, keep_dims): +# class Min(nn.Module): +# def forward(self, x): +# return torch.min(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randn(*input_shape)] - self.run_test( - Min(), - inputs, - expected_ops={torch.ops.aten.min.dim}, - ) +# inputs = [torch.randn(*input_shape)] +# self.run_test( +# Min(), +# inputs, +# expected_ops={torch.ops.aten.min.dim}, +# ) - @parameterized.expand( - [ - ((3, 2, 4), 1, True, torch.int, 0, 5), - ((2, 3, 4, 5), None, True, torch.int, -10, 10), - ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), - ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), - ] - ) - def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): - class Min(nn.Module): - def forward(self, x): - return torch.min(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), 1, True, torch.int, 0, 5), +# ((2, 3, 4, 5), None, True, torch.int, -10, 10), +# ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), +# ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), +# ] +# ) +# def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): +# class Min(nn.Module): +# def forward(self, x): +# return torch.min(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randint(low, high, input_shape, dtype=dtype)] - self.run_test( - Min(), - inputs, - expected_ops={torch.ops.aten.min.dim}, - check_dtype=False, - ) +# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] +# self.run_test( +# Min(), +# inputs, +# expected_ops={torch.ops.aten.min.dim}, +# check_dtype=False, +# ) - @parameterized.expand( - [ - ((3, 2, 4), [1], True, torch.int, 0, 5), - ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), - ((2, 3, 4, 5), None, False, torch.int32, -5, 0), - ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), - ] - ) - def test_min_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): - class Min(nn.Module): - def forward(self, x): - return torch.min(x, dim=dim, keepdim=keep_dims) +# @parameterized.expand( +# [ +# ((3, 2, 4), [1], True, torch.int, 0, 5), +# ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), +# ((2, 3, 4, 5), None, False, torch.int32, -5, 0), +# ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), +# ] +# ) +# def test_min_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): +# class Min(nn.Module): +# def forward(self, x): +# return torch.min(x, dim=dim, keepdim=keep_dims) - inputs = [torch.randint(low, high, input_shape, dtype=dtype)] - self.run_test( - Min(), - inputs, - expected_ops={torch.ops.aten.min.dim}, - check_dtype=False, - ) +# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] +# self.run_test( +# Min(), +# inputs, +# expected_ops={torch.ops.aten.min.dim}, +# check_dtype=False, +# ) -if __name__ == "__main__": - run_tests() +# if __name__ == "__main__": +# run_tests() From 63236884860de9ba3e87349ec81bcb7eb2e63caf Mon Sep 17 00:00:00 2001 From: Evan Li Date: Tue, 3 Oct 2023 21:33:38 -0700 Subject: [PATCH 4/9] fix bugs --- .../dynamo/conversion/impl/reduce.py | 20 ++++----- tests/py/dynamo/conversion/test_max_aten.py | 43 ------------------ tests/py/dynamo/conversion/test_min_aten.py | 43 ------------------ tests/py/dynamo/conversion/test_prod_aten.py | 44 ------------------- 4 files changed, 10 insertions(+), 140 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py index 17dc2c988f..bd32ad153b 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py @@ -43,8 +43,8 @@ def sum( source_ir: Optional[SourceIR], name: str, input_val: TRTTensor, - dim: Optional[Union[int, Sequence[int]]] = None, - keepdim: bool = False, + dim: Optional[Union[int, Sequence[int]]], + keepdim: bool, ) -> TRTTensor: if (isinstance(input_val, TRTTensor)) and ( input_val.dtype == trt.int8 or input_val.dtype == trt.int32 @@ -70,8 +70,8 @@ def prod( source_ir: Optional[SourceIR], name: str, input_val: TRTTensor, - dim: Optional[Union[int, Sequence[int]]] = None, - keepdim: bool = False, + dim: int, + keepdim: bool, ) -> TRTTensor: if (isinstance(input_val, TRTTensor)) and ( input_val.dtype == trt.int8 or input_val.dtype == trt.int32 @@ -97,8 +97,8 @@ def prod( # source_ir: Optional[SourceIR], # name: str, # input_val: TRTTensor, -# dim: Optional[Union[int, Sequence[int]]] = None, -# keepdim: bool = False, +# dim: Optional[int], +# keepdim: bool, # ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: # if (isinstance(input_val, TRTTensor)) and ( # input_val.dtype == trt.int8 or input_val.dtype == trt.int32 @@ -124,8 +124,8 @@ def prod( # source_ir: Optional[SourceIR], # name: str, # input_val: TRTTensor, -# dim: Optional[Union[int, Sequence[int]]] = None, -# keepdim: bool = False, +# dim: Optional[int], +# keepdim: bool, # ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: # if (isinstance(input_val, TRTTensor)) and ( # input_val.dtype == trt.int8 or input_val.dtype == trt.int32 @@ -151,8 +151,8 @@ def mean( source_ir: Optional[SourceIR], name: str, input_val: TRTTensor, - dim: Optional[Union[int, Sequence[int]]] = None, - keepdim: bool = False, + dim: Optional[Union[int, Sequence[int]]], + keepdim: bool, ) -> TRTTensor: if (isinstance(input_val, TRTTensor)) and ( input_val.dtype == trt.int8 or input_val.dtype == trt.int32 diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index e9b66987dd..7da7610d94 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -31,7 +31,6 @@ # [ # ((3, 2, 4), 1, True), # ((2, 3, 4, 5), 3, True), -# ((2, 3, 4, 5), None, False), # ((6, 7, 5, 4, 5), 4, False), # ((1, 5, 2, 1), -3, False), # ((1, 5, 2, 3), -2, True), @@ -49,31 +48,10 @@ # expected_ops={torch.ops.aten.max.dim}, # ) -# @parameterized.expand( -# [ -# ((3, 2, 4), [1], True), -# ((2, 1, 4, 5), None, True), -# ((2, 3, 4, 5), [0, 1, 2, 3], False), -# ((6, 7, 5, 4, 5), [1, 3, 4], False), -# ((6, 7, 5, 4, 5), [-5, -4, -2], False), -# ] -# ) -# def test_max_dim_tuple(self, input_shape, dim, keep_dims): -# class Max(nn.Module): -# def forward(self, x): -# return torch.max(x, dim=dim, keepdim=keep_dims) - -# inputs = [torch.randn(*input_shape)] -# self.run_test( -# Max(), -# inputs, -# expected_ops={torch.ops.aten.max.dim}, -# ) # @parameterized.expand( # [ # ((3, 2, 4), 1, True, torch.int, 0, 5), -# ((2, 3, 4, 5), None, True, torch.int, -10, 10), # ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), # ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), # ] @@ -91,27 +69,6 @@ # check_dtype=False, # ) -# @parameterized.expand( -# [ -# ((3, 2, 4), [1], True, torch.int, 0, 5), -# ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), -# ((2, 3, 4, 5), None, False, torch.int32, -5, 0), -# ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), -# ] -# ) -# def test_max_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): -# class Max(nn.Module): -# def forward(self, x): -# return torch.max(x, dim=dim, keepdim=keep_dims) - -# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] -# self.run_test( -# Max(), -# inputs, -# expected_ops={torch.ops.aten.max.dim}, -# check_dtype=False, -# ) - # if __name__ == "__main__": # run_tests() diff --git a/tests/py/dynamo/conversion/test_min_aten.py b/tests/py/dynamo/conversion/test_min_aten.py index 7507407263..728c3e632d 100644 --- a/tests/py/dynamo/conversion/test_min_aten.py +++ b/tests/py/dynamo/conversion/test_min_aten.py @@ -31,7 +31,6 @@ # [ # ((3, 2, 4), 1, True), # ((2, 3, 4, 5), 3, True), -# ((2, 3, 4, 5), None, False), # ((6, 7, 5, 4, 5), 4, False), # ((1, 5, 2, 1), -3, False), # ((1, 5, 2, 3), -2, True), @@ -51,31 +50,10 @@ # disable_passes=True, # ) -# @parameterized.expand( -# [ -# ((3, 2, 4), [1], True), -# ((2, 1, 4, 5), None, True), -# ((2, 3, 4, 5), [0, 1, 2, 3], False), -# ((6, 7, 5, 4, 5), [1, 3, 4], False), -# ((6, 7, 5, 4, 5), [-5, -4, -2], False), -# ] -# ) -# def test_min_dim_tuple(self, input_shape, dim, keep_dims): -# class Min(nn.Module): -# def forward(self, x): -# return torch.min(x, dim=dim, keepdim=keep_dims) - -# inputs = [torch.randn(*input_shape)] -# self.run_test( -# Min(), -# inputs, -# expected_ops={torch.ops.aten.min.dim}, -# ) # @parameterized.expand( # [ # ((3, 2, 4), 1, True, torch.int, 0, 5), -# ((2, 3, 4, 5), None, True, torch.int, -10, 10), # ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), # ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), # ] @@ -93,27 +71,6 @@ # check_dtype=False, # ) -# @parameterized.expand( -# [ -# ((3, 2, 4), [1], True, torch.int, 0, 5), -# ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), -# ((2, 3, 4, 5), None, False, torch.int32, -5, 0), -# ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), -# ] -# ) -# def test_min_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): -# class Min(nn.Module): -# def forward(self, x): -# return torch.min(x, dim=dim, keepdim=keep_dims) - -# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] -# self.run_test( -# Min(), -# inputs, -# expected_ops={torch.ops.aten.min.dim}, -# check_dtype=False, -# ) - # if __name__ == "__main__": # run_tests() diff --git a/tests/py/dynamo/conversion/test_prod_aten.py b/tests/py/dynamo/conversion/test_prod_aten.py index ada30d6bea..f83464e63a 100644 --- a/tests/py/dynamo/conversion/test_prod_aten.py +++ b/tests/py/dynamo/conversion/test_prod_aten.py @@ -31,7 +31,6 @@ def forward(self, x): [ ((3, 2, 4), 1, True), ((2, 3, 4, 5), 3, True), - ((2, 3, 4, 5), None, False), ((6, 7, 5, 4, 5), 4, False), ((1, 5, 2, 1), -3, False), ((1, 5, 2, 3), -2, True), @@ -49,31 +48,9 @@ def forward(self, x): expected_ops={torch.ops.aten.prod.dim_int}, ) - @parameterized.expand( - [ - ((3, 2, 4), [1], True), - ((2, 1, 4, 5), None, True), - ((2, 3, 4, 5), [0, 1, 2, 3], False), - ((6, 7, 5, 4, 5), [1, 3, 4], False), - ((6, 7, 5, 4, 5), [-5, -4, -2], False), - ] - ) - def test_prod_dim_tuple(self, input_shape, dim, keep_dims): - class Prod(nn.Module): - def forward(self, x): - return torch.prod(x, dim=dim, keepdim=keep_dims) - - inputs = [torch.randn(*input_shape)] - self.run_test( - Prod(), - inputs, - expected_ops={torch.ops.aten.prod.dim_int}, - ) - @parameterized.expand( [ ((3, 2, 4), 1, True, torch.int, 0, 5), - ((2, 3, 4, 5), None, True, torch.int, -10, 10), ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), ] @@ -91,27 +68,6 @@ def forward(self, x): check_dtype=False, ) - @parameterized.expand( - [ - ((3, 2, 4), [1], True, torch.int, 0, 5), - ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), - ((2, 3, 4, 5), None, False, torch.int32, -5, 0), - ((6, 7, 5, 4, 5), [1, 3, 4], False, torch.int32, -5, 5), - ] - ) - def test_prod_dim_tuple_int(self, input_shape, dim, keep_dims, dtype, low, high): - class Prod(nn.Module): - def forward(self, x): - return torch.prod(x, dim=dim, keepdim=keep_dims) - - inputs = [torch.randint(low, high, input_shape, dtype=dtype)] - self.run_test( - Prod(), - inputs, - expected_ops={torch.ops.aten.prod.dim_int}, - check_dtype=False, - ) - if __name__ == "__main__": run_tests() From f0091e8ac5a5a1371dc3a84578ae85416d4dc208 Mon Sep 17 00:00:00 2001 From: Evan Li Date: Wed, 4 Oct 2023 15:02:39 -0700 Subject: [PATCH 5/9] filter out dead nodes, and update min and max converters --- .../dynamo/conversion/aten_ops_converters.py | 133 ++++++++++++------ .../dynamo/conversion/impl/reduce.py | 118 +++++++++------- tests/py/dynamo/conversion/harness.py | 41 ++++-- tests/py/dynamo/conversion/test_max_aten.py | 130 ++++++++--------- tests/py/dynamo/conversion/test_min_aten.py | 131 +++++++++-------- 5 files changed, 311 insertions(+), 242 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 0609ddcde3..08330343c5 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1,4 +1,5 @@ import logging +import operator from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union import numpy as np @@ -712,51 +713,93 @@ def aten_ops_prod( ) -# def one_user_validator(min_node: Node) -> bool: -# # Validate only one user, which is a getitem node that accesses the first element in the list -# return (len(min_node.users) == 1 and -# list(min_node.users)[0].target == operator.getitem and -# list(min_node.users)[0].args[1] == 0) - - -# @dynamo_tensorrt_converter(torch.ops.aten.max.default) # type: ignore[misc] -# @dynamo_tensorrt_converter(torch.ops.aten.max.dim) # type: ignore[misc] -# def aten_ops_max( -# ctx: ConversionContext, -# target: Target, -# args: Tuple[Argument, ...], -# kwargs: Dict[str, Argument], -# name: str, -# ) -> Union[TRTTensor, Sequence[TRTTensor]]: -# return impl.reduce.max( -# ctx, -# target, -# SourceIR.ATEN, -# name, -# args[0], -# args_bounds_check(args, 1, replacement=None), -# args_bounds_check(args, 2, replacement=False), -# ) - - -# @dynamo_tensorrt_converter(torch.ops.aten.min.default) # type: ignore[misc] -# @dynamo_tensorrt_converter(torch.ops.aten.min.dim) # type: ignore[misc] -# def aten_ops_min( -# ctx: ConversionContext, -# target: Target, -# args: Tuple[Argument, ...], -# kwargs: Dict[str, Argument], -# name: str, -# ) -> Union[TRTTensor, Sequence[TRTTensor]]: -# return impl.reduce.min( -# ctx, -# target, -# SourceIR.ATEN, -# name, -# args[0], -# args_bounds_check(args, 1, replacement=None), -# args_bounds_check(args, 2, replacement=False), -# ) +def one_user_validator(node: Node) -> bool: + # Validate only one user, which is a getitem node that accesses the first element in the list + return ( + len(node.users) == 1 + and list(node.users)[0].target == operator.getitem + and list(node.users)[0].args[1] == 0 + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.max.default) # type: ignore[misc] +def aten_ops_max( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.max( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + dim=None, + keepdim=False, + return_indices=False, + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.max.dim, capability_validator=one_user_validator) # type: ignore[misc] +def aten_ops_max_dim( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.max( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args[1], + args_bounds_check(args, 2, replacement=False), + return_indices=True, + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.min.default) # type: ignore[misc] +def aten_ops_min( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.min( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + dim=None, + keepdim=False, + return_indices=False, + ) + + +@dynamo_tensorrt_converter(torch.ops.aten.min.dim, capability_validator=one_user_validator) # type: ignore[misc] +def aten_ops_min_dim( + ctx: ConversionContext, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + return impl.reduce.min( + ctx, + target, + SourceIR.ATEN, + name, + args[0], + args[1], + args_bounds_check(args, 2, replacement=False), + return_indices=True, + ) @dynamo_tensorrt_converter(torch.ops.aten.mean.default) # type: ignore[misc] diff --git a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py index bd32ad153b..7e001d8e10 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py @@ -1,4 +1,4 @@ -from typing import Optional, Sequence, Union +from typing import Optional, Sequence, Tuple, Union import tensorrt as trt from torch.fx.node import Target @@ -70,7 +70,7 @@ def prod( source_ir: Optional[SourceIR], name: str, input_val: TRTTensor, - dim: int, + dim: Optional[Union[int, Sequence[int]]], keepdim: bool, ) -> TRTTensor: if (isinstance(input_val, TRTTensor)) and ( @@ -91,58 +91,68 @@ def prod( return layer.get_output(0) -# def max( -# ctx: ConversionContext, -# target: Target, -# source_ir: Optional[SourceIR], -# name: str, -# input_val: TRTTensor, -# dim: Optional[int], -# keepdim: bool, -# ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: -# if (isinstance(input_val, TRTTensor)) and ( -# input_val.dtype == trt.int8 or input_val.dtype == trt.int32 -# ): -# input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) - -# if dim is None: -# dim = tuple(range(len(input_val.shape))) - -# layer = ctx.net.add_reduce( -# input_val, -# trt.ReduceOperation.MAX, -# axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), -# keep_dims=keepdim, -# ) -# set_layer_name(layer, target, name, source_ir) -# return layer.get_output(0), torch.tensor(0) - - -# def min( -# ctx: ConversionContext, -# target: Target, -# source_ir: Optional[SourceIR], -# name: str, -# input_val: TRTTensor, -# dim: Optional[int], -# keepdim: bool, -# ) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: -# if (isinstance(input_val, TRTTensor)) and ( -# input_val.dtype == trt.int8 or input_val.dtype == trt.int32 -# ): -# input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) - -# if dim is None: -# dim = tuple(range(len(input_val.shape))) - -# layer = ctx.net.add_reduce( -# input_val, -# trt.ReduceOperation.MIN, -# axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), -# keep_dims=keepdim, -# ) -# set_layer_name(layer, target, name, source_ir) -# return layer.get_output(0), torch.randint(0, 3, list(layer.get_output(0).shape)) +def max( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_val: TRTTensor, + dim: Optional[Union[int, Sequence[int]]], + keepdim: bool, + return_indices: bool, +) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: + if (isinstance(input_val, TRTTensor)) and ( + input_val.dtype == trt.int8 or input_val.dtype == trt.int32 + ): + input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + + if dim is None: + dim = tuple(range(len(input_val.shape))) + + layer = ctx.net.add_reduce( + input_val, + trt.ReduceOperation.MAX, + axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), + keep_dims=keepdim, + ) + set_layer_name(layer, target, name, source_ir) + + if return_indices: + return layer.get_output(0), None + + return layer.get_output(0) + + +def min( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + input_val: TRTTensor, + dim: Optional[Union[int, Sequence[int]]], + keepdim: bool, + return_indices: bool, +) -> Union[TRTTensor, Tuple[TRTTensor, TRTTensor]]: + if (isinstance(input_val, TRTTensor)) and ( + input_val.dtype == trt.int8 or input_val.dtype == trt.int32 + ): + input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + + if dim is None: + dim = tuple(range(len(input_val.shape))) + + layer = ctx.net.add_reduce( + input_val, + trt.ReduceOperation.MIN, + axes=get_axes_for_reduce_op(get_positive_dim(dim, len(input_val.shape))), + keep_dims=keepdim, + ) + set_layer_name(layer, target, name, source_ir) + + if return_indices: + return layer.get_output(0), None + + return layer.get_output(0) def mean( diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py index be13f7d2c1..6aefd0ffb2 100644 --- a/tests/py/dynamo/conversion/harness.py +++ b/tests/py/dynamo/conversion/harness.py @@ -197,18 +197,35 @@ def generate_graph( use_dynamo_tracer: bool, enable_passes: bool, ): - if use_dynamo_tracer: - fx_module = torch._dynamo.export( - mod, - *original_inputs, - aten_graph=True, - assume_static_by_default=True, - tracing_mode="real", - ).graph_module - else: - fx_module = torch.fx.symbolic_trace(mod) - if enable_passes: - fx_module = apply_lowering_passes(fx_module, original_inputs) + # Torchdynamo+aot proxytensor tracer + # Below are common passes + passes_list = [ + compose_bmm, + compose_chunk, + compose_getitem_slice, + replace_aten_reshape_alias_with_replace, + replace_aten_op_with_indices, + replace_transpose_mm_op_with_linear, # after compose_bmm + replace_native_layernorm_with_layernorm, + remove_ops, + replace_builtin_ops, # after replace_native_layernorm_with_layernorm + ] + # Combine with customized passes specific to any model + if customized_passes: + passes_list.extend(customized_passes) + + if disable_passes: + passes_list = [] + + fx_module, _ = aten_tracer.trace(mod, original_inputs) + for passes in passes_list: + pr: PassResult = passes(fx_module) + fx_module = pr.graph_module + fx_module(*original_inputs) + + fx_module = run_const_fold(fx_module) + fx_module.graph.eliminate_dead_code() + _LOGGER.info(f"FX graph= {fx_module.graph}") return fx_module diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index 7da7610d94..bb37172458 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -1,74 +1,74 @@ -# import torch -# import torch.nn as nn -# from parameterized import parameterized -# from torch.testing._internal.common_utils import run_tests +import torch +import torch.nn as nn +from harness import DispatchTestCase +from parameterized import parameterized +from torch.testing._internal.common_utils import run_tests -# from .harness import DispatchTestCase +class TestMaxConverter(DispatchTestCase): + @parameterized.expand( + [ + ((3, 2, 4),), + ((2, 3, 4, 5),), + ((2, 3, 4, 5),), + ((6, 7, 5, 4, 5),), + ] + ) + def test_max_dim_int_default(self, input_shape): + class Max(nn.Module): + def forward(self, x): + return torch.max(x) -# class TestMaxConverter(DispatchTestCase): -# @parameterized.expand( -# [ -# ((3, 2, 4),), -# ((2, 3, 4, 5),), -# ((2, 3, 4, 5),), -# ((6, 7, 5, 4, 5),), -# ] -# ) -# def test_max_dim_int_default(self, input_shape): -# class Max(nn.Module): -# def forward(self, x): -# return torch.max(x) + inputs = [torch.randn(*input_shape)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.default}, + ) -# inputs = [torch.randn(*input_shape)] -# self.run_test( -# Max(), -# inputs, -# expected_ops={torch.ops.aten.max.default}, -# ) + @parameterized.expand( + [ + ((3, 2, 4), 1, True), + ((2, 3, 4, 5), 3, True), + ((6, 7, 5, 4, 5), 4, False), + ((1, 5, 2, 1), -3, False), + ((1, 5, 2, 3), -2, True), + ] + ) + def test_max_dim_int(self, input_shape, dim, keep_dims): + class Max(nn.Module): + def forward(self, x): + return torch.ops.aten.max.dim(x, dim, keep_dims)[0] -# @parameterized.expand( -# [ -# ((3, 2, 4), 1, True), -# ((2, 3, 4, 5), 3, True), -# ((6, 7, 5, 4, 5), 4, False), -# ((1, 5, 2, 1), -3, False), -# ((1, 5, 2, 3), -2, True), -# ] -# ) -# def test_max_dim_int(self, input_shape, dim, keep_dims): -# class Max(nn.Module): -# def forward(self, x): -# return torch.max(x, dim=dim, keepdim=keep_dims) + inputs = [torch.randn(*input_shape)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.dim}, + disable_passes=True, + ) -# inputs = [torch.randn(*input_shape)] -# self.run_test( -# Max(), -# inputs, -# expected_ops={torch.ops.aten.max.dim}, -# ) + @parameterized.expand( + [ + ((3, 2, 4), 1, True, torch.int, 0, 5), + ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), + ] + ) + def test_max_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Max(nn.Module): + def forward(self, x): + return torch.ops.aten.max.dim(x, dim, keep_dims)[0] + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Max(), + inputs, + expected_ops={torch.ops.aten.max.dim}, + check_dtype=False, + disable_passes=True, + ) -# @parameterized.expand( -# [ -# ((3, 2, 4), 1, True, torch.int, 0, 5), -# ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), -# ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), -# ] -# ) -# def test_max_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): -# class Max(nn.Module): -# def forward(self, x): -# return torch.max(x, dim=dim, keepdim=keep_dims) -# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] -# self.run_test( -# Max(), -# inputs, -# expected_ops={torch.ops.aten.max.dim}, -# check_dtype=False, -# ) - - -# if __name__ == "__main__": -# run_tests() +if __name__ == "__main__": + run_tests() diff --git a/tests/py/dynamo/conversion/test_min_aten.py b/tests/py/dynamo/conversion/test_min_aten.py index 728c3e632d..6c19b23b58 100644 --- a/tests/py/dynamo/conversion/test_min_aten.py +++ b/tests/py/dynamo/conversion/test_min_aten.py @@ -1,76 +1,75 @@ -# import torch -# import torch.nn as nn -# from parameterized import parameterized -# from torch.testing._internal.common_utils import run_tests +import torch +import torch.nn as nn +from parameterized import parameterized +from torch.testing._internal.common_utils import run_tests -# from .harness import DispatchTestCase +from .harness import DispatchTestCase -# class TestMinConverter(DispatchTestCase): -# @parameterized.expand( -# [ -# ((3, 2, 4),), -# ((2, 3, 4, 5),), -# ((2, 3, 4, 5),), -# ((6, 7, 5, 4, 5),), -# ] -# ) -# def test_min_dim_int_default(self, input_shape): -# class Min(nn.Module): -# def forward(self, x): -# return torch.min(x) +class TestMinConverter(DispatchTestCase): + @parameterized.expand( + [ + ((3, 2, 4),), + ((2, 3, 4, 5),), + ((2, 3, 4, 5),), + ((6, 7, 5, 4, 5),), + ] + ) + def test_min_dim_int_default(self, input_shape): + class Min(nn.Module): + def forward(self, x): + return torch.ops.aten.min.default(x) -# inputs = [torch.randn(*input_shape)] -# self.run_test( -# Min(), -# inputs, -# expected_ops={torch.ops.aten.min.default}, -# ) + inputs = [torch.randn(*input_shape)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.default}, + ) -# @parameterized.expand( -# [ -# ((3, 2, 4), 1, True), -# ((2, 3, 4, 5), 3, True), -# ((6, 7, 5, 4, 5), 4, False), -# ((1, 5, 2, 1), -3, False), -# ((1, 5, 2, 3), -2, True), -# ] -# ) -# def test_min_dim_int(self, input_shape, dim, keep_dims): -# class Min(nn.Module): -# def forward(self, x): -# return torch.min(x, dim=dim, keepdim=keep_dims) -# # return torch.ops.aten.min.dim(x, dim, keep_dims) + @parameterized.expand( + [ + ((3, 2, 4), 1, True), + ((2, 3, 4, 5), 3, True), + ((6, 7, 5, 4, 5), 4, False), + ((1, 5, 2, 1), -3, False), + ((1, 5, 2, 3), -2, True), + ] + ) + def test_min_dim_int(self, input_shape, dim, keep_dims): + class Min(nn.Module): + def forward(self, x): + return torch.min(x, dim=dim, keepdim=keep_dims)[0] -# inputs = [torch.randn(*input_shape)] -# self.run_test( -# Min(), -# inputs, -# expected_ops={torch.ops.aten.min.dim}, -# disable_passes=True, -# ) + inputs = [torch.randn(*input_shape)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.dim}, + disable_passes=True, + ) + @parameterized.expand( + [ + ((3, 2, 4), 1, True, torch.int, 0, 5), + ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), + ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), + ] + ) + def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): + class Min(nn.Module): + def forward(self, x): + return torch.ops.aten.min.dim(x, dim=dim, keepdim=keep_dims)[0] -# @parameterized.expand( -# [ -# ((3, 2, 4), 1, True, torch.int, 0, 5), -# ((2, 3, 4, 5), 2, False, torch.int32, -5, 0), -# ((6, 7, 5, 4, 5), 4, False, torch.int32, -5, 5), -# ] -# ) -# def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): -# class Min(nn.Module): -# def forward(self, x): -# return torch.min(x, dim=dim, keepdim=keep_dims) + inputs = [torch.randint(low, high, input_shape, dtype=dtype)] + self.run_test( + Min(), + inputs, + expected_ops={torch.ops.aten.min.dim}, + check_dtype=False, + disable_passes=True, + ) -# inputs = [torch.randint(low, high, input_shape, dtype=dtype)] -# self.run_test( -# Min(), -# inputs, -# expected_ops={torch.ops.aten.min.dim}, -# check_dtype=False, -# ) - -# if __name__ == "__main__": -# run_tests() +if __name__ == "__main__": + run_tests() From 754c3b57eec5bfcd2758b9994a04910d68c80b8f Mon Sep 17 00:00:00 2001 From: Evan Li Date: Wed, 4 Oct 2023 16:54:11 -0700 Subject: [PATCH 6/9] fix a max test bug --- tests/py/dynamo/conversion/test_max_aten.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index bb37172458..36ed06e126 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -1,9 +1,10 @@ import torch import torch.nn as nn -from harness import DispatchTestCase from parameterized import parameterized from torch.testing._internal.common_utils import run_tests +from .harness import DispatchTestCase + class TestMaxConverter(DispatchTestCase): @parameterized.expand( From 732de3b659b417e271e75d4fd030f88d6c29faa1 Mon Sep 17 00:00:00 2001 From: Evan Li Date: Wed, 4 Oct 2023 18:14:53 -0700 Subject: [PATCH 7/9] fix dim bugs and update tests --- py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py | 2 +- py/torch_tensorrt/dynamo/conversion/impl/reduce.py | 7 +++++-- tests/py/dynamo/conversion/test_amax_aten.py | 2 ++ tests/py/dynamo/conversion/test_max_aten.py | 2 +- tests/py/dynamo/conversion/test_min_aten.py | 2 +- tests/py/dynamo/conversion/test_prod_aten.py | 2 +- tests/py/dynamo/conversion/test_sum_aten.py | 4 +++- 7 files changed, 14 insertions(+), 7 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 08330343c5..7e88f20ae6 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -668,7 +668,7 @@ def aten_ops_amax( SourceIR.ATEN, name, args[0], - args[1], + args_bounds_check(args, 1, replacement=[]), args_bounds_check(args, 2, replacement=False), ) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py index 7e001d8e10..692a52b0df 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/reduce.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/reduce.py @@ -27,6 +27,9 @@ def amax( ): input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) + if dim is None or (isinstance(dim, (tuple, list)) and len(dim) == 0): + dim = tuple(range(len(input_val.shape))) + layer = ctx.net.add_reduce( input_val, trt.ReduceOperation.MAX, @@ -51,7 +54,7 @@ def sum( ): input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) - if dim is None: + if dim is None or (isinstance(dim, (tuple, list)) and len(dim) == 0): dim = tuple(range(len(input_val.shape))) layer = ctx.net.add_reduce( @@ -169,7 +172,7 @@ def mean( ): input_val = cast_trt_tensor(ctx, input_val, trt.float32, name) - if dim is None: + if dim is None or (isinstance(dim, (tuple, list)) and len(dim) == 0): dim = tuple(range(len(input_val.shape))) layer = ctx.net.add_reduce( diff --git a/tests/py/dynamo/conversion/test_amax_aten.py b/tests/py/dynamo/conversion/test_amax_aten.py index 9ac95dfdd0..bdb0db5779 100644 --- a/tests/py/dynamo/conversion/test_amax_aten.py +++ b/tests/py/dynamo/conversion/test_amax_aten.py @@ -29,6 +29,7 @@ def forward(self, x): @parameterized.expand( [ + ((1, 2, 4), [], True), ((3, 2, 4), [1], True), ((2, 1, 4, 5), [0, 3], True), ((2, 3, 4, 5), [0, 1, 2, 3], False), @@ -69,6 +70,7 @@ def forward(self, x): @parameterized.expand( [ + ((1, 2, 4), [], True, torch.int, 0, 5), ((3, 2, 4), [1], True, torch.int, 0, 5), ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), ((2, 3, 4, 5), [0, 1, 2, 3], False, torch.int32, -5, 0), diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index 36ed06e126..a1865c008b 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -9,9 +9,9 @@ class TestMaxConverter(DispatchTestCase): @parameterized.expand( [ + ((1, 2),), ((3, 2, 4),), ((2, 3, 4, 5),), - ((2, 3, 4, 5),), ((6, 7, 5, 4, 5),), ] ) diff --git a/tests/py/dynamo/conversion/test_min_aten.py b/tests/py/dynamo/conversion/test_min_aten.py index 6c19b23b58..9f0bd4e864 100644 --- a/tests/py/dynamo/conversion/test_min_aten.py +++ b/tests/py/dynamo/conversion/test_min_aten.py @@ -9,9 +9,9 @@ class TestMinConverter(DispatchTestCase): @parameterized.expand( [ + ((1, 2),), ((3, 2, 4),), ((2, 3, 4, 5),), - ((2, 3, 4, 5),), ((6, 7, 5, 4, 5),), ] ) diff --git a/tests/py/dynamo/conversion/test_prod_aten.py b/tests/py/dynamo/conversion/test_prod_aten.py index f83464e63a..9346de6149 100644 --- a/tests/py/dynamo/conversion/test_prod_aten.py +++ b/tests/py/dynamo/conversion/test_prod_aten.py @@ -9,9 +9,9 @@ class TestProdConverter(DispatchTestCase): @parameterized.expand( [ + ((1, 2),), ((3, 2, 4),), ((2, 3, 4, 5),), - ((2, 3, 4, 5),), ((6, 7, 5, 4, 5),), ] ) diff --git a/tests/py/dynamo/conversion/test_sum_aten.py b/tests/py/dynamo/conversion/test_sum_aten.py index b279bed43e..2562c641f4 100644 --- a/tests/py/dynamo/conversion/test_sum_aten.py +++ b/tests/py/dynamo/conversion/test_sum_aten.py @@ -9,9 +9,9 @@ class TestSumConverter(DispatchTestCase): @parameterized.expand( [ + ((1, 2),), ((3, 2, 4),), ((2, 3, 4, 5),), - ((2, 3, 4, 5),), ((6, 7, 5, 4, 5),), ] ) @@ -49,6 +49,7 @@ def forward(self, x): @parameterized.expand( [ + ((1, 2, 4), [], True), ((3, 2, 4), [1], True), ((2, 1, 4, 5), None, True), ((2, 3, 4, 5), [0, 1, 2, 3], False), @@ -89,6 +90,7 @@ def forward(self, x): @parameterized.expand( [ + ((1, 2, 4), [], True, torch.int, 0, 5), ((3, 2, 4), [1], True, torch.int, 0, 5), ((2, 1, 4, 5), [0, 3], True, torch.int, -10, 10), ((2, 3, 4, 5), None, False, torch.int32, -5, 0), From 1717018dbeba270a8bba7de341018ba205b6ad54 Mon Sep 17 00:00:00 2001 From: Evan Li Date: Thu, 5 Oct 2023 17:54:28 -0700 Subject: [PATCH 8/9] rebase --- .../dynamo/conversion/aten_ops_converters.py | 58 ++++--------------- tests/py/dynamo/conversion/harness.py | 41 ++++--------- tests/py/dynamo/conversion/test_max_aten.py | 7 +-- tests/py/dynamo/conversion/test_min_aten.py | 9 +-- 4 files changed, 25 insertions(+), 90 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 7e88f20ae6..392dd216da 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -138,12 +138,12 @@ def aten_ops_sigmoid( ) -@dynamo_tensorrt_converter(torch.ops.aten.index.Tensor) +@dynamo_tensorrt_converter(torch.ops.aten.index.Tensor) # type: ignore[misc] @enforce_tensor_types( { 0: (TRTTensor,), } -) +) # type: ignore[misc] def aten_ops_index( ctx: ConversionContext, target: Target, @@ -723,27 +723,8 @@ def one_user_validator(node: Node) -> bool: @dynamo_tensorrt_converter(torch.ops.aten.max.default) # type: ignore[misc] -def aten_ops_max( - ctx: ConversionContext, - target: Target, - args: Tuple[Argument, ...], - kwargs: Dict[str, Argument], - name: str, -) -> Union[TRTTensor, Sequence[TRTTensor]]: - return impl.reduce.max( - ctx, - target, - SourceIR.ATEN, - name, - args[0], - dim=None, - keepdim=False, - return_indices=False, - ) - - @dynamo_tensorrt_converter(torch.ops.aten.max.dim, capability_validator=one_user_validator) # type: ignore[misc] -def aten_ops_max_dim( +def aten_ops_max( ctx: ConversionContext, target: Target, args: Tuple[Argument, ...], @@ -756,34 +737,15 @@ def aten_ops_max_dim( SourceIR.ATEN, name, args[0], - args[1], - args_bounds_check(args, 2, replacement=False), - return_indices=True, + dim=args_bounds_check(args, 1, replacement=None), + keepdim=args_bounds_check(args, 2, replacement=False), + return_indices=(target == torch.ops.aten.max.dim), ) @dynamo_tensorrt_converter(torch.ops.aten.min.default) # type: ignore[misc] -def aten_ops_min( - ctx: ConversionContext, - target: Target, - args: Tuple[Argument, ...], - kwargs: Dict[str, Argument], - name: str, -) -> Union[TRTTensor, Sequence[TRTTensor]]: - return impl.reduce.min( - ctx, - target, - SourceIR.ATEN, - name, - args[0], - dim=None, - keepdim=False, - return_indices=False, - ) - - @dynamo_tensorrt_converter(torch.ops.aten.min.dim, capability_validator=one_user_validator) # type: ignore[misc] -def aten_ops_min_dim( +def aten_ops_min( ctx: ConversionContext, target: Target, args: Tuple[Argument, ...], @@ -796,9 +758,9 @@ def aten_ops_min_dim( SourceIR.ATEN, name, args[0], - args[1], - args_bounds_check(args, 2, replacement=False), - return_indices=True, + dim=args_bounds_check(args, 1, replacement=None), + keepdim=args_bounds_check(args, 2, replacement=False), + return_indices=(target == torch.ops.aten.min.dim), ) diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py index 6aefd0ffb2..be13f7d2c1 100644 --- a/tests/py/dynamo/conversion/harness.py +++ b/tests/py/dynamo/conversion/harness.py @@ -197,35 +197,18 @@ def generate_graph( use_dynamo_tracer: bool, enable_passes: bool, ): - # Torchdynamo+aot proxytensor tracer - # Below are common passes - passes_list = [ - compose_bmm, - compose_chunk, - compose_getitem_slice, - replace_aten_reshape_alias_with_replace, - replace_aten_op_with_indices, - replace_transpose_mm_op_with_linear, # after compose_bmm - replace_native_layernorm_with_layernorm, - remove_ops, - replace_builtin_ops, # after replace_native_layernorm_with_layernorm - ] - # Combine with customized passes specific to any model - if customized_passes: - passes_list.extend(customized_passes) - - if disable_passes: - passes_list = [] - - fx_module, _ = aten_tracer.trace(mod, original_inputs) - for passes in passes_list: - pr: PassResult = passes(fx_module) - fx_module = pr.graph_module - fx_module(*original_inputs) - - fx_module = run_const_fold(fx_module) - fx_module.graph.eliminate_dead_code() - + if use_dynamo_tracer: + fx_module = torch._dynamo.export( + mod, + *original_inputs, + aten_graph=True, + assume_static_by_default=True, + tracing_mode="real", + ).graph_module + else: + fx_module = torch.fx.symbolic_trace(mod) + if enable_passes: + fx_module = apply_lowering_passes(fx_module, original_inputs) _LOGGER.info(f"FX graph= {fx_module.graph}") return fx_module diff --git a/tests/py/dynamo/conversion/test_max_aten.py b/tests/py/dynamo/conversion/test_max_aten.py index a1865c008b..7839bc4113 100644 --- a/tests/py/dynamo/conversion/test_max_aten.py +++ b/tests/py/dynamo/conversion/test_max_aten.py @@ -18,13 +18,12 @@ class TestMaxConverter(DispatchTestCase): def test_max_dim_int_default(self, input_shape): class Max(nn.Module): def forward(self, x): - return torch.max(x) + return torch.ops.aten.max.default(x) inputs = [torch.randn(*input_shape)] self.run_test( Max(), inputs, - expected_ops={torch.ops.aten.max.default}, ) @parameterized.expand( @@ -45,8 +44,6 @@ def forward(self, x): self.run_test( Max(), inputs, - expected_ops={torch.ops.aten.max.dim}, - disable_passes=True, ) @parameterized.expand( @@ -65,9 +62,7 @@ def forward(self, x): self.run_test( Max(), inputs, - expected_ops={torch.ops.aten.max.dim}, check_dtype=False, - disable_passes=True, ) diff --git a/tests/py/dynamo/conversion/test_min_aten.py b/tests/py/dynamo/conversion/test_min_aten.py index 9f0bd4e864..3d0cd29923 100644 --- a/tests/py/dynamo/conversion/test_min_aten.py +++ b/tests/py/dynamo/conversion/test_min_aten.py @@ -24,7 +24,6 @@ def forward(self, x): self.run_test( Min(), inputs, - expected_ops={torch.ops.aten.min.default}, ) @parameterized.expand( @@ -39,14 +38,12 @@ def forward(self, x): def test_min_dim_int(self, input_shape, dim, keep_dims): class Min(nn.Module): def forward(self, x): - return torch.min(x, dim=dim, keepdim=keep_dims)[0] + return torch.ops.aten.min.dim(x, dim, keep_dims)[0] inputs = [torch.randn(*input_shape)] self.run_test( Min(), inputs, - expected_ops={torch.ops.aten.min.dim}, - disable_passes=True, ) @parameterized.expand( @@ -59,15 +56,13 @@ def forward(self, x): def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): class Min(nn.Module): def forward(self, x): - return torch.ops.aten.min.dim(x, dim=dim, keepdim=keep_dims)[0] + return torch.ops.aten.min.dim(x, dim, keep_dims)[0] inputs = [torch.randint(low, high, input_shape, dtype=dtype)] self.run_test( Min(), inputs, - expected_ops={torch.ops.aten.min.dim}, check_dtype=False, - disable_passes=True, ) From 62389a0d8e837007f310cf0c97ece1a3326fcb1f Mon Sep 17 00:00:00 2001 From: Evan Li Date: Fri, 6 Oct 2023 11:55:52 -0700 Subject: [PATCH 9/9] rebase and fix bugs for tests --- tests/py/dynamo/conversion/test_maximum_aten.py | 12 ++++++------ tests/py/dynamo/conversion/test_minimum_aten.py | 12 ++++++------ tests/py/dynamo/conversion/test_prod_aten.py | 14 +++++++------- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/tests/py/dynamo/conversion/test_maximum_aten.py b/tests/py/dynamo/conversion/test_maximum_aten.py index 2be1d9c74b..0ae49f7395 100644 --- a/tests/py/dynamo/conversion/test_maximum_aten.py +++ b/tests/py/dynamo/conversion/test_maximum_aten.py @@ -7,23 +7,23 @@ from .harness import DispatchTestCase -class TestMaxConverter(DispatchTestCase): +class TestMaximumConverter(DispatchTestCase): @parameterized.expand( [ ("2d", (2, 1)), ("3d", (2, 1, 2)), ] ) - def test_max(self, _, shape): - class max(nn.Module): + def test_maximum(self, _, shape): + class Maximum(nn.Module): def forward(self, lhs_val, rhs_val): - return torch.max(lhs_val, rhs_val) + return torch.maximum(lhs_val, rhs_val) inputs = [torch.randn(shape), torch.randn(shape)] self.run_test( - max(), + Maximum(), inputs, - expected_ops={torch.ops.aten.maximum.default}, + use_dynamo_tracer=True, ) diff --git a/tests/py/dynamo/conversion/test_minimum_aten.py b/tests/py/dynamo/conversion/test_minimum_aten.py index 35d0d7163f..4bb70681b1 100644 --- a/tests/py/dynamo/conversion/test_minimum_aten.py +++ b/tests/py/dynamo/conversion/test_minimum_aten.py @@ -7,23 +7,23 @@ from .harness import DispatchTestCase -class TestMinConverter(DispatchTestCase): +class TestMinimumConverter(DispatchTestCase): @parameterized.expand( [ ("2d", (2, 1)), ("3d", (2, 1, 2)), ] ) - def test_min(self, _, shape): - class min(nn.Module): + def test_minimum(self, _, shape): + class Minimum(nn.Module): def forward(self, lhs_val, rhs_val): - return torch.min(lhs_val, rhs_val) + return torch.minimum(lhs_val, rhs_val) inputs = [torch.randn(shape), torch.randn(shape)] self.run_test( - min(), + Minimum(), inputs, - expected_ops={torch.ops.aten.minimum.default}, + use_dynamo_tracer=True, ) diff --git a/tests/py/dynamo/conversion/test_prod_aten.py b/tests/py/dynamo/conversion/test_prod_aten.py index 9346de6149..3fbb602098 100644 --- a/tests/py/dynamo/conversion/test_prod_aten.py +++ b/tests/py/dynamo/conversion/test_prod_aten.py @@ -20,11 +20,11 @@ class Prod(nn.Module): def forward(self, x): return torch.prod(x) - inputs = [torch.randn(*input_shape)] + inputs = [torch.randn(input_shape)] self.run_test( Prod(), inputs, - expected_ops={torch.ops.aten.prod.default}, + use_dynamo_tracer=True, ) @parameterized.expand( @@ -39,13 +39,13 @@ def forward(self, x): def test_prod_dim_int(self, input_shape, dim, keep_dims): class Prod(nn.Module): def forward(self, x): - return torch.prod(x, dim=dim, keepdim=keep_dims) + return torch.prod(x, dim, keep_dims) - inputs = [torch.randn(*input_shape)] + inputs = [torch.randn(input_shape)] self.run_test( Prod(), inputs, - expected_ops={torch.ops.aten.prod.dim_int}, + use_dynamo_tracer=True, ) @parameterized.expand( @@ -58,14 +58,14 @@ def forward(self, x): def test_prod_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high): class Prod(nn.Module): def forward(self, x): - return torch.prod(x, dim=dim, keepdim=keep_dims) + return torch.prod(x, dim, keep_dims) inputs = [torch.randint(low, high, input_shape, dtype=dtype)] self.run_test( Prod(), inputs, - expected_ops={torch.ops.aten.prod.dim_int}, check_dtype=False, + use_dynamo_tracer=True, )