未验证 提交 bcea3b89 编写于 作者: Y YuanRisheng 提交者: GitHub

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

* perfect structure kernel registry

* fix ci bugs
上级 897fb6ab
...@@ -485,14 +485,31 @@ struct OpKernelRegistrarFunctorEx<PlaceType, ...@@ -485,14 +485,31 @@ struct OpKernelRegistrarFunctorEx<PlaceType,
USE_OP_KERNEL(op_type) USE_OP_KERNEL(op_type)
// clang-format on // clang-format on
template <typename StructureKernel, typename enable = void>
struct StructKernelImpl;
template <typename StructureKernel> 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) { static void Compute(phi::KernelContext* ctx) {
auto exe_ctx = static_cast<paddle::framework::ExecutionContext*>(ctx); auto exe_ctx = static_cast<paddle::framework::ExecutionContext*>(ctx);
StructureKernel().Compute(*exe_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(...) \ #define PHI_STRUCTURE_KERNEL(...) \
::paddle::framework::StructKernelImpl<__VA_ARGS__>::Compute ::paddle::framework::StructKernelImpl<__VA_ARGS__>::Compute
#define PHI_STRUCTURE_VARIADIC_KERNEL(...) nullptr #define PHI_STRUCTURE_VARIADIC_KERNEL(...) nullptr
......
...@@ -54,6 +54,7 @@ class GetTensorFromSelectedRowsOp : public framework::OperatorWithKernel { ...@@ -54,6 +54,7 @@ class GetTensorFromSelectedRowsOp : public framework::OperatorWithKernel {
} }
}; };
template <typename T, typename DeviceContext>
class GetTensorFromSelectedRowsKernel { class GetTensorFromSelectedRowsKernel {
public: public:
void operator()(const framework::ExecutionContext &ctx) const { void operator()(const framework::ExecutionContext &ctx) const {
...@@ -101,24 +102,22 @@ REGISTER_OPERATOR(get_tensor_from_selected_rows, ...@@ -101,24 +102,22 @@ REGISTER_OPERATOR(get_tensor_from_selected_rows,
ops::GetTensorFromSelectedRowsOpProtoMaker, ops::GetTensorFromSelectedRowsOpProtoMaker,
ops::GetTensorFromSelectedRowsOpVarTypeInference); ops::GetTensorFromSelectedRowsOpVarTypeInference);
REGISTER_OP_CPU_KERNEL_FUNCTOR(get_tensor_from_selected_rows, PD_REGISTER_STRUCT_KERNEL(get_tensor_from_selected_rows,
float, CPU,
ops::GetTensorFromSelectedRowsKernel, ALL_LAYOUT,
double, ops::GetTensorFromSelectedRowsKernel,
ops::GetTensorFromSelectedRowsKernel, float,
int, double,
ops::GetTensorFromSelectedRowsKernel, int,
int64_t, int64_t) {}
ops::GetTensorFromSelectedRowsKernel);
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
REGISTER_OP_CUDA_KERNEL_FUNCTOR(get_tensor_from_selected_rows, PD_REGISTER_STRUCT_KERNEL(get_tensor_from_selected_rows,
float, GPU,
ops::GetTensorFromSelectedRowsKernel, ALL_LAYOUT,
double, ops::GetTensorFromSelectedRowsKernel,
ops::GetTensorFromSelectedRowsKernel, float,
int, double,
ops::GetTensorFromSelectedRowsKernel, int,
int64_t, int64_t) {}
ops::GetTensorFromSelectedRowsKernel);
#endif #endif
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册