From 208a7cf1ce6084bdfb25cd6e7387d0722f594c5a Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 2 Dec 2021 09:52:40 +0100 Subject: [PATCH 1/5] masked softmax --- src/operator/nn/dnnl/dnnl_base-inl.h | 4 + .../nn/dnnl/dnnl_masked_softmax-inl.h | 86 ++++++++ src/operator/nn/dnnl/dnnl_masked_softmax.cc | 200 ++++++++++++++++++ src/operator/nn/dnnl/dnnl_ops-inl.h | 6 + src/operator/nn/dnnl/dnnl_softmax-inl.h | 40 ---- src/operator/nn/dnnl/dnnl_softmax.cc | 42 ++++ src/operator/nn/masked_softmax.cc | 140 ++++++++++++ src/operator/nn/softmax-inl.h | 16 ++ src/operator/nn/softmax.cc | 64 ------ 9 files changed, 494 insertions(+), 104 deletions(-) create mode 100644 src/operator/nn/dnnl/dnnl_masked_softmax-inl.h create mode 100644 src/operator/nn/dnnl/dnnl_masked_softmax.cc create mode 100644 src/operator/nn/masked_softmax.cc diff --git a/src/operator/nn/dnnl/dnnl_base-inl.h b/src/operator/nn/dnnl/dnnl_base-inl.h index 3b0bda9c90de..54aa1caad1b6 100644 --- a/src/operator/nn/dnnl/dnnl_base-inl.h +++ b/src/operator/nn/dnnl/dnnl_base-inl.h @@ -180,6 +180,7 @@ struct LeakyReLUParam; struct ConvolutionParam; struct DeconvolutionParam; struct SoftmaxParam; +struct MaskedSoftmaxParam; struct SoftmaxOutputParam; struct ReshapeParam; struct LayerNormParam; @@ -192,6 +193,9 @@ bool SupportDNNLConv(const ConvolutionParam& params, const NDArray& input); bool SupportDNNLDeconv(const DeconvolutionParam& params, const NDArray& input); bool SupportDNNLSoftmax(const SoftmaxParam& param, const NDArray& input, const NDArray& output); bool SupportDNNLLogSoftmax(const SoftmaxParam& param, const NDArray& input, const NDArray& output); +bool SupportDNNLMaskedSoftmax(const MaskedSoftmaxParam& param, + const std::vector& input, + const NDArray& output); bool SupportDNNLSoftmaxOutput(const SoftmaxOutputParam& param); bool SupportDNNLTranspose(const NDArray& data); bool SupportDNNLBatchDot(const std::vector& inputs, const NDArray& output); diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h b/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h new file mode 100644 index 000000000000..7a1fcca55f3c --- /dev/null +++ b/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h @@ -0,0 +1,86 @@ +/* + * 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 dnnl_masked_softmax-inl.h + */ +#ifndef MXNET_OPERATOR_NN_DNNL_DNNL_MASKED_SOFTMAX_INL_H_ +#define MXNET_OPERATOR_NN_DNNL_DNNL_MASKED_SOFTMAX_INL_H_ +#include +#include +#include + +#include "./dnnl_base-inl.h" +#include "./dnnl_ops-inl.h" +#include "../softmax-inl.h" +#include "./dnnl_softmax-inl.h" + +#if MXNET_USE_ONEDNN == 1 + +namespace mxnet { +namespace op { + +struct Primitives { + public: + Primitives(){}; + + dnnl::eltwise_forward minusone; + dnnl::eltwise_forward::primitive_desc minusone_pd; + + dnnl::binary mask_input; + dnnl::binary::primitive_desc mask_input_pd; + + dnnl::binary mask_output; + dnnl::binary::primitive_desc mask_output_pd; +}; + +class DNNLMaskedSoftmaxFwd { + public: + struct Tensors { + Tensors(const std::vector& inputs, const std::vector& outputs); + + const NDArray& input; + const NDArray& mask; + const NDArray& output; + }; + + DNNLMaskedSoftmaxFwd(const MaskedSoftmaxParam& param, const Tensors& tensors) + : primitives(CreatePrimitives(param, tensors)) {} + + static DNNLMaskedSoftmaxFwd& GetCached(const MaskedSoftmaxParam& param, const Tensors& tensors); + static std::shared_ptr CreatePrimitives(const MaskedSoftmaxParam& param, + const Tensors& tensors); + + void Execute(const Tensors& tensors, + const std::vector& req, + const MaskedSoftmaxParam& param, + bool is_train) const; + + private: + std::shared_ptr primitives; +}; + +DNNLMaskedSoftmaxFwd::Tensors::Tensors(const std::vector& inputs, + const std::vector& outputs) + : input(inputs[0]), mask(inputs[1]), output(outputs[0]) {} + +} // namespace op +} // namespace mxnet +#endif // MXNET_USE_ONEDNN == 1 +#endif // MXNET_OPERATOR_NN_DNNL_DNNL_MASKED_SOFTMAX_INL_H_ \ No newline at end of file diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax.cc b/src/operator/nn/dnnl/dnnl_masked_softmax.cc new file mode 100644 index 000000000000..2846a44dd8d9 --- /dev/null +++ b/src/operator/nn/dnnl/dnnl_masked_softmax.cc @@ -0,0 +1,200 @@ +/* + * 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 dnnl_maskedsoftmax.cc + */ + +#include "./dnnl_masked_softmax-inl.h" +#include "./dnnl_softmax-inl.h" + +#if MXNET_USE_ONEDNN == 1 +namespace mxnet { +namespace op { + +bool SupportDNNLMaskedSoftmax(const MaskedSoftmaxParam& param, + const std::vector& inputs, + const NDArray& output) { + CHECK_EQ(inputs.size(), 2); + const auto mask = inputs[1]; + SoftmaxParam softmax_param; + softmax_param.axis = param.axis; + softmax_param.dtype = inputs[0].dtype(); + softmax_param.temperature = param.temperature; + return mask.dtype() == mshadow::kBool && SupportDNNLSoftmax(softmax_param, inputs[0], output); +} + +inline static dnnl::memory::dims GetOneDNNDims(const NDArray& arr) { + return dnnl::memory::dims(arr.shape().begin(), arr.shape().end()); +} + +typedef ParamOpSign MaskedSoftmaxSignature; + +DNNLMaskedSoftmaxFwd& DNNLMaskedSoftmaxFwd::GetCached( + const MaskedSoftmaxParam& param, + const DNNLMaskedSoftmaxFwd::Tensors& tensors) { + using maskedsoftmax_fwd_map = + std::unordered_map; +#if DMLC_CXX11_THREAD_LOCAL + static thread_local maskedsoftmax_fwd_map fwds; +#else + static MX_THREAD_LOCAL maskedsoftmax_fwd_map fwds; +#endif + MaskedSoftmaxSignature key(param); + key.AddSign(tensors.input); + key.AddSign(tensors.mask); + key.AddSign(tensors.output); + + auto it = fwds.find(key); + if (it == fwds.end()) { + const DNNLMaskedSoftmaxFwd fwd(param, tensors); + it = AddToCache(&fwds, key, fwd); + } + return it->second; +} + +std::shared_ptr DNNLMaskedSoftmaxFwd::CreatePrimitives(const MaskedSoftmaxParam& param, + const Tensors& tensors) { + const dnnl::engine engine = CpuEngine::Get()->get_engine(); + + auto prim = std::make_shared(); + + auto mem_format = + static_cast(GetDefaultFormat(tensors.input.shape().ndim())); + dnnl::memory::dims input_dims = GetOneDNNDims(tensors.input); + // output_desc is the same + dnnl::memory::desc input_desc = + dnnl::memory::desc(input_dims, dnnl::memory::data_type::f32, mem_format); + + dnnl::memory::dims mask_dims = GetOneDNNDims(tensors.mask); + dnnl::memory::desc mask_desc = + dnnl::memory::desc(mask_dims, dnnl::memory::data_type::s8, mem_format); + + // (mask - 1) + auto minusone_desc = dnnl::eltwise_forward::desc(dnnl::prop_kind::forward_scoring, + dnnl::algorithm::eltwise_linear, + mask_desc, + 1.0f, // multiply factor + -1.0f); // minus one + prim->minusone_pd = dnnl::eltwise_forward::primitive_desc(minusone_desc, engine); + prim->minusone = dnnl::eltwise_forward(prim->minusone_pd); + + // input * mask + auto mask_input_desc = + dnnl::binary::desc(dnnl::algorithm::binary_add, + input_desc, + dnnl::memory::desc(mask_dims, dnnl::memory::data_type::f32, mem_format), + input_desc); + prim->mask_input_pd = dnnl::binary::primitive_desc(mask_input_desc, engine); + prim->mask_input = dnnl::binary(prim->mask_input_pd); + + // output * mask + auto mask_output_desc = + dnnl::binary::desc(dnnl::algorithm::binary_mul, input_desc, mask_desc, input_desc); + prim->mask_output_pd = dnnl::binary::primitive_desc(mask_output_desc, engine); + prim->mask_output = dnnl::binary(prim->mask_output_pd); + + return prim; +} + +void DNNLMaskedSoftmaxForward(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + TmpMemMgr::Get()->Init(ctx.requested[0]); + const auto& param = nnvm::get(attrs.parsed); + const auto tensors = DNNLMaskedSoftmaxFwd::Tensors(inputs, outputs); + const auto& fwd = DNNLMaskedSoftmaxFwd::GetCached(param, tensors); + fwd.Execute(tensors, req, param, ctx.is_train); +} + +void DNNLMaskedSoftmaxFwd::Execute(const Tensors& tensors, + const std::vector& req, + const MaskedSoftmaxParam& param, + bool is_train) const { + using namespace mxnet::op; + using namespace mxnet::op::mxnet_op; + using namespace dnnl; + /* + 1. out = input * [(mask - 1) * inf] + 2. out = softmax(out) + 3. out = out * mask + */ + DNNLStream* stream = DNNLStream::Get(); + auto engine = CpuEngine::Get()->get_engine(); + const NDArray& input = tensors.input; + const NDArray& mask = tensors.mask; + const NDArray& output = tensors.output; + + // 1. A) out = mask - 1 + const memory::desc mask_desc = this->primitives->minusone_pd.src_desc(); + dnnl::memory* mask_minusone_mem = TmpMemMgr::Get()->Alloc(mask_desc); + dnnl::memory mask_mem = + dnnl::memory(mask_desc, engine, reinterpret_cast(mask.data().dptr())); + stream->RegisterPrimArgs(this->primitives->minusone, + { + {DNNL_ARG_SRC, mask_mem}, + {DNNL_ARG_DST, *mask_minusone_mem}, + }); + + // 1. B) out = out * inf + const memory::desc converted_mask_desc = this->primitives->mask_input_pd.src1_desc(); + dnnl::memory* convertedmask_mem = TmpMemMgr::Get()->Alloc(converted_mask_desc); + dnnl::primitive_attr attr; + attr.set_output_scales(0, {mshadow::red::limits::MaxValue()}); + std::unordered_map args( + {{DNNL_ARG_FROM, *mask_minusone_mem}, {DNNL_ARG_TO, *convertedmask_mem}}); + stream->RegisterPrimArgs(dnnl::reorder(*mask_minusone_mem, *convertedmask_mem, attr), args); + + // prepare softmax primitive and memory + SoftmaxParam p; + p.axis = param.axis; + p.temperature = param.temperature; + // int axis = CheckAxis(param.axis, input.shape().ndim()); + // auto softmax = DNNLSoftmaxFwd::GetSoftmaxFwd(p, axis, is_train, output, output); + + auto softmax_tensors = DNNLSoftmaxFwd::Tensors(output, output); + auto softmax_op = DNNLSoftmaxFwd::GetCached(p, softmax_tensors, is_train); + auto softmax_out_mem = output.GetDNNLData(softmax_op.softmax_pd->dst_desc()); + const auto input_mem = input.GetDNNLData(); + + // 1. C) out = input * out + stream->RegisterPrimArgs(this->primitives->mask_input, + {{DNNL_ARG_SRC_0, *input_mem}, + {DNNL_ARG_SRC_1, *convertedmask_mem}, + {DNNL_ARG_DST, *softmax_out_mem}}); + stream->Submit(); + + // 2. out = softmax(out) + // stream->RegisterPrimArgs( + // softmax.GetFwd(), {{DNNL_ARG_SRC, *softmax_out_mem}, {DNNL_ARG_DST, *softmax_out_mem}}); + // stream->Submit(); + softmax_op.Execute(softmax_tensors); + + // 3. out = out * mask + stream->RegisterPrimArgs(this->primitives->mask_output, + {{DNNL_ARG_SRC_0, *softmax_out_mem}, + {DNNL_ARG_SRC_1, mask_mem}, + {DNNL_ARG_DST, *softmax_out_mem}}); + stream->Submit(); +} +} // namespace op +} // namespace mxnet +#endif diff --git a/src/operator/nn/dnnl/dnnl_ops-inl.h b/src/operator/nn/dnnl/dnnl_ops-inl.h index 2eea57424129..40e944939bea 100644 --- a/src/operator/nn/dnnl/dnnl_ops-inl.h +++ b/src/operator/nn/dnnl/dnnl_ops-inl.h @@ -125,6 +125,12 @@ void DNNLLogSoftmaxBackward(const nnvm::NodeAttrs& attrs, const std::vector& req, const std::vector& out_data); +void DNNLMaskedSoftmaxForward(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs); + /* For softmax_output */ void DNNLSoftmaxOutputForward(const nnvm::NodeAttrs& attrs, const OpContext& ctx, diff --git a/src/operator/nn/dnnl/dnnl_softmax-inl.h b/src/operator/nn/dnnl/dnnl_softmax-inl.h index eea520eef88c..c681f3373f57 100644 --- a/src/operator/nn/dnnl/dnnl_softmax-inl.h +++ b/src/operator/nn/dnnl/dnnl_softmax-inl.h @@ -75,31 +75,12 @@ class DNNLSoftmaxFwd { DNNLSoftmaxFwd(const SoftmaxParam& param, const Tensors& tensors, const bool is_train); void Execute(const Tensors& tensors) const; - private: std::shared_ptr softmax_pd; std::shared_ptr softmax_fwd; std::shared_ptr temperature_pd; std::shared_ptr temperature_fwd; }; -DNNLSoftmaxFwd::Tensors::Tensors(const NDArray& data, const NDArray& output) - : data(data), out(output) {} - -DNNLSoftmaxFwd::DNNLSoftmaxFwd(const SoftmaxParam& param, - const Tensors& tensors, - const bool is_train) { - const float temperature = param.temperature.has_value() ? param.temperature.value() : 1.0f; - const int axis = CheckAxis(param.axis, tensors.data.shape().ndim()); - const auto input_mem = tensors.data.GetDNNLData(); - - softmax_pd = std::make_shared(GetSoftmaxFwdPd(*input_mem, axis, is_train)); - softmax_fwd = std::make_shared(*softmax_pd); - - if (temperature != 1.0f) { - temperature_pd = std::make_shared(GetTemperaturePd(*input_mem, temperature)); - temperature_fwd = std::make_shared(*temperature_pd); - } -} class DNNLSoftmaxBwd { public: @@ -126,27 +107,6 @@ class DNNLSoftmaxBwd { std::shared_ptr temperature_fwd; }; -DNNLSoftmaxBwd::Tensors::Tensors(const std::vector& inputs, - const std::vector& outputs) - : out_grad(inputs[0]), out(inputs[1]), data_grad(outputs[0]) {} - -DNNLSoftmaxBwd::DNNLSoftmaxBwd(const SoftmaxParam& param, const Tensors& tensors) { - const float temperature = param.temperature.has_value() ? param.temperature.value() : 1.0f; - const int axis = CheckAxis(param.axis, tensors.out.shape().ndim()); - const auto out_grad_mem = tensors.out_grad.GetDNNLData(); - const auto out_mem = tensors.out.GetDNNLData(); - const auto softmax_fwd_pd = DNNLSoftmaxFwd::GetSoftmaxFwdPd(*out_mem, axis, true); - - softmax_bwd_pd = std::make_shared( - GetSoftmaxBwdPd(*out_grad_mem, *out_mem, axis, softmax_fwd_pd)); - softmax_bwd = std::make_shared(*softmax_bwd_pd); - - if (temperature != 1.0f) { - temperature_pd = - std::make_shared(DNNLSoftmaxFwd::GetTemperaturePd(*out_mem, temperature)); - temperature_fwd = std::make_shared(*temperature_pd); - } -} } // namespace op } // namespace mxnet diff --git a/src/operator/nn/dnnl/dnnl_softmax.cc b/src/operator/nn/dnnl/dnnl_softmax.cc index 50c71e5e488b..b150908f10a7 100644 --- a/src/operator/nn/dnnl/dnnl_softmax.cc +++ b/src/operator/nn/dnnl/dnnl_softmax.cc @@ -69,6 +69,9 @@ void DNNLSoftmaxForward(const nnvm::NodeAttrs& attrs, fwd.Execute(tensors); } +DNNLSoftmaxFwd::Tensors::Tensors(const NDArray& data, const NDArray& output) + : data(data), out(output) {} + typedef ParamOpSign DNNLSoftmaxSignature; DNNLSoftmaxFwd& DNNLSoftmaxFwd::GetCached(const SoftmaxParam& param, const Tensors& tensors, @@ -95,6 +98,22 @@ DNNLSoftmaxFwd& DNNLSoftmaxFwd::GetCached(const SoftmaxParam& param, return it->second; } +DNNLSoftmaxFwd::DNNLSoftmaxFwd(const SoftmaxParam& param, + const Tensors& tensors, + const bool is_train) { + const float temperature = param.temperature.has_value() ? param.temperature.value() : 1.0f; + const int axis = CheckAxis(param.axis, tensors.data.shape().ndim()); + const auto input_mem = tensors.data.GetDNNLData(); + + softmax_pd = std::make_shared(GetSoftmaxFwdPd(*input_mem, axis, is_train)); + softmax_fwd = std::make_shared(*softmax_pd); + + if (temperature != 1.0f) { + temperature_pd = std::make_shared(GetTemperaturePd(*input_mem, temperature)); + temperature_fwd = std::make_shared(*temperature_pd); + } +} + softmax_fwd_pd_t DNNLSoftmaxFwd::GetSoftmaxFwdPd(const dnnl::memory& input_mem, const int axis, const bool is_train) { @@ -172,6 +191,11 @@ void DNNLSoftmaxBackward(const nnvm::NodeAttrs& attrs, bwd.Execute(tensors, req); } + +DNNLSoftmaxBwd::Tensors::Tensors(const std::vector& inputs, + const std::vector& outputs) + : out_grad(inputs[0]), out(inputs[1]), data_grad(outputs[0]) {} + DNNLSoftmaxBwd& DNNLSoftmaxBwd::GetCached(const SoftmaxParam& param, const Tensors& tensors) { #if DMLC_CXX11_THREAD_LOCAL static thread_local std::unordered_map bwds; @@ -195,6 +219,24 @@ DNNLSoftmaxBwd& DNNLSoftmaxBwd::GetCached(const SoftmaxParam& param, const Tenso return it->second; } +DNNLSoftmaxBwd::DNNLSoftmaxBwd(const SoftmaxParam& param, const Tensors& tensors) { + const float temperature = param.temperature.has_value() ? param.temperature.value() : 1.0f; + const int axis = CheckAxis(param.axis, tensors.out.shape().ndim()); + const auto out_grad_mem = tensors.out_grad.GetDNNLData(); + const auto out_mem = tensors.out.GetDNNLData(); + const auto softmax_fwd_pd = DNNLSoftmaxFwd::GetSoftmaxFwdPd(*out_mem, axis, /*is_train*/ true); + + softmax_bwd_pd = std::make_shared( + GetSoftmaxBwdPd(*out_grad_mem, *out_mem, axis, softmax_fwd_pd)); + softmax_bwd = std::make_shared(*softmax_bwd_pd); + + if (temperature != 1.0f) { // avoid dividing by 1 + temperature_pd = + std::make_shared(DNNLSoftmaxFwd::GetTemperaturePd(*out_mem, temperature)); + temperature_fwd = std::make_shared(*temperature_pd); + } +} + void DNNLSoftmaxBwd::Execute(const Tensors& tensors, const std::vector& req) const { DNNLStream* stream = DNNLStream::Get(); diff --git a/src/operator/nn/masked_softmax.cc b/src/operator/nn/masked_softmax.cc new file mode 100644 index 000000000000..6258a2993a8b --- /dev/null +++ b/src/operator/nn/masked_softmax.cc @@ -0,0 +1,140 @@ +/* + * 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 masked_softmax.cc + */ + +#include "./softmax-inl.h" +#include "../tensor/elemwise_unary_op.h" +#include "../tensor/elemwise_binary_op.h" +#include "../operator_common.h" +#if MXNET_USE_ONEDNN == 1 +#include "dnnl/dnnl_base-inl.h" +#include "dnnl/dnnl_ops-inl.h" +#endif + +namespace mxnet { +namespace op { +DMLC_REGISTER_PARAMETER(MaskedSoftmaxParam); + +#if MXNET_USE_ONEDNN == 1 +static void MaskedSoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs, + const OpContext& ctx, + const std::vector& inputs, + const std::vector& req, + const std::vector& outputs) { + if (inputs[0].shape().Size() == 0U) return; + const MaskedSoftmaxParam& param = nnvm::get(attrs.parsed); + if (SupportDNNLMaskedSoftmax(param, inputs, outputs[0])) { + DNNL_OPCHECK_INIT(false, outputs.size(), inputs, outputs); + + DNNLRun(DNNLMaskedSoftmaxForward, attrs, ctx, inputs, req, outputs); + + auto fn = MaskedSoftmaxCompute; + DNNL_OPCHECK_RUN(fn, attrs, ctx, inputs, req, outputs); + return; + } + FallBackCompute(MaskedSoftmaxCompute, attrs, ctx, + inputs, req, outputs); +} + +inline static bool MaskedSoftmaxStorageType(const nnvm::NodeAttrs& attrs, + const int dev_mask, + DispatchMode* dispatch_mode, + std::vector *in_attrs, + std::vector *out_attrs) { + const MaskedSoftmaxParam& param = nnvm::get(attrs.parsed); + CHECK_EQ(in_attrs->size(), 2U); + CHECK_EQ(out_attrs->size(), 1U); + + return DNNLStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs, + out_attrs); +} +#endif + +NNVM_REGISTER_OP(masked_softmax) +.add_alias("_npx_masked_softmax") +.describe(R"code(Applies the softmax function masking elements according to the mask provided)code" ADD_FILELINE) +.set_attr_parser(ParamParser) +.set_attr("FListInputNames", + [](const NodeAttrs& attrs){ + return std::vector{"data", "mask"}; + }) +.set_attr("FListOutputNames", + [](const NodeAttrs& attrs) { + return std::vector{"output"}; + }) +.set_attr("FCompute", MaskedSoftmaxCompute) +.set_attr("FGradient", + [](const nnvm::ObjectPtr& n, const std::vector& ograds) { + auto data_grad = MakeNode("_backward_masked_softmax", n->attrs.name + "_backward_data", + {ograds[0], n->inputs[1], nnvm::NodeEntry(n, 0, 0)}, + &n->attrs.dict, &n); + auto mask_grad = MakeNode("zeros_like", n->attrs.name + "_backward_mask", + {n->inputs[1]}, nullptr, &n); + std::vector ret; + ret.emplace_back(data_grad); + ret.emplace_back(mask_grad); + return ret; + }) +.set_attr("FInferType", MaskedSoftmaxOpType) +.set_num_inputs(2) +.set_num_outputs(1) +.set_attr("FInferShape", MaskedSoftmaxOpShape) +.set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) +.set_attr("FInplaceOption", + [](const NodeAttrs& attrs){ + return std::vector >{{0, 0}}; + }) +#if MXNET_USE_ONEDNN == 1 +.set_attr("TIsDNNL", true) +.set_attr("FComputeEx", MaskedSoftmaxComputeExCPU) +.set_attr("FInferStorageType", MaskedSoftmaxStorageType) +#endif +.add_argument("data", "NDArray-or-Symbol", "The input array.") +.add_argument("mask", "NDArray-or-Symbol", "Mask to apply.") +.add_arguments(MaskedSoftmaxParam::__FIELDS__()); + +NNVM_REGISTER_OP(_backward_masked_softmax) +.set_num_inputs(3) +.set_num_outputs(1) +.set_attr("FListInputNames", + [](const NodeAttrs& attrs){ + return std::vector{"ograd", "mask", "output"}; + }) +.set_attr("FInferShape", MaskedSoftmaxGradOpShape) +.set_attr("FInferType", MaskedSoftmaxGradOpType) +.set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) +.set_attr("FInplaceOption", MaskedSoftmaxGradOpInplaceOption) +.add_argument("args", "NDArray-or-Symbol[]", "Positional input arguments") +.set_attr_parser(ParamParser) +.set_attr("FCompute", MaskedSoftmaxGradCompute); +} // namespace op +} // namespace mxnet \ No newline at end of file diff --git a/src/operator/nn/softmax-inl.h b/src/operator/nn/softmax-inl.h index 71c205539efd..3023cd7177fc 100644 --- a/src/operator/nn/softmax-inl.h +++ b/src/operator/nn/softmax-inl.h @@ -1012,6 +1012,11 @@ struct MaskedSoftmaxParam : public dmlc::Parameter { (*dict)["temperature"] = temperature_s.str(); (*dict)["normalize"] = normalize_s.str(); } + + bool operator==(const MaskedSoftmaxParam& other) const { + return this->axis == other.axis && this->temperature == other.temperature && + this->normalize == other.normalize; + } }; static inline bool softmax_has_dtype_override(const nnvm::NodeAttrs& attrs) { @@ -1545,6 +1550,17 @@ struct hash { return ret; } }; + +template <> +struct hash { + size_t operator()(const mxnet::op::MaskedSoftmaxParam& val) { + size_t ret = 0; + ret = dmlc::HashCombine(ret, val.axis); + ret = dmlc::HashCombine(ret, val.temperature); + ret = dmlc::HashCombine(ret, val.normalize); + return ret; + } +}; } // namespace std #endif // MXNET_OPERATOR_NN_SOFTMAX_INL_H_ diff --git a/src/operator/nn/softmax.cc b/src/operator/nn/softmax.cc index 29f546458c19..004e5a2bb4e9 100644 --- a/src/operator/nn/softmax.cc +++ b/src/operator/nn/softmax.cc @@ -33,7 +33,6 @@ namespace mxnet { namespace op { DMLC_REGISTER_PARAMETER(SoftmaxParam); -DMLC_REGISTER_PARAMETER(MaskedSoftmaxParam); #if MXNET_USE_ONEDNN == 1 static void SoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs, @@ -194,68 +193,5 @@ NNVM_REGISTER_OP(_backward_softmax) .set_attr("FCompute", SoftmaxGradCompute); -NNVM_REGISTER_OP(masked_softmax) - .add_alias("_npx_masked_softmax") - .describe( - R"code(Applies the softmax function masking elements according to the mask provided)code" ADD_FILELINE) - .set_attr_parser(ParamParser) - .set_attr("FListInputNames", - [](const NodeAttrs& attrs) { - return std::vector{"data", "mask"}; - }) - .set_attr("FListOutputNames", - [](const NodeAttrs& attrs) { - return std::vector{"output"}; - }) - .set_attr("FCompute", MaskedSoftmaxCompute) - .set_attr( - "FGradient", - [](const nnvm::ObjectPtr& n, const std::vector& ograds) { - auto data_grad = MakeNode("_backward_masked_softmax", - n->attrs.name + "_backward_data", - {ograds[0], n->inputs[1], nnvm::NodeEntry(n, 0, 0)}, - &n->attrs.dict, - &n); - auto mask_grad = - MakeNode("zeros_like", n->attrs.name + "_backward_mask", {n->inputs[1]}, nullptr, &n); - std::vector ret; - ret.emplace_back(data_grad); - ret.emplace_back(mask_grad); - return ret; - }) - .set_attr("FInferType", MaskedSoftmaxOpType) - .set_num_inputs(2) - .set_num_outputs(1) - .set_attr("FInferShape", MaskedSoftmaxOpShape) - .set_attr("FResourceRequest", - [](const NodeAttrs& attrs) { - return std::vector{ResourceRequest::kTempSpace}; - }) - .set_attr("FInplaceOption", - [](const NodeAttrs& attrs) { - return std::vector >{{0, 0}}; - }) - .add_argument("data", "NDArray-or-Symbol", "The input array.") - .add_argument("mask", "NDArray-or-Symbol", "Mask to apply.") - .add_arguments(MaskedSoftmaxParam::__FIELDS__()); - -NNVM_REGISTER_OP(_backward_masked_softmax) - .set_num_inputs(3) - .set_num_outputs(1) - .set_attr("FListInputNames", - [](const NodeAttrs& attrs) { - return std::vector{"ograd", "mask", "output"}; - }) - .set_attr("FInferShape", MaskedSoftmaxGradOpShape) - .set_attr("FInferType", MaskedSoftmaxGradOpType) - .set_attr("FResourceRequest", - [](const NodeAttrs& attrs) { - return std::vector{ResourceRequest::kTempSpace}; - }) - .set_attr("FInplaceOption", MaskedSoftmaxGradOpInplaceOption) - .add_argument("args", "NDArray-or-Symbol[]", "Positional input arguments") - .set_attr_parser(ParamParser) - .set_attr("FCompute", - MaskedSoftmaxGradCompute); } // namespace op } // namespace mxnet From 719d4f6eff85d16bacd76b63ff74225a23535107 Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 27 Jan 2022 10:38:19 +0100 Subject: [PATCH 2/5] code refactor --- src/operator/nn/dnnl/dnnl_masked_softmax-inl.h | 10 +++++----- src/operator/nn/dnnl/dnnl_masked_softmax.cc | 9 ++------- src/operator/nn/masked_softmax.cc | 2 +- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h b/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h index 7a1fcca55f3c..ccf48097e902 100644 --- a/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h +++ b/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h @@ -26,10 +26,10 @@ #include #include -#include "./dnnl_base-inl.h" -#include "./dnnl_ops-inl.h" -#include "../softmax-inl.h" -#include "./dnnl_softmax-inl.h" +#include "dnnl_base-inl.h" +#include "dnnl_ops-inl.h" +#include "operator/nn/softmax-inl.h" +#include "dnnl_softmax-inl.h" #if MXNET_USE_ONEDNN == 1 @@ -83,4 +83,4 @@ DNNLMaskedSoftmaxFwd::Tensors::Tensors(const std::vector& inputs, } // namespace op } // namespace mxnet #endif // MXNET_USE_ONEDNN == 1 -#endif // MXNET_OPERATOR_NN_DNNL_DNNL_MASKED_SOFTMAX_INL_H_ \ No newline at end of file +#endif // MXNET_OPERATOR_NN_DNNL_DNNL_MASKED_SOFTMAX_INL_H_ diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax.cc b/src/operator/nn/dnnl/dnnl_masked_softmax.cc index 2846a44dd8d9..a92c72f22e09 100644 --- a/src/operator/nn/dnnl/dnnl_masked_softmax.cc +++ b/src/operator/nn/dnnl/dnnl_masked_softmax.cc @@ -21,8 +21,8 @@ * \file dnnl_maskedsoftmax.cc */ -#include "./dnnl_masked_softmax-inl.h" -#include "./dnnl_softmax-inl.h" +#include "dnnl_masked_softmax-inl.h" +#include "dnnl_softmax-inl.h" #if MXNET_USE_ONEDNN == 1 namespace mxnet { @@ -167,8 +167,6 @@ void DNNLMaskedSoftmaxFwd::Execute(const Tensors& tensors, SoftmaxParam p; p.axis = param.axis; p.temperature = param.temperature; - // int axis = CheckAxis(param.axis, input.shape().ndim()); - // auto softmax = DNNLSoftmaxFwd::GetSoftmaxFwd(p, axis, is_train, output, output); auto softmax_tensors = DNNLSoftmaxFwd::Tensors(output, output); auto softmax_op = DNNLSoftmaxFwd::GetCached(p, softmax_tensors, is_train); @@ -183,9 +181,6 @@ void DNNLMaskedSoftmaxFwd::Execute(const Tensors& tensors, stream->Submit(); // 2. out = softmax(out) - // stream->RegisterPrimArgs( - // softmax.GetFwd(), {{DNNL_ARG_SRC, *softmax_out_mem}, {DNNL_ARG_DST, *softmax_out_mem}}); - // stream->Submit(); softmax_op.Execute(softmax_tensors); // 3. out = out * mask diff --git a/src/operator/nn/masked_softmax.cc b/src/operator/nn/masked_softmax.cc index 6258a2993a8b..842e94861b90 100644 --- a/src/operator/nn/masked_softmax.cc +++ b/src/operator/nn/masked_softmax.cc @@ -137,4 +137,4 @@ NNVM_REGISTER_OP(_backward_masked_softmax) .set_attr("FCompute", MaskedSoftmaxGradCompute); } // namespace op -} // namespace mxnet \ No newline at end of file +} // namespace mxnet From ced7e26fa97511ab5008bdccf522ccc8c8566b2b Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 27 Jan 2022 11:20:44 +0100 Subject: [PATCH 3/5] fix sanity --- .../nn/dnnl/dnnl_masked_softmax-inl.h | 2 - src/operator/nn/dnnl/dnnl_softmax.cc | 3 +- src/operator/nn/masked_softmax.cc | 141 +++++++++--------- 3 files changed, 72 insertions(+), 74 deletions(-) diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h b/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h index ccf48097e902..683f743a106d 100644 --- a/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h +++ b/src/operator/nn/dnnl/dnnl_masked_softmax-inl.h @@ -38,8 +38,6 @@ namespace op { struct Primitives { public: - Primitives(){}; - dnnl::eltwise_forward minusone; dnnl::eltwise_forward::primitive_desc minusone_pd; diff --git a/src/operator/nn/dnnl/dnnl_softmax.cc b/src/operator/nn/dnnl/dnnl_softmax.cc index b150908f10a7..93fe5570cbe0 100644 --- a/src/operator/nn/dnnl/dnnl_softmax.cc +++ b/src/operator/nn/dnnl/dnnl_softmax.cc @@ -191,7 +191,6 @@ void DNNLSoftmaxBackward(const nnvm::NodeAttrs& attrs, bwd.Execute(tensors, req); } - DNNLSoftmaxBwd::Tensors::Tensors(const std::vector& inputs, const std::vector& outputs) : out_grad(inputs[0]), out(inputs[1]), data_grad(outputs[0]) {} @@ -230,7 +229,7 @@ DNNLSoftmaxBwd::DNNLSoftmaxBwd(const SoftmaxParam& param, const Tensors& tensors GetSoftmaxBwdPd(*out_grad_mem, *out_mem, axis, softmax_fwd_pd)); softmax_bwd = std::make_shared(*softmax_bwd_pd); - if (temperature != 1.0f) { // avoid dividing by 1 + if (temperature != 1.0f) { // avoid dividing by 1 temperature_pd = std::make_shared(DNNLSoftmaxFwd::GetTemperaturePd(*out_mem, temperature)); temperature_fwd = std::make_shared(*temperature_pd); diff --git a/src/operator/nn/masked_softmax.cc b/src/operator/nn/masked_softmax.cc index 842e94861b90..7dc0dc596da3 100644 --- a/src/operator/nn/masked_softmax.cc +++ b/src/operator/nn/masked_softmax.cc @@ -40,101 +40,102 @@ static void MaskedSoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { - if (inputs[0].shape().Size() == 0U) return; + if (inputs[0].shape().Size() == 0U) + return; const MaskedSoftmaxParam& param = nnvm::get(attrs.parsed); if (SupportDNNLMaskedSoftmax(param, inputs, outputs[0])) { DNNL_OPCHECK_INIT(false, outputs.size(), inputs, outputs); DNNLRun(DNNLMaskedSoftmaxForward, attrs, ctx, inputs, req, outputs); - auto fn = MaskedSoftmaxCompute; + auto fn = MaskedSoftmaxCompute; DNNL_OPCHECK_RUN(fn, attrs, ctx, inputs, req, outputs); return; } - FallBackCompute(MaskedSoftmaxCompute, attrs, ctx, - inputs, req, outputs); + FallBackCompute( + MaskedSoftmaxCompute, attrs, ctx, inputs, req, outputs); } inline static bool MaskedSoftmaxStorageType(const nnvm::NodeAttrs& attrs, const int dev_mask, DispatchMode* dispatch_mode, - std::vector *in_attrs, - std::vector *out_attrs) { + std::vector* in_attrs, + std::vector* out_attrs) { const MaskedSoftmaxParam& param = nnvm::get(attrs.parsed); CHECK_EQ(in_attrs->size(), 2U); CHECK_EQ(out_attrs->size(), 1U); - return DNNLStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs, - out_attrs); + return DNNLStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs, out_attrs); } #endif NNVM_REGISTER_OP(masked_softmax) -.add_alias("_npx_masked_softmax") -.describe(R"code(Applies the softmax function masking elements according to the mask provided)code" ADD_FILELINE) -.set_attr_parser(ParamParser) -.set_attr("FListInputNames", - [](const NodeAttrs& attrs){ - return std::vector{"data", "mask"}; - }) -.set_attr("FListOutputNames", - [](const NodeAttrs& attrs) { - return std::vector{"output"}; - }) -.set_attr("FCompute", MaskedSoftmaxCompute) -.set_attr("FGradient", - [](const nnvm::ObjectPtr& n, const std::vector& ograds) { - auto data_grad = MakeNode("_backward_masked_softmax", n->attrs.name + "_backward_data", - {ograds[0], n->inputs[1], nnvm::NodeEntry(n, 0, 0)}, - &n->attrs.dict, &n); - auto mask_grad = MakeNode("zeros_like", n->attrs.name + "_backward_mask", - {n->inputs[1]}, nullptr, &n); - std::vector ret; - ret.emplace_back(data_grad); - ret.emplace_back(mask_grad); - return ret; - }) -.set_attr("FInferType", MaskedSoftmaxOpType) -.set_num_inputs(2) -.set_num_outputs(1) -.set_attr("FInferShape", MaskedSoftmaxOpShape) -.set_attr("FResourceRequest", - [](const NodeAttrs& attrs) { - return std::vector{ResourceRequest::kTempSpace}; - }) -.set_attr("FInplaceOption", - [](const NodeAttrs& attrs){ - return std::vector >{{0, 0}}; - }) + .add_alias("_npx_masked_softmax") + .describe( + R"code(Applies the softmax function masking elements according to the mask provided)code" ADD_FILELINE) + .set_attr_parser(ParamParser) + .set_attr("FListInputNames", + [](const NodeAttrs& attrs) { + return std::vector{"data", "mask"}; + }) + .set_attr("FListOutputNames", + [](const NodeAttrs& attrs) { + return std::vector{"output"}; + }) + .set_attr("FCompute", MaskedSoftmaxCompute) + .set_attr( + "FGradient", + [](const nnvm::ObjectPtr& n, const std::vector& ograds) { + auto data_grad = MakeNode("_backward_masked_softmax", + n->attrs.name + "_backward_data", + {ograds[0], n->inputs[1], nnvm::NodeEntry(n, 0, 0)}, + &n->attrs.dict, + &n); + auto mask_grad = + MakeNode("zeros_like", n->attrs.name + "_backward_mask", {n->inputs[1]}, nullptr, &n); + std::vector ret; + ret.emplace_back(data_grad); + ret.emplace_back(mask_grad); + return ret; + }) + .set_attr("FInferType", MaskedSoftmaxOpType) + .set_num_inputs(2) + .set_num_outputs(1) + .set_attr("FInferShape", MaskedSoftmaxOpShape) + .set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) + .set_attr("FInplaceOption", + [](const NodeAttrs& attrs) { + return std::vector >{{0, 0}}; + }) #if MXNET_USE_ONEDNN == 1 -.set_attr("TIsDNNL", true) -.set_attr("FComputeEx", MaskedSoftmaxComputeExCPU) -.set_attr("FInferStorageType", MaskedSoftmaxStorageType) + .set_attr("TIsDNNL", true) + .set_attr("FComputeEx", MaskedSoftmaxComputeExCPU) + .set_attr("FInferStorageType", MaskedSoftmaxStorageType) #endif -.add_argument("data", "NDArray-or-Symbol", "The input array.") -.add_argument("mask", "NDArray-or-Symbol", "Mask to apply.") -.add_arguments(MaskedSoftmaxParam::__FIELDS__()); + .add_argument("data", "NDArray-or-Symbol", "The input array.") + .add_argument("mask", "NDArray-or-Symbol", "Mask to apply.") + .add_arguments(MaskedSoftmaxParam::__FIELDS__()); NNVM_REGISTER_OP(_backward_masked_softmax) -.set_num_inputs(3) -.set_num_outputs(1) -.set_attr("FListInputNames", - [](const NodeAttrs& attrs){ - return std::vector{"ograd", "mask", "output"}; - }) -.set_attr("FInferShape", MaskedSoftmaxGradOpShape) -.set_attr("FInferType", MaskedSoftmaxGradOpType) -.set_attr("FResourceRequest", - [](const NodeAttrs& attrs) { - return std::vector{ResourceRequest::kTempSpace}; - }) -.set_attr("FInplaceOption", MaskedSoftmaxGradOpInplaceOption) -.add_argument("args", "NDArray-or-Symbol[]", "Positional input arguments") -.set_attr_parser(ParamParser) -.set_attr("FCompute", MaskedSoftmaxGradCompute); + .set_num_inputs(3) + .set_num_outputs(1) + .set_attr("FListInputNames", + [](const NodeAttrs& attrs) { + return std::vector{"ograd", "mask", "output"}; + }) + .set_attr("FInferShape", MaskedSoftmaxGradOpShape) + .set_attr("FInferType", MaskedSoftmaxGradOpType) + .set_attr("FResourceRequest", + [](const NodeAttrs& attrs) { + return std::vector{ResourceRequest::kTempSpace}; + }) + .set_attr("FInplaceOption", MaskedSoftmaxGradOpInplaceOption) + .add_argument("args", "NDArray-or-Symbol[]", "Positional input arguments") + .set_attr_parser(ParamParser) + .set_attr("FCompute", + MaskedSoftmaxGradCompute); } // namespace op } // namespace mxnet From 45f631cbab646fe6e4c7a058fc460e5c409cb9f8 Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 3 Feb 2022 14:35:27 +0100 Subject: [PATCH 4/5] apply review --- src/operator/nn/dnnl/dnnl_masked_softmax.cc | 1 + src/operator/nn/masked_softmax.cc | 12 ++++++------ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax.cc b/src/operator/nn/dnnl/dnnl_masked_softmax.cc index a92c72f22e09..20935f53dc51 100644 --- a/src/operator/nn/dnnl/dnnl_masked_softmax.cc +++ b/src/operator/nn/dnnl/dnnl_masked_softmax.cc @@ -133,6 +133,7 @@ void DNNLMaskedSoftmaxFwd::Execute(const Tensors& tensors, using namespace mxnet::op::mxnet_op; using namespace dnnl; /* + Three steps of masked softmax: 1. out = input * [(mask - 1) * inf] 2. out = softmax(out) 3. out = out * mask diff --git a/src/operator/nn/masked_softmax.cc b/src/operator/nn/masked_softmax.cc index 7dc0dc596da3..386b53e8d968 100644 --- a/src/operator/nn/masked_softmax.cc +++ b/src/operator/nn/masked_softmax.cc @@ -21,13 +21,13 @@ * \file masked_softmax.cc */ -#include "./softmax-inl.h" -#include "../tensor/elemwise_unary_op.h" -#include "../tensor/elemwise_binary_op.h" -#include "../operator_common.h" +#include "softmax-inl.h" +#include "operator/tensor/elemwise_unary_op.h" +#include "operator/tensor/elemwise_binary_op.h" +#include "operator/operator_common.h" #if MXNET_USE_ONEDNN == 1 -#include "dnnl/dnnl_base-inl.h" -#include "dnnl/dnnl_ops-inl.h" +#include "operator/nn/dnnl/dnnl_base-inl.h" +#include "operator/nn/dnnl/dnnl_ops-inl.h" #endif namespace mxnet { From 105677a62a00c8cc3e85658259bd4b1bfb3b7729 Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Fri, 11 Feb 2022 15:30:42 +0100 Subject: [PATCH 5/5] Change variable name --- src/operator/nn/dnnl/dnnl_masked_softmax.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/operator/nn/dnnl/dnnl_masked_softmax.cc b/src/operator/nn/dnnl/dnnl_masked_softmax.cc index 20935f53dc51..789a6bfe7125 100644 --- a/src/operator/nn/dnnl/dnnl_masked_softmax.cc +++ b/src/operator/nn/dnnl/dnnl_masked_softmax.cc @@ -157,12 +157,12 @@ void DNNLMaskedSoftmaxFwd::Execute(const Tensors& tensors, // 1. B) out = out * inf const memory::desc converted_mask_desc = this->primitives->mask_input_pd.src1_desc(); - dnnl::memory* convertedmask_mem = TmpMemMgr::Get()->Alloc(converted_mask_desc); + dnnl::memory* converted_mask_mem = TmpMemMgr::Get()->Alloc(converted_mask_desc); dnnl::primitive_attr attr; attr.set_output_scales(0, {mshadow::red::limits::MaxValue()}); std::unordered_map args( - {{DNNL_ARG_FROM, *mask_minusone_mem}, {DNNL_ARG_TO, *convertedmask_mem}}); - stream->RegisterPrimArgs(dnnl::reorder(*mask_minusone_mem, *convertedmask_mem, attr), args); + {{DNNL_ARG_FROM, *mask_minusone_mem}, {DNNL_ARG_TO, *converted_mask_mem}}); + stream->RegisterPrimArgs(dnnl::reorder(*mask_minusone_mem, *converted_mask_mem, attr), args); // prepare softmax primitive and memory SoftmaxParam p; @@ -177,7 +177,7 @@ void DNNLMaskedSoftmaxFwd::Execute(const Tensors& tensors, // 1. C) out = input * out stream->RegisterPrimArgs(this->primitives->mask_input, {{DNNL_ARG_SRC_0, *input_mem}, - {DNNL_ARG_SRC_1, *convertedmask_mem}, + {DNNL_ARG_SRC_1, *converted_mask_mem}, {DNNL_ARG_DST, *softmax_out_mem}}); stream->Submit();