TaoLv commented on a change in pull request #18320:
URL: https://github.com/apache/incubator-mxnet/pull/18320#discussion_r426401872
##########
File path: src/operator/nn/log_softmax.cc
##########
@@ -26,10 +26,79 @@
#include "../tensor/elemwise_unary_op.h"
#include "../tensor/elemwise_binary_op.h"
#include "../operator_common.h"
+#if MXNET_USE_MKLDNN == 1
+#include "mkldnn/mkldnn_base-inl.h"
+#include "mkldnn/mkldnn_ops-inl.h"
+#endif
namespace mxnet {
namespace op {
+#if MXNET_USE_MKLDNN == 1
+static void LogSoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const SoftmaxParam& param = nnvm::get<SoftmaxParam>(attrs.parsed);
+ if (SupportMKLDNNLogSoftmax(param, inputs[0], outputs[0])) {
+ MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ MKLDNNRun(MKLDNNLogSoftmaxForward, attrs, ctx, inputs[0], req[0],
outputs[0]);
+ auto fn = SoftmaxCompute<cpu, mxnet_op::log_softmax_fwd>;
+ MKLDNN_OPCHECK_RUN(fn, attrs, ctx, inputs, req, outputs);
+ return;
+ }
+ FallBackCompute(SoftmaxCompute<cpu, mxnet_op::log_softmax_fwd>, attrs, ctx,
+ inputs, req, outputs);
+}
+
+static void LogSoftmaxGradComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const SoftmaxParam& param = nnvm::get<SoftmaxParam>(attrs.parsed);
+ if (SupportMKLDNNLogSoftmax(param, inputs[1], outputs[0])) {
+ MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ MKLDNNRun(MKLDNNLogSoftmaxBackward, attrs, ctx, inputs, req, outputs);
+ auto fn = SoftmaxGradCompute<cpu, op::mshadow_op::left,
mxnet_op::log_softmax_bwd>;
+ MKLDNN_OPCHECK_RUN(fn, attrs, ctx, inputs, req, outputs);
+ return;
+ }
+ FallBackCompute(SoftmaxGradCompute<cpu, op::mshadow_op::left,
mxnet_op::log_softmax_bwd>,
+ attrs, ctx, inputs, req, outputs);
+}
+
+inline static bool LogSoftmaxStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
Review comment:
Please fix indent.
##########
File path: src/operator/nn/mkldnn/mkldnn_log_softmax.cc
##########
@@ -0,0 +1,228 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * \file mkldnn_log_softmax.cc
+ * \brief Implementation of log_softmax function with MKLDNN support
+*/
+
+#include "../softmax-inl.h"
+#include "./mkldnn_ops-inl.h"
+#include "./mkldnn_base-inl.h"
+
+#if MXNET_USE_MKLDNN == 1
+namespace mxnet {
+namespace op {
+
+static mkldnn::logsoftmax_forward::primitive_desc GetLogSoftmaxFwdPd(
+ bool is_train,
+ const int axis,
+ const mkldnn::memory &input_mem) {
+ mkldnn::memory::desc data_md = input_mem.get_desc();
+ auto cpu_engine = CpuEngine::Get()->get_engine();
+ auto prop = is_train ? mkldnn::prop_kind::forward_training
+ : mkldnn::prop_kind::forward_scoring;
+ auto desc = mkldnn::logsoftmax_forward::desc(prop, data_md, axis);
+ return mkldnn::logsoftmax_forward::primitive_desc(desc, cpu_engine);
+}
+
+static mkldnn::logsoftmax_backward::primitive_desc GetLogSoftmaxBwdPd(
+ const mkldnn::memory &diff_mem,
+ const mkldnn::memory &data_mem,
+ const int axis,
+ const
mkldnn::logsoftmax_forward::primitive_desc &hint_fwd_pd) {
+ mkldnn::memory::desc diff_md = diff_mem.get_desc();
+ mkldnn::memory::desc data_md = data_mem.get_desc();
+ auto cpu_engine = CpuEngine::Get()->get_engine();
+ auto desc = mkldnn::logsoftmax_backward::desc(diff_md, data_md, axis);
+ return mkldnn::logsoftmax_backward::primitive_desc(desc, cpu_engine,
hint_fwd_pd);
+}
+
+
+bool SupportMKLDNNLogSoftmax(const SoftmaxParam ¶m,
+ const NDArray &data,
+ const NDArray &output) {
+ // MKLDNN does not support temperature argument in their softmax function
+ // now. Need update this once they start to support it.
Review comment:
Duplicated comments. Also change softmax to log_softmax.
##########
File path: src/operator/nn/log_softmax.cc
##########
@@ -49,12 +118,21 @@ Examples::
)code")
.set_attr_parser(ParamParser<SoftmaxParam>)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs){
+ return std::vector<std::string>{"data"};
+})
.set_attr<FCompute>("FCompute<cpu>", SoftmaxCompute<cpu,
mxnet_op::log_softmax_fwd>)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<bool>("TIsMKLDNN", true)
+.set_attr<FComputeEx>("FComputeEx<cpu>", LogSoftmaxComputeExCPU)
+.set_attr<FInferStorageType>("FInferStorageType", LogSoftmaxStorageType)
+#endif
.set_attr<nnvm::FGradient>("FGradient",
SoftmaxFGradient{"_backward_log_softmax"})
.set_attr<nnvm::FInferType>("FInferType", SoftmaxOpType)
.set_num_inputs(1)
.set_num_outputs(1)
-.set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<1, 1>)
+.set_attr<mxnet::FInferShape>("FInferShape", SoftmaxOpShape)
Review comment:
Could you please elaborate?
##########
File path: src/operator/nn/mkldnn/mkldnn_log_softmax.cc
##########
@@ -0,0 +1,228 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * \file mkldnn_log_softmax.cc
+ * \brief Implementation of log_softmax function with MKLDNN support
+*/
+
+#include "../softmax-inl.h"
+#include "./mkldnn_ops-inl.h"
+#include "./mkldnn_base-inl.h"
+
+#if MXNET_USE_MKLDNN == 1
+namespace mxnet {
+namespace op {
+
+static mkldnn::logsoftmax_forward::primitive_desc GetLogSoftmaxFwdPd(
+ bool is_train,
+ const int axis,
+ const mkldnn::memory &input_mem) {
+ mkldnn::memory::desc data_md = input_mem.get_desc();
+ auto cpu_engine = CpuEngine::Get()->get_engine();
+ auto prop = is_train ? mkldnn::prop_kind::forward_training
+ : mkldnn::prop_kind::forward_scoring;
+ auto desc = mkldnn::logsoftmax_forward::desc(prop, data_md, axis);
+ return mkldnn::logsoftmax_forward::primitive_desc(desc, cpu_engine);
+}
+
+static mkldnn::logsoftmax_backward::primitive_desc GetLogSoftmaxBwdPd(
+ const mkldnn::memory &diff_mem,
+ const mkldnn::memory &data_mem,
+ const int axis,
+ const
mkldnn::logsoftmax_forward::primitive_desc &hint_fwd_pd) {
+ mkldnn::memory::desc diff_md = diff_mem.get_desc();
+ mkldnn::memory::desc data_md = data_mem.get_desc();
+ auto cpu_engine = CpuEngine::Get()->get_engine();
+ auto desc = mkldnn::logsoftmax_backward::desc(diff_md, data_md, axis);
+ return mkldnn::logsoftmax_backward::primitive_desc(desc, cpu_engine,
hint_fwd_pd);
+}
+
+
+bool SupportMKLDNNLogSoftmax(const SoftmaxParam ¶m,
+ const NDArray &data,
+ const NDArray &output) {
+ // MKLDNN does not support temperature argument in their softmax function
+ // now. Need update this once they start to support it.
+ const int ndim = data.shape().ndim();
+ const int in_dtype = data.dtype();
+ const int out_dtype = output.dtype();
+ const int axis = CheckAxis(param.axis, ndim);
+ // MKLDNN does not support temperature argument in their softmax function
+ // now. Need update this once they start to support it.
+ // Currently, MKLDNN shows bad performance when softmax is not performed on
the last dimension
+ if (param.temperature.has_value() ||
+ in_dtype != mshadow::kFloat32 ||
+ in_dtype != out_dtype ||
+ axis != (ndim - 1)) {
+ return false;
+ }
+
+ // only supports ndim = 1, 2, 3, 4 for now
+ return (ndim >= 1 && ndim <= 4);
+}
+
+class MKLDNNLogSoftmaxFwd {
+ public:
+ mkldnn::logsoftmax_forward::primitive_desc pd;
+
+ MKLDNNLogSoftmaxFwd(const bool is_train,
+ const int axis,
+ const mkldnn::memory &input) :
pd(GetLogSoftmaxFwdPd(is_train, axis, input)) {
+ fwd_ = std::make_shared<mkldnn::logsoftmax_forward>(pd);
+ }
+
+ const mkldnn::logsoftmax_forward &GetFwd() const {
+ return *fwd_;
+ }
+
+ private:
+ std::shared_ptr<mkldnn::logsoftmax_forward> fwd_;
+};
+
+typedef ParamOpSign<SoftmaxParam> MKLDNNSoftmaxSignature;
+
+static MKLDNNLogSoftmaxFwd &GetLogSoftmaxFwd(const SoftmaxParam ¶m,
+ const int real_axis,
+ const bool is_train,
+ const NDArray &data,
+ const NDArray &output) {
+#if DMLC_CXX11_THREAD_LOCAL
+ static thread_local std::unordered_map<MKLDNNSoftmaxSignature,
+ MKLDNNLogSoftmaxFwd,
+ OpHash> fwds;
+#else
+ static MX_THREAD_LOCAL std::unordered_map<MKLDNNSoftmaxSignature,
+ MKLDNNLogSoftmaxFwd,
+ OpHash> fwds;
+#endif
+
+ MKLDNNSoftmaxSignature key(param);
+ key.AddSign(real_axis);
+ key.AddSign(is_train);
+ key.AddSign(data);
+ key.AddSign(output);
+
+ auto it = fwds.find(key);
+ if (it == fwds.end()) {
+ MKLDNNLogSoftmaxFwd fwd(is_train, real_axis, *(data.GetMKLDNNData()));
+ it = AddToCache(&fwds, key, fwd);
+ }
+ return it->second;
+}
+
+void MKLDNNLogSoftmaxForward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const NDArray &in_data,
+ const OpReqType &req,
+ const NDArray &out_data) {
+ if (req == kNullOp) return;
+ // same as the FCompute path, softmax only supports kWriteTo and
kWriteInplace for now.
Review comment:
Fix the comment - should it be log_softmax?
##########
File path: src/operator/nn/log_softmax.cc
##########
@@ -26,10 +26,79 @@
#include "../tensor/elemwise_unary_op.h"
#include "../tensor/elemwise_binary_op.h"
#include "../operator_common.h"
+#if MXNET_USE_MKLDNN == 1
+#include "mkldnn/mkldnn_base-inl.h"
+#include "mkldnn/mkldnn_ops-inl.h"
+#endif
namespace mxnet {
namespace op {
+#if MXNET_USE_MKLDNN == 1
+static void LogSoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const SoftmaxParam& param = nnvm::get<SoftmaxParam>(attrs.parsed);
+ if (SupportMKLDNNLogSoftmax(param, inputs[0], outputs[0])) {
+ MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ MKLDNNRun(MKLDNNLogSoftmaxForward, attrs, ctx, inputs[0], req[0],
outputs[0]);
+ auto fn = SoftmaxCompute<cpu, mxnet_op::log_softmax_fwd>;
+ MKLDNN_OPCHECK_RUN(fn, attrs, ctx, inputs, req, outputs);
+ return;
+ }
+ FallBackCompute(SoftmaxCompute<cpu, mxnet_op::log_softmax_fwd>, attrs, ctx,
+ inputs, req, outputs);
+}
+
+static void LogSoftmaxGradComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ const SoftmaxParam& param = nnvm::get<SoftmaxParam>(attrs.parsed);
+ if (SupportMKLDNNLogSoftmax(param, inputs[1], outputs[0])) {
+ MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ MKLDNNRun(MKLDNNLogSoftmaxBackward, attrs, ctx, inputs, req, outputs);
+ auto fn = SoftmaxGradCompute<cpu, op::mshadow_op::left,
mxnet_op::log_softmax_bwd>;
+ MKLDNN_OPCHECK_RUN(fn, attrs, ctx, inputs, req, outputs);
+ return;
+ }
+ FallBackCompute(SoftmaxGradCompute<cpu, op::mshadow_op::left,
mxnet_op::log_softmax_bwd>,
+ attrs, ctx, inputs, req, outputs);
+}
+
+inline static bool LogSoftmaxStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+
+ return MKLDNNStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs,
+ out_attrs);
+}
+
+inline static bool LogSoftmaxGradStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
Review comment:
Ditto.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
