You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/10/24 09:27:09 UTC
[GitHub] huangzhiyuan closed pull request #12957: Add reshape op supported
by MKL-DNN
huangzhiyuan closed pull request #12957: Add reshape op supported by MKL-DNN
URL: https://github.com/apache/incubator-mxnet/pull/12957
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/src/operator/nn/mkldnn/mkldnn_base.cc b/src/operator/nn/mkldnn/mkldnn_base.cc
index a60d6555c74..c43399e1ac1 100644
--- a/src/operator/nn/mkldnn/mkldnn_base.cc
+++ b/src/operator/nn/mkldnn/mkldnn_base.cc
@@ -106,17 +106,13 @@ void MKLDNNCopy(const mkldnn::memory &mem, const mkldnn::memory* this_mem) {
} else if (!same_shape(this_desc, from_desc)) {
// In this case, the source memory stores data in a customized layout. We
// need to reorganize the data in memory before we can reshape.
- mkldnn::memory::primitive_desc def_pd = GetPrimitiveDesc(from_pd, from_def_format);
- mkldnn::memory *def_mem = TmpMemMgr::Get()->Alloc(def_pd);
- stream->RegisterPrim(mkldnn::reorder(mem, *def_mem));
- // Now we can reshape it
mkldnn::memory::dims dims(this_desc.data.dims,
this_desc.data.dims + this_desc.data.ndims);
auto this_dtype = static_cast<mkldnn::memory::data_type>(this_desc.data.data_type);
auto this_format = static_cast<mkldnn::memory::format>(GetDefaultFormat(this_desc));
mkldnn::memory::desc data_md(dims, this_dtype, this_format);
mkldnn::memory::primitive_desc pd(data_md, from_pd.get_engine());
- mkldnn_mem_ptr tmp_mem(new mkldnn::memory(pd, def_mem->get_data_handle()));
+ mkldnn_mem_ptr tmp_mem(new mkldnn::memory(pd, mem.get_data_handle()));
stream->RegisterMem(tmp_mem);
stream->RegisterPrim(mkldnn::reorder(*tmp_mem, *this_mem));
} else if (from_pd == this_pd) {
diff --git a/src/operator/tensor/matrix_op.cc b/src/operator/tensor/matrix_op.cc
index 77d9bf06e2d..c56c0a39712 100644
--- a/src/operator/tensor/matrix_op.cc
+++ b/src/operator/tensor/matrix_op.cc
@@ -103,6 +103,45 @@ DMLC_REGISTER_PARAMETER(StackParam);
DMLC_REGISTER_PARAMETER(SqueezeParam);
DMLC_REGISTER_PARAMETER(DepthToSpaceParam);
+static void ReshapeComputeExCPU(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);
+#if MXNET_USE_MKLDNN == 1
+ if (inputs[0].IsMKLDNNData()) {
+ MKLDNNCopy(attrs, ctx, inputs[0], req[0], outputs[0]);
+ } else {
+ // 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.
+ FallBackCompute(UnaryOp::IdentityCompute<cpu>, attrs, ctx, inputs, req, outputs);
+ return;
+ }
+#endif
+}
+
+inline static bool ReshapeStorageType(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
+ if (dev_mask == mshadow::cpu::kDevMask
+ && in_attrs->at(0) == kDefaultStorage
+ && out_attrs->at(0) == kDefaultStorage) {
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ }
+#endif
+ return ret;
+}
+
NNVM_REGISTER_OP(Reshape)
.add_alias("reshape")
.describe(R"code(Reshapes the input array.
@@ -171,9 +210,19 @@ If the argument `reverse` is set to 1, then the special values are inferred from
.set_num_outputs(1)
.set_attr_parser(ParamParser<ReshapeParam>)
.set_attr<nnvm::FInferShape>("FInferShape", ReshapeShape)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<FInferStorageType>("FInferStorageType", ReshapeStorageType)
+#endif
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_copy"})
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::IdentityCompute<cpu>)
+.set_attr<FComputeEx>("FComputeEx<cpu>", ReshapeComputeExCPU)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<bool>("TIsMKLDNN", true)
+.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& n) {
+ return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+})
+#else
.set_attr<nnvm::FInplaceOption>("FInplaceOption",
[](const NodeAttrs& attrs) {
return std::vector<std::pair<int, int> >{{0, 0}};
@@ -182,6 +231,7 @@ If the argument `reverse` is set to 1, then the special values are inferred from
[](const NodeAttrs& attrs){
return std::vector<bool>{true};
})
+#endif
.add_argument("data", "NDArray-or-Symbol", "Input data to reshape.")
.add_arguments(ReshapeParam::__FIELDS__());
@@ -914,7 +964,7 @@ NNVM_REGISTER_OP(depth_to_space)
.describe(R"code(Rearranges(permutes) data from depth into blocks of spatial data.
Similar to ONNX DepthToSpace operator:
https://github.com/onnx/onnx/blob/master/docs/Operators.md#DepthToSpace.
-The output is a new tensor where the values from depth dimension are moved in spatial blocks
+The output is a new tensor where the values from depth dimension are moved in spatial blocks
to height and width dimension. The reverse of this operation is ``space_to_depth``.
.. math::
@@ -925,7 +975,7 @@ to height and width dimension. The reverse of this operation is ``space_to_depth
y = reshape(x \prime \prime, [N, C / (block\_size ^ 2), H * block\_size, W * block\_size])
\end{gather*}
-where :math:`x` is an input tensor with default layout as :math:`[N, C, H, W]`: [batch, channels, height, width]
+where :math:`x` is an input tensor with default layout as :math:`[N, C, H, W]`: [batch, channels, height, width]
and :math:`y` is the output tensor of layout :math:`[N, C / (block\_size ^ 2), H * block\_size, W * block\_size]`
Example::
@@ -965,9 +1015,9 @@ Example::
NNVM_REGISTER_OP(space_to_depth)
.describe(R"code(Rearranges(permutes) blocks of spatial data into depth.
Similar to ONNX SpaceToDepth operator:
-https://github.com/onnx/onnx/blob/master/docs/Operators.md#SpaceToDepth
+https://github.com/onnx/onnx/blob/master/docs/Operators.md#SpaceToDepth
-The output is a new tensor where the values from height and width dimension are
+The output is a new tensor where the values from height and width dimension are
moved to the depth dimension. The reverse of this operation is ``depth_to_space``.
.. math::
@@ -978,7 +1028,7 @@ moved to the depth dimension. The reverse of this operation is ``depth_to_space`
y = reshape(x \prime \prime, [N, C * (block\_size ^ 2), H / block\_size, W / block\_size])
\end{gather*}
-where :math:`x` is an input tensor with default layout as :math:`[N, C, H, W]`: [batch, channels, height, width]
+where :math:`x` is an input tensor with default layout as :math:`[N, C, H, W]`: [batch, channels, height, width]
and :math:`y` is the output tensor of layout :math:`[N, C * (block\_size ^ 2), H / block\_size, W / block\_size]`
Example::
@@ -987,8 +1037,8 @@ Example::
[12, 18, 13, 19, 14, 20],
[3, 9, 4, 10, 5, 11],
[15, 21, 16, 22, 17, 23]]]]
-
-
+
+
space_to_depth(x, 2) = [[[[0, 1, 2],
[3, 4, 5]],
[[6, 7, 8],
----------------------------------------------------------------
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:
users@infra.apache.org
With regards,
Apache Git Services