chore: import upstream snapshot with attribution

paddle/phi/kernels/selected_rows/impl/ftrl_kernel_impl.h

@@ -0,0 +1,176 @@
+// 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.
+
+#pragma once
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/for_range.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, int MajorType = Eigen::RowMajor>
+using EigenVector = EigenVector<T, MajorType>;
+
+template <typename T>
+class SparseFTRLFunctor {
+ private:
+  const T* g_;
+  const T* p_;
+  const T* s_acc_;
+  const T* l_acc_;
+  const T* lr_;
+  const T l1_;
+  const T l2_;
+  const T lr_power_;
+  const int64_t* rows_;
+  const int64_t row_numel_;
+  T* p_out_;
+  T* s_acc_out_;
+  T* l_acc_out_;
+
+ public:
+  SparseFTRLFunctor(const T* g,
+                    const T* p,
+                    const T* s_acc,
+                    const T* lr,
+                    const T l1,
+                    const T l2,
+                    const T lr_power,
+                    const int64_t* rows,
+                    int64_t row_numel,
+                    T* p_out,
+                    T* s_acc_out,
+                    T* l_acc_out)
+      : g_(g),
+        p_(p),
+        s_acc_(s_acc),
+        lr_(lr),
+        l1_(l1),
+        l2_(l2),
+        lr_power_(lr_power),
+        rows_(rows),
+        row_numel_(row_numel),
+        p_out_(p_out),
+        s_acc_out_(s_acc_out),
+        l_acc_out_(l_acc_out) {}
+
+  inline HOSTDEVICE void operator()(size_t i) {
+    auto j = rows_[i / row_numel_] * row_numel_ + i % row_numel_;
+    const T g = g_[i];
+    const T p = p_[j];
+    const T s_acc = s_acc_[j];
+    const T lr = lr_[0];
+
+    auto new_acc = s_acc + g * g;
+
+    if (lr_power_ == static_cast<T>(-0.5)) {
+      l_acc_out_[j] += g - (std::sqrt(new_acc) - std::sqrt(s_acc)) / lr * p;
+    } else {
+      l_acc_out_[j] +=
+          g - (std::pow(new_acc, -lr_power_) - std::pow(s_acc, -lr_power_)) /
+                  lr * p;
+    }
+
+    auto l_acc = l_acc_out_[j];
+
+    if (std::fabs(l_acc) > l1_) {
+      auto x = -l_acc;
+      if (l_acc >= static_cast<T>(0)) {
+        x += l1_;
+      } else {
+        x -= l1_;
+      }
+
+      auto y = static_cast<T>(2) * l2_;
+      if (lr_power_ == static_cast<T>(-0.5)) {
+        y += std::sqrt(new_acc) / lr;
+      } else {
+        y += std::pow(new_acc, -lr_power_) / lr;
+      }
+
+      auto pre_shrink = x / y;
+      p_out_[j] = pre_shrink;
+    } else {
+      p_out_[j] = static_cast<T>(0);
+    }
+
+    s_acc_out_[j] += g * g;
+  }
+};
+
+template <typename T, typename Context>
+void FTRLOpKernel(const Context& dev_ctx,
+                  const DenseTensor& param,
+                  const DenseTensor& squared_accumulator,
+                  const DenseTensor& linear_accumulator,
+                  const SelectedRows& grad_in,
+                  const DenseTensor& learning_rate,
+                  float l1_in,
+                  float l2_in,
+                  float lr_power_in,
+                  DenseTensor* param_out,
+                  DenseTensor* squared_accum_out,
+                  DenseTensor* linear_accum_out) {
+  auto* lr_in = &learning_rate;
+
+  auto* param_in = &param;
+  auto* sq_accum_in = &squared_accumulator;
+
+  auto* sq_accum_out = squared_accum_out;
+  auto* lin_accum_out = linear_accum_out;
+
+  dev_ctx.template Alloc<T>(param_out);
+  dev_ctx.template Alloc<T>(sq_accum_out);
+  dev_ctx.template Alloc<T>(lin_accum_out);
+
+  auto l1 = static_cast<T>(l1_in) + static_cast<T>(1e-10);
+  auto l2 = static_cast<T>(l2_in) + static_cast<T>(1e-10);
+  auto lr_power = static_cast<T>(lr_power_in);
+
+  auto grad = &grad_in;
+
+  SelectedRows tmp_merged_grad;
+  SelectedRows* merged_grad = &tmp_merged_grad;
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, *grad, merged_grad);
+
+  auto* merged_rows = merged_grad->mutable_rows();
+  phi::MixVector<int64_t> mixv_merged_rows(merged_rows);
+  const int64_t* rows = mixv_merged_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = static_cast<int64_t>(merged_grad->value().dims()[1]);
+  auto row_height = static_cast<int64_t>(merged_grad->rows().size());
+
+  funcs::ForRange<Context> for_range(static_cast<const Context&>(dev_ctx),
+                                     row_numel * row_height);
+
+  SparseFTRLFunctor<T> functor(merged_grad->value().data<T>(),
+                               param_in->data<T>(),
+                               sq_accum_in->data<T>(),
+                               lr_in->data<T>(),
+                               l1,
+                               l2,
+                               lr_power,
+                               rows,
+                               row_numel,
+                               dev_ctx.template Alloc<T>(param_out),
+                               dev_ctx.template Alloc<T>(sq_accum_out),
+                               dev_ctx.template Alloc<T>(lin_accum_out));
+  for_range(functor);
+}
+
+}  // namespace sr
+}  // namespace phi