paddlepaddle--paddle/paddle/phi/kernels/gpu/adamax_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/adamax_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"
namespace phi {
template <typename T, typename MT>
__global__ void AdamaxGPUKernel(const T* param,
                                const T* grad,
                                const MT* learning_rate,
                                const MT* moment,
                                const MT* inf_norm,
                                const MT* beta1_pow,
                                const MT* master_param,
                                MT d_beta1,
                                MT d_beta2,
                                MT d_epsilon,
                                int64_t num,
                                T* param_out,
                                MT* moment_out,
                                MT* inf_norm_out,
                                MT* master_param_out) {
  int64_t idx =
      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
      static_cast<int64_t>(threadIdx.x);

  MT lr = static_cast<MT>(learning_rate[0]);
  MT d_pow = static_cast<MT>(beta1_pow[0]);
  MT one = static_cast<MT>(1.0f);
  auto l_r = lr / (one - d_pow);

  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_moment = static_cast<MT>(moment[index]);
    MT d_inf = static_cast<MT>(inf_norm[index]);
    // compute
    auto mom_out = d_beta1 * d_moment + (one - d_beta1) * d_grad;
    auto norm_out = std::max(std::abs(d_grad), d_beta2 * d_inf + d_epsilon);
    auto out_data = d_param - l_r * (mom_out / norm_out);
    // store
    param_out[index] = static_cast<T>(out_data);
    moment_out[index] = static_cast<T>(mom_out);
    inf_norm_out[index] = static_cast<T>(norm_out);

    if (master_param_out) {
      master_param_out[index] = out_data;
    }
  }
}

template <typename T, typename Context>
void AdamaxKernel(const Context& dev_ctx,
                  const DenseTensor& param,
                  const DenseTensor& grad,
                  const DenseTensor& learning_rate,
                  const DenseTensor& moment,
                  const DenseTensor& inf_norm,
                  const DenseTensor& beta1_pow,
                  const optional<DenseTensor>& master_param,
                  float beta1,
                  float beta2,
                  float epsilon,
                  bool multi_precision,
                  DenseTensor* param_out,
                  DenseTensor* moment_out,
                  DenseTensor* inf_norm_out,
                  DenseTensor* master_param_outs) {
  using MT = typename dtype::template MPTypeTrait<T>::Type;
  T* param_out_data = dev_ctx.template Alloc<T>(param_out);
  MT* moment_out_data = dev_ctx.template Alloc<MT>(moment_out);
  MT* inf_norm_out_data = dev_ctx.template Alloc<MT>(inf_norm_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_outs) : nullptr;
  PADDLE_ENFORCE_EQ(
      beta1_pow.numel(),
      1,
      errors::InvalidArgument("beta1 pow's size should be 1, but received "
                              "value is:%d.",
                              beta1_pow.numel()));

  MT beta1_ = static_cast<MT>(beta1);
  MT beta2_ = static_cast<MT>(beta2);
  MT epsilon_ = static_cast<MT>(epsilon);

  int64_t numel = param.numel();
  auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, 1);
  int grid = config.block_per_grid.x;
  int block = config.thread_per_block.x;
  auto stream = dev_ctx.stream();

  AdamaxGPUKernel<T, MT><<<block, grid, 0, stream>>>(param.data<T>(),
                                                     grad.data<T>(),
                                                     learning_rate.data<MT>(),
                                                     moment.data<MT>(),
                                                     inf_norm.data<MT>(),
                                                     beta1_pow.data<MT>(),
                                                     master_in_data,
                                                     beta1_,
                                                     beta2_,
                                                     epsilon_,
                                                     numel,
                                                     param_out_data,
                                                     moment_out_data,
                                                     inf_norm_out_data,
                                                     master_out_data);
}
}  // namespace phi
PD_REGISTER_KERNEL(
    adamax, GPU, ALL_LAYOUT, phi::AdamaxKernel, float, double, phi::float16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
  }
}