chore: import upstream snapshot with attribution

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) {}