Repository: pytorch
Issue: masked_select is x3 slower than reshaping and index_select
Labels: module: performance,triaged,module: advanced indexing,module: viewing and reshaping

## 🐛 Bug

masked_select is about x3 slower than index select of the corresponding nonzero'ed mask
(this is true for nightly, I have not checked with stable)

## To Reproduce
```
import torch                                                                                                                                        
import torch.utils.benchmark as benchmark                                                                                                           
                                                                                                                                                    
device = torch.device('cuda')                                                                                                                       
size = torch.Size([100, 100, 100])                                                                                                                 
mask = torch.randint(low=0, high=2, size=size, device=device).bool()                                                                                              
x = torch.randn(size=size, device=device, dtype=torch.float)                                                                                        
                                                                                                                                                      
out0 = torch.masked_select(x, mask)                                                                                                                 
out1 = x.view(-1).index_select(0, mask.view(-1).nonzero().squeeze())                                                                                
                                                                                                                                                      
assert (out0 == out1).all()                                                                                                                         
                                                                                                                                                      
t0 = benchmark.Timer(stmt='torch.masked_select(x, mask)', globals={'x':x, 'mask':mask})                                                             
t1 = benchmark.Timer(stmt='x.view(-1).index_select(0, mask.view(-1).nonzero().squeeze())', globals={'x':x, 'mask':mask})                            
                                                                                                                                                      
print(t0.timeit(100))                                                                                                                               
print(t1.timeit(100)) 
```
On my system this takes **140.64 us** for ``masked_select``, and **57.36 us** for the equivalent ``index_select``.

If the initial sizes are different but broadcastable and "broadcast_tensors" is used before index_select, a comparable slowdown is observed. This can be easily checked with a few modifications of the script.

If instead of Size([100, 100, 100]), Size([1000 ** 3]) is used, then this prints **59.90 us** for masked_select, and **56.86 us** us for index_select, so index_select is still faster, but the difference is much smaller. 

## Expected behavior

Since masked_select is a more specific operator than the hand made implementation with index select, it should be faster. I couldn't find any use case where masked_select was faster.

## Environment

Collecting environment information...
PyTorch version: 1.9.0.dev20210223
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 8.4.0-3ubuntu2) 8.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.16.3

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: Quadro RTX 6000
Nvidia driver version: 460.32.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.1
[pip3] pytorch-memlab==0.2.2
[pip3] torch==1.9.0.dev20210223
[pip3] torchaudio==0.9.0.dev20210415
[pip3] torchvision==0.10.0.dev20210415
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.0.221             h6bb024c_0  
[conda] mkl                       2020.1                      217  
[conda] mkl-service               2.3.0            py37he8ac12f_0  
[conda] mkl_fft                   1.2.0            py37h23d657b_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] numpy                     1.18.1                   pypi_0    pypi
[conda] numpy-base                1.19.2           py37hfa32c7d_0  
[conda] pytorch-memlab            0.2.2                    pypi_0    pypi
[conda] torch                     1.9.0.dev20210223          pypi_0    pypi
[conda] torchaudio                0.9.0.dev20210415          pypi_0    pypi
[conda] torchvision               0.10.0.dev20210415          pypi_0    pypi

## Additional context
Suggestion: if masked_select is slower in all cases then maybe it should just be implemented under the hood as 
```
x, mask = torch.broadcast_tensors(x, mask)
out = x.reshape(-1).index_select(0, mask.reshape(-1).nonzero().flatten())
```
to prevent unintentional slowdown in code. 
I run into this issue when I replaced some calls to something like the above code with `masked_select` since I 
realized it was essentially doing the same thing, but I suffered a dramatic slowdown after this.

cc @VitalyFedyunin @ngimel