add gard_api inplace version (#71121)

xiaoguoguo626807 · web-flow · commit af6288318da4 · 2025-02-14T12:43:29.000+08:00
diff --git a/paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py b/paddle/fluid/eager/auto_code_generator/generator/codegen_utils.py
@@ -179,6 +179,19 @@ def IsGradName(string):
     return string.endswith("_grad")
 
 
+def FindRenameForwardName(string):
+    # when op has double_grad and double_grad api has same output of grad api,
+    # double_grad's forward yaml is different from input/output name of grad api
+    # this func find the rename name in double_grad's forward_yaml.
+    # eg acos_grad x_grad -> grad_x, out_grad -> grad_out
+    if string.endswith('_grad'):
+        base_part = string[:-5]
+        transformed_string = 'grad_' + base_part
+        return transformed_string
+    else:
+        raise Exception(f"{string} is not a grad name")
+
+
 def IsPlainTensorType(string):
     plain_tensor_types = ['Tensor&', 'Tensor', 'const Tensor&', 'const Tensor']
     if string in plain_tensor_types:
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -20,6 +20,7 @@
 from codegen_utils import (
     AssertMessage,
     FindForwardName,
+    FindRenameForwardName,
     FunctionGeneratorBase,
     GeneratorBase,
     GetAutoGradMetaName,
@@ -1245,7 +1246,6 @@ def GenerateNodeCreationCodes(self, for_backward=False, is_inplaced=False):
         ) in backward_forward_inputs_map.items():
             is_optional = name in optional_inputs
             is_inplace_input = is_inplaced and name in self.forward_inplace_map
-
             if is_fwd_input:
                 if is_optional:
                     if is_inplace_input:
@@ -1619,7 +1619,7 @@ def GenerateForwardLayoutAutotune(
 
         return layout_logic_str
 
-    def GenerateForwardDefinitionAndDeclaration(self, is_inplaced):
+    def GenerateForwardDefinitionAndDeclaration(self, is_inplaced, grad_flag):
         namespace = self.namespace
         if self.forward_api_name[-1] == '_' and not is_inplaced:
             return
@@ -1637,7 +1637,26 @@ def GenerateForwardDefinitionAndDeclaration(self, is_inplaced):
 
         optional_inputs = self.optional_inputs
         intermediate_outputs = self.intermediate_outputs
-        forward_inplace_map = self.forward_inplace_map if is_inplaced else {}
+        if is_inplaced:
+            if grad_flag and self.grad_api_contents is not None:
+                forward_inplace_map = {}
+                for key, value in self.forward_inplace_map.items():
+                    if key not in self.forward_inputs_position_map:
+                        key = FindRenameForwardName(key)
+                        assert (
+                            key in self.forward_inputs_position_map
+                        ), f"{key} not in {self.forward_api_name} forward_inputs_position_map"
+                    if value not in self.forward_outputs_position_map:
+                        value = FindRenameForwardName(value)
+                        assert (
+                            value in self.forward_outputs_position_map
+                        ), f"{value} not in {self.forward_api_name} forward_outputs_position_map"
+                    forward_inplace_map[key] = value
+                self.forward_inplace_map = forward_inplace_map
+            else:
+                forward_inplace_map = self.forward_inplace_map
+        else:
+            forward_inplace_map = {}
         indent = GetIndent(1)
 
         # Get Function Args
@@ -2186,14 +2205,16 @@ def GenerateForwardDefinitionAndDeclaration(self, is_inplaced):
 
         self.forward_declaration_str += f"TEST_API {returns_type_str} {forward_ad_function_name}({inputs_args_declaration_str});\n"
 
-    def GenerateInplacedForwardDygraphFunctions(self):
+    def GenerateInplacedForwardDygraphFunctions(self, grad_flag):
         # Inplaced Version Dygraph Function Generation
         forward_api_name = self.forward_api_name
         forward_api_contents = self.forward_api_contents
 
         if forward_api_name != "sum" and "inplace" in forward_api_contents:
             # Function Definition and Declaration Generation
-            self.GenerateForwardDefinitionAndDeclaration(is_inplaced=True)
+            self.GenerateForwardDefinitionAndDeclaration(
+                is_inplaced=True, grad_flag=grad_flag
+            )
             self.UpdateCoreOpsInformation(is_inplaced=True)
 
     def UpdateCoreOpsInformation(self, is_inplaced):
@@ -2236,12 +2257,13 @@ def run(self, grad_flag=False):
         ###################
 
         # Definition And Declaration
-        self.GenerateForwardDefinitionAndDeclaration(is_inplaced=False)
+        self.GenerateForwardDefinitionAndDeclaration(
+            is_inplaced=False, grad_flag=grad_flag
+        )
 
         self.UpdateCoreOpsInformation(is_inplaced=False)
 
-        if not grad_flag:
-            self.GenerateInplacedForwardDygraphFunctions()
+        self.GenerateInplacedForwardDygraphFunctions(grad_flag)
 
 
 class DygraphNodeGenerator(DygraphFunctionGeneratorBase):
diff --git a/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py b/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py
@@ -615,7 +615,7 @@ def GeneratePythonCFunction(self):
                 )
                 # Generate Python-C Function Registration
                 self.python_c_function_reg_str = python_c_inplace_func_reg_str
-            elif "backward_op" not in self.forward_api_contents:
+            else:
                 self.python_c_function_str += python_c_inplace_func_str
                 self.python_c_function_declare_str += (
                     python_c_function_declare_str
@@ -670,7 +670,7 @@ def GeneratePythonCFunctions(self):
         for forward_api_content in forward_api_list:
             if "backward_op" in forward_api_content and forward_api_content[
                 "backward_op"
-            ].endswith(('double_grad', 'triple_grad', 'grad_grad', '_grad_')):
+            ].endswith(('double_grad', 'triple_grad', 'grad_grad')):
                 continue
             f_generator = PythonCSingleFunctionGenerator(
                 forward_api_content, namespace
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -254,8 +254,10 @@ def _need_skip(self, op_info, op_name):
             if op_name[:-5] in NO_NEED_GEN_STATIC_ONLY_APIS:
                 return True
         if op_name.endswith("_grad_"):
-            return True
-
+            if op_name.endswith(
+                ("double_grad_", "_grad_grad_", "triple_grad_")
+            ):
+                return True
         return (
             super()._need_skip(op_info, op_name)
             or op_name.endswith('xpu')
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2929,16 +2929,16 @@
   optional : pos_weight
 
 - backward_op : sigmoid_double_grad
-  forward : sigmoid_grad (Tensor out, Tensor fwd_grad_out) -> Tensor(grad_x)
-  args : (Tensor out, Tensor fwd_grad_out, Tensor grad_x_grad)
-  output : Tensor(out_grad), Tensor(fwd_grad_out_grad)
+  forward : sigmoid_grad (Tensor out, Tensor grad_out) -> Tensor(grad_x)
+  args : (Tensor out, Tensor grad_out, Tensor grad_x_grad)
+  output : Tensor(out_grad), Tensor(grad_out_grad)
   infer_meta :
     func : GeneralBinaryGradInferMeta
-    param : [out, fwd_grad_out]
+    param : [out, grad_out]
   kernel :
     func : sigmoid_double_grad
   backward : sigmoid_triple_grad
-  inplace : (grad_x_grad -> fwd_grad_out_grad)
+  inplace : (grad_x_grad -> grad_out_grad)
 
 - backward_op : sigmoid_grad
   forward : sigmoid (Tensor x) -> Tensor(out)
diff --git a/test/legacy_test/test_grad_api.py b/test/legacy_test/test_grad_api.py
@@ -84,6 +84,29 @@ def backward(self, k_grad):
         return x_grad, z_grad, y_grad, w1_grad, b1_grad
 
 
+class dygraph_inplace(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self._param_attr = base.ParamAttr(
+            initializer=paddle.nn.initializer.Constant(value=0.1)
+        )
+        self.w1 = self.create_parameter(
+            attr=self._param_attr, shape=[2, 2], dtype='float32', is_bias=False
+        )
+        self.b1 = self.create_parameter(
+            attr=self._param_attr, shape=[2, 2], dtype='float32', is_bias=False
+        )
+
+    def forward(self, x):
+        self.x = x
+        self.k = paddle.tanh(self.x)
+        return self.k
+
+    def backward(self, k_grad):
+        z_grad = paddle._C_ops.tanh_grad_(self.k, k_grad)
+        return z_grad
+
+
 class TestBaseLayer(unittest.TestCase):
     def test_dy_to_st(self):
         layer = dy_to_st()
@@ -100,9 +123,9 @@ def test_dy_to_st(self):
         x_grad_check = x.grad
         w1_grad_check = layer.w1.grad
         b1_grad_check = layer.b1.grad
-        assert np.allclose(x_grad.numpy(), x_grad_check.numpy())
-        assert np.allclose(w1_grad.numpy(), w1_grad_check.numpy())
-        assert np.allclose(b1_grad.numpy(), b1_grad_check.numpy())
+        np.testing.assert_allclose(x_grad.numpy(), x_grad_check.numpy())
+        np.testing.assert_allclose(w1_grad.numpy(), w1_grad_check.numpy())
+        np.testing.assert_allclose(b1_grad.numpy(), b1_grad_check.numpy())
 
     def test_dygraph(self):
         layer = dygraph()
@@ -117,9 +140,20 @@ def test_dygraph(self):
         x_grad_check = x.grad
         w1_grad_check = layer.w1.grad
         b1_grad_check = layer.b1.grad
-        assert np.allclose(x_grad.numpy(), x_grad_check.numpy())
-        assert np.allclose(w1_grad.numpy(), w1_grad_check.numpy())
-        assert np.allclose(b1_grad.numpy(), b1_grad_check.numpy())
+        np.testing.assert_allclose(x_grad.numpy(), x_grad_check.numpy())
+        np.testing.assert_allclose(w1_grad.numpy(), w1_grad_check.numpy())
+        np.testing.assert_allclose(b1_grad.numpy(), b1_grad_check.numpy())
+
+    def test_dygraph_inplace(self):
+        layer = dygraph_inplace()
+        x = paddle.to_tensor([[1.0, 2.0], [3.0, 4.0]], dtype='float32')
+        out_grad = paddle.to_tensor([[1.0, 1.0], [1.0, 1.0]], dtype='float32')
+        x.stop_gradient = False
+
+        out = layer(x)
+        x_grad = layer.backward(out_grad)
+
+        np.testing.assert_allclose(out_grad.numpy(), x_grad.numpy())
 
 
 if __name__ == '__main__':
