[XNNPACK] Add support for Linear fused BatchNorm

[keyprocedure]

Summary

These changes implement a fusion pass in the XNNPACK partitioner to support linear + batchnorm operations. This pass involves identifying and combining each linear node that exclusively precedes a batchnorm node in the Export IR graph. Fusion occurs by updating the linear node's weight and bias with a new fused weight and bias computed from the linear and batchnorm parameters. In the case of linear nodes without bias, the fused bias is added as a new parameter to the linear node. All users of the batchnorm output are then redirected to the fused linear node, and the batchnorm node is removed from the graph.

This linear + batchnorm pass follows the existing implementation pattern of the convolution + batchnorm pass. These fusion passes fold batchnorm into the preceding convolution or linear ops during export. This allows the XNNPACK backend to run a single fused operation at inference, reducing memory usage and latency without affecting model accuracy.

Note: The current linear + batchnorm fusion implementation supports FP32 only. Quantized support is planned for a future PR to TorchAO.

Fixes #11587

Test plan

Tests were added to verify that linear + batchnorm fusion occurs for FP32 models when the linear layer has a single user, and is skipped for linear layers with multiple users. Both linear cases, with and without bias, are tested. A separate test ensures that standalone batchnorm layers are not partitioned, since XNNPACK does not currently support them.

Tests run and passed via:
python -m unittest executorch.backends.xnnpack.test.passes.test_batch_norm_fusion

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 3 Unrelated Failures

As of commit 6c92d24 with merge base a8d7298 ():

NEW FAILURES - The following jobs have failed:

• pull / test-eval_llama-mmlu-linux / linux-job (gh)
  RuntimeError: Command docker exec -t ffcd01d7e75a2a833efe96cc61ddd80e586410c48b612d54c96af89e1bb3fd16 /exec failed with exit code 1
• trunk / test-arm-cortex-m-size-test (bare_metal) / linux-job (gh)
  RuntimeError: Command docker exec -t dae6a96f6841507247d92d0bb60f0bb4cc2d76382659cab3a852159c73ca48aa /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• trunk / test-models-linux-aarch64 (mobilebert, xnnpack-quantization-delegation, linux.arm64.2xlarge) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• trunk / unittest-release / linux / linux-job (gh) (trunk failure)
  [ FAILED ] TensorPtrMakerTest.FailedCreateTensorUsingFromBlobWithIllegalStrides
• trunk / unittest-release / macos / macos-job (gh) (trunk failure)
  [ FAILED ] TensorPtrMakerTest.FailedCreateTensorUsingFromBlobWithIllegalStrides

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

[keyprocedure]

@pytorchbot label "release notes: none"

[mcr229]

awesome! Do you mind fusing these two passes? Something like BatchNormFusion pass? that way we can just generalize these?

[keyprocedure]

Good call, combining the two passes cleaned up a lot of the duplication. Let me know if anything else needs to be changed.

[digantdesai]

Thanks. LGTM. I will let Max stamp it.

[mcr229]

Love this thank you!

[keyprocedure]

Awesome, I'm glad!

Should we support fusion for Linear with bias=False as well? I can add it to this PR.

[digantdesai]

@keyprocedure let us know once this is good to go. Me or @mcr229 can merge this. Also you might need to rebase, rerun ci.

[keyprocedure]

@digantdesai I added support for Linear with bias=False and rebased. Everything is good to go from my end - ready for CI. Thanks for reviewing and following up earlier!

[keyprocedure]

The CI failures look unrelated to this PR. I noticed they were being discussed on Discord as known trunk issues. Just wanted to share in case it's helpful. Let me know if there's anything you'd like me to update.

Here's a summary of the failures:

• test-eval_llama-mmlu-linux: RuntimeError: Dataset scripts are no longer supported, but found mmlu_no_train.py
• test-arm-cortex-m-size-test (bare_metal): size check failed in cmake-out/test/size_test
• unittest-release (linux & macos): TensorPtrMakerTest.FailedCreateTensorUsingFromBlobWithIllegalStrides