OP primitive cache: use memory as signature for MKLDNN storage type

zheng-da · Jan 11, 2018 · 7fa4d0a · 7fa4d0a
1 parent 982796b
commit 7fa4d0a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/operator/nn/mkldnn/mkldnn_act.cc b/src/operator/nn/mkldnn/mkldnn_act.cc
@@ -141,7 +141,7 @@ static MKLDNNActForward &GetActForward(const ActivationParam& param,
   MKLDNNActSignature key(param);
   key.AddSign(ctx.is_train);
   key.AddSign(param.act_type);
-  key.AddSign(in_mem);
+  key.AddSign(in_data);
 
   auto it = fwds.find(key);
   if (it == fwds.end()) {

diff --git a/src/operator/nn/mkldnn/mkldnn_base-inl.h b/src/operator/nn/mkldnn/mkldnn_base-inl.h
@@ -318,9 +318,13 @@ class MKLDNNOpSignature {
   }
 
   void AddSign(const NDArray &arr) {
-    hash = hash * 2 + arr.dtype();
-    eles.push_back(arr.dtype());
-    AddSign(arr.shape());
+    if (arr.IsMKLDNN()) {
+      AddSign(*(arr.GetMKLDNNData()));
+    } else {
+      hash = hash * 2 + arr.dtype();
+      eles.push_back(arr.dtype());
+      AddSign(arr.shape());
+    }
   }
 
   void AddSign(const TShape &shape) {