NXP backend: Use zero_point to pad quantized average_pool.

Author: MartinPavella

backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py

@@ -3,11 +3,16 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import numpy as np
+
 from executorch.backends.nxp.backend.ir.converter.conversion import (
     aten_translator,
     common,
 )
 from executorch.backends.nxp.backend.ir.converter.conversion.common import OpsList
+from executorch.backends.nxp.backend.ir.converter.conversion.translator import (
+    tf_lite_type_to_numpy,
+)
 from executorch.backends.nxp.backend.ir.converter.node_converter import NodeConverter
 from executorch.backends.nxp.backend.ir.tflite_generator import tflite_model
 from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options import (
@@ -57,9 +62,20 @@ def _convert_2d_avg_pool(
         )
 
         if explicit_padding is not None:
-            # Need to prepend a 'Pad' operator, which adds 0s. But these will be included in the computation!
+            # Need to prepend a 'Pad' operator, which adds 0s (or `zero_point` for the quantized case). But these will
+            #  be included in the computation!
+            input_quantization = t_op.tmp_inputs[0].quantization
+            pad_value = (
+                None
+                if input_quantization is None
+                else np.array(input_quantization.zero_point[0]).astype(
+                    tf_lite_type_to_numpy(t_op.tmp_inputs[0].type)
+                )
+            )
             ops.add_pre(
-                self.builder.create_pad_operator_before(t_op, 0, explicit_padding)
+                self.builder.create_pad_operator_before(
+                    t_op, 0, explicit_padding, pad_value
+                )
             )
 
         return ops.flatten()

backends/nxp/backend/ir/converter/node_converters/ops_converters/convolution_converter.py

@@ -5,8 +5,6 @@
 
 import numpy as np
 import torch
-from torch.fx import Node
-from torch.nn import Parameter
 
 from executorch.backends.nxp.backend.edge_helper import (
     input_tensor,
@@ -42,6 +40,8 @@
     conv_2d_options,
     depthwise_conv_2d_options,
 )
+from torch.fx import Node
+from torch.nn import Parameter
 
 
 class ConvolutionConverter(NodeConverter):

backends/nxp/tests/executorch_pipeline.py

@@ -48,7 +48,7 @@ def get_random_float_data(input_shapes: tuple[int] | list[tuple[int]]):
 
 def to_quantized_edge_program(
     model: torch.nn.Module,
-    input_shapes: tuple[int] | list[tuple[int]],
+    input_shapes: tuple[int, ...] | list[tuple[int, ...]],
     operators_not_to_delegate: list[str] = None,
     target="imxrt700",
     neutron_converter_flavor="SDK_25_03",

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

@@ -10,6 +10,12 @@
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
 )
+from executorch.backends.nxp.backend.ir.converter.builder.model_builder import (
+    ModelBuilder,
+)
+from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import (
+    BuiltinOperator,
+)
 from executorch.backends.nxp.tests.executorch_pipeline import (
     to_edge_program,
     to_quantized_edge_program,
@@ -156,3 +162,49 @@ def test_avg_pool_2d_quant_conversion(mocker, input_shape, padding, count_includ
         tflite_output_preprocess=ToNCHWPreprocess(),
         input_data=input_data,
     )
+
+
+def test_avg_pool_2d_quant_conversion__padded(mocker):
+    input_shape = (1, 8, 8, 8)
+    model = AvgPool2dModule(True, 1)
+
+    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+    ops_spy = mocker.spy(ModelBuilder, "finish")
+
+    # Run conversion
+    _ = to_quantized_edge_program(model, input_shape)
+
+    # Capture the converter operators.
+    ops = ops_spy.spy_return.sub_graphs[0].operators.vector
+
+    # Capture generated model
+    tflite_flatbuffers_model, io_formats = converter_spy.spy_return
+
+    # Capture converted program
+    exported_program: ExportedProgram = converter_spy.call_args.args[1]
+
+    input_data = (np.random.random(input_shape).astype(np.float32) * 50).astype(np.int8)
+
+    convert_run_compare(
+        exported_program,
+        tflite_input_preprocess=ToNHWCPreprocess(),
+        tfl_model=tflite_flatbuffers_model,
+        tflite_output_preprocess=ToNCHWPreprocess(),
+        input_data=input_data,
+    )
+
+    assert len(ops) == 2
+    assert ops[0].builtin_options.operator_type == BuiltinOperator.PADV2
+    assert ops[1].builtin_options.operator_type == BuiltinOperator.AVERAGE_POOL_2D
+
+    # Make sure the padding used the `zero-point`.
+    pad_value = ops[0].tmp_inputs[2].tmp_buffer.data.item()
+    assert (
+        pad_value == ops[0].tmp_inputs[0].quantization.zero_point[0]
+    )  # `Pad` input zp.
+    assert (
+        pad_value == ops[0].tmp_outputs[0].quantization.zero_point[0]
+    )  # `Pad` output zp.
+    assert (
+        pad_value == ops[1].tmp_inputs[0].quantization.zero_point[0]
+    )  # `AvgPool` input zp.