Skip to content

Commit 4147869

Browse files
authored
NXP backend: [fix] check the conv/linear node's user count in BatchNorm fusion (pytorch#20601)
## Summary The `_is_conv` / `_is_linear` closures in the NXP BatchNorm fusion passes check `len(node.users)` instead of `len(node_.users)`. `node_` is the closure's parameter (the conv/linear node being tested), but `node` resolves to the enclosing `for node in graph_module.graph.nodes` loop variable — which at the call site (`_is_conv(bn_node.args[0])`) is always the **BatchNorm** node. The single-user guard exists to prevent fusion when the conv/linear output feeds consumers *other* than the BatchNorm, because folding the BatchNorm into the conv/linear weights changes that output for **all** consumers. Due to the typo the guard inspected the BatchNorm's user count rather than the conv/linear's, so a conv/linear with multiple consumers was still fused — corrupting the other consumers' inputs. Affected files: - `backends/nxp/aten_passes/fuse_batch_norm_with_conv_pass.py` - `backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py` ## Test plan Reproduced on a small `torch.fx` graph where the conv has two consumers (BatchNorm + a relu), so fusion must be skipped: - **Before:** the pass fuses anyway → model output max-abs-diff vs. reference ≈ **76.2** (corrupted). - **After:** the pass is correctly a no-op → diff **0.0**. - **After**, single-consumer conv: still fuses correctly, numerics preserved (diff ≈ 3.8e-6). The linear pass contains the identical typo and the identical one-character fix.
1 parent a6d812a commit 4147869

3 files changed

Lines changed: 114 additions & 2 deletions

File tree

backends/nxp/aten_passes/fuse_batch_norm_with_conv_pass.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ def _is_conv(node_: Node):
6666
torch.ops.aten.conv1d.default,
6767
torch.ops.aten.conv2d.default,
6868
)
69-
has_single_user = len(node.users) == 1
69+
has_single_user = len(node_.users) == 1
7070

7171
return is_conv and has_single_user
7272

backends/nxp/aten_passes/fuse_batch_norm_with_linear_pass.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ def _is_linear(node_: Node):
8181
node_.op == "call_function"
8282
and node_.target == torch.ops.aten.linear.default
8383
)
84-
has_single_user = len(node.users) == 1
84+
has_single_user = len(node_.users) == 1
8585

8686
return is_linear and has_single_user
8787

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# Copyright 2025-2026 NXP
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import torch
7+
from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_conv_pass import (
8+
FuseBatchNormWithConvPass,
9+
)
10+
from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import (
11+
FuseBatchNormWithLinearPass,
12+
)
13+
14+
15+
def _graph_has_batch_norm(graph_module: torch.fx.GraphModule) -> bool:
16+
return any(
17+
node.op == "call_function" and node.target == torch.ops.aten.batch_norm.default
18+
for node in graph_module.graph.nodes
19+
)
20+
21+
22+
class ConvBatchNorm(torch.nn.Module):
23+
def __init__(self, share_conv_output: bool):
24+
super().__init__()
25+
self.conv = torch.nn.Conv2d(3, 4, kernel_size=3)
26+
self.bn = torch.nn.BatchNorm2d(4)
27+
self.share_conv_output = share_conv_output
28+
29+
def forward(self, x):
30+
y = self.conv(x)
31+
out = self.bn(y)
32+
if self.share_conv_output:
33+
# The conv output now feeds another consumer, so the BatchNorm must
34+
# not be fused into the conv weights (that would corrupt this branch).
35+
out = out + torch.relu(y)
36+
return out
37+
38+
39+
class LinearBatchNorm(torch.nn.Module):
40+
def __init__(self, share_linear_output: bool):
41+
super().__init__()
42+
self.linear = torch.nn.Linear(4, 4)
43+
self.bn = torch.nn.BatchNorm1d(4)
44+
self.share_linear_output = share_linear_output
45+
46+
def forward(self, x):
47+
y = self.linear(x)
48+
out = self.bn(y)
49+
if self.share_linear_output:
50+
out = out + torch.relu(y)
51+
return out
52+
53+
54+
def _exported_module(module, example_input):
55+
module.eval()
56+
return torch.export.export(module, example_input, strict=True).module()
57+
58+
59+
def test_conv_batch_norm_fused_when_single_user():
60+
gm = _exported_module(
61+
ConvBatchNorm(share_conv_output=False), (torch.randn(1, 3, 8, 8),)
62+
)
63+
assert _graph_has_batch_norm(gm)
64+
FuseBatchNormWithConvPass().call(gm)
65+
gm.graph.eliminate_dead_code()
66+
gm.recompile()
67+
# Conv has a single user (the BatchNorm), so the fusion removes the BatchNorm.
68+
assert not _graph_has_batch_norm(gm)
69+
70+
71+
def test_conv_batch_norm_not_fused_when_conv_has_multiple_users():
72+
example_input = (torch.randn(1, 3, 8, 8),)
73+
module = ConvBatchNorm(share_conv_output=True)
74+
reference = _exported_module(module, example_input)
75+
gm = _exported_module(module, example_input)
76+
77+
assert _graph_has_batch_norm(gm)
78+
FuseBatchNormWithConvPass().call(gm)
79+
gm.graph.eliminate_dead_code()
80+
gm.recompile()
81+
82+
# Conv feeds two consumers, so the BatchNorm must be left untouched.
83+
assert _graph_has_batch_norm(gm)
84+
x = torch.randn(1, 3, 8, 8)
85+
torch.testing.assert_close(reference(x), gm(x))
86+
87+
88+
def test_linear_batch_norm_fused_when_single_user():
89+
gm = _exported_module(
90+
LinearBatchNorm(share_linear_output=False), (torch.randn(4, 4),)
91+
)
92+
assert _graph_has_batch_norm(gm)
93+
FuseBatchNormWithLinearPass().call(gm)
94+
gm.graph.eliminate_dead_code()
95+
gm.recompile()
96+
assert not _graph_has_batch_norm(gm)
97+
98+
99+
def test_linear_batch_norm_not_fused_when_linear_has_multiple_users():
100+
example_input = (torch.randn(4, 4),)
101+
module = LinearBatchNorm(share_linear_output=True)
102+
reference = _exported_module(module, example_input)
103+
gm = _exported_module(module, example_input)
104+
105+
assert _graph_has_batch_norm(gm)
106+
FuseBatchNormWithLinearPass().call(gm)
107+
gm.graph.eliminate_dead_code()
108+
gm.recompile()
109+
110+
assert _graph_has_batch_norm(gm)
111+
x = torch.randn(4, 4)
112+
torch.testing.assert_close(reference(x), gm(x))

0 commit comments

Comments
 (0)