NXP backend: Add support for `softmax` with the new Neutron flow. by irtrukhina · Pull Request #19739 · pytorch/executorch

irtrukhina · 2026-05-22T08:08:26Z

Summary

Add support for the softmax to the Neutron backend using the new Neutron MLIR flow. Add new tests and update old tests for the softmax as well.

Test plan

New tests are included in the PR.

cc @robert-kalmar @JakeStevens @digantdesai @rascani

pytorch-bot · 2026-05-22T08:08:31Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19739

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 9 Pending

As of commit 79b5ba5 with merge base d29706a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

linux-foundation-easycla · 2026-05-22T08:08:34Z

❌ - login: @irtrukhina / name: Irina Korchakova. The commit (0d13bba, 79b5ba5) is not authorized under a signed CLA. Please click here to be authorized. For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

MartinPavella · 2026-06-05T13:35:46Z

            1. Input and Output must be INT8/UINT8
            2. Channels <= 2040
-            3. Total spatial size (H*W) <= 4096
+            3. Total spatial size <= 4096


Nit: As far as I know total spatial size is not a universally known term. Before your change, it was supposed to represent H*W, whereas now it means N*H*W. Please keep the explicit definition here.

Kept the explicit definition (NHW)

MartinPavella · 2026-06-08T06:29:39Z

Please take a look at the canceled pull / unittest-nxp-neutron / linux-job (pull_request). Otherwise, LGTM

irtrukhina · 2026-07-31T12:31:30Z

@MartinPavella, could you please take a final look at the PR? All tests have passed, and the documentation has been fixed to meet our checks.

MartinPavella

Just a couple of observations to consider. No blockers.

MartinPavella · 2026-07-31T12:46:58Z

+    # Allow MSE up to the theoretical error introduced by 1-bit quantization (1/256).
    comparator = NumericalStatsOutputComparator(
-        max_mse_error=1.0e-3, is_classification_task=True
+        max_mse_error=0.00390625,


I don't see how the dequantized value of a single-bit error directly dictates the MSE threshold. So the comment doesn't make sense to me. (same for the following test)

Also Nit: Using 1/256 would be cleaner.

Sorry. This issue was caused by a combination of my brain overheating in the summer heat and a copy-paste error. The MSE should be the squared one-bit quantization error. I fixed the comparator settings and removed the comment.

MartinPavella · 2026-07-31T12:51:33Z

+            expected_delegated_ops={Softmax: 1},
+            expected_non_delegated_ops={},
+        )
+        output_comparator = NumericalStatsOutputComparator(


Why are you using the NumericalStatsOutputComparator?
For other ops we use the AllCloseOutputComparator with a quantized dataset and atol=1. So it is clear only single-bit errors are tollerated.

It is a very old PR and was created before we add remove_quant_io_ops option into lower_run_compare.
Replaced NumericalStatsOutputComparator with AllCloseOutputComparator(atol=1).

roman-janik-nxp

I recommend to fix the issues mentioned by Martin.

roman-janik-nxp · 2026-08-03T14:17:42Z

-    ],
-)
-def test_softmax_delegation__channel_first(input_shape, dim: int, mocker):
-    model = ConvSoftmaxModule(dim, input_shape[1])


Is there are reason why no longer test channel first variant?

It doesn't make sense to me. Neutron only supports channels-last format. Any errors in the channels-first version are more likely caused by incorrect dimension ordering than by the softmax operation.
These tests are specifically focused on validating softmax.
Do you agree?

Well spotted @roman-janik-nxp, we should definitely have a test with a softmax after a convolution/maxpool/... node to make sure our SoftmaxConverter handles that edge case correctly.
@irtrukhina please add such a test.

irtrukhina requested a review from robert-kalmar as a code owner May 22, 2026 08:08

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 22, 2026

robert-kalmar added module: nxp Issues related to NXP Neutron NPU delegation and code under backends/nxp/ release notes: nxp Changes to the NXP Neutron backend delegate labels May 22, 2026

robert-kalmar requested review from MartinPavella and roman-janik-nxp May 22, 2026 09:01

MartinPavella reviewed May 22, 2026

View reviewed changes

Comment thread backends/nxp/tests/generic_tests/test_integration.py Outdated

Comment thread backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py Outdated

Comment thread backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py Outdated

irtrukhina force-pushed the feature/EIEX-728_revisit_softmax_support branch from 154ffff to 56a1296 Compare June 1, 2026 16:09

MartinPavella reviewed Jun 3, 2026

View reviewed changes

Comment thread backends/nxp/backend/ir/converter/node_converters/ops_converters/softmax_converter.py Outdated

Comment thread backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py Outdated

irtrukhina force-pushed the feature/EIEX-728_revisit_softmax_support branch from c8324ed to c122421 Compare June 4, 2026 08:54

MartinPavella reviewed Jun 4, 2026

View reviewed changes

MartinPavella reviewed Jun 8, 2026

View reviewed changes

irtrukhina force-pushed the feature/EIEX-728_revisit_softmax_support branch from 502f045 to f504223 Compare June 9, 2026 08:09

robert-kalmar marked this pull request as draft June 23, 2026 16:50

irtrukhina force-pushed the feature/EIEX-728_revisit_softmax_support branch from f504223 to 8cc126d Compare July 14, 2026 14:23

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

0d13bba

irtrukhina force-pushed the feature/EIEX-728_revisit_softmax_support branch from 8cc126d to 0d13bba Compare July 31, 2026 10:26

MartinPavella approved these changes Jul 31, 2026

View reviewed changes

roman-janik-nxp reviewed Aug 3, 2026

View reviewed changes

Update comparators in tests

79b5ba5

Uh oh!

Conversation

irtrukhina commented May 22, 2026 • edited by pytorch-bot Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

pytorch-bot Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19739

⏳ No Failures, 9 Pending

Uh oh!

linux-foundation-easycla Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

MartinPavella commented Jun 8, 2026

Uh oh!

irtrukhina commented Jul 31, 2026

Uh oh!

MartinPavella left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

roman-janik-nxp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

irtrukhina commented May 22, 2026 •

edited by pytorch-bot Bot

Loading

pytorch-bot Bot commented May 22, 2026 •

edited

Loading

linux-foundation-easycla Bot commented May 22, 2026 •

edited

Loading