cmake_minimum_required(VERSION 3.24...3.31)

option(TMOL_ENABLE_CUDA "Build CUDA extensions (disable for CPU-only)" ON)

# ═══════════════════════════════════════════════════════════════════════════════
# 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()

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

find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)

# Save Python_INCLUDE_DIRS NOW — find_package(Torch) / Caffe2 may overwrite it.
set(_TMOL_PYTHON_INCLUDE_DIRS "${Python_INCLUDE_DIRS}")
if(NOT _TMOL_PYTHON_INCLUDE_DIRS)
  execute_process(
    COMMAND ${Python_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_path('include'))"
    OUTPUT_VARIABLE _TMOL_PYTHON_INCLUDE_DIRS
    OUTPUT_STRIP_TRAILING_WHITESPACE
  )
endif()
message(STATUS "tmol: Python include dirs = ${_TMOL_PYTHON_INCLUDE_DIRS}")

# PyTorch 2.13 and newer require C++20 for third-party extensions. Keep the
# older release lanes on C++17 so their established CUDA toolchains remain
# unchanged.
execute_process(
  COMMAND ${Python_EXECUTABLE} -c "import torch; v=torch.__version__.split('.'); print(f'{v[0]};{v[1]}')"
  OUTPUT_VARIABLE _TORCH_VERSION
  OUTPUT_STRIP_TRAILING_WHITESPACE
)
list(GET _TORCH_VERSION 0 TORCH_VERSION_MAJOR)
list(GET _TORCH_VERSION 1 TORCH_VERSION_MINOR)

if(TORCH_VERSION_MAJOR GREATER 2
    OR (TORCH_VERSION_MAJOR EQUAL 2 AND TORCH_VERSION_MINOR GREATER_EQUAL 13))
  set(_TMOL_CXX_STANDARD 20)
else()
  set(_TMOL_CXX_STANDARD 17)
endif()
set(CMAKE_CXX_STANDARD ${_TMOL_CXX_STANDARD})
set(CMAKE_CXX_STANDARD_REQUIRED ON)
if(TMOL_HAS_CUDA)
  set(CMAKE_CUDA_STANDARD ${_TMOL_CXX_STANDARD})
  set(CMAKE_CUDA_STANDARD_REQUIRED ON)
endif()
message(STATUS "tmol: PyTorch ${TORCH_VERSION_MAJOR}.${TORCH_VERSION_MINOR}, using C++${_TMOL_CXX_STANDARD}")

if(TMOL_HAS_CUDA)
  # Set TORCH_CUDA_ARCH_LIST env var BEFORE find_package(Torch) so that
  # Caffe2's cuda.cmake uses our architectures instead of defaulting to
  # all archs nvcc supports (which includes sm_52 that lacks double atomicAdd).
  # CMake resolves the symbolic `native` value while enabling CUDA; use that
  # numeric result rather than passing an empty architecture list to Torch.
  if(CMAKE_CUDA_ARCHITECTURES STREQUAL "native")
    set(_TMOL_TORCH_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES_NATIVE})
  else()
    set(_TMOL_TORCH_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES})
  endif()
  set(_TORCH_ARCH_STR "")
  foreach(_A ${_TMOL_TORCH_CUDA_ARCHITECTURES})
    string(REPLACE "-real" "" _A "${_A}")
    string(REPLACE "-virtual" "" _A "${_A}")
    string(LENGTH "${_A}" _ALEN)
    if(_ALEN EQUAL 3)
      string(SUBSTRING "${_A}" 0 2 _MAJOR)
      string(SUBSTRING "${_A}" 2 -1 _MINOR)
      set(_TORCH_ARCH_STR "${_TORCH_ARCH_STR} ${_MAJOR}.${_MINOR}")
    elseif(_ALEN EQUAL 2)
      string(SUBSTRING "${_A}" 0 1 _MAJOR)
      string(SUBSTRING "${_A}" 1 -1 _MINOR)
      set(_TORCH_ARCH_STR "${_TORCH_ARCH_STR} ${_MAJOR}.${_MINOR}")
    endif()
  endforeach()
  string(STRIP "${_TORCH_ARCH_STR}" _TORCH_ARCH_STR)
  set(ENV{TORCH_CUDA_ARCH_LIST} "${_TORCH_ARCH_STR}")
  message(STATUS "tmol: Setting TORCH_CUDA_ARCH_LIST=${_TORCH_ARCH_STR} for find_package(Torch)")
endif()

# A Python wheel installs TorchConfig.cmake below the torch package rather than
# a system CMake prefix. Discover it from the interpreter selected above so a
# normal source build does not require a hand-written CMAKE_PREFIX_PATH.
execute_process(
  COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
  OUTPUT_VARIABLE _TMOL_TORCH_CMAKE_PREFIX_PATH
  RESULT_VARIABLE _TMOL_TORCH_PREFIX_RESULT
  ERROR_VARIABLE _TMOL_TORCH_PREFIX_ERROR
  OUTPUT_STRIP_TRAILING_WHITESPACE
)
if(NOT _TMOL_TORCH_PREFIX_RESULT EQUAL 0 OR NOT _TMOL_TORCH_CMAKE_PREFIX_PATH)
  message(FATAL_ERROR
    "tmol: could not locate Torch's CMake package with ${Python_EXECUTABLE}: "
    "${_TMOL_TORCH_PREFIX_ERROR}")
endif()
list(PREPEND CMAKE_PREFIX_PATH "${_TMOL_TORCH_CMAKE_PREFIX_PATH}")

execute_process(
  COMMAND ${Python_EXECUTABLE} -m pybind11 --cmakedir
  OUTPUT_VARIABLE _TMOL_PYBIND11_CMAKE_PREFIX_PATH
  RESULT_VARIABLE _TMOL_PYBIND11_PREFIX_RESULT
  ERROR_VARIABLE _TMOL_PYBIND11_PREFIX_ERROR
  OUTPUT_STRIP_TRAILING_WHITESPACE
)
if(NOT _TMOL_PYBIND11_PREFIX_RESULT EQUAL 0 OR NOT _TMOL_PYBIND11_CMAKE_PREFIX_PATH)
  message(FATAL_ERROR
    "tmol: could not locate pybind11's CMake package with ${Python_EXECUTABLE}: "
    "${_TMOL_PYBIND11_PREFIX_ERROR}")
endif()
list(PREPEND CMAKE_PREFIX_PATH "${_TMOL_PYBIND11_CMAKE_PREFIX_PATH}")

find_package(Torch REQUIRED)
find_package(pybind11 CONFIG REQUIRED)
# Linux Torch wheels require OpenMP compiler flags for at::parallel_for;
# platforms using Torch's native thread pool do not.
find_package(OpenMP COMPONENTS CXX QUIET)

find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib" NO_DEFAULT_PATH)
message(STATUS "tmol: TORCH_PYTHON_LIBRARY = ${TORCH_PYTHON_LIBRARY}")

if(TMOL_HAS_CUDA)
  # Caffe2's TorchConfig.cmake can set CMAKE_CUDA_ARCHITECTURES=OFF or inherit
  # PyTorch's full -gencode list (including arches this nvcc cannot compile).
  set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_CUDA_ARCHITECTURES}" CACHE STRING "CUDA architectures" FORCE)
  set(CMAKE_CUDA_ARCHITECTURES "${_TMOL_CUDA_ARCHITECTURES}")
  message(STATUS "tmol: Restored CMAKE_CUDA_ARCHITECTURES=${CMAKE_CUDA_ARCHITECTURES}")

  # Caffe2's TorchConfig.cmake unconditionally appends the -gencode flags that
  # PyTorch was originally built with. Replace them with our architectures.
  string(REGEX REPLACE "-gencode[;= ]arch=compute_[0-9]+,code=sm_[0-9]+" "" CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
  string(REGEX REPLACE "-gencode[;= ]arch=compute_[0-9]+,code=compute_[0-9]+" "" CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
  string(REGEX REPLACE " +" " " CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
  string(STRIP "${CMAKE_CUDA_FLAGS}" CMAKE_CUDA_FLAGS)

  # CMake emits the requested architecture flags for each CUDA target. Do not
  # append another copy here: doing so compiles every device image twice.
  message(STATUS "tmol: CMAKE_CUDA_FLAGS after removing Torch gencode = ${CMAKE_CUDA_FLAGS}")
endif()

# ═══════════════════════════════════════════════════════════════════════════════
# Compiler flags — match PyTorch's ABI and set tmol defaults
# ═══════════════════════════════════════════════════════════════════════════════

# TORCH_CXX_FLAGS contains -D_GLIBCXX_USE_CXX11_ABI=0 or =1
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
message(STATUS "tmol: TORCH_CXX_FLAGS = ${TORCH_CXX_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")
  endif()
endfunction()

# Common compile options for CUDA extensions
function(tmol_set_cuda_flags TARGET)
  target_compile_options(${TARGET} PRIVATE
    $<$<COMPILE_LANGUAGE:CXX>:-O3 -w -DWITH_CUDA -DWITH_NVTX>
    $<$<COMPILE_LANGUAGE:CUDA>:
      -O3 -w
      --expt-extended-lambda
      --expt-relaxed-constexpr
      -DWITH_NVTX -DWITH_CUDA
      --threads=${TMOL_NVCC_THREADS}
      -DTORCH_VERSION_MAJOR=${TORCH_VERSION_MAJOR}
      -DTORCH_VERSION_MINOR=${TORCH_VERSION_MINOR}
    >
  )
  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)
  target_compile_options(${TARGET} PRIVATE
    $<$<COMPILE_LANGUAGE:CXX>:-O3 -w -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})
  target_link_libraries(${TARGET_NAME} PRIVATE ${TORCH_PYTHON_LIBRARY})
  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})
  target_link_libraries(${TARGET_NAME} PRIVATE ${TORCH_PYTHON_LIBRARY})
  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. tmol._C — monolithic TORCH_LIBRARY extension (14 op namespaces)
# ═══════════════════════════════════════════════════════════════════════════════

set(_C_SOURCES
  # 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
  # pose/compiled/apsp
  tmol/pose/compiled/apsp_vestibule.ops.cpp
  tmol/pose/compiled/apsp.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
  # 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/na_torsion (CUDA implementation registered separately below)
  tmol/score/na_torsion/potentials/compiled.ops.cpp
)

if(TMOL_HAS_CUDA)
  list(APPEND _C_SOURCES
    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/pose/compiled/apsp.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/lk_ball/potentials/lk_ball_pose_score.cuda.cu
    tmol/score/lk_ball/potentials/gen_pose_waters.cuda.cu
    tmol/score/na_torsion/potentials/na_torsion_pose_score.cuda.cu
  )
endif()

add_library(_C MODULE ${_C_SOURCES})

# _C is a MODULE library loaded by torch, not a regular Python extension.
# No pybind11 entry point — ops register via TORCH_LIBRARY at dlopen time.
set_target_properties(_C PROPERTIES
  PREFIX ""
  OUTPUT_NAME "_C"
  LIBRARY_OUTPUT_DIRECTORY "${CMAKE_SOURCE_DIR}/tmol"
)
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()
if(OpenMP_CXX_FOUND)
  target_link_libraries(_C PRIVATE OpenMP::OpenMP_CXX)
endif()
install(TARGETS _C DESTINATION tmol)


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

# 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
)


# ═══════════════════════════════════════════════════════════════════════════════
# 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_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()
