From 5e268097917825ddaa00a86ee95a4a6c4f50124b Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton@arm.com>
Date: Thu, 12 Jan 2023 22:20:53 +0000
Subject: Add FFT2d and RFFT2d serialization support

Signed-off-by: Luke Hutton <luke.hutton@arm.com>
Change-Id: I79b3684ae1cf24e71ec7098ed7db5be5bc2f6fb2
---
 python/serializer/tosa_serializer.py |  8 +++++++
 python/tosa/Attribute.py             |  1 +
 python/tosa/FFTAttribute.py          | 46 ++++++++++++++++++++++++++++++++++++
 python/tosa/Op.py                    |  2 ++
 4 files changed, 57 insertions(+)
 create mode 100644 python/tosa/FFTAttribute.py

(limited to 'python')

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
-- 
cgit v1.2.1