/*************************************************************************
 * 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 TEST_KERNELS_PYTEST_REGISTER_API_H_
#define TEST_KERNELS_PYTEST_REGISTER_API_H_
#include <torch/extension.h>
#include <torch/torch.h>

namespace tmo {
namespace kernel_test_api {
at::Tensor embedding_test(const at::Tensor &input,
                          const at::Tensor &weight,
                          const int vocab_offset,
                          const int vocab_part);

at::Tensor alibi_slope_test(const at::Tensor &true_seq_lens,
                            const int batch,
                            const int tp_head_num,
                            const int tp_num,
                            const bool use_dynamic,
                            const int max_sequence_length);

at::Tensor create_cos_sin_table_test(at::Tensor &rotary_emb_alpha_cached,
                                     const at::Tensor &seq_lens,
                                     const int max_position_embeddings,
                                     const int batch_stride,
                                     const int rotary_seq_len,
                                     const int rotary_dim,
                                     const int rotary_stride,
                                     const float rotary_scaling,
                                     const int rotary_scaling_type,
                                     const float rotary_base,
                                     const bool interleaved);

void slice_cu_seq_lens_test(const at::Tensor &cu_seq_lens,
                            at::Tensor &sliced_cu_seq_lens,
                            int batch,
                            int parallel_num);

void generate_cu_seq_lens_test(at::Tensor &gen_cu_seq_lens,
                               int seq_len,
                               int parallel_num,
                               bool is_causal_mask,
                               bool is_kv_seq_len);

void per_head_quantize_test(at::Tensor &dst,
                            at::Tensor &scale,
                            const at::Tensor &src,
                            int bs,
                            int seq_len,
                            int head_num,
                            int head_size,
                            int dst_bs_stride,
                            int dst_seq_stride,
                            int dst_head_stride,
                            int src_bs_stride,
                            int src_seq_stride,
                            int src_head_stride);

at::Tensor rotary_embedding_test(const at::Tensor &input,
                                 const at::Tensor &sin_table,
                                 const at::Tensor &cos_table,
                                 const c10::optional<at::Tensor> &seq_offsets,
                                 const c10::optional<at::Tensor> &cu_seq_lens,
                                 const bool &interleaved,
                                 const bool &discrete,
                                 const bool &dynamic_ntk,
                                 const bool &rope_2d,
                                 const int &max_seq_len,
                                 const bool &no_offset);

}  // namespace kernel_test_api
}  // namespace tmo
#endif  // TEST_KERNELS_PYTEST_REGISTER_API_H_
