optimize fusion gru kernel at size 8

paddle/fluid/operators/math/jit_kernel_rnn.cc

@@ -136,6 +136,21 @@ static std::shared_ptr<const VActKernel<T>> GetActKernel(
   return nullptr;
 }
 
+template <jit::cpu_isa_t isa>
+static std::unique_ptr<AVXAct> GetAVXAct(const std::string& type) {
+  if (type == "sigmoid") {
+    return std::unique_ptr<AVXAct>(new AVXActImpl<kSigmoid, isa>());
+  } else if (type == "relu") {
+    return std::unique_ptr<AVXAct>(new AVXActImpl<kRelu, isa>());
+  } else if (type == "tanh") {
+    return std::unique_ptr<AVXAct>(new AVXActImpl<kTanh, isa>());
+  } else if (type == "identity" || type == "") {
+    return std::unique_ptr<AVXAct>(new AVXActImpl<kIdentity, isa>());
+  }
+  PADDLE_THROW("Not support type: %s", type);
+  return nullptr;
+}
+
 /* LSTM JitKernel */
 template <typename T, jit::cpu_isa_t isa, jit_block>
 class LSTMKernelImpl : public LSTMKernel<T> {
@@ -198,21 +213,9 @@ class LSTMKernelImpl : public LSTMKernel<T> {
       const std::string& act_gate, const std::string& act_cand,              \
       const std::string& act_cell, int d)                                    \
       : LSTMKernel<float>() {                                                \
-    auto GetAVXAct = [&](const std::string& type) -> std::unique_ptr<AVXAct> { \
-      if (type == "sigmoid") {                                                 \
-        return std::unique_ptr<AVXAct>(new AVXActImpl<kSigmoid, isa>());       \
-      } else if (type == "relu") {                                             \
-        return std::unique_ptr<AVXAct>(new AVXActImpl<kRelu, isa>());          \
-      } else if (type == "tanh") {                                             \
-        return std::unique_ptr<AVXAct>(new AVXActImpl<kTanh, isa>());          \
-      } else if (type == "identity" || type == "") {                           \
-        return std::unique_ptr<AVXAct>(new AVXActImpl<kIdentity, isa>());      \
-      }                                                                        \
-      PADDLE_THROW("Not support type: %s", type);                              \
-    };                                                                         \
-    avx_act_gate_ = GetAVXAct(act_gate);                                       \
-    avx_act_cand_ = GetAVXAct(act_cand);                                       \
-    avx_act_cell_ = GetAVXAct(act_cell);                                       \
+    avx_act_gate_ = GetAVXAct<isa>(act_gate);                                \
+    avx_act_cand_ = GetAVXAct<isa>(act_cand);                                \
+    avx_act_cell_ = GetAVXAct<isa>(act_cell);                                \
   }                                                                          \
   template <>                                                                \
   void LSTMKernelImpl<float, isa, kEQ8>::ComputeCtHt(                        \
@@ -375,6 +378,7 @@ class GRUKernelImpl : public GRUKernel<T> {
     act_state_d_->Compute(gates + d2_, gates + d2_);
     vmul_d_->Compute(gates, gates + d2_, ht);
   }
+
   void ComputeHtPart1(T* gates, const T* ht_1, T* ht) const override {
     // W: {W_update, W_reset; W_state}
     act_gate_d2_->Compute(gates, gates);
@@ -394,8 +398,65 @@ class GRUKernelImpl : public GRUKernel<T> {
   int d_, d2_;
   std::shared_ptr<const VActKernel<T>> act_gate_d2_, act_gate_d_, act_state_d_;
   std::shared_ptr<const VMulKernel<T>> vmul_d_;
+#ifdef __AVX__
+  std::unique_ptr<const AVXAct> avx_act_gate_, avx_act_state_;
+#endif
 };
 
+#define INTRI8_FLOAT(isa)                                                     \
+  template <>                                                                 \
+  GRUKernelImpl<float, isa, kEQ8>::GRUKernelImpl(                             \
+      const std::string& act_gate, const std::string& act_state, int d)       \
+      : GRUKernel<float>() {                                                  \
+    avx_act_gate_ = GetAVXAct<isa>(act_gate);                                 \
+    avx_act_state_ = GetAVXAct<isa>(act_state);                               \
+  }                                                                           \
+  template <>                                                                 \
+  void GRUKernelImpl<float, isa, kEQ8>::ComputeH1(float* gates, float* ht)    \
+      const {                                                                 \
+    __m256 u, s;                                                              \
+    /* W: {W_update, W_reset; W_state} */                                     \
+    u = _mm256_loadu_ps(gates);                                               \
+    s = _mm256_loadu_ps(gates + 16);                                          \
+    s = _mm256_mul_ps(avx_act_gate_->Compute(u), avx_act_state_->Compute(s)); \
+    _mm256_storeu_ps(ht, s);                                                  \
+  }                                                                           \
+  template <>                                                                 \
+  void GRUKernelImpl<float, isa, kEQ8>::ComputeHtPart1(                       \
+      float* gates, const float* ht_1, float* ht) const {                     \
+    /* not exactly equal the any implementation */                            \
+    __m256 r, ht0;                                                            \
+    r = _mm256_loadu_ps(gates + 8);                                           \
+    ht0 = _mm256_loadu_ps(ht_1);                                              \
+    r = _mm256_mul_ps(avx_act_gate_->Compute(r), ht0);                        \
+    _mm256_storeu_ps(ht, r);                                                  \
+  }                                                                           \
+  template <>                                                                 \
+  void GRUKernelImpl<float, isa, kEQ8>::ComputeHtPart2(                       \
+      float* gates, const float* ht_1, float* ht) const {                     \
+    /* not exactly equal the any implementation */                            \
+    __m256 u, s, ht0;                                                         \
+    u = _mm256_loadu_ps(gates);                                               \
+    s = _mm256_loadu_ps(gates + 16);                                          \
+    ht0 = _mm256_loadu_ps(ht_1);                                              \
+    u = avx_act_gate_->Compute(u);                                            \
+    s = _mm256_mul_ps(u, avx_act_state_->Compute(s));                         \
+    u = _mm256_sub_ps(_mm256_set1_ps(1.f), u);                                \
+    u = _mm256_mul_ps(u, ht0);                                                \
+    u = _mm256_add_ps(s, u);                                                  \
+    _mm256_storeu_ps(ht, u);                                                  \
+  }
+
+#ifdef __AVX__
+INTRI8_FLOAT(jit::avx);
+#endif
+#ifdef __AVX2__
+INTRI8_FLOAT(jit::avx2);
+#endif
+#ifdef __AVX512F__
+INTRI8_FLOAT(jit::avx512f);
+#endif
+
 #define JITKERNEL_DECLARE_GRU(ker_class, ker_dtype)                       \
   template <>                                                             \
   std::shared_ptr<const GRUKernel<ker_dtype>> KernelPool::Get<            \
@@ -412,6 +473,7 @@ class GRUKernelImpl : public GRUKernel<T> {
 REGISTER_JITKERNEL_ARGS(gru, GRUKernel, JITKERNEL_DECLARE_GRU,
                         JITKERNEL_KEY_GRU, JITKERNEL_NEW_GRU_IMPL);
 
+#undef INTRI8_FLOAT
 #undef JITKERNEL_NEW_GRU_IMPL
 #undef JITKERNEL_KEY_GRU
 #undef JITKERNEL_DECLARE_GRU

python/paddle/fluid/tests/unittests/test_fusion_gru_op.py

@@ -125,6 +125,12 @@ class TestFusionGRUOpMD2(TestFusionGRUOp):
         self.D = 8
 
 
+class TestFusionGRUOpMD3(TestFusionGRUOp):
+    def set_confs(self):
+        self.M = 17
+        self.D = 15
+
+
 class TestFusionGRUOpBS1(TestFusionGRUOp):
     def set_confs(self):
         self.lod = [[3]]