From af8c018824aeae958e26eca757909ecfdcb82e03 Mon Sep 17 00:00:00 2001
From: Sapozhnikov Andrey <andrey.sapozhnikov@intel.com>
Date: Fri, 16 Jul 2021 16:23:35 +0300
Subject: [PATCH] [GNA] KSOFunction test fix

---
 .../src/gna_plugin/gna_plugin.cpp             |  5 +++--
 .../gna_plugin/optimizer/gna_pass_manager.cpp | 20 +++++++++----------
 .../skip_tests_config.cpp                     |  2 --
 .../eltwise_reshape_activation.cpp            |  3 ++-
 4 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/inference-engine/src/gna_plugin/gna_plugin.cpp b/inference-engine/src/gna_plugin/gna_plugin.cpp
index c40b97209e1075..4d565ca0c2d975 100644
--- a/inference-engine/src/gna_plugin/gna_plugin.cpp
+++ b/inference-engine/src/gna_plugin/gna_plugin.cpp
@@ -752,12 +752,14 @@ void GNAPlugin::LoadNetwork(CNNNetwork & _network) {
         passes->registerPass<FuseFQIntoWeightsPass>();
         passes->registerPass<MoveFakeQuantizeLayerIntoQuantParamsPass>();
 
+        passes->registerPass<SubstituteScaleShiftBroadCastPass>();
+        passes->registerPass<BroadcastConstPass>();
+
         passes->registerPass<TransposeWeightsFromNCHWToNHWCPass>();
 
         passes->registerPass<SubstitutePReluPass>();
         passes->registerPass<SubstituteSoftSignPass>();
 
-        passes->registerPass<BroadcastConstPass>();
         passes->registerPass<ReorderMaxPoolPass>();
         passes->registerPass<EltwiseSplitOverChannelsPass>();
         passes->registerPass<InsertSplitAligningFilterPass>();
@@ -775,7 +777,6 @@ void GNAPlugin::LoadNetwork(CNNNetwork & _network) {
 #if GNA_LIB_VER == 2
         passes->registerPass<ForbidActivationFusingPass>();
 #endif
-        passes->registerPass<SubstituteScaleShiftBroadCastPass>();
         passes->registerPass<FuseMultipleIdentitiesPass>();
         passIdx = passes->run(passIdx);
     };
diff --git a/inference-engine/src/gna_plugin/optimizer/gna_pass_manager.cpp b/inference-engine/src/gna_plugin/optimizer/gna_pass_manager.cpp
index ae731465025e05..a21affeae88ffe 100644
--- a/inference-engine/src/gna_plugin/optimizer/gna_pass_manager.cpp
+++ b/inference-engine/src/gna_plugin/optimizer/gna_pass_manager.cpp
@@ -1530,16 +1530,7 @@ void SubstituteScaleShiftBroadCastPass::run() {
             continue;
         }
 
-        // only 3d scaleshift supported where number of c is arbitrary
-        auto lastD = reshape_batch ? dataDims[1] : dataDims.back();
-        if (lastD != weightsElements) {
-            THROW_GNA_EXCEPTION << "Unsupported layer: " << l->name
-                                << " should have last dim(" << lastD << ") equal to weights(" << weightsElements << ") length";
-        }
-        if (dataDims.size() == 2) {
-            THROW_GNA_EXCEPTION << "For layer: " << l->name
-                                << " weights size(" << weightsElements<< ") invalid: should match input size of(" << lastD << ")";
-        }
+        // TODO: add broadcasting rules checks
 
         gnalog() << "Substitution ScaleShift broadcast for layer: " << l->name << "\n";
         if (nElements % scaleShift->_weights->size()) {
@@ -2237,6 +2228,15 @@ void TransposeWeightsFromNCHWToNHWCPass::run() {
                 }
                 auto weightable = dynamic_cast<WeightableLayer*>(l.get());
                 IE_ASSERT(weightable != nullptr);
+                size_t totalElements = 0;
+                auto totalWeights = weightable->_weights->size();
+                for (auto && transpositionInfoPart : transpositionInfo) {
+                        totalElements += transpositionInfoPart.num_transpose_rows * transpositionInfoPart.num_transpose_columns;
+                }
+                if (totalElements != totalWeights) {
+                        THROW_GNA_EXCEPTION << l->name << " weights columns from transposition info (" << totalElements
+                                            << ") don't match input dimensions (" << totalWeights << ")";
+                    }
                 ConvertTensorFromNCHWToNHWC(weightable->precision.size(), 1, weightable->_weights->size(),
                     weightable->_weights->cbuffer().as<uint8_t*>(), true, transpositionInfo);
                 if (weightable->_biases) {
diff --git a/inference-engine/tests/functional/plugin/gna/shared_tests_instances/skip_tests_config.cpp b/inference-engine/tests/functional/plugin/gna/shared_tests_instances/skip_tests_config.cpp
index 589ec41c83dd8d..9d9436db568c5d 100644
--- a/inference-engine/tests/functional/plugin/gna/shared_tests_instances/skip_tests_config.cpp
+++ b/inference-engine/tests/functional/plugin/gna/shared_tests_instances/skip_tests_config.cpp
@@ -44,8 +44,6 @@ std::vector<std::string> disabledTestPatterns() {
         R"(.*ConstantResultSubgraphTest.*inPrc=(U8|I8|I32|U64|I64|BOOL).*)",
         // TODO: Issue 51528
         R"(.*CachingSupport.*_(u8|i16)_.*)",
-        // TODO: Issue 51525
-        R"(.*CachingSupport.*KSOFunction.*)",
         // TODO: Issue 57363 (Param -> Result subgraphs)
         R"(.*smoke_MemoryTest.*LOW_LATENCY.*iteration_count=1_.*)",
         // TODO: Issue 57368 (accuracy)
diff --git a/inference-engine/tests/functional/plugin/gna/shared_tests_instances/subgraph_tests/eltwise_reshape_activation.cpp b/inference-engine/tests/functional/plugin/gna/shared_tests_instances/subgraph_tests/eltwise_reshape_activation.cpp
index 8e7af93c892d00..3e97b5b16f6b3d 100644
--- a/inference-engine/tests/functional/plugin/gna/shared_tests_instances/subgraph_tests/eltwise_reshape_activation.cpp
+++ b/inference-engine/tests/functional/plugin/gna/shared_tests_instances/subgraph_tests/eltwise_reshape_activation.cpp
@@ -13,7 +13,8 @@ const std::vector<std::vector<std::vector<size_t>>> shapes = {
     {{1, 64}, {64, 1}},
     {{8, 256}, {16, 128}},
     {{6, 384}, {18, 128}},
-    {{8, 2048}, {32, 512}}
+    {{8, 2048}, {32, 512}},
+    {{2, 4, 64, 64}, {1, 8, 64, 64}}
 };
 
 const std::vector<InferenceEngine::Precision> netPrecisions = {