[MXNET-107]Fused GRU implementation for CPU

python/mxnet/gluon/rnn/rnn_layer.py

@@ -186,7 +186,7 @@ def forward(self, inputs, states=None):
                 self.i2h_weight[i].shape = (self._gates*self._hidden_size, inputs.shape[2])
                 self.i2h_weight[i]._finish_deferred_init()
         if inputs.context.device_type == 'gpu' or \
-           self._mode == 'lstm' and not self._dropout:
+           self._mode in ['lstm', 'gru'] and not self._dropout:
             out = self._forward_kernel(inputs, states)
         else:
             out = self._forward(inputs, states)

src/operator/rnn-inl.h

@@ -101,13 +101,15 @@ inline size_t GetRNNWorkspaceSize(int seq_length,
   switch (mode) {
     case rnn_enum::kRnnRelu:
     case rnn_enum::kRnnTanh:
-    case rnn_enum::kGru:
-      LOG(FATAL) << "Only LSTM is supported at the moment";
+      LOG(FATAL) << "Only LSTM and GRU are supported at the moment";
       break;
     case rnn_enum::kLstm:
       size = (seq_length + 1) * batch_size * hidden_size * 4 + batch_size * hidden_size * 2
              + seq_length * batch_size * hidden_size * direction;
       break;
+    case rnn_enum::kGru:
+      size = seq_length * batch_size * hidden_size * direction * 4 + batch_size * hidden_size * 8;
+      break;
     default:
       LOG(FATAL) << "unknown RNN mode " << mode;
       break;
@@ -125,12 +127,16 @@ inline size_t GetRNNReserveSpaceSize(int num_layer,
   switch (mode) {
     case rnn_enum::kRnnRelu:
     case rnn_enum::kRnnTanh:
-    case rnn_enum::kGru:
-      LOG(FATAL) << "Only LSTM is supported at the moment";
+      LOG(FATAL) << "Only LSTM and GRU are supported at the moment";
       break;
     case rnn_enum::kLstm:
       size = num_layer * direction * seq_length * batch_size * hidden_size * 6;
       break;
+    case rnn_enum::kGru:
+      size = seq_length * batch_size * hidden_size * direction * num_layer * 8 +
+          batch_size * hidden_size * direction * 9 +
+          seq_length * batch_size * 7 * hidden_size * direction;
+      break;
     default:
       LOG(FATAL) << "unknown RNN mode " << mode;
       break;
@@ -221,14 +227,18 @@ void RNNForwardTraining(DType* ws,
   switch (mode) {
     case rnn_enum::kRnnTanh:
     case rnn_enum::kRnnRelu:
-    case rnn_enum::kGru:
-      LOG(FATAL) << "Only LSTM is supported at the moment";
+      LOG(FATAL) << "Only LSTM and GRU are supported at the moment";
       break;
     case rnn_enum::kLstm:
       LstmForwardTraining<DType>(ws, rs, state_outputs, num_layers, direction, seq_length,
                                  batch_size, input_size, state_size, x_ptr, hx_ptr, cx_ptr,
                                  w_ptr, b_ptr, y_ptr, hy_ptr, cy_ptr);
       break;
+    case rnn_enum::kGru:
+      GruForwardTraining<DType>(ws, rs, state_outputs, num_layers, direction, seq_length,
+                                batch_size, input_size, state_size, x_ptr, hx_ptr,
+                                w_ptr, y_ptr, hy_ptr);
+      break;
     default:
       LOG(FATAL) << "unknown RNN mode " << mode;
       break;
@@ -256,14 +266,18 @@ void RNNForwardInference(DType* ws,
   switch (mode) {
     case rnn_enum::kRnnRelu:
     case rnn_enum::kRnnTanh:
-    case rnn_enum::kGru:
-      LOG(FATAL) << "Only LSTM is supported at the moment";
+      LOG(FATAL) << "Only LSTM and GRU are supported at the moment";
       break;
     case rnn_enum::kLstm:
       LstmForwardInference<DType>(ws, state_outputs, num_layers, direction, seq_length,
                                   batch_size, input_size, state_size, x_ptr, hx_ptr, cx_ptr,
                                   w_ptr, b_ptr, y_ptr, hy_ptr, cy_ptr);
       break;
+    case rnn_enum::kGru:
+      GruForwardInference<DType>(ws, state_outputs, num_layers, direction, seq_length,
+                                 batch_size, input_size, state_size, x_ptr, hx_ptr,
+                                 w_ptr, y_ptr, hy_ptr);
+      break;
     default:
       LOG(FATAL) << "unknown RNN mode" << mode;
       break;
@@ -296,13 +310,17 @@ void RNNBackward(DType* ws,
   switch (mode) {
     case rnn_enum::kRnnRelu:
     case rnn_enum::kRnnTanh:
-    case rnn_enum::kGru:
       break;
     case rnn_enum::kLstm:
       LstmBackward<DType>(ws, rs, num_layers, direction, seq_length, batch_size,
                           input_size, state_size, x_ptr, hx_ptr, cx_ptr, w_ptr, y_ptr,
                           dy_ptr, dhy_ptr, dcy_ptr, dx_ptr, dhx_ptr, dcx_ptr, dw_ptr, db_ptr);
       break;
+    case rnn_enum::kGru:
+      GruBackward<DType>(ws, rs, num_layers, direction, seq_length, batch_size,
+                         input_size, state_size, x_ptr, hx_ptr, w_ptr,
+                         dy_ptr, dhy_ptr, dx_ptr, dhx_ptr, dw_ptr);
+      break;
     default:
       LOG(FATAL) << "unknown RNN mode" << mode;
       break;
@@ -330,7 +348,8 @@ class RNNOp : public Operator{
                        const std::vector<TBlob> &aux_args) {
     using namespace mshadow;
     using namespace mshadow::expr;
-    CHECK_EQ(param_.mode, rnn_enum::kLstm) << "Only lstm mode is supported at the moment.";
+    CHECK(param_.mode == rnn_enum::kLstm || param_.mode == rnn_enum::kGru)
+        << "Only lstm and gru mode are supported at the moment.";
     CHECK_EQ(param_.p, 0) << "Dropout is not supported at the moment.";
 
     size_t in_expected = (param_.mode == rnn_enum::kLstm) ? 4 : 3;
@@ -442,8 +461,10 @@ class RNNOp : public Operator{
                         const std::vector<TBlob> &aux_args) {
     using namespace mshadow;
     using namespace mshadow::expr;
-    CHECK_EQ(param_.mode, rnn_enum::kLstm) << "Only lstm mode is supported at the moment.";
+    CHECK(param_.mode == rnn_enum::kLstm || param_.mode == rnn_enum::kGru)
+        << "Only lstm and gru mode are supported at the moment.";
     CHECK_EQ(param_.p, 0) << "Dropout is not supported at the moment.";
+
     size_t in_expected = (param_.mode == rnn_enum::kLstm) ? 4 : 3;
     size_t out_expected = (param_.mode == rnn_enum::kLstm) ? 3 : 2;
     if (!param_.state_outputs) {