diff options
Diffstat (limited to 'src/backends/neon/NeonBackend.cpp')
| -rw-r--r-- | src/backends/neon/NeonBackend.cpp | 21 |
1 files changed, 5 insertions, 16 deletions
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; } |