From 0fe86582a7f42cb1086b9e067d133402ad8af911 Mon Sep 17 00:00:00 2001
From: Zhennan Qin <zhennan.qin@intel.com>
Date: Thu, 24 Oct 2019 16:10:36 +0800
Subject: [PATCH] Fix reorder2default

Change-Id: Iae5043fd1c105af38c877be47559383b35dd9ae1
---
 src/ndarray/ndarray.cc | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
index 2fe9c1f05015..e11665d72ab0 100644
--- a/src/ndarray/ndarray.cc
+++ b/src/ndarray/ndarray.cc
@@ -1615,12 +1615,13 @@ void NDArray::Save(dmlc::Stream *strm) const {
     nd_cpu.WaitToRead();
     save_data = nd_cpu.data();
   } else {
+#if MXNET_USE_MKLDNN == 1
+    // For mkldnn, a copy of *this can ensure no write access pending on *this.
+    nd_cpu = this->Copy(Context::CPU());
+    nd_cpu.WaitToRead();
+#else
     this->WaitToRead();
     nd_cpu = *this;
-#if MXNET_USE_MKLDNN == 1
-    if (nd_cpu.IsMKLDNNData()) {
-      nd_cpu = nd_cpu.Reorder2Default();
-    }
 #endif
     save_data = nd_cpu.data();
   }
@@ -1869,10 +1870,17 @@ void NDArray::SyncCopyFromCPU(const void *data, size_t size) const {
   TBlob src((void*)data, dshape, cpu::kDevMask, this->dtype_, 0); // NOLINT(*)
 
   if (this->ctx().dev_mask() == cpu::kDevMask) {
+    Engine::Get()->PushAsync(
+        [&](RunContext rctx, Engine::CallbackOnComplete on_complete) {
+          NDArray src = *this;
+#if MXNET_USE_MKLDNN == 1
+          src = this->Reorder2Default();
+#endif
+          ndarray::Copy<cpu, cpu>(src.data(), &dst, Context::CPU(), Context::CPU(), rctx);
+          on_complete();
+        },
+        this->ctx(), {this->var()}, {}, FnProperty::kNormal, 0, "SyncCopyCPU2CPU");
     this->WaitToWrite();
-    RunContext rctx{this->ctx(), nullptr, nullptr, false};
-    TBlob dst = this->data();
-    ndarray::Copy<cpu, cpu>(src, &dst, Context::CPU(), Context::CPU(), rctx);
   } else {
 #if MXNET_USE_CUDA
     Engine::Get()->PushAsync(