migrate squeeze kernel to phi, test=kunlun (#45968)

paddle/fluid/operators/squeeze_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/squeeze_op.h"
-#ifdef PADDLE_WITH_XPU
-
-namespace ops = paddle::operators;
-namespace plat = paddle::platform;
-
-REGISTER_OP_XPU_KERNEL(
-    squeeze,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::SqueezeKernel<paddle::platform::XPUDeviceContext, int64_t>);
-REGISTER_OP_XPU_KERNEL(
-    squeeze_grad,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::SqueezeGradKernel<paddle::platform::XPUDeviceContext, int64_t>);
-REGISTER_OP_XPU_KERNEL(
-    squeeze2,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, float>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, double>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, int>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::Squeeze2Kernel<paddle::platform::XPUDeviceContext, int64_t>);
-REGISTER_OP_XPU_KERNEL(
-    squeeze2_grad,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, float>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, double>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, plat::float16>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, bool>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, int>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, int8_t>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, uint8_t>,
-    ops::Squeeze2GradKernel<paddle::platform::XPUDeviceContext, int64_t>);
-
-#endif

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

paddle/phi/kernels/cpu/squeeze_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/squeeze_kernel.h"
-
-#include "paddle/phi/backends/cpu/cpu_context.h"
-#include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/squeeze_kernel_impl.h"
-
-PD_REGISTER_KERNEL(squeeze,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::SqueezeKernel,
-                   float,
-                   double,
-                   phi::dtype::bfloat16,
-                   bool,
-                   int,
-                   uint8_t,
-                   int8_t,
-                   int64_t,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>) {}
-
-PD_REGISTER_KERNEL(squeeze_with_xshape,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::SqueezeWithXShapeKernel,
-                   float,
-                   double,
-                   phi::dtype::bfloat16,
-                   bool,
-                   int,
-                   uint8_t,
-                   int8_t,
-                   int64_t,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>) {}

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

paddle/phi/kernels/impl/squeeze_grad_kernel_impl.h

-// 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.
-
-#pragma once
-#include "paddle/phi/common/int_array.h"
-#include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/core/tensor_utils.h"
-
-namespace phi {
-template <typename T, typename Context>
-void SqueezeGradKernel(const Context& dev_ctx,
-                       const DenseTensor& xshape,
-                       const DenseTensor& dout,
-                       const IntArray& axes,
-                       DenseTensor* dx) {
-  auto xshape_dims = xshape.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(), false, dx);
-  dx->Resize(x_dims);
-}
-}  // namespace phi

paddle/phi/kernels/gpu/squeeze_kernel.cu → paddle/phi/kernels/squeeze_grad_kernel.cc

@@ -12,20 +12,35 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/squeeze_kernel.h"
+#include "paddle/phi/kernels/squeeze_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/squeeze_kernel_impl.h"
+#include "paddle/phi/core/tensor_utils.h"
 
-PD_REGISTER_KERNEL(squeeze,
-                   GPU,
+namespace phi {
+template <typename T, typename Context>
+void SqueezeGradKernel(const Context& dev_ctx,
+                       const DenseTensor& xshape,
+                       const DenseTensor& dout,
+                       const IntArray& axes,
+                       DenseTensor* dx) {
+  auto xshape_dims = xshape.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(), false, dx);
+  dx->Resize(x_dims);
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(squeeze_grad,
+                   CPU,
                    ALL_LAYOUT,
-                   phi::SqueezeKernel,
+                   phi::SqueezeGradKernel,
                    float,
                    double,
                    phi::dtype::bfloat16,
-                   phi::dtype::float16,
                    bool,
                    int,
                    uint8_t,
@@ -34,14 +49,15 @@ PD_REGISTER_KERNEL(squeeze,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
 
-PD_REGISTER_KERNEL(squeeze_with_xshape,
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(squeeze_grad,
                    GPU,
                    ALL_LAYOUT,
-                   phi::SqueezeWithXShapeKernel,
+                   phi::SqueezeGradKernel,
                    float,
                    double,
-                   phi::dtype::bfloat16,
                    phi::dtype::float16,
+                   phi::dtype::bfloat16,
                    bool,
                    int,
                    uint8_t,
@@ -49,3 +65,21 @@ PD_REGISTER_KERNEL(squeeze_with_xshape,
                    int64_t,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
+
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL(squeeze_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeGradKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t) {}
+
+#endif

paddle/phi/kernels/impl/squeeze_kernel_impl.h → paddle/phi/kernels/squeeze_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/squeeze_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"
 
@@ -43,3 +45,94 @@ void SqueezeWithXShapeKernel(const Context& dev_ctx,
 }
 
 }  // namespace phi
+
+PD_REGISTER_KERNEL(squeeze,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeKernel,
+                   float,
+                   double,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
+PD_REGISTER_KERNEL(squeeze_with_xshape,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeWithXShapeKernel,
+                   float,
+                   double,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   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(squeeze,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
+PD_REGISTER_KERNEL(squeeze_with_xshape,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeWithXShapeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL(squeeze,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(squeeze_with_xshape,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::SqueezeWithXShapeKernel,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   bool,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int64_t) {}
+#endif