/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#ifndef CSRC_KERNELS_GENERATE_ALIBI_SLOPE_MLUH_
#define CSRC_KERNELS_GENERATE_ALIBI_SLOPE_MLUH_

#include "cnnl.h"
#include "kernel_utils.h"
namespace tmo {
/**
 * @brief Generate causal mask for context satge.
 * @param queue: The queue for mlu.
 * @param alibi_slopes: Output. Pointer to the MLU memory that stores the output, the shape must be
 * [batch_num, head_num].
 * @param true_seq_lens: Input. Pointer to the MLU memory that stores the actual sequence length of
 * each batch, the shape must be [batch_num].
 * @param batch_num: Batch number.
 * @param head_start: The index of first head.
 * @param head_num: Head number in this card.
 * @param head_num_total: Total head number in all cards.
 * @param max_sequence_length: The maximum sequence length used during training.
 * @param use_dynamic: A boolean value indicates whether to use dynamic NTK.
 */
KernelStatus invokeGenerateAlibiSlope(cnrtQueue_t queue,
                                      void *alibi_slopes,
                                      void *true_seq_lens,
                                      int batch_num,
                                      int head_start,
                                      int head_num,
                                      int head_num_total,
                                      int max_sequence_length,
                                      bool use_dynamic);
}  // namespace tmo

#endif  // CSRC_KERNELS_GENERATE_ALIBI_SLOPE_MLUH_