From 8a84a3e1e6921caea7ab73a6e4a7e8872b9d05ac Mon Sep 17 00:00:00 2001
From: XiaotaoChen <xiaotao.chen@intel.com>
Date: Thu, 1 Nov 2018 09:09:10 +0800
Subject: [PATCH] support mkl log when dtype is fp32 or fp64

---
 src/operator/tensor/elemwise_unary_op.h       | 37 +++++++++++++++++++
 .../tensor/elemwise_unary_op_basic.cc         |  3 +-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/operator/tensor/elemwise_unary_op.h b/src/operator/tensor/elemwise_unary_op.h
index 83b86bf1d94c..961d776c7a29 100644
--- a/src/operator/tensor/elemwise_unary_op.h
+++ b/src/operator/tensor/elemwise_unary_op.h
@@ -34,6 +34,9 @@
 #include "../mxnet_op.h"
 #include "../elemwise_op_common.h"
 #include "../../ndarray/ndarray_function.h"
+#if MSHADOW_USE_MKL == 1
+#include "mkl.h"
+#endif
 
 namespace mxnet {
 namespace op {
@@ -348,6 +351,40 @@ class UnaryOp : public OpBase {
       LogUnimplementedOp(attrs, ctx, inputs, req, outputs);
     }
   }
+
+#if MSHADOW_USE_MKL == 1
+#define MKLLOG(fname, DType) \
+static void MKLLog(size_t size, const DType* pIn, DType* pOut) { \
+  fname(size, pIn, pOut); \
+}
+
+MKLLOG(vsLn, float)
+MKLLOG(vdLn, double)
+#endif
+
+  template<typename xpu, typename OP>
+  static void LogCompute(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;
+    // if defined MSHADOW_USE_MKL then call mkl log when req is KWriteTo and type_flag
+    // is mshadow::kFloat32 or mshadow::kFloat64
+#if MSHADOW_USE_MKL == 1
+    auto type_flag = inputs[0].type_flag_;
+    if (req[0] == kWriteTo && (type_flag == mshadow::kFloat32
+          || type_flag == mshadow::kFloat64)) {
+      MSHADOW_SGL_DBL_TYPE_SWITCH(type_flag, DType, {
+        MKLLog(inputs[0].Size(), inputs[0].dptr<DType>(), outputs[0].dptr<DType>());
+      })
+    } 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 */
diff --git a/src/operator/tensor/elemwise_unary_op_basic.cc b/src/operator/tensor/elemwise_unary_op_basic.cc
index 301fc48d2128..9730d0096e58 100644
--- a/src/operator/tensor/elemwise_unary_op_basic.cc
+++ b/src/operator/tensor/elemwise_unary_op_basic.cc
@@ -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)
 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