This is an automated email from the ASF dual-hosted git repository.
bgawrych pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new c70ea50014 [FEATURE] Integrate oneDNN support for scalar power()
operator (#20904)
c70ea50014 is described below
commit c70ea500142576439e1eb1e8725be6bed063457b
Author: AdamGrabowski <[email protected]>
AuthorDate: Tue Apr 26 10:31:34 2022 +0200
[FEATURE] Integrate oneDNN support for scalar power() operator (#20904)
* Integrate oneDNN support for power() operator from MXNet NumPy and
NDArray modules
* Move DNNLPowerForward function to dnnl file
* Fix include files, declaration, SupportDNNL function and OPCHECK macros
* Support only fp32
---
src/operator/nn/dnnl/dnnl_base-inl.h | 1 +
src/operator/nn/dnnl/dnnl_ops-inl.h | 6 ++
src/operator/nn/dnnl/dnnl_power_scalar-inl.h | 61 ++++++++++++++
src/operator/nn/dnnl/dnnl_power_scalar.cc | 94 ++++++++++++++++++++++
.../numpy/np_elemwise_broadcast_op_scalar.cc | 4 +
src/operator/tensor/elemwise_binary_scalar_op.h | 17 ++++
.../tensor/elemwise_binary_scalar_op_extended.cc | 33 ++++++++
7 files changed, 216 insertions(+)
diff --git a/src/operator/nn/dnnl/dnnl_base-inl.h
b/src/operator/nn/dnnl/dnnl_base-inl.h
index c1766d9355..8e3d4835d1 100644
--- a/src/operator/nn/dnnl/dnnl_base-inl.h
+++ b/src/operator/nn/dnnl/dnnl_base-inl.h
@@ -206,6 +206,7 @@ bool SupportDNNLSplit(const NDArray& input);
bool SupportDNNLStack(const std::vector<NDArray>& inputs);
bool SupportDNNLBinary(const std::vector<NDArray>& inputs);
bool SupportDNNLEltwise(const NDArray& input, const NDArray& output);
+bool SupportDNNLPower(const NDArray& input);
} // namespace op
static int GetTypeSize(int dtype) {
diff --git a/src/operator/nn/dnnl/dnnl_ops-inl.h
b/src/operator/nn/dnnl/dnnl_ops-inl.h
index 06ed1e0f26..8ccdaaceda 100644
--- a/src/operator/nn/dnnl/dnnl_ops-inl.h
+++ b/src/operator/nn/dnnl/dnnl_ops-inl.h
@@ -217,6 +217,12 @@ void DNNLWhereForward(const nnvm::NodeAttrs& attrs,
const std::vector<OpReqType>& req,
const std::vector<NDArray>& outputs);
+void DNNLPowerForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const NDArray& input,
+ const OpReqType& req,
+ const NDArray& output);
+
} // namespace op
} // namespace mxnet
diff --git a/src/operator/nn/dnnl/dnnl_power_scalar-inl.h
b/src/operator/nn/dnnl/dnnl_power_scalar-inl.h
new file mode 100644
index 0000000000..eddbffffff
--- /dev/null
+++ b/src/operator/nn/dnnl/dnnl_power_scalar-inl.h
@@ -0,0 +1,61 @@
+/*
+ * 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_power_scalar-inl.h
+ * \author: Adam Grabowski, [email protected]
+ */
+
+#ifndef MXNET_OPERATOR_NN_DNNL_DNNL_POWER_SCALAR_INL_H_
+#define MXNET_OPERATOR_NN_DNNL_DNNL_POWER_SCALAR_INL_H_
+
+#if MXNET_USE_ONEDNN == 1
+
+#include "dnnl_base-inl.h"
+#include "dnnl_ops-inl.h"
+#include "operator/tensor/elemwise_binary_scalar_op.h"
+
+namespace mxnet {
+namespace op {
+
+using eltwise_fwd_t = dnnl::eltwise_forward;
+using eltwise_fwd_pd_t = dnnl::eltwise_forward::primitive_desc;
+
+class DNNLPowerFwd {
+ public:
+ static DNNLPowerFwd& GetPowerForward(const nnvm::NodeAttrs& attrs,
+ const NDArray& input,
+ const NDArray& outputs);
+
+ DNNLPowerFwd(const NDArray& input, const float exponent);
+
+ void Execute(const NDArray& input, const OpReqType& req, const NDArray&
output);
+
+ private:
+ std::shared_ptr<eltwise_fwd_t> fwd;
+ std::shared_ptr<eltwise_fwd_pd_t> fwd_pd;
+};
+
+typedef OpSignature DNNLPowerSignature;
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_ONEDNN == 1
+#endif // MXNET_OPERATOR_NN_DNNL_DNNL_POWER_SCALAR_INL_H_
diff --git a/src/operator/nn/dnnl/dnnl_power_scalar.cc
b/src/operator/nn/dnnl/dnnl_power_scalar.cc
new file mode 100644
index 0000000000..223ab21c1c
--- /dev/null
+++ b/src/operator/nn/dnnl/dnnl_power_scalar.cc
@@ -0,0 +1,94 @@
+/*
+ * 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_power_scalar.cc
+ * \author: Adam Grabowski, [email protected]
+ */
+
+#if MXNET_USE_ONEDNN == 1
+
+#include "dnnl_power_scalar-inl.h"
+
+namespace mxnet {
+namespace op {
+
+DNNLPowerFwd& DNNLPowerFwd::GetPowerForward(const nnvm::NodeAttrs& attrs,
+ const NDArray& input,
+ const NDArray& output) {
+ const NumpyBinaryScalarParam& param =
nnvm::get<NumpyBinaryScalarParam>(attrs.parsed);
+#if DMLC_CXX11_THREAD_LOCAL
+ static thread_local std::unordered_map<DNNLPowerSignature, DNNLPowerFwd,
OpHash> fwds;
+#else
+ static MX_THREAD_LOCAL std::unordered_map<DNNLPowerSignature, DNNLPowerFwd,
OpHash> fwds;
+#endif
+ DNNLPowerSignature key;
+ key.AddSign(static_cast<float>(param.scalar));
+ key.AddSign(input);
+ key.AddSign(output);
+
+ auto it = fwds.find(key);
+ if (it == fwds.end()) {
+ const DNNLPowerFwd fwd(input, static_cast<float>(param.scalar));
+ it = AddToCache(&fwds, key, fwd);
+ }
+ return it->second;
+}
+
+DNNLPowerFwd::DNNLPowerFwd(const NDArray& input, const float exponent) {
+ auto src_desc = input.GetDNNLData()->get_desc();
+ dnnl::eltwise_forward::desc fwd_desc(
+ dnnl::prop_kind::forward_scoring, dnnl::algorithm::eltwise_pow,
src_desc, 1, exponent);
+ fwd_pd = std::make_shared<eltwise_fwd_pd_t>(fwd_desc,
mxnet::CpuEngine::Get()->get_engine());
+ fwd = std::make_shared<eltwise_fwd_t>(*fwd_pd);
+}
+
+void DNNLPowerFwd::Execute(const NDArray& input, const OpReqType& req, const
NDArray& output) {
+ auto engine = mxnet::CpuEngine::Get()->get_engine();
+ auto src = input.GetDNNLData();
+ dnnl_output_t out_mem = CreateDNNLMem(output, fwd_pd->dst_desc(), req,
&input);
+
+ dnnl_args_map_t args = {
+ {DNNL_ARG_SRC, *src},
+ {DNNL_ARG_DST, *out_mem.second},
+ };
+
+ DNNLStream::Get()->RegisterPrimArgs(*fwd, args);
+ CommitOutput(output, out_mem);
+ DNNLStream::Get()->Submit();
+}
+
+void DNNLPowerForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const NDArray& input,
+ const OpReqType& req,
+ const NDArray& output) {
+ DNNLPowerFwd& fwd = DNNLPowerFwd::GetPowerForward(attrs, input, output);
+ fwd.Execute(input, req, output);
+}
+
+bool SupportDNNLPower(const NDArray& input) {
+ return input.shape().Size() != 0 && input.shape().ndim() > 0 &&
input.shape().ndim() <= 6 &&
+ input.dtype() == mshadow::kFloat32;
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_ONEDNN == 1
diff --git a/src/operator/numpy/np_elemwise_broadcast_op_scalar.cc
b/src/operator/numpy/np_elemwise_broadcast_op_scalar.cc
index e4e61d1226..0e7cdaf3df 100644
--- a/src/operator/numpy/np_elemwise_broadcast_op_scalar.cc
+++ b/src/operator/numpy/np_elemwise_broadcast_op_scalar.cc
@@ -55,6 +55,10 @@ MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_rmod_scalar)
MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_power_scalar)
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu,
mshadow_op::power>)
+#if MXNET_USE_ONEDNN == 1
+ .set_attr<FComputeEx>("FComputeEx<cpu>", PowerComputeExCPU)
+ .set_attr<FInferStorageType>("FInferStorageType", PowerStorageType)
+#endif
.set_attr<nnvm::FGradient>("FGradient",
ElemwiseGradUseIn{"_backward_power_scalar"});
MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_rpower_scalar)
diff --git a/src/operator/tensor/elemwise_binary_scalar_op.h
b/src/operator/tensor/elemwise_binary_scalar_op.h
index 8c025ef9ec..f18cc16c44 100644
--- a/src/operator/tensor/elemwise_binary_scalar_op.h
+++ b/src/operator/tensor/elemwise_binary_scalar_op.h
@@ -33,6 +33,9 @@
#include "../elemwise_op_common.h"
#include "../../common/alm.h"
#include "elemwise_unary_op.h"
+#if MXNET_USE_ONEDNN == 1
+#include "operator/nn/dnnl/dnnl_power_scalar-inl.h"
+#endif
namespace mxnet {
namespace op {
@@ -441,6 +444,20 @@ class BinaryScalarOp : public UnaryOp {
}
};
+#if MXNET_USE_ONEDNN == 1
+bool PowerStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* inputs,
+ std::vector<int>* outputs);
+
+void PowerComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<mxnet::NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<mxnet::NDArray>& outputs);
+#endif
+
#define MXNET_OPERATOR_REGISTER_BINARY_SCALAR(name)
\
NNVM_REGISTER_OP(name)
\
.set_num_inputs(1)
\
diff --git a/src/operator/tensor/elemwise_binary_scalar_op_extended.cc
b/src/operator/tensor/elemwise_binary_scalar_op_extended.cc
index 3c4938074f..564951138a 100644
--- a/src/operator/tensor/elemwise_binary_scalar_op_extended.cc
+++ b/src/operator/tensor/elemwise_binary_scalar_op_extended.cc
@@ -49,8 +49,41 @@ MXNET_OPERATOR_REGISTER_BINARY(_backward_minimum_scalar)
.set_attr_parser(ParamParser<NumpyBinaryScalarParam>)
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Backward<cpu,
mshadow_op::le>);
+#if MXNET_USE_ONEDNN == 1
+bool PowerStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* inputs,
+ std::vector<int>* outputs) {
+ CHECK_EQ(inputs->size(), 1);
+ CHECK_EQ(outputs->size(), 1);
+
+ return DNNLStorageType(attrs, dev_mask, true, dispatch_mode, inputs,
outputs);
+}
+
+void PowerComputeExCPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<mxnet::NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<mxnet::NDArray>& outputs) {
+ if (SupportDNNLPower(inputs[0])) {
+ DNNL_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ DNNLRun(DNNLPowerForward, attrs, ctx, inputs[0], req[0], outputs[0]);
+ DNNL_OPCHECK_RUN(
+ (BinaryScalarOp::Compute<cpu, mshadow_op::power>), attrs, ctx, inputs,
req, outputs);
+ } else {
+ FallBackCompute(
+ BinaryScalarOp::Compute<cpu, mshadow_op::power>, attrs, ctx, inputs,
req, outputs);
+ }
+}
+#endif // MXNET_USE_ONEDNN == 1
+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_power_scalar)
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu,
mshadow_op::power>)
+#if MXNET_USE_ONEDNN == 1
+ .set_attr<FComputeEx>("FComputeEx<cpu>", PowerComputeExCPU)
+ .set_attr<FInferStorageType>("FInferStorageType", PowerStorageType)
+#endif
.set_attr<nnvm::FGradient>("FGradient",
ElemwiseGradUseIn{"_backward_power_scalar"})
.add_alias("_PowerScalar");
