You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by ha...@apache.org on 2018/06/10 17:45:50 UTC
[incubator-mxnet] branch master updated: [MXNET-394] concat of CSR
NDArrays on first dimension (#11024)
This is an automated email from the ASF dual-hosted git repository.
haibin pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new 04eb7f1 [MXNET-394] concat of CSR NDArrays on first dimension (#11024)
04eb7f1 is described below
commit 04eb7f1600371f1d0db91d33e82d45ceda7b8d02
Author: Ziyue Huang <zy...@gmail.com>
AuthorDate: Mon Jun 11 01:45:41 2018 +0800
[MXNET-394] concat of CSR NDArrays on first dimension (#11024)
* concat of csr ndarrays on first dimension
* update
* update
* trigger
* CR comments
* update docs and sparse op alias
* CR comments
* update
* update
---
docs/api/python/ndarray/sparse.md | 9 +++
docs/api/python/symbol/sparse.md | 9 +++
src/operator/nn/concat-inl.h | 85 ++++++++++++++++++++++++++++
src/operator/nn/concat.cc | 57 +++++++++++++------
src/operator/nn/concat.cu | 20 ++++++-
tests/python/unittest/test_sparse_ndarray.py | 17 ++++++
6 files changed, 178 insertions(+), 19 deletions(-)
diff --git a/docs/api/python/ndarray/sparse.md b/docs/api/python/ndarray/sparse.md
index 581a74f..2f1b89c 100644
--- a/docs/api/python/ndarray/sparse.md
+++ b/docs/api/python/ndarray/sparse.md
@@ -169,6 +169,15 @@ We summarize the interface for each class in the following sections.
CSRNDArray.square
```
+### Joining arrays
+
+```eval_rst
+.. autosummary::
+ :nosignatures:
+
+ concat
+```
+
### Indexing
```eval_rst
diff --git a/docs/api/python/symbol/sparse.md b/docs/api/python/symbol/sparse.md
index 86191e3..d26ba07 100644
--- a/docs/api/python/symbol/sparse.md
+++ b/docs/api/python/symbol/sparse.md
@@ -76,6 +76,15 @@ In the rest of this document, we list sparse related routines provided by the
cast_storage
```
+### Joining arrays
+
+```eval_rst
+.. autosummary::
+ :nosignatures:
+
+ concat
+```
+
### Indexing routines
```eval_rst
diff --git a/src/operator/nn/concat-inl.h b/src/operator/nn/concat-inl.h
index a7f1fa8..7a58ae6 100644
--- a/src/operator/nn/concat-inl.h
+++ b/src/operator/nn/concat-inl.h
@@ -154,6 +154,91 @@ void ConcatGradCompute(const nnvm::NodeAttrs& attrs, const OpContext& ctx,
});
}
+/*!
+ * \brief concat CSRNDArray on the first dimension.
+ */
+struct concat_csr_first_dim {
+ /*!
+ * \param i the i-th row of the input ndarray
+ * \param out_idx output csr ndarray column indices
+ * \param out_data output csr ndarray data
+ * \param out_indptr output csr ndarray row index pointer
+ * \param in_idx input csr ndarray column indices
+ * \param in_data input csr ndarray data
+ * \param in_indptr input csr ndarray row index pointer
+ * \param indptr_offset offset for ouput ndarray row index pointer
+ * \param idx_offset offset for ouput ndarray column indices
+ */
+ template<typename DType, typename RType, typename IType>
+ MSHADOW_XINLINE static void Map(int i, const OpReqType req,
+ DType* out_data, const DType* in_data,
+ RType* out_indptr, const RType* in_indptr,
+ IType* out_idx, const IType* in_idx,
+ const nnvm::dim_t indptr_offset,
+ const nnvm::dim_t idx_offset) {
+ if (i == 0) out_indptr[0] = 0;
+ out_indptr[i+1+indptr_offset] = in_indptr[i+1] + idx_offset;
+ for (nnvm::dim_t j = in_indptr[i]; j < in_indptr[i+1]; ++j) {
+ KERNEL_ASSIGN(out_idx[j+idx_offset], req, in_idx[j]);
+ KERNEL_ASSIGN(out_data[j+idx_offset], req, in_data[j]);
+ }
+ }
+};
+
+template<typename xpu>
+void ConcatCSRImpl(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ using namespace csr;
+ const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
+ int num_args = param.num_args;
+ int concat_dim = param.dim;
+ CHECK_EQ(inputs.size(), num_args);
+ CHECK_EQ(outputs.size(), 1);
+ int axis = CheckAxis(concat_dim, inputs[0].shape().ndim());
+ CHECK_EQ(axis, 0) << "concat of csr ndarrays on axis 1 is not supported.";
+ if (req[0] == kNullOp) return;
+ Stream<xpu>* s = ctx.get_stream<xpu>();
+ nnvm::dim_t nnz = 0;
+ for (int i=0; i < num_args; i++) {
+ nnz += inputs[i].aux_shape(kIdx)[0];
+ }
+ const NDArray& out = outputs[0];
+ if (nnz == 0) {
+ FillZerosCsrImpl(s, out);
+ return;
+ }
+ const nnvm::dim_t num_rows = out.shape()[0];
+ out.CheckAndAllocAuxData(kIndPtr, Shape1(num_rows+1));
+
+ MSHADOW_IDX_TYPE_SWITCH(inputs[0].aux_type(kIndPtr), RType, {
+ MSHADOW_IDX_TYPE_SWITCH(inputs[0].aux_type(kIdx), IType, {
+ MSHADOW_TYPE_SWITCH(inputs[0].dtype(), DType, {
+ RType* out_indptr = out.aux_data(kIndPtr).dptr<RType>();
+ out.CheckAndAllocAuxData(kIdx, Shape1(nnz));
+ out.CheckAndAllocData(Shape1(nnz));
+ IType* out_idx = out.aux_data(kIdx).dptr<IType>();
+ DType* out_data = out.data().dptr<DType>();
+ nnvm::dim_t indptr_offset = 0;
+ nnvm::dim_t idx_offset = 0;
+ for (const auto& in : inputs) {
+ const RType* in_indptr = in.aux_data(kIndPtr).dptr<RType>();
+ const IType* in_idx = in.aux_data(kIdx).dptr<IType>();
+ const DType* in_data = in.data().dptr<DType>();
+ Kernel<concat_csr_first_dim, xpu>::Launch(s, in.shape()[0], req[0], out_data,
+ in_data, out_indptr, in_indptr, out_idx, in_idx, indptr_offset, idx_offset);
+ indptr_offset += in.shape()[0];
+ idx_offset += in.aux_shape(kIdx)[0];
+ }
+ });
+ });
+ });
+}
+
} // namespace op
} // namespace mxnet
diff --git a/src/operator/nn/concat.cc b/src/operator/nn/concat.cc
index a7fcb1c..0433245 100644
--- a/src/operator/nn/concat.cc
+++ b/src/operator/nn/concat.cc
@@ -112,18 +112,30 @@ inline static bool ConcatForwardInferStorageType(const nnvm::NodeAttrs& attrs,
std::vector<int> *out_attrs) {
CHECK(!in_attrs->empty());
CHECK_EQ(out_attrs->size(), 1U);
- DispatchMode wanted_mode;
-#if MXNET_USE_MKLDNN == 1
+ auto& out_stype = out_attrs->at(0);
+ bool dispatched = false;
const ConcatParam& param = nnvm::get<ConcatParam>(attrs.parsed);
- if (dev_mask == mshadow::cpu::kDevMask
+ if (!dispatched && common::ContainsOnlyStorage(*in_attrs, kCSRStorage)
+ && param.dim == 0) {
+ dispatched = storage_type_assign(&out_stype, kCSRStorage,
+ dispatch_mode, DispatchMode::kFComputeEx);
+ }
+#if MXNET_USE_MKLDNN == 1
+ if (!dispatched && dev_mask == mshadow::cpu::kDevMask
&& common::ContainsOnlyStorage(*in_attrs, kDefaultStorage)
- && param.dim > 0)
- wanted_mode = DispatchMode::kFComputeEx;
- else
+ && param.dim > 0) {
+ dispatched = storage_type_assign(&out_stype, kDefaultStorage,
+ dispatch_mode, DispatchMode::kFComputeEx);
+ }
#endif
- wanted_mode = DispatchMode::kFCompute;
- return storage_type_assign(out_attrs, mxnet::kDefaultStorage,
- dispatch_mode, wanted_mode);
+ if (!dispatched && common::ContainsOnlyStorage(*in_attrs, kDefaultStorage)) {
+ dispatched = storage_type_assign(&out_stype, kDefaultStorage,
+ dispatch_mode, DispatchMode::kFCompute);
+ }
+ if (!dispatched) {
+ dispatched = dispatch_fallback(out_attrs, dispatch_mode);
+ }
+ return dispatched;
}
inline static bool BackwardConcatStorageType(const nnvm::NodeAttrs& attrs,
@@ -146,7 +158,6 @@ inline static bool BackwardConcatStorageType(const nnvm::NodeAttrs& attrs,
dispatch_mode, wanted_mode);
}
-#if MXNET_USE_MKLDNN == 1
static void ConcatComputeExCPU(const nnvm::NodeAttrs& attrs,
const OpContext& op_ctx,
const std::vector<NDArray>& inputs,
@@ -156,17 +167,24 @@ static void ConcatComputeExCPU(const nnvm::NodeAttrs& attrs,
CHECK_EQ(outputs.size(), 1U);
CHECK_EQ(req.size(), 1U);
if (req[0] == kNullOp) return;
- // MKLDNN support 2D and 4D concat
- if ((inputs[0].shape().ndim() == 2 || inputs[0].shape().ndim() == 4)
+ if (common::ContainsOnlyStorage(inputs, kCSRStorage) &&
+ outputs[0].storage_type() == kCSRStorage) {
+ ConcatCSRImpl<cpu>(attrs, op_ctx, inputs, req, outputs);
+#if MXNET_USE_MKLDNN == 1
+ } else if ((inputs[0].shape().ndim() == 2 || inputs[0].shape().ndim() == 4)
&& inputs[0].dtype() == mshadow::kFloat32) {
MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
MKLDNNConcatForward(attrs, op_ctx, inputs, req, outputs);
MKLDNN_OPCHECK_RUN(ConcatCompute<cpu>, attrs, op_ctx, inputs, req, outputs);
- return;
+ } else if (common::ContainsOnlyStorage(inputs, kDefaultStorage)) {
+ FallBackCompute(ConcatCompute<cpu>, attrs, op_ctx, inputs, req, outputs);
+#endif
+ } else {
+ LogUnimplementedOp(attrs, op_ctx, inputs, req, outputs);
}
- FallBackCompute(ConcatCompute<cpu>, attrs, op_ctx, inputs, req, outputs);
}
+#if MXNET_USE_MKLDNN == 1
static void ConcatGradComputeExCPU(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<NDArray>& inputs,
@@ -201,6 +219,8 @@ struct ConcatGrad {
DMLC_REGISTER_PARAMETER(ConcatParam);
NNVM_REGISTER_OP(Concat)
+MXNET_ADD_SPARSE_OP_ALIAS(concat)
+.add_alias("concat")
.describe(R"code(Joins input arrays along a given axis.
.. note:: `Concat` is deprecated. Use `concat` instead.
@@ -210,6 +230,11 @@ which they will be concatenated.
The dimension of the output array along the concatenated axis will be equal
to the sum of the corresponding dimensions of the input arrays.
+The storage type of ``concat`` output depends on storage types of inputs
+
+- concat(csr, csr, ..., csr, dim=0) = csr
+- otherwise, ``concat`` generates output with default storage
+
Example::
x = [[1,1],[2,2]]
@@ -261,16 +286,12 @@ Example::
.set_attr<nnvm::FInferType>("FInferType", ConcatType)
.set_attr<FInferStorageType>("FInferStorageType", ConcatForwardInferStorageType)
.set_attr<FCompute>("FCompute<cpu>", ConcatCompute<cpu>)
-#if MXNET_USE_MKLDNN == 1
.set_attr<FComputeEx>("FComputeEx<cpu>", ConcatComputeExCPU)
-#endif
.set_attr<nnvm::FGradient>("FGradient", ConcatGrad{"_backward_Concat"})
.set_attr<std::string>("key_var_num_args", "num_args")
.add_argument("data", "NDArray-or-Symbol[]", "List of arrays to concatenate")
.add_arguments(ConcatParam::__FIELDS__());
-NNVM_REGISTER_OP(Concat).add_alias("concat");
-
NNVM_REGISTER_OP(_backward_Concat)
.set_num_outputs([](const NodeAttrs& attrs) {
const ConcatParam& params = nnvm::get<ConcatParam>(attrs.parsed);
diff --git a/src/operator/nn/concat.cu b/src/operator/nn/concat.cu
index f6bf5ec..4f6b8fc 100644
--- a/src/operator/nn/concat.cu
+++ b/src/operator/nn/concat.cu
@@ -29,8 +29,26 @@
namespace mxnet {
namespace op {
+static void ConcatComputeExGPU(const nnvm::NodeAttrs& attrs,
+ const OpContext& op_ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ CHECK(!inputs.empty());
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ if (req[0] == kNullOp) return;
+ if (common::ContainsOnlyStorage(inputs, kCSRStorage) &&
+ outputs[0].storage_type() == kCSRStorage) {
+ ConcatCSRImpl<gpu>(attrs, op_ctx, inputs, req, outputs);
+ } else {
+ LogUnimplementedOp(attrs, op_ctx, inputs, req, outputs);
+ }
+}
+
NNVM_REGISTER_OP(Concat)
-.set_attr<FCompute>("FCompute<gpu>", ConcatCompute<gpu>);
+.set_attr<FCompute>("FCompute<gpu>", ConcatCompute<gpu>)
+.set_attr<FComputeEx>("FComputeEx<gpu>", ConcatComputeExGPU);
NNVM_REGISTER_OP(_backward_Concat)
.set_attr<FCompute>("FCompute<gpu>", ConcatGradCompute<gpu>);
diff --git a/tests/python/unittest/test_sparse_ndarray.py b/tests/python/unittest/test_sparse_ndarray.py
index 1ed5080..c90fb13 100644
--- a/tests/python/unittest/test_sparse_ndarray.py
+++ b/tests/python/unittest/test_sparse_ndarray.py
@@ -156,6 +156,23 @@ def test_sparse_nd_slice():
@with_seed()
+def test_sparse_nd_concat():
+ def check_concat(arrays):
+ ret = np.concatenate([arr.asnumpy() for arr in arrays], axis=0)
+ same(mx.nd.concat(*arrays, dim=0).asnumpy(), ret)
+ nds = []
+ zero_nds = []
+ ncols = rnd.randint(2, 10)
+ for i in range(3):
+ shape = (rnd.randint(2, 10), ncols)
+ A, _ = rand_sparse_ndarray(shape, 'csr')
+ nds.append(A)
+ zero_nds.append(mx.nd.zeros(shape).tostype('csr'))
+ check_concat(nds)
+ check_concat(zero_nds)
+
+
+@with_seed()
def test_sparse_nd_equal():
for stype in ['row_sparse', 'csr']:
shape = rand_shape_2d()
--
To stop receiving notification emails like this one, please contact
haibin@apache.org.