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

PYBIND11_MODULE(btunittests, m) {
  m.def("embedding_test", &tmo::kernel_test_api::embedding_test, "Test embedding kernel function.");
  m.def("alibi_slope_test", &tmo::kernel_test_api::alibi_slope_test,
        "Test alibi_slope kernel function.");
  m.def("create_cos_sin_table_test", &tmo::kernel_test_api::create_cos_sin_table_test,
        "Test create_cos_sin_table_test kernel function.");
  m.def("slice_cu_seq_lens_test", &tmo::kernel_test_api::slice_cu_seq_lens_test,
        "Test slice_cu_seq_lens_test kernel function.");
  m.def("generate_cu_seq_lens_test", &tmo::kernel_test_api::generate_cu_seq_lens_test,
        "Test generate_cu_seq_lens_test kernel function.");
  m.def("per_head_quantize_test", &tmo::kernel_test_api::per_head_quantize_test,
        "Test per_head_quantize_test kernel function.");
  m.def("rotary_embedding_test", &tmo::kernel_test_api::rotary_embedding_test,
        "Test rotary_embedding kernel function.");
}
