diff --git a/backends/arm/operator_support/pool_2d_support.py b/backends/arm/operator_support/pool_2d_support.py index 750fab2730d..f4ada36de80 100644 --- a/backends/arm/operator_support/pool_2d_support.py +++ b/backends/arm/operator_support/pool_2d_support.py @@ -54,8 +54,11 @@ 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}" ) diff --git a/backends/arm/test/ops/test_avg_pool2d.py b/backends/arm/test/ops/test_avg_pool2d.py index 2a50ef38834..c48595aec7f 100644 --- a/backends/arm/test/ops/test_avg_pool2d.py +++ b/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.is_option_enabled("tosa_ref_model"): + 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.is_option_enabled("tosa_ref_model"): + pipeline.change_args("run_method_and_compare_outputs", qtol=1, atol=1, rtol=1) + pipeline.run() @common.parametrize("test_module", test_modules) diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 832dcb3286c..acb27f13798 100644 --- a/backends/arm/test/targets.bzl +++ b/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",