[PHI]Support register functor kernel into PHI (#51914)

* perfect structure kernel registry * fix ci bugs
PaddlePaddle · Mar 27, 2023 · bcea3b8 · bcea3b8
1 parent 897fb6a
commit bcea3b8
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 19 deletions.
diff --git a/paddle/fluid/framework/op_registry.h b/paddle/fluid/framework/op_registry.h
@@ -485,14 +485,31 @@ struct OpKernelRegistrarFunctorEx<PlaceType,
   USE_OP_KERNEL(op_type)
 // clang-format on
 
+template <typename StructureKernel, typename enable = void>
+struct StructKernelImpl;
+
 template <typename StructureKernel>
-struct StructKernelImpl {
+struct StructKernelImpl<
+    StructureKernel,
+    typename std::enable_if<std::is_base_of<paddle::framework::OpKernelBase,
+                                            StructureKernel>::value>::type> {
   static void Compute(phi::KernelContext* ctx) {
     auto exe_ctx = static_cast<paddle::framework::ExecutionContext*>(ctx);
     StructureKernel().Compute(*exe_ctx);
   }
 };
 
+template <typename StructureKernel>
+struct StructKernelImpl<
+    StructureKernel,
+    typename std::enable_if<!std::is_base_of<paddle::framework::OpKernelBase,
+                                             StructureKernel>::value>::type> {
+  static void Compute(phi::KernelContext* ctx) {
+    auto exe_ctx = static_cast<paddle::framework::ExecutionContext*>(ctx);
+    StructureKernel()(*exe_ctx);
+  }
+};
+
 #define PHI_STRUCTURE_KERNEL(...) \
   ::paddle::framework::StructKernelImpl<__VA_ARGS__>::Compute
 #define PHI_STRUCTURE_VARIADIC_KERNEL(...) nullptr

diff --git a/paddle/fluid/operators/get_tensor_from_selected_rows_op.cc b/paddle/fluid/operators/get_tensor_from_selected_rows_op.cc
@@ -54,6 +54,7 @@ class GetTensorFromSelectedRowsOp : public framework::OperatorWithKernel {
   }
 };
 
+template <typename T, typename DeviceContext>
 class GetTensorFromSelectedRowsKernel {
  public:
   void operator()(const framework::ExecutionContext &ctx) const {
@@ -101,24 +102,22 @@ REGISTER_OPERATOR(get_tensor_from_selected_rows,
                   ops::GetTensorFromSelectedRowsOpProtoMaker,
                   ops::GetTensorFromSelectedRowsOpVarTypeInference);
 
-REGISTER_OP_CPU_KERNEL_FUNCTOR(get_tensor_from_selected_rows,
-                               float,
-                               ops::GetTensorFromSelectedRowsKernel,
-                               double,
-                               ops::GetTensorFromSelectedRowsKernel,
-                               int,
-                               ops::GetTensorFromSelectedRowsKernel,
-                               int64_t,
-                               ops::GetTensorFromSelectedRowsKernel);
+PD_REGISTER_STRUCT_KERNEL(get_tensor_from_selected_rows,
+                          CPU,
+                          ALL_LAYOUT,
+                          ops::GetTensorFromSelectedRowsKernel,
+                          float,
+                          double,
+                          int,
+                          int64_t) {}
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-REGISTER_OP_CUDA_KERNEL_FUNCTOR(get_tensor_from_selected_rows,
-                                float,
-                                ops::GetTensorFromSelectedRowsKernel,
-                                double,
-                                ops::GetTensorFromSelectedRowsKernel,
-                                int,
-                                ops::GetTensorFromSelectedRowsKernel,
-                                int64_t,
-                                ops::GetTensorFromSelectedRowsKernel);
+PD_REGISTER_STRUCT_KERNEL(get_tensor_from_selected_rows,
+                          GPU,
+                          ALL_LAYOUT,
+                          ops::GetTensorFromSelectedRowsKernel,
+                          float,
+                          double,
+                          int,
+                          int64_t) {}
 #endif