From 4b0f1b0cb8566f5b93a7ad412ca38a2789a1fffd Mon Sep 17 00:00:00 2001
From: ronnywang <ronny1996@163.com>
Date: Tue, 8 Nov 2022 21:46:50 +0800
Subject: [PATCH] [CustomDevice] fix the not ready kernel can not register.
 (#47758)

---
 paddle/phi/core/custom_kernel.cc | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/paddle/phi/core/custom_kernel.cc b/paddle/phi/core/custom_kernel.cc
index 356dd1482c9..75c651c5402 100644
--- a/paddle/phi/core/custom_kernel.cc
+++ b/paddle/phi/core/custom_kernel.cc
@@ -16,6 +16,9 @@
 
 #include "glog/logging.h"
 
+static std::vector<std::string> gpu_exclusive_kernels({"sync_batch_norm",
+                                                       "sync_batch_norm_grad"});
+
 namespace phi {
 
 void CustomKernelMap::RegisterCustomKernel(const std::string& name,
@@ -41,12 +44,15 @@ void CustomKernelMap::RegisterCustomKernels() {
   }
   auto& kernels = KernelFactory::Instance().kernels();
   for (auto& pair : kernels_) {
-    PADDLE_ENFORCE_NE(
-        kernels.find(pair.first),
-        kernels.end(),
-        phi::errors::InvalidArgument(
+    if (kernels.find(pair.first) == kernels.cend()) {
+      if (std::find(gpu_exclusive_kernels.cbegin(),
+                    gpu_exclusive_kernels.cend(),
+                    pair.first) == gpu_exclusive_kernels.cend()) {
+        PADDLE_THROW(phi::errors::InvalidArgument(
             "The kernel %s is not ready for custom kernel registering.",
             pair.first));
+      }
+    }
 
     for (auto& info_pair : pair.second) {
       PADDLE_ENFORCE_EQ(
-- 
GitLab