enable dil_layer_norm

Author: jiayisunx

scripts/cpu/gen-dense-cpu-ops.py

@@ -73,8 +73,8 @@
     'aten::view(Tensor(a) self, int[] size) -> Tensor(a)',
     'aten::index_select(Tensor self, int dim, Tensor index) -> Tensor',
     'aten::_unsafe_view(Tensor self, int[] size) -> Tensor',
-    #'aten::native_layer_norm(Tensor input, Tensor? weight, Tensor? bias, int M, int N, float eps) -> (Tensor, Tensor, Tensor)',
-    #'aten::native_layer_norm_backward(Tensor grad_out, Tensor input, Tensor mean, Tensor rstd, Tensor? weight, int M, int N, bool[3] output_mask) -> (Tensor, Tensor, Tensor)'
+    'aten::native_layer_norm(Tensor input, Tensor? weight, Tensor? bias, int M, int N, float eps) -> (Tensor, Tensor, Tensor)',
+    'aten::native_layer_norm_backward(Tensor grad_out, Tensor input, Tensor mean, Tensor rstd, Tensor? weight, int M, int N, bool[3] output_mask) -> (Tensor, Tensor, Tensor)'
 ]
 
 _FN_IPEX_FUNCS_WITH_SIMPLE_ATEN_SIG = [

tests/cpu/test_lazy_reorder.py

@@ -998,7 +998,11 @@ def test_layer_norm(self):
             input = torch.randn(2, 5, 10, 10, dtype=torch.float32)
             input_dpcpp=input.to(device=device)
             m = torch.nn.LayerNorm([10, 10])
-            self.assertEqual(m(input), m(input_dpcpp))
+            m_dpcpp = copy.deepcopy(m).to(device=device)
+            output = m(input)
+            output_dpcpp = m_dpcpp(input_dpcpp)
+            self.assertTrue(ipex.core.is_dil_tensor(output_dpcpp))
+            self.assertEqual(output, output_dpcpp)
 
     def test_layer_norm_backward(self):
         with AutoDNNL(True):

torch_ipex/csrc/cpu/DevOPs.cpp

@@ -1937,6 +1937,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> AtenIpexCPUDev::dil_native_layer_
     double eps) {
   DEBUG("AtenIpexCPUDev::dil_native_layer_norm\n");
   CHECK_DNNL_OP_PRE_COND(X);
+  //It's a temporary solution to fall back to fp32 since bf16 layer_norm is not ready for dnnl path now.
+  dbl::comm::reorder_to_dtype(X, at::kFloat);
   dil::tensor x = dbl::comm::try_gen_dil_tensor(X);
   const dil::tensor scale = dbl::comm::try_gen_dil_tensor(gamma);
   const dil::tensor shift = dbl::comm::try_gen_dil_tensor(beta);
@@ -1976,6 +1978,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> AtenIpexCPUDev::dil_native_layer_
   DEBUG("AtenIpexCPUDev::dil_native_layer_norm_backward\n");
   CHECK_DNNL_OP_PRE_COND(dY);
   CHECK_DNNL_OP_PRE_COND(X);
+  //it's a temporary solution to fall back to fp32 since bf16 layer_norm is not ready for dnnl path now.
+  dbl::comm::reorder_to_dtype(dY, at::kFloat);
+  dbl::comm::reorder_to_dtype(X, at::kFloat);
   dil::tensor dy = dbl::comm::try_gen_dil_tensor(dY);
   dil::tensor x = dbl::comm::try_gen_dil_tensor(X);
   dil::tensor m = dbl::comm::try_gen_dil_tensor(mean);