/*************************************************************************
 * 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_GET_GLM_POSITION_ID_MLUH_
#define CSRC_KERNELS_GET_GLM_POSITION_ID_MLUH_

#include "cnnl.h"
#include "kernel_utils.h"
namespace tmo {
/**
 * @brief Get generate position id from context position id, when position id is 2D,
 *        increase block position id by one.
 * @example
 * in GLM network, context_pos_id shape is [batch, 2, context_seq_len], data is
 * [[[0, 1, 2, 2, 2, 2, 2], [0, 0, 0, 1, 1, 1, 1]],
 *  [[0, 1, 2, 3, 4, 5, 5], [0, 0, 0, 0, 0, 0, 1]]]
 * after invoke this kernel, the data is
 * [[[2], [2]],
 *  [[5], [2]]]
 * @param queue: The queue for mlu.
 * @param context_pos_id: Input. Pointer to the MLU memory that stores the position id of
 * context.
 * @param generate_pos_id: Output. Pointer to the MLU memory that stores the position id of
 * generate.
 * @param batch: Batch size.
 * @param context_seq_len: The sequence length of context.
 * @param pos_id_dimension: The dimension of position id, 1 for 1D, 2 for 2D.
 */
KernelStatus invokeSliceIndividualPosId(cnrtQueue_t queue,
                                        int *context_pos_id,
                                        int *generate_pos_id,
                                        int batch,
                                        int context_seq_len,
                                        int pos_id_dimension);

/**
 * @brief Increase block position id by one in generate stage.
 * @example
 * in GLM network, generate_pos_id shape is [batch, 2, 1], data is
 * [[[2], [1]], [[5], [1]]]
 * after invoke this kernel, the data is
 * [[[2], [2]], [[5], [2]]]
 * @param queue: The queue for mlu.
 * @param generate_pos_id: Output/Input. Pointer to the MLU memory that stores the position id of
 * generate.
 * @param batch: Batch size.
 */
KernelStatus invokeIncrement2DPosId(cnrtQueue_t queue, int *generate_pos_id, int batch);
}  // namespace tmo

#endif  // CSRC_KERNELS_GET_GLM_POSITION_ID_MLUH_
