Main Compliance: RFFT2D support

Correct ref model to produce imaginery values of zero as specification
indicates at certain output positions.
Fix up precise and abs modes for RFFT2D in ref model to produce correct
results and bounds using abs weights.

Signed-off-by: Jeremy Johnson <jeremy.johnson@arm.com>
Change-Id: I33767e4219a260278f7933f28b1799223a95a3cc

7 files changed, 202 insertions, 11 deletions

diff --git a/reference_model/src/generate/generate_dot_product.cc b/reference_model/src/generate/generate_dot_product.cc
index 7337969..117d49d 100644
--- a/reference_model/src/generate/generate_dot_product.cc
+++ b/reference_model/src/generate/generate_dot_product.cc
@@ -963,6 +963,63 @@ bool generateFFT2D(const TosaReference::GenerateConfig& cfg,
 
     return true;
 }
+//---------------------------------------------------------------------------//
+//                              RFFT2D                                       //
+//---------------------------------------------------------------------------//
+
+template <typename DataType>
+bool generateRFFT2DReal(const TosaReference::GenerateConfig& cfg,
+                        TosaReference::IDotProductGenerator& generator,
+                        DataType* data,
+                        size_t size)
+{
+    const int64_t T  = TosaReference::numElementsFromShape(cfg.shape);
+    const uint32_t H = cfg.shape[1];
+    const uint32_t W = cfg.shape[2];
+
+    for (int64_t t = 0; t < T; ++t)
+    {
+        uint32_t x = t % W;
+        uint32_t y = (t / W) % H;
+        uint32_t k = y * W + x;
+
+        data[t] = static_cast<DataType>(generator(k));
+    }
+    return true;
+}
+
+bool generateRFFT2D(const TosaReference::GenerateConfig& cfg,
+                    TosaReference::IDotProductGenerator& generator,
+                    void* data,
+                    size_t size)
+{
+    if (cfg.shape.size() != 3)
+    {
+        WARNING("[Generator][DP][RFFT2D] Tensor shape expected 3 dimensions.");
+        return false;
+    }
+
+    switch (cfg.dataType)
+    {
+        case DType::DType_FP32: {
+            float* outData = reinterpret_cast<float*>(data);
+            switch (cfg.inputPos)
+            {
+                case 0:
+                    return generateRFFT2DReal(cfg, generator, outData, size);
+                default:
+                    WARNING("[Generator][DP][RFFT2D] Invalid input tensor slot position to operator.");
+                    return false;
+            }
+            break;
+        }
+        default:
+            WARNING("[Generator][DP][RFFT2D] Only supports FP32.");
+            return false;
+    }
+
+    return true;
+}
 }    // namespace
 
 namespace TosaReference
@@ -1003,6 +1060,8 @@ bool generateDotProduct(const GenerateConfig& cfg, void* data, size_t size)
             return generateConv3D(cfg, *generator, data, size);
         case tosa::Op_FFT2D:
             return generateFFT2D(cfg, *generator, data, size);
+        case tosa::Op_RFFT2D:
+            return generateRFFT2D(cfg, *generator, data, size);
         default:
             WARNING("[Generator][DP] Unsupported operator.");
             return false;
diff --git a/reference_model/src/generate/generate_utils.cc b/reference_model/src/generate/generate_utils.cc
index c495fb6..cf5308b 100644
--- a/reference_model/src/generate/generate_utils.cc
+++ b/reference_model/src/generate/generate_utils.cc
@@ -78,6 +78,7 @@ NLOHMANN_JSON_SERIALIZE_ENUM(Op,
                                  { Op::Op_REDUCE_PRODUCT, "REDUCE_PRODUCT" },
                                  { Op::Op_REDUCE_SUM, "REDUCE_SUM" },
                                  { Op::Op_REVERSE, "REVERSE" },
+                                 { Op::Op_RFFT2D, "RFFT2D" },
                                  { Op::Op_SCATTER, "SCATTER" },
                                  { Op::Op_SELECT, "SELECT" },
                                  { Op::Op_SIGMOID, "SIGMOID" },
diff --git a/reference_model/src/ops/tensor_ops.cc b/reference_model/src/ops/tensor_ops.cc
index 8d8dac7..dd66f79 100644
--- a/reference_model/src/ops/tensor_ops.cc
+++ b/reference_model/src/ops/tensor_ops.cc
@@ -1820,6 +1820,9 @@ int OpRFFT2d<Dtype>::eval()
     int32_t out_imag_height = out_imag->getShape()[1];
     int32_t out_imag_width  = out_imag->getShape()[2];
 
+    int32_t half_in_height = in_height / 2;
+    int32_t half_in_width  = in_width / 2;
+
     // Check Tosa Level
     auto tosa_level = g_func_config.tosa_level;
     LEVEL_CHECK(in_height <= tosa_level.MAX_KERNEL, "H should be smaller than or equal to MAX_KERNEL");
@@ -1831,7 +1834,8 @@ int OpRFFT2d<Dtype>::eval()
                in_batch, in_height, in_width, out_real_batch, out_real_height, out_real_width, out_imag_batch,
                out_imag_height, out_imag_width);
 
-    OutEigenType sum_real, sum_imag, a;
+    OutEigenType sum_real, sum_imag;
+    OutEigenType a, a_cos, a_sin, v_ir;
 
     TIn in_val = this->in->getTensor();
 
@@ -1853,10 +1857,41 @@ int OpRFFT2d<Dtype>::eval()
                 {
                     for (int ix = 0; ix < in_width; ix++)
                     {
-                        // Use explicit cast to ensure intermmediate calculations are completed using OutEigenType
-                        a = 2 * M_PI * ((iy * (OutEigenType)oy) / in_height + (ix * (OutEigenType)ox) / in_width);
-                        sum_real += in_val(n, iy, ix) * cos(a);
-                        sum_imag += -in_val(n, iy, ix) * sin(a);
+                        OutEigenType val = in_val(n, iy, ix);
+                        // Perform the periodic calculation in integer maths to keep
+                        // the accuracy of the co-efficients similar for FP32 normal
+                        // and FP64 precise mode
+                        int32_t ay = (static_cast<int64_t>(iy) * static_cast<int64_t>(oy)) % in_height;
+                        int32_t ax = (static_cast<int64_t>(ix) * static_cast<int64_t>(ox)) % in_width;
+
+                        // Use explicit cast to ensure intermediate calculations are completed using OutEigenType
+                        a = 2 * M_PI * ((OutEigenType)ay / in_height + (OutEigenType)ax / in_width);
+
+                        // Calculate weight values (co-efficients)
+                        a_cos = cos(a);
+                        a_sin = sin(a);
+
+                        if (g_func_config.abs_mode)
+                        {
+                            // Bounded op - Use abs weight values
+                            a_cos = std::abs(a_cos);
+                            a_sin = std::abs(a_sin);
+                            // Bounded op - Use abs real value for imaginary calc
+                            v_ir = val;
+                        }
+                        else
+                        {
+                            // Normal op - Use negative real value for imaginary calc
+                            v_ir = -val;
+                        }
+                        sum_real += val * a_cos;
+                        // Imaginary values with locations (0,0), (0,W/2), (H/2,0) and (H/2,W/2) are zero.
+                        // But due to sin(M_PI) not returning 0 because of M_PI being approximate, only
+                        // add to the imaginary sum when not processing these locations.
+                        if ((ay % (half_in_height)) + (ax % (half_in_width)) > 0)
+                        {
+                            sum_imag += v_ir * a_sin;
+                        }
                     }
                 }
                 this->out_real->getTensor()(n, oy, ox) = sum_real;
diff --git a/reference_model/test/generate_tests.cpp b/reference_model/test/generate_tests.cpp
index 3be402c..564af4a 100644
--- a/reference_model/test/generate_tests.cpp
+++ b/reference_model/test/generate_tests.cpp
@@ -1500,4 +1500,60 @@ TEST_CASE("positive - FP32 fft2d dot product (values -8, -7 & -6 from the end)")
     }
 }
 
+TEST_CASE("positive - FP32 rfft2d dot product (values -8, -7 & -6 from the end)")
+{
+    std::string templateJsonCfg = R"({
+        "tensors" : {
+            "real" : {
+                "generator": "DOT_PRODUCT",
+                "data_type": "FP32",
+                "input_type": "VARIABLE",
+                "shape" : [ 4, 2, 4 ],
+                "input_pos": 0,
+                "op" : "FFT2D",
+                "dot_product_info": {
+                    "s": _SET_,
+                    "ks": 8,
+                    "acc_type": "FP32"
+                }
+            }
+        }
+    })";
+
+    const std::string tosaNameReal = "real";
+    const size_t tosaElements      = 4 * 2 * 4;
+
+    SUBCASE("fft2d, set 0, real")
+    {
+        std::vector<uint32_t> expected = { 0xbe14f2f5, 0xbdb6fe4d, 0x3f30b473 };
+        fft2d_test_FP32(tosaNameReal, tosaElements, templateJsonCfg, "0", expected);
+    }
+    SUBCASE("fft2d, set 1, real")
+    {
+        // NOTE: Python test script produced 0x5e7219eb - so off by 1
+        std::vector<uint32_t> expected = { 0x5e490017, 0x5e57dd30, 0x5e992496 };
+        fft2d_test_FP32(tosaNameReal, tosaElements, templateJsonCfg, "1", expected);
+    }
+    SUBCASE("fft2d, set 2, real")
+    {
+        std::vector<uint32_t> expected = { 0x3f800000, 0xbe7f1cd4, 0xbdfc67ff };
+        fft2d_test_FP32(tosaNameReal, tosaElements, templateJsonCfg, "2", expected);
+    }
+    SUBCASE("fft2d, set 3, real")
+    {
+        std::vector<uint32_t> expected = { 0x41800000, 0xbf6d219b, 0x3f2bd153 };
+        fft2d_test_FP32(tosaNameReal, tosaElements, templateJsonCfg, "3", expected);
+    }
+    SUBCASE("fft2d, set 4, real")
+    {
+        std::vector<uint32_t> expected = { 0x0, 0x0, 0x0 };
+        fft2d_test_FP32(tosaNameReal, tosaElements, templateJsonCfg, "4", expected);
+    }
+    SUBCASE("fft2d, set 5, real")
+    {
+        std::vector<uint32_t> expected = { 0xdd3f6b86, 0xde49ecfd, 0x5e0be03d };
+        fft2d_test_FP32(tosaNameReal, tosaElements, templateJsonCfg, "5", expected);
+    }
+}
+
 TEST_SUITE_END();    // generate
diff --git a/verif/conformance/tosa_main_profile_ops_info.json b/verif/conformance/tosa_main_profile_ops_info.json
index a53d0c7..b8efd35 100644
--- a/verif/conformance/tosa_main_profile_ops_info.json
+++ b/verif/conformance/tosa_main_profile_ops_info.json
@@ -2702,6 +2702,7 @@
         "profile": [
             "tosa-mi"
         ],
+        "support_for": [ "lazy_data_gen" ],
         "generation": {
             "standard": {
                 "generator_args": [
@@ -2709,13 +2710,13 @@
                         "--target-dtype",
                         "fp32",
                         "--fp-values-range",
-                        "-2.0,2.0"
+                        "-max,max"
                     ],
                     [
                         "--target-dtype",
                         "fp32",
                         "--fp-values-range",
-                        "-2.0,2.0",
+                        "-max,max",
                         "--target-shape",
                         "1,16,512",
                         "--target-shape",
diff --git a/verif/generator/tosa_arg_gen.py b/verif/generator/tosa_arg_gen.py
index f6a46b4..a4bced3 100644
--- a/verif/generator/tosa_arg_gen.py
+++ b/verif/generator/tosa_arg_gen.py
@@ -2821,6 +2821,31 @@ class TosaArgGen:
         # Return list of tuples: (arg_str, args_dict)
         return arg_list
 
+    @staticmethod
+    def agRFFT2d(testGen, opName, shapeList, dtype, error_name=None):
+        arg_list = []
+
+        shape = shapeList[0]
+        dot_products = gtu.product(shape)
+        ks = shape[1] * shape[2]  # H*W
+        args_dict = {
+            "dot_products": dot_products,
+            "shape": shape,
+            "ks": ks,
+            "acc_type": dtype,
+        }
+        arg_list.append(("", args_dict))
+
+        arg_list = TosaArgGen._add_data_generators(
+            testGen,
+            opName,
+            dtype,
+            arg_list,
+            error_name,
+        )
+        # Return list of tuples: (arg_str, args_dict)
+        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_test_gen.py b/verif/generator/tosa_test_gen.py
index 9c3cd32..d82f919 100644
--- a/verif/generator/tosa_test_gen.py
+++ b/verif/generator/tosa_test_gen.py
@@ -2588,10 +2588,14 @@ class TosaTestGen:
     def build_rfft2d(
         self,
         op,
-        val,
+        inputs,
+        args_dict,
         validator_fcns=None,
         error_name=None,
+        qinfo=None,
     ):
+        assert len(inputs) == 1
+        val = inputs[0]
         results = OutputShaper.rfft2dOp(self.ser, self.rng, val, error_name)
 
         input_names = [val.name]
@@ -2629,7 +2633,14 @@ class TosaTestGen:
         attr.RFFTAttribute(local_bound)
 
         self.ser.addOperator(op["op"], input_names, output_names, attr)
-        return results
+
+        compliance = []
+        for res in results:
+            compliance.append(
+                self.tensorComplianceMetaData(op, val.dtype, args_dict, res, error_name)
+            )
+
+        return TosaTestGen.BuildInfo(results, compliance)
 
     def build_shape_op(
         self, op, inputs, args_dict, validator_fcns=None, error_name=None, qinfo=None
@@ -4781,8 +4792,8 @@ class TosaTestGen:
             "build_fcn": (
                 build_rfft2d,
                 TosaTensorGen.tgRFFT2d,
-                TosaTensorValuesGen.tvgDefault,
-                None,
+                TosaTensorValuesGen.tvgLazyGenDefault,
+                TosaArgGen.agRFFT2d,
             ),
             "types": [DType.FP32],
             "error_if_validators": (
@@ -4795,6 +4806,9 @@ class TosaTestGen:
                 TosaErrorValidator.evKernelNotPowerOfTwo,
                 TosaErrorValidator.evFFTOutputShapeMismatch,
             ),
+            "data_gen": {
+                "fp": (gtu.DataGenType.DOT_PRODUCT,),
+            },
         },
         # Shape
         "add_shape": {