chore: import upstream snapshot with attribution

paddle/phi/kernels/xpu/complex_grad_kernel.cc

@@ -0,0 +1,201 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifdef PADDLE_WITH_XPU_FFT
+#include "paddle/phi/kernels/complex_grad_kernel.h"
+#include "fft/cuComplex.h"
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/common/type_traits.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/expand_grad_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace xfft_internal::xpu {
+// just for declaration here, the real implementation is in libcufft.so
+template <typename T, typename TComplex>
+int combine_as_complex(
+    const XPUStream stream, int N, const T* real, const T* imag, TComplex* out);
+template <>
+int combine_as_complex(const XPUStream stream,
+                       int N,
+                       const float* real,
+                       const float* imag,
+                       float2* out);
+template <>
+int combine_as_complex(const XPUStream stream,
+                       int N,
+                       const double* real,
+                       const double* imag,
+                       double2* out);
+
+template <typename TComplex, typename T>
+int complex_spilt(
+    const XPUStream stream, int N, const TComplex* in, T* real, T* imag);
+template <>
+int complex_spilt(
+    const XPUStream stream, int N, const float2* in, float* real, float* imag);
+template <>
+int complex_spilt(const XPUStream stream,
+                  int N,
+                  const double2* in,
+                  double* real,
+                  double* imag);
+}  // namespace xfft_internal::xpu
+
+namespace phi {
+
+template <class T, class Context>
+static DenseTensor Fill(const Context& dev_ctx,
+                        std::vector<int> shape,
+                        T fill_value) {
+  DenseTensor ret;
+  ret.Resize(shape);
+  dev_ctx.template Alloc<T>(&ret);
+  funcs::SetConstant<Context, T>()(dev_ctx, &ret, fill_value);
+  return ret;
+}
+
+template <typename T, typename Context>
+void RealGradKernel(const Context& dev_ctx,
+                    const DenseTensor& dout,
+                    DenseTensor* dx) {
+  using XPUComplexType =
+      typename XPUComplexTypeTrait<phi::dtype::Real<T>>::Type;
+  if (dx && dx->numel() == 0) {
+    dev_ctx.template Alloc<T>(dx);
+    return;
+  }
+  auto numel = dout.numel();
+  auto* dx_data =
+      dev_ctx.template Alloc<T>(dx, static_cast<size_t>(numel * sizeof(T)));
+  DenseTensor imag = Fill<phi::dtype::Real<T>, Context>(
+      dev_ctx, vectorize<int>(dout.dims()), phi::dtype::Real<T>(0.0));
+  int r = xfft_internal::xpu::combine_as_complex(
+      dev_ctx.x_context()->xpu_stream,
+      numel,
+      reinterpret_cast<const phi::dtype::Real<T>*>(
+          dout.data<phi::dtype::Real<T>>()),
+      imag.data<phi::dtype::Real<T>>(),
+      reinterpret_cast<XPUComplexType*>(dx_data));
+  PADDLE_ENFORCE_XPU_SUCCESS(r);
+}
+
+template <typename T, typename Context>
+void ImagGradKernel(const Context& dev_ctx,
+                    const DenseTensor& dout,
+                    DenseTensor* dx) {
+  using XPUComplexType =
+      typename XPUComplexTypeTrait<phi::dtype::Real<T>>::Type;
+  if (dx && dx->numel() == 0) {
+    dev_ctx.template Alloc<T>(dx);
+    return;
+  }
+  auto numel = dout.numel();
+  auto* dx_data =
+      dev_ctx.template Alloc<T>(dx, static_cast<size_t>(numel * sizeof(T)));
+  DenseTensor real = Fill<phi::dtype::Real<T>, Context>(
+      dev_ctx, vectorize<int>(dout.dims()), phi::dtype::Real<T>(0.0));
+  int r = xfft_internal::xpu::combine_as_complex(
+      dev_ctx.x_context()->xpu_stream,
+      numel,
+      real.data<phi::dtype::Real<T>>(),
+      reinterpret_cast<const phi::dtype::Real<T>*>(
+          dout.data<phi::dtype::Real<T>>()),
+      reinterpret_cast<XPUComplexType*>(dx_data));
+  PADDLE_ENFORCE_XPU_SUCCESS(r);
+}
+
+template <typename T, typename Context>
+void ComplexGradKernel(const Context& dev_ctx,
+                       const DenseTensor& x,
+                       const DenseTensor& y,
+                       const DenseTensor& dout,
+                       DenseTensor* dx,
+                       DenseTensor* dy) {
+  using C = phi::dtype::complex<T>;
+  using XPUComplexType = typename XPUComplexTypeTrait<T>::Type;
+  if (dout.numel() == 0) {
+    if (dx) {
+      if (dx->numel() == 0) {
+        dev_ctx.template Alloc<T>(dx);
+      } else {
+        Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
+      }
+    }
+    if (dy) {
+      if (dy->numel() == 0) {
+        dev_ctx.template Alloc<T>(dy);
+      } else {
+        Full<T, Context>(dev_ctx, dy->dims(), 0, dy);
+      }
+    }
+    return;
+  }
+  auto numel = dout.numel();
+  DenseTensor real_dout, imag_dout;
+  real_dout.Resize(dout.dims());
+  imag_dout.Resize(dout.dims());
+  T* real_data = dev_ctx.template Alloc<T>(&real_dout);
+  T* imag_data = dev_ctx.template Alloc<T>(&imag_dout);
+  int r = xfft_internal::xpu::complex_spilt(
+      dev_ctx.x_context()->xpu_stream,
+      numel,
+      reinterpret_cast<const XPUComplexType*>(dout.data<C>()),
+      real_data,
+      imag_data);
+  PADDLE_ENFORCE_XPU_SUCCESS(r);
+  if (dx) {
+    if (x.dims() == dout.dims()) {
+      dx->ShareDataWith(real_dout);
+    } else {
+      ExpandGradKernel<T, Context>(
+          dev_ctx, x, real_dout, phi::IntArray(vectorize(x.dims())), dx);
+    }
+  }
+
+  if (dy) {
+    if (y.dims() == dout.dims()) {
+      dy->ShareDataWith(imag_dout);
+    } else {
+      ExpandGradKernel<T, Context>(
+          dev_ctx, y, imag_dout, phi::IntArray(vectorize(y.dims())), dy);
+    }
+  }
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(imag_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::ImagGradKernel,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
+}
+
+PD_REGISTER_KERNEL(real_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::RealGradKernel,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetDataType(phi::dtype::ToReal(kernel_key.dtype()));
+}
+
+PD_REGISTER_KERNEL(
+    complex_grad, XPU, ALL_LAYOUT, phi::ComplexGradKernel, float, double) {
+  kernel->InputAt(2).SetDataType(phi::dtype::ToComplex(kernel_key.dtype()));
+}
+#endif  // PADDLE_WITH_XPU_FFT