[Prim] enable whitelist and blacklist for custom_vjp

python/paddle/fluid/core.py

@@ -446,6 +446,10 @@ def __sync_stat_with_flag(flag):
         )
 
 
+def _is_all_prim_enabled():
+    return _is_fwd_prim_enabled() and _is_bwd_prim_enabled()
+
+
 # Alert!!! This method is only for test coveraget, user should never use it directly, this may cause serious system errors.
 def _test_use_sync(value):
     __sync_stat_with_flag(value)

python/paddle/fluid/tests/unittests/autograd/test_primapi.py

@@ -1046,14 +1046,51 @@ def tearDown(self):
         paddle.disable_static()
 
     @param.parameterized.expand((({'dropout'},),))
-    def test_exclude(self, exclude):
+    def test_blacklist(self, blacklist):
         program = paddle.static.Program()
         with paddle.static.program_guard(program):
-            x = paddle.rand((1,))
-            y = paddle.nn.functional.dropout(x)
-        primapi.to_prim(program.blocks, exclude)
+            paddle.nn.functional.softmax(
+                paddle.nn.functional.dropout(paddle.rand((1,)))
+            )
+        primapi.to_prim(program.blocks, blacklist=blacklist)
+        ops = tuple(op.type for op in program.block(0).ops)
+        self.assertTrue(all(tuple(op in ops for op in blacklist)))
+
+    @param.parameterized.expand((({'dropout'},),))
+    def test_whitelist(self, whitelist):
+        program = paddle.static.Program()
+        with paddle.static.program_guard(program):
+            paddle.nn.functional.softmax(
+                paddle.nn.functional.dropout(paddle.rand((1,)))
+            )
+        primapi.to_prim(program.blocks, whitelist=whitelist)
         ops = tuple(op.type for op in program.block(0).ops)
-        self.assertTrue(all(tuple(op in ops for op in exclude)))
+        self.assertTrue(all(tuple(op not in ops for op in whitelist)))
+
+    @param.parameterized.expand((({'softmax'}, {'softmax', 'dropout'}),))
+    def test_both_not_empty(self, blacklist, whitelist):
+        program = paddle.static.Program()
+        with paddle.static.program_guard(program):
+            paddle.nn.functional.softmax(
+                paddle.nn.functional.dropout(paddle.rand((1,)))
+            )
+        primapi.to_prim(
+            program.blocks, blacklist=blacklist, whitelist=whitelist
+        )
+        ops = tuple(op.type for op in program.block(0).ops)
+        self.assertTrue(all(tuple(op in ops for op in blacklist)))
+
+    @param.parameterized.expand(((('dropout',), 'softmax'),))
+    def test_type_error(self, blacklist, whitelist):
+        program = paddle.static.Program()
+        with paddle.static.program_guard(program):
+            paddle.nn.functional.softmax(
+                paddle.nn.functional.dropout(paddle.rand((1,)))
+            )
+        with self.assertRaises(TypeError):
+            primapi.to_prim(
+                program.blocks, blacklist=blacklist, whitelist=whitelist
+            )
 
 
 if __name__ == '__main__':

python/paddle/fluid/tests/unittests/dygraph_to_static/test_cinn_prim.py

@@ -77,6 +77,8 @@ def train(self, use_prim):
     def check_prim(self, net, use_prim):
         if not use_prim:
             return
+        # Please use PartialProgramLayer(second output parameter of get_concrete_program) rather than
+        # main_program here, as main_program is original program before to_prim.
         fwd_ops = [
             op.type
             for op in net.forward.get_concrete_program(self.x)[1]

python/paddle/fluid/tests/unittests/dygraph_to_static/test_cinn_prim_gelu.py

@@ -93,6 +93,8 @@ def _train(self, use_prim, approximate, data):
     def check_prim(self, net, use_prim):
         if not use_prim:
             return
+        # Please use PartialProgramLayer(second output parameter of get_concrete_program) rather than
+        # main_program here, as main_program is original program before to_prim.
         fwd_ops = [
             op.type
             for op in net.forward.get_concrete_program(self.x)[1]

python/paddle/fluid/tests/unittests/dygraph_to_static/test_cinn_prim_layer_norm.py

@@ -89,6 +89,8 @@ def train(self, use_prim):
     def check_prim(self, net, use_prim):
         if not use_prim:
             return
+        # Please use PartialProgramLayer(second output parameter of get_concrete_program) rather than
+        # main_program here, as main_program is original program before to_prim.
         fwd_ops = [
             op.type
             for op in net.forward.get_concrete_program(self.x, self.w, self.b)[

python/paddle/fluid/tests/unittests/dygraph_to_static/test_cinn_prim_mean.py

@@ -100,6 +100,8 @@ def _train(self, use_prim, data, axis, keep_dim):
     def check_prim(self, net, use_prim):
         if not use_prim:
             return
+        # Please use PartialProgramLayer(second output parameter of get_concrete_program) rather than
+        # main_program here, as main_program is original program before to_prim.
         fwd_ops = [
             op.type
             for op in net.forward.get_concrete_program(self.x)[1]

python/paddle/fluid/tests/unittests/dygraph_to_static/test_partial_program_hook.py

@@ -0,0 +1,71 @@
+# 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.fluid import core
+from paddle.jit.dy2static import partial_program, program_translator
+
+
+class TestPartiaProgramLayerHook(unittest.TestCase):
+    def setUp(self):
+        self._hook = partial_program.PartialProgramLayerHook()
+
+    def test_before_append_backward(self):
+        self.assertIsNone(self._hook.before_append_backward(None))
+
+    def test_after_append_backward(self):
+        self.assertIsNone(self._hook.after_append_backward(None, 0))
+
+    def test_after_infer(self):
+        self.assertIsNone(self._hook.after_infer(None))
+
+
+class TestPrimHook(unittest.TestCase):
+    def setUp(self):
+        core._set_prim_all_enabled(False)
+
+        def f():
+            return paddle.nn.functional.dropout(paddle.rand((1,)))
+
+        concrete_program, partial_program = paddle.jit.to_static(
+            f
+        ).get_concrete_program()
+        self._hook = program_translator.PrimHooker(
+            concrete_program.main_program
+        )
+        self._forward = partial_program.forward_program
+        self._whole = partial_program._train_program
+
+        core._set_prim_all_enabled(True)
+
+    def tearDown(self):
+        core._set_prim_all_enabled(False)
+
+    def test_before_append_backward(self):
+        self._hook.before_append_backward(self._forward)
+        self.assertNotIn(
+            'dropout', tuple(op.type for op in self._forward.blocks[0].ops)
+        )
+
+    def test_after_append_backward(self):
+        self._hook.after_append_backward(self._whole, 0)
+        self.assertNotIn(
+            'dropout_grad', tuple(op.type for op in self._whole.blocks[0].ops)
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_batch_norm.py

@@ -21,6 +21,7 @@
 import paddle.nn as nn
 import paddle.nn.functional as F
 from paddle.fluid import core, framework
+from paddle.incubate.autograd import primapi
 from paddle.nn import BatchNorm
 from paddle.tensor import ones  # noqa: F401
 from paddle.tensor import zeros  # noqa: F401
@@ -183,7 +184,7 @@ def cal_composite(
                 attrs.use_global_stats,
             )
             blocks = main_program.blocks
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
         exe = paddle.static.Executor()
         exe.run(startup_program)

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_batch_norm_grad.py

@@ -20,6 +20,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 np.random.seed(2023)
 
@@ -190,7 +191,7 @@ def cal_composite(
                 attrs.use_global_stats,
             )
             blocks = main_program.blocks
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             z = paddle.static.gradients([y], [x1])
 

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_dropout.py

@@ -19,6 +19,7 @@
 
 import paddle
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 np.random.seed(2023)
 
@@ -154,7 +155,7 @@ def dropout(x, p, is_test, mode, seed=0):
                     input_, p, training=(not is_test), mode=mode
                 )
                 if core._is_fwd_prim_enabled():
-                    paddle.incubate.autograd.to_prim(mp.blocks)
+                    primapi.to_prim(mp.blocks)
                 grad = paddle.static.gradients(output, input_)[0]
             exe = paddle.static.Executor(self.place)
             exe.run(sp)

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_gelu.py

@@ -22,6 +22,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 
 def generate_data(shape, dtype="float32"):
@@ -89,7 +90,7 @@ def cal_composite(self, inputs):
             # Ensure that gelu in original block
             self.assertTrue('gelu' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that gelu is splitted into small ops

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_gelu_grad.py

@@ -22,6 +22,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 
 def generate_data(shape, dtype="float32"):
@@ -97,7 +98,7 @@ def cal_composite_grad(self, inputs):
             # Ensure that gelu in original block
             self.assertTrue('gelu' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that gelu is splitted into small ops
@@ -164,7 +165,7 @@ def cal_composite_grad(self, inputs):
             x.stop_gradient = False
             y = fn(x)
             blocks = main_program.blocks
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
             z = paddle.static.gradients([y], x)
 
         exe = paddle.static.Executor()

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_layer_norm.py

@@ -20,6 +20,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 
 def generate_data(shape1, shape2, shape3, dtype="float32"):
@@ -98,7 +99,7 @@ def cal_composite(self, inputs, norm_shape, weight, bias):
             # Ensure that layer_norm in original block
             self.assertTrue('layer_norm' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that layer_norm is splitted into small ops
@@ -137,7 +138,7 @@ def cal2_composite(self, inputs, norm_shape, weight, bias):
             # Ensure that layer_norm in original block
             self.assertTrue('layer_norm' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that layer_norm is splitted into small ops

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_layer_norm_grad.py

@@ -22,6 +22,7 @@
 import paddle
 import paddle.nn.functional as F
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 TOLERANCE_NUMPY = {
     "float32": {"rtol": 2e-5, "atol": 2e-5},
@@ -196,7 +197,7 @@ def cal_composite_backward(self, inputs, norm_shape, weight, bias):
             # Ensure that layer_norm in original block
             self.assertTrue('layer_norm' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that layer_norm is splitted into small ops
@@ -242,7 +243,7 @@ def cal2_composite_backward(self, inputs, norm_shape, weight, bias):
             # Ensure that layer_norm in original block
             self.assertTrue('layer_norm' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that layer_norm is splitted into small ops
@@ -341,7 +342,7 @@ def cal_composite_backward(self, inputs, norm_shape, weight, bias):
             y = fn(x, norm_shape, w, b)
 
             blocks = main_program.blocks
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
             z = paddle.static.gradients([y], x)
 
         exe = paddle.static.Executor()
@@ -374,7 +375,7 @@ def cal2_composite_backward(self, inputs, norm_shape, weight, bias):
             y = fn(x, norm_shape, weight, bias)
 
             blocks = main_program.blocks
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
             z = paddle.static.gradients([y], x)
 
         exe = paddle.static.Executor()
@@ -480,7 +481,7 @@ def cal_composite_backward(self, inputs, norm_shape, weight, bias, y_grad):
             # Ensure that layer_norm in original block
             self.assertTrue('layer_norm' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that layer_norm is splitted into small ops
@@ -532,7 +533,7 @@ def cal_composite_backward_prim(
             )
 
             blocks = main_program.blocks
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
             z = paddle.static.gradients([y], x)
 
         exe = paddle.static.Executor()

python/paddle/fluid/tests/unittests/prim/composite_ops/test_composite_mean.py

@@ -20,6 +20,7 @@
 import paddle
 import paddle.tensor as tensor
 from paddle.fluid import core
+from paddle.incubate.autograd import primapi
 
 
 def generate_data(shape, dtype="float32"):
@@ -93,7 +94,7 @@ def cal_composite(self, inputs):
             # Ensure that reduce_mean in original block
             self.assertTrue('reduce_mean' in fwd_ops)
 
-            paddle.incubate.autograd.to_prim(blocks)
+            primapi.to_prim(blocks)
 
             fwd_ops_new = [op.type for op in blocks[0].ops]
             # Ensure that reduce_mean is splitted into small ops