From 23be07e855c066f192c1007d529064462853a27c Mon Sep 17 00:00:00 2001
From: jimfly01 <jim.flynn@arm.com>
Date: Tue, 4 Dec 2018 17:47:22 +0000
Subject: IVGCVSW-2333 Add ParseSub method to TfParser

* Also added unit test Sub.cpp

Change-Id: I6d23c11ae894ee433cd28ffdf0248b14e01b0131
---
 src/armnnTfParser/TfParser.cpp | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

(limited to 'src/armnnTfParser/TfParser.cpp')

diff --git a/src/armnnTfParser/TfParser.cpp b/src/armnnTfParser/TfParser.cpp
index 0d425257e8..8f6352c6e7 100644
--- a/src/armnnTfParser/TfParser.cpp
+++ b/src/armnnTfParser/TfParser.cpp
@@ -385,6 +385,7 @@ const std::map<std::string, TfParser::OperationParsingFunction> TfParser::ms_Ope
     { "Maximum",               &TfParser::ParseMaximum },
     { "Minimum",               &TfParser::ParseMinimum },
     { "Pad",                   &TfParser::ParsePad },
+    { "Sub",                   &TfParser::ParseSub },
 };
 
 ITfParser* ITfParser::CreateRaw()
@@ -1612,6 +1613,45 @@ ParsedTfOperationPtr TfParser::ParseMinimum(const tensorflow::NodeDef& nodeDef,
     return std::make_unique<SingleLayerParsedTfOperation>(this, nodeDef, layer);
 }
 
+ParsedTfOperationPtr TfParser::ParseSub(const tensorflow::NodeDef& nodeDef, const tensorflow::GraphDef& graphDef)
+{
+    std::vector<OutputOfParsedTfOperation> inputs = GetInputParsedTfOperationsChecked(nodeDef, 2);
+
+    IOutputSlot* input0Slot = &inputs[0].m_IndexedValue->ResolveArmnnOutputSlot(inputs[0].m_Index);
+    IOutputSlot* input1Slot = &inputs[1].m_IndexedValue->ResolveArmnnOutputSlot(inputs[1].m_Index);
+
+    const TensorInfo& input0Info = input0Slot->GetTensorInfo();
+    const TensorInfo& input1Info = input1Slot->GetTensorInfo();
+
+    if (input0Info.GetNumDimensions() == 1)
+    {
+        const bool isNHWC = true;
+        input0Slot = AddBroadcastReshapeLayer(input1Slot, input0Slot, isNHWC, *m_Network, nodeDef);
+    }
+
+    if (input1Info.GetNumDimensions() == 1)
+    {
+        const bool isNHWC = true;
+        input1Slot = AddBroadcastReshapeLayer(input0Slot, input1Slot, isNHWC, *m_Network, nodeDef);
+    }
+
+    IConnectableLayer* const layer = m_Network->AddSubtractionLayer(nodeDef.name().c_str());
+
+    input0Slot->Connect(layer->GetInputSlot(0));
+    input1Slot->Connect(layer->GetInputSlot(1));
+
+    if (input0Info.GetNumDimensions() == 1)
+    {
+        layer->GetOutputSlot(0).SetTensorInfo(input1Slot->GetTensorInfo());
+    }
+    else
+    {
+        layer->GetOutputSlot(0).SetTensorInfo(input0Slot->GetTensorInfo());
+    }
+
+    return std::make_unique<SingleLayerParsedTfOperation>(this, nodeDef, layer);
+}
+
 unsigned int CheckPaddingTensor(const ConstTensor& paddingTensor,
                                 const TensorInfo& inputTensorInfo,
                                 const std::string& nodeName)
-- 
cgit v1.2.1