Skip to content

Commit

Permalink
Merge pull request #6 from juliusshufan/erf
Browse files Browse the repository at this point in the history 
Fix review comments
  • Loading branch information
@TaoLv
TaoLv committed Apr 25, 2019
2 parents 672be6a + 06c51e9 commit acd7b56
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 105 deletions.
101 changes: 42 additions & 59 deletions src/operator/tensor/elemwise_unary_op.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,27 @@ class UnaryOp : public OpBase {
}

#if MSHADOW_USE_MKL == 1
template<typename OP, typename MKL_OP>
static void MKL_Compute(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs) {
if (req[0] == kNullOp) return;
auto type_flag = inputs[0].type_flag_;
size_t input_size = inputs[0].Size();
if ((req[0] == kWriteTo || req[0] == kWriteInplace) &&
mkl_func::check_size(input_size) &&
mkl_func::check_type(type_flag)) {
// set DType as float or double according to type_flag
MSHADOW_SGL_DBL_TYPE_SWITCH(type_flag, DType, {
MKL_OP::Vectorize(input_size, inputs[0].dptr<DType>(), outputs[0].dptr<DType>());
});
} else {
Compute<cpu, OP>(attrs, ctx, inputs, req, outputs);
}
}

template<typename OP, typename MKL_OP>
static void MKL_ComputeEx(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
Expand Down Expand Up @@ -375,28 +396,6 @@ class UnaryOp : public OpBase {
}
}

#if MSHADOW_USE_MKL == 1
template<typename OP, typename MKL_OP>
static void MKL_Compute(const nnvm::NodeAttrs& attrs,
const OpContext& ctx,
const std::vector<TBlob>& inputs,
const std::vector<OpReqType>& req,
const std::vector<TBlob>& outputs) {
if (req[0] == kNullOp) return;
auto type_flag = inputs[0].type_flag_;
size_t input_size = inputs[0].Size();
if ((req[0] == kWriteTo || req[0] == kWriteInplace) &&
mkl_func::check_size(input_size) &&
mkl_func::check_type(type_flag)) {
// set DType as float or double according to type_flag
MSHADOW_SGL_DBL_TYPE_SWITCH(type_flag, DType, {
MKL_OP::Vectorize(input_size, inputs[0].dptr<DType>(), outputs[0].dptr<DType>());
});
} else {
Compute<cpu, OP>(attrs, ctx, inputs, req, outputs);
}
}
#endif // MSHADOW_USE_MKL == 1
};

/*! \brief Map legacy unary_bwd to backward_grad */
Expand Down Expand Up @@ -570,48 +569,32 @@ struct ReshapeLikeParam : public dmlc::Parameter<ReshapeLikeParam> {
}) \
.add_argument("data", "NDArray-or-Symbol", "The input array.")

/*! \bried MKL Unary compute.
* With this macro means mxnet compile with MKL to accelerate math function with mkl.
* Will Register FCompute with UnaryOp::MKL_Compute() to compelet the math function.
*/
#define MXNET_MKL_OPERATOR_REGISTER_UNARY(__name$) \
NNVM_REGISTER_OP(__name$) \
.set_num_inputs(1) \
.set_num_outputs(1) \
.set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<1, 1>) \
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) \
.set_attr<nnvm::FInplaceOption>("FInplaceOption", \
[](const NodeAttrs& attrs){ \
return std::vector<std::pair<int, int> >{{0, 0}}; \
}) \
.add_argument("data", "NDArray-or-Symbol", "The input array.")

#if MSHADOW_USE_MKL == 1
/*! \bried MKL Unary compute.
* * With this macro means mxnet compile with MKL to accelerate math function with mkl.
* * Will Register FCompute with UnaryOp::MKL_Compute() to compelet the math function.
*/
#define MXNET_MKL_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(__name$, __xpu$, __kernel$, __mkl_kernel$) \
MXNET_MKL_OPERATOR_REGISTER_UNARY(__name$) \
MXNET_ADD_SPARSE_OP_ALIAS(__name$) \
.set_attr<FInferStorageType>("FInferStorageType", ElemwiseStorageType<1, 1, false, true, true>) \
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::MKL_Compute<__kernel$, __mkl_kernel$>) \
.set_attr<FComputeEx>("FComputeEx<" #__xpu$ ">", UnaryOp::MKL_ComputeEx<__kernel$, __mkl_kernel$>)

/*! \bried MKL Unary compute.
* * With this macro means mxnet compile with MKL to accelerate math function with mkl.
* * Will Register FCompute with UnaryOp::MKL_Compute() to compelet the math function.
*/
#define MXNET_MKL_OPERATOR_REGISTER_UNARY_WITH_RSP(__name$, __xpu$, __kernel$, __mkl_kernel$) \
MXNET_MKL_OPERATOR_REGISTER_UNARY(__name$) \
MXNET_ADD_SPARSE_OP_ALIAS(__name$) \
.set_attr<FInferStorageType>("FInferStorageType", ElemwiseStorageType<1, 1, false, true, false>)\
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::MKL_Compute<__kernel$, __mkl_kernel$>) \
.set_attr<FComputeEx>("FComputeEx<" #__xpu$ ">", UnaryOp::MKL_ComputeEx<__kernel$, __mkl_kerbel$>)

#define MXNET_MKL_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(__name$, __xpu$, __kernel$, __mkl_kernel$)\
MXNET_MKL_OPERATOR_REGISTER_UNARY(__name$) \
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::MKL_Compute<__kernel$, __mkl_kernel$>)
#define MXNET_MKL_OPERATOR_REGISTER_UNARY_WITH_RSP_CSR(__name$, __xpu$, __kernel$, __mkl_kernel$) \
MXNET_OPERATOR_REGISTER_UNARY(__name$) \
MXNET_ADD_SPARSE_OP_ALIAS(__name$) \
.set_attr<FInferStorageType>("FInferStorageType", ElemwiseStorageType<1, 1, false, true, true>) \
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::MKL_Compute<__kernel$, __mkl_kernel$>) \
.set_attr<FComputeEx>("FComputeEx<" #__xpu$ ">", UnaryOp::MKL_ComputeEx<__kernel$, __mkl_kernel$>)

/*! \bried MKL Unary compute.
* * With this macro means mxnet compile with MKL to accelerate math function with mkl.
* * Will Register FCompute with UnaryOp::MKL_Compute() to compelet the math function.
*/
#define MXNET_MKL_OPERATOR_REGISTER_UNARY_WITH_RSP(__name$, __xpu$, __kernel$, __mkl_kernel$) \
MXNET_OPERATOR_REGISTER_UNARY(__name$) \
MXNET_ADD_SPARSE_OP_ALIAS(__name$) \
.set_attr<FInferStorageType>("FInferStorageType", ElemwiseStorageType<1, 1, false, true, false>) \
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::MKL_Compute<__kernel$, __mkl_kernel$>) \
.set_attr<FComputeEx>("FComputeEx<" #__xpu$ ">", UnaryOp::MKL_ComputeEx<__kernel$, __mkl_kerbel$>)

#define MXNET_MKL_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(__name$, __xpu$, __kernel$, __mkl_kernel$) \
MXNET_OPERATOR_REGISTER_UNARY(__name$) \
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::MKL_Compute<__kernel$, __mkl_kernel$>)
#endif

/*! \brief Unary compute, with FComputeEx for csr and rsp available */
Expand Down
35 changes: 9 additions & 26 deletions src/operator/tensor/elemwise_unary_op_basic.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -929,30 +929,23 @@ The storage type of ``cbrt`` output depends upon the input storage type:

MXNET_OPERATOR_REGISTER_BINARY_WITH_SPARSE_CPU_DR(_backward_cbrt,
unary_bwd<mshadow_op::cube_root_grad>);

// erf
#if MSHADOW_USE_MKL == 1
MXNET_MKL_OPERATOR_REGISTER_UNARY(erf)
MXNET_OPERATOR_REGISTER_UNARY(erf)
.describe(R"code(Returns element-wise gauss error function of the input.
Example::
erf([0, -1., 10.]) = [0., -0.8427, 1.]
)code" ADD_FILELINE)
#if MSHADOW_USE_MKL == 1
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::MKL_Compute<mshadow_op::erf, mkl_func::erf>)
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_erf"});
#else
MXNET_OPERATOR_REGISTER_UNARY(erf)
.describe(R"code(Returns element-wise gauss error function of the input.
Example::
erf([0, -1., 10.]) = [0., -0.8427, 1.]
)code" ADD_FILELINE)
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::Compute<cpu, mshadow_op::erf>)
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_erf"});
#endif // MSHADOW_USE_MKL == 1
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_erf"});


MXNET_OPERATOR_REGISTER_BINARY(_backward_erf)
.set_attr<FCompute>("FCompute<cpu>",
Expand Down Expand Up @@ -1030,18 +1023,6 @@ The storage type of ``exp`` output is always dense
#endif

// log
#if MSHADOW_USE_MKL == 1
MXNET_MKL_OPERATOR_REGISTER_UNARY(log)
.describe(R"code(Returns element-wise Natural logarithmic value of the input.
The natural logarithm is logarithm in base *e*, so that ``log(exp(x)) = x``
The storage type of ``log`` output is always dense
)code" ADD_FILELINE)
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::MKL_Compute<mshadow_op::log, mkl_func::log>)
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_log"});
#else
MXNET_OPERATOR_REGISTER_UNARY(log)
MXNET_ADD_SPARSE_OP_ALIAS(log)
.describe(R"code(Returns element-wise Natural logarithmic value of the input.
Expand All @@ -1051,10 +1032,12 @@ The natural logarithm is logarithm in base *e*, so that ``log(exp(x)) = x``
The storage type of ``log`` output is always dense
)code" ADD_FILELINE)
#if MSHADOW_USE_MKL == 1
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::MKL_Compute<mshadow_op::log, mkl_func::log>)
#else
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::Compute<cpu, mshadow_op::log>)
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_log"});
#endif // MSHADOW_USE_MKL == 1

.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_log"});

// log10
MXNET_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(log10, cpu, mshadow_op::log10)
Expand Down
25 changes: 5 additions & 20 deletions tests/python/gpu/test_operator_gpu.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2200,42 +2200,30 @@ def test_context_num_gpus():
assert mx.context.num_gpus() > 0

def math_log(shape, dtype, check_value):
np_x = np.random.rand(shape[0], shape[1])
np_x = np.random.rand(*tuple(shape))
x = mx.nd.array(np_x, dtype=dtype)
mx.nd.waitall()
y = mx.nd.log(data=x)
y.wait_to_read()
if check_value:
x_ = x.as_in_context(mx.cpu())
mx.nd.waitall()
y_ = mx.nd.log(data=x_)
y_.wait_to_read()
assert_almost_equal(y.asnumpy(), y_.asnumpy())

def math_erf(shape, dtype, check_value):
np_x = np.random.rand(shape[0], shape[1])
np_x = np.random.rand(*tuple(shape))
x = mx.nd.array(np_x, dtype=dtype)
mx.nd.waitall()
y = mx.nd.erf(data=x)
y.wait_to_read()
if check_value:
x_ = x.as_in_context(mx.cpu())
mx.nd.waitall()
y_ = mx.nd.erf(data=x_)
y_.wait_to_read()
assert_almost_equal(y.asnumpy(), y_.asnumpy())

def math_square(shape, dtype, check_value):
np_x = np.random.rand(shape[0], shape[1])
np_x = np.random.rand(*tuple(shape))
x = mx.nd.array(np_x, dtype=dtype)
mx.nd.waitall()
y = mx.nd.square(data=x)
y.wait_to_read()
if check_value:
x_ = x.as_in_context(mx.cpu())
mx.nd.waitall()
y_ = mx.nd.square(data=x_)
y_.wait_to_read()
assert_almost_equal(y.asnumpy(), y_.asnumpy())

def run_math(op, shape, dtype="float32", check_value=True):
Expand All @@ -2252,12 +2240,9 @@ def run_math(op, shape, dtype="float32", check_value=True):
def test_math():
ops = ['log', 'erf', 'square']
check_value= True
lshape = 1000
rshapes = [1, 10, 100, 1000, 10000]
shape_lst = [[1000], [100,1000], [10,100,100], [10,100,100,100]]
dtypes = ["float32", "float64"]
for rshape in rshapes:
shape = (lshape, rshape)
print("shape:(%d, %d), " % (lshape, rshape), end="")
for shape in shape_lst:
for dtype in dtypes:
for op in ops:
run_math(op, shape, dtype, check_value=check_value)
Expand Down

0 comments on commit acd7b56

Please sign in to comment.