support custom vjp trait

paddle/fluid/pir/dialect/op_generator/op_gen.py

@@ -15,6 +15,8 @@
 import argparse
 import logging
 import os
+import pathlib
+import sys
 
 import yaml
 from op_build_gen import gen_build_func_str
@@ -30,6 +32,12 @@
     vjp_interface_implementation_gen_op_list,
 )
 
+# import from paddle/fluid/primitive/code_gen/gen.py
+sys.path.append(
+    str(pathlib.Path(__file__).resolve().parents[3] / 'primitive/codegen')
+)
+import gen as vjp_gen
+
 # =====================================
 # String Template for h file code gen
 # =====================================
@@ -54,6 +62,7 @@
 #include "paddle/fluid/pir/dialect/operator/interface/infermeta.h"
 #include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
 #include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
+#include "paddle/fluid/pir/dialect/operator/trait/custom_vjp.h"
 #include "paddle/fluid/framework/infershape_utils.h"
 #include "paddle/phi/core/infermeta_utils.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
@@ -818,6 +827,12 @@ def OpGenerator(
     ops_declare_list = []  # all op class declare store in this list
     ops_defined_list = []  # all op class defined store in this list
     ops_vjp_defined_list = []  # all op vjp static interface defination
+
+    # (4) parse name of ops which have custom vjp rules
+    custom_vjp_op_name_list = []
+    for custom_vjp in vjp_gen.CUSTOM_VJP:
+        custom_vjp_op_name_list.append(custom_vjp[:-5])  # cut _grad
+
     for key, op_info in op_info_items.items():
         # get op inputs info
         op_input_name_list = op_info.input_name_list
@@ -873,6 +888,10 @@ def OpGenerator(
             op_interfaces += ["paddle::dialect::VjpInterface"]
         exclusive_interface_str = gen_exclusive_interface_str(op_info)
 
+        # if op has custom vjp rule, then append a CustomVjpTrait to it
+        if op_info.op_phi_name[0] in custom_vjp_op_name_list:
+            op_traits += ["paddle::dialect::CustomVjpTrait"]
+
         # check op inputs and mutable_attributes grad semantics
         input_grad_semantics = get_input_grad_semantic(op_info, op_info_items)
         mutable_attribute_grad_semantics = get_mutable_attribute_grad_semantic(

paddle/fluid/pir/dialect/operator/trait/custom_vjp.h

@@ -0,0 +1,38 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*
+Custom VJP stands for manually implemented backward rules for composite
+operators. CustomVjpTrait will be added for those composite operators that
+defines custom vjp rules. Finally, by calling has_custom_vjp(op), users can
+check whether an operator has a CustomVjpTrait, and thus check whether a custom
+vjp rule is defined for that operator.
+*/
+
+#pragma once
+
+#include "paddle/pir/core/op_base.h"
+
+namespace paddle {
+namespace dialect {
+class CustomVjpTrait : public pir::OpTraitBase<CustomVjpTrait> {
+ public:
+  explicit CustomVjpTrait(pir::Operation *op)
+      : pir::OpTraitBase<CustomVjpTrait>(op) {}
+};
+
+}  // namespace dialect
+}  // namespace paddle
+
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::CustomVjpTrait)

paddle/fluid/pir/dialect/operator/trait/trait.cc

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/fluid/pir/dialect/operator/trait/custom_vjp.h"
 #include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::InplaceTrait)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::CustomVjpTrait)

paddle/fluid/primitive/codegen/gen.py

@@ -31,7 +31,7 @@
 
 # import from paddle/fluid/pir/dialect/op_generator/api_gen.py
 sys.path.append(
-    str(pathlib.Path(__file__).resolve().parents[2] / 'ir/dialect/op_generator')
+    str(pathlib.Path(__file__).resolve().parents[2] / 'pir/dialect/op_generator')
 )
 
 # fmt: on
@@ -67,7 +67,12 @@
     'slice_double_grad',
     'layer_norm_grad',
 ]
-VJP_COMPS = ['divide_grad', 'sum_grad', 'gelu_grad']
+
+
+PRIM_VJP = ['divide_grad', 'sum_grad']  # vjp list of primitive op
+CUSTOM_VJP = ['gelu_grad']  # custom vjp list of composite op
+VJP_COMPS = PRIM_VJP + CUSTOM_VJP
+
 BACKENDS = [
     'add_n',
     'mean',

paddle/fluid/pybind/pybind.cc

@@ -196,6 +196,7 @@ limitations under the License. */
 #include "paddle/fluid/eager/nan_inf_utils.h"
 #include "paddle/fluid/imperative/layout_autotune.h"
 #include "paddle/fluid/pir/dialect/operator/interface/vjp.h"
+#include "paddle/fluid/pir/dialect/operator/trait/custom_vjp.h"
 #include "paddle/fluid/prim/utils/eager/eager_tensor_operants.h"
 #include "paddle/fluid/prim/utils/static/static_tensor_operants.h"
 #include "paddle/fluid/pybind/eager_utils.h"
@@ -751,6 +752,21 @@ void BindVjp(pybind11::module *m) {
     if (vjp_interface_impl == nullptr) return false;
     return true;
   });
+
+  m->def(
+      "has_custom_vjp",
+      [](pir::Operation &op) -> py::bool_ {
+        return op.info().HasTrait<paddle::dialect::CustomVjpTrait>();
+      },
+      R"DOC(
+           Return whether an op has custom vjp rules.
+
+           Args:
+               op (pir::Operation): op to be checked
+
+           Returns:
+               out (bool): True means that the op has custom vjp rules, False means it does not.
+           )DOC");
 }
 PYBIND11_MODULE(libpaddle, m) {
   BindImperative(&m);

test/prim/new_ir_prim/test_custom_vjp_trait.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import paddle
+from paddle import ir, nn
+from paddle.base.core import has_custom_vjp
+
+paddle.enable_static()
+
+
+def get_gelu_program_new_ir():
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x = paddle.static.data('x', [2, 3, 3], dtype='float32')
+        net = nn.GELU()
+        out = net(x)
+    newir_program = ir.translate_to_new_ir(main_program.desc)
+    return newir_program
+
+
+def get_multiply_program_new_ir():
+    main_program, start_program = (
+        paddle.static.Program(),
+        paddle.static.Program(),
+    )
+    with paddle.static.program_guard(main_program, start_program):
+        x = paddle.static.data('x', [2, 3, 3], dtype='float32')
+        y = paddle.static.data('y', [2, 3, 3], dtype='float32')
+        out = paddle.multiply(x, y)
+    newir_program = ir.translate_to_new_ir(main_program.desc)
+    return newir_program
+
+
+class TestCustomVjpTrait(unittest.TestCase):
+    def test_gelu_op_custom_vjp_trait(self):
+        newir_program = get_gelu_program_new_ir()
+        op = newir_program.global_block().ops[-1]
+        self.assertEqual(op.name(), "pd_op.gelu")
+        self.assertEqual(has_custom_vjp(op), True)
+
+    def test_multiply_op_custom_vjp_trait(self):
+        newir_program = get_multiply_program_new_ir()
+        op = newir_program.global_block().ops[-1]
+        self.assertEqual(op.name(), "pd_op.multiply")
+        self.assertEqual(has_custom_vjp(op), False)
+
+
+if __name__ == "__main__":
+    unittest.main()