未验证 提交 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,
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
......
......@@ -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
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册