IVGCVSW-6161 ConstTensorsAsInput: Optimizer - Redirect ConstTensor layer members

 * Optimization that searches for layers with ConstantLayers as inputs.
 * The layer member variables are then redirected to these ConstantLayers.

Signed-off-by: Matthew Sloyan <matthew.sloyan@arm.com>
Change-Id: I24a2bf0e8575b808343e0bbe3897b344e94796ad

6 files changed, 175 insertions, 0 deletions

diff --git a/Android.mk b/Android.mk
index f3e4f40534..d472e1a5cc 100644
--- a/Android.mk
+++ b/Android.mk
@@ -395,6 +395,7 @@ LOCAL_SRC_FILES := \
         src/armnn/test/optimizations/OptimizeInversePermutesTests.cpp \
         src/armnn/test/optimizations/PermuteAndBatchToSpaceAsDepthToSpaceTests.cpp \
         src/armnn/test/optimizations/PermuteAsReshapeTests.cpp \
+        src/armnn/test/optimizations/RedirectMembersToConstantInputsTests.cpp \
         src/armnn/test/optimizations/ReduceMultipleAxesTests.cpp \
         src/armnn/test/optimizations/SquashEqualSiblingsTests.cpp \
         src/armnn/test/optimizations/TransposeAsReshapeTests.cpp \
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9ac7ecffaa..1ccc80df1e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -593,6 +593,7 @@ if(BUILD_UNIT_TESTS)
         src/armnn/test/optimizations/OptimizeInversePermutesTests.cpp
         src/armnn/test/optimizations/PermuteAndBatchToSpaceAsDepthToSpaceTests.cpp
         src/armnn/test/optimizations/PermuteAsReshapeTests.cpp
+        src/armnn/test/optimizations/RedirectMembersToConstantInputsTests.cpp
         src/armnn/test/optimizations/ReduceMultipleAxesTests.cpp
         src/armnn/test/optimizations/SquashEqualSiblingsTests.cpp
         src/armnn/test/optimizations/TransposeAsReshapeTests.cpp
diff --git a/src/armnn/Network.cpp b/src/armnn/Network.cpp
index 74c195f676..d340f021e2 100644
--- a/src/armnn/Network.cpp
+++ b/src/armnn/Network.cpp
@@ -1618,6 +1618,7 @@ IOptimizedNetworkPtr Optimize(const INetwork& inNetwork,
                                                 PermuteAsReshape(),
                                                 TransposeAsReshape(),
                                                 OptimizeConsecutiveReshapes(),
+                                                RedirectMembersToConstantInputs(),
                                                 FoldPadIntoConvolution2d(),
                                                 FoldPadIntoDepthwiseConvolution2d(),
                                                 FoldPadIntoPooling2d(),
diff --git a/src/armnn/optimizations/All.hpp b/src/armnn/optimizations/All.hpp
index 5decc7c969..2bc54d993d 100644
--- a/src/armnn/optimizations/All.hpp
+++ b/src/armnn/optimizations/All.hpp
@@ -18,5 +18,6 @@
 #include "OptimizeInversePermutes.hpp"
 #include "PermuteAsReshape.hpp"
 #include "PermuteAndBatchToSpaceAsDepthToSpace.hpp"
+#include "RedirectMembersToConstantInputs.hpp"
 #include "SquashEqualSiblings.hpp"
 #include "TransposeAsReshape.hpp"
\ No newline at end of file
diff --git a/src/armnn/optimizations/RedirectMembersToConstantInputs.hpp b/src/armnn/optimizations/RedirectMembersToConstantInputs.hpp
new file mode 100644
index 0000000000..5bf5ae5460
--- /dev/null
+++ b/src/armnn/optimizations/RedirectMembersToConstantInputs.hpp
@@ -0,0 +1,86 @@
+//
+// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include "IgnoreUnused.hpp"
+#include "Optimization.hpp"
+
+#include <armnn/utility/PolymorphicDowncast.hpp>
+
+namespace armnn
+{
+namespace optimizations
+{
+
+class RedirectMembersToConstantInputsImpl
+{
+public:
+    /// Search for layers with ConstantLayers as inputs. If the inputs are constant redirect the layers member
+    /// variable for ConstTensors (e.g. m_weights) to the data stored in the ConstantLayer it is connected to.
+    void Run(Graph& graph, Layer& layer) const
+    {
+        IgnoreUnused(graph);
+
+        switch (layer.GetType())
+        {
+            case LayerType::BatchNormalization:
+                break;
+            case LayerType::Convolution2d:
+                break;
+            case LayerType::DepthwiseConvolution2d:
+                break;
+            case LayerType::DetectionPostProcess:
+                break;
+            case LayerType::FullyConnected:
+                RedirectWeightsAndBiases<FullyConnectedLayer>(&layer);
+                break;
+            case LayerType::Lstm:
+                break;
+            case LayerType::TransposeConvolution2d:
+                break;
+            default:
+                break;
+        }
+    }
+
+protected:
+    RedirectMembersToConstantInputsImpl()  = default;
+    ~RedirectMembersToConstantInputsImpl() = default;
+
+private:
+    template <typename LayerT>
+    static LayerT* RedirectWeightsAndBiases(Layer* layer)
+    {
+        LayerT* layerPtr = PolymorphicDowncast<LayerT*>(layer);
+
+        // Loop through input slots to check for constant weights and biases layers.
+        // Weights index = 1, Biases index = 2.
+        for (unsigned int inputSlotIndex = 1; inputSlotIndex != layerPtr->GetNumInputSlots(); ++inputSlotIndex)
+        {
+            OutputSlot* outputSlot = layerPtr->GetInputSlot(inputSlotIndex).GetConnectedOutputSlot();
+            if (outputSlot->GetOwningLayer().GetType() == LayerType::Constant)
+            {
+                // Get constant layer and redirect base layer member variables.
+                ConstantLayer& constantLayer = dynamic_cast<ConstantLayer&>(outputSlot->GetOwningLayer());
+                if (inputSlotIndex == 1)
+                {
+                    layerPtr->m_Weight = constantLayer.m_LayerOutput;
+                }
+                else if (inputSlotIndex == 2)
+                {
+                    layerPtr->m_Bias = constantLayer.m_LayerOutput;
+                }
+            }
+        }
+
+        return layerPtr;
+    }
+};
+
+using RedirectMembersToConstantInputs = OptimizeForType<FullyConnectedLayer, RedirectMembersToConstantInputsImpl>;
+
+} // namespace optimizations
+} // namespace armnn
diff --git a/src/armnn/test/optimizations/RedirectMembersToConstantInputsTests.cpp b/src/armnn/test/optimizations/RedirectMembersToConstantInputsTests.cpp
new file mode 100644
index 0000000000..46b06a55c7
--- /dev/null
+++ b/src/armnn/test/optimizations/RedirectMembersToConstantInputsTests.cpp
@@ -0,0 +1,85 @@
+//
+// Copyright © 2021 Arm Ltd and Contributors. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#include "../TestUtils.hpp"
+
+#include <Optimizer.hpp>
+
+#include <doctest/doctest.h>
+
+TEST_SUITE("Optimizer")
+{
+using namespace armnn::optimizations;
+
+TEST_CASE("RedirectMembersToConstantInputsFullyConnectedTest")
+{
+    armnn::Graph graph;
+
+    const armnn::TensorInfo inputInfo  ({ 1, 2, 2, 3 }, armnn::DataType::Float32);
+    const armnn::TensorInfo outputInfo ({ 1, 2, 2, 3 }, armnn::DataType::Float32);
+    const armnn::TensorInfo weightsInfo({ 4 }, armnn::DataType::Float32, 0.0f, 0, true);
+    const armnn::TensorInfo biasesInfo ({ 2 }, armnn::DataType::Float32, 0.0f, 0, true);
+
+    // Check if isConstant is enabled for weights and biases tensor info.
+    CHECK(weightsInfo.IsConstant());
+    CHECK(biasesInfo.IsConstant());
+
+    armnn::FullyConnectedDescriptor desc;
+    desc.m_BiasEnabled = true;
+    desc.m_ConstantWeights = false;
+
+    // Create the simple test network with Weights and Biases as inputs to a FullyConnected layer.
+    auto input   = graph.AddLayer<armnn::InputLayer>(0, "Input");
+    auto weights = graph.AddLayer<armnn::ConstantLayer>("Weights");
+    auto biases  = graph.AddLayer<armnn::ConstantLayer>("Biases");
+    auto fcLayer = graph.AddLayer<armnn::FullyConnectedLayer>(desc, "FullyConnected");
+    auto output  = graph.AddLayer<armnn::OutputLayer>(1, "Output");
+
+    float expectedWeightsData[] = { 1.0f, 1.0f, 1.0f, 1.0f };
+    float expectedBiasesData[]  = { 2.0f, 2.0f };
+
+    // Set the m_LayerOutput for the optimizer to point to.
+    armnn::ConstTensor weightsTensor(weightsInfo, &expectedWeightsData);
+    armnn::ConstTensor biasesTensor(biasesInfo, &expectedBiasesData);
+    weights->m_LayerOutput = std::make_unique<armnn::ScopedTensorHandle>(weightsTensor);
+    biases->m_LayerOutput  = std::make_unique<armnn::ScopedTensorHandle>(biasesTensor);
+
+    input->GetOutputSlot().SetTensorInfo(inputInfo);
+    weights->GetOutputSlot().SetTensorInfo(weightsInfo);
+    biases->GetOutputSlot().SetTensorInfo(biasesInfo);
+    fcLayer->GetOutputSlot().SetTensorInfo(outputInfo);
+
+    // Connect up the layers
+    input->GetOutputSlot(0).Connect(fcLayer->GetInputSlot(0));
+    weights->GetOutputSlot(0).Connect(fcLayer->GetInputSlot(1));
+    biases->GetOutputSlot(0).Connect(fcLayer->GetInputSlot(2));
+    fcLayer->GetOutputSlot(0).Connect(output->GetInputSlot(0));
+
+    // Member variables should be null before optimization.
+    CHECK(fcLayer->m_Weight == nullptr);
+    CHECK(fcLayer->m_Bias == nullptr);
+
+    // Run the optimizer
+    armnn::Optimizer::Pass(graph, armnn::MakeOptimizations(RedirectMembersToConstantInputs()));
+
+    // Check if member variables are not null and shape is set correctly.
+    CHECK(fcLayer->m_Weight != nullptr);
+    CHECK(fcLayer->m_Bias != nullptr);
+    CHECK(fcLayer->m_Weight->GetTensorInfo().GetShape() == weightsInfo.GetShape());
+    CHECK(fcLayer->m_Bias->GetTensorInfo().GetShape() == biasesInfo.GetShape());
+
+    // Check whether data matches expected float data
+    const float* weightsData = fcLayer->m_Weight->GetConstTensor<float>();
+    CHECK(weightsData[0] == expectedWeightsData[0]);
+    CHECK(weightsData[1] == expectedWeightsData[1]);
+    CHECK(weightsData[2] == expectedWeightsData[2]);
+    CHECK(weightsData[3] == expectedWeightsData[3]);
+
+    const float* biasesData = fcLayer->m_Bias->GetConstTensor<float>();
+    CHECK(biasesData[0] == expectedBiasesData[0]);
+    CHECK(biasesData[1] == expectedBiasesData[1]);
+}
+
+}
\ No newline at end of file