Arm backend: Support for avg_pool2d for TOSA 1.0

Add support for AVG_POOL2D TOSA 1.0

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

Author: oscarandersson8218

backends/arm/operators/op_avg_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,
@@ -36,14 +34,16 @@ def __init__(self, *args):
     def _build_generic_avgpool2d(
         self,
         node: torch.fx.Node,
-        tosa_graph: ts.TosaSerializer,
+        tosa_graph: Any,
         inputs: List[TosaArg],
         output: TosaArg,
         input_zp: int,
         output_zp: int,
-        accumulator_type: ts.DType,
+        accumulator_type: Any,
     ) -> None:
 
+        import tosa_tools.v0_80.serializer.tosa_serializer as ts  # type: ignore
+
         input_tensor = inputs[0]
         kernel_size_list = inputs[1].special
         stride_size_list = inputs[2].special
@@ -79,10 +79,12 @@ def _build_generic_avgpool2d(
     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]
         assert input_tensor.dtype == ts.DType.INT8
 
@@ -110,10 +112,135 @@ class AvgPool2dVisitor_0_80_MI(AvgPool2dVisitor_0_80_BI):
     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
+
+        assert (
+            inputs[0].dtype == ts.DType.INT8 or inputs[0].dtype == ts.DType.FP32
+        ), "Only FP32 and INT8 supported"
+
+        if inputs[0].dtype == ts.DType.INT8:
+            super().define_node(node, tosa_graph, inputs, output)
+
+        if inputs[0].dtype == ts.DType.FP32:
+            accumulator_type = ts.DType.FP32
+            # Initilize zero point to zero.
+            input_zp = 0
+            output_zp = 0
+
+            self._build_generic_avgpool2d(
+                node, tosa_graph, inputs, output, input_zp, output_zp, accumulator_type
+            )
+
+
+@register_node_visitor
+class AvgPool2dVisitor(NodeVisitor):
+    target = "aten.avg_pool2d.default"
+
+    tosa_specs = [
+        TosaSpecification.create_from_string("TOSA-1.0+INT"),
+    ]
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def _build_generic_avgpool2d(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+        input_zp: int,
+        output_zp: int,
+        accumulator_type: Any,
+    ) -> None:
+
+        import serializer.tosa_serializer as ts  # type: ignore
+
+        input_tensor = inputs[0]
+        kernel_size_list = inputs[1].special
+        stride_size_list = 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.AvgPool2dAttribute(
+            kernel=kernel_size_list,
+            stride=stride_size_list,
+            pad=pad_size_list,
+            acc_type=accumulator_type,
+        )
+        input_zp_tensor = tosa_graph.addConst(
+            shape=[1], dtype=output.dtype, vals=[input_zp]
+        )
+        output_zp_tensor = tosa_graph.addConst(
+            shape=[1], dtype=output.dtype, vals=[output_zp]
+        )
+
+        tosa_graph.addOperator(
+            ts.TosaOp.Op().AVG_POOL2D,
+            [input_tensor.name, input_zp_tensor.name, output_zp_tensor.name],
+            [output.name],
+            attr,
+        )
+
+    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]
+        assert input_tensor.dtype == ts.DType.INT8
+
+        accumulator_type = ts.DType.INT32
+
+        input_qargs = get_input_qparams(node)
+        input_zp = input_qargs[0].zp
+
+        output_qargs = get_output_qparams(node)
+        output_zp = output_qargs[0].zp
+
+        self._build_generic_avgpool2d(
+            node, tosa_graph, inputs, output, input_zp, output_zp, accumulator_type
+        )
+
+
+@register_node_visitor
+class AvgPool2dVisitor_FP(AvgPool2dVisitor):
+    target = "aten.avg_pool2d.default"
+
+    tosa_specs = [
+        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
+
         assert (
             inputs[0].dtype == ts.DType.INT8 or inputs[0].dtype == ts.DType.FP32
         ), "Only FP32 and INT8 supported"