move dgc_momentum kernel to phi

paddle/fluid/operators/optimizers/dgc_momentum_op.cc

@@ -12,10 +12,10 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/optimizers/dgc_momentum_op.h"
-
 #include <string>
 
+#include "paddle/fluid/framework/op_registry.h"
+
 namespace paddle {
 namespace operators {
 
@@ -203,6 +203,3 @@ namespace ops = paddle::operators;
 REGISTER_OP_WITHOUT_GRADIENT(dgc_momentum,
                              ops::DGCMomentumOp,
                              ops::DGCMomentumOpMaker);
-
-PD_REGISTER_STRUCT_KERNEL(
-    dgc_momentum, CPU, ALL_LAYOUT, ops::DGCMomentumKernel, float) {}

paddle/fluid/operators/optimizers/dgc_momentum_op.cu → paddle/phi/kernels/cpu/dgc_momentum_kernel.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2023 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.
@@ -12,9 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/optimizers/dgc_momentum_op.h"
+#include "paddle/phi/kernels/dgc_momentum_kernel.h"
 
-namespace ops = paddle::operators;
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/dgc_momentum_kernel_impl.h"
 
-PD_REGISTER_STRUCT_KERNEL(
-    dgc_momentum, GPU, ALL_LAYOUT, ops::DGCMomentumKernel, float) {}
+PD_REGISTER_KERNEL(
+    dgc_momentum, CPU, ALL_LAYOUT, phi::DGCMomentumKernel, float) {}

paddle/phi/kernels/dgc_momentum_kernel.h

@@ -0,0 +1,42 @@
+// Copyright (c) 2023 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 DGCMomentumKernel(const Context& dev_ctx,
+                       const DenseTensor& param,
+                       const DenseTensor& grad,
+                       const DenseTensor& velocity,
+                       const DenseTensor& learning_rate,
+                       const DenseTensor& master_param,
+                       const DenseTensor& current_step_tensor,
+                       const DenseTensor& nranks_tensor,
+                       float mu,
+                       bool use_nesterov,
+                       const std::string& regularization_method,
+                       float regularization_coeff,
+                       bool multi_precision,
+                       float rescale_grad,
+                       float rampup_begin_step,
+                       DenseTensor* param_out,
+                       DenseTensor* velocity_out,
+                       DenseTensor* master_param_out,
+                       DenseTensor* grad_out);
+
+}  // namespace phi

paddle/phi/kernels/gpu/dgc_momentum_kernel.cu

@@ -0,0 +1,22 @@
+// Copyright (c) 2023 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/dgc_momentum_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/dgc_momentum_kernel_impl.h"
+
+PD_REGISTER_KERNEL(
+    dgc_momentum, GPU, ALL_LAYOUT, phi::DGCMomentumKernel, float) {}