Skip to content

Commit

Permalink
Sparse retain improvement (#138)
Browse files Browse the repository at this point in the history 
* Add one more kernel for sparse retain

* Fix compile

* Change STORAGE_TYPE_ASSIGN_CHECK to type_assign for fallback

* Fix

* Add gpu compile
  • Loading branch information
@reminisce @eric-haibin-lin
reminisce authored and eric-haibin-lin committed Aug 5, 2017
1 parent 325f4db commit 5866b2b
Show file tree
Hide file tree
Showing 9 changed files with 410 additions and 252 deletions.
2 changes: 1 addition & 1 deletion src/c_api/c_api_ndarray.cc
View file
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
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
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
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

0 comments on commit 5866b2b

Please sign in to comment.