add support for AveragePool, GlobalAveragePool, and GlobalMaxPool operators for Generic target

Author: marchioa

Deeploy/Targets/Generic/Bindings.py

@@ -13,8 +13,9 @@
 from Deeploy.FutureExtension.CodeTransformationPasses.FutureCodeTransformation import FutureGeneration
 from Deeploy.Targets.Generic.Templates import AddTemplate, BatchNormalizationTemplate, ConcatTemplate, ConvTemplate, \
     ConvTransposeTemplate, DebugPrintTemplate, DequantTemplate, DummyTemplate, DWConvTemplate, FloatAddTemplate, \
-    FloatCeilTemplate, FloatClipTemplate, FloatConvTemplate, FloatDivTemplate, FloatDWConvTemplate, FloatExpTemplate, \
-    FloatFloorTemplate, FloatGELUTemplate, FloatGemmTemplate, FloatGroupNormTemplate, FloatHardSigmoidTemplate, \
+    FloatAveragePoolTemplate, FloatCeilTemplate, FloatClipTemplate, FloatConvTemplate, FloatDivTemplate, \
+    FloatDWConvTemplate, FloatExpTemplate, FloatFloorTemplate, FloatGELUTemplate, FloatGemmTemplate, \
+    FloatGlobalAveragePoolTemplate, FloatGlobalMaxPoolTemplate, FloatGroupNormTemplate, FloatHardSigmoidTemplate, \
     FloatHardSwishTemplate, FloatInstanceNormTemplate, FloatLayernormTemplate, FloatMatMulTemplate, \
     FloatMaxPoolTemplate, FloatMulTemplate, FloatPadTemplate, FloatPowTemplate, FloatReduceMeanTemplate, \
     FloatReluTemplate, FloatSigmoidTemplate, FloatSoftmaxTemplate, FloatSqrtTemplate, FloatSubTemplate, \
@@ -399,3 +400,23 @@
              PointerClass(float32_t)], [PointerClass(float32_t)]), FloatGroupNormTemplate.referenceTemplate,
         BasicTransformer),
 ]
+
+BasicAveragePool1DBindings = [
+    NodeBinding(DummyChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatAveragePoolTemplate.referenceTemplate1d, BasicTransformer)
+]
+
+BasicAveragePool2DBindings = [
+    NodeBinding(DummyChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatAveragePoolTemplate.referenceTemplate2d, BasicTransformer)
+]
+
+BasicGlobalAveragePoolBindings = [
+    NodeBinding(DummyChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatGlobalAveragePoolTemplate.referenceTemplate, BasicTransformer)
+]
+
+BasicGlobalMaxPoolBindings = [
+    NodeBinding(DummyChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatGlobalMaxPoolTemplate.referenceTemplate, BasicTransformer)
+]

Deeploy/Targets/Generic/Layers.py

@@ -752,3 +752,15 @@ class InstanceNormLayer(ONNXLayer):
 
 class GroupNormLayer(ONNXLayer):
     pass
+
+
+class AveragePoolLayer(ONNXLayer):
+    pass
+
+
+class GlobalAveragePoolLayer(ONNXLayer):
+    pass
+
+
+class GlobalMaxPoolLayer(ONNXLayer):
+    pass

Deeploy/Targets/Generic/Parsers.py

@@ -3000,3 +3000,120 @@ def parseNode(self, node: gs.Node) -> bool:
             return False
         self.operatorRepresentation['num_groups'] = node.attrs['num_groups']
         return True
+
+
+class AveragePoolParser(NodeParser):
+
+    def parseNode(self, node: gs.Node) -> bool:
+
+        if not all([
+                node.op == 'AveragePool',
+                len(node.inputs) == 1,
+                len(node.outputs) == 1,
+                'kernel_shape' in node.attrs,
+        ]):
+            return False
+
+        kernel_shape = node.attrs['kernel_shape']
+        spatial_ndim = len(kernel_shape)
+
+        auto_pad = node.attrs.get('auto_pad', 'NOTSET')
+        ceil_mode = node.attrs.get('ceil_mode', 0)
+        count_include_pad = node.attrs.get('count_include_pad ', 0)
+        dilations = node.attrs.get('dilations', (1,) * spatial_ndim)
+        strides = node.attrs.get('strides', (1,) * spatial_ndim)
+        pads = node.attrs.get('pads', (0,) * (2 * spatial_ndim))
+
+        if not all([
+                auto_pad == 'NOTSET',  # TODO: implement other values
+                ceil_mode == 0,  # TODO: implement other values
+                count_include_pad == 0,  # TODO: implement other values
+                all([d == 1 for d in dilations]),  # TODO: implement other values
+                len(dilations) == spatial_ndim,
+                len(strides) == spatial_ndim,
+                len(pads) == 2 * spatial_ndim,
+                all([s > 0 for s in strides]),
+        ]):
+            return False
+
+        self.operatorRepresentation['kernel_shape'] = kernel_shape
+        self.operatorRepresentation['auto_pad'] = auto_pad
+        self.operatorRepresentation['ceil_mode'] = ceil_mode
+        self.operatorRepresentation['count_include_pad'] = count_include_pad
+        self.operatorRepresentation['dilations'] = dilations
+        self.operatorRepresentation['strides'] = strides
+        self.operatorRepresentation['pads'] = pads
+
+        return True
+
+    def parseNodeCtxt(self,
+                      ctxt: NetworkContext,
+                      node: gs.Node,
+                      channels_first: bool = True) -> Tuple[NetworkContext, bool]:
+
+        data_in = ctxt.lookup(node.inputs[0].name)
+        data_out = ctxt.lookup(node.outputs[0].name)
+        self.operatorRepresentation['data_in'] = data_in.name
+        self.operatorRepresentation['data_out'] = data_out.name
+
+        self.operatorRepresentation['batch_size'] = data_in.shape[0]
+        self.operatorRepresentation['num_channels'] = data_in.shape[1]
+
+        spatial_shape = data_in.shape[2:]
+        if len(self.operatorRepresentation['kernel_shape']) != len(spatial_shape):
+            return ctxt, False
+
+        if len(spatial_shape) == 1:
+            self.operatorRepresentation['length'] = spatial_shape[0]
+        elif len(spatial_shape) == 2:
+            self.operatorRepresentation['height'] = spatial_shape[0]
+            self.operatorRepresentation['width'] = spatial_shape[1]
+        else:
+            return ctxt, False
+
+        return ctxt, True
+
+
+class AveragePool1DParser(AveragePoolParser):
+
+    def parseNode(self, node: gs.Node) -> bool:
+        return super().parseNode(node) and len(node.attrs['kernel_shape']) == 1
+
+
+class AveragePool2DParser(AveragePoolParser):
+
+    def parseNode(self, node: gs.Node) -> bool:
+        return super().parseNode(node) and len(node.attrs['kernel_shape']) == 2
+
+
+class GlobalPoolParser(NodeParser):
+
+    def parseNode(self, node: gs.Node) -> bool:
+        return len(node.inputs) == 1 and len(node.outputs) == 1
+
+    def parseNodeCtxt(self,
+                      ctxt: NetworkContext,
+                      node: gs.Node,
+                      channels_first: bool = True) -> Tuple[NetworkContext, bool]:
+
+        data_in = ctxt.lookup(node.inputs[0].name)
+        data_out = ctxt.lookup(node.outputs[0].name)
+        self.operatorRepresentation['data_in'] = data_in.name
+        self.operatorRepresentation['data_out'] = data_out.name
+        self.operatorRepresentation['batch_size'] = data_in.shape[0]
+        self.operatorRepresentation['num_channels'] = data_in.shape[1]
+        self.operatorRepresentation['spatial_size'] = np.prod(data_in.shape[2:])
+
+        return ctxt, True
+
+
+class GlobalAveragePoolParser(GlobalPoolParser):
+
+    def parseNode(self, node: gs.Node) -> bool:
+        return super().parseNode(node) and node.op == 'GlobalAveragePool'
+
+
+class GlobalMaxPoolParser(GlobalPoolParser):
+
+    def parseNode(self, node: gs.Node) -> bool:
+        return super().parseNode(node) and node.op == 'GlobalMaxPool'

Deeploy/Targets/Generic/Platform.py

@@ -6,32 +6,33 @@
     RemoveEmptyConvBiasPass, RemoveOnlySingletonReduceMeanPass
 from Deeploy.DeeployTypes import ConstantBuffer, DeploymentEngine, DeploymentPlatform, NodeMapper, NodeTemplate, \
     StructBuffer, TopologyOptimizer, TransientBuffer, VariableBuffer
-from Deeploy.Targets.Generic.Bindings import BasicAddBindings, BasicBatchNormBindings, BasicCeilBindings, \
-    BasicClipBindings, BasicConcatBindings, BasicConv1DBindings, BasicConv2DBindings, BasicConvTransposeBindings, \
-    BasicDebugPrintBindings, BasicDequantBindings, BasicDivBindings, BasicDWConv1DBinding, BasicDWConv2DBindings, \
-    BasicExpBindings, BasicFloorBindings, BasicGatherBindings, BasicGELUBindings, BasicGEMMBindings, \
+from Deeploy.Targets.Generic.Bindings import BasicAddBindings, BasicAveragePool1DBindings, BasicAveragePool2DBindings, \
+    BasicBatchNormBindings, BasicCeilBindings, BasicClipBindings, BasicConcatBindings, BasicConv1DBindings, \
+    BasicConv2DBindings, BasicConvTransposeBindings, BasicDebugPrintBindings, BasicDequantBindings, BasicDivBindings, \
+    BasicDWConv1DBinding, BasicDWConv2DBindings, BasicExpBindings, BasicFloorBindings, BasicGatherBindings, \
+    BasicGELUBindings, BasicGEMMBindings, BasicGlobalAveragePoolBindings, BasicGlobalMaxPoolBindings, \
     BasicGroupNormBindings, BasicHardSigmoidBindings, BasicHardSwishBindings, BasicInstanceNormBindings, \
     BasicITAPartialSoftmaxBinding, BasicITASoftmaxBinding, BasicLayerNormBindings, BasicMatMulBindings, \
     BasicMaxPool1DBindings, BasicMaxPool2DBindings, BasicMulBindings, BasicPad1DBindings, BasicPad2DBindings, \
     BasicPowBindings, BasicQuantBindings, BasicReduceMeanBindings, BasicReduceSumBindings, BasicReluBinding, \
     BasicReshapeBindings, BasicRQIntegerDivBinding, BasicRQSBindings, BasicRQSGELUBinding, BasicSigmoidBindings, \
     BasicSliceBindings, BasicSoftmaxBindings, BasicSqrtBindings, BasicSubBindings, BasicSwishBindings, \
     BasicTransposeBindings, DummyBinding
-from Deeploy.Targets.Generic.Layers import AddLayer, BatchNormalizationLayer, CeilLayer, ClipLayer, ConcatLayer, \
-    ConvLayer, ConvTransposeLayer, DebugPrintLayer, DequantLayer, DivLayer, ExpLayer, FloorLayer, GatherLayer, \
-    GELULayer, GEMMLayer, GroupNormLayer, InstanceNormLayer, ITAMaxLayer, LayerNormLayer, MatMulLayer, MaxPoolLayer, \
-    MulLayer, PadLayer, PowLayer, QuantLayer, ReduceMeanLayer, ReduceSumLayer, ReluLayer, RequantShiftLayer, \
-    ReshapeLayer, RQIntegerDivLayer, RQSiGELULayer, SigmoidLayer, SliceLayer, SoftmaxLayer, SqrtLayer, SubLayer, \
-    SwishLayer, TransposeLayer
-from Deeploy.Targets.Generic.Parsers import AddParser, BatchNormParser, CeilParser, ClipParser, ConcatParser, \
-    ConvTranspose1DParser, DebugParser, DequantParser, DivParser, DummyParser, ExpParser, FlattenParser, FloorParser, \
-    GatherParser, GELUParser, GenericConv1DParser, GenericConv2DParser, GenericDWConv1DParser, GenericDWConv2DParser, \
-    GenericGEMMParser, GenericMaxPool2DParser, GroupNormParser, HardSigmoidParser, HardSwishParser, \
-    InstanceNormParser, IntegerDivParser, ITAMaxParser, ITAPartialMaxParser, LayerNormParser, MatMulParser, \
-    MaxPool1DParser, MulParser, Pad1DParser, Pad2DParser, PowParser, QuantParser, ReduceMeanParser, ReduceSumParser, \
-    ReluParser, RequantShiftParser, ReshapeParser, RQIntegerDivParser, RQSiGELUParser, SigmoidParser, SliceParser, \
-    SoftmaxParser, SqrtParser, SubParser, SwishParser, TransposeParser, UnsqueezeParser, iLayerNormParser, \
-    iSoftmaxParser
+from Deeploy.Targets.Generic.Layers import AddLayer, AveragePoolLayer, BatchNormalizationLayer, CeilLayer, ClipLayer, \
+    ConcatLayer, ConvLayer, ConvTransposeLayer, DebugPrintLayer, DequantLayer, DivLayer, ExpLayer, FloorLayer, \
+    GatherLayer, GELULayer, GEMMLayer, GlobalAveragePoolLayer, GlobalMaxPoolLayer, GroupNormLayer, InstanceNormLayer, \
+    ITAMaxLayer, LayerNormLayer, MatMulLayer, MaxPoolLayer, MulLayer, PadLayer, PowLayer, QuantLayer, ReduceMeanLayer, \
+    ReduceSumLayer, ReluLayer, RequantShiftLayer, ReshapeLayer, RQIntegerDivLayer, RQSiGELULayer, SigmoidLayer, \
+    SliceLayer, SoftmaxLayer, SqrtLayer, SubLayer, SwishLayer, TransposeLayer
+from Deeploy.Targets.Generic.Parsers import AddParser, AveragePool1DParser, AveragePool2DParser, BatchNormParser, \
+    CeilParser, ClipParser, ConcatParser, ConvTranspose1DParser, DebugParser, DequantParser, DivParser, DummyParser, \
+    ExpParser, FlattenParser, FloorParser, GatherParser, GELUParser, GenericConv1DParser, GenericConv2DParser, \
+    GenericDWConv1DParser, GenericDWConv2DParser, GenericGEMMParser, GenericMaxPool2DParser, GlobalAveragePoolParser, \
+    GlobalMaxPoolParser, GroupNormParser, HardSigmoidParser, HardSwishParser, InstanceNormParser, IntegerDivParser, \
+    ITAMaxParser, ITAPartialMaxParser, LayerNormParser, MatMulParser, MaxPool1DParser, MulParser, Pad1DParser, \
+    Pad2DParser, PowParser, QuantParser, ReduceMeanParser, ReduceSumParser, ReluParser, RequantShiftParser, \
+    ReshapeParser, RQIntegerDivParser, RQSiGELUParser, SigmoidParser, SliceParser, SoftmaxParser, SqrtParser, \
+    SubParser, SwishParser, TransposeParser, UnsqueezeParser, iLayerNormParser, iSoftmaxParser
 from Deeploy.Targets.Generic.Templates import AllocateTemplate, FreeTemplate
 from Deeploy.Targets.Generic.TopologyOptimizationPasses.Passes import DequantPatternPass, ExtractPaddingFromConvPass, \
     ExtractPaddingFromPoolPass, MatMulAddMergePass, MergeConstAddAndRequantPass, QuantPatternPass, \
@@ -89,6 +90,10 @@
 HardSwishMapper = NodeMapper(HardSwishParser(), BasicHardSwishBindings)
 InstanceNormMapper = NodeMapper(InstanceNormParser(), BasicInstanceNormBindings)
 GroupNormMapper = NodeMapper(GroupNormParser(), BasicGroupNormBindings)
+AveragePool1DMapper = NodeMapper(AveragePool1DParser(), BasicAveragePool1DBindings)
+AveragePool2DMapper = NodeMapper(AveragePool2DParser(), BasicAveragePool2DBindings)
+GlobalAveragePoolMapper = NodeMapper(GlobalAveragePoolParser(), BasicGlobalAveragePoolBindings)
+GlobalMaxPoolMapper = NodeMapper(GlobalMaxPoolParser(), BasicGlobalMaxPoolBindings)
 
 # Dummy nodes are intended for development purposes only!
 # They should always generate compiler errors to not accidentally end up in production code
@@ -146,6 +151,9 @@
     'HardSwish': SwishLayer([HardSwishMapper]),
     'InstanceNormalization': InstanceNormLayer([InstanceNormMapper]),
     'GroupNormalization': GroupNormLayer([GroupNormMapper]),
+    'AveragePool': AveragePoolLayer([AveragePool1DMapper, AveragePool2DMapper]),
+    'GlobalAveragePool': GlobalAveragePoolLayer([GlobalAveragePoolMapper]),
+    'GlobalMaxPool': GlobalMaxPoolLayer([GlobalMaxPoolMapper]),
     # # For example, you can use the DummpyMapper, in case you want to test
     # # deployment or optimizations with GlobalAveragePool nodes but did not yet
     # # implement the corresponding kernel

Deeploy/Targets/Generic/Templates/FloatAveragePoolTemplate.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: 2023 ETH Zurich and University of Bologna
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from Deeploy.DeeployTypes import NetworkContext, NodeTemplate, OperatorRepresentation
+
+
+class _AveragePoolTemplate(NodeTemplate):
+
+    def alignToContext(self, ctxt: NetworkContext,
+                       operatorRepresentation: OperatorRepresentation) -> tuple[NetworkContext, dict, list[str]]:
+
+        data_in = ctxt.lookup(operatorRepresentation['data_in'])
+        operatorRepresentation['type_width'] = data_in._type.referencedType.typeWidth
+        return ctxt, operatorRepresentation, []
+
+
+referenceTemplate1d = _AveragePoolTemplate("""
+// Average Pool 1D (Name: ${nodeName}, Op: ${nodeOp})
+AveragePool1d_fp${type_width}_fp${type_width}(
+    ${data_in}, ${data_out}, ${batch_size}, ${num_channels}, ${length}, ${kernel_shape[0]},
+    ${strides[0]}, ${pads[0]}, ${pads[1]});
+""")
+
+referenceTemplate2d = _AveragePoolTemplate("""
+// Average Pool 2D (Name: ${nodeName}, Op: ${nodeOp})
+AveragePool2d_fp${type_width}_fp${type_width}(
+    ${data_in}, ${data_out}, ${batch_size}, ${num_channels}, ${height}, ${width},
+    ${kernel_shape[0]}, ${kernel_shape[1]}, ${strides[0]}, ${strides[1]},
+    ${pads[0]}, ${pads[1]}, ${pads[2]}, ${pads[3]});
+""")

Deeploy/Targets/Generic/Templates/FloatGlobalAveragePoolTemplate.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: 2023 ETH Zurich and University of Bologna
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from Deeploy.DeeployTypes import NetworkContext, NodeTemplate, OperatorRepresentation
+
+
+class _GlobalAveragePoolTemplate(NodeTemplate):
+
+    def alignToContext(self, ctxt: NetworkContext,
+                       operatorRepresentation: OperatorRepresentation) -> tuple[NetworkContext, dict, list[str]]:
+
+        data_in = ctxt.lookup(operatorRepresentation['data_in'])
+        operatorRepresentation['type_width'] = data_in._type.referencedType.typeWidth
+        return ctxt, operatorRepresentation, []
+
+
+referenceTemplate = _GlobalAveragePoolTemplate("""
+// Global Average Pool 1D (Name: ${nodeName}, Op: ${nodeOp})
+GlobalAveragePool_fp${type_width}_fp${type_width}(
+    ${data_in}, ${data_out}, ${batch_size}, ${num_channels}, ${spatial_size});
+""")

Deeploy/Targets/Generic/Templates/FloatGlobalMaxPoolTemplate.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: 2023 ETH Zurich and University of Bologna
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from Deeploy.DeeployTypes import NetworkContext, NodeTemplate, OperatorRepresentation
+
+
+class _GlobalMaxPoolTemplate(NodeTemplate):
+
+    def alignToContext(self, ctxt: NetworkContext,
+                       operatorRepresentation: OperatorRepresentation) -> tuple[NetworkContext, dict, list[str]]:
+
+        data_in = ctxt.lookup(operatorRepresentation['data_in'])
+        operatorRepresentation['type_width'] = data_in._type.referencedType.typeWidth
+        return ctxt, operatorRepresentation, []
+
+
+referenceTemplate = _GlobalMaxPoolTemplate("""
+// Global Max Pool 1D (Name: ${nodeName}, Op: ${nodeOp})
+GlobalMaxPool_fp${type_width}_fp${type_width}(
+    ${data_in}, ${data_out}, ${batch_size}, ${num_channels}, ${spatial_size});
+""")