diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py index 8f2cdb20a633..8d7d1d454ec2 100644 --- a/python/mxnet/ndarray/numpy/_op.py +++ b/python/mxnet/ndarray/numpy/_op.py @@ -1027,7 +1027,7 @@ def add(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.add(x1, x2, out=out) - return _api_internal.add(x1, x2, out) + return _api_internal.add(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1063,7 +1063,7 @@ def subtract(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.subtract(x1, x2, out=out) - return _api_internal.subtract(x1, x2, out) + return _api_internal.subtract(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1099,7 +1099,7 @@ def multiply(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.multiply(x1, x2, out=out) - return _api_internal.multiply(x1, x2, out) + return _api_internal.multiply(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1138,7 +1138,7 @@ def divide(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.divide(x1, x2, out=out) - return _api_internal.true_divide(x1, x2, out) + return _api_internal.true_divide(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1179,7 +1179,7 @@ def true_divide(x1, x2, out=None): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.true_divide(x1, x2, out=out) - return _api_internal.true_divide(x1, x2, out) + return _api_internal.true_divide(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1218,7 +1218,7 @@ def floor_divide(x1, x2, out=None): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.floor_divide(x1, x2, out=out) - return _api_internal.floor_divide(x1, x2, out) + return _api_internal.floor_divide(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1247,7 +1247,7 @@ def mod(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.mod(x1, x2, out=out) - return _api_internal.mod(x1, x2, out) + return _api_internal.mod(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1437,7 +1437,7 @@ def matmul(a, b, out=None): ... mxnet.base.MXNetError: ... : Multiplication by scalars is not allowed. """ - return _api_internal.matmul(a, b, out) + return _api_internal.matmul(a, b, out, False) @set_module('mxnet.ndarray.numpy') @@ -1466,7 +1466,7 @@ def remainder(x1, x2, out=None): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): _np.mod(x1, x2, out=out) - return _api_internal.mod(x1, x2, out) + return _api_internal.mod(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -1496,7 +1496,7 @@ def power(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.power(x1, x2, out=out) - return _api_internal.power(x1, x2, out) + return _api_internal.power(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -6883,7 +6883,7 @@ def bitwise_and(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.bitwise_and(x1, x2, out=out) - return _api_internal.bitwise_and(x1, x2, out) + return _api_internal.bitwise_and(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -6923,7 +6923,7 @@ def bitwise_xor(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.bitwise_xor(x1, x2, out=out) - return _api_internal.bitwise_xor(x1, x2, out) + return _api_internal.bitwise_xor(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -6963,7 +6963,7 @@ def bitwise_or(x1, x2, out=None, **kwargs): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.bitwise_or(x1, x2, out=out) - return _api_internal.bitwise_or(x1, x2, out) + return _api_internal.bitwise_or(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -10089,7 +10089,7 @@ def bitwise_left_shift(x1, x2, out=None): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.left_shift(x1, x2, out=out) - return _api_internal.bitwise_left_shift(x1, x2, out) + return _api_internal.bitwise_left_shift(x1, x2, out, False) @set_module('mxnet.ndarray.numpy') @@ -10128,4 +10128,4 @@ def bitwise_right_shift(x1, x2, out=None): """ if isinstance(x1, numeric_types) and isinstance(x2, numeric_types): return _np.right_shift(x1, x2, out=out) - return _api_internal.bitwise_right_shift(x1, x2, out) + return _api_internal.bitwise_right_shift(x1, x2, out, False) diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 6f2789c3603a..783c6f961aa8 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -52,7 +52,7 @@ wrap_sort_functions from ..device import current_device from ..ndarray import numpy as _mx_nd_np -from ..ndarray.numpy import _internal as _npi +from ..ndarray.numpy import _internal as _npi, _api_internal from ..ndarray.ndarray import _storage_type from ..dlpack import ndarray_from_numpy, ndarray_to_dlpack_for_write, DLDeviceType,\ ndarray_from_dlpack @@ -1091,7 +1091,7 @@ def __iadd__(self, other): """x.__iadd__(y) <=> x += y""" if not self.writable: raise ValueError('trying to add to a readonly ndarray') - return add(self, other, out=self) + return _api_internal.add(self, other, self, True) @wrap_mxnp_np_ufunc def __radd__(self, other): @@ -1145,27 +1145,27 @@ def __rshift__(self, other): @wrap_mxnp_np_ufunc def __iand__(self, other): """x.__iand__(y) <=> x &= y""" - return bitwise_and(self, other, out=self) + return _api_internal.bitwise_and(self, other, self, True) @wrap_mxnp_np_ufunc def __ior__(self, other): r"""x.__ior__(y) <=> x \|= y""" - return bitwise_or(self, other, out=self) + return _api_internal.bitwise_or(self, other, self, True) @wrap_mxnp_np_ufunc def __ixor__(self, other): """x.__ixor__(y) <=> x ^= y""" - return bitwise_xor(self, other, out=self) + return _api_internal.bitwise_xor(self, other, self, True) @wrap_mxnp_np_ufunc def __ilshift__(self, other): """x.__ilshift__(y) <=> x <<= y""" - return bitwise_left_shift(self, other, out=self) + return _api_internal.bitwise_left_shift(self, other, self, True) @wrap_mxnp_np_ufunc def __irshift__(self, other): """x.__irshift__(y) <=> x >>= y""" - return bitwise_right_shift(self, other, out=self) + return _api_internal.bitwise_right_shift(self, other, self, True) @wrap_mxnp_np_ufunc def __rlshift__(self, other): @@ -1207,7 +1207,7 @@ def __isub__(self, other): """x.__isub__(y) <=> x -= y""" if not self.writable: raise ValueError('trying to subtract from a readonly ndarray') - return subtract(self, other, out=self) + return _api_internal.subtract(self, other, self, True) @wrap_mxnp_np_ufunc def __rsub__(self, other): @@ -1229,7 +1229,7 @@ def __ifloordiv__(self, other): """x.__ifloordiv__(y) <=> x //= y""" if not self.writable: raise ValueError('trying to divide from a readonly ndarray') - return floor_divide(self, other, out=self) + return _api_internal.floor_divide(self, other, self, True) @wrap_mxnp_np_ufunc def __rfloordiv__(self, other): @@ -1249,7 +1249,7 @@ def __imul__(self, other): r"""x.__imul__(y) <=> x \*= y""" if not self.writable: raise ValueError('trying to add to a readonly ndarray') - return multiply(self, other, out=self) + return _api_internal.multiply(self, other, self, True) @wrap_mxnp_np_ufunc def __rmul__(self, other): @@ -1269,7 +1269,7 @@ def __rdiv__(self, other): @wrap_mxnp_np_ufunc def __idiv__(self, other): """x.__idiv__(y) <=> x /= y""" - return divide(self, other, out=self) + return _api_internal.true_divide(self, other, self, True) @wrap_mxnp_np_ufunc def __truediv__(self, other): @@ -1284,7 +1284,7 @@ def __rtruediv__(self, other): @wrap_mxnp_np_ufunc def __itruediv__(self, other): """x.__itruediv__(y) <=> x /= y""" - return divide(self, other, out=self) + return _api_internal.true_divide(self, other, self, True) @wrap_mxnp_np_ufunc def __mod__(self, other): @@ -1299,7 +1299,7 @@ def __rmod__(self, other): @wrap_mxnp_np_ufunc def __imod__(self, other): """x.__imod__(y) <=> x %= y""" - return mod(self, other, out=self) + return _api_internal.mod(self, other, self, True) @wrap_mxnp_np_ufunc def __pow__(self, other): @@ -1314,7 +1314,7 @@ def __rpow__(self, other): @wrap_mxnp_np_ufunc def __ipow__(self, other): """x.__ipow__(y) <=> x **= y""" - return power(self, other, out=self) + return _api_internal.power(self, other, self, True) @wrap_mxnp_np_ufunc def __eq__(self, other): @@ -1362,7 +1362,7 @@ def __rmatmul__(self, other): @wrap_mxnp_np_ufunc def __imatmul__(self, other): """x.__imatmul__(y) <=> x @= y""" - return matmul(self, other, out=self) + return _api_internal.matmul(self, other, self, True) def __bool__(self): num_elements = self.size diff --git a/src/api/operator/ufunc_helper.cc b/src/api/operator/ufunc_helper.cc index 927253150f9a..b4d10049df34 100644 --- a/src/api/operator/ufunc_helper.cc +++ b/src/api/operator/ufunc_helper.cc @@ -25,6 +25,7 @@ #include "utils.h" #include "../../imperative/imperative_utils.h" #include "../../operator/tensor/elemwise_binary_scalar_op.h" +#include "../../operator/numpy/np_elemwise_broadcast_op.h" namespace mxnet { @@ -55,6 +56,31 @@ void UFuncHelper(NDArray* lhs, } } +void UFuncHelper(NDArray* lhs, + NDArray* rhs, + NDArray* out, + runtime::MXNetRetValue* ret, + const nnvm::Op* op, + bool in_place) { + using namespace runtime; + nnvm::NodeAttrs attrs; + op::NumpyBinaryParam param = {}; + param.in_place = in_place; + attrs.op = op; + attrs.parsed = param; + SetAttrDict(&attrs); + NDArray* inputs[] = {lhs, rhs}; + int num_inputs = 2; + NDArray** outputs = out == nullptr ? nullptr : &out; + int num_outputs = out != nullptr; + auto ndoutputs = Invoke(op, &attrs, num_inputs, inputs, &num_outputs, outputs); + if (outputs) { + *ret = PythonArg(2); + } else { + *ret = reinterpret_cast(ndoutputs[0]); + } +} + void UFuncHelper(NDArray* lhs, int64_t rhs, NDArray* out, @@ -164,7 +190,19 @@ void UFuncHelper(runtime::MXNetArgs args, NDArray* out = args[2].operator NDArray*(); if (args[0].type_code() == kNDArrayHandle) { if (args[1].type_code() == kNDArrayHandle) { - UFuncHelper(args[0].operator NDArray*(), args[1].operator NDArray*(), out, ret, fn_array); + int args_size = args.size(); + if (args_size == 4) { + bool in_place = args[3].operator bool(); + if (in_place) { + UFuncHelper( + args[0].operator NDArray*(), args[1].operator NDArray*(), out, ret, fn_array, true); + } else { + UFuncHelper( + args[0].operator NDArray*(), args[1].operator NDArray*(), out, ret, fn_array, false); + } + } else { + UFuncHelper(args[0].operator NDArray*(), args[1].operator NDArray*(), out, ret, fn_array); + } } else if (args[1].type_code() == kDLInt) { UFuncHelper(args[0].operator NDArray*(), args[1].operator int64_t(), out, ret, lfn_scalar); } else { diff --git a/src/operator/numpy/np_elemwise_broadcast_op.h b/src/operator/numpy/np_elemwise_broadcast_op.h index fa329bf248d5..8bc9851c6144 100644 --- a/src/operator/numpy/np_elemwise_broadcast_op.h +++ b/src/operator/numpy/np_elemwise_broadcast_op.h @@ -35,6 +35,20 @@ namespace mxnet { namespace op { +struct NumpyBinaryParam : public dmlc::Parameter { + bool in_place; + DMLC_DECLARE_PARAMETER(NumpyBinaryParam) { + DMLC_DECLARE_FIELD(in_place).set_default(false).describe( + "Indicate whether this binary operation is in-place"); + } + + void SetAttrDict(std::unordered_map* dict) { + std::ostringstream in_place_s; + in_place_s << in_place; + (*dict)["in_place"] = in_place_s.str(); + } +}; + inline void PrintErrorMessage(const std::string& op_name, const int dtype1, const int dtype2) { LOG(FATAL) << "Operator " << op_name << " does not support combination of " << mshadow::dtype_string(dtype1) << " with " << mshadow::dtype_string(dtype2) @@ -450,6 +464,22 @@ void NumpyBinaryBroadcastCompute(const nnvm::NodeAttrs& attrs, return; } + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + TBlob temp_tblob; + mshadow::Stream* s = ctx.get_stream(); + MSHADOW_TYPE_SWITCH_EXT(lhs.type_flag_, LType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + BinaryBroadcastCompute( + attrs, ctx, {temp_tblob.reshape(rhs.shape_), lhs}, req, outputs); + return; + } + MixedBinaryBroadcastCompute(attrs, ctx, inputs, req, outputs); } @@ -475,7 +505,22 @@ void NumpyBinaryBroadcastComputeWithBool(const nnvm::NodeAttrs& attrs, BinaryBroadcastComputeWithBool(attrs, ctx, inputs, req, outputs); return; } - if (!common::is_float(lhs.type_flag_) && !common::is_float(rhs.type_flag_)) { + + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + Stream* s = ctx.get_stream(); + TBlob temp_tblob; + MSHADOW_TYPE_SWITCH_EXT_WITH_BOOL(lhs.type_flag_, LType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + BinaryBroadcastComputeWithBool( + attrs, ctx, {temp_tblob.reshape(rhs.shape_), lhs}, req, outputs); + return; + } else if (!common::is_float(lhs.type_flag_) && !common::is_float(rhs.type_flag_)) { Stream* s = ctx.get_stream(); TBlob temp_tblob; if (lhs.type_flag_ == out.type_flag_) { @@ -545,7 +590,18 @@ void NumpyBinaryBroadcastIntComputeWithBool(const nnvm::NodeAttrs& attrs, } Stream* s = ctx.get_stream(); TBlob temp_tblob; - if (lhs.type_flag_ == out.type_flag_) { + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + MXNET_INT_TYPE_SWITCH_EXT_WITH_BOOL(lhs.type_flag_, LType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + BinaryBroadcastIntComputeWithBool( + attrs, ctx, {temp_tblob.reshape(rhs.shape_), lhs}, req, outputs); + } else if (lhs.type_flag_ == out.type_flag_) { MXNET_INT_TYPE_SWITCH_EXT_WITH_BOOL(lhs.type_flag_, LType, { Tensor temp_tensor = ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); @@ -873,6 +929,12 @@ inline bool NumpyBinaryMixedPrecisionType(const nnvm::NodeAttrs& attrs, CHECK_EQ(out_attrs->size(), 1U); const int ltype = in_attrs->at(0); const int rtype = in_attrs->at(1); + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + TYPE_ASSIGN_CHECK(*out_attrs, 0, ltype); + return true; + } if (ltype != -1 && rtype != -1 && (ltype != rtype)) { // Only when both input types are known and not the same, we enter the mixed-precision mode TYPE_ASSIGN_CHECK(*out_attrs, 0, common::type_promotion(ltype, rtype)); @@ -886,6 +948,7 @@ inline bool NumpyBinaryMixedPrecisionType(const nnvm::NodeAttrs& attrs, NNVM_REGISTER_OP(name) \ .set_num_inputs(2) \ .set_num_outputs(1) \ + .set_attr_parser(ParamParser) \ .set_attr("FListInputNames", \ [](const NodeAttrs& attrs) { \ return std::vector{"lhs", "rhs"}; \ @@ -915,6 +978,12 @@ inline bool NumpyBinaryMixedIntPrecisionTypeWithBool(const nnvm::NodeAttrs& attr << "1st input only supports integer types or bool types."; CHECK(common::is_int(rtype) || rtype == mshadow::kBool) << "2nd input only supports integer types or bool types."; + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + TYPE_ASSIGN_CHECK(*out_attrs, 0, ltype); + return true; + } if (ltype != -1 && rtype != -1 && (ltype != rtype)) { // Only when both input types are known and not the same, we enter the mixed-precision mode TYPE_ASSIGN_CHECK(*out_attrs, 0, common::type_promotion(ltype, rtype)); @@ -928,6 +997,7 @@ inline bool NumpyBinaryMixedIntPrecisionTypeWithBool(const nnvm::NodeAttrs& attr NNVM_REGISTER_OP(name) \ .set_num_inputs(2) \ .set_num_outputs(1) \ + .set_attr_parser(ParamParser) \ .set_attr("FListInputNames", \ [](const NodeAttrs& attrs) { \ return std::vector{"lhs", "rhs"}; \ @@ -955,6 +1025,12 @@ inline bool NumpyBinaryMixedIntPrecisionType(const nnvm::NodeAttrs& attrs, const int rtype = in_attrs->at(1); CHECK(common::is_int(ltype)) << "1st input only supports integer types."; CHECK(common::is_int(rtype)) << "2nd input only supports integer types."; + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + TYPE_ASSIGN_CHECK(*out_attrs, 0, ltype); + return true; + } if (ltype != -1 && rtype != -1 && (ltype != rtype)) { // Only when both input types are known and not the same, we enter the mixed-precision mode TYPE_ASSIGN_CHECK(*out_attrs, 0, common::type_promotion(ltype, rtype)); @@ -968,6 +1044,7 @@ inline bool NumpyBinaryMixedIntPrecisionType(const nnvm::NodeAttrs& attrs, NNVM_REGISTER_OP(name) \ .set_num_inputs(2) \ .set_num_outputs(1) \ + .set_attr_parser(ParamParser) \ .set_attr("FListInputNames", \ [](const NodeAttrs& attrs) { \ return std::vector{"lhs", "rhs"}; \ diff --git a/src/operator/numpy/np_elemwise_broadcast_op_add.cc b/src/operator/numpy/np_elemwise_broadcast_op_add.cc index 50a79ab5dc2f..253f35fae099 100644 --- a/src/operator/numpy/np_elemwise_broadcast_op_add.cc +++ b/src/operator/numpy/np_elemwise_broadcast_op_add.cc @@ -27,6 +27,8 @@ namespace mxnet { namespace op { +DMLC_REGISTER_PARAMETER(NumpyBinaryParam); + MXNET_OPERATOR_REGISTER_NP_BINARY_MIXED_PRECISION(_npi_add) .set_attr("FCompute", NumpyBinaryBroadcastComputeWithBool("FCompute", - NumpyBinaryBroadcastCompute) + NumpyBinaryBroadcastComputeWithBool) .set_attr("FGradient", ElemwiseGradUseIn{"_backward_npi_broadcast_sub"}); NNVM_REGISTER_OP(_backward_npi_broadcast_sub) diff --git a/src/operator/numpy/np_true_divide-inl.h b/src/operator/numpy/np_true_divide-inl.h index 047489f648cc..2f81faf551f8 100644 --- a/src/operator/numpy/np_true_divide-inl.h +++ b/src/operator/numpy/np_true_divide-inl.h @@ -90,89 +90,108 @@ void TrueDivideElemwiseCompute(const nnvm::NodeAttrs& attrs, const TBlob& lhs = inputs[0]; const TBlob& rhs = inputs[1]; const TBlob& out = outputs[0]; - if (lhs.type_flag_ == rhs.type_flag_) { - // Case when types of the 2 input tensors are the same - if (common::is_float(lhs.type_flag_)) { - // If both are the same floats, normal launch - MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { - MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, DType, { - Kernel, xpu>::Launch( - s, out.Size(), out.dptr(), lhs.dptr(), rhs.dptr()); - }); - }); - } else { - // If both are the same integers, output is float32 or float64 - CHECK_EQ(out.type_flag_, mxnet::common::GetDefaultDtype()) - << "true_divide only supports float32 and float64" - " output when input's dtype is " - << type_string(lhs.type_flag_); - MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { - MXNET_INT_TYPE_SWITCH(lhs.type_flag_, DType, { - Kernel, xpu>::Launch( - s, out.Size(), out.dptr(), lhs.dptr(), rhs.dptr()); - }); + + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + TBlob temp_tblob; + MSHADOW_TYPE_SWITCH_EXT(lhs.type_flag_, LType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { + MSHADOW_TYPE_SWITCH_EXT(out.type_flag_, DType, { + Kernel, xpu>::Launch( + s, out.Size(), out.dptr(), lhs.dptr(), temp_tblob.dptr()); }); - } + }); } else { - // Case when types of the 2 input tensors are different - if (common::is_float(lhs.type_flag_) && common::is_float(rhs.type_flag_)) { - // both lhs and rhs are float types, output type is the more precise one - TBlob temp_tblob; - if (lhs.type_flag_ == out.type_flag_) { - MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { - Tensor temp_tensor = - ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); - temp_tblob = TBlob(temp_tensor); - }); - CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + if (lhs.type_flag_ == rhs.type_flag_) { + // Case when types of the 2 input tensors are the same + if (common::is_float(lhs.type_flag_)) { + // If both are the same floats, normal launch MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { - MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { + MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, DType, { Kernel, xpu>::Launch( - s, out.Size(), out.dptr(), lhs.dptr(), temp_tblob.dptr()); + s, out.Size(), out.dptr(), lhs.dptr(), rhs.dptr()); }); }); } else { - MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { - Tensor temp_tensor = - ctx.requested[0].get_space_typed(Shape1(lhs.Size()), s); - temp_tblob = TBlob(temp_tensor); - }); - CastCompute(attrs, ctx, {lhs}, {kWriteTo}, {temp_tblob}); + // If both are the same integers, output is float32 or float64 + CHECK_EQ(out.type_flag_, mxnet::common::GetDefaultDtype()) + << "true_divide only supports float32 and float64" + " output when input's dtype is " + << type_string(lhs.type_flag_); MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { - MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { + MXNET_INT_TYPE_SWITCH(lhs.type_flag_, DType, { Kernel, xpu>::Launch( - s, out.Size(), out.dptr(), temp_tblob.dptr(), rhs.dptr()); + s, out.Size(), out.dptr(), lhs.dptr(), rhs.dptr()); }); }); } - } else if (common::is_float(lhs.type_flag_) || common::is_float(rhs.type_flag_)) { - // one is float type, the other is integer type, the output type should be the same as float - CHECK_EQ(out.type_flag_, common::is_float(lhs.type_flag_) ? lhs.type_flag_ : rhs.type_flag_) - << "This case out type should be same as the float type"; - if (common::is_float(lhs.type_flag_)) { - // lhs is the float one - MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { + } else { + // Case when types of the 2 input tensors are different + if (common::is_float(lhs.type_flag_) && common::is_float(rhs.type_flag_)) { + // both lhs and rhs are float types, output type is the more precise one + TBlob temp_tblob; + if (lhs.type_flag_ == out.type_flag_) { MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { - MXNET_INT_TYPE_SWITCH(rhs.type_flag_, RType, { - Kernel, xpu>::Launch( - s, out.Size(), out.dptr(), rhs.dptr(), lhs.dptr()); + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { + MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { + Kernel, xpu>::Launch( + s, out.Size(), out.dptr(), lhs.dptr(), temp_tblob.dptr()); }); }); - }); - } else { - // rhs is the float one - MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { - MXNET_INT_TYPE_SWITCH(lhs.type_flag_, LType, { - MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { + } else { + MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(lhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {lhs}, {kWriteTo}, {temp_tblob}); + MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { + MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { Kernel, xpu>::Launch( - s, out.Size(), out.dptr(), lhs.dptr(), rhs.dptr()); + s, out.Size(), out.dptr(), temp_tblob.dptr(), rhs.dptr()); }); }); - }); + } + } else if (common::is_float(lhs.type_flag_) || common::is_float(rhs.type_flag_)) { + // one is float type, the other is integer type, the output type should be the same as float + CHECK_EQ(out.type_flag_, common::is_float(lhs.type_flag_) ? lhs.type_flag_ : rhs.type_flag_) + << "This case out type should be same as the float type"; + if (common::is_float(lhs.type_flag_)) { + // lhs is the float one + MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { + MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { + MXNET_INT_TYPE_SWITCH(rhs.type_flag_, RType, { + Kernel, xpu>::Launch( + s, out.Size(), out.dptr(), rhs.dptr(), lhs.dptr()); + }); + }); + }); + } else { + // rhs is the float one + MXNET_ASSIGN_REQ_SWITCH(req[0], Req, { + MXNET_INT_TYPE_SWITCH(lhs.type_flag_, LType, { + MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { + Kernel, xpu>::Launch( + s, out.Size(), out.dptr(), lhs.dptr(), rhs.dptr()); + }); + }); + }); + } + } else { + // lhs is integer type, rhs is integer type, output type should be float + LOG(FATAL) << "not implemented yet..."; } - } else { - // lhs is integer type, rhs is integer type, output type should be float - LOG(FATAL) << "not implemented yet..."; } } } @@ -190,6 +209,8 @@ void TrueDivideBroadcastCompute(const nnvm::NodeAttrs& attrs, mxnet::TShape new_lshape, new_rshape, new_oshape; int ndim = BinaryBroadcastShapeCompact( inputs[0].shape_, inputs[1].shape_, outputs[0].shape_, &new_lshape, &new_rshape, &new_oshape); + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; if (!ndim) { TrueDivideElemwiseCompute(attrs, ctx, inputs, req, outputs); } else { @@ -203,52 +224,32 @@ void TrueDivideBroadcastCompute(const nnvm::NodeAttrs& attrs, mshadow::Shape oshape = new_oshape.get(); mshadow::Shape lstride = calc_stride(new_lshape.get()); mshadow::Shape rstride = calc_stride(new_rshape.get()); - if (lhs.type_flag_ == rhs.type_flag_) { - // When the both inputs have the same data types - if (common::is_float(lhs.type_flag_)) { - // If both inputs are the same float types, output is the same float type - MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, DType, { - Kernel, xpu>::template LaunchEx( - s, - new_oshape.Size(), - req[0], - lstride, - rstride, - oshape, - lhs.dptr(), - rhs.dptr(), - out.dptr()); - }); - } else { - CHECK_EQ(out.type_flag_, mxnet::common::GetDefaultDtype()) - << "true_divide only supports float32 and float64 output when input's dtype is " - << type_string(lhs.type_flag_); - MXNET_INT_TYPE_SWITCH(lhs.type_flag_, DType, { - // If both inputs are the same integer types, output is float type - Kernel, xpu>::template LaunchEx( - s, - new_oshape.Size(), - req[0], - lstride, - rstride, - oshape, - lhs.dptr(), - rhs.dptr(), - out.dptr()); - }); - } + if (is_inplace) { + TBlob temp_tblob; + MSHADOW_TYPE_SWITCH_EXT(lhs.type_flag_, LType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + MSHADOW_TYPE_SWITCH_EXT(out.type_flag_, DType, { + Kernel, xpu>::template LaunchEx( + s, + new_oshape.Size(), + req[0], + lstride, + rstride, + oshape, + lhs.dptr(), + temp_tblob.dptr(), + out.dptr()); + }); } else { - if (common::is_float(lhs.type_flag_) && common::is_float(rhs.type_flag_)) { - // lhs and rhs have different float types, the output is the more precise one - TBlob temp_tblob; - if (lhs.type_flag_ == out.type_flag_) { - MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { - Tensor temp_tensor = - ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); - temp_tblob = TBlob(temp_tensor); - }); - CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); - MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { + if (lhs.type_flag_ == rhs.type_flag_) { + // When the both inputs have the same data types + if (common::is_float(lhs.type_flag_)) { + // If both inputs are the same float types, output is the same float type + MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, DType, { Kernel, xpu>::template LaunchEx(s, new_oshape.Size(), @@ -257,17 +258,15 @@ void TrueDivideBroadcastCompute(const nnvm::NodeAttrs& attrs, rstride, oshape, lhs.dptr(), - temp_tblob.dptr(), + rhs.dptr(), out.dptr()); }); } else { - MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { - Tensor temp_tensor = - ctx.requested[0].get_space_typed(Shape1(lhs.Size()), s); - temp_tblob = TBlob(temp_tensor); - }); - CastCompute(attrs, ctx, {lhs}, {kWriteTo}, {temp_tblob}); - MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { + CHECK_EQ(out.type_flag_, mxnet::common::GetDefaultDtype()) + << "true_divide only supports float32 and float64 output when input's dtype is " + << type_string(lhs.type_flag_); + MXNET_INT_TYPE_SWITCH(lhs.type_flag_, DType, { + // If both inputs are the same integer types, output is float type Kernel, xpu>::template LaunchEx(s, new_oshape.Size(), @@ -275,37 +274,42 @@ void TrueDivideBroadcastCompute(const nnvm::NodeAttrs& attrs, lstride, rstride, oshape, - temp_tblob.dptr(), + lhs.dptr(), rhs.dptr(), - out.dptr()); + out.dptr()); }); } - } else if (common::is_float(lhs.type_flag_) || common::is_float(rhs.type_flag_)) { - // one of lhs and rhs is float, the output is the same type as the float one - if (common::is_float(lhs.type_flag_)) { - // lhs is float type, output will be the same float type - CHECK_EQ(lhs.type_flag_, out.type_flag_) - << "lhs should have the same type as out, infer type broken?"; - MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { - MXNET_INT_TYPE_SWITCH(rhs.type_flag_, RType, { - Kernel, + } else { + if (common::is_float(lhs.type_flag_) && common::is_float(rhs.type_flag_)) { + // lhs and rhs have different float types, the output is the more precise one + TBlob temp_tblob; + if (lhs.type_flag_ == out.type_flag_) { + MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(rhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {rhs}, {kWriteTo}, {temp_tblob}); + MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { + Kernel, xpu>::template LaunchEx(s, new_oshape.Size(), req[0], - rstride, lstride, + rstride, oshape, - rhs.dptr(), - lhs.dptr(), - out.dptr()); + lhs.dptr(), + temp_tblob.dptr(), + out.dptr()); }); - }); - } else { - // rhs is float type, output will be the same float type - CHECK_EQ(rhs.type_flag_, out.type_flag_) - << "rhs should have the same type as out, infer type broken?"; - MXNET_INT_TYPE_SWITCH(lhs.type_flag_, LType, { + } else { MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { + Tensor temp_tensor = + ctx.requested[0].get_space_typed(Shape1(lhs.Size()), s); + temp_tblob = TBlob(temp_tensor); + }); + CastCompute(attrs, ctx, {lhs}, {kWriteTo}, {temp_tblob}); + MSHADOW_REAL_TYPE_SWITCH(out.type_flag_, DType, { Kernel, xpu>::template LaunchEx(s, new_oshape.Size(), @@ -313,15 +317,54 @@ void TrueDivideBroadcastCompute(const nnvm::NodeAttrs& attrs, lstride, rstride, oshape, - lhs.dptr(), - rhs.dptr(), - out.dptr()); + temp_tblob.dptr(), + rhs.dptr(), + out.dptr()); }); - }); + } + } else if (common::is_float(lhs.type_flag_) || common::is_float(rhs.type_flag_)) { + // one of lhs and rhs is float, the output is the same type as the float one + if (common::is_float(lhs.type_flag_)) { + // lhs is float type, output will be the same float type + CHECK_EQ(lhs.type_flag_, out.type_flag_) + << "lhs should have the same type as out, infer type broken?"; + MSHADOW_REAL_TYPE_SWITCH(lhs.type_flag_, LType, { + MXNET_INT_TYPE_SWITCH(rhs.type_flag_, RType, { + Kernel, + xpu>::template LaunchEx(s, + new_oshape.Size(), + req[0], + rstride, + lstride, + oshape, + rhs.dptr(), + lhs.dptr(), + out.dptr()); + }); + }); + } else { + // rhs is float type, output will be the same float type + CHECK_EQ(rhs.type_flag_, out.type_flag_) + << "rhs should have the same type as out, infer type broken?"; + MXNET_INT_TYPE_SWITCH(lhs.type_flag_, LType, { + MSHADOW_REAL_TYPE_SWITCH(rhs.type_flag_, RType, { + Kernel, + xpu>::template LaunchEx(s, + new_oshape.Size(), + req[0], + lstride, + rstride, + oshape, + lhs.dptr(), + rhs.dptr(), + out.dptr()); + }); + }); + } + } else { + // lhs and rhs have different integer types, the output is float type + LOG(FATAL) << "not implemented yet..."; } - } else { - // lhs and rhs have different integer types, the output is float type - LOG(FATAL) << "not implemented yet..."; } } }); diff --git a/src/operator/numpy/np_true_divide.cc b/src/operator/numpy/np_true_divide.cc index 639379d36cd0..33570e280a61 100644 --- a/src/operator/numpy/np_true_divide.cc +++ b/src/operator/numpy/np_true_divide.cc @@ -51,8 +51,17 @@ bool TrueDivideType(const nnvm::NodeAttrs& attrs, if (dtype == -1) return false; } - const int lhs_dtype = in_attrs->at(0); + + if (num_inputs == 2) { + const NumpyBinaryParam& param = nnvm::get(attrs.parsed); + bool is_inplace = param.in_place; + if (is_inplace) { + TYPE_ASSIGN_CHECK(*out_attrs, 0, lhs_dtype); + return true; + } + } + const int rhs_dtype = (num_inputs == 2) ? in_attrs->at(1) : @@ -64,6 +73,7 @@ bool TrueDivideType(const nnvm::NodeAttrs& attrs, NNVM_REGISTER_OP(_npi_true_divide) .set_num_inputs(2) .set_num_outputs(1) + .set_attr_parser(ParamParser) .set_attr("FListInputNames", [](const NodeAttrs& attrs) { return std::vector{"lhs", "rhs"}; diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py index 99aacbf0e6fb..2356b6033211 100644 --- a/tests/python/unittest/test_numpy_op.py +++ b/tests/python/unittest/test_numpy_op.py @@ -11759,3 +11759,29 @@ def array_values(low, high, shape): assert_almost_equal(mx_out.asnumpy(), np_out.astype(mx_out.dtype), rtol=rtol, atol=atol, use_broadcast=False, equal_nan=True) + +@use_np +@pytest.mark.parametrize('a_dtype', np.numeric_dtypes) +@pytest.mark.parametrize('b_dtype', np.numeric_dtypes) +@pytest.mark.parametrize('shape', [ + (), + (2, 0, 2, 2), + (5, 5) +]) +@pytest.mark.parametrize('op', [ + '__iadd__', '__iand__', '__ior__', '__ixor__', '__isub__', '__imul__', '__imod__', + '__itruediv__', '__idiv__', '__ifloordiv__', '__ipow__', '__ilshift__', '__irshift__']) +def test_in_place_dtype(a_dtype, b_dtype, shape, op): + try: + promote_type = np.result_type(a_dtype, b_dtype) + except Exception as e: + # Unkown type promotion between two types + return + if op in ('__ilshift__', '__irshift__', '__iand__', '__ior__', '__ixor__') and \ + (a_dtype not in np.integer_dtypes or b_dtype not in np.integer_dtypes): + return + a = np.array(onp.random.uniform(1, 5, shape), dtype=a_dtype) + b = np.array(onp.random.uniform(1, 5, shape), dtype=b_dtype) + getattr(a, op)(b) + assert a.dtype == a_dtype + npx.waitall()