From 92201d3956a4f64615baf5bc9e979bcfc6bd09bd Mon Sep 17 00:00:00 2001
From: tensor-tang <tangjian03@baidu.com>
Date: Fri, 4 Jan 2019 06:41:40 +0000
Subject: [PATCH] support avg and sqrt pool and add  mkl impl

test=develop
---
 .../operators/jit/more/mkl/CMakeLists.txt     |  1 +
 paddle/fluid/operators/jit/more/mkl/mkl.cc    | 31 +++++++++++++++++++
 paddle/fluid/operators/jit/more/mkl/mkl.h     | 26 ++++++++++++++++
 paddle/fluid/operators/jit/refer/refer.h      |  9 ++++++
 4 files changed, 67 insertions(+)

diff --git a/paddle/fluid/operators/jit/more/mkl/CMakeLists.txt b/paddle/fluid/operators/jit/more/mkl/CMakeLists.txt
index 863cc720d6..f5ed2f0572 100644
--- a/paddle/fluid/operators/jit/more/mkl/CMakeLists.txt
+++ b/paddle/fluid/operators/jit/more/mkl/CMakeLists.txt
@@ -9,3 +9,4 @@ USE_JITKERNEL_MORE(kVScal, mkl)
 USE_JITKERNEL_MORE(kVExp, mkl)
 USE_JITKERNEL_MORE(kVSigmoid, mkl)
 USE_JITKERNEL_MORE(kVTanh, mkl)
+USE_JITKERNEL_MORE(kSeqPool, mkl)
diff --git a/paddle/fluid/operators/jit/more/mkl/mkl.cc b/paddle/fluid/operators/jit/more/mkl/mkl.cc
index a5b088d481..5a499ac2c0 100644
--- a/paddle/fluid/operators/jit/more/mkl/mkl.cc
+++ b/paddle/fluid/operators/jit/more/mkl/mkl.cc
@@ -72,6 +72,26 @@ void VExp<double>(const double* x, double* y, int n) {
   platform::dynload::vdExp(n, x, y);
 }
 
+template <>
+void VCopy<float>(const float* x, float* y, int n) {
+  platform::dynload::cblas_scopy(n, x, 1, y, 1);
+}
+
+template <>
+void VCopy<double>(const double* x, double* y, int n) {
+  platform::dynload::cblas_dcopy(n, x, 1, y, 1);
+}
+
+template <>
+void VAXPY<float>(float a, const float* x, float* y, int n) {
+  platform::dynload::cblas_saxpy(n, a, x, 1, y, 1);
+}
+
+template <>
+void VAXPY<double>(double a, const double* x, double* y, int n) {
+  platform::dynload::cblas_daxpy(n, a, x, 1, y, 1);
+}
+
 // TODO(TJ): tuning me carefully on AVX, AVX2 and AVX512
 template <>
 bool VMulKernel<float>::UseMe(const int& d) const {
@@ -103,6 +123,16 @@ bool VTanhKernel<float>::UseMe(const int& d) const {
   return d > 7;
 }
 
+template <>
+bool SeqPoolKernel<float>::UseMe(const seq_pool_attr_t& attr) const {
+  return true;
+}
+
+template <>
+bool SeqPoolKernel<double>::UseMe(const seq_pool_attr_t& attr) const {
+  return true;
+}
+
 #define AWALYS_USE_ME_WITH_DOUBLE(func)                  \
   template <>                                            \
   bool func##Kernel<double>::UseMe(const int& d) const { \
@@ -135,5 +165,6 @@ REGISTER_MKL_KERNEL(kVScal, VScal);
 REGISTER_MKL_KERNEL(kVExp, VExp);
 REGISTER_MKL_KERNEL(kVSigmoid, VSigmoid);
 REGISTER_MKL_KERNEL(kVTanh, VTanh);
+REGISTER_MKL_KERNEL(kSeqPool, SeqPool);
 
 #undef REGISTER_MKL_KERNEL
diff --git a/paddle/fluid/operators/jit/more/mkl/mkl.h b/paddle/fluid/operators/jit/more/mkl/mkl.h
index ee1031c028..0a3816db24 100644
--- a/paddle/fluid/operators/jit/more/mkl/mkl.h
+++ b/paddle/fluid/operators/jit/more/mkl/mkl.h
@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <cmath>
 #include <type_traits>
 #include "paddle/fluid/operators/jit/kernel_base.h"
 
@@ -35,6 +36,12 @@ void VScal(const T* a, const T* x, T* y, int n);
 template <typename T>
 void VExp(const T* x, T* y, int n);
 
+template <typename T>
+void VCopy(const T* x, T* y, int n);
+
+template <typename T>
+void VAXPY(T a, const T* x, T* y, int n);
+
 template <typename T>
 void VSigmoid(const T* x, T* y, int n) {
   const T min = SIGMOID_THRESHOLD_MIN;
@@ -60,6 +67,23 @@ void VTanh(const T* x, T* y, int n) {
   }
 }
 
+template <typename T>
+void SeqPool(const T* x, T* y, const seq_pool_attr_t* attr) {
+  VCopy<T>(x, y, attr->w);
+  for (int h = 1; h != attr->h; ++h) {
+    VAXPY<T>(static_cast<T>(1), x + h * attr->w, y, attr->w);
+  }
+  if (attr->type == SeqPoolType::kAvg || attr->type == SeqPoolType::kSqrt) {
+    T scalar = static_cast<T>(1);
+    if (attr->type == SeqPoolType::kAvg) {
+      scalar = scalar / static_cast<T>(attr->h);
+    } else {
+      scalar = scalar / std::sqrt(static_cast<T>(attr->h));
+    }
+    VScal<T>(&scalar, y, y, attr->w);
+  }
+}
+
 #define DECLARE_MKL_KERNEL(name, tuples)                             \
   template <typename T>                                              \
   class name##Kernel : public KernelMore<tuples<T>> {                \
@@ -81,6 +105,8 @@ DECLARE_MKL_KERNEL(VExp, XYNTuples);
 DECLARE_MKL_KERNEL(VSigmoid, XYNTuples);
 DECLARE_MKL_KERNEL(VTanh, XYNTuples);
 
+DECLARE_MKL_KERNEL(SeqPool, SeqPoolTuples);
+
 #undef DECLARE_MKL_KERNEL
 
 }  // namespace mkl
diff --git a/paddle/fluid/operators/jit/refer/refer.h b/paddle/fluid/operators/jit/refer/refer.h
index c2aa922528..4e19783c86 100644
--- a/paddle/fluid/operators/jit/refer/refer.h
+++ b/paddle/fluid/operators/jit/refer/refer.h
@@ -344,6 +344,15 @@ void SeqPool(const T* x, T* y, const seq_pool_attr_t* attr) {
       src += attr->w;
     }
   }
+  if (attr->type == SeqPoolType::kAvg || attr->type == SeqPoolType::kSqrt) {
+    T scalar = static_cast<T>(1);
+    if (attr->type == SeqPoolType::kAvg) {
+      scalar = scalar / static_cast<T>(attr->h);
+    } else {
+      scalar = scalar / std::sqrt(static_cast<T>(attr->h));
+    }
+    VScal<T>(&scalar, y, y, attr->w);
+  }
 }
 
 #define DECLARE_REFER_KERNEL(name, tuples)             \
-- 
GitLab