NXP backend: Condition Linear+BatchNorm fusing passes to valid combination only

[StrycekSimon]

Summary

Introduce shape conditioning to Linear+BN fusion related passes.
Fix RemoveSimulatedLinearBatchNormFusionQATPass pass recompiling graph when not necessary.

Problem description:
Linear and BatchNorm layers are fusable only in cases where both layers work with the same dimension. Linear layer works strictly with only the last dimension of it's input. BatchNorm on the other hand works in channels dimension (the second one in most cases) except for one case when using BatchNorm1d and input shape [N, C] (see official PyTorch documentation for more info). In this one case both BatchNorm and Linear works on the same dimension and are correctly fusable.

Example:

inp = torch.randn((2, 2, 2)) # (N, C, L)
lin = nn.Linear(in_features=3, out_features=5) # Weight of shape (5, 3)
bn = nn.BatchNorm2d(num_features=2) # Running var shape (2, )

# With the following fusing equation:
# linear_w_fused = linear_w * (gamma / sqrt(var + eps))
# And the following shapes:
# linear_w.shape: (5, 3); var.shape: (2, )
# We cannot correctly perform the multiplication.

Test plan

New unit test cases were added + relevant older ones were adjusted.

cc @robert-kalmar @JakeStevens @digantdesai

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 8 New Failures

As of commit da6cae3 with merge base f78535d ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner-mypy (gh)
  >>> Lint for test/end2end/test_end2end.py:
• pull / test-qnn-wheel-packages-linux (3.10) / linux-job (gh)
  RuntimeError: Command docker exec -t 063b2eaa136e4d14261dd5cf1c1c1e517a495082a40ae118ddfac1502475e16c /exec failed with exit code 1
• pull / test-qnn-wheel-packages-linux (3.11) / linux-job (gh)
  RuntimeError: Command docker exec -t 55f825c7e48bf7522fca152d30cec8ecb9c8f22984d2c88dc238dafc8e5867df /exec failed with exit code 1
• pull / test-qnn-wheel-packages-linux (3.12) / linux-job (gh)
  RuntimeError: Command docker exec -t 1c170441104761481534868899560154e1e903c492dabc67472d697f60c313a9 /exec failed with exit code 1
• pull / test-qnn-wheel-packages-linux (3.13) / linux-job (gh)
  RuntimeError: Command docker exec -t c521ce0733574c4f6a153d8e5a0bf1848fa61f6a01858c0a578cb0970c30e4ce /exec failed with exit code 1
• pull / test-samsung-models-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 045ad7f1a6cfd444ef865017c0ce77e5bbcd0ed2cfc7d84dcc6388860bef66a8 /exec failed with exit code 1
• pull / test-samsung-quantmodels-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 82551a82d47238a6152b484c2d2ed007c3c145eb0ddcade3bfa3ab928dbeefbd /exec failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh)
  backends/xnnpack/test/ops/test_conv2d.py::TestConv2d::test_fp16_conv2d

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

[Copilot]

Pull request overview

This PR tightens the NXP backend’s Linear+BatchNorm fusion logic by conditioning fusion passes on valid shape combinations (only fuse when Linear and BN operate on the same dimension), and avoids unnecessary graph recompilation in the QAT “remove simulated fusion” pass.

Changes:

• Add shape-based gating (via FX tensor_meta shape) to Linear+BN fusion passes to prevent invalid fusions.
• Track whether graph edits were made in RemoveSimulatedLinearBatchNormFusionQATPass and only recompile when needed.
• Update/refactor NXP backend tests and test models; add new negative tests ensuring incompatible Linear+BN are not fused.

Reviewed changes

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

Show a summary per file

File	Description
backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py	Adds output-shape conditioning before fusing Linear+BN.
backends/nxp/aten_passes/simulated_linear_bn_fusion_passes/add_simulated_linear_bn_fusion_qat_pass.py	Adds output-shape conditioning to simulated Linear+BN QAT fusion insertion.
backends/nxp/aten_passes/simulated_linear_bn_fusion_passes/remove_simulated_linear_bn_fusion_qat_pass.py	Returns/propagates a “made changes” flag to avoid unnecessary recompile.
backends/nxp/backend/graph_utils.py	Adds get_output_shape helper for retrieving tensor_meta shapes.
backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py	Updates model usage and adds negative tests for non-fusable shapes.
backends/nxp/tests/models.py	Refactors BN-related helper modules used by tests (but currently introduces a problematic import / rename compatibility issue).
backends/nxp/tests/test_batch_norm_fusion.py	Updates fusion tests to use shared helpers and adds incompatible-shape full-pipeline coverage.

---

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