From 4b0f1b0cb8566f5b93a7ad412ca38a2789a1fffd Mon Sep 17 00:00:00 2001 From: ronnywang 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 356dd1482c..75c651c540 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 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