Repository: pytorch
Issue: Bug in Conv/BN fuser with torch.fx
Labels: triaged,module: fx,oncall: fx

### 🐛 Describe the bug

Conv/BN fuser, as introduced in https://github.com/pytorch/pytorch/pull/47657 , only checks if the output of conv is used in multiple places. It does not check every usage of conv/bn modules.

A failure example here:

```python
import torch
from torch import nn

class ToyModel(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.conv1 = nn.Conv2d(1, 1, 1)
        self.bn1 = nn.BatchNorm2d(1)
        self.bn1.weight.data.normal_()
    
    def forward(self, x, y):
        return self.conv1(x) + self.bn1(self.conv1(y))

model = ToyModel()
model.eval()

a = torch.rand(64, 1, 32, 32)
b = torch.rand(64, 1, 32, 32)

from torch.fx.experimental.optimization import fuse
output = model(a, b)

print(torch.fx.symbolic_trace(model).graph)
fused_model = fuse(model)
print(fused_model.graph)
output2 = fused_model(a, b)
print((output - output2).abs().max())
```

Output:
```
graph():
    %x : [#users=1] = placeholder[target=x]
    %y : [#users=1] = placeholder[target=y]
    %conv1 : [#users=1] = call_module[target=conv1](args = (%x,), kwargs = {})
    %conv1_1 : [#users=1] = call_module[target=conv1](args = (%y,), kwargs = {})
    %bn1 : [#users=1] = call_module[target=bn1](args = (%conv1_1,), kwargs = {})
    %add : [#users=1] = call_function[target=operator.add](args = (%conv1, %bn1), kwargs = {})
    return add

graph():
    %x : [#users=1] = placeholder[target=x]
    %y : [#users=1] = placeholder[target=y]
    %conv1 : [#users=1] = call_module[target=conv1](args = (%x,), kwargs = {})
    %conv1_1 : [#users=1] = call_module[target=conv1](args = (%y,), kwargs = {})
    %add : [#users=1] = call_function[target=operator.add](args = (%conv1, %conv1_1), kwargs = {})
    return add

tensor(0.3768, grad_fn=<MaxBackward1>)
```

The solution is to check the global usage count of conv/bn modules, and to only fuse modules that are used once.

### Versions

```
Collecting environment information...
PyTorch version: 1.12.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.4.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.25.1
Libc version: N/A

Python version: 3.8.13 (default, Mar 28 2022, 06:13:39)  [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-13.4.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.22.4
[pip3] torch==1.12.0
[pip3] torchvision==0.13.0
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] torch                     1.12.0                   pypi_0    pypi
[conda] torchvision               0.13.0                   pypi_0    pypi
```

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv