Add FFT2d and RFFT2d serialization support

Signed-off-by: Luke Hutton <luke.hutton@arm.com>
Change-Id: I79b3684ae1cf24e71ec7098ed7db5be5bc2f6fb2

7 files changed, 149 insertions, 39 deletions

diff --git a/include/attribute.def b/include/attribute.def
index 121bd89..86361a4 100644
--- a/include/attribute.def
+++ b/include/attribute.def
@@ -127,3 +127,6 @@ DEF_ATTRIBUTE(Custom, 3,
               string, S, identifier,
               string, S, config,
               uint8_t, V, implementation_attrs)
+
+DEF_ATTRIBUTE(FFT, 1,
+              bool, S, inverse)
diff --git a/include/tosa_generated.h b/include/tosa_generated.h
index bb501be..1a79453 100644
--- a/include/tosa_generated.h
+++ b/include/tosa_generated.h
@@ -71,6 +71,9 @@ struct NegateAttributeBuilder;
 struct CustomAttribute;
 struct CustomAttributeBuilder;
 
+struct FFTAttribute;
+struct FFTAttributeBuilder;
+
 struct Version;
 struct VersionBuilder;
 
@@ -252,11 +255,13 @@ enum Op : uint32_t {
   Op_CUSTOM = 66,
   Op_COND_IF = 67,
   Op_WHILE_LOOP = 68,
+  Op_FFT2D = 69,
+  Op_RFFT2D = 70,
   Op_MIN = Op_UNKNOWN,
-  Op_MAX = Op_WHILE_LOOP
+  Op_MAX = Op_RFFT2D
 };
 
-inline const Op (&EnumValuesOp())[69] {
+inline const Op (&EnumValuesOp())[71] {
   static const Op values[] = {
     Op_UNKNOWN,
     Op_ARGMAX,
@@ -326,13 +331,15 @@ inline const Op (&EnumValuesOp())[69] {
     Op_IDENTITY,
     Op_CUSTOM,
     Op_COND_IF,
-    Op_WHILE_LOOP
+    Op_WHILE_LOOP,
+    Op_FFT2D,
+    Op_RFFT2D
   };
   return values;
 }
 
 inline const char * const *EnumNamesOp() {
-  static const char * const names[70] = {
+  static const char * const names[72] = {
     "UNKNOWN",
     "ARGMAX",
     "AVG_POOL2D",
@@ -402,13 +409,15 @@ inline const char * const *EnumNamesOp() {
     "CUSTOM",
     "COND_IF",
     "WHILE_LOOP",
+    "FFT2D",
+    "RFFT2D",
     nullptr
   };
   return names;
 }
 
 inline const char *EnumNameOp(Op e) {
-  if (flatbuffers::IsOutRange(e, Op_UNKNOWN, Op_WHILE_LOOP)) return "";
+  if (flatbuffers::IsOutRange(e, Op_UNKNOWN, Op_RFFT2D)) return "";
   const size_t index = static_cast<size_t>(e);
   return EnumNamesOp()[index];
 }
@@ -436,11 +445,12 @@ enum Attribute : uint8_t {
   Attribute_FullyConnectedAttribute = 19,
   Attribute_NegateAttribute = 20,
   Attribute_CustomAttribute = 21,
+  Attribute_FFTAttribute = 22,
   Attribute_MIN = Attribute_NONE,
-  Attribute_MAX = Attribute_CustomAttribute
+  Attribute_MAX = Attribute_FFTAttribute
 };
 
-inline const Attribute (&EnumValuesAttribute())[22] {
+inline const Attribute (&EnumValuesAttribute())[23] {
   static const Attribute values[] = {
     Attribute_NONE,
     Attribute_PoolAttribute,
@@ -463,13 +473,14 @@ inline const Attribute (&EnumValuesAttribute())[22] {
     Attribute_MatMulAttribute,
     Attribute_FullyConnectedAttribute,
     Attribute_NegateAttribute,
-    Attribute_CustomAttribute
+    Attribute_CustomAttribute,
+    Attribute_FFTAttribute
   };
   return values;
 }
 
 inline const char * const *EnumNamesAttribute() {
-  static const char * const names[23] = {
+  static const char * const names[24] = {
     "NONE",
     "PoolAttribute",
     "ConvAttribute",
@@ -492,13 +503,14 @@ inline const char * const *EnumNamesAttribute() {
     "FullyConnectedAttribute",
     "NegateAttribute",
     "CustomAttribute",
+    "FFTAttribute",
     nullptr
   };
   return names;
 }
 
 inline const char *EnumNameAttribute(Attribute e) {
-  if (flatbuffers::IsOutRange(e, Attribute_NONE, Attribute_CustomAttribute)) return "";
+  if (flatbuffers::IsOutRange(e, Attribute_NONE, Attribute_FFTAttribute)) return "";
   const size_t index = static_cast<size_t>(e);
   return EnumNamesAttribute()[index];
 }
@@ -591,6 +603,10 @@ template<> struct AttributeTraits<tosa::CustomAttribute> {
   static const Attribute enum_value = Attribute_CustomAttribute;
 };
 
+template<> struct AttributeTraits<tosa::FFTAttribute> {
+  static const Attribute enum_value = Attribute_FFTAttribute;
+};
+
 bool VerifyAttribute(flatbuffers::Verifier &verifier, const void *obj, Attribute type);
 bool VerifyAttributeVector(flatbuffers::Verifier &verifier, const flatbuffers::Vector<flatbuffers::Offset<void>> *values, const flatbuffers::Vector<uint8_t> *types);
 
@@ -2075,6 +2091,47 @@ inline flatbuffers::Offset<CustomAttribute> CreateCustomAttributeDirect(
       implementation_attrs__);
 }
 
+struct FFTAttribute FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
+  typedef FFTAttributeBuilder Builder;
+  enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
+    VT_INVERSE = 4
+  };
+  bool inverse() const {
+    return GetField<uint8_t>(VT_INVERSE, 0) != 0;
+  }
+  bool Verify(flatbuffers::Verifier &verifier) const {
+    return VerifyTableStart(verifier) &&
+           VerifyField<uint8_t>(verifier, VT_INVERSE, 1) &&
+           verifier.EndTable();
+  }
+};
+
+struct FFTAttributeBuilder {
+  typedef FFTAttribute Table;
+  flatbuffers::FlatBufferBuilder &fbb_;
+  flatbuffers::uoffset_t start_;
+  void add_inverse(bool inverse) {
+    fbb_.AddElement<uint8_t>(FFTAttribute::VT_INVERSE, static_cast<uint8_t>(inverse), 0);
+  }
+  explicit FFTAttributeBuilder(flatbuffers::FlatBufferBuilder &_fbb)
+        : fbb_(_fbb) {
+    start_ = fbb_.StartTable();
+  }
+  flatbuffers::Offset<FFTAttribute> Finish() {
+    const auto end = fbb_.EndTable(start_);
+    auto o = flatbuffers::Offset<FFTAttribute>(end);
+    return o;
+  }
+};
+
+inline flatbuffers::Offset<FFTAttribute> CreateFFTAttribute(
+    flatbuffers::FlatBufferBuilder &_fbb,
+    bool inverse = false) {
+  FFTAttributeBuilder builder_(_fbb);
+  builder_.add_inverse(inverse);
+  return builder_.Finish();
+}
+
 struct Version FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   typedef VersionBuilder Builder;
   enum FlatBuffersVTableOffset FLATBUFFERS_VTABLE_UNDERLYING_TYPE {
@@ -2320,6 +2377,9 @@ struct TosaOperator FLATBUFFERS_FINAL_CLASS : private flatbuffers::Table {
   const tosa::CustomAttribute *attribute_as_CustomAttribute() const {
     return attribute_type() == tosa::Attribute_CustomAttribute ? static_cast<const tosa::CustomAttribute *>(attribute()) : nullptr;
   }
+  const tosa::FFTAttribute *attribute_as_FFTAttribute() const {
+    return attribute_type() == tosa::Attribute_FFTAttribute ? static_cast<const tosa::FFTAttribute *>(attribute()) : nullptr;
+  }
   const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>> *inputs() const {
     return GetPointer<const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>> *>(VT_INPUTS);
   }
@@ -2426,6 +2486,10 @@ template<> inline const tosa::CustomAttribute *TosaOperator::attribute_as<tosa::
   return attribute_as_CustomAttribute();
 }
 
+template<> inline const tosa::FFTAttribute *TosaOperator::attribute_as<tosa::FFTAttribute>() const {
+  return attribute_as_FFTAttribute();
+}
+
 struct TosaOperatorBuilder {
   typedef TosaOperator Table;
   flatbuffers::FlatBufferBuilder &fbb_;
@@ -2821,6 +2885,10 @@ inline bool VerifyAttribute(flatbuffers::Verifier &verifier, const void *obj, At
       auto ptr = reinterpret_cast<const tosa::CustomAttribute *>(obj);
       return verifier.VerifyTable(ptr);
     }
+    case Attribute_FFTAttribute: {
+      auto ptr = reinterpret_cast<const tosa::FFTAttribute *>(obj);
+      return verifier.VerifyTable(ptr);
+    }
     default: return true;
   }
 }
diff --git a/python/serializer/tosa_serializer.py b/python/serializer/tosa_serializer.py
index 8f70fb0..e8311ce 100644
--- a/python/serializer/tosa_serializer.py
+++ b/python/serializer/tosa_serializer.py
@@ -361,6 +361,14 @@ class TosaSerializerAttribute(TosaSerializerUnion):
         self.ints.append((a.AddInput1Zp, input1_zp))
         self.ints.append((a.AddOutputZp, output_zp))
 
+    def FFTAttribute(self, inverse):
+        from tosa import FFTAttribute as a, Attribute
+
+        self.utype = Attribute.Attribute().FFTAttribute
+        self.optFcns = (a.Start, a.End)
+
+        self.bools.append((a.AddInverse, inverse))
+
 
 class TosaSerializerTensor:
     def __init__(
diff --git a/python/tosa/Attribute.py b/python/tosa/Attribute.py
index 7ada553..25ade44 100644
--- a/python/tosa/Attribute.py
+++ b/python/tosa/Attribute.py
@@ -25,3 +25,4 @@ class Attribute(object):
     FullyConnectedAttribute = 19
     NegateAttribute = 20
     CustomAttribute = 21
+    FFTAttribute = 22
diff --git a/python/tosa/FFTAttribute.py b/python/tosa/FFTAttribute.py
new file mode 100644
index 0000000..3dd4083
--- /dev/null
+++ b/python/tosa/FFTAttribute.py
@@ -0,0 +1,46 @@
+# automatically generated by the FlatBuffers compiler, do not modify
+
+# namespace: tosa
+
+import flatbuffers
+from flatbuffers.compat import import_numpy
+np = import_numpy()
+
+class FFTAttribute(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def GetRootAs(cls, buf, offset=0):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = FFTAttribute()
+        x.Init(buf, n + offset)
+        return x
+
+    @classmethod
+    def GetRootAsFFTAttribute(cls, buf, offset=0):
+        """This method is deprecated. Please switch to GetRootAs."""
+        return cls.GetRootAs(buf, offset)
+    @classmethod
+    def FFTAttributeBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
+        return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x54\x4F\x53\x41", size_prefixed=size_prefixed)
+
+    # FFTAttribute
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+    # FFTAttribute
+    def Inverse(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
+        if o != 0:
+            return bool(self._tab.Get(flatbuffers.number_types.BoolFlags, o + self._tab.Pos))
+        return False
+
+def FFTAttributeStart(builder): builder.StartObject(1)
+def Start(builder):
+    return FFTAttributeStart(builder)
+def FFTAttributeAddInverse(builder, inverse): builder.PrependBoolSlot(0, inverse, 0)
+def AddInverse(builder, inverse):
+    return FFTAttributeAddInverse(builder, inverse)
+def FFTAttributeEnd(builder): return builder.EndObject()
+def End(builder):
+    return FFTAttributeEnd(builder)
\ No newline at end of file
diff --git a/python/tosa/Op.py b/python/tosa/Op.py
index 6b4c7b0..b8e68c0 100644
--- a/python/tosa/Op.py
+++ b/python/tosa/Op.py
@@ -72,3 +72,5 @@ class Op(object):
     CUSTOM = 66
     COND_IF = 67
     WHILE_LOOP = 68
+    FFT2D = 69
+    RFFT2D = 70
diff --git a/schema/tosa.fbs b/schema/tosa.fbs
index 093c235..ff8697b 100644
--- a/schema/tosa.fbs
+++ b/schema/tosa.fbs
@@ -20,6 +20,9 @@ file_identifier "TOSA";
 // File extension of any written files.
 file_extension "tosa";
 
+// NOTE: New values added to the schema should be placed
+// at the end of the list in order to keep schema stable.
+
 enum DType:uint32 {
   UNKNOWN = 0,
   BOOL,
@@ -43,8 +46,6 @@ enum ResizeMode:uint32 {
 
 enum Op:uint32 {
   UNKNOWN = 0,
-
-  // Tensor Operator
   ARGMAX,
   AVG_POOL2D,
   CONV2D,
@@ -54,14 +55,10 @@ enum Op:uint32 {
   MATMUL,
   MAX_POOL2D,
   TRANSPOSE_CONV2D,
-
-  // Activation
   CLAMP,
   RESERVED,
   SIGMOID,
   TANH,
-
-  // Elementwise-Binary
   ADD,
   ARITHMETIC_RIGHT_SHIFT,
   BITWISE_AND,
@@ -79,8 +76,6 @@ enum Op:uint32 {
   POW,
   SUB,
   TABLE,
-
-  // Elementwise-Unary
   ABS,
   BITWISE_NOT,
   CEIL,
@@ -92,24 +87,16 @@ enum Op:uint32 {
   NEGATE,
   RECIPROCAL,
   RSQRT,
-
-  // Elementwise-Ternary
   SELECT,
-
-  // Logical
   EQUAL,
   GREATER,
   GREATER_EQUAL,
-
-  // Reduction
   REDUCE_ANY,
   REDUCE_ALL,
   REDUCE_MAX,
   REDUCE_MIN,
   REDUCE_PRODUCT,
   REDUCE_SUM,
-
-  // Data layout operation
   CONCAT,
   PAD,
   RESHAPE,
@@ -117,28 +104,18 @@ enum Op:uint32 {
   SLICE,
   TILE,
   TRANSPOSE,
-
-  // Gather/scatter operation
   GATHER,
   SCATTER,
-
-  // Image
   RESIZE,
-
-  // Type conversion
   CAST,
   RESCALE,
-
-  // Data Nodes
   CONST,
   IDENTITY,
-
-  // Custom operations
   CUSTOM,
-
-  // Control flow operators
   COND_IF,
   WHILE_LOOP,
+  FFT2D,
+  RFFT2D,
 }
 
 union Attribute {
@@ -162,7 +139,8 @@ union Attribute {
   MatMulAttribute,
   FullyConnectedAttribute,
   NegateAttribute,
-  CustomAttribute
+  CustomAttribute,
+  FFTAttribute,
 }
 
 table PoolAttribute {
@@ -288,6 +266,10 @@ table CustomAttribute {
   implementation_attrs:[ubyte];
 }
 
+table FFTAttribute {
+  inverse: bool;
+}
+
 table Version {
   _major: int32 = 0;
   _minor: int32 = 51;