Add FFT2d to the reference model

Includes:
* FFT2d reference implementation
* Basic TOSA tests

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

3 files changed, 219 insertions, 4 deletions

diff --git a/verif/generator/tosa_arg_gen.py b/verif/generator/tosa_arg_gen.py
index 05a7d2b..370570c 100644
--- a/verif/generator/tosa_arg_gen.py
+++ b/verif/generator/tosa_arg_gen.py
@@ -417,6 +417,45 @@ class TosaTensorGen:
         return [ifm_shape, filter_shape, bias_shape]
 
     @staticmethod
+    def tgFFT2d(testGen, op, rank, error_name=None):
+        pl, const = op["operands"]
+
+        if error_name != ErrorIf.WrongRank:
+            assert rank == 3
+        assert pl == 2 and const == 0
+
+        # IFM dimensions are NHW
+        ifm_shape = testGen.makeShape(rank)
+
+        # Select nearest lower power of two from input height and width
+        ifm_shape[1] = 2 ** int(math.log(ifm_shape[1], 2))
+        ifm_shape[2] = 2 ** int(math.log(ifm_shape[2], 2))
+
+        # Constrict the overall size of the shape when creating ERROR_IF tests
+        if error_name:
+            ifm_shape = TosaErrorIfArgGen.eiRestrictDimensions(ifm_shape)
+
+        # Generate an invalid kernel that is not a power of two
+        if error_name == ErrorIf.KernelNotPowerOfTwo:
+            inc_h = 2 if ifm_shape[1] == 1 else 1
+            inc_w = 2 if ifm_shape[2] == 1 else 1
+            inc_choices = [(inc_h, 0), (0, inc_w), (inc_h, inc_w)]
+            selected_inc = testGen.rng.choice(inc_choices)
+            ifm_shape[1] += selected_inc[0]
+            ifm_shape[2] += selected_inc[1]
+
+        ifm_shape = testGen.constrictBatchSize(ifm_shape)
+
+        ifm_shapes = [ifm_shape.copy(), ifm_shape.copy()]
+        if error_name == ErrorIf.FFTInputShapeMismatch:
+            modify_shape = testGen.rng.choice([0, 1])
+            # Only modify kernel (H, W)
+            modify_dim = testGen.rng.choice([1, 2])
+            ifm_shapes[modify_shape][modify_dim] *= 2
+
+        return [ifm_shapes[0], ifm_shapes[1]]
+
+    @staticmethod
     def tgRFFT2d(testGen, op, rank, error_name=None):
         pl, const = op["operands"]
 
@@ -1613,6 +1652,15 @@ class TosaArgGen:
 
         return arg_list
 
+    @staticmethod
+    def agFFT2d(testGen, opName, shapeList, dtype, error_name=None):
+        arg_list = []
+
+        arg_list.append(("inverseTrue", [True]))
+        arg_list.append(("inverseFalse", [False]))
+
+        return arg_list
+
     # Helper function for reshape.  Gets some factors of a larger number.
     @staticmethod
     def getFactors(val, start=1):
diff --git a/verif/generator/tosa_error_if.py b/verif/generator/tosa_error_if.py
index 93f975d..ee227b3 100644
--- a/verif/generator/tosa_error_if.py
+++ b/verif/generator/tosa_error_if.py
@@ -79,6 +79,8 @@ class ErrorIf(object):
     CondIfCondShapeNotSizeOne = "CondIfCondShapeNotSizeOne"
     CondGraphOutputShapeNotSizeOne = "CondGraphOutputShapeNotSizeOne"
     KernelNotPowerOfTwo = "KernelNotPowerOfTwo"
+    FFTInputShapeMismatch = "FFTInputShapeMismatch"
+    FFTOutputShapeMismatch = "FFTOutputShapeMismatch"
 
 
 class TosaErrorIfArgGen:
@@ -562,7 +564,7 @@ class TosaErrorValidator:
                 ):
                     error_result = True
 
-            elif op["op"] == Op.RFFT2D:
+            elif op["op"] in [Op.FFT2D, Op.RFFT2D]:
                 if not all([ty == input_dtype for ty in output_dtype]):
                     error_result = True
 
@@ -686,7 +688,7 @@ class TosaErrorValidator:
             op = kwargs["op"]
             output_list = kwargs["output_list"]
             expected_length = 1
-            if op["op"] == Op.RFFT2D:
+            if op["op"] in [Op.FFT2D, Op.RFFT2D]:
                 expected_length = 2
 
             if len(output_list) != expected_length:
@@ -2446,6 +2448,64 @@ class TosaErrorValidator:
         }
         return info_dict
 
+    @staticmethod
+    def evFFTInputShapeMismatch(check=False, **kwargs):
+        error_name = ErrorIf.FFTInputShapeMismatch
+        param_reqs = {"rank": None, "dtype": None, "shape": None}
+        error_result = False
+        error_reason = "Mismatch between real and imaginary input shapes"
+
+        if check:
+            input1 = kwargs["input1"]
+            input2 = kwargs["input2"]
+
+            if input1.shape != input2.shape:
+                error_result = True
+
+        info_dict = {
+            "error_name": error_name,
+            "error_result": error_result,
+            "error_reason": error_reason,
+            "param_reqs": param_reqs,
+        }
+        return info_dict
+
+    @staticmethod
+    def evFFTOutputShapeMismatch(check=False, **kwargs):
+        error_name = ErrorIf.FFTOutputShapeMismatch
+        param_reqs = {"rank": None, "dtype": None, "shape": None}
+        error_result = False
+        error_reason = (
+            "Mismatch between provided and expected output kernel (H, W) shape"
+        )
+
+        if check:
+            op = kwargs["op"]
+            input_shape = kwargs["input_shape"]
+
+            if len(input_shape) == 3:
+                output_shapes = kwargs["output_shape"]
+
+                # Ignoring batch size (N) from input shape
+                expected_shape = input_shape[1:]
+                if op["op"] == Op.RFFT2D:
+                    expected_shape[1] = expected_shape[1] // 2 + 1
+
+                # Ignoring batch size (N) from output shapes
+                output_shape_0 = output_shapes[0][1:]
+                output_shape_1 = output_shapes[1][1:]
+                # Ensure sure the kernel sizes (H, W) of both outputs match the expected
+                if output_shape_0 != output_shape_1 or output_shape_0 != expected_shape:
+                    error_result = True
+
+        info_dict = {
+            "error_name": error_name,
+            "error_result": error_result,
+            "error_reason": error_reason,
+            "param_reqs": param_reqs,
+        }
+        return info_dict
+
 
 class TosaInvalidValidator:
     @staticmethod
diff --git a/verif/generator/tosa_test_gen.py b/verif/generator/tosa_test_gen.py
index 5f9e2c1..2b762aa 100644
--- a/verif/generator/tosa_test_gen.py
+++ b/verif/generator/tosa_test_gen.py
@@ -213,6 +213,12 @@ class TosaTestGen:
         else:
             raise Exception(f"Unknown dtype, cannot determine width: {dtype}")
 
+    def constrictBatchSize(self, shape):
+        # Limit the batch size unless an explicit target shape set
+        if self.args.max_batch_size and not self.args.target_shapes:
+            shape[0] = min(shape[0], self.args.max_batch_size)
+        return shape
+
     # Argument generators
     # Returns a list of tuples (stringDescriptor, [build_fcn_arg_list])
     # Where the string descriptor is used to generate the test name and
@@ -2081,6 +2087,48 @@ class TosaTestGen:
 
         return acc_out
 
+    def build_fft2d(
+        self, op, val1, val2, inverse, validator_fcns=None, error_name=None
+    ):
+        results = OutputShaper.fft2dOp(self.ser, self.rng, val1, val2, error_name)
+
+        input_names = [val1.name, val2.name]
+        pCount, cCount = op["operands"]
+        num_operands = pCount + cCount
+
+        output_names = [res.name for res in results]
+        output_shapes = [res.shape for res in results]
+        output_dtypes = [res.dtype for res in results]
+
+        input_names, output_names = TosaErrorIfArgGen.eiInvalidateInputOutputList(
+            self, error_name, input_names, output_names
+        )
+
+        if not TosaErrorValidator.evValidateErrorIfs(
+            self.ser,
+            validator_fcns,
+            error_name,
+            op=op,
+            inverse=inverse,
+            input1=val1,
+            input2=val2,
+            input_shape=val1.shape,
+            input_dtype=val1.dtype,
+            output_shape=output_shapes,
+            output_dtype=output_dtypes,
+            result_tensors=results,
+            input_list=input_names,
+            output_list=output_names,
+            num_operands=num_operands,
+        ):
+            return None
+
+        attr = ts.TosaSerializerAttribute()
+        attr.FFTAttribute(inverse)
+
+        self.ser.addOperator(op["op"], input_names, output_names, attr)
+        return results
+
     def build_rfft2d(self, op, val, validator_fcns=None, error_name=None):
         results = OutputShaper.rfft2dOp(self.ser, self.rng, val, error_name)
 
@@ -2089,6 +2137,7 @@ class TosaTestGen:
         num_operands = pCount + cCount
 
         output_names = [res.name for res in results]
+        output_shapes = [res.shape for res in results]
         output_dtypes = [res.dtype for res in results]
 
         input_names, output_names = TosaErrorIfArgGen.eiInvalidateInputOutputList(
@@ -2102,6 +2151,7 @@ class TosaTestGen:
             op=op,
             input_shape=val.shape,
             input_dtype=val.dtype,
+            output_shape=output_shapes,
             output_dtype=output_dtypes,
             result_tensors=results,
             input_list=input_names,
@@ -3927,6 +3977,29 @@ class TosaTestGen:
                 TosaErrorValidator.evCondGraphOutputShapeNotSizeOne,
             ),
         },
+        "fft2d": {
+            "op": Op.FFT2D,
+            "operands": (2, 0),
+            "rank": (3, 3),
+            "build_fcn": (
+                build_fft2d,
+                TosaTensorGen.tgFFT2d,
+                TosaTensorValuesGen.tvgDefault,
+                TosaArgGen.agFFT2d,
+            ),
+            "types": [DType.FP32],
+            "error_if_validators": (
+                TosaErrorValidator.evWrongInputType,
+                TosaErrorValidator.evWrongOutputType,
+                TosaErrorValidator.evWrongInputList,
+                TosaErrorValidator.evWrongOutputList,
+                TosaErrorValidator.evWrongRank,
+                TosaErrorValidator.evBatchMismatch,
+                TosaErrorValidator.evKernelNotPowerOfTwo,
+                TosaErrorValidator.evFFTInputShapeMismatch,
+                TosaErrorValidator.evFFTOutputShapeMismatch,
+            ),
+        },
         "rfft2d": {
             "op": Op.RFFT2D,
             "operands": (1, 0),
@@ -3946,6 +4019,7 @@ class TosaTestGen:
                 TosaErrorValidator.evWrongRank,
                 TosaErrorValidator.evBatchMismatch,
                 TosaErrorValidator.evKernelNotPowerOfTwo,
+                TosaErrorValidator.evFFTOutputShapeMismatch,
             ),
         },
     }
@@ -4770,6 +4844,37 @@ class OutputShaper:
         return ser.addOutput(output_shape, out_dtype)
 
     @staticmethod
+    def fft2dOp(serializer, rng, ifm1, ifm2, error_name=None):
+        outputs = []
+
+        assert ifm1.dtype == ifm2.dtype
+        input_dtype = ifm1.dtype
+
+        if error_name != ErrorIf.FFTInputShapeMismatch:
+            assert ifm1.shape == ifm2.shape
+
+        input_shape = ifm1.shape
+        if error_name != ErrorIf.WrongRank:
+            assert len(input_shape) == 3
+
+        output_shape = input_shape.copy()
+        output_dtype = input_dtype
+
+        if error_name == ErrorIf.WrongOutputType:
+            excludes = [DType.FP32]
+            wrong_dtypes = list(usableDTypes(excludes=excludes))
+            output_dtype = rng.choice(wrong_dtypes)
+        elif error_name == ErrorIf.BatchMismatch:
+            output_shape[0] += rng.integers(1, 10)
+        elif error_name == ErrorIf.FFTOutputShapeMismatch:
+            modify_dim = rng.choice([1, 2])
+            output_shape[modify_dim] += rng.integers(1, 10)
+
+        outputs.append(serializer.addOutput(output_shape, output_dtype))
+        outputs.append(serializer.addOutput(output_shape, output_dtype))
+        return outputs
+
+    @staticmethod
     def rfft2dOp(serializer, rng, value, error_name=None):
         outputs = []
 
@@ -4785,8 +4890,10 @@ class OutputShaper:
             wrong_dtypes = list(usableDTypes(excludes=excludes))
             output_dtype = rng.choice(wrong_dtypes)
         elif error_name == ErrorIf.BatchMismatch:
-            incorrect_batch = input_shape[0] + rng.integers(1, 10)
-            output_shape = [incorrect_batch, *input_shape[1:]]
+            output_shape[0] += rng.integers(1, 10)
+        elif error_name == ErrorIf.FFTOutputShapeMismatch:
+            modify_dim = rng.choice([1, 2])
+            output_shape[modify_dim] += rng.integers(1, 10)
 
         outputs.append(serializer.addOutput(output_shape, output_dtype))
         outputs.append(serializer.addOutput(output_shape, output_dtype))