NXP backend: Enable prelu with new Neutron flow

Author: roman-janik-nxp

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

@@ -3,7 +3,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from executorch.backends.nxp.backend.data_format import NXP_NODE_FORMAT
+import torch
+
 from executorch.backends.nxp.backend.ir.converter.node_converter import (
     CustomDelegationOptions,
     NodeConverter,
@@ -24,38 +25,17 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        node_shape = node.meta["val"].shape
-        rank = len(node_shape)
-
-        # According to Neutron spec., PReLU can be done only on 4D tensors
-        if rank != 4:
-            return False
-
-        ch_idx, h_idx, w_idx = PReLUConverter._get_channel_spatial_indices(node)
-        # According to Neutron spec., size of channels must be divisible by num_macs.
-        num_macs = neutron_target_spec.get_num_macs()
-        if node_shape[ch_idx] % num_macs != 0:
+        if not NodeConverter.inputs_satisfy_broadcast_condition(node):
             return False
 
-        # According to Neutron spec., height * width cannot be greater than a given constant.
-        if node_shape[w_idx] * node_shape[h_idx] > 4096:
+        supported_types = [torch.int8, torch.uint8]
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, supported_types, [0, 1], [0]
+        ):
             return False
 
         return True
 
-    @staticmethod
-    def _get_channel_spatial_indices(node: Node):
-        if node.meta[NXP_NODE_FORMAT].is_channels_first():
-            ch_idx = 1
-            h_idx = 2
-            w_idx = 3
-        else:
-            ch_idx = 3
-            h_idx = 1
-            w_idx = 2
-
-        return ch_idx, h_idx, w_idx
-
     @staticmethod
     def _is_supported_in_IR(
         node: Node,

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

@@ -106,7 +106,7 @@ def test__basic_nsys_inference_qat(self, mocker, request):
             ),
             pytest.param(
                 [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
-                id="2 inputs 2D, num_elems of input == num_elems of output",
+                id="2 inputs 1D + 2D, num_elems of input == num_elems of output",
             ),
         ],
     )

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

@@ -92,7 +92,7 @@ def test__basic_nsys_inference_qat(self, mocker, request, x_input_shape):
             ),
             pytest.param(
                 [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
-                id="2 inputs 2D, num_elems of input == num_elems of output",
+                id="2 inputs 1D + 2D, num_elems of input == num_elems of output",
             ),
         ],
     )

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

@@ -9,18 +9,32 @@
 from executorch.backends.nxp.backend.edge_program_converter import (
     EdgeProgramToIRConverter,
 )
-from executorch.backends.nxp.tests.executors import (
-    convert_run_compare,
-    graph_contains_any_of_ops,
+from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops
+from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier
+from executorch.backends.nxp.tests.model_output_comparator import (
+    AllCloseOutputComparator,
 )
 from executorch.backends.nxp.tests.models import (
+    ConvPReLUModule,
     LinearPReLUModule,
     TwoPartitionPReLUModel,
 )
+
+from executorch.backends.nxp.tests.nsys_testing import lower_run_compare
+from executorch.backends.nxp.tests.ops_aliases import (
+    AddMm,
+    Convolution,
+    ExecutorchDelegateCall,
+    GtScalar,
+    MulTensor,
+    PermuteCopy,
+    Prelu,
+    ViewCopy,
+    WhereSelf,
+)
 from torch.export import ExportedProgram
 from executorch.backends.nxp.tests.use_qat import *  # noqa F403
 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program
-from executorch.exir.dialects._ops import ops as exir_ops
 
 
 @pytest.fixture(autouse=True)
@@ -29,123 +43,148 @@ def reseed_model_per_test_run():
     np.random.seed(23)
 
 
-# noinspection PyProtectedMember
-ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate
-
-
-@pytest.mark.parametrize(
-    "input_shape",
-    [
-        pytest.param((1, 8, 24, 32), id="4D."),
-    ],
-)
-def test_prelu_with_linear_quant_conversion(mocker, input_shape):
-    converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
-
-    # Run conversion
-    channels = input_shape[-1]
-    edge_program = to_quantized_edge_program(
-        LinearPReLUModule(in_features=channels, out_features=channels),
+class TestPreluConverter:
+    # noinspection PyMethodMayBeStatic
+    def assert_delegated(
+        self,
+        model,
         input_shape,
-    ).exported_program()
-
-    # Capture generated entities
-    neutron_ir_model, *_ = converter_spy.spy_return
-    exported_program: ExportedProgram = converter_spy.call_args.args[1]
-
-    # Check `prelu` was not decomposed into simpler edge operators
-    assert not graph_contains_any_of_ops(
-        exported_program.graph,
+        mocker,
+        request,
+        expected_delegated_ops=None,
+        use_qat=False,
+    ):
+        rank = len(input_shape)
+        graph_verifier = DetailedGraphVerifier(
+            mocker,
+            expected_delegated_ops=expected_delegated_ops
+            or {
+                Prelu: 1,
+                AddMm: 1,
+                PermuteCopy: 1,
+                ViewCopy: 0 if rank == 2 else 2,
+            },
+            expected_non_delegated_ops={},
+        )
+        comparator = AllCloseOutputComparator(atol=1)
+
+        lower_run_compare(
+            model,
+            input_shape,
+            graph_verifier,
+            request,
+            output_comparator=comparator,
+            remove_quant_io_ops=True,
+            use_qat=use_qat,
+        )
+
+    @pytest.mark.parametrize(
+        "input_shape",
         [
-            exir_ops.edge.aten.gt.Scalar,
-            exir_ops.edge.aten.mul.Tensor,
-            exir_ops.edge.aten.where.self,
+            pytest.param((1,), id="1D."),
+            pytest.param(
+                (36, 487),
+                id="2D incorrect results.",
+                marks=pytest.mark.xfail(
+                    reason="AIR-14737: incorrect results", strict=True
+                ),
+            ),
+            pytest.param(
+                (87, 842),
+                id="2D incorrect results alt.",
+                marks=pytest.mark.xfail(
+                    reason="AIR-14737: incorrect results", strict=True
+                ),
+            ),
+            pytest.param((7, 83), id="2D."),
+            pytest.param(
+                (1, 43, 183),
+                id="3D incorrect results alt.",
+                marks=pytest.mark.xfail(
+                    reason="AIR-14737: incorrect results", strict=True
+                ),
+            ),
+            pytest.param((1, 43, 93), id="3D."),
+            pytest.param((1, 4, 7, 8), id="4D."),
+            pytest.param((1, 4, 3, 4, 14), id="5D."),
         ],
     )
-
-    assert graph_contains_any_of_ops(
-        exported_program.graph,
-        [exir_ops.edge.aten.prelu.default],
-    )
-
-    # Check `prelu` was delegated
-    assert not graph_contains_any_of_ops(
-        edge_program.graph,
-        [exir_ops.edge.aten.prelu.default],
-    )
-
-    input_data = (
-        (2 * np.random.random(input_shape).astype(np.float32) - 1) * 50
-    ).astype(np.int8)
-
-    convert_run_compare(exported_program, input_data, tfl_model=neutron_ir_model)
-
-
-@pytest.mark.parametrize(
-    "input_shape",
-    [
-        pytest.param((1, 8, 24, 32), id="4D."),
-    ],
-)
-def test_prelu_2_partitions(mocker, input_shape):
-    # TODO (Martin) Add a channels last dim order variant of this test to verify correct partitioning.
-    # Run conversion
-    edge_program = to_quantized_edge_program(
-        TwoPartitionPReLUModel(), [input_shape, input_shape]
-    ).exported_program()
-
-    # Check `prelu` was delegated
-    assert not graph_contains_any_of_ops(
-        edge_program.graph,
-        [exir_ops.edge.aten.prelu.default],
-    )
-
-    # Check there are two partitions
-    edge_nodes = list(edge_program.graph.nodes)
-    assert sum(n.target == ExecutorchDelegateCall for n in edge_nodes) == 2
-
-
-@pytest.mark.parametrize(
-    "input_shape",
-    [
-        pytest.param((1,), id="1D not supported."),
-        pytest.param((1, 8), id="2D not supported."),
-        pytest.param((1, 8, 16), id="3D not supported."),
-        pytest.param((1, 8, 16, 32, 64), id="5D not supported."),
-        pytest.param((1, 8, 16, 31), id="channels must be divisible by NUM_MACS"),
-        pytest.param((1, 8, 1024, 8), id="width*height is too big (limit 4096)"),
-    ],
-)
-def test_prelu__no_delegation__unsupported_conversion(mocker, input_shape):
-    # Run conversion
-    channels = input_shape[-1]
-    edge_program = to_quantized_edge_program(
-        LinearPReLUModule(in_features=channels, out_features=channels),
-        input_shape,
-    ).exported_program()
-
-    # Check `prelu` was not delegated (only `linear` was)
-    edge_nodes = list(edge_program.graph.nodes)
-    assert sum(n.target == ExecutorchDelegateCall for n in edge_nodes) == 1
-
-    # Check `prelu` was decomposed into simpler edge operators
-    assert graph_contains_any_of_ops(
-        edge_program.graph,
-        [
-            exir_ops.edge.aten.gt.Scalar,
-        ],
-    )
-
-    assert graph_contains_any_of_ops(
-        edge_program.graph,
-        [
-            exir_ops.edge.aten.mul.Tensor,
-        ],
-    )
-
-    assert graph_contains_any_of_ops(
-        edge_program.graph,
+    def test__basic_nsys_inference(self, mocker, request, input_shape):
+        channels = input_shape[-1]
+        model = LinearPReLUModule(in_features=channels, out_features=channels)
+
+        converter_spy = mocker.spy(EdgeProgramToIRConverter, "convert_program")
+
+        self.assert_delegated(model, input_shape, mocker, request)
+
+        # Capture generated entities
+        exported_program: ExportedProgram = converter_spy.call_args.args[1]
+
+        # Check `prelu` was not decomposed into simpler edge operators
+        assert not graph_contains_any_of_ops(
+            exported_program.graph,
+            [
+                GtScalar,
+                MulTensor,
+                WhereSelf,
+            ],
+        )
+
+    def test__basic_nsys_inference_qat(self, mocker, request):
+        input_shape = (2, 4, 6, 7)
+        channels = input_shape[-1]
+        model = LinearPReLUModule(in_features=channels, out_features=channels)
+
+        self.assert_delegated(model, input_shape, mocker, request, use_qat=True)
+
+    def test__num_parameters_param(self, mocker, request):
+        input_shape = (1, 43, 93)
+        channels = input_shape[-1]
+        num_parameters = input_shape[1]
+        model = LinearPReLUModule(
+            in_features=channels, out_features=channels, num_parameters=num_parameters
+        )
+
+        self.assert_delegated(model, input_shape, mocker, request)
+
+    def test_prelu_2_partitions(self):
+        input_shape = (1, 8, 24, 32)
+        # Run conversion
+        edge_program = to_quantized_edge_program(
+            TwoPartitionPReLUModel(), [input_shape, input_shape]
+        ).exported_program()
+
+        # Check `prelu` was delegated
+        assert not graph_contains_any_of_ops(
+            edge_program.graph,
+            [Prelu],
+        )
+
+        # Check there are two partitions
+        edge_nodes = list(edge_program.graph.nodes)
+        assert sum(n.target == ExecutorchDelegateCall for n in edge_nodes) == 2
+
+    @pytest.mark.parametrize(
+        "input_shape",
         [
-            exir_ops.edge.aten.where.self,
+            pytest.param((1, 8, 42, 24), id="4D."),
+            pytest.param(
+                (1, 8, 42, 21),
+                id="4D incorrect results.",
+                marks=pytest.mark.xfail(
+                    reason="AIR-14737: incorrect results", strict=True
+                ),
+            ),
         ],
     )
+    def test__w_conv(self, mocker, request, input_shape):
+        channels = input_shape[1]
+        model = ConvPReLUModule(in_channels=channels)
+        expected_delegated_ops = {
+            Prelu: 1,
+            Convolution: 1,
+        }
+
+        self.assert_delegated(
+            model, input_shape, mocker, request, expected_delegated_ops
+        )

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

@@ -106,7 +106,7 @@ def test__basic_nsys_inference_qat(self, mocker, request):
             ),
             pytest.param(
                 [ModelInputSpec((10,)), ModelInputSpec((1, 1))],
-                id="2 inputs 2D, num_elems of input == num_elems of output",
+                id="2 inputs 1D + 2D, num_elems of input == num_elems of output",
             ),
         ],
     )