Act change.

luobao-intel · Jul 31, 2018 · 8455702 · 8455702
1 parent 7ffb252
commit 8455702
Showing 1 changed file with 10 additions and 37 deletions.
diff --git a/src/operator/nn/activation.cc b/src/operator/nn/activation.cc
@@ -101,66 +101,35 @@ void ActivationGradComputeExCPU(const nnvm::NodeAttrs& attrs,
 }
 #endif
 
+#if MXNET_USE_MKLDNN == 1
 inline static bool ActivationStorageType(const nnvm::NodeAttrs& attrs,
                                          const int dev_mask,
                                          DispatchMode* dispatch_mode,
                                          std::vector<int> *in_attrs,
                                          std::vector<int> *out_attrs) {
   CHECK_EQ(in_attrs->size(), 1);
   CHECK_EQ(out_attrs->size(), 1);
-  bool ret = ElemwiseStorageType<1, 1, false, false, false>(attrs, dev_mask,
-                                                            dispatch_mode,
-                                                            in_attrs, out_attrs);
-#if MXNET_USE_MKLDNN == 1
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
-  if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNAct(param)) {
-    *dispatch_mode = DispatchMode::kFComputeEx;
-  }
-#endif
-  return ret;
+  return MKLDNNStorageType(attrs, dev_mask, SupportMKLDNNAct(param),
+                           dispatch_mode, in_attrs, out_attrs);
 }
 
 inline static bool BackwardActStorageType(const nnvm::NodeAttrs& attrs,
                                           const int dev_mask,
                                           DispatchMode* dispatch_mode,
                                           std::vector<int> *in_attrs,
                                           std::vector<int> *out_attrs) {
-  bool ret = false;
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
-#if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1)
   if (param.act_type != activation::kReLU) {
     CHECK_EQ(in_attrs->size(), 3U);
-    ret = ElemwiseStorageType<3, 1, false, false, false>(attrs, dev_mask,
-                                                         dispatch_mode,
-                                                         in_attrs, out_attrs);
   } else {
     // for ReLU activation, the backward pass only needs ograd and output
     CHECK_EQ(in_attrs->size(), 2U);
-    ret = ElemwiseStorageType<2, 1, false, false, false>(attrs, dev_mask,
-                                                         dispatch_mode,
-                                                         in_attrs, out_attrs);
-  }
-#else
-  if (param.act_type == activation::kSoftSign) {
-    CHECK_EQ(in_attrs->size(), 3U);
-    ret = ElemwiseStorageType<3, 1, false, false, false>(attrs, dev_mask,
-                                                         dispatch_mode,
-                                                         in_attrs, out_attrs);
-  } else {
-    CHECK_EQ(in_attrs->size(), 2U);
-    ret = ElemwiseStorageType<2, 1, false, false, false>(attrs, dev_mask,
-                                                         dispatch_mode,
-                                                         in_attrs, out_attrs);
   }
-#endif
-  CHECK_EQ(out_attrs->size(), 1U);
-#if MXNET_USE_MKLDNN == 1
-  if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNAct(param)) {
-    *dispatch_mode = DispatchMode::kFComputeEx;
-  }
-#endif
-  return ret;
+  return MKLDNNStorageType(attrs, dev_mask, SupportMKLDNNAct(param),
+                           dispatch_mode, in_attrs, out_attrs);
 }
+#endif
 
 MXNET_OPERATOR_REGISTER_UNARY(Activation)
 .describe(R"code(Applies an activation function element-wise to the input.
@@ -175,7 +144,9 @@ The following activation functions are supported:
 
 )code" ADD_FILELINE)
 .set_attr_parser(ParamParser<ActivationParam>)
+#if MXNET_USE_MKLDNN == 1
 .set_attr<FInferStorageType>("FInferStorageType", ActivationStorageType)
+#endif
 .set_attr<nnvm::FListOutputNames>("FListOutputNames",
     [](const NodeAttrs& attrs) {
     return std::vector<std::string>{"output"};
@@ -196,7 +167,9 @@ NNVM_REGISTER_OP(_backward_Activation)
   })
 .set_num_outputs(1)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
+#if MXNET_USE_MKLDNN == 1
 .set_attr<FInferStorageType>("FInferStorageType", BackwardActStorageType)
+#endif
 .set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<3, 1>)
 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<3, 1>)
 .set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs){