[PIR]Migrate layer_norm into pir (#57324)

0x45f · lxd-cumt · web-flow · commit d29524ffeb9f · 2023-09-19T12:40:03.000+08:00
* Migrate layer_norm into pir

* add default constructor for IrMetaTensor

* fix bugs of generating Op::Build when Op has optional tensor

* Fix layer norm backward

* Enable ut

---------

Co-authored-by: lxd-cumt &lt;1141051934@qq.com&gt;
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -422,13 +422,16 @@ def make_input_stopgradient(op):
         ):
             if not grad_semantic:
                 continue
-            if input.get_defining_op().name() == "builtin.combine":
+            if (
+                input.get_defining_op() is not None
+                and input.get_defining_op().name() == "builtin.combine"
+            ):
                 stop_gradient = make_input_stopgradient(input.get_defining_op())
                 input_grad_stopgradients.append(
                     [info[0] for info in stop_gradient]
                 )
             else:
-                if input in no_grad_set:
+                if input.get_defining_op() is None or input in no_grad_set:
                     input_grad_stopgradients.append([True])
                 else:
                     input_grad_stopgradients.append([False])
@@ -445,7 +448,10 @@ def update_input_grad_map(op, input_grads):
         ):
             if not grad_semantic:
                 continue
-            if input.get_defining_op().name() == "builtin.combine":
+            if (
+                input.get_defining_op() is not None
+                and input.get_defining_op().name() == "builtin.combine"
+            ):
                 update_input_grad_map(input.get_defining_op(), input_grads[i])
             else:
                 input_grad = input_grads[i]
diff --git a/python/paddle/nn/functional/norm.py b/python/paddle/nn/functional/norm.py
@@ -17,7 +17,7 @@
 # TODO: define normalization api
 import paddle
 from paddle import _C_ops, base, in_dynamic_mode
-from paddle.base.framework import in_dygraph_mode
+from paddle.base.framework import in_dygraph_mode, in_dynamic_or_pir_mode
 
 from ...base.data_feeder import check_type, check_variable_and_dtype
 from ...base.layer_helper import LayerHelper
@@ -342,7 +342,7 @@ def layer_norm(
             + str(input_shape)
         )
 
-    if in_dygraph_mode():
+    if in_dynamic_or_pir_mode():
         out = _C_ops.layer_norm(x, weight, bias, epsilon, begin_norm_axis)
         return out
 
diff --git a/test/legacy_test/test_layer_norm_op.py b/test/legacy_test/test_layer_norm_op.py
@@ -142,6 +142,7 @@ def test_check_output(self):
             atol=self.ori_atol,
             rtol=self.ori_rtol,
             check_prim=True,
+            check_new_ir=True,
         )
 
     def test_check_grad(self):
@@ -150,6 +151,7 @@ def test_check_grad(self):
             ['Y'],
             max_relative_error=self.max_relative_error,
             check_prim=True,
+            check_new_ir=True,
         )
 
     def initConfig(self):
@@ -239,6 +241,7 @@ def test_check_output(self):
             atol=self.ori_atol,
             rtol=self.ori_rtol,
             check_prim=True,
+            check_new_ir=True,
         )
 
     def test_check_grad(self):
@@ -248,6 +251,7 @@ def test_check_grad(self):
             ['Y'],
             max_relative_error=self.max_relative_error,
             check_prim=True,
+            check_new_ir=True,
         )
 
     def initConfig(self):
@@ -519,6 +523,7 @@ def initConfig(self):
 class TestLayerNormOp(unittest.TestCase):
     def setUp(self):
         self.use_cudnn = True
+        paddle.enable_static()
 
     def __assert_close(self, tensor, np_array, msg, atol=1e-4):
         np.testing.assert_allclose(
