From 249d57f37eab12a656612c8d9815e7740c6e6c4f Mon Sep 17 00:00:00 2001
From: Vladislav Golubev <vladislav.golubev@intel.com>
Date: Fri, 17 Mar 2023 11:43:32 +0100
Subject: [PATCH] [Transformations] CompressQuantizeWeights: fp16 weights
 support (#16323)

* [Transformations] CompressQuantizeWeights: fp16 weights support

* Code style fix

* Code style fix
---
 .../src/compress_quantize_weigths.cpp         | 15 +++++++++++----
 .../tests/utils/compress_quantize_weights.cpp | 19 ++++++++++++++-----
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/common/offline_transformations/src/compress_quantize_weigths.cpp b/src/common/offline_transformations/src/compress_quantize_weigths.cpp
index 2718c1462b8..819ca0752ee 100644
--- a/src/common/offline_transformations/src/compress_quantize_weigths.cpp
+++ b/src/common/offline_transformations/src/compress_quantize_weigths.cpp
@@ -4,6 +4,7 @@
 
 #include <compress_quantize_weights.hpp>
 #include <ngraph/opsets/opset8.hpp>
+#include <ngraph/pattern/op/or.hpp>
 #include <ngraph/pattern/op/wrap_type.hpp>
 #include <ngraph/rt_info.hpp>
 #include <ngraph/validation_util.hpp>
@@ -36,7 +37,10 @@ static bool has_dequantization_subgraph(const std::shared_ptr<ngraph::Node>& fir
 }
 
 ngraph::pass::CompressQuantizeWeights::CompressQuantizeWeights() {
-    auto weights_pattern = pattern::wrap_type<opset8::Constant>();
+    auto weights_const_pattern = pattern::wrap_type<opset8::Constant>();
+    auto weigths_convert_pattern = pattern::wrap_type<opset8::Convert>({weights_const_pattern});
+    OutputVector weights_options{weights_const_pattern, weigths_convert_pattern};
+    auto weights_pattern = std::make_shared<pattern::op::Or>(weights_options);
     auto input_low_pattern = pattern::wrap_type<opset8::Constant>();
     auto input_high_pattern = pattern::wrap_type<opset8::Constant>();
     auto output_low_pattern = pattern::wrap_type<opset8::Constant>();
@@ -93,11 +97,14 @@ ngraph::pass::CompressQuantizeWeights::CompressQuantizeWeights() {
             auto new_output_low = op::Constant::create(input_type, Shape{}, {-static_cast<float>(levels / 2)});
             auto new_output_high =
                 std::make_shared<opset8::Add>(new_output_low, op::Constant::create(input_type, Shape{}, {levels - 1}));
-            const auto& weights = pattern_value_map.at(weights_pattern);
+            const auto& weights_const = pattern_value_map.at(weights_const_pattern);
             const auto& input_low = pattern_value_map.at(input_low_pattern);
             const auto& input_high = pattern_value_map.at(input_high_pattern);
+            const auto& fq_data_input = pattern_value_map.count(weigths_convert_pattern)
+                                            ? pattern_value_map.at(weigths_convert_pattern)
+                                            : weights_const;
             auto quantize =
-                fq->clone_with_new_inputs({weights, input_low, input_high, new_output_low, new_output_high});
+                fq->clone_with_new_inputs({fq_data_input, input_low, input_high, new_output_low, new_output_high});
             // Convert quantized weights to low precision type
             std::shared_ptr<Node> new_weights = std::make_shared<opset8::Convert>(quantize, quantized_type);
             // Constant fold quantized weights
@@ -106,7 +113,7 @@ ngraph::pass::CompressQuantizeWeights::CompressQuantizeWeights() {
             } else {
                 return false;
             }
-            new_weights->set_friendly_name(weights.get_node()->get_friendly_name());
+            new_weights->set_friendly_name(weights_const.get_node()->get_friendly_name());
 
             /*
                Dequantize part is performed by Convert(from low to high precision)->Subtract->Multiply subgraph.
diff --git a/src/common/transformations/tests/utils/compress_quantize_weights.cpp b/src/common/transformations/tests/utils/compress_quantize_weights.cpp
index e4ac431fa38..c71a9754a2f 100644
--- a/src/common/transformations/tests/utils/compress_quantize_weights.cpp
+++ b/src/common/transformations/tests/utils/compress_quantize_weights.cpp
@@ -31,13 +31,18 @@ struct CompressQuantizeWeightsParams {
     float zero_point_val;
 };
 
-class CompressQuantizeWeightsTests : public testing::WithParamInterface<CompressQuantizeWeightsParams>,
-                                     public TransformationTestsF {
+class CompressQuantizeWeightsTests
+    : public testing::WithParamInterface<std::tuple<CompressQuantizeWeightsParams, element::Type>>,
+      public TransformationTestsF {
     void SetUp() override {
         TransformationTestsF::SetUp();
-        auto param = GetParam();
+        CompressQuantizeWeightsParams param;
+        ov::element::Type data_prc;
+        std::tie(param, data_prc) = GetParam();
         {
-            auto data = opset8::Constant::create(element::f32, param.shape, param.weights);
+            std::shared_ptr<Node> data = opset8::Constant::create(data_prc, param.shape, param.weights);
+            if (data_prc == element::f16)
+                data = std::make_shared<opset8::Convert>(data, element::f32);
             auto input_low = opset8::Constant::create(element::f32, Shape{}, {param.in_low});
             auto input_high = opset8::Constant::create(element::f32, Shape{}, {param.in_high});
             auto output_low = opset8::Constant::create(element::f32, Shape{}, {param.out_low});
@@ -116,7 +121,11 @@ static std::vector<CompressQuantizeWeightsParams> params = {
      -64.25f},
 };
 
-INSTANTIATE_TEST_SUITE_P(TransformationTests, CompressQuantizeWeightsTests, ::testing::ValuesIn(params));
+static element::TypeVector data_precisions = {element::f32, element::f16};
+
+INSTANTIATE_TEST_SUITE_P(TransformationTests,
+                         CompressQuantizeWeightsTests,
+                         ::testing::Combine(::testing::ValuesIn(params), ::testing::ValuesIn(data_precisions)));
 
 TEST_F(TransformationTestsF, CompressQuantizeWeightsWithDequantizationSubgraph) {
     {