Move real and imag op to phi (#39777)

* Move Real OP to phi * Move Imag OP to phi * Move Real and Imag InferShape to phi * Move Real and Imag to complex_kernel * Change PT_REGISTER_XXX to PD_REGISTER_XXX

Author: From00

paddle/fluid/operators/imag_op.cc

@@ -12,23 +12,17 @@ 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. */
 
-#include "paddle/fluid/operators/imag_op.h"
+#include "paddle/fluid/framework/infershape_utils.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace operators {
 
 class ImagOp : public framework::OperatorWithKernel {
  public:
  using framework::OperatorWithKernel::OperatorWithKernel;
-
- void InferShape(framework::InferShapeContext* ctx) const override {
- OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Imag");
- OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Imag");
-
- auto x_dims = ctx->GetInputDim("X");
- ctx->SetOutputDim("Out", x_dims);
- ctx->ShareLoD("X", "Out");
- }
 };
 
 class ImagOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -88,19 +82,13 @@ DECLARE_INPLACE_OP_INFERER(ImagGradOpInplaceInferer,
 } // namespace operators
 } // namespace paddle
 
+DELCARE_INFER_SHAPE_FUNCTOR(imag, ImagInferShapeFunctor,
+ PT_INFER_META(phi::UnchangedInferMeta));
+
 namespace ops = paddle::operators;
 
 REGISTER_OPERATOR(imag, ops::ImagOp, ops::ImagOpMaker,
  ops::ImagGradOpMaker<paddle::framework::OpDesc>,
- ops::ImagGradOpMaker<paddle::imperative::OpBase>);
+ ops::ImagGradOpMaker<paddle::imperative::OpBase>,
+ ImagInferShapeFunctor);
 REGISTER_OPERATOR(imag_grad, ops::ImagGradOp);
-
-REGISTER_OP_CPU_KERNEL(imag, ops::ImagKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<float>>,
- ops::ImagKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<double>>);
-REGISTER_OP_CPU_KERNEL(imag_grad,
- ops::ImagGradKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<float>>,
- ops::ImagGradKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<double>>);

paddle/fluid/operators/imag_op.cu

@@ -12,22 +12,17 @@ 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. */
 
-#include "paddle/fluid/operators/real_op.h"
+#include "paddle/fluid/framework/infershape_utils.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/phi/core/infermeta_utils.h"
+#include "paddle/phi/infermeta/unary.h"
 
 namespace paddle {
 namespace operators {
 
 class RealOp : public framework::OperatorWithKernel {
  public:
  using framework::OperatorWithKernel::OperatorWithKernel;
- void InferShape(framework::InferShapeContext* ctx) const override {
- OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Real");
- OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Real");
-
- auto x_dims = ctx->GetInputDim("X");
- ctx->SetOutputDim("Out", x_dims);
- ctx->ShareLoD("X", "Out");
- }
 };
 
 class RealOpMaker : public framework::OpProtoAndCheckerMaker {
@@ -87,19 +82,13 @@ DECLARE_INPLACE_OP_INFERER(RealGradOpInplaceInferer,
 } // namespace operators
 } // namespace paddle
 
+DELCARE_INFER_SHAPE_FUNCTOR(real, RealInferShapeFunctor,
+ PT_INFER_META(phi::UnchangedInferMeta));
+
 namespace ops = paddle::operators;
 
 REGISTER_OPERATOR(real, ops::RealOp, ops::RealOpMaker,
  ops::RealGradOpMaker<::paddle::framework::OpDesc>,
- ops::RealGradOpMaker<::paddle::imperative::OpBase>);
+ ops::RealGradOpMaker<::paddle::imperative::OpBase>,
+ RealInferShapeFunctor);
 REGISTER_OPERATOR(real_grad, ops::RealGradOp);
-
-REGISTER_OP_CPU_KERNEL(real, ops::RealKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<float>>,
- ops::RealKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<double>>);
-REGISTER_OP_CPU_KERNEL(real_grad,
- ops::RealGradKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<float>>,
- ops::RealGradKernel<paddle::platform::CPUDeviceContext,
- paddle::platform::complex<double>>);

paddle/fluid/operators/imag_op.h

@@ -0,0 +1,31 @@
+/* Copyright (c) 2021 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. */
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RealGradKernel(const Context& dev_ctx,
+ const DenseTensor& dout,
+ DenseTensor* dx);
+
+template <typename T, typename Context>
+void ImagGradKernel(const Context& dev_ctx,
+ const DenseTensor& dout,
+ DenseTensor* dx);
+
+} // namespace phi

paddle/fluid/operators/real_op.cc

@@ -50,4 +50,14 @@ DenseTensor Conj(const Context& dev_ctx, const DenseTensor& x) {
  return x;
 }
 
+template <typename T, typename DeviceContext>
+void RealKernel(const DeviceContext& dev_ctx,
+ const DenseTensor& x,
+ DenseTensor* out);
+
+template <typename T, typename DeviceContext>
+void ImagKernel(const DeviceContext& dev_ctx,
+ const DenseTensor& x,
+ DenseTensor* out);
+
 } // namespace phi

paddle/fluid/operators/real_op.cu

@@ -0,0 +1,33 @@
+// Copyright (c) 2022 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.
+
+#include "paddle/phi/kernels/complex_grad_kernel.h"
+#include "paddle/phi/kernels/impl/complex_grad_kernel_impl.h"
+
+#include "paddle/phi/common/complex.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(real_grad,
+ CPU,
+ ALL_LAYOUT,
+ phi::RealGradKernel,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>) {}
+
+PD_REGISTER_KERNEL(imag_grad,
+ CPU,
+ ALL_LAYOUT,
+ phi::ImagGradKernel,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>) {}

paddle/fluid/operators/real_op.h

@@ -31,3 +31,17 @@ PD_REGISTER_KERNEL(conj,
  double,
  int,
  int64_t) {}
+
+PD_REGISTER_KERNEL(real,
+ CPU,
+ ALL_LAYOUT,
+ phi::RealKernel,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>) {}
+
+PD_REGISTER_KERNEL(imag,
+ CPU,
+ ALL_LAYOUT,
+ phi::ImagKernel,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>) {}