jitkernel lstm refer support peephole

tensor-tang · tensor-tang · commit f91386087378 · 2018-11-21T08:47:12.000Z
test=develop
diff --git a/paddle/fluid/operators/fused/fusion_lstm_op.cc b/paddle/fluid/operators/fused/fusion_lstm_op.cc
@@ -236,27 +236,31 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
   const int D = wh_dims[0];                                 \
   const int D4 = wh_dims[1]
 
-#define INIT_OTHER_DEFINES                                                  \
-  const T* x_data = x->data<T>();                                           \
-  const T* wx_data = wx->data<T>();                                         \
-  const T* wh_data = wh->data<T>();                                         \
-  /* diagonal weight*/                                                      \
-  const T* wp_data = bias->data<T>() + D4;                                  \
-  /* for peephole only*/                                                    \
-  T* checked_cell_data = nullptr;                                           \
-  auto place = ctx.GetPlace();                                              \
-  if (use_peepholes) {                                                      \
-    /* w_ic * Ct-1, w_fc * Ct-1  ; w_oc * Ct => ih*/                        \
-    auto* checked_cell = ctx.Output<Tensor>("CheckedCell");                 \
-    checked_cell_data = checked_cell->mutable_data<T>(place);               \
-  }                                                                         \
-  const auto& ker =                                                         \
-      math::jitkernel::KernelPool::Instance()                               \
-          .template Get<math::jitkernel::LSTMKernel<T>, const std::string&, \
-                        const std::string&, const std::string&>(            \
-              ctx.Attr<std::string>("gate_activation"),                     \
-              ctx.Attr<std::string>("candidate_activation"),                \
-              ctx.Attr<std::string>("cell_activation"), D, use_peepholes)
+#define INIT_OTHER_DEFINES                                      \
+  const T* x_data = x->data<T>();                               \
+  const T* wx_data = wx->data<T>();                             \
+  const T* wh_data = wh->data<T>();                             \
+  /* diagonal weight*/                                          \
+  const T* wp_data = bias->data<T>() + D4;                      \
+  /* for peephole only*/                                        \
+  T* checked_cell_data = nullptr;                               \
+  auto place = ctx.GetPlace();                                  \
+  if (use_peepholes) {                                          \
+    /* w_ic * Ct-1, w_fc * Ct-1  ; w_oc * Ct => ih*/            \
+    auto* checked_cell = ctx.Output<Tensor>("CheckedCell");     \
+    checked_cell_data = checked_cell->mutable_data<T>(place);   \
+  }                                                             \
+  const math::jitkernel::lstm_attr_t attr(                      \
+      D, ctx.Attr<std::string>("gate_activation"),              \
+      ctx.Attr<std::string>("candidate_activation"),            \
+      ctx.Attr<std::string>("cell_activation"), use_peepholes); \
+  math::jitkernel::lstm_t one_step;                             \
+  one_step.wp = wp_data;                                        \
+  one_step.checked = checked_cell_data;                         \
+  const auto& ker =                                             \
+      math::jitkernel::KernelPool::Instance()                   \
+          .template Get<math::jitkernel::LSTMKernel<T>,         \
+                        const math::jitkernel::lstm_attr_t&>(attr)
 
 // Wh GEMM
 #define GEMM_WH_ADDON(bs, prev, out)                                           \
@@ -299,7 +303,10 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
         prev_h_data = h0_data + bid * D;
         prev_c_data = c0_data + bid * D;
       } else {
-        ker->ComputeC1H1(xx_data, c_out_data, h_out_data, wp_data);
+        one_step.gates = xx_data;
+        one_step.ct = c_out_data;
+        one_step.ht = h_out_data;
+        ker->ComputeC1H1(&one_step, &attr);
         tstart = 1;
         // move one step
         prev_h_data = h_out_data;
@@ -310,8 +317,12 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
       }
       for (int step = tstart; step < seq_len; ++step) {
         GEMM_WH_ADDON(1, prev_h_data, xx_data);
-        ker->ComputeCtHt(xx_data, prev_c_data, c_out_data, h_out_data, wp_data,
-                         checked_cell_data);
+
+        one_step.gates = xx_data;
+        one_step.ct_1 = prev_c_data;
+        one_step.ct = c_out_data;
+        one_step.ht = h_out_data;
+        ker->ComputeCtHt(&one_step, &attr);
         // move one step
         prev_h_data = h_out_data;
         prev_c_data = c_out_data;
@@ -388,7 +399,11 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
       T* cur_h_out_data = batched_h_out_data;
       T* cur_c_out_data = batched_c_out_data;
       for (int i = 0; i < max_bs; ++i) {
-        ker->ComputeC1H1(cur_in_data, cur_c_out_data, cur_h_out_data, wp_data);
+        one_step.gates = cur_in_data;
+        one_step.ct = cur_c_out_data;
+        one_step.ht = cur_h_out_data;
+        ker->ComputeC1H1(&one_step, &attr);
+
         cur_in_data += D4;
         cur_c_out_data += D;
         cur_h_out_data += D;
@@ -413,8 +428,12 @@ class FuisonLSTMKernel : public framework::OpKernel<T> {
       T* cur_c_out_data = batched_c_out_data;
       T* cur_h_out_data = batched_h_out_data;
       for (int i = 0; i < cur_bs; ++i) {
-        ker->ComputeCtHt(cur_in_data, cur_prev_c_data, cur_c_out_data,
-                         cur_h_out_data, wp_data, checked_cell_data);
+        one_step.gates = cur_in_data;
+        one_step.ct_1 = cur_prev_c_data;
+        one_step.ct = cur_c_out_data;
+        one_step.ht = cur_h_out_data;
+        ker->ComputeCtHt(&one_step, &attr);
+
         // move one batch
         cur_in_data += D4;
         cur_prev_c_data += D;
diff --git a/paddle/fluid/operators/math/jit_code.cc b/paddle/fluid/operators/math/jit_code.cc
@@ -233,7 +233,7 @@ void LSTMJitCode::generate() {
     vmovups(ymm_src, ptr[reg_ptr_gates + offset + num_]);
     act<ymm_t>(ymm_i, ymm_src, act_gate_);
     vmulps(ymm_c, ymm_c, ymm_i);
-    if (first_) {
+    if (!compute_c1h1_) {
       // f
       vmovups(ymm_src, ptr[reg_ptr_gates + offset + 2 * num_]);
       act<ymm_t>(ymm_f, ymm_src, act_gate_);
@@ -242,8 +242,8 @@ void LSTMJitCode::generate() {
       vaddps(ymm_f, ymm_f, ymm_c);
     }
     /* H_t = act_cell(C_t) * ogated */
-    ymm_t ymm_ct = first_ ? ymm_c : ymm_f;
-    ymm_t ymm_o = first_ ? ymm_f : ymm_c;
+    ymm_t ymm_ct = compute_c1h1_ ? ymm_c : ymm_f;
+    ymm_t ymm_o = compute_c1h1_ ? ymm_f : ymm_c;
     ymm_t ymm_tmp = ymm_i;
     act<ymm_t>(ymm_tmp, ymm_ct, act_cell_);
     vmovups(ymm_src, ptr[reg_ptr_gates + offset + 3 * num_]);
diff --git a/paddle/fluid/operators/math/jit_code.h b/paddle/fluid/operators/math/jit_code.h
@@ -319,6 +319,12 @@ class LSTMJitCode : public VActJitCode {
  public:
   const char* name() const override {
     std::string base = "LSTMJitCode";
+    if (use_peephole_) {
+      base += "_Peephole";
+    }
+    if (compute_c1h1_) {
+      base += "_C1H1";
+    }
     auto AddTypeStr = [&](operand_type type) {
       switch (type) {
         case operand_type::relu:
@@ -340,30 +346,42 @@ class LSTMJitCode : public VActJitCode {
           break;
       }
     };
-    if (first_) {
-      base += "_C1H1";
-    }
     AddTypeStr(act_gate_);
     AddTypeStr(act_cand_);
     AddTypeStr(act_cell_);
     return base.c_str();
   }
 
-  explicit LSTMJitCode(int d, bool first, operand_type act_gate,
-                       operand_type act_cand, operand_type act_cell,
+  explicit LSTMJitCode(bool compute_c1h1, const lstm_attr_t& attr,
                        size_t code_size = 256 * 1024, void* code_ptr = nullptr)
-      : VActJitCode(d, act_gate, code_size, code_ptr),
-        num_(d),
-        first_(first),
-        act_gate_(act_gate),
-        act_cand_(act_cand),
-        act_cell_(act_cell) {}
+      : VActJitCode(attr.d, operand_type::sigmoid /* this is bugy*/, code_size,
+                    code_ptr),
+        compute_c1h1_(compute_c1h1) {
+    auto typeExchange = [](const std::string& type) -> gen::operand_type {
+      if (type == "sigmoid") {
+        return operand_type::sigmoid;
+      } else if (type == "relu") {
+        return operand_type::relu;
+      } else if (type == "tanh") {
+        return operand_type::tanh;
+      } else if (type == "identity" || type == "") {
+        return operand_type::identity;
+      }  // else throw error
+      return operand_type::identity;
+    };
+    num_ = attr.d;
+    use_peephole_ = attr.use_peephole;
+    act_gate_ = typeExchange(attr.act_gate);
+    act_cand_ = typeExchange(attr.act_cand);
+    act_cell_ = typeExchange(attr.act_cell);
+  }
   static bool init(int d);
   void generate() override;
 
  protected:
   int num_;
-  bool first_;
+  bool compute_c1h1_;
+  bool use_peephole_;
   operand_type act_gate_;
   operand_type act_cand_;
   operand_type act_cell_;
diff --git a/paddle/fluid/operators/math/jit_kernel.h b/paddle/fluid/operators/math/jit_kernel.h
@@ -122,18 +122,9 @@ class VTanhKernel : public VActKernel<T> {};
 template <typename T>
 class LSTMKernel : public Kernel {
  public:
-  virtual void ComputeCtHt(T *gates, const T *ct_1, T *ct, T *ht,
-                           /* below only used in peephole*/
-                           const T *wp_data = nullptr,
-                           T *checked = nullptr) const = 0;
-
-  virtual void ComputeC1H1(T *gates, T *ct, T *ht,
-                           /* below only used in peephole*/
-                           const T *wp_data = nullptr) const = 0;
-
-  //  void (*ComputeCtHt)(lstm_t *);
-  // // compute c1 and h1 without c0 or h0
-  //  void (*ComputeC1H1)(lstm_t *);
+  void (*ComputeCtHt)(lstm_t *, const lstm_attr_t *);
+  // compute c1 and h1 without c0 or h0
+  void (*ComputeC1H1)(lstm_t *, const lstm_attr_t *);
 };
 
 template <typename T>
diff --git a/paddle/fluid/operators/math/jit_kernel_impl.h b/paddle/fluid/operators/math/jit_kernel_impl.h
@@ -33,18 +33,24 @@ typedef struct {
   const void* ct_1;
   void* ct;
   void* ht;
-  /* below only used in peephole*/
-  const void* wp_data{nullptr};
+  /* weight_peephole and checked data are only used in peephole*/
+  const void* wp{nullptr};
   void* checked{nullptr};
 } lstm_t;
 
 typedef struct lstm_attr_s {
+  bool use_peephole;
   int d;
   std::string act_gate, act_cand, act_cell;
   lstm_attr_s() = default;
   lstm_attr_s(int _d, const std::string& _act_gate,
-              const std::string& _act_cand, const std::string& _act_cell)
-      : d(_d), act_gate(_act_gate), act_cand(_act_cand), act_cell(_act_cell) {}
+              const std::string& _act_cand, const std::string& _act_cell,
+              bool _use_peephole = false)
+      : use_peephole(_use_peephole),
+        d(_d),
+        act_gate(_act_gate),
+        act_cand(_act_cand),
+        act_cell(_act_cell) {}
 } lstm_attr_t;
 
 }  // namespace jitkernel
diff --git a/paddle/fluid/operators/math/jit_kernel_macro.h b/paddle/fluid/operators/math/jit_kernel_macro.h
@@ -82,10 +82,10 @@ namespace jitkernel {
 #define REGISTER_JITKERNEL_ARGS(ker_key, ker_class, marco_define_name,     \
                                 marco_declare, macro_find_key, macro_impl) \
   marco_define_name(ker_key, ker_class);                                   \
-  REGISTER_JITKERNEL_WITH_DTYPE(ker_class, float, JITKERNEL_DECLARE,       \
-                                JITKERNEL_FIND_KEY, JITKERNEL_IMPL);       \
-  REGISTER_JITKERNEL_WITH_DTYPE(ker_class, double, JITKERNEL_DECLARE,      \
-                                JITKERNEL_FIND_KEY, JITKERNEL_IMPL)
+  REGISTER_JITKERNEL_WITH_DTYPE(ker_class, float, marco_declare,           \
+                                macro_find_key, macro_impl);               \
+  REGISTER_JITKERNEL_WITH_DTYPE(ker_class, double, marco_declare,          \
+                                macro_find_key, macro_impl)
 
 #define REGISTER_JITKERNEL(ker_key, ker_class)                       \
   REGISTER_JITKERNEL_ARGS(ker_key, ker_class, JITKERNEL_DEFINE_NAME, \
diff --git a/paddle/fluid/operators/math/jit_kernel_refer.h b/paddle/fluid/operators/math/jit_kernel_refer.h
@@ -117,35 +117,50 @@ void (*getActFunc(const std::string& type))(const T*, T*, int) {  // NOLINT
 }
 
 template <typename T>
-void LSTMCtHt(lstm_t* step, lstm_attr_t* attr) {
+void LSTMCtHt(lstm_t* step, const lstm_attr_t* attr) {
   T* gates = reinterpret_cast<T*>(step->gates);
   const T* ct_1 = reinterpret_cast<const T*>(step->ct_1);
   T* ct = reinterpret_cast<T*>(step->ct);
   T* ht = reinterpret_cast<T*>(step->ht);
+  const T* wp = reinterpret_cast<const T*>(step->wp);
+  T* checked = reinterpret_cast<T*>(step->checked);
   auto act_gate = getActFunc<T>(attr->act_gate);
   auto act_cand = getActFunc<T>(attr->act_cand);
   auto act_cell = getActFunc<T>(attr->act_cell);
   int d = attr->d;
   int d2 = d * 2;
   int d3 = d * 3;
   // gates: W_ch, W_ih, W_fh, W_oh
-  act_gate(gates + d, gates + d, d3);
+  if (attr->use_peephole) {
+    VMul(wp, ct_1, checked, d);
+    VMul(wp + d, ct_1, checked + d, d);
+    VAdd(checked, gates + d, gates + d, d2);
+    act_gate(gates + d, gates + d, d2);
+  } else {
+    act_gate(gates + d, gates + d, d3);
+  }
 
-  /* C_t = C_t-1 * fgated + cand_gated * igated */
+  // C_t = C_t-1 * fgated + cand_gated * igated
   act_cand(gates, gates, d);
   VMul(gates, gates + d, gates + d, d);
   VMul(ct_1, gates + d2, gates + d2, d);
   VAdd(gates + d, gates + d2, ct, d);
 
-  /* H_t = act_cell(C_t) * ogated */
+  if (attr->use_peephole) {
+    // get ogated
+    VMul(wp + d2, ct, gates + d, d);
+    VAdd(gates + d, gates + d3, gates + d3, d);
+    act_gate(gates + d3, gates + d3, d);
+  }
+  // H_t = act_cell(C_t) * ogated
   act_cell(ct, gates + d2, d);
   VMul(gates + d2, gates + d3, ht, d);
 }
 
+// compute c1 and h1 without c0 or h0
 template <typename T>
-void LSTMC1H1(lstm_t* step, lstm_attr_t* attr) {
+void LSTMC1H1(lstm_t* step, const lstm_attr_t* attr) {
   T* gates = reinterpret_cast<T*>(step->gates);
-  const T* ct_1 = reinterpret_cast<const T*>(step->ct_1);
   T* ct = reinterpret_cast<T*>(step->ct);
   T* ht = reinterpret_cast<T*>(step->ht);
   auto act_gate = getActFunc<T>(attr->act_gate);
@@ -158,10 +173,16 @@ void LSTMC1H1(lstm_t* step, lstm_attr_t* attr) {
   act_gate(gates + d, gates + d, d);
   act_cand(gates, gates, d);
   VMul(gates, gates + d, ct, d);
+  if (attr->use_peephole) {
+    // get outgated, put W_oc * C_t on igated
+    const T* wp = reinterpret_cast<const T*>(step->wp);
+    VMul(wp + d2, ct, gates + d, d);
+    VAdd(gates + d, gates + d3, gates + d3, d);
+  }
   /* H_t = act_cell(C_t) * ogated */
   act_gate(gates + d3, gates + d3, d);
   act_cell(ct, gates + d2, d);
-  Vmul(gates + d2, gates + d3, ht, d);
+  VMul(gates + d2, gates + d3, ht, d);
 }
 
 }  // namespace refer
diff --git a/paddle/fluid/operators/math/jit_kernel_rnn.cc b/paddle/fluid/operators/math/jit_kernel_rnn.cc
diff --git a/paddle/fluid/operators/math/jit_kernel_test.cc b/paddle/fluid/operators/math/jit_kernel_test.cc
