Fix the stride issue in DecomposeAvgPool2d and add a test for it. (#13152)

Author: tirwu01

backends/arm/_passes/decompose_avg_pool2d.py

@@ -45,7 +45,10 @@ def call_operator(self, op, args, kwargs, meta):
         x = args[0]
         kernel_h, kernel_w = args[1]
         kernel_size = kernel_h * kernel_w
-        stride_h, stride_w = args[2]
+        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)
         ceil_mode = args[4] if len(args) > 4 else False
         count_include_pad = args[5] if len(args) > 5 else True
@@ -108,7 +111,14 @@ def call_operator(self, op, args, kwargs, meta):
             x = super().call_operator(cat_op, (cat_nodes, 2), kwargs, meta)
             new_pad_h = 0
 
-        avgpool_args = (x, args[1], args[2], [new_pad_h, new_pad_w], ceil_mode, False)
+        avgpool_args = (
+            x,
+            args[1],
+            [stride_h, stride_w],
+            [new_pad_h, new_pad_w],
+            ceil_mode,
+            False,
+        )
         x = super().call_operator(avgpool_op, avgpool_args, kwargs, meta)
 
         # Multiply by factor (kernel_size / divisor_override) if divisor_override

backends/arm/test/passes/test_decompose_avg_pool2d_pass.py

@@ -0,0 +1,75 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Tuple
+
+import torch
+from executorch.backends.arm._passes.decompose_avg_pool2d import DecomposeAvgPool2d
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+
+input_t = Tuple[torch.Tensor]  # Input x
+
+
+class AvgPool2dWithStride(torch.nn.Module):
+    """
+    avg_pool2d model with explicit stride parameter
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(1, 3, 8, 8),)
+
+    def forward(self, x):
+        return torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+
+
+class AvgPool2dWithoutStride(torch.nn.Module):
+    """
+    avg_pool2d model without stride parameter (should default to kernel_size)
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(1, 3, 8, 8),)
+
+    def forward(self, x):
+        return torch.nn.functional.avg_pool2d(x, kernel_size=3)
+
+
+class AvgPool2dListKernel(torch.nn.Module):
+    """
+    avg_pool2d model with list kernel_size and no stride
+    """
+
+    def get_inputs(self) -> input_t:
+        return (torch.rand(1, 3, 8, 8),)
+
+    def forward(self, x):
+        return torch.nn.functional.avg_pool2d(x, kernel_size=[2, 3])
+
+
+modules = {
+    "avg_pool2d_with_stride": AvgPool2dWithStride(),
+    "avg_pool2d_without_stride": AvgPool2dWithoutStride(),
+    "avg_pool2d_list_kernel": AvgPool2dListKernel(),
+}
+
+
+@common.parametrize("module", modules)
+def test_decompose_avg_pool2d_tosa_MI(module):
+    """Test that DecomposeAvgPool2d pass works correctly with and without stride parameters."""
+    pipeline = PassPipeline[input_t](
+        module,
+        module.get_inputs(),
+        quantize=False,
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
+        },
+        ops_after_pass={
+            # After decomposition, we should still see avg_pool2d (transformed)
+            "executorch_exir_dialects_edge__ops_aten_avg_pool2d_default": 1,
+        },
+        pass_list=[DecomposeAvgPool2d],
+    )
+    pipeline.run()