improve hashing.

rahul003 · Mar 29, 2018 · 58854be · 58854be
1 parent 8f3194f
commit 58854be
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/include/mxnet/ndarray.h b/include/mxnet/ndarray.h
@@ -588,6 +588,9 @@ class NDArray {
   bool IsDefaultData() const {
     return ptr_->IsDefault();
   }
+
+  const mkldnn::memory::desc &GetMKLDNNDesc() const;
+
   /*
    * All functions below return a raw pointer to mkldnn memory. Actually there
    * is a shared pointer that hold the memory either in NDArray or in MKLDNN

diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
@@ -168,6 +168,11 @@ nnvm::Symbol NDArray::get_autograd_symbol() const {
 
 #if MXNET_USE_MKLDNN == 1
 
+const mkldnn::memory::desc &NDArray::GetMKLDNNDesc() const {
+  CHECK(ptr_->mkl_mem_);
+  return ptr_->mkl_mem_->GetDesc();
+}
+
 NDArray NDArray::MKLDNNDataReshape(const TShape &shape) const {
   CHECK(!is_none()) << "NDArray is not initialized";
   CHECK_GE(shape_.Size(), shape.Size())

diff --git a/src/operator/nn/mkldnn/mkldnn_base-inl.h b/src/operator/nn/mkldnn/mkldnn_base-inl.h
@@ -392,6 +392,10 @@ class MKLDNNMemory {
     return size;
   }
 
+  const mkldnn::memory::desc &GetDesc() const {
+    return desc;
+  }
+
   mkldnn::memory::primitive_desc GetPrimitiveDesc() const {
     return mem->get_primitive_desc();
   }

diff --git a/src/operator/operator_common.h b/src/operator/operator_common.h
@@ -519,8 +519,7 @@ class OpSignature {
    */
 
 #if MXNET_USE_MKLDNN == 1
-  void AddSign(const mkldnn::memory &mem) {
-    auto desc = mem.get_primitive_desc().desc();
+  void AddSign(const mkldnn::memory::desc &desc) {
     hash = hash * 2 + desc.data.format;
     eles.push_back(desc.data.format);
     hash = hash * 2 + desc.data.data_type;
@@ -541,7 +540,7 @@ class OpSignature {
   void AddSign(const NDArray &arr) {
 #if MXNET_USE_MKLDNN == 1
     if (arr.IsMKLDNNData()) {
-      AddSign(*(arr.GetMKLDNNData()));
+      AddSign(arr.GetMKLDNNDesc());
     } else {
 #endif
       hash = hash * 2 + arr.dtype();