From 5291ecc11c10c7a87aa99d00bcfbd8c579698ebe Mon Sep 17 00:00:00 2001 From: Roman Kazantsev Date: Mon, 10 Jun 2024 12:28:01 +0400 Subject: [PATCH] [MOC] Implement LSTMCell fusion for Keras LSTM layers (#24873) **Details:** Implement LSTMCell fusion for Keras LSTM layers. Such sub-graph with split weights appears for bi-directional LSTM layer. For custom model, we managed to fuse 6 LSTMCell sub-graphs and obtained 19% performance gain on CPU. **Ticket:** 142924 --------- Signed-off-by: Kazantsev, Roman --- .../common_optimizations/lstm_cell_fusion.hpp | 36 +- .../common_optimizations/lstm_cell_fusion.cpp | 316 +++++++++++++++++- .../common_optimizations/lstm_cell_fusion.cpp | 186 +++++++++++ 3 files changed, 523 insertions(+), 15 deletions(-) diff --git a/src/common/transformations/include/transformations/common_optimizations/lstm_cell_fusion.hpp b/src/common/transformations/include/transformations/common_optimizations/lstm_cell_fusion.hpp index a545f9ed80e..947e39edcfd 100644 --- a/src/common/transformations/include/transformations/common_optimizations/lstm_cell_fusion.hpp +++ b/src/common/transformations/include/transformations/common_optimizations/lstm_cell_fusion.hpp @@ -1,4 +1,4 @@ -// Copyright (C) 2023 Intel Corporation +// Copyright (C) 2023-2024 Intel Corporation // SPDX-License-Identifier: Apache-2.0 // @@ -11,17 +11,43 @@ namespace ov { namespace pass { class TRANSFORMATIONS_API LSTMCellFusion; +class TRANSFORMATIONS_API LSTMCellFusionWithJointWeights; +class TRANSFORMATIONS_API LSTMCellFusionWithSplitWeights; } // namespace pass } // namespace ov /** * @ingroup ov_transformation_common_api - * @brief LSTMCellFusion transformation replaces a sequence of - * operations with LSTMCell op. + * @brief LSTMCellFusionWithJointWeights transformation converts + * a sequence of operations with merged weights into LSTMCell op. */ -class ov::pass::LSTMCellFusion : public ov::pass::MatcherPass { +class ov::pass::LSTMCellFusionWithJointWeights : public ov::pass::MatcherPass { +public: + OPENVINO_RTTI("LSTMCellFusionWithJointWeights", "0"); + LSTMCellFusionWithJointWeights(); +}; + +/** + * @ingroup ov_transformation_common_api + * @brief LSTMCellFusionWithSplitWeights transformation converts + * a sequence of operations with split weights into LSTMCell op. + */ +class ov::pass::LSTMCellFusionWithSplitWeights : public ov::pass::MatcherPass { +public: + OPENVINO_RTTI("LSTMCellFusionWithSplitWeights", "0"); + LSTMCellFusionWithSplitWeights(); +}; + +/** + * @ingroup ov_transformation_common_api + * @brief LSTMCellFusion transformation replaces various sub-graphs with a LSTMCell op. + */ +class ov::pass::LSTMCellFusion : public ov::pass::GraphRewrite { public: OPENVINO_RTTI("LSTMCellFusion", "0"); - LSTMCellFusion(); + LSTMCellFusion() { + add_matcher(); + add_matcher(); + } }; diff --git a/src/common/transformations/src/transformations/common_optimizations/lstm_cell_fusion.cpp b/src/common/transformations/src/transformations/common_optimizations/lstm_cell_fusion.cpp index b7ad38dc74e..763f1e57cd1 100644 --- a/src/common/transformations/src/transformations/common_optimizations/lstm_cell_fusion.cpp +++ b/src/common/transformations/src/transformations/common_optimizations/lstm_cell_fusion.cpp @@ -1,4 +1,4 @@ -// Copyright (C) 2023 Intel Corporation +// Copyright (C) 2023-2024 Intel Corporation // SPDX-License-Identifier: Apache-2.0 // @@ -23,6 +23,107 @@ #include "openvino/pass/pattern/op/wrap_type.hpp" #include "transformations/utils/utils.hpp" +using namespace ov; +using namespace ov::pass; + +namespace { +static std::string get_activation_name(const std::shared_ptr& node) { + std::string name = node->get_type_name(); + name[0] = std::tolower(name[0]); + return name; +} + +void generate_gate_pattern(const std::shared_ptr& x, + const std::shared_ptr& h, + std::shared_ptr& it, + std::shared_ptr& wi, + std::shared_ptr& ri, + std::shared_ptr& bi, + std::shared_ptr& x_by_wi, + std::shared_ptr& h_by_ri) { + wi = pattern::any_input([](const Output& output) { + return pattern::has_static_shape()(output) && pattern::rank_equals(2)(output); + }); + ri = pattern::any_input([](const Output& output) { + return pattern::has_static_shape()(output) && pattern::rank_equals(2)(output); + }); + bi = pattern::any_input([](const Output& output) { + return pattern::has_static_shape()(output) && pattern::rank_equals(1)(output); + }); + + x_by_wi = pattern::wrap_type({x, wi}); + auto x_by_wi_biased = pattern::wrap_type({x_by_wi, bi}); + h_by_ri = pattern::wrap_type({h, ri}); + it = pattern::wrap_type({x_by_wi_biased, h_by_ri}); + it = pattern::wrap_type({it}); +} + +bool check_weights_format(const ov::Output& w, + const ov::Output& r, + const ov::Output& b, + size_t input_size, + size_t hidden_size, + const ov::pass::pattern::PatternValueMap& pattern_map, + const std::shared_ptr& x_by_w_label, + const std::shared_ptr& h_by_r_label) { + // w must be of a shape [input_size, hidden_size] + // r must be of a shape [hidden_size, hidden_size] + // b must be of a shape [hidden_size] + if (w.get_shape() != ov::Shape{input_size, hidden_size}) { + return false; + } + + if (r.get_shape() != ov::Shape{hidden_size, hidden_size}) { + return false; + } + + if (b.get_shape() != ov::Shape{hidden_size}) { + return false; + } + + // check transpose attributes for MatMul operations + if (const auto& matmul = ov::as_type_ptr(pattern_map.at(x_by_w_label).get_node_shared_ptr())) { + if (matmul->get_transpose_a() || matmul->get_transpose_b()) { + return false; + } + } else { + return false; + } + + if (const auto& matmul = ov::as_type_ptr(pattern_map.at(h_by_r_label).get_node_shared_ptr())) { + if (matmul->get_transpose_a() || matmul->get_transpose_b()) { + return false; + } + } else { + return false; + } + + return true; +} + +ov::Output prepare_weight_fico(const ov::Output& f, + const ov::Output& i, + const ov::Output& c, + const ov::Output& o, + NodeRegistry& rg) { + // at this point input weights are of shape [input_size, hidden_size] + // before concatenation it needs to transpose them + // to get a shape equal to [hidden_size, input_size] + auto transpose_order = rg.make(element::i32, ov::Shape{2}, std::vector{1, 0}); + auto f_tr = rg.make(f, transpose_order); + auto i_tr = rg.make(i, transpose_order); + auto c_tr = rg.make(c, transpose_order); + auto o_tr = rg.make(o, transpose_order); + + ov::Output w = rg.make(ov::OutputVector{f_tr, i_tr, c_tr, o_tr}, 0); + if (const auto& constant = ov::util::constantfold_subgraph(w)) { + w = constant; + } + + return w; +} +} // namespace + /* The following graph is fused to LSTMCell @@ -107,15 +208,8 @@ +----------+ */ - -static std::string get_activation_name(const std::shared_ptr& node) { - std::string name = node->get_type_name(); - name[0] = std::tolower(name[0]); - return name; -} - -ov::pass::LSTMCellFusion::LSTMCellFusion() { - MATCHER_SCOPE(LSTMCellFusion); +ov::pass::LSTMCellFusionWithJointWeights::LSTMCellFusionWithJointWeights() { + MATCHER_SCOPE(LSTMCellFusionWithJointWeights); auto x_label = pattern::any_input(pattern::rank_equals(2)); auto h_label = pattern::any_input(pattern::rank_equals(2)); @@ -352,3 +446,205 @@ ov::pass::LSTMCellFusion::LSTMCellFusion() { auto m = std::make_shared(Ho_label, matcher_name); this->register_matcher(m, callback); } + +/* + The following graph is fused to LSTMCell + + +----+ +----+ +----+ +----+ + | It | | Ct | | Ft | | Ot | + +----+ +----+ +----+ +----+ + | | | | + v | | v + +------+-----+ | | +------+-----+ + | Activation | | | | Activation | + | (i_t) | | | | (o_t) | + +------+-----+ | | +------+-----+ + | v | | + | +------+-----+ | | + | | Activation | | | + | | (c_t) | | | + | +------+-----+ | | + | | | | + | | v | + +---+ +---+ +------+-----+ | + | | | Activation | +-----+ | + v v | (f_t) | | C | | + +--+---+---+ +------------+ +-----+ | + | Multiply | | | | + +----+-----+ | +--------+ | + | | | | + | v v | + | +---+---+--+ | + | | Multiply | | + | +----+-----+ | + | | | + | | | + +---------+ +--------+ | + | | | + v v | + +-+-----+-+ | + | Add | | + | (C out) | | + +----+----+ | + | | + v | + +-----+------+ | + | Activation | | + +-----+------+ | + | | + | | + +----------+ +-------------------+ + | | + v v + +--+----+--+ + | Multiply | + | (H out) | + +----------+ + + where each of It, Ct, Ft, and Ot represents a separate graph: + + +-----+ +-----+ + | X | | W | + +--+--+ +--+--+ + | | + +---+ +---+ + | | + v v + +--+--+--+ +------+ +-----+ +----+ + | MatMul | | Bias | | H | | R | + +----+---+ +---+--+ +-----+ +----+ + | | | | + | +--------+ +---+ +---+ + | | | | + v v v v + +--+--+--+ +--------+ + | Add | | MatMul | + +----+---+ +--+-----+ + | | + | +-----------------------------+ + | | + v v + +--+---+--+ + | Add | + +----+----+ + | + + */ +ov::pass::LSTMCellFusionWithSplitWeights::LSTMCellFusionWithSplitWeights() { + MATCHER_SCOPE(LSTMCellFusionWithSplitWeights); + + auto x_label = pattern::any_input(pattern::rank_equals(2)); + auto h_label = pattern::any_input(pattern::rank_equals(2)); + auto c_label = pattern::any_input(pattern::rank_equals(2)); + + // it expects a pattern with split weights (input, recurrent and bias) for each gate + std::shared_ptr it_label, wi_label, ri_label, bi_label, x_by_wi_label, h_by_ri_label; + generate_gate_pattern(x_label, h_label, it_label, wi_label, ri_label, bi_label, x_by_wi_label, h_by_ri_label); + std::shared_ptr ft_label, wf_label, rf_label, bf_label, x_by_wf_label, h_by_rf_label; + generate_gate_pattern(x_label, h_label, ft_label, wf_label, rf_label, bf_label, x_by_wf_label, h_by_rf_label); + std::shared_ptr ot_label, wo_label, ro_label, bo_label, x_by_wo_label, h_by_ro_label; + generate_gate_pattern(x_label, h_label, ot_label, wo_label, ro_label, bo_label, x_by_wo_label, h_by_ro_label); + std::shared_ptr c1t_label, wc_label, rc_label, bc_label, x_by_wc_label, h_by_rc_label; + generate_gate_pattern(x_label, h_label, c1t_label, wc_label, rc_label, bc_label, x_by_wc_label, h_by_rc_label); + + auto it_mul_c1t_label = pattern::wrap_type({it_label, c1t_label}); + auto ft_mul_c_label = pattern::wrap_type({ft_label, c_label}); + auto ct_label = pattern::wrap_type({ft_mul_c_label, it_mul_c1t_label}); + + auto ct_activated_label = pattern::wrap_type({ct_label}); + auto ht_label = pattern::wrap_type({ct_activated_label, ot_label}); + + matcher_pass_callback callback = [=](pattern::Matcher& m) { + NodeRegistry rg; + + const auto& pattern_map = m.get_pattern_value_map(); + const auto& x = pattern_map.at(x_label); + const auto& h = pattern_map.at(h_label); + const auto& c = pattern_map.at(c_label); + + const auto& wi = pattern_map.at(wi_label); + const auto& ri = pattern_map.at(ri_label); + const auto& bi = pattern_map.at(bi_label); + + const auto& wf = pattern_map.at(wf_label); + const auto& rf = pattern_map.at(rf_label); + const auto& bf = pattern_map.at(bf_label); + + const auto& wo = pattern_map.at(wo_label); + const auto& ro = pattern_map.at(ro_label); + const auto& bo = pattern_map.at(bo_label); + + const auto& wc = pattern_map.at(wc_label); + const auto& rc = pattern_map.at(rc_label); + const auto& bc = pattern_map.at(bc_label); + + const auto& it = pattern_map.at(it_label).get_node_shared_ptr(); + const auto& ft = pattern_map.at(ft_label).get_node_shared_ptr(); + const auto& ot = pattern_map.at(ot_label).get_node_shared_ptr(); + const auto& c1t = pattern_map.at(c1t_label).get_node_shared_ptr(); + const auto& ct_activated = pattern_map.at(ct_activated_label).get_node_shared_ptr(); + + auto ct = pattern_map.at(ct_label); + auto ht = pattern_map.at(ht_label); + + // check weights format + auto input_size = static_cast(wi.get_shape()[0]); + auto hidden_size = static_cast(wi.get_shape()[1]); + bool weights_format_ok = + check_weights_format(wi, ri, bi, input_size, hidden_size, pattern_map, x_by_wi_label, h_by_ri_label); + weights_format_ok &= + check_weights_format(wf, rf, bf, input_size, hidden_size, pattern_map, x_by_wf_label, h_by_rf_label); + weights_format_ok &= + check_weights_format(wo, ro, bo, input_size, hidden_size, pattern_map, x_by_wo_label, h_by_ro_label); + weights_format_ok &= + check_weights_format(wc, rc, bc, input_size, hidden_size, pattern_map, x_by_wc_label, h_by_rc_label); + if (!weights_format_ok) { + return false; + } + + // check activation functions + std::string f_activation_name = it->get_type_name(); + if (f_activation_name != ft->get_type_name() || f_activation_name != ot->get_type_name()) { + return false; + } + f_activation_name[0] = std::tolower(f_activation_name[0]); + + std::string g_activation_name = get_activation_name(c1t); + std::string h_activation_name = get_activation_name(ct_activated); + + // prepare weights in a format required by LSTMCell + // W weights must be in fico format + ov::Output W = prepare_weight_fico(wf, wi, wc, wo, rg); + ov::Output R = prepare_weight_fico(rf, ri, rc, ro, rg); + + ov::Output B = rg.make(ov::OutputVector{bf, bi, bc, bo}, 0); + if (const auto& constant = ov::util::constantfold_subgraph(B)) + B = constant; + + auto lstm_cell = rg.make( + x, + h, + c, + W, + R, + B, + static_cast(hidden_size), + std::vector{f_activation_name, g_activation_name, h_activation_name}); + + if (transformation_callback(lstm_cell)) { + return false; + } + + lstm_cell->set_friendly_name(m.get_match_root()->get_friendly_name()); + + copy_runtime_info(m.get_matched_nodes(), rg.get()); + + ht.replace(lstm_cell->output(0)); + ct.replace(lstm_cell->output(1)); + + return true; + }; + + auto m = std::make_shared(ht_label, matcher_name); + this->register_matcher(m, callback); +} diff --git a/src/common/transformations/tests/common_optimizations/lstm_cell_fusion.cpp b/src/common/transformations/tests/common_optimizations/lstm_cell_fusion.cpp index 72da74a1545..4326601be61 100644 --- a/src/common/transformations/tests/common_optimizations/lstm_cell_fusion.cpp +++ b/src/common/transformations/tests/common_optimizations/lstm_cell_fusion.cpp @@ -4,7 +4,10 @@ #include "transformations/common_optimizations/lstm_cell_fusion.hpp" +#include + #include "common_test_utils/ov_test_utils.hpp" +#include "openvino/core/validation_util.hpp" #include "openvino/op/abs.hpp" #include "openvino/op/add.hpp" #include "openvino/op/concat.hpp" @@ -12,10 +15,13 @@ #include "openvino/op/lstm_cell.hpp" #include "openvino/op/matmul.hpp" #include "openvino/op/multiply.hpp" +#include "openvino/op/negative.hpp" #include "openvino/op/parameter.hpp" +#include "openvino/op/relu.hpp" #include "openvino/op/sigmoid.hpp" #include "openvino/op/split.hpp" #include "openvino/op/tanh.hpp" +#include "openvino/op/transpose.hpp" using namespace ov; @@ -108,3 +114,183 @@ TEST_P(LSTMCellFusionTestSuite, SubgraphFusedToLSTMCell) { INSTANTIATE_TEST_SUITE_P(LSTMCellFusion, LSTMCellFusionTestSuite, testing::Combine(testing::Values(false, true), testing::Values(1, 2), testing::Values(1, -1))); + +using LSTMCellTFKerasFusionParam = std::tuple; // hidden size + +class LSTMCellFusionWithSplitWeights : public testing::WithParamInterface, + public TransformationTestsF {}; + +namespace { +void generate_weights_value(std::vector& weights_value, const Shape& weights_shape) { + weights_value.resize(shape_size(weights_shape)); + std::mt19937 rng(9812); + std::uniform_real_distribution distribution(-300, 300); + for (size_t i = 0; i < weights_value.size(); ++i) { + weights_value[i] = distribution(rng); + } +} + +void generate_weights(std::vector& w, + std::vector& r, + std::vector& b, + size_t input_size, + size_t hidden_size) { + Shape w_shape({hidden_size, input_size}); + Shape r_shape({hidden_size, hidden_size}); + Shape b_shape({hidden_size}); + + generate_weights_value(w, w_shape); + generate_weights_value(r, r_shape); + generate_weights_value(b, b_shape); +} + +ov::Output get_activation_function(const std::string& activation_name, const ov::Output& input) { + if (activation_name == "relu") { + return std::make_shared(input); + } else if (activation_name == "tanh") { + return std::make_shared(input); + } else if (activation_name == "sigmoid") { + return std::make_shared(input); + } + + throw "unsupported activation function"; +} + +ov::Output generate_gate_subgraph(const std::shared_ptr& x, + const std::shared_ptr& h, + const std::vector& w, + const std::vector& r, + const std::vector& b, + size_t input_size, + size_t hidden_size, + const std::string& f_activation) { + // w must be of a shape [input_size, hidden_size] + // r must be of a shape [hidden_size, hidden_size] + // b must be of a shape [hidden_size] + Shape w_shape({input_size, hidden_size}); + Shape r_shape({hidden_size, hidden_size}); + Shape b_shape({hidden_size}); + + auto w_const = op::v0::Constant::create(element::f32, w_shape, w); + auto r_const = op::v0::Constant::create(element::f32, r_shape, r); + auto b_const = op::v0::Constant::create(element::f32, b_shape, b); + + auto x_by_wi = std::make_shared(x, w_const); + auto x_by_wi_biased = std::make_shared(x_by_wi, b_const); + auto h_by_ri = std::make_shared(h, r_const); + auto it = std::make_shared(x_by_wi_biased, h_by_ri); + return get_activation_function(f_activation, it); +} + +ov::Output prepare_weight_fico(const std::vector& f_val, + const std::vector& i_val, + const std::vector& c_val, + const std::vector& o_val, + Shape w_shape) { + auto f = std::make_shared(element::f32, w_shape, f_val); + auto i = std::make_shared(element::f32, w_shape, i_val); + auto c = std::make_shared(element::f32, w_shape, c_val); + auto o = std::make_shared(element::f32, w_shape, o_val); + + auto tr_order = std::make_shared(element::i32, ov::Shape{2}, std::vector{1, 0}); + auto f_tr = std::make_shared(f, tr_order); + auto i_tr = std::make_shared(i, tr_order); + auto c_tr = std::make_shared(c, tr_order); + auto o_tr = std::make_shared(o, tr_order); + + ov::Output w = std::make_shared(ov::OutputVector{f_tr, i_tr, c_tr, o_tr}, 0); + if (const auto& constant = ov::util::constantfold_subgraph(w)) { + w = constant; + } + + return w; +} +} // namespace + +TEST_P(LSTMCellFusionWithSplitWeights, SubgraphFusedToLSTMCell) { + const auto& param = GetParam(); + const std::string& f_activation = std::get<0>(param); + const std::string& g_activation = std::get<1>(param); + const std::string& h_activation = std::get<2>(param); + size_t input_size = std::get<3>(param); + size_t hidden_size = std::get<4>(param); + size_t batch_size = 2; + + // generate weights values + std::vector wi, ri, bi, wf, rf, bf, wo, ro, bo, wc, rc, bc; + generate_weights(wi, ri, bi, input_size, hidden_size); + generate_weights(wf, rf, bf, input_size, hidden_size); + generate_weights(wo, ro, bo, input_size, hidden_size); + generate_weights(wc, rc, bc, input_size, hidden_size); + + { + auto x = std::make_shared(element::f32, Shape{batch_size, input_size}); + auto h = std::make_shared(element::f32, Shape{batch_size, hidden_size}); + auto c = std::make_shared(element::f32, Shape{batch_size, hidden_size}); + + auto it = generate_gate_subgraph(x, h, wi, ri, bi, input_size, hidden_size, f_activation); + auto ft = generate_gate_subgraph(x, h, wf, rf, bf, input_size, hidden_size, f_activation); + auto ot = generate_gate_subgraph(x, h, wo, ro, bo, input_size, hidden_size, f_activation); + auto c1t = generate_gate_subgraph(x, h, wc, rc, bc, input_size, hidden_size, g_activation); + + auto it_mul_c1t = std::make_shared(it, c1t); + auto ft_mul_c = std::make_shared(ft, c); + auto ct = std::make_shared(ft_mul_c, it_mul_c1t); + + auto ct_activated = get_activation_function(h_activation, ct); + auto ht = std::make_shared(ct_activated, ot); + + auto c_neg = std::make_shared(ct); + auto h_abs = std::make_shared(ht); + + model = std::make_shared(NodeVector{h_abs, c_neg}, ParameterVector{x, h, c}); + manager.register_pass(); + } + + { + auto x = std::make_shared(element::f32, Shape{batch_size, input_size}); + auto h = std::make_shared(element::f32, Shape{batch_size, hidden_size}); + auto c = std::make_shared(element::f32, Shape{batch_size, hidden_size}); + + // concatenate weights in fico format + auto w = prepare_weight_fico(wf, wi, wc, wo, ov::Shape{input_size, hidden_size}); + auto r = prepare_weight_fico(rf, ri, rc, ro, ov::Shape{hidden_size, hidden_size}); + + std::vector b_value = bf; + b_value.insert(b_value.end(), bi.begin(), bi.end()); + b_value.insert(b_value.end(), bc.begin(), bc.end()); + b_value.insert(b_value.end(), bo.begin(), bo.end()); + auto b = op::v0::Constant::create(element::f32, Shape{4 * hidden_size}, b_value); + + auto lstm_cell = + std::make_shared(x, + h, + c, + w, + r, + b, + hidden_size, + std::vector{f_activation, g_activation, h_activation}); + + auto c_neg = std::make_shared(lstm_cell->output(1)); + auto h_abs = std::make_shared(lstm_cell->output(0)); + + model_ref = std::make_shared(NodeVector{h_abs, c_neg}, ParameterVector{x, h, c}); + } + + comparator.enable(FunctionsComparator::CmpValues::CONST_VALUES); + comparator.enable(FunctionsComparator::CmpValues::ATTRIBUTES); + comparator.enable(FunctionsComparator::CmpValues::ACCURACY); +} + +INSTANTIATE_TEST_SUITE_P(LSTMCellFusion, + LSTMCellFusionWithSplitWeights, + testing::Combine(testing::Values("sigmoid", "tanh", "relu"), + testing::Values("sigmoid", "relu"), + testing::Values("tanh", "relu"), + testing::Values(2, 3), + testing::Values(3, 4)));