Enable fft_c2r

Author: CuiYifeng

src/ATen/native/xpu/SpectralOps.cpp

@@ -42,4 +42,25 @@ Tensor& _fft_c2c_xpu_out(
 #endif // USE_ONEMKL
 }
 
+Tensor _fft_c2r_xpu(
+    const Tensor& self,
+    IntArrayRef dim,
+    int64_t normalization,
+    int64_t last_dim_size) {
+  TORCH_CHECK(self.is_complex());
+
+  return native::xpu::_fft_c2r_mkl(self, dim, normalization, last_dim_size);
+}
+
+Tensor _fft_c2r_xpu_out(
+    const Tensor& self,
+    IntArrayRef dim,
+    int64_t normalization,
+    int64_t last_dim_size,
+    Tensor& out) {
+  TORCH_CHECK(self.is_complex());
+
+  return native::xpu::_fft_c2r_mkl_out(self, dim, normalization, last_dim_size, out);
+}
+
 } // namespace at::native

src/ATen/native/xpu/XPUFallback.template

@@ -190,7 +190,6 @@ TORCH_LIBRARY_IMPL(aten, XPU, m) {
     "_cholesky_solve_helper",
     "dot",
     "_efficient_attention_forward",
-    "_fft_c2r",
     "_fft_r2c",
     "_flash_attention_forward",
     "geqrf",

src/ATen/native/xpu/mkl/SpectralOps.cpp

@@ -2,7 +2,11 @@
 #include <ATen/native/Resize.h>
 #include <ATen/native/SpectralOpsUtils.h>
 #include <ATen/native/xpu/mkl/SpectralOps.h>
+#include <ATen/ops/complex.h>
+#include <ATen/ops/imag.h>
 #include <ATen/ops/mul.h>
+#include <ATen/ops/real.h>
+#include <ATen/ops/zeros_like.h>
 #include <comm/SYCLContext.h>
 #include <comm/TensorInfo.h>
 #include <oneapi/mkl.hpp>
@@ -84,8 +88,7 @@ void _mkl_dft(
   }
 
   if (!complex_input || !complex_output) {
-    desc.set_value(
-        config_param::CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX);
+    desc.set_value(config_param::CONJUGATE_EVEN_STORAGE, DFTI_COMPLEX_COMPLEX);
   }
 
   desc.set_value(
@@ -398,5 +401,81 @@ Tensor& _fft_c2c_mkl_out(
       out, result, normalization, result.sizes(), dim);
 }
 
+void HermitSymmImpl(Tensor& input, int64_t dim, int pos) {
+  std::vector<at::indexing::TensorIndex> indices(
+      input.dim(), at::indexing::Slice());
+
+  indices[dim] = pos;
+
+  Tensor values = at::complex(
+      at::real(input.index(indices)),
+      at::zeros_like(at::imag(input.index(indices))));
+
+  input.index_put_(indices, values);
+}
+
+void HermitSymm(Tensor& input, int64_t dim, int64_t out_size) {
+  HermitSymmImpl(input, dim, 0);
+
+  if (out_size % 2 == 0)
+    HermitSymmImpl(input, dim, -1);
+}
+
+Tensor _fft_c2r_mkl(
+    const Tensor& self,
+    IntArrayRef dim,
+    int64_t normalization,
+    int64_t last_dim_size) {
+  if (dim.empty()) {
+    return self.clone();
+  }
+
+  auto input = self;
+
+  if (dim.size() > 1) {
+    auto c2c_dims = dim.slice(0, dim.size() - 1);
+    input = _fft_c2c_mkl(
+        self,
+        c2c_dims,
+        static_cast<int64_t>(fft_norm_mode::none),
+        /*forward=*/false);
+  }
+
+  auto in_sizes = input.sizes();
+  DimVector out_sizes(in_sizes.begin(), in_sizes.end());
+  out_sizes[dim.back()] = last_dim_size;
+
+  auto out = at::empty(
+      out_sizes,
+      self.options().dtype(c10::toRealValueType(self.scalar_type())));
+
+  input = input.clone(MemoryFormat::Contiguous);
+
+  HermitSymm(input, dim.back(), out_sizes[dim.back()]);
+
+  impl::_exec_fft(
+      out,
+      input,
+      out_sizes,
+      dim.back(),
+      /*onesided=*/true,
+      /*forward=*/false);
+
+  return impl::_fft_apply_normalization(out, normalization, out_sizes, dim);
+}
+
+Tensor& _fft_c2r_mkl_out(
+    const Tensor& self,
+    IntArrayRef dim,
+    int64_t normalization,
+    int64_t last_dim_size,
+    Tensor& out) {
+  auto result = _fft_c2c_mkl(
+      self, dim, static_cast<int64_t>(fft_norm_mode::none), last_dim_size);
+  at::native::resize_output(out, result.sizes());
+  return impl::_fft_apply_normalization_out(
+      out, result, normalization, result.sizes(), dim);
+}
+
 } // namespace at::native::xpu
 #endif // USE_ONEMKL

src/ATen/native/xpu/mkl/SpectralOps.h

@@ -17,4 +17,17 @@ TORCH_XPU_API Tensor& _fft_c2c_mkl_out(
     bool forward,
     Tensor& out);
 
+TORCH_XPU_API Tensor _fft_c2r_mkl(
+    const Tensor& self,
+    IntArrayRef dim,
+    int64_t normalization,
+    int64_t last_dim_size);
+
+TORCH_XPU_API Tensor& _fft_c2r_mkl_out(
+    const Tensor& self,
+    IntArrayRef dim,
+    int64_t normalization,
+    int64_t last_dim_size,
+    Tensor& out);
+
 } // namespace at::native::xpu

yaml/native/native_functions.yaml

@@ -9322,3 +9322,14 @@
   variants: function
   dispatch:
     XPU: _fft_c2c_xpu_out
+
+# Complex to real inverse FFT
+- func: _fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor
+  variants: function
+  dispatch:
+    XPU: _fft_c2r_xpu
+
+- func: _fft_c2r.out(Tensor self, int[] dim, int normalization, SymInt last_dim_size, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  dispatch:
+    XPU: _fft_c2r_xpu_out