reminisce commented on a change in pull request #15940: Add boolean ndarray
URL: https://github.com/apache/incubator-mxnet/pull/15940#discussion_r315054507
##########
File path: src/operator/numpy/np_elemwise_broadcast_op.cc
##########
@@ -182,5 +114,211 @@
MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_rpower_scalar)
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu,
mshadow_op::rpower>)
.set_attr<nnvm::FGradient>("FGradient",
ElemwiseGradUseOut{"_backward_rpower_scalar"});
+#if MXNET_USE_TVM_OP
+static constexpr char func_equal_cpu[] = "equal_cpu";
+static constexpr char func_equal_gpu[] = "equal_gpu";
+static constexpr char func_not_equal_cpu[] = "not_equal_cpu";
+static constexpr char func_not_equal_gpu[] = "not_equal_gpu";
+static constexpr char func_greater_cpu[] = "greater_cpu";
+static constexpr char func_greater_gpu[] = "greater_gpu";
+static constexpr char func_less_cpu[] = "less_cpu";
+static constexpr char func_less_gpu[] = "less_gpu";
+static constexpr char func_greater_equal_cpu[] = "greater_equal_cpu";
+static constexpr char func_greater_equal_gpu[] = "greater_equal_gpu";
+static constexpr char func_less_equal_cpu[] = "less_equal_cpu";
+static constexpr char func_less_equal_gpu[] = "less_equal_gpu";
+
+bool NumpyBinaryLogicOpType(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);
+ if (in_attrs->at(0) == -1 && in_attrs->at(1) == -1) return false;
+ TYPE_ASSIGN_CHECK(*in_attrs, 0, in_attrs->at(1));
+ TYPE_ASSIGN_CHECK(*in_attrs, 1, in_attrs->at(0));
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kBool);
+ return true;
+}
+
+static TBlob PrependAxes(const TBlob& src, const int dst_ndim) {
+ CHECK_LE(src.shape_.ndim(), dst_ndim);
+ const int src_ndim = src.shape_.ndim();
+ if (src_ndim == dst_ndim) return src;
+ mxnet::TShape dst_shape(dst_ndim, 1);
+ for (int i = dst_ndim - src_ndim; i < dst_ndim; ++i) {
+ dst_shape[i] = src.shape_[i - dst_ndim + src_ndim];
+ }
+ return src.reshape(dst_shape);
+}
+
+template<const char* func>
+void TVMBinaryBroadcastCompute(const nnvm::NodeAttrs& attrs,
+ const mxnet::OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ CHECK_EQ(inputs.size(), 2U);
+ CHECK_EQ(outputs.size(), 1U);
+ if (outputs[0].shape_.Size() == 0U) return; // skip zero-size tensor
+
+ // prepare tblobs and TVMArgs
+ std::vector<TBlob> tblobs = {inputs[0], inputs[1], outputs[0]};
+ std::vector<int> type_codes;
+ std::vector<TVMValue> values;
+
+ const int ondim = outputs[0].shape_.ndim();
+ const size_t num_args = inputs.size() + outputs.size();
+ type_codes.resize(num_args);
+ values.resize(num_args);
+ for (size_t i = 0; i < num_args; ++i) {
+ tblobs[i] = PrependAxes(tblobs[i], ondim);
+ type_codes[i] = kArrayHandle;
+ values[i].v_handle = const_cast<DLTensor*>(&(tblobs[i].dltensor()));
+ }
+ tvm::runtime::TVMArgs tvm_args(&values[0], &type_codes[0], tblobs.size());
+ tvm::runtime::TVMOpModule::Get()->CallEx(func, ctx, tblobs, tvm_args);
+}
+
+#define MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(name)
\
+ NNVM_REGISTER_OP(_npi_##name)
\
+ .set_num_inputs(2)
\
+ .set_num_outputs(1)
\
+ .set_attr<nnvm::FListInputNames>("FListInputNames",
\
+ [](const NodeAttrs& attrs) {
\
+ return std::vector<std::string>{"lhs", "rhs"};
\
+ })
\
+ .set_attr<mxnet::FInferShape>("FInferShape", BinaryBroadcastShape)
\
+ .set_attr<nnvm::FInferType>("FInferType", NumpyBinaryLogicOpType)
\
+ .set_attr<nnvm::FInplaceOption>("FInplaceOption",
\
+ [](const NodeAttrs& attrs) {
\
+ return std::vector<std::pair<int, int> >{{0, 0}, {1, 0}};
\
+ })
\
+ .set_attr<FCompute>("FCompute<cpu>",
TVMBinaryBroadcastCompute<func_##name##_cpu>) \
+ .set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
\
+ .add_argument("lhs", "NDArray-or-Symbol", "First input to the function")
\
+ .add_argument("rhs", "NDArray-or-Symbol", "Second input to the function")
+
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(equal);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(not_equal);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(greater);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(less);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(greater_equal);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC(less_equal);
+
+#if MXNET_USE_CUDA
+#define MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(name)
\
+ NNVM_REGISTER_OP(_npi_##name)
\
+ .set_attr<FCompute>("FCompute<gpu>",
TVMBinaryBroadcastCompute<func_##name##_gpu>)
+
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(equal);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(not_equal);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(greater);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(less);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(greater_equal);
+MXNET_OPERATOR_REGISTER_NP_BINARY_LOGIC_GPU(less_equal);
+#endif // MXNET_USE_CUDA
+
+bool NumpyBinaryScalarLogicOpType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1U);
+ CHECK_EQ(out_attrs->size(), 1U);
+ if (in_attrs->at(0) == -1) return false;
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kBool);
+ return true;
+}
+
+template<const char* func>
+void TVMBinaryBroadcastScalarCompute(const nnvm::NodeAttrs& attrs,
+ const mxnet::OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
+ CHECK_EQ(inputs.size(), 1U);
+ CHECK_EQ(outputs.size(), 1U);
+ if (outputs[0].shape_.Size() == 0U) return; // skip zero-size tensor
+
+ // prepare tblobs and TVMArgs
+ std::vector<TBlob> tblobs = {inputs[0], outputs[0]};
+ std::vector<int> type_codes;
+ std::vector<TVMValue> values;
+
+ const size_t num_args = 3; // one input tensor, one scalar param, and one
output
+ type_codes.resize(num_args);
+ values.resize(num_args);
+
+ // input tensor setup
+ type_codes[0] = kArrayHandle;
+ values[0].v_handle = const_cast<DLTensor*>(&(tblobs[0].dltensor()));
+
+ // scalar param
+ type_codes[1] = kDLFloat;
Review comment:
@yzhliu Since I need to pass a `double` param to the op func generated by
TVM, I cannot use the `Call` function defined in the `TVMOpModule`. I moved the
logic of preparing `TVMArgs` up here from the `Call` function to MXNet op's
`FCompute` function and added an independent `CallEx` in `TVMOpModule` to just
invoke the kernel. We can discuss the change of the API to cater for more use
cases.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
