/*************************************************************************
 * 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_UPDATE_OUT_AND_LSE_MLUH_
#define CSRC_KERNELS_UPDATE_OUT_AND_LSE_MLUH_

#include "cnnl.h"
#include "kernel_utils.h"
namespace tmo {
/**
 * @brief Update out and log-sum-exp(lse) according to block out and block lse.
 * @param queue: The queue for mlu.
 * @param out: Input/Output. Pointer to the MLU memory that stores the origin out.
 *        In pad mode, the shape must be [batch, max_seq_len, head_num, head_size].
 *        In pack mode, the shape must be [total_seq_len, head_num, head_size].
 *        Dim of seq_len and head_num may have stride.
 * @param lse: Input/Output. Pointer to the MLU memory that stores the origin lse.
 *        The shape must be [batch, head_num, max_seq_len]. Lse must be continuous.
 * @param block_out: Input. Pointer to the MLU memory that stores the block out.
 *        In pad mode, the shape must be [batch, block_seq_len, head_num, head_size].
 *        In pack mode, the shape must be [total_block_seq_len, head_num, head_size].
 *        Dim of seq_len and head_num may have stride.
 * @param block_lse: Input. Pointer to the MLU memory that stores the origin lse.
 *        The shape must be [batch, head_num, block_seq_len]. Block_lse must be continuous.
 * @param seq_offsets: Input. Pointer to the MLU memory that stores the origin out
 *        and lse sequence offset. The shape must be [batch].
 *        Seq_offsets must be continuous, and could be nullptr.
 * @param cu_seqs: Input. Pointer to the MLU memory that stores the cumulative sum of out seq_lens,
 *        In pack mode, the shape must be [batch + 1].
 *        In pad mode. cu_seqs does not work, could be nullptr.
 * @param block_cu_seqs: Input. Pointer to the MLU memory that stores the cumulative sum of block
 * out seq_lens, In pack mode, the shape must be [batch + 1]. In pad mode. block_cu_seqs does not
 * work, could be nullptr.
 * @param dtype: Data type.
 * @param batch: Batch size.
 * @param head_num: Head number.
 * @param head_size: Head size.
 * @param max_seq_len: The sequence length of origin out.
 * @param block_seq_len: The sequence length of block out.
 * @param bs_stride: The stride of batch in origin out, does not work when packed is true.
 * @param seq_stride: The stride of seq_len in origin out.
 * @param head_stride: The stride of head_num in origin out.
 * @param block_bs_stride: The stride of batch in block out, does not work when packed is true.
 * @param block_seq_stride: The stride of seq_len in block out.
 * @param block_head_stride: The stride of head_num in block out.
 * @param packed: A boolean value indicates whether to use pack mode.
 * @note All seq_lens in block out should be less than or equal to origin out.
 */
KernelStatus invokeUpdateOutAndLse(cnrtQueue_t queue,
                                   void *out,
                                   float *lse,
                                   const void *block_out,
                                   const float *block_lse,
                                   const int *seq_offsets,
                                   const int *cu_seqs,
                                   const int *block_cu_seqs,
                                   const int batch,
                                   const int head_num,
                                   const int head_size,
                                   const int max_seq_len,
                                   const int block_seq_len,
                                   const int64_t bs_stride,
                                   const int64_t seq_stride,
                                   const int64_t head_stride,
                                   const int64_t block_bs_stride,
                                   const int64_t block_seq_stride,
                                   const int64_t block_head_stride,
                                   const bool packed,
                                   const cnnlDataType_t dtype);
}  // namespace tmo

#endif  // CSRC_KERNELS_RESHAPE_LINEAR_CACHE_MLUH_
