/*************************************************************************
 * 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_OPERATE_CU_SEQ_LENS_MLUH_
#define CSRC_KERNELS_OPERATE_CU_SEQ_LENS_MLUH_

#include "cnnl.h"
#include "kernel_utils.h"

namespace tmo {

/**
 * @brief slice cu_seq_lens and cu_k_seq_lens for parallel context when attention split batch;
 * @example
 * cu_seq_lens: [0, 2, 5, 10, 20, 33, 46, 51, 77]
 * batch: 8, parallel_num: 3
 * sliced_cu_seq_lens: [0, 2, 5, 10, 0, 10, 23, 36, 0, 5, 31]
 * @param queue: The queue for mlu.
 * @param cu_seq_lens: Input. Pointer to the MLU memory that stores the current seq lens, the shape
 * is [batch + 1].
 * @param sliced_cu_seq_lens: Output. Pointer to the MLU memory that stores the sliced current seq
 * lens, the shape is [batch + loop_time].
 * @param batch: Batch size.
 * @param parallel_num: Parallel num of batch.
 */
KernelStatus invokeSliceCuSeqlens(cnrtQueue_t queue,
                                  int *cu_seq_lens,
                                  int *sliced_cu_seq_lens,
                                  int batch,
                                  int parallel_num);

/**
 * @brief generate cu_seq_lens and cu_k_seq_lens for parallel context when attention split seq;
 * @example
 * seq_len: 11, parallel_num: 3
 * gen_cu_seq_lens for q: [0, 4, 0, 4, 0, 3]
 * @example
 * seq_len: 11, parallel_num: 3
 * is_causal_mask false, gen_cu_seq_lens for kv: [0, 11, 0, 11, 0, 11]
 * is_causal_mask true , gen_cu_seq_lens for kv: [0,  4, 0,  8, 0, 11]
 * @param queue: The queue for mlu.
 * @param gen_cu_seq_lens: Output. Pointer to the MLU memory that stores the generated current seq
 * lens, the shape is [2 * loop_time].
 * @param seq_len: Sequence length.
 * @param parallel_num: Parallel num of sequence length.
 * @param is_causal_mask: Whether self attention use causal mask.
 * @param is_kv_seq_len: The gen_cu_seq_lens is for q or kv.
 */
KernelStatus invokeGenerateCuSeqlens(cnrtQueue_t queue,
                                     int *gen_cu_seq_lens,
                                     int seq_len,
                                     int parallel_num,
                                     bool is_causal_mask,
                                     bool is_kv_seq_len);

}  // namespace tmo

#endif  // CSRC_KERNELS_OPERATE_CU_SEQ_LENS_MLUH_
