feat: support 1D, 2D, and 3D avg and max pooling dynamo converters (#2317)

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1413,3 +1413,90 @@ def aten_ops_linear(
         weight=args[1],
         bias=args_bounds_check(args, 2, None),
     )
+
+
+def avg_pool_param_validator(pool_node: Node) -> bool:
+    ceil_mode = args_bounds_check(pool_node.args, 4, False)
+    divisor_override = args_bounds_check(pool_node.args, 6)
+
+    if ceil_mode is not False:
+        _LOGGER.debug(
+            f"Currently we don't support specifying ceil_mode, got ceil_mode={ceil_mode}."
+        )
+        return False
+
+    if divisor_override is not None:
+        _LOGGER.debug(
+            f"Currently we don't support divisor_override, got divisor_override={divisor_override}."
+        )
+        return False
+
+    return True
+
+
+# Note: AvgPool1d uses avg_pool2d as it converts to 2D first.
+@dynamo_tensorrt_converter(torch.ops.aten.avg_pool1d.default, capability_validator=avg_pool_param_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.avg_pool2d.default, capability_validator=avg_pool_param_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.avg_pool3d.default, capability_validator=avg_pool_param_validator)  # type: ignore[misc]
+def aten_ops_avg_pool(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pool.avg_poolNd(
+        network,
+        target,
+        source_ir=SourceIR.ATEN,
+        name=name,
+        input=args[0],
+        kernel_size=args[1],
+        stride=args_bounds_check(args, 2, replacement=[]),
+        padding=args_bounds_check(args, 3, replacement=0),
+        ceil_mode=args_bounds_check(args, 4, replacement=False),
+        count_include_pad=args_bounds_check(args, 5, replacement=True),
+        divisor_override=args_bounds_check(args, 6, replacement=None),
+    )
+
+
+def max_pool_param_validator(pool_node: Node) -> bool:
+    dilation = args_bounds_check(pool_node.args, 4, 1)
+    ceil_mode = args_bounds_check(pool_node.args, 5, False)
+
+    if dilation != 1:
+        _LOGGER.debug(f"Currently we don't support dilation, got dilation={dilation}.")
+        return False
+
+    if ceil_mode is not False:
+        _LOGGER.debug(
+            f"Currently we don't support specifying ceil_mode, got ceil_mode={ceil_mode}."
+        )
+        return False
+
+    return True
+
+
+# Note: MaxPool1d uses max_pool2d as it converts to 2D first.
+@dynamo_tensorrt_converter(torch.ops.aten.max_pool1d.default, capability_validator=max_pool_param_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.max_pool2d.default, capability_validator=max_pool_param_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.max_pool3d.default, capability_validator=max_pool_param_validator)  # type: ignore[misc]
+def aten_ops_max_pool(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.pool.max_poolNd(
+        network,
+        target,
+        source_ir=SourceIR.ATEN,
+        name=name,
+        input=args[0],
+        kernel_size=args[1],
+        stride=args_bounds_check(args, 2, replacement=[]),
+        padding=args_bounds_check(args, 3, replacement=0),
+        dilation=args_bounds_check(args, 4, replacement=1),
+        ceil_mode=args_bounds_check(args, 5, replacement=False),
+    )

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -12,6 +12,7 @@
     matmul,
     normalization,
     permutation,
+    pool,
     reduce,
     select,
     shape,

py/torch_tensorrt/dynamo/conversion/impl/pool.py

@@ -0,0 +1,105 @@
+from typing import Optional, Sequence, Union
+
+import tensorrt as trt
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.converter_utils import extend_attr_to_tuple
+from torch_tensorrt.fx.converters.converter_utils import (
+    has_dynamic_shape,
+    set_layer_name,
+)
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+
+def avg_poolNd(
+    network: TRTNetwork,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    kernel_size: Sequence[int],
+    stride: Union[int, Sequence[int]],
+    padding: Union[int, Sequence[int]] = 0,
+    ceil_mode: bool = False,
+    count_include_pad: bool = True,
+    divisor_override: Optional[int] = None,
+) -> TRTTensor:
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for pooling."
+
+    if ceil_mode is not False:
+        raise RuntimeError("ceil_mode is not yet supported!")
+
+    if divisor_override is not None:
+        raise RuntimeError("divisor_override is not yet supported!")
+
+    dim = len(kernel_size)
+
+    kernel_size = extend_attr_to_tuple(kernel_size, dim)
+
+    if stride == []:
+        stride = kernel_size
+    else:
+        stride = extend_attr_to_tuple(stride, dim)
+
+    padding = extend_attr_to_tuple(padding, dim)
+
+    # add average pooling layer
+    pool_layer = network.add_pooling_nd(
+        input=input,
+        type=trt.PoolingType.AVERAGE,
+        window_size=kernel_size,
+    )
+
+    pool_layer.stride_nd = stride
+    pool_layer.padding_nd = padding
+    pool_layer.average_count_excludes_padding = not count_include_pad
+
+    set_layer_name(pool_layer, target, name, source_ir)
+    return pool_layer.get_output(0)
+
+
+def max_poolNd(
+    network: TRTNetwork,
+    target: Union[Target, str],
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    kernel_size: Sequence[int],
+    stride: Union[int, Sequence[int]],
+    padding: Union[int, Sequence[int]] = 0,
+    dilation: Union[int, Sequence[int]] = 1,
+    ceil_mode: bool = False,
+) -> TRTTensor:
+    if has_dynamic_shape(input.shape):
+        assert input.shape[1] != -1, "Channel dim can't be dynamic for pooling."
+
+    if dilation != 1:
+        raise RuntimeError("dilation is not yet supported!")
+
+    if ceil_mode is not False:
+        raise RuntimeError("ceil_mode is not yet supported!")
+
+    dim = len(kernel_size)
+
+    kernel_size = extend_attr_to_tuple(kernel_size, dim)
+
+    if stride == []:
+        stride = kernel_size
+    else:
+        stride = extend_attr_to_tuple(stride, dim)
+
+    padding = extend_attr_to_tuple(padding, dim)
+
+    # add max pooling layer
+    pool_layer = network.add_pooling_nd(
+        input=input,
+        type=trt.PoolingType.MAX,
+        window_size=kernel_size,
+    )
+
+    pool_layer.stride_nd = stride
+    pool_layer.padding_nd = padding
+
+    set_layer_name(pool_layer, target, name, source_ir)
+    return pool_layer.get_output(0)