cmake_policy(SET CMP0079 NEW)

function(sd_cuda_sources VENDOR LANGUAGE)
    set(SOURCES)
    foreach(OP rope int8 w4 convrot affine modulate bf16 rms_affine registration)
        set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
        string(TOUPPER "${VENDOR}" UPPER_VENDOR)
        file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
        set_source_files_properties("${SOURCE}" PROPERTIES LANGUAGE ${LANGUAGE})
        list(APPEND SOURCES "${SOURCE}")
    endforeach()
    set(TARGET "sd-${VENDOR}")
    add_library(${TARGET} OBJECT ${SOURCES})
    target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src")
    target_link_libraries(${TARGET} PRIVATE ggml)
    set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
    target_compile_features(${TARGET} PRIVATE cxx_std_17)
    target_link_libraries(${SD_LIB} PRIVATE ${TARGET})
    set(SD_CUDA_TARGET ${TARGET} PARENT_SCOPE)
    set(SD_CUDA_SOURCES ${SOURCES} PARENT_SCOPE)
endfunction()

if(SD_CUDA)
    enable_language(CUDA)
    sd_cuda_sources(cuda CUDA)
    target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu sol_attention.cu)
    set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
    if(SD_CUDA_CUTLASS_PATH)
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_W4_CUTLASS)
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include")
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS --expt-relaxed-constexpr)
        if(MSVC)
            set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
        endif()
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS)
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${PROJECT_SOURCE_DIR}/src/core/cuda")
    endif()
    set_target_properties(${SD_CUDA_TARGET} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
    target_link_libraries(${SD_CUDA_TARGET} PRIVATE CUDA::cudart CUDA::cublas)
    target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
endif()

if(SD_HIPBLAS)
    list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}" /opt/rocm /usr)
    find_package(hip REQUIRED)
    find_package(hipblas REQUIRED)
    if(WIN32 OR CMAKE_CXX_COMPILER MATCHES "hipcc")
        sd_cuda_sources(hip CXX)
        target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::device)
    else()
        sd_cuda_sources(hip HIP)
        set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
    endif()
    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
    # W4A8 WMMA uses the wave64 fragment layout on both GFX11 and GFX12.
    set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_hip.cu" APPEND PROPERTY COMPILE_OPTIONS "-mwavefrontsize64")
    target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
    target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
endif()

if(SD_MUSA)
    if(NOT MUSA_PATH)
        if(DEFINED ENV{MUSA_PATH})
            set(MUSA_PATH "$ENV{MUSA_PATH}")
        elseif(EXISTS /opt/musa)
            set(MUSA_PATH /opt/musa)
        else()
            set(MUSA_PATH /usr/local/musa)
        endif()
    endif()
    list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
    find_package(MUSAToolkit REQUIRED)
    set(CMAKE_CXX_COMPILER "${MUSA_PATH}/bin/clang++")
    sd_cuda_sources(musa CXX)
    set(FLAGS -Od3 -fno-strict-aliasing -ffast-math -fsigned-char -x musa -mtgpu -fmusa-flush-denormals-to-zero)
    if(NOT DEFINED MUSA_ARCHITECTURES)
        set(MUSA_ARCHITECTURES "22;31")
    endif()
    foreach(ARCH ${MUSA_ARCHITECTURES})
        list(APPEND FLAGS "--cuda-gpu-arch=mp_${ARCH}")
    endforeach()
    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "${FLAGS}")
    if(GGML_STATIC)
        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart_static MUSA::mublas_static)
        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart_static MUSA::mublas_static)
    else()
        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart MUSA::mublas)
        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart MUSA::mublas)
    endif()
endif()