From fc884b4141a28fbd3c62f665341ec88158fcd332 Mon Sep 17 00:00:00 2001
From: Finn Williams <Finn.Williams@Arm.com>
Date: Thu, 11 Jun 2020 17:35:44 +0100
Subject: Add support for Rank and scalar tensors

!armnn:3330

Signed-off-by: Finn Williams <Finn.Williams@Arm.com>
Change-Id: Icc429d9fabb570193d12bffef0e00dda7b51032f
---
 ConversionUtils_1_3.hpp | 47 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

(limited to 'ConversionUtils_1_3.hpp')

diff --git a/ConversionUtils_1_3.hpp b/ConversionUtils_1_3.hpp
index 3acb49a7..d5d89df1 100644
--- a/ConversionUtils_1_3.hpp
+++ b/ConversionUtils_1_3.hpp
@@ -642,4 +642,49 @@ bool ConvertQuantizedLstm(const HalOperation& operation, const HalModel& model,
              SetupAndTrackLayerOutputSlot<HalPolicy>(operation, 2, *layer, 2, model, data, &constOutputInfo));
 }
 
-} // armnn_driver namespace
\ No newline at end of file
+template<typename HalPolicy,
+         typename HalOperation = typename HalPolicy::Operation,
+         typename HalModel     = typename HalPolicy::Model>
+bool ConvertRank(const HalOperation& operation, const HalModel& model, ConversionData& data)
+{
+    using HalOperand = typename HalPolicy::Operand;
+
+    const HalOperand* inputOperand = GetInputOperand<HalPolicy>(operation, 0, model);
+    const HalOperand* outputOperand = GetOutputOperand<HalPolicy>(operation, 0, model);
+
+    if (inputOperand == nullptr || outputOperand == nullptr)
+    {
+        return Fail("%s: Operation has invalid inputs", __func__);
+    }
+
+    const Shape inputOperandShape = GetOperandShape(*inputOperand);
+    const Shape outputOperandShape = GetOperandShape(*outputOperand);
+
+    LayerInputHandle input = ConvertToLayerInputHandle<HalPolicy>(operation, 0, model, data);
+    if (!input.IsValid())
+    {
+        return Fail("%s: Could not read input 0", __func__);
+    }
+
+    armnn::TensorInfo outInfo = GetTensorInfoForOperand(*outputOperand);
+
+    bool isSupported = false;
+    FORWARD_LAYER_SUPPORT_FUNC(__func__,
+                               IsRankSupported,
+                               data.m_Backends,
+                               isSupported,
+                               input.GetTensorInfo(),
+                               outInfo);
+    if (!isSupported)
+    {
+        return false;
+    }
+
+    armnn::IConnectableLayer* layer = data.m_Network->AddRankLayer();
+    assert(layer != nullptr);
+    input.Connect(layer->GetInputSlot(0));
+
+    return SetupAndTrackLayerOutputSlot<HalPolicy>(operation, 0, *layer, model, data, &outInfo);
+}
+
+} // armnn_driver namespace
-- 
cgit v1.2.1