Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 71 additions & 23 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,15 @@ include_directories(${gflags_SOURCE_DIR}/include)
# glog
set(WITH_GFLAGS OFF CACHE BOOL "Disable glog finding system gflags" FORCE)
set(WITH_GTEST OFF CACHE BOOL "Disable glog finding system gtest" FORCE)
add_subdirectory(third_party/glog)
get_property(_infinitrain_build_testing_was_cached CACHE BUILD_TESTING PROPERTY TYPE SET)
block()
set(BUILD_TESTING OFF)
add_subdirectory(third_party/glog)
endblock()
if(NOT _infinitrain_build_testing_was_cached)
unset(BUILD_TESTING CACHE)
endif()
unset(_infinitrain_build_testing_was_cached)
include_directories(${glog_SOURCE_DIR}/src)

# eigen
Expand Down Expand Up @@ -131,19 +139,27 @@ endif()
# ------------------------------------------------------------------------------

add_library(infini_train STATIC ${SRC})
add_library(InfiniTrain::infini_train ALIAS infini_train)
Comment thread
kilinchange marked this conversation as resolved.
add_library(InfiniTrain::cpu_kernels ALIAS infini_train_cpu_kernels)
target_compile_features(infini_train PUBLIC cxx_std_20)
target_include_directories(infini_train PUBLIC
"$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>"
)
target_link_libraries(infini_train
PUBLIC
glog
gflags
infini_train_cpu_kernels
Eigen3::Eigen
)
if(USE_OMP)
target_link_libraries(infini_train PUBLIC OpenMP::OpenMP_CXX)
endif()

if(USE_CUDA)
# infini_train contains cuda runtime wrappers (*.cc) like cuda_blas_handle.cc/cuda_guard.cc
# Those may need CUDA runtime/driver/cublas symbols at final link, so attach them here too.
target_link_libraries(infini_train
PUBLIC
infini_train_cuda_kernels
CUDA::cudart
CUDA::cublas
CUDA::cuda_driver
Expand All @@ -157,37 +173,66 @@ if(USE_CUDA)
endif()

# ------------------------------------------------------------------------------
# Helper: link libraries in a group to fix static lib one-pass resolution
# (THIS is what fixes "undefined reference" from cuda_kernels -> core symbols)
# Final executable link interface
# ------------------------------------------------------------------------------

set(_infini_train_executable_archives
infini_train
infini_train_cpu_kernels
)
if(USE_CUDA)
list(APPEND _infini_train_executable_archives infini_train_cuda_kernels)
endif()
set_property(GLOBAL PROPERTY INFINI_TRAIN_EXECUTABLE_ARCHIVES
${_infini_train_executable_archives})

function(link_infini_train_exe target_name)
if(USE_CUDA)
target_link_libraries(${target_name} PRIVATE
"-Wl,--start-group"
"-Wl,--whole-archive"
infini_train
infini_train_cpu_kernels
infini_train_cuda_kernels
"-Wl,--no-whole-archive"
"-Wl,--end-group"
)
cmake_parse_arguments(LINK "" "" "EXTRA_ARCHIVES" ${ARGN})
if(LINK_KEYWORDS_MISSING_VALUES)
message(FATAL_ERROR
"link_infini_train_exe: missing values for ${LINK_KEYWORDS_MISSING_VALUES}")
endif()
if(LINK_UNPARSED_ARGUMENTS)
message(FATAL_ERROR
"link_infini_train_exe: unexpected arguments: ${LINK_UNPARSED_ARGUMENTS}")
endif()

get_target_property(_target_type ${target_name} TYPE)
if(_target_type STREQUAL "INTERFACE_LIBRARY")
set(_link_scope INTERFACE)
elseif(_target_type STREQUAL "EXECUTABLE")
set(_link_scope PRIVATE)
else()
target_link_libraries(${target_name} PRIVATE
"-Wl,--start-group"
"-Wl,--whole-archive"
infini_train
infini_train_cpu_kernels
"-Wl,--no-whole-archive"
"-Wl,--end-group"
)
message(FATAL_ERROR
"link_infini_train_exe requires an executable or interface library target")
endif()

get_property(_framework_archives GLOBAL PROPERTY INFINI_TRAIN_EXECUTABLE_ARCHIVES)
target_link_libraries(${target_name} ${_link_scope}
"-Wl,--start-group"
"-Wl,--whole-archive"
${_framework_archives}
${LINK_EXTRA_ARCHIVES}
"-Wl,--no-whole-archive"
"-Wl,--end-group"
)
endfunction()

# Libraries and providers link InfiniTrain::infini_train for the core API. Final
# executables must link InfiniTrain::infini_train_executable or a provider-specific
# executable interface; linking only InfiniTrain::infini_train may discard static
# registration objects.
# EXTRA_ARCHIVES is reserved for archives that still rely on static initialization.
add_library(infini_train_executable INTERFACE)
link_infini_train_exe(infini_train_executable)
add_library(InfiniTrain::infini_train_executable ALIAS infini_train_executable)
unset(_infini_train_executable_archives)

# ------------------------------------------------------------------------------
# Examples
# ------------------------------------------------------------------------------

if(PROJECT_IS_TOP_LEVEL)
add_executable(mnist
example/mnist/main.cc
example/mnist/dataset.cc
Expand Down Expand Up @@ -220,10 +265,13 @@ add_executable(llama3
example/llama3/checkpoint_loader.cc
)
link_infini_train_exe(llama3)
endif()

# Tools
if(PROJECT_IS_TOP_LEVEL)
add_subdirectory(tools/infini_run)
set_target_properties(infini_run PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR})
endif()

# Tests
if(BUILD_TEST)
Expand Down
Loading
Loading