| 142 | |
| 143 | template <typename T, typename Context> |
| 144 | void ComplexKernel(const Context& dev_ctx, |
| 145 | const DenseTensor& x, |
| 146 | const DenseTensor& y, |
| 147 | DenseTensor* out) { |
| 148 | using C = phi::dtype::complex<T>; |
| 149 | using XPUComplexType = typename XPUComplexTypeTrait<T>::Type; |
| 150 | if (out->numel() == 0) { |
| 151 | dev_ctx.template Alloc<C>(out); |
| 152 | return; |
| 153 | } |
| 154 | auto x_dims = x.dims(); |
| 155 | auto y_dims = y.dims(); |
| 156 | auto out_dims = funcs::BroadcastTwoDims(x_dims, y_dims); |
| 157 | std::vector<int64_t> out_dims_vec = vectorize(out_dims); |
| 158 | |
| 159 | DenseTensor broadcasted_x, broadcasted_y; |
| 160 | const T* x_data = nullptr; |
| 161 | const T* y_data = nullptr; |
| 162 | |
| 163 | if (x_dims == out_dims) { |
| 164 | x_data = x.data<T>(); |
| 165 | } else { |
| 166 | broadcasted_x.Resize(out_dims); |
| 167 | dev_ctx.template Alloc<T>(&broadcasted_x); |
| 168 | ExpandKernel<T, Context>( |
| 169 | dev_ctx, x, phi::IntArray(out_dims_vec), &broadcasted_x); |
| 170 | x_data = broadcasted_x.data<T>(); |
| 171 | } |
| 172 | |
| 173 | if (y_dims == out_dims) { |
| 174 | y_data = y.data<T>(); |
| 175 | } else { |
| 176 | broadcasted_y.Resize(out_dims); |
| 177 | dev_ctx.template Alloc<T>(&broadcasted_y); |
| 178 | ExpandKernel<T, Context>( |
| 179 | dev_ctx, y, phi::IntArray(out_dims_vec), &broadcasted_y); |
| 180 | y_data = broadcasted_y.data<T>(); |
| 181 | } |
| 182 | |
| 183 | dev_ctx.template Alloc<C>(out); |
| 184 | int r = xfft_internal::xpu::combine_as_complex( |
| 185 | dev_ctx.x_context()->xpu_stream, |
| 186 | out->numel(), |
| 187 | x_data, |
| 188 | y_data, |
| 189 | reinterpret_cast<XPUComplexType*>(out->data<C>())); |
| 190 | PADDLE_ENFORCE_XPU_SUCCESS(r); |
| 191 | } |
| 192 | } // namespace phi |
| 193 | |
| 194 | PD_REGISTER_KERNEL(conj, |