| 119 | |
| 120 | template <typename T, typename Context> |
| 121 | void ComplexGradKernel(const Context& dev_ctx, |
| 122 | const DenseTensor& x, |
| 123 | const DenseTensor& y, |
| 124 | const DenseTensor& dout, |
| 125 | DenseTensor* dx, |
| 126 | DenseTensor* dy) { |
| 127 | using C = phi::dtype::complex<T>; |
| 128 | using XPUComplexType = typename XPUComplexTypeTrait<T>::Type; |
| 129 | if (dout.numel() == 0) { |
| 130 | if (dx) { |
| 131 | if (dx->numel() == 0) { |
| 132 | dev_ctx.template Alloc<T>(dx); |
| 133 | } else { |
| 134 | Full<T, Context>(dev_ctx, dx->dims(), 0, dx); |
| 135 | } |
| 136 | } |
| 137 | if (dy) { |
| 138 | if (dy->numel() == 0) { |
| 139 | dev_ctx.template Alloc<T>(dy); |
| 140 | } else { |
| 141 | Full<T, Context>(dev_ctx, dy->dims(), 0, dy); |
| 142 | } |
| 143 | } |
| 144 | return; |
| 145 | } |
| 146 | auto numel = dout.numel(); |
| 147 | DenseTensor real_dout, imag_dout; |
| 148 | real_dout.Resize(dout.dims()); |
| 149 | imag_dout.Resize(dout.dims()); |
| 150 | T* real_data = dev_ctx.template Alloc<T>(&real_dout); |
| 151 | T* imag_data = dev_ctx.template Alloc<T>(&imag_dout); |
| 152 | int r = xfft_internal::xpu::complex_spilt( |
| 153 | dev_ctx.x_context()->xpu_stream, |
| 154 | numel, |
| 155 | reinterpret_cast<const XPUComplexType*>(dout.data<C>()), |
| 156 | real_data, |
| 157 | imag_data); |
| 158 | PADDLE_ENFORCE_XPU_SUCCESS(r); |
| 159 | if (dx) { |
| 160 | if (x.dims() == dout.dims()) { |
| 161 | dx->ShareDataWith(real_dout); |
| 162 | } else { |
| 163 | ExpandGradKernel<T, Context>( |
| 164 | dev_ctx, x, real_dout, phi::IntArray(vectorize(x.dims())), dx); |
| 165 | } |
| 166 | } |
| 167 | |
| 168 | if (dy) { |
| 169 | if (y.dims() == dout.dims()) { |
| 170 | dy->ShareDataWith(imag_dout); |
| 171 | } else { |
| 172 | ExpandGradKernel<T, Context>( |
| 173 | dev_ctx, y, imag_dout, phi::IntArray(vectorize(y.dims())), dy); |
| 174 | } |
| 175 | } |
| 176 | } |
| 177 | } // namespace phi |
| 178 | |