From cb4bbf4fe49ec08b16c348b0db988199de789500 Mon Sep 17 00:00:00 2001
From: Won Jeon <won.jeon@arm.com>
Date: Thu, 10 Aug 2023 08:50:15 +0000
Subject: Add DIM operator to serialization library

Signed-off-by: Won Jeon <won.jeon@arm.com>
Change-Id: I7c68f94a088e206c99d34d152601cf342bff5541
---
 include/tosa_generated.h           | 26 ++++++++++++++++----------
 python/tosa/DType.py               |  1 +
 python/tosa/Op.py                  |  1 +
 schema/tosa.fbs                    |  2 ++
 src/tosa_serialization_handler.cpp |  2 +-
 5 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/include/tosa_generated.h b/include/tosa_generated.h
index 1cddf98..22819f1 100644
--- a/include/tosa_generated.h
+++ b/include/tosa_generated.h
@@ -112,11 +112,12 @@ enum DType : uint32_t {
   DType_UINT16 = 9,
   DType_FP16 = 10,
   DType_BF16 = 11,
+  DType_SHAPE = 12,
   DType_MIN = DType_UNKNOWN,
-  DType_MAX = DType_BF16
+  DType_MAX = DType_SHAPE
 };
 
-inline const DType (&EnumValuesDType())[12] {
+inline const DType (&EnumValuesDType())[13] {
   static const DType values[] = {
     DType_UNKNOWN,
     DType_BOOL,
@@ -129,13 +130,14 @@ inline const DType (&EnumValuesDType())[12] {
     DType_FP32,
     DType_UINT16,
     DType_FP16,
-    DType_BF16
+    DType_BF16,
+    DType_SHAPE
   };
   return values;
 }
 
 inline const char * const *EnumNamesDType() {
-  static const char * const names[13] = {
+  static const char * const names[14] = {
     "UNKNOWN",
     "BOOL",
     "UINT8",
@@ -148,13 +150,14 @@ inline const char * const *EnumNamesDType() {
     "UINT16",
     "FP16",
     "BF16",
+    "SHAPE",
     nullptr
   };
   return names;
 }
 
 inline const char *EnumNameDType(DType e) {
-  if (::flatbuffers::IsOutRange(e, DType_UNKNOWN, DType_BF16)) return "";
+  if (::flatbuffers::IsOutRange(e, DType_UNKNOWN, DType_SHAPE)) return "";
   const size_t index = static_cast<size_t>(e);
   return EnumNamesDType()[index];
 }
@@ -265,11 +268,12 @@ enum Op : uint32_t {
   Op_FFT2D = 69,
   Op_RFFT2D = 70,
   Op_ERF = 71,
+  Op_DIM = 72,
   Op_MIN = Op_UNKNOWN,
-  Op_MAX = Op_ERF
+  Op_MAX = Op_DIM
 };
 
-inline const Op (&EnumValuesOp())[72] {
+inline const Op (&EnumValuesOp())[73] {
   static const Op values[] = {
     Op_UNKNOWN,
     Op_ARGMAX,
@@ -342,13 +346,14 @@ inline const Op (&EnumValuesOp())[72] {
     Op_WHILE_LOOP,
     Op_FFT2D,
     Op_RFFT2D,
-    Op_ERF
+    Op_ERF,
+    Op_DIM
   };
   return values;
 }
 
 inline const char * const *EnumNamesOp() {
-  static const char * const names[73] = {
+  static const char * const names[74] = {
     "UNKNOWN",
     "ARGMAX",
     "AVG_POOL2D",
@@ -421,13 +426,14 @@ inline const char * const *EnumNamesOp() {
     "FFT2D",
     "RFFT2D",
     "ERF",
+    "DIM",
     nullptr
   };
   return names;
 }
 
 inline const char *EnumNameOp(Op e) {
-  if (::flatbuffers::IsOutRange(e, Op_UNKNOWN, Op_ERF)) return "";
+  if (::flatbuffers::IsOutRange(e, Op_UNKNOWN, Op_DIM)) return "";
   const size_t index = static_cast<size_t>(e);
   return EnumNamesOp()[index];
 }
diff --git a/python/tosa/DType.py b/python/tosa/DType.py
index 15da2f6..6df2dcb 100644
--- a/python/tosa/DType.py
+++ b/python/tosa/DType.py
@@ -15,3 +15,4 @@ class DType(object):
     UINT16 = 9
     FP16 = 10
     BF16 = 11
+    SHAPE = 12
diff --git a/python/tosa/Op.py b/python/tosa/Op.py
index b310311..d628479 100644
--- a/python/tosa/Op.py
+++ b/python/tosa/Op.py
@@ -75,3 +75,4 @@ class Op(object):
     FFT2D = 69
     RFFT2D = 70
     ERF = 71
+    DIM = 72
diff --git a/schema/tosa.fbs b/schema/tosa.fbs
index 057933c..f101fa3 100644
--- a/schema/tosa.fbs
+++ b/schema/tosa.fbs
@@ -36,6 +36,7 @@ enum DType:uint32 {
   UINT16,
   FP16,
   BF16,
+  SHAPE,
 }
 
 enum ResizeMode:uint32 {
@@ -117,6 +118,7 @@ enum Op:uint32 {
   FFT2D,
   RFFT2D,
   ERF,
+  DIM,
 }
 
 union Attribute {
diff --git a/src/tosa_serialization_handler.cpp b/src/tosa_serialization_handler.cpp
index de7efb1..cbb862f 100644
--- a/src/tosa_serialization_handler.cpp
+++ b/src/tosa_serialization_handler.cpp
@@ -682,7 +682,7 @@ tosa_err_t TosaSerializationHandler::Serialize()
             }
             auto fb_block_tensors = _builder.CreateVector(fboffset_block_tensors);
             auto fboffset_block   = CreateTosaBasicBlock(_builder, block_name, fb_block_operators, fb_block_tensors,
-                                                       fb_block_inputs, fb_block_outputs);
+                                                         fb_block_inputs, fb_block_outputs);
             fboffset_blocks.push_back(fboffset_block);
         }    // end block for_loop
         auto fb_blocks = _builder.CreateVector(fboffset_blocks);
-- 
cgit v1.2.1