migrate unsqueeze kernels to phi, test=kunlun (#45673)

paddle/fluid/operators/unsqueeze_op_xpu.cc

-/* Copyright (c) 2019 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/fluid/operators/unsqueeze_op.h"
-#ifdef PADDLE_WITH_XPU
-namespace ops = paddle::operators;
-namespace plat = paddle::platform;
-
-REGISTER_OP_XPU_KERNEL(
-    unsqueeze,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, int64_t>);
-REGISTER_OP_XPU_KERNEL(
-    unsqueeze_grad,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::UnsqueezeGradKernel<paddle::platform::XPUDeviceContext, int64_t>);
-REGISTER_OP_XPU_KERNEL(
-    unsqueeze2,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::UnsqueezeKernel<paddle::platform::XPUDeviceContext, int64_t>);
-REGISTER_OP_XPU_KERNEL(
-    unsqueeze2_grad,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext,
-                              plat::float16>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::Unsqueeze2GradKernel<paddle::platform::XPUDeviceContext, int64_t>);
-
-#endif

paddle/fluid/operators/where_op_xpu.cc

-// Copyright (c) 2020 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
-
-#include "paddle/fluid/framework/op_registry.h"
-
-namespace paddle {
-namespace operators {
-
-template <typename DeviceContext, typename T>
-class WhereXPUKernel : public framework::OpKernel<T> {
- public:
-  void Compute(const framework::ExecutionContext& context) const override {
-    auto* condition = context.Input<framework::Tensor>("Condition");
-    auto* X = context.Input<framework::Tensor>("X");
-    auto* Y = context.Input<framework::Tensor>("Y");
-    auto* out = context.Output<framework::Tensor>("Out");
-
-    const bool* cond_data = condition->data<bool>();
-    const T* x_data = X->data<T>();
-    const T* y_data = Y->data<T>();
-    T* out_data = out->mutable_data<T>(context.GetPlace());
-
-    auto cond_dims = phi::vectorize<int>(condition->dims());
-    auto input_dims = phi::vectorize<int>(X->dims());
-
-    auto& dev_ctx = context.template device_context<DeviceContext>();
-    int ret = xpu::select(dev_ctx.x_context(),
-                          cond_data,
-                          x_data,
-                          y_data,
-                          out_data,
-                          cond_dims,
-                          input_dims);
-    PADDLE_ENFORCE_EQ(ret,
-                      XPU_SUCCESS,
-                      platform::errors::External(
-                          "XPU select kernel return wrong value[%d %s]",
-                          ret,
-                          XPUAPIErrorMsg[ret]));
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-
-REGISTER_OP_XPU_KERNEL(
-    where,
-    ops::WhereXPUKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::WhereXPUKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::WhereXPUKernel<paddle::platform::XPUDeviceContext, int64_t>);
-#endif

paddle/phi/kernels/cpu/unsqueeze_grad_kernel.cc

-// 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/unsqueeze_grad_kernel.h"
-
-#include "paddle/phi/backends/cpu/cpu_context.h"
-#include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/unsqueeze_grad_kernel_impl.h"
-
-PD_REGISTER_KERNEL(unsqueeze_grad,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::UnsqueezeGradKernel,
-                   phi::dtype::bfloat16,
-                   bool,
-                   int,
-                   int16_t,
-                   uint8_t,
-                   int8_t,
-                   int64_t,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>,
-                   float,
-                   double) {}

paddle/phi/kernels/cpu/unsqueeze_kernel.cc

-// 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/unsqueeze_kernel.h"
-
-#include "paddle/phi/backends/cpu/cpu_context.h"
-#include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/unsqueeze_kernel_impl.h"
-
-PD_REGISTER_KERNEL(unsqueeze,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::UnsqueezeKernel,
-                   float,
-                   double,
-                   phi::dtype::bfloat16,
-                   bool,
-                   int,
-                   int16_t,
-                   uint8_t,
-                   int8_t,
-                   int64_t,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>) {}
-
-PD_REGISTER_KERNEL(unsqueeze_with_xshape,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::UnsqueezeWithXShapeKernel,
-                   float,
-                   double,
-                   phi::dtype::bfloat16,
-                   bool,
-                   int,
-                   int16_t,
-                   uint8_t,
-                   int8_t,
-                   int64_t,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/unsqueeze_grad_kernel.cu

-// 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/unsqueeze_grad_kernel.h"
-
-#include "paddle/phi/backends/gpu/gpu_context.h"
-#include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/unsqueeze_grad_kernel_impl.h"
-
-PD_REGISTER_KERNEL(unsqueeze_grad,
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::UnsqueezeGradKernel,
-                   phi::dtype::bfloat16,
-                   phi::dtype::float16,
-                   bool,
-                   int,
-                   int16_t,
-                   uint8_t,
-                   int8_t,
-                   int64_t,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>,
-                   float,
-                   double) {}

paddle/phi/kernels/gpu/unsqueeze_kernel.cu → paddle/phi/kernels/unsqueeze_grad_kernel.cc

@@ -12,19 +12,33 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/unsqueeze_kernel.h"
+#include "paddle/phi/kernels/unsqueeze_grad_kernel.h"
 
-#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/all_context.h"
 #include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/unsqueeze_kernel_impl.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/unsqueeze.h"
 
-PD_REGISTER_KERNEL(unsqueeze,
-                   GPU,
+namespace phi {
+template <typename T, typename Context>
+void UnsqueezeGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x_shape,
+                         const DenseTensor& dout,
+                         DenseTensor* dx) {
+  auto xshape_dims = x_shape.dims();
+  auto x_dims = phi::slice_ddim(xshape_dims, 1, xshape_dims.size());
+  dev_ctx.template Alloc<T>(dx);
+  phi::Copy(dev_ctx, dout, dev_ctx.GetPlace(), true, dx);
+  dx->Resize(x_dims);
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(unsqueeze_grad,
+                   CPU,
                    ALL_LAYOUT,
-                   phi::UnsqueezeKernel,
+                   phi::UnsqueezeGradKernel,
                    float,
                    double,
-                   phi::dtype::float16,
                    phi::dtype::bfloat16,
                    bool,
                    int,
@@ -35,10 +49,11 @@ PD_REGISTER_KERNEL(unsqueeze,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
 
-PD_REGISTER_KERNEL(unsqueeze_with_xshape,
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(unsqueeze_grad,
                    GPU,
                    ALL_LAYOUT,
-                   phi::UnsqueezeWithXShapeKernel,
+                   phi::UnsqueezeGradKernel,
                    float,
                    double,
                    phi::dtype::float16,
@@ -51,3 +66,21 @@ PD_REGISTER_KERNEL(unsqueeze_with_xshape,
                    int64_t,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
+
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL(unsqueeze_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeGradKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t) {}
+
+#endif

paddle/phi/kernels/impl/unsqueeze_kernel_impl.h → paddle/phi/kernels/unsqueeze_kernel.cc

@@ -12,8 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#pragma once
-#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/unsqueeze_kernel.h"
+
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_utils.h"
 #include "paddle/phi/kernels/funcs/unsqueeze.h"
 
@@ -48,3 +50,98 @@ void UnsqueezeWithXShapeKernel(const Context& dev_ctx,
   UnsqueezeKernel<T, Context>(dev_ctx, x, axes, out);
 }
 }  // namespace phi
+
+PD_REGISTER_KERNEL(unsqueeze,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeKernel,
+                   float,
+                   double,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   int16_t,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
+PD_REGISTER_KERNEL(unsqueeze_with_xshape,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeWithXShapeKernel,
+                   float,
+                   double,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   int16_t,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(unsqueeze,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   int16_t,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
+PD_REGISTER_KERNEL(unsqueeze_with_xshape,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeWithXShapeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   int16_t,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL(unsqueeze,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(unsqueeze_with_xshape,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::UnsqueezeWithXShapeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t) {}
+#endif

paddle/phi/kernels/impl/unsqueeze_grad_kernel_impl.h → paddle/phi/kernels/xpu/where_kernel.cc

@@ -12,20 +12,38 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#pragma once
-#include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/where_kernel.h"
+
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/kernel_registry.h"
 
 namespace phi {
+
 template <typename T, typename Context>
-void UnsqueezeGradKernel(const Context& dev_ctx,
-                         const DenseTensor& x_shape,
-                         const DenseTensor& dout,
-                         DenseTensor* dx) {
-  auto xshape_dims = x_shape.dims();
-  auto x_dims = phi::slice_ddim(xshape_dims, 1, xshape_dims.size());
-  dev_ctx.template Alloc<T>(dx);
-  phi::Copy(dev_ctx, dout, dev_ctx.GetPlace(), true, dx);
-  dx->Resize(x_dims);
+void WhereKernel(const Context& ctx,
+                 const DenseTensor& condition,
+                 const DenseTensor& x,
+                 const DenseTensor& y,
+                 DenseTensor* out) {
+  const bool* cond_data = condition.data<bool>();
+  const T* x_data = x.data<T>();
+  const T* y_data = y.data<T>();
+  T* out_data = ctx.template Alloc<T>(out);
+
+  auto cond_dims = phi::vectorize<int>(condition.dims());
+  auto input_dims = phi::vectorize<int>(x.dims());
+
+  int ret = xpu::select(ctx.x_context(),
+                        cond_data,
+                        x_data,
+                        y_data,
+                        out_data,
+                        cond_dims,
+                        input_dims);
+  PADDLE_ENFORCE_XDNN_SUCCESS(ret, "select");
 }
+
 }  // namespace phi
+
+PD_REGISTER_KERNEL(
+    where, XPU, ALL_LAYOUT, phi::WhereKernel, float, int, int64_t) {}