cmake_minimum_required(VERSION 3.27)

project(gpuastar_cu VERSION ${SKBUILD_PROJECT_VERSION} LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)

if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
  set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
endif()

set(Python_FIND_UNVERSIONED_NAMES LAST CACHE STRING "Search unversioned Python names after versioned names")
find_package(Python 3.10 REQUIRED COMPONENTS Interpreter)

set(_gpuastar_cu_with_cuda_default ON)
if(APPLE)
  # macOS source-tree development installs contain Python sources only.
  # NVIDIA CUDA is not available on macOS.
  set(_gpuastar_cu_with_cuda_default OFF)
endif()
option(GPUASTAR_CU_WITH_CUDA "Build the CUDA extension module" ${_gpuastar_cu_with_cuda_default})
unset(_gpuastar_cu_with_cuda_default)

if(NOT GPUASTAR_CU_WITH_CUDA)
  message(STATUS "GPUASTAR_CU_WITH_CUDA is OFF. Installing gpuastar-cu Python package without the CUDA extension.")
  install(FILES src/gpuastar_cu/py.typed DESTINATION gpuastar_cu)
  return()
endif()

find_package(Python 3.10 REQUIRED COMPONENTS Interpreter Development.Module)

include("${CMAKE_CURRENT_LIST_DIR}/cmake/GPUAStarPythonCuda.cmake")
gpuastar_configure_python_cuda_toolkit()

if(DEFINED ENV{GPUASTAR_CU_CUDA_ARCHITECTURES} AND NOT "$ENV{GPUASTAR_CU_CUDA_ARCHITECTURES}" STREQUAL "")
  set(_gpuastar_cu_cuda_architectures "$ENV{GPUASTAR_CU_CUDA_ARCHITECTURES}")
else()
  set(_gpuastar_cu_cuda_architectures "75-real;80-real;86-real;87-real;89-real;90")
endif()
set(CMAKE_CUDA_ARCHITECTURES "${_gpuastar_cu_cuda_architectures}" CACHE STRING "CUDA architectures" FORCE)
if(
  DEFINED ENV{GPUASTAR_CU_CUDA_ARCHITECTURES}
  AND NOT "$ENV{GPUASTAR_CU_CUDA_ARCHITECTURES}" STREQUAL ""
  AND NOT DEFINED ENV{GPUASTAR_CU_TORCH_CUDA_ARCH_LIST}
)
  message(STATUS "GPUASTAR_CU_CUDA_ARCHITECTURES is set. Set GPUASTAR_CU_TORCH_CUDA_ARCH_LIST too if Torch CMake needs different arch names.")
endif()
if(DEFINED ENV{GPUASTAR_CU_TORCH_CUDA_ARCH_LIST} AND NOT "$ENV{GPUASTAR_CU_TORCH_CUDA_ARCH_LIST}" STREQUAL "")
  set(
    GPUASTAR_CU_TORCH_CUDA_ARCH_LIST
    "$ENV{GPUASTAR_CU_TORCH_CUDA_ARCH_LIST}"
    CACHE STRING
    "CUDA architectures for Torch CMake"
    FORCE
  )
else()
  set(GPUASTAR_CU_TORCH_CUDA_ARCH_LIST "7.5;8.0;8.6;8.7;8.9;9.0+PTX" CACHE STRING "CUDA architectures for Torch CMake")
endif()
set(ENV{TORCH_CUDA_ARCH_LIST} "${GPUASTAR_CU_TORCH_CUDA_ARCH_LIST}")
set(TORCH_CUDA_ARCH_LIST "${GPUASTAR_CU_TORCH_CUDA_ARCH_LIST}" CACHE STRING "CUDA architectures for Torch CMake" FORCE)

enable_language(CUDA)

set(CMAKE_CUDA_STANDARD 17)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)

execute_process(
  COMMAND "${Python_EXECUTABLE}" -c "import torch.utils,sys; sys.stdout.write(torch.utils.cmake_prefix_path)"
  OUTPUT_VARIABLE TORCH_CMAKE_PREFIX
  OUTPUT_STRIP_TRAILING_WHITESPACE
  RESULT_VARIABLE TORCH_PREFIX_RESULT
)
if(NOT TORCH_PREFIX_RESULT EQUAL 0)
  message(FATAL_ERROR "Could not import torch to discover torch.utils.cmake_prefix_path")
endif()
list(PREPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX}")

find_package(CUDAToolkit REQUIRED)

if(DEFINED ENV{GPUASTAR_CU_EXPECTED_CUDA_VERSION} AND NOT "$ENV{GPUASTAR_CU_EXPECTED_CUDA_VERSION}" STREQUAL "")
  set(_gpuastar_cu_expected_cuda_version "$ENV{GPUASTAR_CU_EXPECTED_CUDA_VERSION}")
  set(_gpuastar_cu_toolkit_major_minor "${CUDAToolkit_VERSION_MAJOR}.${CUDAToolkit_VERSION_MINOR}")
  string(REGEX MATCH "^[0-9]+\\.[0-9]+" _gpuastar_cu_compiler_major_minor "${CMAKE_CUDA_COMPILER_VERSION}")
  if(NOT "${_gpuastar_cu_toolkit_major_minor}" STREQUAL "${_gpuastar_cu_expected_cuda_version}")
    message(
      FATAL_ERROR
      "Expected CUDA toolkit ${_gpuastar_cu_expected_cuda_version}, but CMake found ${CUDAToolkit_VERSION}"
    )
  endif()
  if(NOT "${_gpuastar_cu_compiler_major_minor}" STREQUAL "${_gpuastar_cu_expected_cuda_version}")
    message(
      FATAL_ERROR
      "Expected CUDA compiler ${_gpuastar_cu_expected_cuda_version}, but CMake found ${CMAKE_CUDA_COMPILER_VERSION}"
    )
  endif()
  unset(_gpuastar_cu_compiler_major_minor)
  unset(_gpuastar_cu_expected_cuda_version)
  unset(_gpuastar_cu_toolkit_major_minor)
endif()

if(NOT TARGET torch::cudart)
  add_library(torch::cudart INTERFACE IMPORTED)
  target_link_libraries(torch::cudart INTERFACE CUDA::cudart)
endif()

if(NOT TARGET CUDA::nvToolsExt)
  set(_gpuastar_cu_nvtoolsext_roots "")
  if(DEFINED CUDAToolkit_ROOT)
    list(APPEND _gpuastar_cu_nvtoolsext_roots "${CUDAToolkit_ROOT}")
  endif()
  if(DEFINED CUDA_TOOLKIT_ROOT_DIR)
    list(APPEND _gpuastar_cu_nvtoolsext_roots "${CUDA_TOOLKIT_ROOT_DIR}")
  endif()
  find_library(
    _gpuastar_cu_nvtoolsext_lib
    NAMES nvToolsExt nvToolsExt64
    PATHS ${_gpuastar_cu_nvtoolsext_roots}
    PATH_SUFFIXES lib lib64 lib/x64 bin
    NO_DEFAULT_PATH
  )
  if(NOT _gpuastar_cu_nvtoolsext_lib)
    find_file(
      _gpuastar_cu_nvtoolsext_lib
      NAMES libnvToolsExt.so.1 libnvToolsExt.so.1.0.0
      PATHS ${_gpuastar_cu_nvtoolsext_roots}
      PATH_SUFFIXES lib lib64 lib/x64 bin
      NO_DEFAULT_PATH
    )
  endif()
  if(_gpuastar_cu_nvtoolsext_lib)
    add_library(CUDA::nvToolsExt UNKNOWN IMPORTED)
    set_target_properties(CUDA::nvToolsExt PROPERTIES
      IMPORTED_LOCATION "${_gpuastar_cu_nvtoolsext_lib}"
    )
  endif()
  unset(_gpuastar_cu_nvtoolsext_lib)
  unset(_gpuastar_cu_nvtoolsext_roots)
endif()

find_package(Torch REQUIRED)

add_library(_C MODULE
  csrc/bwas.cpp
  csrc/bwas_kernel.cu
)
set_property(TARGET _C PROPERTY CUDA_ARCHITECTURES "${_gpuastar_cu_cuda_architectures}")

target_include_directories(_C PRIVATE csrc ${TORCH_INCLUDE_DIRS})
target_link_libraries(_C PRIVATE ${TORCH_LIBRARIES} Python::Module CUDA::cudart)
target_compile_definitions(_C PRIVATE TORCH_EXTENSION_NAME=_C GPUASTAR_USE_PY_PRINT_REDIRECT=1)

if(MSVC)
  target_compile_options(_C PRIVATE $<$<COMPILE_LANGUAGE:CXX>:/O2>)
else()
  target_compile_options(_C PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-O3>)
endif()
target_compile_options(_C PRIVATE
  $<$<COMPILE_LANGUAGE:CUDA>:-O3>
  $<$<COMPILE_LANGUAGE:CUDA>:--extended-lambda>
  $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
  $<$<COMPILE_LANGUAGE:CUDA>:-lineinfo>
  $<$<COMPILE_LANGUAGE:CUDA>:-Xptxas=-O3>
)
if(DEFINED ENV{GPUASTAR_CU_USE_FAST_MATH} AND "$ENV{GPUASTAR_CU_USE_FAST_MATH}" STREQUAL "1")
  # Opt-in only: fast-math can change floating-point priority comparisons and
  # therefore may affect search tie behavior.
  target_compile_options(_C PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:--use_fast_math>)
endif()

set_target_properties(_C PROPERTIES
  PREFIX ""
  OUTPUT_NAME "_C"
  CUDA_SEPARABLE_COMPILATION OFF
)
if(WIN32)
  set_target_properties(_C PROPERTIES SUFFIX ".pyd")
endif()

set(_gpuastar_cu_build_rpath "")
set(_gpuastar_cu_install_rpath "")
if(UNIX AND NOT APPLE)
  list(APPEND _gpuastar_cu_build_rpath "$ORIGIN")
  list(APPEND
    _gpuastar_cu_install_rpath
    "$ORIGIN"
    "$ORIGIN/../torch/lib"
    "$ORIGIN/../nvidia/cu13/lib"
    "$ORIGIN/../nvidia/cuda_runtime/lib"
    "$ORIGIN/../nvidia/cuda_nvrtc/lib"
  )
  if(SKBUILD_STATE STREQUAL "editable")
    list(APPEND _gpuastar_cu_install_rpath "$ORIGIN/../../..")
  endif()
elseif(APPLE)
  list(APPEND _gpuastar_cu_build_rpath "@loader_path")
  list(APPEND _gpuastar_cu_install_rpath "@loader_path" "@loader_path/../torch/lib")
endif()
if(_gpuastar_cu_build_rpath OR _gpuastar_cu_install_rpath)
  set_target_properties(_C PROPERTIES
    BUILD_RPATH "${_gpuastar_cu_build_rpath}"
    INSTALL_RPATH "${_gpuastar_cu_install_rpath}"
    INSTALL_REMOVE_ENVIRONMENT_RPATH TRUE
    INSTALL_RPATH_USE_LINK_PATH FALSE
  )
endif()

install(TARGETS _C LIBRARY DESTINATION gpuastar_cu RUNTIME DESTINATION gpuastar_cu)
install(FILES src/gpuastar_cu/py.typed DESTINATION gpuastar_cu)
