Added logsoftmax BF16/FP32 FWD oneDNN kernel

paddle/fluid/operators/log_softmax_op.cc

@@ -31,9 +31,17 @@ class LogSoftmaxOp : public framework::OperatorWithKernel {
  protected:
   framework::OpKernelType GetExpectedKernelType(
       const framework::ExecutionContext& ctx) const override {
-    return framework::OpKernelType(
-        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
-        ctx.device_context());
+    auto input_data_type =
+        framework::OperatorWithKernel::IndicateVarDataType(ctx, "X");
+
+#ifdef PADDLE_WITH_MKLDNN
+    if (this->CanMKLDNNBeUsed(ctx, input_data_type)) {
+      return framework::OpKernelType(input_data_type, ctx.GetPlace(),
+                                     framework::DataLayout::kMKLDNN,
+                                     framework::LibraryType::kMKLDNN);
+    }
+#endif
+    return framework::OpKernelType(input_data_type, ctx.GetPlace());
   }
 };
 
@@ -48,6 +56,10 @@ class LogSoftmaxOpMaker : public framework::OpProtoAndCheckerMaker {
                  "The dimension index of Input(x) to perform log_softmax,"
                  "default -1 for last dimension")
         .SetDefault(-1);
+    AddAttr<bool>("use_mkldnn",
+                  "(bool, default false) Only used in mkldnn kernel")
+        .SetDefault(false)
+        .AsExtra();
     AddComment(R"DOC(
 LogSoftmax Operator.
 

paddle/fluid/operators/mkldnn/log_softmax_mkldnn_op.cc

@@ -0,0 +1,78 @@
+/* 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/fluid/operators/softmax_op.h"
+#include "paddle/fluid/platform/mkldnn_reuse.h"
+
+namespace paddle {
+namespace operators {
+
+using framework::Tensor;
+
+template <typename T>
+class LogSoftmaxMKLDNNHandler
+    : public platform::MKLDNNHandlerNoCachingT<T, dnnl::logsoftmax_forward> {
+ public:
+  LogSoftmaxMKLDNNHandler(const dnnl::engine mkldnn_engine,
+                          platform::Place cpu_place, const Tensor* x,
+                          const int axis)
+      : platform::MKLDNNHandlerNoCachingT<T, dnnl::logsoftmax_forward>(
+            mkldnn_engine, cpu_place) {
+    const auto logsoftmax_tz = phi::vectorize(x->dims());
+    const auto md = dnnl::memory::desc(
+        logsoftmax_tz, platform::MKLDNNGetDataType<T>(), x->format());
+
+    this->AcquireForwardPrimitiveDescriptor(dnnl::prop_kind::forward_inference,
+                                            md, axis);
+  }
+};
+
+template <typename T>
+class LogSoftmaxMKLDNNKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    auto& dev_ctx =
+        ctx.template device_context<platform::MKLDNNDeviceContext>();
+    const auto& mkldnn_engine = dev_ctx.GetEngine();
+
+    const Tensor* x = ctx.Input<Tensor>("X");
+    Tensor* out = ctx.Output<Tensor>("Out");
+
+    int axis = ctx.Attr<int>("axis");
+    axis = axis >= 0 ? axis : x->dims().size() + axis;
+
+    LogSoftmaxMKLDNNHandler<T> handler(mkldnn_engine, ctx.GetPlace(), x, axis);
+
+    auto src_memory_p = handler.AcquireSrcMemory(x);
+    auto dst_memory_p = handler.AcquireDstMemory(out);
+
+    auto logsoftmax_p = handler.AcquireForwardPrimitive();
+
+    auto& astream = platform::MKLDNNDeviceContext::tls().get_stream();
+    logsoftmax_p->execute(astream, {{DNNL_ARG_SRC, *src_memory_p},
+                                    {DNNL_ARG_DST, *dst_memory_p}});
+    astream.wait();
+
+    out->set_layout(framework::DataLayout::kMKLDNN);
+    out->set_format(x->format());
+  }
+};
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+
+REGISTER_OP_KERNEL(log_softmax, MKLDNN, ::paddle::platform::CPUPlace,
+                   ops::LogSoftmaxMKLDNNKernel<float>,
+                   ops::LogSoftmaxMKLDNNKernel<paddle::platform::bfloat16>);

python/paddle/fluid/tests/unittests/ir/inference/test_mkldnn_log_softmax_op.py

@@ -0,0 +1,63 @@
+# 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.
+
+from auto_scan_test import MkldnnAutoScanTest
+from program_config import TensorConfig, ProgramConfig, OpConfig
+import numpy as np
+from functools import partial
+import unittest
+from hypothesis import given
+import hypothesis.strategies as st
+
+
+class TestMKLDNNLogSoftmaxOp(MkldnnAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def sample_program_configs(self, *args, **kwargs):
+        def generate_input(*args, **kwargs):
+            return np.random.random(kwargs['in_shape']).astype(np.float32)
+
+        logsoftmax_op = OpConfig(
+            type="log_softmax",
+            inputs={"X": ["input_data"]},
+            outputs={"Out": ["output_data"]},
+            attrs={"axis": kwargs['axis']})
+
+        program_config = ProgramConfig(
+            ops=[logsoftmax_op],
+            weights={},
+            inputs={
+                "input_data": TensorConfig(data_gen=partial(generate_input,
+                                                            *args, **kwargs)),
+            },
+            outputs=["output_data"])
+
+        yield program_config
+
+    def sample_predictor_configs(self, program_config):
+        config = self.create_inference_config(use_mkldnn=True)
+        yield config, (1e-5, 1e-5)
+
+    @given(
+        axis=st.sampled_from([-2, -1, 0, 1]),
+        in_shape=st.lists(
+            st.integers(
+                min_value=2, max_value=5), min_size=3, max_size=5))
+    def test(self, *args, **kwargs):
+        self.run_test(quant=False, *args, **kwargs)
+
+
+if __name__ == "__main__":
+    unittest.main()

python/paddle/fluid/tests/unittests/mkldnn/test_log_softmax_mkldnn_op.py

@@ -0,0 +1,93 @@
+#   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.
+
+import unittest
+import numpy as np
+import paddle
+from paddle.fluid import core
+from paddle.fluid.tests.unittests.test_log_softmax import ref_log_softmax
+from paddle.fluid.tests.unittests.op_test import OpTest, OpTestTool, convert_float_to_uint16
+
+
+@OpTestTool.skip_if_not_cpu_bf16()
+class TestLogSoftmaxOneDNNOp(OpTest):
+    def setUp(self):
+        self.op_type = 'log_softmax'
+        self.set_dtype()
+        self.set_shape()
+        self.set_axis()
+
+        x = np.random.uniform(0.1, 1.0, self.shape).astype(np.float32)
+        out = np.apply_along_axis(ref_log_softmax, self.axis, x)
+
+        if self.dtype == np.uint16:
+            x = convert_float_to_uint16(x)
+
+        self.inputs = {'X': x}
+        self.outputs = {'Out': out}
+        self.attrs = {'axis': self.axis, 'use_mkldnn': True}
+
+    def set_dtype(self):
+        self.dtype = np.float32
+
+    def set_shape(self):
+        self.shape = [2, 3, 4, 5]
+
+    def set_axis(self):
+        self.axis = -1
+
+    def test_check_output(self):
+        self.check_output_with_place(core.CPUPlace())
+
+
+class TestLogSoftmax1DOneDNNOp(TestLogSoftmaxOneDNNOp):
+    def set_shape(self):
+        self.shape = [100]
+
+
+class TestLogSoftmax3DOneDNNOp(TestLogSoftmaxOneDNNOp):
+    def set_shape(self):
+        self.shape = [12, 10, 3]
+
+
+class TestLogSoftmax5DOneDNNOp(TestLogSoftmaxOneDNNOp):
+    def set_shape(self):
+        self.shape = [2, 3, 4, 5, 6]
+
+
+class TestLogSoftmaxPositiveAxisOneDNNOp(TestLogSoftmaxOneDNNOp):
+    def set_axis(self):
+        self.axis = 2
+
+
+# BF16 TESTS
+class TestLogSoftmax1DBF16OneDNNOp(TestLogSoftmax1DOneDNNOp):
+    def set_dtype(self):
+        self.dtype = np.uint16
+
+
+class TestLogSoftmaxPositiveAxisBF16OneDNNOp(
+        TestLogSoftmaxPositiveAxisOneDNNOp):
+    def set_dtype(self):
+        self.dtype = np.uint16
+
+
+class TestLogSoftmax5DBF16OneDNNOp(TestLogSoftmax5DOneDNNOp):
+    def set_shape(self):
+        self.shape = [2, 3, 4, 5, 6]
+
+
+if __name__ == "__main__":
+    paddle.enable_static()
+    unittest.main()

python/paddle/fluid/tests/unittests/test_log_softmax.py

@@ -14,7 +14,7 @@
 
 import unittest
 import numpy as np
-from op_test import OpTest
+from paddle.fluid.tests.unittests.op_test import OpTest
 import paddle
 import paddle.nn.functional as F