Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Sparse retain improvement #138

Merged
merged 5 commits into from
Aug 5, 2017
Merged

Sparse retain improvement #138

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e8617fd
Add one more kernel for sparse retain
reminisce Aug 1, 2017
c81defb
Fix compile
reminisce Aug 2, 2017
0a72988
Change STORAGE_TYPE_ASSIGN_CHECK to type_assign for fallback
reminisce Aug 3, 2017
5c4e9c1
Fix
reminisce Aug 3, 2017
9da68f3
Add gpu compile
reminisce Aug 3, 2017
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
2 changes: 1 addition & 1 deletion src/c_api/c_api_ndarray.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -393,7 +393,7 @@ void PushOperator(const OpStatePtr& state,
#if MXNET_USE_CUDA
CastNonDefaultStorage<gpu>(temp_in_src, temp_in_dst, opctx);
fcompute(state, opctx, input_blobs, req, output_blobs);
CastNonDefaultStorage<gpu>(temp_our_dst, temp_out_src, opctx);
CastNonDefaultStorage<gpu>(temp_out_dst, temp_out_src, opctx);
#else
LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
#endif
Expand Down
2 changes: 1 addition & 1 deletion src/kvstore/comm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
#include <thread>
#include "mxnet/ndarray.h"
#include "../ndarray/ndarray_function.h"
#include "../operator/tensor/indexing_op.h"
#include "../operator/tensor/sparse_retain-inl.h"
namespace mxnet {
namespace kvstore {
/**
Expand Down
41 changes: 0 additions & 41 deletions src/operator/tensor/indexing_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,46 +229,5 @@ Examples::
.add_argument("indices", "NDArray-or-Symbol", "array of locations where to set on_value")
.add_arguments(OneHotParam::__FIELDS__());

NNVM_REGISTER_OP(sparse_retain)
.describe(R"code(pick rows specified by user input index array from a row sparse matrix
and save them in the output sparse matrix.

Example::

data = [[1, 2], [3, 4], [5, 6]]
indices = [0, 1, 3]
shape = (4, 2)
rsp_in = row_sparse(data, indices)
to_retain = [0, 3]
rsp_out = sparse_retain(rsp_in, to_retain)
rsp_out.values = [[1, 2], [5, 6]]
rsp_out.indices = [0, 3]

)code" ADD_FILELINE)
.set_num_inputs(2)
.set_num_outputs(1)
.set_attr<nnvm::FListInputNames>("FListInputNames",
[](const NodeAttrs& attrs) {
return std::vector<std::string>{"data", "indices"};
})
.set_attr<nnvm::FInferShape>("FInferShape", SparseRetainOpShape)
.set_attr<nnvm::FInferType>("FInferType", SparseRetainOpType)
.set_attr<FInferStorageType>("FInferStorageType", SparseRetainForwardInferStorageType)
.set_attr<FComputeEx>("FComputeEx<cpu>", SparseRetainOpForwardEx<cpu>)
.set_attr<nnvm::FGradient>("FGradient",
[](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
return MakeNonlossGradNode("_backward_sparse_retain", n, ograds,
{n->inputs[sr::kIdx]}, n->attrs.dict);
})
.add_argument("data", "NDArray-or-Symbol", "The input array for sparse_retain operator.")
.add_argument("indices", "NDArray-or-Symbol", "The index array of rows ids that will be retained.");

NNVM_REGISTER_OP(_backward_sparse_retain)
.set_num_inputs(2)
.set_num_outputs(2)
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
.set_attr<FInferStorageType>("FInferStorageType", SparseRetainBackwardInferStorageType)
.set_attr<FComputeEx>("FComputeEx<cpu>", SparseRetainOpBackwardEx<cpu>);

} // namespace op
} // namespace mxnet
203 changes: 0 additions & 203 deletions src/operator/tensor/indexing_op.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -670,209 +670,6 @@ void OneHotOpForward(const nnvm::NodeAttrs& attrs,
});
}

/*!
* \brief sparse retain namespace
*/
namespace sr {
enum SparseRetainOpInputs {kArr, kIdx};
enum SparseRetainOpOutputs {kOut};
} // namespace sr

inline bool SparseRetainOpShape(const nnvm::NodeAttrs& attrs,
std::vector<TShape> *in_attrs,
std::vector<TShape> *out_attrs) {
CHECK_EQ(in_attrs->size(), 2U)
<< "sparse_retain operator takes 2 arguments (" << in_attrs->size() << " given)";
CHECK_EQ(out_attrs->size(), 1U);

TShape tshape((*in_attrs)[sr::kArr]);
shape_assign(&tshape, (*out_attrs)[sr::kOut]);
SHAPE_ASSIGN_CHECK(*in_attrs, sr::kArr, tshape);
SHAPE_ASSIGN_CHECK(*out_attrs, sr::kOut, tshape);
return true;
}

inline bool SparseRetainOpType(const nnvm::NodeAttrs& attrs,
std::vector<int> *in_attrs,
std::vector<int> *out_attrs) {
CHECK_EQ(in_attrs->size(), 2U);
CHECK_EQ(out_attrs->size(), 1U);
CHECK_NE((*in_attrs)[sr::kIdx], -1) << "Index type must be set for sparse_retain operator";

TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[sr::kArr]);
TYPE_ASSIGN_CHECK(*in_attrs, 0, (*out_attrs)[sr::kOut]);
return (*in_attrs)[0] != -1;
}

inline bool SparseRetainForwardInferStorageType(const nnvm::NodeAttrs& attrs,
const Context& ctx,
std::vector<int> *in_attrs,
std::vector<int> *out_attrs) {
CHECK_EQ(in_attrs->size(), 2U);
CHECK_EQ(out_attrs->size(), 1U);
STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, sr::kArr, kRowSparseStorage);
STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, sr::kOut, kRowSparseStorage);
return true;
}

inline bool SparseRetainBackwardInferStorageType(const nnvm::NodeAttrs& attrs,
const Context& ctx,
std::vector<int> *in_attrs,
std::vector<int> *out_attrs) {
CHECK_EQ(in_attrs->size(), 2U);
CHECK_EQ(out_attrs->size(), 2U);
STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, sr::kOut, kDefaultStorage);
STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, sr::kIdx, kDefaultStorage);
STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, sr::kArr, kRowSparseStorage);
STORAGE_TYPE_ASSIGN_CHECK(*out_attrs, sr::kIdx, kDefaultStorage);
return true;
}

struct SparseRetainRspForward {
template<typename DType, typename RType, typename IType>
MSHADOW_XINLINE static void Map(int i, DType* out_data, RType* out_idx,
const DType* in_data, const RType* in_idx,
const IType* idx, const size_t nnr,
const size_t num_cols) {
const RType irow = idx[i];
int j = -1, left = 0, right = nnr - 1;
while (left <= right) {
int m = left + (right - left) / 2;
const auto in_idx_m = in_idx[m];
if (in_idx_m == irow) {
j = m;
break;
} else if (in_idx_m < irow) {
left = m + 1;
} else {
right = m - 1;
}
}
out_idx[i] = idx[i];
if (j >= 0) {
const size_t in_offset = j * num_cols;
const size_t out_offset = i * num_cols;
for (size_t k = 0; k < num_cols; ++k) {
out_data[out_offset+k] = in_data[in_offset+k];
}
}
}
};

template<typename xpu>
void SparseRetainOpForwardRspImpl(mshadow::Stream<xpu> *s, const NDArray &input,
const TBlob &indices, OpReqType req,
NDArray *output) {
using namespace rowsparse;
if (req == kNullOp || !input.storage_initialized() || indices.Size() == 0U) {
FillZerosRspImpl(s, output);
return;
}
const TBlob input_data = input.data();
const TBlob input_idx = input.aux_data(rowsparse::kIdx);
output->CheckAndAlloc({mshadow::Shape1(indices.Size())});
TBlob output_data = output->data();
TBlob output_idx = output->aux_data(rowsparse::kIdx);
const index_t row_length = input_data.shape_.ProdShape(1, input_data.shape_.ndim());

using namespace mxnet_op;
MSHADOW_TYPE_SWITCH(output_data.type_flag_, DType, { // output data type
MSHADOW_IDX_TYPE_SWITCH(output_idx.type_flag_, RType, { // row index data type
MSHADOW_TYPE_SWITCH(indices.type_flag_, IType, { // index array data type
Kernel<set_zero, xpu>::Launch(s, output_data.Size(), output_data.dptr<DType>());
Kernel<SparseRetainRspForward, xpu>::Launch(s, indices.Size(), output_data.dptr<DType>(),
output_idx.dptr<RType>(), input_data.dptr<DType>(), input_idx.dptr<RType>(),
indices.dptr<IType>(), input_data.shape_[0], row_length);
});
});
});
}

template<typename xpu>
void SparseRetainOpForwardEx(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(), 2U);
CHECK_EQ(outputs.size(), 1U);
CHECK_EQ(req.size(), 1U);
CHECK_EQ(req[sr::kOut], kWriteTo) << "sparse_retain only supports req=\'write\'";

CHECK_EQ(inputs[sr::kArr].storage_type(), kRowSparseStorage)
<< "sparse_retain operator only takes row sparse NDArray as input";
CHECK_EQ(inputs[sr::kIdx].storage_type(), kDefaultStorage)
<< "sparse_retain operator only takes default NDArray as its index array";
CHECK_EQ(outputs[sr::kOut].storage_type(), kRowSparseStorage)
<< "sparse_retain operator only outputs row sparse NDArray";

const NDArray& input_nd = inputs[sr::kArr];
const TBlob idx_data = inputs[sr::kIdx].data();
NDArray output_nd = outputs[sr::kOut];
mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
SparseRetainOpForwardRspImpl<xpu>(s, input_nd, idx_data, req[sr::kOut], &output_nd);
}

template<int req>
struct SparseRetainRspBackward {
template<typename DType, typename RType, typename IType>
MSHADOW_XINLINE static void Map(int i, DType* in_grad, RType* in_grad_idx,
const DType* out_grad, const IType* idx,
const size_t num_cols) {
const RType irow = idx[i];
in_grad_idx[i] = irow;
const size_t out_offset = irow * num_cols;
const size_t in_offset = i * num_cols;
for (size_t j = 0; j < num_cols; ++j) {
KERNEL_ASSIGN(in_grad[in_offset+j], req, out_grad[out_offset+j]);
}
}
};

template<typename xpu>
void SparseRetainOpBackwardEx(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(), 2U);
CHECK_EQ(outputs.size(), 2U);
CHECK_EQ(req.size(), 2U);
CHECK_NE(req[sr::kArr], kWriteInplace);
CHECK_EQ(req[sr::kIdx], kNullOp)
<< "sparse_retain does not support calculating gradients of indices";

CHECK_EQ(inputs[sr::kOut].storage_type(), kDefaultStorage)
<< "sparse_retain backward only takes default NDArray as ograd";
CHECK_EQ(inputs[sr::kIdx].storage_type(), kDefaultStorage)
<< "sparse_retain backward only takes default NDArray as its index array";
CHECK_EQ(outputs[sr::kArr].storage_type(), kRowSparseStorage)
<< "sparse_retain backward only outputs row sparse NDArray as grad of input";

const TBlob out_grad_data = inputs[sr::kOut].data();
const TBlob idx_data = inputs[sr::kIdx].data();

NDArray in_grad_nd = outputs[sr::kArr];
in_grad_nd.CheckAndAlloc({mshadow::Shape1(idx_data.Size())});
TBlob in_grad_data = in_grad_nd.data();
TBlob in_grad_idx = in_grad_nd.aux_data(rowsparse::kIdx);
const auto row_length = out_grad_data.shape_.ProdShape(1, out_grad_data.shape_.ndim());

using namespace mxnet_op;
Stream<xpu> *s = ctx.get_stream<xpu>();
MSHADOW_TYPE_SWITCH(out_grad_data.type_flag_, DType, { // output data type
MSHADOW_IDX_TYPE_SWITCH(in_grad_idx.type_flag_, RType, { // row index data type
MSHADOW_TYPE_SWITCH(idx_data.type_flag_, IType, { // index array data type
MXNET_ASSIGN_REQ_SWITCH(req[sr::kArr], req_type, {
Kernel<SparseRetainRspBackward<req_type>, xpu>::Launch(
s, in_grad_idx.Size(), in_grad_data.dptr<DType>(), in_grad_idx.dptr<RType>(),
out_grad_data.dptr<DType>(), idx_data.dptr<IType>(), row_length);
});
});
});
});
}

} // namespace op
} // namespace mxnet
#ifdef __CUDACC__
Expand Down
Loading