Add tp.MaxPool (#204)

Signed-off-by: yizhuoz004 <[email protected]>
Co-authored-by: pranavm-nvidia <[email protected]>

Author: yizhuoz004

tripy/tests/frontend/trace/ops/test_convolution.py

@@ -66,7 +66,7 @@ def test_invalid_rank_fails(self, conv_func):
     @pytest.mark.parametrize(
         "padding, err, expect_input_stack_info",
         [
-            (((2, 2),), r"expects padding-entries to have same dimension-size as size of window dimensions", True),
+            (((2, 2),), r"Padding must have the same length as kernel_dims.", False),
             (((2, 2, 2), (2, 2, 2)), r"Padding must be provided as a sequence of pairs of integers.", False),
             (((1, 2), (-3, 1)), r"Negative padding is not supported.", False),
         ],
@@ -84,7 +84,7 @@ def test_invalid_padding(self, conv_func, padding, err, expect_input_stack_info)
         "stride, err, expect_input_stack_info",
         [
             ((-1, 0), r"Non-positive stride is not supported.", False),
-            ((2, 2, 2), r"expects window-strides to have same dimension-size as size of window dimensions", True),
+            ((2, 2, 2), r"Stride must have the same length as kernel_dims.", False),
         ],
     )
     def test_invalid_stride(self, conv_func, stride, err, expect_input_stack_info):
@@ -136,11 +136,7 @@ def test_infer_rank(self, conv_func):
         "dilation, err, expect_input_stack_info",
         [
             ((-1, 0), r"Non-positive dilation is not supported.", False),
-            (
-                (2, 2, 2),
-                r"expects window-dilation factors to have same dimension-size as size of window dimensions.",
-                True,
-            ),
+            ((2, 2, 2), r"Dilation must have the same length as kernel_dims.", False),
         ],
     )
     def test_invalid_rhs_dilation(self, conv_func, dilation, err, expect_input_stack_info):

tripy/tests/frontend/trace/ops/test_pooling.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from tests import helper
+import tripy as tp
+
+
+class TestPooling:
+    def test_invalid_kernel_dims(self):
+        a = tp.ones((1, 1, 4, 4))
+        with helper.raises(tp.TripyException, "Unsupported kernel_dims, must be 2D or 3D."):
+            tp.maxpool(a, (2,))
+
+    def test_invalid_stride(self):
+        a = tp.ones((1, 1, 4, 4))
+        with helper.raises(tp.TripyException, "Stride must have the same length as kernel_dims."):
+            tp.maxpool(a, (2, 2), stride=(1,))
+
+    def test_invalid_padding_length(self):
+        a = tp.ones((1, 1, 4, 4))
+        with helper.raises(tp.TripyException, "Padding must have the same length as kernel_dims."):
+            tp.maxpool(a, (2, 2), padding=((1, 1),))
+
+    def test_invalid_padding_contents(self):
+        a = tp.ones((1, 1, 4, 4))
+        with helper.raises(tp.TripyException, "Padding must be provided as a sequence of pairs of integers."):
+            tp.maxpool(a, (2, 2), padding=((1, 1, 1), (1, 1, 1)))

tripy/tests/integration/test_pooling.py

@@ -0,0 +1,61 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pytest
+import torch
+
+import tripy as tp
+
+
+class TestPooling:
+
+    @pytest.mark.parametrize(
+        "kernel_dims, stride, padding",
+        [
+            ((3, 3), (1, 1), ((0, 0), (0, 0))),
+            ((4, 4), (2, 2), ((1, 1), (2, 2))),
+        ],
+    )
+    @pytest.mark.parametrize("dtype", [tp.float32, tp.float16, tp.int8])
+    def test_maxpool_2d(self, kernel_dims, stride, padding, dtype):
+        inp_tp = tp.reshape(tp.arange(64, dtype=dtype), (1, 1, 8, 8))
+        out = tp.maxpool(inp_tp, kernel_dims=kernel_dims, stride=stride, padding=padding)
+        out_torch = torch.from_dlpack(out).to("cpu")
+
+        torch_padding = (padding[0][0], padding[1][0])
+        pool_torch = torch.nn.MaxPool2d(kernel_size=kernel_dims, stride=stride, padding=torch_padding)
+        expected = pool_torch(torch.from_dlpack(inp_tp).to("cpu"))
+        assert torch.allclose(expected, out_torch)
+        assert expected.shape == out_torch.shape
+
+    @pytest.mark.parametrize(
+        "kernel_dims, stride, padding",
+        [
+            ((2, 2, 2), (2, 2, 2), ((0, 0), (1, 1), (1, 1))),
+        ],
+    )
+    @pytest.mark.parametrize("dtype", [tp.float32, tp.float16])
+    def test_maxpool_3d(self, kernel_dims, stride, padding, dtype):
+        inp_tp = tp.reshape(tp.arange(512, dtype=dtype), (1, 1, 8, 8, 8))
+        out = tp.maxpool(inp_tp, kernel_dims=kernel_dims, stride=stride, padding=padding)
+        out_torch = torch.from_dlpack(out).to("cpu")
+
+        torch_padding = (padding[0][0], padding[1][0], padding[2][0])
+        pool_torch = torch.nn.MaxPool3d(kernel_size=kernel_dims, stride=stride, padding=torch_padding)
+        expected = pool_torch(torch.from_dlpack(inp_tp).to("cpu"))
+        assert torch.allclose(expected, out_torch)
+        assert expected.shape == out_torch.shape

tripy/tripy/backend/mlir/utils.py

@@ -16,6 +16,7 @@
 #
 
 import contextlib
+import numbers
 import os
 import re
 import sys
@@ -84,13 +85,20 @@ def get_mlir_dtype(dtype: "tripy.dtype"):
     }[dtype.name]
 
 
-def get_mlir_scalar_attr(dtype: "tripy.dtype", value):
-    from tripy.common.datatype import floating
-
+def get_mlir_scalar_attr(mlir_dtype, value):
     # MLIR represents float dtypes as FloatAttr
     # and non-float dtypes as IntegerAttr
-    attr_func = ir.FloatAttr.get if issubclass(dtype, floating) else ir.IntegerAttr.get
-    return attr_func(get_mlir_dtype(dtype), value)
+    attr_func = ir.IntegerAttr.get if isinstance(mlir_dtype, ir.IntegerType) else ir.FloatAttr.get
+    return attr_func(mlir_dtype, value)
+
+
+def list_to_dense_attr(data: List, mlir_dtype):
+    if isinstance(data, numbers.Number):
+        return [get_mlir_scalar_attr(mlir_dtype, data)]
+    attrs = []
+    for element in data:
+        attrs.extend(list_to_dense_attr(element, mlir_dtype))
+    return attrs
 
 
 def get_mlir_quant_dtype(

tripy/tripy/flat_ir/ops/__init__.py

@@ -42,6 +42,7 @@
 from tripy.flat_ir.ops.plugin import PluginOp
 from tripy.flat_ir.ops.pow import PowOp
 from tripy.flat_ir.ops.reduce import ArgMinMaxOp, ReduceOp
+from tripy.flat_ir.ops.reduce_window import ReduceWindowOp
 from tripy.flat_ir.ops.reshape import DynamicReshapeOp
 from tripy.flat_ir.ops.round_nearest_even import RoundNearestEvenOp
 from tripy.flat_ir.ops.rsqrt import RsqrtOp

tripy/tripy/flat_ir/ops/constant.py

@@ -77,15 +77,9 @@ def to_mlir(self, operands):
             )
         else:
             out_dtype = self.outputs[0].dtype
-
-            def to_attrs(data):
-                if isinstance(data, numbers.Number):
-                    return [mlir_utils.get_mlir_scalar_attr(out_dtype, data)]
-                attrs = []
-                for element in data:
-                    attrs.extend(to_attrs(element))
-                return attrs
-
-            attr = ir.DenseElementsAttr.get(attrs=to_attrs(self.data), type=self.outputs[0].to_mlir())
+            attr = ir.DenseElementsAttr.get(
+                attrs=mlir_utils.list_to_dense_attr(self.data, mlir_utils.get_mlir_dtype(out_dtype)),
+                type=self.outputs[0].to_mlir(),
+            )
 
         return [stablehlo.ConstantOp(attr)]

tripy/tripy/flat_ir/ops/reduce_window.py

@@ -0,0 +1,76 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from dataclasses import dataclass
+from typing import Sequence, Tuple
+
+from mlir_tensorrt.compiler import ir
+from mlir_tensorrt.compiler.dialects import stablehlo
+
+from tripy.backend.mlir import utils as mlir_utils
+from tripy.flat_ir.ops.base import BaseFlatIROp
+
+
+@dataclass(repr=False)
+class ReduceWindowOp(BaseFlatIROp):
+
+    reduce_mode: str
+    window_dims: Sequence[int]
+    window_strides: Sequence[int]
+    padding: Sequence[Tuple[int]]
+
+    def _get_reduce_func(self):
+        if self.reduce_mode == "max":
+            return stablehlo.MaxOp
+        else:
+            raise NotImplementedError()
+
+    def to_mlir(self, operands):
+        input_dtype = self.inputs[0].dtype
+        out_type = self.outputs[0].to_mlir()
+
+        window_dims_attr = ir.DenseI64ArrayAttr.get(self.window_dims)
+        window_strides_attr = ir.DenseI64ArrayAttr.get(self.window_strides)
+        padding_attr_type = ir.RankedTensorType.get(
+            [len(self.padding), 2],
+            ir.IntegerType.get_signless(64),
+        )
+        padding_attr = ir.DenseElementsAttr.get(
+            attrs=mlir_utils.list_to_dense_attr(self.padding, ir.IntegerType.get_signless(64)),
+            type=padding_attr_type,
+        )
+
+        reduce = stablehlo.ReduceWindowOp(
+            result=[out_type],
+            inputs=[operands[0]],
+            init_values=[operands[1]],
+            window_dimensions=window_dims_attr,
+            window_strides=window_strides_attr,
+            padding=padding_attr,
+        )
+
+        reduce_arg_type = ir.RankedTensorType.get(
+            [],
+            mlir_utils.get_mlir_dtype(input_dtype),
+        )
+        reduce_block = ir.Block.create_at_start(reduce.regions[0], [reduce_arg_type, reduce_arg_type])
+        reduce_func = self._get_reduce_func()
+        with ir.InsertionPoint(reduce_block):
+            out = reduce_func(*reduce_block.arguments)
+            stablehlo.ReturnOp([out])
+
+        return [reduce]

tripy/tripy/frontend/module/convolution.py

@@ -23,6 +23,7 @@
 from tripy.common import datatype
 from tripy.frontend.module.module import Module
 from tripy.frontend.module.parameter import Parameter, DefaultParameter
+from tripy.frontend.trace.ops import utils as op_utils
 
 from tripy.common.exception import raise_error
 
@@ -70,37 +71,12 @@ def __init__(
                 ],
             )
 
+        op_utils.check_conv_pooling_args(kernel_dims, stride, padding, dilation)
         rank = len(kernel_dims) + 2
         self.padding = utils.default(padding, tuple(((0, 0) for _ in range(rank - 2))))
-
-        if not all(len(pad) == 2 for pad in self.padding):
-            raise_error(
-                f"Padding must be provided as a sequence of pairs of integers.",
-                details=[f"Supplied padding attribute: {self.padding} contains sequences that are not of length 2."],
-            )
-
-        if not all(p1 >= 0 and p2 >= 0 for p1, p2 in self.padding):
-            raise_error(
-                "Negative padding is not supported.",
-                details=[f"Got padding: {self.padding} but all values must be non-negative integers."],
-            )
-
         self.stride = utils.default(stride, (1,) * (rank - 2))
-
-        if not all(s > 0 for s in self.stride):
-            raise_error(
-                "Non-positive stride is not supported.",
-                details=[f"Got stride: {self.stride} but all values must be integers greater than 0."],
-            )
-
         self.dilation = utils.default(dilation, (1,) * (rank - 2))
 
-        if not all(isinstance(d, int) and d > 0 for d in self.dilation):
-            raise_error(
-                "Non-positive dilation is not supported.",
-                details=[f"Got dilation: {self.dilation} but all values must be integers greater than 0."],
-            )
-
         if bias:
             self.bias = DefaultParameter((out_channels,), dtype=dtype)