fix fp32 flatten issue

src/operator/nn/mkldnn/mkldnn_flatten.cc

@@ -0,0 +1,87 @@
+/*
+ * 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_flatten.cc
+ * \brief Implement flatten operator by using mkldnn reorder primitive
+ * \author Wuxun Zhang
+*/
+
+#if MXNET_USE_MKLDNN == 1
+
+#include "mkldnn_reshape-inl.h"
+
+namespace mxnet {
+namespace op {
+
+class MKLDNNFlattenFwd : public MKLDNNReshapeFwd {
+ public:
+  explicit MKLDNNFlattenFwd(const OpReqType &req,
+                            const NDArray &input,
+                            const NDArray &output)
+    : MKLDNNReshapeFwd(req, input, output) {}
+};
+
+static MKLDNNFlattenFwd &GetFlattenForward(const OpReqType &req,
+                                           const NDArray &input,
+                                           const NDArray &output) {
+#if DMLC_CXX11_THREAD_LOCAL
+  static thread_local std::unordered_map<OpSignature,
+                                         MKLDNNFlattenFwd, OpHash> fwds;
+#else
+  static MX_THREAD_LOCAL std::unordered_map<OpSignature,
+                                            MKLDNNFlattenFwd, OpHash> fwds;
+#endif
+  OpSignature key;
+  key.AddSign(req);
+  key.AddSign(input);
+
+  auto it = fwds.find(key);
+  if (it == fwds.end()) {
+    MKLDNNFlattenFwd fwd(req, input, output);
+    it = AddToCache(&fwds, key, fwd);
+  }
+  return it->second;
+}
+
+void MKLDNNFlattenForward(const nnvm::NodeAttrs &attrs,
+                          const OpContext &ctx,
+                          const NDArray &input,
+                          const OpReqType &req,
+                          const NDArray &output) {
+  if (req == kNullOp) return;
+  CHECK_NE(req, kAddTo) << "kAddTo is not supported yet";
+
+  auto fwd = GetFlattenForward(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, ws_ptr);
+}
+
+}  // namespace op
+}  // namespace mxnet
+
+#endif

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

@@ -119,12 +119,11 @@ void MKLDNNTransposeForward(const nnvm::NodeAttrs& attrs,
                             const OpReqType &req,
                             const NDArray &output);
 
-void MKLDNNReshapeForward(const nnvm::NodeAttrs &attrs,
+void MKLDNNFlattenForward(const nnvm::NodeAttrs &attrs,
                           const OpContext &ctx,
-                          const NDArray &data,
+                          const NDArray &input,
                           const OpReqType &req,
                           const NDArray &output);
-
 }  // namespace op
 }  // namespace mxnet
 #endif  // MXNET_USE_MKLDNN == 1

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

@@ -0,0 +1,84 @@
+/*
+ * 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.
+ */
+
+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file mkldnn_reshape-inl.h
+ * \brief Function definition of mkldnn reshape operator
+ */
+
+#ifndef MXNET_OPERATOR_NN_MKLDNN_MKLDNN_RESHAPE_INL_H_
+#define MXNET_OPERATOR_NN_MKLDNN_MKLDNN_RESHAPE_INL_H_
+
+#if MXNET_USE_MKLDNN == 1
+#include <vector>
+#include "mkldnn_base-inl.h"
+#include "../../tensor/matrix_op-inl.h"
+
+namespace mxnet {
+namespace op {
+
+inline bool SupportMKLDNNReshape(const ReshapeParam &param,
+                                 const NDArray &data) {
+  auto data_ndim = data.shape().ndim();
+
+  if (data_ndim > 4 ||
+      data.dtype() != mshadow::kFloat32 ||
+      param.shape.ndim() > 4)
+    return false;
+
+  return true;
+}
+
+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_;
+  bool needInvalidateInput = false;
+
+ public:
+  MKLDNNReshapeFwd(const OpReqType &req,
+                   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,
+               void* workspace = nullptr);
+};
+
+typedef ParamOpSign<ReshapeParam> MKLDNNReshapeSignature;
+MKLDNNReshapeFwd &GetReshapeForward(const ReshapeParam& param,
+                                    const OpReqType &req,
+                                    const NDArray &input,
+                                    const NDArray &output);
+void MKLDNNReshapeForward(const nnvm::NodeAttrs& attrs,
+                          const OpContext &ctx,
+                          const NDArray &input,
+                          const OpReqType &req,
+                          const NDArray &output);
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // MXNET_USE_MKLDNN == 1
+#endif  // MXNET_OPERATOR_NN_MKLDNN_MKLDNN_RESHAPE_INL_H_

src/operator/nn/mkldnn/mkldnn_reshape.cc

@@ -27,37 +27,14 @@
 
 #include <mkldnn.hpp>
 #include "../../tensor/matrix_op-inl.h"
+#include "mkldnn_reshape-inl.h"
 
 namespace mxnet {
 namespace op {
 
-bool SupportMKLDNNReshape(const ReshapeParam &param,
-                          const NDArray &data) {
-  auto data_ndim = data.shape().ndim();
-
-  if (data_ndim > 4 ||
-      data.dtype() != mshadow::kFloat32 ||
-      param.shape.ndim() > 4)
-    return false;
-
-  return true;
-}
-
-typedef ParamOpSign<ReshapeParam> MKLDNNReshapeSignature;
-
-class MKLDNNReshapeForward {
-  std::shared_ptr<mkldnn::memory> data_;
-  std::shared_ptr<mkldnn::memory> out_;
-  std::shared_ptr<mkldnn::memory> temp_;
-  std::vector<mkldnn::primitive> prims_;
-
-  bool needInvalidateInput = false;
-
- public:
-  MKLDNNReshapeForward(const ReshapeParam &param,
-                       const OpReqType &req,
-                       const NDArray &input,
-                       const NDArray &output) {
+MKLDNNReshapeFwd::MKLDNNReshapeFwd(const OpReqType &req,
+                                   const NDArray &input,
+                                   const NDArray &output) {
     auto engine = CpuEngine::Get()->get_engine();
 
     // data_
@@ -98,62 +75,63 @@ class MKLDNNReshapeForward {
     } else {
       LOG(FATAL) << "not supported req type: " << req;
     }
-  }
+}
 
-  int GetWorkspaceSize() {
-    return temp_ ? temp_->get_primitive_desc().get_size() : 0;
-  }
+int MKLDNNReshapeFwd::GetWorkspaceSize() {
+  return temp_ ? temp_->get_primitive_desc().get_size() : 0;
+}
 
-  void SetNewMem(const NDArray &input, const NDArray &output, void* workspace = nullptr) {
-    if (input.IsMKLDNNData()) {
-      this->data_->set_data_handle(input.GetMKLDNNData()->get_data_handle());
-    } else {
-      MSHADOW_TYPE_SWITCH(input.dtype(), DTYPE, {
-        this->data_->set_data_handle(input.data().dptr<DTYPE>());
-      })
-    }
+void MKLDNNReshapeFwd::SetNewMem(const NDArray &input,
+                                 const NDArray &output,
+                                 void* workspace) {
+  if (input.IsMKLDNNData()) {
+    this->data_->set_data_handle(input.GetMKLDNNData()->get_data_handle());
+  } else {
+    MSHADOW_TYPE_SWITCH(input.dtype(), DTYPE, {
+      this->data_->set_data_handle(input.data().dptr<DTYPE>());
+    })
+  }
 
-    if (output.IsMKLDNNData()) {
-      this->out_->set_data_handle(output.GetMKLDNNData()->get_data_handle());
-    } else {
-      MSHADOW_TYPE_SWITCH(output.dtype(), DTYPE, {
-        this->out_->set_data_handle(output.data().dptr<DTYPE>());
-      })
-    }
+  if (output.IsMKLDNNData()) {
+    this->out_->set_data_handle(output.GetMKLDNNData()->get_data_handle());
+  } else {
+    MSHADOW_TYPE_SWITCH(output.dtype(), DTYPE, {
+      this->out_->set_data_handle(output.data().dptr<DTYPE>());
+    })
+  }
 
-    if (workspace) {
-      this->temp_->set_data_handle(workspace);
-    }
+  if (workspace) {
+    this->temp_->set_data_handle(workspace);
   }
+}
 
-  void Execute(const NDArray &input,
-               const NDArray &output,
-               void* workspace = nullptr) {
-    // set memory handles
-    SetNewMem(input, output, workspace);
-    // register primitives
-    auto stream = MKLDNNStream::Get();
-    for (auto &v : this->prims_) {
-      stream->RegisterPrim(v);
-    }
-    stream->Submit();
-    // invalidate mkldnn memory in input
-    if (needInvalidateInput) {
-      const_cast<NDArray &>(input).InvalidateMKLDNNData();
-    }
+void MKLDNNReshapeFwd::Execute(const NDArray &input,
+                               const NDArray &output,
+                               void* workspace) {
+  // set memory handles
+  SetNewMem(input, output, workspace);
+  // register primitives
+  auto stream = MKLDNNStream::Get();
+  for (auto &v : this->prims_) {
+    stream->RegisterPrim(v);
   }
-};
+  stream->Submit();
+  // invalidate mkldnn memory in input
+  if (needInvalidateInput) {
+    const_cast<NDArray &>(input).InvalidateMKLDNNData();
+  }
+}
 
-static MKLDNNReshapeForward &GetReshapeForward(const ReshapeParam& param,
-                                               const OpReqType &req,
-                                               const NDArray &input,
-                                               const NDArray &output) {
+MKLDNNReshapeFwd &GetReshapeForward(const ReshapeParam& param,
+                                    const OpReqType &req,
+                                    const NDArray &input,
+                                    const NDArray &output) {
 #if DMLC_CXX11_THREAD_LOCAL
   static thread_local std::unordered_map<MKLDNNReshapeSignature,
-                                         MKLDNNReshapeForward, OpHash> fwds;
+                                         MKLDNNReshapeFwd, OpHash> fwds;
 #else
   static MX_THREAD_LOCAL std::unordered_map<MKLDNNReshapeSignature,
-                                            MKLDNNReshapeForward, OpHash> fwds;
+                                            MKLDNNReshapeFwd, OpHash> fwds;
 #endif
   MKLDNNReshapeSignature key(param);
   key.AddSign(req);
@@ -162,7 +140,7 @@ static MKLDNNReshapeForward &GetReshapeForward(const ReshapeParam& param,
 
   auto it = fwds.find(key);
   if (it == fwds.end()) {
-    MKLDNNReshapeForward fwd(param, req, input, output);
+    MKLDNNReshapeFwd fwd(req, input, output);
     it = AddToCache(&fwds, key, fwd);
   }
   return it->second;

src/operator/tensor/matrix_op.cc

@@ -28,6 +28,7 @@
 #include "../nn/mkldnn/mkldnn_ops-inl.h"
 #include "../nn/mkldnn/mkldnn_base-inl.h"
 #include "../nn/mkldnn/mkldnn_slice-inl.h"
+#include "../nn/mkldnn/mkldnn_reshape-inl.h"
 
 namespace mxnet {
 namespace op {
@@ -233,12 +234,8 @@ static void FlattenEx(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(inputs.size(), 1U);
   CHECK_EQ(outputs.size(), 1U);
 #if MXNET_USE_MKLDNN == 1
-  if (inputs[0].IsMKLDNNData()) {
-    MKLDNNCopy(attrs, ctx, inputs[0], req[0], outputs[0]);
-    // If the output is a special MKLDNN layout and the number of dimensions
-    // is larger than 2, we should use the default layout.
-    if (outputs[0].IsMKLDNNData() && inputs[0].shape().ndim() > 2)
-      const_cast<NDArray &>(outputs[0]).Reorder2Default();
+  if (SupportMKLDNNArray(inputs[0].dtype(), inputs[0].shape())) {
+    MKLDNNFlattenForward(attrs, ctx, inputs[0], req[0], outputs[0]);
     return;
   } else {
     // This happens if inputs are supposed to be in MKLDNN format
@@ -252,10 +249,10 @@ static void FlattenEx(const nnvm::NodeAttrs& attrs,
 
 #if MXNET_USE_MKLDNN == 1
 static inline bool FlattenStorageType(const nnvm::NodeAttrs& attrs,
-                                   const int dev_mask,
-                                   DispatchMode* dispatch_mode,
-                                   std::vector<int> *in_attrs,
-                                   std::vector<int> *out_attrs) {
+                                      const int dev_mask,
+                                      DispatchMode* dispatch_mode,
+                                      std::vector<int> *in_attrs,
+                                      std::vector<int> *out_attrs) {
   CHECK_EQ(in_attrs->size(), 1);
   CHECK_EQ(out_attrs->size(), 1);
   return MKLDNNStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs,