chore: import upstream snapshot with attribution

paddle/phi/kernels/cpu/rprop_kernel.cc

@@ -0,0 +1,138 @@
+// 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/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/jit/kernels.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void RpropKernelCPUImpl(const Context& dev_ctx,
+                        const DenseTensor& param,
+                        const DenseTensor& grad,
+                        const DenseTensor& prev,
+                        const DenseTensor& learning_rate,
+                        const DenseTensor& learning_rate_range,
+                        const DenseTensor& etas,
+                        DenseTensor* param_out,
+                        DenseTensor* prev_out,
+                        DenseTensor* learning_rate_out) {
+  auto param_eigen = EigenVector<T>::Flatten(param);
+  auto prev_eigen = EigenVector<T>::Flatten(prev);
+  auto param_out_eigen = EigenVector<T>::Flatten(*param_out);
+  auto prev_out_eigen = EigenVector<T>::Flatten(*prev_out);
+  auto learning_rate_out_eigen = EigenVector<T>::Flatten(*learning_rate_out);
+  auto learning_rate_min = learning_rate_range.data<T>()[0];
+  auto learning_rate_max = learning_rate_range.data<T>()[1];
+  auto eta_negative = etas.data<T>()[0];
+  auto eta_positive = etas.data<T>()[1];
+
+  DenseTensor grad_tensor;
+  grad_tensor.Resize(grad.dims());
+  dev_ctx.template Alloc<T>(&grad_tensor);
+  Copy<Context>(dev_ctx, grad, dev_ctx.GetPlace(), true, &grad_tensor);
+  auto grad_eigen = EigenVector<T>::Flatten(grad_tensor);
+
+  DenseTensor product_tensor;
+  product_tensor.Resize(grad.dims());
+  dev_ctx.template Alloc<T>(&product_tensor);
+  auto product_eigen = EigenVector<T>::Flatten(product_tensor);
+
+  DenseTensor learning_rate_tensor;
+  learning_rate_tensor.Resize(learning_rate.dims());
+  dev_ctx.template Alloc<T>(&learning_rate_tensor);
+  Copy<Context>(
+      dev_ctx, learning_rate, dev_ctx.GetPlace(), true, &learning_rate_tensor);
+  auto learning_rate_eigen = EigenVector<T>::Flatten(learning_rate_tensor);
+
+  DenseTensor eta_tensor;
+  eta_tensor.Resize(learning_rate.dims());
+  dev_ctx.template Alloc<T>(&eta_tensor);
+  auto eta_eigen = EigenVector<T>::Flatten(eta_tensor);
+
+  product_eigen = grad_eigen * prev_eigen;
+  T* product_data = product_tensor.data<T>();
+  T* grad_data = grad_tensor.data<T>();
+  T* eta_data = eta_tensor.data<T>();
+  T zero = static_cast<T>(0);
+  T one = static_cast<T>(1);
+  for (int64_t i = 0, n = product_tensor.numel(); i < n; i++) {
+    if (product_data[i] > zero) {
+      eta_data[i] = eta_positive;
+    } else if (product_data[i] == zero) {
+      eta_data[i] = one;
+    } else if (product_data[i] < zero) {
+      grad_data[i] = zero;
+      eta_data[i] = eta_negative;
+    }
+  }
+
+  learning_rate_eigen = learning_rate_eigen * eta_eigen;
+  T* learning_rate_data = learning_rate_tensor.data<T>();
+  for (int64_t i = 0, n = learning_rate_tensor.numel(); i < n; i++) {
+    if (learning_rate_data[i] > learning_rate_max) {
+      learning_rate_data[i] = learning_rate_max;
+    } else if (learning_rate_data[i] < learning_rate_min) {
+      learning_rate_data[i] = learning_rate_min;
+    }
+  }
+
+  param_out_eigen = param_eigen - grad_eigen.sign() * learning_rate_eigen;
+  prev_out_eigen = grad_eigen;
+  learning_rate_out_eigen = learning_rate_eigen;
+  Copy<Context>(dev_ctx, grad_tensor, dev_ctx.GetPlace(), true, prev_out);
+  Copy<Context>(dev_ctx,
+                learning_rate_tensor,
+                dev_ctx.GetPlace(),
+                true,
+                learning_rate_out);
+}
+
+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 UNUSED,
+                 const DenseTensor& learning_rate_range,
+                 const DenseTensor& etas,
+                 bool multi_precision UNUSED,
+                 DenseTensor* param_out,
+                 DenseTensor* prev_out,
+                 DenseTensor* learning_rate_out,
+                 DenseTensor* master_param_out UNUSED) {
+  dev_ctx.template Alloc<T>(param_out);
+  dev_ctx.template Alloc<T>(prev_out);
+  dev_ctx.template Alloc<T>(learning_rate_out);
+  RpropKernelCPUImpl<T, Context>(dev_ctx,
+                                 param,
+                                 grad,
+                                 prev,
+                                 learning_rate,
+                                 learning_rate_range,
+                                 etas,
+                                 param_out,
+                                 prev_out,
+                                 learning_rate_out);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    rprop, CPU, ALL_LAYOUT, phi::RpropKernel, phi::bfloat16, float, double) {}