Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

[mkldnn-1.0]enable mkldnn elemwise_sum #16521

Merged
merged 4 commits into from
Oct 19, 2019
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/operator/tensor/elemwise_sum.cc
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ bool ElementWiseSumForwardInferStorageType(const nnvm::NodeAttrs& attrs,
CHECK_EQ(out_attrs->size(), 1U);
bool ret = ElemwiseStorageAttr<false, true, false>(attrs, dev_mask, dispatch_mode,
in_attrs, out_attrs);
#if MXNET_USE_MKLDNN == 1
#if MXNET_USE_MKLDNN == 100
// We should always use FComputeEx.
if (dev_mask == mshadow::cpu::kDevMask
&& common::ContainsOnlyStorage(*in_attrs, kDefaultStorage)
Expand All @@ -94,7 +94,7 @@ bool ElementWiseSumForwardInferStorageType(const nnvm::NodeAttrs& attrs,
return ret;
}

#if MXNET_USE_MKLDNN == 1
#if MXNET_USE_MKLDNN == 100
static inline bool IsMKLDNNData(const std::vector<NDArray> &arrs) {
for (auto &arr : arrs) {
if (!arr.IsMKLDNNData())
Expand Down Expand Up @@ -123,7 +123,7 @@ void ElementWiseSumComputeExCPU(const nnvm::NodeAttrs& attrs,
ResourceRequest(ResourceRequest::kTempSpace));
NDArray out_nd = outputs[0];
mxnet::ndarray::ElementwiseSum<cpu>(s, rsc, inputs, &out_nd);
#if MXNET_USE_MKLDNN == 1
#if MXNET_USE_MKLDNN == 100
} else if (IsMKLDNNData(inputs)) {
MKLDNNSumForward(attrs, ctx, inputs, req[0], outputs[0]);
} else if (common::ContainsOnlyStorage(inputs, kDefaultStorage)) {
Expand Down Expand Up @@ -178,7 +178,7 @@ The storage type of ``add_n`` output depends on storage types of inputs
[](const NodeAttrs& attrs) {
return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
})
#if MXNET_USE_MKLDNN == 1
#if MXNET_USE_MKLDNN == 100
.set_attr<bool>("TIsMKLDNN", true)
#endif
.set_attr<mxnet::FInferShape>("FInferShape", ElementWiseSumShape)
Expand Down