From 2c378d49d8aace5eca0b5862b6ddb02a67cd904e Mon Sep 17 00:00:00 2001 From: zyfncg Date: Tue, 26 Jul 2022 13:13:14 +0000 Subject: [PATCH] support auto fallback to cpu kernel for cusom device --- paddle/phi/core/kernel_factory.cc | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/paddle/phi/core/kernel_factory.cc b/paddle/phi/core/kernel_factory.cc index d6f6e60fe2d3d..ae9c16e0cc710 100644 --- a/paddle/phi/core/kernel_factory.cc +++ b/paddle/phi/core/kernel_factory.cc @@ -110,6 +110,7 @@ const Kernel& KernelFactory::SelectKernelOrThrowError( << "] is not registered."; } #endif + auto kernel_iter = iter->second.find(kernel_key); // TODO(chenweihang): polish refind impl here if (kernel_iter == iter->second.end() && @@ -118,6 +119,22 @@ const Kernel& KernelFactory::SelectKernelOrThrowError( kernel_key.backend(), phi::DataLayout::ALL_LAYOUT, kernel_key.dtype()); kernel_iter = iter->second.find(any_layout_kernel_key); } + +#ifdef PADDLE_WITH_CUSTOM_DEVICE + if (kernel_iter == iter->second.end()) { + // Fallback CPU backend + phi::KernelKey cpu_kernel_key( + phi::Backend::CPU, kernel_key.layout(), kernel_key.dtype()); + kernel_iter = iter->second.find(cpu_kernel_key); + if (kernel_iter == iter->second.end() && + kernel_key.layout() != phi::DataLayout::ALL_LAYOUT) { + phi::KernelKey any_layout_kernel_key( + phi::Backend::CPU, phi::DataLayout::ALL_LAYOUT, kernel_key.dtype()); + kernel_iter = iter->second.find(any_layout_kernel_key); + } + } +#endif + PADDLE_ENFORCE_NE( kernel_iter, iter->second.end(),