Repository: pytorch/TensorRT

Issue:
Expose IReduceLayer in dynamo.conversion.impl



Discussion:

Resolving patch:
diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
--- 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
@@ -137,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,
@@ -667,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),
     )
 
@@ -692,6 +693,97 @@ 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),
+    )
+
+
+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]
+@dynamo_tensorrt_converter(torch.ops.aten.max.dim, capability_validator=one_user_validator)  # 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=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]
+@dynamo_tensorrt_converter(torch.ops.aten.min.dim, capability_validator=one_user_validator)  # 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=args_bounds_check(args, 1, replacement=None),
+        keepdim=args_bounds_check(args, 2, replacement=False),
+        return_indices=(target == torch.ops.aten.min.dim),
+    )
+
+
+@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 +1210,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 +1228,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
--- 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
@@ -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,
@@ -43,16 +46,17 @@ 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
     ):
         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(
         input_val,
         trt.ReduceOperation.SUM,
@@ -61,3 +65,121 @@ 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]]],
+    keepdim: bool,
+) -> 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]]],
+    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(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_val: TRTTensor,
+    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
+    ):
+        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.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)


Test patch:
diff --git a/tests/py/dynamo/conversion/test_amax_aten.py b/tests/py/dynamo/conversion/test_amax_aten.py
--- 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
--- 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,60 @@
 class TestMaxConverter(DispatchTestCase):
     @parameterized.expand(
         [
-            ("2d", (2, 1)),
-            ("3d", (2, 1, 2)),
+            ((1, 2),),
+            ((3, 2, 4),),
+            ((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.ops.aten.max.default(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},
+        )
+
+    @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]
+
+        inputs = [torch.randn(*input_shape)]
+        self.run_test(
+            Max(),
+            inputs,
+        )
+
+    @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,
+            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
--- /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 TestMaximumConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d", (2, 1)),
+            ("3d", (2, 1, 2)),
+        ]
+    )
+    def test_maximum(self, _, shape):
+        class Maximum(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.maximum(lhs_val, rhs_val)
+
+        inputs = [torch.randn(shape), torch.randn(shape)]
+        self.run_test(
+            Maximum(),
+            inputs,
+            use_dynamo_tracer=True,
+        )
+
+
+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
--- 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,21 +9,62 @@
 class TestMinConverter(DispatchTestCase):
     @parameterized.expand(
         [
-            ("2d", (2, 1)),
-            ("3d", (2, 1, 2)),
+            ((1, 2),),
+            ((3, 2, 4),),
+            ((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.ops.aten.min.default(x)
 
-        inputs = [torch.randn(shape), torch.randn(shape)]
+        inputs = [torch.randn(*input_shape)]
         self.run_test(
-            min(),
+            Min(),
             inputs,
         )
 
+    @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.ops.aten.min.dim(x, dim, keep_dims)[0]
+
+        inputs = [torch.randn(*input_shape)]
+        self.run_test(
+            Min(),
+            inputs,
+        )
+
+    @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, keep_dims)[0]
+
+        inputs = [torch.randint(low, high, input_shape, dtype=dtype)]
+        self.run_test(
+            Min(),
+            inputs,
+            check_dtype=False,
+        )
+
 
 if __name__ == "__main__":
     run_tests()
diff --git a/tests/py/dynamo/conversion/test_minimum_aten.py b/tests/py/dynamo/conversion/test_minimum_aten.py
new file mode 100644
--- /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 TestMinimumConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d", (2, 1)),
+            ("3d", (2, 1, 2)),
+        ]
+    )
+    def test_minimum(self, _, shape):
+        class Minimum(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.minimum(lhs_val, rhs_val)
+
+        inputs = [torch.randn(shape), torch.randn(shape)]
+        self.run_test(
+            Minimum(),
+            inputs,
+            use_dynamo_tracer=True,
+        )
+
+
+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
new file mode 100644
--- /dev/null
+++ b/tests/py/dynamo/conversion/test_prod_aten.py
@@ -0,0 +1,73 @@
+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(
+        [
+            ((1, 2),),
+            ((3, 2, 4),),
+            ((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,
+            use_dynamo_tracer=True,
+        )
+
+    @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_prod_dim_int(self, input_shape, dim, keep_dims):
+        class Prod(nn.Module):
+            def forward(self, x):
+                return torch.prod(x, dim, keep_dims)
+
+        inputs = [torch.randn(input_shape)]
+        self.run_test(
+            Prod(),
+            inputs,
+            use_dynamo_tracer=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_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, keep_dims)
+
+        inputs = [torch.randint(low, high, input_shape, dtype=dtype)]
+        self.run_test(
+            Prod(),
+            inputs,
+            check_dtype=False,
+            use_dynamo_tracer=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()
diff --git a/tests/py/dynamo/conversion/test_sum_aten.py b/tests/py/dynamo/conversion/test_sum_aten.py
--- 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),