polish kernel signature details

PaddlePaddle · Oct 26, 2021 · 8f100da · 8f100da · paddle-bot-old · Oct 26, 2021
1 parent 72910fa
commit 8f100da
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 17 deletions.
diff --git a/paddle/fluid/framework/operator.cc b/paddle/fluid/framework/operator.cc
@@ -1763,14 +1763,12 @@ OpKernelType OperatorWithKernel::GetKernelTypeForVar(
 
 KernelSignature OperatorWithKernel::GetExpectedPtenKernelArgs(
     const ExecutionContext& ctx) const {
-  if (KernelSignatureMap::Instance().Has(Type())) {
-    return *(KernelSignatureMap::Instance().GetNullable(Type()));
-  } else {
+  if (!KernelSignatureMap::Instance().Has(Type())) {
     KernelArgsNameMakerByOpProto maker(Info().proto_);
-    auto signature = std::move(maker.GetKernelSignature());
-    KernelSignatureMap::Instance().Insert(Type(), signature);
-    return signature;
+    KernelSignatureMap::Instance().Emplace(
+        Type(), std::move(maker.GetKernelSignature()));
   }
+  return KernelSignatureMap::Instance().Get(Type());
 }
 
 pten::KernelContext OperatorWithKernel::BuildPtenKernelContext(

diff --git a/paddle/fluid/framework/pten_utils.cc b/paddle/fluid/framework/pten_utils.cc
@@ -59,11 +59,6 @@ pten::KernelKey TransOpKernelTypeToPtenKernelKey(
   return pten::KernelKey(backend, layout, dtype);
 }
 
-KernelSignatureMap& KernelSignatureMap::Instance() {
-  static KernelSignatureMap g_kernel_signature_map;
-  return g_kernel_signature_map;
-}
-
 const paddle::SmallVector<std::string>&
 KernelArgsNameMakerByOpProto::GetInputArgsNames() {
   for (int i = 0; i < op_proto_->inputs_size(); ++i) {

diff --git a/paddle/fluid/framework/pten_utils.h b/paddle/fluid/framework/pten_utils.h
@@ -44,7 +44,10 @@ pten::KernelKey TransOpKernelTypeToPtenKernelKey(
 // TODO(chenweihang): we can generate this map by proto info in compile time
 class KernelSignatureMap {
  public:
-  static KernelSignatureMap& Instance();
+  static KernelSignatureMap& Instance() {
+    static KernelSignatureMap g_kernel_signature_map;
+    return g_kernel_signature_map;
+  }
 
   bool Has(const std::string& op_type) const {
     return map_.find(op_type) != map_.end();
@@ -56,6 +59,12 @@ class KernelSignatureMap {
     }
   }
 
+  void Emplace(const std::string& op_type, KernelSignature&& signature) {
+    if (!Has(op_type)) {
+      map_.emplace(op_type, signature);
+    }
+  }
+
   const KernelSignature* GetNullable(const std::string& op_type) const {
     auto it = map_.find(op_type);
     if (it == map_.end()) {
@@ -65,6 +74,15 @@ class KernelSignatureMap {
     }
   }
 
+  const KernelSignature& Get(const std::string& op_type) const {
+    auto it = map_.find(op_type);
+    PADDLE_ENFORCE_NE(
+        it, map_.end(),
+        platform::errors::NotFound(
+            "Operator `%s`'s kernel signature is not registered.", op_type));
+    return it->second;
+  }
+
  private:
   KernelSignatureMap() = default;
   paddle::flat_hash_map<std::string, KernelSignature> map_;

diff --git a/paddle/fluid/platform/flags.cc b/paddle/fluid/platform/flags.cc
@@ -684,16 +684,14 @@ PADDLE_DEFINE_EXPORTED_bool(
 /**
  * Pt kernel related FLAG
  * Name: FLAGS_run_pten_kernel
- * Since Version: 2.2.0
+ * Since Version: 2.3.0
  * Value Range: bool, default=false
  * Example: FLAGS_run_pten_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(run_pten_kernel, true,
-                            "It controls whether to use pt kernel");
+                            "It controls whether to use pten kernel");
 
 /**
  * Distributed related FLAG

diff --git a/python/paddle/fluid/tests/unittests/test_fuse_all_reduce_pass.py b/python/paddle/fluid/tests/unittests/test_fuse_all_reduce_pass.py
@@ -83,7 +83,7 @@ def _decorate_compare_fused_all_reduce(self, model, use_device):
             use_device,
             init_feed_dict=init_data,
             optimizer=self.optimizer,
-            fuse_all_optimizer_ops=False)
+            fuse_all_optimizer_ops=True)
 
     def test_simple_fc_with_fuse_all_reduce(self):
         self._decorate_compare_fused_all_reduce(simple_fc_net, DeviceType.CUDA)