include(ExternalProject)

find_package(Threads)

# NOTE: A git version of 1.6.5 or later is required if this download method is used.
find_package(Git QUIET REQUIRED)

set(FLASH_ATTENTION_PROJECT flash_attention)

set(FLASH_ATTENTION_URL
    https://oneflow-static.oss-cn-beijing.aliyuncs.com/third_party_mirror/flash-attention-v2-eed2e82b880e06237af3e50ceac4cf6728b15645.zip
)

set(FLASH_ATTENTION_INSTALL_DIR ${THIRD_PARTY_DIR}/flash_attention)
set(FLASH_ATTENTION_INCLUDE_DIR ${FLASH_ATTENTION_INSTALL_DIR}/include CACHE PATH "" FORCE)
set(FLASH_ATTENTION_LIBRARY_DIR ${FLASH_ATTENTION_INSTALL_DIR}/lib CACHE PATH "" FORCE)
set(FLASH_ATTENTION_LIBRARIES ${FLASH_ATTENTION_LIBRARY_DIR}/libflash_attention.so)

if(THIRD_PARTY)
  ExternalProject_Add(
    ${FLASH_ATTENTION_PROJECT}
    PREFIX flash_attention
    URL ${FLASH_ATTENTION_URL}
    URL_HASH MD5=63192a05973f614aff594a8bd11813ce
    UPDATE_COMMAND ""
    BUILD_BYPRODUCTS ${FLASH_ATTENTION_LIBRARIES}
    CMAKE_ARGS -DCMAKE_BUILD_TYPE:STRING=${CMAKE_BUILD_TYPE}
               -DCMAKE_CXX_FLAGS:STRING=${CMAKE_CXX_FLAGS}
               -DCMAKE_CXX_FLAGS_DEBUG:STRING=${CMAKE_CXX_FLAGS_DEBUG}
               -DCMAKE_CXX_FLAGS_RELEASE:STRING=${CMAKE_CXX_FLAGS_RELEASE}
               -DCMAKE_CUDA_ARCHITECTURES:STRING=${CMAKE_CUDA_ARCHITECTURES}
    CMAKE_CACHE_ARGS
      -DCMAKE_CUDA_COMPILER:STRING=${CUDAToolkit_NVCC_EXECUTABLE}
      -DCMAKE_C_COMPILER_LAUNCHER:STRING=${CMAKE_C_COMPILER_LAUNCHER}
      -DCMAKE_CXX_COMPILER_LAUNCHER:STRING=${CMAKE_CXX_COMPILER_LAUNCHER}
      -DCMAKE_INSTALL_PREFIX:PATH=${FLASH_ATTENTION_INSTALL_DIR}
      -DCMAKE_INSTALL_LIBDIR:PATH=${FLASH_ATTENTION_LIBRARY_DIR}
      -DCMAKE_INSTALL_MESSAGE:STRING=${CMAKE_INSTALL_MESSAGE}
      -DCMAKE_BUILD_TYPE:STRING=${CMAKE_BUILD_TYPE})
endif(THIRD_PARTY)