You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by jx...@apache.org on 2018/03/26 16:34:53 UTC
[incubator-mxnet] branch master updated: [MXNET-36] Update ndarray
binary ops to use kernel launch instead of mshadow operations (#9029)
This is an automated email from the ASF dual-hosted git repository.
jxie 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 cfc0941 [MXNET-36] Update ndarray binary ops to use kernel launch instead of mshadow operations (#9029)
cfc0941 is described below
commit cfc09410ddbcc194e1c6743be814cfa761af2d00
Author: Rahul Huilgol <ra...@gmail.com>
AuthorDate: Mon Mar 26 09:34:48 2018 -0700
[MXNET-36] Update ndarray binary ops to use kernel launch instead of mshadow operations (#9029)
* update binary ops to use kernel launch instead of mshadow eval
* change return of vector
* return const-vars
* address review comments
* comments added
---
include/mxnet/op_attr_types.h | 2 +-
src/kvstore/kvstore_dist_server.h | 5 +--
src/ndarray/ndarray.cc | 90 +++++++++++++++++++++++++++++++-------
src/ndarray/ndarray_function-inl.h | 30 +++++++++++--
src/ndarray/ndarray_function.h | 20 +++++----
5 files changed, 114 insertions(+), 33 deletions(-)
diff --git a/include/mxnet/op_attr_types.h b/include/mxnet/op_attr_types.h
index 820eb1a..800872b 100644
--- a/include/mxnet/op_attr_types.h
+++ b/include/mxnet/op_attr_types.h
@@ -49,8 +49,8 @@ enum OpReqType {
kWriteTo,
/*!
* \brief perform an inplace write,
- * Target shares memory with one of input arguments.
* This option only happen when
+ * Target shares memory with one of input arguments.
*/
kWriteInplace,
/*! \brief add to the provided space */
diff --git a/src/kvstore/kvstore_dist_server.h b/src/kvstore/kvstore_dist_server.h
index f1637c4..c2ddcd8 100644
--- a/src/kvstore/kvstore_dist_server.h
+++ b/src/kvstore/kvstore_dist_server.h
@@ -297,9 +297,6 @@ class KVStoreDistServer {
CopyFromTo(recved, &merged.array, 0);
} else {
NDArray out(kRowSparseStorage, stored.shape(), Context());
- std::vector<Engine::VarHandle> const_vars;
- const_vars.push_back(recved.var());
- const_vars.push_back(merged.array.var());
// accumulate row_sparse gradients
// TODO(haibin) override + operator for row_sparse NDArray
// instead of calling BinaryComputeRspRsp directly
@@ -309,7 +306,7 @@ class KVStoreDistServer {
op::ElemwiseBinaryOp::ComputeEx<cpu, op::mshadow_op::plus>(
{}, {}, {recved, merged.array}, {kWriteTo}, {out});
on_complete();
- }, recved.ctx(), const_vars, {out.var()},
+ }, recved.ctx(), {recved.var(), merged.array.var()}, {out.var()},
FnProperty::kNormal, 0, PROFILER_MESSAGE_FUNCNAME);
CopyFromTo(out, &merged.array, 0);
}
diff --git a/src/ndarray/ndarray.cc b/src/ndarray/ndarray.cc
index d4a6583..35dd445 100644
--- a/src/ndarray/ndarray.cc
+++ b/src/ndarray/ndarray.cc
@@ -824,16 +824,18 @@ void TernaryOp(const NDArray &lhs,
}
/*!
- * \brief run a binary operation
- * \param lhs left operand
- * \param rhs right operand
- * \param out the output ndarray
- * \param binary_op the real
- */
+* \brief Performs some preparation required to apply binary operators.
+* Checks context and shape of ndarrays, allocates space for output
+* and prepares const variables for engine
+* \param lhs left operand
+* \param rhs right operand
+* \param out the output ndarray
+* \param binary_op the real operation
+*/
template<typename OP>
-void BinaryOp(const NDArray &lhs,
- const NDArray &rhs,
- NDArray *out) {
+std::vector<Engine::VarHandle> BinaryOpPrepare(const NDArray &lhs,
+ const NDArray &rhs,
+ NDArray *out) {
// no check if both of them are on cpu
if (lhs.ctx().dev_mask() != cpu::kDevMask || rhs.ctx().dev_mask() != cpu::kDevMask) {
CHECK(lhs.ctx() == rhs.ctx()) << "operands context mismatch";
@@ -848,15 +850,71 @@ void BinaryOp(const NDArray &lhs,
CHECK(out->ctx() == lhs.ctx()) << "target context mismatch";
}
CHECK(out->shape() == OP::GetShape(lhs.shape(), rhs.shape()))
- << "target shape mismatch";
+ << "target shape mismatch";
}
+ std::vector<Engine::VarHandle> const_vars;
+ // prepare const variables for engine
+ if (lhs.var() != out->var()) const_vars.push_back(lhs.var());
+ if (rhs.var() != out->var()) const_vars.push_back(rhs.var());
+ return const_vars;
+}
+
+/*!
+* \brief run a binary operation using the kernel launch method
+* \param lhs left operand
+* \param rhs right operand
+* \param out the output ndarray
+* \param binary_op the real operation
+*/
+template<typename OP>
+void BinaryOpKernel(const NDArray &lhs,
+ const NDArray &rhs,
+ NDArray *out) {
+ std::vector<Engine::VarHandle> const_vars = BinaryOpPrepare<OP>(lhs, rhs, out);
// important: callback must always capture by value
NDArray ret = *out;
- // get the const variables
- std::vector<Engine::VarHandle> const_vars;
- if (lhs.var() != ret.var()) const_vars.push_back(lhs.var());
- if (rhs.var() != ret.var()) const_vars.push_back(rhs.var());
+ switch (lhs.ctx().dev_mask()) {
+ case cpu::kDevMask: {
+ Engine::Get()->PushSync([lhs, rhs, ret](RunContext ctx) {
+ TBlob tmp = ret.data();
+ mshadow::Stream<cpu>* s = ctx.get_stream<cpu>();
+ ndarray::BinaryOpKernelImpl<OP>(s, lhs.data(), rhs.data(), &tmp);
+ },
+ lhs.ctx(), const_vars, {ret.var()},
+ FnProperty::kNormal, 0, PROFILER_MESSAGE_FUNCNAME);
+ break;
+ }
+#if MXNET_USE_CUDA
+ case gpu::kDevMask: {
+ Engine::Get()->PushSync([lhs, rhs, ret](RunContext ctx) {
+ TBlob tmp = ret.data();
+ mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+ ndarray::BinaryOpKernelImpl<OP>(s, lhs.data(), rhs.data(), &tmp);
+ // Wait GPU kernel to complete
+ ctx.get_stream<gpu>()->Wait();
+ }, lhs.ctx(), const_vars, {ret.var()},
+ FnProperty::kNormal, 0, PROFILER_MESSAGE_FUNCNAME);
+ break;
+}
+#endif
+ default: LOG(FATAL) << MXNET_GPU_NOT_ENABLED_ERROR;
+ }
+}
+/*!
+ * \brief run a binary operation using mshadow operations
+ * \param lhs left operand
+ * \param rhs right operand
+ * \param out the output ndarray
+ * \param binary_op the real operation
+ */
+template<typename OP>
+void BinaryOp(const NDArray &lhs,
+ const NDArray &rhs,
+ NDArray *out) {
+ std::vector<Engine::VarHandle> const_vars = BinaryOpPrepare<OP>(lhs, rhs, out);
+ // important: callback must always capture by value
+ NDArray ret = *out;
// redirect everything to mshadow operations
switch (lhs.ctx().dev_mask()) {
case cpu::kDevMask: {
@@ -1421,7 +1479,7 @@ template<typename OP>
inline NDArray BinaryOpRet(const NDArray &lhs,
const NDArray &rhs) {
NDArray ret;
- BinaryOp<OP>(lhs, rhs, &ret);
+ BinaryOpKernel<OP>(lhs, rhs, &ret);
return ret;
}
@@ -1436,7 +1494,7 @@ inline NDArray ScalarOpRet(const NDArray &lhs,
template<typename OP>
inline NDArray &BinaryOpApply(NDArray *dst,
const NDArray &src) {
- BinaryOp<OP>(*dst, src, dst);
+ BinaryOpKernel<OP>(*dst, src, dst);
return *dst;
}
diff --git a/src/ndarray/ndarray_function-inl.h b/src/ndarray/ndarray_function-inl.h
index a80d9db..d494f08 100644
--- a/src/ndarray/ndarray_function-inl.h
+++ b/src/ndarray/ndarray_function-inl.h
@@ -47,6 +47,15 @@
}
#endif
+#ifndef DECL_BINARY_LAUNCH
+#define DECL_BINARY_LAUNCH(XPU, OP) \
+ template <> \
+ void BinaryOpKernelImpl<OP, XPU>(mshadow::Stream<XPU> *s, \
+ const TBlob& lhs, const TBlob& rhs, TBlob *out) { \
+ BinaryOpKernelLaunch<OP>(s, lhs, rhs, out); \
+ }
+#endif
+
#ifndef DECL_SCALAR
#define DECL_SCALAR(XPU, OP, FUN, REVERSE) \
template<> \
@@ -433,18 +442,31 @@ void EvalBroadcast<DEVICE>(TBlob const& src, TBlob* ret, int size, RunContext ct
out = mshadow::expr::broadcast_with_axis(in, 0, size);
}
+template<typename OP, typename xpu>
+void BinaryOpKernelLaunch(mshadow::Stream<xpu>* s, const TBlob& lhs, const TBlob& rhs, TBlob *out) {
+ using namespace op::mxnet_op;
+ using namespace mshadow;
+ MSHADOW_TYPE_SWITCH(out->type_flag_, DType, {
+ Kernel<op_with_req<OP, kWriteInplace>, xpu >::
+ Launch(s,
+ lhs.Size(),
+ out->dptr<DType>(),
+ lhs.dptr<DType>(),
+ rhs.dptr<DType>());
+ });
+}
// declarations
DECL_BINARY(DEVICE, MatChooseRowElem, EvalMatChooseRowElem_)
DECL_TERNARY(DEVICE, MatFillRowElem, EvalMatFillRowElem_)
DECL_BINARY(DEVICE, OneHotEncode, EvalOneHot_)
-DECL_BINARY(DEVICE, Plus, EvalBinary_)
-DECL_BINARY(DEVICE, Minus, EvalBinary_)
-DECL_BINARY(DEVICE, Mul, EvalBinary_)
-DECL_BINARY(DEVICE, Div, EvalBinary_)
DECL_SCALAR(DEVICE, Plus, EvalScalar_, true)
DECL_SCALAR(DEVICE, Minus, EvalScalar_, true)
DECL_SCALAR(DEVICE, Mul, EvalScalar_, true)
DECL_SCALAR(DEVICE, Div, EvalScalar_, true)
+DECL_BINARY_LAUNCH(DEVICE, Plus)
+DECL_BINARY_LAUNCH(DEVICE, Minus)
+DECL_BINARY_LAUNCH(DEVICE, Mul)
+DECL_BINARY_LAUNCH(DEVICE, Div)
// for reverse seq
DECL_SCALAR(DEVICE, Plus, EvalScalar_, false)
diff --git a/src/ndarray/ndarray_function.h b/src/ndarray/ndarray_function.h
index 518bb77..97c23b6 100644
--- a/src/ndarray/ndarray_function.h
+++ b/src/ndarray/ndarray_function.h
@@ -46,20 +46,20 @@ struct BinaryBase {
};
// operators
-struct Plus : public BinaryBase {
- typedef op::mshadow_op::plus mshadow_op;
+struct Plus : public BinaryBase, public mshadow::op::plus {
+ typedef mshadow::op::plus mshadow_op;
};
-struct Minus : public BinaryBase {
- typedef op::mshadow_op::minus mshadow_op;
+struct Minus : public BinaryBase, public mshadow::op::minus {
+ typedef mshadow::op::minus mshadow_op;
};
-struct Mul : public BinaryBase {
- typedef op::mshadow_op::mul mshadow_op;
+struct Mul : public BinaryBase, public mshadow::op::mul {
+ typedef mshadow::op::mul mshadow_op;
};
-struct Div : public BinaryBase {
- typedef op::mshadow_op::div mshadow_op;
+struct Div : public BinaryBase, public mshadow::op::div {
+ typedef mshadow::op::div mshadow_op;
};
struct Mod : public BinaryBase {
@@ -208,6 +208,10 @@ void Eval(mshadow::Stream<xpu> *s,
template <typename Device>
void EvalBroadcast(TBlob const& src, TBlob* ret, int size, RunContext ctx);
+template <typename OP, typename xpu>
+void BinaryOpKernelImpl(mshadow::Stream<xpu> *s, const TBlob& lhs,
+ const TBlob& rhs, TBlob *out);
+
} // namespace ndarray
} // namespace mxnet
#endif // MXNET_NDARRAY_NDARRAY_FUNCTION_H_
--
To stop receiving notification emails like this one, please contact
jxie@apache.org.