StepMesh/CMakeLists.txt at main · stepfun-ai/StepMesh · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
cmake_minimum_required(VERSION 3.22 FATAL_ERROR)

project(af LANGUAGES C CXX)

set(CMAKE_CXX_STANDARD 17)
execute_process(COMMAND ${Python_EXECUTABLE}
        -c "import torch; print(int(torch.compiled_with_cxx11_abi()))"
        OUTPUT_VARIABLE TORCH_CXX11_ABI OUTPUT_STRIP_TRAILING_WHITESPACE)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++17 -fPIC -O3 -Wall -finline-functions -msse2 -D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI} ")


# set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++17 -fPIC -O3 -Wall -finline-functions  -msse2 ")

# import pytorch library
find_package (Python COMPONENTS Interpreter Development)
execute_process(COMMAND ${Python_EXECUTABLE}
        -c "import torch; print(torch.utils.cmake_prefix_path)"
        OUTPUT_VARIABLE PYTORCH_CMAKE_PREFIX_PATH OUTPUT_STRIP_TRAILING_WHITESPACE)
message("MY PYTHON_EXECUTABLE ${Python_EXECUTABLE}")
message("MY PYTORCH_CMAKE_PREFIX_PATH ${PYTORCH_CMAKE_PREFIX_PATH}")

list(APPEND CMAKE_PREFIX_PATH "${PYTORCH_CMAKE_PREFIX_PATH}/Torch")

find_package(Torch REQUIRED CONFIG)
message("MY TORCH_INCLUDE_DIRS ${TORCH_INCLUDE_DIRS}")
message("MY CUDA_INCLUDE_DIRS ${CUDA_INCLUDE_DIRS}")
include_directories(${TORCH_INCLUDE_DIRS})
# Save ABI setting before adding TORCH_CXX_FLAGS (which might override it)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
# Ensure ABI setting is preserved after TORCH_CXX_FLAGS

list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules)


if("$ENV{USE_CUDA}" STREQUAL "0")
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DDMLC_USE_ZMQ -DDMLC_USE_RDMA -DSTEPMESH_USE_GDR")
else()
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DDMLC_USE_ZMQ -DDMLC_USE_CUDA -DSTEPMESH_USE_GDR -DDMLC_USE_RDMA -DSTEPMESH_ENABLE_TRACE")
endif()

link_directories("${PROJECT_SOURCE_DIR}/deps/lib")

list(APPEND STEPAF_LIBS pthread dl zmq ibverbs rdmacm rt)

include_directories("${PROJECT_SOURCE_DIR}/include/")
include_directories("${PROJECT_SOURCE_DIR}/deps/include/")
include_directories("${PROJECT_BINARY_DIR}/include/")
include_directories("${PROJECT_BINARY_DIR}/src/")

if("${CUDA_INCLUDE_DIRS}" STREQUAL "")
    message(WARNING "CUDA_INCLUDE_DIRS is empty")
else()
    include_directories("${CUDA_INCLUDE_DIRS}")
endif()

FILE(GLOB SOURCE "src/*.cc")
FILE(GLOB BACKEND_SOURCE "src/backend/*.cc")
message("MY SOURCES ${SOURCE}")
message("MY BACKEND_SOURCE ${BACKEND_SOURCE}")

add_library(af ${SOURCE} ${BACKEND_SOURCE})

target_link_libraries(af
        ${STEPAF_LIBS}
        ${TORCH_LIBRARIES}
        ${TORCH_PYTHON_LIBRARY})

add_subdirectory(tests)