-
Notifications
You must be signed in to change notification settings - Fork 6.8k
support mkl log when dtype is fp32 or fp64 #13150
Changes from 3 commits
96f7ab3
49b89e4
2ba108f
e70998e
9206a72
faf4fef
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -34,6 +34,10 @@ | |
#include "../mxnet_op.h" | ||
#include "../elemwise_op_common.h" | ||
#include "../../ndarray/ndarray_function.h" | ||
#if MSHADOW_USE_MKL == 1 | ||
#include "mkl.h" | ||
#endif | ||
#include<climits> | ||
|
||
namespace mxnet { | ||
namespace op { | ||
|
@@ -348,6 +352,42 @@ class UnaryOp : public OpBase { | |
LogUnimplementedOp(attrs, ctx, inputs, req, outputs); | ||
} | ||
} | ||
|
||
#if MSHADOW_USE_MKL == 1 | ||
static inline void MKLLog(MKL_INT size, const float* pIn, float* pOut) { | ||
vsLn(size, pIn, pOut); | ||
} | ||
|
||
static inline void MKLLog(MKL_INT size, const double* pIn, double* pOut) { | ||
vdLn(size, pIn, pOut); | ||
} | ||
#endif | ||
|
||
template<typename xpu, typename OP> | ||
static void LogCompute(const nnvm::NodeAttrs& attrs, | ||
TaoLv marked this conversation as resolved.
Show resolved
Hide resolved
|
||
const OpContext& ctx, | ||
const std::vector<TBlob>& inputs, | ||
const std::vector<OpReqType>& req, | ||
const std::vector<TBlob>& outputs) { | ||
if (req[0] == kNullOp) return; | ||
// if defined MSHADOW_USE_MKL then call mkl log when req is KWriteTo, type_flag is | ||
// mshadow::kFloat32 or mshadow::kFloat64 and data size less than or equal MKL_INT_MAX | ||
#if MSHADOW_USE_MKL == 1 | ||
auto type_flag = inputs[0].type_flag_; | ||
const size_t MKL_INT_MAX = (sizeof(MKL_INT) == sizeof(int)) ? INT_MAX : LLONG_MAX; | ||
size_t input_size = inputs[0].Size(); | ||
if (req[0] == kWriteTo && (type_flag == mshadow::kFloat32 | ||
|| type_flag == mshadow::kFloat64) && input_size <= MKL_INT_MAX) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. fix indent. |
||
MSHADOW_SGL_DBL_TYPE_SWITCH(type_flag, DType, { | ||
MKLLog(input_size, inputs[0].dptr<DType>(), outputs[0].dptr<DType>()); | ||
}) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. need |
||
} else { | ||
Compute<xpu, OP>(attrs, ctx, inputs, req, outputs); | ||
} | ||
#else | ||
Compute<xpu, OP>(attrs, ctx, inputs, req, outputs); | ||
#endif | ||
} | ||
}; | ||
|
||
/*! \brief Map legacy unary_bwd to backward_grad */ | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -940,7 +940,7 @@ The storage type of ``exp`` output is always dense | |
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseOut{"_mul"}); | ||
|
||
// log | ||
MXNET_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(log, cpu, mshadow_op::log) | ||
MXNET_OPERATOR_REGISTER_UNARY(log) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Sorry, I do not understand this change. Can you please explain? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. the defination of #define MXNET_OPERATOR_REGISTER_UNARY_WITH_SPARSE_DR(__name$, __xpu$, __kernel$) \
MXNET_OPERATOR_REGISTER_UNARY(__name$) \
.set_attr<FCompute>("FCompute<" #__xpu$ ">", UnaryOp::Compute<__xpu$, __kernel$>) it will set FCompute attribute to MXNET_OPERATOR_REGISTER_UNARY(log)
.set_attr<FCompute>("FCompute<cpu>", UnaryOp::LogCompute<cpu, mshadow_op::log>) |
||
MXNET_ADD_SPARSE_OP_ALIAS(log) | ||
.describe(R"code(Returns element-wise Natural logarithmic value of the input. | ||
|
||
|
@@ -949,6 +949,7 @@ 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::LogCompute<cpu, mshadow_op::log>) | ||
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_log"}); | ||
|
||
// log10 | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
need space before
<
and put this include before#include "./cast_storage-inl.h"
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
all is done