zheng-da commented on a change in pull request #10025: Language model with
Google's billion words dataset
URL: https://github.com/apache/incubator-mxnet/pull/10025#discussion_r173305447
##########
File path: src/operator/nn/fully_connected.cc
##########
@@ -73,22 +76,47 @@ static bool FullyConnectedShape(const nnvm::NodeAttrs&
attrs,
return true;
}
-#if MXNET_USE_MKLDNN == 1
void FullyConnectedComputeExCPU(const nnvm::NodeAttrs& attrs,
const OpContext &ctx,
const std::vector<NDArray> &inputs,
const std::vector<OpReqType> &req,
const std::vector<NDArray> &outputs) {
- if (SupportMKLDNN(inputs[0])) {
- MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
- MKLDNNFCForward(attrs, ctx, inputs, req, outputs);
- MKLDNN_OPCHECK_RUN(FullyConnectedCompute<cpu>, attrs, ctx, inputs, req,
- outputs);
+ const FullyConnectedParam& param =
nnvm::get<FullyConnectedParam>(attrs.parsed);
+ const bool valid_data = inputs[0].storage_type() == kDefaultStorage;
+ const bool valid_weight = inputs[1].storage_type() == kDefaultStorage ||
+ inputs[1].storage_type() == kRowSparseStorage;
+ const bool valid_out = outputs[0].storage_type() == kDefaultStorage;
+ bool valid_bias = true;
+ if (!param.no_bias) {
+ valid_bias = inputs[2].storage_type() == kDefaultStorage ||
+ inputs[2].storage_type() == kRowSparseStorage;
+ }
+#if MXNET_USE_MKLDNN == 1
+ if (common::ContainsOnlyStorage(inputs, kDefaultStorage) &&
+ common::ContainsOnlyStorage(outputs, kDefaultStorage)) {
+ if (SupportMKLDNN(inputs[0])) {
+ MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ MKLDNNFCForward(attrs, ctx, inputs, req, outputs);
+ MKLDNN_OPCHECK_RUN(FullyConnectedCompute<cpu>, attrs, ctx, inputs, req,
+ outputs);
+ } else {
+ FallBackCompute(FullyConnectedCompute<cpu>, attrs, ctx, inputs, req,
outputs);
+ }
return;
}
- FallBackCompute(FullyConnectedCompute<cpu>, attrs, ctx, inputs, req,
outputs);
+#endif
+ if (valid_data && valid_weight && valid_bias && valid_out) {
+ std::vector<TBlob> in_blobs(inputs.size());
+ for (size_t i = 0; i < in_blobs.size(); i++) in_blobs[i] =
inputs[i].data();
Review comment:
Actually, if an NDArray uses MKLDNN format, data() will convert its layout
inside the array. This caused a race condition, I just fixed. You shouldn't
call data() on an NDArray with MKLDNN format. Please check FallBackCompute to
see how it is handled correctly.
I also don't understand what happens if the input array is a row sparse. If
the row sparse array doesn't have zero-entry rows, it works fine. What is the
row sparse array have zero-entry rows, isn't the memory returned from data() is
smaller than we expect?
----------------------------------------------------------------
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
