1 // 2 // Copyright © 2017 Arm Ltd. All rights reserved. 3 // SPDX-License-Identifier: MIT 4 // 5 6 #include <TestUtils.hpp> 7 8 #include <Optimizer.hpp> 9 10 #include <doctest/doctest.h> 11 12 TEST_SUITE("Optimizer") 13 { 14 using namespace armnn::optimizations; 15 16 TEST_CASE("OptimizeConsecutiveReshapesTest") 17 { 18 armnn::Graph graph; 19 20 const armnn::TensorInfo info0({ 1, 2, 3, 5 }, armnn::DataType::Float32); 21 22 auto output = graph.AddLayer<armnn::OutputLayer>(0, "output"); 23 auto input = graph.InsertNewLayer<armnn::InputLayer>(output->GetInputSlot(0), 0, "input"); 24 25 input->GetOutputHandler().SetTensorInfo(info0); 26 27 { 28 // Inserts two reshapes. 29 const armnn::TensorInfo info1({ 1, 30, 1, 1 }, armnn::DataType::Float32); 30 const armnn::TensorInfo info2({ 1, 2, 1, 15 }, armnn::DataType::Float32); 31 32 std::string reshape1Name = "reshape1"; 33 std::string reshape2Name = "reshape2"; 34 35 auto reshape1 = graph.InsertNewLayer<armnn::ReshapeLayer>( 36 output->GetInputSlot(0), armnn::ReshapeDescriptor{ info1.GetShape() }, reshape1Name.c_str()); 37 auto reshape2 = graph.InsertNewLayer<armnn::ReshapeLayer>( 38 output->GetInputSlot(0), armnn::ReshapeDescriptor{ info2.GetShape() }, reshape2Name.c_str()); 39 40 reshape1->GetOutputHandler().SetTensorInfo(info1); 41 reshape2->GetOutputHandler().SetTensorInfo(info2); 42 43 CHECK(CheckSequence(graph.cbegin(), graph.cend(), &IsLayerOfType<armnn::InputLayer>, 44 &IsLayerOfType<armnn::ReshapeLayer>, &IsLayerOfType<armnn::ReshapeLayer>, 45 &IsLayerOfType<armnn::OutputLayer>)); 46 47 armnn::Optimizer::Pass(graph, armnn::MakeOptimizations(OptimizeConsecutiveReshapes())); 48 __anon32618ce20102(const armnn::Layer* const layer) 49 auto checkReshape = [&info2](const armnn::Layer* const layer) -> bool { 50 const auto reshapeLayer = static_cast<const armnn::ReshapeLayer*>(layer); 51 return IsLayerOfType<armnn::ReshapeLayer>(layer) && 52 (reshapeLayer->GetParameters().m_TargetShape == info2.GetShape()) && 53 (reshapeLayer->GetOutputHandler().GetTensorInfo().GetShape() == info2.GetShape()); 54 }; 55 56 // The two reshapes are replaced by a single equivalent reshape. 57 CHECK(CheckSequence(graph.cbegin(), graph.cend(), &IsLayerOfType<armnn::InputLayer>, checkReshape, 58 &IsLayerOfType<armnn::OutputLayer>)); 59 60 // Check the new reshape layer has the other two reshapes as related layers 61 std::list<std::string> testRelatedLayers = { reshape2Name, reshape1Name }; 62 63 CHECK(CheckRelatedLayers<armnn::ReshapeLayer>(graph, testRelatedLayers)); 64 } 65 66 { 67 // Inserts a reshape to the input shape. 68 auto reshapeToIn = graph.InsertNewLayer<armnn::ReshapeLayer>( 69 output->GetInputSlot(0), armnn::ReshapeDescriptor{ info0.GetShape() }, "reshapeToIn"); 70 71 reshapeToIn->GetOutputHandler().SetTensorInfo(info0); 72 73 armnn::Optimizer::Pass(graph, armnn::MakeOptimizations(OptimizeConsecutiveReshapes())); 74 75 // The two reshapes are removed. 76 CHECK(CheckSequence(graph.cbegin(), graph.cend(), &IsLayerOfType<armnn::InputLayer>, 77 &IsLayerOfType<armnn::OutputLayer>)); 78 } 79 } 80 81 }