NXP backend: Linear + BatchNorm QAT fusing

[StrycekSimon]

Summary

Adds two passes for inserting/removing simulated BatchNorm fusion for QAT training, similarly to how _fuse_conv_bn_qat adds simulated Conv+BatchNorm fusion in prepare_qat_pt2e function from TorchAO.

Test plan

Added integration tests that covers newly added implementation,

cc @robert-kalmar

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 11 New Failures, 1 Pending

As of commit d70a29e with merge base 7a9fb3f ():

NEW FAILURES - The following jobs have failed:

• pull / android / run-emulator (gh)
  The process '/opt/android/sdk/platform-tools/adb' failed with exit code 224
• pull / test-coreml-bc-macos (macos-m1-stable) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• pull / test-coreml-bc-macos (macos-m2-stable) / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• pull / test-moshi-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 356f9e228c7c2eed5c844b4f081538b521d601a2cd9c956bead1ec55f0295e5c /exec failed with exit code 1
• pull / test-parakeet-xnnpack-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 27447cc15993d4e1ba88b48f2e3c7676362f6c97ceeecc0b750a24cc24d4b567 /exec failed with exit code 1
• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 105f918dab59f09bae28b9e3f29fb3a9b15b00075bec1bf4daf3d1ab39783022 /exec failed with exit code 127
• pull / test-samsung-quantmodels-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 43694dda6a9247090aba44bb0ff0a528517e822a5b60e9ae26af59e995c72b76 /exec failed with exit code 127
• pull / test-vulkan-operators-linux / linux-job (gh)
  RuntimeError: Command docker exec -t e9041bddeeec1c45af8a2920b1e09a1c9d300e54c83430546db49c66c994205a /exec failed with exit code 127
• pull / unittest-arm-backend-with-no-deps (test_pytest_models_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t 8d551438624b32083139b41df38683299fd06134873d8bcf079e3612f1d83aca /exec failed with exit code 1
• pull / unittest-arm-backend-with-no-deps (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t 9fd20b195736dc8a912c8b01825a6e26c320554febf1bc43a11be8655283af00 /exec failed with exit code 1
• Test Metal Backend / test-metal-backend-modules / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 127

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

[Copilot]

Pull request overview

This pull request adds support for Linear + BatchNorm QAT (Quantization Aware Training) fusion in the NXP backend, following a similar approach to the existing Conv + BatchNorm fusion in TorchAO.

Changes:

• Adds AddSimulatedLinearBatchNormFusionQATPass to insert simulated fusion during QAT training by normalizing linear weights and denormalizing outputs using BatchNorm statistics
• Adds RemoveSimulatedLinearBatchNormFusionQATPass to remove the simulated fusion artifacts after QAT training is complete
• Updates FuseBatchNormWithLinearPass to handle fake-quantized nodes during fusion
• Updates LinearPattern quantizer to skip quantization between linear and batch norm nodes during QAT
• Adds comprehensive integration tests covering the fusion workflow

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
backends/nxp/tests/models.py	Adds LinearBNModule test model for testing Linear + BatchNorm patterns
backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py	Adds integration tests for simulated fusion, complete fusion pipeline, and graph equivalence
backends/nxp/aten_passes/add_simulated_linear_bn_fusion_qat_pass.py	Implements pass to add simulated Linear + BatchNorm fusion for QAT by normalizing weights and denormalizing outputs
backends/nxp/aten_passes/remove_simulated_linear_bn_fusion_qat_pass.py	Implements pass to remove simulated fusion artifacts after QAT training
backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py	Adds _unwrap_if_fq helper function to handle fake-quantized nodes during fusion
backends/nxp/quantizer/patterns.py	Updates LinearPattern to skip output quantization when followed by BatchNorm during QAT

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 8 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[robert-kalmar]

Only change I really would like to see is integration into the executorch_pipeline.py and to the aot_neutron_convert.py if it make sense (at this moment). Other my comments are nice to have.

[StrycekSimon]

Only change I really would like to see is integration into the executorch_pipeline.py and to the aot_neutron_convert.py if it make sense (at this moment). Other my comments are nice to have.

I will add it to the calibrate_and_quantize function just to include it in tests and cover that there are no conflicts (failing tests) with the rest of the implementation. Just note that the only model in aot_neutron_compile.py that uses QAT so far is CifarNet on which it is pretty pointless to explicitly use it as it does not contain the Linear+BN pattern.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[StrycekSimon]

After the recent push, one of PTQ test cases for mm op conversion fails. I will investigate.

Update: It's a mismatch by 1 in one of the output values, much like what can be seen in some other tests (eg. test_clone_converter.py) probably caused by float arithmetics. Increasing atol to 1 like in similar tests resolved this issue.

[Copilot]

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.