RuntimeError: Undefined type BFloat16 from matmul above certain sizes (maybe MPS only)

[Vargol]

🐛 Describe the bug

When using torch.matmul (and the @ operator) on bfloat16 tensors above a certain size I'm getting the error
RuntimeError: Undefined type BFloat16

Example code

import torch

sdxl_latent_rgb_factors = torch.tensor(
            [
                #   R        G        B
                [0.3816, 0.4930, 0.5320],
                [-0.3753, 0.1631, 0.1739],
                [0.1770, 0.3588, -0.2048],
                [-0.4350, -0.2644, -0.4289],
            ],
            dtype=torch.bfloat16,
            device='mps',
        )

print('-------181,181,4--works----------')

samples = torch.zeros(  181 , 181, 4, 
                     dtype=torch.bfloat16,
                    device='mps',
)

x = torch.matmul(samples, sdxl_latent_rgb_factors)

print('-------182,182,4--fails--------')

samples2 = torch.zeros(  182 , 182, 4, 
                     dtype=torch.bfloat16,
                    device='mps',
)

x = torch.matmul(samples2, sdxl_latent_rgb_factors)

results in

 python ../Diffusers/matmul2.py 
-------181,181,4--works----------
-------182,182,4--fails--------
Traceback (most recent call last):
  File "/Volumes/SSD2TB/AI/Cascade/../Diffusers/matmul2.py", line 31, in <module>
    x = torch.matmul(samples2, sdxl_latent_rgb_factors)
RuntimeError: Undefined type BFloat16

Versions

Collecting environment information...
PyTorch version: 2.3.0.dev20240309
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.2.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.24.4
Libc version: N/A

Python version: 3.10.13 (main, Nov  9 2023, 13:59:31) [Clang 15.0.0 (clang-1500.0.40.1)] (64-bit runtime)
Python platform: macOS-14.2.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

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.3.0.dev20240309
[pip3] torchaudio==2.2.0.dev20240301
[pip3] torchvision==0.18.0.dev20240309
[conda] Could not collect

cc @kulinseth @albanD @malfet @DenisVieriu97 @razarmehr

[malfet]

Ok, I know the problem: MPS supports BFloat16, but metal is not, and as MPS matmul is buggy for large matrices, there is a naive Metal matmul implementation, that yet not support bf16

[atalman]

validated with 2.3:

-------181,181,4--works----------
-------182,182,4--fails--------