Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Fix GetMKLDNNData for delay alloc #16618

Merged
merged 6 commits into from
Oct 28, 2019
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/ndarray/ndarray.cc
Original file line number Diff line number Diff line change
Expand Up @@ -647,6 +647,7 @@ const mkldnn::memory *NDArray::GetMKLDNNData() const {
// If this is a view, we can't create a MKLDNN memory for the chunk
// because we don't have the complete data type and shape information for
// the chunk.
CheckAndAlloc();
void *off_addr = static_cast<char *>(ptr_->shandle.dptr) + byte_offset_;
// Create the primitive desc for the new mkldnn memory.
mkldnn::memory::dims dims(shape().ndim());
Expand All @@ -665,6 +666,7 @@ const mkldnn::memory *NDArray::GetMKLDNNData() const {
} else {
// If this isn't a view, we can create a MKLDNN memory and store it in the
// chunk.
CheckAndAlloc();
ptr_->SetMKLMem(shape_, dtype_);
MKLDNNStream::Get()->RegisterMem(ptr_->mkl_mem_->GetMem());
return ptr_->mkl_mem_->GetRaw();
Expand Down