Arm backend: Add TOSA AVG_POOL2D op (#18972)

### Summary

Adds new TOSA dialect op for AVG_POOL2D. aten.AvgPool2d nodes are
replaced by tosa.AVG_POOL2D in RewriteAvgPool2dPass. op_avg_pool2d node
visitor is replaced by a simpler node vistior for tosa.avg_pool2d.


Signed-off-by: Saoirse Stewart <saoirse.stewart@arm.com>

Author: SaoirseARM

backends/arm/_passes/__init__.py

@@ -143,6 +143,7 @@
 from .replace_scalar_with_tensor_pass import (  # noqa
     ReplaceScalarWithTensorByProfilePass,
 )
+from .rewrite_avg_pool2d_pass import RewriteAvgPool2dPass  # noqa
 from .rewrite_bool_bitwise_to_logical_pass import (  # noqa
     RewriteBoolBitwiseToLogicalPass,
 )

backends/arm/_passes/arm_pass_manager.py

@@ -10,7 +10,6 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 
-import executorch.backends.arm.tosa.dialect  # noqa: unused
 from executorch.backends.arm._passes import (
     AccumulateIndexPutPass,
     AnnotateOutputDimOrderPass,
@@ -126,6 +125,7 @@
     RemoveNoopPass,
     ReplaceInfAndLimitValuesPass,
     ReplaceScalarWithTensorByProfilePass,
+    RewriteAvgPool2dPass,
     RewriteBoolBitwiseToLogicalPass,
     RewriteBoolToFp32CastViaInt8Pass,
     RewriteConvPass,
@@ -144,7 +144,6 @@
     UnsqueezeBeforeRepeatPass,
     UnsqueezeScalarPlaceholdersPass,
 )
-
 from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.common.pipeline_config import (
@@ -463,6 +462,8 @@ def _tosa_pipeline(
                 DecomposeSliceScatterPass(),
                 AccumulateIndexPutPass(),
                 DecomposeIndexTensorToGatherPass(),
+                DecomposeAdaptiveAvgPool2dPass(),
+                DecomposeAvgPool2dPass(),
                 Conv1dUnsqueezePass(),
             ]
         )
@@ -499,17 +500,16 @@ def _tosa_pipeline(
                 DecomposeSoftmaxPass(),
                 ConvertMinMaxPass(),
                 DecomposeAnyPass(),
-                DecomposeAdaptiveAvgPool2dPass(),
-                DecomposeAvgPool2dPass(),
                 DecorateFp32toInt32CastingPass(),
-                ComputeConstantOpsAOTPass(exported_program),
-                FuseConstantArgsPass(exported_program),
                 ConvertExpandCopyToRepeatPass(),
                 UnsqueezeBeforeRepeatPass(),
                 DecomposeCumsumPass(exported_program),
                 DecomposeAsStridedCopyPass(),
                 DecomposeMaxPool2dPass(),
                 SizeAdjustInputPass(),
+                RewriteAvgPool2dPass(),
+                ComputeConstantOpsAOTPass(exported_program),
+                FuseConstantArgsPass(exported_program),
                 DecomposeSelectPass(),
                 ConvertSqueezesToViewPass(),
                 CastToInt32Pass(),
@@ -605,6 +605,8 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     DecomposeLayerNormPass(tfa_pass=True),
                     DecomposeVarPass(tfa_pass=True),
                     DecomposeMeanDimPass(graph_module, self.tosa_spec, tfa_pass=True),
+                    DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
+                    DecomposeAvgPool2dPass(tfa_pass=True),
                 ]
             )
 
@@ -630,8 +632,6 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     DecomposeDivPass(tfa_pass=True),
                     DecomposeLinalgVectorNormPass(tfa_pass=True),
                     DecomposeSqrtPass(tfa_pass=True),
-                    DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
-                    DecomposeAvgPool2dPass(tfa_pass=True),
                     DecomposeSoftmaxPass(
                         tfa_pass=True,
                     ),

backends/arm/_passes/decompose_avg_pool2d_pass.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
-from typing import Set, Type
+from typing import Any, Set, Type
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
@@ -17,138 +17,162 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
-edge_div_ops = (exir_ops.edge.aten.avg_pool2d.default,)
-aten_div_ops = (torch.ops.aten.avg_pool2d.default,)
+edge_avg_pool2d = (exir_ops.edge.aten.avg_pool2d.default,)
+aten_avg_pool2d = (torch.ops.aten.avg_pool2d.default,)
 
 
 def get_decomposition(op) -> tuple:
-    if op in edge_div_ops:
+    if op in edge_avg_pool2d:
         return (
-            exir_ops.edge.aten.full.default,
-            exir_ops.edge.aten.cat.default,
+            exir_ops.edge.aten.constant_pad_nd.default,
             exir_ops.edge.aten.avg_pool2d.default,
             exir_ops.edge.aten.mul.Tensor,
         )
-    if op in aten_div_ops:
+    if op in aten_avg_pool2d:
         return (
-            torch.ops.aten.full.default,
-            torch.ops.aten.cat.default,
+            torch.ops.aten.pad.default,
             torch.ops.aten.avg_pool2d.default,
             torch.ops.aten.mul.Tensor,
         )
     raise RuntimeError(f"Can't get avg_pool2d decomposition for op {op}")
 
 
+def _compute_post_pad(
+    size: int,
+    kernel: int,
+    stride: int,
+    pad: int,
+    ceil_mode: bool,
+    divisor_override,
+) -> int:
+
+    if pad == 0:
+        return pad
+    if ceil_mode and divisor_override is None:
+        return pad
+
+    pad_adjust = adjust_pooling_pad_if_needed(size, kernel, stride, pad, ceil_mode)
+
+    # Padding must always be above 0, the above adjustment may return -1
+    if pad_adjust > 0:
+        return pad_adjust
+    return pad
+
+
+def _get_avgpool_post_pad(
+    h,
+    w,
+    kernel: tuple,
+    stride_h,
+    stride_w,
+    pad_h,
+    pad_w,
+    ceil_mode,
+    count_include_pad,
+    divisor_override,
+) -> tuple[list[Any], list[int]]:
+    """Compute the post-padding configuration for avg_pool2d when pre-
+    materializing explicit zero padding ahead of the pooling operation.
+
+    Given the original spatial dimensions (h, w), pooling kernel size, stride,
+    and explicit pre-padding amounts (pad_h, pad_w), this function returns the
+    additional padding to apply on the right and bottom edges so that avg_pool2d
+    with count_include_pad and/or divisor_override produces the equivalent
+    result without built-in padding.
+
+    """
+
+    k_h, k_w = kernel
+    post_h, post_w = (0, 0)
+    new_pad_h, new_pad_w = pad_h, pad_w
+
+    if not count_include_pad:
+        return [new_pad_h, new_pad_w], [new_pad_h, new_pad_w]
+
+    post_h = _compute_post_pad(h, k_h, stride_h, pad_h, ceil_mode, divisor_override)
+    post_w = _compute_post_pad(w, k_w, stride_w, pad_w, ceil_mode, divisor_override)
+
+    # Return our pre-padding calculation. Turn off built-in padding.
+    return [pad_w, post_w, pad_h, post_h], [0, 0]
+
+
 class DecomposeAvgPool2dPass(ArmPass):
     _passes_required_after: Set[Type[ExportPass]] = {ComputeConstantOpsAOTPass}
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in (edge_div_ops + aten_div_ops) or not self.allowed_to_transform(
-            meta
-        ):
+        if op not in (
+            edge_avg_pool2d + aten_avg_pool2d
+        ) or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
-        full_op, cat_op, avgpool_op, mul_op = get_decomposition(op)
+        pad_op, avgpool_op, mul_op = get_decomposition(op)
 
         x = args[0]
-        full_kwargs = {"device": x.data.device, "dtype": x.data.dtype}
         kernel_h, kernel_w = args[1]
         kernel_size = kernel_h * kernel_w
+
         if len(args) > 2 and args[2] is not None:
             stride_h, stride_w = args[2]
         else:
             stride_h, stride_w = kernel_h, kernel_w
-        pad_h, pad_w = new_pad_h, new_pad_w = args[3] if len(args) > 3 else (0, 0)
+        pad_h, pad_w = args[3] if len(args) > 3 else (0, 0)
         ceil_mode = args[4] if len(args) > 4 else False
         count_include_pad = args[5] if len(args) > 5 else True
         divisor_override = args[6] if len(args) > 6 else None
 
         n, c, h, w = x.data.shape
-        post_pad_w, post_pad_h = (0, 0)
 
         # Count_include_pad == False means that we use a different divisor for edge elements
         # When divisor_override is set, this will be overriden anyways.
         # It is easier to replace a constant divisor, so set count_include_pad == True
         if divisor_override is not None:
             count_include_pad = True
 
-        # Add width padding manually if count_include_pad
-        if count_include_pad and pad_w > 0:
-            pre_pad_shape = [n, c, h, pad_w]
-            pre_pad = super().call_operator(
-                full_op, (pre_pad_shape, 0.0), full_kwargs, meta, updated=True
-            )
-
-            if ceil_mode and divisor_override is None:
-                post_pad_w = pad_w
-            else:
-                post_pad_w = adjust_pooling_pad_if_needed(
-                    w, kernel_w, stride_w, pad_w, ceil_mode
-                )
-
-            if post_pad_w > 0:
-                post_pad_shape = [n, c, h, post_pad_w]
-                post_pad = super().call_operator(
-                    full_op, (post_pad_shape, 0.0), full_kwargs, meta, updated=True
-                )
-                cat_nodes = [pre_pad, x, post_pad]
-            else:
-                cat_nodes = [pre_pad, x]
-
-            x = super().call_operator(
-                cat_op, (cat_nodes, 3), kwargs, meta, updated=True
-            )
-            new_pad_w = 0
-
-        # Add height padding manually if count_include_pad
-        if count_include_pad and pad_h > 0:
-            pre_pad_shape = [n, c, pad_h, w + pad_w + post_pad_w]
-            pre_pad = super().call_operator(
-                full_op, (pre_pad_shape, 0.0), full_kwargs, meta, updated=True
-            )
+        pad, new_pad = _get_avgpool_post_pad(
+            h,
+            w,
+            args[1],
+            stride_h,
+            stride_w,
+            pad_h,
+            pad_w,
+            ceil_mode,
+            count_include_pad,
+            divisor_override,
+        )
 
-            if ceil_mode and divisor_override is None:
-                post_pad_h = pad_h
-            else:
-                post_pad_h = adjust_pooling_pad_if_needed(
-                    h, kernel_h, stride_h, pad_h, ceil_mode
-                )
-
-            if post_pad_h > 0:
-                post_pad_shape = [n, c, post_pad_h, w + pad_w + post_pad_w]
-                post_pad = super().call_operator(
-                    full_op, (post_pad_shape, 0.0), full_kwargs, meta, updated=True
-                )
-                cat_nodes = [pre_pad, x, post_pad]
+        if count_include_pad and (pad_h > 0 or pad_w > 0):
+            if op in aten_avg_pool2d:
+                pad_args = (x, pad, "constant", 0.0)
             else:
-                cat_nodes = [pre_pad, x]
+                pad_args = (x, pad, 0.0)
 
             x = super().call_operator(
-                cat_op, (cat_nodes, 2), kwargs, meta, updated=True
+                pad_op,
+                pad_args,
+                {},
+                meta,
+                updated=True,
             )
-            new_pad_h = 0
 
         avgpool_args = (
             x,
             args[1],
             [stride_h, stride_w],
-            [new_pad_h, new_pad_w],
+            new_pad,
             ceil_mode,
             False,
         )
+
         x = super().call_operator(avgpool_op, avgpool_args, kwargs, meta, updated=True)
 
-        # Multiply by factor (kernel_size / divisor_override) if divisor_override
         if divisor_override is not None and divisor_override != kernel_size:
-            override_multiplier = super().call_operator(
-                full_op,
-                ([1, 1, 1, 1], kernel_size / divisor_override),
-                full_kwargs,
+            x = super().call_operator(
+                mul_op,
+                (x, super().call_scalar(kernel_size / divisor_override, meta)),
+                {},
                 meta,
                 updated=True,
             )
-            x = super().call_operator(
-                mul_op, (x, override_multiplier), kwargs, meta, updated=True
-            )
 
         return x