[PHI] Migrate relu6 and abs kernels (#45397)

* abs relu6 fwd

* abs bwd

* gaussian_random_kernel and mkldnn-onednn renaming

* scale kernel

* whitespace

* whitespace

* revert scale migration

* whitespaces

* revert changes to gaussian kernel

* whitespaces

paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc

@@ -144,13 +144,6 @@ void eltwise_grad_use_out(const framework::ExecutionContext &ctx,
   dx->set_mem_desc(diff_src_memory_p->get_desc());
 }
 
-template <typename T, dnnl::algorithm algorithm>
-struct MKLDNNActivationFunc : public BaseActivationFunctor<T> {
-  void operator()(const framework::ExecutionContext &ctx) const {
-    eltwise_forward<T>(ctx, algorithm);
-  }
-};
-
 template <typename T, dnnl::algorithm algorithm>
 struct MKLDNNActivationGradFunc : public BaseActivationFunctor<T> {
   void operator()(const framework::ExecutionContext &ctx) const {
@@ -158,13 +151,6 @@ struct MKLDNNActivationGradFunc : public BaseActivationFunctor<T> {
   }
 };
 
-template <typename T, dnnl::algorithm algorithm>
-struct MKLDNNActivationGradUseOutFunc : public BaseActivationFunctor<T> {
-  void operator()(const framework::ExecutionContext &ctx) const {
-    eltwise_grad_use_out<T>(ctx, algorithm);
-  }
-};
-
 template <typename T>
 struct GeluMKLDNNFunctor : public BaseActivationFunctor<T> {
   void operator()(const framework::ExecutionContext &ctx) const {
@@ -196,33 +182,24 @@ struct SoftplusMKLDNNFunctor : public BaseActivationFunctor<T> {
   }
 };
 
-template <typename T>
-using Relu6MKLDNNFunctor =
-    MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_bounded_relu>;
-
-template <typename T>
-using AbsMKLDNNFunctor = MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_abs>;
-
 template <typename T>
 using Relu6MKLDNNGradFunctor =
     MKLDNNActivationGradFunc<T, dnnl::algorithm::eltwise_bounded_relu>;
 
-template <typename T>
-using AbsMKLDNNGradFunctor =
-    MKLDNNActivationGradFunc<T, dnnl::algorithm::eltwise_abs>;
-
 }  // namespace operators
 }  // namespace paddle
 
 namespace ops = paddle::operators;
 
-#define REGISTER_ACTIVATION_MKLDNN_KERNEL(act_type, functor, grad_functor)    \
+#define REGISTER_FWD_ACTIVATION_MKLDNN_KERNEL(act_type, functor) \
   REGISTER_OP_KERNEL(                                            \
       act_type,                                                  \
       MKLDNN,                                                    \
       ::paddle::platform::CPUPlace,                              \
       ops::MKLDNNActivationKernel<ops::functor<float>>,          \
-      ops::MKLDNNActivationKernel<ops::functor<paddle::platform::bfloat16>>); \
+      ops::MKLDNNActivationKernel<ops::functor<paddle::platform::bfloat16>>);
+
+#define REGISTER_GRAD_ACTIVATION_MKLDNN_KERNEL(act_type, grad_functor) \
   REGISTER_OP_KERNEL(                                                  \
       act_type##_grad,                                                 \
       MKLDNN,                                                          \
@@ -231,24 +208,7 @@ namespace ops = paddle::operators;
       ops::MKLDNNActivationGradKernel<                                 \
           ops::grad_functor<paddle::platform::bfloat16>>);
 
-#define REGISTER_ACTIVATION_MKLDNN_KERNEL_FWD_ONLY(act_type, functor) \
-  REGISTER_OP_KERNEL(act_type,                                        \
-                     MKLDNN,                                          \
-                     ::paddle::platform::CPUPlace,                    \
-                     ops::MKLDNNActivationKernel<ops::functor<float>>);
-
-#define FOR_EACH_MKLDNN_KERNEL_FUNCTOR(__macro)            \
-  __macro(abs, AbsMKLDNNFunctor, AbsMKLDNNGradFunctor);    \
-  __macro(gelu, GeluMKLDNNFunctor, GeluMKLDNNGradFunctor); \
-  __macro(relu6, Relu6MKLDNNFunctor, Relu6MKLDNNGradFunctor);
-
-FOR_EACH_MKLDNN_KERNEL_FUNCTOR(REGISTER_ACTIVATION_MKLDNN_KERNEL);
-
-namespace ops = paddle::operators;
-REGISTER_OP_KERNEL(
-    softplus,
-    MKLDNN,
-    paddle::platform::CPUPlace,
-    ops::MKLDNNActivationKernel<ops::SoftplusMKLDNNFunctor<float>>,
-    ops::MKLDNNActivationKernel<
-        ops::SoftplusMKLDNNFunctor<paddle::platform::bfloat16>>);
+REGISTER_FWD_ACTIVATION_MKLDNN_KERNEL(softplus, SoftplusMKLDNNFunctor);
+REGISTER_FWD_ACTIVATION_MKLDNN_KERNEL(gelu, GeluMKLDNNFunctor);
+REGISTER_GRAD_ACTIVATION_MKLDNN_KERNEL(gelu, GeluMKLDNNGradFunctor);
+REGISTER_GRAD_ACTIVATION_MKLDNN_KERNEL(relu6, Relu6MKLDNNGradFunctor);

paddle/fluid/platform/mkldnn_reuse.h

@@ -25,7 +25,7 @@ limitations under the License. */
 #include "paddle/fluid/operators/pool_op.h"
 #include "paddle/fluid/platform/mkldnn_helper.h"
 #include "paddle/fluid/platform/place.h"
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_reuse.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_reuse.h"
 
 namespace paddle {
 namespace platform {

paddle/phi/kernels/funcs/data_layout_transform.cc

@@ -25,8 +25,8 @@
 #include "paddle/phi/core/dense_tensor.h"
 
 #ifdef PADDLE_WITH_MKLDNN
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_helper.h"
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_reuse.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_helper.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_reuse.h"
 #endif
 
 namespace phi {

paddle/phi/kernels/funcs/onednn/mkldnn_reuse.h → paddle/phi/kernels/funcs/onednn/onednn_reuse.h

@@ -25,7 +25,7 @@ limitations under the License. */
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_helper.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_helper.h"
 
 namespace phi {
 namespace funcs {

paddle/phi/kernels/onednn/activation_grad_kernel.cc

@@ -19,7 +19,7 @@
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/activation_functor.h"
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_reuse.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_reuse.h"
 
 namespace phi {
 
@@ -147,6 +147,10 @@ struct MKLDNNActivationGradUseOutFunc : public funcs::BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+using AbsMKLDNNGradFunctor =
+    MKLDNNActivationGradFunc<T, dnnl::algorithm::eltwise_abs>;
+
 template <typename T>
 using ReluMKLDNNGradFunctor =
     MKLDNNActivationGradFunc<T, dnnl::algorithm::eltwise_relu>;
@@ -193,6 +197,7 @@ DEFINE_ONEDNN_ACTIVATION_GRAD_KERNEL_DEPOUT(Sqrt, SqrtMKLDNNGradUseOutFunctor);
 DEFINE_ONEDNN_ACTIVATION_GRAD_KERNEL_DEPOUT(Sigmoid,
                                             SigmoidMKLDNNGradUseOutFunctor);
 DEFINE_ONEDNN_ACTIVATION_GRAD_KERNEL_DEPOUT(Exp, ExpMKLDNNGradUseOutFunctor);
+DEFINE_ONEDNN_ACTIVATION_GRAD_KERNEL_DEPOUT(Abs, AbsMKLDNNGradFunctor);
 DEFINE_ONEDNN_ACTIVATION_GRAD_KERNEL_DEPOUT(Relu, ReluMKLDNNGradFunctor);
 
 DEFINE_ONEDNN_ACT_GRAD_KERNEL_WITH_ONE_ATTRS_DEPX(LeakyRelu,
@@ -240,6 +245,7 @@ PD_REGISTER_KERNEL(relu_grad,
   PD_REGISTER_KERNEL(                                  \
       name, OneDNN, ALL_LAYOUT, phi::func, float, phi::dtype::bfloat16) {}
 
+PD_REGISTER_ACTIVATION_GRAD_KERNEL(abs_grad, AbsGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(elu_grad, EluGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(exp_grad, ExpGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(hard_swish_grad, HardSwishGradKernel)

paddle/phi/kernels/onednn/activation_kernel.cc

@@ -19,7 +19,7 @@
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/activation_functor.h"
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_reuse.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_reuse.h"
 
 namespace phi {
 
@@ -88,10 +88,17 @@ struct MKLDNNActivationFunc : public funcs::BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+using AbsMKLDNNFunctor = MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_abs>;
+
 template <typename T>
 using ReluMKLDNNFunctor =
     MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_relu>;
 
+template <typename T>
+using Relu6MKLDNNFunctor =
+    MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_bounded_relu>;
+
 template <typename T>
 using SwishMKLDNNFunctor =
     MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_swish>;
@@ -126,6 +133,7 @@ template <typename T>
 using RoundMKLDNNFunctor =
     MKLDNNActivationFunc<T, dnnl::algorithm::eltwise_round>;
 
+DEFINE_ONEDNN_ACTIVATION_KERNEL(Abs, AbsMKLDNNFunctor)
 DEFINE_ONEDNN_ACTIVATION_KERNEL(Relu, ReluMKLDNNFunctor)
 DEFINE_ONEDNN_ACTIVATION_KERNEL(Tanh, TanhMKLDNNFunctor)
 DEFINE_ONEDNN_ACTIVATION_KERNEL(Exp, ExpMKLDNNFunctor)
@@ -137,6 +145,7 @@ DEFINE_ONEDNN_ACTIVATION_KERNEL(Round, RoundMKLDNNFunctor)
 DEFINE_ONEDNN_ACT_KERNEL_WITH_ONE_ATTRS(LeakyRelu, ReluMKLDNNFunctor, alpha)
 DEFINE_ONEDNN_ACT_KERNEL_WITH_ONE_ATTRS(Mish, MishMKLDNNFunctor, threshold)
 DEFINE_ONEDNN_ACT_KERNEL_WITH_ONE_ATTRS(Elu, EluMKLDNNFunctor, alpha)
+DEFINE_ONEDNN_ACT_KERNEL_WITH_ONE_ATTRS(Relu6, Relu6MKLDNNFunctor, threshold)
 DEFINE_ONEDNN_ACT_KERNEL_WITH_ONE_ATTRS(Swish, SwishMKLDNNFunctor, beta)
 
 template <typename T, typename Context>
@@ -158,13 +167,15 @@ PD_REGISTER_KERNEL(round, OneDNN, ALL_LAYOUT, phi::RoundKernel, float) {}
   PD_REGISTER_KERNEL(                             \
       name, OneDNN, ALL_LAYOUT, phi::func, float, phi::dtype::bfloat16) {}
 
+PD_REGISTER_ACTIVATION_KERNEL(abs, AbsKernel)
 PD_REGISTER_ACTIVATION_KERNEL(elu, EluKernel)
 PD_REGISTER_ACTIVATION_KERNEL(exp, ExpKernel)
 PD_REGISTER_ACTIVATION_KERNEL(hard_swish, HardSwishKernel)
 PD_REGISTER_ACTIVATION_KERNEL(leaky_relu, LeakyReluKernel)
 PD_REGISTER_ACTIVATION_KERNEL(mish, MishKernel)
+PD_REGISTER_ACTIVATION_KERNEL(relu, ReluKernel)
+PD_REGISTER_ACTIVATION_KERNEL(relu6, Relu6Kernel)
 PD_REGISTER_ACTIVATION_KERNEL(sigmoid, SigmoidKernel)
 PD_REGISTER_ACTIVATION_KERNEL(sqrt, SqrtKernel)
 PD_REGISTER_ACTIVATION_KERNEL(swish, SwishKernel)
 PD_REGISTER_ACTIVATION_KERNEL(tanh, TanhKernel)
-PD_REGISTER_ACTIVATION_KERNEL(relu, ReluKernel)

paddle/phi/kernels/transfer_layout_kernel.cc

@@ -23,7 +23,7 @@ limitations under the License. */
 #include "paddle/phi/kernels/funcs/data_layout_transform.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 #ifdef PADDLE_WITH_MKLDNN
-#include "paddle/phi/kernels/funcs/onednn/mkldnn_helper.h"
+#include "paddle/phi/kernels/funcs/onednn/onednn_helper.h"
 #endif
 namespace phi {