cmake_minimum_required(VERSION 3.16)
project(mindspeed_ops_C)

set(CMAKE_CXX_STANDARD 17)

include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake)

set(PYTHON_SUPPORTED_VERSIONS "3.10" "3.11" "3.12")

find_package(pybind11 REQUIRED)

# Detect torch CMake config via Python.
# This adds torch's cmake prefix path so find_package(Torch) can locate it.
append_cmake_prefix_path("torch" "import torch; print(torch.utils.cmake_prefix_path)")

find_package(Torch REQUIRED)

# Verify torch version
run_python(TORCH_VERSION
  "import torch; print(torch.__version__)" "Failed to get torch version")
message(STATUS "Detected torch version: ${TORCH_VERSION}")

set(ASCEND_CANN_PACKAGE_PATH ${ASCEND_HOME_PATH})

set(MINDSPEED_CSRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/mindspeed_ops/csrc)

file(GLOB MINDSPEED_SRC
    ${MINDSPEED_CSRC_DIR}/torch_binding.cpp
    ${MINDSPEED_CSRC_DIR}/aclnn_torch_adapter/*.cpp
)

include_directories(
    ${pybind11_INCLUDE_DIRS}
    ${PYTHON_INCLUDE_PATH}
    ${TORCH_INCLUDE_DIRS}
    ${TORCH_NPU_PATH}/include
    ${ASCEND_HOME_PATH}/include
    ${MINDSPEED_CSRC_DIR}
    ${MINDSPEED_CSRC_DIR}/cann
)

pybind11_add_module(mindspeed_ops_C ${MINDSPEED_SRC})

target_link_directories(
    mindspeed_ops_C
    PRIVATE
    ${TORCH_LIBRARY_DIRS}
    ${TORCH_NPU_PATH}/lib/
    ${ASCEND_HOME_PATH}/lib64
)

target_link_libraries(
    mindspeed_ops_C
    PUBLIC
    ${TORCH_LIBRARIES}
    torch_npu
    ascendcl
)

target_link_options(mindspeed_ops_C PRIVATE "-Wl,-rpath,$ORIGIN:$ORIGIN/lib")

set(MINDSPEED_INSTALL_PATH "${CMAKE_INSTALL_PREFIX}")

install(TARGETS mindspeed_ops_C LIBRARY DESTINATION ${MINDSPEED_INSTALL_PATH})