Skip to content

Commit 96c1fe1

Browse files
eric-haibin-linChaiBapchya
eric-haibin-lin
authored and
ChaiBapchya
committed
Sparse support for logic ops (apache#12860)
* remove check * fix lint * fix gpu build
1 parent 0cd8b51 commit 96c1fe1

4 files changed

Lines changed: 107 additions & 20 deletions

File tree

‎src/operator/tensor/elemwise_binary_scalar_op.h‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,28 @@ class BinaryScalarOp : public UnaryOp {
273273
}
274274
}
275275

276+
template<typename xpu, typename OP>
277+
static void LogicComputeEx(const nnvm::NodeAttrs &attrs,
278+
const OpContext &ctx,
279+
const std::vector<NDArray> &inputs,
280+
const std::vector<OpReqType> &req,
281+
const std::vector<NDArray> &outputs) {
282+
DCHECK_EQ(inputs.size(), 1);
283+
DCHECK_EQ(outputs.size(), 1);
284+
const auto in_stype = inputs[0].storage_type();
285+
const auto out_stype = outputs[0].storage_type();
286+
if (req[0] == kNullOp) {
287+
return;
288+
}
289+
if ((in_stype == kRowSparseStorage && out_stype == kRowSparseStorage) ||
290+
(in_stype == kCSRStorage && out_stype == kCSRStorage)) {
291+
// csr -> csr, or rsp -> rsp
292+
UnaryOp::MapToFCompute<xpu>(attrs, ctx, inputs, req, outputs, Compute<xpu, OP>);
293+
} else {
294+
LogUnimplementedOp(attrs, ctx, inputs, req, outputs);
295+
}
296+
}
297+
276298
template<typename xpu, typename OP>
277299
static void Backward(const nnvm::NodeAttrs &attrs,
278300
const OpContext &ctx,

‎src/operator/tensor/elemwise_binary_scalar_op_logic.cc‎

Lines changed: 47 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -29,33 +29,68 @@
2929
namespace mxnet {
3030
namespace op {
3131

32-
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_equal_scalar)
33-
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::eq>)
32+
#define MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(__name$, __kernel$) \
33+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(__name$) \
34+
.set_attr<FInferStorageType>("FInferStorageType", BinaryScalarLogicStorageType<__kernel$>) \
35+
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, __kernel$>) \
36+
.set_attr<FComputeEx>("FComputeEx<cpu>", BinaryScalarOp::LogicComputeEx<cpu, __kernel$>)
37+
38+
template<typename OP>
39+
static bool BinaryScalarLogicStorageType(const nnvm::NodeAttrs& attrs,
40+
const int dev_mask,
41+
DispatchMode* dispatch_mode,
42+
std::vector<int> *in_attrs,
43+
std::vector<int> *out_attrs) {
44+
CHECK_EQ(in_attrs->size(), 1);
45+
CHECK_EQ(out_attrs->size(), 1);
46+
const auto in_stype = in_attrs->at(0);
47+
auto &out_stype = out_attrs->at(0);
48+
bool dispatched = false;
49+
const double alpha = nnvm::get<double>(attrs.parsed);
50+
bool is_sparse = OP::Map(static_cast<double>(0), alpha) == 0;
51+
if (!dispatched && in_stype == kDefaultStorage) {
52+
// dns -> dns
53+
dispatched = storage_type_assign(&out_stype, kDefaultStorage,
54+
dispatch_mode, DispatchMode::kFCompute);
55+
}
56+
if (!dispatched && in_stype == kRowSparseStorage && is_sparse) {
57+
// rsp -> rsp
58+
dispatched = storage_type_assign(&out_stype, kRowSparseStorage,
59+
dispatch_mode, DispatchMode::kFComputeEx);
60+
}
61+
if (!dispatched && in_stype == kCSRStorage && is_sparse) {
62+
// csr -> csr
63+
dispatched = storage_type_assign(&out_stype, kCSRStorage,
64+
dispatch_mode, DispatchMode::kFComputeEx);
65+
}
66+
if (!dispatched) {
67+
dispatched = dispatch_fallback(out_attrs, dispatch_mode);
68+
}
69+
return dispatched;
70+
}
71+
72+
73+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(_equal_scalar, mshadow_op::eq)
3474
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
3575
.add_alias("_EqualScalar");
3676

37-
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_not_equal_scalar)
38-
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::ne>)
77+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(_not_equal_scalar, mshadow_op::ne)
3978
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
4079
.add_alias("_NotEqualScalar");
4180

42-
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_greater_scalar)
43-
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::gt>)
81+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(_greater_scalar, mshadow_op::gt)
4482
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
4583
.add_alias("_GreaterScalar");
4684

47-
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_greater_equal_scalar)
48-
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::ge>)
85+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(_greater_equal_scalar, mshadow_op::ge)
4986
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
5087
.add_alias("_GreaterEqualScalar");
5188

52-
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_lesser_scalar)
53-
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::lt>)
89+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(_lesser_scalar, mshadow_op::lt)
5490
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
5591
.add_alias("_LesserScalar");
5692

57-
MXNET_OPERATOR_REGISTER_BINARY_SCALAR(_lesser_equal_scalar)
58-
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::le>)
93+
MXNET_OPERATOR_REGISTER_BINARY_SCALAR_LOGIC(_lesser_equal_scalar, mshadow_op::le)
5994
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
6095
.add_alias("_LesserEqualScalar");
6196

‎src/operator/tensor/elemwise_binary_scalar_op_logic.cu‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,22 +28,28 @@ namespace mxnet {
2828
namespace op {
2929

3030
NNVM_REGISTER_OP(_equal_scalar)
31-
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::eq>);
31+
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::eq>)
32+
.set_attr<FComputeEx>("FComputeEx<gpu>", BinaryScalarOp::LogicComputeEx<gpu, mshadow_op::eq>);
3233

3334
NNVM_REGISTER_OP(_not_equal_scalar)
34-
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::ne>);
35+
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::ne>)
36+
.set_attr<FComputeEx>("FComputeEx<gpu>", BinaryScalarOp::LogicComputeEx<gpu, mshadow_op::ne>);
3537

3638
NNVM_REGISTER_OP(_greater_scalar)
37-
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::gt>);
39+
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::gt>)
40+
.set_attr<FComputeEx>("FComputeEx<gpu>", BinaryScalarOp::LogicComputeEx<gpu, mshadow_op::gt>);
3841

3942
NNVM_REGISTER_OP(_greater_equal_scalar)
40-
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::ge>);
43+
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::ge>)
44+
.set_attr<FComputeEx>("FComputeEx<gpu>", BinaryScalarOp::LogicComputeEx<gpu, mshadow_op::ge>);
4145

4246
NNVM_REGISTER_OP(_lesser_scalar)
43-
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::lt>);
47+
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::lt>)
48+
.set_attr<FComputeEx>("FComputeEx<gpu>", BinaryScalarOp::LogicComputeEx<gpu, mshadow_op::lt>);
4449

4550
NNVM_REGISTER_OP(_lesser_equal_scalar)
46-
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::le>);
51+
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::le>)
52+
.set_attr<FComputeEx>("FComputeEx<gpu>", BinaryScalarOp::LogicComputeEx<gpu, mshadow_op::le>);
4753

4854
NNVM_REGISTER_OP(_logical_and_scalar)
4955
.set_attr<FCompute>("FCompute<gpu>", BinaryScalarOp::Compute<gpu, mshadow_op::logical_and>);

‎tests/python/unittest/test_sparse_ndarray.py‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,12 @@ def test_sparse_nd_equal():
180180
y = sparse_nd_ones(shape, stype)
181181
z = x == y
182182
assert (z.asnumpy() == np.zeros(shape)).all()
183-
z = 0 == x
183+
z = 0 == y
184+
assert (z.asnumpy() == np.zeros(shape)).all()
185+
assert z.stype == 'default'
186+
z = 1 == y
184187
assert (z.asnumpy() == np.ones(shape)).all()
188+
assert z.stype == stype
185189

186190

187191
@with_seed()
@@ -192,8 +196,12 @@ def test_sparse_nd_not_equal():
192196
y = sparse_nd_ones(shape, stype)
193197
z = x != y
194198
assert (z.asnumpy() == np.ones(shape)).all()
195-
z = 0 != x
199+
z = 0 != y
200+
assert (z.asnumpy() == np.ones(shape)).all()
201+
assert z.stype == stype
202+
z = 1 != y
196203
assert (z.asnumpy() == np.zeros(shape)).all()
204+
assert z.stype == 'default'
197205

198206

199207
@with_seed()
@@ -206,8 +214,13 @@ def test_sparse_nd_greater():
206214
assert (z.asnumpy() == np.zeros(shape)).all()
207215
z = y > 0
208216
assert (z.asnumpy() == np.ones(shape)).all()
217+
assert z.stype == stype
209218
z = 0 > y
210219
assert (z.asnumpy() == np.zeros(shape)).all()
220+
assert z.stype == stype
221+
z = y > 1
222+
assert (z.asnumpy() == np.zeros(shape)).all()
223+
assert z.stype == stype
211224

212225

213226
@with_seed()
@@ -220,10 +233,13 @@ def test_sparse_nd_greater_equal():
220233
assert (z.asnumpy() == np.zeros(shape)).all()
221234
z = y >= 0
222235
assert (z.asnumpy() == np.ones(shape)).all()
236+
assert z.stype == 'default'
223237
z = 0 >= y
224238
assert (z.asnumpy() == np.zeros(shape)).all()
239+
assert z.stype == 'default'
225240
z = y >= 1
226241
assert (z.asnumpy() == np.ones(shape)).all()
242+
assert z.stype == stype
227243

228244

229245
@with_seed()
@@ -236,8 +252,13 @@ def test_sparse_nd_lesser():
236252
assert (z.asnumpy() == np.zeros(shape)).all()
237253
z = 0 < y
238254
assert (z.asnumpy() == np.ones(shape)).all()
255+
assert z.stype == stype
239256
z = y < 0
240257
assert (z.asnumpy() == np.zeros(shape)).all()
258+
assert z.stype == stype
259+
z = y < 1
260+
assert (z.asnumpy() == np.zeros(shape)).all()
261+
assert z.stype == 'default'
241262

242263

243264
@with_seed()
@@ -250,10 +271,13 @@ def test_sparse_nd_lesser_equal():
250271
assert (z.asnumpy() == np.zeros(shape)).all()
251272
z = 0 <= y
252273
assert (z.asnumpy() == np.ones(shape)).all()
274+
assert z.stype == 'default'
253275
z = y <= 0
254276
assert (z.asnumpy() == np.zeros(shape)).all()
277+
assert z.stype == 'default'
255278
z = 1 <= y
256279
assert (z.asnumpy() == np.ones(shape)).all()
280+
assert z.stype == stype
257281

258282

259283
@with_seed()

0 commit comments

Comments
 (0)