调整channel_shuffle及其梯度的核函数的位置

paddle/phi/kernels/channel_shuffle_grad_kernel.h

@@ -20,7 +20,7 @@
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleGradKernel(const Context& ctx,
+void ChannelShuffleGradKernel(const Context& dev_ctx,
                               const DenseTensor& out_grad,
                               int groups,
                               const std::string& data_format,

paddle/phi/kernels/channel_shuffle_kernel.h

@@ -20,7 +20,7 @@
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleKernel(const Context& ctx,
+void ChannelShuffleKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           int groups,
                           const std::string& data_format,

paddle/phi/kernels/cpu/channel_shuffle_grad_kernel.cc

@@ -0,0 +1,26 @@
+// 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/channel_shuffle_grad_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/cpu/channel_shuffle_kernel.cc

@@ -0,0 +1,26 @@
+// 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/channel_shuffle_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/channel_shuffle_grad_kernel.cu

@@ -0,0 +1,26 @@
+// 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/channel_shuffle_grad_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/gpu/channel_shuffle_kernel.cu

@@ -0,0 +1,26 @@
+// 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/channel_shuffle_kernel.h"
+#include "paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+PD_REGISTER_KERNEL(channel_shuffle,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::ChannelShuffleKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/channel_shuffle_grad_kernel.cc → paddle/phi/kernels/impl/channel_shuffle_grad_kernel_impl.h

@@ -12,25 +12,25 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/channel_shuffle_grad_kernel.h"
+#pragma once
+
 #include <string>
 #include <vector>
-#include "paddle/phi/backends/all_context.h"
+
 #include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleGradKernel(const Context& ctx,
+void ChannelShuffleGradKernel(const Context& dev_ctx,
                               const DenseTensor& out_grad,
                               int groups,
                               const std::string& data_format,
                               DenseTensor* x_grad) {
   auto* dout = &out_grad;
   auto* dx = x_grad;
-  ctx.template Alloc<T>(dx);
+  dev_ctx.template Alloc<T>(dx);
   bool channel_last = (data_format == "NHWC");
   auto do_dims = dout->dims();
   auto dx_dims = dx->dims();
@@ -51,24 +51,8 @@ void ChannelShuffleGradKernel(const Context& ctx,
     o.Resize({dx_dims[0], dx_dims[1], dx_dims[2], groups, dx_dims[3] / groups});
   }
   phi::funcs::Transpose<Context, T, 5> trans;
-  trans(ctx, t, &o, axis);
+  trans(dev_ctx, t, &o, axis);
   dx->Resize(dx_dims);
 }
 
 }  // namespace phi
-
-PD_REGISTER_KERNEL(channel_shuffle_grad,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleGradKernel,
-                   float,
-                   double) {}
-
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(channel_shuffle_grad,
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleGradKernel,
-                   float,
-                   double) {}
-#endif

paddle/phi/kernels/channel_shuffle_kernel.cc → paddle/phi/kernels/impl/channel_shuffle_kernel_impl.h

@@ -12,24 +12,24 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/phi/kernels/channel_shuffle_kernel.h"
+#pragma once
+
 #include <string>
 #include <vector>
-#include "paddle/phi/backends/all_context.h"
+
 #include "paddle/phi/core/dense_tensor.h"
-#include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
 
 namespace phi {
 
 template <typename T, typename Context>
-void ChannelShuffleKernel(const Context& ctx,
+void ChannelShuffleKernel(const Context& dev_ctx,
                           const DenseTensor& x,
                           int groups,
                           const std::string& data_format,
                           DenseTensor* out) {
   auto* in = &x;
-  ctx.template Alloc<T>(out);
+  dev_ctx.template Alloc<T>(out);
   bool channel_last = (data_format == "NHWC");
   auto in_dims = in->dims();
   auto o_dims = out->dims();
@@ -50,24 +50,8 @@ void ChannelShuffleKernel(const Context& ctx,
     o.Resize({in_dims[0], in_dims[1], in_dims[2], in_dims[3] / groups, groups});
   }
   phi::funcs::Transpose<Context, T, 5> trans;
-  trans(ctx, t, &o, axis);
+  trans(dev_ctx, t, &o, axis);
   out->Resize(o_dims);
 }
 
 }  // namespace phi
-
-PD_REGISTER_KERNEL(channel_shuffle,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleKernel,
-                   float,
-                   double) {}
-
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(channel_shuffle,
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::ChannelShuffleKernel,
-                   float,
-                   double) {}
-#endif