From be06f10f79ccbdb9b110342c89c1d70238cc141c Mon Sep 17 00:00:00 2001
From: Mike Kelly <mike.kelly@arm.com>
Date: Mon, 17 Jul 2023 17:49:55 +0100
Subject: IVGCVSW-7891 Failure in Nightly tests

 * Added check to ensure that Reshapes are not removed on Neon if they are
   before or after a SplitterLayer and have more than 4 dimensions.
 * Moved NCHW check to a function to reduce clutter.

Signed-off-by: Mike Kelly <mike.kelly@arm.com>
Change-Id: I45d97634484e8dc0ca7675c23481caf84eb3fe90
---
 src/backends/backendsCommon/SubgraphUtils.hpp | 41 +++++++++++++++++++++++++++
 src/backends/neon/NeonBackend.cpp             | 21 ++++----------
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/src/backends/backendsCommon/SubgraphUtils.hpp b/src/backends/backendsCommon/SubgraphUtils.hpp
index ade4b63976..823da76f29 100644
--- a/src/backends/backendsCommon/SubgraphUtils.hpp
+++ b/src/backends/backendsCommon/SubgraphUtils.hpp
@@ -199,6 +199,47 @@ LayerType* FoldPadLayer(OptimizationViews& optimizationViews,
     return replacementLayer;
 }
 
+/// Checks if the Layer is connected to any Layer that has an NCHW layout.
+inline bool ConnectedToLayerWithNCHW(Layer* baseLayer)
+{
+    Layer& parentLayer = baseLayer->GetInputSlot(0).GetConnectedOutputSlot()->GetOwningLayer();
+
+    if (IsNCHW(parentLayer))
+    {
+        return true;
+    }
+    for (unsigned int i = 0; i < baseLayer->GetOutputSlot(0).GetNumConnections(); ++i)
+    {
+        Layer& nextLayer = baseLayer->GetOutputSlot(0).GetConnection(i)->GetOwningLayer();
+        if (IsNCHW(nextLayer))
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
+/// Checks if the Layer is connected to a Splitter Layer through a Tensor that has more than 4 dimensions.
+inline bool ConnectedToSplitterWithMoreThan4Dims(Layer* baseLayer)
+{
+    Layer& parentLayer = baseLayer->GetInputSlot(0).GetConnectedOutputSlot()->GetOwningLayer();
+    TensorInfo parentTensorInfo = baseLayer->GetInputSlot(0).GetConnectedOutputSlot()->GetTensorInfo();
+    if (parentTensorInfo.GetNumDimensions() > 4 && parentLayer.GetType() == LayerType::Splitter)
+    {
+        return true;
+    }
+    for (unsigned int i = 0; i < baseLayer->GetOutputSlot(0).GetNumConnections(); ++i)
+    {
+        Layer& nextLayer = baseLayer->GetOutputSlot(0).GetConnection(i)->GetOwningLayer();
+        TensorInfo nextTensorInfo = baseLayer->GetOutputSlot(0).GetConnection(i)->GetTensorInfo();
+        if (nextTensorInfo.GetNumDimensions() > 4 && nextLayer.GetType() == LayerType::Splitter)
+        {
+            return true;
+        }
+    }
+    return false;
+}
+
 inline void RemoveReshapeLayer(ReshapeLayer* baseLayer,
                                std::map<LayerGuid, Layer*>& untouched,
                                OptimizationViews& optimizationViews)
diff --git a/src/backends/neon/NeonBackend.cpp b/src/backends/neon/NeonBackend.cpp
index 098b1ff109..60e25672ae 100644
--- a/src/backends/neon/NeonBackend.cpp
+++ b/src/backends/neon/NeonBackend.cpp
@@ -510,26 +510,15 @@ OptimizationViews NeonBackend::OptimizeSubgraphView(const SubgraphView& subgraph
         if (base.GetType() == LayerType::Reshape)
         {
             ReshapeLayer* baseLayer = PolymorphicDowncast<ReshapeLayer*>(&base);
-            Layer& parentLayer = baseLayer->GetInputSlot(0).GetConnectedOutputSlot()->GetOwningLayer();
 
-            // Cannot currently remove the Reshape if it's connected to any layer that has an NCHW layout
-            if (IsNCHW(parentLayer))
+            // Cannot remove a Reshape if it's connected to any layer that has an NCHW layout
+            if (ConnectedToLayerWithNCHW(baseLayer))
             {
                 continue;
             }
-            bool isNCHW = false;
-
-            for (unsigned int i = 0; i < baseLayer->GetOutputSlot(0).GetNumConnections(); ++i)
-            {
-                Layer& nextLayer = baseLayer->GetOutputSlot(0).GetConnection(i)->GetOwningLayer();
-
-                if (IsNCHW(nextLayer))
-                {
-                    isNCHW = true;
-                    break;
-                }
-            }
-            if (isNCHW)
+            // Cannot remove a Reshape if it's connected to a SplitterLayer through a Tensor that has more than
+            // 4 dimensions
+            if (ConnectedToSplitterWithMoreThan4Dims(baseLayer))
             {
                 continue;
             }
-- 
cgit v1.2.1