Skip to content

Commit

Permalink
[MKLDNN]Fix reorder2default (apache#16602)
Browse files Browse the repository at this point in the history
* Fix reorder2default

Change-Id: I74c87af9535f6264e6d1ea7eaed089a6480a3358

* fix

Change-Id: I6d07b43b520a47e7c78bd4b4b6390f5fb95e6957

* Fix

Change-Id: Id72f25c34291be4711f55569c6d61467edd6113d

* Fix CI

Change-Id: I8c33a82555d5ace2d0b682c1e3eefa13f3a44768

* Run CI

Change-Id: Ie8a6dab80ef91c0337cafbae4e3db277e0c7ebf7
  • Loading branch information
ZhennanQin authored and pengzhao-intel committed Oct 24, 2019
1 parent bde443e commit dd4eaf5
Showing 1 changed file with 15 additions and 10 deletions.
25 changes: 15 additions & 10 deletions src/ndarray/ndarray.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1635,11 +1635,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();
}
Expand Down Expand Up @@ -2024,15 +2026,18 @@ void NDArray::SyncCopyToCPU(void *data, size_t size) const {
TBlob dst(data, dshape, cpu::kDevMask, this->dtype_, 0); // NOLINT(*)

if (this->ctx().dev_mask() == cpu::kDevMask) {
this->WaitToRead();
RunContext rctx{this->ctx(), nullptr, nullptr, false};
NDArray src = *this;
Engine::Get()->PushAsync(
[&](RunContext rctx, Engine::CallbackOnComplete on_complete) {
RunContext ctx{this->ctx(), nullptr, nullptr, false};
NDArray src = *this;
#if MXNET_USE_MKLDNN == 1
if (src.IsMKLDNNData())
src = this->Reorder2Default();
src = this->Reorder2Default();
#endif
ndarray::Copy<cpu, cpu>(src.data(), &dst,
Context::CPU(), Context::CPU(), rctx);
ndarray::Copy<cpu, cpu>(src.data(), &dst, Context::CPU(), Context::CPU(), ctx);
on_complete();
},
this->ctx(), {this->var()}, {}, FnProperty::kNormal, 0, "SyncCopyCPU2CPU");
this->WaitToWrite();
} else {
#if MXNET_USE_CUDA
Engine::Get()->PushAsync(
Expand Down

0 comments on commit dd4eaf5

Please sign in to comment.