NXP backend: Add support for aten.avg_pool1d.default.

Author: MartinPavella

backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py

@@ -13,15 +13,18 @@
 
 # Operator aliases for better readability.
 AddMM = exir_ops.edge.aten.addmm.default
-ViewCopy = exir_ops.edge.aten.view_copy.default
-MM = exir_ops.edge.aten.mm.default
+AvgPool2D = exir_ops.edge.aten.avg_pool2d.default
 Conv = exir_ops.edge.aten.convolution.default
+Clone = exir_ops.edge.aten.clone.default
+CloneDimOrder = exir_ops.edge.dim_order_ops._clone_dim_order.default
 HardTanh = exir_ops.edge.aten.hardtanh.default
+MM = exir_ops.edge.aten.mm.default
 Relu = exir_ops.edge.aten.relu.default
 Sigmoid = exir_ops.edge.aten.sigmoid.default
+SqueezeCopy = exir_ops.edge.aten.squeeze_copy.dims
 Tanh = exir_ops.edge.aten.tanh.default
-Clone = exir_ops.edge.aten.clone.default
-CloneDimOrder = exir_ops.edge.dim_order_ops._clone_dim_order.default
+UnsqueezeCopy = exir_ops.edge.aten.unsqueeze_copy.default
+ViewCopy = exir_ops.edge.aten.view_copy.default
 
 
 def insert_qdq_pair_after_node(
@@ -106,7 +109,11 @@ class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass):
         ],
         ViewCopy: [Clone, CloneDimOrder],
         Conv: [
-            ViewCopy,  # For 1D conv.
+            ViewCopy,  # For 1D conv
+        ],
+        AvgPool2D: [
+            ViewCopy,  # For 1D AvgPool
+            UnsqueezeCopy,
         ],
     }
 
@@ -206,6 +213,10 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass):
             ViewCopy,  # For 1D conv.
         ],
         ViewCopy: [Clone, CloneDimOrder],
+        AvgPool2D: [
+            ViewCopy,
+            SqueezeCopy,
+        ],
     }
 
     def run(self, graph_module: torch.fx.GraphModule) -> PassResult:

backends/nxp/neutron_partitioner.py

@@ -75,13 +75,15 @@ class QDQCluster:
 
     AUXILIARY_OPS = [
         operator.getitem,
-        exir_ops.edge.aten.view_copy.default,
-        exir_ops.edge.aten.permute_copy.default,
+        exir_ops.edge.aten.clone.default,
         exir_ops.edge.aten.hardtanh.default,
+        exir_ops.edge.aten.permute_copy.default,
         exir_ops.edge.aten.relu.default,
         exir_ops.edge.aten.sigmoid.default,
+        exir_ops.edge.aten.squeeze_copy.dims,
         exir_ops.edge.aten.tanh.default,
-        exir_ops.edge.aten.clone.default,
+        exir_ops.edge.aten.unsqueeze_copy.default,
+        exir_ops.edge.aten.view_copy.default,
         exir_ops.edge.dim_order_ops._clone_dim_order.default,
     ]
 

backends/nxp/quantizer/neutron_quantizer.py

@@ -17,7 +17,8 @@
     AdaptiveAvgPoolPattern,
     AddmmPattern,
     AddTensorPattern,
-    AvgPoolPattern,
+    AvgPool1DPattern,
+    AvgPool2DPattern,
     BatchNormPattern,
     CatPattern,
     Conv1dPattern,
@@ -252,7 +253,8 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False)
                 OpQuantizer(AdaptiveAvgPoolPattern(is_qat=is_qat), static_qconfig),
                 OpQuantizer(AddTensorPattern(is_qat=is_qat), static_qconfig),
                 OpQuantizer(AddmmPattern(self, is_qat=is_qat), static_fc_qconfig),
-                OpQuantizer(AvgPoolPattern(is_qat=is_qat), static_qconfig),
+                OpQuantizer(AvgPool1DPattern(is_qat=is_qat), static_qconfig),
+                OpQuantizer(AvgPool2DPattern(is_qat=is_qat), static_qconfig),
                 OpQuantizer(BatchNormPattern(is_qat=is_qat), static_qconfig),
                 OpQuantizer(CatPattern(is_qat=is_qat), static_qconfig),
                 OpQuantizer(Conv1dPattern(is_qat=is_qat), static_qconfig),

backends/nxp/quantizer/patterns.py

@@ -324,7 +324,16 @@ def get_anchors(
         )
 
 
-class AvgPoolPattern(SharedSpecPattern):
+class AvgPool1DPattern(SharedSpecPattern):
+    """
+    Quantizer for AvgPool1D operator.
+    """
+
+    def partition_types(self):
+        return [torch.ops.aten.avg_pool1d.default]
+
+
+class AvgPool2DPattern(SharedSpecPattern):
     """
     Quantizer for AvgPool2D operator.
     """

backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py

@@ -9,7 +9,6 @@
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
 )
-
 from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig
 from executorch.backends.nxp.backend.ir.converter.builder.model_builder import (
     ModelBuilder,
@@ -23,12 +22,25 @@
 )
 from executorch.backends.nxp.tests.executors import (
     convert_run_compare,
+    graph_contains_any_of_ops,
+    ToChannelFirstPreprocess,
+    ToChannelLastPreprocess,
     ToNCHWPreprocess,
     ToNHWCPreprocess,
 )
 from executorch.backends.nxp.tests.models import AvgPool2dConvModule, AvgPool2dModule
 from torch.export import ExportedProgram
 from executorch.backends.nxp.tests.use_qat import *  # noqa F403
+from executorch.exir.dialects._ops import ops as exir_ops
+
+# noinspection PyProtectedMember
+AvgPool2D = exir_ops.edge.aten.avg_pool2d.default
+ExecutorchDelegateCall = torch._higher_order_ops.executorch_call_delegate
+Squeeze = exir_ops.edge.aten.squeeze.default
+SqueezeDim = exir_ops.edge.aten.squeeze.dim
+SqueezeDims = exir_ops.edge.aten.squeeze.dims
+Unsqueeze = exir_ops.edge.aten.unsqueeze.default
+ViewCopy = exir_ops.edge.aten.view_copy.default
 
 
 @pytest.fixture(autouse=True)
@@ -219,3 +231,59 @@ def test_avg_pool_2d_quant_conversion__padded(mocker, use_qat):
     assert (
         pad_value == ops[1].tmp_inputs[0].quantization.zero_point[0]
     )  # `AvgPool` input zp.
+
+
+class AvgPool1DModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self.avg_pool = torch.nn.AvgPool1d(
+            kernel_size=3,
+        )
+
+    def forward(self, x):
+        return self.avg_pool(x)
+
+
+def test_from_avg_pool_1d(mocker):
+    model = AvgPool1DModule()
+    input_shape = (
+        1,
+        3,
+        12,
+    )  # Don't use multiples of `num_macs` so the `view_copy` nodes will NOT be deleagted.
+    extended_shape = (1, 3, 1, 12)
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+    delegated_ep = to_quantized_edge_program(
+        model, input_shape, use_neutron_for_format_conversion=False
+    ).exported_program()
+
+    # Make sure the `avg_pool` was delegated.
+    assert graph_contains_any_of_ops(delegated_ep.graph, [ExecutorchDelegateCall])
+    assert not graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D])
+
+    # Make sure both `view_copy` nodes were added, and there is no `squeeze` or `unsqueeze`.
+    assert len([n for n in delegated_ep.graph.nodes if n.target == ViewCopy]) == 2
+    assert not graph_contains_any_of_ops(
+        delegated_ep.graph, [Unsqueeze, Squeeze, SqueezeDim, SqueezeDims]
+    )
+
+    # Verify correct behavior of the converted NeutronIR model.
+    intermediate_ep = converter_spy.call_args.args[1]
+    neutron_ir_model, _ = converter_spy.spy_return
+
+    input_data = (
+        np.random.random(extended_shape).astype(np.float32) * 256.0 - 128.0
+    ).astype(np.int8)
+
+    # Make sure the tested program contains the `avg_pool`.
+    assert graph_contains_any_of_ops(intermediate_ep.graph, [AvgPool2D])
+
+    convert_run_compare(
+        intermediate_ep,
+        tfl_model=neutron_ir_model,
+        input_data=input_data,
+        tflite_input_preprocess=ToChannelLastPreprocess(),
+        tflite_output_preprocess=ToChannelFirstPreprocess(),
+    )

docs/source/backends/nxp/op-support.csv

@@ -3,6 +3,7 @@ aten.abs.default,int8,static int8,
 aten._adaptive_avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False"
 aten.addmm.default,int8,static int8,2D tensor only
 aten.add.Tensor,int8,static int8,"alpha = 1, input tensor of rame rank"
+aten.avg_pool1d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False"
 aten.avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False"
 aten.cat.default,int8,static int8,"input_channels % 8 = 0, output_channels %8 = 0"
 aten.clone.default,int8,static int8,