Merge pull request #187 from ROCmSoftwarePlatform/lstm_operator

Lstm operator

src/eliminate_contiguous.cpp

@@ -27,6 +27,13 @@ void eliminate_contiguous::apply(program& p) const
 {
     for(auto ins : iterator_for(p))
     {
+        // skip the reshape operator for now, since there is a bug
+        // for the transpose followed by a reshape
+        if(ins->name() == "reshape")
+        {
+            continue;
+        }
+
         // Make a copy so we can modify it while we iterate
         auto args = ins->inputs();
         for(auto arg : ins->inputs())

src/include/migraphx/operators.hpp

@@ -1259,6 +1259,57 @@ struct gru
     }
 };
 
+struct lstm
+{
+    std::size_t hidden_size = 1;
+    std::vector<operation> actv_funcs{sigmoid{}, tanh{}, tanh{}};
+    rnn_direction direction = rnn_direction::forward;
+    float clip              = 0.0f;
+    int input_forget        = 0;
+
+    std::string name() const { return "lstm"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        auto in_dims     = inputs[0].lens();
+        auto hidden_dims = inputs[2].lens();
+        if(hidden_size != hidden_dims[2])
+        {
+            MIGRAPHX_THROW("LSTM: hidden size mismatch in attribute and input");
+        }
+
+        std::size_t num_directions = 1;
+        if(direction == rnn_direction::bidirectional)
+        {
+            num_directions = 2;
+        }
+
+        if(num_directions != hidden_dims[0])
+        {
+            MIGRAPHX_THROW("LSTM: num_direction does not match the direction attribute");
+        }
+
+        std::vector<std::size_t> out_dims(in_dims);
+        out_dims.insert(out_dims.begin() + 1, num_directions);
+        out_dims.back() = hidden_size;
+
+        return {inputs[0].type(), out_dims};
+    }
+};
+
+struct lstm_last_cell_output
+{
+    std::string name() const { return "lstm_last_cell_output"; }
+    shape compute_shape(std::vector<shape> inputs) const
+    {
+        check_shapes{inputs, *this}.has(1);
+        auto dims = inputs[0].lens();
+
+        // remove the first dimension, remaing are output shape
+        dims.erase(dims.begin());
+        return {inputs[0].type(), dims};
+    }
+};
+
 struct undefined
 {
     std::string name() const { return "undefined"; }

src/include/migraphx/rewrite_rnn.hpp

@@ -45,6 +45,18 @@ struct rewrite_rnn
                                           const operation& actv_func2) const;
 
     std::vector<operation> gru_actv_funcs(instruction_ref ins) const;
+
+    // for lstm operators
+    void apply_lstm(program& prog, instruction_ref ins) const;
+    std::vector<instruction_ref> lstm_cell(bool is_forward,
+                                           program& prog,
+                                           instruction_ref ins,
+                                           std::vector<instruction_ref> inputs,
+                                           const operation& actv_func1,
+                                           const operation& actv_func2,
+                                           const operation& actv_func3) const;
+
+    std::vector<operation> lstm_actv_funcs(instruction_ref ins) const;
 };
 
 } // namespace MIGRAPHX_INLINE_NS

src/onnx/onnx.cpp

@@ -89,6 +89,7 @@ struct onnx_parser
         add_mem_op("Transpose", &onnx_parser::parse_transpose);
         add_mem_op("RNN", &onnx_parser::parse_rnn);
         add_mem_op("GRU", &onnx_parser::parse_gru);
+        add_mem_op("LSTM", &onnx_parser::parse_lstm);
         add_mem_op("Pad", &onnx_parser::parse_pad);
 
         // init the activation function map
@@ -749,15 +750,17 @@ struct onnx_parser
         {
             auto names = attributes.at("activations").strings();
             vec_names.clear();
-            for_each(names.begin(), names.end(), [&](auto& fn) { vec_names.push_back(fn); });
+            vec_names.resize(names.size());
+            std::copy(names.begin(), names.end(), vec_names.begin());
         }
 
-        for_each(vec_names.begin(), vec_names.end(), [&](auto& fn) {
-            if(map_actv_funcs.count(fn) == 0)
-            {
-                MIGRAPHX_THROW("RNN: activation function " + std::string(fn) + " not supported");
-            }
+        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
+            return (map_actv_funcs.count(name) == 0);
         });
+        if(name_it != vec_names.end())
+        {
+            MIGRAPHX_THROW("RNN: activation function " + std::string(*name_it) + " not supported");
+        }
 
         // bidirectional case should have two activation functions.
         // one is for forward, and the other is for reverse.
@@ -839,8 +842,7 @@ struct onnx_parser
             auto names = attributes.at("activations").strings();
             vec_names.clear();
             vec_names.resize(names.size());
-            std::transform(
-                names.begin(), names.end(), vec_names.begin(), [](auto& str) { return str; });
+            std::copy(names.begin(), names.end(), vec_names.begin());
         }
 
         // need 4 activation functions
@@ -878,12 +880,13 @@ struct onnx_parser
             }
         }
 
-        for_each(vec_names.begin(), vec_names.end(), [&](auto& name) {
-            if(map_actv_funcs.count(name) == 0)
-            {
-                MIGRAPHX_THROW("GRU: activation function " + std::string(name) + " not supported");
-            }
+        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
+            return (map_actv_funcs.count(name) == 0);
         });
+        if(name_it != vec_names.end())
+        {
+            MIGRAPHX_THROW("GRU: activation function " + std::string(*name_it) + " not supported");
+        }
 
         std::vector<operation> vec_actv_funcs(vec_names.size());
         std::transform(vec_names.begin(), vec_names.end(), vec_actv_funcs.begin(), [&](auto& name) {
@@ -920,6 +923,178 @@ struct onnx_parser
         return {hidden_states, last_output};
     }
 
+    std::vector<instruction_ref>
+    parse_lstm(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
+    {
+        migraphx::shape input_shape = args[0]->get_shape();
+        std::size_t hidden_size     = args[2]->get_shape().lens()[2];
+
+        if(contains(attributes, "hidden_size"))
+        {
+            std::size_t hidden_size_att = parse_value(attributes.at("hidden_size")).at<int>();
+            if(hidden_size != hidden_size_att)
+            {
+                MIGRAPHX_THROW("LSTM: hidden size mismatch in input and attribute");
+            }
+        }
+
+        // Handling of direction to be added later
+        std::string direction{"forward"};
+        if(contains(attributes, "direction"))
+        {
+            direction = attributes.at("direction").s();
+        }
+
+        op::rnn_direction dirct = op::rnn_direction::forward;
+        if(direction == "bidirectional")
+        {
+            dirct = op::rnn_direction::bidirectional;
+        }
+        else if(direction == "reverse")
+        {
+            dirct = op::rnn_direction::reverse;
+        }
+        else if(direction == "forward")
+        {
+            dirct = op::rnn_direction::forward;
+        }
+        else
+        {
+            MIGRAPHX_THROW("LSTM: incorrect direction attribute");
+        }
+
+        std::vector<std::string> vec_names = {"sigmoid", "tanh", "tanh"};
+        if(contains(attributes, "activations"))
+        {
+            auto names = attributes.at("activations").strings();
+            vec_names.clear();
+            vec_names.resize(names.size());
+            std::copy(names.begin(), names.end(), vec_names.begin());
+        }
+
+        // need 6 activation functions for bidirectional directions
+        if(dirct == op::rnn_direction::bidirectional)
+        {
+            // 6 activation functions are used in the bidirectional
+            // scenario. No spec is provided in onnx::operator. we
+            // use the algorithm that: if 1 actv function is provided,
+            // repeat 1st six times. If 2 actv functins are provided,
+            // repeat 2nd once, then repeat all three once
+            // if 3 actv funcs are provide, repeat all three once.
+            // the same algorithm is used for 4, 5, and 6 actv funcions
+            // provided. This may need change later
+            switch(vec_names.size())
+            {
+            case 1:
+                vec_names = {vec_names.at(0),
+                             vec_names.at(0),
+                             vec_names.at(0),
+                             vec_names.at(0),
+                             vec_names.at(0),
+                             vec_names.at(0)};
+                break;
+
+            case 2:
+                // repeat the 2nd actv func once, then repeat all three another time
+                vec_names = {vec_names.at(0),
+                             vec_names.at(1),
+                             vec_names.at(1),
+                             vec_names.at(0),
+                             vec_names.at(1),
+                             vec_names.at(1)};
+                break;
+
+            case 3:
+                // repeat all three actv funcs once
+                vec_names = {vec_names.at(0),
+                             vec_names.at(1),
+                             vec_names.at(2),
+                             vec_names.at(0),
+                             vec_names.at(1),
+                             vec_names.at(2)};
+                break;
+
+            case 4:
+                vec_names = {vec_names.at(0),
+                             vec_names.at(1),
+                             vec_names.at(2),
+                             vec_names.at(3),
+                             vec_names.at(3),
+                             vec_names.at(3)};
+                break;
+
+            case 5:
+                vec_names = {vec_names.at(0),
+                             vec_names.at(1),
+                             vec_names.at(2),
+                             vec_names.at(3),
+                             vec_names.at(4),
+                             vec_names.at(4)};
+                break;
+
+            default: break;
+            }
+        }
+        else
+        {
+            switch(vec_names.size())
+            {
+            case 1: vec_names = {vec_names.at(0), vec_names.at(0), vec_names.at(0)}; break;
+
+            case 2:
+                // repeat the 2nd actv func once, so we have 3 actv funcs
+                vec_names = {vec_names.at(0), vec_names.at(1), vec_names.at(1)};
+                break;
+
+            default: break;
+            }
+        }
+
+        auto name_it = std::find_if(vec_names.begin(), vec_names.end(), [&](auto& name) {
+            return (map_actv_funcs.count(name) == 0);
+        });
+        if(name_it != vec_names.end())
+        {
+            MIGRAPHX_THROW("LSTM: activation function " + std::string(*name_it) + " not supported");
+        }
+
+        std::vector<operation> vec_actv_funcs(vec_names.size());
+        std::transform(vec_names.begin(), vec_names.end(), vec_actv_funcs.begin(), [&](auto& name) {
+            return map_actv_funcs[name];
+        });
+
+        float clip = 0.0;
+        if(contains(attributes, "clip"))
+        {
+            clip = parse_value(attributes.at("clip")).at<float>();
+        }
+
+        int input_forget = 0;
+        if(contains(attributes, "input_forget"))
+        {
+            input_forget = parse_value(attributes.at("input_forget")).at<int>();
+        }
+
+        // append undefined opeator to make 6 arguments
+        if(args.size() < 8)
+        {
+            auto ins = prog.add_instruction(op::undefined{});
+            args.insert(args.end(), 8 - args.size(), ins);
+        }
+
+        // first output for concatenation of hidden states
+        auto hidden_states = prog.add_instruction(
+            op::lstm{hidden_size, vec_actv_funcs, dirct, clip, input_forget}, std::move(args));
+
+        // second output for last lstm output
+        auto last_output = prog.add_instruction(op::rnn_last_output{}, hidden_states);
+
+        // third output for last cell output
+        auto last_cell_output = prog.add_instruction(op::lstm_last_cell_output{}, hidden_states);
+
+        return {hidden_states, last_output, last_cell_output};
+    }
+
     void parse_from(std::istream& is)
     {
         onnx::ModelProto model;

test/CMakeLists.txt

@@ -113,13 +113,18 @@ endif()
 
 # Onnx test
 set(TEST_ONNX_DIR ${CMAKE_CURRENT_SOURCE_DIR}/onnx)
-add_executable(test_onnx ${TEST_ONNX_DIR}/onnx_test.cpp)
-rocm_clang_tidy_check(test_onnx)
-target_link_libraries(test_onnx migraphx_onnx)
-target_include_directories(test_onnx PUBLIC include)
-add_test(NAME test_onnx COMMAND $<TARGET_FILE:test_onnx> WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/onnx) 
-add_dependencies(tests test_onnx)
-add_dependencies(check test_onnx)
+file (GLOB ONNX_TESTS ${TEST_ONNX_DIR}/*.cpp)
+foreach(ONNX_TEST ${ONNX_TESTS})
+    get_filename_component(BASE_NAME ${ONNX_TEST} NAME_WE)
+    set(TEST_NAME test_${BASE_NAME})
+    add_executable(${TEST_NAME} ${TES_ONNX_DIR}/${ONNX_TEST})
+    rocm_clang_tidy_check(${TEST_NAME})
+    target_link_libraries(${TEST_NAME} migraphx_onnx)
+    target_include_directories(${TEST_NAME} PUBLIC include)
+    add_test(NAME ${TEST_NAME} COMMAND $<TARGET_FILE:${TEST_NAME}> WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/onnx) 
+    add_dependencies(tests ${TEST_NAME})
+    add_dependencies(check ${TEST_NAME})
+endforeach()
 
 if(MIGRAPHX_ENABLE_PYTHON)
 add_subdirectory(py)