[mkldnn-v1.0] Add MKL-DNN reshape&flatten&expand_dims (#16258)

* Add mkldnn 1.0 support for reshape/flatten/expanddims ops

* improve log & modify definition location of args_map_

* fix comments

* rebase code

* trigger CI

* trigger CI

* trigger CI

* trigger CI

src/operator/nn/mkldnn/mkldnn_base-inl.h

@@ -189,7 +189,7 @@ bool SupportMKLDNNDeconv(const DeconvolutionParam& params, const NDArray &input)
 bool SupportMKLDNNSoftmax(const SoftmaxParam& param, const NDArray &input, const NDArray &output);
 bool SupportMKLDNNSoftmaxOutput(const SoftmaxOutputParam &param);
 bool SupportMKLDNNTranspose(const TransposeParam& param, const NDArray &data);
-bool SupportMKLDNNReshape(const ReshapeParam &param, const NDArray &data);
+bool SupportMKLDNNReshape(const NDArray &in_data, const NDArray &out_data);
 } // namespace op
 
 static int GetTypeSize(int dtype) {

src/operator/nn/mkldnn/mkldnn_expand_dims.cc

@@ -0,0 +1,70 @@
+/*
+ * 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
+ *
+ * https://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_expand_dims.cc
+ * \brief Implement expand_dims operator via MKL-DNN reorder primitive
+ * \author Wuxun Zhang
+*/
+
+#if MXNET_USE_MKLDNN == 100
+
+#include "mkldnn_reshape-inl.h"
+
+namespace mxnet {
+namespace op {
+
+class MKLDNNExpandDimsFwd : public MKLDNNReshapeFwd {
+ public:
+ explicit MKLDNNExpandDimsFwd(const OpReqType &req,
+ const NDArray &input,
+ const NDArray &output)
+ : MKLDNNReshapeFwd(req, input, output) {}
+};
+
+typedef ParamOpSign<ExpandDimParam> MKLDNNExpandDimsSignature;
+
+void MKLDNNExpandDimsForward(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const NDArray &input,
+ const OpReqType &req,
+ const NDArray &output) {
+ const ExpandDimParam& param = nnvm::get<ExpandDimParam>(attrs.parsed);
+ if (req == kNullOp) return;
+ CHECK_NE(req, kAddTo) << "kAddTo is not supported yet";
+
+ auto fwd = GetCachedForward<MKLDNNExpandDimsFwd, ExpandDimParam,
+ MKLDNNExpandDimsSignature>(param, req, input, output);
+
+ auto ws_size = fwd.GetWorkspaceSize();
+ void* ws_ptr = nullptr;
+ if (ws_size) {
+ mshadow::Stream<cpu> *s = ctx.get_stream<cpu>();
+ mshadow::Tensor<cpu, 1, char> ws = ctx.requested[0]
+ .get_space_typed<cpu, 1, char>(mshadow::Shape1(ws_size), s);
+ ws_ptr = reinterpret_cast<void*>(ws.dptr_);
+ }
+
+ fwd.Execute(input, output, req, ws_ptr);
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif

src/operator/nn/mkldnn/mkldnn_flatten-inl.h

@@ -25,7 +25,7 @@
 
 #ifndef MXNET_OPERATOR_NN_MKLDNN_MKLDNN_FLATTEN_INL_H_
 #define MXNET_OPERATOR_NN_MKLDNN_MKLDNN_FLATTEN_INL_H_
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
 
 #include "mkldnn_reshape-inl.h"
 

src/operator/nn/mkldnn/mkldnn_flatten.cc

@@ -19,11 +19,11 @@
 
 /*!
  * \file mkldnn_flatten.cc
- * \brief Implement flatten operator by using mkldnn reorder primitive
+ * \brief Implement flatten operator via using MKL-DNN reorder primitive
  * \author Wuxun Zhang
 */
 
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
 
 #include "mkldnn_flatten-inl.h"
 
@@ -70,7 +70,7 @@ void MKLDNNFlattenForward(const nnvm::NodeAttrs &attrs,
  ws_ptr = reinterpret_cast<void*>(ws.dptr_);
  }
 
- fwd.Execute(input, output, ws_ptr);
+ fwd.Execute(input, output, req, ws_ptr);
 }
 
 } // namespace op

src/operator/nn/mkldnn/mkldnn_ops-inl.h

@@ -63,18 +63,6 @@ void MKLDNNConcatBackward(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
  const std::vector<NDArray>& inputs,
  const std::vector<OpReqType>& req,
  const std::vector<NDArray>& outputs);
-
-void MKLDNNReshapeForward(const nnvm::NodeAttrs& attrs,
- const OpContext &ctx,
- const NDArray &input,
- const OpReqType &req,
- const NDArray &output);
-
-void MKLDNNFlattenForward(const nnvm::NodeAttrs &attrs,
- const OpContext &ctx,
- const NDArray &input,
- const OpReqType &req,
- const NDArray &output);
 #endif
 
 #if MXNET_USE_MKLDNN == 100
@@ -142,6 +130,22 @@ void MKLDNNTransposeForward(const nnvm::NodeAttrs& attrs,
  const NDArray &data,
  const OpReqType &req,
  const NDArray &output);
+
+void MKLDNNReshapeForward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const NDArray &input,
+ const OpReqType &req,
+ const NDArray &output);
+void MKLDNNFlattenForward(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const NDArray &input,
+ const OpReqType &req,
+ const NDArray &output);
+void MKLDNNExpandDimsForward(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const NDArray &input,
+ const OpReqType &req,
+ const NDArray &output);
 #endif
 
 } // namespace op

src/operator/nn/mkldnn/mkldnn_reshape-inl.h

@@ -26,7 +26,7 @@
 #ifndef MXNET_OPERATOR_NN_MKLDNN_MKLDNN_RESHAPE_INL_H_
 #define MXNET_OPERATOR_NN_MKLDNN_MKLDNN_RESHAPE_INL_H_
 
-#if MXNET_USE_MKLDNN == 1
+#if MXNET_USE_MKLDNN == 100
 #include <vector>
 #include "mkldnn_base-inl.h"
 #include "../../tensor/matrix_op-inl.h"
@@ -36,7 +36,6 @@ namespace op {
 
 class MKLDNNReshapeFwd {
  protected:
- std::shared_ptr<mkldnn::memory> data_;
  std::shared_ptr<mkldnn::memory> out_;
  std::shared_ptr<mkldnn::memory> temp_;
  std::vector<mkldnn::primitive> prims_;
@@ -47,15 +46,39 @@ class MKLDNNReshapeFwd {
  const NDArray &input,
  const NDArray &output);
  int GetWorkspaceSize();
- void SetNewMem(const NDArray &input,
- const NDArray &output,
- void* workspace = nullptr);
  void Execute(const NDArray &input,
  const NDArray &output,
+ const OpReqType &req,
  void* workspace = nullptr);
 };
 
 typedef ParamOpSign<ReshapeParam> MKLDNNReshapeSignature;
+
+template<typename MKLDNNOpFwdType, typename ParamType, typename MKLDNNSigatureType>
+MKLDNNOpFwdType &GetCachedForward(const ParamType& param,
+ const OpReqType &req,
+ const NDArray &input,
+ const NDArray &output) {
+#if DMLC_CXX11_THREAD_LOCAL
+ static thread_local std::unordered_map<MKLDNNSigatureType,
+ MKLDNNOpFwdType, OpHash> fwds;
+#else
+ static MX_THREAD_LOCAL std::unordered_map<MKLDNNSigatureType,
+ MKLDNNOpFwdType, OpHash> fwds;
+#endif
+ MKLDNNSigatureType key(param);
+ key.AddSign(req);
+ key.AddSign(input);
+ key.AddSign(output);
+
+ auto it = fwds.find(key);
+ if (it == fwds.end()) {
+ MKLDNNOpFwdType fwd(req, input, output);
+ it = AddToCache(&fwds, key, fwd);
+ }
+ return it->second;
+}
+
 MKLDNNReshapeFwd &GetReshapeForward(const ReshapeParam& param,
  const OpReqType &req,
  const NDArray &input,