add use_pt_kernel Flags to control whether to use pt kernel (#13)

* add use_pt_kernel Flags to control whether to use pt kernel * change the default value to true for cheking pt kernels
chenwhql · Oct 12, 2021 · aa6ed57 · aa6ed57
1 parent 9b33270
commit aa6ed57
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 2 deletions.
diff --git a/paddle/fluid/framework/operator.cc b/paddle/fluid/framework/operator.cc
@@ -51,6 +51,7 @@ DECLARE_bool(check_nan_inf);
 DECLARE_bool(enable_unused_var_check);
 PADDLE_DEFINE_EXPORTED_int32(inner_op_parallelism, 0,
                              "number of threads for inner op");
+DECLARE_bool(use_pt_kernel);
 
 namespace paddle {
 namespace framework {
@@ -1155,7 +1156,8 @@ void OperatorWithKernel::RunImpl(const Scope& scope,
   // phase
 
   // VLOG(1) << "Pt KernelFactory: " << pt::KernelFactory::Instance();
-  if (pt::KernelFactory::Instance().ContainsKernel(type_.c_str())) {
+  if (FLAGS_use_pt_kernel &&
+      pt::KernelFactory::Instance().ContainsKernel(type_.c_str())) {
     if (pt_kernel_key_.get() == nullptr || pt_kernel_.get() == nullptr) {
       ChoosePtKernel(*runtime_ctx, *dev_ctx);
     }

diff --git a/paddle/fluid/imperative/prepared_operator.cc b/paddle/fluid/imperative/prepared_operator.cc
@@ -22,6 +22,7 @@
 #include "paddle/fluid/platform/xpu/xpu_op_list.h"
 #endif
 DECLARE_bool(check_nan_inf);
+DECLARE_bool(use_pt_kernel);
 
 namespace paddle {
 namespace imperative {
@@ -205,7 +206,8 @@ PreparedOp PrepareImpl(const NameVarMap<VarType>& ins,
 #endif
 
   // 1. get expected kernel key
-  if (pt::KernelFactory::Instance().ContainsKernel(op.Type().c_str())) {
+  if (FLAGS_use_pt_kernel &&
+      pt::KernelFactory::Instance().ContainsKernel(op.Type().c_str())) {
     auto kernel_name =
         ConstructPtKernelName<VarType>(op.Type(), (*op.Info().proto_), ins);
     auto inputs = BuildInputMap<VarType>(ins);

diff --git a/paddle/fluid/platform/flags.cc b/paddle/fluid/platform/flags.cc
@@ -673,3 +673,17 @@ PADDLE_DEFINE_EXPORTED_int32(get_host_by_name_time, 120,
 PADDLE_DEFINE_EXPORTED_bool(
     apply_pass_to_program, false,
     "It controls whether to apply IR pass to program when using Fleet APIs");
+
+/**
+ * Pt kernel related FLAG
+ * Name: FLAGS_use_pt_kernel
+ * Since Version: 2.2.0
+ * Value Range: bool, default=false
+ * Example: FLAGS_use_pt_kernel=true would use the pt kernel to compute in the
+ * Op.
+ * Note:
+ */
+// TODO(chentianyu03): change default value to false before merge into develop
+// branch
+PADDLE_DEFINE_EXPORTED_bool(use_pt_kernel, true,
+                            "It controls whether to use pt kernel");