NXP backend: Add support for softmax with the new Neutron flow.

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

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import numpy as np
+import torch
 
 from executorch.backends.nxp.backend.custom_delegation_options import (
     CustomDelegationOptions,
@@ -58,40 +59,31 @@ def _is_supported_on_target(
         parameters_mapping: dict[str, Parameter],
         custom_delegation_options: CustomDelegationOptions,
     ) -> bool:
-        """Check if the softmax operation can be executed on Neutron hardware.
-
-        Hardware constraints:
-        1. Input rank must be >= 2 (Neutron does not support 1D)
-        2. Channels must be a multiple of num_macs
-        3. Channels < 4096 / num_pipes * 4
-        4. Total spatial size (N*H*W) <= 4096
-        5. (channels * spatial_size) / num_macs <= 65536
+        """Hardware constraints:
+        1. Input and Output must be INT8/UINT8
+        2. Channels <= 2040
+        3. Total spatial size (N*H*W) <= 4096
+        4. Total size (channels * spatial_size) <= 524288
         """
-        input_shape = node.meta["val"].shape
-
-        # Constraint 1: Neutron does not support 1D SoftMax
-        if len(input_shape) == 1:
+        # Constraint 1: Input and Output must be INT8/UINT8.
+        supported_types = [torch.int8, torch.uint8]
+        if not NodeConverter.uses_quantization_type_for_io(
+            node, supported_types, [0], [0]
+        ):
             return False
 
-        num_macs = neutron_target_spec.get_num_macs()
-        num_pipes = neutron_target_spec.get_num_pipes()
+        # Constraint 2: Channel size limit
         channels = SoftmaxConverter._get_channels(node)
-        total_spatial_size = SoftmaxConverter._get_total_spatial_size(node)
-
-        # Constraint 2: Channels must be a multiple of num_macs
-        if channels % num_macs != 0:
+        if channels > 2040:
             return False
 
-        # Constraint 3: Channel size limit
-        if channels >= 4096 / num_pipes * 4:
-            return False
-
-        # Constraint 4: Spatial size limit
+        # Constraint 3: Spatial size limit
+        total_spatial_size = SoftmaxConverter._get_total_spatial_size(node)
         if total_spatial_size > 4096:
             return False
 
-        # Constraint 5: Total processing size limit
-        if channels * total_spatial_size / num_macs > 65536:
+        # Constraint 4: Total processing size limit
+        if channels * total_spatial_size > 524288:
             return False
 
         return True

backends/nxp/tests/generic_tests/test_cifarnet.py

@@ -11,10 +11,7 @@
 from executorch.backends.nxp.tests.config_importer import test_config
 from executorch.backends.nxp.tests.dataset_creator import CopyDatasetCreator
 from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec
-from executorch.backends.nxp.tests.graph_verifier import (
-    BaseGraphVerifier,
-    NonDelegatedNode,
-)
+from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier
 from executorch.backends.nxp.tests.model_output_comparator import (
     NumericalStatsOutputComparator,
 )
@@ -56,15 +53,14 @@ def test_cifarnet(mocker, request, cifar_test_files, channels_last):
         model.to(memory_format=torch.channels_last)
         input_spec.dim_order = torch.channels_last
 
-    non_dlg_nodes = [NonDelegatedNode("aten__softmax_default", 1)]
-
     comparator = NumericalStatsOutputComparator(
-        max_mse_error=1.0e-3, is_classification_task=True
+        max_mse_error=1.53e-5,
+        is_classification_task=True,
     )
     lower_run_compare(
         model,
         [input_spec],
-        BaseGraphVerifier(1, non_dlg_nodes),
+        BaseGraphVerifier(1, []),
         request,
         dataset_creator=CopyDatasetCreator(cifar_test_files),
         output_comparator=comparator,
@@ -84,18 +80,15 @@ def test_cifarnet_qat(mocker, request, cifar_test_files):
     model = CifarNet().get_eager_model().eval()
 
     input_shape = (1, 3, 32, 32)
-    non_dlg_nodes = [NonDelegatedNode("aten__softmax_default", 1)]
 
-    # The higher MSE threshold is due to using weaker "MovingAbs" observers instead of "MinMax" observers.
-    # The "MovingAbs" observers capture only limited number of past calibration samples compared to "MinMax",
-    # which uses statistics from the whole calibration set.
     comparator = NumericalStatsOutputComparator(
-        max_mse_error=8e-2, is_classification_task=True
+        max_mse_error=1.53e-5,
+        is_classification_task=True,
     )
     lower_run_compare(
         model,
         input_shape,
-        BaseGraphVerifier(1, non_dlg_nodes),
+        BaseGraphVerifier(1, []),
         request,
         dataset_creator=CopyDatasetCreator(cifar_test_files),
         output_comparator=comparator,

backends/nxp/tests/generic_tests/test_integration.py

@@ -30,8 +30,8 @@ def test_conv_fc_softmax__to_executorch_program(use_qat):
 
     delegation_info = get_delegation_info(program.graph_module)
     assert delegation_info.num_delegated_subgraphs == 1
-    assert delegation_info.num_non_delegated_nodes == 11
-    assert delegation_info.num_delegated_nodes == 15
+    assert delegation_info.num_non_delegated_nodes == 5
+    assert delegation_info.num_delegated_nodes == 18
 
     # Make sure Convolution and AddMM are delegated.
     assert not graph_contains_any_of_ops(program.graph, [Convolution, AddMM])
@@ -46,8 +46,8 @@ def test_cifarnet(use_qat):
 
     delegation_info = get_delegation_info(exec_prog.exported_program().graph_module)
     assert delegation_info.num_delegated_subgraphs == 1
-    assert delegation_info.num_non_delegated_nodes == 11
-    assert delegation_info.num_delegated_nodes == 47
+    assert delegation_info.num_non_delegated_nodes == 5
+    assert delegation_info.num_delegated_nodes == 50
 
     nodes = list(exec_prog.exported_program().graph.nodes)
     # `nodes[2].target` is an OpOverload (not and EdgeOpOverload that we usually test against), so just check the name.

backends/nxp/tests/generic_tests/test_neutron_backend_executor.py

@@ -89,7 +89,7 @@ def test_conv_fc__lowered_program_and_tflite_output_match(mocker):
     # No Transpose ops in produced TFLite model
     tflite_subgraph = Model.GetRootAs(tflite_flatbuffers_model).Subgraphs(0)
 
-    assert tflite_subgraph.OperatorsLength() == 3
+    assert tflite_subgraph.OperatorsLength() == 4
     assert (
         tflite_subgraph.Operators(0).BuiltinOptionsType()
         == BuiltinOptions.Conv2DOptions
@@ -102,6 +102,10 @@ def test_conv_fc__lowered_program_and_tflite_output_match(mocker):
         tflite_subgraph.Operators(2).BuiltinOptionsType()
         == BuiltinOptions.FullyConnectedOptions
     )
+    assert (
+        tflite_subgraph.Operators(3).BuiltinOptionsType()
+        == BuiltinOptions.SoftmaxOptions
+    )
 
     # Verify outputs of program and TFLite model
     input_data = (

backends/nxp/tests/generic_tests/test_profiling.py

@@ -186,10 +186,6 @@ def forward(self, x):
 
 
 class TestProfiling:
-    @pytest.mark.xfail(
-        reason="Profiling support for cmodel and SoftMax fix will be available in Neutron SW 3.2.",
-        strict=True,
-    )
     def test__softmax(self, caplog, request):
         caplog.set_level(logging.INFO)
         model = SoftmaxModule(-1)
@@ -240,7 +236,6 @@ def test__simple_parallel_pool(self, caplog, request):
             10: (),  # Neutron Dump
         }
 
-    @pytest.mark.xfail(reason="SoftMax support PR is not merged so far.", strict=True)
     def test__cifar(self, caplog, request):
         caplog.set_level(logging.INFO)
         input_shape = (1, 3, 32, 32)