From 967670890baa54717811f8e72e011a479cca6d6b Mon Sep 17 00:00:00 2001
From: Tom Allsop <tom.allsop@arm.com>
Date: Tue, 6 May 2025 11:38:52 +0100
Subject: [PATCH] Arm backend: Adjust MaxPool2d padding when window is not
 divisible by stride

* MaxPool2dVisitor will adjust padding if the pooling window is not divisible by the stride

Signed-off-by: Tom Allsop <tom.allsop@arm.com>
Change-Id: I92c4f714067b8498f4df328dc971991a2585ec6c
---
 backends/arm/operators/op_max_pool2d.py | 46 +++++++++++++++++++++++++
 backends/arm/test/ops/test_max_pool.py  |  1 +
 2 files changed, 47 insertions(+)

diff --git a/backends/arm/operators/op_max_pool2d.py b/backends/arm/operators/op_max_pool2d.py
index 40f48d3896f..6ba85789914 100644
--- a/backends/arm/operators/op_max_pool2d.py
+++ b/backends/arm/operators/op_max_pool2d.py
@@ -23,6 +23,24 @@
 from executorch.backends.arm.tosa_specification import TosaSpecification
 
 
+# Similarly to Conv2d, the TOSA spec requires that following is exactly divisible:
+# `(input + 2 * pad - kernel_size) / stride`
+# PyTorch however, does not require this, so as needed, we must adjust the padding.
+def adjust_pad_if_needed(
+    input_size: int, kernel_size: int, stride: int, pad: int
+) -> int:
+    if pad == 0:
+        return pad
+
+    mod_remainder = (input_size + 2 * pad - kernel_size) % stride
+
+    # No need to adjust
+    if mod_remainder == 0:
+        return pad
+
+    return pad - mod_remainder
+
+
 @register_node_visitor
 class MaxPool2dVisitor_0_80(NodeVisitor):
     target = "aten.max_pool2d.default"
@@ -61,6 +79,20 @@ def define_node(
         except IndexError:
             pad_size_list = [0, 0, 0, 0]
 
+        # Adjust the padding as necessary
+        pad_size_list[1] = adjust_pad_if_needed(
+            input_tensor.shape[2],
+            kernel_size[0],
+            stride[0],
+            pad_size_list[1],
+        )
+        pad_size_list[3] = adjust_pad_if_needed(
+            input_tensor.shape[3],
+            kernel_size[1],
+            stride[1],
+            pad_size_list[3],
+        )
+
         accumulator_type = output.dtype
 
         # Initilize zero point to zero.
@@ -131,6 +163,20 @@ def define_node(
         except IndexError:
             pad_size_list = [0, 0, 0, 0]
 
+        # Adjust the padding as necessary
+        pad_size_list[1] = adjust_pad_if_needed(
+            input_tensor.shape[2],
+            kernel_size[0],
+            stride[0],
+            pad_size_list[1],
+        )
+        pad_size_list[3] = adjust_pad_if_needed(
+            input_tensor.shape[3],
+            kernel_size[1],
+            stride[1],
+            pad_size_list[3],
+        )
+
         attr = ts.TosaSerializerAttribute()
         attr.MaxPool2dAttribute(
             kernel=kernel_size, stride=stride, pad=pad_size_list, nan_mode=1
diff --git a/backends/arm/test/ops/test_max_pool.py b/backends/arm/test/ops/test_max_pool.py
index 657597f9058..4db8c62bd88 100644
--- a/backends/arm/test/ops/test_max_pool.py
+++ b/backends/arm/test/ops/test_max_pool.py
@@ -31,6 +31,7 @@
     ("zeros", torch.zeros(1, 1, 4, 8), [2, 2, 1]),
     ("ones", torch.ones(1, 16, 50, 32), [4, 2, 0]),
     ("rand", torch.rand(1, 16, 52, 16), [4, 3, 0]),
+    ("non_divisible", torch.rand(1, 16, 112, 112), [3, 2, 1]),
 ]
 
 test_data_suite_mult_batches = [