chore: import upstream snapshot with attribution

paddle/phi/kernels/funcs/math/CMakeLists.txt

@@ -0,0 +1,21 @@
+file(
+  GLOB func_cc_srcs
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "*.cc")
+list(FILTER func_cc_srcs EXCLUDE REGEX ".*_xpu\\.cc$")
+
+if(WITH_GPU OR WITH_ROCM)
+  file(
+    GLOB func_cu_srcs
+    RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+    "*.cu")
+endif()
+
+if(WITH_XPU)
+  file(
+    GLOB func_xpu_srcs
+    RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+    "*_xpu.cc")
+endif()
+
+collect_srcs(kernels_srcs SRCS ${func_cc_srcs} ${func_cu_srcs} ${func_xpu_srcs})

paddle/phi/kernels/funcs/math/beam_search.cc

@@ -0,0 +1,311 @@
+// 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/funcs/math/beam_search.h"
+#include "glog/logging.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+
+namespace phi {
+namespace math {
+
+template <typename T>
+class BeamSearchFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext &dev_ctx UNUSED,
+                  const DenseTensor *pre_ids,
+                  const DenseTensor *pre_scores,
+                  const DenseTensor *ids,
+                  const DenseTensor *scores,
+                  DenseTensor *selected_ids,
+                  DenseTensor *selected_scores,
+                  DenseTensor *parent_idx,
+                  size_t level,
+                  size_t beam_size,
+                  int end_id,
+                  bool is_accumulated) {
+    auto abs_lod = ToAbsOffset(scores->lod());
+    auto &high_level = abs_lod[level];
+
+    auto items = SelectTopBeamSizeItems(pre_ids,
+                                        pre_scores,
+                                        ids,
+                                        scores,
+                                        level,
+                                        beam_size,
+                                        end_id,
+                                        is_accumulated);
+    auto selected_items = ToMap(items, high_level.back());
+    if (FLAGS_v == 3) {
+      VLOG(3) << "selected_items:";
+      for (size_t i = 0; i < selected_items.size(); ++i) {
+        VLOG(3) << "offset: " << i;
+        for (auto &item : selected_items[i]) {
+          VLOG(3) << item.ToString();
+        }
+      }
+    }
+
+    PruneEndBeams(pre_ids, abs_lod, &selected_items, level, end_id);
+    // calculate the output tensor's height
+    size_t num_instances = std::accumulate(
+        std::begin(selected_items),
+        std::end(selected_items),
+        0,
+        [](size_t a, std::vector<Item> &b) { return a + b.size(); });
+    // the output tensor shape should be [num_instances, 1]
+    auto dims =
+        make_ddim(std::vector<int64_t>({static_cast<int>(num_instances), 1}));
+    selected_ids->Resize(dims);
+    auto *selected_ids_data = dev_ctx.template Alloc<int64_t>(selected_ids);
+    selected_scores->Resize(dims);
+    auto *selected_scores_data = dev_ctx.template Alloc<float>(selected_scores);
+    if (parent_idx != nullptr) {
+      parent_idx->Resize({static_cast<int64_t>(num_instances)});
+    }
+    auto *parent_idx_data =
+        parent_idx ? dev_ctx.template Alloc<int>(parent_idx) : nullptr;
+
+    // fill in data
+    std::vector<size_t> low_level;
+    size_t low_offset = 0;
+    for (auto &items : selected_items) {
+      low_level.push_back(low_offset);
+      for (auto &item : items) {
+        if (parent_idx) {
+          parent_idx_data[low_offset] = static_cast<int>(low_level.size() - 1);
+        }
+        selected_ids_data[low_offset] = item.id;
+        selected_scores_data[low_offset] = item.score;
+        low_offset++;
+      }
+    }
+    low_level.push_back(low_offset);
+
+    // fill lod
+    LegacyLoD lod(2);
+    lod[0].assign(high_level.begin(), high_level.end());
+    lod[1].assign(low_level.begin(), low_level.end());
+    if (!CheckLegacyLoD(lod)) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "lod %s is not right in"
+          " beam_search, please check your code.",
+          LoDToString(lod)));
+    }
+    selected_ids->set_lod(lod);
+    selected_scores->set_lod(lod);
+  }
+
+  /*
+   * The basic items help to sort.
+   */
+  struct Item {
+    Item() = default;
+    Item(size_t offset, size_t id, float score)
+        : offset(offset), id(id), score(score) {}
+    // offset in the higher lod level.
+    size_t offset;
+    // prefix id in the lower lod level.
+    // size_t prefix;
+    // the candidate id
+    size_t id;
+    // the corresponding score
+    float score;
+
+    inline bool operator<(const Item &in) const {
+      return (score < in.score) ||
+             ((score == in.score) && (offset < in.offset));
+    }
+
+    inline Item &operator=(const Item &in) {
+      if (this != &in) {
+        this->offset = in.offset;
+        this->id = in.id;
+        this->score = in.score;
+        return *this;
+      }
+      return *this;
+    }
+
+    std::string ToString() {
+      std::ostringstream os;
+      os << "{";
+      os << "offset: " << offset << ", ";
+      os << "id: " << id << ", ";
+      os << "score: " << score << "";
+      os << "}";
+      return os.str();
+    }
+  };
+
+ protected:
+  /*
+   * Prune the source sentences all branches finished, and it is optional.
+   * Pruning must one step later than finishing (thus pre_ids is needed here),
+   * since the end tokens must be written out.
+   */
+  void PruneEndBeams(const DenseTensor *pre_ids,
+                     const LegacyLoD &abs_lod,
+                     std::vector<std::vector<Item>> *items,
+                     size_t lod_level,
+                     int end_id) {
+    auto *pre_ids_data = pre_ids->data<int64_t>();
+    auto &high_level = abs_lod[lod_level];
+    for (size_t src_idx = 0; src_idx < high_level.size() - 1; ++src_idx) {
+      size_t src_prefix_start = high_level[src_idx];
+      size_t src_prefix_end = high_level[src_idx + 1];
+      bool finish_flag = true;
+      for (size_t offset = src_prefix_start; offset < src_prefix_end;
+           offset++) {
+        for (auto &item : items->at(offset)) {
+          if (item.id != static_cast<size_t>(end_id) ||
+              pre_ids_data[offset] != end_id) {
+            finish_flag = false;
+            break;
+          }
+        }
+        if (!finish_flag) break;
+      }
+      if (finish_flag) {  // all branches of the beam (source sentence) end and
+                          // prune this beam
+        for (size_t offset = src_prefix_start; offset < src_prefix_end;
+             offset++)
+          items->at(offset).clear();
+      }
+    }
+  }
+
+  /*
+   * Transform the items into a map whose key is offset, value is the items.
+   * NOTE low performance.
+   */
+  std::vector<std::vector<Item>> ToMap(
+      const std::vector<std::vector<Item>> &items, size_t element_num) {
+    std::vector<std::vector<Item>> result;
+    result.resize(element_num);
+    for (auto &entries : items) {
+      for (const auto &item : entries) {
+        result[item.offset].push_back(item);
+      }
+    }
+    return result;
+  }
+
+  void Insert(std::vector<Item> *top_beam_ptr,
+              const Item &item,
+              size_t beam_size) {
+    std::vector<Item> &top_beam = *top_beam_ptr;
+
+    size_t num_beams = top_beam.size();
+    if (num_beams < beam_size) {
+      top_beam.resize(num_beams + 1);
+      num_beams++;
+    } else {
+      if (item < top_beam[beam_size - 1]) {
+        return;
+      }
+    }
+
+    for (int k = static_cast<int>(num_beams) - 2; k >= 0; --k) {
+      if (top_beam[k] < item) {
+        top_beam[k + 1] = top_beam[k];
+      } else {
+        top_beam[k + 1] = item;
+        return;
+      }
+    }
+    top_beam[0] = item;
+  }
+
+  /*
+   * For each source, select top beam_size records.
+   */
+  std::vector<std::vector<Item>> SelectTopBeamSizeItems(
+      const DenseTensor *pre_ids,
+      const DenseTensor *pre_scores,
+      const DenseTensor *ids,
+      const DenseTensor *scores,
+      size_t lod_level,
+      size_t beam_size,
+      int end_id,
+      bool is_accumulated) {
+    std::vector<std::vector<Item>> result;
+
+    // find the current candidates
+    auto abs_lod = ToAbsOffset(scores->lod());
+
+    auto *pre_ids_data = pre_ids->data<int64_t>();
+    auto *pre_scores_data = pre_scores->data<float>();
+
+    auto *ids_data = ids ? ids->data<int64_t>() : nullptr;
+    auto *scores_data = scores->data<float>();
+
+    size_t num_seqs = scores->NumElements(lod_level);
+    size_t seq_width = 1;
+    for (int i = 1; i < scores->dims().size(); i++) {
+      seq_width *= scores->dims()[i];
+    }
+
+    for (size_t seq_id = 0; seq_id < num_seqs; ++seq_id) {
+      size_t seq_offset_start = abs_lod[lod_level][seq_id];
+      size_t seq_offset_end = abs_lod[lod_level][seq_id + 1];
+
+      std::vector<Item> top_beam;
+      top_beam.reserve(beam_size);
+
+      for (size_t offset = seq_offset_start; offset < seq_offset_end;
+           ++offset) {
+        auto pre_id = pre_ids_data[offset];
+        auto pre_score = pre_scores_data[offset];
+        if (pre_id == end_id) {
+          // Allocate all probability mass to end_id for finished branches and
+          // the other candidate ids can be ignored.
+          Item item(offset, end_id, pre_score);
+          Insert(&top_beam, item, beam_size);
+        } else {
+          size_t index = offset * seq_width;
+          for (size_t d = 0; d < seq_width; d++, index++) {
+            int64_t id = ids_data ? ids_data[index] : static_cast<int64_t>(d);
+            float score = is_accumulated
+                              ? scores_data[index]
+                              : pre_score + std::log(scores_data[index]);
+            Item item(offset, id, score);
+            Insert(&top_beam, item, beam_size);
+          }
+        }
+      }
+
+      result.emplace_back(top_beam);
+    }
+
+    if (FLAGS_v == 3) {
+      VLOG(3) << "SelectTopBeamSizeItems result size " << result.size();
+      for (auto &items : result) {
+        VLOG(3) << "item set:";
+        for (auto &item : items) {
+          VLOG(3) << item.ToString();
+        }
+      }
+    }
+
+    return result;
+  }
+};
+
+template class PADDLE_API BeamSearchFunctor<CPUContext, int>;
+template class PADDLE_API BeamSearchFunctor<CPUContext, int64_t>;
+template class PADDLE_API BeamSearchFunctor<CPUContext, float>;
+template class PADDLE_API BeamSearchFunctor<CPUContext, double>;
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/beam_search.cu

@@ -0,0 +1,541 @@
+// 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/funcs/math/beam_search.h"
+#include "paddle/phi/backends/gpu/gpu_device_function.h"
+#include "paddle/phi/backends/gpu/gpu_launch_config.h"
+
+namespace phi {
+namespace math {
+
+struct Triple {
+  __device__ __forceinline__ Triple() = default;
+  __device__ __forceinline__ Triple(int o, int i, float s)
+      : offset(o), id(i), score(s) {}
+
+  __device__ __forceinline__ void set(int o, int i, float s) {
+    offset = o;
+    id = i;
+    score = s;
+  }
+
+  __device__ __forceinline__ void operator=(const Triple& in) {
+    offset = in.offset;
+    id = in.id;
+    score = in.score;
+  }
+
+  __device__ __forceinline__ bool operator<(const float s) const {
+    return score < s;
+  }
+
+  __device__ __forceinline__ bool operator<(const Triple& in) const {
+    return (score < in.score) || ((score == in.score) && (offset < in.offset));
+  }
+
+  int offset;
+  int id;
+  float score;
+};
+
+__device__ __forceinline__ void Insert(Triple* top_beam,
+                                       const Triple& p,
+                                       int beam_size) {
+  if (p < top_beam[beam_size - 1]) {
+    return;
+  }
+  for (int k = beam_size - 2; k >= 0; --k) {
+    if (top_beam[k] < p) {
+      top_beam[k + 1] = top_beam[k];
+    } else {
+      top_beam[k + 1] = p;
+      return;
+    }
+  }
+  top_beam[0] = p;
+}
+
+template <int MaxThreadsPerSeq, bool IsAccumulated = true>
+__device__ __forceinline__ int SelectTopBeam(Triple* top_beam,
+                                             const int64_t* pre_ids,
+                                             const float* pre_scores,
+                                             const int64_t* ids,
+                                             const float* scores,
+                                             const int seq_offset_start,
+                                             const int seq_offset_end,
+                                             const int seq_width,
+                                             int beam_size,
+                                             int end_id,
+                                             int used_threads) {
+  // top_beam is shared memory
+  const int tid = threadIdx.x;
+  const int tid_of_seq = threadIdx.x % MaxThreadsPerSeq;
+
+  int num_used_threads = used_threads;
+
+  Triple* top_beam_local = top_beam + tid * beam_size;
+  if (tid_of_seq < num_used_threads) {
+    for (int i = 0; i < beam_size; ++i) {
+      top_beam_local[i].set(-1, -1, -INFINITY);
+    }
+
+    for (int offset = seq_offset_start; offset < seq_offset_end; ++offset) {
+      int pre_id = static_cast<int>(pre_ids[offset]);
+      if (pre_id == end_id) {
+        if (tid_of_seq == 0) {
+          Triple tmp(offset, end_id, pre_scores[offset]);
+          Insert(top_beam_local, tmp, beam_size);
+        }
+      } else {
+        int64_t index = static_cast<int64_t>(offset) * seq_width + tid_of_seq;
+        if (!IsAccumulated) {
+          float pre_score = pre_scores[offset];
+          for (int i = tid_of_seq; i < seq_width; i += num_used_threads) {
+            float score = pre_score + __logf(scores[index]);
+            int id = ids ? static_cast<int>(ids[index]) : i;
+            Triple tmp(offset, id, score);
+            Insert(top_beam_local, tmp, beam_size);
+            index += num_used_threads;
+          }
+        } else {
+          for (int i = tid_of_seq; i < seq_width; i += num_used_threads) {
+            int id = ids ? static_cast<int>(ids[index]) : i;
+            float score = scores[index];
+            Triple tmp(offset, id, score);
+            Insert(top_beam_local, tmp, beam_size);
+            index += num_used_threads;
+          }
+        }
+      }
+    }
+  }
+
+  while (num_used_threads > 1) {
+    if (num_used_threads > 16) {
+      __syncthreads();
+    }
+
+    if ((num_used_threads & 0x1) != 0) {
+      // If num_used_threads is a odd number, merge local top_beam of thread 0
+      // and num_used_threads - 1
+      if (tid_of_seq == 0) {
+        int index_in_sh = (num_used_threads - 1 + tid) * beam_size;
+        for (int i = 0; i < beam_size; i++) {
+          Insert(top_beam_local, top_beam[index_in_sh], beam_size);
+          index_in_sh++;
+        }
+      }
+    }
+
+    num_used_threads = num_used_threads >> 1;
+    if (tid_of_seq < num_used_threads) {
+      int index_in_sh = (num_used_threads + tid) * beam_size;
+      for (int i = 0; i < beam_size; i++) {
+        Insert(top_beam_local, top_beam[index_in_sh], beam_size);
+        index_in_sh++;
+      }
+    }
+  }
+
+  if (tid_of_seq == 0) {
+    int num_items = 0;
+    for (int i = 0; i < beam_size; ++i) {
+      num_items =
+          (top_beam_local[i].score > -INFINITY) ? num_items + 1 : num_items;
+    }
+    return num_items;
+  }
+
+  return 0;
+}
+
+__device__ __forceinline__ bool PruneEndBeams(Triple* top_beam_local,
+                                              const int64_t* pre_ids,
+                                              const int end_id,
+                                              int num_items) {
+  bool finish_flag = true;
+  for (int i = 0; i < num_items; ++i) {
+    int offset = top_beam_local[i].offset;
+    if (top_beam_local[i].id != end_id ||
+        static_cast<int>(pre_ids[offset]) != end_id) {
+      finish_flag = false;
+      break;
+    }
+  }
+  return finish_flag;
+}
+
+template <bool ReturnParentIdx = false>
+__device__ __forceinline__ void WriteBack(int64_t* selected_ids,
+                                          float* selected_scores,
+                                          int* parent_idx,
+                                          size_t* selected_offsets,
+                                          Triple* top_beam_local,
+                                          const int seq_offset_start,
+                                          const int seq_offset_end,
+                                          const int selected_seq_start,
+                                          const int selected_seq_length) {
+  const int tid = threadIdx.x;  // use 1 thread only for each sequence
+  int global_index = selected_seq_start;
+  for (int global_offset = seq_offset_start; global_offset < seq_offset_end;
+       ++global_offset) {
+    for (int local_index = 0; local_index < selected_seq_length;
+         ++local_index) {
+      if (top_beam_local[local_index].offset == global_offset) {
+        selected_ids[global_index] =
+            static_cast<int64_t>(top_beam_local[local_index].id);
+        selected_scores[global_index] = top_beam_local[local_index].score;
+        if (ReturnParentIdx) {
+          parent_idx[global_index] = static_cast<int>(global_offset);
+        }
+        global_index++;
+      }
+    }
+    selected_offsets[global_offset + 1] = static_cast<size_t>(global_index);
+  }
+}
+
+template <int MaxLength, int MaxThreadsPerSeq, int MaxSeqs>
+__device__ void BeamSearchDetails(int64_t* selected_ids,
+                                  float* selected_scores,
+                                  int* parent_idx,
+                                  size_t* selected_offsets,
+                                  const int64_t* pre_ids,
+                                  const float* pre_scores,
+                                  const int64_t* ids,
+                                  const float* scores,
+                                  const int seq_offset_start,
+                                  const int seq_offset_end,
+                                  const int seq_width,
+                                  int beam_size,
+                                  int end_id,
+                                  bool is_accumulated,
+                                  int num_used_threads) {
+  __shared__ Triple top_beam[MaxLength];
+
+  int num_items = 0;
+  if (is_accumulated) {
+    num_items = SelectTopBeam<MaxThreadsPerSeq, true>(top_beam,
+                                                      pre_ids,
+                                                      pre_scores,
+                                                      ids,
+                                                      scores,
+                                                      seq_offset_start,
+                                                      seq_offset_end,
+                                                      seq_width,
+                                                      beam_size,
+                                                      end_id,
+                                                      num_used_threads);
+  } else {
+    num_items = SelectTopBeam<MaxThreadsPerSeq, false>(top_beam,
+                                                       pre_ids,
+                                                       pre_scores,
+                                                       ids,
+                                                       scores,
+                                                       seq_offset_start,
+                                                       seq_offset_end,
+                                                       seq_width,
+                                                       beam_size,
+                                                       end_id,
+                                                       num_used_threads);
+  }
+
+  const int tid = threadIdx.x;  // use 1 thread only for each sequence
+  const int tid_of_seq = tid % MaxThreadsPerSeq;
+  if (tid_of_seq == 0) {
+    // Use 1 thread for each sequence.
+    Triple* top_beam_local = top_beam + tid * beam_size;
+    bool finish_flag =
+        PruneEndBeams(top_beam_local, pre_ids, end_id, num_items);
+
+    int selected_seq_start = 0;
+    int selected_seq_length = finish_flag ? 0 : num_items;
+
+    if (MaxSeqs > 1) {
+      const int seq_id = (MaxSeqs > 1) ? tid / MaxThreadsPerSeq : tid;
+      __shared__ int shared_mem[MaxSeqs];
+
+      // [0, MaxSeqs - 1], length of each sequences
+      shared_mem[seq_id] = selected_seq_length;
+      __syncthreads();
+
+      for (int s = 0; s < seq_id; ++s) {
+        selected_seq_start += shared_mem[s];
+      }
+
+      if (seq_id == 0) {
+        selected_offsets[0] = 0;
+      }
+    } else {
+      selected_offsets[0] = 0;
+    }
+
+    if (parent_idx) {
+      WriteBack<true>(selected_ids,
+                      selected_scores,
+                      parent_idx,
+                      selected_offsets,
+                      top_beam_local,
+                      seq_offset_start,
+                      seq_offset_end,
+                      selected_seq_start,
+                      selected_seq_length);
+    } else {
+      WriteBack<false>(selected_ids,
+                       selected_scores,
+                       parent_idx,
+                       selected_offsets,
+                       top_beam_local,
+                       seq_offset_start,
+                       seq_offset_end,
+                       selected_seq_start,
+                       selected_seq_length);
+    }
+  }
+}
+
+template <int MaxLength, int MaxThreadsPerSeq, int MaxSeqs>
+__global__ void BeamSearchKernel(int64_t* selected_ids,
+                                 float* selected_scores,
+                                 int* parent_idx,
+                                 size_t* selected_offsets,
+                                 const int64_t* pre_ids,
+                                 const float* pre_scores,
+                                 const int64_t* ids,
+                                 const float* scores,
+                                 const size_t* seq_offsets,
+                                 const int num_seqs,
+                                 const int seq_width,
+                                 int beam_size,
+                                 int end_id,
+                                 bool is_accumulated,
+                                 int num_used_threads) {
+  const int tid = threadIdx.x;
+  const int seq_id = (MaxSeqs > 1) ? tid / MaxThreadsPerSeq : tid;
+
+  int seq_offset_start = static_cast<int>(seq_offsets[seq_id]);
+  int seq_offset_end = static_cast<int>(seq_offsets[seq_id + 1]);
+
+  BeamSearchDetails<MaxLength, MaxThreadsPerSeq, MaxSeqs>(selected_ids,
+                                                          selected_scores,
+                                                          parent_idx,
+                                                          selected_offsets,
+                                                          pre_ids,
+                                                          pre_scores,
+                                                          ids,
+                                                          scores,
+                                                          seq_offset_start,
+                                                          seq_offset_end,
+                                                          seq_width,
+                                                          beam_size,
+                                                          end_id,
+                                                          is_accumulated,
+                                                          num_used_threads);
+}
+
+template <int MaxLength, int MaxThreadsPerSeq>
+__global__ void BeamSearchKernelSingle(int64_t* selected_ids,
+                                       float* selected_scores,
+                                       int* parent_idx,
+                                       size_t* selected_offsets,
+                                       const int64_t* pre_ids,
+                                       const float* pre_scores,
+                                       const int64_t* ids,
+                                       const float* scores,
+                                       const int seq_length,
+                                       const int seq_width,
+                                       int beam_size,
+                                       int end_id,
+                                       bool is_accumulated,
+                                       int num_used_threads) {
+  const int seq_offset_start = 0;
+  const int seq_offset_end = seq_length;
+
+  BeamSearchDetails<MaxLength, MaxThreadsPerSeq, 1>(selected_ids,
+                                                    selected_scores,
+                                                    parent_idx,
+                                                    selected_offsets,
+                                                    pre_ids,
+                                                    pre_scores,
+                                                    ids,
+                                                    scores,
+                                                    seq_offset_start,
+                                                    seq_offset_end,
+                                                    seq_width,
+                                                    beam_size,
+                                                    end_id,
+                                                    is_accumulated,
+                                                    num_used_threads);
+}
+
+static inline int GetNumUsedThreads(const int max_threads_per_seq,
+                                    const int seq_width,
+                                    int beam_size) {
+  int num_used_threads = (seq_width + beam_size - 1) / beam_size;
+  num_used_threads = max_threads_per_seq < num_used_threads
+                         ? max_threads_per_seq
+                         : num_used_threads;
+
+  num_used_threads =
+      num_used_threads > 32
+          ? (num_used_threads >> 5) << 5
+          : (num_used_threads > 16
+                 ? 32
+                 : (num_used_threads > 8
+                        ? 16
+                        : (num_used_threads > 4
+                               ? 8
+                               : (num_used_threads > 2 ? 4
+                                                       : num_used_threads))));
+  return num_used_threads;
+}
+
+template <typename T>
+class BeamSearchFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& dev_ctx,
+                  const DenseTensor* pre_ids,
+                  const DenseTensor* pre_scores,
+                  const DenseTensor* ids,
+                  const DenseTensor* scores,
+                  DenseTensor* selected_ids,
+                  DenseTensor* selected_scores,
+                  DenseTensor* parent_idx,
+                  size_t level,
+                  size_t beam_size,
+                  int end_id,
+                  bool is_accumulated) {
+    auto abs_lod = ToAbsOffset(scores->lod());
+
+    const int64_t* pre_ids_data = pre_ids->data<int64_t>();
+    const float* pre_scores_data = pre_scores->data<float>();
+    const int64_t* ids_data = ids ? ids->data<int64_t>() : nullptr;
+    const float* scores_data = scores->data<float>();
+
+    const size_t num_seqs = abs_lod[level].size() - 1;
+    size_t seq_width = 1;
+    for (int i = 1; i < scores->dims().size(); i++) {
+      seq_width *= scores->dims()[i];
+    }
+
+    // Reserve a big enough memory.
+    auto selected_dims =
+        make_ddim({static_cast<int64_t>(num_seqs * beam_size), 1});
+    selected_ids->Resize(selected_dims);
+    int64_t* selected_ids_data = dev_ctx.template Alloc<int64_t>(selected_ids);
+    selected_scores->Resize(selected_dims);
+    float* selected_scores_data =
+        dev_ctx.template Alloc<float>(selected_scores);
+    if (parent_idx != nullptr) {
+      parent_idx->Resize({static_cast<int64_t>(num_seqs * beam_size)});
+    }
+    int* parent_idx_data =
+        parent_idx ? dev_ctx.template Alloc<int>(parent_idx) : nullptr;
+
+    LegacyLoD selected_lod(2);
+    selected_lod[0].assign(abs_lod[level].begin(), abs_lod[level].end());
+    selected_lod[1].resize(scores->dims()[0] + 1);
+    phi::MixVector<size_t> mix_vector(&selected_lod[1]);
+    phi::MixVector<size_t> mixv_abs(&abs_lod[level]);
+    size_t* selected_offsets = mix_vector.CUDAMutableData(dev_ctx.GetPlace());
+
+    if (num_seqs == 1) {
+      const int seq_length = static_cast<int>(abs_lod[level][1]);
+      const int kMaxThreadsPerSeq = 1024;
+      int num_used_threads = GetNumUsedThreads(kMaxThreadsPerSeq,
+                                               static_cast<int>(seq_width),
+                                               static_cast<int>(beam_size));
+      switch (phi::backends::gpu::RoundToPowerOfTwo(beam_size * seq_width)) {
+        CUDA_LAUNCH_KERNEL_HELPER(
+            BeamSearchKernelSingle<kPowerOfTwoDim, kMaxThreadsPerSeq>
+            <<<1, kMaxThreadsPerSeq, 0, dev_ctx.stream()>>>(
+                selected_ids_data,
+                selected_scores_data,
+                parent_idx_data,
+                selected_offsets,
+                pre_ids_data,
+                pre_scores_data,
+                ids_data,
+                scores_data,
+                seq_length,
+                static_cast<int>(seq_width),
+                static_cast<int>(beam_size),
+                static_cast<int>(end_id),
+                is_accumulated,
+                num_used_threads));
+      }
+    } else if (num_seqs <= 4) {
+      const size_t* seq_offsets = mixv_abs.CUDAData(dev_ctx.GetPlace());
+      // Use only 1 block
+      const int kMaxThreadsPerSeq = 32;
+      const int kMaxSeqs = 4;
+      int num_used_threads = GetNumUsedThreads(kMaxThreadsPerSeq,
+                                               static_cast<int>(seq_width),
+                                               static_cast<int>(beam_size));
+      switch (
+          phi::backends::gpu::RoundToPowerOfTwo(beam_size * num_seqs * 32)) {
+        CUDA_LAUNCH_KERNEL_HELPER(
+            BeamSearchKernel<kPowerOfTwoDim, kMaxThreadsPerSeq, kMaxSeqs>
+            <<<1, num_seqs * kMaxThreadsPerSeq, 0, dev_ctx.stream()>>>(
+                selected_ids_data,
+                selected_scores_data,
+                parent_idx_data,
+                selected_offsets,
+                pre_ids_data,
+                pre_scores_data,
+                ids_data,
+                scores_data,
+                seq_offsets,
+                static_cast<int>(num_seqs),
+                static_cast<int>(seq_width),
+                static_cast<int>(beam_size),
+                end_id,
+                is_accumulated,
+                num_used_threads));
+      }
+    } else {
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Not implemented other number of sequences yet."));
+    }
+
+    dev_ctx.Wait();
+    mix_vector.CopyToCPU();
+    if (!CheckLegacyLoD(selected_lod)) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "lod %s is not right in"
+          " beam_search, please check your code.",
+          LoDToString(selected_lod)));
+    }
+
+    selected_ids->set_lod(selected_lod);
+    selected_scores->set_lod(selected_lod);
+    if (selected_lod[1].back() < num_seqs * beam_size) {
+      auto final_selected_dims =
+          make_ddim({static_cast<int64_t>(selected_lod[1].back()), 1});
+      selected_ids->Resize(final_selected_dims);
+      selected_scores->Resize(final_selected_dims);
+      if (parent_idx) {
+        parent_idx->Resize({static_cast<int64_t>(selected_lod[1].back())});
+      }
+    }
+  }
+};
+
+template class BeamSearchFunctor<GPUContext, int>;
+template class BeamSearchFunctor<GPUContext, int64_t>;
+template class PADDLE_API BeamSearchFunctor<GPUContext, float>;
+template class BeamSearchFunctor<GPUContext, double>;
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/beam_search.h

@@ -0,0 +1,163 @@
+// 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 <vector>
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/lod_utils.h"
+#include "paddle/phi/core/mixed_vector.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+namespace math {
+
+static inline std::string LoDToString(const LegacyLoD& lod) {
+  std::ostringstream stream;
+  for (const auto& row : lod) {
+    for (const auto& element : row) {
+      stream << element << " ";
+    }
+    stream << "\n";
+  }
+  return stream.str();
+}
+
+static inline bool CheckLegacyLoD(const LoD& in, int tensor_height = -1) {
+  if (in.empty()) return true;
+  for (const auto& level : in) {
+    // check: there should be more than 2 offsets existing in each level.
+    if (level.size() < 2) return false;
+    // check: the first offset(the begin offset) of each level should be 0.
+    if (level.front() != 0) return false;
+    // check: all the offsets in a level should be non-descending
+    if (!std::is_sorted(level.begin(), level.end())) {
+      return false;
+    }
+  }
+  // check: the lowest level's last offset should equals `tensor_height` if
+  //        tensor_height>0.
+  if (tensor_height > 0 &&
+      static_cast<size_t>(tensor_height) != in.back().back())
+    return false;
+
+  // check: the higher level's last offset should equals the lower level's
+  // size-1.
+  // NOTE LoD store the levels from top to bottom, so the higher level goes
+  // first.
+  for (size_t level = 0; level < in.size() - 1; level++) {
+    if (in[level].back() != in[level + 1].size() - 1) return false;
+  }
+  return true;
+}
+
+/*
+ * This is an implementation of beam search.
+ *
+ * To explain the details, lets take machine translation task for example, in
+ * this task, one source sentence is translated to multiple target sentences,
+ * during this period, one sentence will be translated to multiple translation
+ * prefixes(target sentence that have not ended), in each time step a prefix
+ * will have some candidates, input the candidate ids and their corresponding
+ * scores (probabilities), it will sort and select the top beam_size candidates
+ * for each source sentence, and store the selected candidates's score and their
+ * corresponding ids to DenseTensors.
+ *
+ * A detailed example:
+ *
+ *  Input
+ *
+ *    ids:
+ *      - LoD (should have 2 levels)
+ *        - first level: [0, 1, 4]
+ *        - second level: [0, 1, 2, 3, 4]
+ *      - tensor's data:
+ *          [[4, 2, 5]
+ *           [2, 1, 3]
+ *           [3, 5, 2]
+ *           [8, 2, 1]]
+ *
+ *    scores:
+ *      - LoD same as `ids`
+ *      - tensor's data
+ *          [[0.5, 0.3, 0.2]
+ *           [0.6, 0.3, 0.1]
+ *           [0.9, 0.5, 0.1]
+ *           [0.7, 0.5, 0.1]]
+ *
+ * The inputs means that there are 2 source sentences to translate, and the
+ * first source has 1 prefix, the second source has 2 prefix.
+ *
+ * Lets assume beam size is 2, and the beam search's output should be
+ *      - LoD
+ *        - first level: [0, 1, 2]
+ *        - second level: [0, 2, 4]
+ *      - id tensor's data
+ *          [[4,
+ *            1,
+ *            3,
+ *            8]]
+ *      - score tensor's data
+ *          [[0.5,
+ *            0.3,
+ *            0.9,
+ *            0.7]]
+ *
+ * TODO all the prune operations should be in the beam search, so it is better
+ * to split the beam search algorithm into a sequence of smaller operators, and
+ * the prune operators can be inserted in this sequence.
+ */
+template <typename DeviceContext, typename T>
+class BeamSearchFunctor {
+ public:
+  /*
+   * The main function of beam search.
+   *
+   * @selected_ids: a [None, 1]-shaped tensor with LoD.
+   *   In a machine translation model, it might be the candidate term id sets,
+   *   each set stored as a variance-length sequence.
+   *   The format might be described with a two-level LoD
+   *   - [[0 1],
+   *      [0 1 2]]
+   *   - [[]
+   *      [0 1]]
+   *   the first level of LoD tells that there are two source sentences. The
+   *   second level describes the details of the candidate id set's offsets in
+   * the source sentences.
+   *
+   *  @selected_scores: a LoD tensor with the same shape and LoD with
+   * selected_ids.
+   *   It stores the corresponding scores of candidate ids in selected_ids.
+   *
+   * Return false if all the input tensor is empty, in machine translation task
+   * that means no candidates is provided, and the task will stop running.
+   */
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor* pre_ids,
+                  const DenseTensor* pre_scores,
+                  const DenseTensor* ids,
+                  const DenseTensor* scores,
+                  DenseTensor* selected_ids,
+                  DenseTensor* selected_scores,
+                  DenseTensor* parent_idx,
+                  size_t level,
+                  size_t beam_size,
+                  int end_id,
+                  bool is_accumulated);
+};
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/beam_search_xpu.cc

@@ -0,0 +1,364 @@
+// 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 "glog/logging.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/kernels/funcs/math/beam_search.h"
+#ifdef PADDLE_WITH_XPU
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/backends/xpu/xpu_header.h"
+#endif
+
+namespace phi {
+namespace math {
+template <typename T>
+int CopyData(const T *x, T **y, int len, const Place &place) {
+  if (nullptr == x || nullptr == y || len <= 0)
+    return xpu::Error_t::INVALID_PARAM;
+
+  *y = reinterpret_cast<T *>(malloc(sizeof(T) * len));
+
+  phi::memory_utils::Copy(CPUPlace(), *y, place, x, len * sizeof(T));
+  return xpu::Error_t::SUCCESS;
+}
+
+template <typename T>
+void CopyDataByCondition(const T *x, T **y, int len, const Place &place) {
+  if (x != nullptr) {
+    int r = CopyData(x, y, len, place);
+    PADDLE_ENFORCE_EQ(
+        r,
+        xpu::Error_t::SUCCESS,
+        common::errors::External("Copy data form xpu to cpu failed"));
+  }
+}
+
+template <typename T>
+class BeamSearchFunctor<XPUContext, T> {
+ public:
+  void operator()(const XPUContext &dev_ctx,
+                  const DenseTensor *pre_ids,
+                  const DenseTensor *pre_scores,
+                  const DenseTensor *ids,
+                  const DenseTensor *scores,
+                  DenseTensor *selected_ids,
+                  DenseTensor *selected_scores,
+                  DenseTensor *parent_idx,
+                  size_t level,
+                  size_t beam_size,
+                  int end_id,
+                  bool is_accumulated) {
+    auto abs_lod = ToAbsOffset(scores->lod());
+    auto &high_level = abs_lod[level];
+
+    auto items = SelectTopBeamSizeItems(pre_ids,
+                                        pre_scores,
+                                        ids,
+                                        scores,
+                                        level,
+                                        beam_size,
+                                        end_id,
+                                        is_accumulated,
+                                        ids->place());
+    auto selected_items = ToMap(items, high_level.back());
+    if (FLAGS_v == 3) {
+      VLOG(3) << "selected_items:";
+      for (size_t i = 0; i < selected_items.size(); ++i) {
+        VLOG(3) << "offset: " << i;
+        for (auto &item : selected_items[i]) {
+          VLOG(3) << item.ToString();
+        }
+      }
+    }
+
+    PruneEndBeams(
+        pre_ids, abs_lod, &selected_items, level, end_id, ids->place());
+    // calculate the output tensor's height
+    size_t num_instances = std::accumulate(
+        std::begin(selected_items),
+        std::end(selected_items),
+        0,
+        [](size_t a, std::vector<Item> &b) { return a + b.size(); });
+    // the output tensor shape should be [num_instances, 1]
+    auto dims =
+        make_ddim(std::vector<int64_t>({static_cast<int>(num_instances), 1}));
+    selected_ids->Resize(dims);
+    auto *selected_ids_data = dev_ctx.template HostAlloc<int64_t>(selected_ids);
+    selected_scores->Resize(dims);
+    auto *selected_scores_data =
+        dev_ctx.template HostAlloc<float>(selected_scores);
+    if (parent_idx != nullptr) {
+      parent_idx->Resize({static_cast<int64_t>(num_instances)});
+    }
+    auto *parent_idx_data =
+        parent_idx ? dev_ctx.template HostAlloc<int>(parent_idx) : nullptr;
+
+    // fill in data
+    std::vector<size_t> low_level;
+    size_t low_offset = 0;
+    for (auto &items : selected_items) {
+      low_level.push_back(low_offset);
+      for (auto &item : items) {
+        if (parent_idx) {
+          parent_idx_data[low_offset] = static_cast<int>(low_level.size() - 1);
+        }
+        selected_ids_data[low_offset] = item.id;
+        selected_scores_data[low_offset] = item.score;
+        low_offset++;
+      }
+    }
+    low_level.push_back(low_offset);
+
+    // fill lod
+    LegacyLoD lod(2);
+    lod[0].assign(high_level.begin(), high_level.end());
+    lod[1].assign(low_level.begin(), low_level.end());
+    if (!CheckLegacyLoD(lod)) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "lod %s is not right in"
+          " beam_search, please check your code.",
+          LoDToString(lod)));
+    }
+    selected_ids->set_lod(lod);
+    selected_scores->set_lod(lod);
+  }
+
+  /*
+   * The basic items help to sort.
+   */
+  struct Item {
+    Item() {}
+    Item(size_t offset, size_t id, float score)
+        : offset(offset), id(id), score(score) {}
+    // offset in the higher lod level.
+    size_t offset;
+    // prefix id in the lower lod level.
+    // size_t prefix;
+    // the candidate id
+    size_t id;
+    // the corresponding score
+    float score;
+
+    inline bool operator<(const Item &in) const {
+      return (score < in.score) ||
+             ((score == in.score) && (offset < in.offset));
+    }
+
+    inline void operator=(const Item &in) {
+      offset = in.offset;
+      id = in.id;
+      score = in.score;
+    }
+
+    std::string ToString() {
+      std::ostringstream os;
+      os << "{";
+      os << "offset: " << offset << ", ";
+      os << "id: " << id << ", ";
+      os << "score: " << score << "";
+      os << "}";
+      return os.str();
+    }
+  };
+
+ protected:
+  /*
+   * Prune the source sentences all branches finished, and it is optional.
+   * Pruning must one step later than finishing (thus pre_ids is needed here),
+   * since the end tokens must be written out.
+   */
+  void PruneEndBeams(const DenseTensor *pre_ids,
+                     const LegacyLoD &abs_lod,
+                     std::vector<std::vector<Item>> *items,
+                     size_t lod_level,
+                     int end_id,
+                     const Place &place) {
+    auto *pre_ids_data_xpu = pre_ids->data<int64_t>();
+    int64_t *pre_ids_data = nullptr;
+    CopyDataByCondition<int64_t>(
+        pre_ids_data_xpu, &pre_ids_data, pre_ids->numel(), place);
+
+    auto &high_level = abs_lod[lod_level];
+    for (size_t src_idx = 0; src_idx < high_level.size() - 1; ++src_idx) {
+      size_t src_prefix_start = high_level[src_idx];
+      size_t src_prefix_end = high_level[src_idx + 1];
+      bool finish_flag = true;
+      for (size_t offset = src_prefix_start; offset < src_prefix_end;
+           offset++) {
+        for (auto &item : items->at(offset)) {
+          if (item.id != static_cast<size_t>(end_id) ||
+              pre_ids_data[offset] != end_id) {
+            finish_flag = false;
+            break;
+          }
+        }
+        if (!finish_flag) break;
+      }
+      if (finish_flag) {  // all branches of the beam (source sentence) end and
+                          // prune this beam
+        for (size_t offset = src_prefix_start; offset < src_prefix_end;
+             offset++)
+          items->at(offset).clear();
+      }
+    }
+    free(pre_ids_data);
+  }
+
+  /*
+   * Transform the items into a map whose key is offset, value is the items.
+   * NOTE low performance.
+   */
+  std::vector<std::vector<Item>> ToMap(
+      const std::vector<std::vector<Item>> &items, size_t element_num) {
+    std::vector<std::vector<Item>> result;
+    result.resize(element_num);
+    for (auto &entries : items) {
+      for (const auto &item : entries) {
+        result[item.offset].push_back(item);
+      }
+    }
+    return result;
+  }
+
+  void Insert(std::vector<Item> *top_beam_ptr,
+              const Item &item,
+              size_t beam_size) {
+    std::vector<Item> &top_beam = *top_beam_ptr;
+
+    size_t num_beams = top_beam.size();
+
+    if (num_beams < beam_size) {
+      top_beam.resize(num_beams + 1);
+      num_beams++;
+    } else {
+      if (item < top_beam[beam_size - 1]) {
+        return;
+      }
+    }
+
+    for (int k = static_cast<int>(num_beams) - 2; k >= 0; --k) {
+      if (top_beam[k] < item) {
+        top_beam[k + 1] = top_beam[k];
+      } else {
+        top_beam[k + 1] = item;
+        return;
+      }
+    }
+    top_beam[0] = item;
+  }
+
+  /*
+   * For each source, select top beam_size records.
+   */
+  std::vector<std::vector<Item>> SelectTopBeamSizeItems(
+      const DenseTensor *pre_ids,
+      const DenseTensor *pre_scores,
+      const DenseTensor *ids,
+      const DenseTensor *scores,
+      size_t lod_level,
+      size_t beam_size,
+      int end_id,
+      bool is_accumulated,
+      const Place &place) {
+    std::vector<std::vector<Item>> result;
+
+    // find the current candidates
+    auto abs_lod = ToAbsOffset(scores->lod());
+
+    auto *pre_ids_data_xpu = pre_ids->data<int64_t>();
+    int64_t *pre_ids_data = nullptr;
+    CopyDataByCondition<int64_t>(
+        pre_ids_data_xpu, &pre_ids_data, pre_ids->numel(), place);
+
+    auto *pre_scores_data_xpu = pre_scores->data<float>();
+    float *pre_scores_data = nullptr;
+    CopyDataByCondition<float>(
+        pre_scores_data_xpu, &pre_scores_data, pre_scores->numel(), place);
+
+    auto *ids_data_xpu = ids ? ids->data<int64_t>() : nullptr;
+    int64_t *ids_data = nullptr;
+    CopyDataByCondition<int64_t>(ids_data_xpu, &ids_data, ids->numel(), place);
+
+    auto *scores_data_xpu = scores->data<float>();
+    float *scores_data = nullptr;
+    CopyDataByCondition<float>(
+        scores_data_xpu, &scores_data, scores->numel(), place);
+
+    size_t num_seqs = scores->NumElements(lod_level);
+    size_t seq_width = 1;
+    for (int i = 1; i < scores->dims().size(); i++) {
+      seq_width *= scores->dims()[i];
+    }
+
+    for (size_t seq_id = 0; seq_id < num_seqs; ++seq_id) {
+      size_t seq_offset_start = abs_lod[lod_level][seq_id];
+      size_t seq_offset_end = abs_lod[lod_level][seq_id + 1];
+
+      std::vector<Item> top_beam;
+      top_beam.reserve(beam_size);
+
+      for (size_t offset = seq_offset_start; offset < seq_offset_end;
+           ++offset) {
+        auto pre_id = pre_ids_data[offset];
+        auto pre_score = pre_scores_data[offset];
+
+        if (pre_id == end_id) {
+          // Allocate all probability mass to end_id for finished branches and
+          // the other candidate ids can be ignored.
+          Item item(offset, end_id, pre_score);
+          Insert(&top_beam, item, beam_size);
+        } else {
+          size_t index = offset * seq_width;
+          for (size_t d = 0; d < seq_width; d++, index++) {
+            int64_t id = ids_data ? ids_data[index] : static_cast<int64_t>(d);
+            float score = is_accumulated
+                              ? scores_data[index]
+                              : pre_score + std::log(scores_data[index]);
+
+            Item item(offset, id, score);
+            Insert(&top_beam, item, beam_size);
+          }
+        }
+      }
+
+      result.emplace_back(top_beam);
+    }
+
+    if (FLAGS_v == 3) {
+      VLOG(3) << "SelectTopBeamSizeItems result size " << result.size();
+      for (auto &items : result) {
+        VLOG(3) << "item set:";
+        for (auto &item : items) {
+          VLOG(3) << item.ToString();
+        }
+      }
+    }
+
+    free(pre_ids_data);
+    free(pre_scores_data);
+    free(ids_data);
+    free(scores_data);
+
+    return result;
+  }
+};
+
+template class BeamSearchFunctor<XPUContext, int>;
+template class BeamSearchFunctor<XPUContext, int64_t>;
+template class BeamSearchFunctor<XPUContext, float>;
+template class BeamSearchFunctor<XPUContext, double>;
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/bert_encoder_functor.cu

@@ -0,0 +1,415 @@
+// 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/funcs/math/bert_encoder_functor.h"
+#include <algorithm>
+#include <type_traits>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+namespace phi {
+namespace math {
+
+// NOTE(chenfeiyu): explicitly use operator+ for float2
+// since float2 is not in namespace phi::funcs, ADL won't help
+using funcs::operator+;
+
+template <typename T>
+__device__ __forceinline__ T local_rsqrt(T num) {
+  return rsqrt(static_cast<float>(num));
+}
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+__device__ __forceinline__ half local_rsqrt(half num) { return hrsqrt(num); }
+#endif
+
+template <typename T, int TPB>
+__device__ inline void LayerNormSmall(T val,
+                                      const funcs::kvp<T> &thread_data,
+                                      const int ld,
+                                      const int idx,
+                                      const T *bias,
+                                      const T *scale,
+                                      T *output,
+                                      T eps) {
+  using BlockReduce = cub::BlockReduce<funcs::kvp<T>, TPB>;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ T mu;      // mean
+  __shared__ T rsigma;  // 1 / std.dev.
+
+  const auto sum_kv = BlockReduce(temp_storage).Reduce(thread_data, cub::Sum());
+
+  if (threadIdx.x == 0) {
+    mu = sum_kv.key;
+    rsigma = local_rsqrt(sum_kv.value - mu * mu + eps);
+  }
+  __syncthreads();
+
+  if (threadIdx.x < ld) {
+    const T g(scale[threadIdx.x]);
+    const T b(bias[threadIdx.x]);
+    output[idx] = g * (val - mu) * rsigma + b;
+  }
+}
+
+template <typename T, int TPB>
+__device__ inline void LayerNorm(const funcs::kvp<T> &thread_data,
+                                 const int ld,
+                                 const int offset,
+                                 const T *bias,
+                                 const T *scale,
+                                 T *output,
+                                 T eps) {
+  using BlockReduce = cub::BlockReduce<funcs::kvp<T>, TPB>;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ T mu;      // mean
+  __shared__ T rsigma;  // 1 / std.dev.
+
+  const auto sum_kv = BlockReduce(temp_storage).Reduce(thread_data, cub::Sum());
+
+  if (threadIdx.x == 0) {
+    mu = sum_kv.key;
+    rsigma = local_rsqrt(sum_kv.value - mu * mu + eps);
+  }
+  __syncthreads();
+
+  for (int i = threadIdx.x; i < ld; i += TPB) {
+    const int idx = offset + i;
+    const T val = output[idx];
+    const T g(scale[i]);
+    const T b(bias[i]);
+    output[idx] = g * (val - mu) * rsigma + b;
+  }
+}
+
+template <typename T, typename T2, int TPB>
+__device__ inline void LayerNorm2(const funcs::kvp<T> &thread_data,
+                                  const int ld,
+                                  const int offset,
+                                  const T2 *bias,
+                                  const T2 *scale,
+                                  T2 *output,
+                                  T eps) {
+  using BlockReduce = cub::BlockReduce<funcs::kvp<T>, TPB>;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ T mu;      // mean
+  __shared__ T rsigma;  // 1 / std.dev.
+
+  const auto sum_kv = BlockReduce(temp_storage).Reduce(thread_data, cub::Sum());
+
+  if (threadIdx.x == 0) {
+    mu = sum_kv.key;
+    rsigma = local_rsqrt(sum_kv.value - mu * mu + eps);
+  }
+  __syncthreads();
+
+  for (int i = threadIdx.x; i < ld; i += TPB) {
+    const int idx = offset + i;
+    T2 val = output[idx];
+    const T2 g = scale[i];
+    const T2 b = bias[i];
+    val.x = T(g.x) * (val.x - mu) * rsigma + T(b.x);
+    val.y = T(g.y) * (val.y - mu) * rsigma + T(b.y);
+    output[idx] = val;
+  }
+}
+
+template <typename T, unsigned TPB>
+__global__ void SkipLayerNormSmallKernel(int num,
+                                         int hidden,
+                                         const T *input1,
+                                         const T *input2,
+                                         T *output,
+                                         const T *scale,
+                                         const T *bias,
+                                         T eps) {
+  const T rld = T(1) / T(hidden);
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<T> thread_data(0, 0);
+  const int idx = offset + threadIdx.x;
+  T val = 0;
+  if (threadIdx.x < hidden) {
+    val = input1[idx] + input2[idx];
+    const T rldval = rld * val;
+    thread_data = pair_sum(thread_data, funcs::kvp<T>(rldval, rldval * val));
+  }
+  LayerNormSmall<T, TPB>(
+      val, thread_data, hidden, idx, bias, scale, output, eps);
+}
+
+// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
+#ifndef __HIPCC__  // @{ Half kernel: SkipLayerNormSmallKernel
+template <>
+__global__ void SkipLayerNormSmallKernel<half, 32>(int num,
+                                                   int hidden,
+                                                   const half *input1,
+                                                   const half *input2,
+                                                   half *output,
+                                                   const half *scale,
+                                                   const half *bias,
+                                                   half eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const half rld = half(1) / half(hidden);
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<half> thread_data(0, 0);
+  const int idx = offset + threadIdx.x;
+  half val = 0;
+  if (threadIdx.x < hidden) {
+    val = input1[idx] + input2[idx];
+    const half rldval = rld * val;
+    thread_data = pair_sum(thread_data, funcs::kvp<half>(rldval, rldval * val));
+  }
+  LayerNormSmall<half, 32>(
+      val, thread_data, hidden, idx, bias, scale, output, eps);
+#endif
+}
+
+template <>
+__global__ void SkipLayerNormSmallKernel<half, 128>(int num,
+                                                    int hidden,
+                                                    const half *input1,
+                                                    const half *input2,
+                                                    half *output,
+                                                    const half *scale,
+                                                    const half *bias,
+                                                    half eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const half rld = half(1) / half(hidden);
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<half> thread_data(0, 0);
+  const int idx = offset + threadIdx.x;
+  half val = 0;
+  if (threadIdx.x < hidden) {
+    val = input1[idx] + input2[idx];
+    const half rldval = rld * val;
+    thread_data = pair_sum(thread_data, funcs::kvp<half>(rldval, rldval * val));
+  }
+  LayerNormSmall<half, 128>(
+      val, thread_data, hidden, idx, bias, scale, output, eps);
+#endif
+}
+
+template <>
+__global__ void SkipLayerNormSmallKernel<half, 384>(int num,
+                                                    int hidden,
+                                                    const half *input1,
+                                                    const half *input2,
+                                                    half *output,
+                                                    const half *scale,
+                                                    const half *bias,
+                                                    half eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const half rld = half(1) / half(hidden);
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<half> thread_data(0, 0);
+  const int idx = offset + threadIdx.x;
+  half val = 0;
+  if (threadIdx.x < hidden) {
+    val = input1[idx] + input2[idx];
+    const half rldval = rld * val;
+    thread_data = pair_sum(thread_data, funcs::kvp<half>(rldval, rldval * val));
+  }
+  LayerNormSmall<half, 384>(
+      val, thread_data, hidden, idx, bias, scale, output, eps);
+#endif
+}
+#endif  // @} End Half kernel: SkipLayerNormSmallKernel
+
+template <typename T, unsigned TPB>
+__global__ void SkipLayerNormKernel(int num,
+                                    int hidden,
+                                    const T *input1,
+                                    const T *input2,
+                                    T *output,
+                                    const T *scale,
+                                    const T *bias,
+                                    T eps) {
+  const T rld = T(1) / T(hidden);
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<T> thread_data(0, 0);
+
+  for (int it = threadIdx.x; it < hidden; it += TPB) {
+    const int idx = offset + it;
+    const T val = input1[idx] + input2[idx];
+    const T rldval = rld * val;
+    thread_data = pair_sum(thread_data, funcs::kvp<T>(rldval, rldval * val));
+    output[idx] = val;
+  }
+  LayerNorm<T, TPB>(thread_data, hidden, offset, bias, scale, output, eps);
+}
+
+// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
+#ifndef __HIPCC__  // @{ Half kernel: SkipLayerNormKernel
+template <>
+__global__ void SkipLayerNormKernel<half, 256>(int num,
+                                               int hidden,
+                                               const half *input1,
+                                               const half *input2,
+                                               half *output,
+                                               const half *scale,
+                                               const half *bias,
+                                               half eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const half rld = half(1) / half(hidden);
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<half> thread_data(0, 0);
+
+  for (int it = threadIdx.x; it < hidden; it += 256) {
+    const int idx = offset + it;
+    const half val = input1[idx] + input2[idx];
+    const half rldval = rld * val;
+    thread_data = pair_sum(thread_data, funcs::kvp<half>(rldval, rldval * val));
+    output[idx] = val;
+  }
+  LayerNorm<half, 256>(thread_data, hidden, offset, bias, scale, output, eps);
+#endif
+}
+#endif  // @} End Half kernel: SkipLayerNormKernel
+
+template <typename T, typename T2, unsigned TPB>
+__global__ void SkipLayerNormKernel2(int num,
+                                     int hidden,
+                                     const T2 *input1,
+                                     const T2 *input2,
+                                     T2 *output,
+                                     const T2 *scale,
+                                     const T2 *bias,
+                                     float eps) {
+  const T rld = T(0.5f / hidden);  // because hidden is hidden/2
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<T> thread_data(0, 0);
+
+  for (int it = threadIdx.x; it < hidden; it += TPB) {
+    const int idx = offset + it;
+    const T2 val2 = input1[idx] + input2[idx];
+    thread_data =
+        pair_sum(thread_data,
+                 funcs::kvp<T>(rld * (val2.x + val2.y),
+                               rld * val2.x * val2.x + rld * val2.y * val2.y));
+    output[idx] = val2;
+  }
+  LayerNorm2<T, T2, TPB>(thread_data, hidden, offset, bias, scale, output, eps);
+}
+
+// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
+#ifndef __HIPCC__  // @{ Half kernel: SkipLayerNormKernel2
+template <>
+__global__ void SkipLayerNormKernel2<half, half2, 256>(int num,
+                                                       int hidden,
+                                                       const half2 *input1,
+                                                       const half2 *input2,
+                                                       half2 *output,
+                                                       const half2 *scale,
+                                                       const half2 *bias,
+                                                       float eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const half rld = half(0.5f / hidden);  // because hidden is hidden/2
+  const int offset = blockIdx.x * hidden;
+  cub::Sum pair_sum;
+  funcs::kvp<half> thread_data(0, 0);
+
+  for (int it = threadIdx.x; it < hidden; it += 256) {
+    const int idx = offset + it;
+    const half2 val2 = input1[idx] + input2[idx];
+    thread_data = pair_sum(
+        thread_data,
+        funcs::kvp<half>(rld * (val2.x + val2.y),
+                         rld * val2.x * val2.x + rld * val2.y * val2.y));
+    output[idx] = val2;
+  }
+  LayerNorm2<half, half2, 256>(
+      thread_data, hidden, offset, bias, scale, output, eps);
+#endif
+}
+#endif  // @} End Half kernel: SkipLayerNormKernel2
+
+template <typename T>
+void SkipLayerNormFunctor<T>::operator()(const int num,
+                                         const int hidden,
+                                         const T *input1,
+                                         const T *input2,
+                                         const T *scale,
+                                         const T *bias,
+                                         T *output,
+                                         float eps,
+                                         gpuStream_t stream) {
+  int block = num / hidden;
+  if (hidden <= WARP_SIZE) {
+    const int threads = WARP_SIZE;
+    SkipLayerNormSmallKernel<T, threads><<<block, threads, 0, stream>>>(
+        num, hidden, input1, input2, output, scale, bias, eps);
+  } else if (hidden <= 128) {
+    const int threads = 128;
+    SkipLayerNormSmallKernel<T, threads><<<block, threads, 0, stream>>>(
+        num, hidden, input1, input2, output, scale, bias, eps);
+  } else if (hidden == 384) {
+    const int threads = 384;
+    SkipLayerNormSmallKernel<T, threads><<<block, threads, 0, stream>>>(
+        num, hidden, input1, input2, output, scale, bias, eps);
+  } else {
+    const int threads = 256;
+    if (hidden % 2 == 0) {
+      if (std::is_same<T, float>::value) {
+        SkipLayerNormKernel2<float, float2, threads>
+            <<<block, threads, 0, stream>>>(
+                num,
+                hidden / 2,
+                reinterpret_cast<const float2 *>(input1),
+                reinterpret_cast<const float2 *>(input2),
+                reinterpret_cast<float2 *>(output),
+                reinterpret_cast<const float2 *>(scale),
+                reinterpret_cast<const float2 *>(bias),
+                eps);
+// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
+#ifndef __HIPCC__
+      } else if (std::is_same<T, __half>::value) {
+        SkipLayerNormKernel2<__half, __half2, threads>
+            <<<block, threads, 0, stream>>>(
+                num,
+                hidden / 2,
+                reinterpret_cast<const __half2 *>(input1),
+                reinterpret_cast<const __half2 *>(input2),
+                reinterpret_cast<__half2 *>(output),
+                reinterpret_cast<const __half2 *>(scale),
+                reinterpret_cast<const __half2 *>(bias),
+                eps);
+#endif
+      } else {
+        assert(false);
+        // should not be here
+      }
+    } else {
+      SkipLayerNormKernel<T, threads><<<block, threads, 0, stream>>>(
+          num, hidden, input1, input2, output, scale, bias, eps);
+    }
+  }
+}
+
+template class SkipLayerNormFunctor<float>;
+
+// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
+#if defined(PADDLE_WITH_CUDA)
+template class SkipLayerNormFunctor<half>;
+#endif
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/bert_encoder_functor.h

@@ -0,0 +1,73 @@
+// 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
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/kernels/funcs/cub.h"
+
+namespace phi {
+namespace math {
+
+template <typename T>
+struct CUDATypeTraits;
+
+template <>
+struct CUDATypeTraits<half> {
+  typedef phi::float16 TYPE;
+};
+
+template <>
+struct CUDATypeTraits<float> {
+  typedef float TYPE;
+};
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+// This functor involves a fusion calculation in Ernie or Bert.
+// The fusion mode is as follows:
+//
+// |           |
+// other_op1   other_op2
+//      |           |
+//      |------elementwise_add
+//                  |
+//              layer_norm
+//                  |
+//              other_op3
+//                  |
+
+template <typename T>
+class SkipLayerNormFunctor {
+ public:
+  void operator()(const int num,
+                  const int hidden,
+                  const T *input1,
+                  const T *input2,
+                  const T *scale,
+                  const T *bias,
+                  T *output,
+                  float eps,
+                  gpuStream_t stream);
+};
+#endif
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/bloomfilter.h

@@ -0,0 +1,194 @@
+// 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
+#define BLOOMFILTER_MAGIC_NUM_NEW 17070416
+
+#include <inttypes.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#include <sys/types.h>
+#elif defined(_WIN32)
+#ifndef NOMINMAX
+#define NOMINMAX  // msvc max/min macro conflict with std::min/max
+#endif
+#include <windows.h>
+#else
+#include <unistd.h>
+#endif  // _WIN32
+
+#include <cinttypes>
+
+namespace phi {
+namespace math {
+
+#pragma pack(push, 4)
+struct bloomfilter {
+  uint64_t magic_num;
+  uint64_t m;
+  uint64_t k;
+  uint64_t count;
+  unsigned char bit_vector[1];
+};
+#pragma pack(pop)
+
+inline int bloomfilter_get(const struct bloomfilter *bloomfilter,
+                           const void *key,
+                           size_t len);
+inline int bloomfilter_check(struct bloomfilter *filter);
+
+#define bit_get(v, n) ((v)[(n) >> 3] & (0x1 << (0x7 - ((n)&0x7))))
+#define ROTL64(x, r) (((x) << (r)) | ((x) >> (64 - (r))))
+#define BIG_CONSTANT(x) (x##LLU)
+
+uint64_t fmix64(uint64_t k) {
+  k ^= k >> 33;
+  k *= BIG_CONSTANT(0xff51afd7ed558ccd);
+  k ^= k >> 33;
+  k *= BIG_CONSTANT(0xc4ceb9fe1a85ec53);
+  k ^= k >> 33;
+  return k;
+}
+
+void murmurhash3_x64_128(const void *key,
+                         const int len,
+                         const uint32_t seed,
+                         void *out) {
+  const uint8_t *data = (const uint8_t *)key;
+  const int nblocks = len / 16;
+
+  uint64_t h1 = seed;
+  uint64_t h2 = seed;
+  int i = 0;
+
+  const uint64_t c1 = BIG_CONSTANT(0x87c37b91114253d5);
+  const uint64_t c2 = BIG_CONSTANT(0x4cf5ad432745937f);
+
+  //----------
+  // body
+
+  const uint64_t *blocks = (const uint64_t *)(data);
+
+  uint64_t k1;
+  uint64_t k2;
+
+  for (i = 0; i < nblocks; i++) {
+    k1 = blocks[i * 2 + 0];
+    k2 = blocks[i * 2 + 1];
+
+    k1 *= c1;
+    k1 = ROTL64(k1, 31);
+    k1 *= c2;
+    h1 ^= k1;
+
+    h1 = ROTL64(h1, 27);
+    h1 += h2;
+    h1 = h1 * 5 + 0x52dce729;
+
+    k2 *= c2;
+    k2 = ROTL64(k2, 33);
+    k2 *= c1;
+    h2 ^= k2;
+
+    h2 = ROTL64(h2, 31);
+    h2 += h1;
+    h2 = h2 * 5 + 0x38495ab5;
+  }
+
+  //----------
+  // tail
+
+  const uint8_t *tail = (const uint8_t *)(data + nblocks * 16);
+  uint64_t nk1 = 0;
+  uint64_t nk2 = 0;
+
+  uint64_t tail0_64 = *(uint64_t *)(tail);     // NOLINT
+  uint64_t tail_64 = *(uint64_t *)(tail + 8);  // NOLINT
+  uint64_t mask0 = 0xffffffffffffffff;
+  uint64_t mask = 0x00ffffffffffffff;
+
+  int flag = len & 15;
+  if (flag && flag <= 8) {
+    tail0_64 &= (mask0 >> ((8 - flag) << 3));
+  } else if (flag > 8) {
+    tail_64 &= (mask >> ((15 - flag) << 3));
+    nk2 ^= tail_64;
+    nk2 *= c2;
+    nk2 = ROTL64(nk2, 33);
+    nk2 *= c1;
+    h2 ^= nk2;
+  }
+
+  if (flag) {
+    nk1 ^= tail0_64;
+    nk1 *= c1;
+    nk1 = ROTL64(nk1, 31);
+    nk1 *= c2;
+    h1 ^= nk1;
+  }
+
+  //----------
+  // finalization
+
+  h1 ^= len;
+  h2 ^= len;
+
+  h1 += h2;
+  h2 += h1;
+
+  h1 = fmix64(h1);
+  h2 = fmix64(h2);
+
+  h1 += h2;
+  h2 += h1;
+
+  reinterpret_cast<uint64_t *>(out)[0] = h1;
+  reinterpret_cast<uint64_t *>(out)[1] = h2;
+}
+
+inline int bloomfilter_check(struct bloomfilter *filter) {
+  if (filter->magic_num == BLOOMFILTER_MAGIC_NUM_NEW) {
+    return 1;
+  } else {
+    fprintf(stderr, "error magic_num, %" PRIu64 "\n", filter->magic_num);
+    return 0;
+  }
+}
+
+inline int bloomfilter_get(const struct bloomfilter *bloomfilter,
+                           const void *key,
+                           size_t len) {
+  uint32_t i;
+  uint64_t result[2];
+
+  for (i = 0; i < bloomfilter->k; i++) {
+    murmurhash3_x64_128(key, len, i, &result);
+    result[0] %= bloomfilter->m;
+    result[1] %= bloomfilter->m;
+    if (!bit_get(bloomfilter->bit_vector, result[0])) {
+      return 0;
+    }
+    if (!bit_get(bloomfilter->bit_vector, result[1])) {
+      return 0;
+    }
+  }
+  return 1;
+}
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/context_project.cc

@@ -0,0 +1,25 @@
+// 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/funcs/math/context_project.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+
+namespace phi {
+namespace math {
+
+template class ContextProjectFunctor<CPUContext, float>;
+template class ContextProjectFunctor<CPUContext, double>;
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/context_project.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/funcs/math/context_project.h"
+
+namespace phi {
+namespace math {
+
+template class ContextProjectFunctor<GPUContext, float>;
+template class ContextProjectFunctor<GPUContext, double>;
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/context_project.h

@@ -0,0 +1,351 @@
+// 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 <algorithm>
+#include <vector>
+
+#include "paddle/common/enforce.h"
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/im2col.h"
+
+namespace phi {
+namespace math {
+
+/*
+ * \brief Context projection concatenates features in adjacent time-steps in
+ * a sequence. The i-th row of the output is the concatenation of
+ * context_length rows of the input. The context_length rows are the
+ * consecutive rows from the i+shift_start row.
+ * ContextProjectGradFunctor is the inverse process of ContextProjectFunctor.
+ *
+ * \param in            Input data.
+ * \param Shape         The shape of Input data:
+ *                        [mini-batch, input_hidden_size].
+ *
+ * \param padding_data  Padding data.
+ * \param Shape         The shape of Padding data:
+ *                        [up_pad + down_pad, input_hidden_size].
+ *
+ * \param col           Col data.
+ * \param Shape         The shape of Col data:
+ *                        [mini-batch, context_length * input_hidden_size].
+ *
+ * For a mini-batch of 2 variable lengths sentences, containing 3, and 1
+ * time-steps:
+ *
+ * Assumed input (X) is a [4, M, N] float DenseTensor, and X->lod()[0] =
+ * [0, 3, 4]. Besides, for the sake of simplicity, we assume M=1 and N=2.
+ *
+ * X = [[a1, a2;
+ *       b1, b2;
+ *       c1, c2]
+ *      [d1, d2]]
+ *
+ * This is to say that input (X) has 4 words and the dimension of each word
+ * representation is 2.
+ *
+ * - Case1:
+ *   If context_start is -1 and padding_trainable is false, we use zero to pad
+ *   instead of learned weight to pad,
+ *   and the context_length is 3, the output (Out) is:
+ *
+ *   Out =[[0,  0,  a1, a2, b1, b2;
+ *          a1, a2, b1, b2, c1, c2;
+ *          b1, b2, c1, c2, 0,  0 ]
+ *          [0,  0, d1, d2, 0,  0 ]]
+ *
+ * - Case2:
+ *   If context_start is -1 and padding_trainable is true, we use learned weight
+ *   to pad,
+ *   and the context_length is 3, the output (Out) is:
+ *
+ *   Out = [[w1, w2, a1, a2, b1, b2;
+ *           a1, a2, b1, b2, c1, c2;
+ *           b1, b2, c1, c2, w3, w4]
+ *          [w1, w2, d1, d2, w3, w4]]
+ *
+ */
+
+template <typename DeviceContext, typename T>
+class ContextProjectFunctor {
+ public:
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor& in,
+                  const DenseTensor* padding_data,
+                  bool padding_trainable,
+                  const int context_start,
+                  const int context_length,
+                  const int context_stride,
+                  const int up_pad,
+                  const int down_pad,
+                  DenseTensor* col) {
+    auto lod_level_0 = in.lod()[0];
+
+    funcs::Im2ColFunctor<funcs::ColFormat::OCF, DeviceContext, float>
+        im2col_ocf;
+
+    std::vector<int> dilation({1, 1});
+    std::vector<int> padding({up_pad, 0, down_pad, 0});
+    std::vector<int> stride({context_stride, 1});
+
+    int input_row_begin, input_row_end;
+    int sequence_height;
+    int64_t sequence_width = in.dims()[1];
+
+    for (int i = 0; i < static_cast<int>(lod_level_0.size()) - 1; ++i) {
+      if (lod_level_0[i] == lod_level_0[i + 1]) continue;
+
+      input_row_begin = (context_start > 0)
+                            ? static_cast<int>(lod_level_0[i]) + context_start
+                            : static_cast<int>(lod_level_0[i]);
+      input_row_end = static_cast<int>(lod_level_0[i + 1]);
+
+      DenseTensor out_t = col->Slice(static_cast<int>(lod_level_0[i]),
+                                     static_cast<int>(lod_level_0[i + 1]));
+
+      sequence_height = static_cast<int>(out_t.dims()[0]);
+
+      if (input_row_begin < input_row_end) {
+        DenseTensor in_t = in.Slice(input_row_begin, input_row_end);
+
+        std::vector<int64_t> output_shape(
+            {sequence_height,
+             1,
+             1,
+             context_length,
+             sequence_width});  // output_height, output_width,
+        // input_channels, filter_height, filter_width
+        out_t.Resize(output_shape);
+
+        std::vector<int64_t> input_shape(
+            {1,
+             input_row_end - input_row_begin,
+             sequence_width});  // input_channels, input_height, input_width
+        in_t.Resize(input_shape);
+        im2col_ocf(dev_ctx, in_t, dilation, stride, padding, &out_t);
+        out_t.Resize({sequence_height, context_length * sequence_width});
+      }
+    }
+    if (padding_trainable) {
+      PADDLE_ENFORCE_NOT_NULL(
+          padding_data,
+          common::errors::InvalidArgument(
+              "The input tensor 'padding_data' should not be NULL."));
+      for (int i = 0; i < static_cast<int>(lod_level_0.size()) - 1; ++i) {
+        if (lod_level_0[i] == lod_level_0[i + 1]) continue;
+
+        DenseTensor out_t = col->Slice(static_cast<int>(lod_level_0[i]),
+                                       static_cast<int>(lod_level_0[i + 1]));
+
+        sequence_height = static_cast<int>(out_t.dims()[0]);
+
+        // add up trainable data
+        out_t.Resize({static_cast<int64_t>(sequence_height) * context_length,
+                      sequence_width});
+
+        if (up_pad > 0) {  // add up pad
+          int padding_rows = std::min(
+              up_pad, static_cast<int>(lod_level_0[i + 1] - lod_level_0[i]));
+
+          for (int k = 0; k < padding_rows; ++k) {
+            int padding_size =
+                k + context_length < up_pad ? context_length : up_pad - k;
+            DenseTensor out_t_sub = out_t.Slice(
+                k * context_length, k * context_length + padding_size);
+            DenseTensor w_sub = padding_data->Slice(k, k + padding_size);
+            phi::Copy(dev_ctx, w_sub, dev_ctx.GetPlace(), false, &out_t_sub);
+          }
+        }
+        if (down_pad > 0) {  // add down pad
+          int down_pad_begin_row =
+              std::max(0,
+                       (sequence_height - context_start - context_length) + 1) +
+              1;
+          int padding_begin = std::max(0, context_start - sequence_height);
+          int padding_size =
+              sequence_height - context_start >= context_length
+                  ? 1
+                  : context_length - (sequence_height - context_start);
+          if (context_start >= sequence_height) padding_size = context_length;
+          int padding_idx = padding_begin;
+          for (int t = 0; t + down_pad_begin_row <= sequence_height;
+               ++t, ++padding_size) {
+            if (context_start >= sequence_height) padding_size = context_length;
+            if (padding_size > context_length) {
+              padding_size = context_length;
+              padding_idx++;
+            }
+            if (padding_begin > 0 || sequence_height == context_start)
+              padding_idx = padding_begin + t;
+
+            DenseTensor out_t_sub = out_t.Slice(
+                (down_pad_begin_row + t) * context_length - padding_size,
+                (down_pad_begin_row + t) * context_length);
+            DenseTensor w_sub = padding_data->Slice(
+                up_pad + padding_idx, up_pad + padding_idx + padding_size);
+            phi::Copy(dev_ctx, w_sub, dev_ctx.GetPlace(), false, &out_t_sub);
+          }
+        }
+        out_t.Resize({sequence_height,
+                      static_cast<int64_t>(context_length) * sequence_width});
+      }
+    }
+  }
+};
+
+template <typename DeviceContext, typename T>
+class ContextProjectGradFunctor {
+ public:
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor& in,
+                  bool padding_trainable,
+                  const int context_start,
+                  const int context_length,
+                  const int context_stride,
+                  const int up_pad,
+                  const int down_pad,
+                  bool pad_grad,
+                  bool input_grad,
+                  DenseTensor* padding_data,
+                  DenseTensor* col) {
+    auto lod_level_0 = in.lod()[0];
+
+    funcs::Col2ImFunctor<funcs::ColFormat::OCF, DeviceContext, float>
+        col2im_ocf;
+
+    std::vector<int> dilation({1, 1});
+    std::vector<int> padding({up_pad, 0, down_pad, 0});
+    std::vector<int> stride({context_stride, 1});
+
+    int input_row_begin, input_row_end;
+    int sequence_height;
+    int64_t sequence_width = in.dims()[1];
+    auto blas = funcs::GetBlas<DeviceContext, T>(dev_ctx);
+
+    if (input_grad) {
+      for (int i = 0; i < static_cast<int>(lod_level_0.size()) - 1; ++i) {
+        if (lod_level_0[i] == lod_level_0[i + 1]) continue;
+
+        input_row_begin = (context_start > 0)
+                              ? static_cast<int>(lod_level_0[i]) + context_start
+                              : static_cast<int>(lod_level_0[i]);
+        input_row_end = static_cast<int>(lod_level_0[i + 1]);
+
+        DenseTensor out_t = col->Slice(static_cast<int>(lod_level_0[i]),
+                                       static_cast<int>(lod_level_0[i + 1]));
+
+        sequence_height = static_cast<int>(out_t.dims()[0]);
+
+        if (input_row_begin < input_row_end) {
+          DenseTensor in_t = in.Slice(input_row_begin, input_row_end);
+
+          std::vector<int64_t> output_shape(
+              {sequence_height,
+               1,
+               1,
+               context_length,
+               sequence_width});  // output_height, output_width,
+          // input_channels, filter_height, filter_width
+          out_t.Resize(output_shape);
+
+          std::vector<int64_t> input_shape(
+              {1,
+               input_row_end - input_row_begin,
+               sequence_width});  // input_channels, input_height, input_width
+          in_t.Resize(input_shape);
+
+          col2im_ocf(dev_ctx, out_t, dilation, stride, padding, &in_t);
+          out_t.Resize({sequence_height, context_length * sequence_width});
+        }
+      }
+    }
+    if (pad_grad) {
+      if (padding_trainable) {
+        for (int i = 0; i < static_cast<int>(lod_level_0.size()) - 1; ++i) {
+          if (lod_level_0[i] == lod_level_0[i + 1]) continue;
+
+          DenseTensor out_t = col->Slice(static_cast<int>(lod_level_0[i]),
+                                         static_cast<int>(lod_level_0[i + 1]));
+
+          sequence_height = static_cast<int>(out_t.dims()[0]);
+          out_t.Resize({static_cast<int64_t>(sequence_height) * context_length,
+                        sequence_width});
+
+          if (up_pad > 0) {
+            int padding_rows = std::min(
+                up_pad, static_cast<int>(lod_level_0[i + 1] - lod_level_0[i]));
+
+            for (int k = 0; k < padding_rows; ++k) {
+              int padding_size =
+                  k + context_length < up_pad ? context_length : up_pad - k;
+              DenseTensor out_t_sub = out_t.Slice(
+                  k * context_length, k * context_length + padding_size);
+              DenseTensor w_sub = padding_data->Slice(k, k + padding_size);
+              PADDLE_ENFORCE_LE_INT_MAX(w_sub.numel(),
+                                        "context_project AXPY size");
+              blas.AXPY(static_cast<int>(w_sub.numel()),
+                        static_cast<T>(1),
+                        out_t_sub.data<T>(),
+                        w_sub.data<T>());
+            }
+          }
+          if (down_pad > 0) {
+            int down_pad_begin_row =
+                std::max(
+                    0, (sequence_height - context_start - context_length) + 1) +
+                1;
+            int padding_begin = std::max(0, context_start - sequence_height);
+            int padding_size =
+                sequence_height - context_start >= context_length
+                    ? 1
+                    : context_length - (sequence_height - context_start);
+            if (context_start >= sequence_height) padding_size = context_length;
+            int padding_idx = padding_begin;
+            for (int t = 0; t + down_pad_begin_row <= sequence_height;
+                 ++t, ++padding_size) {
+              if (context_start >= sequence_height)
+                padding_size = context_length;
+              if (padding_size > context_length) {
+                padding_size = context_length;
+                padding_idx++;
+              }
+              if (padding_begin > 0 || sequence_height == context_start)
+                padding_idx = padding_begin + t;
+
+              DenseTensor out_t_sub = out_t.Slice(
+                  (down_pad_begin_row + t) * context_length - padding_size,
+                  (down_pad_begin_row + t) * context_length);
+              DenseTensor w_sub = padding_data->Slice(
+                  up_pad + padding_idx, up_pad + padding_idx + padding_size);
+              PADDLE_ENFORCE_LE_INT_MAX(w_sub.numel(),
+                                        "context_project AXPY size");
+              blas.AXPY(static_cast<int>(w_sub.numel()),
+                        static_cast<T>(1),
+                        out_t_sub.data<T>(),
+                        w_sub.data<T>());
+            }
+          }
+          out_t.Resize({sequence_height,
+                        static_cast<int64_t>(context_length) * sequence_width});
+        }
+      }
+    }
+  }
+};
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/cos_sim_functor.cc

@@ -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.
+
+#include "paddle/phi/kernels/funcs/math/cos_sim_functor.h"
+
+namespace phi {
+namespace math {
+template <typename T>
+struct CosSimDyFunctor<CPUContext, T> {
+  void operator()(const CPUContext& dev_ctx,
+                  const T* x_norm,
+                  const T* y_norm,
+                  const T* x,
+                  const T* y,
+                  const T* z,
+                  const T* dz,
+                  const size_t rows,
+                  const size_t cols,
+                  T* dy) const {
+    for (size_t row_id = 0; row_id < rows; ++row_id) {
+      auto xy_norm_prod = x_norm[row_id] * y_norm[0];
+      auto dz_data = dz[row_id];
+      auto z_data = z[row_id];
+      auto* x_data = x + cols * row_id;
+      auto reciprocal_xy_norm_prod = 1 / xy_norm_prod;
+
+      auto y_norm_square = y_norm[0] * y_norm[0];
+      auto reciprocal_y_norm_square = 1 / y_norm_square;
+      for (size_t i = 0; i < cols; ++i) {
+        dy[i] += dz_data * (x_data[i] * reciprocal_xy_norm_prod -
+                            z_data * y[i] * reciprocal_y_norm_square);
+      }
+    }
+  }
+};
+
+template struct CosSimDyFunctor<CPUContext, float>;
+template struct CosSimDyFunctor<CPUContext, double>;
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/cos_sim_functor.cu

@@ -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/kernels/funcs/math/cos_sim_functor.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+
+namespace phi {
+namespace math {
+
+template <typename T>
+__global__ void CosSimDyKernel(const T* x_norm,
+                               const T* y_norm,
+                               const T* x,
+                               const T* y,
+                               const T* z,
+                               const T* dz,
+                               const size_t rows,
+                               const size_t cols,
+                               T* dy) {
+  int grid_size = blockDim.x * gridDim.x;
+  T y_norm_data = y_norm[0];
+  for (size_t row_id =
+           static_cast<size_t>(blockIdx.x) * static_cast<size_t>(blockDim.x) +
+           static_cast<size_t>(threadIdx.x);
+       row_id < rows;
+       row_id += grid_size) {
+    T xy_norm_prod = x_norm[row_id] * y_norm_data;
+    T dz_data = dz[row_id];
+    T z_data = z[row_id];
+    const T* x_data = x + cols * row_id;
+    T reciprocal_xy_norm_prod = 1 / xy_norm_prod;
+
+    T y_norm_square = y_norm_data * y_norm_data;
+    T reciprocal_y_norm_square = 1 / y_norm_square;
+    for (size_t i = 0; i < cols; ++i) {
+      T dy_data = dz_data * (x_data[i] * reciprocal_xy_norm_prod -
+                             z_data * y[i] * reciprocal_y_norm_square);
+      CudaAtomicAdd(dy + i, dy_data);
+    }
+  }
+}
+
+template <typename T>
+struct CosSimDyFunctor<GPUContext, T> {
+  void operator()(const GPUContext& dev_ctx,
+                  const T* x_norm,
+                  const T* y_norm,
+                  const T* x,
+                  const T* y,
+                  const T* z,
+                  const T* dz,
+                  const size_t rows,
+                  const size_t cols,
+                  T* dy) const {
+    const int block_size = 512;
+    dim3 threads(block_size, 1);
+    dim3 grid((rows + block_size - 1) / block_size, 1);
+    CosSimDyKernel<T><<<grid, threads, 0, dev_ctx.stream()>>>(
+        x_norm, y_norm, x, y, z, dz, rows, cols, dy);
+  }
+};
+
+template struct CosSimDyFunctor<GPUContext, float>;
+template struct CosSimDyFunctor<GPUContext, double>;
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/cos_sim_functor.h

@@ -0,0 +1,190 @@
+// 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 <math.h>
+#include <stdlib.h>
+
+#include "paddle/common/hostdevice.h"
+#include "paddle/phi/backends/all_context.h"
+
+namespace phi {
+namespace math {
+
+template <typename T, bool same_row>
+struct CosSimFunctor {
+  CosSimFunctor(const T* x, const T* y, T* x_norm, T* y_norm, T* z, int cols)
+      : x_norm_(x_norm),
+        y_norm_(y_norm),
+        x_(x),
+        y_(y),
+        z_(z),
+        cols_(static_cast<size_t>(cols)) {}
+
+  inline HOSTDEVICE void operator()(size_t row_id) const {
+    auto* x = x_ + cols_ * row_id;
+    T xx = 0, xy = 0, yy = 0;
+    T eps = 1e-8;
+    if (same_row) {
+      auto* y = y_ + cols_ * row_id;
+      T tep_x, tep_y;
+      for (size_t i = 0; i < cols_; ++i) {
+        tep_x = x[i];
+        tep_y = y[i];
+        xx += tep_x * tep_x;
+
+        yy += tep_y * tep_y;
+        xy += tep_x * tep_y;
+      }
+      xx = xx > eps ? xx : eps;
+      yy = yy > eps ? yy : eps;
+      xx = sqrt(xx);
+      yy = sqrt(yy);
+      y_norm_[row_id] = yy;
+      x_norm_[row_id] = xx;
+      z_[row_id] = xy / (xx * yy);
+    } else {  // This can be wrote in a better way.
+      T tep_x, tep_y;
+      for (size_t i = 0; i < cols_; ++i) {
+        tep_x = x[i];
+        tep_y = y_[i];
+        xx += tep_x * tep_x;
+        yy += tep_y * tep_y;
+        xy += tep_x * tep_y;
+      }
+      xx = xx > eps ? xx : eps;
+      yy = yy > eps ? yy : eps;
+      xx = sqrt(xx);
+      yy = sqrt(yy);
+      if (row_id == 0) y_norm_[0] = yy;
+      x_norm_[row_id] = xx;
+      z_[row_id] = xy / (xx * yy);
+    }
+  }
+
+  T* x_norm_;
+  T* y_norm_;
+  const T* x_;
+  const T* y_;
+  T* z_;
+  const size_t cols_;
+};
+
+template <typename T>
+struct CosSimGradFunctor {
+  CosSimGradFunctor(const T* x_norm,
+                    const T* y_norm,
+                    const T* x,
+                    const T* y,
+                    const T* z,
+                    const T* dz,
+                    T* dx,
+                    int cols)
+      : x_norm_(x_norm),
+        y_norm_(y_norm),
+        x_(x),
+        y_(y),
+        z_(z),
+        dz_(dz),
+        dx_(dx),
+        cols_(static_cast<size_t>(cols)) {}
+
+  inline HOSTDEVICE void operator()(size_t row_id) const {
+    auto x_norm_square = x_norm_[row_id] * x_norm_[row_id];
+    auto xy_norm_prod = x_norm_[row_id] * y_norm_[row_id];
+    auto dz = dz_[row_id];
+    auto z = z_[row_id];
+
+    auto* dx = dx_ + cols_ * row_id;
+    auto* x = x_ + cols_ * row_id;
+    auto* y = y_ + cols_ * row_id;
+
+    auto reciprocal_xy_norm_prod = 1 / xy_norm_prod;
+    auto reciprocal_x_norm_square = 1 / x_norm_square;
+    for (size_t i = 0; i < cols_; ++i) {
+      dx[i] = dz * (y[i] * reciprocal_xy_norm_prod -
+                    z * x[i] * reciprocal_x_norm_square);
+    }
+  }
+
+  const T* x_norm_;
+  const T* y_norm_;
+  const T* x_;
+  const T* y_;
+  const T* z_;
+  const T* dz_;
+  T* dx_;
+  const size_t cols_;
+};
+
+template <typename T>
+struct CosSimDxFunctor {
+  CosSimDxFunctor(const T* x_norm,
+                  const T* y_norm,
+                  const T* x,
+                  const T* y,
+                  const T* z,
+                  const T* dz,
+                  T* dx,
+                  int cols)
+      : x_norm_(x_norm),
+        y_norm_(y_norm),
+        x_(x),
+        y_(y),
+        z_(z),
+        dz_(dz),
+        dx_(dx),
+        cols_(static_cast<size_t>(cols)) {}
+
+  inline HOSTDEVICE void operator()(size_t row_id) const {
+    auto xy_norm_prod = x_norm_[row_id] * y_norm_[0];
+    auto dz = dz_[row_id];
+    auto z = z_[row_id];
+    auto* x = x_ + cols_ * row_id;
+    auto reciprocal_xy_norm_prod = 1 / xy_norm_prod;
+    auto x_norm_square = x_norm_[row_id] * x_norm_[row_id];
+    auto* dx = dx_ + cols_ * row_id;
+    auto reciprocal_x_norm_square = 1 / x_norm_square;
+
+    for (size_t i = 0; i < cols_; ++i) {
+      dx[i] = dz * (y_[i] * reciprocal_xy_norm_prod -
+                    z * x[i] * reciprocal_x_norm_square);
+    }
+  }
+  const T* x_norm_;
+  const T* y_norm_;
+  const T* x_;
+  const T* y_;
+  const T* z_;
+  const T* dz_;
+  T* dx_;
+  const size_t cols_;
+};
+
+template <typename DeviceContext, typename T>
+struct CosSimDyFunctor {
+  void operator()(const DeviceContext& dev_ctx,
+                  const T* x_norm,
+                  const T* y_norm,
+                  const T* x,
+                  const T* y,
+                  const T* z,
+                  const T* dz,
+                  const size_t rows,
+                  const size_t cols,
+                  T* dy) const;
+};
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/prelu.cu

@@ -0,0 +1,150 @@
+// 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/funcs/math/prelu.h"
+
+namespace phi {
+namespace math {
+
+#define CUDA_NUM_THREADS 1024
+
+inline static int PADDLE_GET_BLOCKS(const int N) {
+  return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;
+}
+
+template <typename T>
+__global__ void PReluChannelFirstWiseKernel(const T *input,
+                                            const T *alpha,
+                                            T *output,
+                                            size_t channel_num,
+                                            size_t plane_size,
+                                            size_t numel) {
+  CUDA_KERNEL_LOOP(index, numel) {
+    size_t temp = index / plane_size;
+    size_t channel_index = temp % channel_num;
+    T scale = alpha[channel_index];
+    T x = input[index];
+    T zero = static_cast<T>(0);
+    output[index] = (x > zero) ? x : scale * x;
+  }
+}
+
+template <typename T>
+__global__ void PReluChannelLastWiseKernel(const T *input,
+                                           const T *alpha,
+                                           T *output,
+                                           size_t channel_num,
+                                           size_t numel) {
+  CUDA_KERNEL_LOOP(index, numel) {
+    size_t channel_index = index % channel_num;
+    T scale = alpha[channel_index];
+    T x = input[index];
+    T zero = static_cast<T>(0);
+    output[index] = (x > zero) ? x : scale * x;
+  }
+}
+
+template <typename T>
+__global__ void PReluElementWiseKernel(const T *input,
+                                       const T *alpha,
+                                       T *output,
+                                       size_t spatial_size,
+                                       size_t numel) {
+  CUDA_KERNEL_LOOP(index, numel) {
+    size_t element_index = index % spatial_size;
+    T scale = alpha[element_index];
+    T x = input[index];
+    T zero = static_cast<T>(0);
+    output[index] = (x > zero) ? x : scale * x;
+  }
+}
+
+template <typename T>
+__global__ void PReluScalarKernel(const T *input,
+                                  const T *alpha,
+                                  T *output,
+                                  size_t numel) {
+  T scale = alpha[0];
+  CUDA_KERNEL_LOOP(index, numel) {
+    T x = input[index];
+    T zero = static_cast<T>(0);
+    output[index] = (x > zero) ? x : scale * x;
+  }
+}
+
+template <typename T>
+void PreluChannelWiseDirectCUDAFunctor<T>::operator()(gpuStream_t stream,
+                                                      const T *input,
+                                                      const T *alpha,
+                                                      T *output,
+                                                      size_t batch_size,
+                                                      size_t channel,
+                                                      bool channel_last,
+                                                      size_t numel) {
+  if (channel_last) {
+    PReluChannelLastWiseKernel<<<PADDLE_GET_BLOCKS(numel),
+                                 CUDA_NUM_THREADS,
+                                 0,
+                                 stream>>>(
+        input, alpha, output, channel, numel);
+  } else {
+    PReluChannelFirstWiseKernel<<<PADDLE_GET_BLOCKS(numel),
+                                  CUDA_NUM_THREADS,
+                                  0,
+                                  stream>>>(
+        input, alpha, output, channel, numel / batch_size / channel, numel);
+  }
+}
+
+template <typename T>
+void PreluElementWiseDirectCUDAFunctor<T>::operator()(gpuStream_t stream,
+                                                      const T *input,
+                                                      const T *alpha,
+                                                      T *output,
+                                                      size_t batch_size,
+                                                      size_t numel) {
+  PReluElementWiseKernel<<<PADDLE_GET_BLOCKS(numel),
+                           CUDA_NUM_THREADS,
+                           0,
+                           stream>>>(
+      input, alpha, output, numel / batch_size, numel);
+}
+
+template <typename T>
+void PreluScalarDirectCUDAFunctor<T>::operator()(gpuStream_t stream,
+                                                 const T *input,
+                                                 const T *alpha,
+                                                 T *output,
+                                                 size_t numel) {
+  PReluScalarKernel<<<PADDLE_GET_BLOCKS(numel), CUDA_NUM_THREADS, 0, stream>>>(
+      input, alpha, output, numel);
+}
+
+template class PreluChannelWiseDirectCUDAFunctor<float>;
+template class PreluChannelWiseDirectCUDAFunctor<phi::float16>;
+template class PreluChannelWiseDirectCUDAFunctor<phi::bfloat16>;
+template class PreluChannelWiseDirectCUDAFunctor<double>;
+
+template class PreluElementWiseDirectCUDAFunctor<float>;
+template class PreluElementWiseDirectCUDAFunctor<phi::float16>;
+template class PreluElementWiseDirectCUDAFunctor<phi::bfloat16>;
+template class PreluElementWiseDirectCUDAFunctor<double>;
+
+template class PreluScalarDirectCUDAFunctor<float>;
+template class PreluScalarDirectCUDAFunctor<phi::float16>;
+template class PreluScalarDirectCUDAFunctor<phi::bfloat16>;
+template class PreluScalarDirectCUDAFunctor<double>;
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/prelu.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 <vector>
+
+#include "paddle/phi/api/include/tensor.h"
+#include "paddle/phi/backends/gpu/gpu_dnn.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+namespace math {
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+template <typename T>
+class PreluChannelWiseDirectCUDAFunctor {
+ public:
+  void operator()(gpuStream_t stream,
+                  const T *input,
+                  const T *alpha,
+                  T *output,
+                  size_t batch_size,
+                  size_t channel,
+                  bool channel_last,
+                  size_t numel);
+};
+
+template <typename T>
+class PreluElementWiseDirectCUDAFunctor {
+ public:
+  void operator()(gpuStream_t stream,
+                  const T *input,
+                  const T *alpha,
+                  T *output,
+                  size_t batch_size,
+                  size_t numel);
+};
+
+template <typename T>
+class PreluScalarDirectCUDAFunctor {
+ public:
+  void operator()(gpuStream_t stream,
+                  const T *input,
+                  const T *alpha,
+                  T *output,
+                  size_t numel);
+};
+
+#endif
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/sampler.cc

@@ -0,0 +1,95 @@
+// 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/funcs/math/sampler.h"
+
+#include <glog/logging.h>
+
+#include "paddle/phi/core/generator.h"
+
+namespace phi::math {
+
+Sampler::~Sampler() = default;
+
+UniformSampler::UniformSampler(int64_t range, unsigned int seed)
+    : Sampler(range, seed), inv_range_(1.0f / (range + 1)) {  // NOLINT
+  random_engine_ = phi::GetCPURandomEngine(seed_);
+  dist_ = std::make_shared<std::uniform_int_distribution<>>(0, range);
+}
+
+int64_t UniformSampler::Sample() const { return (*dist_)(*random_engine_); }
+
+float UniformSampler::Probability(int64_t value) const { return inv_range_; }
+
+LogUniformSampler::LogUniformSampler(int64_t range, unsigned int seed)
+    : Sampler(range, seed), log_range_(log(range + 1)) {  // NOLINT
+  random_engine_ = phi::GetCPURandomEngine(seed_);
+  dist_ = std::make_shared<std::uniform_real_distribution<>>(0, 1);
+}
+
+int64_t LogUniformSampler::Sample() const {
+  // Got Log Uniform distribution from uniform distribution by
+  // inverse_transform_sampling method
+  // More details:
+  // https://wanghaoshuang.github.io/2017/11/Log-uniform-distribution-sampler/
+  auto cur_random = (*dist_)(*random_engine_);
+  const int64_t value = static_cast<int64_t>(exp(cur_random * log_range_)) - 1;
+  // Mathematically, value should be <= range_, but might not be due to some
+  // floating point roundoff, so we mod by range_.
+  return value % range_;
+}
+
+float LogUniformSampler::Probability(int64_t value) const {
+  // Given f(x) = 1/[(x+1) * log_range_]
+  // The value's  probability  is integral of f(x) from value to (value + 1)
+  // More details:
+  // https://wanghaoshuang.github.io/2017/11/Log-uniform-distribution-sampler
+  return (log((value + 2.0) / (value + 1.0))) / log_range_;  // NOLINT
+}
+
+CustomSampler::CustomSampler(int64_t range,
+                             const float *probabilities,
+                             const int *alias,
+                             const float *alias_probabilities,
+                             unsigned int seed)
+    : Sampler(range, seed) {
+  random_engine_ = phi::GetCPURandomEngine(seed_);
+  real_dist_ = std::make_shared<std::uniform_real_distribution<>>(0, 1);
+  int_dist_ = std::make_shared<std::uniform_int_distribution<>>(0, range);
+
+  alias_probs_ = alias_probabilities;
+  probs_ = probabilities;
+  alias_ = alias;
+}
+
+int64_t CustomSampler::Sample() const {
+  auto index = (*int_dist_)(*random_engine_);
+  auto p = (*real_dist_)(*random_engine_);
+  if (p > alias_probs_[index]) {
+    int alias = alias_[index];
+
+    if (alias == exceptional_val) {
+      LOG(WARNING) << "WARNING: CustomSampler get alias " << exceptional_val;
+      return index;
+    }
+
+    return alias;
+  } else {
+    return index;
+  }
+}
+
+float CustomSampler::Probability(int64_t value) const { return probs_[value]; }
+
+}  // namespace phi::math

paddle/phi/kernels/funcs/math/sampler.h

@@ -0,0 +1,133 @@
+// 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 <cstdint>
+#include <memory>
+#include <random>
+#include <vector>
+
+#include "paddle/phi/core/enforce.h"
+
+namespace phi {
+namespace math {
+
+// TODO(wanghaoshuang): Support for GPU
+
+/**
+ * Sample integers from [0, range).
+ */
+class Sampler {
+ public:
+  explicit Sampler(int64_t range, unsigned int seed = 0UL) : range_(range) {
+    PADDLE_ENFORCE_GT(
+        range,
+        0,
+        common::errors::InvalidArgument(
+            "Range should be greater than 0, but received %d.", range));
+    if (seed == 0) {
+      std::random_device r;
+      seed_ = r();
+    } else {
+      seed_ = seed;
+    }
+  }
+
+  virtual ~Sampler();
+
+  // Sample a single value
+  virtual int64_t Sample() const = 0;
+
+  // The probability that a single call to Sample() returns the given value.
+  virtual float Probability(int64_t value) const = 0;
+
+  int64_t range() { return range_; }
+
+ protected:
+  const int64_t range_;
+  unsigned int seed_;
+};
+
+/**
+ * Sample integers from [0, range).
+ * And the distribution function is:
+ * P(x) = 1 / range
+ */
+class UniformSampler : public Sampler {
+ public:
+  explicit UniformSampler(int64_t range, unsigned int seed = 0UL);
+
+  ~UniformSampler() override {}
+
+  int64_t Sample() const override;
+
+  float Probability(int64_t value) const override;
+
+ private:
+  const float inv_range_;
+  std::shared_ptr<std::mt19937_64> random_engine_;
+  std::shared_ptr<std::uniform_int_distribution<>> dist_;
+};
+
+/**
+ * Sample integers from [0, range).
+ * And the distribution function is:
+ * P(x) = (1/ln(range+1)) * ln(1 + 1/(x + 1))
+ */
+class LogUniformSampler : public Sampler {
+ public:
+  explicit LogUniformSampler(int64_t range, unsigned int seed = 0UL);
+
+  ~LogUniformSampler() override {}
+
+  int64_t Sample() const override;
+
+  float Probability(int64_t value) const override;
+
+ private:
+  const float log_range_;
+  std::shared_ptr<std::mt19937_64> random_engine_;
+  std::shared_ptr<std::uniform_real_distribution<>> dist_;
+};
+
+/**
+ * Sample integers from [0, range) from custom distribution.
+ */
+class CustomSampler : public Sampler {
+ public:
+  explicit CustomSampler(int64_t range,
+                         const float* probabilities,
+                         const int* alias,
+                         const float* alias_probabilities,
+                         unsigned int seed = 0UL);
+
+  ~CustomSampler() override {}
+
+  int64_t Sample() const override;
+
+  float Probability(int64_t value) const override;
+
+ private:
+  const float* alias_probs_;
+  const int* alias_;
+  const float* probs_;
+  const int exceptional_val = -1;
+  std::shared_ptr<std::mt19937_64> random_engine_;
+  std::shared_ptr<std::uniform_real_distribution<>> real_dist_;
+  std::shared_ptr<std::uniform_int_distribution<>> int_dist_;
+};
+
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/tree2col.cc

@@ -0,0 +1,199 @@
+// 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/funcs/math/tree2col.h"
+
+#include <deque>
+#include <stack>
+
+namespace phi {
+namespace math {
+std::vector<TreeNode> Tree2ColUtil::construct_patch(
+    size_t root, int max_depth, const std::vector<std::vector<int>> &tr) {
+  std::stack<TreeNode, std::deque<TreeNode>> stack;
+  std::unordered_map<int, bool> visited;
+  std::vector<TreeNode> patch;
+
+  stack.emplace(root, 1, 1, 0);
+  patch.emplace_back(root, 1, 1, 0);
+  visited[static_cast<int>(root)] = true;
+
+  while (!stack.empty()) {
+    TreeNode &u = stack.top();
+    bool end = true;
+    size_t node = u.get_node(), sz = tr[node].size();
+    visited[static_cast<int>(node)] = true;
+    for (size_t i = 0; i < sz; i++) {
+      size_t v = tr[node][i];
+      if (!visited[static_cast<int>(v)] &&
+          static_cast<int>(u.get_depth()) + 1 < max_depth) {
+        visited[static_cast<int>(v)] = true;
+        stack.emplace(v, i, sz, u.get_depth() + 1);
+        patch.emplace_back(v, i + 1, sz, u.get_depth() + 1);
+        end = false;
+      }
+    }
+    if (end) {
+      stack.pop();
+    }
+  }
+  return patch;
+}
+
+void Tree2ColUtil::construct_tree(const DenseTensor &EdgeSet,
+                                  std::vector<std::vector<int>> *tr,
+                                  size_t *node_count) {
+  const auto &edge_set_dims = EdgeSet.dims();
+  PADDLE_ENFORCE_EQ(edge_set_dims[1],
+                    2,
+                    common::errors::InvalidArgument(
+                        "The second dimension of the EdgeSet shall be 2, but "
+                        "got %ld != 2. Please check the input value.",
+                        edge_set_dims[1]));
+  int64_t edge_count = EdgeSet.numel();
+
+  const int *edge_data = EdgeSet.data<int>();
+
+  for (int64_t i = 0; i < edge_count; i += 2) {
+    int u = edge_data[i], v = edge_data[i + 1];
+    if (u != 0 && v != 0) (*node_count)++;
+  }
+  (*node_count)++;
+
+  tr->resize(static_cast<size_t>(*node_count + 1));
+
+  for (int64_t i = 0; i < edge_count; i += 2) {
+    int u = edge_data[i], v = edge_data[i + 1];
+    if (u != 0 && v != 0) {
+      tr->at(u).push_back(v);
+    } else {
+      break;
+    }
+  }
+}
+
+template <typename T>
+class Tree2ColFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext &dev_ctx,
+                  const DenseTensor &EdgeSet,
+                  const DenseTensor &node_features,
+                  DenseTensor *patch,
+                  int max_depth) {
+    std::vector<std::vector<int>> tr;
+    const auto &feature_dims = node_features.dims();
+    funcs::SetConstant<CPUContext, T> constant;
+    int64_t feature_size = feature_dims[1];
+    size_t patch_elem_size = 3 * static_cast<size_t>(feature_size);
+    size_t node_count = 0, patch_count = 0, patch_size = 0;
+    Tree2ColUtil::construct_tree(EdgeSet, &tr, &node_count);
+    std::vector<std::vector<TreeNode>> processing_list;
+    for (size_t u = 1; u <= node_count; u++) {
+      std::vector<TreeNode> temp_patch =
+          Tree2ColUtil::construct_patch(u, max_depth, tr);
+      if (!temp_patch.empty()) {
+        processing_list.emplace_back(temp_patch);
+      }
+    }
+    patch_size = processing_list.size();
+
+    patch->Resize({static_cast<int64_t>(patch_size),
+                   static_cast<int64_t>(patch_elem_size)});
+    T *patch_data = dev_ctx.template Alloc<T>(patch);
+    constant(dev_ctx, patch, 0);
+    const T *features = node_features.data<T>();
+
+    for (auto &patch_item : processing_list) {
+      size_t pointer_base = patch_count * patch_elem_size;
+      for (auto &v : patch_item) {
+        T eta_l = v.eta_l<T>(max_depth), eta_r = v.eta_r<T>(max_depth),
+          eta_t = v.eta_t<T>(max_depth);
+        size_t id = v.get_node() - 1;
+        for (int i = 0; i < feature_size; i++) {
+          patch_data[pointer_base + i * 3] +=
+              eta_l * features[id * feature_size + i];
+          patch_data[pointer_base + i * 3 + 1] +=
+              eta_r * features[id * feature_size + i];
+          patch_data[pointer_base + i * 3 + 2] +=
+              eta_t * features[id * feature_size + i];
+        }
+      }
+      patch_count++;
+    }
+    patch->Resize({static_cast<int64_t>(patch_count),
+                   static_cast<int64_t>(patch_elem_size)});
+  }
+};
+template <typename T>
+class Col2TreeFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext &dev_ctx,
+                  const DenseTensor &EdgeSet,
+                  const DenseTensor &out_grad,
+                  DenseTensor *in_grad,
+                  int max_depth) {
+    std::vector<std::vector<int>> tr;
+    const auto &output_dims = out_grad.dims();
+    funcs::SetConstant<CPUContext, T> constant;
+    int64_t output_size = output_dims[1];
+    size_t grad_elem_size = 3 * static_cast<size_t>(output_size);
+    size_t node_count = 0, grad_count = 0;
+    Tree2ColUtil::construct_tree(EdgeSet, &tr, &node_count);
+    std::vector<std::vector<TreeNode>> processing_list;
+    std::vector<std::vector<TreeNode>> grad_list;
+    grad_list.resize(node_count);
+    for (size_t u = 1; u <= node_count; u++) {
+      std::vector<TreeNode> tmp =
+          Tree2ColUtil::construct_patch(u, max_depth, tr);
+      if (!tmp.empty()) {
+        processing_list.push_back(tmp);
+      }
+    }
+    for (size_t patch_id = 0; patch_id < processing_list.size(); patch_id++) {
+      for (auto v : processing_list[patch_id]) {
+        grad_list[v.get_node() - 1].push_back(v.change_node(patch_id + 1));
+      }
+    }
+    in_grad->Resize({static_cast<int64_t>(node_count),
+                     static_cast<int64_t>(grad_elem_size)});
+    T *grad_data = dev_ctx.template Alloc<T>(in_grad);
+
+    constant(dev_ctx, in_grad, 0);
+    const T *out_g = out_grad.data<T>();
+    for (auto &patch_item : grad_list) {
+      size_t pointer_base = grad_count * grad_elem_size;
+      for (auto &v : patch_item) {
+        T eta_l = v.eta_l<T>(max_depth), eta_r = v.eta_r<T>(max_depth),
+          eta_t = v.eta_t<T>(max_depth);
+        size_t id = v.get_node() - 1;
+        for (int i = 0; i < output_size; i++) {
+          grad_data[pointer_base + i * 3] +=
+              eta_l * out_g[id * output_size + i];
+          grad_data[pointer_base + i * 3 + 1] +=
+              eta_r * out_g[id * output_size + i];
+          grad_data[pointer_base + i * 3 + 2] +=
+              eta_t * out_g[id * output_size + i];
+        }
+      }
+      grad_count++;
+    }
+  }
+};
+
+template class Tree2ColFunctor<CPUContext, float>;
+template class Tree2ColFunctor<CPUContext, double>;
+template class Col2TreeFunctor<CPUContext, float>;
+template class Col2TreeFunctor<CPUContext, double>;
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/tree2col.cu

@@ -0,0 +1,239 @@
+// 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 <stack>
+
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/math/tree2col.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+namespace math {
+using Node = phi::math::TreeNode;
+template <typename T>
+__global__ void tree2col(const T* eta,
+                         const int* node,
+                         const int* index,
+                         const T* vectors,
+                         T* result,
+                         int feature_size,
+                         int n) {
+  const int thread_id =
+      (blockIdx.x * gridDim.y + blockIdx.y) * blockDim.x + threadIdx.x;
+  const int patch_id = thread_id / feature_size;
+  const int j = thread_id % feature_size;
+  if (patch_id < n) {
+    const int begin_o = patch_id * 3 * feature_size;
+    const int begin = index[patch_id * 2], end = index[patch_id * 2 + 1];
+    T res_l = 0, res_r = 0, res_t = 0;
+    for (int i = begin; i < end; i++) {
+      const int id = node[i];
+      const T vec = vectors[id * feature_size + j];
+      res_l += eta[i * 3] * vec;
+      res_r += eta[i * 3 + 1] * vec;
+      res_t += eta[i * 3 + 2] * vec;
+    }
+    result[begin_o + j * 3] = res_l;
+    result[begin_o + j * 3 + 1] = res_r;
+    result[begin_o + j * 3 + 2] = res_t;
+  }
+}
+template <typename T>
+class Tree2ColFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& dev_ctx,
+                  const DenseTensor& EdgeSet,
+                  const DenseTensor& node_features,
+                  DenseTensor* patch,
+                  int max_depth) {
+    std::vector<std::vector<int>> tr;
+    auto gpu_place = dev_ctx.GetPlace();
+    auto cpu_place = CPUPlace();
+    auto stream = dev_ctx.stream();
+    auto feature_dims = node_features.dims();
+    funcs::SetConstant<GPUContext, T> constant;
+
+    DenseTensor EdgeSet_cpu;
+    phi::Copy(dev_ctx, EdgeSet, cpu_place, false, &EdgeSet_cpu);
+    int64_t feature_size = feature_dims[1];
+    size_t patch_elem_size = 3 * static_cast<size_t>(feature_size);
+    size_t node_count = 0, patch_count = 0, total_size = 0;
+    size_t max_size = feature_dims[0];
+    Tree2ColUtil::construct_tree(EdgeSet_cpu, &tr, &node_count);
+
+    std::vector<std::vector<Node>> processing_list;
+    for (size_t u = 1; u <= node_count; u++) {
+      std::vector<Node> tmp = Tree2ColUtil::construct_patch(u, max_depth, tr);
+      if (!tmp.empty()) {
+        processing_list.push_back(tmp);
+        total_size += tmp.size();
+      }
+    }
+
+    size_t patch_size = processing_list.size();
+    DenseTensor node_cpu, node_gpu, eta_cpu, eta_gpu, index_cpu, index_gpu;
+    node_cpu.Resize({static_cast<int64_t>(total_size)});
+    int* node = dev_ctx.template Alloc<int>(&node_cpu);
+    eta_cpu.Resize({static_cast<int64_t>(total_size * 3)});
+    T* eta = dev_ctx.template Alloc<T>(&eta_cpu);
+    index_cpu.Resize({static_cast<int64_t>(patch_size * 2)});
+    int* index = dev_ctx.template Alloc<int>(&index_cpu);
+
+    int idx = 0, index_idx = 0;
+    for (auto& tmp : processing_list) {
+      index[index_idx++] = idx;
+      for (auto& v : tmp) {
+        node[idx] = static_cast<int>(v.node - 1);
+        eta[idx * 3] = v.eta_l<T>(max_depth);
+        eta[idx * 3 + 1] = v.eta_r<T>(max_depth);
+        eta[idx * 3 + 2] = v.eta_t<T>(max_depth);
+        idx++;
+      }
+      index[index_idx++] = idx;
+    }
+    phi::Copy(dev_ctx, node_cpu, gpu_place, false, &node_gpu);
+    phi::Copy(dev_ctx, eta_cpu, gpu_place, false, &eta_gpu);
+    phi::Copy(dev_ctx, index_cpu, gpu_place, false, &index_gpu);
+
+    PADDLE_ENFORCE_LE_INT_MAX(feature_size, "tree2col kernel feature_size");
+    PADDLE_ENFORCE_LE_INT_MAX(static_cast<int64_t>(patch_size),
+                              "tree2col kernel patch_size");
+    const int feature_size_int = static_cast<int>(feature_size);
+    const int patch_size_int = static_cast<int>(patch_size);
+
+    int64_t elem_size = static_cast<int64_t>(patch_size) * feature_size;
+    int64_t blocks64 = (elem_size + 1024 - 1) / 1024;
+    PADDLE_ENFORCE_LE_INT_MAX(blocks64, "CUDA launch grid blocks");
+    int block_x = 512;
+    int64_t block_y64 = (blocks64 + block_x - 1) / block_x;
+    PADDLE_ENFORCE_LE_INT_MAX(block_y64, "CUDA launch grid.y");
+    int block_y = static_cast<int>(block_y64);
+    dim3 threads(1024, 1);
+    dim3 grid(block_x, block_y);
+
+    patch->Resize({static_cast<int64_t>(max_size),
+                   static_cast<int64_t>(patch_elem_size)});
+    dev_ctx.template Alloc<T>(patch);
+    constant(dev_ctx, patch, 0);
+    tree2col<T><<<grid, threads, 0, stream>>>(eta_gpu.data<T>(),
+                                              node_gpu.data<int>(),
+                                              index_gpu.data<int>(),
+                                              node_features.data<T>(),
+                                              patch->data<T>(),
+                                              feature_size_int,
+                                              patch_size_int);
+  }
+};
+template <typename T>
+class Col2TreeFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& dev_ctx,
+                  const DenseTensor& EdgeSet,
+                  const DenseTensor& patch_grad,
+                  DenseTensor* embedding_grad,
+                  int max_depth) {
+    std::vector<std::vector<int>> tr;
+    auto gpu_place = dev_ctx.GetPlace();
+    auto cpu_place = CPUPlace();
+    auto stream = dev_ctx.stream();
+    auto output_dims = patch_grad.dims();
+    funcs::SetConstant<GPUContext, T> constant;
+
+    DenseTensor EdgeSet_cpu;
+    phi::Copy(dev_ctx, EdgeSet, cpu_place, false, &EdgeSet_cpu);
+    int64_t output_size = output_dims[1];
+    size_t patch_elem_size = 3 * static_cast<size_t>(output_size);
+    size_t node_count = 0, patch_count = 0;
+    size_t max_size = output_dims[0];
+    Tree2ColUtil::construct_tree(EdgeSet_cpu, &tr, &node_count);
+    std::vector<std::vector<Node>> processing_list;
+    std::vector<std::vector<Node>> grad_list;
+    grad_list.resize(node_count);
+    size_t total_size = 0, grad_size = node_count;
+    for (size_t u = 1; u <= node_count; u++) {
+      std::vector<Node> tmp = Tree2ColUtil::construct_patch(u, max_depth, tr);
+      if (!tmp.empty()) {
+        processing_list.push_back(tmp);
+      }
+    }
+    for (size_t patch_id = 0; patch_id < processing_list.size(); patch_id++) {
+      for (auto v : processing_list[patch_id]) {
+        grad_list[v.get_node() - 1].push_back(v.change_node(patch_id + 1));
+      }
+    }
+    for (auto& tmp : grad_list) {
+      total_size += tmp.size();
+    }
+
+    DenseTensor node_cpu, node_gpu, eta_cpu, eta_gpu, index_cpu, index_gpu;
+    node_cpu.Resize({static_cast<int64_t>(total_size)});
+    int* node = dev_ctx.template Alloc<int>(&node_cpu);
+    eta_cpu.Resize({static_cast<int64_t>(total_size * 3)});
+    T* eta = dev_ctx.template Alloc<T>(&eta_cpu);
+    index_cpu.Resize({static_cast<int64_t>(grad_size * 2)});
+    int* index = dev_ctx.template Alloc<int>(&index_cpu);
+
+    size_t idx = 0, index_idx = 0;
+    for (auto& tmp : grad_list) {
+      index[index_idx++] = idx;
+      for (auto& v : tmp) {
+        node[idx] = static_cast<int>(v.node - 1);
+        eta[idx * 3] = v.eta_l<T>(max_depth);
+        eta[idx * 3 + 1] = v.eta_r<T>(max_depth);
+        eta[idx * 3 + 2] = v.eta_t<T>(max_depth);
+        idx++;
+      }
+      index[index_idx++] = idx;
+    }
+    phi::Copy(dev_ctx, node_cpu, gpu_place, false, &node_gpu);
+    phi::Copy(dev_ctx, eta_cpu, gpu_place, false, &eta_gpu);
+    phi::Copy(dev_ctx, index_cpu, gpu_place, false, &index_gpu);
+
+    PADDLE_ENFORCE_LE_INT_MAX(output_size, "tree2col kernel output_size");
+    PADDLE_ENFORCE_LE_INT_MAX(static_cast<int64_t>(grad_size),
+                              "tree2col kernel grad_size");
+    const int output_size_int = static_cast<int>(output_size);
+    const int grad_size_int = static_cast<int>(grad_size);
+
+    int64_t elem_size = static_cast<int64_t>(output_size) * grad_size;
+    int64_t blocks64 = (elem_size + 1024 - 1) / 1024;
+    PADDLE_ENFORCE_LE_INT_MAX(blocks64, "CUDA launch grid blocks");
+    int block_x = 512;
+    int64_t block_y64 = (blocks64 + block_x - 1) / block_x;
+    PADDLE_ENFORCE_LE_INT_MAX(block_y64, "CUDA launch grid.y");
+    int block_y = static_cast<int>(block_y64);
+    dim3 threads(1024, 1);
+    dim3 grid(block_x, block_y);
+
+    embedding_grad->Resize({static_cast<int64_t>(max_size),
+                            static_cast<int64_t>(patch_elem_size)});
+    dev_ctx.template Alloc<T>(embedding_grad);
+
+    constant(dev_ctx, embedding_grad, 0);
+    tree2col<T><<<grid, threads, 0, stream>>>(eta_gpu.data<T>(),
+                                              node_gpu.data<int>(),
+                                              index_gpu.data<int>(),
+                                              patch_grad.data<T>(),
+                                              embedding_grad->data<T>(),
+                                              output_size_int,
+                                              grad_size_int);
+  }
+};
+
+template class Tree2ColFunctor<GPUContext, float>;
+template class Tree2ColFunctor<GPUContext, double>;
+template class Col2TreeFunctor<GPUContext, float>;
+template class Col2TreeFunctor<GPUContext, double>;
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/tree2col.h

@@ -0,0 +1,91 @@
+// 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 <array>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+namespace math {
+class TreeNode {
+ public:
+  size_t node;
+  explicit TreeNode(size_t node = 0,
+                    size_t index = 0,
+                    size_t pclen = 0,
+                    size_t depth = 0)
+      : node(node), index(index), pclen(pclen), depth(depth) {}
+  template <typename T>
+  T eta_t(T filter_depth) {
+    return ((filter_depth - this->depth) / filter_depth);
+  }
+  template <typename T>
+  T eta_l(T filter_depth) {
+    T temp;
+    if (this->pclen == 1) {
+      temp = 0.5;
+    } else {
+      temp = (this->index - 1.0) / (this->pclen - 1.0);
+    }
+    return (1.0 - this->eta_t<T>(filter_depth)) * temp;
+  }
+  template <typename T>
+  T eta_r(T filter_depth) {
+    return (1.0 - this->eta_t<T>(filter_depth)) *
+           (1.0 - this->eta_l<T>(filter_depth));
+  }
+  TreeNode change_node(size_t v) {
+    return TreeNode(v, this->index, this->pclen, this->depth);
+  }
+  size_t get_node() { return this->node; }
+  size_t get_depth() { return this->depth; }
+
+ private:
+  size_t index, pclen, depth;
+};
+class Tree2ColUtil {
+ public:
+  static std::vector<TreeNode> construct_patch(
+      size_t root, int max_depth, const std::vector<std::vector<int>> &tr);
+
+  static void construct_tree(const DenseTensor &EdgeSet,
+                             std::vector<std::vector<int>> *tr,
+                             size_t *node_count);
+};
+
+template <typename DeviceContext, typename T>
+class Tree2ColFunctor {
+ public:
+  void operator()(const DeviceContext &dev_ctx,
+                  const DenseTensor &EdgeSet,
+                  const DenseTensor &node_features,
+                  DenseTensor *patch,
+                  int max_depth);
+};
+template <typename DeviceContext, typename T>
+class Col2TreeFunctor {
+ public:
+  void operator()(const DeviceContext &dev_ctx,
+                  const DenseTensor &EdgeSet,
+                  const DenseTensor &out_grad,
+                  DenseTensor *in_grad,
+                  int max_depth);
+};
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/unpooling.cc

@@ -0,0 +1,216 @@
+// 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/funcs/math/unpooling.h"
+
+namespace phi {
+namespace math {
+template <typename T>
+class Unpool2dMaxFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  DenseTensor* output) {
+    const int batch_size = static_cast<int>(input.dims()[0]);
+    const int input_height = static_cast<int>(input.dims()[2]);
+    const int input_width = static_cast<int>(input.dims()[3]);
+    const int output_channels = static_cast<int>(output->dims()[1]);
+    const int output_height = static_cast<int>(output->dims()[2]);
+    const int output_width = static_cast<int>(output->dims()[3]);
+    int64_t input_feasize = static_cast<int64_t>(input_height) * input_width;
+    int64_t output_feasize = static_cast<int64_t>(output_height) * output_width;
+    const T* input_data = input.data<T>();
+    const int* indices_data = indices.data<int>();
+    T* output_data = context.template Alloc<T>(output);
+    for (int b = 0; b < batch_size; ++b) {
+      for (int c = 0; c < output_channels; ++c) {
+        for (int64_t i = 0; i < input_feasize; ++i) {
+          int index = indices_data[i];
+
+          PADDLE_ENFORCE_LT(
+              index,
+              output_feasize,
+              common::errors::InvalidArgument(
+                  "index should less than output tensor height * output tensor "
+                  "width. Expected %ld < %ld, but got "
+                  "%ld >= %ld. Please check input value.",
+                  index,
+                  output_feasize,
+                  index,
+                  output_feasize));
+          output_data[index] = input_data[i];
+        }
+        input_data += input_feasize;
+        indices_data += input_feasize;
+        output_data += output_feasize;
+      }
+    }
+  }
+};
+template <class T>
+class Unpool2dMaxGradFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  const DenseTensor& output,
+                  const DenseTensor& output_grad,
+                  DenseTensor* input_grad) {
+    const int batch_size = static_cast<int>(input.dims()[0]);
+    const int input_height = static_cast<int>(input.dims()[2]);
+    const int input_width = static_cast<int>(input.dims()[3]);
+    const int output_channels = static_cast<int>(output.dims()[1]);
+    const int output_height = static_cast<int>(output.dims()[2]);
+    const int output_width = static_cast<int>(output.dims()[3]);
+    int64_t input_feasize = static_cast<int64_t>(input_height) * input_width;
+    int64_t output_feasize = static_cast<int64_t>(output_height) * output_width;
+    const int* indices_data = indices.data<int>();
+    const T* output_grad_data = output_grad.data<T>();
+    T* input_grad_data = context.template Alloc<T>(input_grad);
+
+    for (int b = 0; b < batch_size; ++b) {
+      for (int c = 0; c < output_channels; ++c) {
+        for (int64_t i = 0; i < input_feasize; ++i) {
+          int index = indices_data[i];
+          PADDLE_ENFORCE_LT(
+              index,
+              output_feasize,
+              common::errors::InvalidArgument(
+                  "index should less than output tensor height * output tensor "
+                  "width. Expected %ld < %ld, but got "
+                  "%ld >= %ld. Please check input value.",
+                  index,
+                  output_feasize,
+                  index,
+                  output_feasize));
+          input_grad_data[i] = output_grad_data[index];
+        }
+        input_grad_data += input_feasize;
+        indices_data += input_feasize;
+        output_grad_data += output_feasize;
+      }
+    }
+  }
+};
+
+template <typename T>
+class Unpool3dMaxFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  DenseTensor* output) {
+    const int batch_size = static_cast<int>(input.dims()[0]);
+    const int input_depth = static_cast<int>(input.dims()[2]);
+    const int input_height = static_cast<int>(input.dims()[3]);
+    const int input_width = static_cast<int>(input.dims()[4]);
+    const int output_channels = static_cast<int>(output->dims()[1]);
+    const int output_depth = static_cast<int>(output->dims()[2]);
+    const int output_height = static_cast<int>(output->dims()[3]);
+    const int output_width = static_cast<int>(output->dims()[4]);
+    int64_t input_feasize =
+        static_cast<int64_t>(input_depth) * input_height * input_width;
+    int64_t output_feasize =
+        static_cast<int64_t>(output_depth) * output_height * output_width;
+    const T* input_data = input.data<T>();
+    const int* indices_data = indices.data<int>();
+    T* output_data = context.template Alloc<T>(output);
+    for (int b = 0; b < batch_size; ++b) {
+      for (int c = 0; c < output_channels; ++c) {
+        for (int64_t i = 0; i < input_feasize; ++i) {
+          int index = indices_data[i];
+
+          PADDLE_ENFORCE_LT(
+              index,
+              output_feasize,
+              common::errors::InvalidArgument(
+                  "index should less than output tensor depth * output tensor "
+                  "height "
+                  "* output tensor width. Expected %ld < %ld, but got "
+                  "%ld >= %ld. Please check input value.",
+                  index,
+                  output_feasize,
+                  index,
+                  output_feasize));
+          output_data[index] = input_data[i];
+        }
+        input_data += input_feasize;
+        indices_data += input_feasize;
+        output_data += output_feasize;
+      }
+    }
+  }
+};
+template <class T>
+class Unpool3dMaxGradFunctor<CPUContext, T> {
+ public:
+  void operator()(const CPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  const DenseTensor& output,
+                  const DenseTensor& output_grad,
+                  DenseTensor* input_grad) {
+    const int batch_size = static_cast<int>(input.dims()[0]);
+    const int input_depth = static_cast<int>(input.dims()[2]);
+    const int input_height = static_cast<int>(input.dims()[3]);
+    const int input_width = static_cast<int>(input.dims()[4]);
+    const int output_channels = static_cast<int>(output.dims()[1]);
+    const int output_depth = static_cast<int>(output.dims()[2]);
+    const int output_height = static_cast<int>(output.dims()[3]);
+    const int output_width = static_cast<int>(output.dims()[4]);
+    int64_t input_feasize =
+        static_cast<int64_t>(input_depth) * input_height * input_width;
+    int64_t output_feasize =
+        static_cast<int64_t>(output_depth) * output_height * output_width;
+    const int* indices_data = indices.data<int>();
+    const T* output_grad_data = output_grad.data<T>();
+    T* input_grad_data = context.template Alloc<T>(input_grad);
+
+    for (int b = 0; b < batch_size; ++b) {
+      for (int c = 0; c < output_channels; ++c) {
+        for (int64_t i = 0; i < input_feasize; ++i) {
+          int index = indices_data[i];
+          PADDLE_ENFORCE_LT(
+              index,
+              output_feasize,
+              common::errors::InvalidArgument(
+                  "index should less than output tensor depth * output tensor "
+                  "height "
+                  "* output tensor width. Expected %ld < %ld, but got "
+                  "%ld >= %ld. Please check input value.",
+                  index,
+                  output_feasize,
+                  index,
+                  output_feasize));
+          input_grad_data[i] = output_grad_data[index];
+        }
+        input_grad_data += input_feasize;
+        indices_data += input_feasize;
+        output_grad_data += output_feasize;
+      }
+    }
+  }
+};
+
+template class Unpool2dMaxGradFunctor<CPUContext, float>;
+template class Unpool2dMaxGradFunctor<CPUContext, double>;
+template class Unpool2dMaxFunctor<CPUContext, float>;
+template class Unpool2dMaxFunctor<CPUContext, double>;
+template class Unpool3dMaxGradFunctor<CPUContext, float>;
+template class Unpool3dMaxGradFunctor<CPUContext, double>;
+template class Unpool3dMaxFunctor<CPUContext, float>;
+template class Unpool3dMaxFunctor<CPUContext, double>;
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/unpooling.cu

@@ -0,0 +1,361 @@
+// 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/funcs/math/unpooling.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+
+namespace phi {
+namespace math {
+template <typename T>
+__global__ void KernelUnpool2dMax(const int64_t nthreads,
+                                  const T* input_data,
+                                  const int* indices_data,
+                                  const int input_height,
+                                  const int input_width,
+                                  const int channels,
+                                  T* output_data,
+                                  const int output_height,
+                                  const int output_width) {
+  CUDA_KERNEL_LOOP_TYPE(linearIndex, nthreads, int64_t) {
+    int64_t c = (linearIndex / input_width / input_height) % channels;
+    int64_t n = linearIndex / input_width / input_height / channels;
+    output_data += (n * channels + c) * output_height * output_width;
+    int maxind = indices_data[linearIndex];
+    output_data[maxind] = input_data[linearIndex];
+  }
+}
+
+template <typename T>
+__global__ void KernelUnpool2dMaxGrad(const int64_t nthreads,
+                                      const T* input_data,
+                                      const int* indices_data,
+                                      const int input_height,
+                                      const int input_width,
+                                      const int channels,
+                                      const T* output_data,
+                                      const T* output_grad,
+                                      const int output_height,
+                                      const int output_width,
+                                      T* input_grad) {
+  CUDA_KERNEL_LOOP_TYPE(linearIndex, nthreads, int64_t) {
+    int64_t c = (linearIndex / input_width / input_height) % channels;
+    int64_t n = linearIndex / input_width / input_height / channels;
+    output_grad += (n * channels + c) * output_height * output_width;
+    int maxind = indices_data[linearIndex];
+    input_grad[linearIndex] = output_grad[maxind];
+  }
+}
+/*
+ * All tensors are in NCHW format.
+ */
+
+template <typename T>
+__global__ void KernelUnpool3dMax(const int64_t nthreads,
+                                  const T* input_data,
+                                  const int* indices_data,
+                                  const int input_depth,
+                                  const int input_height,
+                                  const int input_width,
+                                  const int channels,
+                                  T* output_data,
+                                  const int output_depth,
+                                  const int output_height,
+                                  const int output_width) {
+  CUDA_KERNEL_LOOP_TYPE(linearIndex, nthreads, int64_t) {
+    int64_t c =
+        (linearIndex / input_depth / input_width / input_height) % channels;
+    int64_t n =
+        linearIndex / input_depth / input_width / input_height / channels;
+    output_data +=
+        (n * channels + c) * output_depth * output_height * output_width;
+    int maxind = indices_data[linearIndex];
+    output_data[maxind] = input_data[linearIndex];
+  }
+}
+
+template <typename T>
+__global__ void KernelUnpool3dMaxGrad(const int64_t nthreads,
+                                      const T* input_data,
+                                      const int* indices_data,
+                                      const int input_depth,
+                                      const int input_height,
+                                      const int input_width,
+                                      const int channels,
+                                      const T* output_data,
+                                      const T* output_grad,
+                                      const int output_depth,
+                                      const int output_height,
+                                      const int output_width,
+                                      T* input_grad) {
+  CUDA_KERNEL_LOOP_TYPE(linearIndex, nthreads, int64_t) {
+    int64_t c =
+        (linearIndex / input_depth / input_width / input_height) % channels;
+    int64_t n =
+        linearIndex / input_depth / input_width / input_height / channels;
+    output_grad +=
+        (n * channels + c) * output_depth * output_height * output_width;
+    int maxind = indices_data[linearIndex];
+    input_grad[linearIndex] = output_grad[maxind];
+  }
+}
+/*
+ * All tensors are in NCDHW format.
+ */
+
+template <typename T>
+class Unpool2dMaxFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  DenseTensor* output) {
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t batch_size = input.dims()[0];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_height = input.dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_width = input.dims()[3];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_channels = output->dims()[1];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_height = output->dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_width = output->dims()[3];
+
+    const T* input_data = input.data<T>();
+    const int* indices_data = indices.data<int>();
+    T* output_data = context.template Alloc<T>(output);
+    int threads = 1024;
+    int64_t max_grid = context.GetCUDAMaxGridDimSize()[0];
+    int grid = std::min((input.numel() + threads - 1) / threads, max_grid);
+    KernelUnpool2dMax<T>
+        <<<grid, threads, 0, context.stream()>>>(input.numel(),
+                                                 input_data,
+                                                 indices_data,
+                                                 input_height,
+                                                 input_width,
+                                                 output_channels,
+                                                 output_data,
+                                                 output_height,
+                                                 output_width);
+  }
+};
+/*
+ * All tensors are in NCHW format.
+ */
+template <typename T>
+class Unpool2dMaxGradFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  const DenseTensor& output,
+                  const DenseTensor& output_grad,
+                  DenseTensor* input_grad) {
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t batch_size = input.dims()[0];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_height = input.dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_width = input.dims()[3];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_channels = output.dims()[1];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_height = output.dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_width = output.dims()[3];
+
+    const T* input_data = input.data<T>();
+    const int* indices_data = indices.data<int>();
+    const T* output_data = output.data<T>();
+    const T* output_grad_data = output_grad.data<T>();
+    T* input_grad_data = context.template Alloc<T>(input_grad);
+    int threads = 1024;
+    int64_t max_grid = context.GetCUDAMaxGridDimSize()[0];
+    int grid = std::min((input.numel() + threads - 1) / threads, max_grid);
+    KernelUnpool2dMaxGrad<T>
+        <<<grid, threads, 0, context.stream()>>>(input.numel(),
+                                                 input_data,
+                                                 indices_data,
+                                                 input_height,
+                                                 input_width,
+                                                 output_channels,
+                                                 output_data,
+                                                 output_grad_data,
+                                                 output_height,
+                                                 output_width,
+                                                 input_grad_data);
+  }
+};
+
+template <typename T>
+class Unpool3dMaxFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  DenseTensor* output) {
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t batch_size = input.dims()[0];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_depth = input.dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_height = input.dims()[3];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_width = input.dims()[4];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_channels = output->dims()[1];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_depth = output->dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_height = output->dims()[3];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_width = output->dims()[4];
+
+    const T* input_data = input.data<T>();
+    const int* indices_data = indices.data<int>();
+    T* output_data = context.template Alloc<T>(output);
+    int threads = 1024;
+    int64_t max_grid = context.GetCUDAMaxGridDimSize()[0];
+    int grid = std::min((input.numel() + threads - 1) / threads, max_grid);
+    KernelUnpool3dMax<T>
+        <<<grid, threads, 0, context.stream()>>>(input.numel(),
+                                                 input_data,
+                                                 indices_data,
+                                                 input_depth,
+                                                 input_height,
+                                                 input_width,
+                                                 output_channels,
+                                                 output_data,
+                                                 output_depth,
+                                                 output_height,
+                                                 output_width);
+  }
+};
+/*
+ * All tensors are in NCDHW format.
+ */
+template <typename T>
+class Unpool3dMaxGradFunctor<GPUContext, T> {
+ public:
+  void operator()(const GPUContext& context,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  const DenseTensor& output,
+                  const DenseTensor& output_grad,
+                  DenseTensor* input_grad) {
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t batch_size = input.dims()[0];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_depth = input.dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_height = input.dims()[3];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t input_width = input.dims()[4];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_channels = output.dims()[1];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_depth = output.dims()[2];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_height = output.dims()[3];
+
+    // TODO(large-tensor): downstream functors may still use int; guard until
+    // upgraded.
+    int64_t output_width = output.dims()[4];
+
+    const T* input_data = input.data<T>();
+    const int* indices_data = indices.data<int>();
+    const T* output_data = output.data<T>();
+    const T* output_grad_data = output_grad.data<T>();
+    T* input_grad_data = context.template Alloc<T>(input_grad);
+    int threads = 1024;
+    int64_t max_grid = context.GetCUDAMaxGridDimSize()[0];
+    int grid = std::min((input.numel() + threads - 1) / threads, max_grid);
+    KernelUnpool3dMaxGrad<T>
+        <<<grid, threads, 0, context.stream()>>>(input.numel(),
+                                                 input_data,
+                                                 indices_data,
+                                                 input_depth,
+                                                 input_height,
+                                                 input_width,
+                                                 output_channels,
+                                                 output_data,
+                                                 output_grad_data,
+                                                 output_depth,
+                                                 output_height,
+                                                 output_width,
+                                                 input_grad_data);
+  }
+};
+
+template class Unpool2dMaxGradFunctor<GPUContext, float>;
+template class Unpool2dMaxGradFunctor<GPUContext, double>;
+template class Unpool2dMaxFunctor<GPUContext, float>;
+template class Unpool2dMaxFunctor<GPUContext, double>;
+template class Unpool3dMaxGradFunctor<GPUContext, float>;
+template class Unpool3dMaxGradFunctor<GPUContext, double>;
+template class Unpool3dMaxFunctor<GPUContext, float>;
+template class Unpool3dMaxFunctor<GPUContext, double>;
+}  // namespace math
+}  // namespace phi

paddle/phi/kernels/funcs/math/unpooling.h

@@ -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.
+
+#pragma once
+
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+namespace math {
+template <typename DeviceContext, typename T>
+class Unpool2dMaxFunctor {
+ public:
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  DenseTensor* output);
+};
+template <typename DeviceContext, class T>
+class Unpool2dMaxGradFunctor {
+ public:
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  const DenseTensor& output,
+                  const DenseTensor& output_grad,
+                  DenseTensor* input_grad);
+};
+
+template <typename DeviceContext, typename T>
+class Unpool3dMaxFunctor {
+ public:
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  DenseTensor* output);
+};
+template <typename DeviceContext, class T>
+class Unpool3dMaxGradFunctor {
+ public:
+  void operator()(const DeviceContext& dev_ctx,
+                  const DenseTensor& input,
+                  const DenseTensor& indices,
+                  const DenseTensor& output,
+                  const DenseTensor& output_grad,
+                  DenseTensor* input_grad);
+};
+}  // namespace math
+}  // namespace phi