Remove redundant memory copy for HF multi-attention submodule for cpu path using MKL prepack
Labels: triaged, bug, oncall: pt2, module: cpu inductor

### 🐛 Describe the bug

There has a redundant memory copy for HF multi-attention submodule for CPU path.

### Error logs

output code:

```
from ctypes import c_void_p, c_long
import torch
import random
from torch import empty_strided, as_strided, device
from torch._inductor.codecache import AsyncCompile

aten = torch.ops.aten
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
async_compile = AsyncCompile()


kernel_cpp_0 = async_compile.cpp('''
#include "/tmp/torchinductor_xiaobing/77/c7773nj5pwikpmm2pwa62rcudlf7p3if7eyqb5k4sjsvewwje4le.h"
extern "C" void kernel(float* __restrict__ in_out_ptr0,
                       float* __restrict__ in_out_ptr1)
{
    for(long i0=0; i0<384; i0+=1)
    {
        auto tmp0 = at::vec::Vectorized<float>::loadu(in_out_ptr0 + 16*i0);
        tmp0.store(in_out_ptr0 + 16*i0);
    }
    #pragma omp simd simdlen(8)
    for(long i0=6144; i0<6144; i0+=1)
    {
        auto tmp0 = in_out_ptr0[i0];
        in_out_ptr0[i0] = tmp0;
    }
    for(long i0=0; i0<384; i0+=1)
    {
        auto tmp0 = at::vec::Vectorized<float>::loadu(in_out_ptr1 + 16*i0);
        tmp0.store(in_out_ptr1 + 16*i0);
    }
    #pragma omp simd simdlen(8)
    for(long i0=6144; i0<6144; i0+=1)
    {
        auto tmp0 = in_out_ptr1[i0];
        in_out_ptr1[i0] = tmp0;
    }
}
''')


async_compile.wait(globals())
del async_compile

def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1 = args
    args.clear()
    buf0 = torch.ops.mkl._mkl_linear(arg9_1, arg2_1, arg0_1, arg1_1, 8)
    del arg0_1
    del arg1_1
    del arg2_1
    buf1 = torch.ops.mkl._mkl_linear(arg9_1, arg5_1, arg3_1, arg4_1, 8)
    del arg3_1
    del arg4_1
    del arg5_1
    del arg9_1
    buf3 = as_strided(buf0, (12, 8, 64), (64, 768, 1)); del buf0  # reuse
    buf4 = as_strided(buf1, (12, 64, 8), (64, 1, 768)); del buf1  # reuse
    kernel_cpp_0(c_void_p(buf3.data_ptr()), c_void_p(buf4.data_ptr()))
    buf5 = empty_strided((12, 8, 8), (64, 8, 1), device='cpu', dtype=torch.float32)
    aten.bmm.out(buf3, buf4, out=buf5)
    return (as_strided(buf5, (1, 12, 8, 8), (768, 64, 8, 1)), )


if __name__ == "__main__":
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    arg0_1 = rand_strided((768, 128), (128, 1), device='cpu', dtype=torch.float32)
    arg1_1 = rand_strided((768, ), (1, ), device='cpu', dtype=torch.float32)
    arg2_1 = rand_strided((2310369, 1), (1, 0), device='cpu', dtype=torch.float32)
    arg3_1 = rand_strided((768, 128), (128, 1), device='cpu', dtype=torch.float32)
    arg4_1 = rand_strided((768, ), (1, ), device='cpu', dtype=torch.float32)
    arg5_1 = rand_strided((2310369, 1), (1, 0), device='cpu', dtype=torch.float32)
    arg6_1 = rand_strided((768, 128), (128, 1), device='cpu', dtype=torch.float32)
    arg7_1 = rand_strided((768, ), (1, ), device='cpu', dtype=torch.float32)
    arg8_1 = rand_strided((2310369, 1), (1, 0), device='cpu', dtype=torch.float32)
    arg9_1 = rand_strided((1, 8, 128), (1024, 128, 1), device='cpu', dtype=torch.float32)
    print_performance(lambda: call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1]))
```

### Minified repro

```
import torch._dynamo
import torch.fx.experimental.optimization as optimization
import copy

from typing import Dict, List, Optional

import time
import torch.profiler as profiler
from torch.fx import symbolic_trace
from torch._inductor import config
config.debug = True

class Model(torch.nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.num_attention_heads = 12
        self.attention_head_size = 64
        self.query = torch.nn.Linear(128, 768)
        self.key = torch.nn.Linear(128, 768)
        self.value = torch.nn.Linear(128, 768)


    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden_states):
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        #value_layer = self.transpose_for_scores(mixed_value_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        return attention_scores


x = torch.randn(1, 8, 128)

model = Model().eval()
y = model(x)

with torch.no_grad():
    opt_model = torch._dynamo.optimize('inductor')(model)

with torch.no_grad():
    for i in range(2):
        y1 = opt_model(x)

```

cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh