enable bf16 layernorm

Author: jiayisunx

tests/cpu/test_bf16_lazy_reorder.py

@@ -486,6 +486,90 @@ def test_batch_norm3d_backward(self):
                 self.assertTrue(ipex.core.is_bf16_dil_tensor(x_auto_mix_bf16.grad))
                 self.assertEqual(x_man_bf16.grad.float(), x_auto_mix_bf16.grad)
 
+class TestLayerNorm(TestCase):
+    def test_layer_norm(self):
+        rand_seed = int(get_rand_seed())
+        print("{} rand sed: {}".format(sys._getframe().f_code.co_name, rand_seed))
+
+        x_cpu, x_auto_mix_inference, x_auto_mix_train, x_man_bf16, x_auto_mix_train_bf16 = _gen_tensor(
+            rand_seed, (2, 5, 10, 10))
+
+        op_cpu, op_auto_mix_inference, op_auto_mix_train, op_man_bf16, op_auto_mix_train_bf16 = _gen_op(
+            rand_seed, torch.nn.LayerNorm([10, 10]), is_bn=True)
+
+        ref_cpu = op_cpu(x_cpu)
+        with AutoDNNL(True), AutoMixPrecision(False):
+            res_bf16 = op_man_bf16(x_man_bf16)
+            self.assertEqual(res_bf16.dtype, torch.bfloat16)
+
+            # FW inference
+            with AutoMixPrecision(True, train=False):
+                self.assertEqual(x_auto_mix_inference.dtype, torch.float)
+                self.assertFalse(ipex.core.is_bf16_dil_tensor(x_auto_mix_inference))
+                res_auto_mix_inference = op_auto_mix_inference(x_auto_mix_inference)
+                self.assertEqual(res_auto_mix_inference.dtype, torch.float)
+                self.assertEqual(x_auto_mix_inference.dtype, torch.float)
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(res_auto_mix_inference))
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(x_auto_mix_inference))
+                self.assertEqual(res_bf16.float(), res_auto_mix_inference)
+
+            # FW train (input is not bf16 dil tensor)
+            with AutoMixPrecision(True, train=True):
+                self.assertEqual(x_auto_mix_train.dtype, torch.float)
+                self.assertFalse(ipex.core.is_bf16_dil_tensor(x_auto_mix_train))
+                res_auto_mix_train = op_auto_mix_train(x_auto_mix_train)
+                self.assertEqual(res_auto_mix_train.dtype, torch.float)
+                self.assertEqual(x_auto_mix_train.dtype, torch.float)
+                self.assertFalse(ipex.core.is_bf16_dil_tensor(res_auto_mix_train))
+                self.assertFalse(ipex.core.is_bf16_dil_tensor(x_auto_mix_train))
+                self.assertEqual(ref_cpu, res_auto_mix_train)
+
+            # FW train (input is bf16 dil tensor)
+            with AutoMixPrecision(True, train=True):
+                self.assertEqual(x_auto_mix_train_bf16.dtype, torch.float)
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(x_auto_mix_train_bf16))
+                res_auto_mix_train_bf16 = op_auto_mix_train_bf16(x_auto_mix_train_bf16)
+                self.assertEqual(res_auto_mix_train_bf16.dtype, torch.float)
+                self.assertEqual(x_auto_mix_train_bf16.dtype, torch.float)
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(res_auto_mix_train_bf16))
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(x_auto_mix_train_bf16))
+                self.assertEqual(res_bf16.float(), res_auto_mix_train_bf16)
+
+    def test_layer_norm_backward(self):
+        rand_seed = int(get_rand_seed())
+        print("{} rand sed: {}".format(sys._getframe().f_code.co_name, rand_seed))
+        x_cpu, _, x_auto_mix, x_man_bf16, x_auto_mix_bf16 = _gen_tensor(rand_seed, (2, 5, 10, 10), is_forward=False)
+
+        op_cpu, _, op_auto_mix, op_man_bf16, op_auto_mix_bf16 = _gen_op(rand_seed, torch.nn.LayerNorm([10, 10]), is_bn=True, is_forward=False)
+
+        out_cpu = op_cpu(x_cpu).sum()
+        out_cpu.backward()
+        with AutoDNNL(True), AutoMixPrecision(False, train=True):
+            out_man_bf16 = op_man_bf16(x_man_bf16).sum()
+            out_man_bf16.backward()
+            self.assertEqual(x_man_bf16.grad.dtype, torch.bfloat16)
+            self.assertEqual(x_cpu.grad.bfloat16().float(), x_man_bf16.grad, 1e-2)
+
+            # BW train (input is not bf16 dil tensor)
+            with AutoMixPrecision(True, train=True):
+                self.assertEqual(x_auto_mix.dtype, torch.float)
+                self.assertFalse(ipex.core.is_bf16_dil_tensor(x_auto_mix))
+                out_auto_mix = op_auto_mix(x_auto_mix).sum()
+                out_auto_mix.backward()
+                self.assertEqual(x_auto_mix.grad.dtype, torch.float)
+                self.assertFalse(ipex.core.is_bf16_dil_tensor(x_auto_mix.grad))
+                self.assertEqual(x_cpu.grad, x_auto_mix.grad)
+
+             # BW train (input is bf16 dil tensor)
+            with AutoMixPrecision(True, train=True):
+                self.assertEqual(x_auto_mix_bf16.dtype, torch.float)
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(x_auto_mix_bf16))
+                out_auto_mix_bf16 = op_auto_mix_bf16(x_auto_mix_bf16).sum()
+                out_auto_mix_bf16.backward()
+                self.assertEqual(x_auto_mix_bf16.grad.dtype, torch.float)
+                self.assertTrue(ipex.core.is_bf16_dil_tensor(x_auto_mix_bf16.grad))
+                self.assertEqual(x_man_bf16.grad.float(), x_auto_mix_bf16.grad)
+
 class TestRelu(TestCase):
     def test_relu(self):
         rand_seed = int(get_rand_seed())

torch_ipex/csrc/cpu/DevOPs.cpp

@@ -2467,8 +2467,7 @@ 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);
+  dbl::comm::reorder_to_bf16_for_mix_prec(X, true);
   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);
@@ -2508,9 +2507,8 @@ 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);
+  dbl::comm::reorder_to_bf16_for_mix_prec(dY, true);
+  dbl::comm::reorder_to_bf16_for_mix_prec(X, true);
   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);