/* 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 framework {} // namespace framework namespace platform { class MKLDNNDeviceContext; } // namespace platform } // namespace paddle namespace paddle { namespace operators { using dnnl::memory; using dnnl::primitive; using dnnl::stream; using framework::DataLayout; using framework::Tensor; using platform::GetMKLDNNFormat; using platform::MKLDNNDeviceContext; using platform::to_void_cast; template class MKLDNNActivationKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext &ctx) const override { Functor functor; functor(ctx); } }; template class MKLDNNActivationGradKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext &ctx) const override { Functor functor; functor(ctx); } }; template void eltwise_forward(const framework::ExecutionContext &ctx, dnnl::algorithm algorithm) { PADDLE_ENFORCE_EQ(platform::is_cpu_place(ctx.GetPlace()), true, paddle::platform::errors::PreconditionNotMet( "Operator DNNL eletwise_forward must use CPUPlace")); auto &dev_ctx = ctx.template device_context(); const auto &mkldnn_engine = dev_ctx.GetEngine(); const auto *x = ctx.Input("X"); auto *out = ctx.Output("Out"); bool is_inplaced = x->IsSharedBufferWith(*out); platform::ActivationMKLDNNHandler handler( algorithm, ctx, mkldnn_engine, ctx.GetPlace(), x); auto src_memory_p = handler.AcquireSrcMemory(x); std::shared_ptr dst_memory_p = nullptr; if (is_inplaced) { dst_memory_p = src_memory_p; out->mutable_data(ctx.GetPlace()); } else { dst_memory_p = handler.AcquireDstMemory(out); } auto activation_p = handler.AcquireForwardPrimitive(); auto &astream = paddle::platform::MKLDNNDeviceContext::tls().get_stream(); activation_p->execute( astream, {{DNNL_ARG_FROM, *src_memory_p}, {DNNL_ARG_TO, *dst_memory_p}}); astream.wait(); out->set_mem_desc(dst_memory_p->get_desc()); } template void eltwise_grad(const framework::ExecutionContext &ctx, dnnl::algorithm algorithm) { auto &dev_ctx = ctx.template device_context(); const auto &mkldnn_engine = dev_ctx.GetEngine(); const auto *x = ctx.Input("X"); const auto *dout = ctx.Input(framework::GradVarName("Out")); auto *dx = ctx.Output(framework::GradVarName("X")); platform::ActivationMKLDNNHandler handler( algorithm, ctx, mkldnn_engine, ctx.GetPlace(), x, dout); auto src_memory_p = handler.AcquireBackwardSrcMemory(x); auto diff_dst_memory_p = handler.AcquireDiffDstMemory(dout); auto diff_src_memory_p = handler.AcquireDiffSrcMemory(dx); auto activation_backward_p = handler.AcquireBackwardPrimitive(); auto &astream = paddle::platform::MKLDNNDeviceContext::tls().get_stream(); activation_backward_p->execute(astream, {{DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DIFF_DST, *diff_dst_memory_p}, {DNNL_ARG_DIFF_SRC, *diff_src_memory_p}}); astream.wait(); dx->set_mem_desc(diff_src_memory_p->get_desc()); } template void eltwise_grad_use_out(const framework::ExecutionContext &ctx, dnnl::algorithm algorithm) { auto &dev_ctx = ctx.template device_context(); const auto &mkldnn_engine = dev_ctx.GetEngine(); const auto *out = ctx.Input("Out"); const auto *dout = ctx.Input(framework::GradVarName("Out")); auto *dx = ctx.Output(framework::GradVarName("X")); platform::ActivationMKLDNNHandler handler( algorithm, ctx, mkldnn_engine, ctx.GetPlace(), out, dout); auto dst_memory_p = handler.AcquireBackwardSrcMemory(out); auto diff_dst_memory_p = handler.AcquireDiffDstMemory(dout); auto diff_src_memory_p = handler.AcquireDiffSrcMemory(dx); auto activation_backward_p = handler.AcquireBackwardPrimitive(); auto &astream = paddle::platform::MKLDNNDeviceContext::tls().get_stream(); activation_backward_p->execute(astream, {{DNNL_ARG_DST, *dst_memory_p}, {DNNL_ARG_DIFF_DST, *diff_dst_memory_p}, {DNNL_ARG_DIFF_SRC, *diff_src_memory_p}}); astream.wait(); dx->set_mem_desc(diff_src_memory_p->get_desc()); } template struct MKLDNNActivationFunc : public BaseActivationFunctor { void operator()(const framework::ExecutionContext &ctx) const { eltwise_forward(ctx, algorithm); } }; template struct MKLDNNActivationGradFunc : public BaseActivationFunctor { void operator()(const framework::ExecutionContext &ctx) const { eltwise_grad(ctx, algorithm); } }; template struct MKLDNNActivationGradUseOutFunc : public BaseActivationFunctor { void operator()(const framework::ExecutionContext &ctx) const { eltwise_grad_use_out(ctx, algorithm); } }; template struct GeluMKLDNNFunctor : public BaseActivationFunctor { void operator()(const framework::ExecutionContext &ctx) const { const bool approximate = ctx.Attr("approximate"); if (approximate) { eltwise_forward(ctx, dnnl::algorithm::eltwise_gelu_tanh); } else { eltwise_forward(ctx, dnnl::algorithm::eltwise_gelu_erf); } } }; template struct GeluMKLDNNGradFunctor : public BaseActivationFunctor { void operator()(const framework::ExecutionContext &ctx) const { const bool approximate = ctx.Attr("approximate"); if (approximate) { eltwise_grad(ctx, dnnl::algorithm::eltwise_gelu_tanh); } else { eltwise_grad(ctx, dnnl::algorithm::eltwise_gelu_erf); } } }; template struct SoftplusMKLDNNFunctor : public BaseActivationFunctor { void operator()(const framework::ExecutionContext &ctx) const { custom_softplus_eltwise_forward(ctx); } }; template using ReluMKLDNNFunctor = MKLDNNActivationFunc; template using Relu6MKLDNNFunctor = MKLDNNActivationFunc; template using SwishMKLDNNFunctor = MKLDNNActivationFunc; template using HardSwishMKLDNNFunctor = MKLDNNActivationFunc; template using MishMKLDNNFunctor = MKLDNNActivationFunc; template using SigmoidMKLDNNFunctor = MKLDNNActivationFunc; template using TanhMKLDNNFunctor = MKLDNNActivationFunc; template using SqrtMKLDNNFunctor = MKLDNNActivationFunc; template using AbsMKLDNNFunctor = MKLDNNActivationFunc; template using EluMKLDNNFunctor = MKLDNNActivationFunc; template using ExpMKLDNNFunctor = MKLDNNActivationFunc; template using RoundMKLDNNFunctor = MKLDNNActivationFunc; template using ReluMKLDNNGradFunctor = MKLDNNActivationGradFunc; template using Relu6MKLDNNGradFunctor = MKLDNNActivationGradFunc; template using SwishMKLDNNGradFunctor = MKLDNNActivationGradFunc; template using HardSwishMKLDNNGradFunctor = MKLDNNActivationGradFunc; template using MishMKLDNNGradFunctor = MKLDNNActivationGradFunc; template using SigmoidMKLDNNGradUseOutFunctor = MKLDNNActivationGradUseOutFunc< T, dnnl::algorithm::eltwise_logistic_use_dst_for_bwd>; template using TanhMKLDNNGradUseOutFunctor = MKLDNNActivationGradUseOutFunc< T, dnnl::algorithm::eltwise_tanh_use_dst_for_bwd>; template using SqrtMKLDNNGradUseOutFunctor = MKLDNNActivationGradUseOutFunc< T, dnnl::algorithm::eltwise_sqrt_use_dst_for_bwd>; template using AbsMKLDNNGradFunctor = MKLDNNActivationGradFunc; template using EluMKLDNNGradUseOutFunctor = MKLDNNActivationGradUseOutFunc< T, dnnl::algorithm::eltwise_elu_use_dst_for_bwd>; template using ExpMKLDNNGradUseOutFunctor = MKLDNNActivationGradUseOutFunc< T, dnnl::algorithm::eltwise_exp_use_dst_for_bwd>; } // namespace operators } // namespace paddle namespace ops = paddle::operators; #define REGISTER_ACTIVATION_MKLDNN_KERNEL(act_type, functor, grad_functor) \ REGISTER_OP_KERNEL( \ act_type, \ MKLDNN, \ ::paddle::platform::CPUPlace, \ ops::MKLDNNActivationKernel>, \ ops::MKLDNNActivationKernel>); \ REGISTER_OP_KERNEL( \ act_type##_grad, \ MKLDNN, \ ::paddle::platform::CPUPlace, \ ops::MKLDNNActivationGradKernel>, \ ops::MKLDNNActivationGradKernel< \ ops::grad_functor>); #define REGISTER_ACTIVATION_MKLDNN_KERNEL_FWD_ONLY(act_type, functor) \ REGISTER_OP_KERNEL(act_type, \ MKLDNN, \ ::paddle::platform::CPUPlace, \ ops::MKLDNNActivationKernel>); #define FOR_EACH_MKLDNN_KERNEL_FUNCTOR(__macro) \ __macro(abs, AbsMKLDNNFunctor, AbsMKLDNNGradFunctor); \ __macro(elu, EluMKLDNNFunctor, EluMKLDNNGradUseOutFunctor); \ __macro(exp, ExpMKLDNNFunctor, ExpMKLDNNGradUseOutFunctor); \ __macro(gelu, GeluMKLDNNFunctor, GeluMKLDNNGradFunctor); \ __macro(hard_swish, HardSwishMKLDNNFunctor, HardSwishMKLDNNGradFunctor); \ __macro(leaky_relu, ReluMKLDNNFunctor, ReluMKLDNNGradFunctor); \ __macro(mish, MishMKLDNNFunctor, MishMKLDNNGradFunctor); \ __macro(relu, ReluMKLDNNFunctor, ReluMKLDNNGradFunctor); \ __macro(relu6, Relu6MKLDNNFunctor, Relu6MKLDNNGradFunctor); \ __macro(sigmoid, SigmoidMKLDNNFunctor, SigmoidMKLDNNGradUseOutFunctor); \ __macro(sqrt, SqrtMKLDNNFunctor, SqrtMKLDNNGradUseOutFunctor); \ __macro(swish, SwishMKLDNNFunctor, SwishMKLDNNGradFunctor); \ __macro(tanh, TanhMKLDNNFunctor, TanhMKLDNNGradUseOutFunctor); FOR_EACH_MKLDNN_KERNEL_FUNCTOR(REGISTER_ACTIVATION_MKLDNN_KERNEL); // round eltwise primitive doesn't support BF16, nor does it support grad REGISTER_ACTIVATION_MKLDNN_KERNEL_FWD_ONLY(round, RoundMKLDNNFunctor); namespace ops = paddle::operators; REGISTER_OP_KERNEL( softplus, MKLDNN, paddle::platform::CPUPlace, ops::MKLDNNActivationKernel>, ops::MKLDNNActivationKernel< ops::SoftplusMKLDNNFunctor>);