fix inferstorage for sparse in FullyConnected

apache · Jul 10, 2018 · 198cae3 · 198cae3
1 parent 7a6e8de
commit 198cae3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 9 deletions.
diff --git a/src/operator/nn/fully_connected-inl.h b/src/operator/nn/fully_connected-inl.h
@@ -35,6 +35,7 @@
 #include "../operator_common.h"
 #include "../elemwise_op_common.h"
 #include "../linalg.h"
+#include "../../common/utils.h"
 
 namespace mxnet {
 namespace op {

diff --git a/src/operator/nn/fully_connected.cc b/src/operator/nn/fully_connected.cc
@@ -210,17 +210,21 @@ inline static bool BackwardFCStorageType(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(in_attrs->size(), 3U);
   CHECK_EQ(out_attrs->size(), out_expected);
 
-  DispatchMode wanted_mode;
-#if 0
+  bool dispatched = false;
   // TODO(zhengda) let's disable MKLDNN for FullyConnected for now.
   // It seems there is a bug.
-  if (dev_mask == mshadow::cpu::kDevMask)
-    *dispatch_mode = DispatchMode::kFComputeEx;
-  else
-#endif
-    wanted_mode = DispatchMode::kFCompute;
-  return storage_type_assign(out_attrs, mxnet::kDefaultStorage,
-                             dispatch_mode, wanted_mode);
+  if (!dispatched && common::ContainsOnlyStorage(*in_attrs, mxnet::kDefaultStorage)) {
+    storage_type_assign(out_attrs, mxnet::kDefaultStorage,
+                        dispatch_mode, DispatchMode::kFCompute);
+  }
+  if (!dispatched && common::ContainsStorageType(*in_attrs, mxnet::kRowSparseStorage)) {
+    dispatched = dispatch_fallback(out_attrs, dispatch_mode);
+  }
+  if (!dispatched) {
+    dispatched = storage_type_assign(out_attrs, mxnet::kDefaultStorage,
+                                     dispatch_mode, DispatchMode::kFCompute);
+  }
+  return dispatched;
 }
 
 DMLC_REGISTER_PARAMETER(FullyConnectedParam);