chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/rprop_kernel.cu

@@ -0,0 +1,152 @@
+// 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/rprop_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_helper.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/mixed_vector.h"
+
+namespace phi {
+
+template <typename T, typename MT>
+__global__ void RpropKernelGPUImpl(const T* param,
+                                   const T* grad,
+                                   const T* prev,
+                                   const T* learning_rate,
+                                   const MT* master_param,
+                                   const T* learning_rate_range,
+                                   const T* etas,
+                                   int64_t num,
+                                   T* param_out,
+                                   T* prev_out,
+                                   T* learning_rate_out,
+                                   MT* master_param_out) {
+  MT learning_rate_min_data = static_cast<MT>(learning_rate_range[0]);
+  MT learning_rate_max_data = static_cast<MT>(learning_rate_range[1]);
+  MT eta_negative_data = static_cast<MT>(etas[0]);
+  MT eta_positive_data = static_cast<MT>(etas[1]);
+  MT zero_data = static_cast<MT>(0);
+  MT one_data = static_cast<MT>(1);
+  MT negative_one_data = static_cast<MT>(-1);
+
+  CUDA_KERNEL_LOOP_TYPE(i, num, int64_t) {
+    MT param_data = master_param ? master_param[i] : static_cast<MT>(param[i]);
+    MT grad_data = static_cast<MT>(grad[i]);
+    MT prev_data = static_cast<MT>(prev[i]);
+    MT learning_rate_data = static_cast<MT>(learning_rate[i]);
+    MT product_data = grad_data * prev_data;
+
+    MT eta_data = one_data;
+    if (product_data > zero_data) {
+      eta_data = eta_positive_data;
+    } else if (product_data < zero_data) {
+      grad_data = zero_data;
+      eta_data = eta_negative_data;
+    }
+
+    learning_rate_data = learning_rate_data * eta_data;
+    if (learning_rate_data > learning_rate_max_data) {
+      learning_rate_data = learning_rate_max_data;
+    } else if (learning_rate_data < learning_rate_min_data) {
+      learning_rate_data = learning_rate_min_data;
+    }
+
+    MT grad_sign_data = zero_data;
+    if (grad_data > zero_data) {
+      grad_sign_data = one_data;
+    } else if (grad_data < zero_data) {
+      grad_sign_data = negative_one_data;
+    }
+
+    param_data = param_data - grad_sign_data * learning_rate_data;
+    prev_data = grad_data;
+
+    param_out[i] = static_cast<T>(param_data);
+    prev_out[i] = static_cast<T>(prev_data);
+    learning_rate_out[i] = static_cast<T>(learning_rate_data);
+    if (master_param_out) {
+      master_param_out[i] = param_data;
+    }
+  }
+}
+
+template <typename T, typename Context>
+void RpropKernel(const Context& dev_ctx,
+                 const DenseTensor& param,
+                 const DenseTensor& grad,
+                 const DenseTensor& prev,
+                 const DenseTensor& learning_rate,
+                 const optional<DenseTensor>& master_param,
+                 const DenseTensor& learning_rate_range,
+                 const DenseTensor& etas,
+                 bool multi_precision,
+                 DenseTensor* param_out,
+                 DenseTensor* prev_out,
+                 DenseTensor* learning_rate_out,
+                 DenseTensor* master_param_out) {
+  using MT = typename MPTypeTrait<T>::Type;
+  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;
+
+  int block = 512;
+  int64_t grid_max = dev_ctx.GetCUDAMaxGridDimSize()[0];
+  int grid = std::min((param.numel() + block - 1) / block, grid_max);
+
+  RpropKernelGPUImpl<T, MT><<<grid, block, 0, dev_ctx.stream()>>>(
+      param.data<T>(),
+      grad.data<T>(),
+      prev.data<T>(),
+      learning_rate.data<T>(),
+      master_in_data,
+      learning_rate_range.data<T>(),
+      etas.data<T>(),
+      param.numel(),
+      dev_ctx.template Alloc<T>(param_out),
+      dev_ctx.template Alloc<T>(prev_out),
+      dev_ctx.template Alloc<T>(learning_rate_out),
+      master_out_data);
+}
+
+}  // namespace phi
+
+#ifdef PADDLE_WITH_CUDA
+PD_REGISTER_KERNEL(rprop,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::RpropKernel,
+                   phi::float16,
+                   phi::bfloat16,
+                   float,
+                   double) {
+  if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
+      kernel_key.dtype() == phi::DataType::BFLOAT16) {
+    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+  }
+}
+#endif
+
+#ifdef PADDLE_WITH_HIP
+PD_REGISTER_KERNEL(
+    rprop, GPU, ALL_LAYOUT, phi::RpropKernel, phi::float16, float, double) {
+  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
+    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
+  }
+}
+#endif