diff --git a/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc index dec2fa1836081..0f75b36f71360 100644 --- a/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc +++ b/paddle/fluid/operators/mkldnn/activation_mkldnn_op.cc @@ -262,6 +262,10 @@ using EluMKLDNNFunctor = MKLDNNActivationFunc; template using ExpMKLDNNFunctor = MKLDNNActivationFunc; +template +using RoundMKLDNNFunctor = + MKLDNNActivationFunc; + template using ReluMKLDNNGradFunctor = MKLDNNActivationGradFunc; @@ -330,6 +334,10 @@ namespace ops = paddle::operators; 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(relu6, Relu6MKLDNNFunctor, Relu6MKLDNNGradFunctor); \ __macro(leaky_relu, ReluMKLDNNFunctor, ReluMKLDNNGradFunctor); \ @@ -341,6 +349,8 @@ namespace ops = paddle::operators; __macro(exp, ExpMKLDNNFunctor, ExpMKLDNNGradUseOutFunctor); FOR_EACH_MKLDNN_KERNEL_FUNCTOR(REGISTER_ACTIVATION_MKLDNN_KERNEL); +REGISTER_ACTIVATION_MKLDNN_KERNEL_FWD_ONLY(round, RoundMKLDNNFunctor); + REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(relu, ReluMKLDNNFunctor, ReluMKLDNNGradFunctor); REGISTER_ACTIVATION_MKLDNN_BF16_KERNEL(gelu, GeluMKLDNNFunctor, diff --git a/python/paddle/fluid/tests/unittests/mkldnn/test_activation_mkldnn_op.py b/python/paddle/fluid/tests/unittests/mkldnn/test_activation_mkldnn_op.py index e2d50fc853887..4e4fe69d914fa 100644 --- a/python/paddle/fluid/tests/unittests/mkldnn/test_activation_mkldnn_op.py +++ b/python/paddle/fluid/tests/unittests/mkldnn/test_activation_mkldnn_op.py @@ -328,6 +328,18 @@ def setUp(self): self.attrs = {"use_mkldnn": True} +class TestMKLDNNRound(TestActivation): + def setUp(self): + self.op_type = "round" + + x = np.random.uniform(0.1, 1, [2, 4, 3, 5]).astype(np.float32) + out = np.round(x) + + self.inputs = {'X': x} + self.outputs = {'Out': out} + self.attrs = {"use_mkldnn": True} + + class TestMKLDNNSigmoidDim4(TestSigmoid): def setUp(self): super(TestMKLDNNSigmoidDim4, self).setUp()