# Math library

# NOTES
#
# 1. Use `target_link_libraries(xxx)` instead of `list(APPEND THIRD_PARTY_LIBS xxx)`. Global flag
#    THIRD_PARTY_LIBS is no longer in use.
#
# 2. Perform
#      set_env add CMAKE_PREFIX_PATH /hpc/software/intel/oneapi/mkl/2023.1.0
#      set_env add CMAKE_MODULE_PATH /hpc/software/intel/oneapi/mkl/2023.1.0/lib/cmake/mkl
#    if MKLConfig.cmake is not found on H3C hpc.
#
# 3. Reference for MKL:
#      https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl-link-line-advisor.html

# Summary of CMake flags and compile definitions:
#
# Overall CMake flags for math library vendors:
#   TBPLAS_USE_MKL, TBPLAS_USE_AOCL, TBPLAS_USE_OPENBLAS, TBPLAS_USE_NETLIB
#
# Interface CMake flags:
#   INTERFACE_ILP64: whether to use ilp64 interface, should be OFF for Eigen/Diag and ON for TBPM
#
# Misc:
#   TBPLAS_MATH_DEF: CMake flag for holding shared compile definitions
#
# Compile definitions for Eigen:
#   EIGEN_USE_MKL_ALL, EIGEN_USE_BLAS, EIGEN_USE_LAPACKE
#
# Compile definitions for Diag:
#   DIAG_USE_MKL: whether diag module uses MKL
#   WITH_LAPACKE: whether to enable LAPACKE support
#
# Compile definitions for TBPM:
#   TBPM_USE_MKL: whether to use sparse library from MKL
#   TBPM_USE_AOCL: whether to use sparse library from AOCL

# Reset all flags
set(TBPLAS_USE_MKL OFF)
set(TBPLAS_USE_AOCL OFF)
set(TBPLAS_USE_OPENBLAS OFF)
set(TBPLAS_USE_NETLIB OFF)
set(INTERFACE_ILP64 OFF)
set(TBPLAS_MATH_DEF "")

# Macros and functions
function(check_math_vendor)
    # Allowed math library vendors for each component
    set(eigen_vendors "^(mkl|aocl|openblas|netlib|default)$")
    set(diag_vendors "^(mkl|aocl|openblas|netlib|default)$")
    set(tbpm_vendors "^(mkl|aocl|default)$")

    # Check for illegal math library vendor
    if(NOT (${EIGEN_BACKEND} MATCHES ${eigen_vendors}))
        message(FATAL_ERROR "Unknown EIGEN_BACKEND: ${EIGEN_BACKEND}")
    endif()
    if(NOT (${DIAG_BACKEND} MATCHES ${diag_vendors}))
        message(FATAL_ERROR "Unknown DIAG_BACKEND: ${DIAG_BACKEND}")
    endif()
    if(NOT (${TBPM_BACKEND} MATCHES ${tbpm_vendors}))
        message(FATAL_ERROR "Unknown TBPM_BACKEND: ${TBPM_BACKEND}")
    endif()

    # Check for incompatible combinations
    if(NOT ("${EIGEN_BACKEND}" STREQUAL "default"))
        # If eigen_backend is non-default, diag_backend should be either default or
        # equal to eigen_backend. And tbpm_backend must be default.
        if((NOT ("${DIAG_BACKEND}" STREQUAL "${EIGEN_BACKEND}")) AND
            (NOT ("${DIAG_BACKEND}" STREQUAL "default")))
            message(FATAL_ERROR "Mixed usage of ${EIGEN_BACKEND} and ${DIAG_BACKEND}")
        endif()
        if(NOT ("${TBPM_BACKEND}" STREQUAL "default"))
            message(FATAL_ERROR "TBPM_BACKEND must be default when EIGEN_BACKEND is ${EIGEN_BACKEND}")
        endif()
    elseif(NOT ("${DIAG_BACKEND}" STREQUAL "default"))
        # If eigen_backend is default, then the choice of diag_backend is free.
        # But tbpm_backend cannot conflict with diag_backend.
        if(NOT ("${TBPM_BACKEND}" STREQUAL "default"))
            message(FATAL_ERROR "TBPM_BACKEND must be default when DIAG_BACKEND is ${DIAG_BACKEND}")
        endif()
    endif()
endfunction()

function(check_aocl)
    # Check if AOCL is properly installed
    if("$ENV{AOCL_ROOT}" STREQUAL "")
        message(FATAL_ERROR "AOCL_ROOT not set")
    else()
        message(STATUS "AOCL_ROOT: " "$ENV{AOCL_ROOT}")
    endif()

    # Check targets of symbolic links under AOCL_ROOT
    execute_process(
        COMMAND readlink "$ENV{AOCL_ROOT}/lib"
        OUTPUT_VARIABLE link_type
        OUTPUT_STRIP_TRAILING_WHITESPACE
    )
    if(INTERFACE_ILP64)
        if(NOT (${link_type} MATCHES "lib_ILP64$"))
            message(FATAL_ERROR "AOCL links not set to ilp64")
        endif()
    else()
        if(NOT (${link_type} MATCHES "lib_LP64$"))
            message(FATAL_ERROR "AOCL links not set to lp64")
        endif()
    endif()
endfunction()

function(check_feast)
    # Require LAPACK
    if(NOT (TBPLAS_USE_MKL OR TBPLAS_USE_AOCL OR TBPLAS_USE_OPENBLAS OR TBPLAS_USE_NETLIB))
        message(FATAL_ERROR "FEAST requires a working LAPACK library")
    endif()

    # Require LP64 interface of MKL/AOCL
    if(INTERFACE_ILP64)
        message(FATAL_ERROR "FEAST requires MKL/AOCL LP64 interface")
    endif()

    # Require FEAST to be installed
    if("$ENV{FEASTROOT}" STREQUAL "")
        message(FATAL_ERROR "FEASTROOT not set")
    else()
        message(STATUS "FEASTROOT: " "$ENV{FEASTROOT}")
    endif()
endfunction()

# Backend for Eigen
#
# 1. For MKL and AOCL, we force using LP64 interface to avoid potential bugs.
# 2. Eigen is used throughout the project, so we make the compile definitions global.
check_math_vendor()
if(${EIGEN_BACKEND} MATCHES "^mkl$")
    set(TBPLAS_USE_MKL ON)
    add_compile_definitions(EIGEN_USE_MKL_ALL)
elseif(${EIGEN_BACKEND} MATCHES "^aocl$")
    set(TBPLAS_USE_AOCL ON)
    add_compile_definitions(EIGEN_USE_BLAS)
    add_compile_definitions(EIGEN_USE_LAPACKE)
elseif(${EIGEN_BACKEND} MATCHES "^openblas$")
    set(TBPLAS_USE_OPENBLAS ON)
    add_compile_definitions(EIGEN_USE_BLAS)
    add_compile_definitions(EIGEN_USE_LAPACKE)
elseif(${EIGEN_BACKEND} MATCHES "^netlib$")
    set(TBPLAS_USE_NETLIB ON)
    add_compile_definitions(EIGEN_USE_BLAS)
    add_compile_definitions(EIGEN_USE_LAPACKE)
endif()

# Backend for diagonalization
#
# 1. For MKL and AOCL, we force using LP64 interface to avoid potential bugs.
#
# 2. Header of LAPACKE is named mkl_lapacke.h in MKL and lapacke.h in other vendors,
#    so we need to pass DIAG_USE_MKL to the compiler to distinguish that.
#
# 3. LAPACKE solver is enabled automatically if DIAG_BACKEND supports it.
if(${DIAG_BACKEND} MATCHES "^mkl$")
    set(TBPLAS_USE_MKL ON)
    list(APPEND TBPLAS_MATH_DEF DIAG_USE_MKL)
    list(APPEND TBPLAS_MATH_DEF WITH_LAPACKE)
elseif(${DIAG_BACKEND} MATCHES "^aocl$")
    set(TBPLAS_USE_AOCL ON)
    list(APPEND TBPLAS_MATH_DEF WITH_LAPACKE)
elseif(${DIAG_BACKEND} MATCHES "^openblas$")
    set(TBPLAS_USE_OPENBLAS ON)
    list(APPEND TBPLAS_MATH_DEF WITH_LAPACKE)
elseif(${DIAG_BACKEND} MATCHES "^netlib$")
    set(TBPLAS_USE_NETLIB ON)
    list(APPEND TBPLAS_MATH_DEF WITH_LAPACKE)
endif()

# Backend for TBPM
#
# MKL/AOCL support for Eigen/Diag and for TBPM cannot be enabled in the same time, since Eigen/Diag
# requires LP64 interface, while TBPM requires ILP64.
#
# Although there are some workarounds for the issues, e.g., using g++ instead of icpx will allow the use
# of ILP64 MKL interface for Eigen/Diag, it is bug-prone and does not make much sense. So we require that
# MKL/AOCL support can only be enabled for either Eigen/Diag or TBPM.
if(${TBPM_BACKEND} MATCHES "^mkl$")
    set(TBPLAS_USE_MKL ON)
    set(INTERFACE_ILP64 ON)
    list(APPEND TBPLAS_MATH_DEF TBPM_USE_MKL)
elseif(${TBPM_BACKEND} MATCHES "^aocl$")
    set(TBPLAS_USE_AOCL ON)
    set(INTERFACE_ILP64 ON)
    list(APPEND TBPLAS_MATH_DEF TBPM_USE_AOCL)
endif()

# Math library vendor setting
if(TBPLAS_USE_MKL)
    # MKL_INTERFACE must go first
    if(INTERFACE_ILP64)
        set(MKL_INTERFACE "ilp64")
    else()
        set(MKL_INTERFACE "lp64")
    endif()
    find_package(MKL CONFIG REQUIRED)
elseif(TBPLAS_USE_AOCL)
    check_aocl()
elseif(TBPLAS_USE_OPENBLAS)
    find_package(OpenBLAS)
elseif(TBPLAS_USE_NETLIB)
    find_package(cblas)
    find_package(lapacke)
endif()

# FEAST
if(WITH_FEAST)
    check_feast()
endif()

# Compile definitions and header directories should be defined PUBLIC,
# otherwise there may be inconsistencies in headers, especially for C++
# which uses templates.
function(tbp_set_math_lib arg)
    # Add shared definitions
    target_compile_definitions(${arg} PUBLIC ${TBPLAS_MATH_DEF})

    if(TBPLAS_USE_MKL)
        target_compile_options(${arg} PUBLIC $<TARGET_PROPERTY:MKL::MKL,INTERFACE_COMPILE_OPTIONS>)
        target_include_directories(${arg} PUBLIC $<TARGET_PROPERTY:MKL::MKL,INTERFACE_INCLUDE_DIRECTORIES>)
        target_link_libraries(${arg} PUBLIC $<LINK_ONLY:MKL::MKL>)
        # For Intel compiler, add proper mkl options.
        if(TBPLAS_USE_INTEL)
            if(WITH_OPENMP)
                target_compile_options(${arg} PUBLIC -qmkl=parallel)
            else()
                target_compile_options(${arg} PUBLIC -qmkl=sequential)
            endif()
        endif()
    elseif(TBPLAS_USE_AOCL)
        # AOCL does not provide proper cmake files. Instead, it provides an initialization script
        # which sets the relevant environment variables and pkgconfig files for linking. So we
        # can link blis(BLAS) and flame(LAPACK) libraries according to guidelines in pkgconfig files.
        target_link_options(${arg} PUBLIC -fopenmp)
        target_link_libraries(${arg} PUBLIC blis-mt m pthread rt)
        target_link_libraries(${arg} PUBLIC flame pthread aoclutils)
        # If ilp64 then tbpm uses sparse library from aocl.
        if(INTERFACE_ILP64)
            # blas, lapack
            target_compile_definitions(${arg} PUBLIC AOCLDA_ILP64)
            # sparse
            target_compile_definitions(${arg} PUBLIC aoclsparse_ILP64)
            target_link_libraries(${arg} PUBLIC aoclsparse)
        endif()
    elseif(TBPLAS_USE_OPENBLAS)
        target_include_directories(${arg} PUBLIC ${OpenBLAS_INCLUDE_DIRS})
        target_link_libraries(${arg} PUBLIC ${OpenBLAS_LIBRARIES})
    elseif(TBPLAS_USE_NETLIB)
        target_include_directories(${arg} PUBLIC ${CBLAS_INCLUDE_DIRS})
        target_include_directories(${arg} PUBLIC ${LAPACKE_INCLUDE_DIRS})
        target_link_libraries(${arg} PUBLIC ${CBLAS_LIBRARIES})
        target_link_libraries(${arg} PUBLIC ${LAPACKE_LIBRARIES})
    endif()

    if(WITH_FEAST)
        target_compile_definitions(${arg} PUBLIC WITH_FEAST)
        target_include_directories(${arg} PUBLIC "$ENV{FEASTROOT}/include")
        target_link_directories(${arg} PUBLIC "$ENV{FEASTROOT}/lib/x64")

        # Link compiler-dependent fortran library
        # FEAST Eigenvalue Solver v4.0 User Guide, pp.11
        # FEAST does not support flang, so we need to compile it using gfortran and link
        # to its runtime library.
        if(TBPLAS_USE_INTEL)
            target_link_libraries(${arg} PUBLIC feast ifcore ifcoremt)
        elseif(TBPLAS_USE_CLANG)
            target_link_libraries(${arg} PUBLIC feast gfortran)
        elseif(TBPLAS_USE_GCC)
            target_link_libraries(${arg} PUBLIC feast gfortran)
        else()
            message(FATAL_ERROR "Fortran library for FEAST not set")
        endif()
    endif()
endfunction()