【prim】update Prim optest about customvjp test, add FLAGS_comp_skip_default_ops (#71335)

xiaoguoguo626807 · web-flow · commit cd8ae37b50bc · 2025-02-28T20:39:50.000+08:00
* modify setup wrong path

* update Prim optest about customvjp test, add FLAGS_comp_skip_default_ops
diff --git a/paddle/common/flags.cc b/paddle/common/flags.cc
@@ -1646,6 +1646,11 @@ PHI_DEFINE_EXPORTED_int32(
 
 PHI_DEFINE_EXPORTED_bool(print_ir, false, "Whether print ir debug str.");
 
+PHI_DEFINE_EXPORTED_bool(
+    comp_skip_default_ops,
+    true,
+    "Whether to skip decomposing comp op in default list (decomp_trans.cc).");
+
 PHI_DEFINE_EXPORTED_bool(
     prim_skip_dynamic,
     true,
diff --git a/paddle/fluid/primitive/base/decomp_trans.cc b/paddle/fluid/primitive/base/decomp_trans.cc
@@ -29,6 +29,7 @@
 COMMON_DECLARE_bool(prim_check_ops);
 COMMON_DECLARE_bool(prim_enable_dynamic);
 COMMON_DECLARE_string(prim_forward_blacklist);
+COMMON_DECLARE_bool(comp_skip_default_ops);
 
 using paddle::dialect::DenseTensorType;
 using paddle::dialect::SelectedRowsType;
@@ -368,9 +369,18 @@ bool DecompProgram::enable_decomp_by_filter(const std::string& op_name) {
       flag = false;
     }
   }
+  std::set<std::string> default_comp_blacklist = {"pd_op.embedding",
+                                                  "pd_op.dropout"};
+
   auto from_flag_blacklist = StringSplit(FLAGS_prim_forward_blacklist);
   if (!from_flag_blacklist.empty())
     blacklist_.insert(from_flag_blacklist.begin(), from_flag_blacklist.end());
+
+  if (FLAGS_comp_skip_default_ops) {
+    blacklist_.insert(default_comp_blacklist.begin(),
+                      default_comp_blacklist.end());
+  }
+
   if (!blacklist_.empty() && blacklist_.find(op_name) != blacklist_.end())
     flag = false;
   return flag;
diff --git a/test/dygraph_to_static/CMakeLists.txt b/test/dygraph_to_static/CMakeLists.txt
@@ -41,7 +41,8 @@ set_tests_properties(test_loop PROPERTIES TIMEOUT 180)
 set_tests_properties(test_mnist_amp PROPERTIES TIMEOUT 240)
 set_tests_properties(test_sentiment PROPERTIES TIMEOUT 150)
 set_tests_properties(test_incubate_jit_inference PROPERTIES TIMEOUT 150)
-
+set_tests_properties(test_partial_program_hook
+                     PROPERTIES ENVIRONMENT "FLAGS_comp_skip_default_ops=False")
 if(TEST test_resnet_amp)
   set_tests_properties(test_resnet_amp PROPERTIES TIMEOUT 360)
 endif()
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -818,7 +818,12 @@ endif()
 set_tests_properties(test_tensordot PROPERTIES TIMEOUT 200)
 set_tests_properties(test_partial_eager_deletion_transformer PROPERTIES TIMEOUT
                                                                         120)
-set_tests_properties(test_dropout_op PROPERTIES TIMEOUT 120)
+set_tests_properties(
+  test_dropout_op PROPERTIES TIMEOUT 120 ENVIRONMENT
+                             "FLAGS_comp_skip_default_ops=False")
+set_tests_properties(test_lookup_table_v2_op
+                     PROPERTIES ENVIRONMENT "FLAGS_comp_skip_default_ops=False")
+
 set_tests_properties(test_imperative_ocr_attention_model PROPERTIES TIMEOUT 120)
 set_tests_properties(test_nn_grad PROPERTIES TIMEOUT 180)
 set_tests_properties(test_elementwise_sub_op PROPERTIES TIMEOUT 120)
@@ -1086,7 +1091,7 @@ endif()
 foreach(STATIC_BUILD_TEST ${STATIC_BUILD_TESTS})
   py_test_modules(
     ${STATIC_BUILD_TEST}_static_build MODULES ${STATIC_BUILD_TEST} ENVS
-    FLAGS_new_executor_static_build=true)
+    FLAGS_new_executor_static_build=true FLAGS_comp_skip_default_ops=False)
 endforeach()
 
 py_test_modules(test_adam_op_multi_thread MODULES test_adam_op ENVS
diff --git a/test/legacy_test/prim_op_test.py b/test/legacy_test/prim_op_test.py
@@ -1119,7 +1119,15 @@ def check_static_comp(self):
                 if not in_pir_mode():
                     primapi.to_prim(main_program.blocks)
                 else:
-                    fw_outs = decompose(main_program, fw_outs)
+                    blacklist = set()
+                    for op in main_program.global_block().ops:
+                        if core.has_custom_vjp(op):
+                            blacklist.add(op.name())
+                    fw_outs = decompose(
+                        main_program,
+                        fw_outs,
+                        blacklist=blacklist,
+                    )
                 outputs_dict = self.get_output_dict(
                     self.outputs, fw_outs, outputs_sig
                 )
diff --git a/test/prim/model/test_bert_prim.py b/test/prim/model/test_bert_prim.py
@@ -34,16 +34,16 @@
 
 
 DY2ST_PRIM_GT = [
-    11.144556999206543,
-    10.343620300292969,
-    10.330279350280762,
-    10.276118278503418,
-    10.222086906433105,
-    10.194628715515137,
-    10.14902114868164,
-    10.096250534057617,
-    10.104615211486816,
-    9.985644340515137,
+    10.649632453918457,
+    10.333406448364258,
+    10.33541202545166,
+    10.260543823242188,
+    10.219606399536133,
+    10.176884651184082,
+    10.124699592590332,
+    10.072620391845703,
+    10.112163543701172,
+    9.969392776489258,
 ]
 
 if core.is_compiled_with_cuda():
diff --git a/test/prim/pir_prim/CMakeLists.txt b/test/prim/pir_prim/CMakeLists.txt
@@ -33,6 +33,7 @@ foreach(target ${TEST_PRIM_PURE_PIR_CASES})
     ENVS
     GLOG_v=1
     FLAGS_enable_pir_api=true
+    FLAGS_comp_skip_default_ops=0
     FLAGS_prim_enable_dynamic=true)
 endforeach()
 
@@ -53,6 +54,7 @@ foreach(target ${TEST_PRIM_DYNAMIC_SHAPE_BACKWARD_PIR_CASES})
     GLOG_v=1
     FLAGS_prim_vjp_skip_default_ops=0
     FLAGS_enable_pir_api=true
+    FLAGS_comp_skip_default_ops=0
     FLAGS_prim_enable_dynamic=true)
 endforeach()
 
