Repository: pytorch
Issue: unique(return_counts=True) fails on MPS for unsorted tensors with 1M+ elements
Labels: triaged,module: correctness (silent),module: mps

Trying to count many unique integers on an M2 Max fails. When trying to reproduce, it starts failing for me somewhere between half a million and about a million (unsorted) integers.

Whether it fails is not deterministic. Sometimes large tensors will work (even within the same session), but at least on my machine, for 1M elements it fails much more often than it works.

Reproduce:
```python
import torch

torch.unique(torch.randint(0, 2, (1_000_000,), device="mps"), return_counts=True)
```

Output:
```
Error: command buffer exited with error status.
	The Metal Performance Shaders operations encoded on it may not have completed.
	Error: 
	(null)
	Internal Error (0000000e:Internal Error)
	<AGXG14XFamilyCommandBuffer: 0x366e63060>
    label = <none> 
    device = <AGXG14CDevice: 0x112011200>
        name = Apple M2 Max 
    commandQueue = <AGXG14XFamilyCommandQueue: 0x111037e00>
        label = <none> 
        device = <AGXG14CDevice: 0x112011200>
            name = Apple M2 Max 
    retainedReferences = 1
```

The resulting counting tensor may have a few correct values at the beginning, and then too few for the rest of the elements, or (incorrect) zero values.

Adding `sorted=True` does not fix the issue.

Test to find the failure point:
```python
for elem_power in range(15, 33):
    num_elems = 2 ** elem_power
    rand_ints = torch.randint(0, 2, (num_elems,), device="mps")
    assert torch.unique(rand_ints, return_counts=True)[1].sum() == num_elems, print(f"2**{elem_power}={num_elems} failed!")
```

I've seen it failing at 2^19, 2^20 and 2^21 elements.

### Versions

I can reproduce the bug on this machine (Mac Studio / M2 Max / 32 GB) on 2.1.0 and on current nightly.

#### 2.1.0

Collecting environment information...
PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: version 3.27.4
Libc version: N/A

Python version: 3.10.13 (main, Sep 11 2023, 08:24:56) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.0-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 M2 Max

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] torch==2.1.0
[pip3] torchvision==0.15.2a0
[conda] Could not collect

---

#### Nightly

Collecting environment information...
PyTorch version: 2.2.0.dev20231012
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: version 3.27.4
Libc version: N/A

Python version: 3.10.13 (main, Sep 11 2023, 08:24:56) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.0-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 M2 Max

Versions of relevant libraries:
[pip3] numpy==1.24.1
[pip3] torch==2.2.0.dev20231012
[pip3] torchaudio==2.2.0.dev20231012
[pip3] torchvision==0.17.0.dev20231012
[conda] Could not collect


cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen @razarmehr @abhudev