Enhanced fft_r2c overhead and reactivate fft_r2c on XPU (#1720)

# Motivation
Revert `fft_r2c` fallback after root causing the `TestInductorOpInfoXPU`
failure. There is a memory layout mismatching between `aten::fft_r2c`
and Inductor meta deducing. We will have another PR to correct Inductor
meta deducing for XPU backend.

# Solution
Inconsistent strides should be resolved in `fft_r2c` meta calculation
instead of `fft_r2c` XPU kernel. The change of this PR includes:
- Reduced `_fft_r2c_mkl` overhead and simplified kernel code
- Reactivate `fft_r2c` on XPU

Author: CuiYifeng

src/ATen/native/xpu/SpectralOps.cpp

@@ -1,10 +1,10 @@
-#include <ATen/native/Resize.h>
-#include <ATen/ops/_fft_r2c_native.h>
 #if defined(USE_ONEMKL_XPU)
 #include <ATen/native/xpu/mkl/SpectralOps.h>
 #else
+#include <ATen/native/Resize.h>
 #include <ATen/ops/_fft_c2c_native.h>
 #include <ATen/ops/_fft_c2r_native.h>
+#include <ATen/ops/_fft_r2c_native.h>
 #endif // USE_ONEMKL_XPU
 
 namespace at::native {
@@ -87,9 +87,13 @@ Tensor _fft_r2c_xpu(
     bool onesided) {
   TORCH_CHECK(self.is_floating_point());
 
+#if defined(USE_ONEMKL_XPU)
+  return native::xpu::_fft_r2c_mkl(self, dim, normalization, onesided);
+#else
   Tensor out_cpu = native::_fft_r2c_mkl(
       self.to(Device(at::kCPU)), dim, normalization, onesided);
   return out_cpu.to(Device(at::kXPU));
+#endif // USE_ONEMKL_XPU
 }
 
 Tensor& _fft_r2c_xpu_out(
@@ -100,11 +104,15 @@ Tensor& _fft_r2c_xpu_out(
     Tensor& out) {
   TORCH_CHECK(self.is_floating_point());
 
+#if defined(USE_ONEMKL_XPU)
+  return native::xpu::_fft_r2c_mkl_out(self, dim, normalization, onesided, out);
+#else
   Tensor out_cpu = native::_fft_r2c_mkl(
       self.to(Device(at::kCPU)), dim, normalization, onesided);
   at::native::resize_output(out, out_cpu.sizes());
   out.copy_(out_cpu);
   return out;
+#endif // USE_ONEMKL_XPU
 }
 
 } // namespace at::native

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

@@ -499,11 +499,10 @@ Tensor _fft_r2c_mkl(
 
   IntArrayRef out_sizes = onesided ? onesided_sizes : input_sizes;
 
-  auto sorted_dims = impl::_sort_dims(self, dim, /*exclude_last=*/true);
   auto out = at::empty(
       out_sizes, self.options().dtype(c10::toComplexType(self.scalar_type())));
 
-  auto working_tensor = self.clone(MemoryFormat::Contiguous);
+  auto working_tensor = self.contiguous();
 
   // First do the R2C transform on the last dimension
   impl::_exec_fft(
@@ -515,17 +514,12 @@ Tensor _fft_r2c_mkl(
         self.options().dtype(c10::toComplexType(self.scalar_type())));
   }
 
-  sorted_dims.resize(sorted_dims.size() - 1);
+  DimVector sorted_dims(dim.begin(), dim.end() - 1);
 
   while (!sorted_dims.empty()) {
-    if (working_tensor.is_same(self)) {
-      working_tensor = std::move(out);
-      out = at::empty(
-          out_sizes,
-          self.options().dtype(c10::toComplexType(self.scalar_type())));
-    } else {
-      std::swap(out, working_tensor);
-    }
+    sorted_dims = impl::_sort_dims(self, sorted_dims);
+
+    std::swap(out, working_tensor);
 
     const auto max_dims =
         std::min(static_cast<size_t>(impl::mkl_max_ndim), sorted_dims.size());
@@ -539,18 +533,13 @@ Tensor _fft_r2c_mkl(
         onesided,
         /*forward=*/true);
     sorted_dims.resize(sorted_dims.size() - max_dims);
-
-    if (sorted_dims.empty()) {
-      break;
-    }
-
-    sorted_dims = impl::_sort_dims(self, sorted_dims);
   }
 
   // Only need to normalize the onesided slice since data in the other half is
   // overwritten
   auto out_slice = out.slice(last_dim, 0, last_dim_halfsize);
-  working_tensor = self;
+  impl::_fft_apply_normalization(out_slice, normalization, input_sizes, dim);
+
   if (!onesided) {
     if (out.sizes()[last_dim] != out_sizes[last_dim]) {
       working_tensor.resize_(out_sizes, MemoryFormat::Contiguous);
@@ -560,7 +549,7 @@ Tensor _fft_r2c_mkl(
     at::native::_fft_fill_with_conjugate_symmetry_(out, dim);
   }
 
-  return impl::_fft_apply_normalization(out, normalization, input_sizes, dim);
+  return out;
 }
 
 Tensor& _fft_r2c_mkl_out(