Function ComputeImpl

paddle/phi/kernels/impl/lamb_kernel_impl.h:122–312

template <typename T, typename MT, typename Context, bool IsMultiPrecision>
void ComputeImpl(const Context& dev_ctx,
                 const DenseTensor& param,
                 const DenseTensor& grad,
                 const DenseTensor& lr,
                 const DenseTensor& mom1,
                 const DenseTensor& mom2,
                 const DenseTensor& beta1_pow,
                 const DenseTensor& beta2_pow,
                 const optional<DenseTensor>& master_param_opt,
                 const optional<DenseTensor>& skip_update_opt,
                 float weight_decay_f,
                 float beta1_f,
                 float beta2_f,
                 float epsilon_f,
                 bool always_adapt,
                 bool multi_precision UNUSED,
                 DenseTensor* param_out,
                 DenseTensor* mom1_out,
                 DenseTensor* mom2_out,
                 DenseTensor* beta1_pow_out,
                 DenseTensor* beta2_pow_out,
                 DenseTensor* master_param_out) {
  if (!IsMultiPrecision) {
    constexpr auto kIsSameType = std::is_same<T, MT>::value;
    PADDLE_ENFORCE_EQ(
        kIsSameType,
        true,
        common::errors::InvalidArgument(
            "When multi_precision=False, T and MT must be the same type."));
  }

  const auto* master_param =
      IsMultiPrecision ? master_param_opt.get_ptr() : nullptr;
  const auto* skip_update = skip_update_opt.get_ptr();
  const bool* skip_update_flag = skip_update && skip_update->IsInitialized()
                                     ? skip_update->data<bool>()
                                     : nullptr;
  if (skip_update_flag &&
      skip_update->place().GetType() == AllocationType::CPU &&
      (*skip_update_flag)) {
    return;
  }

  auto weight_decay = static_cast<MT>(weight_decay_f);
  auto beta1 = static_cast<MT>(beta1_f);
  auto beta2 = static_cast<MT>(beta2_f);
  auto epsilon = static_cast<MT>(epsilon_f);
  auto numel = param.numel();
  funcs::ForRange<Context> for_range(dev_ctx, numel);
  DenseTensor trust_ratio_div;
  trust_ratio_div.Resize(param.dims());
  auto* trust_ratio_div_ptr = dev_ctx.template Alloc<MT>(&trust_ratio_div);

  const void* param_ptr = param.data();
  const void* master_param_ptr = master_param ? master_param->data() : nullptr;
  void* param_out_ptr = dev_ctx.template Alloc<T>(param_out);
  void* master_param_out_ptr =
      master_param_out ? dev_ctx.template Alloc<MT>(master_param_out) : nullptr;