more test examples are added for the RNN operator.

src/include/migraphx/operators.hpp

@@ -1098,7 +1098,7 @@ struct rnn
 
         if(num_directions != hidden_dims[0])
         {
-            MIGRAPHX_THROW("RNN: num_direction does not match the direction attribute");
+            MIGRAPHX_THROW("RNN: num_direction mismatch in attribute and input");
         }
 
         std::vector<std::size_t> out_dims(in_dims);

src/onnx/onnx.cpp

@@ -732,7 +732,7 @@ struct onnx_parser
                                                   std::move(args));
         result.push_back(hidden_states);
 
-        // second out for the last hidden state
+        // second output for the last hidden state
         auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
         result.push_back(last_output);
 

src/rewrite_rnn.cpp

@@ -186,6 +186,7 @@ std::vector<instruction_ref> rewrite_rnn::rnn_cell(bool is_forward,
     long seq_index      = is_forward ? 0 : seq_len - 1;
     for(std::size_t i = 0; i < seq_len; i++)
     {
+        seq_index = is_forward ? i : (seq_len - 1 - i);
         auto xt = prog.insert_instruction(ins, op::slice{{0}, {seq_index}, {seq_index + 1}}, input);
         xt      = prog.insert_instruction(ins, op::squeeze{{0}}, xt);
         auto xt_wi = prog.insert_instruction(ins, op::dot{}, xt, tran_sw);
@@ -220,7 +221,6 @@ std::vector<instruction_ref> rewrite_rnn::rnn_cell(bool is_forward,
                              ? last_out
                              : prog.insert_instruction(ins, op::concat{0}, last_out, hidden_out);
         }
-        seq_index = is_forward ? (seq_index + 1) : (seq_index - 1);
     }
 
     std::vector<instruction_ref> out_args;

test/gpu/miopen.cpp

@@ -1049,6 +1049,126 @@ struct test_conv_bn_relu_pooling2
     }
 };
 
+struct test_rnn_forward
+{
+    migraphx::program create_program() const
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 1;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        auto seq = p.add_parameter("seq", in_shape);
+        auto ih = p.add_parameter("ih", ih_shape);
+        auto w = p.add_parameter("w", w_shape);
+        auto r = p.add_parameter("r", r_shape);
+        auto bias = p.add_parameter("bias", b_shape);
+
+        auto output = p.add_instruction(migraphx::op::rnn{hidden_size,
+                                            {migraphx::op::tanh{}, migraphx::op::tanh{}},
+                                            migraphx::op::rnn::forward,
+                                            clip},
+                          seq,
+                          w,
+                          r,
+                          bias,
+                          ih);
+        auto last = p.add_instruction(migraphx::op::rnn_last_output{}, output);
+        p.add_instruction(migraphx::op::add{}, last, last);
+
+        return p;
+    }
+};
+
+struct test_rnn_reverse
+{
+    migraphx::program create_program() const
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 1;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        auto seq = p.add_parameter("seq", in_shape);
+        auto ih = p.add_parameter("ih", ih_shape);
+        auto w = p.add_parameter("w", w_shape);
+        auto r = p.add_parameter("r", r_shape);
+        auto bias = p.add_parameter("bias", b_shape);
+
+        auto output = p.add_instruction(migraphx::op::rnn{hidden_size,
+                                            {migraphx::op::tanh{}, migraphx::op::tanh{}},
+                                            migraphx::op::rnn::reverse,
+                                            clip},
+                          seq,
+                          w,
+                          r,
+                          bias,
+                          ih);
+        auto last = p.add_instruction(migraphx::op::rnn_last_output{}, output);
+        p.add_instruction(migraphx::op::add{}, last, last);
+        
+        return p;
+    }
+};
+
+struct test_rnn_bidirectional
+{
+    migraphx::program create_program() const
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 2;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        auto seq = p.add_parameter("seq", in_shape);
+        auto ih = p.add_parameter("ih", ih_shape);
+        auto w = p.add_parameter("w", w_shape);
+        auto r = p.add_parameter("r", r_shape);
+        auto bias = p.add_parameter("bias", b_shape);
+
+        auto output = p.add_instruction(migraphx::op::rnn{hidden_size,
+                                            {migraphx::op::tanh{}, migraphx::op::tanh{}},
+                                            migraphx::op::rnn::bidirectional,
+                                            clip},
+                          seq,
+                          w,
+                          r,
+                          bias,
+                          ih);
+        auto last = p.add_instruction(migraphx::op::rnn_last_output{}, output);
+        p.add_instruction(migraphx::op::add{}, last, last);
+        
+        return p;
+    }
+};
+
 int main()
 {
     verify_program<test_abs>();
@@ -1108,4 +1228,5 @@ int main()
     verify_program<test_slice>();
     verify_program<test_gather>();
     verify_program<test_gather_neg_axis>();
+    verify_program<test_rnn_forward>();
 }

test/op_shape_test.cpp

@@ -249,4 +249,136 @@ TEST_CASE(gather)
     }
 }
 
+TEST_CASE(rnn)
+{
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 1;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        expect_shape(migraphx::shape{migraphx::shape::float_type, {seq_len, num_dirct, batch_size, hidden_size}},
+                     migraphx::op::rnn{hidden_size,
+                                      {migraphx::op::tanh{}},
+                                      migraphx::op::rnn::forward, clip},
+                                      in_shape, w_shape, r_shape, b_shape, ih_shape);        
+    }
+
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 1;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        expect_shape(migraphx::shape{migraphx::shape::float_type, {seq_len, num_dirct, batch_size, hidden_size}},
+                     migraphx::op::rnn{hidden_size,
+                                      {migraphx::op::tanh{}},
+                                      migraphx::op::rnn::reverse, clip},
+                                      in_shape, w_shape, r_shape, b_shape, ih_shape);        
+    }
+
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 2;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        expect_shape(migraphx::shape{migraphx::shape::float_type, {seq_len, num_dirct, batch_size, hidden_size}},
+                     migraphx::op::rnn{hidden_size,
+                                      {migraphx::op::tanh{}},
+                                      migraphx::op::rnn::bidirectional, clip},
+                                      in_shape, w_shape, r_shape, b_shape, ih_shape);        
+    }
+
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 1;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        throws_shape(migraphx::op::rnn{hidden_size + 1,
+                                      {migraphx::op::tanh{}},
+                                      migraphx::op::rnn::forward, clip},
+                                      in_shape, w_shape, r_shape, b_shape, ih_shape);        
+    }
+
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 1;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        throws_shape(migraphx::op::rnn{hidden_size,
+                                      {migraphx::op::tanh{}},
+                                      migraphx::op::rnn::bidirectional, clip},
+                                      in_shape, w_shape, r_shape, b_shape, ih_shape);        
+    }
+
+    {
+        std::size_t batch_size  = 2;
+        std::size_t seq_len     = 2;
+        std::size_t hidden_size = 4;
+        std::size_t input_size  = 3;
+        std::size_t num_dirct   = 2;
+        float clip = 0.0f;
+
+        migraphx::program p;
+        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
+        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
+        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
+        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
+        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};
+
+        throws_shape(migraphx::op::rnn{hidden_size,
+                                      {migraphx::op::tanh{}},
+                                      migraphx::op::rnn::forward, clip},
+                                      in_shape, w_shape, r_shape, b_shape, ih_shape);        
+    }
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }