Repository: pytorch
Issue: Faster dropout (aten.bernoulli) on cpu
Labels: module: performance,triaged,module: mps

### 🚀 The feature, motivation and pitch

I use torch 2.3.1 on MacOS 15.1 beta 

I profiling my model and observed that aten::bernoulli is taking 34% of the runtime on cpu while  dropout time is marginal for MPS (GPU) device.


Can we optimize bernoulli computation for cpu ?

### Alternatives

_No response_

### Additional context


Here the stats:

cpu using torch.profiler:
````
STAGE:2024-10-17 11:38:42 87880:56463086 ActivityProfilerController.cpp:314] Completed Stage: Warm Up
STAGE:2024-10-17 11:38:42 87880:56463086 ActivityProfilerController.cpp:320] Completed Stage: Collection
STAGE:2024-10-17 11:38:42 87880:56463086 ActivityProfilerController.cpp:324] Completed Stage: Post Processing
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                          aten::dropout         0.05%      56.000us        35.83%      44.111ms       6.302ms             7  
                                       aten::bernoulli_        33.98%      41.829ms        33.98%      41.829ms       5.976ms             7  
    autograd::engine::evaluate_function: IndexBackward0         0.11%     140.000us        14.23%      17.515ms     547.344us            32  
                                         IndexBackward0         0.08%     101.000us        14.11%      17.375ms     542.969us            32  
                                 aten::_index_put_impl_        14.00%      17.238ms        14.05%      17.298ms     455.211us            38  
````

cpu using cProfile

````
         1823 function calls (1767 primitive calls) in 0.117 seconds

   Ordered by: cumulative time
   List reduced from 168 to 100 due to restriction <100>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     24/2    0.000    0.000    0.069    0.034 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/module.py:1528(_wrapped_call_impl)
     24/2    0.000    0.000    0.069    0.034 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/module.py:1534(_call_impl)
        1    0.004    0.004    0.069    0.069 /Users/tgg/Github/mlx-graphs-last/AttentiveFP.py:432(forward)
        1    0.000    0.000    0.045    0.045 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/_tensor.py:466(backward)
        1    0.000    0.000    0.045    0.045 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/autograd/__init__.py:165(backward)
        1    0.000    0.000    0.045    0.045 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/autograd/graph.py:739(_engine_run_backward)
        1    0.045    0.045    0.045    0.045 {method 'run_backward' of 'torch._C._EngineBase' objects}
        7    0.000    0.000    0.041    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/dropout.py:58(forward)
        7    0.000    0.000    0.041    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/functional.py:1279(dropout)
        7    0.041    0.006    0.041    0.006 {built-in method torch.dropout}

````

mps using cProfile 
````
        1823 function calls (1767 primitive calls) in 0.093 seconds

   Ordered by: cumulative time
   List reduced from 168 to 100 due to restriction <100>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     24/2    0.000    0.000    0.080    0.040 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/module.py:1528(_wrapped_call_impl)
     24/2    0.000    0.000    0.080    0.040 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/module.py:1534(_call_impl)
        1    0.068    0.068    0.080    0.080 /Users/tgg/Github/mlx-graphs-last/AttentiveFP.py:432(forward)
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/_tensor.py:466(backward)
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/autograd/__init__.py:165(backward)
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/autograd/graph.py:739(_engine_run_backward)
        1    0.006    0.006    0.006    0.006 {method 'run_backward' of 'torch._C._EngineBase' objects}
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/optim/optimizer.py:374(wrapper)
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/optim/optimizer.py:58(_use_grad)
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/optim/adam.py:135(step)
        1    0.000    0.000    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/optim/adam.py:260(adam)
        1    0.001    0.001    0.006    0.006 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/optim/adam.py:338(_single_tensor_adam)
        3    0.002    0.001    0.002    0.001 {built-in method torch.stack}
        4    0.000    0.000    0.001    0.000 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/rnn.py:1457(forward)
        4    0.001    0.000    0.001    0.000 {built-in method torch.gru_cell}
        1    0.001    0.001    0.001    0.001 /Users/tgg/Github/mlx-graphs-last/AttentiveFP.py:441(<listcomp>)
        3    0.001    0.000    0.001    0.000 {method 'type' of 'torch._C.TensorBase' objects}
        1    0.001    0.001    0.001    0.001 /Users/tgg/Github/mlx-graphs-last/AttentiveFP.py:444(<listcomp>)
        7    0.000    0.000    0.001    0.000 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/dropout.py:58(forward)
        7    0.000    0.000    0.001    0.000 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/functional.py:1279(dropout)
       11    0.000    0.000    0.001    0.000 /Users/tgg/miniforge3/envs/mlxgraphenv-py311/lib/python3.11/site-packages/torch/nn/modules/linear.py:115(forward)
        7    0.001    0.000    0.001    0.000 {built-in method torch.dropout}
````

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