cmake_minimum_required(VERSION 3.24...3.31)

set(TMOL_OPS_CUDA "" CACHE STRING "Build an optional CUDA ops package: 12 or 13")
option(TMOL_ENABLE_CUDA "Build CUDA extensions (disable for CPU-only)" OFF)
set(_TMOL_NATIVE_PACKAGE tmol)
if(TMOL_OPS_CUDA)
  if(NOT TMOL_OPS_CUDA MATCHES "^(12|13)$")
    message(FATAL_ERROR "TMOL_OPS_CUDA must be 12 or 13")
  endif()
  set(TMOL_ENABLE_CUDA ON)
  set(_TMOL_NATIVE_PACKAGE "tmol_ops_cu${TMOL_OPS_CUDA}")
endif()

# ═══════════════════════════════════════════════════════════════════════════════
# Detect CUDA and choose languages
# ═══════════════════════════════════════════════════════════════════════════════

include(CheckLanguage)
if(TMOL_ENABLE_CUDA AND NOT CMAKE_CUDA_COMPILER)
  # check_language() does not reliably discover toolkit installs under
  # /usr/local/cuda when CMake is launched outside an activated CUDA env.
  find_program(
    _TMOL_NVCC_EXECUTABLE
    NAMES nvcc
    HINTS "$ENV{CUDA_HOME}" "$ENV{CUDA_PATH}" /usr/local/cuda
    PATH_SUFFIXES bin
  )
  if(_TMOL_NVCC_EXECUTABLE)
    set(CMAKE_CUDA_COMPILER "${_TMOL_NVCC_EXECUTABLE}" CACHE FILEPATH "CUDA compiler")
  endif()
endif()
check_language(CUDA)

if(CMAKE_CUDA_COMPILER AND TMOL_ENABLE_CUDA)
  # ── CUDA architecture selection — MUST be set before project() so that
  # enable_language(CUDA) uses our architectures, not the nvcc defaults.
  # CMAKE_CUDA_ARCHITECTURES can be set via:
  #   -DCMAKE_CUDA_ARCHITECTURES="80;90;103"  (from scikit-build-core -C flags)
  #   TORCH_CUDA_ARCH_LIST env var (PyTorch format: "7.0 8.0 9.0+PTX")
  #   or falls back to a sensible default.
  #
  # Ask nvcc instead of maintaining a version-to-architecture table. CUDA 13.1,
  # for example, can emit sm_103 and sm_121 even though a numeric "max arch"
  # test would neither validate the former nor discover the latter. Keep the
  # project floor at Turing (sm_75), matching tmol's supported Colab wheels.
  execute_process(
    COMMAND ${CMAKE_CUDA_COMPILER} --list-gpu-code
    OUTPUT_VARIABLE _TMOL_NVCC_GPU_CODE
    ERROR_VARIABLE _TMOL_NVCC_GPU_CODE_ERROR
    RESULT_VARIABLE _TMOL_NVCC_GPU_CODE_RESULT
    OUTPUT_STRIP_TRAILING_WHITESPACE
  )
  string(REGEX MATCHALL "sm_[0-9]+" _TMOL_NVCC_SM_NAMES "${_TMOL_NVCC_GPU_CODE}")
  set(_TMOL_NVCC_SUPPORTED_ARCHS "")
  foreach(_SM_NAME ${_TMOL_NVCC_SM_NAMES})
    string(REPLACE "sm_" "" _SM_ARCH "${_SM_NAME}")
    if(_SM_ARCH GREATER_EQUAL 75)
      list(APPEND _TMOL_NVCC_SUPPORTED_ARCHS "${_SM_ARCH}")
    endif()
  endforeach()
  list(REMOVE_DUPLICATES _TMOL_NVCC_SUPPORTED_ARCHS)
  list(SORT _TMOL_NVCC_SUPPORTED_ARCHS COMPARE NATURAL ORDER ASCENDING)
  if(NOT _TMOL_NVCC_SUPPORTED_ARCHS)
    message(FATAL_ERROR
      "tmol: nvcc --list-gpu-code did not report any supported sm_75+ targets "
      "(exit ${_TMOL_NVCC_GPU_CODE_RESULT}): ${_TMOL_NVCC_GPU_CODE_ERROR}")
  endif()
  # Ship SASS for every release target and a single forward-compatible PTX
  # image at the newest target. Unsuffixed numeric CMake architectures would
  # otherwise embed redundant PTX for every generation.
  set(_TMOL_NVCC_RELEASE_ARCHS "")
  list(LENGTH _TMOL_NVCC_SUPPORTED_ARCHS _TMOL_NVCC_ARCH_COUNT)
  math(EXPR _TMOL_NVCC_LAST_ARCH_INDEX "${_TMOL_NVCC_ARCH_COUNT} - 1")
  foreach(_ARCH_INDEX RANGE ${_TMOL_NVCC_LAST_ARCH_INDEX})
    list(GET _TMOL_NVCC_SUPPORTED_ARCHS ${_ARCH_INDEX} _SM_ARCH)
    if(_ARCH_INDEX EQUAL _TMOL_NVCC_LAST_ARCH_INDEX)
      list(APPEND _TMOL_NVCC_RELEASE_ARCHS "${_SM_ARCH}")
    else()
      list(APPEND _TMOL_NVCC_RELEASE_ARCHS "${_SM_ARCH}-real")
    endif()
  endforeach()

  if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
    if(DEFINED ENV{CMAKE_CUDA_ARCHITECTURES})
      set(CMAKE_CUDA_ARCHITECTURES "$ENV{CMAKE_CUDA_ARCHITECTURES}" CACHE STRING "CUDA architectures")
    elseif(DEFINED ENV{TORCH_CUDA_ARCH_LIST})
      set(_ARCH_STR "$ENV{TORCH_CUDA_ARCH_LIST}")
      string(REPLACE " " ";" _ARCH_LIST "${_ARCH_STR}")
      set(_CLEAN_ARCHS "")
      foreach(_A ${_ARCH_LIST})
        string(FIND "${_A}" "+PTX" _HAS_PTX)
        string(REPLACE "+PTX" "" _A "${_A}")
        string(REPLACE "." "" _A "${_A}")
        if(_HAS_PTX GREATER_EQUAL 0)
          list(APPEND _CLEAN_ARCHS "${_A}")
        else()
          list(APPEND _CLEAN_ARCHS "${_A}-real")
        endif()
      endforeach()
      set(CMAKE_CUDA_ARCHITECTURES ${_CLEAN_ARCHS} CACHE STRING "CUDA architectures")
    else()
      set(CMAKE_CUDA_ARCHITECTURES "native" CACHE STRING "CUDA architectures")
    endif()
  endif()
  if(CMAKE_CUDA_ARCHITECTURES STREQUAL "all-major" OR
      CMAKE_CUDA_ARCHITECTURES STREQUAL "all")
    # CMake's symbolic all values are compiler-policy dependent. Resolve them
    # to nvcc's exact sm_75+ list so release wheels cannot silently omit an
    # architecture discovered above.
    set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_NVCC_RELEASE_ARCHS}" CACHE STRING "CUDA architectures" FORCE)
  endif()

  if(CMAKE_CUDA_ARCHITECTURES STREQUAL "native")
    # CMake 3.24+ queries the GPU(s) visible at configure time and emits only
    # their native device code. This is the fast default for local builds.
    set(_TMOL_CUDA_ARCHITECTURES "native")
  else()
    # Filter requested numeric arches by exact compiler support. Architectures
    # are not a simple compatibility interval, so a maximum clamp is wrong.
    set(_TMOL_FILTERED_ARCHS "")
    set(_TMOL_REJECTED_ARCHS "")
    foreach(_A ${CMAKE_CUDA_ARCHITECTURES})
      string(REPLACE "-real" "" _A_NUM "${_A}")
      string(REPLACE "-virtual" "" _A_NUM "${_A_NUM}")
      if(_A_NUM IN_LIST _TMOL_NVCC_SUPPORTED_ARCHS)
        list(APPEND _TMOL_FILTERED_ARCHS "${_A}")
      else()
        list(APPEND _TMOL_REJECTED_ARCHS "${_A}")
      endif()
    endforeach()
    if(NOT _TMOL_FILTERED_ARCHS)
      message(FATAL_ERROR
        "tmol: none of the requested CUDA architectures "
        "(${CMAKE_CUDA_ARCHITECTURES}) are supported by ${CMAKE_CUDA_COMPILER}; "
        "supported sm_75+ targets: ${_TMOL_NVCC_SUPPORTED_ARCHS}")
    endif()
    if(_TMOL_REJECTED_ARCHS)
      message(WARNING
        "tmol: ignoring CUDA architectures not supported by this nvcc: "
        "${_TMOL_REJECTED_ARCHS}")
    endif()
    set(_TMOL_CUDA_ARCHITECTURES "${_TMOL_FILTERED_ARCHS}")
  endif()
  # Preserve the selected value independently: project() may overwrite the
  # cache while enabling CUDA. `all` is numeric here; `native` intentionally
  # remains symbolic so nvcc resolves the GPU visible at compile time.
  set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_CUDA_ARCHITECTURES}" CACHE STRING "CUDA architectures" FORCE)
  set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_CUDA_ARCHITECTURES}")
  message(STATUS "tmol: nvcc supports sm_75+ targets ${_TMOL_NVCC_SUPPORTED_ARCHS}")
  message(STATUS "tmol: using CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}")

  project(tmol LANGUAGES CXX CUDA)
  set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_CUDA_ARCHITECTURES}" CACHE STRING "CUDA architectures" FORCE)
  set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_CUDA_ARCHITECTURES}")
  set(TMOL_HAS_CUDA TRUE)
  message(STATUS "tmol: CUDA enabled, CMAKE_CUDA_ARCHITECTURES = ${CMAKE_CUDA_ARCHITECTURES}")
else()
  project(tmol LANGUAGES CXX)
  set(TMOL_HAS_CUDA FALSE)
  message(STATUS "tmol: CUDA not found or disabled; building CPU-only extensions")
endif()

if(TMOL_OPS_CUDA)
  if(NOT TMOL_HAS_CUDA OR NOT CMAKE_CUDA_COMPILER_VERSION MATCHES "^${TMOL_OPS_CUDA}\\.")
    message(FATAL_ERROR "tmol-ops-cu${TMOL_OPS_CUDA} requires a CUDA ${TMOL_OPS_CUDA} toolkit")
  endif()
endif()

# ═══════════════════════════════════════════════════════════════════════════════
# Dependencies
# ═══════════════════════════════════════════════════════════════════════════════

find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)

# Build with current stable headers, target the oldest supported stable C ABI.
# Production modules avoid TorchConfig.cmake and ATen C++ objects, and do not
# link libtorch_python. C++17 retains the manylinux baseline C++ runtime.
set(_TMOL_PYTHON_INCLUDE_DIRS "${Python_INCLUDE_DIRS}")
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
  # LTO can resolve STL calls to Torch's unversioned template exports, even
  # with libstdc++ first. Keep those calls bound to the system C++ runtime.
  set(CMAKE_INTERPROCEDURAL_OPTIMIZATION OFF)
endif()
if(TMOL_HAS_CUDA)
  set(CMAKE_CUDA_STANDARD 17)
  set(CMAKE_CUDA_STANDARD_REQUIRED ON)
endif()
execute_process(
  COMMAND ${Python_EXECUTABLE} -c "import pathlib, torch; print(pathlib.Path(torch.__file__).parent)"
  OUTPUT_VARIABLE _TMOL_TORCH_DIR OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY
)
set(TMOL_STABLE_HEADERS "${_TMOL_TORCH_DIR}/include" CACHE PATH "PyTorch 2.14+ headers for production stable ABI modules")
set(TORCH_INCLUDE_DIRS "${_TMOL_TORCH_DIR}/include" "${_TMOL_TORCH_DIR}/include/torch/csrc/api/include")
find_library(TORCH_LIBRARY torch PATHS "${_TMOL_TORCH_DIR}/lib" NO_DEFAULT_PATH REQUIRED)
find_library(TORCH_CPU_LIBRARY torch_cpu PATHS "${_TMOL_TORCH_DIR}/lib" NO_DEFAULT_PATH REQUIRED)
set(TORCH_LIBRARIES "${TORCH_LIBRARY}" "${TORCH_CPU_LIBRARY}")
if(TMOL_HAS_CUDA)
  find_library(TORCH_CUDA_LIBRARY torch_cuda PATHS "${_TMOL_TORCH_DIR}/lib" NO_DEFAULT_PATH REQUIRED)
  list(APPEND TORCH_LIBRARIES "${TORCH_CUDA_LIBRARY}")
endif()
execute_process(
  COMMAND ${Python_EXECUTABLE} -m pybind11 --cmakedir
  OUTPUT_VARIABLE _TMOL_PYBIND11_CMAKE_PREFIX_PATH OUTPUT_STRIP_TRAILING_WHITESPACE
  COMMAND_ERROR_IS_FATAL ANY
)
list(PREPEND CMAKE_PREFIX_PATH "${_TMOL_PYBIND11_CMAKE_PREFIX_PATH}")
find_package(pybind11 CONFIG REQUIRED)

# ═══════════════════════════════════════════════════════════════════════════════
# Compiler flags
# ═══════════════════════════════════════════════════════════════════════════════


if(TMOL_HAS_CUDA)
  set(TMOL_NVCC_THREADS "4" CACHE STRING "Threads per nvcc invocation")
endif()

# ═══════════════════════════════════════════════════════════════════════════════
# Include directories
# ═══════════════════════════════════════════════════════════════════════════════

set(TMOL_INCLUDE_DIRS
  ${CMAKE_SOURCE_DIR}
  ${CMAKE_SOURCE_DIR}/tmol/extern
  ${TORCH_INCLUDE_DIRS}
  ${_TMOL_PYTHON_INCLUDE_DIRS}
)

# NVTX includes (from pip nvidia-nvtx package or CUDA toolkit)
if(TMOL_HAS_CUDA)
  execute_process(
    COMMAND ${Python_EXECUTABLE} -c
      "import nvidia.nvtx, os; print(os.path.join(nvidia.nvtx.__path__[0], 'include'))"
    OUTPUT_VARIABLE _NVTX_INCLUDE
    OUTPUT_STRIP_TRAILING_WHITESPACE
    ERROR_QUIET
    RESULT_VARIABLE _NVTX_RESULT
  )
  if(_NVTX_RESULT EQUAL 0 AND EXISTS "${_NVTX_INCLUDE}")
    list(APPEND TMOL_INCLUDE_DIRS "${_NVTX_INCLUDE}")
  else()
    find_package(CUDAToolkit QUIET)
    if(CUDAToolkit_FOUND)
      list(APPEND TMOL_INCLUDE_DIRS "${CUDAToolkit_INCLUDE_DIRS}")
    endif()
  endif()
endif()

# ═══════════════════════════════════════════════════════════════════════════════
# Helper functions
# ═══════════════════════════════════════════════════════════════════════════════

# Torch's CMake targets advertise optional CUDA libraries such as NVRTC even
# when an extension does not use them.  Keep those libraries out of DT_NEEDED
# so wheels remain loadable with PyTorch distributions that do not ship every
# optional CUDA component (notably the monolithic aarch64 wheels).
function(tmol_set_portable_link_flags TARGET)
  if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
    target_link_options(${TARGET} PRIVATE "-Wl,--as-needed")
    # Resolve the C++ runtime before Torch's incidental template exports.
    # Otherwise an extension can silently depend on its build Torch's STL.
    if("stdc++" IN_LIST CMAKE_CXX_IMPLICIT_LINK_LIBRARIES)
      target_link_libraries(${TARGET} PRIVATE stdc++)
    endif()
  endif()
endfunction()

if(TMOL_BUILD_TESTS)
  find_package(OpenMP COMPONENTS CXX QUIET)
endif()

function(tmol_set_abi TARGET)
  if("${TARGET}" MATCHES "^test_")
    set_target_properties(${TARGET} PROPERTIES CXX_STANDARD 20 CUDA_STANDARD 20)
    find_library(_TMOL_TORCH_PYTHON torch_python PATHS "${_TMOL_TORCH_DIR}/lib" NO_DEFAULT_PATH REQUIRED)
    find_library(_TMOL_C10 c10 PATHS "${_TMOL_TORCH_DIR}/lib" NO_DEFAULT_PATH REQUIRED)
    target_link_libraries(${TARGET} PRIVATE "${_TMOL_TORCH_PYTHON}" "${_TMOL_C10}")
    if(OpenMP_CXX_FOUND)
      target_link_libraries(${TARGET} PRIVATE OpenMP::OpenMP_CXX)
    endif()
  else()
    target_compile_definitions(${TARGET} PRIVATE
      TORCH_TARGET_VERSION=0x020b000000000000
      TMOL_INTEROP_NAMESPACE="tmol_native_${TARGET}")
    if("${TARGET}" STREQUAL "_C")
      target_compile_definitions(${TARGET} PRIVATE TMOL_MAIN_MODULE)
    else()
      target_sources(${TARGET} PRIVATE tmol/native.cpp)
    endif()
    target_include_directories(${TARGET} BEFORE PRIVATE "${TMOL_STABLE_HEADERS}")
  endif()
endfunction()

# Common compile options for CUDA extensions
function(tmol_set_cuda_flags TARGET)
  tmol_set_abi(${TARGET})
  target_compile_options(${TARGET} PRIVATE
    $<$<COMPILE_LANGUAGE:CXX>:-O3 -w -fno-math-errno -DWITH_CUDA -DWITH_NVTX -DUSE_CUDA>
    $<$<COMPILE_LANGUAGE:CUDA>:
      -O3 -w
      --expt-extended-lambda
      --expt-relaxed-constexpr
      -DWITH_NVTX -DWITH_CUDA -DUSE_CUDA
      --threads=${TMOL_NVCC_THREADS}
    >
  )
  target_include_directories(${TARGET} PRIVATE ${TMOL_INCLUDE_DIRS})
  tmol_set_portable_link_flags(${TARGET})
  target_link_libraries(${TARGET} PRIVATE ${TORCH_LIBRARIES})
endfunction()

# Common compile options for CPU-only extensions
function(tmol_set_cpp_flags TARGET)
  tmol_set_abi(${TARGET})
  target_compile_options(${TARGET} PRIVATE
    $<$<COMPILE_LANGUAGE:CXX>:-O3 -w -fno-math-errno -DWITH_NVTX>
  )
  target_include_directories(${TARGET} PRIVATE ${TMOL_INCLUDE_DIRS})
  tmol_set_portable_link_flags(${TARGET})
  target_link_libraries(${TARGET} PRIVATE ${TORCH_LIBRARIES})
endfunction()

# Create a pybind11 CUDA extension and install it
function(tmol_add_pybind_cuda_ext TARGET_NAME INSTALL_DIR EXT_NAME)
  pybind11_add_module(${TARGET_NAME} MODULE ${ARGN})
  target_compile_definitions(${TARGET_NAME} PRIVATE TORCH_EXTENSION_NAME=${EXT_NAME})
  tmol_set_cuda_flags(${TARGET_NAME})
  set_target_properties(${TARGET_NAME} PROPERTIES
    OUTPUT_NAME "${EXT_NAME}"
    LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/${INSTALL_DIR}"
  )
  install(TARGETS ${TARGET_NAME} DESTINATION "${INSTALL_DIR}")
endfunction()

# Create a pybind11 CPU-only extension and install it
function(tmol_add_pybind_cpp_ext TARGET_NAME INSTALL_DIR EXT_NAME)
  pybind11_add_module(${TARGET_NAME} MODULE ${ARGN})
  target_compile_definitions(${TARGET_NAME} PRIVATE TORCH_EXTENSION_NAME=${EXT_NAME})
  tmol_set_cpp_flags(${TARGET_NAME})
  set_target_properties(${TARGET_NAME} PROPERTIES
    OUTPUT_NAME "${EXT_NAME}"
    LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/${INSTALL_DIR}"
  )
  install(TARGETS ${TARGET_NAME} DESTINATION "${INSTALL_DIR}")
endfunction()

# ═══════════════════════════════════════════════════════════════════════════════
# 1. Native CPU or CUDA operator module
# ═══════════════════════════════════════════════════════════════════════════════

set(_C_SOURCES
  tmol/native.cpp
  # optimization
  tmol/optimization/armijo_compiled.ops.cpp
  tmol/optimization/armijo_compiled.cpu.cpp
  # io/details/compiled
  tmol/io/details/compiled/compiled.ops.cpp
  tmol/io/details/compiled/gen_pose_leaf_atoms.cpu.cpp
  tmol/io/details/compiled/resolve_his_taut.cpu.cpp
  # kinematics/compiled
  tmol/kinematics/compiled/compiled_ops.cpp
  tmol/kinematics/compiled/compiled.cpu.cpp
  # pack/compiled
  tmol/pack/compiled/compiled.ops.cpp
  tmol/pack/compiled/compiled.cpu.cpp
  # pack/rotamer/dunbrack
  tmol/pack/rotamer/dunbrack/compiled.ops.cpp
  tmol/pack/rotamer/dunbrack/compiled.cpu.cpp
  # score/backbone_torsion
  tmol/score/backbone_torsion/potentials/compiled.ops.cpp
  tmol/score/backbone_torsion/potentials/backbone_torsion_pose_score.cpu.cpp
  # score/cartbonded
  tmol/score/cartbonded/potentials/compiled.ops.cpp
  tmol/score/cartbonded/potentials/cartbonded_pose_score.cpu.cpp
  # score/genbonded
  tmol/score/genbonded/potentials/compiled.ops.cpp
  tmol/score/genbonded/potentials/genbonded_pose_score.cpu.cpp
  # score/constraint
  tmol/score/constraint/potentials/compiled.ops.cpp
  tmol/score/constraint/potentials/constraint_score.cpu.cpp
  # score/disulfide
  tmol/score/disulfide/potentials/compiled.ops.cpp
  tmol/score/disulfide/potentials/disulfide_pose_score.cpu.cpp
  # score/dunbrack
  tmol/score/dunbrack/potentials/compiled.ops.cpp
  tmol/score/dunbrack/potentials/dunbrack_pose_score.cpu.cpp
  # score/elec
  tmol/score/elec/potentials/compiled.ops.cpp
  tmol/score/elec/potentials/elec_pose_score.cpu.cpp
  # score/hbond
  tmol/score/hbond/potentials/compiled.ops.cpp
  tmol/score/hbond/potentials/hbond_pose_score.cpu.cpp
  tmol/score/hbond/potentials/gen_hbond_bases.cpu.cpp
  # score/ljlk
  tmol/score/ljlk/potentials/compiled.ops.cpp
  tmol/score/ljlk/potentials/ljlk_pose_score.cpu.cpp
  tmol/score/ljlk/potentials/ljlk_elec_pose_score.cpu.cpp
  # score/lk_ball
  tmol/score/lk_ball/potentials/compiled.ops.cpp
  tmol/score/lk_ball/potentials/lk_ball_pose_score.cpu.cpp
  tmol/score/lk_ball/potentials/gen_pose_waters.cpu.cpp
  # score/metal
  tmol/score/metal/potentials/compiled.ops.cpp
  tmol/score/metal/potentials/metal_coordination_pose_score.cpu.cpp
  # score/na_torsion (CUDA implementation registered separately below)
  tmol/score/na_torsion/potentials/compiled.ops.cpp
)

if(TMOL_HAS_CUDA)
  list(APPEND _C_SOURCES
    tmol/score/common/whole_pose_scoring.cuda.cu
    tmol/optimization/armijo_compiled.cuda.cu
    tmol/io/details/compiled/gen_pose_leaf_atoms.cuda.cu
    tmol/io/details/compiled/resolve_his_taut.cuda.cu
    tmol/kinematics/compiled/compiled.cuda.cu
    tmol/pack/compiled/compiled.cuda.cu
    tmol/pack/rotamer/dunbrack/compiled.cuda.cu
    tmol/score/backbone_torsion/potentials/backbone_torsion_pose_score.cuda.cu
    tmol/score/cartbonded/potentials/cartbonded_pose_score.cuda.cu
    tmol/score/genbonded/potentials/genbonded_pose_score.cuda.cu
    tmol/score/constraint/potentials/constraint_score.cuda.cu
    tmol/score/disulfide/potentials/disulfide_pose_score.cuda.cu
    tmol/score/dunbrack/potentials/dunbrack_pose_score.cuda.cu
    tmol/score/elec/potentials/elec_pose_score.cuda.cu
    tmol/score/hbond/potentials/hbond_pose_score.cuda.cu
    tmol/score/hbond/potentials/gen_hbond_bases.cuda.cu
    tmol/score/ljlk/potentials/ljlk_pose_score.cuda.cu
    tmol/score/ljlk/potentials/ljlk_elec_pose_score.cuda.cu
    tmol/score/lk_ball/potentials/lk_ball_pose_score.cuda.cu
    tmol/score/lk_ball/potentials/gen_pose_waters.cuda.cu
    tmol/score/metal/potentials/metal_coordination_pose_score.cuda.cu
    tmol/score/na_torsion/potentials/na_torsion_pose_score.cuda.cu
  )
  # Approximate CUDA math is accurate for the pairwise nonbonded terms and
  # materially improves their throughput. Keep every other term exact:
  # torsion-heavy gradients are substantially more sensitive to approximation.
  set_source_files_properties(
    tmol/score/elec/potentials/elec_pose_score.cuda.cu
    tmol/score/hbond/potentials/hbond_pose_score.cuda.cu
    tmol/score/ljlk/potentials/ljlk_pose_score.cuda.cu
    tmol/score/ljlk/potentials/ljlk_elec_pose_score.cuda.cu
    tmol/score/lk_ball/potentials/lk_ball_pose_score.cuda.cu
    PROPERTIES COMPILE_OPTIONS "--use_fast_math"
  )
endif()

pybind11_add_module(_C MODULE ${_C_SOURCES})

# Register native operators through the stable dispatcher ABI.
set_target_properties(_C PROPERTIES
  PREFIX ""
  OUTPUT_NAME "_C"
  LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/${_TMOL_NATIVE_PACKAGE}"
)
execute_process(
  COMMAND ${Python_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX'))"
  OUTPUT_VARIABLE _PY_EXT_SUFFIX
  OUTPUT_STRIP_TRAILING_WHITESPACE
)
set_target_properties(_C PROPERTIES SUFFIX "${_PY_EXT_SUFFIX}")

if(TMOL_HAS_CUDA)
  tmol_set_cuda_flags(_C)
else()
  tmol_set_cpp_flags(_C)
endif()
install(TARGETS _C DESTINATION "${_TMOL_NATIVE_PACKAGE}")


# ═══════════════════════════════════════════════════════════════════════════════
# 2. Production pybind11 modules
# ═══════════════════════════════════════════════════════════════════════════════

# CPU helpers belong to the base package.
if(NOT TMOL_OPS_CUDA)
# bspline (CPU-only)
tmol_add_pybind_cpp_ext(
  bspline_compiled "tmol/numeric/bspline_compiled" "_compiled"
  tmol/numeric/bspline_compiled/bspline.pybind.cpp
)

# cubic hermite polynomial (CPU-only)
tmol_add_pybind_cpp_ext(
  cubic_hermite "tmol/score/common" "_cubic_hermite_polynomial"
  tmol/score/common/_cubic_hermite_polynomial.cpp
)


endif()

# ═══════════════════════════════════════════════════════════════════════════════
# 3. Test extensions (optional, behind TMOL_BUILD_TESTS)
# ═══════════════════════════════════════════════════════════════════════════════

option(TMOL_BUILD_TESTS "Build test C++/CUDA extensions" OFF)

if(TMOL_BUILD_TESTS)

  # --- pybind CPU-only test extensions ---

  tmol_add_pybind_cpp_ext(
    test_hbond_ext "tmol/tests/score/hbond/potentials" "_ext"
    tmol/tests/score/hbond/potentials/compiled.pybind.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_ljlk_ext "tmol/tests/score/ljlk/potentials" "_ext"
    tmol/tests/score/ljlk/potentials/compiled.pybind.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_lkball_ext "tmol/tests/score/lk_ball/potentials" "_ext"
    tmol/tests/score/lk_ball/potentials/compiled.pybind.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_polynomial_ext "tmol/tests/score/common/polynomial" "_ext"
    tmol/tests/score/common/polynomial/polynomial.pybind.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_geom_ext "tmol/tests/score/common/geom" "_ext"
    tmol/tests/score/common/geom/geom.pybind.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_uaid_util "tmol/tests/score/common" "_uaid_util"
    tmol/tests/score/common/uaid_util.pybind.cc
  )

  tmol_add_pybind_cpp_ext(
    test_tensor_struct "tmol/tests/utility/tensor" "_tensor_struct"
    tmol/tests/utility/tensor/tensor_struct.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_tensor_accessor "tmol/tests/utility/tensor" "_tensor_accessor"
    tmol/tests/utility/tensor/tensor_accessor.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_tensor_collection "tmol/tests/utility/tensor" "_tensor_collection"
    tmol/tests/utility/tensor/tensor_collection.cpp
  )

  tmol_add_pybind_cpp_ext(
    test_in_place_heap "tmol/tests/utility/datastructures" "_in_place_heap"
    tmol/tests/utility/datastructures/in_place_heap.cpp
  )

  # --- mixed-backend and CUDA-only test extensions ---

  if(TMOL_HAS_CUDA)
    tmol_add_pybind_cuda_ext(
      test_bspline_cuda "tmol/tests/numeric" "_bspline_cuda"
      tmol/tests/numeric/bspline.cuda.cu
    )

    tmol_add_pybind_cuda_ext(
      test_geom_cuda "tmol/tests/score/common/geom" "_ext_cuda"
      tmol/tests/score/common/geom/geom.cu
    )

    tmol_add_pybind_cuda_ext(
      test_bonded_atom "tmol/tests/score/bonded_atom" "_ext"
      tmol/tests/score/bonded_atom/test.pybind.cpp
      tmol/tests/score/bonded_atom/test_cpu.cpp
      tmol/tests/score/bonded_atom/test_cuda.cu
    )

    tmol_add_pybind_cuda_ext(
      test_device_operations "tmol/tests/score/common/device_operations" "_ext"
      tmol/tests/score/common/device_operations/test.pybind.cpp
      tmol/tests/score/common/device_operations/test_cpu.cpp
      tmol/tests/score/common/device_operations/test_cuda.cu
    )

    tmol_add_pybind_cuda_ext(
      test_warp_segreduce "tmol/tests/score/common" "_warp_segreduce"
      tmol/tests/score/common/warp_segreduce.cpp
      tmol/tests/score/common/warp_segreduce.cuda.cu
    )

    tmol_add_pybind_cuda_ext(
      test_warp_stride_reduce "tmol/tests/score/common" "_warp_stride_reduce"
      tmol/tests/score/common/warp_stride_reduce.cpp
      tmol/tests/score/common/warp_stride_reduce.cuda.cu
    )

    tmol_add_pybind_cuda_ext(
      test_segscan "tmol/tests/kinematics/segscan" "_ext"
      tmol/tests/kinematics/segscan/segscan.cu
    )

    tmol_add_pybind_cuda_ext(
      test_dunbrack_ext "tmol/tests/pack/rotamer/dunbrack" "_ext"
      tmol/tests/pack/rotamer/dunbrack/compiled.pybind.cpp
      tmol/tests/pack/rotamer/dunbrack/test_cpu.cpp
      tmol/tests/pack/rotamer/dunbrack/test_cuda.cu
    )
  else()
    # These extensions exercise both backends. Keep their CPU tests available
    # when CI intentionally builds without a CUDA toolchain.
    tmol_add_pybind_cpp_ext(
      test_bonded_atom "tmol/tests/score/bonded_atom" "_ext"
      tmol/tests/score/bonded_atom/test.pybind.cpp
      tmol/tests/score/bonded_atom/test_cpu.cpp
    )

    tmol_add_pybind_cpp_ext(
      test_device_operations "tmol/tests/score/common/device_operations" "_ext"
      tmol/tests/score/common/device_operations/test.pybind.cpp
      tmol/tests/score/common/device_operations/test_cpu.cpp
    )

    tmol_add_pybind_cpp_ext(
      test_dunbrack_ext "tmol/tests/pack/rotamer/dunbrack" "_ext"
      tmol/tests/pack/rotamer/dunbrack/compiled.pybind.cpp
      tmol/tests/pack/rotamer/dunbrack/test_cpu.cpp
    )
  endif()

endif()
