Repository: pytorch
Issue: [MPS] Incorrect output from convolution ops with large dimensions
Labels: high priority,triaged,module: regression,module: correctness (silent),module: mps

### 🐛 Describe the bug

#140726 attempted to reintroduce support for convolutions with `out_channels` > 2**16. While this appeared to work for `Conv1d`, it introduced a regression in `Conv2d` (see https://github.com/pytorch/pytorch/issues/142515#issuecomment-2533495363). It remains unclear whether `Conv3d` is affected.

The issue results in a silent correctness bug. A guard was previously added in #129484 to prevent this (lifted in #140726 for macOS >= 15.1 which caused the regression), and we should reinstate the guard for the affected operations.

#### Minimal repro

```python
import torch
import torch.nn.functional as F

weight_cpu = torch.randn(10, 10, 1, 65536, device="cpu")
weight_mps = weight_cpu.detach().clone().to("mps")

x_cpu = torch.randn(1, 10, 1, 65536, device="cpu")
x_mps = x_cpu.detach().clone().to("mps")

y_cpu = F.conv2d(x_cpu, weight_cpu)
y_mps = F.conv2d(x_mps, weight_mps)

print(y_cpu)
print(y_mps)
```

outputs

```
tensor([[[[-1725.4941]],

         [[  497.0602]],

         [[ -475.0262]],

         [[-1376.2498]],

         [[-1118.9232]],

         [[  566.6099]],

         [[ -195.5112]],

         [[-1461.4688]],

         [[ -119.6234]],

         [[ -679.2219]]]])
tensor([[[[0.]],

         [[0.]],

         [[0.]],

         [[0.]],

         [[0.]],

         [[0.]],

         [[0.]],

         [[0.]],

         [[0.]],

         [[0.]]]], device='mps:0')
```

### TODO

- [ ] Investigate which ops are currently broken.
- [ ] Add additional tests with large dimensions for all operations on the affected code path.
- [ ] Reinstate the guard for broken operations on all macOS versions.
- [ ] Ensure the fix is included in the v2.6.0 release.

Sorry for the breakage!

xref #129207 #134416 #140722 huggingface/parler-tts#148 haoheliu/versatile_audio_super_resolution#70

### Versions

PyTorch version: 2.6.0a0+gite1196df
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.1.1 (arm64)
GCC version: Could not collect
Clang version: 16.0.0 (clang-1600.0.26.4)
CMake version: version 3.31.1
Libc version: N/A

Python version: 3.12.8 | packaged by conda-forge | (main, Dec  5 2024, 14:19:53) [Clang 18.1.8 ] (64-bit runtime)
Python platform: macOS-15.1.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 M3 Max

Versions of relevant libraries:
[pip3] mypy==1.13.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.2
[pip3] onnx==1.17.0
[pip3] onnxscript==0.1.0.dev20240817
[pip3] optree==0.13.0
[pip3] torch==2.6.0a0+gite1196df
[pip3] torchaudio==2.5.0a0+ba696ea
[pip3] torchvision==0.20.0a0+e9a3213
[conda] numpy                     1.26.2                   pypi_0    pypi
[conda] optree                    0.13.0                   pypi_0    pypi
[conda] torch                     2.6.0a0+gite1196df           dev_0    <develop>
[conda] torchaudio                2.5.0a0+ba696ea           dev_0    <develop>
[conda] torchvision               0.20.0a0+e9a3213           dev_0    <develop>

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen