Repository: pytorch
Issue: modified CPUReproTests.test_two_local_buffers_in_outer_loop_fusion fails Codegen on all CPU platforms
Labels: module: cpu,triaged

### 🐛 Describe the bug

Whilst investigating another bug, I discovered this happens on aarch64 and x86

```
    scalar_kernel = codegen_kernel(CppKernel)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/codegen/cpp.py", line 3830, in codegen_kernel
    run(kernel)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/codegen/cpp.py", line 3842, in run
    fn(vars, reduction_vars)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/codegen/cpp.py", line 4019, in fn
    return node.codegen(index_vars)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/scheduler.py", line 1053, in codegen
    self._body(*index_vars)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/loop_body.py", line 404, in __call__
    result = self.root_block()
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/loop_body.py", line 638, in __call__
    return InterpreterShim(graph, submodules).run(V.get_ops_handler())
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/loop_body.py", line 60, in run
    return super().run(*args, **kwargs)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/fx/interpreter.py", line 167, in run
    self.env[node] = self.run_node(node)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/loop_body.py", line 56, in run_node
    return super().run_node(n)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/fx/interpreter.py", line 228, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/fx/interpreter.py", line 332, in call_method
    return getattr(self_obj, target)(*args_tail, **kwargs)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/sizevars.py", line 915, in load
    return self._inner.load(name, self._simplify(index))
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/codegen/common.py", line 1973, in load
    out = self.load(name, index)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/codegen/cpp.py", line 1973, in load
    var = self.args.input(name)
  File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/_inductor/codegen/common.py", line 1138, in input
    assert name not in V.graph.removed_buffers, name
torch._dynamo.exc.BackendCompilerFailed: backend='compile_fx_wrapper' raised:
AssertionError: buf1
```

full stacktrace attached in text file.
[bug.txt](https://github.com/user-attachments/files/18038947/bug.txt)


This is a modified version of fn from **CPUReproTests.test_two_local_buffers_in_outer_loop_fusion**

```
import torch

def fn(x):
    softmax = torch.nn.functional.softmax(x, dim=-1)
    sum = torch.sum(softmax, dim=-1)
    sum_broadcast = torch.broadcast_to(
        sum.unsqueeze(-1), [*(sum.size()[0:3]), 256]
    )
    sum_exp = torch.exp(sum_broadcast)
    return torch.sum(sum_exp, dim=-1)

x = torch.randn(4, 12, 1023, 1022)
foo = torch.compile(fn)
foo(x)
```

This should be reproducible on any platform , eager mode seems to work fine.

### Versions

Collecting environment information...
PyTorch version: 2.6.0a0+gitcc64ad6
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (aarch64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.31.1
Libc version: glibc-2.35

Python version: 3.10.16 | packaged by conda-forge | (main, Dec  5 2024, 14:08:42) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.5.0-1018-aws-aarch64-with-glibc2.35
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:
Architecture:                       aarch64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
CPU(s):                             48
On-line CPU(s) list:                0-47
Vendor ID:                          ARM
Model name:                         Neoverse-N1
Model:                              1
Thread(s) per core:                 1
Core(s) per cluster:                48
Socket(s):                          -
Cluster(s):                         1
Stepping:                           r3p1
BogoMIPS:                           243.75
Flags:                              fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm lrcpc dcpop asimddp ssbs
L1d cache:                          3 MiB (48 instances)
L1i cache:                          3 MiB (48 instances)
L2 cache:                           48 MiB (48 instances)
L3 cache:                           32 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-47
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:           Mitigation; __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; CSV2, BHB
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] mypy==1.13.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.22.4
[pip3] onnx==1.17.0
[pip3] onnxscript==0.1.0.dev20240817
[pip3] optree==0.13.0
[pip3] torch==2.6.0a0+gitcc64ad6
[conda] No relevant packages

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10