Optimize LayerNorm with explicit vectorization using Vec256 (#29104)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/29104
We would like to provide the vectorized implementation for layer norm. This PR reuses https://github.com/pytorch/pytorch/pull/23349.
Test Plan:
buck test mode/dev-nosan //caffe2/test:nn -- "LayerNorm"
buck test mode/dev-nosan //caffe2/test:nn -- "test_LayerNorm_1d_no_elementwise_affine_eval"
python run_test.py -i nn -- TestNN.test_LayerNorm_1d_no_elementwise_affine_eval
Differential Revision: D18293522
fbshipit-source-id: f4cfed6e62bac1b43ee00c32b495ecc836bd9ec5
