[PHI] Migrate softplus kernel (#47406)

* add extra attr property set

* add type_info for all context

* add onednn context to all context

* fix context compile error

* simplify conv kernel args

* pass runtime attr into dev_ctx

* fix marco error

* clear conv_grad_kernel extra args

* merge conv_grad_grad into conv_grad

* clear conv2d_grad_grad extra attrs

* remove redundant imports

* migrate softmax

* clear yaml and eager extra attr

* fix conv1d error

* change to thread local

* fix npu compile failed

* try to fix windows compile failed

* add conv2d onednn phi kernel

* fix ci bugs (#36)

* fix compile bugs (#38)

* fix extra input transform bug (#39)

* support dynamic created attr (#40)

* reset extra info gen code

* rm conv_grad_grad kernel

* reimpl pass attr adapting

* add int attr support

* remove vector inputnames creating

* merge dev

* fix map at error

* adjust attribute

* adapt funcs to PHI

* init

* adjust imports

* support postops

* format codeblocks

* revert changes to softmax
Co-authored-by: NChen Weihang <chenweihang@baidu.com>
Co-authored-by: NYuanRisheng <yuanrisheng@baidu.com>

paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc

-/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-   http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License. */
-
-#include "paddle/fluid/operators/activation_op.h"
-#include "paddle/fluid/operators/mkldnn/softplus_mkldnn_op.h"
-#include "paddle/fluid/platform/mkldnn_reuse.h"
-
-namespace phi {
-class DenseTensor;
-}  // namespace phi
-
-namespace paddle {
-namespace operators {
-
-using dnnl::memory;
-using dnnl::primitive;
-using dnnl::stream;
-using phi::DataLayout;
-
-using platform::GetMKLDNNFormat;
-using platform::MKLDNNDeviceContext;
-using platform::to_void_cast;
-
-template <typename Functor>
-class MKLDNNActivationKernel
-    : public framework::OpKernel<typename Functor::ELEMENT_TYPE> {
- public:
-  void Compute(const framework::ExecutionContext &ctx) const override {
-    Functor functor;
-    functor(ctx);
-  }
-};
-
-template <typename T>
-struct SoftplusMKLDNNFunctor : public BaseActivationFunctor<T> {
-  void operator()(const framework::ExecutionContext &ctx) const {
-    custom_softplus_eltwise_forward<T>(ctx);
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-
-#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>>);
-
-REGISTER_FWD_ACTIVATION_MKLDNN_KERNEL(softplus, SoftplusMKLDNNFunctor);

paddle/fluid/operators/mkldnn/softplus_mkldnn_op.h → paddle/phi/kernels/onednn/softplus_kernel.cc

-/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
-
-#pragma once
-#include "paddle/fluid/platform/mkldnn_reuse.h"
-
-namespace paddle {
-namespace operators {
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/activation_kernel.h"
+
+#include "paddle/phi/backends/onednn/onednn_reuse.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
 
 template <typename T>
-class SoftplusMKLDNNHandler
-    : public platform::MKLDNNHandlerNoCachingT<T, dnnl::binary> {
+class SoftplusOneDNNHandler
+    : public funcs::OneDNNHandlerNoCachingT<T, dnnl::binary> {
  public:
-  SoftplusMKLDNNHandler(const framework::ExecutionContext& ctx,
+  SoftplusOneDNNHandler(const OneDNNContext& dev_ctx,
                         const phi::DenseTensor* x,
-                        const float beta,
-                        const dnnl::engine engine)
-      : platform::MKLDNNHandlerNoCachingT<T, dnnl::binary>(engine,
-                                                           ctx.GetPlace()) {
-    auto x_tz = phi::vectorize(x->dims());
-
-    auto beta_tz = std::vector<int64_t>(x_tz.size(), 1);
-    auto beta_md =
-        dnnl::memory::desc(beta_tz,
-                           platform::MKLDNNGetDataType<T>(),
-                           platform::GetPlainMKLDNNFormat(x_tz.size()));
-
+                        const float beta)
+      : funcs::OneDNNHandlerNoCachingT<T, dnnl::binary>(dev_ctx.GetEngine(),
+                                                        dev_ctx.GetPlace()) {
     dnnl::post_ops post_ops;
     post_ops.append_eltwise(
         1.0f, dnnl::algorithm::eltwise_soft_relu, 0.0f, 0.0f);
@@ -43,12 +35,16 @@ class SoftplusMKLDNNHandler
       post_ops.append_eltwise(
           1.0f, dnnl::algorithm::eltwise_linear, 1.0f / beta, 0.0f);
     }
-
-    platform::AppendActivation(ctx, post_ops);
-
+    funcs::AppendActivation(dev_ctx, post_ops);
     dnnl::primitive_attr attrs;
     attrs.set_post_ops(post_ops);
 
+    auto x_tz = phi::vectorize(x->dims());
+    auto beta_tz = std::vector<int64_t>(x_tz.size(), 1);
+    auto beta_md = dnnl::memory::desc(beta_tz,
+                                      funcs::OneDNNGetDataType<T>(),
+                                      funcs::GetPlainOneDNNFormat(x_tz.size()));
+
     this->AcquireForwardPrimitiveDescriptor(attrs,
                                             dnnl::algorithm::binary_mul,
                                             x->mem_desc(),
@@ -57,39 +53,31 @@ class SoftplusMKLDNNHandler
   }
 
   std::shared_ptr<dnnl::memory> AcquireBetaMemory(const float* beta) {
-    return this->AcquireMemoryFromPrimitive(
-        this->fwd_pd_->src1_desc(), platform::to_void_cast<float>(beta));
+    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->src1_desc(),
+                                            funcs::to_void_cast<float>(beta));
   }
 };
 
-template <typename T>
-void custom_softplus_eltwise_forward(const framework::ExecutionContext& ctx) {
-  const auto& dev_ctx =
-      ctx.template device_context<platform::MKLDNNDeviceContext>();
-  const auto& mkldnn_engine = dev_ctx.GetEngine();
-
-  const auto* x = ctx.Input<phi::DenseTensor>("X");
-  auto* out = ctx.Output<phi::DenseTensor>("Out");
-
-  bool is_inplaced = x->IsSharedBufferWith(*out);
-
-  const float beta = ctx.Attr<float>("beta");
-
-  SoftplusMKLDNNHandler<T> handler(ctx, x, beta, mkldnn_engine);
-
-  auto src_memory_p = handler.AcquireSrcMemory(x);
+template <typename T, typename Context>
+void SoftplusKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    float beta,
+                    float threshold,
+                    DenseTensor* out) {
+  SoftplusOneDNNHandler<T> handler(dev_ctx, &x, beta);
 
+  auto src_memory_p = handler.AcquireSrcMemory(&x);
   auto beta_memory_p = handler.AcquireBetaMemory(&beta);
   std::shared_ptr<dnnl::memory> dst_memory_p = nullptr;
-  if (is_inplaced) {
+  if (x.IsSharedBufferWith(*out)) {
     dst_memory_p = src_memory_p;
-    out->mutable_data<T>(ctx.GetPlace());
+    dev_ctx.template Alloc<T>(out);
   } else {
     dst_memory_p = handler.AcquireDstMemory(out);
   }
   auto binary_p = handler.AcquireForwardPrimitive();
 
-  auto& astream = paddle::platform::MKLDNNDeviceContext::tls().get_stream();
+  auto& astream = OneDNNContext::tls().get_stream();
 
   const std::unordered_map<int, dnnl::memory> args = {
       {DNNL_ARG_SRC_0, *src_memory_p},
@@ -101,5 +89,12 @@ void custom_softplus_eltwise_forward(const framework::ExecutionContext& ctx) {
 
   out->set_mem_desc(dst_memory_p->get_desc());
 }
-}  // namespace operators
-}  // namespace paddle
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(softplus,
+                   OneDNN,
+                   ONEDNN,
+                   phi::SoftplusKernel,
+                   float,
+                   phi::dtype::bfloat16) {}