Fix RFFT bugs in VMLA. (#5308)

The method used in op_kernels_fft.h should be resize, not reserve.
reserve() does not modify the size variable in a vector, which causes
undefined-behavior.

Also adapt the RFFT python tests to e2e/xla_ops lit tests. The input and
expected output are derived from python tests.

note: I gave it a shot to IRFFT, but it did not work. And the purpose is
to decouple testing from python for #5299,
so this change focus on RFFT support.

integrations/tensorflow/e2e/BUILD

@@ -109,7 +109,6 @@ VMLA_FAILING = [
     "einsum_dynamic_test.py",
     "einsum_static_test.py",
     "einsum_vector_test.py",
-    "fft_test.py",  # TODO(natashaknk): Get diffs within tolerance
     "mandelbrot_test.py",  # TODO(silvasean): Get this working on IREE.
     "ring_buffer_test.py",  # TODO(b/148747011)
     "strings_test.py",

integrations/tensorflow/e2e/CMakeLists.txt

@@ -40,7 +40,6 @@ iree_e2e_cartesian_product_test_suite(
     "einsum_dynamic_test.py,iree_vmla,"
     "einsum_static_test.py,iree_vmla,"
     "einsum_vector_test.py,iree_vmla,"
-    "fft_test.py,iree_vmla,"
     "mandelbrot_test.py,iree_vmla,"
     "ring_buffer_test.py,iree_vmla,"
     "strings_test.py,iree_vmla,"

iree/modules/vmla/op_kernels_fft.h

@@ -127,7 +127,7 @@ struct Rfft {
     int element_count = real_src_buffer.size() / 2 + 1;
 
     std::vector<T> complex_output;
-    complex_output.reserve(element_count * 4);
+    complex_output.resize(element_count * 4);
 
     pffft_transform_ordered(fft_state, &real_src_buffer[0], &complex_output[0],
                             NULL, PFFFT_FORWARD);

iree/test/e2e/xla_ops/BUILD

@@ -95,6 +95,7 @@ iree_check_single_backend_test_suite(
             "dot_general.mlir",
             "exponential.mlir",
             "exponential_minus_one.mlir",
+            "fft.mlir",
             "finite.mlir",
             "floor.mlir",
             "gather.mlir",

iree/test/e2e/xla_ops/CMakeLists.txt

@@ -32,6 +32,7 @@ iree_check_single_backend_test_suite(
     "dot_general.mlir"
     "exponential.mlir"
     "exponential_minus_one.mlir"
+    "fft.mlir"
     "finite.mlir"
     "floor.mlir"
     "gather.mlir"

iree/test/e2e/xla_ops/fft.mlir

@@ -0,0 +1,27 @@
+// TODO(hanchung): Add other types of fft tests, e.g. fft, ifft, irfft.
+
+func @rfft_real() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[
+    9.0, 1.0, 4.5, -0.3, 10.0, -1.0, 5.5, 0.3, 299.0, 3.5, -0.777, 2.0, 1.7,
+    3.5, -4.5, 0.0, 9.0, 1.0, 4.5, -0.3, 10.0, -1.0, 5.5, 0.3, 299.0, 3.5,
+    -0.777, 2.0, 1.7, 3.5, -4.5, 0.0]> : tensor<32xf32>
+  %0 = "mhlo.fft"(%input) {
+    fft_length = dense<32> : tensor<1xi64>, fft_type = "RFFT"
+  } : (tensor<32xf32>) -> tensor<17xcomplex<f32>>
+  %real = "mhlo.real"(%0) : (tensor<17xcomplex<f32>>) -> tensor<17xf32>
+  check.expect_almost_eq_const(%real, dense<[666.8460, 0.0, -590.16925, 0.0, 593.4485, 0.0, -579.52875, 0.0, 629.95404, 0.0, -567.1126, 0.0, 591.75146, 0.0, -583.1894, 0.0, 630.846]> : tensor<17xf32>) : tensor<17xf32>
+  return
+}
+
+func @rfft_imag() attributes { iree.module.export } {
+  %input = iree.unfoldable_constant dense<[
+    9.0, 1.0, 4.5, -0.3, 10.0, -1.0, 5.5, 0.3, 299.0, 3.5, -0.777, 2.0, 1.7,
+    3.5, -4.5, 0.0, 9.0, 1.0, 4.5, -0.3, 10.0, -1.0, 5.5, 0.3, 299.0, 3.5,
+    -0.777, 2.0, 1.7, 3.5, -4.5, 0.0]> : tensor<32xf32>
+  %0 = "mhlo.fft"(%input) {
+    fft_length = dense<32> : tensor<1xi64>, fft_type = "RFFT"
+  } : (tensor<32xf32>) -> tensor<17xcomplex<f32>>
+  %imag = "mhlo.imag"(%0) : (tensor<17xcomplex<f32>>) -> tensor<17xf32>
+  check.expect_almost_eq_const(%imag, dense<[0.0, 0.0, -23.956373, 0.0, -10.254326, 0.0, -6.1443653, 0.0, -10.0, 0.0, 3.865515, 0.0, 0.63767385, 0.0, 52.453506, 0.0, 0.0]> : tensor<17xf32>) : tensor<17xf32>
+  return
+}