/*************************************************************************
 * 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_CREATE_COS_SIN_TABLE_MLUH_
#define CSRC_KERNELS_CREATE_COS_SIN_TABLE_MLUH_

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

namespace tmo {

/**
 * @brief Create cos and sin table for rotary embedding.
 * @param queue: The queue for mlu.
 * @param cos_sin_table: Output. Pointer to the MLU memory that stores the cos and sin table.
 *        If rotary_scaling_type is linear, the shape is [rotary_seq_len, rotary_stride].
 *        If rotary_scaling_type is dynamic ntk, the shape is [batch, rotary_seq_len,
 * rotary_stride].
 * @param rotary_emb_alpha_cached: Output/Input. Pointer to the MLU memory that
 *        stores the ntk alpha cache. Only used in dynamic ntk, the shape is [batch].
 * @param seq_lens: Input. Pointer to the MLU memory that stores the true sequence len.
 *        The shape is [batch].
 * @param max_position_embeddings: The maximum rotary embedding positions.
 * @param batch: Batch size.
 * @param batch_stride: The stride for batch dim of cos_sin_table.
 *        Only used in dynamic ntk, the value is rotary_seq_len * rotary_stride.
 * @param rotary_seq_len: The rotary sequence length of cos and sin table.
 * @param rotary_dim: The rotary dim value of cos and sin table.
 * @param rotary_stride: The stride of rotary_seq_len dim for cos and sin table.
 * @param rotary_base: The rotary base, value is usually 10000.
 * @param rotary_scaling: The rotary scaling, value is usually 1.
 * @param rotary_scaling_type: The rotary scaling type, value is linear or dynamic ntk.
 * @param interleaved: A boolean value indicates compute mode of rotary embedding.
 * @param dtype: Data type of cos and sin table generated.
 */
KernelStatus invokeCreateCosSinTable(cnrtQueue_t queue,
                                     void *cos_sin_table,
                                     float *rotary_emb_alpha_cached,
                                     int *seq_lens,
                                     int max_position_embeddings,
                                     int batch,
                                     int batch_stride,
                                     int rotary_seq_len,
                                     int rotary_dim,
                                     int rotary_stride,
                                     float rotary_base,
                                     float rotary_scaling,
                                     int rotary_scaling_type,
                                     bool interleaved,
                                     cnnlDataType_t dtype);

}  // namespace tmo

#endif  // CSRC_KERNELS_CREATE_COS_SIN_TABLE_MLUH_
