Repository: pytorch
Issue: cpu faster than mps for both GRUCell and LSTMCell on apple silicon M3 max: MPS Metal
Labels: module: performance,triaged,module: mps

### 🐛 Describe the bug

I compare the GRUCell & LTSMCell on both cpu and mps device on apple silicon:

````
# Example usage of the CustomminGRUCell
hidden_size = 128
batch_size = 32
input_size = 128

device = torch.device('cpu')
# Initialize cell and example data
gru_cell = nn.GRUCell(input_size,hidden_size).to(device)
h = torch.randn(batch_size, hidden_size).to(device)  # Initial hidden state
inputs = torch.randn(batch_size, input_size).to(device)  # Input tensor
````
**result with CPU using %timeit gru_cell(inputs, h)
117 µs ± 3.59 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

**result with MPS using %timeit gru_cell(inputs, h)
249 µs ± 4.36 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

*Similar with LSTMCell:
````
# Example usage of the CustomminGRUCell
hidden_size = 128
batch_size = 32
input_size = 128

device = torch.device('cpu')
# Initialize cell and example data
lstm_cell = nn.LSTMCell(input_size,hidden_size).to(device)
h = torch.randn(batch_size, hidden_size).to(device)  # Initial hidden state
c = torch.randn(batch_size, hidden_size).to(device)  # Initial hidden state

inputs = torch.randn(batch_size, input_size).to(device)  # Input tensor
`````
**Result on CPU using %timeit lstm_cell(inputs, (h, c))
150 µs ± 3.43 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

**Result on MPS using %timeit lstm_cell(inputs, (h, c))
236 µs ± 3.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

I have also tested the minGRUCell case:
````
import torch
import torch.nn as nn
import torch.nn.functional as F

class minGRUCell(nn.Module):
    def __init__(self, hidden_features):
        super(minGRUCell, self).__init__()
        self.hidden_features = hidden_features

        # Linear layers for input and hidden state transformation
        self.dense_i = nn.Linear(hidden_features, hidden_features)
        self.dense_h = nn.Linear(hidden_features, hidden_features)

    def forward(self, inputs, h):
        # GRU Gates
        htilde = self.dense_i(inputs)
        z = torch.sigmoid( self.dense_h(inputs))  # Update gate ie 0 to 1 value!

        # Update hidden state
        new_h = torch.lerp(h, htilde, z)

        return new_h  # Return updated hidden state


# Example usage of the CustomminGRUCell
hidden_size = 128
batch_size = 32
input_size = 128

device = torch.device('cpu')
# Initialize cell and example data
mingru_cell = minGRUCell(hidden_size).to(device)
h = torch.randn(batch_size, hidden_size).to(device)  # Initial hidden state
inputs = torch.randn(batch_size, input_size).to(device)  # Input tensor

# Forward pass
new_h = mingru_cell(inputs, h)

`````
**Result on cpu : %timeit mingru_cell(inputs, h)
24.4 µs ± 269 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

**Result on mps: %timeit mingru_cell(inputs, h)
84.3 µs ± 1.23 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Any idea how to improve this ???

### Versions

Collecting environment information...
PyTorch version: 2.3.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.1 (arm64)
GCC version: Could not collect
Clang version: 18.1.3
CMake version: version 3.30.3
Libc version: N/A

Python version: 3.11.8 | packaged by conda-forge | (main, Feb 16 2024, 20:49:36) [Clang 16.0.6 ] (64-bit runtime)
Python platform: macOS-15.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==0.800
[pip3] mypy_extensions==0.4.4
[pip3] numpy==1.26.3
[pip3] onnx==1.16.1
[pip3] onnx-tf==1.10.0
[pip3] onnx2keras==0.0.24
[pip3] onnxruntime==1.18.0
[pip3] optree==0.12.1
[pip3] torch==2.3.1
[pip3] torch_cluster==1.6.3
[pip3] torch_geometric==2.5.0
[pip3] torch_scatter==2.1.2
[pip3] torch_sparse==0.6.18
[pip3] torch_spline_conv==1.2.2
[pip3] torchdata==0.7.1
[pip3] torchvision==0.18.1
[conda] nomkl                     1.0                  h5ca1d4c_0    conda-forge
[conda] numpy                     1.26.3                   pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] torch                     2.3.1                    pypi_0    pypi
[conda] torch-cluster             1.6.3                    pypi_0    pypi
[conda] torch-geometric           2.5.0                    pypi_0    pypi
[conda] torch-scatter             2.1.2                    pypi_0    pypi
[conda] torch-sparse              0.6.18                   pypi_0    pypi
[conda] torch-spline-conv         1.2.2                    pypi_0    pypi
[conda] torchdata                 0.7.1                    pypi_0    pypi
[conda] torchvision               0.18.1                   pypi_0    pypi

cc @msaroufim @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen