chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/radam_kernel.cu

@@ -0,0 +1,181 @@
+// Copyright (c) 2024 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/radam_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+template <typename T, typename MT>
+__global__ void RAdamGPUKernel(const T* param,
+                               const T* grad,
+                               const MT* learning_rate,
+                               const MT* beta1_pow,
+                               const MT* beta2_pow,
+                               const MT* rho,
+                               const MT* moment1,
+                               const MT* moment2,
+                               const MT* master_param,
+                               MT beta1,
+                               MT beta2,
+                               MT epsilon,
+                               MT rho_inf,
+                               int64_t num,
+                               T* param_out,
+                               MT* beta1_pow_out,
+                               MT* beta2_pow_out,
+                               MT* rho_out,
+                               MT* moment1_out,
+                               MT* moment2_out,
+                               MT* master_param_out) {
+  MT lr_scalar = static_cast<MT>(learning_rate[0]);
+
+  int64_t idx =
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
+      static_cast<int64_t>(threadIdx.x);
+
+  for (int64_t index = idx; index < num; index += gridDim.x * blockDim.x) {
+    // load and cast input to MT
+    MT d_param =
+        master_param ? master_param[index] : static_cast<MT>(param[index]);
+    MT d_grad = static_cast<MT>(grad[index]);
+    MT d_beta1_pow = static_cast<MT>(beta1_pow[index]);
+    MT d_beta2_pow = static_cast<MT>(beta2_pow[index]);
+    MT d_rho = static_cast<MT>(rho[index]);
+    MT d_moment1 = static_cast<MT>(moment1[index]);
+    MT d_moment2 = static_cast<MT>(moment2[index]);
+
+    // compute
+    MT beta1_pow_scalar = d_beta1_pow * beta1;
+    MT beta2_pow_scalar = d_beta2_pow * beta2;
+    MT rho_scalar = (d_rho * (beta2 - beta2_pow_scalar) + beta2_pow_scalar) /
+                    (static_cast<MT>(1) - beta2_pow_scalar);
+
+    MT m1_out = beta1 * d_moment1 + (static_cast<MT>(1) - beta1) * d_grad;
+    MT m2_out =
+        beta2 * d_moment2 + (static_cast<MT>(1) - beta2) * d_grad * d_grad;
+
+    MT m1_hat = m1_out / (static_cast<MT>(1) - beta1_pow_scalar);
+
+    MT rho_t = rho_inf - static_cast<MT>(2) * rho_scalar;
+    MT p_out = static_cast<MT>(0);
+
+    if (rho_t > static_cast<MT>(5)) {
+      MT l_t = std::sqrt((static_cast<MT>(1) - beta2_pow_scalar)) /
+               (std::sqrt(m2_out) + epsilon);
+      MT r_t = std::sqrt(((rho_t - static_cast<MT>(4)) *
+                          (rho_t - static_cast<MT>(2)) * rho_inf) /
+                         ((rho_inf - static_cast<MT>(4)) *
+                          (rho_inf - static_cast<MT>(2)) * rho_t));
+
+      p_out = d_param - lr_scalar * m1_hat * r_t * l_t;
+    } else {
+      p_out = d_param - lr_scalar * m1_hat;
+    }
+
+    // store
+    param_out[index] = static_cast<T>(p_out);
+    beta1_pow_out[index] = static_cast<T>(beta1_pow_scalar);
+    beta2_pow_out[index] = static_cast<T>(beta2_pow_scalar);
+    rho_out[index] = static_cast<T>(rho_scalar);
+    moment1_out[index] = static_cast<MT>(m1_out);
+    moment2_out[index] = static_cast<MT>(m2_out);
+
+    if (master_param_out) {
+      master_param_out[index] = p_out;
+    }
+  }
+}
+
+template <typename T, typename Context>
+void RAdamKernel(const Context& dev_ctx,
+                 const DenseTensor& param,
+                 const DenseTensor& grad,
+                 const DenseTensor& learning_rate,
+                 const DenseTensor& beta1_pow,
+                 const DenseTensor& beta2_pow,
+                 const DenseTensor& rho,
+                 const DenseTensor& moment1,
+                 const DenseTensor& moment2,
+                 const optional<DenseTensor>& master_param,
+                 float beta1,
+                 float beta2,
+                 float epsilon,
+                 bool multi_precision,
+                 DenseTensor* param_out,
+                 DenseTensor* beta1_pow_out,
+                 DenseTensor* beta2_pow_out,
+                 DenseTensor* rho_out,
+                 DenseTensor* moment1_out,
+                 DenseTensor* moment2_out,
+                 DenseTensor* master_param_out) {
+  using MT = typename dtype::template MPTypeTrait<T>::Type;
+  T* param_out_data = dev_ctx.template Alloc<T>(param_out);
+
+  MT* beta1_pow_out_data = dev_ctx.template Alloc<MT>(beta1_pow_out);
+  MT* beta2_pow_out_data = dev_ctx.template Alloc<MT>(beta2_pow_out);
+  MT* rho_out_data = dev_ctx.template Alloc<MT>(rho_out);
+
+  MT* moment1_out_data = dev_ctx.template Alloc<MT>(moment1_out);
+  MT* moment2_out_data = dev_ctx.template Alloc<MT>(moment2_out);
+
+  const MT* master_in_data =
+      multi_precision ? master_param->data<MT>() : nullptr;
+  MT* master_out_data =
+      multi_precision ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;
+
+  MT beta1_ = static_cast<MT>(beta1);
+  MT beta2_ = static_cast<MT>(beta2);
+  MT epsilon_ = static_cast<MT>(epsilon);
+
+  MT rho_inf =
+      static_cast<MT>(2) / (static_cast<MT>(1) - beta2_) - static_cast<MT>(1);
+
+  int64_t numel = param.numel();
+  int block = 512;
+  int64_t grid_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
+  int grid = std::min((param.numel() + block - 1) / block, grid_max);
+  auto stream = dev_ctx.stream();
+
+  RAdamGPUKernel<T, MT><<<block, grid, 0, stream>>>(param.data<T>(),
+                                                    grad.data<T>(),
+                                                    learning_rate.data<MT>(),
+                                                    beta1_pow.data<MT>(),
+                                                    beta2_pow.data<MT>(),
+                                                    rho.data<MT>(),
+                                                    moment1.data<MT>(),
+                                                    moment2.data<MT>(),
+                                                    master_in_data,
+                                                    beta1_,
+                                                    beta2_,
+                                                    epsilon_,
+                                                    rho_inf,
+                                                    numel,
+                                                    param_out_data,
+                                                    beta1_pow_out_data,
+                                                    beta2_pow_out_data,
+                                                    rho_out_data,
+                                                    moment1_out_data,
+                                                    moment2_out_data,
+                                                    master_out_data);
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    radam, GPU, ALL_LAYOUT, phi::RAdamKernel, float, double, phi::float16) {}