[LPT] Assign + ReadValue transformation (#8690)

* [LPT] Assign + ReadValue transformation

* [LPT] Assign+ReadValue: applied review comments

* fixed some misses

src/common/low_precision_transformations/include/low_precision/assign_and_read_value.hpp

@@ -0,0 +1,27 @@
+// Copyright (C) 2018-2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include <ngraph/ngraph.hpp>
+#include "layer_transformation.hpp"
+
+namespace ngraph {
+namespace pass {
+namespace low_precision {
+
+class LP_TRANSFORMATIONS_API AssignAndReadValueTransformation : public LayerTransformation {
+public:
+    NGRAPH_RTTI_DECLARATION;
+    AssignAndReadValueTransformation(const std::shared_ptr<ngraph::Function> function, const Params& params = Params());
+    bool transform(TransformationContext& context, ngraph::pattern::Matcher& m) override;
+    bool canBeTransformed(const TransformationContext& context, std::shared_ptr<Node> op) const override;
+    bool isPrecisionPreserved(std::shared_ptr<Node> layer) const noexcept override;
+private:
+    std::shared_ptr<ngraph::Function> function;
+};
+
+} // namespace low_precision
+} // namespace pass
+} // namespace ngraph

src/common/low_precision_transformations/include/low_precision/common/fake_quantize_dequantization.hpp

@@ -33,6 +33,7 @@ public:
     bool multiplyHasZeroOrDenormal() const;
     bool isShared() const;
     bool isLowPrecision() const;
+    std::shared_ptr<Node> copyWithNewInput(const std::shared_ptr<Node>& input) const;
 
     static bool checkElementwise(const std::shared_ptr<ngraph::Node>& elementwise);
 

src/common/low_precision_transformations/include/low_precision/fake_quantize.hpp

@@ -22,11 +22,10 @@ public:
 
     static bool checkElementwise(const std::shared_ptr<Node>& eltwise);
 
-private:
-    std::shared_ptr<opset1::FakeQuantize> fuseElementwise(
+    static std::shared_ptr<opset1::FakeQuantize> fuseElementwise(
             TransformationContext& context,
             MatcherPass* matcherPass,
-            const std::shared_ptr<opset1::FakeQuantize>& fakeQuantize) const;
+            const std::shared_ptr<opset1::FakeQuantize>& fakeQuantize);
 };
 
 } // namespace low_precision

src/common/low_precision_transformations/include/low_precision/network_helper.hpp

@@ -41,17 +41,12 @@ public:
 
     static std::vector<Input<Node>> consumer_inputs(std::shared_ptr<Node> node);
 
-    // returns true if at least one child is not FQ
-    static bool notAllChildrensAreFQ(const NodeVector& layer);
-
     // Collect and return a vector with all nodes that consumes any of the `node` output
     static std::vector<std::shared_ptr<Node>> consumers(std::shared_ptr<Node> node);
 
     // return true if op is on a constant path
     static bool isConstantPath(const std::shared_ptr<Node>& op);
 
-    static Shape alignShapeForChannelDim(const Shape& shape, Rank rank);
-
     template <typename OperationType>
     static std::shared_ptr<Node> setOutDataPrecisionForTypeRelaxed(std::shared_ptr<OperationType> operation, const element::Type& precision);
 
@@ -215,87 +210,6 @@ public:
         const std::shared_ptr<Node>& dequantization,
         const std::shared_ptr<Node>& newNode);
 
-    static void replaceAttributeInNodes(
-        std::shared_ptr<ngraph::Function> f,
-        const std::string& name,
-        const ov::Any& newAttribute,
-        const ov::Any& oldAttribute,
-        const std::shared_ptr<ngraph::Node>& initialNode) {
-        std::set<std::shared_ptr<Node>> visited;
-        std::deque<std::shared_ptr<Node>> nodes;
-        nodes.emplace_back(initialNode);
-
-        while (!nodes.empty()) {
-            auto node = nodes.front();
-            nodes.pop_front();
-
-            if (visited.count(node) || ov::is_type<op::Constant>(node)) {
-                continue;
-            }
-
-            visited.insert(node);
-
-            bool handleConnectedNodes = false;
-            if (NetworkHelper::isPrecisionPreserved(node) || ov::is_type<opset1::FakeQuantize>(node)) {
-                auto& rt = node->get_rt_info();
-
-                if (node == initialNode) {
-                    rt[name] = newAttribute;
-                    handleConnectedNodes = true;
-                } else {
-                    auto it = rt.find(name);
-                    if (it != rt.end()) {
-                        const auto currentAttribute = it->second;
-                        if (oldAttribute == currentAttribute) {
-                            rt[name] = newAttribute;
-                        }
-                        handleConnectedNodes = true;
-                    }
-                }
-            }
-
-            if (!handleConnectedNodes) {
-                continue;
-            }
-
-            if (!ov::is_type<opset1::FakeQuantize>(node)) {
-                for (size_t index = 0ul; index < node->get_input_size(); ++index) {
-                    auto getInput = [](const std::shared_ptr<ngraph::Node>& node, const size_t index) {
-                        const auto dequantization = NetworkHelper::getDequantization(node, index);
-                        if (!dequantization.empty() &&
-                            (ov::is_type<opset1::Convert>(dequantization.data.get_node())) &&
-                            ov::is_type<opset1::FakeQuantize>(dequantization.data.get_node()->get_input_node_ptr(0))) {
-                            const auto input = dequantization.data.get_node()->input(0);
-                            return input;
-                        }
-                        return node->input(index);
-                    };
-
-                    const auto& input = getInput(node, index);
-                    const auto& input_node = input.get_source_output().get_node_shared_ptr();
-
-                    //const auto& input_node = input.get_source_output().get_node_shared_ptr();
-                    if (visited.count(input_node) || ov::is_type<op::Constant>(input_node)) {
-                        continue;
-                    }
-
-                    nodes.push_front(input_node);
-                }
-            }
-
-            for (auto& output : node->outputs()) {
-                for (auto& input_value : output.get_target_inputs()) {
-                    const auto& output_node = input_value.get_node()->shared_from_this();
-                    if (visited.count(output_node) || ov::is_type<op::Constant>(output_node)) {
-                        continue;
-                    }
-
-                    nodes.push_front(output_node);
-                }
-            }
-        }
-    }
-
     template <typename SharedAttribute>
     static void reassign(
         const std::shared_ptr<typename SharedAttribute::SharedValueAttribute::SharedValue>& sharedValue,
@@ -370,14 +284,6 @@ std::shared_ptr<Node> make_op_pattern(const ngraph::NodeVector& args) {
     return std::make_shared<ngraph::pattern::op::Any>(element::undefined, PartialShape{}, [](std::shared_ptr<Node> n) {return !!ov::as_type_ptr<T>(n); }, args);
 }
 
-template <typename T>
-std::shared_ptr<Node> make_op_label() {
-    return std::make_shared<ngraph::pattern::op::Label>(
-            element::undefined,
-            PartialShape{},
-            [](std::shared_ptr<Node> n) {return !!ov::as_type_ptr<T>(n); });
-}
-
 template <typename T, typename... Args>
 std::shared_ptr<Node> fold(Args&&... args) {
     auto node = std::make_shared<T>(std::forward<Args>(args)...);

src/common/low_precision_transformations/src/assign_and_read_value.cpp

@@ -0,0 +1,132 @@
+// Copyright (C) 2021 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "low_precision/assign_and_read_value.hpp"
+#include <ngraph/ngraph.hpp>
+
+#include <ngraph/pattern/op/wrap_type.hpp>
+#include "low_precision/network_helper.hpp"
+#include <ngraph/opsets/opset6.hpp>
+#include <ngraph/pattern/op/or.hpp>
+#include <openvino/op/util/assign_base.hpp>
+#include "low_precision/fake_quantize.hpp"
+
+namespace ngraph {
+namespace pass {
+namespace low_precision {
+
+NGRAPH_RTTI_DEFINITION(ngraph::pass::low_precision::AssignAndReadValueTransformation, "AssignAndReadValueTransformation", 0);
+
+AssignAndReadValueTransformation::AssignAndReadValueTransformation(const std::shared_ptr<ngraph::Function> function, const Params& params) :
+    LayerTransformation(params), function(function) {
+    auto assign3 = pattern::wrap_type<opset3::Assign>({ pattern::wrap_type<opset1::Multiply>() });
+    auto assign6 = pattern::wrap_type<opset6::Assign>({ pattern::wrap_type<opset1::Multiply>() });
+
+    ngraph::graph_rewrite_callback callback = [=](pattern::Matcher& m) {
+        const auto& opsMap = m.get_pattern_value_map();
+        auto op = m.get_match_root();
+        auto assignIt = opsMap.find(assign3);
+        if (assignIt == opsMap.end()) {
+            assignIt = opsMap.find(assign6);
+        }
+        const auto assign = assignIt->second.get_node_shared_ptr();
+        // check that we have ReadValue as the first dependency
+        if (assign->get_control_dependencies().empty()) {
+            return false;
+        }
+
+        if (transformation_callback(op)) {
+            return false;
+        }
+        return transform(*context, m);
+    };
+
+    auto m = std::make_shared<ngraph::pattern::Matcher>(
+        std::make_shared<pattern::op::Or>(OutputVector{ assign3, assign6 }),
+        "AssignAndReadValueTransformation");
+    this->register_matcher(m, callback);
+}
+
+bool AssignAndReadValueTransformation::transform(TransformationContext& context, ngraph::pattern::Matcher& m) {
+    if (!canBeTransformed(context, m.get_match_root())) {
+        return false;
+    }
+
+    const auto oldAssign = m.get_match_root();
+    const auto readValue = oldAssign->get_control_dependencies()[0];
+    oldAssign->remove_control_dependency(readValue);
+
+    const auto assign = NetworkHelper::separateInStandaloneBranch(oldAssign);
+    const auto dequantization = NetworkHelper::getDequantization(assign);
+
+    auto oldVar = ov::as_type_ptr<op::ReadValueBase>(readValue)->get_variable();
+    auto variableInfo = oldVar->get_info();
+    // set new precision for oldVar to update precision in newReadValue
+    oldVar->update({variableInfo.data_shape, dequantization.data.get_element_type(), variableInfo.variable_id});
+    // transform ReadValue part
+    const auto newConstant = foldConvert(readValue->get_input_node_shared_ptr(0), dequantization.data.get_element_type());
+    const auto newReadValue = readValue->copy_with_new_inputs({newConstant});
+    const auto newDequantization = dequantization.copyWithNewInput(newReadValue);
+    replace_node(readValue, newDequantization);
+
+    // transform Assign part
+
+    const auto newAssign = assign->copy_with_new_inputs({dequantization.data});
+    function->remove_sink(as_type_ptr<op::Sink>(oldAssign));
+    function->add_sinks({as_type_ptr<op::Sink>(newAssign)});
+
+    NetworkHelper::copyInfo(assign, newAssign);
+    replace_node(assign, newAssign);
+    newAssign->add_control_dependency(newReadValue);
+
+    // fuse dequantization multiply with FQ after ReadValue if possible
+    const auto nextLayers = newDequantization->get_output_target_inputs(0);
+    if (nextLayers.size() > 1) {
+        return true;
+    }
+    const auto fakeQuantize = as_type_ptr<opset1::FakeQuantize>(nextLayers.begin()->get_node()->shared_from_this());
+
+    if (fakeQuantize == nullptr) {
+        return true;
+    }
+    auto fakeQuantizeInputs = fakeQuantize->input_values();
+
+    const auto inputLow = as_type_ptr<opset1::Constant>(fakeQuantizeInputs[1].get_node_shared_ptr());
+    const auto inputHigh = as_type_ptr<opset1::Constant>(fakeQuantizeInputs[2].get_node_shared_ptr());
+
+    if (inputLow == nullptr || inputHigh == nullptr) {
+        return true;
+    }
+
+    FakeQuantizeTransformation::fuseElementwise(context, this, fakeQuantize);
+
+    return true;
+}
+
+bool AssignAndReadValueTransformation::canBeTransformed(const TransformationContext& context, std::shared_ptr<Node> op) const {
+    if (!LayerTransformation::canBeTransformed(context, op)) {
+        return false;
+    }
+
+    const auto readValue = std::dynamic_pointer_cast<op::ReadValueBase>(op->get_control_dependencies()[0]);
+    if (!readValue) {
+        return false;
+    }
+
+    // TODO: remove this limitation and change the transformation when this constant will be accepted to be non-zero
+    if (!NetworkHelper::isZeroConst(readValue->get_input_node_shared_ptr(0))) {
+        return false;
+    }
+
+    const auto dequantization = NetworkHelper::getDequantization(op);
+    return dequantization.subtract == nullptr && dequantization.multiply != nullptr;
+}
+
+bool AssignAndReadValueTransformation::isPrecisionPreserved(std::shared_ptr<Node> layer) const noexcept {
+    return false;
+}
+
+} // namespace low_precision
+} // namespace pass
+} // namespace ngraph

src/common/low_precision_transformations/src/fake_quantize.cpp

@@ -129,17 +129,15 @@ bool FakeQuantizeTransformation::checkElementwise(const std::shared_ptr<Node>& e
 std::shared_ptr<opset1::FakeQuantize> FakeQuantizeTransformation::fuseElementwise(
     TransformationContext& context,
     MatcherPass* matcherPass,
-    const std::shared_ptr<opset1::FakeQuantize>& fakeQuantize) const {
+    const std::shared_ptr<opset1::FakeQuantize>& fakeQuantize) {
     const std::shared_ptr<Node> eltwise = fakeQuantize->get_input_node_shared_ptr(0);
 
-    std::shared_ptr<Node> inputLowConst_f32 = foldConvert(fakeQuantize->input_value(1), deqPrecision);
-    std::shared_ptr<Node> inputHighConst_f32 = foldConvert(fakeQuantize->input_value(2), deqPrecision);
+    std::shared_ptr<Node> inputLowConst_f32 = foldConvert(fakeQuantize->input_value(1), element::f32);
+    std::shared_ptr<Node> inputHighConst_f32 = foldConvert(fakeQuantize->input_value(2), element::f32);
 
     std::shared_ptr<opset1::Constant> constant = fq::getConstant(eltwise);
     if (ov::is_type<opset1::Multiply>(eltwise) && checkElementwise(eltwise)) {
-        const auto value = constant->get_output_element_type(0) == deqPrecision ?
-            constant :
-            foldConvert(constant, deqPrecision);
+        const auto value = foldConvert(constant, element::f32);
 
         const auto valueVec = ov::as_type_ptr<opset1::Constant>(value)->cast_vector<float>();
 
@@ -159,9 +157,7 @@ std::shared_ptr<opset1::FakeQuantize> FakeQuantizeTransformation::fuseElementwis
         inputLowConst_f32 = fq::updateShape(inputLowConst_f32, fakeQuantize->get_output_partial_shape(0));
         inputHighConst_f32 =  fq::updateShape(inputHighConst_f32, fakeQuantize->get_output_partial_shape(0));
     } else if (ov::is_type<opset1::Subtract>(eltwise) && checkElementwise(eltwise)) {
-        const auto value = constant->get_output_element_type(0) == deqPrecision ?
-            constant :
-            foldConvert(constant, deqPrecision);
+        const auto value = foldConvert(constant, element::f32);
 
         inputLowConst_f32 = fq::updateShape(fold<opset1::Add>(inputLowConst_f32, value), fakeQuantize->get_output_partial_shape(0));
         inputHighConst_f32 = fq::updateShape(fold<opset1::Add>(inputHighConst_f32, value), fakeQuantize->get_output_partial_shape(0));
@@ -173,9 +169,7 @@ std::shared_ptr<opset1::FakeQuantize> FakeQuantizeTransformation::fuseElementwis
             return nullptr;
         }
 
-        const auto value = constant->get_output_element_type(0) == deqPrecision ?
-            constant :
-            foldConvert(constant, deqPrecision);
+        const auto value = foldConvert(constant, element::f32);
 
         inputLowConst_f32 = fq::updateShape(fold<opset1::Subtract>(inputLowConst_f32, value), fakeQuantize->get_output_partial_shape(0));
         inputHighConst_f32 = fq::updateShape(fold<opset1::Subtract>(inputHighConst_f32, value), fakeQuantize->get_output_partial_shape(0));
@@ -196,8 +190,8 @@ std::shared_ptr<opset1::FakeQuantize> FakeQuantizeTransformation::fuseElementwis
         data->output(outputIdx),
         inputLowConst_f32,
         inputHighConst_f32,
-        foldConvert(fakeQuantize->input_value(3), deqPrecision),
-        foldConvert(fakeQuantize->input_value(4), deqPrecision) }));
+        foldConvert(fakeQuantize->input_value(3), element::f32),
+        foldConvert(fakeQuantize->input_value(4), element::f32) }));
 
     matcherPass->register_new_node(newFakeQuantize);
 

src/common/low_precision_transformations/src/fake_quantize_dequantization.cpp

@@ -155,6 +155,26 @@ bool FakeQuantizeDequantization::checkElementwise(const std::shared_ptr<ngraph::
     return true;
 }
 
+std::shared_ptr<Node> FakeQuantizeDequantization::copyWithNewInput(const std::shared_ptr<Node>& input) const {
+    auto lastNode = input;
+    if (convert) {
+        lastNode = convert->copy_with_new_inputs({lastNode});
+    }
+    if (subtract) {
+        std::shared_ptr<Node> input1 = nullptr;
+        if (subtractConvert) {
+            input1 = subtractConvert;
+        } else {
+            input1 = subtractConstant;
+        }
+        lastNode = subtract->copy_with_new_inputs({lastNode, input1});
+    }
+    if (multiply) {
+        lastNode = multiply->copy_with_new_inputs({lastNode, multiplyConstant});
+    }
+    return lastNode;
+}
+
 int FakeQuantizeDequantization::fillDequantizationParams(
     const std::shared_ptr<ngraph::Node>& elementwise,
     std::shared_ptr<ngraph::opset1::Convert>& convert,

src/common/low_precision_transformations/src/low_precision.cpp

@@ -38,6 +38,7 @@
 
 // general transformations
 #include "low_precision/add.hpp"
+#include "low_precision/assign_and_read_value.hpp"
 #include "low_precision/avg_pool.hpp"
 #include "low_precision/clamp.hpp"
 #include "low_precision/convolution.hpp"
@@ -207,6 +208,7 @@ bool ngraph::pass::low_precision::LowPrecision::run_on_model(const std::shared_p
 
     std::shared_ptr<ngraph::pass::GraphRewrite> common = manager.register_pass<ngraph::pass::GraphRewrite>();
     common->add_matcher<ngraph::pass::low_precision::AddTransformation>(params);
+    common->add_matcher<ngraph::pass::low_precision::AssignAndReadValueTransformation>(f, params);
     common->add_matcher<ngraph::pass::low_precision::AvgPoolTransformation>(params);
     common->add_matcher<ngraph::pass::low_precision::ClampTransformation>(params);
     common->add_matcher<ngraph::pass::low_precision::ConcatTransformation>(params);

src/common/low_precision_transformations/src/network_helper.cpp

@@ -37,17 +37,6 @@ bool NetworkHelper::is_castable_to_one_of(NodeTypeInfo type, const std::unordere
     return false;
 }
 
-bool NetworkHelper::notAllChildrensAreFQ(const NodeVector& childrens) {
-    // NOTE: This check was added for models that don't have FQ after AvgPool
-    //       They will have transparent precision as it was in old LPT.
-    for (const auto& child : childrens) {
-        if (!ov::is_type<opset1::FakeQuantize>(child)) {
-            return true;
-        }
-    }
-    return false;
-}
-
 // Collect and return a vector with all nodes that consumes any of the `node` output
 std::vector<Input<Node>> NetworkHelper::consumer_inputs(std::shared_ptr<Node> node) {
     std::vector<Input<Node>> result;
@@ -199,15 +188,6 @@ size_t NetworkHelper::getGroupsCount(std::shared_ptr<Node> layer) {
     }
 }
 
-// Assumin tensor in NC... layout, append necessary number of 1s to shape to align it to a give rank
-Shape NetworkHelper::alignShapeForChannelDim(const Shape& shape, Rank rank) {
-    assert(shape.size() == 1);
-    assert(rank.is_static());
-    Shape result = shape;
-    result.resize(rank.get_length() - 1, 1);
-    return result;
-}
-
 void NetworkHelper::removeLayer(std::shared_ptr<Node> layer) {
     ngraph::replace_output_update_name(layer->output(0), layer->input_value(0));
 }
@@ -1359,9 +1339,12 @@ FakeQuantizeDequantization NetworkHelper::getDequantization(const std::shared_pt
 
     const std::shared_ptr<opset1::Convert> convert = ov::as_type_ptr<opset1::Convert>(dataNode.get_node_shared_ptr());
     if (convert != nullptr) {
-        if ((convert->input(0).get_element_type() != element::i8) && (convert->input(0).get_element_type() != element::u8) &&
-            (convert->input(0).get_element_type() != element::i4) && (convert->input(0).get_element_type() != element::u4) &&
-            (convert->output(0).get_element_type() != element::f32)) {
+        auto defaultPrecisions = LayerTransformation::getDefaultPrecisions();
+        auto el_type = convert->input(0).get_element_type();
+        auto foundIt = std::find(defaultPrecisions.begin(), defaultPrecisions.end(), el_type);
+        if (foundIt == defaultPrecisions.end() &&
+            el_type != element::i4  && el_type != element::u4 &&
+            el_type != element::f32 && el_type != element::f16) {
             return FakeQuantizeDequantization(dataNode, nullptr, subtract, subtractConvert, subtractConstant, multiply, multiplyConstant);
         }
         dataNode = convert->get_input_source_output(0);