Repository: pytorch
Issue: torch.nn.LayerNorm is very slow on GPU (much slower than a custom LayerNorm version in the ConvNext model)
Labels: module: nn,module: cuda,triaged

### 🐛 Describe the bug

I found that for a (B, C, H, W) tensor, `nn.LayerNorm` is much slower (0.088s w/o permute and 0.14s with necessary permute) than the custom LayerNorm version for the ConvNext model (0.04s) `https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119`, while both results are equivalent (up to numerical errors).

```
import numpy as np
import torch
import torch.nn as nn

import time

ln_time = 0.0
ln_ln_time = 0.0
ln_permute1_time = 0.0
ln_permute2_time = 0.0
custom_ln_time = 0.0
bn_time = 0.0

T = 8
device = 'cuda:0'
size=(16,256,512,128)
weight = torch.rand(size[1]).to(device)
bias = torch.randn(size[1]).to(device)

class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. 
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with 
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs 
    with shape (batch_size, channels, height, width).
    """
    def __init__(self, normalized_shape, eps=1e-8, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(weight.clone())
        self.bias = nn.Parameter(bias.clone())
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError 
        self.normalized_shape = (normalized_shape, )
    
    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


ln = nn.LayerNorm(size[1], eps=1e-8).to(device)
ln.weight = nn.Parameter(weight.clone())
ln.bias = nn.Parameter(bias.clone())
ln_custom = LayerNorm(size[1], data_format="channels_first").to(device)

a = torch.randn(size).to(device)
b = a.permute(0,2,3,1).contiguous()
ln_a = ln(b)
ln_a = ln_a.permute(0,3,1,2).contiguous()
ln_custom_a = ln_custom(a)
print("|ln - ln_custom| max", torch.max(torch.abs(ln_a - ln_custom_a)))

del a,b



ln_custom = LayerNorm(size[1], data_format="channels_first").to(device)
for t in range(T):
    print(t)
    a = torch.randn(size).to(device)
    torch.cuda.synchronize()
    tt = time.time()
    b=ln_custom(a)
    torch.cuda.synchronize()
    custom_ln_time += time.time() - tt
    del b
print("custom_ln_time", custom_ln_time / T)

ln = nn.LayerNorm(size[1]).to(device)
ln.weight = nn.Parameter(weight.clone())
ln.bias = nn.Parameter(bias.clone())
for t in range(T):
    print(t)
    a = torch.randn(size).to(device)
    torch.cuda.synchronize()

    tt = time.time()
    b=a.permute(0,2,3,1).contiguous()
    torch.cuda.synchronize()
    ln_permute1_time += time.time() - tt

    tt=time.time()
    b=ln(b)
    torch.cuda.synchronize()
    ln_permute2_time += time.time() - tt

    tt=time.time()
    b=b.permute(0,3,1,2).contiguous()
    torch.cuda.synchronize()
    ln_ln_time += time.time() - tt
    del b

print("pytorch ln_ln_time", ln_ln_time / T, "ln_permute1_time", ln_permute1_time / T, "ln_permute2_time", ln_permute2_time / T)
print("pytorch ln_total_time", (ln_ln_time + ln_permute1_time + ln_permute2_time) / T)
```

Output:

```
|ln - ln_custom| max tensor(1.4305e-06, device='cuda:0', grad_fn=<MaxBackward1>)
0
1
2
3
4
5
6
7
custom_ln_time 0.0418052077293396
0
1
2
3
4
5
6
7
pytorch ln_ln_time 0.08832478523254395 ln_permute1_time 0.017529428005218506 ln_permute2_time 0.03777831792831421
pytorch ln_total_time 0.14363253116607666
```

### Versions

PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.13.0-39-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.3.58
GPU models and configuration: 
GPU 0: NVIDIA GeForce GTX 1080 Ti
GPU 1: NVIDIA GeForce GTX 1080 Ti

Nvidia driver version: 470.103.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.3
[pip3] numpy-quaternion==2022.4.1
[pip3] pytorch3d==0.3.0
[pip3] torch==1.11.0
[pip3] torchaudio==0.11.0
[pip3] torchsparse==1.4.0.dev0
[pip3] torchvision==0.12.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.1            py38hd3c417c_0  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] numpy                     1.22.3                   pypi_0    pypi
[conda] numpy-base                1.21.2           py38h79a1101_0  
[conda] numpy-quaternion          2022.4.1                 pypi_0    pypi
[conda] pytorch                   1.11.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] pytorch3d                 0.3.0                    pypi_0    pypi
[conda] torchaudio                0.11.0               py38_cu113    pytorch
[conda] torchsparse               1.4.0.dev0               pypi_0    pypi
[conda] torchvision               0.12.0               py38_cu113    pytorch

cc @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @ngimel