[Prim][Pir] Add FLAG_prim_backward_blacklist (#69156)

* add FLAG_prim_backward_blacklist

Author: cubehan3

paddle/common/flags.cc

@@ -1689,6 +1689,14 @@ 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)
 /**

paddle/fluid/pybind/pybind.cc

@@ -233,6 +233,7 @@ limitations under the License. */
 #endif
 
 COMMON_DECLARE_bool(use_mkldnn);
+COMMON_DECLARE_string(prim_backward_blacklist);
 
 // disable auto conversion to list in Python
 PYBIND11_MAKE_OPAQUE(phi::TensorArray);
@@ -847,6 +848,25 @@ 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",
@@ -878,6 +898,10 @@ 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);
         }

test/prim/pir_prim/test_pir_prim_flags.py

@@ -133,6 +133,14 @@ def test_prim_backward_blacklist(self):
         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"}
+        )
+        self.train()
+        core._set_prim_all_enabled(False)
+
 
 if __name__ == '__main__':
     unittest.main()