FFT, RFFT: Improve reference pseudo-code

Add types for the angle caclulations.
Add conformance appendix section for RFFT2D.

Signed-off-by: Dominic Symes <dominic.symes@arm.com>
Change-Id: Ic0ab79c0784c6359d2d66290e124dfb670a879cc

3 files changed, 28 insertions, 4 deletions

diff --git a/chapters/appendix_a.adoc b/chapters/appendix_a.adoc
index b162738..6d85204 100644
--- a/chapters/appendix_a.adoc
+++ b/chapters/appendix_a.adoc
@@ -275,6 +275,24 @@ for (0 <= y < H, 0 <= x < W, 0 <= m < H, 0 <= n < W) {
 }
 ----
 
+==== RFFT2D
+
+The following generates input test data for test set S.
+For compliant implementation, the test must pass whenever the attributes satisfy:
+`N*H*W >= MIN_DOT_PRODUCTS`
+
+[source,c++]
+----
+KS = H*W;
+for (0 <= n < N, 0 <= y < H, 0 <= x < W) {
+  input_real[n, y, x] = tosa_mi_data(S, KS, 0, y*W+x, ((0*N+n)*H+y)*IW+x);
+}
+for (0 <= y < H, 0 <= x < W, 0 <= m < H, 0 <= n < W) {
+  weight_real[y, x, m, n] = real(exp(2*pi*i*((m*h/H) + (n*w/W))));
+  weight_imag[y, x, m, n] = imag(exp(2*pi*i*((m*h/H) + (n*w/W))));
+}
+----
+
 ==== REDUCE_SUM
 
 The following generates input test data for test set S.
diff --git a/pseudocode/operators/FFT2D.tosac b/pseudocode/operators/FFT2D.tosac
index a958aa4..e8f0243 100644
--- a/pseudocode/operators/FFT2D.tosac
+++ b/pseudocode/operators/FFT2D.tosac
@@ -22,8 +22,10 @@ for_each(0 <= n < N, 0 <= oy < H, 0 <= ox < W) {
     for_each(0 <= iy < H, 0 <= ix < W) {
         in_out_t val_real = tensor_read<in_out_t>(input_real, [N,H,W], [n,iy,ix]);
         in_out_t val_imag = tensor_read<in_out_t>(input_imag, [N,H,W], [n,iy,ix]);
-        float_t a = sign_val * 2 * pi() * ((iy * oy) / H + (ix * ox) / W);
-        sum_real += val_real * cos(a) + val_imag * sin(a);
+        int32_t ay = (static_cast<size_t>(iy) * static_cast<size_t>(oy)) % static_cast<int32_t>(H);
+        int32_t ax = (static_cast<size_t>(ix) * static_cast<size_t>(ox)) % static_cast<int32_t>(W);
+        in_out_t a = sign_val * 2 * pi() * (static_cast<in_out_t>(ay) / H + static_cast<in_out_t>(ax) / W);
+        sum_real +=  val_real * cos(a) + val_imag * sin(a);
         sum_imag += -val_real * sin(a) + val_imag * cos(a);
     }
     tensor_write<in_out_t>(output_real, [N,H,W], [n,oy,ox], sum_real);
diff --git a/pseudocode/operators/RFFT2D.tosac b/pseudocode/operators/RFFT2D.tosac
index f664826..c4dfab2 100644
--- a/pseudocode/operators/RFFT2D.tosac
+++ b/pseudocode/operators/RFFT2D.tosac
@@ -15,9 +15,13 @@ for_each(0 <= n < N, 0 <= oy < H, 0 <= ox < W/2 + 1) {
     in_out_t sum_imag = 0.0;
     for_each(0 <= iy < H, 0 <= ix < W) {
         in_out_t val_real = tensor_read<in_out_t>(input_real, [N,H,W], [n,iy,ix]);
-        float_t a = 2 * pi() * ((iy * oy) / H + (ix * ox) / W);
+        int32_t ay = (static_cast<size_t>(iy) * static_cast<size_t>(oy)) % static_cast<int32_t>(H);
+        int32_t ax = (static_cast<size_t>(ix) * static_cast<size_t>(ox)) % static_cast<int32_t>(W);
+        in_out_t a = sign_val * 2 * pi() * (static_cast<in_out_t>(ay) / H + static_cast<in_out_t>(ax) / W);
         sum_real += val_real * cos(a);
-        sum_imag += -val_real * sin(a);
+        if ((ay % (H/2)) + (ax % (W/2)) > 0) {
+            sum_imag += -val_real * sin(a);
+        }
     }
     tensor_write<in_out_t>(output_real, [N,H,W], [n,oy,ox], sum_real);
     tensor_write<in_out_t>(output_imag, [N,H,W], [n,oy,ox], sum_imag);