Handle avg_pool2d with padding == 0 as no padding (#10697)

Summary: Pull Request resolved: #10697

Differential Revision: D74117402

Author: 3l1

backends/arm/operator_support/pool_2d_support.py

@@ -54,8 +54,9 @@ def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
         kernel = cast(tuple[int, int], node.args[1])
         stride = cast(tuple[int, int], node.args[2])
         if len(node.args) > 3:
+            padding = cast(tuple[int, int], node.args[3])
             # Padding case
-            if not all(1 <= k <= 8 for k in kernel):
+            if not all(1 <= k <= 8 for k in kernel) and not all(v == 0 for v in padding):
                 self.reporter.report_reject(
                     node, f"Avgpool2d with padding needs kernel dims < 8, got {kernel}"
                 )

backends/arm/test/ops/test_avg_pool2d.py

@@ -9,9 +9,11 @@
 
 from typing import Tuple
 
+import pytest
+
 import torch
 
-from executorch.backends.arm.test import common
+from executorch.backends.arm.test import common, conftest
 
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineBI,
@@ -64,15 +66,24 @@ def forward(self, x):
 
 
 @common.parametrize("test_module", test_modules)
+@pytest.mark.tosa_ref_model
 def test_avgpool2d_tosa_MI(test_module):
     model, input_tensor = test_module
 
-    pipeline = TosaPipelineMI[input_t](model, input_tensor, aten_op, exir_op)
-    pipeline.change_args("run_method_and_compare_outputs", qtol=1, atol=1, rtol=1)
-    pipeline.run()
+    pipeline = TosaPipelineMI[input_t](
+        model,
+        input_tensor,
+        aten_op,
+        exir_op,
+        run_on_tosa_ref_model=conftest.is_option_enabled("tosa_ref_model"),
+    )
+    if conftest.get_option("tosa_version") == "1.0":
+        pipeline.change_args("run_method_and_compare_outputs", qtol=1, atol=1, rtol=1)
+        pipeline.run()
 
 
 @common.parametrize("test_module", test_modules)
+@pytest.mark.tosa_ref_model
 def test_avgpool2d_tosa_BI(test_module):
     model, input_tensor = test_module
 
@@ -82,9 +93,11 @@ def test_avgpool2d_tosa_BI(test_module):
         aten_op,
         exir_op,
         symmetric_io_quantization=True,
+        run_on_tosa_ref_model=conftest.is_option_enabled("tosa_ref_model"),
     )
-    pipeline.change_args("run_method_and_compare_outputs", qtol=1, atol=1, rtol=1)
-    pipeline.run()
+    if conftest.get_option("tosa_version") == "0.80":
+        pipeline.change_args("run_method_and_compare_outputs", qtol=1, atol=1, rtol=1)
+        pipeline.run()
 
 
 @common.parametrize("test_module", test_modules)

backends/arm/test/targets.bzl

@@ -13,6 +13,7 @@ def define_arm_tests():
 
     # Operators
     test_files += [
+        "ops/test_avg_pool2d.py",
         "ops/test_linear.py", 
         "ops/test_slice.py",
         "ops/test_sigmoid.py",