What should do to add MKLDNN kernel
Created by: jacquesqiao
Base class that need to implement:
-
MKLDNNLoDTensor
derived from LoDTensor. in formal PaddlePaddle. Like MKLDNNMatrix.h for v2 -
MKLDNNKernelBase
derived from KernelBase. Like MKLDNNLayer.h in formal PaddlePaddle for v2 -
Tensor transform functions used to tranform between LodTensor and MKLDNNKernel
Steps to add new kernel
Take pool_cudnn_op
as example:
- add
pool_cudnn_op.cu.cc
- implement
PoolCUDNNOpKernel
- register kernel with library type
CUDNN
REGISTER_OP_KERNEL(pool2d, CUDNN, ::paddle::platform::CUDAPlace,
ops::PoolCUDNNOpKernel<float>,
ops::PoolCUDNNOpKernel<double>);
If we want to add pool_mkldnn
kernel, the step should be:
- add new file
pool_mkldnn_op.cc
- implement PoolMKLDNNOpKernel
- register PoolMKLDNNOpKernel with library
MKLDNN
REGISTER_OP_KERNEL(pool2d, MKLDNN, ::paddle::platform::MKLDNNPlace,
ops::PoolMKLDNNOpKernel<float>,
ops::PoolMKLDNNOpKernel<double>);