[Prim][PIR] Polish backward_blacklist usage (#69356)

cubehan3 · web-flow · commit 7aae19454108 · 2024-11-18T20:08:35.000+08:00
* polish backward_blacklist usage
diff --git a/paddle/common/flags.cc b/paddle/common/flags.cc
@@ -1685,14 +1685,6 @@ PHI_DEFINE_EXPORTED_string(
     "",
     "It controls the forward blacklist ops not to be decomposed.");
 
-// PIR and prim related FLAG
-// Example: If prim_backward_blacklist="relu_grad;mean_grad",
-// it will block the decompsitions of `relu` and `mean` backward grads.
-PHI_DEFINE_EXPORTED_string(
-    prim_backward_blacklist,
-    "",
-    "It controls the forward blacklist ops not to be decomposed.");
-
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL) || \
     defined(PADDLE_WITH_XPU_BKCL)
 /**
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -848,25 +848,6 @@ static std::vector<std::vector<pir::Value>> GenerateBackwardBlockForPyLayerOp(
   return res;
 }
 
-namespace {
-std::unordered_set<std::string> StringSplit(const std::string &str) {
-  std::istringstream iss(str);
-  std::unordered_set<std::string> tokens;
-  std::string token;
-  while (std::getline(iss, token, ';')) {
-    size_t startpos = token.find_first_not_of(' ');
-    size_t endpos = token.find_last_not_of(' ');
-    if ((startpos != std::string::npos) && (endpos != std::string::npos)) {
-      token = token.substr(startpos, endpos - startpos + 1);
-    } else if (startpos != std::string::npos) {
-      token = token.substr(startpos);
-    }
-    tokens.insert(token);
-  }
-  return tokens;
-}
-}  // namespace
-
 void BindVjp(pybind11::module *m) {
   m->def(
       "call_vjp",
@@ -898,10 +879,6 @@ void BindVjp(pybind11::module *m) {
                          common::errors::InvalidArgument(
                              "The vjp function is not registered in %s op ",
                              fwd_op.name()));
-          const std::unordered_set<std::string> backward_blacklist_ops =
-              StringSplit(FLAGS_prim_backward_blacklist);
-          paddle::prim::PrimCommonUtils::SetPrimBackwardBlacklist(
-              backward_blacklist_ops);
           vjp_res = vjp_interface.Vjp(
               &fwd_op, inputs, outputs, out_grads, stop_gradients);
         }
diff --git a/python/paddle/base/core.py b/python/paddle/base/core.py
@@ -594,17 +594,23 @@ def _set_prim_forward_blacklist(*args):
             prim_config["forward_blacklist"].add(item)
 
 
-def _reset_prim_forward_blacklist():
-    prim_config["forward_blacklist"] = set()
+# Currently, this function is not utilized anywhere in the codebase.
+# It may be intended for future use or could be removed if unnecessary.
+# def _reset_prim_forward_blacklist():
+#     prim_config["forward_blacklist"] = set()
 
 
 def _set_prim_backward_blacklist(*args):
     ops = set(args)
+    new_ops = set()
     for item in ops:
-        prim_config["backward_blacklist"].add(item)
         if not isinstance(item, str):
-            raise TypeError("all items in set must belong to string")
-    _set_bwd_prim_blacklist(ops)
+            raise TypeError("All items in set must be strings.")
+        if item.startswith("pd_op."):
+            item = item[6:]
+        prim_config["backward_blacklist"].add(item)
+        new_ops.add(item)
+    _set_bwd_prim_blacklist(new_ops)
 
 
 def _set_prim_backward_enabled(value):
diff --git a/test/prim/pir_prim/CMakeLists.txt b/test/prim/pir_prim/CMakeLists.txt
@@ -93,3 +93,15 @@ foreach(target ${TEST_PRIM_TRANS_PIR_CASES})
   py_test_modules(${target} MODULES ${target} ENVS GLOG_v=1
                   FLAGS_enable_pir_in_executor=true)
 endforeach()
+
+set(TEST_PRIM_BACKWARD_BLACKLIST "test_pir_prim_flags_v4")
+
+py_test_modules(
+  ${TEST_PRIM_BACKWARD_BLACKLIST}
+  MODULES
+  ${TEST_PRIM_BACKWARD_BLACKLIST}
+  ENVS
+  GLOG_v=1
+  FLAGS_enable_pir_api=true
+  FLAGS_prim_enable_dynamic=true
+  FLAGS_prim_backward_blacklist=pd_op.tanh_grad)
diff --git a/test/prim/pir_prim/test_pir_prim_flags.py b/test/prim/pir_prim/test_pir_prim_flags.py
@@ -129,15 +129,7 @@ def check_prim(self, net):
 
     def test_prim_backward_blacklist(self):
         core._set_prim_all_enabled(True)
-        core._set_prim_backward_blacklist("tanh_grad", "exp_grad")
-        self.train()
-        core._set_prim_all_enabled(False)
-
-    def test_prim_backward_blacklist_flag(self):
-        core._set_prim_all_enabled(True)
-        paddle.set_flags(
-            {"FLAGS_prim_backward_blacklist": "tanh_grad;exp_grad"}
-        )
+        core._set_prim_backward_blacklist("pd_op.tanh_grad", "pd_op.exp_grad")
         self.train()
         core._set_prim_all_enabled(False)
 
diff --git a/test/prim/pir_prim/test_pir_prim_flags_v4.py b/test/prim/pir_prim/test_pir_prim_flags_v4.py
@@ -0,0 +1,63 @@
+# Copyright (c) 2024 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
+#
+#     https://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.base import core
+
+
+class PrimeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        x1 = paddle.tanh(x)
+        x2 = paddle.exp(x)
+        x3 = x1 + x2
+        res = paddle.nn.functional.gelu(x3)
+        return res
+
+
+class TestPrimBackwardBlacklistFlags(unittest.TestCase):
+    def train(self):
+        x = paddle.randn([2, 4])
+        x.stop_gradient = False
+        net = PrimeNet()
+        net.forward = paddle.jit.to_static(full_graph=True)(net.forward)
+        out = net(x)
+        loss = paddle.mean(out)
+        loss.backward()
+        self.check_prim(net)
+
+    def check_prim(self, net):
+        program = net.forward.program_cache.last()[-1][-1].train_program
+        if isinstance(
+            program, paddle.jit.dy2static.pir_partial_program.RunnableProgram
+        ):
+            program = program.program
+        block = program.global_block()
+        ops = [op.name() for op in block.ops]
+        self.assertTrue('pd_op.tanh_grad' in ops)
+        self.assertTrue('pd_op.exp_grad' not in ops)
+        self.assertTrue('pd_op.gelu_grad' not in ops)
+
+    def test_prim_backward_blacklist_flag(self):
+        core._set_prim_all_enabled(True)
+        self.train()
+        core._set_prim_all_enabled(False)
+
+
+if __name__ == '__main__':
+    unittest.main()
