[FEATURE] Add _npi_power_scalar and _npi_multiply_scalar fuse (#20976)

* [FEATURE] Add _npi_power_scalar and _npi_multiply_scalar fuse

* Merge _npi_power_scalar implementation with implementation of this fuse

* Fix clang

* Fix CI

* Fix review and simplify the implementation

* Add checks for the amount of inputs and outputs

* Fix CI

* Add Reset() function

* Fix DNNLPowMulScalarShape and Type functions

* Fix DNNLPowMulScalarType

* Fix DNNLPowMulScalarType

* Add generic implementation for sq_pow_mul_scalar operator

* Fix sanity

* Fix req

* Add Filter method to property

* Add new line

* Fix gpu CI

* Add '_sg_pow_mul_scalar' to symbol_fp16.py

* Fix CI on MacOS

* Fix SupportDNNL*

* Make PowMulScalarCompute more readable

* Fix PowMulScalarCompute

* Fix memory usage

* Fix build

Author: bartekkuncer

include/mxnet/op_attr_types.h

@@ -58,7 +58,7 @@ enum OpReqType {
 };
 
 /*!
- * \brief All the possible information needed by Operator.Forward and Backward
+ * \brief All the possible information needed by Operator.
  *  This is the superset of RunContext.
  *  We use this data structure to bookkeep everything needed by Forward and Backward.
  * \sa Resource

include/mxnet/tensor_blob.h

@@ -210,7 +210,7 @@ class TBlob {
     CHECK(Device::kDevMask == this->dev_mask())
         << "TBlob.get: device type do not match specified type";
     CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
-        << "TBlob.get_with_shape: data type do not match specified type."
+        << "TBlob.get_with_shape: data type do not match specified type. "
         << "Expected: " << mshadow::dtype_string(type_flag_) << " v.s. given "
         << mshadow::dtype_string(mshadow::DataType<DType>::kFlag);
     return mshadow::Tensor<Device, 2, DType>(static_cast<DType*>(dptr_), shape_.FlatTo2D(), stream);
@@ -248,7 +248,7 @@ class TBlob {
   template <typename DType>
   inline DType* dptr() const {
     CHECK(mshadow::DataType<DType>::kFlag == type_flag_)
-        << "TBlob.get_with_shape: data type do not match specified type."
+        << "TBlob.get_with_shape: data type do not match specified type. "
         << "Expected: " << mshadow::dtype_string(type_flag_) << " v.s. given "
         << mshadow::dtype_string(mshadow::DataType<DType>::kFlag);
     return static_cast<DType*>(dptr_);

python/mxnet/amp/lists/symbol_fp16.py

@@ -636,7 +636,8 @@
         '_sg_onednn_fully_connected',
         '_sg_onednn_selfatt_qk',
         '_sg_onednn_selfatt_valatt',
-        '_sg_onednn_batch_dot'
+        '_sg_onednn_batch_dot',
+        '_sg_pow_mul_scalar'
     ])
 
 # Functions that have to be cast to FP32 only for

src/operator/nn/dnnl/dnnl_base-inl.h

@@ -254,7 +254,6 @@ bool SupportDNNLLeakyRelu(const LeakyReLUParam& param);
 bool SupportDNNLLeakyRelu(const LeakyReLUParam& param, const NDArray& input);
 bool SupportDNNLLogSoftmax(const SoftmaxParam& param, const NDArray& input);
 bool SupportDNNLMaskedSoftmax(const MaskedSoftmaxParam& param, const std::vector<NDArray>& input);
-bool SupportDNNLPower(const NDArray& input);
 bool SupportDNNLQuantizedAct(const ActivationParam& param);
 bool SupportDNNLReshape(const NDArray& input);
 bool SupportDNNLSlice(const SliceParam& param, const NDArray& input, const NDArray& output);

src/operator/nn/dnnl/dnnl_pow_mul_scalar-inl.h

@@ -0,0 +1,100 @@
+/*
+ * 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_pow_mul_scalar-inl.h
+ */
+
+#ifndef MXNET_OPERATOR_NN_DNNL_DNNL_POW_MUL_SCALAR_INL_H_
+#define MXNET_OPERATOR_NN_DNNL_DNNL_POW_MUL_SCALAR_INL_H_
+
+#if MXNET_USE_ONEDNN == 1
+
+#include <vector>
+
+#include "operator/tensor/elemwise_binary_scalar_op.h"
+
+namespace mxnet {
+namespace op {
+
+struct DNNLPowMulScalarParam : public dmlc::Parameter<DNNLPowMulScalarParam> {
+  float exponent;
+  float multiplier;
+  bool exp_is_int;
+  bool mul_is_int;
+
+  DMLC_DECLARE_PARAMETER(DNNLPowMulScalarParam) {
+    DMLC_DECLARE_FIELD(exponent).describe("Exponent for power operation.").set_default(1);
+    DMLC_DECLARE_FIELD(multiplier).describe("Multiplier for multiply operation.").set_default(1);
+    DMLC_DECLARE_FIELD(exp_is_int)
+        .describe("Indicate whether exponent is int type.")
+        .set_default(true);
+    DMLC_DECLARE_FIELD(mul_is_int)
+        .describe("Indicate whether multiplier is int type.")
+        .set_default(true);
+  }
+
+  bool operator==(const DNNLPowMulScalarParam& other) const {
+    return this->exponent == other.exponent && this->multiplier == other.multiplier &&
+           this->exp_is_int == other.exp_is_int && this->mul_is_int == other.mul_is_int;
+  }
+};
+
+using eltwise_fwd_t    = dnnl::eltwise_forward;
+using eltwise_fwd_pd_t = dnnl::eltwise_forward::primitive_desc;
+
+typedef ParamOpSign<DNNLPowMulScalarParam> DNNLPowMulScalarSignature;
+
+class DNNLPowMulScalarFwd {
+ public:
+  static DNNLPowMulScalarFwd& GetCached(const DNNLPowMulScalarParam& param,
+                                        const NDArray& input,
+                                        const NDArray& output);
+
+  DNNLPowMulScalarFwd(const DNNLPowMulScalarParam& param, const NDArray& input);
+
+  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;
+};
+
+template <bool subgraph>
+inline void DNNLPowMulScalarForward(const nnvm::NodeAttrs& attrs,
+                                    const OpContext& ctx,
+                                    const std::vector<NDArray>& inputs,
+                                    const std::vector<OpReqType>& req,
+                                    const std::vector<NDArray>& outputs) {
+  DNNLPowMulScalarParam param;
+  if (subgraph) {
+    param = nnvm::get<DNNLPowMulScalarParam>(attrs.parsed);
+  } else {
+    param.multiplier = 1;
+    param.exponent   = nnvm::get<NumpyBinaryScalarParam>(attrs.parsed).scalar;
+  }
+  DNNLPowMulScalarFwd& fwd = DNNLPowMulScalarFwd::GetCached(param, inputs[0], outputs[0]);
+  fwd.Execute(inputs[0], req[0], outputs[0]);
+}
+
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // MXNET_USE_ONEDNN == 1
+#endif  // MXNET_OPERATOR_NN_DNNL_DNNL_POW_MUL_SCALAR_INL_H_

src/operator/nn/dnnl/dnnl_power_scalar.cc …/operator/nn/dnnl/dnnl_pow_mul_scalar.ccsrc/operator/nn/dnnl/dnnl_power_scalar.cc renamed to src/operator/nn/dnnl/dnnl_pow_mul_scalar.cc src/operator/nn/dnnl/dnnl_power_scalar.cc renamed to src/operator/nn/dnnl/dnnl_pow_mul_scalar.cc

@@ -18,49 +18,54 @@
  */
 
 /*!
- * \file dnnl_power_scalar.cc
- * \author: Adam Grabowski, adam.grabowski@intel.com
+ * \file dnnl_pow_mul_scalar.cc
  */
 
 #if MXNET_USE_ONEDNN == 1
 
-#include "dnnl_power_scalar-inl.h"
+#include "dnnl_pow_mul_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);
+DMLC_REGISTER_PARAMETER(DNNLPowMulScalarParam);
+
+DNNLPowMulScalarFwd& DNNLPowMulScalarFwd::GetCached(const DNNLPowMulScalarParam& param,
+                                                    const NDArray& input,
+                                                    const NDArray& output) {
 #if DMLC_CXX11_THREAD_LOCAL
-  static thread_local std::unordered_map<DNNLPowerSignature, DNNLPowerFwd, OpHash> fwds;
+  static thread_local std::unordered_map<DNNLPowMulScalarSignature, DNNLPowMulScalarFwd, OpHash>
+      fwds;
 #else
-  static MX_THREAD_LOCAL std::unordered_map<DNNLPowerSignature, DNNLPowerFwd, OpHash> fwds;
+  static MX_THREAD_LOCAL std::unordered_map<DNNLPowMulScalarSignature, DNNLPowMulScalarFwd, OpHash>
+      fwds;
 #endif
-  DNNLPowerSignature key;
-  key.AddSign(static_cast<float>(param.scalar));
+  DNNLPowMulScalarSignature key(param);
   key.AddSign(input);
   key.AddSign(output);
 
   auto it = fwds.find(key);
   if (it == fwds.end()) {
-    const DNNLPowerFwd fwd(input, static_cast<float>(param.scalar));
+    const DNNLPowMulScalarFwd fwd(param, input);
     it = AddToCache(&fwds, key, fwd);
   }
   return it->second;
 }
 
-DNNLPowerFwd::DNNLPowerFwd(const NDArray& input, const float exponent) {
+DNNLPowMulScalarFwd::DNNLPowMulScalarFwd(const DNNLPowMulScalarParam& param, const NDArray& input) {
   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);
+  dnnl::eltwise_forward::desc fwd_desc(dnnl::prop_kind::forward_scoring,
+                                       dnnl::algorithm::eltwise_pow,
+                                       src_desc,
+                                       param.multiplier,
+                                       param.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();
+void DNNLPowMulScalarFwd::Execute(const NDArray& input,
+                                  const OpReqType& req,
+                                  const NDArray& output) {
   auto src              = input.GetDNNLData();
   dnnl_output_t out_mem = CreateDNNLMem(output, fwd_pd->dst_desc(), req, &input);
 
@@ -73,22 +78,18 @@ void DNNLPowerFwd::Execute(const NDArray& input, const OpReqType& req, const NDA
   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
 
+namespace std {
+template <>
+struct hash<mxnet::op::DNNLPowMulScalarParam> {
+  size_t operator()(const mxnet::op::DNNLPowMulScalarParam& val) {
+    size_t ret = 0;
+    ret        = dmlc::HashCombine(ret, val.exponent);
+    ret        = dmlc::HashCombine(ret, val.multiplier);
+    return ret;
+  }
+};
+}  // namespace std
 #endif  // MXNET_USE_ONEDNN == 1

src/operator/nn/dnnl/dnnl_power_scalar-inl.h

@@ -29,6 +29,8 @@
 namespace mxnet {
 namespace op {
 
+enum SelectStatus { kFail = 0, kStart, kSuccess };
+
 inline uint32_t DefaultSubgraphOpNumInputs(const nnvm::NodeAttrs& attrs) {
   const nnvm::Symbol& sym = *attrs.subgraphs[0];
   return sym.ListInputNames(nnvm::Symbol::kAll).size();

src/operator/subgraph/common.h

@@ -35,8 +35,6 @@ namespace op {
 
 class SgDNNLBNReLUSelector : public SubgraphSelector {
  public:
-  enum SelectStatus { kStart, kSuccess, kFail };
-
   explicit SgDNNLBNReLUSelector(const bool disable_bn_relu)
       : disable_bn_relu_(disable_bn_relu), status_(kStart) {}
 

src/operator/subgraph/dnnl/dnnl_bn_relu_property.h

@@ -37,7 +37,7 @@ namespace op {
 class SgDNNLConvSelector : public SubgraphSelector {
  public:
   /*! \brief pattern match status_ */
-  enum SelectStatus {
+  enum SelectStatusConv {
     kFail = 0,
     kStart,
     kBN,
@@ -51,7 +51,7 @@ class SgDNNLConvSelector : public SubgraphSelector {
   bool disable_conv_act_;
   bool disable_conv_sum_;
   bool quantize_;
-  SelectStatus status_;
+  SelectStatusConv status_;
   std::vector<const nnvm::Node*> matched_list_;
 
  public: