chore: import upstream snapshot with attribution

paddle/phi/kernels/selected_rows/activation_kernel.cc

@@ -0,0 +1,66 @@
+// 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/selected_rows/activation_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/activation_kernel.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void SquareKernel(const Context& dev_ctx,
+                  const SelectedRows& x,
+                  SelectedRows* out) {
+  out->set_rows(x.rows());
+  out->set_height(x.height());
+  phi::SquareKernel<T, Context>(dev_ctx, x.value(), out->mutable_value());
+}
+
+template <typename T, typename Context>
+void SqrtKernel(const Context& dev_ctx,
+                const SelectedRows& x,
+                SelectedRows* out) {
+  out->set_rows(x.rows());
+  out->set_height(x.height());
+  phi::SqrtKernel<T, Context>(dev_ctx, x.value(), out->mutable_value());
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    square_sr, CPU, ALL_LAYOUT, phi::sr::SquareKernel, float, double) {}
+
+PD_REGISTER_KERNEL(
+    sqrt_sr, CPU, ALL_LAYOUT, phi::sr::SqrtKernel, float, double) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+PD_REGISTER_KERNEL(square_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::SquareKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(
+    sqrt_sr, GPU, ALL_LAYOUT, phi::sr::SqrtKernel, float, double) {}
+
+#endif

paddle/phi/kernels/selected_rows/activation_kernel.h

@@ -0,0 +1,34 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void SquareKernel(const Context& dev_ctx,
+                  const SelectedRows& x,
+                  SelectedRows* out);
+
+template <typename T, typename Context>
+void SqrtKernel(const Context& dev_ctx,
+                const SelectedRows& x,
+                SelectedRows* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/adam_kernel.h

@@ -0,0 +1,53 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void AdamDenseParamSparseGradKernel(const Context& dev_ctx,
+                                    const DenseTensor& param,
+                                    const SelectedRows& grad,
+                                    const DenseTensor& learning_rate,
+                                    const DenseTensor& moment1,
+                                    const DenseTensor& moment2,
+                                    const optional<DenseTensor>& moment2_max,
+                                    const DenseTensor& beta1_pow,
+                                    const DenseTensor& beta2_pow,
+                                    const optional<DenseTensor>& master_param,
+                                    const optional<DenseTensor>& skip_update,
+                                    const Scalar& beta1,
+                                    const Scalar& beta2,
+                                    const Scalar& epsilon,
+                                    bool lazy_mode,
+                                    int64_t min_row_size_to_use_multithread,
+                                    bool multi_precision,
+                                    bool use_global_beta_pow,
+                                    bool amsgrad,
+                                    DenseTensor* param_out,
+                                    DenseTensor* moment1_out,
+                                    DenseTensor* moment2_out,
+                                    DenseTensor* moment2_max_out,
+                                    DenseTensor* beta1_pow_out,
+                                    DenseTensor* beta2_pow_out,
+                                    DenseTensor* master_param_outs);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/adamw_kernel.h

@@ -0,0 +1,56 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void AdamwDenseParamSparseGradKernel(const Context& dev_ctx,
+                                     const DenseTensor& param,
+                                     const SelectedRows& grad,
+                                     const DenseTensor& learning_rate,
+                                     const DenseTensor& moment1,
+                                     const DenseTensor& moment2,
+                                     const optional<DenseTensor>& moment2_max,
+                                     const DenseTensor& beta1_pow,
+                                     const DenseTensor& beta2_pow,
+                                     const optional<DenseTensor>& master_param,
+                                     const optional<DenseTensor>& skip_update,
+                                     const Scalar& beta1,
+                                     const Scalar& beta2,
+                                     const Scalar& epsilon,
+                                     float lr_ratio,
+                                     float coeff,
+                                     bool with_decay,
+                                     bool lazy_mode,
+                                     int64_t min_row_size_to_use_multithread,
+                                     bool multi_precision,
+                                     bool use_global_beta_pow,
+                                     bool amsgrad,
+                                     DenseTensor* param_out,
+                                     DenseTensor* moment1_out,
+                                     DenseTensor* moment2_out,
+                                     DenseTensor* moment2_max_out,
+                                     DenseTensor* beta1_pow_out,
+                                     DenseTensor* beta2_pow_out,
+                                     DenseTensor* master_param_outs);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/add_n_kernel.h

@@ -0,0 +1,27 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void AddNKernel(const Context& dev_ctx,
+                const std::vector<const SelectedRows*>& x,
+                SelectedRows* out);
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/assign_kernel.cc

@@ -0,0 +1,58 @@
+// 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/selected_rows/assign_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/assign_kernel.h"
+
+namespace phi::sr {
+
+// Note: use `const optional<SelectedRows>& x`
+// as input if needed
+template <typename Context>
+void AssignKernel(const Context& dev_ctx,
+                  const SelectedRows& x,
+                  SelectedRows* out) {
+  out->set_rows(x.rows());
+  out->set_height(x.height());
+  phi::AssignKernel<Context>(dev_ctx, x.value(), out->mutable_value());
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL_FOR_ALL_DTYPE(assign_sr,
+                                 CPU,
+                                 ALL_LAYOUT,
+                                 phi::sr::AssignKernel<phi::CPUContext>) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL_FOR_ALL_DTYPE(assign_sr,
+                                 GPU,
+                                 ALL_LAYOUT,
+                                 phi::sr::AssignKernel<phi::GPUContext>) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
+#endif
+
+#ifdef PADDLE_WITH_XPU
+PD_REGISTER_KERNEL_FOR_ALL_DTYPE(assign_sr,
+                                 XPU,
+                                 ALL_LAYOUT,
+                                 phi::sr::AssignKernel<phi::XPUContext>) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}
+#endif

paddle/phi/kernels/selected_rows/assign_kernel.h

@@ -0,0 +1,28 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename Context>
+void AssignKernel(const Context& dev_ctx,
+                  const SelectedRows& x,
+                  SelectedRows* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/clip_by_norm_kernel.h

@@ -0,0 +1,29 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ClipByNormKernel(const Context& dev_ctx,
+                      const SelectedRows& x,
+                      float max_norm,
+                      SelectedRows* out);
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/clip_kernel.h

@@ -0,0 +1,33 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/impl/clip_kernel_impl.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ClipSparseKernel(const Context& dev_ctx,
+                      const SelectedRows& x,
+                      const Scalar& min,
+                      const Scalar& max,
+                      SelectedRows* out);
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/cpu/adam_kernel.cc

@@ -0,0 +1,260 @@
+// 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/selected_rows/adam_kernel.h"
+
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/core/threadpool.h"
+#include "paddle/phi/kernels/funcs/adam_functors.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+PD_DECLARE_int32(inner_op_parallelism);
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void AdamDenseParamSparseGradKernel(const Context& dev_ctx,
+                                    const DenseTensor& param,
+                                    const SelectedRows& grad,
+                                    const DenseTensor& learning_rate,
+                                    const DenseTensor& moment1,
+                                    const DenseTensor& moment2,
+                                    const optional<DenseTensor>& moment2_max,
+                                    const DenseTensor& beta1_pow,
+                                    const DenseTensor& beta2_pow,
+                                    const optional<DenseTensor>& master_param
+                                        UNUSED,
+                                    const optional<DenseTensor>& skip_update,
+                                    const Scalar& beta1,
+                                    const Scalar& beta2,
+                                    const Scalar& epsilon,
+                                    bool lazy_mode,
+                                    int64_t min_row_size_to_use_multithread,
+                                    bool multi_precision UNUSED,
+                                    bool use_global_beta_pow,
+                                    bool amsgrad,
+                                    DenseTensor* param_out,
+                                    DenseTensor* moment1_out,
+                                    DenseTensor* moment2_out,
+                                    DenseTensor* moment2_max_out,
+                                    DenseTensor* beta1_pow_out,
+                                    DenseTensor* beta2_pow_out,
+                                    DenseTensor* master_param_outs UNUSED) {
+  VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
+
+  bool skip_update_ = false;
+  if (skip_update.is_initialized()) {
+    PADDLE_ENFORCE_EQ(
+        skip_update->numel(),
+        1,
+        errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
+                                skip_update->numel()));
+    std::vector<bool> skip_update_vec;
+    TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
+    skip_update_ = skip_update_vec[0];
+  }
+  // skip_update=true, just copy input to output, and TensorCopy will call
+  // mutable_data
+  if (skip_update_) {
+    VLOG(4) << "Adam skip update";
+    phi::Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
+    phi::Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
+    phi::Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
+    if (amsgrad) {
+      phi::Copy(dev_ctx,
+                moment2_max.get(),
+                dev_ctx.GetPlace(),
+                false,
+                moment2_max_out);
+    }
+    if (!use_global_beta_pow) {
+      phi::Copy(dev_ctx, beta1_pow, dev_ctx.GetPlace(), false, beta1_pow_out);
+      phi::Copy(dev_ctx, beta2_pow, dev_ctx.GetPlace(), false, beta2_pow_out);
+    }
+    return;
+  }
+
+  T beta1_ = beta1.to<T>();
+  T beta2_ = beta2.to<T>();
+  T epsilon_ = epsilon.to<T>();
+
+  VLOG(3) << "beta1_pow.numel() : " << beta1_pow.numel();
+  VLOG(3) << "beta2_pow.numel() : " << beta2_pow.numel();
+  VLOG(3) << "param.numel(): " << param.numel();
+
+  PADDLE_ENFORCE_EQ(
+      beta1_pow_out->numel(),
+      1,
+      errors::InvalidArgument("beta1 pow output size should be 1, but received "
+                              "value is:%d.",
+                              beta1_pow_out->numel()));
+
+  PADDLE_ENFORCE_EQ(
+      beta2_pow_out->numel(),
+      1,
+      errors::InvalidArgument("beta2 pow output size should be 1, but received "
+                              "value is:%d.",
+                              beta2_pow_out->numel()));
+
+  if (grad.rows().empty()) {
+    VLOG(3) << "grad row size is 0!!";
+    return;
+  }
+
+  std::vector<int64_t> cpu_rows(grad.rows().begin(), grad.rows().end());
+  bool is_strict_sorted = true;
+  for (size_t i = 1; i < cpu_rows.size(); ++i) {
+    if (cpu_rows[i - 1] >= cpu_rows[i]) {
+      is_strict_sorted = false;
+      break;
+    }
+  }
+
+  SelectedRows tmp_grad_merge;
+  const SelectedRows* grad_merge_ptr = nullptr;
+  if (is_strict_sorted) {
+    grad_merge_ptr = &grad;
+  } else {
+    // merge duplicated rows if any.
+    // The rows of grad_merge have been sorted inside MergeAdd functor
+    funcs::scatter::MergeAdd<Context, T> merge_func;
+    merge_func(dev_ctx, grad, &tmp_grad_merge, true);
+    grad_merge_ptr = &tmp_grad_merge;
+  }
+
+  auto& grad_merge = *grad_merge_ptr;
+  auto& grad_tensor = grad_merge.value();
+  const T* grad_data = grad_tensor.template data<T>();
+  auto* grad_merge_rows = &grad_merge.rows();
+  phi::MixVector<int64_t> mixv_grad_merge_rows(grad_merge_rows);
+  const int64_t* rows = mixv_grad_merge_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = grad_tensor.numel() / grad_merge.rows().size();
+
+  T lr_value = static_cast<T>(learning_rate.data<double>()[0]);
+  funcs::SparseAdamFunctor<T, funcs::CPUAdam> functor(
+      beta1_,
+      beta2_,
+      epsilon_,
+      beta1_pow.data<T>(),
+      beta2_pow.data<T>(),
+      moment1.data<T>(),
+      dev_ctx.template Alloc<T>(moment1_out),
+      moment2.data<T>(),
+      dev_ctx.template Alloc<T>(moment2_out),
+      amsgrad ? moment2_max.get().data<T>() : nullptr,
+      amsgrad ? dev_ctx.template Alloc<T>(moment2_max_out) : nullptr,
+      &lr_value,
+      grad_data,
+      param.data<T>(),
+      dev_ctx.template Alloc<T>(param_out),
+      rows,
+      row_numel,
+      grad_merge.rows().size(),
+      lazy_mode,
+      amsgrad);
+  // update beta1 and beta2
+  if (!use_global_beta_pow) {
+    dev_ctx.template Alloc<T>(beta1_pow_out)[0] =
+        beta1_ * beta1_pow.data<T>()[0];
+    dev_ctx.template Alloc<T>(beta2_pow_out)[0] =
+        beta2_ * beta2_pow.data<T>()[0];
+  }
+  if (lazy_mode) {
+    VLOG(3) << "run cpu lazy mode";
+    size_t row_count = grad_merge.rows().size();
+    std::vector<int64_t> cpu_rows(grad_merge.rows());
+    for (size_t row_index = 0; row_index < row_count; ++row_index) {
+      for (size_t offset = 0; offset < row_numel; ++offset) {
+        size_t i = cpu_rows[row_index] * row_numel + offset;
+        functor.adam_update(i, grad_data[row_index * row_numel + offset]);
+      }
+    }
+  }
+#ifndef _WIN32
+  else if (FLAGS_inner_op_parallelism > 1 &&  // NOLINT
+           min_row_size_to_use_multithread > 0 &&
+           param.dims()[0] > min_row_size_to_use_multithread) {
+    VLOG(3) << "use multi thread, inner_op_parallelism="
+            << FLAGS_inner_op_parallelism << " min_row_size_to_use_multithread="
+            << min_row_size_to_use_multithread;
+    if (FLAGS_inner_op_parallelism > 10) {
+      VLOG(1) << "FLAGS_inner_op_parallelism " << FLAGS_inner_op_parallelism
+              << " is two large!";
+    }
+    auto& grad_rows = grad_merge.rows();
+    std::unordered_map<size_t, int> row_id_to_grad_row_offset;
+    size_t param_row_count = param.numel() / row_numel;
+    if (param_row_count < 1000) {
+      VLOG(1) << "param_row_count should be larger then 1000 to use "
+                 "multi thread, currently "
+              << param_row_count;
+    }
+    for (int i = 0; i < static_cast<int>(grad_rows.size()); ++i) {
+      row_id_to_grad_row_offset[grad_rows[i]] = i;
+    }
+    std::vector<std::future<void>> fs;
+    int64_t line_in_each_thread = static_cast<int64_t>(
+        param_row_count / FLAGS_inner_op_parallelism + static_cast<int64_t>(1));
+    for (int i = 0; i < FLAGS_inner_op_parallelism; ++i) {
+      int64_t start = i * line_in_each_thread;
+      int64_t end = (i + 1) * line_in_each_thread;
+      if (start >= static_cast<int64_t>(param_row_count)) {
+        break;
+      }
+      if (end > static_cast<int64_t>(param_row_count)) {
+        end = static_cast<int64_t>(param_row_count);
+      }
+      fs.push_back(phi::Async([&functor,
+                               &row_id_to_grad_row_offset,
+                               &grad_data,
+                               row_numel,
+                               start,
+                               end]() {
+        for (int64_t row_id = start; row_id < end; ++row_id) {
+          auto iter = row_id_to_grad_row_offset.find(row_id);
+          if (iter != row_id_to_grad_row_offset.end()) {
+            for (size_t row_offset = 0U; row_offset < row_numel; ++row_offset) {
+              functor.adam_update(
+                  row_id * row_numel + row_offset,
+                  grad_data[iter->second * row_numel + row_offset]);
+            }
+          } else {
+            for (size_t row_offset = 0U; row_offset < row_numel; ++row_offset) {
+              functor.adam_update(row_id * row_numel + row_offset, 0);
+            }
+          }
+        }
+      }));
+    }
+    for (auto& item : fs) item.wait();
+  }
+#endif    // !_WIN32
+  else {  // NOLINT
+    functor(param.numel());
+  }
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(adam_dense_param_sparse_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::AdamDenseParamSparseGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/selected_rows/cpu/adamw_kernel.cc

@@ -0,0 +1,147 @@
+// 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/selected_rows/adamw_kernel.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/adam_kernel.h"
+#include "paddle/phi/kernels/funcs/adam_functors.h"
+#include "paddle/phi/kernels/selected_rows/adam_kernel.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void AdamwDenseParamSparseGradKernel(const Context& dev_ctx,
+                                     const DenseTensor& param,
+                                     const SelectedRows& grad,
+                                     const DenseTensor& learning_rate,
+                                     const DenseTensor& moment1,
+                                     const DenseTensor& moment2,
+                                     const optional<DenseTensor>& moment2_max,
+                                     const DenseTensor& beta1_pow,
+                                     const DenseTensor& beta2_pow,
+                                     const optional<DenseTensor>& master_param,
+                                     const optional<DenseTensor>& skip_update,
+                                     const Scalar& beta1,
+                                     const Scalar& beta2,
+                                     const Scalar& epsilon,
+                                     float lr_ratio,
+                                     float coeff,
+                                     bool with_decay,
+                                     bool lazy_mode,
+                                     int64_t min_row_size_to_use_multithread,
+                                     bool multi_precision,
+                                     bool use_global_beta_pow,
+                                     bool amsgrad,
+                                     DenseTensor* param_out,
+                                     DenseTensor* moment1_out,
+                                     DenseTensor* moment2_out,
+                                     DenseTensor* moment2_max_out,
+                                     DenseTensor* beta1_pow_out,
+                                     DenseTensor* beta2_pow_out,
+                                     DenseTensor* master_param_outs) {
+  bool skip_update_ = false;
+  if (skip_update.is_initialized()) {
+    PADDLE_ENFORCE_EQ(
+        skip_update->numel(),
+        1,
+        errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
+                                skip_update->numel()));
+    std::vector<bool> skip_update_vec;
+    TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
+    skip_update_ = skip_update_vec[0];
+  }
+  VLOG(3) << "Skip update" << skip_update_;
+
+  if (skip_update_ || !with_decay) {
+    AdamDenseParamSparseGradKernel<T, Context>(dev_ctx,
+                                               param,
+                                               grad,
+                                               learning_rate,
+                                               moment1,
+                                               moment2,
+                                               moment2_max,
+                                               beta1_pow,
+                                               beta2_pow,
+                                               master_param,
+                                               skip_update,
+                                               beta1,
+                                               beta2,
+                                               epsilon,
+                                               lazy_mode,
+                                               min_row_size_to_use_multithread,
+                                               multi_precision,
+                                               use_global_beta_pow,
+                                               amsgrad,
+                                               param_out,
+                                               moment1_out,
+                                               moment2_out,
+                                               moment2_max_out,
+                                               beta1_pow_out,
+                                               beta2_pow_out,
+                                               master_param_outs);
+    return;
+  }
+
+  auto* param_ =
+      master_param.is_initialized() ? master_param.get_ptr() : &param;
+  T coeff_ = static_cast<T>(coeff);
+  T lr_ratio_ = static_cast<T>(lr_ratio);
+  funcs::AdamWFunctor<T, funcs::CPUAdamW> functor(
+      coeff_,
+      lr_ratio_,
+      learning_rate.data<T>(),
+      const_cast<T*>(param_->data<T>()));
+  functor(param_->numel());
+
+  AdamDenseParamSparseGradKernel<T, Context>(dev_ctx,
+                                             param,
+                                             grad,
+                                             learning_rate,
+                                             moment1,
+                                             moment2,
+                                             moment2_max,
+                                             beta1_pow,
+                                             beta2_pow,
+                                             master_param,
+                                             skip_update,
+                                             beta1,
+                                             beta2,
+                                             epsilon,
+                                             lazy_mode,
+                                             min_row_size_to_use_multithread,
+                                             multi_precision,
+                                             use_global_beta_pow,
+                                             amsgrad,
+                                             param_out,
+                                             moment1_out,
+                                             moment2_out,
+                                             moment2_max_out,
+                                             beta1_pow_out,
+                                             beta2_pow_out,
+                                             master_param_outs);
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(adamw_dense_param_sparse_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::AdamwDenseParamSparseGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/selected_rows/cpu/add_n_kernel.cc

@@ -0,0 +1,25 @@
+// 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/selected_rows/impl/add_n_kernel_impl.h"
+
+PD_REGISTER_KERNEL(add_n_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::AddNKernel,
+                   float,
+                   double,
+                   int,
+                   phi::bfloat16,
+                   int64_t) {}

paddle/phi/kernels/selected_rows/cpu/clip_by_norm_kernel.cc

@@ -0,0 +1,22 @@
+// 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/selected_rows/clip_by_norm_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/clip_by_norm_kernel_impl.h"
+
+PD_REGISTER_KERNEL(
+    clip_by_norm_sr, CPU, ALL_LAYOUT, phi::sr::ClipByNormKernel, float) {}

paddle/phi/kernels/selected_rows/cpu/clip_kernel.cc

@@ -0,0 +1,28 @@
+// 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/selected_rows/clip_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/clip_kernel_impl.h"
+
+PD_REGISTER_KERNEL(clip_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::ClipSparseKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/phi/kernels/selected_rows/cpu/dgc_clip_by_norm_kernel.cc

@@ -0,0 +1,19 @@
+// 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.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/dgc_clip_by_norm_kernel_impl.h"
+PD_REGISTER_KERNEL(
+    dgc_clip_by_norm_sr, CPU, ALL_LAYOUT, phi::sr::DGCClipByNormKernel, float) {
+}

paddle/phi/kernels/selected_rows/cpu/ftrl_kernel.cc

@@ -0,0 +1,18 @@
+// 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.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/ftrl_kernel_impl.h"
+
+PD_REGISTER_KERNEL(ftrl_sr, CPU, ALL_LAYOUT, phi::sr::FTRLOpKernel, float) {}

paddle/phi/kernels/selected_rows/cpu/get_tensor_from_selected_rows_kernel.cc

@@ -0,0 +1,24 @@
+// 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.
+
+#include "paddle/phi/kernels/selected_rows/impl/get_tensor_from_selected_rows_kernel_impl.h"
+
+PD_REGISTER_KERNEL(get_tensor_from_selected_rows,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::GetTensorFromSelectedRowsKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/phi/kernels/selected_rows/cpu/lamb_kernel.cc

@@ -0,0 +1,21 @@
+// 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/selected_rows/lamb_kernel.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/lamb_kernel_impl.h"
+
+PD_REGISTER_KERNEL(
+    lamb_sr, CPU, ALL_LAYOUT, phi::sr::LambKernel, float, double) {}

paddle/phi/kernels/selected_rows/cpu/load_kernel.cc

@@ -0,0 +1,18 @@
+// 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.
+
+#include "paddle/phi/kernels/selected_rows/impl/load_kernel_impl.h"
+
+PD_REGISTER_KERNEL(
+    load_sr, CPU, ALL_LAYOUT, phi::sr::LoadSelectedRowsKernel, float) {}

paddle/phi/kernels/selected_rows/cpu/lookup_table_grad_kernel.cc

@@ -0,0 +1,161 @@
+// 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.
+
+#include <string>
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+namespace sr {
+
+constexpr int64_t kNoPadding = -1;
+
+template <typename T, typename Context>
+void LookupTableGradKernel(const Context &dev_ctx,
+                           const SelectedRows &w,
+                           const DenseTensor &ids_in,
+                           const DenseTensor &out_grad,
+                           bool is_sparse,
+                           int64_t padding_idx,
+                           bool is_test,
+                           DenseTensor *w_grad) {
+  DDim table_dim;
+
+  auto *table_t = &w;
+  table_dim = table_t->value().dims();
+
+  // Since paddings are not trainable and fixed in forward, the gradient of
+  // paddings makes no sense and we don't deal with it in backward.
+
+  auto *ids = &ids_in;
+  auto *d_output = &out_grad;
+  auto *d_table = w_grad;
+
+  auto *ids_data = ids->data<int64_t>();
+
+  int64_t N = table_dim[0];
+  int64_t D = table_dim[1];
+
+  auto *d_output_data = d_output->data<T>();
+  auto *d_table_data = dev_ctx.template Alloc<T>(d_table);
+
+  memset(d_table_data, 0, d_table->numel() * sizeof(T));
+
+  for (int64_t i = 0; i < ids->numel(); ++i) {
+    if (padding_idx != kNoPadding && ids_data[i] == padding_idx) {
+      // the gradient of padding_idx should be 0, already done by memset, so
+      // do nothing.
+    } else {
+      PADDLE_ENFORCE_LT(
+          ids_data[i],
+          N,
+          common::errors::InvalidArgument(
+              "Variable value (input) of OP(lookup_table_grad) "
+              "expected >= 0 and < %ld, but got %ld. Please check input "
+              "value.",
+              N,
+              ids_data[i]));
+      PADDLE_ENFORCE_GE(
+          ids_data[i],
+          0,
+          common::errors::InvalidArgument(
+              "Variable value (input) of OP(lookup_table_grad) "
+              "expected >= 0 and < %ld, but got %ld. Please check input "
+              "value.",
+              N,
+              ids_data[i]));
+      for (int j = 0; j < D; ++j) {
+        d_table_data[ids_data[i] * D + j] += d_output_data[i * D + j];
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void LookupTableSparseGradKernel(const Context &dev_ctx,
+                                 const SelectedRows &w,
+                                 const DenseTensor &ids_in,
+                                 const DenseTensor &out_grad,
+                                 bool is_sparse,
+                                 int64_t padding_idx,
+                                 bool is_test,
+                                 SelectedRows *w_grad) {
+  DDim table_dim;
+
+  auto *table_t = &w;
+  table_dim = table_t->value().dims();
+
+  // Since paddings are not trainable and fixed in forward, the gradient of
+  // paddings makes no sense and we don't deal with it in backward.
+
+  auto *ids = &ids_in;
+  auto *d_output = &out_grad;
+  auto *d_table = w_grad;
+
+  auto *ids_data = ids->data<int64_t>();
+  int64_t ids_num = ids->numel();
+
+  std::vector<int64_t> new_rows;
+  new_rows.resize(ids_num);
+  std::memcpy(&new_rows[0], ids_data, ids_num * sizeof(int64_t));
+  d_table->set_rows(new_rows);
+
+  auto *d_table_value = d_table->mutable_value();
+  d_table_value->Resize({ids_num, table_dim[1]});
+  dev_ctx.template Alloc<T>(d_table_value);
+  d_table->set_height(table_dim[0]);
+
+  auto *d_output_data = d_output->data<T>();
+  auto *d_table_data = d_table_value->data<T>();
+
+  auto d_output_dims = d_output->dims();
+  auto d_output_dims_2d =
+      common::flatten_to_2d(d_output_dims, d_output_dims.size() - 1);
+  PADDLE_ENFORCE_EQ(d_table_value->dims(),
+                    d_output_dims_2d,
+                    common::errors::InvalidArgument(
+                        "ShapeError: The shape of lookup_table@Grad and "
+                        "output@Grad should be same. "
+                        "But received lookup_table@Grad's shape = [%s], "
+                        "output@Grad's shape = [%s].",
+                        d_table_value->dims(),
+                        d_output_dims_2d));
+  memcpy(d_table_data, d_output_data, sizeof(T) * d_output->numel());
+}
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(lookup_table_grad_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::LookupTableGradKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}
+
+PD_REGISTER_KERNEL(lookup_table_sparse_grad_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::LookupTableSparseGradKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}

paddle/phi/kernels/selected_rows/cpu/lookup_table_kernel.cc

@@ -0,0 +1,131 @@
+// 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.
+
+#include <string>
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+namespace sr {
+
+constexpr int64_t kNoPadding = -1;
+
+template <typename T, typename Context>
+void LookupTableKernel(const Context &dev_ctx,
+                       const SelectedRows &w,
+                       const DenseTensor &ids_in,
+                       bool is_sparse,
+                       bool is_distributed UNUSED,
+                       int64_t padding_idx,
+                       bool remote_prefetch UNUSED,
+                       const std::string &entry_config UNUSED,
+                       bool is_test,
+                       const std::string &entry UNUSED,
+                       const std::string &table_class UNUSED,
+                       const std::vector<std::string> &table_names UNUSED,
+                       int trainer_id UNUSED,
+                       bool grad_inplace UNUSED,
+                       const std::vector<std::string> &epmap UNUSED,
+                       const std::vector<int64_t> &height_sections UNUSED,
+                       SelectedRows *out) {
+  auto *ids_t = &ids_in;  // int tensor
+  auto *output_t = out;   // float tensor
+
+  int64_t *ids = const_cast<int64_t *>(ids_t->data<int64_t>());
+  int64_t ids_numel = ids_t->numel();
+
+  const auto &table_t = w;
+  int64_t row_width = table_t.value().dims()[1];
+  const auto *table = table_t.value().data<T>();
+  auto *output = dev_ctx.template Alloc<T>(output_t);
+  auto input_data_type = table_t.value().dtype();
+  for (int64_t i = 0; i < ids_numel; ++i) {
+    if (padding_idx != kNoPadding && ids[i] == padding_idx) {
+      memset(output + i * row_width, 0, row_width * sizeof(T));
+    } else {
+      PADDLE_ENFORCE_GE(ids[i],
+                        0,
+                        common::errors::InvalidArgument(
+                            "Variable value (input) of OP(lookup_table) "
+                            "expected >= 0. But received %ld",
+                            ids[i]));
+      if (is_test) {
+        auto id_index = table_t.GetIndexFromId(ids[i]);
+
+        if (id_index != -1) {
+          if (input_data_type == phi::DataType::INT8 ||
+              input_data_type == phi::DataType::INT16 ||
+              input_data_type == phi::DataType::BFLOAT16) {
+            memcpy(output + i * row_width,
+                   table + id_index * row_width,
+                   row_width * sizeof(T));
+          } else {
+            auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
+            blas.VCOPY(row_width,
+                       table + id_index * row_width,
+                       output + i * row_width);
+          }
+        } else {
+          memset(output + i * row_width, 0, row_width * sizeof(T));
+        }
+      } else {
+        auto id_index = table_t.Index(ids[i]);
+        PADDLE_ENFORCE_GE(ids[i],
+                          0,
+                          common::errors::InvalidArgument(
+                              "Variable value (input) of OP(lookup_table) "
+                              "expected >= 0. But received %ld",
+                              ids[i]));
+        PADDLE_ENFORCE_GE(
+            id_index,
+            0,
+            common::errors::InvalidArgument(
+                "the input key should be exists. But received %d.", id_index));
+
+        if (input_data_type == phi::DataType::INT8 ||
+            input_data_type == phi::DataType::INT16 ||
+            input_data_type == phi::DataType::BFLOAT16) {
+          memcpy(output + i * row_width,
+                 table + id_index * row_width,
+                 row_width * sizeof(T));
+        } else {
+          auto blas = funcs::GetBlas<CPUContext, T>(dev_ctx);
+          blas.VCOPY(
+              row_width, table + id_index * row_width, output + i * row_width);
+        }
+      }
+    }
+  }
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(lookup_table_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::LookupTableKernel,
+                   float,
+                   double,
+                   int8_t,
+                   int16_t,
+                   phi::bfloat16) {}

paddle/phi/kernels/selected_rows/cpu/nce_grad_kernel.cc

@@ -0,0 +1,220 @@
+// 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.
+
+#include <math.h>
+
+#include <iterator>
+#include <random>
+#include <set>
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/math/sampler.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+#include "paddle/utils/optional.h"
+
+namespace phi {
+namespace sr {
+
+using Sampler = phi::math::Sampler;
+
+template <typename T, typename Context>
+void NCEGradKernel(const Context &dev_ctx,
+                   const DenseTensor &input_in,
+                   const DenseTensor &label_in,
+                   const optional<DenseTensor> &bias_in,
+                   const DenseTensor &weight_in,
+                   const DenseTensor &sample_logits_in,
+                   const DenseTensor &sample_labels_in,
+                   const optional<DenseTensor> &sample_weight_in,
+                   const optional<DenseTensor> &custom_dist_probs,
+                   const optional<DenseTensor> &custom_dist_alias,
+                   const optional<DenseTensor> &custom_dist_alias_probs,
+                   const DenseTensor &cost_grad,
+                   int num_total_classes,
+                   const std::vector<int> &custom_neg_classes,
+                   int num_neg_samples,
+                   int sampler_in,
+                   int seed,
+                   bool is_sparse,
+                   bool remote_prefetch,
+                   bool is_test,
+                   DenseTensor *input_grad,
+                   DenseTensor *bias_grad,
+                   SelectedRows *weight_grad) {
+  auto d_out = &cost_grad;
+  const T *d_out_data = d_out->data<T>();
+  auto label = &label_in;
+  auto sample_out = &sample_logits_in;
+  const T *sample_out_data = sample_out->data<T>();
+  auto sample_labels = &sample_labels_in;
+  const int64_t *sample_labels_data = sample_labels->data<int64_t>();
+  auto sample_weight = sample_weight_in.get_ptr();
+  const T *sample_weight_data = nullptr;
+  if (sample_weight != nullptr) {
+    sample_weight_data = sample_weight->data<T>();
+  }
+  int num_true_class = 1;
+  if (label != nullptr) {
+    num_true_class = label->dims()[1];
+  }
+
+  int sampler_type = sampler_in;
+  Sampler *sampler;
+  switch (sampler_type) {
+    case 0: {
+      sampler = new phi::math::UniformSampler(num_total_classes - 1, seed);
+      break;
+    }
+    case 1: {
+      sampler = new phi::math::LogUniformSampler(num_total_classes - 1, seed);
+      break;
+    }
+    case 2: {
+      auto dist_probs = custom_dist_probs.get_ptr();
+      auto dist_alias = custom_dist_alias.get_ptr();
+      auto dist_alias_probs = custom_dist_alias_probs.get_ptr();
+
+      PADDLE_ENFORCE_EQ(
+          dist_probs->numel(),
+          num_total_classes,
+          common::errors::InvalidArgument(
+              "ShapeError: The number of elements in Input(CustomDistProbs) "
+              "should be equal to the number of total classes. But Received: "
+              "Input(CustomDistProbs).numel() = %d, Attr(num_total_classes) "
+              "= %d.",
+              dist_probs->numel(),
+              num_total_classes));
+      PADDLE_ENFORCE_EQ(
+          dist_alias->numel(),
+          num_total_classes,
+          common::errors::InvalidArgument(
+              "ShapeError: The number of elements in Input(CustomDistAlias) "
+              "should be equal to the number of total classes. But Received: "
+              "Input(CustomDistAlias).numel() = %d, Attr(num_total_classes) "
+              "= %d.",
+              dist_alias->numel(),
+              num_total_classes));
+      PADDLE_ENFORCE_EQ(
+          dist_alias_probs->numel(),
+          num_total_classes,
+          common::errors::InvalidArgument(
+              "ShapeError: The number of elements in "
+              "Input(CustomDistAliasProbs) "
+              "should be equal to the number of total classes. But Received: "
+              "Input(CustomDistAliasProbs).numel() = %d, "
+              "Attr(num_total_classes) = %d.",
+              dist_alias_probs->numel(),
+              num_total_classes));
+
+      const float *probs_data = dist_probs->data<float>();
+      const int *alias_data = dist_alias->data<int>();
+      const float *alias_probs_data = dist_alias_probs->data<float>();
+      sampler = new phi::math::CustomSampler(num_total_classes - 1,
+                                             probs_data,
+                                             alias_data,
+                                             alias_probs_data,
+                                             seed);
+      break;
+    }
+    default: {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Unsupported SamplerType. SamplerType should be 0: Uniform, "
+          "1: LogUniform or 2: CustomDist. Received SamplerType: %d",
+          sampler_type));
+    }
+  }
+
+  //    T b = 1. / num_total_classes * num_neg_samples;
+  DenseTensor sample_grad;  // tmp tensor
+  sample_grad.Resize(sample_labels->dims());
+  T *sample_grad_data = dev_ctx.template Alloc<T>(&sample_grad);
+
+  // backward cost
+  for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+    int64_t label_idx = i % sample_labels->dims()[1];
+    int64_t sample_idx = i / sample_labels->dims()[1];
+    float b = sampler->Probability(sample_labels_data[i]) * num_neg_samples;
+    T o = sample_out_data[i];
+    T w = sample_weight == nullptr ? 1 : sample_weight_data[sample_idx];
+    sample_grad_data[i] = label_idx < num_true_class
+                              ? w * (b / (o + b)) * (o - 1)
+                              : w * (o * (1 - o) / (o + b));
+    sample_grad_data[i] *= d_out_data[sample_idx];
+  }
+
+  // get d_bias
+  auto d_bias = bias_grad;
+  if (d_bias != nullptr) {
+    T *d_bias_data = dev_ctx.template Alloc<T>(d_bias);
+    std::fill(d_bias_data, d_bias_data + d_bias->numel(), 0.0);
+    for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+      d_bias_data[sample_labels_data[i]] += sample_grad_data[i];
+    }
+  }
+
+  if (!is_sparse) {
+    PADDLE_THROW(
+        common::errors::InvalidArgument("The parameter weight_grad of a NCE_OP "
+                                        "must be DenseTensor"));
+  } else {
+    std::vector<int64_t> labels;
+    for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+      labels.push_back(sample_labels_data[i]);
+    }
+    std::set<T> st(labels.begin(), labels.end());
+    labels.assign(st.begin(), st.end());
+
+    DDim table_dim = weight_in.dims();
+
+    auto d_w = weight_grad;
+
+    d_w->set_rows(labels);
+    d_w->set_height(table_dim[0]);
+
+    auto *d_table_value = d_w->mutable_value();
+    d_table_value->Resize({static_cast<int64_t>(labels.size()), table_dim[1]});
+    auto d_w_data = dev_ctx.template Alloc<T>(d_table_value);
+    std::fill(d_w_data, d_w_data + d_table_value->numel(), 0.0);
+
+    auto d_w_matrix = EigenMatrix<T>::From(*d_table_value);
+    auto x_matrix = EigenMatrix<T>::From(input_in);
+    for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+      d_w_matrix.chip(d_w->Index(sample_labels_data[i]), 0) +=
+          x_matrix.chip(static_cast<int>(i / sample_labels->dims()[1]), 0) *
+          sample_grad_data[i];
+    }
+  }
+
+  // get d_x
+  auto d_x = input_grad;
+  if (d_x != nullptr) {
+    auto *d_x_data = dev_ctx.template Alloc<T>(d_x);
+    std::fill(d_x_data, d_x_data + d_x->numel(), 0.0);
+    auto d_x_matrix = EigenMatrix<T>::From(*d_x);
+    auto w_matrix = EigenMatrix<T>::From(weight_in);
+    for (int64_t i = 0; i < sample_labels->numel(); ++i) {
+      d_x_matrix.chip(static_cast<int>(i / sample_labels->dims()[1]), 0) +=
+          w_matrix.chip(sample_labels_data[i], 0) * sample_grad_data[i];
+    }
+  }
+
+  delete sampler;
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(
+    nce_sr_grad, CPU, ALL_LAYOUT, phi::sr::NCEGradKernel, float, double) {}

paddle/phi/kernels/selected_rows/cpu/save_kernel.cc

@@ -0,0 +1,31 @@
+// 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.
+
+#include "paddle/phi/kernels/selected_rows/impl/save_kernel_impl.h"
+
+PD_REGISTER_KERNEL(save_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::SaveSelectedRowsKernel,
+                   float,
+                   double,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int64_t,
+                   phi::float16,
+                   phi::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}

paddle/phi/kernels/selected_rows/cpu/share_data_kernel.cc

@@ -0,0 +1,29 @@
+// 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.
+
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/kernels/selected_rows/impl/share_data_kernel_impl.h"
+
+PD_REGISTER_KERNEL(share_data_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::ShareDataKernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::float16) {}

paddle/phi/kernels/selected_rows/cpu/uniform_random_batch_size_like_kernel.cc

@@ -0,0 +1,77 @@
+// 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.
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/uniform_random_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void CPUUniformRandomKernel(const Context& dev_ctx,
+                            const SelectedRows& input,
+                            const std::vector<int>& shape,
+                            int input_dim_idx,
+                            int output_dim_idx,
+                            float min,
+                            float max,
+                            int seed,
+                            int diag_num,
+                            int diag_step,
+                            float diag_val,
+                            DataType dtype UNUSED,
+                            SelectedRows* out) {
+  out->set_rows(input.rows());
+  out->set_height(input.height());
+  DenseTensor* tensor = out->mutable_value();
+  T* data = dev_ctx.template Alloc<T>(tensor);
+  int64_t size = tensor->numel();
+
+  funcs::UniformRealDistribution<T>(
+      data, size, min, max, static_cast<unsigned int>(seed));
+
+  unsigned int diag_num_tmp = static_cast<unsigned int>(diag_num);
+  unsigned int diag_step_tmp = static_cast<unsigned int>(diag_step);
+  auto diag_val_tmp = static_cast<T>(diag_val);
+  if (diag_num_tmp > 0) {
+    PADDLE_ENFORCE_GT(
+        size,
+        (diag_num_tmp - 1) * (diag_step_tmp + 1),
+        common::errors::InvalidArgument(
+            "ShapeInvalid: the diagonal's elements is equal (num-1) "
+            "* (step-1) with num %d, step %d,"
+            "It should be smaller than %d, but received %d",
+            diag_num_tmp,
+            diag_step_tmp,
+            (diag_num_tmp - 1) * (diag_step_tmp + 1),
+            size));
+    for (int64_t i = 0; i < diag_num_tmp; ++i) {
+      int64_t pos = i * diag_step_tmp + i;
+      data[pos] = diag_val_tmp;
+    }
+  }
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(uniform_random_batch_size_like_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::CPUUniformRandomKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}

paddle/phi/kernels/selected_rows/elementwise_multiply_kernel.cc

@@ -0,0 +1,106 @@
+/* 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/selected_rows/elementwise_multiply_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/elementwise_multiply_kernel.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void MultiplyRawKernel(const Context& dev_ctx,
+                       const SelectedRows& x,
+                       const DenseTensor& y,
+                       int axis,
+                       SelectedRows* out) {
+  PADDLE_ENFORCE_EQ(common::product(y.dims()),
+                    1,
+                    common::errors::InvalidArgument(
+                        "For MultiplyKernel, if X is Sparse, Y must "
+                        "contain only one element."));
+  out->set_rows(x.rows());
+  out->set_height(x.height());
+  auto z = out->mutable_value();
+  z->Resize(x.value().dims());
+  dev_ctx.Alloc(z, x.value().dtype());
+  MultiplyRawKernel<T, Context>(dev_ctx, x.value(), y, axis, z);
+}
+
+template <typename T, typename Context>
+void MultiplyKernel(const Context& dev_ctx,
+                    const SelectedRows& x,
+                    const DenseTensor& y,
+                    SelectedRows* out) {
+  int axis = -1;
+  MultiplyRawKernel<T, Context>(dev_ctx, x, y, axis, out);
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(multiply_raw_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::MultiplyRawKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::bfloat16,
+                   phi::complex64,
+                   phi::complex128) {}
+PD_REGISTER_KERNEL(multiply_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::MultiplyKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::bfloat16,
+                   phi::complex64,
+                   phi::complex128) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(multiply_raw_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::MultiplyRawKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::bfloat16,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {}
+PD_REGISTER_KERNEL(multiply_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::MultiplyKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::bfloat16,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {}
+#endif

paddle/phi/kernels/selected_rows/elementwise_multiply_kernel.h

@@ -0,0 +1,38 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/legacy/elementwise_multiply_kernel.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void MultiplyRawKernel(const Context& dev_ctx,
+                       const SelectedRows& x,
+                       const DenseTensor& y,
+                       int axis,
+                       SelectedRows* out);
+
+template <typename T, typename Context>
+void MultiplyKernel(const Context& dev_ctx,
+                    const SelectedRows& x,
+                    const DenseTensor& y,
+                    SelectedRows* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/full_kernel.cc

@@ -0,0 +1,145 @@
+/* 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/selected_rows/full_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#endif
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void FullKernel(const Context& dev_ctx,
+                const IntArray& shape,
+                const Scalar& val,
+                DataType dtype,
+                SelectedRows* out) {
+  phi::FullKernel<T>(dev_ctx, shape, val, dtype, out->mutable_value());
+}
+
+template <typename T, typename Context>
+void FullWithTensorKernel(const Context& dev_ctx,
+                          const DenseTensor& value,
+                          const IntArray& shape,
+                          DataType dtype,
+                          SelectedRows* out) {
+  phi::FullWithTensorKernel<T>(
+      dev_ctx, value, shape, dtype, out->mutable_value());
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(full_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::FullKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::float16,
+                   phi::bfloat16,
+                   phi::complex64,
+                   phi::complex128) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(full_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::FullKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {}
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+PD_REGISTER_KERNEL(full_sr,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::sr::FullKernel,
+                   float,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::float16) {}
+#endif
+
+PD_REGISTER_KERNEL(full_with_tensor_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::FullWithTensorKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::float16,
+                   phi::bfloat16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::CPU);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(full_with_tensor_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::FullWithTensorKernel,
+                   float,
+                   double,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::float16,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::CPU);
+}
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+PD_REGISTER_KERNEL(full_with_tensor_sr,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::sr::FullWithTensorKernel,
+                   float,
+                   uint8_t,
+                   int16_t,
+                   int,
+                   int64_t,
+                   bool,
+                   phi::float16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::CPU);
+}
+#endif

paddle/phi/kernels/selected_rows/full_kernel.h

@@ -0,0 +1,38 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void FullKernel(const Context& dev_ctx,
+                const IntArray& shape,
+                const Scalar& val,
+                DataType dtype,
+                SelectedRows* out);
+
+template <typename T, typename Context>
+void FullWithTensorKernel(const Context& dev_ctx,
+                          const DenseTensor& value,
+                          const IntArray& shape,
+                          DataType dtype,
+                          SelectedRows* out);
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/gpu/adam_kernel.cu

@@ -0,0 +1,342 @@
+// 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/selected_rows/adam_kernel.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/adam_functors.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>
+__global__ void UpdateBetaPow(T beta1,
+                              T beta2,
+                              const T* beta1_pow_,
+                              const T* beta2_pow_,
+                              T* beta1_pow_out,
+                              T* beta2_pow_out) {
+  *beta1_pow_out = beta1 * beta1_pow_[0];
+  *beta2_pow_out = beta2 * beta2_pow_[0];
+}
+
+template <typename T, typename MT>
+__global__ void SparseAdamCUDAKernelREG(MT beta1,
+                                        MT beta2,
+                                        MT epsilon,
+                                        const MT beta1_pow,
+                                        const MT beta2_pow,
+                                        const MT* mom1_,
+                                        MT* mom1_out_,
+                                        const MT* mom2_,
+                                        MT* mom2_out_,
+                                        const MT* mom2_max_,
+                                        MT* mom2_max_out_,
+                                        const double* lr_,
+                                        const T* grad_,
+                                        const T* param_,
+                                        T* param_out_,
+                                        const MT* master_param,
+                                        MT* master_param_out,
+                                        const int64_t* rows_,
+                                        int64_t row_numel,
+                                        int64_t row_count,
+                                        bool lazy_mode,
+                                        int ndim,
+                                        bool amsgrad) {
+  int64_t id =
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
+      static_cast<int64_t>(threadIdx.x);
+  MT lr = static_cast<MT>(*lr_);
+
+  for (; id < ndim; id += blockDim.x * gridDim.x) {
+    auto row_idx =
+        funcs::BinarySearch<int64_t>(rows_, row_count, id / row_numel);
+    if (lazy_mode && row_idx < 0) {
+      return;
+    } else {
+      MT mom1 = mom1_[id];
+      MT mom2 = mom2_[id];
+
+      MT p = master_param ? master_param[id] : static_cast<MT>(param_[id]);
+      MT g = row_idx >= 0
+                 ? static_cast<MT>(grad_[row_idx * row_numel + id % row_numel])
+                 : static_cast<MT>(0);
+      mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
+      mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;
+
+      MT denom;
+      if (amsgrad) {
+        MT mom2_max = mom2_max_[id];
+        MT moment2_max_ = std::max(mom2, mom2_max);
+        mom2_max_out_[id] = moment2_max_;
+
+        denom = (sqrt(moment2_max_) / sqrt(static_cast<MT>(1.0) - beta2_pow)) +
+                epsilon;
+      } else {
+        denom = (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
+      }
+
+      p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));
+
+      // Write back to global memory
+      mom1_out_[id] = mom1;
+      mom2_out_[id] = mom2;
+      param_out_[id] = static_cast<T>(p);
+      if (master_param_out) {
+        master_param_out[id] = p;
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void AdamDenseParamSparseGradKernel(const Context& dev_ctx,
+                                    const DenseTensor& param,
+                                    const SelectedRows& grad,
+                                    const DenseTensor& learning_rate,
+                                    const DenseTensor& moment1,
+                                    const DenseTensor& moment2,
+                                    const optional<DenseTensor>& moment2_max,
+                                    const DenseTensor& beta1_pow,
+                                    const DenseTensor& beta2_pow,
+                                    const optional<DenseTensor>& master_param,
+                                    const optional<DenseTensor>& skip_update,
+                                    const Scalar& beta1,
+                                    const Scalar& beta2,
+                                    const Scalar& epsilon,
+                                    bool lazy_mode,
+                                    int64_t min_row_size_to_use_multithread,
+                                    bool multi_precision,
+                                    bool use_global_beta_pow,
+                                    bool amsgrad,
+                                    DenseTensor* param_out,
+                                    DenseTensor* moment1_out,
+                                    DenseTensor* moment2_out,
+                                    DenseTensor* moment2_max_out,
+                                    DenseTensor* beta1_pow_out,
+                                    DenseTensor* beta2_pow_out,
+                                    DenseTensor* master_param_outs) {
+  using MT = typename phi::dtype::MPTypeTrait<T>::Type;
+
+  VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
+
+  bool skip_update_ = false;
+  if (skip_update.is_initialized()) {
+    PADDLE_ENFORCE_EQ(
+        skip_update->numel(),
+        1,
+        errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
+                                skip_update->numel()));
+    std::vector<bool> skip_update_vec;
+    TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
+    skip_update_ = skip_update_vec[0];
+  }
+  // skip_update=true, just copy input to output, and TensorCopy will call
+  // mutable_data
+  if (skip_update_) {
+    VLOG(4) << "Adam skip update";
+    phi::Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
+    phi::Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
+    phi::Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
+    if (amsgrad) {
+      phi::Copy(dev_ctx,
+                moment2_max.get(),
+                dev_ctx.GetPlace(),
+                false,
+                moment2_max_out);
+    }
+    if (!use_global_beta_pow) {
+      phi::Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
+      phi::Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
+    }
+    return;
+  }
+
+  MT beta1_ = beta1.to<MT>();
+  MT beta2_ = beta2.to<MT>();
+  MT epsilon_ = epsilon.to<MT>();
+  VLOG(3) << "beta1_pow.numel() : " << beta1_pow.numel()
+          << "beta2_pow.numel() : " << beta2_pow.numel();
+  VLOG(3) << "param.numel(): " << param.numel();
+  PADDLE_ENFORCE_EQ(
+      beta1_pow_out->numel(),
+      1,
+      errors::InvalidArgument("beta1 pow output size should be 1, but received "
+                              "value is:%d.",
+                              beta1_pow_out->numel()));
+
+  PADDLE_ENFORCE_EQ(
+      beta2_pow_out->numel(),
+      1,
+      errors::InvalidArgument("beta2 pow output size should be 1, but received "
+                              "value is:%d.",
+                              beta2_pow_out->numel()));
+
+  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;
+
+  const MT* moment2_max_in_data =
+      amsgrad ? moment2_max.get().data<MT>() : nullptr;
+  MT* moment2_max_out_data =
+      amsgrad ? dev_ctx.template Alloc<MT>(moment2_max_out) : nullptr;
+
+  if (grad.rows().size() == 0) {
+    VLOG(3) << "grad row size is 0!!";
+    return;
+  }
+
+  std::vector<int64_t> cpu_rows(grad.rows().begin(), grad.rows().end());
+  bool is_strict_sorted = true;
+  for (size_t i = 1; i < cpu_rows.size(); ++i) {
+    if (cpu_rows[i - 1] >= cpu_rows[i]) {
+      is_strict_sorted = false;
+      break;
+    }
+  }
+
+  SelectedRows tmp_grad_merge;
+  const SelectedRows* grad_merge_ptr;
+  if (is_strict_sorted) {
+    grad_merge_ptr = &grad;
+  } else {
+    // merge duplicated rows if any.
+    // The rows of grad_merge have been sorted inside MergeAdd functor
+    funcs::scatter::MergeAdd<Context, T> merge_func;
+    merge_func(dev_ctx, grad, &tmp_grad_merge, true);
+    grad_merge_ptr = &tmp_grad_merge;
+  }
+  auto& grad_merge = *grad_merge_ptr;
+  auto& grad_tensor = grad_merge.value();
+  const T* grad_data = grad_tensor.template data<T>();
+  auto* grad_merge_rows = &grad_merge.rows();
+  phi::MixVector<int64_t> mixv_grad_merge_rows(grad_merge_rows);
+  const int64_t* rows = mixv_grad_merge_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = grad_tensor.numel() / grad_merge.rows().size();
+
+  if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
+    int threads = 512;
+    int64_t ndim = param.numel();
+    int64_t blocks = (ndim + threads - 1) / threads;
+
+    // NOTE(large-tensor): Kernel launch requires int type for grid dimension
+    PADDLE_ENFORCE_LE_INT_MAX(blocks, "blocks");
+    SparseAdamCUDAKernelREG<T, MT>
+        <<<static_cast<int>(blocks), threads, 0, dev_ctx.stream()>>>(
+            beta1_,
+            beta2_,
+            epsilon_,
+            *beta1_pow.data<MT>(),
+            *beta2_pow.data<MT>(),
+            moment1.data<MT>(),
+            dev_ctx.template Alloc<MT>(moment1_out),
+            moment2.data<MT>(),
+            dev_ctx.template Alloc<MT>(moment2_out),
+            moment2_max_in_data,
+            moment2_max_out_data,
+            learning_rate.data<double>(),
+            grad_data,
+            param.data<T>(),
+            dev_ctx.template Alloc<T>(param_out),
+            master_in_data,
+            master_out_data,
+            rows,
+            row_numel,
+            grad_merge.rows().size(),
+            lazy_mode,
+            ndim,
+            amsgrad);
+    if (!use_global_beta_pow) {
+      // Update with cpu
+      dev_ctx.template HostAlloc<MT>(beta1_pow_out)[0] =
+          beta1_ * beta1_pow.data<MT>()[0];
+      dev_ctx.template HostAlloc<MT>(beta2_pow_out)[0] =
+          beta2_ * beta2_pow.data<MT>()[0];
+    }
+  } else {
+    funcs::SparseAdamFunctor<T, funcs::GPUAdam, MT> functor(
+        beta1_,
+        beta2_,
+        epsilon_,
+        beta1_pow.data<MT>(),
+        beta2_pow.data<MT>(),
+        moment1.data<MT>(),
+        dev_ctx.template Alloc<MT>(moment1_out),
+        moment2.data<MT>(),
+        dev_ctx.template Alloc<MT>(moment2_out),
+        moment2_max_in_data,
+        moment2_max_out_data,
+        learning_rate.data<double>(),
+        grad_data,
+        param.data<T>(),
+        dev_ctx.template Alloc<T>(param_out),
+        master_in_data,
+        master_out_data,
+        rows,
+        row_numel,
+        grad_merge.rows().size(),
+        lazy_mode,
+        amsgrad);
+
+    // FIXME(minqiyang): remove BinarySearch in GPU later
+    funcs::ForRange<Context> for_range(dev_ctx, param.numel());
+    for_range(functor);
+    if (!use_global_beta_pow) {
+      // update beta1 and beta2
+      UpdateBetaPow<MT><<<1, 32, 0, dev_ctx.stream()>>>(
+          beta1_,
+          beta2_,
+          beta1_pow.data<MT>(),
+          beta2_pow.data<MT>(),
+          dev_ctx.template Alloc<MT>(beta1_pow_out),
+          dev_ctx.template Alloc<MT>(beta2_pow_out));
+    }
+  }
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(adam_dense_param_sparse_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::AdamDenseParamSparseGradKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  // Skip beta1_pow, beta2_pow, skip_update data transform
+  kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(9).SetBackend(phi::Backend::ALL_BACKEND);
+
+  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);
+    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
+  }
+  kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
+  kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
+}

paddle/phi/kernels/selected_rows/gpu/adamw_kernel.cu

@@ -0,0 +1,367 @@
+// 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/selected_rows/adamw_kernel.h"
+
+#include <math.h>  // for sqrt in CPU and CUDA
+
+#include <vector>
+
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/amp_type_traits.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/adam_functors.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>
+__global__ void UpdateAdamWBetaPow(T beta1,
+                                   T beta2,
+                                   const T* beta1_pow_,
+                                   const T* beta2_pow_,
+                                   T* beta1_pow_out,
+                                   T* beta2_pow_out) {
+  *beta1_pow_out = beta1 * beta1_pow_[0];
+  *beta2_pow_out = beta2 * beta2_pow_[0];
+}
+
+template <typename T, typename MT>
+__global__ void SparseAdamWCUDAKernelREG(MT beta1,
+                                         MT beta2,
+                                         MT epsilon,
+                                         MT coeff,
+                                         MT lr_ratio,
+                                         const MT beta1_pow,
+                                         const MT beta2_pow,
+                                         const MT* mom1_,
+                                         MT* mom1_out_,
+                                         const MT* mom2_,
+                                         MT* mom2_out_,
+                                         const MT* mom2_max_,
+                                         MT* mom2_max_out_,
+                                         const MT* lr_,
+                                         const T* grad_,
+                                         const T* param_,
+                                         T* param_out_,
+                                         const MT* master_param,
+                                         MT* master_param_out,
+                                         const int64_t* rows_,
+                                         int64_t row_numel,
+                                         int64_t row_count,
+                                         bool lazy_mode,
+                                         int ndim,
+                                         bool amsgrad) {
+  int64_t id =
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
+      static_cast<int64_t>(threadIdx.x);
+  MT lr = *lr_ * lr_ratio;
+
+  for (; id < ndim; id += blockDim.x * gridDim.x) {
+    auto row_idx =
+        funcs::BinarySearch<int64_t>(rows_, row_count, id / row_numel);
+    if (lazy_mode && row_idx < 0) {
+      return;
+    } else {
+      MT mom1 = static_cast<MT>(mom1_[id]);
+      MT mom2 = static_cast<MT>(mom2_[id]);
+
+      MT p = master_param ? master_param[id] : static_cast<MT>(param_[id]);
+      MT g = row_idx >= 0
+                 ? static_cast<MT>(grad_[row_idx * row_numel + id % row_numel])
+                 : static_cast<MT>(0);
+
+      p *= (static_cast<MT>(1.0) - lr * coeff);
+
+      mom1 = beta1 * mom1 + (static_cast<MT>(1.0) - beta1) * g;
+      mom2 = beta2 * mom2 + (static_cast<MT>(1.0) - beta2) * g * g;
+
+      MT denom;
+      if (amsgrad) {
+        MT mom2_max = static_cast<MT>(mom2_max_[id]);
+        MT mom2_max_ = std::max(mom2, mom2_max);
+        mom2_max_out_[id] = mom2_max_;
+
+        denom = (sqrt(mom2_max_) / sqrt(static_cast<MT>(1.0) - beta2_pow)) +
+                epsilon;
+      } else {
+        denom = (sqrt(mom2) / sqrt(static_cast<MT>(1.0) - beta2_pow)) + epsilon;
+      }
+
+      p += (mom1 / denom) * (-(lr / (static_cast<MT>(1.0) - beta1_pow)));
+
+      // Write back to global memory
+      mom1_out_[id] = mom1;
+      mom2_out_[id] = mom2;
+      param_out_[id] = static_cast<T>(p);
+      if (master_param_out) {
+        master_param_out[id] = p;
+      }
+    }
+  }
+}
+
+template <typename T, typename Context>
+void AdamwDenseParamSparseGradKernel(const Context& dev_ctx,
+                                     const DenseTensor& param,
+                                     const SelectedRows& grad,
+                                     const DenseTensor& learning_rate,
+                                     const DenseTensor& moment1,
+                                     const DenseTensor& moment2,
+                                     const optional<DenseTensor>& moment2_max,
+                                     const DenseTensor& beta1_pow,
+                                     const DenseTensor& beta2_pow,
+                                     const optional<DenseTensor>& master_param,
+                                     const optional<DenseTensor>& skip_update,
+                                     const Scalar& beta1,
+                                     const Scalar& beta2,
+                                     const Scalar& epsilon,
+                                     float lr_ratio,
+                                     float coeff,
+                                     bool with_decay,
+                                     bool lazy_mode,
+                                     int64_t min_row_size_to_use_multithread,
+                                     bool multi_precision,
+                                     bool use_global_beta_pow,
+                                     bool amsgrad,
+                                     DenseTensor* param_out,
+                                     DenseTensor* moment1_out,
+                                     DenseTensor* moment2_out,
+                                     DenseTensor* moment2_max_out,
+                                     DenseTensor* beta1_pow_out,
+                                     DenseTensor* beta2_pow_out,
+                                     DenseTensor* master_param_outs) {
+  using MT = typename phi::dtype::MPTypeTrait<T>::Type;
+
+  VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
+
+  MT coeff_ = static_cast<MT>(coeff);
+  MT lr_ratio_ = static_cast<MT>(lr_ratio);
+
+  bool skip_update_ = false;
+  if (skip_update.is_initialized()) {
+    PADDLE_ENFORCE_EQ(
+        skip_update->numel(),
+        1,
+        errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
+                                skip_update->numel()));
+    std::vector<bool> skip_update_vec;
+    TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
+    skip_update_ = skip_update_vec[0];
+  }
+
+  // skip_update=true, just copy input to output, and TensorCopy will call
+  // mutable_data
+  if (skip_update_) {
+    VLOG(4) << "Adamw skip update";
+    phi::Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
+    phi::Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
+    phi::Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
+    if (amsgrad) {
+      phi::Copy(dev_ctx,
+                moment2_max.get(),
+                dev_ctx.GetPlace(),
+                false,
+                moment2_max_out);
+    }
+    if (!use_global_beta_pow) {
+      phi::Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
+      phi::Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
+    }
+    return;
+  }
+
+  // if with_decay = false, coeff = 0
+  if (!with_decay) {
+    coeff_ = static_cast<MT>(0.0);
+  }
+
+  MT beta1_ = beta1.to<MT>();
+  MT beta2_ = beta2.to<MT>();
+  MT epsilon_ = epsilon.to<MT>();
+  VLOG(3) << "beta1_pow.numel() : " << beta1_pow.numel()
+          << "beta2_pow.numel() : " << beta2_pow.numel();
+  VLOG(3) << "param.numel(): " << param.numel();
+  PADDLE_ENFORCE_EQ(
+      beta1_pow_out->numel(),
+      1,
+      errors::InvalidArgument("beta1 pow output size should be 1, but received "
+                              "value is:%d.",
+                              beta1_pow_out->numel()));
+
+  PADDLE_ENFORCE_EQ(
+      beta2_pow_out->numel(),
+      1,
+      errors::InvalidArgument("beta2 pow output size should be 1, but received "
+                              "value is:%d.",
+                              beta2_pow_out->numel()));
+
+  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;
+
+  const MT* moment2_max_in_data =
+      amsgrad ? moment2_max.get().data<MT>() : nullptr;
+  MT* moment2_max_out_data =
+      amsgrad ? dev_ctx.template Alloc<MT>(moment2_max_out) : nullptr;
+
+  if (grad.rows().size() == 0) {
+    VLOG(3) << "grad row size is 0!!";
+    return;
+  }
+
+  std::vector<int64_t> cpu_rows(grad.rows().begin(), grad.rows().end());
+  bool is_strict_sorted = true;
+  for (size_t i = 1; i < cpu_rows.size(); ++i) {
+    if (cpu_rows[i - 1] >= cpu_rows[i]) {
+      is_strict_sorted = false;
+      break;
+    }
+  }
+
+  SelectedRows tmp_grad_merge;
+  const SelectedRows* grad_merge_ptr;
+  if (is_strict_sorted) {
+    grad_merge_ptr = &grad;
+  } else {
+    // merge duplicated rows if any.
+    // The rows of grad_merge have been sorted inside MergeAdd functor
+    funcs::scatter::MergeAdd<Context, T> merge_func;
+    merge_func(dev_ctx, grad, &tmp_grad_merge, true);
+    grad_merge_ptr = &tmp_grad_merge;
+  }
+  auto& grad_merge = *grad_merge_ptr;
+  auto& grad_tensor = grad_merge.value();
+  const T* grad_data = grad_tensor.template data<T>();
+  auto* grad_merge_rows = &grad_merge.rows();
+  phi::MixVector<int64_t> mixv_grad_merge_rows(grad_merge_rows);
+  const int64_t* rows = mixv_grad_merge_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = grad_tensor.numel() / grad_merge.rows().size();
+
+  if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
+    int threads = 512;
+    int64_t ndim = param.numel();
+    int64_t blocks = (ndim + threads - 1) / threads;
+
+    // NOTE(large-tensor): Kernel launch requires int type for grid dimension
+    PADDLE_ENFORCE_LE_INT_MAX(blocks, "blocks");
+    SparseAdamWCUDAKernelREG<T, MT>
+        <<<static_cast<int>(blocks), threads, 0, dev_ctx.stream()>>>(
+            beta1_,
+            beta2_,
+            epsilon_,
+            coeff_,
+            lr_ratio_,
+            *beta1_pow.data<MT>(),
+            *beta2_pow.data<MT>(),
+            moment1.data<MT>(),
+            dev_ctx.template Alloc<MT>(moment1_out),
+            moment2.data<MT>(),
+            dev_ctx.template Alloc<MT>(moment2_out),
+            moment2_max_in_data,
+            moment2_max_out_data,
+            learning_rate.data<MT>(),
+            grad_data,
+            param.data<T>(),
+            dev_ctx.template Alloc<T>(param_out),
+            master_in_data,
+            master_out_data,
+            rows,
+            row_numel,
+            grad_merge.rows().size(),
+            lazy_mode,
+            ndim,
+            amsgrad);
+    if (!use_global_beta_pow) {
+      // Update with cpu
+      dev_ctx.template HostAlloc<MT>(beta1_pow_out)[0] =
+          beta1_ * beta1_pow.data<MT>()[0];
+      dev_ctx.template HostAlloc<MT>(beta2_pow_out)[0] =
+          beta2_ * beta2_pow.data<MT>()[0];
+    }
+  } else {
+    funcs::SparseAdamWFunctor<T, funcs::GPUAdamW, MT> functor(
+        beta1_,
+        beta2_,
+        epsilon_,
+        coeff_,
+        lr_ratio_,
+        beta1_pow.data<MT>(),
+        beta2_pow.data<MT>(),
+        moment1.data<MT>(),
+        dev_ctx.template Alloc<MT>(moment1_out),
+        moment2.data<MT>(),
+        dev_ctx.template Alloc<MT>(moment2_out),
+        moment2_max_in_data,
+        moment2_max_out_data,
+        learning_rate.data<MT>(),
+        grad_data,
+        param.data<T>(),
+        dev_ctx.template Alloc<T>(param_out),
+        master_in_data,
+        master_out_data,
+        rows,
+        row_numel,
+        grad_merge.rows().size(),
+        lazy_mode,
+        amsgrad);
+
+    // FIXME(minqiyang): remove BinarySearch in GPU later
+    funcs::ForRange<Context> for_range(dev_ctx, param.numel());
+    for_range(functor);
+    if (!use_global_beta_pow) {
+      // update beta1 and beta2
+      UpdateAdamWBetaPow<MT><<<1, 32, 0, dev_ctx.stream()>>>(
+          beta1_,
+          beta2_,
+          beta1_pow.data<MT>(),
+          beta2_pow.data<MT>(),
+          dev_ctx.template Alloc<MT>(beta1_pow_out),
+          dev_ctx.template Alloc<MT>(beta2_pow_out));
+    }
+  }
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(adamw_dense_param_sparse_grad,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::AdamwDenseParamSparseGradKernel,
+                   float,
+                   double,
+                   phi::float16) {
+  // Skip beta1_pow, beta2_pow, skip_update data transform
+  kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(9).SetBackend(phi::Backend::ALL_BACKEND);
+
+  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);
+    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(5).SetDataType(phi::DataType::FLOAT32);
+    kernel->OutputAt(6).SetDataType(phi::DataType::FLOAT32);
+  }
+  kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
+  kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
+}

paddle/phi/kernels/selected_rows/gpu/add_n_kernel.cu

@@ -0,0 +1,26 @@
+// 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/selected_rows/impl/add_n_kernel_impl.h"
+
+PD_REGISTER_KERNEL(add_n_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::AddNKernel,
+                   float,
+                   double,
+                   int,
+                   phi::bfloat16,
+                   phi::float16,
+                   int64_t) {}

paddle/phi/kernels/selected_rows/gpu/clip_by_norm_kernel.cu

@@ -0,0 +1,26 @@
+// 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/selected_rows/clip_by_norm_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/clip_by_norm_kernel_impl.h"
+
+PD_REGISTER_KERNEL(clip_by_norm_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::ClipByNormKernel,
+                   float,
+                   phi::float16) {}

paddle/phi/kernels/selected_rows/gpu/clip_kernel.cu

@@ -0,0 +1,29 @@
+// 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/selected_rows/clip_kernel.h"
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/clip_kernel_impl.h"
+
+PD_REGISTER_KERNEL(clip_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::ClipSparseKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   phi::float16) {}

paddle/phi/kernels/selected_rows/gpu/dgc_clip_by_norm_kernel.cu

@@ -0,0 +1,19 @@
+// 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.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/dgc_clip_by_norm_kernel_impl.h"
+PD_REGISTER_KERNEL(
+    dgc_clip_by_norm_sr, GPU, ALL_LAYOUT, phi::sr::DGCClipByNormKernel, float) {
+}

paddle/phi/kernels/selected_rows/gpu/ftrl_kernel.cu

@@ -0,0 +1,18 @@
+// 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.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/ftrl_kernel_impl.h"
+
+PD_REGISTER_KERNEL(ftrl_sr, GPU, ALL_LAYOUT, phi::sr::FTRLOpKernel, float) {}

paddle/phi/kernels/selected_rows/gpu/get_tensor_from_selected_rows_kernel.cu

@@ -0,0 +1,24 @@
+// 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.
+
+#include "paddle/phi/kernels/selected_rows/impl/get_tensor_from_selected_rows_kernel_impl.h"
+
+PD_REGISTER_KERNEL(get_tensor_from_selected_rows,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::GetTensorFromSelectedRowsKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}

paddle/phi/kernels/selected_rows/gpu/lamb_kernel.cu

@@ -0,0 +1,29 @@
+// 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/selected_rows/lamb_kernel.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/lamb_kernel_impl.h"
+
+PD_REGISTER_KERNEL(lamb_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::LambKernel,
+                   phi::float16,
+                   float,
+                   double) {
+  kernel->InputAt(5).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
+}

paddle/phi/kernels/selected_rows/gpu/load_kernel.cu

@@ -0,0 +1,18 @@
+// 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.
+
+#include "paddle/phi/kernels/selected_rows/impl/load_kernel_impl.h"
+
+PD_REGISTER_KERNEL(
+    load_sr, GPU, ALL_LAYOUT, phi::sr::LoadSelectedRowsKernel, float) {}

paddle/phi/kernels/selected_rows/gpu/lookup_table_grad_kernel.cu

@@ -0,0 +1,201 @@
+// 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.
+
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, int BlockDimX, int BlockDimY, int GridDimX>
+__global__ void LookupTableGrad(T *table,
+                                const T *output,
+                                const int64_t *ids,
+                                const int64_t N,
+                                const int64_t K,
+                                const int64_t D) {
+  int idx = threadIdx.x;
+  int64_t idy = static_cast<int64_t>(blockIdx.x) +
+                static_cast<int64_t>(threadIdx.y) * GridDimX;
+
+  while (idy < K) {
+    int64_t id = ids[idy];
+    PADDLE_ENFORCE(
+        id >= 0,
+        "Variable value (input) of OP(lookup_table_grad) "
+        "expected >= 0 and < %ld, but got %ld. Please check input value.",
+        N,
+        id);
+    PADDLE_ENFORCE(
+        id < N,
+        "Variable value (input) of OP(lookup_table_grad) "
+        "expected >= 0 and < %ld, but got %ld. Please check input value.",
+        N,
+        id);
+    const T *out = output + idy * D;
+    T *tab = table + id * D;
+    for (int i = idx; i < D; i += BlockDimX) {
+      CudaAtomicAdd(&tab[i], out[i]);
+    }
+    idy += BlockDimY * GridDimX;
+  }
+}
+
+template <typename T, typename Context>
+void LookupTableGradCUDAKernel(
+    const Context &dev_ctx,
+    const SelectedRows &w,
+    const DenseTensor &ids_in,
+    const DenseTensor &out_grad,
+    bool is_sparse,
+    bool is_distributed UNUSED,
+    int64_t padding_idx,
+    bool remote_prefetch UNUSED,
+    const std::string &entry_config UNUSED,
+    bool is_test,
+    const std::string &entry UNUSED,
+    const std::string &table_class UNUSED,
+    const std::vector<std::string> &table_names UNUSED,
+    int trainer_id UNUSED,
+    bool grad_inplace UNUSED,
+    const std::vector<std::string> &epmap UNUSED,
+    const std::vector<int64_t> &height_sections UNUSED,
+    DenseTensor *w_grad) {
+  // Since paddings are not trainable and fixed in forward, the gradient of
+  // paddings makes no sense and we don't deal with it in backward.
+
+  auto ids_t = &ids_in;
+  auto d_output_t = &out_grad;
+  auto d_table_t = w_grad;
+
+  int64_t N = d_table_t->dims()[0];
+  int64_t D = d_table_t->dims()[1];
+  int64_t K = ids_t->numel();
+  const int64_t *ids = ids_t->data<int64_t>();
+  const T *d_output = d_output_t->data<T>();
+  T *d_table = dev_ctx.template Alloc<T>(d_table_t);
+
+  auto t = EigenVector<T>::Flatten(*d_table_t);
+  t.device(*dev_ctx.eigen_device()) = t.constant(static_cast<T>(0));
+
+#ifdef PADDLE_WITH_HIP
+  dim3 threads(64, 4);
+#else
+  dim3 threads(128, 8);
+#endif  // PADDLE_WITH_HIP
+  dim3 grids(8, 1);
+
+#ifdef PADDLE_WITH_HIP
+  LookupTableGrad<T, 64, 4, 8><<<grids, threads, 0, dev_ctx.stream()>>>(
+      d_table, d_output, ids, N, K, D);
+#else
+  LookupTableGrad<T, 128, 8, 8><<<grids, threads, 0, dev_ctx.stream()>>>(
+      d_table, d_output, ids, N, K, D);
+#endif  // PADDLE_WITH_HIP
+}
+
+template <typename T, typename Context>
+void LookupTableSparseGradCUDAKernel(
+    const Context &dev_ctx,
+    const SelectedRows &w,
+    const DenseTensor &ids_in,
+    const DenseTensor &out_grad,
+    bool is_sparse,
+    bool is_distributed UNUSED,
+    int64_t padding_idx,
+    bool remote_prefetch UNUSED,
+    const std::string &entry_config UNUSED,
+    bool is_test,
+    const std::string &entry UNUSED,
+    const std::string &table_class UNUSED,
+    const std::vector<std::string> &table_names UNUSED,
+    int trainer_id UNUSED,
+    bool grad_inplace UNUSED,
+    const std::vector<std::string> &epmap UNUSED,
+    const std::vector<int64_t> &height_sections UNUSED,
+    SelectedRows *w_grad) {
+  // Since paddings are not trainable and fixed in forward, the gradient of
+  // paddings makes no sense and we don't deal with it in backward.
+  auto *ids = &ids_in;
+  auto *table = &w;
+  auto *d_output = &out_grad;
+  auto *d_table = w_grad;
+
+  auto *ids_data = ids->data<int64_t>();
+  int64_t ids_num = ids->numel();
+
+  auto stream = dev_ctx.stream();
+  // copy GPU memory to CPU pinned memory
+  Vector<int64_t> new_rows;
+  new_rows.resize(ids_num);
+  auto gpu_place = dev_ctx.GetPlace();
+
+  // TODO(yuyang18): Strange code here.
+  phi::MixVector<int64_t> mixv_new_rows(&new_rows);
+  phi::memory_utils::Copy(gpu_place,
+                          mixv_new_rows.CUDAMutableData(dev_ctx.GetPlace()),
+                          gpu_place,
+                          ids_data,
+                          ids_num * sizeof(int64_t),
+                          stream);
+  mixv_new_rows.CopyToCPU();
+  d_table->set_rows(new_rows);
+
+  auto *d_table_value = d_table->mutable_value();
+  d_table_value->Resize({ids_num, table->dims()[1]});
+  dev_ctx.template Alloc<T>(d_table_value);
+
+  auto *d_table_data = d_table_value->data<T>();
+  auto *d_output_data = d_output->data<T>();
+  auto d_output_dims = d_output->dims();
+  auto d_output_dims_2d =
+      common::flatten_to_2d(d_output_dims, d_output_dims.size() - 1);
+  PADDLE_ENFORCE_EQ(d_table_value->dims(),
+                    d_output_dims_2d,
+                    common::errors::InvalidArgument(
+                        "ShapeError: The shape of lookup_table@Grad and "
+                        "output@Grad should be same. "
+                        "But received lookup_table@Grad's shape = [%s], "
+                        "output@Grad's shape = [%s].",
+                        d_table_value->dims(),
+                        d_output_dims_2d));
+  phi::memory_utils::Copy(gpu_place,
+                          d_table_data,
+                          gpu_place,
+                          d_output_data,
+                          d_output->numel() * sizeof(T),
+                          stream);
+}
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(lookup_table_grad_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::LookupTableGradCUDAKernel,
+                   float,
+                   double,
+                   phi::float16) {}
+
+PD_REGISTER_KERNEL(lookup_table_sparse_grad_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::LookupTableSparseGradCUDAKernel,
+                   float,
+                   double,
+                   phi::float16) {}

paddle/phi/kernels/selected_rows/gpu/lookup_table_kernel.cu

@@ -0,0 +1,134 @@
+// 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.
+
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/funcs/eigen/common.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+template <typename T,
+          int BlockDimX,
+          int BlockDimY,
+          int GridDimX,
+          bool PaddingFlag>
+__global__ void LookupTable(T *output,
+                            const T *table,
+                            const int64_t *ids,
+                            const int64_t N,
+                            const int64_t K,
+                            const int64_t D,
+                            const int64_t padding_idx) {
+  int idx = threadIdx.x;
+  int64_t idy = static_cast<int64_t>(blockIdx.x) +
+                static_cast<int64_t>(threadIdx.y) * GridDimX;
+
+  while (idy < K) {
+    int64_t id = ids[idy];
+    PADDLE_ENFORCE(
+        id >= 0,
+        "Variable value (input) of OP(lookup_table) "
+        "expected >= 0 and < %ld, but got %ld. Please check input value.",
+        N,
+        id);
+    PADDLE_ENFORCE(
+        id < N,
+        "Variable value (input) of OP(lookup_table) "
+        "expected >= 0 and < %ld, but got %ld. Please check input value.",
+        N,
+        id);
+    T *out = output + idy * D;
+    const T *tab = table + id * D;
+    for (int i = idx; i < D; i += BlockDimX) {
+      if (PaddingFlag) {
+        if (id == padding_idx)
+          out[i] = static_cast<T>(0);
+        else
+          out[i] = tab[i];
+      } else {
+        out[i] = tab[i];
+      }
+    }
+    idy += BlockDimY * GridDimX;
+  }
+}
+
+template <typename T, typename Context>
+void LookupTableCUDAKernel(const Context &dev_ctx,
+                           const SelectedRows &w,
+                           const DenseTensor &ids_in,
+                           bool is_sparse,
+                           bool is_distributed UNUSED,
+                           int64_t padding_idx,
+                           bool remote_prefetch UNUSED,
+                           const std::string &entry_config UNUSED,
+                           bool is_test,
+                           const std::string &entry UNUSED,
+                           const std::string &table_class UNUSED,
+                           const std::vector<std::string> &table_names UNUSED,
+                           int trainer_id UNUSED,
+                           bool grad_inplace UNUSED,
+                           const std::vector<std::string> &epmap UNUSED,
+                           const std::vector<int64_t> &height_sections UNUSED,
+                           SelectedRows *out) {
+  auto *table_t = &w;
+  auto *ids_t = &ids_in;
+  auto *output_t = out;
+
+  size_t N = table_t->dims()[0];
+  size_t D = table_t->dims()[1];
+  size_t K = ids_t->numel();
+
+  auto *ids = ids_t->data<int64_t>();
+  auto *table = table_t->value().data<T>();
+  auto *output = dev_ctx.template Alloc<T>(output_t);
+
+#ifdef PADDLE_WITH_HIP
+  dim3 threads(64, 4);
+#else
+  dim3 threads(128, 8);
+#endif  // PADDLE_WITH_HIP
+  dim3 grids(8, 1);
+#ifdef PADDLE_WITH_HIP
+  if (padding_idx == -1)
+    LookupTable<T, 64, 4, 8, false><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+  else
+    LookupTable<T, 64, 4, 8, true><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+#else
+  if (padding_idx == -1)
+    LookupTable<T, 128, 8, 8, false><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+  else
+    LookupTable<T, 128, 8, 8, true><<<grids, threads, 0, dev_ctx.stream()>>>(
+        output, table, ids, N, K, D, padding_idx);
+#endif  // PADDLE_WITH_HIP
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(lookup_table_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::LookupTableCUDAKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   int8_t,
+                   int16_t) {}

paddle/phi/kernels/selected_rows/gpu/save_kernel.cu

@@ -0,0 +1,31 @@
+// 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.
+
+#include "paddle/phi/kernels/selected_rows/impl/save_kernel_impl.h"
+
+PD_REGISTER_KERNEL(save_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::SaveSelectedRowsKernel,
+                   float,
+                   double,
+                   int,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int64_t,
+                   phi::float16,
+                   phi::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+}

paddle/phi/kernels/selected_rows/gpu/share_data_kernel.cu

@@ -0,0 +1,29 @@
+// 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.
+
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/kernels/selected_rows/impl/share_data_kernel_impl.h"
+
+PD_REGISTER_KERNEL(share_data_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::ShareDataKernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::float16) {}

paddle/phi/kernels/selected_rows/gpu/uniform_random_batch_size_like_kernel.cu

@@ -0,0 +1,60 @@
+// 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.
+
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/uniform_random_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void GPUUniformRandomKernel(const Context& dev_ctx,
+                            const SelectedRows& input,
+                            const std::vector<int>& shape,
+                            int input_dim_idx,
+                            int output_dim_idx,
+                            float min,
+                            float max,
+                            int seed,
+                            int diag_num,
+                            int diag_step,
+                            float diag_val,
+                            DataType dtype UNUSED,
+                            SelectedRows* out) {
+  out->set_rows(input.rows());
+  out->set_height(input.height());
+  DenseTensor* tensor = out->mutable_value();
+  dev_ctx.template Alloc<T>(tensor);
+  funcs::UniformRandom<T>(reinterpret_cast<const GPUContext&>(dev_ctx),
+                          tensor,
+                          seed,
+                          min,
+                          max,
+                          diag_num,
+                          diag_step,
+                          diag_val);
+}
+
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(uniform_random_batch_size_like_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::GPUUniformRandomKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}

paddle/phi/kernels/selected_rows/hsigmoid_loss_grad_kernel.cc

@@ -0,0 +1,89 @@
+// 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/selected_rows/hsigmoid_loss_grad_kernel.h"
+
+#include <set>
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/mixed_vector.h"
+#include "paddle/phi/kernels/cpu/hsigmoid_loss_grad.h"
+
+namespace phi::sr {
+
+static std::vector<int64_t> PathToRows(const DenseTensor& path) {
+  std::set<int64_t> rows;
+  const int64_t* paths = path.data<int64_t>();
+  for (int64_t i = 0; i < path.numel(); ++i) {
+    int64_t row = paths[i];
+    if (row < 0) {
+      continue;
+    }
+    rows.emplace(row);
+  }
+  return std::vector<int64_t>(rows.begin(), rows.end());
+}
+
+template <typename T, typename Context>
+void HSigmoidLossGradKernel(const Context& dev_ctx,
+                            const DenseTensor& x,
+                            const DenseTensor& w,
+                            const DenseTensor& label,
+                            const optional<DenseTensor>& path,
+                            const optional<DenseTensor>& code,
+                            const optional<DenseTensor>& bias,
+                            const DenseTensor& pre_out,
+                            const DenseTensor& out_grad,
+                            int num_classes,
+                            bool is_sparse,
+                            DenseTensor* x_grad,
+                            SelectedRows* w_grad,
+                            DenseTensor* bias_grad) {
+  PADDLE_ENFORCE_NOT_NULL(
+      path.get_ptr(),
+      errors::NotFound("Custom tree must be set for sparse mode!"));
+  Vector<int64_t> real_rows = PathToRows(*path);
+  w_grad->set_rows(real_rows);
+  // Build a map of id -> row_index to speed up finding the index of one id
+  w_grad->set_height(w.dims()[0]);
+  auto* w_grad_value = w_grad->mutable_value();
+  DDim temp_dim(w.dims());
+  temp_dim[0] = static_cast<int>(real_rows.size());
+  w_grad_value->Resize(temp_dim);
+  phi::HSigmoidLossGradKernelImpl<T>(dev_ctx,
+                                     x,
+                                     w,
+                                     label,
+                                     path,
+                                     code,
+                                     bias,
+                                     pre_out,
+                                     out_grad,
+                                     num_classes,
+                                     is_sparse,
+                                     x_grad,
+                                     w_grad_value,
+                                     bias_grad,
+                                     w_grad);
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(hsigmoid_loss_grad_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::HSigmoidLossGradKernel,
+                   float,
+                   double) {}

paddle/phi/kernels/selected_rows/hsigmoid_loss_grad_kernel.h

@@ -0,0 +1,40 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void HSigmoidLossGradKernel(const Context& dev_ctx,
+                            const DenseTensor& x,
+                            const DenseTensor& w,
+                            const DenseTensor& label,
+                            const optional<DenseTensor>& path,
+                            const optional<DenseTensor>& code,
+                            const optional<DenseTensor>& bias,
+                            const DenseTensor& pre_out,
+                            const DenseTensor& out_grad,
+                            int num_classes,
+                            bool is_sparse,
+                            DenseTensor* x_grad,
+                            SelectedRows* w_grad,
+                            DenseTensor* bias_grad);
+
+}  // namespace sr
+}  // namespace phi

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

@@ -0,0 +1,90 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/kernels/selected_rows/add_n_kernel.h"
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void AddNKernel(const Context &dev_ctx,
+                const std::vector<const SelectedRows *> &x,
+                SelectedRows *out) {
+  dev_ctx.template Alloc<T>(out->mutable_value());
+
+  bool in_place = false;
+  if (x.size() > 0 && x[0]->value().Holder() == out->value().Holder()) {
+    in_place = true;
+  }
+
+  if (in_place && x.size() < 2) {
+    return;
+  }
+
+  std::vector<const SelectedRows *> inputs;
+  SelectedRows temp_in0;
+
+  if (in_place) {
+    auto &in0 = *x[0];
+    temp_in0.set_height(in0.height());
+    temp_in0.set_rows(in0.rows());
+    Copy<Context>(
+        dev_ctx, in0.value(), in0.place(), false, temp_in0.mutable_value());
+    inputs.push_back(&temp_in0);
+    for (size_t i = 1; i < x.size(); ++i) {
+      auto &in = *x[i];
+      if (in.rows().size() > 0) {
+        inputs.push_back(&in);
+      }
+    }
+  } else {
+    for (auto in_var : x) {
+      auto &in = *in_var;
+      if (in.rows().size() > 0) {
+        inputs.push_back(in_var);
+      }
+    }
+  }
+
+  out->mutable_rows()->clear();
+
+  bool has_data = false;
+  for (auto &in : inputs) {
+    if (in->rows().size() > 0) {
+      has_data = true;
+      break;
+    }
+  }
+  if (has_data) {
+    funcs::scatter::MergeAdd<Context, T> merge_add;
+    merge_add(dev_ctx, inputs, out);
+
+    out->SyncIndex();
+
+  } else {
+    // no data, just set a empty out tensor.
+    auto *out_dense = out->mutable_value();
+    out_dense->clear();
+    out_dense->Resize({0});
+    dev_ctx.template Alloc<T>(out_dense);
+  }
+}
+}  // namespace sr
+}  // namespace phi

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

@@ -0,0 +1,45 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/clip_by_norm_kernel.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+#include "paddle/phi/kernels/selected_rows/clip_by_norm_kernel.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ClipByNormKernel(const Context& dev_ctx,
+                      const SelectedRows& x,
+                      float max_norm,
+                      SelectedRows* out) {
+  SelectedRows merged_input;
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, x, &merged_input);
+  auto input = &(merged_input.value());
+  out->set_rows(merged_input.rows());
+  out->set_height(merged_input.height());
+  auto out_tensor = out->mutable_value();
+  out_tensor->Resize(merged_input.value().dims());
+  return phi::ClipByNormKernel<T, Context>(
+      dev_ctx, *input, max_norm, out_tensor);
+}
+
+}  // namespace sr
+}  // namespace phi

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

@@ -0,0 +1,61 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+#include "paddle/phi/kernels/selected_rows/clip_kernel.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ClipSparseKernel(const Context& dev_ctx,
+                      const SelectedRows& x,
+                      const Scalar& min,
+                      const Scalar& max,
+                      SelectedRows* out) {
+  auto max_ = max.to<T>();
+  auto min_ = min.to<T>();
+
+  PADDLE_ENFORCE_LE(
+      min_,
+      max_,
+      errors::InvalidArgument("max should be greater than or equal to min. "
+                              "But received min = %f, max = %f",
+                              static_cast<float>(min_),
+                              static_cast<float>(max_)));
+
+  PADDLE_ENFORCE_NE(&x,
+                    out,
+                    errors::InvalidArgument("Inplace clip is not allowed "
+                                            "when x is SelectedRows"));
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, x, out);
+  auto* out_tensor = out->mutable_value();
+  auto* out_data = out_tensor->data<T>();
+  int64_t numel = out_tensor->numel();
+  phi::Transform<Context> trans;
+  trans(dev_ctx,
+        out_data,
+        out_data + numel,
+        out_data,
+        ClipFunctor<T>(min_, max_));
+}
+}  // namespace sr
+}  // namespace phi

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

@@ -0,0 +1,51 @@
+// 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 "glog/logging.h"
+#include "paddle/phi/kernels/selected_rows/clip_by_norm_kernel.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void DGCClipByNormKernel(const Context& dev_ctx,
+                         const SelectedRows& x_in,
+                         const DenseTensor& current_step_in,
+                         float max_norm,
+                         float rampup_begin_step,
+                         SelectedRows* out) {
+  if (static_cast<int>(rampup_begin_step) < 0) {
+    return;
+  }
+
+  auto current_step_tensor = &current_step_in;
+  auto* current_step = current_step_tensor->data<T>();
+
+  VLOG(10) << "current_step:" << *current_step
+           << ", rampup_begin_step:" << rampup_begin_step;
+
+  if (static_cast<int>(*current_step) < static_cast<int>(rampup_begin_step)) {
+    VLOG(10) << "current_step:" << *current_step
+             << " < rampup_begin_step:" << rampup_begin_step
+             << " so doesn't use dgc_clip_by_norm";
+    return;
+  }
+
+  auto* x = &x_in;
+  SelectedRows* output_selected_rows = out;
+  return phi::sr::ClipByNormKernel<T>(
+      dev_ctx, *x, max_norm, output_selected_rows);
+}
+}  // namespace sr
+}  // namespace phi

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

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

@@ -0,0 +1,32 @@
+// 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/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void GetTensorFromSelectedRowsKernel(const Context& dev_ctx,
+                                     const SelectedRows& x,
+                                     DenseTensor* out) {
+  out->Resize(x.value().dims());
+  dev_ctx.template Alloc<T>(out);
+  phi::Copy(dev_ctx, x.value(), dev_ctx.GetPlace(), false, out);
+}
+}  // namespace sr
+}  // namespace phi

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

@@ -0,0 +1,39 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/kernels/funcs/isfinite_functor.h"
+#include "paddle/phi/kernels/selected_rows/isfinite_kernel.h"
+
+namespace phi {
+
+#define DEFINE_ISFINITE_SR(isfinite)                                      \
+  template <typename T, typename Context>                                 \
+  void isfinite##SR(                                                      \
+      const Context& dev_ctx, const SelectedRows& x, SelectedRows* out) { \
+    if (out && out->numel() == 0) {                                       \
+      dev_ctx.template Alloc<bool>(out);                                  \
+      return;                                                             \
+    }                                                                     \
+    dev_ctx.template Alloc<bool>(out);                                    \
+    Isinf##Kernel<T, Context>(dev_ctx, x.value(), out->mutable_value());  \
+  }
+
+DEFINE_ISFINITE_SR(Isinf)
+DEFINE_ISFINITE_SR(Isnan)
+DEFINE_ISFINITE_SR(Isfinite)
+#undef DEFINE_ISFINITE_SR
+
+}  // namespace phi

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

@@ -0,0 +1,371 @@
+// 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.
+#pragma once
+#include "glog/logging.h"
+
+#include "paddle/phi/common/data_type.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/funcs/lamb_functors.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename MT, typename Context, bool IsMultiPrecision>
+void ComputeRowImpl(const Context& dev_ctx,
+                    const DenseTensor& param,
+                    const SelectedRows& 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,
+                    DenseTensor* param_out,
+                    DenseTensor* mom1_out,
+                    DenseTensor* mom2_out,
+                    DenseTensor* beta1_pow_out,
+                    DenseTensor* beta2_pow_out,
+                    DenseTensor* master_param_out);
+
+template <typename T, typename Context>
+void LambKernel(const Context& dev_ctx,
+                const DenseTensor& param,
+                const SelectedRows& grad,
+                const DenseTensor& learning_rate,
+                const DenseTensor& moment1,
+                const DenseTensor& moment2,
+                const DenseTensor& beta1_pow,
+                const DenseTensor& beta2_pow,
+                const optional<DenseTensor>& master_param,
+                const optional<DenseTensor>& skip_update,
+                float weight_decay,
+                float beta1,
+                float beta2,
+                float epsilon,
+                bool always_adapt,
+                bool multi_precision,
+                DenseTensor* param_out,
+                DenseTensor* moment1_out,
+                DenseTensor* moment2_out,
+                DenseTensor* beta1_pow_out,
+                DenseTensor* beta2_pow_out,
+                DenseTensor* master_param_outs) {
+  using MT = typename phi::dtype::MPTypeTrait<T>::Type;
+  if (multi_precision) {
+    ComputeRowImpl<T, MT, Context, true>(dev_ctx,
+                                         param,
+                                         grad,
+                                         learning_rate,
+                                         moment1,
+                                         moment2,
+                                         beta1_pow,
+                                         beta2_pow,
+                                         master_param,
+                                         skip_update,
+                                         weight_decay,
+                                         beta1,
+                                         beta2,
+                                         epsilon,
+                                         always_adapt,
+                                         multi_precision,
+                                         param_out,
+                                         moment1_out,
+                                         moment2_out,
+                                         beta1_pow_out,
+                                         beta2_pow_out,
+                                         master_param_outs);
+  } else {
+    ComputeRowImpl<T, T, Context, false>(dev_ctx,
+                                         param,
+                                         grad,
+                                         learning_rate,
+                                         moment1,
+                                         moment2,
+                                         beta1_pow,
+                                         beta2_pow,
+                                         master_param,
+                                         skip_update,
+                                         weight_decay,
+                                         beta1,
+                                         beta2,
+                                         epsilon,
+                                         always_adapt,
+                                         multi_precision,
+                                         param_out,
+                                         moment1_out,
+                                         moment2_out,
+                                         beta1_pow_out,
+                                         beta2_pow_out,
+                                         master_param_outs);
+  }
+}
+
+template <typename T, typename MT, typename Context, bool IsMultiPrecision>
+void ComputeRowImpl(const Context& dev_ctx,
+                    const DenseTensor& param,
+                    const SelectedRows& 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 =
+      dev_ctx.AllocateTmpTensor<MT, DeviceContext>(param.dims(), dev_ctx);*/
+  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;
+  // Update moments
+  bool should_update_beta_pow_later = false;
+  const MT *beta1_pow_ptr = nullptr, *beta2_pow_ptr = nullptr;
+  MT *beta1_pow_out_ptr = nullptr, *beta2_pow_out_ptr = nullptr;
+  VLOG(10) << "Beta1Pow place: " << beta1_pow.place()
+           << " , Beta2Pow place: " << beta2_pow.place();
+  // Diff from here
+  PADDLE_ENFORCE_EQ(IsMultiPrecision,
+                    false,
+                    common::errors::Unimplemented(
+                        "SelectedRows gradient is not supported when "
+                        "multi_precision=True."));
+  constexpr bool kIsSameType = std::is_same<T, MT>::value;
+  PADDLE_ENFORCE_EQ(kIsSameType,
+                    true,
+                    common::errors::Unimplemented(
+                        "SelectedRows gradient is not supported when "
+                        "multi_precision=True."));
+  if (grad.rows().size() == 0) {
+    VLOG(3) << "grad row size is 0!!";
+    return;
+  }
+
+  std::vector<int64_t> cpu_rows(grad.rows().begin(), grad.rows().end());
+  bool is_strict_sorted = true;
+  for (size_t i = 1; i < cpu_rows.size(); ++i) {
+    if (cpu_rows[i - 1] >= cpu_rows[i]) {
+      is_strict_sorted = false;
+      break;
+    }
+  }
+
+  SelectedRows tmp_grad_merge;
+  const SelectedRows* grad_merge_ptr;
+  if (is_strict_sorted) {
+    grad_merge_ptr = &grad;
+  } else {
+    // merge duplicated rows if any.
+    // The rows of grad_merge have been sorted inside MergeAdd functor
+    funcs::scatter::MergeAdd<Context, T> merge_func;
+    merge_func(dev_ctx, grad, &tmp_grad_merge, true);
+    grad_merge_ptr = &tmp_grad_merge;
+  }
+
+  auto& grad_merge = *grad_merge_ptr;
+  auto& grad_tensor = grad_merge.value();
+  const T* grad_data = grad_tensor.template data<T>();
+  auto* grad_merge_rows = &grad_merge.rows();
+  phi::MixVector<int64_t> mixv_grad_merge_rows(grad_merge_rows);
+  const int64_t* rows = mixv_grad_merge_rows.Data(dev_ctx.GetPlace());
+  auto row_numel = grad_tensor.numel() / grad_merge.rows().size();
+  if (dev_ctx.GetPlace().GetType() == AllocationType::GPU &&
+      beta1_pow.place() == phi::CPUPlace() &&
+      beta2_pow.place() == phi::CPUPlace()) {
+    SparseLambMomentREGUpdateFunctor<T> moment_update_functor(
+        static_cast<T>(weight_decay),
+        static_cast<T>(beta1),
+        static_cast<T>(beta2),
+        static_cast<T>(epsilon),
+        *beta1_pow.template data<T>(),
+        *beta2_pow.template data<T>(),
+        mom1.template data<T>(),
+        dev_ctx.template Alloc<T>(mom1_out),
+        mom2.template data<T>(),
+        dev_ctx.template Alloc<T>(mom2_out),
+        grad_data,
+        param.template data<T>(),
+        trust_ratio_div.template data<T>(),
+        rows,
+        row_numel,
+        grad_merge.rows().size(),
+        skip_update_flag);
+    for_range(moment_update_functor);
+    T* beta1_pow_out_data = dev_ctx.template HostAlloc<T>(beta1_pow_out);
+    beta1_pow_out_data[0] =
+        static_cast<T>(beta1) * beta1_pow.template data<T>()[0];
+    T* beta2_pow_out_data = dev_ctx.template HostAlloc<T>(beta2_pow_out);
+    beta2_pow_out_data[0] =
+        static_cast<T>(beta2) * beta2_pow.template data<T>()[0];
+  } else {
+    beta1_pow_ptr = beta1_pow.template data<MT>();
+    beta2_pow_ptr = beta2_pow.template data<MT>();
+    beta1_pow_out_ptr = dev_ctx.template Alloc<MT>(beta1_pow_out);
+    beta2_pow_out_ptr = dev_ctx.template Alloc<MT>(beta2_pow_out);
+    should_update_beta_pow_later = true;
+    SparseLambMomentMENUpdateFunctor<T> moment_update_functor(
+        static_cast<T>(weight_decay),
+        static_cast<T>(beta1),
+        static_cast<T>(beta2),
+        static_cast<T>(epsilon),
+        reinterpret_cast<const T*>(beta1_pow_ptr),
+        reinterpret_cast<const T*>(beta2_pow_ptr),
+        mom1.template data<T>(),
+        dev_ctx.template Alloc<T>(mom1_out),
+        mom2.template data<T>(),
+        dev_ctx.template Alloc<T>(mom2_out),
+        grad_data,
+        param.template data<T>(),
+        trust_ratio_div.template data<T>(),
+        rows,
+        row_numel,
+        grad_merge.rows().size(),
+        skip_update_flag);
+    for_range(moment_update_functor);
+  }
+  // Same from here
+  // Update parameter
+  // The code in the following part is exactly the same as that in
+  // paddle/phi/kernels/impl/lamb_kernel_impl.h Please modify it together
+
+  // DenseTensor p_norm_t;
+  // p_norm_t.Resize({1});
+  // auto* p_norm_ptr = dev_ctx.template Alloc<MT>(&p_norm_t);
+
+  // DenseTensor trust_ratio_div_norm_t;
+  // trust_ratio_div_norm_t.Resize({1});
+  // auto* trust_ratio_div_norm_ptr =
+  //     dev_ctx.template Alloc<MT>(&trust_ratio_div_norm_t);
+
+  DenseTensor p_norm_t;
+  DataType dtype = phi::CppTypeToDataType<MT>::Type();
+  FullKernel<MT, Context>(
+      dev_ctx, std::vector<int64_t>({1}), 0, dtype, &p_norm_t);
+  auto* p_norm_ptr = p_norm_t.data<MT>();
+
+  DenseTensor trust_ratio_div_norm_t;
+  FullKernel<MT, Context>(
+      dev_ctx, std::vector<int64_t>({1}), 0, dtype, &trust_ratio_div_norm_t);
+  auto* trust_ratio_div_norm_ptr = trust_ratio_div_norm_t.data<MT>();
+
+  // TODO(zengjinle): remove the following Eigen operations when
+  // *skip_update == true.
+  if (weight_decay > static_cast<MT>(0) || always_adapt) {
+    memory_utils::Buffer buffer(dev_ctx.GetPlace());
+    funcs::SquaredL2Norm(dev_ctx,
+                         reinterpret_cast<const MT*>(
+                             IsMultiPrecision ? master_param_ptr : param_ptr),
+                         p_norm_ptr,
+                         numel,
+                         &buffer);
+    funcs::SquaredL2Norm(
+        dev_ctx, trust_ratio_div_ptr, trust_ratio_div_norm_ptr, numel, &buffer);
+  }
+
+  if (VLOG_IS_ON(1)) {
+    const auto& name = "Param";
+    auto pn = funcs::ToVector(p_norm_ptr, 1, dev_ctx.GetPlace());
+    auto tn = funcs::ToVector(trust_ratio_div_norm_ptr, 1, dev_ctx.GetPlace());
+    auto dtype = DataTypeToString(CppTypeToDataType<T>::Type());
+    VLOG(1) << "Param " << dtype << " " << name << " pn = " << pn[0]
+            << " , tn = " << tn[0];
+  }
+
+#define CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC(__should_update_beta_pow)          \
+  do {                                                                        \
+    LambParamUpdateFunctor<T, MT, IsMultiPrecision, __should_update_beta_pow> \
+        param_update_functor(lr.template data<MT>(),                          \
+                             static_cast<const T*>(param_ptr),                \
+                             static_cast<const MT*>(master_param_ptr),        \
+                             p_norm_ptr,                                      \
+                             trust_ratio_div_ptr,                             \
+                             trust_ratio_div_norm_ptr,                        \
+                             static_cast<T*>(param_out_ptr),                  \
+                             static_cast<MT*>(master_param_out_ptr),          \
+                             skip_update_flag);                               \
+    if (__should_update_beta_pow) {                                           \
+      param_update_functor.SetBetaPows(beta1_pow_ptr,                         \
+                                       beta2_pow_ptr,                         \
+                                       beta1_pow_out_ptr,                     \
+                                       beta2_pow_out_ptr,                     \
+                                       beta1,                                 \
+                                       beta2);                                \
+    }                                                                         \
+    for_range(param_update_functor);                                          \
+  } while (0)
+
+  if (should_update_beta_pow_later) {
+    CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC(true);
+  } else {
+    CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC(false);
+  }
+
+#undef CALL_PADDLE_UPDATE_LAMB_PARAM_FUNC
+}
+
+}  // namespace sr
+}  // namespace phi

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

@@ -0,0 +1,47 @@
+// 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 <string>
+
+#include "paddle/phi/core/framework/selected_rows_serialize.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void LoadSelectedRowsKernel(const Context& dev_ctx,
+                            const std::string& file_path,
+                            int64_t seek,
+                            const std::vector<int64_t>& shape,
+                            bool load_as_fp16,
+                            SelectedRows* out) {
+  // FIXME(yuyang18): We save variable to local file now, but we should change
+  // it to save an output stream.
+  std::ifstream fin(file_path, std::ios::binary);
+  PADDLE_ENFORCE_EQ(
+      static_cast<bool>(fin),
+      true,
+      errors::Unavailable("Load operator fail to open file %s, please check "
+                          "whether the model file is complete or damaged.",
+                          file_path));
+  PADDLE_ENFORCE_NOT_NULL(
+      out,
+      errors::InvalidArgument("The variable to be loaded cannot be found."));
+
+  phi::DeserializeFromStream(fin, out, dev_ctx);
+}
+}  // namespace phi::sr

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

@@ -0,0 +1,63 @@
+// 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 <stdint.h>
+
+#include <fstream>
+#include <numeric>
+#include <string>
+#include <vector>
+
+#include "paddle/phi/core/framework/dense_tensor_serialize.h"
+#include "paddle/phi/core/framework/selected_rows_serialize.h"
+#include "paddle/phi/core/framework/var_type_helper.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/cast_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void SaveSelectedRowsKernel(const Context& dev_ctx,
+                            const SelectedRows& x,
+                            const std::string& file_path,
+                            bool overwrite,
+                            bool save_as_fp16) {
+  PADDLE_ENFORCE_EQ(
+      FileExists(file_path) && !overwrite,
+      false,
+      common::errors::PreconditionNotMet(
+          "%s exists!, cannot save to it when overwrite is set to false.",
+          file_path));
+  PADDLE_ENFORCE_EQ(save_as_fp16,
+                    false,
+                    common::errors::Unimplemented(
+                        "SelectedRows is not supported to save as float16."));
+
+  MkDirRecursively(DirName(file_path).c_str());
+
+  // FIXME(yuyang18): We save variable to local file now, but we should change
+  // it to save an output stream.
+  std::ofstream fout(file_path, std::ios::binary);
+  PADDLE_ENFORCE_EQ(static_cast<bool>(fout),
+                    true,
+                    common::errors::Unavailable(
+                        "Cannot open %s to save variables.", file_path));
+  phi::SerializeToStream(fout, x, dev_ctx);
+  fout.close();
+}
+
+}  // namespace phi

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

@@ -0,0 +1,29 @@
+// 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/kernel_registry.h"
+
+namespace phi {
+namespace sr {
+template <typename T, typename Context>
+void ShareDataKernel(const Context &dev_ctx,
+                     const SelectedRows &x,
+                     SelectedRows *out) {
+  out->set_rows(x.rows());
+  out->set_height(x.height());
+  out->mutable_value()->ShareDataWith(x.value());
+}
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/isfinite_kernel.cc

@@ -0,0 +1,84 @@
+// 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/selected_rows/isfinite_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#endif
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/selected_rows/impl/isfinite_kernel_impl.h"
+
+PD_REGISTER_KERNEL(isinf_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::IsinfSR,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(isnan_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::IsnanSR,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(isfinite_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::IsfiniteSR,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(isinf_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::IsinfSR,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(isnan_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::IsnanSR,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {}
+
+PD_REGISTER_KERNEL(isfinite_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::IsfiniteSR,
+                   float,
+                   double,
+                   phi::float16,
+                   int,
+                   int64_t) {}
+#endif

paddle/phi/kernels/selected_rows/isfinite_kernel.h

@@ -0,0 +1,32 @@
+/* Copyright (c) 2021 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/selected_rows.h"
+#include "paddle/phi/kernels/isfinite_kernel.h"
+
+namespace phi {
+
+#define DEFINE_ISFINITE_SR(isfinite_sr)   \
+  template <typename T, typename Context> \
+  void isfinite_sr(                       \
+      const Context& dev_ctx, const SelectedRows& x, SelectedRows* out);
+
+DEFINE_ISFINITE_SR(IsinfSR)
+DEFINE_ISFINITE_SR(IsnanSR)
+DEFINE_ISFINITE_SR(IsfiniteSR)
+#undef DEFINE_ISFINITE_SR
+
+}  // namespace phi

paddle/phi/kernels/selected_rows/lamb_kernel.h

@@ -0,0 +1,47 @@
+// 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.
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void LambKernel(const Context& dev_ctx,
+                const DenseTensor& param,
+                const SelectedRows& grad,
+                const DenseTensor& learning_rate,
+                const DenseTensor& moment1,
+                const DenseTensor& moment2,
+                const DenseTensor& beta1_pow,
+                const DenseTensor& beta2_pow,
+                const optional<DenseTensor>& master_param,
+                const optional<DenseTensor>& skip_update,
+                float weight_decay,
+                float beta1,
+                float beta2,
+                float epsilon,
+                bool always_adapt,
+                bool multi_precision,
+                DenseTensor* param_out,
+                DenseTensor* moment1_out,
+                DenseTensor* moment2_out,
+                DenseTensor* beta1_pow_out,
+                DenseTensor* beta2_pow_out,
+                DenseTensor* master_param_outs);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/merge_selected_rows_kernel.cc

@@ -0,0 +1,49 @@
+// 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/selected_rows/merge_selected_rows_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void MergeSelectedRowsKernel(const Context& dev_ctx,
+                             const SelectedRows& x,
+                             SelectedRows* out) {
+  funcs::scatter::MergeAdd<Context, T> merge_func;
+  merge_func(dev_ctx, x, out);
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(merge_selected_rows,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::MergeSelectedRowsKernel,
+                   float,
+                   double) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(merge_selected_rows,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::MergeSelectedRowsKernel,
+                   float,
+                   double) {}
+#endif

paddle/phi/kernels/selected_rows/merge_selected_rows_kernel.h

@@ -0,0 +1,28 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void MergeSelectedRowsKernel(const Context& dev_ctx,
+                             const SelectedRows& x,
+                             SelectedRows* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/scale_kernel.cc

@@ -0,0 +1,67 @@
+/* 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/selected_rows/scale_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/scale_kernel.h"
+namespace phi::sr {
+
+template <typename T, typename Context>
+void ScaleKernel(const Context& dev_ctx,
+                 const SelectedRows& x,
+                 const Scalar& scale,
+                 const Scalar& bias,
+                 bool bias_after_scale,
+                 SelectedRows* out) {
+  if (x.value().Holder() != out->value().Holder() ||
+      x.value().data() != out->value().data()) {
+    out->set_rows(x.rows());
+    out->set_height(x.height());
+  }
+  phi::ScaleKernel<T>(
+      dev_ctx, x.value(), scale, bias, bias_after_scale, out->mutable_value());
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(scale_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::ScaleKernel,
+                   float,
+                   double,
+                   phi::bfloat16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(scale_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::ScaleKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   uint8_t,
+                   int8_t,
+                   int16_t,
+                   int,
+                   int64_t) {}
+#endif

paddle/phi/kernels/selected_rows/scale_kernel.h

@@ -0,0 +1,32 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ScaleKernel(const Context& dev_ctx,
+                 const SelectedRows& x,
+                 const Scalar& scale,
+                 const Scalar& bias,
+                 bool bias_after_scale,
+                 SelectedRows* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/shape_kernel.cc

@@ -0,0 +1,191 @@
+/* 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/selected_rows/shape_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/shape_kernel.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void ShapeKernel(const Context& dev_ctx,
+                 const SelectedRows& input,
+                 DenseTensor* out) {
+  phi::ShapeKernel<T, Context>(dev_ctx, input.value(), out);
+}
+
+template <typename T, typename Context>
+void Shape64Kernel(const Context& dev_ctx,
+                   const SelectedRows& input,
+                   DenseTensor* out) {
+  phi::Shape64Kernel<T, Context>(dev_ctx, input.value(), out);
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(shape_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::ShapeKernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(shape_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::ShapeKernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
+}
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+PD_REGISTER_KERNEL(shape_sr,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::sr::ShapeKernel,
+                   bool,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::float16,
+                   phi::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
+}
+#endif
+
+#if defined(PADDLE_WITH_CUSTOM_DEVICE) && !defined(PADDLE_WITH_CUDA)
+PD_REGISTER_KERNEL(shape_sr,
+                   Custom,
+                   ALL_LAYOUT,
+                   phi::sr::ShapeKernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
+}
+#endif
+
+PD_REGISTER_KERNEL(shape64_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::Shape64Kernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(shape64_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::Shape64Kernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
+}
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+PD_REGISTER_KERNEL(shape64_sr,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::sr::Shape64Kernel,
+                   bool,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::float16,
+                   phi::bfloat16) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
+}
+#endif
+
+#if defined(PADDLE_WITH_CUSTOM_DEVICE) && !defined(PADDLE_WITH_CUDA)
+PD_REGISTER_KERNEL(shape64_sr,
+                   Custom,
+                   ALL_LAYOUT,
+                   phi::sr::Shape64Kernel,
+                   bool,
+                   int,
+                   int8_t,
+                   uint8_t,
+                   int64_t,
+                   float,
+                   double,
+                   phi::complex64,
+                   phi::complex128) {
+  kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
+  kernel->OutputAt(0).SetDataType(phi::DataType::INT64);
+}
+#endif

paddle/phi/kernels/selected_rows/shape_kernel.h

@@ -0,0 +1,33 @@
+/* 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. */
+
+#pragma once
+
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void ShapeKernel(const Context& dev_ctx,
+                 const SelectedRows& input,
+                 DenseTensor* out);
+
+template <typename T, typename Context>
+void Shape64Kernel(const Context& dev_ctx,
+                   const SelectedRows& input,
+                   DenseTensor* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/uniform_kernel.cc

@@ -0,0 +1,106 @@
+/* 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/selected_rows/uniform_kernel.h"
+#include "paddle/phi/kernels/legacy/uniform_kernel.h"
+
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/uniform_kernel.h"
+
+namespace phi::sr {
+
+template <typename T, typename Context>
+void UniformRawKernel(const Context& dev_ctx,
+                      const IntArray& shape,
+                      DataType dtype,
+                      const Scalar& min,
+                      const Scalar& max,
+                      int seed,
+                      int diag_num,
+                      int diag_step,
+                      float diag_val,
+                      SelectedRows* out) {
+  phi::UniformRawKernel<T>(dev_ctx,
+                           shape,
+                           dtype,
+                           min,
+                           max,
+                           seed,
+                           diag_num,
+                           diag_step,
+                           diag_val,
+                           out->mutable_value());
+}
+
+template <typename T, typename Context>
+void UniformKernel(const Context& dev_ctx,
+                   const IntArray& shape,
+                   DataType dtype,
+                   const Scalar& min,
+                   const Scalar& max,
+                   int seed,
+                   SelectedRows* out) {
+  phi::UniformKernel<T>(
+      dev_ctx, shape, dtype, min, max, seed, out->mutable_value());
+}
+
+}  // namespace phi::sr
+
+PD_REGISTER_KERNEL(uniform_raw_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::UniformRawKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}
+
+PD_REGISTER_KERNEL(uniform_sr,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::sr::UniformKernel,
+                   float,
+                   double,
+                   phi::bfloat16) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+PD_REGISTER_KERNEL(uniform_raw_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::UniformRawKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   phi::bfloat16) {}
+
+PD_REGISTER_KERNEL(uniform_sr,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::sr::UniformKernel,
+                   float,
+                   double,
+                   phi::float16,
+                   phi::bfloat16) {}
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+
+PD_REGISTER_KERNEL(
+    uniform_raw_sr, XPU, ALL_LAYOUT, phi::sr::UniformRawKernel, float) {}
+
+PD_REGISTER_KERNEL(uniform_sr, XPU, ALL_LAYOUT, phi::sr::UniformKernel, float) {
+}
+#endif

paddle/phi/kernels/selected_rows/uniform_kernel.h

@@ -0,0 +1,46 @@
+// 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.
+
+#pragma once
+
+#include "paddle/phi/common/int_array.h"
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/selected_rows.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void UniformRawKernel(const Context& dev_ctx,
+                      const IntArray& shape,
+                      DataType dtype,
+                      const Scalar& min,
+                      const Scalar& max,
+                      int seed,
+                      int diag_num,
+                      int diag_step,
+                      float diag_val,
+                      SelectedRows* out);
+
+template <typename T, typename Context>
+void UniformKernel(const Context& dev_ctx,
+                   const IntArray& shape,
+                   DataType dtype,
+                   const Scalar& min,
+                   const Scalar& max,
+                   int seed,
+                   SelectedRows* out);
+
+}  // namespace sr
+}  // namespace phi

paddle/phi/kernels/selected_rows/xpu/adam_kernel.cc

@@ -0,0 +1,362 @@
+// 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/selected_rows/adam_kernel.h"
+
+#include "glog/logging.h"
+
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/adam_functors.h"
+#include "paddle/phi/kernels/funcs/selected_rows_functor.h"
+
+namespace phi {
+namespace sr {
+
+template <typename T, typename Context>
+void AdamDenseParamSparseGradKernel(
+    const Context& dev_ctx,
+    const DenseTensor& param,
+    const SelectedRows& grad,
+    const DenseTensor& learning_rate,
+    const DenseTensor& moment1,
+    const DenseTensor& moment2,
+    const optional<DenseTensor>& moment2_max,  // UNUSED
+    const DenseTensor& beta1_pow,
+    const DenseTensor& beta2_pow,
+    const optional<DenseTensor>& master_param,
+    const optional<DenseTensor>& skip_update,
+    const Scalar& beta1,
+    const Scalar& beta2,
+    const Scalar& epsilon,
+    bool lazy_mode,
+    int64_t min_row_size_to_use_multithread,
+    bool multi_precision,
+    bool use_global_beta_pow,
+    bool amsgrad,  // UNUSED
+    DenseTensor* param_out,
+    DenseTensor* moment1_out,
+    DenseTensor* moment2_out,
+    DenseTensor* moment2_max_out,  // UNUSED
+    DenseTensor* beta1_pow_out,
+    DenseTensor* beta2_pow_out,
+    DenseTensor* master_param_outs) {
+  PADDLE_ENFORCE_NE(
+      amsgrad,
+      true,
+      common::errors::Unimplemented("Operation amsgrad is not supported yet."));
+
+  using XPUType = typename XPUTypeTrait<T>::Type;
+  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+  float* param_ptr = nullptr;
+  funcs::GetDataPointer<Context, float>(
+      param, &param_ptr, dev_ctx, &RAII_GUARD);
+
+  float* mom1_ptr = nullptr;
+  funcs::GetDataPointer<Context, float>(
+      moment1, &mom1_ptr, dev_ctx, &RAII_GUARD);
+
+  float* mom2_ptr = nullptr;
+  funcs::GetDataPointer<Context, float>(
+      moment2, &mom2_ptr, dev_ctx, &RAII_GUARD);
+
+  float* lr_ptr = nullptr;
+  funcs::GetDataPointer<Context, float>(
+      learning_rate, &lr_ptr, dev_ctx, &RAII_GUARD);
+
+  float* beta1_pow_ptr = nullptr;
+  const float* beta1_const_pow_ptr = nullptr;
+
+  if (beta1_pow.place() == CPUPlace()) {
+    if (beta1_pow.dtype() == DataType::FLOAT16) {
+      XPUType* beta1_pow_t =
+          RAII_GUARD.alloc_l3_or_gm<XPUType>(beta1_pow.numel());
+      memory_utils::Copy(param.place(),
+                         beta1_pow_t,
+                         beta1_pow.place(),
+                         beta1_pow.data<T>(),
+                         sizeof(T) * beta1_pow.numel());
+
+      int r = xpu::cast<XPUType, float>(
+          dev_ctx.x_context(), beta1_pow_t, beta1_pow_ptr, beta1_pow.numel());
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
+    } else {
+      beta1_pow_ptr = RAII_GUARD.alloc_l3_or_gm<float>(beta1_pow.numel());
+      memory_utils::Copy(param.place(),
+                         beta1_pow_ptr,
+                         beta1_pow.place(),
+                         beta1_pow.data<T>(),
+                         sizeof(T) * beta1_pow.numel());
+    }
+
+  } else {
+    if (beta1_pow.dtype() == DataType::FLOAT16)
+      funcs::GetDataPointer<Context, float>(
+          beta1_pow, &beta1_pow_ptr, dev_ctx, &RAII_GUARD);
+    else
+      beta1_const_pow_ptr = beta1_pow.template data<float>();
+  }
+
+  float* beta2_pow_ptr = nullptr;
+  const float* beta2_const_pow_ptr = nullptr;
+
+  if (beta2_pow.place() == CPUPlace()) {
+    if (beta2_pow.dtype() == DataType::FLOAT16) {
+      XPUType* beta2_pow_t =
+          RAII_GUARD.alloc_l3_or_gm<XPUType>(beta2_pow.numel());
+      memory_utils::Copy(param.place(),
+                         beta2_pow_t,
+                         beta2_pow.place(),
+                         beta2_pow.data<T>(),
+                         sizeof(T) * beta2_pow.numel());
+
+      int r = xpu::cast<XPUType, float>(
+          dev_ctx.x_context(), beta2_pow_t, beta2_pow_ptr, beta2_pow.numel());
+      PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
+    } else {
+      beta2_pow_ptr = RAII_GUARD.alloc_l3_or_gm<float>(beta2_pow.numel());
+      memory_utils::Copy(param.place(),
+                         beta2_pow_ptr,
+                         beta2_pow.place(),
+                         beta2_pow.data<T>(),
+                         sizeof(T) * beta2_pow.numel());
+    }
+  } else {
+    if (beta2_pow.dtype() == DataType::FLOAT16)
+      funcs::GetDataPointer<Context, float>(
+          beta2_pow, &beta2_pow_ptr, dev_ctx, &RAII_GUARD);
+    else
+      beta2_const_pow_ptr = beta2_pow.template data<float>();
+  }
+
+  DenseTensor xpu_param_out;
+  float* param_out_ptr = nullptr;
+  const DenseTensorMeta meta_param(DataType::FLOAT32, param_out->dims());
+  xpu_param_out.set_meta(meta_param);
+  funcs::GetOutDataPointer<Context, float>(
+      param_out, &xpu_param_out, &param_out_ptr, dev_ctx);
+
+  DenseTensor xpu_mom1_out;
+  float* mom1_out_ptr = nullptr;
+  const DenseTensorMeta meta_mom1(DataType::FLOAT32, moment1_out->dims());
+  xpu_mom1_out.set_meta(meta_mom1);
+  funcs::GetOutDataPointer<Context, float>(
+      moment1_out, &xpu_mom1_out, &mom1_out_ptr, dev_ctx);
+
+  DenseTensor xpu_mom2_out;
+  float* mom2_out_ptr = nullptr;
+  const DenseTensorMeta meta_mom2(DataType::FLOAT32, moment2_out->dims());
+  xpu_mom2_out.set_meta(meta_mom2);
+  funcs::GetOutDataPointer<Context, float>(
+      moment2_out, &xpu_mom2_out, &mom2_out_ptr, dev_ctx);
+
+  bool skip_update_ = false;
+  if (skip_update.is_initialized()) {
+    PADDLE_ENFORCE_EQ(
+        skip_update->numel(),
+        1,
+        errors::InvalidArgument("Input(SkipUpdate) size must be 1, but get %d",
+                                skip_update->numel()));
+    std::vector<bool> skip_update_vec;
+    TensorToVector(*skip_update, dev_ctx, &skip_update_vec);
+    skip_update_ = skip_update_vec[0];
+  }
+
+  if (skip_update_) {
+    VLOG(4) << "Adam skip update";
+    phi::Copy(dev_ctx, param, dev_ctx.GetPlace(), false, param_out);
+    phi::Copy(dev_ctx, moment1, dev_ctx.GetPlace(), false, moment1_out);
+    phi::Copy(dev_ctx, moment2, dev_ctx.GetPlace(), false, moment2_out);
+    if (!use_global_beta_pow) {
+      phi::Copy(dev_ctx, beta1_pow, beta1_pow.place(), false, beta1_pow_out);
+      phi::Copy(dev_ctx, beta2_pow, beta2_pow.place(), false, beta2_pow_out);
+    }
+    return;
+  }
+
+  PADDLE_ENFORCE_EQ(
+      beta1_pow_out->numel(),
+      1,
+      errors::InvalidArgument("Tensor holds the wrong size, Expected beta1 pow "
+                              "output size is 1, but received "
+                              "value is:%d.",
+                              beta1_pow_out->numel()));
+
+  PADDLE_ENFORCE_EQ(
+      beta2_pow_out->numel(),
+      1,
+      errors::InvalidArgument("Tensor holds the wrong size, Expected beta2 pow "
+                              "output size is 1, but received "
+                              "value is:%d.",
+                              beta2_pow_out->numel()));
+
+  VLOG(4) << "use_global_beta_pow:" << use_global_beta_pow;
+
+  auto beta1_ = beta1.to<float>();
+  auto beta2_ = beta2.to<float>();
+  auto epsilon_ = epsilon.to<float>();
+
+  float* grad_c = nullptr;
+  if (grad.rows().size() == 0) {
+    VLOG(3) << "grad row size is 0!!";
+    return;
+  }
+
+  std::vector<int64_t> cpu_rows(grad.rows().begin(), grad.rows().end());
+  bool is_strict_sorted = true;
+  for (size_t i = 1; i < cpu_rows.size(); ++i) {
+    if (cpu_rows[i - 1] >= cpu_rows[i]) {
+      is_strict_sorted = false;
+      break;
+    }
+  }
+
+  SelectedRows tmp_grad_merge;
+  const SelectedRows* grad_merge_ptr;
+  if (is_strict_sorted) {
+    grad_merge_ptr = &grad;
+  } else {
+    funcs::scatter::MergeAdd<Context, float> merge_func;
+    merge_func(dev_ctx, grad, &tmp_grad_merge, true);
+
+    xpu_wait(dev_ctx.x_context()->xpu_stream);
+    grad_merge_ptr = &tmp_grad_merge;
+  }
+
+  auto& grad_merge = *grad_merge_ptr;
+  auto& grad_tensor = grad_merge.value();
+
+  funcs::GetDataPointer<Context, float>(
+      grad_tensor, &grad_c, dev_ctx, &RAII_GUARD);
+
+  int row_count = grad_merge.rows().size();
+  std::vector<int> rows(row_count);
+  int* xpu_rows = RAII_GUARD.alloc_l3_or_gm<int>(row_count);
+  std::vector<int64_t> merge_rows(grad_merge.rows().begin(),
+                                  grad_merge.rows().end());
+  for (size_t i = 0; i < grad_merge.rows().size(); ++i) {
+    rows[i] = static_cast<int>(merge_rows[i]);
+  }
+  xpu_wait(dev_ctx.x_context()->xpu_stream);
+  memory_utils::Copy(dev_ctx.GetPlace(),
+                     xpu_rows,
+                     CPUPlace(),
+                     rows.data(),
+                     row_count * sizeof(int));
+  auto row_numel = grad_tensor.numel() / grad_merge.rows().size();
+  auto ori_rows = param.numel() / row_numel;
+
+  int r = xpu::sparse_adam(
+      dev_ctx.x_context(),
+      grad_c != nullptr ? grad_c : grad_tensor.template data<float>(),
+      mom1_ptr != nullptr ? mom1_ptr : moment1.template data<float>(),
+      mom2_ptr != nullptr ? mom2_ptr : moment2.template data<float>(),
+      param_ptr != nullptr ? param_ptr : param.template data<float>(),
+      beta1_pow_ptr != nullptr ? beta1_pow_ptr : beta1_const_pow_ptr,
+      beta2_pow_ptr != nullptr ? beta2_pow_ptr : beta2_const_pow_ptr,
+      lr_ptr != nullptr ? lr_ptr : learning_rate.template data<float>(),
+      mom1_out_ptr,
+      mom2_out_ptr,
+      param_out_ptr,
+      beta1_,
+      beta2_,
+      epsilon_,
+      ori_rows,
+      xpu_rows,
+      row_numel,
+      grad_merge.rows().size(),
+      lazy_mode);
+
+  PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
+
+  funcs::CopyOutData<Context, float>(
+      xpu_mom1_out, moment1_out, dev_ctx, &RAII_GUARD);
+  funcs::CopyOutData<Context, float>(
+      xpu_mom2_out, moment1_out, dev_ctx, &RAII_GUARD);
+  funcs::CopyOutData<Context, float>(
+      xpu_param_out, moment1_out, dev_ctx, &RAII_GUARD);
+
+  if (!use_global_beta_pow) {
+    // update in cpu and then copy to xpu
+    if (beta1_pow.place() == CPUPlace() && beta2_pow.place() == CPUPlace()) {
+      funcs::SetBetaData<Context, float>(
+          beta1_pow, beta1_pow_out, beta1_, dev_ctx);
+
+      funcs::SetBetaData<Context, float>(
+          beta2_pow, beta2_pow_out, beta2_, dev_ctx);
+    } else {
+      float* beta1_pow_out_p1 = nullptr;
+
+      if (beta1_pow_out->dtype() == DataType::FLOAT16) {
+        funcs::Scale<Context, float>(beta1_pow_out,
+                                     beta1_pow,
+                                     beta1_pow_ptr,
+                                     beta1_,
+                                     dev_ctx,
+                                     &RAII_GUARD);
+      } else {
+        const float* beta1_pow_data = beta1_pow.template data<float>();
+        beta1_pow_out_p1 = dev_ctx.template Alloc<float>(beta1_pow_out);
+        r = xpu::scale(dev_ctx.x_context(),
+                       beta1_pow_data,
+                       beta1_pow_out_p1,
+                       beta1_pow.numel(),
+                       false,
+                       beta1_,
+                       0.0f);
+        xpu_wait(dev_ctx.x_context()->xpu_stream);
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
+      }
+
+      float* beta2_pow_out_p1 = nullptr;
+      if (beta2_pow_out->dtype() == DataType::FLOAT16) {
+        funcs::Scale<Context, float>(beta2_pow_out,
+                                     beta2_pow,
+                                     beta2_pow_ptr,
+                                     beta2_,
+                                     dev_ctx,
+                                     &RAII_GUARD);
+      } else {
+        const float* beta2_pow_data = beta2_pow.template data<float>();
+        beta2_pow_out_p1 = dev_ctx.template Alloc<float>(beta2_pow_out);
+        r = xpu::scale(dev_ctx.x_context(),
+                       beta2_pow_data,
+                       beta2_pow_out_p1,
+                       beta2_pow.numel(),
+                       false,
+                       beta2_,
+                       0.0f);
+        xpu_wait(dev_ctx.x_context()->xpu_stream);
+        PADDLE_ENFORCE_XDNN_SUCCESS(r, "adam");
+      }
+    }
+  }
+}
+}  // namespace sr
+}  // namespace phi
+
+PD_REGISTER_KERNEL(adam_dense_param_sparse_grad,
+                   XPU,
+                   ALL_LAYOUT,
+                   phi::sr::AdamDenseParamSparseGradKernel,
+                   float,
+                   phi::float16) {
+  // Skip beta1_pow, beta2_pow, skip_update data transform
+  kernel->InputAt(6).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->InputAt(9).SetBackend(phi::Backend::ALL_BACKEND);
+  kernel->OutputAt(4).SetBackend(phi::Backend::UNDEFINED);
+  kernel->OutputAt(5).SetBackend(phi::Backend::UNDEFINED);
+}