zheng-da commented on a change in pull request #8302: Refactor operators &
MKLDNN
URL: https://github.com/apache/incubator-mxnet/pull/8302#discussion_r162501336
##########
File path: src/operator/tensor/elemwise_unary_op_basic.cc
##########
@@ -108,12 +109,66 @@
MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU(_backward_sigmoid,
unary_bwd<mshadow_op::sigmoid_grad>);
// copy
+static void CopyEx(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ CHECK_EQ(inputs.size(), 1U);
+ CHECK_EQ(outputs.size(), 1U);
+ const auto in_stype = inputs[0].storage_type();
+ const auto out_stype = outputs[0].storage_type();
+#if MXNET_USE_MKLDNN == 1
+ if (inputs[0].IsMKLDNNData()) {
+ MKLDNNCopy(attrs, ctx, inputs[0], req[0], outputs[0]);
+ return;
+ } else if (in_stype == kDefaultStorage && out_stype == kDefaultStorage) {
+ // This happens if inputs are supposed to be in MKLDNN format
+ // but MKLDNN doesn't support the data type or the shape. We're
+ // forced to convert it to the default format.
+ std::vector<TBlob> in_blobs(1);
+ std::vector<TBlob> out_blobs(1);
+ in_blobs[0] = inputs[0].data();
+ out_blobs[0] = outputs[0].data();
+ UnaryOp::IdentityCompute<cpu>(attrs, ctx, in_blobs, req, out_blobs);
+ return;
+ }
+#endif
+ UnaryOp::IdentityComputeEx<cpu>(attrs, ctx, inputs, req, outputs);
+}
+
+static inline bool CopyStorageType(const nnvm::NodeAttrs& attrs,
Review comment:
I'm not sure if ElemwiseStorageType should be used for other mkldnn ops. is
it a right thing to do?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
With regards,
Apache Git Services
