Repository: pytorch
Issue: 144x less efficient CPU usage when training NN past a certain width
Labels: module: performance,triaged,module: arm

### 🐛 Describe the bug

The code below is a minimal NN training loop with a fully connected NN of shape (10->width->width->10).
When width is 45, everything is fine, the code takes about 1 second, and only uses 1 CPU.
When width is 46, the code takes 9 seconds, and uses all 16 cpus.
So it's 144x less efficient (what are all those cycles doing?)
I'm guessing it switches over to a multi-threaded implementation when multiplying matrices of a certain size.  But something doesn't seem right.

I also tried the "mps" backend, which surprisingly has enough overhead that it isn't faster until the network is very wide.  

```
import time
import os

import torch
from torch import nn, optim
print(f"{torch.__version__=}, {os.uname()=}")


batch_size = 128
all_inputs = torch.randn((batch_size * 100, 10))
all_targets = all_inputs + 0.01 * torch.randn((batch_size * 100, 10))
for device, omp_num_threads in [("cpu", None), ("cpu", 1), ("mps", 1)]:
    if omp_num_threads is not None:
        torch.set_num_threads(omp_num_threads)
    for width in [32, 45, 46, 64, 128, 256, 512, 1024, 2048, 4096]:
        if device == "cpu" and width > 256: break  # too slow, don't bother
        network = nn.Sequential(nn.Linear(10, width), nn.Linear(width, width), nn.Linear(width, 10)).to(device)
        optimizer = optim.Adam(network.parameters(), lr=3e-4)
        t_start = time.time()
        for epoch in range(50):
            for offset in range(0, len(all_inputs), batch_size):
                inputs = all_inputs[offset:offset+batch_size].to(device)
                targets = all_targets[offset:offset+batch_size].to(device)
                optimizer.zero_grad()
                ((network(inputs) - targets) ** 2).mean().backward()
                optimizer.step()
        final_loss = ((network(all_inputs.to(device)) - all_targets.to(device)) ** 2).mean()
        print(f"{torch.get_num_threads()=}, device={device}, nn_width={width}, final_loss={final_loss:2.5f}, took {time.time() - t_start:2.1f} secs")
```

output on my machine is:

```
torch.__version__='2.3.1.post100', os.uname()=posix.uname_result(sysname='Darwin', nodename='username.macbook.pro.m3.lan', release='23.5.0', version='Darwin Kernel Version 23.5.0: Wed May  1 20:17:33 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T6031', machine='arm64')
torch.get_num_threads()=16, device=cpu, nn_width=32, final_loss=0.00010, took 0.8 secs
torch.get_num_threads()=16, device=cpu, nn_width=45, final_loss=0.00010, took 1.0 secs
torch.get_num_threads()=16, device=cpu, nn_width=46, final_loss=0.00010, took 8.8 secs  <---- 16 cpus, and 9x slower!
torch.get_num_threads()=16, device=cpu, nn_width=64, final_loss=0.00011, took 6.8 secs
torch.get_num_threads()=16, device=cpu, nn_width=128, final_loss=0.00012, took 19.9 secs
torch.get_num_threads()=16, device=cpu, nn_width=256, final_loss=0.00015, took 65.6 secs
# everything is way faster with just 1 thread (below)
torch.get_num_threads()=1, device=cpu, nn_width=32, final_loss=0.00010, took 0.9 secs 
torch.get_num_threads()=1, device=cpu, nn_width=45, final_loss=0.00010, took 1.0 secs
torch.get_num_threads()=1, device=cpu, nn_width=46, final_loss=0.00010, took 2.5 secs  <---- 1 cpu, and faster!
torch.get_num_threads()=1, device=cpu, nn_width=64, final_loss=0.00011, took 1.9 secs 
torch.get_num_threads()=1, device=cpu, nn_width=128, final_loss=0.00012, took 2.5 secs
torch.get_num_threads()=1, device=cpu, nn_width=256, final_loss=0.00015, took 4.2 secs
# mps has a lot of overhead, but eventually is faster
torch.get_num_threads()=1, device=mps, nn_width=32, final_loss=0.00010, took 8.7 secs
torch.get_num_threads()=1, device=mps, nn_width=45, final_loss=0.00010, took 8.7 secs
torch.get_num_threads()=1, device=mps, nn_width=46, final_loss=0.00010, took 8.9 secs
torch.get_num_threads()=1, device=mps, nn_width=64, final_loss=0.00011, took 8.4 secs
torch.get_num_threads()=1, device=mps, nn_width=128, final_loss=0.00012, took 11.8 secs
torch.get_num_threads()=1, device=mps, nn_width=256, final_loss=0.00015, took 8.4 secs
torch.get_num_threads()=1, device=mps, nn_width=512, final_loss=0.00019, took 11.2 secs
torch.get_num_threads()=1, device=mps, nn_width=1024, final_loss=0.00027, took 9.2 secs
torch.get_num_threads()=1, device=mps, nn_width=2048, final_loss=0.00033, took 10.0 secs
torch.get_num_threads()=1, device=mps, nn_width=4096, final_loss=0.00032, took 27.9 secs. <-- quadratic runtime starts here as expected
```

### Versions


```
Collecting environment information...
PyTorch version: 2.3.1.post100
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.5 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.0 | packaged by conda-forge | (main, Jan 14 2023, 12:26:40) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.5-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 M3 Max

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.3.1.post100
[conda] numpy                     1.26.4          py311he598dae_0  
[conda] numpy-base                1.26.4          py311hfbfe69c_0  
[conda] pytorch                   2.3.1           gpu_mps_py311h7b7e308_100  
```

cc @msaroufim @malfet @snadampal @milpuz01