Arm backend: Support for max_pool2d for TOSA 1.0

Add support for MAX_POOL2D TOSA 1.0

Signed-off-by: Oscar Andersson <[email protected]>
Change-Id: I3432622ede7b029e78ca8af5c9c71f17f551d4b3

Author: oscarandersson8218

backends/arm/operators/op_max_pool2d.py

@@ -4,12 +4,10 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-unsafe
-from typing import List
+from typing import Any, List
 
 import torch
 
-import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
-
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     get_input_qparams,
     get_output_qparams,
@@ -19,22 +17,29 @@
     register_node_visitor,
 )
 from executorch.backends.arm.tosa_mapping import TosaArg
+from executorch.backends.arm.tosa_specification import TosaSpecification
 
 
 @register_node_visitor
-class MaxPool2dVisitor(NodeVisitor):
+class MaxPool2dVisitor_0_80(NodeVisitor):
     target = "aten.max_pool2d.default"
 
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-0.80+BI"),
+        TosaSpecification.create_from_string("TOSA-0.80+MI"),
+    ]
+
     def __init__(self, *args):
         super().__init__(*args)
 
     def define_node(
         self,
         node: torch.fx.Node,
-        tosa_graph: ts.TosaSerializer,
+        tosa_graph: Any,
         inputs: List[TosaArg],
         output: TosaArg,
     ) -> None:
+        import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
 
         input_tensor = inputs[0]
         kernel_size = inputs[1].special
@@ -80,3 +85,53 @@ def define_node(
             [output.name],
             attr,
         )
+
+
+@register_node_visitor
+class MaxPool2dVisitor(NodeVisitor):
+    target = "aten.max_pool2d.default"
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-1.0+INT"),
+        TosaSpecification.create_from_string("TOSA-1.0+FP"),
+    ]
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        import serializer.tosa_serializer as ts  # type: ignore
+
+        input_tensor = inputs[0]
+        kernel_size = inputs[1].special
+        stride = inputs[2].special
+
+        try:
+            pad_size_list = inputs[3].special
+            pad_size_list = [
+                pad_size_list[0],
+                pad_size_list[0],
+                pad_size_list[1],
+                pad_size_list[1],
+            ]
+        except IndexError:
+            pad_size_list = [0, 0, 0, 0]
+
+        attr = ts.TosaSerializerAttribute()
+        attr.MaxPool2dAttribute(
+            kernel=kernel_size, stride=stride, pad=pad_size_list, nan_mode=1
+        )
+
+        tosa_graph.addOperator(
+            ts.TosaOp.Op().MAX_POOL2D,
+            [input_tensor.name],
+            [output.name],
+            attr,
+        )