Support passthrough ops for batch norm fusion (#4776)

[DrJessop]

Summary:

Prior to this diff, the fuse batch norm with conv function in prepare only supported conv immediately followed by batch norm. Now, we also support fusing through a non-channel slice. We may consider adding more passthrough ops when needed.

Differential Revision: D116357358

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4776

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit a427d97 with merge base fa7b7eb ():

NEW FAILURE - The following job has failed:

• Run PT2E GPU Tests / test-rocm (ROCM Nightly, linux.rocm.gpu.ecosystem.mi350.1, --pre torch torchvision --index-url ht... / linux-job (gh)
  Could not load credentials from any providers

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

[meta-codesync]

@DrJessop has exported this pull request. If you are a Meta employee, you can view the originating Diff in D116357358.

[vkuzo]

@claude review

[claude]

Claude finished @vkuzo's task in 1m 53s —— View job

---

Claude Code Review

• Read changed files (torchao/quantization/pt2e/utils.py, test file)
• Analyze fusion logic changes
• Analyze test coverage
• Post review feedback

---

Review Summary

Nice, well-scoped change. The core idea — allowing conv→BN fusion to "see through" a spatial (non-channel) slice and an optional contiguous/clone — is mathematically sound: BatchNorm is a per-channel affine applied elementwise, so it commutes with a slice on any non-channel dim (bn(slice(y)) == slice(bn(y))). Folding the BN into the conv weights and leaving the slice in place therefore preserves numerics. The refactor of the ad-hoc conv-lookup into _get_conv_bn_fusion_nodes is a readability win, and the guard rails (single-user checks, channel-dim rejection) are the right ones.

The three new tests cover the important cases well: the positive spatial-slice case, the rejected channel slice, and the rejected shared-conv-output case.

Correctness — looks good

• _is_non_channel_slice_node correctly resolves negative dim via dim + input_val.dim() and rejects dim == 1. Returning False when the val/meta is unavailable is appropriately conservative.
• The single-user constraints on the conv, the slice, and the contiguous/clone node are all enforced, so the shared-output case (test_..._shared_conv_output_is_skipped) is correctly rejected.
• replacement_node is set to the direct BN input (the contiguous/clone node when present, else the slice), so BN's users are rewired to the correct downstream node after the BN is erased. 👍

Suggestions (non-blocking)

1. sym_size users may block fusion unnecessarily (utils.py:996). The check len(conv_node.users) != 1 or len(slice_node.users) != 1 counts all users. Under dynamic shapes, the conv/slice output can also feed aten.sym_size/sym_numel nodes, which would make this bail out even though fusion is still valid. This file already has _filter_sym_size_users for exactly this situation — consider using it here so dynamic-shape graphs still fuse:

   if len(_filter_sym_size_users(conv_node)) != 1 or len(_filter_sym_size_users(slice_node)) != 1:
       return None

   This is a missed-optimization, not a correctness bug (it only ever skips fusion), so it's safe to defer. Fix this →

2. contiguous/clone only handled directly before BN. _get_conv_bn_fusion_nodes steps through at most one contiguous/clone node, and only when it sits between the slice and the BN. A graph like conv → contiguous → bn (no slice) won't fuse, and neither will conv → slice → clone → contiguous → bn. That matches the PR's stated scope ("fuse through a non-channel slice"), so this is fine — just flagging it as a known boundary in case follow-up patterns are needed.

3. Minor: consistent arg access in _is_non_channel_slice_node (utils.py:951 vs 955). The positive-dim branch reads dim via the normalized get_arg(n, "dim", int), while the negative-dim branch reaches into n.args[0] directly for the slice input. Both are correct, but pulling the input through get_arg/normalized args too would keep the accessor style uniform.

Test nits (optional)

• The three new tests are eval()-only (PTQ path, batch_norm/_native_batch_norm_legit_no_training). A QAT (prepare_qat_pt2e) variant that hits the training BN ops would broaden coverage, since the fusion helper also supports the training targets. Not required for this PR.

Overall this is a clean, correct, and well-tested addition. The suggestions above are all optional polish.
· branch export-D116357358

[vkuzo]

lgtm. Could you make sure OSS CI is green, especially the ruff piece

[DrJessop]

https://github.com/pytorch/ao/actions/runs/32752092932/job/97511173340?pr=4776 Keeps failing with "Node 20 is being deprecated", which is unrelated to my diff. Any thoughts on how to move forward to get a green CI?