<commit_before>/* Copyright (c) 2016 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/fluid/operators/math/beam_search.h"
#include "paddle/fluid/platform/cuda_device_function.h"

namespace paddle {
namespace operators {
namespace math {

struct Triple {
  __device__ __forceinline__ Triple() {}
  __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 {
        int index = 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();
    }

    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;<commit_msg>Fix error in CUDA kernel of beam_search. (#15957)<commit_after>/* Copyright (c) 2016 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/fluid/operators/math/beam_search.h"
#include "paddle/fluid/platform/cuda_device_function.h"

namespace paddle {
namespace operators {
namespace math {

struct Triple {
  __device__ __forceinline__ Triple() {}
  __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 {
        int index = 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;
{'fraction_non_alphanumeric': 0.10781303309450699, 'fraction_numerical': 0.008074604799272148, 'mean_word_length': 2.662640566430654, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5439582', 'n_tokens_mistral': 3091, 'n_tokens_neox': 2930, 'n_words': 845}