From ae622b7553d3d3f49447160655f2feb7aa3b0e17 Mon Sep 17 00:00:00 2001
From: Matteo Martincigh <matteo.martincigh@arm.com>
Date: Tue, 23 Oct 2018 18:25:38 +0100
Subject: IVGCVSW-2049 Create the CL Mean Float workload

 * Refactored HalPolicy::ConvertMean
 * Added sanity checks when parsing the Mean operands

Change-Id: Iea19556adb957b200e7fd01022a87230aaaf7157
---
 1.1/HalPolicy.cpp | 48 +++++++++++++++++++++++++-----------------------
 1 file changed, 25 insertions(+), 23 deletions(-)

(limited to '1.1')

diff --git a/1.1/HalPolicy.cpp b/1.1/HalPolicy.cpp
index 9ca37ec0..c647812d 100644
--- a/1.1/HalPolicy.cpp
+++ b/1.1/HalPolicy.cpp
@@ -149,40 +149,43 @@ bool HalPolicy::ConvertSub(const Operation& operation, const Model& model, Conve
 bool HalPolicy::ConvertMean(const Operation& operation, const Model& model, ConversionData& data)
 {
     LayerInputHandle input = ConvertToLayerInputHandle(operation, 0, model, data);
-
     if (!input.IsValid())
     {
         return Fail("%s: Operation has invalid inputs", __func__);
     }
 
-    const armnn::TensorInfo& inputInfo  = input.GetTensorInfo();
-
-    armnn::MeanDescriptor descriptor;
-
     const Operand* axisOperand = GetInputOperand(operation, 1, model);
-    if (axisOperand)
+    if (!axisOperand)
     {
-        std::vector<int32_t> axis;
-        GetTensorInt32Values(*axisOperand, axis, model, data);
-        unsigned int rank = inputInfo.GetNumDimensions();
-        // convert the axis to unsigned int.
-        for (auto& i : axis)
-        {
-            unsigned int unsignedAxis = (i + rank) % rank;
-            if (std::find(descriptor.m_Axis.begin(), descriptor.m_Axis.end(), unsignedAxis) == descriptor.m_Axis.end())
-            {
-                descriptor.m_Axis.push_back(unsignedAxis);
-            }
-        }
+        return Fail("%s: Could not read input 1", __func__);
     }
 
-    int32_t keepDims;
-    GetInputInt32(operation, 2, keepDims, model, data);
-    if (keepDims > 0)
+    std::vector<int32_t> axis;
+    if (!GetTensorInt32Values(*axisOperand, axis, model, data))
     {
-        descriptor.m_KeepDims = true;
+        return Fail("%s: Input 1 has invalid values", __func__);
     }
 
+    const armnn::TensorInfo& inputInfo = input.GetTensorInfo();
+
+    // Convert the axis to unsigned int and remove duplicates.
+    unsigned int rank = inputInfo.GetNumDimensions();
+    std::set<unsigned int> uniqueAxis;
+    std::transform(axis.begin(), axis.end(),
+                   std::inserter(uniqueAxis, uniqueAxis.begin()),
+                   [rank](int i) -> unsigned int { return (i + rank) % rank; });
+
+    // Get the "keep dims" flag.
+    int32_t keepDims = 0;
+    if (!GetInputInt32(operation, 2, keepDims, model, data))
+    {
+        return Fail("%s: Could not read input 2", __func__);
+    }
+
+    armnn::MeanDescriptor descriptor;
+    descriptor.m_Axis.assign(uniqueAxis.begin(), uniqueAxis.end());
+    descriptor.m_KeepDims = keepDims > 0;
+
     const Operand* output = GetOutputOperand(operation, 0, model);
     if (!output)
     {
@@ -204,7 +207,6 @@ bool HalPolicy::ConvertMean(const Operation& operation, const Model& model, Conv
     armnn::IConnectableLayer* const layer = data.m_Network->AddMeanLayer(descriptor);
     assert(layer != nullptr);
     input.Connect(layer->GetInputSlot(0));
-    layer->GetOutputSlot(0).SetTensorInfo(outputInfo);
 
     return SetupAndTrackLayerOutputSlot(operation, 0, *layer, model, data);
 }
-- 
cgit v1.2.1