cmake_minimum_required(VERSION 3.8)
project(qcu LANGUAGES C CXX CUDA)
set(CMAKE_CXX_STANDARD 11)
option(QCU_ENABLE_NVSHMEM "Enable the optional NVSHMEM coarse-grid transport" OFF)
message(STATUS "PATH is: $ENV{PATH}")
message(STATUS "LD_LIBRARY_PATH is: $ENV{LD_LIBRARY_PATH}")
# Architecture precedence: explicit CMake/environment override, env.sh's
# detected sm_arch, then the compute capability of the active GPU.  Auto
# detected values are emitted as cubin + PTX so they remain portable across a
# minor driver/toolkit mismatch; an explicit override is passed through
# unchanged.
set(_QCU_AUTO_CUDA_ARCHITECTURES "")
if(DEFINED QCU_CUDA_ARCHITECTURES AND NOT "${QCU_CUDA_ARCHITECTURES}" STREQUAL "")
  set(_QCU_CUDA_ARCHITECTURES "${QCU_CUDA_ARCHITECTURES}")
elseif(DEFINED CMAKE_CUDA_ARCHITECTURES AND NOT "${CMAKE_CUDA_ARCHITECTURES}" STREQUAL "")
  set(_QCU_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}")
elseif(DEFINED ENV{QCU_CUDA_ARCHITECTURES} AND NOT "$ENV{QCU_CUDA_ARCHITECTURES}" STREQUAL "")
  set(_QCU_CUDA_ARCHITECTURES "$ENV{QCU_CUDA_ARCHITECTURES}")
elseif(DEFINED ENV{sm_arch} AND NOT "$ENV{sm_arch}" STREQUAL "")
  set(_QCU_AUTO_CUDA_ARCHITECTURES "$ENV{sm_arch}")
  message(STATUS "QCU CUDA architecture auto-detected from env.sh: sm_${_QCU_AUTO_CUDA_ARCHITECTURES}")
else()
  find_program(_QCU_NVIDIA_SMI NAMES nvidia-smi)
  if(_QCU_NVIDIA_SMI)
    execute_process(
      COMMAND "${_QCU_NVIDIA_SMI}" --query-gpu=compute_cap --format=csv,noheader
      RESULT_VARIABLE _QCU_NVIDIA_SMI_RESULT
      OUTPUT_VARIABLE _QCU_NVIDIA_SMI_OUTPUT
      ERROR_QUIET
      OUTPUT_STRIP_TRAILING_WHITESPACE)
    if(_QCU_NVIDIA_SMI_RESULT EQUAL 0 AND NOT "${_QCU_NVIDIA_SMI_OUTPUT}" STREQUAL "")
      string(REPLACE "\n" ";" _QCU_GPU_CAPS "${_QCU_NVIDIA_SMI_OUTPUT}")
      list(GET _QCU_GPU_CAPS 0 _QCU_GPU_CAP)
      string(STRIP "${_QCU_GPU_CAP}" _QCU_GPU_CAP)
      string(REGEX MATCH "^([0-9]+)\\.([0-9]+)$" _QCU_GPU_CAP_MATCH "${_QCU_GPU_CAP}")
      if(_QCU_GPU_CAP_MATCH)
        set(_QCU_AUTO_CUDA_ARCHITECTURES "${CMAKE_MATCH_1}${CMAKE_MATCH_2}")
        message(STATUS "QCU CUDA architecture auto-detected from nvidia-smi: sm_${_QCU_AUTO_CUDA_ARCHITECTURES}")
      endif()
    endif()
  endif()
endif()

if(_QCU_AUTO_CUDA_ARCHITECTURES)
  set(_QCU_CUDA_ARCHITECTURES
      "${_QCU_AUTO_CUDA_ARCHITECTURES}-real;${_QCU_AUTO_CUDA_ARCHITECTURES}-virtual")
endif()

if(NOT DEFINED _QCU_CUDA_ARCHITECTURES OR "${_QCU_CUDA_ARCHITECTURES}" STREQUAL "")
  message(FATAL_ERROR
    "Unable to determine CUDA architecture. Set QCU_CUDA_ARCHITECTURES "
    "(for example 70-real;70-virtual), pass -DCMAKE_CUDA_ARCHITECTURES, "
    "or source env.sh on a machine with nvidia-smi.")
endif()

set(QCU_CUDA_ARCHITECTURES "${_QCU_CUDA_ARCHITECTURES}")
set(CMAKE_CUDA_ARCHITECTURES "${_QCU_CUDA_ARCHITECTURES}")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -O3 --expt-relaxed-constexpr")
message(STATUS "QCU CUDA architectures: ${QCU_CUDA_ARCHITECTURES}")
set(MPI_INCLUDE_PATH $ENV{mpi_home}/include)
set(MPI_CXX_LIBRARIES $ENV{mpi_home}/lib/libmpi.so)
include(GNUInstallDirs)
include_directories(${LD_LIBRARY_PATH})
include_directories(${MPI_INCLUDE_PATH})
include_directories(./include)
aux_source_directory(./src SOURCE_FILE)
add_library(qcu SHARED ${SOURCE_FILE})
find_package(MPI REQUIRED)
find_package(CUDA REQUIRED)
target_link_libraries(qcu ${MPI_C_LIBRARIES})
target_link_libraries(qcu ${CUDA_LIBRARIES})
target_link_libraries(qcu curand)
target_link_libraries(qcu cublas)

# NVSHMEM is deliberately opt-in.  A normal CUDA/MPI build must not need the
# headers or library, while an explicit request must fail during configure
# instead of producing a binary whose optional transport silently disappears.
if(QCU_ENABLE_NVSHMEM)
  find_path(QCU_NVSHMEM_INCLUDE_DIR
            NAMES nvshmem.h
            HINTS $ENV{NVSHMEM_HOME} $ENV{NVSHMEM_ROOT}
            PATH_SUFFIXES include)
  find_path(QCU_NVSHMEMX_INCLUDE_DIR
            NAMES nvshmemx.h
            HINTS $ENV{NVSHMEM_HOME} $ENV{NVSHMEM_ROOT}
            PATH_SUFFIXES include)
  find_library(QCU_NVSHMEM_LIBRARY
               NAMES nvshmem
               HINTS $ENV{NVSHMEM_HOME} $ENV{NVSHMEM_ROOT}
               PATH_SUFFIXES lib lib64)
  if(NOT QCU_NVSHMEM_INCLUDE_DIR OR
     NOT QCU_NVSHMEMX_INCLUDE_DIR OR
     NOT QCU_NVSHMEM_LIBRARY)
    message(FATAL_ERROR
      "QCU_ENABLE_NVSHMEM=ON requires nvshmem.h, nvshmemx.h and libnvshmem")
  endif()
  target_include_directories(qcu PRIVATE
    ${QCU_NVSHMEM_INCLUDE_DIR} ${QCU_NVSHMEMX_INCLUDE_DIR})
  target_compile_definitions(qcu PRIVATE QCU_HAVE_NVSHMEM=1)
  target_link_libraries(qcu ${QCU_NVSHMEM_LIBRARY})
  message(STATUS "NVSHMEM enabled: ${QCU_NVSHMEM_LIBRARY}")
else()
  message(STATUS "NVSHMEM disabled (set -DQCU_ENABLE_NVSHMEM=ON to enable)")
endif()
