/*************************************************************************
 * 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_OFFLINE_QUANT_TO_LINEAR_CACHE_MLUH_
#define CSRC_KERNELS_OFFLINE_QUANT_TO_LINEAR_CACHE_MLUH_

#include "kernel_utils.h"

namespace tmo {
/**
 * @brief Quantize current key and value, Then store key and value to key_cache and value_cache.
 * @param queue: The queue for mlu.
 * @param key_cache: Pointer to the MLU memory that stores the key cache,
 *        the shape must be [max_batch, head_num, cache_mem_len, head_size].
 *        Data type of key_cache must be int8. key_cache could be nullptr.
 * @param value_cache: Pointer to the MLU memory that stores the value cache,
 *        the shape must be [max_batch, head_num, cache_mem_len, head_size].
 *        Data type of value_cache must be int8. value_cache could be nullptr.
 * @param key_cache_scale: Pointer to the MLU memory that stores the key cache scale,
 *        the shape must be [head_num, cache_mem_len] when quant_mode is not zero,
 *        and [head_num, head_size] when quant_mode is zero. Data type of key_cache_scale
 *        must be float. value_cache could be nullptr.
 * @param value_cache_scale: Pointer to the MLU memory that stores the value cache scale,
 *        the shape must be [head_num, cache_mem_len] when quant_mode is not zero,
 *        and [head_num, head_size] when quant_mode is zero. Data type of value_cache_scale
 *        must be float. value_cache_scale could be nullptr.
 * @param cache_bs_offsets: Pointer to the MLU memory that stores the batch
 *        offset of cache, the shape must be [batch], if it's nullptr, the
 *        default value is {0, 1, 2 ... batch - 1}.
 * @param cache_seq_offsets: Pointer to the MLU memory that stores the sequence
 *        offset of cache, the shape must be [batch], if it's nullptr, the
 *        default value is 0 for every batch.
 * @param key: Pointer to the MLU memory that stores the key,
 *        the shape must be [batch, max_contxt_len, head_num, head_size].
 *        Data type of key couble be float/half/bfloat16. key could be nullptr.
 * @param value: Pointer to the MLU memory that stores the value,
 *        the shape must be [batch, max_contxt_len, head_num, head_size].
 *        Data type of value couble be float/half/bfloat16, value could be nullptr.
 * @param context_seq_offsets: Pointer to the MLU memory that stores the
 *        sequence offset of context, the shape must be [batch]. if it's nullptr,
 *        the default value is 0 for every batch. It must be nullptr when packed is true.
 * @param context_lens: Pointer to the MLU memory that stores the sequence length or cumulative
 *        sequence length of context. when packed is false, the shape must be [batch], which
 *        indicates sequence length of context. when packed is true, the shape must be [batch + 1],
 which
 *        indicates cumulative sequence length of context.
 * @param dtype: Data type.
 * @param batch: Batch size.
 * @param head_num: Head number.
 * @param head_size: Head size.
 * @param max_contxt_len: The maximum sequence length of context.
 * @param cache_mem_len: The maximum sequence length of cache.
 * @param contxt_bs_stride: The stride of batch in context, does not work when packed is true.
 * @param contxt_head_stride: The stride of head_num in context.
 * @param contxt_seq_stride: The stride of max_contxt_len in context.
 * @param cache_bs_stride: The stride of batch in cache.
 * @param cache_head_stride: The stride of head_num in cache.
 * @param cache_seq_stride: The stride of cache_mem_len in cache.
 * @param cache_scale_bs_stride: The stride of batch in cache scale.
 * @param cache_scale_head_stride: The stride of head in cache scale.
 * @param packed: A boolean value indicates whether to use pack mode.
 * @param quant_mode: A int value indicates the quantify mode, 0 means quantify by per_channel, and
 others value means quantify by per_head.
 * @note If one of key/key_cache/key_cache_scale is nullptr, nothing todo for key.
         If one of value/value_cache/value_cache_scale is nullptr, nothing todo for value.
 */
KernelStatus invokeOfflineQuantToLinearCache(cnrtQueue_t queue,
                                             void *key_cache,
                                             void *value_cache,
                                             const void *key_cache_scale,
                                             const void *value_cache_scale,
                                             const void *cache_bs_offsets,
                                             const void *cache_seq_offsets,
                                             const void *key,
                                             const void *value,
                                             const void *context_seq_offsets,
                                             const void *context_lens,
                                             const cnnlDataType_t dtype,
                                             const int batch,
                                             const int head_num,
                                             const int head_size,
                                             const int max_context_len,
                                             const int cache_mem_len,
                                             const size_t context_bs_stride,
                                             const size_t context_head_stride,
                                             const size_t context_seq_stride,
                                             const size_t cache_bs_stride,
                                             const size_t cache_head_stride,
                                             const size_t cache_seq_stride,
                                             const size_t cache_scale_head_stride,
                                             const bool packed,
                                             const int quant_mode);
}  // namespace tmo

#endif  // CSRC_KERNELS_OFFLINE_QUANT_TO_LINEAR_CACHE_MLUH_
