chore: import upstream snapshot with attribution

paddle/phi/kernels/impl/beam_search_decode_kernel_impl.h

@@ -0,0 +1,167 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/kernels/funcs/beam_search_decode.h"
+
+namespace phi {
+
+struct BeamSearchDecodeFunctor {
+  BeamSearchDecodeFunctor(const TensorArray& step_ids,
+                          const TensorArray& step_scores,
+                          DenseTensor* id_tensor,
+                          DenseTensor* score_tensor,
+                          size_t beam_size,
+                          int end_id)
+      : beam_size_(beam_size),
+        end_id_(end_id),
+        step_ids_origin_(step_ids),
+        step_scores_origin_(step_scores),
+        id_tensor_(id_tensor),
+        score_tensor_(score_tensor) {
+    tensor_on_gpu_ = false;
+    // First make a copy of GPU data on CPU
+    if (step_ids_origin_[0].place().GetType() == AllocationType::GPU ||
+        step_ids_origin_[0].place().GetType() == AllocationType::CUSTOM) {
+      if (step_ids_origin_[0].place().GetType() == AllocationType::GPU ||
+          step_ids_origin_[0].place().GetType() == AllocationType::CUSTOM) {
+        tensor_on_gpu_ = true;
+      }
+      DeviceContextPool& pool = DeviceContextPool::Instance();
+      auto* dev_ctx = pool.Get(step_ids_origin_[0].place());
+      // Copy all tensors in the input tensor array
+      for (auto& step_id : step_ids_origin_) {
+        DenseTensor out;
+        if (step_id.numel() > 0) {
+          if (tensor_on_gpu_) {
+            dev_ctx->Wait();
+          }
+          Copy(*dev_ctx, step_id, CPUPlace(), false, &out);
+          dev_ctx->Wait();
+        }
+
+        out.set_lod(step_id.lod());
+        step_ids_.push_back(out);
+      }
+    }
+    if (step_scores_origin_[0].place().GetType() == AllocationType::GPU ||
+        step_scores_origin_[0].place().GetType() == AllocationType::CUSTOM) {
+      if (step_scores_origin_[0].place().GetType() == AllocationType::GPU ||
+          step_scores_origin_[0].place().GetType() == AllocationType::CUSTOM) {
+        tensor_on_gpu_ = true;
+      }
+      DeviceContextPool& pool = DeviceContextPool::Instance();
+      auto* dev_ctx = pool.Get(step_scores_origin_[0].place());
+      // Copy all tensors in the input tensor array
+      for (auto& step_score : step_scores_origin_) {
+        DenseTensor out;
+        if (step_score.numel() > 0) {
+          if (tensor_on_gpu_) {
+            dev_ctx->Wait();
+          }
+          Copy(*dev_ctx, step_score, CPUPlace(), false, &out);
+          dev_ctx->Wait();
+        }
+
+        out.set_lod(step_score.lod());
+        step_scores_.push_back(out);
+      }
+    }
+  }
+
+  template <typename T>
+  void apply_mix() const {
+    if (std::is_same<bool, T>::value) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "beam search decode op does not support bool!"));
+
+    } else {
+      funcs::BeamSearchDecoder<T> beam_search_decoder(beam_size_, end_id_);
+      // Check if the tensor is on GPU. If so, use the CPU copy instead
+      if (tensor_on_gpu_) {
+        beam_search_decoder.Backtrace(
+            step_ids_, step_scores_, id_tensor_, score_tensor_);
+      } else {
+        beam_search_decoder.Backtrace(
+            step_ids_origin_, step_scores_origin_, id_tensor_, score_tensor_);
+      }
+    }
+  }
+
+  bool tensor_on_gpu_;
+  size_t beam_size_;
+  int end_id_;
+  // TODO(Superjomn) Here might result serious performance issue in the
+  // concurrency
+  // scenarios.
+  const TensorArray& step_ids_origin_;
+  const TensorArray& step_scores_origin_;
+  TensorArray step_ids_ = TensorArray();
+  TensorArray step_scores_ = TensorArray();
+  DenseTensor* id_tensor_;
+  DenseTensor* score_tensor_;
+};
+
+template <typename T, typename Context>
+void BeamSearchDecodeOpKernel(const Context& dev_ctx,
+                              const TensorArray& ids_in,
+                              const TensorArray& scores_in,
+                              int beam_size,
+                              int end_id,
+                              DenseTensor* sentence_ids,
+                              DenseTensor* sentence_scores) {
+  const TensorArray* ids = &ids_in;
+  const TensorArray* scores = &scores_in;
+  const size_t step_num = ids->size();
+  PADDLE_ENFORCE_GT(
+      step_num,
+      0UL,
+      common::errors::InvalidArgument(
+          "beam search steps, which is the "
+          "size of Input(Ids) TensorArray. beam search steps should "
+          "be larger than 0, but received %d. ",
+          step_num));
+  const size_t source_num = ids->at(0).lod().at(0).size() - 1;
+  PADDLE_ENFORCE_GT(
+      source_num,
+      0UL,
+      common::errors::InvalidArgument(
+          "source_num is the sequence number of the "
+          "first decoding step, indicating by Input(Ids)[0].lod[0].size. "
+          "The number of source_num should be larger than "
+          "0, but received %d. ",
+          source_num));
+
+  for (size_t i = 0; i < step_num; ++i) {
+    size_t tmp = ids->at(i).lod().size();
+    PADDLE_ENFORCE_EQ(
+        tmp,
+        2UL,
+        common::errors::InvalidArgument(
+            "For the i step in beam search steps,"
+            "the size of Input(Ids)[i].lod() should larger than 2,"
+            "but received %d. ",
+            tmp));
+  }
+
+  // prepare output
+  DenseTensor* sentenceIds = sentence_ids;
+  DenseTensor* sentenceScores = sentence_scores;
+  BeamSearchDecodeFunctor bs(
+      *ids, *scores, sentenceIds, sentenceScores, beam_size, end_id);
+  bs.apply_mix<T>();
+}
+}  // namespace phi