From fd0cae3153dd7844b1b14adda8571e4588ed762c Mon Sep 17 00:00:00 2001
From: Sadik Armagan <sadik.armagan@arm.com>
Date: Mon, 8 Nov 2021 17:18:31 +0000
Subject: IVGCVSW-5636 'Implement NNAPI caching functions'

* Serialize/Deserialize axis parameter on Softmax Descriptor.

Signed-off-by: Sadik Armagan <sadik.armagan@arm.com>
Change-Id: I1bba0dd7d78f9f211f77115dcefccf24a4a61b39
---
 src/armnnDeserializer/Deserializer.cpp      |  4 +++-
 src/armnnSerializer/ArmnnSchema.fbs         |  1 +
 src/armnnSerializer/ArmnnSchema_generated.h | 14 ++++++++++++--
 src/armnnSerializer/Serializer.cpp          |  4 +++-
 4 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/src/armnnDeserializer/Deserializer.cpp b/src/armnnDeserializer/Deserializer.cpp
index bfd4f6b560..8b1e9b970c 100644
--- a/src/armnnDeserializer/Deserializer.cpp
+++ b/src/armnnDeserializer/Deserializer.cpp
@@ -1652,7 +1652,8 @@ void IDeserializer::DeserializerImpl::ParseFill(GraphPtr graph, unsigned int lay
     CHECK_VALID_SIZE(outputs.size(), 1);
 
     auto layerName = GetLayerName(graph, layerIndex);
-    armnn::FillDescriptor descriptor(1.0f);
+    armnn::FillDescriptor descriptor;
+    descriptor.m_Value = graph->layers()->Get(layerIndex)->layer_as_FillLayer()->descriptor()->value();
     IConnectableLayer* layer = m_Network->AddFillLayer(descriptor, layerName.c_str());
 
     armnn::TensorInfo outputTensorInfo = ToTensorInfo(outputs[0]);
@@ -2484,6 +2485,7 @@ void IDeserializer::DeserializerImpl::ParseSoftmax(GraphPtr graph, unsigned int
 
     armnn::SoftmaxDescriptor descriptor;
     descriptor.m_Beta = graph->layers()->Get(layerIndex)->layer_as_SoftmaxLayer()->descriptor()->beta();
+    descriptor.m_Axis = graph->layers()->Get(layerIndex)->layer_as_SoftmaxLayer()->descriptor()->axis();
     auto layerName = GetLayerName(graph, layerIndex);
 
     IConnectableLayer* layer = m_Network->AddSoftmaxLayer(descriptor, layerName.c_str());
diff --git a/src/armnnSerializer/ArmnnSchema.fbs b/src/armnnSerializer/ArmnnSchema.fbs
index 40de3496b0..362dd5c653 100644
--- a/src/armnnSerializer/ArmnnSchema.fbs
+++ b/src/armnnSerializer/ArmnnSchema.fbs
@@ -495,6 +495,7 @@ table SoftmaxLayer {
 
 table SoftmaxDescriptor {
     beta:float;
+    axis:int = -1;
 }
 
 table DepthwiseConvolution2dLayer {
diff --git a/src/armnnSerializer/ArmnnSchema_generated.h b/src/armnnSerializer/ArmnnSchema_generated.h
index 7747f9edd9..b66bac69db 100644
--- a/src/armnnSerializer/ArmnnSchema_generated.h
+++ b/src/armnnSerializer/ArmnnSchema_generated.h
@@ -5126,14 +5126,19 @@ inline flatbuffers::Offset<SoftmaxLayer> CreateSoftmaxLayer(
 struct SoftmaxDescriptor FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   typedef SoftmaxDescriptorBuilder Builder;
   enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
-    VT_BETA = 4
+    VT_BETA = 4,
+    VT_AXIS = 6
   };
   float beta() const {
     return GetField<float>(VT_BETA, 0.0f);
   }
+  int32_t axis() const {
+    return GetField<int32_t>(VT_AXIS, -1);
+  }
   bool Verify(flatbuffers::Verifier &verifier) const {
     return VerifyTableStart(verifier) &&
            VerifyField<float>(verifier, VT_BETA) &&
+           VerifyField<int32_t>(verifier, VT_AXIS) &&
            verifier.EndTable();
   }
 };
@@ -5145,6 +5150,9 @@ struct SoftmaxDescriptorBuilder {
   void add_beta(float beta) {
     fbb_.AddElement<float>(SoftmaxDescriptor::VT_BETA, beta, 0.0f);
   }
+  void add_axis(int32_t axis) {
+    fbb_.AddElement<int32_t>(SoftmaxDescriptor::VT_AXIS, axis, -1);
+  }
   explicit SoftmaxDescriptorBuilder(flatbuffers::FlatBufferBuilder &_fbb)
         : fbb_(_fbb) {
     start_ = fbb_.StartTable();
@@ -5159,8 +5167,10 @@ struct SoftmaxDescriptorBuilder {
 
 inline flatbuffers::Offset<SoftmaxDescriptor> CreateSoftmaxDescriptor(
     flatbuffers::FlatBufferBuilder &_fbb,
-    float beta = 0.0f) {
+    float beta = 0.0f,
+    int32_t axis = -1) {
   SoftmaxDescriptorBuilder builder_(_fbb);
+  builder_.add_axis(axis);
   builder_.add_beta(beta);
   return builder_.Finish();
 }
diff --git a/src/armnnSerializer/Serializer.cpp b/src/armnnSerializer/Serializer.cpp
index c08784352d..a9f7b7d307 100644
--- a/src/armnnSerializer/Serializer.cpp
+++ b/src/armnnSerializer/Serializer.cpp
@@ -1037,7 +1037,9 @@ void SerializerStrategy::SerializeSoftmaxLayer(const armnn::IConnectableLayer* l
 
     // Create the FlatBuffer SoftmaxDescriptor
     auto flatBufferSoftmaxDesc =
-        serializer::CreateSoftmaxDescriptor(m_flatBufferBuilder, softmaxDescriptor.m_Beta);
+        serializer::CreateSoftmaxDescriptor(m_flatBufferBuilder,
+                                            softmaxDescriptor.m_Beta,
+                                            softmaxDescriptor.m_Axis);
 
     // Create the FlatBuffer SoftmaxLayer
     auto flatBufferSoftmaxLayer =
-- 
cgit v1.2.1