Update

[ghstack-poisoned]

Author: swolchok

kernels/optimized/cpu/op_fft_r2c.cpp

@@ -11,6 +11,8 @@
 
 #include <pocketfft_hdronly.h>
 
+#include <optional>
+
 namespace torch::executor::native {
 
 // TODO: contents of this anonymous namespace are copy/pasted from

kernels/optimized/cpu/targets.bzl

@@ -27,7 +27,7 @@ _OPTIMIZED_ATEN_OPS = (
     op_target(name = "op_exp"),
     op_target(
         name = "op_fft_r2c",
-        deps = [] if runtime.is_oss else ["fbsource//third-party/pocketfft:pocketfft"],
+        deps = [] if runtime.is_oss else ["fbsource//third-party/pocket_fft:pocketfft"],
     ),
     op_target(name = "op_sigmoid"),
     op_target(

kernels/test/op_fft_r2c_test.cpp

@@ -8,6 +8,7 @@
 
 #include <executorch/kernels/test/FunctionHeaderWrapper.h> // Declares the operator
 #include <executorch/kernels/test/TestUtil.h>
+#include <executorch/kernels/test/supported_features.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_util.h>
@@ -130,6 +131,10 @@ TEST_F(OpFftR2cOutTest, MultipleDims) {
 }
 
 TEST_F(OpFftR2cOutTest, InvalidNorm) {
+  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
+    GTEST_SKIP() << "ATen MKL path does not validate norm";
+    return;
+  }
   auto invalid_norm = [this](int64_t norm) {
     test_dtype<float, ScalarType::Float, /* expect_failure = */ true>(norm);
   };
@@ -140,6 +145,10 @@ TEST_F(OpFftR2cOutTest, InvalidNorm) {
 }
 
 TEST_F(OpFftR2cOutTest, InvalidDim) {
+  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
+    GTEST_SKIP() << "ATen fails UBSAN";
+    return;
+  }
   auto negative_dim = [this]() {
     test_dtype<float, ScalarType::Float, /* expect_failure = */ true>(0, -1);
     test_dtype<float, ScalarType::Float, /* expect_failure = */ true>(0, 3);
@@ -150,6 +159,10 @@ TEST_F(OpFftR2cOutTest, InvalidDim) {
 
 // TODO: support this and patch test accordingly!
 TEST_F(OpFftR2cOutTest, TwoSidedIsNotSupported) {
+  if (torch::executor::testing::SupportedFeatures::get()->is_aten) {
+    GTEST_SKIP() << "ATen supports two-sided";
+    return;
+  }
   auto twosided = [this]() {
     test_dtype<double, ScalarType::Double, /* expect_failure = */ true>(
         0, 1, false);

kernels/test/targets.bzl

@@ -244,6 +244,7 @@ def define_common_targets():
     _common_op_test("op_exp_test", ["aten", "portable", "optimized"])
     _common_op_test("op_expand_copy_test", ["aten", "portable"])
     _common_op_test("op_expm1_test", ["aten", "portable"])
+    _common_op_test("op_fft_r2c_test", ["aten", "optimized"])
     _common_op_test("op_fill_test", ["aten", "portable"])
     _common_op_test("op_flip_test", ["aten", "portable"])
     _common_op_test("op_floor_divide_test", ["aten", "portable"])