[Inductor] Vectorize Embedding Lookup in CPP
Labels: triaged, oncall: pt2, module: cpu inductor

We explicitly disable vectorization on indirect indexing right now:
https://github.com/pytorch/pytorch/blob/7c353eb39559f2c8897a0580700dd0a6f943d34f/torch/_inductor/codegen/cpp.py#L929

It is a conservative design choice and would miss some cases like embedding with large enough vector length. Below is the generated code from DistilBertForQuestionAnswering as an example:
```c++
#include "/tmp/torchinductor_jgong5/uh/cuhop6zhqdvygrhcsy2duu7yusyui3ja4mmjk5kccnhxojwtrwhg.h"
extern "C" void kernel(const long* __restrict__ in_ptr0,
                       const float* __restrict__ in_ptr1,
                       const long* __restrict__ in_ptr2,
                       const float* __restrict__ in_ptr3,
                       const float* __restrict__ in_ptr4,
                       const float* __restrict__ in_ptr5,
                       float* __restrict__ out_ptr0,
                       float* __restrict__ out_ptr1,
                       float* __restrict__ out_ptr2,
                       float* __restrict__ out_ptr3)
{
    #pragma GCC ivdep
    for(long i0=0; i0<128; i0+=1)
    {
        {
            {
                float tmp5 = 0;
                for(long i1=0; i1<768; i1+=1)
                {
                    {
                        auto tmp0 = in_ptr0[i0];
                        auto tmp2 = in_ptr2[i0];
                        auto tmp1 = in_ptr1[i1 + (768*tmp0)];
                        auto tmp3 = in_ptr3[i1 + (768*tmp2)];
                        auto tmp4 = tmp1 + tmp3;
                        tmp5 += tmp4;
                    }
                }
                out_ptr0[i0] = tmp5;
            }
        }
    }
    #pragma GCC ivdep
    for(long i0=0; i0<128; i0+=1)
    {
        {
            {
                float tmp10 = 0;
                float tmp11 = 0;
                for(long i1=0; i1<768; i1+=1)
                {
                    {
                        auto tmp0 = in_ptr0[i0];
                        auto tmp2 = in_ptr2[i0];
                        auto tmp5 = out_ptr0[i0];
                        auto tmp1 = in_ptr1[i1 + (768*tmp0)];
                        auto tmp3 = in_ptr3[i1 + (768*tmp2)];
                        auto tmp4 = tmp1 + tmp3;
                        auto tmp6 = static_cast<float>(768);
                        auto tmp7 = tmp5 / tmp6;
                        auto tmp8 = tmp4 - tmp7;
                        auto tmp9 = tmp8 * tmp8;
                        tmp10 += tmp9;
                        tmp11 += tmp4;
                    }
                }
                out_ptr1[i0] = tmp10;
                out_ptr2[i0] = tmp11;
            }
        }
    }
    #pragma GCC ivdep
    for(long i0=0; i0<128; i0+=1)
    {
        #pragma GCC ivdep
        for(long i1=0; i1<768; i1+=1)
        {
            {
                {
                    auto tmp0 = in_ptr0[i0];
                    auto tmp2 = in_ptr2[i0];
                    auto tmp5 = out_ptr2[i0];
                    auto tmp9 = out_ptr1[i0];
                    auto tmp16 = in_ptr4[i1];
                    auto tmp18 = in_ptr5[i1];
                    auto tmp1 = in_ptr1[i1 + (768*tmp0)];
                    auto tmp3 = in_ptr3[i1 + (768*tmp2)];
                    auto tmp4 = tmp1 + tmp3;
                    auto tmp6 = static_cast<float>(768);
                    auto tmp7 = tmp5 / tmp6;
                    auto tmp8 = tmp4 - tmp7;
                    auto tmp10 = tmp9 / tmp6;
                    auto tmp11 = static_cast<float>(1e-12);
                    auto tmp12 = tmp10 + tmp11;
                    auto tmp13 = std::sqrt(tmp12);
                    auto tmp14 = 1 / tmp13;
                    auto tmp15 = tmp8 * tmp14;
                    auto tmp17 = tmp15 * tmp16;
                    auto tmp19 = tmp17 + tmp18;
                    out_ptr3[i1 + (768*i0)] = tmp19;
                }
            }
        }
    }
}
```

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