You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2017/10/27 17:13:11 UTC
[GitHub] cjolivier01 closed pull request #8340: Optimizations to set_to, copy, fill ops
cjolivier01 closed pull request #8340: Optimizations to set_to, copy, fill ops
URL: https://github.com/apache/incubator-mxnet/pull/8340
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py
index 1cd9f40e52..9eeddda30c 100644
--- a/python/mxnet/ndarray/ndarray.py
+++ b/python/mxnet/ndarray/ndarray.py
@@ -1854,8 +1854,10 @@ def full(shape, val, ctx=None, dtype=mx_real_t, out=None):
>>> mx.nd.full((1, 2), 2.0, dtype='float16').asnumpy()
array([[ 2., 2.]], dtype=float16)
"""
- out = empty(shape, ctx, dtype) if out is None else out
- out[:] = val
+ if ctx is None:
+ ctx = Context.default_ctx
+ dtype = mx_real_t if dtype is None else dtype
+ out = _internal._full(shape=shape, ctx=ctx, dtype=dtype, value=val, out=out)
return out
diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
index 6903db0b0d..4713c1ee1c 100644
--- a/python/mxnet/symbol/symbol.py
+++ b/python/mxnet/symbol/symbol.py
@@ -2759,7 +2759,7 @@ def full(shape, val, dtype=None, **kwargs):
"""
if dtype is None:
dtype = _numpy.float32
- return _internal._MulScalar(ones(shape=shape, dtype=dtype, **kwargs), scalar=val)
+ return _internal._full(shape=shape, dtype=dtype, value=val, **kwargs)
# pylint: disable=redefined-outer-name
def arange(start, stop=None, step=1.0, repeat=1, name=None, dtype=None):
diff --git a/src/operator/mxnet_op.h b/src/operator/mxnet_op.h
index bd3b0f27c5..c74b21605b 100644
--- a/src/operator/mxnet_op.h
+++ b/src/operator/mxnet_op.h
@@ -199,21 +199,25 @@ MSHADOW_XINLINE Shape<ndim> calc_stride(const Shape<ndim>& shape) {
return stride;
}
-
-struct fill {
- template<typename DType>
- MSHADOW_XINLINE static void Map(int i, DType* out, const DType val) {
- out[i] = val;
- }
-};
-
-
-struct set_zero {
- template<typename DType>
- MSHADOW_XINLINE static void Map(int i, DType* out) {
- out[i] = static_cast<DType>(0);
- }
-};
+/*!
+ * \brief Simple copy data from one blob to another
+ * \param to Destination blob
+ * \param from Source blob
+ */
+template <typename xpu>
+MSHADOW_CINLINE void copy(mshadow::Stream<xpu> *s, const TBlob& to, const TBlob& from) {
+ CHECK_EQ(from.Size(), to.Size());
+ CHECK_EQ(from.dev_mask(), to.dev_mask());
+ MSHADOW_TYPE_SWITCH(to.type_flag_, DType, {
+ if (to.type_flag_ == from.type_flag_) {
+ mshadow::Copy(to.FlatTo1D<xpu, DType>(), from.FlatTo1D<xpu, DType>(), s);
+ } else {
+ MSHADOW_TYPE_SWITCH(from.type_flag_, SrcDType, {
+ to.FlatTo1D<xpu, DType>(s) = mshadow::expr::tcast<DType>(from.FlatTo1D<xpu, SrcDType>(s));
+ })
+ }
+ })
+}
/*! \brief Binary op backward gradient OP wrapper */
template<typename GRAD_OP>
@@ -234,22 +238,75 @@ struct backward_grad {
*/
template<typename OP, int req>
struct op_with_req {
+ typedef OP Operation;
+
+ /*! \brief input is one tensor */
template<typename DType>
MSHADOW_XINLINE static void Map(int i, DType *out, const DType *in) {
KERNEL_ASSIGN(out[i], req, OP::Map(in[i]));
}
+ /*! \brief inputs are two tensors */
template<typename DType>
MSHADOW_XINLINE static void Map(int i, DType *out, const DType *lhs, const DType *rhs) {
KERNEL_ASSIGN(out[i], req, OP::Map(lhs[i], rhs[i]));
}
+ /*! \brief input is tensor and a scalar value */
template<typename DType>
MSHADOW_XINLINE static void Map(int i, DType *out, const DType *in, const DType value) {
KERNEL_ASSIGN(out[i], req, OP::Map(in[i], value));
}
+
+ /*! \brief No inputs (ie fill to constant value) */
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType *out) {
+ KERNEL_ASSIGN(out[i], req, OP::Map());
+ }
+
+ /*! \brief input is single scalar value */
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType *out, const DType value) {
+ KERNEL_ASSIGN(out[i], req, OP::Map(value));
+ }
+
+ /*! \brief inputs are two tensors and a scalar value */
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType *out,
+ const DType *input_1, const DType *input_2, const DType value) {
+ KERNEL_ASSIGN(out[i], req, OP::Map(input_1[i], input_2[i], value));
+ }
+
+ /*! \brief inputs are three tensors (ie backward grad with binary grad function) */
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType *out,
+ const DType *input_1,
+ const DType *input_2,
+ const DType *input_3) {
+ KERNEL_ASSIGN(out[i], req, OP::Map(input_1[i], input_2[i], input_3[i]));
+ }
};
+/*!
+ * \brief Set to immediate scalar value kernel
+ * \tparam val Scalar immediate
+ */
+template<int val>
+struct set_to_int {
+ // mxnet_op version (when used directly with Kernel<>::Launch()) */
+ template<typename DType>
+ MSHADOW_XINLINE static void Map(int i, DType* out) {
+ out[i] = DType(val);
+ }
+ // mshadow_op version (when used with op_with_req<>)
+ MSHADOW_XINLINE static int Map() {
+ return val;
+ }
+};
+
+/*! \brief Special-case kernel shortcut for setting to zero */
+using set_zero = set_to_int<0>;
+
template<typename OP, typename xpu>
struct Kernel;
diff --git a/src/operator/tensor/elemwise_binary_op.h b/src/operator/tensor/elemwise_binary_op.h
index 10c38de6f3..6a1cc02c87 100644
--- a/src/operator/tensor/elemwise_binary_op.h
+++ b/src/operator/tensor/elemwise_binary_op.h
@@ -103,7 +103,7 @@ class ElemwiseBinaryOp : public OpBase {
#pragma omp parallel for
for (int i = iter_out; i < index_out_min; ++i) {
MXNET_ASSIGN_REQ_SWITCH(req, Req, {
- SerialLaunchCPU<SetToScalar<Req>>(s, size, (*out)[i].dptr_, zero_input_val);
+ SerialLaunchCPU<OpBase::set_to_scalar<Req>>(s, size, (*out)[i].dptr_, zero_input_val);
});
}
}
diff --git a/src/operator/tensor/elemwise_binary_scalar_op.h b/src/operator/tensor/elemwise_binary_scalar_op.h
index 989ecb2714..b866a296f6 100644
--- a/src/operator/tensor/elemwise_binary_scalar_op.h
+++ b/src/operator/tensor/elemwise_binary_scalar_op.h
@@ -66,7 +66,7 @@ class BinaryScalarOp : public UnaryOp {
const int64_t dense_block_count = next_input_row - output_row;
if (dense_block_count > 0) {
MXNET_ASSIGN_REQ_SWITCH(req, Req, {
- mxnet_op::Kernel<OpBase::SetToScalar<Req>, cpu>::Launch(
+ mxnet_op::Kernel<OpBase::set_to_scalar<Req>, cpu>::Launch(
stream,
items_per_row * dense_block_count,
output_data.dptr_ + items_per_row * output_row,
diff --git a/src/operator/tensor/elemwise_unary_op.h b/src/operator/tensor/elemwise_unary_op.h
index 60bac8c347..d455b7e761 100644
--- a/src/operator/tensor/elemwise_unary_op.h
+++ b/src/operator/tensor/elemwise_unary_op.h
@@ -27,9 +27,11 @@
#include <mxnet/operator_util.h>
#include <vector>
#include <utility>
+#include <algorithm>
#include "../mshadow_op.h"
#include "../mxnet_op.h"
#include "../elemwise_op_common.h"
+#include "../../ndarray/ndarray_function.h"
namespace mxnet {
namespace op {
@@ -51,15 +53,10 @@ class OpBase {
}
}
+ /*! \brief simple kernel to set to a scalar value of arbitrary type */
template<int req>
- struct SetToScalar {
- template<typename DType>
- MSHADOW_XINLINE static void Map(int i, DType *out, const DType value) {
- KERNEL_ASSIGN(out[i], req, value);
- }
- };
+ using set_to_scalar = mxnet_op::op_with_req<mshadow_op::identity, req>;
- protected:
/*! \brief Copy blob data */
template<typename xpu>
static void inline CopyBlob(mshadow::Stream<xpu> *s,
@@ -174,7 +171,7 @@ class OpBase {
const OpReqType req,
DType *out) {
MXNET_ASSIGN_REQ_SWITCH(req, Req, {
- SerialLaunchCPU<SetToScalar<Req>>(s, size, out, val);
+ SerialLaunchCPU<OpBase::set_to_scalar<Req>>(s, size, out, val);
});
}
}; // OpBase
@@ -333,23 +330,31 @@ class UnaryOp : public OpBase {
template<typename xpu>
static void IdentityCompute(const nnvm::NodeAttrs& attrs,
- const OpContext& ctx,
- const std::vector<TBlob>& inputs,
- const std::vector<OpReqType>& req,
- const std::vector<TBlob>& outputs) {
+ const OpContext& ctx,
+ const std::vector<TBlob>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<TBlob>& outputs) {
using namespace mshadow;
using namespace mshadow::expr;
- Stream<xpu> *s = ctx.get_stream<xpu>();
- if (req[0] == kNullOp) return;
- if (req[0] == kWriteInplace) {
- CHECK_EQ(inputs[0].dptr_, outputs[0].dptr_); return;
+ switch (req[0]) {
+ case kWriteTo:
+ CHECK_EQ(outputs[0].dev_mask(), inputs[0].dev_mask());
+ mxnet_op::copy(ctx.get_stream<xpu>(), outputs[0], inputs[0]);
+ break;
+ case kAddTo: {
+ Stream<xpu> *s = ctx.get_stream<xpu>();
+ MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+ mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, kAddTo>, xpu>::Launch(
+ s, inputs[0].Size(), outputs[0].dptr<DType>(), inputs[0].dptr<DType>());
+ });
+ }
+ break;
+ case kWriteInplace:
+ CHECK_EQ(inputs[0].dptr_, outputs[0].dptr_);
+ break;
+ case kNullOp:
+ break;
}
- MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
- MXNET_ASSIGN_REQ_SWITCH(req[0], Req, {
- mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, Req>, xpu>::Launch(
- s, inputs[0].Size(), outputs[0].dptr<DType>(), inputs[0].dptr<DType>());
- });
- });
}
template<typename xpu>
diff --git a/src/operator/tensor/indexing_op.h b/src/operator/tensor/indexing_op.h
index 36b6e76f70..1c731088ee 100644
--- a/src/operator/tensor/indexing_op.h
+++ b/src/operator/tensor/indexing_op.h
@@ -854,7 +854,6 @@ void ScatterNDForward(const nnvm::NodeAttrs& attrs,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs) {
- using namespace mxnet_op;
using namespace mshadow;
CHECK_EQ(inputs.size(), 2U);
CHECK_EQ(outputs.size(), 1U);
@@ -868,14 +867,11 @@ void ScatterNDForward(const nnvm::NodeAttrs& attrs,
mshadow::Shape<10> strides;
for (int i = M-1, stride = K; i >= 0; stride *= oshape[i], --i) strides[i] = stride;
MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, DType, { // output data type switch
- if (req[0] == kWriteTo) {
- Kernel<fill, xpu>::Launch(s, oshape.Size(), outputs[0].dptr<DType>(),
- static_cast<DType>(0));
- }
+ Fill<true>(s, outputs[0], req[0], 0);
MSHADOW_TYPE_SWITCH(inputs[1].type_flag_, IType, { // indices data type switch
- Kernel<scatter_nd, xpu>::Launch(
- s, N, req[0], N, M, K, strides, outputs[0].dptr<DType>(),
- inputs[0].dptr<DType>(), inputs[1].dptr<IType>());
+ mxnet_op::Kernel<scatter_nd, xpu>::Launch(
+ s, N, req[0], N, M, K, strides, outputs[0].dptr<DType>(),
+ inputs[0].dptr<DType>(), inputs[1].dptr<IType>());
});
});
}
diff --git a/src/operator/tensor/init_op.cc b/src/operator/tensor/init_op.cc
index e5d11e4733..52e488acd9 100644
--- a/src/operator/tensor/init_op.cc
+++ b/src/operator/tensor/init_op.cc
@@ -28,6 +28,7 @@ namespace mxnet {
namespace op {
DMLC_REGISTER_PARAMETER(InitOpParam);
+DMLC_REGISTER_PARAMETER(InitOpWithScalarParam);
DMLC_REGISTER_PARAMETER(RangeParam);
@@ -53,6 +54,16 @@ NNVM_REGISTER_OP(_ones)
.set_attr<FCompute>("FCompute<cpu>", FillCompute<cpu, 1>)
.add_arguments(InitOpParam::__FIELDS__());
+NNVM_REGISTER_OP(_full)
+ .describe("fill target with a scalar value")
+ .set_num_inputs(0)
+ .set_num_outputs(1)
+ .set_attr_parser(ParamParser<InitOpWithScalarParam>)
+ .set_attr<nnvm::FInferShape>("FInferShape", InitShape<InitOpWithScalarParam>)
+ .set_attr<nnvm::FInferType>("FInferType", InitType<InitOpWithScalarParam>)
+ .set_attr<FCompute>("FCompute<cpu>", InitFillWithScalarCompute<cpu>)
+.add_arguments(InitOpWithScalarParam::__FIELDS__());
+
NNVM_REGISTER_OP(_arange)
.describe("Return evenly spaced values within a given interval. Similar to Numpy")
.set_num_inputs(0)
diff --git a/src/operator/tensor/init_op.cu b/src/operator/tensor/init_op.cu
index a1982ca236..4c16a94353 100644
--- a/src/operator/tensor/init_op.cu
+++ b/src/operator/tensor/init_op.cu
@@ -49,6 +49,9 @@ NNVM_REGISTER_OP(_zeros)
NNVM_REGISTER_OP(_ones)
.set_attr<FCompute>("FCompute<gpu>", FillCompute<gpu, 1>);
+NNVM_REGISTER_OP(_full)
+.set_attr<FCompute>("FCompute<gpu>", InitFillWithScalarCompute<gpu>);
+
NNVM_REGISTER_OP(_arange)
.set_attr<FCompute>("FCompute<gpu>", RangeCompute<gpu>);
diff --git a/src/operator/tensor/init_op.h b/src/operator/tensor/init_op.h
index ea4243ee91..b2bd15ceec 100644
--- a/src/operator/tensor/init_op.h
+++ b/src/operator/tensor/init_op.h
@@ -34,7 +34,7 @@
#include <limits>
#include "../elemwise_op_common.h"
#include "../mxnet_op.h"
-
+#include "../mshadow_op.h"
namespace mxnet {
namespace op {
@@ -98,6 +98,32 @@ struct RangeParam : public dmlc::Parameter<RangeParam> {
}
};
+/*! \brief Initialize and fill output with an arbitrary value */
+struct InitOpWithScalarParam : dmlc::Parameter<InitOpWithScalarParam> {
+ TShape shape;
+ std::string ctx;
+ int dtype;
+ double value;
+ DMLC_DECLARE_PARAMETER(InitOpWithScalarParam) {
+ DMLC_DECLARE_FIELD(shape)
+ .set_default(TShape())
+ .describe("The shape of the output");
+ DMLC_DECLARE_FIELD(ctx)
+ .set_default("")
+ .describe("Context of output, in format [cpu|gpu|cpu_pinned](n)."
+ "Only used for imperative calls.");
+ DMLC_DECLARE_FIELD(dtype).set_default(mshadow::kFloat32)
+ .add_enum("float32", mshadow::kFloat32)
+ .add_enum("float64", mshadow::kFloat64)
+ .add_enum("float16", mshadow::kFloat16)
+ .add_enum("uint8", mshadow::kUint8)
+ .add_enum("int32", mshadow::kInt32)
+ .describe("Target data type.");
+ DMLC_DECLARE_FIELD(value)
+ .describe("Value with which to fill newly created tensor");
+ }
+};
+
/*! \brief Parse keyword arguments as PType arguments and save to parsed */
inline void RangeParamParser(nnvm::NodeAttrs* attrs) {
RangeParam param;
@@ -165,19 +191,77 @@ inline bool InitStorageType(const nnvm::NodeAttrs& attrs,
return true;
}
+/*!
+ * \brief General-purpose blob value-filling function
+ * \tparam xpu cpu or gpu
+ * \tparam ValueType Data type of supplied value
+ * \tparam is_integer Whether to optimize for an integer value
+ * \param s Stream
+ * \param b The blob to fill with a value
+ * \param req Request type (kNullOp, kWriteTo, etc)
+ * \param val The value to use for the filling operation
+ */
+template <bool is_integer = false, typename ValueType, typename xpu>
+void Fill(mshadow::Stream<xpu> *s, const TBlob& b, const OpReqType req, ValueType val) {
+ if (req != kNullOp) {
+ const size_t size = b.Size();
+ if (val == 0) {
+ if (req != kAddTo) {
+ if (b.dev_mask() == cpu::kDevMask) {
+ MSHADOW_TYPE_SWITCH(b.type_flag_, DType, {
+ memset(b.dptr_, 0, size * sizeof(DType));
+ });
+ } else {
+ // Optimize common use-case of filling with ones
+ MSHADOW_TYPE_SWITCH(b.type_flag_, DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req, Req, {
+ mxnet_op::Kernel<mxnet_op::op_with_req<mxnet_op::set_to_int<0>, Req>, xpu>::Launch(
+ s, b.Size(), b.dptr<DType>());
+ });
+ });
+ }
+ }
+ } else if (is_integer && val == 1) {
+ // Optimize common use-case of filling with ones
+ MSHADOW_TYPE_SWITCH(b.type_flag_, DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req, Req, {
+ mxnet_op::Kernel<mxnet_op::op_with_req<mxnet_op::set_to_int<1>, Req>, xpu>::Launch(
+ s, b.Size(), b.dptr<DType>());
+ });
+ });
+ } else {
+ // Generic fill kernel from variable
+ MSHADOW_TYPE_SWITCH(b.type_flag_, DType, {
+ MXNET_ASSIGN_REQ_SWITCH(req, Req, {
+ mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, Req>, xpu>::Launch(
+ s, b.Size(), b.dptr<DType>(), static_cast<DType>(val));
+ });
+ });
+ }
+ }
+}
+
+/*! \brief Fill output with a scalar integer value */
template<typename xpu, int value>
void FillCompute(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs) {
- using namespace mshadow;
- using namespace mshadow::expr;
- Stream<xpu> *s = ctx.get_stream<xpu>();
- MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
- Tensor<xpu, 1, DType> out = outputs[0].FlatTo1D<xpu, DType>(s);
- ASSIGN_DISPATCH(out, req[0], scalar<DType>(value));
- });
+ Fill<true>(ctx.get_stream<xpu>(), outputs[0], req[0], value);
+}
+
+/*! \brief Fill output with an arbitrary value */
+template<typename xpu>
+void InitFillWithScalarCompute(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const std::vector<TBlob> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<TBlob> &outputs) {
+ CHECK_EQ(inputs.size(), 0);
+ CHECK_EQ(outputs.size(), 1U);
+ const auto& param = nnvm::get<InitOpWithScalarParam>(attrs.parsed);
+ Fill<true>(ctx.get_stream<xpu>(), outputs[0], req[0], param.value);
}
struct PopulateFullIdxRspKernel {
@@ -196,15 +280,12 @@ inline void FillDnsZerosRspImpl(mshadow::Stream<xpu> *s, NDArray *dst) {
using namespace mshadow;
using namespace mxnet_op;
CHECK_EQ(dst->storage_type(), kRowSparseStorage);
- MSHADOW_REAL_TYPE_SWITCH(dst->dtype(), DType, {
- MSHADOW_IDX_TYPE_SWITCH(dst->aux_type(kIdx), IType, {
- auto num_rows = dst->shape()[0];
- dst->CheckAndAlloc({Shape1(num_rows)});
- auto idx = dst->aux_data(kIdx);
- auto val = dst->data();
- Kernel<set_zero, xpu>::Launch(s, val.Size(), val.dptr<DType>());
- Kernel<PopulateFullIdxRspKernel, xpu>::Launch(s, num_rows, idx.dptr<IType>());
- });
+ MSHADOW_IDX_TYPE_SWITCH(dst->aux_type(kIdx), IType, {
+ const index_t num_rows = dst->shape()[0];
+ dst->CheckAndAlloc({Shape1(num_rows)});
+ Fill<true>(s, dst->data(), kWriteTo, 0);
+ auto idx = dst->aux_data(kIdx).FlatTo1D<xpu, IType>(s);
+ Kernel<PopulateFullIdxRspKernel, xpu>::Launch(s, num_rows, idx.dptr_);
});
}
@@ -244,10 +325,7 @@ inline void FillZerosCsrImpl(mshadow::Stream<mshadow::cpu> *s, const NDArray& ds
dst.set_aux_shape(csr::kIdx, mshadow::Shape1(0));
dst.CheckAndAllocAuxData(csr::kIndPtr, mshadow::Shape1(dst.shape()[0] + 1));
TBlob indptr_data = dst.aux_data(csr::kIndPtr);
- MSHADOW_IDX_TYPE_SWITCH(dst.aux_type(csr::kIndPtr), IType, {
- mxnet_op::Kernel<mxnet_op::set_zero, mshadow::cpu>::Launch(
- s, indptr_data.Size(), indptr_data.dptr<IType>());
- });
+ Fill<true>(s, dst.aux_data(csr::kIndPtr), kWriteTo, 0);
}
void FillZerosCsrImpl(mshadow::Stream<mshadow::gpu> *s, const NDArray& dst);
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services